Imagination Lambda Return

Problem

Compute the lambda return of a Dreamer-style imagined trajectory.

Starting from the bootstrap value at the horizon, apply the backward recursion G_t = rewards[t] + gamma * ((1 - lam) * values[t + 1] + lam * G_{t + 1}) for each step. values has one more entry than rewards and the final value initializes the recursion.

Hint

Lambda blends the one-step value estimate with the longer-horizon return.

Test cases

case 1

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

expected: 1.45

case 2

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

expected: 2.8

case 3

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

expected: 4.33

case 4

in: [[],[4],0.9,0.5]

expected: 4

case 5

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

expected: 2.5

Starter codePython

def imagination_lambda_return(rewards, values, gamma, lam):
    # Your code here
    pass

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