Advantage Estimation (GAE) - Reinforcement Learning

Generalized Advantage Estimation (GAE)

What you'll build

GAE(γ, λ) advantages for one trajectory, computed with the standard backward recursion used in every PPO implementation.

Background

The advantage A(s, a) = Q(s, a) - V(s) says how much better an action was than the policy's average. Policy gradients weighted by advantages instead of raw returns have far lower variance. GAE (Schulman et al., 2016) blends one-step TD errors with an exponential weight λ, trading bias (small λ, trusts the critic) against variance (large λ, trusts the sampled rewards).

Formula

δ_t = r_t + γ · V(s_{t+1}) - V(s_t)                 (TD error)
Â_t = Σ_{l=0}^{T-1-t} (γλ)^l · δ_{t+l}
    = δ_t + γλ · Â_{t+1}        with Â_T = 0         (backward recursion)

• λ = 0: Â_t = δ_t (one-step TD: low variance, biased)
• λ = 1: Â_t = G_t - V(s_t) (Monte Carlo return minus baseline: unbiased, high variance)
• Typical: γ = 0.99, λ = 0.95

Inputs

• rewards: (T,) rewards r_0 ... r_{T-1}
• values: (T+1,) critic values V(s_0) ... V(s_T). The last entry is the bootstrap value of the state after the final step: pass 0 if the episode terminated there, or the critic's estimate if the trajectory was truncated.
• gamma, lam: floats in [0, 1]

Expected return format

A NumPy array of shape (T,), rounded to 2 decimals.

Edge cases and conventions

• T = 1 gives Â_0 = δ_0.
• Iterate from t = T-1 down to 0; Â_{T-1} = δ_{T-1}.

Constraints

NumPy only. O(T) time: don't compute the double sum.

Follow-up

• Implementations also return value targets: returns = advantages + values[:-1]. Why is that the TD(λ) return?
• With several episodes in one rollout buffer, how do you stop the recursion from leaking across an episode boundary? (Multiply by (1 - done_t).)

Input: rewards = [1, 1, 1], values = [0, 0, 0, 0], gamma = 0.99, lam = 0.95
Output: [2.83, 1.94, 1.0]

Input: rewards = [1, 2, 3], values = [1, 2, 3, 0], gamma = 0.9, lam = 0.9
Output: [3.99, 2.7, 0.0]

import numpy as np

def compute_gae(rewards: np.ndarray, values: np.ndarray,
                gamma: float, lam: float) -> np.ndarray:
    """
    Compute Generalized Advantage Estimation.

    Args:
        rewards: (T,) rewards r_0 ... r_{T-1}
        values: (T+1,) values V(s_0) ... V(s_T); the last entry is the
                bootstrap value (0 if the episode terminated)
        gamma: discount factor
        lam: GAE lambda

    Returns:
        (T,) advantages A_0 ... A_{T-1}, rounded to 2 decimals
    """
    rewards = np.array(rewards, dtype=float)
    values = np.array(values, dtype=float)
    T = len(rewards)
    advantages = np.zeros(T)

    # Your code here
    pass