Generalized Advantage Estimation

mediumPROPolicy Gradients and Advantage EstimationReinforcement LearningRL 1: Foundations

Problem Statement

Compute GAE(gamma, lambda) advantages and the corresponding value targets for a rollout.

Background

An actor-critic agent needs an advantage estimate $A_t \approx q(s_t,a_t) - v(s_t)$ : how much better this action was than the critic's expectation. The one-step TD error is the cheapest such estimate,

$\delta_t = r_t + \gamma V(s_{t+1}) - V(s_t)$

low variance but biased by whatever the critic gets wrong. The full Monte Carlo advantage is unbiased but high variance. GAE dials between them with exactly the geometric average you already met as the lambda return:

$A_t^{\text{GAE}} = \sum_{l=0}^{\infty} (\gamma\lambda)^l \delta_{t+l}$

which computes in one backward pass:

$A_t = \delta_t + \gamma\lambda\,(1 - d_t)\, A_{t+1}$

Two implementation details do all the damage in practice:

The done mask appears twice. It kills the bootstrap inside $\delta_t$ (a terminal state has no successor value) and it stops the accumulator from leaking advantage across an episode boundary into the previous episode.
**The value targets are **A_t + V(s_t)****, not the raw discounted returns. This is what keeps the critic's regression target consistent with the advantages the actor is being trained on.

lam = 0 recovers the plain TD error and lam = 1 recovers the Monte Carlo advantage — check your implementation against both.

Your Task

Implement:

def gae(rewards, values, dones, last_value, gamma, lam):
    ...

rewards[t], values[t] (the critic's $V(s_t)$ ), dones[t] (bool, True if step t ended an episode).
last_value — the critic's value of the state after the final step of the rollout, used only when the last step is not terminal.
Process the rollout backwards.

Return a **tuple **(advantages, returns)****, both lists of T floats, where returns[t] = advantages[t] + values[t].

Input / Output Format

Lists of floats/bools in; a tuple of two lists of floats out, rounded to 4 decimals by the grader.

Sample

rewards = [1.0, 1.0]
values  = [0.0, 0.0]
dones   = [False, True]
adv, ret = gae(rewards, values, dones, 0.0, 0.9, 0.95)
print([round(x, 4) for x in adv])
print([round(x, 4) for x in ret])

Output:

[1.855, 1.0]
[1.855, 1.0]

Step 1 is terminal so delta_1 = 1.0 and A_1 = 1.0. Step 0 has delta_0 = 1.0 + 0.9*0.0 - 0.0 = 1.0 and A_0 = 1.0 + 0.90.951.0 = 1.855. The values are all zero, so the returns equal the advantages.

Example

Input

gae([1.0, 1.0], [0.0, 0.0], [False, True], 0.0, 0.9, 0.95)

Output

[1.855, 1.0]
[1.855, 1.0]

Reasoning

Working backwards: step 1 is terminal so there is no bootstrap, delta_1 = 1.0 - 0.0 = 1.0 and A_1 = 1.0. Step 0 is not terminal, so delta_0 = 1.0 + 0.9V(s_1) - V(s_0) = 1.0, and A_0 = delta_0 + 0.90.95*A_1 = 1.855. Since every value is 0, returns equal advantages.

Constraints

1 <= T <= 10000
dones[t] is a Python bool.
0.0 <= gamma <= 1.0, 0.0 <= lam <= 1.0
The done flag must zero out both the bootstrap in delta_t and the carried-over gamma*lam*A_{t+1} term.
returns[t] must be advantages[t] + values[t].
Do not round inside the function; return (advantages, returns) in that order.

Learn this in

Reinforcement LearningML Study Plan Policy Gradient MethodsReinforcement Learning Actor-Critic & PPOReinforcement 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

Generalized Advantage Estimation (GAE) is a technique used in Actor-Critic reinforcement learning algorithms to estimate the advantage function $A(s, a) = Q(s, a) - V(s)$ . The advantage function measures how much better a specific action is compared to the average expected return from that state. Estimating this accurately is crucial for reducing the variance of policy gradient updates while maintaining low bias.

The core concept relies on balancing Temporal Difference (TD) errors and Monte Carlo returns. A one-step TD error, $\delta_t = r_t + \gamma V(s_{t+1}) - V(s_t)$ , is low variance but biased because it depends on the critic's estimate $V(s_{t+1})$ . Conversely, the full Monte Carlo return is unbiased but has high variance. GAE introduces a hyperparameter $\lambda$ to create a weighted geometric sum of $k$ -step TD errors. When $\lambda=0$ , GAE reduces to the one-step TD error (high bias, low variance). When $\lambda=1$ , it approximates the Monte Carlo advantage (low bias, high variance).

The recursive formula for GAE allows for efficient computation in a single backward pass through the trajectory: $A_t = \delta_t + \gamma\lambda (1 - d_t) A_{t+1}$ Here, $d_t$ is the done flag. If the episode ends at step $t$ , the bootstrap term $A_{t+1}$ is zeroed out because there is no future state to bootstrap from within the same episode. This prevents "leakage" of advantage estimates across episode boundaries. The final value targets (or returns) used to train the critic are computed as $R_t = A_t + V(s_t)$ , ensuring the critic learns to predict the same returns that the actor uses for advantage estimation.

2. Algorithm Approach

The problem requires implementing the backward recursive calculation of GAE. The approach is iterative and processes the rollout from the last time step $T-1$ down to $0$ .

Initialize: Start with the advantage of the step after the last one, $A_T$ , which is effectively $0$ (or handled via the last_value if not terminal).
Iterate Backwards: For each step $t$ from $T-1$ down to $0$ :

Calculate the TD error $\delta_t$ . This requires the current reward, the current value, and the next value. The "next value" is values[t+1] if $t < T-1$ , or last_value if $t = T-1$ .
Apply the done mask: If the episode ends at $t$ (dones[t] is True), the next value contribution to the TD error is zeroed out (conceptually $V(s_{t+1})=0$ for the delta calculation relative to the end of episode, though strictly $\delta_t = r_t - V(s_t)$ if terminal). More precisely, the bootstrap term in the GAE recursion is multiplied by $(1 - d_t)$ .
Update the Advantage: $A_t = \delta_t + \gamma \lambda (1 - d_t) A_{t+1}$ .
Compute the Return: $R_t = A_t + V(s_t)$ .

Store Results: Keep track of $A_t$ and $R_t$ for each step.

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.

PRO

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