Actor-critic and GAE

The last chapter left me with REINFORCE and a baseline, and a promise: the advantage $A(s,a)$ is the "right" thing to multiply the score function by, and any good estimate of it will do. This chapter makes good on that promise by building the estimate the whole modern stack uses. I will learn the baseline instead of guessing it (that is the "critic"), I will show that the choice of advantage estimator is really a single knob trading bias against variance, and I will derive Generalized Advantage Estimation (GAE), which is the specific setting of that knob PPO and its descendants reach for. The derivation is a small, pretty telescoping sum, and it is worth doing by hand once because it demystifies the one line of code (delta + gamma*lam*last) that every PPO implementation hides the whole idea inside. At the end I price what this costs on a 16GB card, because "learn the baseline" means "carry a second network," and on the baseline machine that is the difference between a run that fits and a run that OOMs.

Theory

From a guessed baseline to a learned critic

Recall where REINFORCE landed. The policy gradient with a state-dependent baseline $b(s)$ is

$${\nabla }_{\theta }J(\theta )={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}\left[\sum _t{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)(G_t-b(s_t))\right],\text{\hspace{1em}(6.1)}$$

where $G_t={\sum }_{l\ge 0}{\gamma }^l{r}_{t+l}$ is the discounted return from step $t$, and subtracting any function of the state alone leaves the gradient unbiased (the baseline term has zero expectation under the score function, which I proved in the REINFORCE chapter). The variance-minimizing baseline is close to the state value $V^{\pi }(s)={\mathbb{E}}_{\pi }[G_t\mid s_t=s]$, so the natural move is to estimate $V^{\pi }$ with a network $V_{\varphi }(s)$ and use it as the baseline. That network is the critic; the policy ${\pi }_{\theta }$ is the actor. The actor proposes actions, the critic grades the states the actor lands in, and the grade feeds back as the baseline in equation (6.1). With $b=V_{\varphi }$, the multiplier $G_t-V_{\varphi }(s_t)$ is an estimate of the advantage

$$A^{\pi }(s_t,a_t)=Q^{\pi }(s_t,a_t)-V^{\pi }(s_t),\text{\hspace{1em}(6.2)}$$

the amount by which taking $a_t$ beats the policy's own average behavior from $s_t$. Positive advantage, push the log-prob up; negative, push it down. That is the entire actor-critic idea in one sentence.

The critic is trained by regression. I have samples of the return $G_t$ (or a bootstrapped target, more on that in a moment), and I fit $V_{\varphi }$ to them by minimizing a squared error,

$${\mathcal{L}}_V(\varphi )=\frac{1}{2}\mathbb{E}[(V_{\varphi }(s_t)-{\hat{G}}_t{)}^2],\text{\hspace{1em}(6.3)}$$

where ${\hat{G}}_t$ is whatever target I have chosen. So there are two losses now, the policy loss (equation 6.1) and the value loss (equation 6.3), optimized together, and this two-headed structure is exactly what PPO carries into language-model training. Everything interesting about actor-critic comes down to what I plug in for ${\hat{G}}_t$, because that single choice sets both the critic's regression target and, through the advantage, the actor's gradient.

The bias-variance dial

Here is the tension at the heart of the whole chapter. I can estimate the advantage in two extreme ways, and they fail in opposite directions.

At one extreme, use the full Monte Carlo return: $A_t\approx G_t-V_{\varphi }(s_t)$. The return $G_t$ is an unbiased sample of $Q^{\pi }$, so this estimator is (nearly) unbiased. But $G_t$ is a sum of many random rewards over the whole rest of the trajectory, so its variance is enormous, and in a language model where the "trajectory" is a 500-token generation, that variance is brutal. Every token's advantage inherits the noise of every reward that came after it.

At the other extreme, bootstrap after one step. Use the one-step temporal-difference (TD) target ${\hat{G}}_t=r_t+\gamma V_{\varphi }(s_{t+1})$, which gives the advantage estimate

$${\delta }_t=r_t+\gamma V_{\varphi }(s_{t+1})-V_{\varphi }(s_t).\text{\hspace{1em}(6.4)}$$

