N-Step Return

Problem

Compute the n-step return from a list of rewards.

Use the first min(n, len(rewards)) rewards and bootstrap with v_next only if n rewards were available; if the episode ended early there is no bootstrap. G = sum of gamma^t * r_t + gamma^n * v_next.

Hint

If fewer than n rewards are available the episode terminated, so drop the bootstrap term.

Test cases

case 1

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

expected: 5.950000000000001

case 2

in: [[1,1],0.5,100,3]

expected: 1.5

case 3

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

expected: 4.7

case 4

in: [[],0.9,7,3]

expected: 0

Starter codePython

def n_step_return(rewards, gamma, v_next, n):
    # Your code here
    pass

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