Trust regions: from TRPO to PPO

I now have a policy gradient (Part V chapter 4), a learned baseline, and a low-variance advantage estimate from GAE (chapter 6). In principle I could just take gradient steps and be done. In practice, if I do that naively, the policy collapses, and it collapses in a way that is not a bug in my code but a fundamental feature of on-policy RL. This chapter is about why, and about the two fixes that define modern policy optimization: TRPO, which solves the problem correctly but expensively with a hard trust-region constraint, and PPO, which approximates that constraint with a clipped objective so simple you can implement it in five lines and so robust it became the default for essentially all of RLHF. I will derive the surrogate objective from the performance-difference lemma, sketch the TRPO constrained problem, and then do the PPO clip in full, including the piecewise case analysis that is the whole reason the clip works. The geometry of that clip is the payoff: once you see which way each piece pushes, PPO stops being a magic formula and becomes obvious.

Theory

Why a naive policy-gradient step destroys the policy

The policy gradient tells me a direction in parameter space, not a distance. The gradient ${\nabla }_{\theta }J$ is a statement about an infinitesimal step; it says nothing about how far I can move before the linear approximation it is built on stops being true. And in policy optimization the approximation breaks down viciously fast, for a reason specific to RL that supervised learning does not share: the data distribution depends on the parameters I am changing. In supervised learning, a too-large step gives a worse model on a fixed dataset, and the next batch corrects it. In RL, a too-large step gives a worse policy, that worse policy generates the next batch of data, and if the policy has moved somewhere terrible it collects terrible data and the estimate of the next gradient is computed on that garbage. There is no fixed dataset to fall back to. A single overlarge step can push the policy into a region where it emits degenerate outputs, those outputs get near-zero or misleading reward, the gradient signal collapses, and the run never recovers. This is the "falling off a cliff" failure, and it is why you cannot just crank the learning rate.

The deeper issue is that the natural distance in parameter space (Euclidean distance on $\theta$) has nothing to do with the distance in policy space (how different the resulting action distributions are). A tiny change in $\theta$ can enormously change ${\pi }_{\theta }$ if you are in a sensitive region, and a large change in $\theta$ can barely move ${\pi }_{\theta }$ elsewhere. So bounding $\Vert \Delta \theta \Vert$ is the wrong constraint. What I actually want to bound is how much the policy distribution moves, and the natural currency for that is the KL divergence $D_{\text{KL}}({\pi }_{\text{old}}\Vert {\pi }_{\theta })$. Keep the new policy inside a small KL ball around the old one and the data distribution cannot lurch, the advantage estimates I computed under ${\pi }_{\text{old}}$ stay approximately valid, and each step is a safe, monotonic-ish improvement. That KL ball is the "trust region," the region where I trust my local model of the objective. Everything below is machinery for optimizing inside it.

The surrogate objective and the performance-difference lemma

To optimize offline (using a batch of data from ${\pi }_{\text{old}}$ to take a step that produces ${\pi }_{\theta }$), I need to express the new policy's performance in terms of old-policy samples. The exact bridge is the performance-difference lemma (Kakade and Langford, 2002):

$$J({\pi }_{\theta })-J({\pi }_{\text{old}})={\mathbb{E}}_{s\sim d^{{\pi }_{\theta }}}{\mathbb{E}}_{a\sim {\pi }_{\theta }(\cdot \mid s)}[A^{{\pi }_{\text{old}}}(s,a)],\text{\hspace{1em}(7.1)}$$

where $d^{{\pi }_{\theta }}$ is the discounted state-visitation distribution under the new policy. The lemma is exact and beautiful, but it has a problem that is the crux of everything: the expectation is over states visited by ${\pi }_{\theta }$, the very policy I am trying to find, so I cannot sample from it yet.

The two approximations I just made (swapping $d^{{\pi }_{\theta }}$ for $d^{{\pi }_{\text{old}}}$, and trusting importance sampling with a ratio that could be anything) are both only valid near ${\pi }_{\text{old}}$. Push ${\pi }_{\theta }$ far and the surrogate $L(\theta )$ stops predicting the true $J({\pi }_{\theta })$, the ratio $r_t$ can blow up, and maximizing $L$ happily walks you off the cliff. So the surrogate is necessary but not sufficient; it must be constrained.

