Chapter 5 · Book pages 142–150

PPO — Proximal Policy Optimization

Main read ≈ 20 minutes · dotted words have hover/tap definitions · ▸ boxes go deeper · finish with the test (100% to complete)

The big picture

Chapter 4 ended with a promise: stage 3 of the pipeline is PPO, and this chapter is the reason. Vanilla policy gradient has a fatal flaw — no constraint on step size. A single unlucky batch can push the policy into a region where it generates garbage; garbage gets low rewards; the next gradient makes things worse; the collapse is unrecoverable. PPO's entire contribution is a cheap way to take safe steps.

PPO is the workhorse behind GPT-4 and Claude — the algorithm that made RLHF practical at frontier scale. Everything here is chapter 3's machinery (policy gradient, advantage, GAE, KL anchor) plus one new idea: clip the probability ratio so no single batch can wreck the policy. Chapter 7 (GRPO) will keep the clip and delete the critic.

Stop 1 — Motivation: from TRPO to PPO

The first serious fix was TRPO (2015): constrain the KL divergence between the old and new policy so updates stay inside a trust region. It works perfectly — but requires expensive second-order optimization: the Fisher information matrix and conjugate gradients.

PPO (2017) achieves similar stability with a simple first-order clipped objective: 10× simpler to implement, works almost as well, and scales to distributed training trivially. That trade — almost-as-good stability for dramatically simpler math — is why PPO, not TRPO, trains frontier models.

Go deeper: why second-order hurts at 70B scale

TRPO's constraint D_KL(π_θold ‖ π_θ) ≤ δ is enforced exactly, but computing it needs curvature information (the Fisher matrix) and iterative solves (conjugate gradients) — machinery that is painful to distribute across hundreds of GPUs. PPO replaces the hard constraint with a soft, per-sample clip that any first-order optimizer (Adam over mini-batches) can apply. Book §5.1, 5.4.4, p. 142–143.

Stop 2 — The clipped objective

The core innovation. Define the probability ratio between the live policy and the frozen snapshot that collected the data, r_t(θ) = π_θ(a_t|s_t) / π_θold(a_t|s_t). PPO's objective:

L^CLIP(θ) = E_t [ min( r_t(θ)Â_t, clip(r_t(θ), 1−ε, 1+ε)Â_t ) ] (5.1)

The min operator creates a pessimistic bound. For a good action (Â > 0) we want to raise its probability — rÂ grows with r, but the clip caps the benefit at r = 1 + ε: "don't get greedy on one good example." For a bad action (Â < 0) we want to lower its probability — rÂ improves as r falls, but the clip caps the benefit at r = 1 − ε: "don't forget too aggressively based on one bad example."

Interactive · Clip explorer

Slide the ratio r and the advantage Â — watch r·Â, clip(r)·Â, and the pessimistic min, and see whether the gradient flows (the four cases of eq. 5.9).

Net effect: the policy changes by at most ±20% per update step (ε = 0.2). That single bound prevents both catastrophic collapse and overconfident specialization.

Go deeper: the four gradient cases of eq. 5.9

With L_t = min(r_tÂ_t, r̄_tÂ_t), the gradient is ∇r_t·Â_t when the unclipped term is smaller, and exactly 0 when the clipped term is smaller. Expanded: Â > 0 and r < 1 + ε → gradient flows (keep pushing up); Â > 0 and r ≥ 1 + ε → gradient zero (already increased enough); Â < 0 and r > 1 − ε → gradient flows (keep pushing down); Â < 0 and r ≤ 1 − ε → gradient zero (already decreased enough). Clipping removes the incentive to move, not the ability — a soft wall, not a hard one. Book §5.2, 5.4.5, p. 142–144.

Stop 3 — The full PPO loss

The clipped policy term is one of three ingredients:

L = L^CLIP − c₁(V_θ(s_t) − V_t^target)² + c₂H[π_θ(·|s_t)] (5.2)

The value loss (c₁ = 0.1) trains the critic to predict returns — clipped too, for stability. The entropy bonus (c₂ = 0.01) rewards the policy for staying spread out, preventing premature convergence to a deterministic policy — critical for exploration.

Go deeper: the full derivation in six steps (§5.4)

