21. Generalized Advantage Estimation (GAE)

Problem

Actor-critic methods need an advantage estimate to tell the policy gradient “how much better than average was this action,” but the two obvious choices both have a flaw: the one-step TD error has low variance but high bias, while the full Monte Carlo return has high variance but low bias. GAE interpolates between these two extremes with a single parameter, lambda, via an exponentially-weighted sum of k-step TD errors. It's the advantage estimator underlying most modern actor-critic training, including the advantage-based updates in DeepMind's own A3C/IMPALA line of actor-critic architectures.

Given per-timestep rewards, a value function's estimates (including one extra bootstrap value for the state just after the last reward), and done flags marking episode boundaries, compute the GAE advantage at every timestep via the backward recursion:
delta_t = r_t + gamma * V(s_{t+1}) * (1 - done_t) - V(s_t)
A_t = delta_t + gamma * lambda * (1 - done_t) * A_{t+1}
with A_T (past the last real timestep) implicitly zero.

Source: src/21_gae_advantage.py

def compute_gae(
    rewards: torch.Tensor, values: torch.Tensor, dones: torch.Tensor,
    gamma: float, lam: float,
) -> torch.Tensor: ...

>>> rewards = torch.tensor([[1.0, 1.0, 1.0]])
>>> values = torch.tensor([[0.5, 0.5, 0.5, 0.0]])
>>> dones = torch.tensor([[0.0, 0.0, 1.0]])
>>> compute_gae(rewards, values, dones, gamma=0.99, lam=0.95).shape
torch.Size([1, 3])

Step-by-Step Approach

1. Note the shape mismatch by design: values has one more timestep than rewards and dones, holding the bootstrap value for the state right after the final reward.
2. Initialize a running gae accumulator to zero (batch,) — this represents A_{t+1} as the loop walks backward, and there is nothing past the last timestep to bootstrap from.
3. Loop over timesteps in reverse. At each step compute the one-step TD error delta_t, using values[:, t+1] as the bootstrap and masking it by (1 - dones[:, t]).
4. Fold delta_t into the running GAE accumulator: gae = delta_t + gamma * lam * (1 - dones[:, t]) * gae — the same done-mask also has to zero the recursive term, not just the bootstrap value.
5. Store gae into the output advantages tensor at position t before moving to t - 1.
6. Sanity-check the lam = 0 edge case: the recursive term vanishes entirely and each advantage should equal exactly its own one-step TD error, nothing from later timesteps.

GAE is really just one accumulator variable threaded backward through time, reused for two purposes at once: it both stores “the advantage so far” and carries “the discounted future correction” into the next (earlier) step's calculation — get the mask on the recursive term wrong and the algorithm silently bootstraps advantage across an episode boundary that shouldn't exist.

Reference solution

import torch


def compute_gae(
    rewards: torch.Tensor, values: torch.Tensor, dones: torch.Tensor,
    gamma: float, lam: float,
) -> torch.Tensor:
    batch, T = rewards.shape
    advantages = torch.zeros_like(rewards)

    # gae carries A_{t+1} backward into the computation of A_t -- it starts
    # at 0 because there is nothing past the last real timestep to bootstrap
    gae = torch.zeros(batch)
    for t in reversed(range(T)):
        # (1 - done_t) zeroes out both the value bootstrap and the recursive
        # term whenever the episode ended at t: no state t+1 exists to look at
        not_done = 1.0 - dones[:, t]
        delta = rewards[:, t] + gamma * values[:, t + 1] * not_done - values[:, t]
        gae = delta + gamma * lam * not_done * gae
        advantages[:, t] = gae

    return advantages

Key Functions & Tricks

• torch.zeros_like(rewards) — allocates the output advantages tensor with the same shape and dtype as the rewards, before filling it in via the backward loop
• reversed(range(T)) — walks timesteps back-to-front so each step's advantage can reuse the already-computed A_{t+1}
• values[:, t + 1] — the one-step-ahead bootstrap value, why values is deliberately one timestep longer than rewards
• (1.0 - dones[:, t]) — applied twice per step — once to the value bootstrap, once to the recursive GAE term — since a terminal transition invalidates both

How to Recognize This Pattern