By the end of §15.4 we had two ways to estimate the advantage $A^\pi(s_t, a_t)$ of an action from one sampled trajectory, and they sit at opposite ends of a trade-off. The Monte Carlo estimate, $G_t - V(s_t)$, waits for every reward that follows and subtracts the critic's prediction. Whatever the critic gets wrong, it is a valid baseline, so the policy gradient stays unbiased, but it carries the noise of an entire future. The TD error, $\delta_t = r_t + \gamma V(s_{t+1}) - V(s_t)$, looks one step ahead and trusts the critic for the rest. It is quiet, but if the critic is wrong, the gradient is wrong.

Is there anything in between? There is a whole family of estimators, and one parameter moves smoothly from one end to the other. Schulman, Moritz, Levine, Jordan, and Abbeel (2016) named it generalized advantage estimation (GAE). It is the advantage estimator inside PPO, and therefore inside RLHF (§16.4), and it is a five-line backward loop that interviewers love to ask for, because it is easy to get almost right. This section derives it, measures its bias and variance against exact advantages, and handles the part that actually breaks in code, the ends of episodes.

Sums of TD errors

Throughout, $V$ is the critic's current estimate (not necessarily the true $V^\pi$), $\gamma$ is the discount, and the TD error at step $t$ is

Add up the next $k$ TD errors, discounting each by how far ahead it is. Writing each one out,

The last two sums are the same list of terms shifted by one. The second runs over $\gamma^1V(s_{t+1}), \dots, \gamma^kV(s_{t+k})$ and the third over $\gamma^0V(s_t), \dots, \gamma^{k-1}V(s_{t+k-1})$, so everything cancels except the two ends. The sum telescopes,

In plain words, the sum of $k$ discounted TD errors is the $k$-step advantage estimate. Take $k$ real rewards, let the critic estimate everything after them, and subtract the critic's estimate of where you started. The part before the subtraction, $G_t^{(k)} = \sum_{l<k}\gamma^lr_{t+l} + \gamma^kV(s_{t+k})$, is called the $k$-step return.

The two estimators we already had are the ends of this family. With $k = 1$ it is the TD error. With $k$ running to the end of the episode (or to infinity, where $\gamma^kV \to 0$), the bootstrap term disappears and it is the Monte Carlo estimate $G_t - V(s_t)$.

Where the bias lives, and how it shrinks

Suppose the critic is off by an error $e(s) = V(s) - V^\pi(s)$. Take the expectation of $\hat A_t^{(k)}$ given $s_t$ and $a_t$. Write $V = V^\pi + e$ in the $k$-step estimate. The part built from the true values has expectation exactly $A^\pi(s_t, a_t)$. To see it, use the Bellman equation $V^\pi(s_{t+k-1}) = \mathbb{E}[r_{t+k-1} + \gamma V^\pi(s_{t+k})\mid s_{t+k-1}]$, which holds because the action at step $t+k-1$ is drawn from $\pi$, to fold the last reward and the bootstrap into $\gamma^{k-1}V^\pi(s_{t+k-1})$, and repeat until one reward is left. What remains is $\mathbb{E}[r_t + \gamma V^\pi(s_{t+1})\mid s_t, a_t] - V^\pi(s_t) = Q^\pi(s_t, a_t) - V^\pi(s_t)$, because the first action is fixed. The errors are left over,

The last term, $-e(s_t)$, depends only on the state, so it is a baseline and does not bias the policy gradient (§15.4). The middle term depends on the action, because the action influences where you are $k$ steps later, so it does bias the gradient. But it is multiplied by $\gamma^k$. In plain words, the further you look ahead before trusting the critic, the less the critic's mistakes matter, and they matter geometrically less.

The variance goes the other way. Each extra step adds one more random reward and one more random transition to the estimate. So $k$ is a bias-variance dial, but an awkward one. It is an integer, the best value depends on the problem, and it throws away the information in all the other horizons.

GAE, an average over every horizon

Instead of choosing one $k$, average all of them, with weights that decay geometrically. For a parameter $\lambda \in [0, 1]$, define

