n-step Returns and the Lambda Return

Problem Statement

Compute the n-step return at a time step, and the lambda return that geometrically averages all of them.

Background

TD(0) bootstraps after one reward; Monte Carlo waits for the whole episode. Those are the two ends of one dial, and the n-step return is the dial:

$G_{t:t+n} = R_{t+1} + \gamma R_{t+2} + \cdots + \gamma^{n-1} R_{t+n} + \gamma^{n} V(S_{t+n})$

Take real rewards for n steps, then bootstrap. If t + n reaches or passes the terminal step, there is nothing to bootstrap from — the return is simply the full remaining discounted reward, with no value term at all. That truncation is the single most common bug here.

The lambda return refuses to choose an n. It averages every n-step return with geometrically decaying weights $(1-\lambda)\lambda^{n-1}$, and gives the entire leftover weight to the full return so the coefficients sum to exactly 1:

$G_t^{\lambda} = (1-\lambda) \sum_{n=1}^{T-t-1} \lambda^{n-1} G_{t:t+n} \; + \; \lambda^{T-t-1} G_{t:T}$

Sanity checks worth remembering: $\lambda = 0$ collapses to the one-step TD return, and $\lambda = 1$ collapses to the Monte Carlo return. Everything interesting is in between, and this same geometric average reappears verbatim as GAE in the actor-critic chapter.

Your Task

The episode has T = len(rewards) steps: rewards[k] is $R_{k+1}$, and values[k] is $V(S_k)$ for k in 0..T-1. The episode terminates after rewards[T-1], so $V(S_T) = 0$.

def n_step_return(rewards, values, t, n, gamma):
    ...

def lambda_return(rewards, values, t, gamma, lam):
    ...

Both return a float.

Input / Output Format

Two lists of floats plus ints/floats in; a float out, rounded to 4 decimals by the grader.

Sample

rewards = [1.0, 2.0, 3.0]
values  = [10.0, 20.0, 30.0]
print(round(n_step_return(rewards, values, 0, 1, 0.5), 4))
print(round(n_step_return(rewards, values, 0, 3, 0.5), 4))

Output:

11.0
2.75

The 1-step return is 1.0 + 0.5V(S_1) = 1.0 + 0.520.0 = 11.0. The 3-step return reaches the terminal step, so it bootstraps from nothing: 1.0 + 0.52.0 + 0.253.0 = 2.75.