n-step and Lambda Returns

hardPROTemporal-Difference Learning and ControlReinforcement LearningRL 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.

Learn this in

Monte Carlo & Temporal-Difference LearningReinforcement Learning

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

Concept primerPREVIEW

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

PRO

Continue the full explanation

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