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 : ```python 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.shape1, 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 :python def train(self): # Sample a batch (a tuple of narray) from the replay buffer # States (BS), actions (B), rewards (B), next_states(BS), 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. 695