📘

n-step and Lambda Returns

HardReinforcement Learning 📚 RL 1: Foundations

Problem Statement

Given a full episode of rewards R_1..R_T and bootstrap value estimates V(s_0)..V(s_T) (length T+1), compute two prediction targets for time 0.

n-step return: $G^{(n)}_0 = \sum_{k=0}^{n-1} \gamma^{k} R_{k+1} + \gamma^{n} V(s_n)$ (if n >= T, use the full Monte-Carlo return with no bootstrap, i.e. gamma^T * V(s_T) added where V(s_T)=0 for a terminal state — but here just apply the formula with V(s_n), and if n >= T set the bootstrap term to gamma^T * values[T]).

lambda return (offline, full episode): $G^{\lambda}_0 = (1-\lambda)\sum_{n=1}^{T-1}\lambda^{n-1} G^{(n)}_0 + \lambda^{T-1} G^{(T)}_0$

Implement n_step_and_lambda(rewards, values, gamma, lam, n) returning the tuple (g_n, g_lambda). rewards has length T, values has length T+1.

Example:

Input:

n_step_and_lambda([1.0, 1.0], [0.0, 0.0, 0.0], 1.0, 0.5, 1)

Output:

(1.0, 1.5)

Reasoning:

Compute the 1-step return ( $G^{(1)}_0$ ): With $n=1$ , we sum the first reward discounted by $\gamma^0$ and add the bootstrap term $\gamma^1 V(s_1)$ . Using the inputs $R_1=1.0$ , $\gamma=1.0$ , and $V(s_1)=0.0$ , the calculation is $1.0 \times 1.0 + 1.0 \times 0.0 = 1.0$ . This is the first element of the output tuple.
Compute the 2-step return ( $G^{(2)}_0$ ): Since the episode length $T=2$ , the 2-step return uses the full horizon. We sum both rewards ( $R_1$ and $R_2$ ) and add the bootstrap term $\gamma^2 V(s_2)$ . The calculation is $(1.0 \times 1.0) + (1.0 \times 1.0) + (1.0 \times 0.0) = 2.0$ .
Calculate the $\lambda$ -return ( $G^{\lambda}_0$ ): The formula combines the 1-step and 2-step returns using $\lambda=0.5$ . It is computed as $(1-\lambda)G^{(1)}_0 + \lambda^{T-1}G^{(T)}_0$ . Substituting the values: $(1 - 0.5) \times 1.0 + 0.5^{2-1} \times 2.0$ .
Finalize the $\lambda$ -return value: Evaluating the expression from the previous step: $0.5 \times 1.0 + 0.5 \times 2.0 = 0.5 + 1.0 = 1.5$ . This is the second element of the output tuple.
The final output is (1.0, 1.5)

Constraints:

len(values) == len(rewards) + 1.
Clamp n to at most T for the n-step term.
Return a tuple (g_n, g_lambda) of floats.

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 Temporal-Difference (TD) learning, prediction targets balance bias and variance by mixing immediate rewards with future value estimates. The n-step return $G^{(n)}_0$ unrolls the Bellman equation for exactly $n$ steps: it accumulates discounted rewards $R_1, \dots, R_n$ and then bootstraps from the estimated value $V(s_n)$ . When $n$ exceeds the episode length $T$ , the bootstrap term collapses to the terminal value, effectively recovering the Monte-Carlo return.

The lambda return $G^{\lambda}_0$ generalizes n-step returns by taking a weighted average over all n-step returns from $n=1$ to $T$ . The weights form a geometric series in $\lambda$ , with the final term (the full Monte-Carlo return) receiving weight $\lambda^{T-1}$ . This construction ensures that $G^{\lambda}$ interpolates smoothly between one-step TD ( $\lambda=0$ ) and Monte-Carlo ( $\lambda=1$ ), providing a tunable bias-variance trade-off.

A critical detail is the discounting convention. Rewards are discounted by $\gamma^k$ where $k$ is the step index relative to time 0. The bootstrap term at step $n$ carries a factor of $\gamma^n$ , not $\gamma^{n-1}$ , because it represents the value of the state after $n$ transitions. Confusing these exponents is the most frequent source of bugs in this problem family.

2. Algorithm Approach

The solution follows a direct computation pattern with two independent passes:

Pass 1 (n-step): Iterate $k$ from $0$ to $\min(n, T)-1$ , accumulating $\gamma^k R_{k+1}$ . Then add the bootstrap term $\gamma^{\min(n,T)} V(s_{\min(n,T)})$ .
Pass 2 (lambda return): Iterate $n$ from $1$ to $T$ , computing each $G^{(n)}_0$ (or reusing a running sum) and accumulating $(1-\lambda)\lambda^{n-1} G^{(n)}_0$ for $n < T$ , plus $\lambda^{T-1} G^{(T)}_0$ for the final term.

A key optimization insight: you can compute all $G^{(n)}_0$ values in a single forward pass by maintaining a running discounted reward sum and updating the bootstrap term incrementally, avoiding redundant recomputation.

3. Step-by-Step Strategy

👑

Continue the full explanation

You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.

🔒

Editor locked

The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.

solution.py

Python 3.13

Test Results

0/0

Run code to see test results.