Trust regions, PPO and RL for language models

Before this lesson: lesson 07 (the policy gradient theorem, REINFORCE and baselines), lesson 08 (advantages, GAE and the soft-greedy policy $\pi \propto \exp(Q/\alpha)$) and the importance ratio of lesson 04. KL divergence: information theory 03 and, for forward vs reverse KL, information theory 06. The Fisher matrix and natural gradient (optional): optimization 10. The full LLM training pipeline (SFT, reward models, DPO) is outside this track. Time: about 95 minutes, plus the notebook.

---

Why this matters #

• PPO is the default on-policy deep RL algorithm. It is simple, works with shared actor-critic networks, and reuses each batch for several epochs. It trained OpenAI Five (Dota 2) and many robotics and game agents.
• RLHF. InstructGPT (Ouyang et al., 2022) fine-tuned GPT-3 with PPO on the score of a reward model learned from human preferences, plus a per-token KL penalty to the supervised model. This recipe made chat assistants follow instructions.
• RL with verifiable rewards. GRPO (DeepSeekMath, 2024) removes the critic and uses a group of sampled answers as the baseline. As of 2025 it and its variants (Dr. GRPO, DAPO and others) are a standard way to train reasoning models such as DeepSeek-R1 on math and code, where a program checks the answer.
• One idea behind all of them: distribution shift. The trust region (TRPO, PPO), the KL penalty (RLHF) and the behaviour constraint of offline RL (lesson 06) all keep the policy where the data, or the reward model, can be trusted.

1. Predict first #

2. The picture #

Why big steps fail. Think of climbing a hill in fog. You can see the slope only near your feet. A small step uphill is safe. A long jump in the same direction can land you lower than before. In RL the "fog" is the data: every sample was collected by the old policy $\pi_{\text{old}}$. To judge a new policy $\pi_\theta$ with old samples, we reweight each sample by the probability ratio

the importance weight of lesson 04. This corrects for the actions. It does not correct for the states: a new policy also visits other states, and the old batch has no information about how often. This error is small for small steps and large for large steps.

A small example where this matters. An agent has two states, $s_1$ (the start) and $s_2$, and three actions. It cannot tell the two states apart (they look the same to it), so one softmax policy $\pi = \mathrm{softmax}(\theta)$ with $\theta \in \R^3$ acts in both. This is the normal situation with function approximation: one parameter vector serves many states. Rewards and next states, with $\gamma = 0.9$:

	$a_1$ (forward)	$a_2$ (wait)	$a_3$ (back)
$s_1$	$+1$, go to $s_2$	$+0.2$, stay in $s_1$	$0$, stay in $s_1$
$s_2$	$-1$, stay in $s_2$	$0$, stay in $s_2$	$0$, go to $s_1$

The old policy is uniform ($\theta_{\text{old}} = 0$). It spends a lot of time in $s_2$, where "back" is very good (advantage $+0.96$). So the policy gradient says: more "back", less "forward". But "back" is slightly bad in $s_1$ (advantage $-0.09$). A policy that almost always goes back never reaches $s_2$ again, so it spends all its time in $s_1$, where "back" earns nothing.

Interactive figure: loading…

This answers question 1. The first small steps help almost as predicted: at $t = 0.25$ the surrogate predicts $0.46$ and the true gain is $0.38$. But the true gain peaks near $t = 0.72$ (about $0.62$) and becomes negative after $t \approx 2.03$, while the surrogate keeps growing to above $2.5$. A trust region of size $\delta = 0.05$ stops at $t = 0.56$, with a true gain of $0.60$, close to the best possible on this line.

PPO's rule, as a picture. PPO does not compute any KL during the update. For each sample it looks at $\rho_t$ and the advantage estimate $\hat A_t$ (from GAE, lesson 08). If $\hat A_t > 0$, the update wants to make the action more likely, so $\rho_t$ grows. Once $\rho_t > 1 + \varepsilon$, PPO stops rewarding further growth: that sample's gradient becomes zero. If $\hat A_t < 0$, the same happens once $\rho_t < 1 - \varepsilon$. Here $\varepsilon$ (typically $0.2$) is the clip range; it is not the exploration rate of earlier lessons.

Interactive figure: loading…

3. Definitions and main results #

3.1 Setup and notation #

We use the notation of the track. $J(\pi) = \mathbb E_{s_0 \sim p_0}[V^\pi(s_0)]$ is the expected discounted return from the start distribution $p_0$. $A^\pi = Q^\pi - V^\pi$ is the advantage. The discounted visitation of a policy is

It is unnormalised; $(1 - \gamma) d^\pi$ is a probability distribution. Three policies appear in this lesson. Keep them apart:

Symbol	What it is	Where it appears
$\pi_\theta$	the policy we are training	everywhere
$\pi_{\text{old}}$	the policy that collected the current batch; fixed during one update	in the ratio $\rho_t$
$\pi_{\text{ref}}$	a frozen copy of the starting model (in LLMs: the supervised model)	in the KL penalty of RLHF

3.2 The performance difference lemma #

The first result says exactly how much better a new policy $\pi'$ is than an old policy $\pi$, using the old advantages.

Theorem (performance difference lemma; Kakade and Langford, 2002). For any two policies $\pi$ and $\pi'$, with bounded rewards and $\gamma < 1$,

In words: run the new policy, and add up the old policy's advantages of the actions it takes.

Why this is true (intuition). $A^\pi(s, a)$ measures how much better action $a$ is than following $\pi$ from $s$. The new policy is a sequence of such "deviations" from $\pi$, one per step. Each deviation changes the return by its advantage, and the changes add up.

Proof.

