DQN Loss and Target Network

Write the two parts of DQN that are easy to get wrong: the temporal difference loss with a target network, and the soft update of that target network.

DQN learns Q(s, a) with a neural network. For a batch of transitions (s, a, r, s', done) from a replay buffer, the loss is:

q      = q_net(s)[a]                                  # the action taken
target = r + gamma * (1 - done) * max_a' target_net(s')[a']
loss   = smooth_l1(q, target)                         # Huber loss, mean

The target uses a second network, a slowly moving copy of q_net. No gradient may flow through the target.

With double=True, use Double DQN instead. The online network chooses the next action, and the target network scores it:

a*     = argmax_a' q_net(s')[a']
target = r + gamma * (1 - done) * target_net(s')[a*]

Then write soft_update, which moves every target parameter part of the way towards the online parameter, in place:

target_param = (1 - tau) * target_param + tau * online_param

The tests run a full DQN training loop with your two functions, on a small corridor task.

Function signature

def dqn_loss(q_net, target_net, states, actions, rewards, next_states,
             dones, gamma, double=False):
def soft_update(target_net, online_net, tau):

Example

Input

q_net(states)            -> [[1.0, 2.0], [0.5, 3.0]]
target_net(next_states)  -> [[4.0, 1.0], [4.0, 1.0]]
actions = [1, 0], rewards = [1.0, 0.0], dones = [0.0, 1.0], gamma = 0.9

Output

tensor(1.1125)

Reasoning

The chosen Q values are 2.0 and 0.5. The first target is 1 + 0.9 * 4 = 4.6. The second transition ends the episode, so its target is the reward 0. The errors are 2.6 and 0.5. The Huber loss is 2.6 - 0.5 = 2.1 for the large error and 0.5 * 0.5^2 = 0.125 for the small one. The mean is 1.1125.

Constraints and notes

• states and next_states have shape (B, obs_dim). actions is int64 of shape (B,). rewards and dones are float of shape (B,), with dones equal to 1.0 where the episode ended.
• Use F.smooth_l1_loss with its defaults (beta=1, mean over the batch).
• After loss.backward(), the target network must have no gradients.
• soft_update returns nothing. It changes the target parameters in place, not the online ones, and must not record anything for autograd.