R2D3 Demonstration Loss

Problem

Combine a standard Q-learning loss with an auxiliary demonstration loss.

Return the mean of the Q losses plus lambda_demo times the mean of the demonstration losses over demonstration samples only. If there are no demonstration samples the second term is 0.0.

Hint

The demonstration term is averaged only over demonstration samples, not the whole batch.

Test cases

case 1

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

expected: 2.525

case 2

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

expected: 2

case 3

in: [[],[],[],0.5]

expected: 0

Starter codePython

def r2d3_demonstration_loss(q_losses, demo_losses, is_demo, lambda_demo):
    # Your code here
    pass

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