TRPO: the hard constraint

TRPO makes the trust region explicit. Maximize the surrogate subject to a hard KL constraint:

$$\underset{\theta }{\max }{\mathbb{E}}_t[r_t(\theta ){\hat{A}}_t]\text{subject to}{\mathbb{E}}_t\left[D_{\text{KL}}\right({\pi }_{\text{old}}(\cdot \mid s_t)\Vert {\pi }_{\theta }(\cdot \mid s_t)\left)\right]\le \delta .\text{\hspace{1em}(7.5)}$$

TRPO works and comes with a monotonic-improvement guarantee, but equation (7.6) is a lot of machinery: conjugate-gradient inner loops, Fisher-vector products, a line search, all per update. For a transformer with billions of parameters and a training loop I want to run on one GPU, this is both a memory and an engineering burden. PPO's entire pitch is: get 90% of the trust-region benefit with a first-order method that needs none of that.

PPO: clipping as a soft trust region

PPO throws away the explicit KL constraint and instead bakes the trust region into the objective itself, so that plain SGD/Adam cannot want to move too far. The clipped surrogate is

$$\boxed{L^{\text{CLIP}}(\theta )={\mathbb{E}}_t\left[\min \right(r_t(\theta ){\hat{A}}_t,\text{clip}(r_t(\theta ),1-ϵ,1+ϵ){\hat{A}}_t\left)\right],}\text{\hspace{1em}(7.7)}$$

where $\text{clip}(x,a,b)=\max (a,\min (x,b))$ and $ϵ$ is a small constant, typically $0.1$ to $0.2$. Two things are happening: the ratio is clipped to the interval $[1-ϵ,1+ϵ]$, and then the objective takes the minimum of the clipped and unclipped terms. The minimum is what makes it a lower bound (a pessimistic surrogate), and the direction of the clip flips depending on the sign of the advantage. That sign-dependence is the whole design, so I will take it apart case by case.

The geometry is worth stating in plain words because it is the mental model I carry into every PPO and GRPO run. Picture $\ell (r)$ as a function of the ratio $r$, with the old policy sitting at $r=1$. For a good action, the objective is a ramp going up and to the right that hits a flat ceiling at $r=1+ϵ$. For a bad action, it is a ramp going up and to the left (toward smaller $r$) that hits a flat floor at $r=1-ϵ$. In both cases there is a flat plateau beyond the boundary in the "desired" direction and a live slope on the near side and in the "undo" direction. Gradient ascent slides up the ramp until it reaches the plateau, then stops, per token, per sample. No Fisher matrix, no line search, no explicit KL, just a clamp and a min. The clip width $ϵ$ is the radius of the trust region measured in ratio space, and $ϵ=0.2$ meaning "don't let any single token's probability move by more than about 20% per update round" is a genuinely useful one-line summary.

Tooling

The tool is again TRL's PPOTrainer, now read through the lens of equation (7.7). One PPO iteration does four things: generate responses from the current policy (this snapshots ${\pi }_{\text{old}}$, whose log-probs are cached), score them with a reward model or verifier and compute per-token rewards including the KL-to-reference penalty, compute advantages ${\hat{A}}_t$ with GAE from the previous chapter, then take several gradient epochs over the batch optimizing $L^{\text{CLIP}}$ plus the clipped value loss plus an entropy bonus. The ratio $r_t(\theta )$ is computed as exp(new_logprob - old_logprob) per token, where old_logprob is the cached value from generation time and new_logprob comes from a fresh forward pass under the current $\theta$; this is why PPO does multiple gradient steps per rollout batch (the "$K$ epochs"), because after the first step $\theta \ne {\theta }_{\text{old}}$ and the ratio genuinely starts to move away from 1, which is the entire point of the clip.

Lab

