📘

n-step Returns and the Lambda Return

HardReinforcement Learning 📚 RL 1: Foundations

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.

Example:

Input:

n_step_return([1.0, 2.0, 3.0], [10.0, 20.0, 30.0], 0, 1, 0.5) and n_step_return([1.0, 2.0, 3.0], [10.0, 20.0, 30.0], 0, 3, 0.5)

Output:

11.0
2.75

Reasoning:

With n=1 the return takes one real reward then bootstraps: 1.0 + 0.5V(S_1) = 1.0 + 0.520.0 = 11.0. With n=3 the window reaches the terminal step, so there is no state left to bootstrap from and the return is the full discounted reward 1.0 + 0.52.0 + 0.253.0 = 2.75.

Constraints:

1 <= T <= 500, 0 <= t < T, 1 <= n
n may exceed the number of remaining steps; in that case the return is the full remaining discounted reward with no bootstrap term.
0.0 <= gamma <= 1.0, 0.0 <= lam <= 1.0
lambda_return must satisfy: lam = 0 gives the 1-step return, lam = 1 gives the full Monte Carlo return.
Do not round inside the functions.

More from RL 1: Foundations

GAE Over a Full TrajectoryHard PPO Clipped Objective (Single Sample)Medium TD Residuals for GAEMedium Return Standardization BaselineMedium REINFORCE Loss TermMedium

1. Background Knowledge

In Reinforcement Learning, estimating the value of a state or action requires balancing bias and variance. Monte Carlo (MC) methods wait for an episode to finish, using the actual total return. This is unbiased but has high variance. Temporal Difference (TD) learning bootstraps from current estimates, updating after every step. This has low variance but high bias. The n-step return generalizes both: it takes $n$ actual rewards and then bootstraps from the estimated value of the state reached after those $n$ steps.

The lambda return ( $G_t^\lambda$ ) is a weighted average of all possible n-step returns. It uses a parameter $\lambda \in [0, 1]$ to control the weighting. When $\lambda=0$ , it weights only the 1-step return (pure TD). When $\lambda=1$ , it weights the full episode return (pure MC). For intermediate values, it provides a smooth trade-off. This concept is foundational for advanced algorithms like Generalized Advantage Estimation (GAE).

2. Algorithm Approach

To solve this, you need to implement two distinct calculations. The first is a direct summation for the n-step return. The second is a weighted sum for the lambda return.

For the n-step return, iterate from time step $t$ up to $t+n-1$ . Accumulate the discounted rewards. If the episode ends before $n$ steps are taken, stop accumulating rewards and do not add a bootstrap value. If the episode continues, add the discounted value estimate $V(S_{t+n})$ .

For the lambda return, you must compute the weighted sum of all valid n-step returns. The weight for the $n$ -step return is $(1-\lambda)\lambda^{n-1}$ . Crucially, the final term (the full episode return) receives the remaining weight $\lambda^{T-t-1}$ to ensure the weights sum to 1. You can implement this by iterating through all possible $n$ values from 1 to $T-t$ , calculating the corresponding n-step return, and applying the geometric weights.

3. Step-by-Step Strategy

Implement n_step_return:

Initialize a variable return_val to 0.
Determine the number of actual steps available: steps = min(n, T - t).
Loop i from 0 to steps - 1:
Add $\gamma^i \times \text{rewards}[t+i]$ to return_val.
If steps < n (meaning we hit the end of the episode), return return_val immediately. Do not bootstrap.
If steps == n (we have exactly $n$ steps and the episode continues), add $\gamma^n \times \text{values}[t+n]$ to return_val.
Return return_val.

Implement lambda_return:

Initialize lambda_ret to 0.
Calculate the total number of steps remaining in the episode: T_minus_t = len(rewards) - t.
Loop n from 1 to T_minus_t - 1:
Calculate the n-step return using your helper function.
Calculate the weight: (1 - lam) * (lam ** (n - 1)).
Add weight * n_step_return to lambda_ret.
Handle the final term (full episode return):
Calculate the full return $G_{t:T}$ (which is effectively the n-step return where $n$ is large enough to reach the end, or simply sum all remaining discounted rewards).
Calculate the final weight: lam ** (T_minus_t - 1).
Add weight * full_return to lambda_ret.
Return lambda_ret.

4. Common Pitfalls

Bootstrapping at the End: The most common error is adding a value term $V(S_{t+n})$ even when $t+n$ exceeds or equals the episode length $T$ . Remember, $V(S_T) = 0$ by definition, but more importantly, if the episode terminates, there is no future state to bootstrap from. The return is just the sum of discounted rewards.
Index Out of Bounds: When accessing values[t+n], ensure $t+n$ is a valid index. If $t+n == T$ , the value is 0, but accessing values[T] will cause an IndexError since values has length $T$ (indices $0$ to $T-1$ ).
Lambda Weight Summation: Ensure the weights sum to 1. The geometric series $(1-\lambda)\sum_{n=1}^{k} \lambda^{n-1}$ sums to $1-\lambda^k$ . The remaining weight $\lambda^k$ must be assigned to the final full return. Missing this final term will result in an incorrect lambda return.
Off-by-One Errors: Be careful with the range of $n$ . The problem defines $G_{t:t+n}$ using $n$ rewards. The lambda return sums from $n=1$ to $T-t-1$ for the intermediate terms, with the last term being the full return.

5. Time & Space Complexity

Time Complexity:
n_step_return: $O(n)$ because it iterates up to $n$ times.
lambda_return: $O(T^2)$ in the naive implementation. It loops $T-t$ times, and each iteration calls n_step_return which takes up to $O(T)$ time. This is acceptable for typical episode lengths in RL problems.
Space Complexity:
$O(1)$ auxiliary space for both functions, as they only use a few variables for accumulation and do not create new data structures proportional to the input size.

solution.py

Python 3.13

Test Results

0/0

Run code to see test results.