DQN Loss (MSE)

Problem

Compute the mean squared DQN loss over a batch.

For each sample the target is reward + gamma * next_q_max * (1 - done), where done is a boolean, and the prediction is q_values[i][actions[i]]. Return the mean of the squared errors; the batch is non-empty.

Hint

Only the Q value of the action that was taken contributes to the loss.

Test cases

case 1

in: [[[1,2],[3,1]],[1,0],[1,0],[2,4],[false,true],0.9]

expected: 4.82

case 2

in: [[[0,0]],[0],[1],[2],[false],0.5]

expected: 4

case 3

in: [[[5,5],[0,0]],[0,1],[5,2],[0,0],[false,false],0.9]

expected: 2

case 4

in: [[[1]],[0],[1],[10],[true],0.99]

expected: 0

Starter codePython

def dqn_loss(q_values, actions, rewards, next_q_max, dones, gamma):
    # Your code here
    pass

4 test cases · run your code against them in the editor