(1) Objective: maximize expected cumulative reward J(θ) (5.3). (2) Policy gradient theorem: ∇J = E[Σ ∇ log π·Â_t] (5.4) — the advantage replaces the full return to cut variance. (3) Importance sampling: PPO collects data with π_θold but updates π_θ; correcting the mismatch turns the objective into the surrogate L^CPI = E[r_t(θ)·Â_t] (5.5–5.7). (4) Unconstrained, one step can push r far from 1.0 — importance weights explode, the policy enters untested regions where the reward model is unreliable, collapse follows; TRPO's answer is the hard KL constraint. (5) PPO clips the surrogate instead (5.8–5.9). (6) Combine: θ_k+1 = θ_k + α·∇[L^CLIP − c₁L_VF + c₂H] (5.10–5.12). The book's summary: importance sampling enables data reuse across epochs, clipping keeps weights from going extreme, the min takes the more conservative of the two — "monotonic improvement with probability 1, using only first-order gradients" (the book's claim, §5.4.6). Book §5.4.1–5.4.6, p. 143–144.

Stop 4 — The rollout buffer

A rollout is one trajectory from the current policy — for RLHF, one prompt answered token by token until end-of-text, each token a "step". The rollout buffer stores, per token:

B = { (s_t, a_t, log π_θold(a_t|s_t), r_t, V(s_t)) } (5.13)

— state, action, the log-probability under the exact policy that generated it (for the ratio), reward, and the value baseline (for GAE). The buffer runs a strict three-phase clockwork: collect (a 70B model with batch 128 × 512 tokens produces up to 65K token-level transitions), train (GAE + K = 3–10 epochs of mini-batch descent on the clipped objective), purge (wipe it completely). Watch the cycle, and contrast it with a replay buffer:

Interactive · Rollout-buffer lifecycle vs replay buffer

Step through collect → train → purge and watch the ratio drift from 1.0; toggle to replay-buffer mode to see the off-policy alternative.

Rollout buffer (PPO, GRPO, on-policy): one batch, a few epochs, discarded — fresh data every cycle. Replay buffer (DQN, SAC, off-policy): millions of transitions stored indefinitely, randomly sampled, reused across many updates. The purge exists because stale data breaks the clip: r_t(θ) is only meaningful near 1.0.

Continuous purging is why PPO needs continuous generation — and why generation is 60–70% of wall-clock time. This is where chapter 2 pays rent: vLLM runs the generation phase with 256+ responses in parallel, memory-efficient batching, and prefix sharing — for GRPO's N = 8 responses per prompt, the prompt KV is computed once and shared across all 8. OpenRLHF and TRL use vLLM as the generation backend, with DeepSpeed/FSDP on the training workers.

Go deeper: the numbers of one collect phase

Batch = 128 prompts × 512 max tokens → up to ~65K token-level transitions per rollout. K epochs typically 3–10 over that single buffer before the purge. The generation bottleneck at 60–70% of wall-clock is the direct consequence of on-policy data going stale — every update cycle must start with fresh rollouts. Book §5.5.1–5.5.3, p. 144–146.

Stop 5 — PPO for RLHF: the full loop

The book's concrete walkthrough, one PPO step for a 70B chat model (Llama-3-70B policy, batch of 128 prompts, 512 max tokens): the reward model scores each response, per-token KL is measured against the frozen reference policy, GAE advantages are whitened to zero mean and unit variance, and the update runs 4 epochs of mini-batches of 16:

PPO end-to-end (book fig. 5.1): generate → score → KL → final reward → GAE → clipped update, repeat (p. 146–147).

Interactive · One PPO step for a 70B chat model

Walk the six steps with the book's actual numbers at each stage.

The final reward combines both signals — R = r_RM − 0.05 × mean_KL, the KL penalty subtracted from the reward model's score, placed only on the last token. The payoff line: the policy improves by ~0.005 win-rate per step; after 10K steps, a 5–10% absolute improvement over SFT. RLHF progress is thousands of small, safe steps — exactly what the clip is for.

Go deeper: tokenization pitfalls

Tokenization decides what a "step" is, and that creates three subtle issues. KL accounting: the same semantic content tokenized differently sums to different total KL — rare words split into more subwords collect a higher total penalty. Credit assignment: GAE assigns advantage per token position, but the model only truly "decides" at the first token of a word; later subword tokens are largely deterministic. Reward placement: reward only at the final token means credit must propagate backward through GAE — longer responses get a more diluted signal. Mitigations: normalize KL by sequence length, word-level reward shaping, or rewards at semantic boundaries. Book §5.6, p. 146.

Stop 6 — Logits and the two-network mechanics

PPO holds two parameter states with the same topology: the live policy π_θ, updated by backpropagation, and the old policy π_θold, a frozen snapshot anchoring one optimization cycle. During rollout, the frozen snapshot emits raw logits z_old (a vector of size |V| = 32K–128K), softmax turns them into probabilities, a token is sampled — and the log-probability log π_θold(a_t|s_t) is stored as a scalar in the buffer.

Why store log-probs? It avoids re-running the frozen network during optimization — saving one full forward pass per mini-batch, which is significant for a 70B model. During training, only the live policy is re-evaluated: z_new = f(s_t; θ), log π_θ = LogSoftmax(z_new)[a_t], and the ratio is recovered in log-space (5.18–5.20):

r_t(θ) = exp( log π_θ(a_t|s_t) − log π_θold(a_t|s_t) ) (5.20)

— exponentiating a difference instead of dividing raw probabilities, which would underflow or overflow numerically.

Go deeper: the weight lifecycle (table 5.1) and continuous actions

Across one cycle: at rollout start θ and θ_old are the same copy, so r = 1.0 by identity; as mini-batch steps accumulate, θ drifts from θ_old and r deviates (e.g. 1.06, 0.94); when clipping activates, r is trapped at 1 ± ε; at optimization end θ_old is discarded; the next cycle copies θ → θ_old and r resets to 1.0. For continuous action spaces (robotics, not LLMs), the network outputs a mean µ and standard deviation σ and log-probs come from the Gaussian log-PDF (5.22) — the ratio and clip then work identically. Book §5.7.1–5.7.5, p. 146–148.

Stop 7 — TRL implementation

The HuggingFace TRL library provides production-ready PPO for LLMs. The valuable part is the config — every hyperparameter from this chapter in one place (book's listing, values verbatim):

ppo_config = PPOConfig(
    learning_rate=1.5e-6,      # low LR for stability
    batch_size=128,            # prompts per step
    mini_batch_size=16,        # gradient accumulation unit
    ppo_epochs=4,              # epochs per batch (reuse data)
    gamma=1.0,                 # no discounting (single turn)
    lam=0.95,                  # GAE lambda
    cliprange=0.2,             # PPO epsilon
    cliprange_value=0.2,       # value function clipping
    vf_coef=0.1,               # value loss coefficient
    init_kl_coef=0.05,         # initial KL penalty
    target_kl=6.0,             # adaptive KL target
    whiten_rewards=True,       # normalize advantages
    gradient_accumulation_steps=4,
    max_grad_norm=1.0,
)

Around it: the model loads as AutoModelForCausalLMWithValueHead (a causal LM plus the critic's value head) with LoRA adapters (r = 64, α = 16, on the q/k/v/o projections), and the training loop per batch is generate (512 tokens, temperature 0.7, top-p 0.9) → score with the reward model → ppo_trainer.step(...), which handles KL, GAE, and clipping internally. Monitoring: stats["ppo/mean_scores"] and stats["ppo/policy/approx_kl"].

Read the config as the chapter in miniature: cliprange 0.2 is stop 2, vf_coef 0.1 is stop 3, ppo_epochs 4 with purge-by-design is stop 4, init_kl_coef 0.05 is stop 5's reward formula, max_grad_norm 1.0 is gradient clipping from the 70B walkthrough.

Go deeper: what the training loop hides

ppo_trainer.step(query_tensors, response_tensors, rewards) runs the whole stop-5 pipeline internally: per-token KL against the reference, final reward R = r_RM − 0.05 × mean_KL, GAE with λ = 0.95 and γ = 1.0, advantage whitening, then 4 epochs × mini-batches of 16 with ε = 0.2. The generation call maps to the vLLM workers in production setups (OpenRLHF, TRL integrations). Book §5.6, 5.8, p. 146, 149–150.

Stop 8 — Critical hyperparameters

The book's danger table — each parameter, its typical value, and the cost of getting it wrong:

Parameter	Typical	Effect of getting it wrong
`cliprange`	0.2	Too low: no learning. Too high: instability.
`init_kl_coef`	0.01–0.1	Too low: reward hacking. Too high: stuck at SFT.
`target_kl`	4–8	Adaptive controller target. Lower = conservative.
`ppo_epochs`	4	Too many: overfits to batch. Too few: wastes generation compute.
`learning_rate`	1–5 × 10⁻⁶	Too high: catastrophic forgetting.
`batch_size`	64–256	Larger = smoother gradients, more generation compute.
`temperature`	0.7–1.0	Lower: less exploration. Higher: noisier advantages.

Interactive · Hyperparameter danger board

Pick a parameter, slide it too low or too high, and see the book's stated consequence.

Two honest tensions to hold: target_kl is 6.0 in the book's TRL config but 4–8 in this table — the config picks one point inside the typical band, the table gives the band. And init_kl_coef 0.05 (config) sits in the middle of the 0.01–0.1 range: low enough to learn, high enough to suppress reward hacking.

The whole chapter cashes out here: PPO is not one algorithm but a dial board, and every dial trades the same two currencies — learning speed and stability. Catastrophic forgetting (LR too high) and "stuck at SFT" (KL too high) are the same failure at opposite ends, which is why the clip, the KL controller, and the tiny learning rate all exist. Book §5.9, p. 150.