The goal of this lab is to make the clip's geometry undeniable by plotting $\ell (r)$ itself, for both advantage signs, and overlaying the unclipped surrogate so you can see exactly where and why the objective flattens. This is the picture from the derivation, rendered from the same arithmetic PPO runs on every token. It is pure math, no GPU, no model, and it runs instantly. Seeing the two ramps-into-plateaus once is worth more than re-reading the min-of-clip formula ten times.

This is a uv project. From the repo root:

uv init labs/ppo-clip-geometry
cd labs/ppo-clip-geometry
uv add numpy matplotlib

"""Plot the PPO-clip per-sample objective as a function of the ratio r.

Reproduces the piecewise picture from the derivation: for A>0 a rising ramp
that ceilings at r=1+eps, for A<0 a rising-to-the-left ramp that floors at
r=1-eps. Overlays the unclipped surrogate r*A to show where the clip bites.
"""
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)
EPS = 0.2

def clipped_obj(r, A, eps=EPS):
    unclipped = r * A
    clipped = np.clip(r, 1 - eps, 1 + eps) * A
    return np.minimum(unclipped, clipped)

def main():
    r = np.linspace(0.0, 2.0, 400)
    fig, axes = plt.subplots(1, 2, figsize=(11, 4.5), sharey=False)

    for ax, A, title in [(axes[0], +1.0, r"$\hat{A} > 0$ (good action)"),
                         (axes[1], -1.0, r"$\hat{A} < 0$ (bad action)")]:
        ax.plot(r, r * A, "--", color="0.6", label=r"unclipped $r\hat{A}$")
        ax.plot(r, clipped_obj(r, A), color="C0", lw=2.4, label=r"$L^{CLIP}$")
        for x in (1 - EPS, 1.0, 1 + EPS):
            ax.axvline(x, color="0.85", lw=1, zorder=0)
        ax.set_title(title); ax.set_xlabel(r"ratio $r_t(\theta)$")
        ax.set_ylabel("per-sample objective"); ax.legend(loc="best")
        ax.annotate(r"$1-\epsilon$", (1 - EPS, ax.get_ylim()[0]),
                    ha="center", va="bottom", fontsize=8, color="0.4")
        ax.annotate(r"$1+\epsilon$", (1 + EPS, ax.get_ylim()[0]),
                    ha="center", va="bottom", fontsize=8, color="0.4")

    fig.suptitle(r"PPO-clip geometry ($\epsilon = 0.2$): the plateau is the trust region")
    fig.tight_layout()
    png = OUT / "clip_geometry.png"
    fig.savefig(png, dpi=130)

    # Print the four regimes as a table so the plateaus are checkable numerically.
    print("A>0:  r=0.9 -> L=%.2f   r=1.3 -> L=%.2f (ceiling at 1+eps)"
          % (clipped_obj(0.9, 1.0), clipped_obj(1.3, 1.0)))
    print("A<0:  r=1.1 -> L=%.2f   r=0.7 -> L=%.2f (floor at 1-eps)"
          % (clipped_obj(1.1, -1.0), clipped_obj(0.7, -1.0)))
    print(f"Artifact: {png.resolve()}")

if __name__ == "__main__":
    main()

Run it:

uv run python clip_geometry.py

What you should see. Two panels. Left ($\hat{A}>0$): a straight line rising with $r$ that abruptly goes flat at $r=1+ϵ=1.2$, with the dashed unclipped line continuing up past it, so the gap between them is the reward PPO is deliberately declining. Right ($\hat{A}<0$): a line rising as $r$ decreases, going flat at $r=1-ϵ=0.8$, again with the unclipped dashed line diverging below. The four printed numbers pin the plateaus: for $\hat{A}>0$, $r=1.3$ gives exactly $1+ϵ=1.2$ (clamped) while $r=0.9$ gives $0.9$ (unclamped); for $\hat{A}<0$, $r=0.7$ gives $-(1-ϵ)=-0.8$ (clamped) while $r=1.1$ gives $-1.1$ (unclamped, still being corrected). That asymmetry, live on the correction side and flat on the overshoot side, is the entire trust-region behavior of PPO in one figure, and it is the picture I want in your head when GRPO reuses this exact clip in the next chapter.