World-Model Discounted Loss

Problem

Compute a discounted multi-step world-model loss.

At each horizon average the squared error over the state dimensions, weight it by gamma^t, and return the discounted sum divided by the number of steps. An empty trajectory returns 0.0.

Hint

Discounting later horizons stops long rollouts from dominating the loss.

Test cases

case 1

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

expected: 2.5

case 2

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

expected: 0.75

case 3

in: [[[2,0]],[[0,0]],0.9]

expected: 2

Starter codePython

def world_model_discounted_loss(predicted, true, gamma):
    # Your code here
    pass

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