GAE One Step

Problem

Compute one step of generalized advantage estimation.

First form the TD residual delta = reward + gamma * next_value - value, then return delta + gamma * lam * next_advantage. Pass next_advantage = 0.0 for the final step of an episode.

Hint

GAE mixes the current TD residual with the discounted advantage of the next step.

Test cases

case 1

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

expected: 1.7000000000000002

case 2

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

expected: 3.41

case 3

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

expected: 0.495

case 4

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

expected: -1.25

Starter codePython

def gae_one_step(reward, value, next_value, gamma, lam, next_advantage):
    # Your code here
    pass

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