# Optimal implementation of vanilla DQN loss in Keras

I've implemented vanilla DQN for continuous/non-images (no CNN) states in keras. But, I'm not sure if my implementation of the loss computation is optimal.

For reminder the loss is defined as : $$loss=(r+\gamma max_{a'} Q(s',a')-Q(s,a))^2 \frac{1}{2}$$

Here is my implementation of the network + loss function :

self.network = Sequential()
self.network.add(Dense(256, activation='relu', input_dim=input_dim))
self.network.add(Dense(32, activation='relu'))
self.network.add(Dense(output_dim))

def q_loss(data, y_pred):
# Extract the concatenated data tensor
action, reward, next_state, done = K.cast(data[:, 0], 'int32'), data[:, 1], data[:, 2:-1], K.cast(data[:, -1], 'bool')
# Compute Q(s,a)
mask = tf.one_hot(action, depth=y_pred.shape, dtype=tf.bool, on_value=True, off_value=False)
q_action = tf.boolean_mask(y_pred, mask)
# Compute the max of values at next state except if done=True
max_q_next = K.max(self.network(next_state), axis=1) * K.cast(tf.logical_not(done), 'float32')
# Compute the TD-error, do not propagate the gradient into the next state value
td_error = reward + 0.95 * K.stop_gradient(max_q_next) - q_action
# Compute the MSE
loss = K.square(td_error) / 2
return loss

self.network.compile(loss=q_loss, optimizer=RMSprop(lr=self.learning_rate))


Here is my train function :

def train(self):
# Sample a batch (a tuple of narray) from the replay buffer
# States (B*S), actions (B), rewards (B), next_states(B*S), dones (B)
# B=batch_size and S=state_size
states, actions, rewards, next_states, dones = self.memory.sample(self.batch_size)
# Concatenate actions, rewards, next_states, dones together because keras loss only accept one tensor
data = np.concatenate([np.expand_dims(actions, axis=1), np.expand_dims(rewards, axis=1), next_states, np.expand_dims(dones, axis=1)], axis=1)
# Train on a batch
self.network.train_on_batch(states, data)


I find the way I compute the DQN td-error and loss being ugly and probably sub-optimal. Do you have a better workaround (maybe with a mix of Keras and Tensorflow combined) ?

I've checked multiple existing Keras implementations (like link), but sadly they mostly compute the td-error outside of keras in full python/numpy, witch is IMO sub-optimal.