The weights $(1-\lambda)\lambda^{k-1}$ are positive and sum to one, so this is a weighted average of the $k$-step estimates, leaning on short horizons when $\lambda$ is small and long ones when $\lambda$ is near 1. It looks expensive, an infinite sum of sums, but it collapses. Substitute each $\hat A^{(k)}_t$ as its sum of TD errors and collect the coefficient of each $\delta_{t+l}$. The error $\delta_{t+l}$ appears (with factor $\gamma^l$) in every $\hat A^{(k)}_t$ with $k \ge l + 1$, so its total weight is

Therefore

In plain words, GAE is a discounted sum of future TD errors, with discount $\gamma\lambda$ instead of $\gamma$. Each TD error is a small surprise, "this step went better (or worse) than the critic expected", and GAE adds up the surprises that follow an action, fading them out with distance.

The sum satisfies a recursion, the same one that returns satisfy, which is how it is always computed,

so one backward pass over a trajectory produces every advantage, in $O(T)$ time.

The two ends are the estimators we started with.

• $\lambda = 0$ gives $\hat A_t = \delta_t$, the one-step TD error. Low variance, biased unless $V = V^\pi$.
• $\lambda = 1$ gives $\hat A_t = \sum_l\gamma^l\delta_{t+l} = G_t - V(s_t)$, by the telescoping above. Unbiased for the gradient whatever $V$ is, high variance.

Everything in between trades one for the other smoothly. The effective horizon of the advantage is $1/(1 - \gamma\lambda)$ steps, about 17 for the common choice $\gamma = 0.99$, $\lambda = 0.95$.

Worked example. A three-step episode ends in a terminal state with rewards $(0, 0, 1)$. The critic predicts $V(s_0) = 0.5$, $V(s_1) = 0.6$, $V(s_2) = 0.8$, and the terminal state has value 0. With $\gamma = 0.9$ the TD errors are

With $\lambda = 0.5$, so $\gamma\lambda = 0.45$, the backward recursion gives $\hat A_2 = 0.2$, then $\hat A_1 = 0.12 + 0.45(0.2) = 0.21$, then $\hat A_0 = 0.04 + 0.45(0.21) = 0.1345$. With $\lambda = 1$ it gives $\hat A_0 = 0.04 + 0.9(0.12) + 0.81(0.2) = 0.31$, which matches the Monte Carlo estimate $G_0 - V(s_0) = 0.81 - 0.5$. With $\lambda = 0$ it is just $0.04$. The three answers differ by a factor of eight, which is the whole point. They make different bets on how much to trust the critic.

import numpy as np

def n_step_advantage(rewards, values, n, gamma):
    raise NotImplementedError

Two views of the same number

GAE has a twin in the older literature. The $\lambda$-return of Sutton's TD($\lambda$) (Sutton 1988; Sutton and Barto, chapter 12) averages $k$-step returns with the same weights, $G_t^\lambda = (1-\lambda)\sum_k\lambda^{k-1}G_t^{(k)}$. Since each $\hat A_t^{(k)} = G^{(k)}_t - V(s_t)$ and the weights sum to one,

So GAE is the $\lambda$-return with the baseline subtracted, and conversely $\hat A_t + V(s_t)$ is the $\lambda$-return. That second form is how PPO trains its critic. The value targets are $\hat R_t = \hat A_t + V_{\text{old}}(s_t)$, regression targets that interpolate between TD targets ($\lambda = 0$) and Monte Carlo returns ($\lambda = 1$) with the same dial.

What is the difference between the two knobs, $\gamma$ and $\lambda$? Both shrink the weight of the distant future and so reduce variance. But they bias different things. A $\gamma < 1$ changes the objective. It makes the agent care less about the far future, and it biases the gradient even with a perfect critic. A $\lambda < 1$ keeps the objective and only adds bias through the critic's errors, so with an accurate critic it is free. Schulman et al. (2016) found that the best $\lambda$ is typically well below the best $\gamma$, which their analysis attributes to $\lambda$ introducing much less bias than $\gamma$ for an accurate value function. On their cart-pole task the fastest learning came from $\lambda$ between 0.92 and 0.98 at $\gamma = 0.99$. PPO's defaults, $\gamma = 0.99$ and $\lambda = 0.95$, come from this line of work.

Bias and variance, measured