1. Fix a start state $s_0$ and a trajectory $s_0, a_0, r_0, s_1, \dots$ of $\pi'$. Because $V^\pi$ is bounded and $\gamma < 1$, the sum $\sum_{t \ge 0} \gamma^t \big(\gamma V^\pi(s_{t+1}) - V^\pi(s_t)\big)$ telescopes (neighbouring terms cancel) and equals $-V^\pi(s_0)$. Write out the first terms to see it: $(\gamma V_1 - V_0) + (\gamma^2 V_2 - \gamma V_1) + (\gamma^3 V_3 - \gamma^2 V_2) + \dots = -V_0$.
2. $V^{\pi'}(s_0) = \mathbb E_{\tau \sim \pi'}[\sum_t \gamma^t r_t]$ is the expected return of $\pi'$. Add the zero-sum identity of step 1 inside the expectation: $V^{\pi'}(s_0) - V^\pi(s_0) = \mathbb E_{\tau \sim \pi'}\big[\sum_{t \ge 0} \gamma^t \big(r_t + \gamma V^\pi(s_{t+1}) - V^\pi(s_t)\big)\big]$.
3. Condition on $(s_t, a_t)$. The next state comes from the environment, $s_{t+1} \sim P(\cdot \mid s_t, a_t)$, which does not depend on the policy. So $\mathbb E[r_t + \gamma V^\pi(s_{t+1}) \mid s_t, a_t] = Q^\pi(s_t, a_t)$, and term $t$ has expectation $\mathbb E[\gamma^t (Q^\pi(s_t, a_t) - V^\pi(s_t))] = \mathbb E[\gamma^t A^\pi(s_t, a_t)]$.
4. Average over $s_0 \sim p_0$. Then group the terms by state: the weight of state $s$ is $\sum_t \gamma^t \Pr_{\pi'}(s_t = s) = d^{\pi'}(s)$, and the action there is drawn from $\pi'(\cdot \mid s)$. $\blacksquare$

The lemma is exact, but it has a problem: it needs states from $d^{\pi'}$, the new policy's visitation. We cannot sample those before we have chosen $\pi'$.

3.3 The surrogate objective and TRPO #

Replace $d^{\pi'}$ by $d^\pi$ and write the inner sum with importance weights. This gives the surrogate objective

Every term can be estimated from the old batch: states and actions come from $\pi$, and $\pi'$ only appears in the ratio. In practice $L$ is the sample mean of $\rho_t \hat A_t$ (up to the constant $1/(1-\gamma)$).

Theorem (first-order agreement). For a differentiable policy $\pi_\theta$ with $\pi = \pi_{\theta_{\text{old}}}$: $L_\pi(\pi_{\theta_{\text{old}}}) = 0 = J(\pi_{\theta_{\text{old}}}) - J(\pi)$, and $\nabla_\theta L_\pi(\pi_\theta)\big|_{\theta_{\text{old}}} = \nabla_\theta J(\pi_\theta)\big|_{\theta_{\text{old}}}$.

Proof.

1. Value: $\sum_a \pi(a \mid s) A^\pi(s, a) = V^\pi(s) - V^\pi(s) = 0$ for every $s$, so $L_\pi(\pi) = 0$.
2. Gradient: $\nabla_\theta L = \sum_s d^\pi(s) \sum_a \nabla_\theta \pi_\theta(a \mid s)\, A^\pi(s, a)$.
3. Since $\sum_a \pi_\theta(a \mid s) = 1$, we have $\sum_a \nabla_\theta \pi_\theta(a \mid s) = 0$, so we may replace $A^\pi$ by $Q^\pi = A^\pi + V^\pi(s)$ without changing the sum.
4. $\sum_s d^\pi(s) \sum_a \nabla_\theta \pi_\theta(a \mid s) Q^\pi(s, a)$ is exactly $\nabla_\theta J$ by the policy gradient theorem of lesson 07. $\blacksquare$

Why replacing $d^{\pi'}$ by $d^\pi$ is a second-order error. The difference between the lemma and the surrogate is $\sum_s \big(d^{\pi'}(s) - d^\pi(s)\big) \sum_a \pi'(a \mid s) A^\pi(s, a)$. Both factors are zero when $\pi' = \pi$, and both grow linearly with a small change of $\theta$. Their product grows like the square of the step. So $L$ is accurate near $\pi$ and can be badly wrong far away, as the first figure showed.

Theorem (TRPO lower bound; Schulman et al., 2015, Theorem 1). With $C = \dfrac{4 \gamma \max_{s,a} |A^\pi(s,a)|}{(1-\gamma)^2}$,

The right side is $0$ at $\pi' = \pi$. So if each update maximizes the right side, the true $J$ never decreases (a minorize-maximize scheme). We do not prove the bound here; the proof is in the TRPO paper (§3 and Appendix A).

In practice $C$ is huge (in the example above $C = 4 \cdot 0.9 \cdot 0.98 / 0.01 \approx 353$), and the penalized steps would be tiny. TRPO therefore solves a constrained problem with a chosen radius $\delta$ (for example $0.01$):

Here $\hat{\mathbb E}_t$ is the average over the batch, and the max over states is replaced by an average.

The KL is locally a quadratic form. Near $\theta_{\text{old}}$, with $\Delta = \theta - \theta_{\text{old}}$,

where $F$ is the Fisher information matrix. The value and the gradient of the KL are zero at $\Delta = 0$, so the first term of its Taylor series is the quadratic one, and its Hessian is $F$. For a softmax over logits, $\nabla_\theta \log \pi(a) = e_a - p$ and $F = \mathrm{diag}(p) - pp^\top$.

The natural gradient, conjugate gradient and line search

Maximize the linear model $g^\top \Delta$ ($g$ is the policy gradient) subject to $\tfrac12 \Delta^\top F \Delta \le \delta$. The Lagrangian condition gives $\Delta \propto F^{-1} g$, the natural gradient, and the constraint fixes the length:

$$\Delta = \sqrt{\frac{2\delta}{g^\top F^{-1} g}}\; F^{-1} g .$$

For a network with millions of parameters, $F$ cannot be stored. TRPO computes $x \approx F^{-1} g$ with about 10 steps of conjugate gradient, which needs only Fisher-vector products $Fv$ (two backward passes, no matrix). Because the quadratic model is only approximate, TRPO then runs a backtracking line search: it halves the step until the true sample KL is at most $\delta$ and the surrogate improves. For a softmax, $F$ is singular along $(1, \dots, 1)$ (adding a constant to all logits does not change $p$), so one uses the pseudo-inverse or a small damping $F + \lambda I$.

Numerical check. $\theta = (0.5, -0.2, 0.1)$, $\Delta = (0.03, -0.02, 0.01)$: the exact KL is $1.918 \cdot 10^{-4}$ and $\tfrac12 \Delta^\top F \Delta = 1.927 \cdot 10^{-4}$. The natural gradient and K-FAC are in optimization 10 and optimization 12.

3.4 PPO: the clipped surrogate #

TRPO needs second-order machinery. PPO (Schulman et al., 2017) gets most of its benefit with first-order SGD. With $\rho_t(\theta) = \pi_\theta(a_t \mid s_t) / \pi_{\text{old}}(a_t \mid s_t)$ and advantage estimates $\hat A_t$:

The derivative of $\rho_t \hat A_t$ with respect to $\log \pi_\theta(a_t \mid s_t)$ is $\rho_t \hat A_t$, because $\rho_t = \exp(\log \pi_\theta - \log \pi_{\text{old}})$. The four cases:

Case	Term chosen by the min	Gradient w.r.t. $\log \pi_\theta(a_t \mid s_t)$
$\hat A > 0$, $\rho \le 1 + \varepsilon$	$\rho \hat A$	$\rho \hat A$ (push up)
$\hat A > 0$, $\rho > 1 + \varepsilon$	$(1 + \varepsilon)\hat A$, a constant	$0$
$\hat A < 0$, $\rho \ge 1 - \varepsilon$	$\rho \hat A$	$\rho \hat A$ (push down)
$\hat A < 0$, $\rho < 1 - \varepsilon$	$(1 - \varepsilon)\hat A$, a constant	$0$

Zero-gradient condition. The gradient is zero only when the sample has already moved by more than $\varepsilon$ in the direction that helps ($\hat A > 0, \rho > 1 + \varepsilon$ or $\hat A < 0, \rho < 1 - \varepsilon$). A sample that moved the wrong way keeps its full gradient.

Pessimistic bound. Since $\min(x, y) \le x$, every term satisfies $\min(\rho \hat A, \operatorname{clip}(\rho)\hat A) \le \rho \hat A$, so $L^{\text{CLIP}} \le \hat{\mathbb E}[\rho \hat A]$. The clipped objective is a lower bound on the unclipped surrogate. It ignores gains beyond the clip range but keeps every loss. At $\rho = 1$ (the start of each update) the two agree, and the gradient of $L^{\text{CLIP}}$ is the ordinary policy gradient $\hat{\mathbb E}[\hat A_t \nabla \log \pi_\theta(a_t \mid s_t)]$.

The full PPO update. A PPO iteration:

1. Run $\pi_{\text{old}}$ for $T$ steps in $N$ parallel environments. Store $\log \pi_{\text{old}}(a_t \mid s_t)$ and $V_\phi(s_t)$.
2. Compute advantages with GAE and value targets $\hat R_t = \hat A_t + V_\phi(s_t)$ (lesson 08).
3. For $K$ epochs (typically 3–10), shuffle the batch into minibatches and take gradient steps on the loss below, with a value-loss weight $c_v$ (for example $0.5$) and an entropy bonus weight $c_e$ (for example $0.01$).
4. Set $\pi_{\text{old}} \leftarrow \pi_\theta$ and repeat.

The loss of step 3 is

Implementation details that matter (Huang et al., The 37 implementation details of PPO):

• Advantage normalization per minibatch: $(\hat A - \text{mean}) / (\text{std} + 10^{-8})$. The step size then does not depend on the reward scale.
• Value loss, often clipped like the policy, with its own weight. Entropy bonus: keeps the policy from becoming deterministic too early.
• Several epochs of shuffled minibatches: this data reuse is the reason to have a trust region.
• Approximate-KL monitoring and early stopping: stop the epochs when the estimated KL passes a target (for example $0.01$–$0.02$). The clip does not bound the KL (see below).
• Global gradient-norm clipping (for example $0.5$) and learning-rate annealing.

Estimating the KL from samples. During the update we only have the sampled actions, drawn from $\pi_{\text{old}}$. With $\rho = \pi_\theta(a \mid s) / \pi_{\text{old}}(a \mid s)$:

Both are unbiased for $\mathrm{KL}(\pi_{\text{old}} \,\|\, \pi_\theta)$. For $k_1$ this is the definition: $\mathrm{KL}(\pi_{\text{old}} \,\|\, \pi_\theta) = \mathbb E_{a \sim \pi_{\text{old}}}[\log(\pi_{\text{old}}/\pi_\theta)]$. For $k_3$, add $\mathbb E_{a \sim \pi_{\text{old}}}[\rho - 1] = \sum_a \pi_\theta(a \mid s) - 1 = 0$. Each $k_1$ term can be negative, but each $k_3$ term is $\ge 0$, because $\log x \le x - 1$. So $k_3$ usually has lower variance (Schulman's note "Approximating KL divergence").

The clip does not bound the KL. The clip only removes the incentive for samples that are already outside the range in the helpful direction. Other samples keep moving the shared parameters, and those changes also move the clipped samples and the actions that were never sampled. Ratios often end far outside $[1 - \varepsilon, 1 + \varepsilon]$, and the KL keeps growing with more epochs or a larger learning rate. The notebook shows this on a 10-action problem: after 10 epochs with $\varepsilon = 0.2$, over half of the ratios are outside $[0.8, 1.2]$, the largest is $1.63$, and the KL is more than ten times its value after the first epoch.

3.5 RL for language models #

The token-level MDP. A language model $\pi_\theta(y_t \mid x, y_{<t})$ generates a response $y$ to a prompt $x$ one token at a time. As an MDP:

• State $s_t = (x, y_{<t})$: the prompt and the tokens so far.
• Action $a_t = y_t$: the next token. The action space is the vocabulary (for example 100 000 tokens).
• Transition: deterministic, append the token. The episode ends at the end-of-sequence token or a length limit.
• Reward: usually one number $r(x, y)$ for the whole response, given at the last token, from a reward model or a checker. Usually $\gamma = 1$.

Because the transition is known and deterministic, the same problem can also be seen as a contextual bandit: the context is $x$, the action is the whole response $y$, and $\pi_\theta(y \mid x) = \prod_t \pi_\theta(y_t \mid x, y_{<t})$.

A reward model from preferences. Humans compare two responses $y_w$ (preferred) and $y_l$ to the same prompt. The Bradley–Terry model says $\Pr(y_w \succ y_l \mid x) = \sigma\big(r_\phi(x, y_w) - r_\phi(x, y_l)\big)$, with the logistic function $\sigma$. The reward model $r_\phi$ (a language model with a scalar head) is trained by minimizing $-\log \sigma\big(r_\phi(x, y_w) - r_\phi(x, y_l)\big)$ over the comparisons. Only reward differences are learned, so $r_\phi$ is defined up to a constant per prompt.

The KL-regularized objective (InstructGPT style):

with a KL coefficient $\beta > 0$. The penalty keeps the policy near $\pi_{\text{ref}}$, where the reward model was trained and is reliable, and keeps the text fluent.

The per-token reward, from the chain rule. Because $\log \pi(y \mid x) = \sum_t \log \pi(y_t \mid s_t)$,

The second line takes, at each step $t$, the expectation over the token $y_t$ given $s_t$ first.

The sequence KL is the expected sum of per-token log-ratios (the chain rule for KL). So PPO-based RLHF gives every token the reward

The sum over the response is $r(x, y) - \beta \log \dfrac{\pi_\theta(y \mid x)}{\pi_{\text{ref}}(y \mid x)}$. A token-level critic and GAE then turn these rewards into per-token advantages.

Interactive figure: loading…

The optimum of the KL-regularized objective has a closed form. Fix a prompt $x$ and allow any distribution $\pi$ over responses.

Theorem. Let $Z = \sum_y \pi_{\text{ref}}(y) \exp(r(y)/\beta)$. The unique maximizer of $F(\pi) = \mathbb E_{y \sim \pi}[r(y)] - \beta\, \mathrm{KL}(\pi \,\|\, \pi_{\text{ref}})$ is

and the maximum value is $F(\pi^*) = \beta \log Z$.

Proof.

1. From the definition of $\pi^*$, take logs: $r(y) = \beta \log \dfrac{\pi^*(y)}{\pi_{\text{ref}}(y)} + \beta \log Z$ for every $y$.
2. Put this into $F$: $F(\pi) = \sum_y \pi(y) \Big[\beta \log \dfrac{\pi^*(y)}{\pi_{\text{ref}}(y)} + \beta \log Z - \beta \log \dfrac{\pi(y)}{\pi_{\text{ref}}(y)}\Big]$.
3. The $\pi_{\text{ref}}$ terms cancel: $F(\pi) = \beta \log Z - \beta \sum_y \pi(y) \log \dfrac{\pi(y)}{\pi^*(y)} = \beta \log Z - \beta\, \mathrm{KL}(\pi \,\|\, \pi^*)$.
4. $\mathrm{KL}(\pi \,\|\, \pi^*) \ge 0$, with equality only for $\pi = \pi^*$ (Gibbs' inequality, information theory 03). $\blacksquare$

So the KL-regularized optimum reweights the reference model by $\exp(r/\beta)$. Small $\beta$: almost all mass on the best response (but only among responses that $\pi_{\text{ref}}$ can produce). Large $\beta$: almost $\pi_{\text{ref}}$. With a uniform reference this is the soft-greedy policy $\pi \propto \exp(Q/\alpha)$ of lesson 08, with $\alpha = \beta$: the KL to a uniform distribution is $\log |\mathcal A| - \mathcal H(\pi)$, so the KL penalty becomes an entropy bonus.

Step 1 also says: at the optimum, the penalized reward $r(y) - \beta \log(\pi^*(y)/\pi_{\text{ref}}(y))$ equals $\beta \log Z$ for every response. All responses look equally good, so every advantage is zero and the gradient vanishes.

Interactive figure: loading…

Direct preference optimization (DPO)

Step 1 of the proof can be read backwards: for the optimal policy, the reward is $r(x, y) = \beta \log \dfrac{\pi^*(y \mid x)}{\pi_{\text{ref}}(y \mid x)} + \beta \log Z(x)$. Put this into the Bradley–Terry model. The term $\beta \log Z(x)$ is the same for both responses and cancels:

$$\Pr(y_w \succ y_l \mid x) = \sigma\Big(\beta \log \frac{\pi^*(y_w \mid x)}{\pi_{\text{ref}}(y_w \mid x)} - \beta \log \frac{\pi^*(y_l \mid x)}{\pi_{\text{ref}}(y_l \mid x)}\Big).$$

DPO (Rafailov et al., 2023) fits a policy $\pi_\theta$ in place of $\pi^*$ directly to the preference data by minimizing $-\log \sigma(\cdot)$ of this expression. It needs no reward model, no sampling and no RL loop, only log-probabilities of fixed pairs. The cost: it learns only from the given pairs (it is an offline method), and it is exact only if the optimum is reachable.

3.6 Critic-free methods: RLOO and GRPO #

PPO for RLHF needs four large models in memory: the policy, the reference, the reward model and a value model about as large as the policy. The value model must assign credit to single tokens from one reward at the end, which is hard. LLMs offer a cheaper baseline: sample $K$ responses $y_1, \dots, y_K$ to the same prompt, with rewards $r_1, \dots, r_K$, and compare them.

• RLOO (REINFORCE leave-one-out; Ahmadian et al., 2024):

The baseline of response $i$ does not use $y_i$; it depends only on the other, independent samples. So $\mathbb E[b_i \nabla \log \pi(y_i)] = \mathbb E[b_i]\, \mathbb E[\nabla \log \pi(y_i)] = 0$: this is unbiased REINFORCE with a per-prompt baseline (lesson 07). In the paper the KL penalty is part of the sequence reward, $r_i = r(x, y_i) - \beta \log(\pi_\theta(y_i)/\pi_{\text{ref}}(y_i))$.

• GRPO (Group Relative Policy Optimization; Shao et al., DeepSeekMath, 2024):

the same advantage for every token of $y_i$. It is used inside a PPO-style clipped ratio per token, and the KL to $\pi_{\text{ref}}$ is a separate loss term, not part of the reward. It is estimated per token with the $k_3$ form, $\frac{\pi_{\text{ref}}}{\pi_\theta} - 1 - \log \frac{\pi_{\text{ref}}}{\pi_\theta}$. Dividing by the std is not unbiased: it gives more weight to prompts whose rewards vary little. As of 2025, variants differ: sample or population std, and Dr. GRPO (Liu et al., 2025) drops both the std and the division by the response length.

Interactive figure: loading…

Why they fit LLMs. The dynamics are known and deterministic, so sampling several responses per prompt is cheap and needs no simulator. The reward comes once per sequence, so a per-prompt Monte Carlo baseline is a good baseline. No value network saves memory and removes a hard regression problem. Ahmadian et al. report RLOO matching or beating PPO for RLHF.

RL with verifiable rewards (RLVR). For math and code, a program can check the final answer: the reward is 1 if correct and 0 if not. There is no reward model to over-optimize, only the risk of gaming the checker (for example, an answer format it accepts by mistake). GRPO with such rewards trained DeepSeek-R1 (2025).

3.7 Reward over-optimization #

The reward model is only a proxy for what people want. If we optimize it hard, the policy finds responses that score high but are not better, for example longer answers or flattering text (reward hacking, an example of Goodhart's law). Gao et al. (2022) measured this with a large "gold" reward model playing the role of humans and smaller proxy reward models. With the distance $d = \sqrt{\mathrm{KL}(\pi \,\|\, \pi_{\text{init}})}$ from the initial policy, the proxy reward keeps rising, while the gold reward rises and then falls. For RL they fit the gold reward as $R(d) = d\,(a - b \log d)$, with constants $a, b$ that depend on the reward-model size.

The KL penalty keeps the policy near $\pi_{\text{ref}}$, where the proxy is more reliable. It does not make the proxy correct. In Gao et al.'s setup, the penalty acted like early stopping: it changed how far training went, but not the gold reward reached at a given KL distance. Other protections: early stopping on held-out human or gold evaluations, larger or ensembled reward models, new preference data from the current policy, and verifiable rewards where possible.

Answers to section 1. (1) No. The true improvement agrees with the surrogate only to first order. In the example, the gain grows more slowly than the step, peaks, and then becomes negative, while the surrogate keeps rising. (2) Yes. Clipping removes the incentive only for samples that are already outside the range in the helpful direction; other samples and shared parameters keep moving the policy. That is why PPO monitors the KL and stops early. (3) Lower. The log-ratio of that response is positive, so $-\beta \log(\pi_\theta/\pi_{\text{ref}}) < 0$ is added to $0.8$.

4. Worked example #

(a) PPO with $\varepsilon = 0.2$ on three samples $(\rho, \hat A)$, with $1 - \varepsilon = 0.8$ and $1 + \varepsilon = 1.2$:

Sample	$\rho \hat A$	$\operatorname{clip}(\rho)\, \hat A$	min	Gradient w.r.t. its $\log \pi$
$(1.5, +1)$	$1.5$	$1.2 \cdot 1 = 1.2$	$1.2$ (clipped constant)	$0$
$(0.6, +1)$	$0.6$	$0.8 \cdot 1 = 0.8$	$0.6$ (unclipped)	$0.6/3 = 0.2$
$(1.1, -2)$	$-2.2$	$1.1 \cdot (-2) = -2.2$	$-2.2$ (inside the range)	$-2.2/3 = -0.733$

$L^{\text{CLIP}} = (1.2 + 0.6 - 2.2)/3 = -0.133$. The unclipped mean is $(1.5 + 0.6 - 2.2)/3 = -0.033$. The clipped value is lower, as the pessimistic bound says. The first sample already moved enough and is ignored. The second sample was good but became less likely; its gradient pushes $\rho$ back up.

(b) Per-token RLHF reward. A 4-token response has log-ratios $\log(\pi_\theta/\pi_{\text{ref}}) = (0.2, -0.1, 0.4, 0.0)$, $\beta = 0.05$, and reward-model score $0.8$.

• Token 1: $-0.05 \cdot 0.2 = -0.01$.
• Token 2: $-0.05 \cdot (-0.1) = +0.005$.
• Token 3: $-0.05 \cdot 0.4 = -0.02$.
• Token 4: $-0.05 \cdot 0 + 0.8 = 0.8$.

Sum: $-0.01 + 0.005 - 0.02 + 0.8 = 0.775$. Check: the sequence log-ratio is $0.2 - 0.1 + 0.4 + 0 = 0.5$, and $0.8 - 0.05 \cdot 0.5 = 0.775$.

(c) RLOO and GRPO. Three responses with rewards $(0.8, 0.2, 0.5)$; mean $\bar r = 0.5$.

• RLOO: $0.8 - \tfrac{0.2 + 0.5}{2} = 0.45$; $\ 0.2 - \tfrac{0.8 + 0.5}{2} = -0.45$; $\ 0.5 - \tfrac{0.8 + 0.2}{2} = 0$. Check: $\tfrac32 (0.3, -0.3, 0) = (0.45, -0.45, 0)$.
• GRPO with the population std $\sqrt{(0.09 + 0.09 + 0)/3} = \sqrt{0.06} = 0.245$: $(1.22, -1.22, 0)$. With the sample std $\sqrt{0.18/2} = 0.3$: $(1, -1, 0)$.

(d) The closed-form optimum. Three responses with $\pi_{\text{ref}} = (0.5, 0.3, 0.2)$, rewards $r = (0, 1, 2)$ and $\beta = 1$.

1. Weights $\pi_{\text{ref}}(y) e^{r(y)}$: $0.5 \cdot 1 = 0.5$, $\ 0.3 \cdot e = 0.8155$, $\ 0.2 \cdot e^2 = 1.4778$.
2. $Z = 0.5 + 0.8155 + 1.4778 = 2.7933$, so $\pi^* = (0.179, 0.292, 0.529)$.
3. $\mathbb E_{\pi^*}[r] = 0.292 + 2 \cdot 0.529 = 1.350$ and $\mathrm{KL}(\pi^* \| \pi_{\text{ref}}) = 0.323$.
4. Objective: $1.350 - 0.323 = 1.027 = \log Z = \log 2.7933$ ✓.
5. Penalized reward of response 1: $0 - \log(0.179/0.5) = 1.027$; of response 3: $2 - \log(0.529/0.2) = 1.027$. The same for all responses, as the theorem says.

With $\beta = 0.1$ the optimum is almost $(0, 0, 1)$ with KL $= 1.609 = \log 5$ (all mass on response 3, which had $\pi_{\text{ref}} = 0.2$). With $\beta = 10$ it is $(0.465, 0.308, 0.227)$, close to $\pi_{\text{ref}}$.

(e) Sequence KL = sum of per-token KLs. A toy model writes two tokens from $\{0, 1\}$. Policy: first token $(0.7, 0.3)$; second token $(0.9, 0.1)$ after a 0 and $(0.4, 0.6)$ after a 1. Reference: first $(0.5, 0.5)$; second $(0.5, 0.5)$ after a 0 and $(0.2, 0.8)$ after a 1.

1. First-token KL: $0.7 \log\tfrac{0.7}{0.5} + 0.3 \log\tfrac{0.3}{0.5} = 0.0823$.
2. Second-token KL after a 0: $0.9 \log\tfrac{0.9}{0.5} + 0.1 \log\tfrac{0.1}{0.5} = 0.3681$; after a 1: $0.4 \log\tfrac{0.4}{0.2} + 0.6 \log\tfrac{0.6}{0.8} = 0.1046$.
3. Chain rule: $0.0823 + 0.7 \cdot 0.3681 + 0.3 \cdot 0.1046 = 0.3713$.
4. Directly over the four sequences, with $\pi_\theta = (0.63, 0.07, 0.12, 0.18)$ and $\pi_{\text{ref}} = (0.25, 0.25, 0.10, 0.40)$: $\sum_y \pi_\theta(y) \log \tfrac{\pi_\theta(y)}{\pi_{\text{ref}}(y)} = 0.3713$ ✓.

5. In code #

NumPy checks of the worked examples. The notebook for this lesson has these cells and the exercises.

import numpy as np
from scipy.special import softmax

# (a) PPO clipped objective and its gradient w.r.t. log pi
rho, A, eps = np.array([1.5, 0.6, 1.1]), np.array([1.0, 1.0, -2.0]), 0.2
terms = np.minimum(rho * A, np.clip(rho, 1 - eps, 1 + eps) * A)
grad = np.where(rho * A <= np.clip(rho, 1 - eps, 1 + eps) * A, rho * A, 0.0) / len(rho)
print(terms.mean(), (rho * A).mean())      # -0.1333 -0.0333
print(grad)                                # [ 0.      0.2    -0.7333]

# (b) per-token RLHF reward
log_ratio, beta, score = np.array([0.2, -0.1, 0.4, 0.0]), 0.05, 0.8
r_tok = -beta * log_ratio
r_tok[-1] += score
print(r_tok, r_tok.sum())                  # [-0.01  0.005 -0.02  0.8 ] 0.775

# (c) RLOO and GRPO
r = np.array([0.8, 0.2, 0.5]); K = len(r)
print(K / (K - 1) * (r - r.mean()))        # [ 0.45 -0.45  0.  ]
print((r - r.mean()) / r.std())            # [ 1.2247 -1.2247  0. ]   (population std)

# (d) closed-form optimum of E[r] - beta KL(pi || pi_ref)
p_ref, rew, beta = np.array([0.5, 0.3, 0.2]), np.array([0.0, 1.0, 2.0]), 1.0
w = p_ref * np.exp(rew / beta); Z = w.sum(); p_star = w / Z
kl = np.sum(p_star * np.log(p_star / p_ref))
print(p_star, p_star @ rew - beta * kl, beta * np.log(Z))   # [0.179 0.2919 0.5291] 1.0272 1.0272
print(rew - beta * np.log(p_star / p_ref))                  # [1.0272 1.0272 1.0272]

# Fisher approximation of the KL for a softmax policy
theta, delta = np.array([0.5, -0.2, 0.1]), np.array([0.03, -0.02, 0.01])
p, q = softmax(theta), softmax(theta + delta)
F = np.diag(p) - np.outer(p, p)
print(np.sum(p * np.log(p / q)), 0.5 * delta @ F @ delta)   # 1.918e-04 1.927e-04

# (e) sequence KL = sum of per-token KLs (toy two-token model)
pol = {"first": np.array([0.7, 0.3]), "second": np.array([[0.9, 0.1], [0.4, 0.6]])}
ref = {"first": np.array([0.5, 0.5]), "second": np.array([[0.5, 0.5], [0.2, 0.8]])}
seq = lambda m, a, b: m["first"][a] * m["second"][a, b]
kl = lambda u, v: np.sum(u * np.log(u / v))
kl_seq = sum(seq(pol, a, b) * np.log(seq(pol, a, b) / seq(ref, a, b)) for a in (0, 1) for b in (0, 1))
kl_tok = kl(pol["first"], ref["first"]) + sum(pol["first"][a] * kl(pol["second"][a], ref["second"][a]) for a in (0, 1))
print(kl_seq, kl_tok)                      # 0.3713 0.3713

6. Where it shows up #

• Deep RL benchmarks and games. PPO with GAE is the default baseline in libraries such as Stable-Baselines3 and CleanRL, and trained OpenAI Five. Its clipped ratio also appears in multi-agent RL (MAPPO).
• RLHF. InstructGPT and many later chat models: a reward model from human comparisons, then PPO with a learned value model and a per-token KL penalty to the supervised model.
• Reasoning models and RLVR. GRPO and its 2025 variants (Dr. GRPO, DAPO) with rewards from answer checkers, unit tests or proof checkers; DeepSeek-R1 is the best-known example.
• Preference optimization without RL. DPO and its variants use the closed-form optimum of the KL-regularized objective to fit preferences directly.
• ML for science. Policy-gradient fine-tuning with a penalty for leaving a prior model steers generative models toward a computed property. REINVENT trains a molecule generator toward a scoring function (for example predicted activity) while keeping it close to the prior generator, and DDPO (Black et al., 2023) fine-tunes diffusion models with a PPO-style clipped ratio. The penalty has the same role as in RLHF: the property score is often a learned proxy, and the prior keeps samples realistic.

7. Common mistakes #

• Mixing up the three policies. $\pi_\theta$ is trained, $\pi_{\text{old}}$ generated the batch and appears in the ratio, $\pi_{\text{ref}}$ is the frozen starting model and appears in the KL penalty. In PPO-RLHF all three exist at the same time.
• "Clipping bounds the KL." It does not. Measure the approximate KL and stop the epochs early.
• Zero-gradient cases reversed. The gradient vanishes for $\hat A > 0, \rho > 1 + \varepsilon$ and for $\hat A < 0, \rho < 1 - \varepsilon$. A sample that moved the wrong way keeps its gradient.
• Forgetting that the surrogate uses the old states. The ratio corrects the action distribution, not the state distribution. This is why big steps are dangerous even with exact ratios.
• Std normalization of an all-equal group. If all $K$ rewards are equal, all GRPO and RLOO advantages are 0 (GRPO needs a small constant in the denominator to avoid $0/0$). Such prompts give no signal.
• Counting the KL twice. PPO-RLHF puts it in the per-token reward; GRPO adds it as a separate loss term. Use one or the other.
• Believing the KL penalty fixes reward hacking. It limits how far the policy moves; it does not make the reward model right.
• Reward timing. We write $r_t$ for the reward after $(s_t, a_t)$; Sutton and Barto write $R_{t+1}$. In the token MDP the score $r(x, y)$ belongs to the last token.

8. Check yourself #

1. Why can the surrogate use $d^\pi$ instead of $d^{\pi'}$?
   Answer

   Because it agrees with the true improvement to first order at $\pi' = \pi$: the same value (0) and the same gradient (the policy gradient). The error is a product of two terms that both vanish at $\pi' = \pi$, so it is second order in the step, and TRPO's theorem bounds it by a multiple of the maximum KL. So the surrogate is accurate inside a small trust region.
2. A sample has $\hat A = -1$ and $\rho = 0.75$, with $\varepsilon = 0.2$. Does it give a gradient?
   Answer

   No. $\rho \hat A = -0.75$ and $\operatorname{clip}(\rho)\hat A = 0.8 \cdot (-1) = -0.8$; the min is $-0.8$, the clipped constant. The probability of this bad action has already fallen by more than $\varepsilon$, in the helpful direction.
3. Show that the RLOO advantage equals $\frac{K}{K-1}(r_i - \bar r)$.
   Answer

   $\sum_{j \ne i} r_j = K\bar r - r_i$, so $r_i - \frac{K\bar r - r_i}{K - 1} = \frac{(K-1) r_i - K\bar r + r_i}{K - 1} = \frac{K(r_i - \bar r)}{K - 1}$.
4. Prove that $k_3 = (\rho - 1) - \log \rho$ is an unbiased estimator of $\mathrm{KL}(\pi_{\text{old}} \| \pi_\theta)$ with $a \sim \pi_{\text{old}}$, and that each term is non-negative.
   Answer

   $\mathbb E_{\pi_{\text{old}}}[\rho] = \sum_a \pi_{\text{old}}(a)\frac{\pi_\theta(a)}{\pi_{\text{old}}(a)} = 1$, so $\mathbb E[\rho - 1] = 0$ and $\mathbb E[k_3] = \mathbb E[-\log\rho] = \mathrm{KL}(\pi_{\text{old}} \| \pi_\theta)$. Each term is $\ge 0$ because $\log x \le x - 1$ for $x > 0$.
5. For the KL-regularized objective with $\pi_{\text{ref}} = (0.5, 0.5)$, rewards $(1, 0)$ and $\beta = 0.5$, find $\pi^*$.
   Answer

   Weights $0.5 e^{2} = 3.695$ and $0.5 e^{0} = 0.5$; $Z = 4.195$; $\pi^* = (0.881, 0.119)$. This is $\sigma(r_1/\beta - r_2/\beta) = \sigma(2)$ for the first response.
6. What does the KL penalty not fix in reward over-optimization?
   Answer

   The proxy reward is still wrong; the penalty only limits how far the policy moves from $\pi_{\text{ref}}$. Gao et al. found that it acted like early stopping, with no better gold reward at a given KL distance.

Recall cardsAnswer in your head first, then reveal. Add the cards to your review deck to see them again just before you would forget.

1. State the performance difference lemma.

   Show answer

2. What is the surrogate objective, and how does it differ from the exact improvement?

   Show answer

3. Write PPO's clipped objective.

   Show answer

4. For which samples is the PPO gradient zero?

   Show answer

5. Why does PPO still monitor the KL?

   Show answer

6. Give the $k_3$ estimator of $\mathrm{KL}(\pi_{\text{old}} \| \pi_\theta)$ and its two properties.

   Show answer

7. What is the per-token reward in PPO-based RLHF?

   Show answer

8. What is the maximizer of $\mathbb E_\pi[r] - \beta\,\mathrm{KL}(\pi \| \pi_{\text{ref}})$?

   Show answer

9. Compute the RLOO advantages for rewards $(1, 0)$.

   Show answer

10. What is the GRPO advantage, and what does Dr. GRPO change?

    Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

• "Walk me through PPO's loss." The ratio $\rho_t$; the clipped surrogate with the min (a pessimistic bound; zero gradient only after enough movement in the helpful direction); the value loss; the entropy bonus; advantage normalization; several epochs of minibatches; KL monitoring and early stopping.
• "Why not just take bigger policy-gradient steps?" The data come from the old policy. The performance difference lemma shows the surrogate uses the old state distribution, so it is only valid near the old policy; large steps can make the policy worse.
• "TRPO vs PPO?" TRPO: a hard KL constraint, natural gradient by conjugate gradient, line search; more principled, harder to implement, awkward with shared actor-critic networks. PPO: first-order, clipped ratio, simpler and usually as good.
• "PPO vs GRPO vs RLOO for LLMs?" PPO needs a value model and GAE for per-token credit. RLOO and GRPO use $K$ samples per prompt as the baseline; less memory, no critic to train, but more samples per prompt. RLOO is unbiased; GRPO's std division reweights prompts.
• "Derive the optimal policy of the KL-regularized objective." Write $F(\pi) = \beta \log Z - \beta\,\mathrm{KL}(\pi \| \pi^*)$ with $\pi^* \propto \pi_{\text{ref}} e^{r/\beta}$. Mention that DPO inverts this to get the reward from the policy.
• "What is reward hacking, and what do you do about it?" Over-optimizing a proxy. KL penalty and early stopping, better or ensembled reward models, fresh human data, verifiable rewards.

10. Practice #

The notebook for this lesson runs in Colab. It has five exercises, each with a check cell and a folded solution:

1. Write ppo_clip_objective and ppo_clip_grad, and check the gradient by finite differences.
2. Write approx_kl ($k_1$ and $k_3$); show by sampling that both are unbiased and that $k_3 \ge 0$.
3. Run PPO for several epochs on one batch of a 10-action softmax problem and track the KL and the fraction of ratios outside the clip range; compare with the unclipped objective.
4. Compute per-token KL rewards on the two-token toy model and check the chain rule (sequence KL = sum of per-token KLs).
5. Write rloo_advantages, grpo_advantages and the closed-form optimum; train the toy model with RLOO plus a KL penalty and check that it converges to $\pi_{\text{ref}} \exp(r/\beta)/Z$.

11. Watch and read #

L4 TRPO and PPO (Foundations of Deep RL Series)Pieter Abbeel · Foundations of Deep RL, lecture 4 · 25 min

The surrogate from importance sampling, the KL trust region and the Fisher approximation, then PPO's clipped objective.

Proximal Policy Optimization (PPO) for LLMs Explained IntuitivelyJulia Turc · 22 min

The token-level view of PPO for language models: value model, GAE, importance ratio and clipping.

DeepSeek's GRPO (Group Relative Policy Optimization) | Reinforcement Learning for LLMsJulia Turc · 23 min

From 14:10, how GRPO replaces the value model by a group baseline, and how it compares with PPO.

Direct Preference Optimization: Your Language Model is Secretly a Reward Model | DPO paper explainedAI Coffee Break with Letitia · 9 min

The closed-form optimum of the KL-regularized objective and how DPO turns it into a classification loss on preference pairs.

• Schulman, Levine, Moritz, Jordan, Abbeel, Trust Region Policy Optimization (2015): §2–3 (the performance difference identity and Theorem 1) and Appendix C (conjugate gradient).
• Schulman, Wolski, Dhariwal, Radford, Klimov, Proximal Policy Optimization Algorithms (2017): §3 (the clipped objective) and §5 (the algorithm).
• Kakade and Langford, Approximately Optimal Approximate Reinforcement Learning (ICML 2002): the original performance difference lemma.
• Huang et al., The 37 Implementation Details of Proximal Policy Optimization (ICLR blog track, 2022).
• John Schulman, Approximating KL Divergence: the $k_1$, $k_2$, $k_3$ estimators.
• OpenAI Spinning Up: TRPO and PPO pages, with pseudocode.
• Ouyang et al., Training language models to follow instructions with human feedback (InstructGPT, 2022): §3.5 (reward model and PPO with a per-token KL penalty).
• Nathan Lambert, Reinforcement Learning from Human Feedback (free online book): chapters on reward models, policy gradients (PPO, RLOO, GRPO) and direct alignment (DPO).
• Ahmadian et al., Back to Basics: Revisiting REINFORCE Style Optimization for Learning from Human Feedback in LLMs (RLOO, 2024).
• Shao et al., DeepSeekMath (GRPO, 2024), §4; Liu et al., Understanding R1-Zero-Like Training: A Critical Perspective (Dr. GRPO, 2025).
• Gao, Schulman, Hilton, Scaling Laws for Reward Model Overoptimization (2022).
• Rafailov et al., Direct Preference Optimization: Your Language Model is Secretly a Reward Model (2023).