This ${\delta }_t$ is the TD residual, and it is the single most important quantity in the GAE derivation, so give it a name and a box in your head. It has tiny variance, because it only involves one real reward $r_t$ and two value lookups. But it is biased: it trusts $V_{\varphi }$, which is wrong early in training, and that bias propagates. If the critic thinks a losing position is fine, ${\delta }_t$ inherits the delusion.

So one-step TD is low-variance and biased; full Monte Carlo is high-variance and unbiased. Everything in between is an $n$-step estimator. Roll out $n$ real rewards, then bootstrap with the critic:

$${\hat{A}}_t^{(n)}=\underset{n\text{-step return}}{\underbrace{r_t+\gamma r_{t+1}+\cdots +{\gamma }^{n-1}{r}_{t+n-1}+{\gamma }^n{V}_{\varphi }(s_{t+n})}}-V_{\varphi }(s_t).\text{\hspace{1em}(6.5)}$$

Small $n$ leans on the critic (low variance, higher bias); large $n$ leans on real rewards (high variance, lower bias); $n\to \infty$ recovers Monte Carlo. The knob is $n$, and picking a single $n$ is an awkward, discrete choice. GAE's contribution is to stop choosing and instead take a smooth, exponentially weighted average over all $n$ at once, controlled by a continuous parameter $\lambda \in [0,1]$. That is the dial done right.

GAE(λ): the telescoping sum

GAE defines the advantage estimate as the exponentially weighted average of every $n$-step estimator, with weight proportional to ${\lambda }^{n-1}$ and a normalizer $(1-\lambda )$ so the weights sum to one:

$${\hat{A}}_t^{\text{GAE}(\gamma ,\lambda )}=(1-\lambda )\sum _{n=1}^{\infty }{\lambda }^{n-1}{\hat{A}}_t^{(n)}.\text{\hspace{1em}(6.7)}$$

The claim, which I will now prove, is that this apparently infinite mixture collapses to a single discounted sum of TD residuals with the discount rate $\gamma \lambda$:

$$\boxed{{\hat{A}}_t^{\text{GAE}(\gamma ,\lambda )}=\sum _{l=0}^{\infty }(\gamma \lambda {)}^l{\delta }_{t+l}.}\text{\hspace{1em}(6.8)}$$

Two endpoints check the formula and fix the intuition. At $\lambda =0$, equation (6.8) keeps only the $l=0$ term, so ${\hat{A}}_t={\delta }_t$: pure one-step TD, minimum variance, maximum reliance on the critic (maximum bias if the critic is wrong). At $\lambda =1$, the discount becomes ${\gamma }^l$ and the sum telescopes back through equation (6.6) into ${\sum }_l{\gamma }^l{r}_{t+l}-V_{\varphi }(s_t)=G_t-V_{\varphi }(s_t)$: pure Monte Carlo, unbiased, maximum variance. So $\lambda$ is literally the bias-variance dial, sliding continuously between the two extremes I set up earlier, and the conventional working value $\lambda \approx 0.95$ sits close to the low-bias end while shaving off the worst of the variance. Note the two discounts do different jobs: $\gamma$ defines the actual objective (how much future reward is worth), while $\lambda$ is purely an estimator knob that does not change what I am optimizing, only how noisily I estimate its gradient.

What this means for a language model