In the gridworld we can do what is impossible in a real problem, compute the true advantage of every state and action exactly and compare every estimate against it. Fix the uniformly random policy. Solve for $V^\pi$ and $Q^\pi$ exactly (§15.2). Then give the agent a wrong critic, $V = V^\pi + \text{noise}$, with independent Gaussian errors of standard deviation 0.2 in each state, sample 2,000 episodes, compute GAE for several values of $\lambda$, and compare each estimate with its target.

What is the right target? By §15.4, the critic's error at the current state is a baseline and harmless, so the honest target is $Q^\pi(s_t,a_t) - V(s_t)$, the true action value minus the critic's own baseline. Deviations from it are what actually reach the gradient. We split the mean squared error, for every visited state and action, into the squared bias of the average estimate and the variance around that average.

The table is the bias-variance trade-off in four columns. At $\lambda = 0$ half the error is bias, the critic's mistakes about the next state leaking straight into the estimate. As $\lambda$ grows the bias falls toward zero (what remains at $\lambda = 1$ is sampling noise in our measurement), while the variance grows more than fourfold. The total error is smallest in the middle, around $\lambda = 0.5$ for this critic. A better critic moves the best $\lambda$ down, toward trusting it more, and a worse one moves it up. The widget below lets you do exactly that.

Episode ends, termination, truncation, and rollout buffers

The recursion $\hat A_t = \delta_t + \gamma\lambda\hat A_{t+1}$ is easy. Getting the edges right is where real implementations break, and there are three different edges.

Termination. If $s_{t+1}$ is a terminal state (the goal, the pit, the end-of-sequence token), its value is zero by definition. So $\delta_t = r_t - V(s_t)$, with no bootstrap, and the recursion must not carry anything back from step $t+1$, which in a rollout buffer belongs to the next episode.

Truncation. If the episode was cut off at step $t$ by a time limit, the world did not end. The agent would have kept collecting reward. So the TD error must bootstrap from the value of the state it actually reached, $\delta_t = r_t + \gamma V(s_{t+1}) - V(s_t)$, but the recursion must still stop, because the next entry in the buffer is a fresh episode. Treating a time-out as a terminal state teaches the critic that the world ends when the clock runs out, a bug that Pardo et al. (2018) studied in detail. Gymnasium separates the two cases into terminated and truncated flags for exactly this reason.

The end of the rollout. PPO collects a fixed number of steps per environment, so the buffer usually ends in the middle of an episode. The last step's TD error bootstraps from the critic's value of the state after the last step, and the recursion starts there with $\hat A_{T} = 0$.

All three are handled by one convention. Store, for every step, the value of the true next state $V(s_{t+1})$ (computed before any reset), a terminated flag, and a truncated flag. Then

In plain words, termination kills the bootstrap and the recursion, truncation kills only the recursion, and the end of the buffer is a truncation that bootstraps from the last next-state value.

import numpy as np

def gae(rewards, values, next_values, terminated, truncated, gamma, lam):
    raise NotImplementedError

import numpy as np

def gae_weights(m, lam):
    raise NotImplementedError

In PyTorch

Production code computes GAE on tensors of shape (T, N) with exactly the backward loop above, vectorized over environments, under torch.no_grad(), because advantages are targets and must never carry gradients. Here it is next to our NumPy version on a random buffer.

Read more

• Schulman, Moritz, Levine, Jordan & Abbeel (2016). High-Dimensional Continuous Control Using Generalized Advantage Estimation. The GAE paper, with the $\gamma$-just analysis of bias and the experiments on $\gamma$ and $\lambda$.
• Sutton (1988). Learning to Predict by the Methods of Temporal Differences. TD($\lambda$) and the $\lambda$-return that GAE is built on.
• Sutton & Barto (2018). Reinforcement Learning: An Introduction, 2nd edition. Chapters 7 and 12 cover n-step returns and eligibility traces in depth.
• Pardo, Tavakoli, Levdik & Kormushev (2018). Time Limits in Reinforcement Learning. Why a time-out is not a terminal state, and what goes wrong when it is treated as one.
• Huang et al. (2022). The 37 Implementation Details of Proximal Policy Optimization. Detail 5 covers GAE, value bootstrapping, and the TD($\lambda$) value targets exactly as PPO implements them.

Watch next

▶

▶