For an LLM, the "trajectory" is a generation and the "steps" are tokens. The state $s_t$ is the prompt plus the tokens emitted so far, the action $a_t$ is the next token, and in the RLHF setting the reward is sparse: zero for every intermediate token and a single scalar (a reward-model score, or a verifier's verdict) at the end of the sequence. With per-token rewards all zero until the end, the TD residual ${\delta }_t=\gamma V_{\varphi }(s_{t+1})-V_{\varphi }(s_t)$ for interior tokens is entirely a difference of critic values, and GAE's job is to smear that terminal reward back across the tokens, credit-assigning it token by token through the learned value function. This is why PPO-for-LLMs genuinely needs a value model: without $V_{\varphi }$ there is nothing to bootstrap from, every interior ${\delta }_t$ would be zero, and all the credit would dump onto the final token. The critic is what turns one scalar at the end into a dense, per-token advantage signal. GRPO, two chapters from now, will refuse to pay for that critic and get its dense signal a different way, and understanding why GAE needs the value model is exactly what makes GRPO's shortcut legible.

Tooling

The tool that embodies GAE is the RLHF trainer, and the concrete reference implementation I lean on is Hugging Face's TRL PPOTrainer together with the AutoModelForCausalLMWithValueHead wrapper. That wrapper is worth dwelling on because it is where the actor-critic structure becomes a memory line item. It takes a causal LM and bolts a second output head onto the final hidden states: a small linear layer $V_{\varphi }:{\mathbb{R}}^{d_{\text{model}}}\to \mathbb{R}$ that reads the same transformer trunk the policy uses and emits a scalar value per token position. So the critic shares the body with the actor and adds only a thin head. That sharing is a deliberate memory economy, but it does not make the critic free, because PPO's optimizer still has to carry gradients and optimizer state for the value pathway, and many setups (and the original InstructGPT recipe) use a separate value network entirely, which doubles the trunk.

Inside the trainer, the GAE computation is exactly equations (6.4), (6.8), and (6.9): after generating responses and scoring them, TRL computes per-token rewards (the reward-model score placed at the final token, plus a per-token KL penalty against the reference policy that I will develop in the PPO and GRPO chapters), runs values = value_head(hidden_states), then loops backward accumulating lastgaelam = delta + gamma * lam * lastgaelam to fill an advantages tensor. The defaults you will see are gamma=1.0 and lam=0.95, and $\gamma =1$ is the usual LLM choice because a generation is short and finite and there is no reason to discount a reward that arrives 200 tokens away relative to one 10 tokens away. The returns tensor advantages + values becomes the regression target for the value loss (equation 6.3), which TRL clips just like the policy loss, a detail I will motivate in the next chapter.

Lab

The goal of this lab is to see the bias-variance dial move. I will build a tiny, fully controllable environment where the true advantage is known in closed form, compute GAE at several $\lambda$ values against a deliberately-wrong critic, and plot estimator bias against estimator variance as $\lambda$ sweeps from 0 to 1. The artifact is a CSV and a PNG showing the classic tradeoff curve, so that when a real PPO run later uses $\lambda =0.95$ you know exactly what that number is buying. No GPU is needed here; this is pure estimator arithmetic and it runs on CPU in a second.

This is a uv project. From the repo root:

uv init labs/gae-dial
cd labs/gae-dial
uv add numpy matplotlib

"""Visualize the GAE(lambda) bias-variance tradeoff on a known-answer MDP.

A short linear-chain MDP with deterministic rewards gives a closed-form true
advantage. We corrupt the critic with a state-dependent error, then measure how GAE's
bias (systematic offset from the true advantage) and variance (spread across
noisy reward rollouts) move as lambda sweeps 0 -> 1.
"""
import csv
from pathlib import Path

import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

OUT = Path("artifacts")
OUT.mkdir(exist_ok=True)

GAMMA = 0.99
T = 20                 # steps per episode
N_EPISODES = 4000      # rollouts to estimate variance
REWARD_NOISE = 1.0     # std of per-step reward noise
CRITIC_ERR = 0.5       # amplitude of the state-dependent error in V_phi

rng = np.random.default_rng(0)

def true_values(mean_rewards):
    """Exact V(s_t) = sum_{k>=0} gamma^k E[r_{t+k}] for the chain."""
    V = np.zeros(T + 1)
    for t in range(T - 1, -1, -1):
        V[t] = mean_rewards[t] + GAMMA * V[t + 1]
    return V

def gae(deltas, lam):
    """Backward recursion A[t] = delta[t] + gamma*lam*A[t+1]  (eq 6.9)."""
    A = np.zeros(len(deltas))
    running = 0.0
    for t in range(len(deltas) - 1, -1, -1):
        running = deltas[t] + GAMMA * lam * running
        A[t] = running
    return A

def main():
    mean_rewards = np.linspace(1.0, 0.2, T)      # deterministic reward means
    V_true = true_values(mean_rewards)
    # A *state-dependent* critic error, not a uniform offset. It is calibrated
    # at the two anchored states -- the start s_0 (every rollout begins there, so
    # it gets the most training signal) and the terminal s_T (value known to be
    # 0) -- and wrong by +CRITIC_ERR everywhere in between. A uniform offset
    # would cancel in every TD residual yet survive as the lone baseline term at
    # the Monte-Carlo limit, making |bias| *grow* with lambda; pinning s_0 (and
    # s_T) to the truth is what makes the MC estimate (lambda=1) unbiased while
    # bootstrap (lambda=0) still leans on the wrong interior values.
    err = np.full(T + 1, CRITIC_ERR)
    err[0] = 0.0
    err[T] = 0.0
    V_phi = V_true + err                          # critic wrong in the interior

    # True advantage at t=0 for the greedy step: A = Q - V. With deterministic
    # dynamics and this policy, the true A_0 is 0 in expectation (the baseline
    # is exactly right when V_true is used), so the *estimator's* mean minus 0
    # is its bias, and its spread across rollouts is its variance.
    lambdas = np.linspace(0.0, 1.0, 21)
    rows, biases, variances = [], [], []

    for lam in lambdas:
        est0 = np.empty(N_EPISODES)
        for e in range(N_EPISODES):
            rewards = mean_rewards + rng.normal(0, REWARD_NOISE, size=T)
            # delta_t = r_t + gamma V(s_{t+1}) - V(s_t), bootstrap with wrong V_phi
            deltas = rewards + GAMMA * V_phi[1:] - V_phi[:-1]
            est0[e] = gae(deltas, lam)[0]
        # Bias: mean estimate minus the true advantage against the *correct* V.
        true_A0 = 0.0
        bias = est0.mean() - true_A0
        var = est0.var()
        biases.append(bias)
        variances.append(var)
        rows.append({"lambda": round(lam, 3),
                     "bias": bias, "variance": var,
                     "abs_bias": abs(bias)})

    csv_path = OUT / "gae_tradeoff.csv"
    with csv_path.open("w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=["lambda", "bias", "variance", "abs_bias"])
        w.writeheader(); w.writerows(rows)

    fig, ax1 = plt.subplots(figsize=(7, 4.5))
    ax1.plot(lambdas, np.abs(biases), "o-", color="C3", label="|bias|")
    ax1.set_xlabel(r"$\lambda$"); ax1.set_ylabel("|bias|", color="C3")
    ax1.tick_params(axis="y", labelcolor="C3")
    ax2 = ax1.twinx()
    ax2.plot(lambdas, variances, "s-", color="C0", label="variance")
    ax2.set_ylabel("variance", color="C0")
    ax2.tick_params(axis="y", labelcolor="C0")
    ax1.set_title(r"GAE($\lambda$): the bias-variance dial")
    fig.tight_layout()
    png_path = OUT / "gae_tradeoff.png"
    fig.savefig(png_path, dpi=130)

    print(f"lambda=0.00  |bias|={abs(biases[0]):.3f}  var={variances[0]:.3f}")
    print(f"lambda=0.95  |bias|={abs(biases[19]):.3f}  var={variances[19]:.3f}")
    print(f"lambda=1.00  |bias|={abs(biases[-1]):.3f}  var={variances[-1]:.3f}")
    print(f"Artifacts: {csv_path.resolve()}  {png_path.resolve()}")

if __name__ == "__main__":
    main()

Run it:

uv run python gae_bias_variance.py

What you should see. As $\lambda$ climbs from 0 to 1, the |bias| curve falls toward zero (the critic's baked-in error drains out of the estimate) while the variance curve climbs steeply (more real, noisy rewards enter the sum). The two curves cross somewhere in the middle, and the region around $\lambda \approx 0.9$ to $0.97$ is visibly the sweet spot: most of the bias is already gone but the variance has not yet exploded, which is the empirical reason the field settled on $\lambda =0.95$. The printed lines give you three points on that curve to sanity-check against the plot, and the CSV lets you re-plot or fit the tradeoff yourself. Read this figure as the picture behind one hyperparameter you will otherwise copy on faith into every PPO and GRPO config in Part VII.