Proximal Policy Optimization for LLMs: Visual walkthrough

We assume only this: an LLM writes text one token at a time, and a separate reward model can look at a finished answer and give it a score. Everything else we earn below.

Step 1 — Text generation is a walk down a tree

WHAT. Before any math, let us agree on what the LLM is doing. It starts with a prompt. Then it picks the next token. Then, with that token added, it picks the next one. And so on until it stops.

WHY picture it as a tree? Because "training the model" will mean nudging the odds at each fork. If we can see the forks, we can see what we are nudging.

PICTURE.

Figure — Proximal Policy Optimization for LLMs

Figure s01 — a branching tree. The white node on the left is the prompt (the starting state). Three arrows fan out to the next possible tokens, labelled with their probabilities 0.6, 0.3, 0.1 (blue, white, pink). From the top token two more arrows branch onward, showing that after each choice a new fork appears. Caption: every circle is a state, every arrow an action (a next token), and the number on the arrow is how likely the model is to take it.

Every circle is a state $s$ — the prompt plus all tokens written so far. Every arrow leaving a circle is a possible action $a$ — the next token. The number on an arrow is the probability the model gives that token.

One answer is a path from the root down to a leaf. Because each choice happens in order, it is natural to number the steps: the token chosen at the first fork is $a_{1}$ from state $s_{1}$ , at the second fork $a_{2}$ from $s_{2}$ , and so on. The little subscript $t$ is just "which fork along the path" — a clock ticking once per token. We will use $t = 1, 2, \dots, T$ , where $T$ is the last token (the answer ends). This finite last step $T$ is why we call an LLM rollout episodic: every generation stops after a bounded number of tokens, so there is no infinite future to worry about.

Step 2 — What we want: a high score, on average

WHAT. A finished answer is the whole path $a_{1}, a_{2}, \dots, a_{T}$ . In general reinforcement learning each step could pay a reward $r_{t}$ , and the total pay of one path is the return $R = \sum_{t = 1}^{T} γ^{t - 1} r_{t} .$

Term by term:

$r_{t}$ — the reward handed out at step $t$ (at the $t$ -th token).
$\sum_{t = 1}^{T}$ — add up the rewards along the whole finite path, from the first token to the last, $T$ .
$γ$ — the discount factor, a number in $[0, 1]$ . It says how much a reward later in the path is worth compared to now: $γ = 1$ counts every step equally, $γ < 1$ shrinks far-future rewards. Because our episodes are short and finite ( $T$ tokens), the sum can't blow up, so RLHF for LLMs almost always uses $γ = 1$ — every token in the answer counts fully.
$R$ — the return: the total score of one complete answer.

The LLM special case. In RLHF the reward model only scores the whole finished answer, not each token. So every $r_{t} = 0$ except at the last token $t = T$ , where $r_{T} = r_{ϕ} (answer)$ . With $γ = 1$ the return collapses to a single number: $R = r_{ϕ} (answer) .$ That is why in the simplified objective below we can write the reward of a path as one score $r$ . We optimize its average over the paths the model actually takes:

$J (θ) = E_{τ \sim π_{θ}} [R (τ)]$

Term by term:

$τ$ (tau) — one whole path (a "trajectory"): the sequence of states and tokens $s_{1}, a_{1}, \dots, s_{T}, a_{T}$ .
$τ \sim π_{θ}$ — " $τ$ is drawn from the policy": we imagine actually letting the model write, following its own arrow-odds fork by fork.
$R (τ)$ — the return of that path (here, the reward-model score of the finished answer).
$E [\cdot]$ — the average (expected value): each path's return, weighted by how likely the model is to walk it.
$J (θ)$ — one number: "how good is this model on average?" Bigger is better.

WHY average, not best-case? Because the model is random — it samples arrows. We cannot control one lucky path; we can only shift the odds so that typical paths score higher.

PICTURE.

Figure s02 — a bar chart of seven possible answers (paths). Each blue bar's height is that path's return $R$ ; the small label under each bar is its probability. The dashed yellow line marks $J$ , the probability-weighted average height. Caption: tall bars are answers the reward model loves; $J$ is the height you'd expect if you let the model write blindfolded. Raising $J$ means fattening the tall bars and starving the short ones.

Step 3 — We cannot differentiate the dice — the log-trick

WHAT. To improve $J$ we want its gradient $\nabla_{θ} J$ : which way to nudge the weights $θ$ to make $J$ go up. But $J$ contains an average over paths the model chooses. You cannot smoothly differentiate "the model rolled this token."

WHY a trick is forced on us. The reward model is a black box — "a human liked this" has no derivative we can chase back through the discrete token choice. We need a way to turn "shift the sampling" into "shift a thing we can differentiate."

The escape is one algebra identity, the log-derivative trick: $\nabla_{θ} π_{θ} = π_{θ} \nabla_{θ} lo g π_{θ}$

Why is $lo g$ the right tool and not something else? Because differentiating a probability weight gives you back the weight itself times $\nabla lo g$ . That factor of "the weight itself" is exactly what turns the sum-over-paths back into an average, so we can estimate it by sampling. Applied per token along the path, it gives the policy-gradient estimate:

$\nabla_{θ} J = E [R (τ) \sum_{t = 1}^{T} \nabla_{θ} lo g π_{θ} (a_{t} ∣ s_{t})]$

Term by term:

$\nabla_{θ} lo g π_{θ} (a_{t} ∣ s_{t})$ — the direction in weight-space that makes the token chosen at step $t$ more likely.
$\sum_{t = 1}^{T}$ — do this at every fork along the path.
multiplied by $R (τ)$ — scaled by how good the whole path turned out.

PICTURE.

Figure s03 — one state node with two arrows leaving it. A thick blue arrow labelled "high $R$ : increase log-prob" points up-right; a thin pink arrow labelled "low $R$ : decrease log-prob" points down-right. Between them sits the formula $\nabla J = E [R \cdot \nabla lo g π_{θ}]$ . Caption: a good path pushes its tokens' probabilities up (thick = strong push), a bad path pushes them down (thin = weak push); the push size is set by the return. This is REINFORCE.

Step 4 — Raw reward is too jumpy — subtract a baseline

WHAT. In Step 3 we scaled every push by the raw return $R$ . Problem: if every answer scores somewhere around $+ 8$ , then every token gets shoved up hard, and the tiny differences that actually matter drown in the noise. We fix this by measuring each token against an expectation. First we need two quantities that describe "how good is my situation":

With those two defined, the advantage is simply their difference:

$A (s, a) = Q (s, a) - V (s)$

Term by term:

$Q (s, a)$ — score expected from taking action $a$ here, then continuing.
$V (s)$ — score expected from this state on average (the "sea level").
$A (s, a)$ — the surplus: how much better than average this action was. Positive = above average, negative = below.

WHY this is allowed and why it helps. Subtracting $V (s)$ does not change the average gradient — a baseline that doesn't depend on the action is mean-zero in the sum over actions. But it slashes the variance: now the signal is "better or worse than usual," a small number centered on zero, instead of "big positive number for everything." (In practice $A$ is estimated with GAE, which is also where the discount $γ$ from Step 2 reappears.) Because each token now carries its own per-step advantage, from here on we index by the step clock $t$ and write $A_{t} = A (s_{t}, a_{t})$ , $s_{t}$ , $a_{t}$ — the same subscript we set up in Step 1:

$\nabla_{θ} J = E [\sum_{t = 1}^{T} A_{t} \nabla_{θ} lo g π_{θ} (a_{t} ∣ s_{t})]$

PICTURE.

Figure s04 — the same seven actions as bars, but now a dashed yellow line cuts across them at the value $V (s)$ , the "sea level." Bars poking above the line are coloured blue (advantage $A_{t} > 0$ , "above average"); bars below are pink ( $A_{t} < 0$ , "below average"). Caption: instead of pushing on the raw height, we push on the height above or below sea level — same shape of learning, far less shaking.

Step 5 — Rollouts are expensive — reuse them with a ratio

WHAT. Generating full LLM answers is slow and costly. We'd like to collect a batch once using the current model — call it the old policy $π_{θ_{old}}$ — and then take several gradient steps on it. But once we step, our new $π_{θ}$ disagrees with the data-collecting $π_{θ_{old}}$ . To stay honest we reweight each token by the probability ratio:

$ρ_{t} (θ) = \frac{π _{θ} ( a _{t} ∣ s _{t} )}{π _{θ_{old}} ( a _{t} ∣ s _{t} )}$

Term by term:

$a_{t}, s_{t}$ — the token and state at step $t$ along a path we already generated (the clock from Step 1).
numerator $π_{θ} (a_{t} ∣ s_{t})$ — how likely the new model is to take the token we already saw.
denominator $π_{θ_{old}} (a_{t} ∣ s_{t})$ — how likely the old (data-collecting) model was.
$ρ_{t}$ — their ratio. $ρ_{t} = 1$ means "no change yet." $ρ_{t} > 1$ means the new model now likes this token more; $ρ_{t} < 1$ means less.

The objective becomes $J = E [ρ_{t} (θ) A_{t}]$ .

WHY the ratio and not the difference? Because this is exactly importance sampling: to average a quantity under one distribution using samples from another, you multiply by the ratio of the two probabilities. It is the only correction that keeps the estimate unbiased.

PICTURE.

Figure s05 — a semicircular dial (a speedometer). The centre needle points straight up at $ρ_{t} = 1.0$ (yellow, "models agree"). The left end reads 0.8 (pink, "abandon token"), the right end reads 1.2 (blue, "over-eager"). Caption: the ratio is a dial of how far the new model has swung away from the old one on this token; dead-centre means no change yet.

Step 6 — The danger, then the clip

WHAT (the danger first). With a good token ( $A_{t} > 0$ ), the objective $ρ_{t} A_{t}$ rewards making $ρ_{t}$ huge — the model can slam a token's probability toward 1 in a single greedy step. For a 7B-parameter LLM that is how you get catastrophic forgetting and reward hacking: the model wrecks its language to chase the score.

THE FIX. PPO clips the ratio to a band $[1 - ϵ, 1 + ϵ]$ and keeps the pessimistic choice:

Term by term:

$clip (ρ_{t}, 0.8, 1.2)$ — force $ρ_{t}$ back inside the band if it strays out.
$min (\cdot, \cdot)$ — take the smaller of greedy vs capped: always the more cautious estimate.
$ϵ$ — the band's half-width; how far one update is allowed to help before the brakes engage.

WHY the $min$ and not just the clip? The $min$ is what makes overshooting pointless in both directions. Watch it split by the sign of $A_{t}$ .

PICTURE — the two shapes.

Figure s06 — a plot with the ratio $ρ_{t}$ on the horizontal axis and the clipped surrogate value on the vertical. The blue curve ( $A_{t} > 0$ , good token) rises with $ρ_{t}$ until $ρ_{t} = 1.2$ , then goes flat — labelled "no reward past 1.2." The pink curve ( $A_{t} < 0$ , bad token) is flat until $ρ_{t} = 0.8$ and then keeps falling — labelled "still falls: correction allowed." Two dotted vertical lines mark the band edges 0.8 and 1.2. Caption: the clip is deliberately asymmetric — it flattens over-optimistic moves but never blocks a correction.

Good action, $A_{t} > 0$ (blue curve): rises with $ρ_{t}$ until $ρ_{t} = 1.2$ , then goes flat. Beyond the band the curve is horizontal → gradient zero → no reward for pushing further. The brake is engaged.
Bad action, $A_{t} < 0$ (pink curve): flat until $ρ_{t} = 0.8$ , then keeps falling. Crucially it does not flatten on the correcting side — the model is still free to keep lowering a bad token's probability.

So the clip is asymmetric on purpose: it stops over-optimistic moves but never blocks a correction.

Step 7 — Both worked cases, read off the curve

Example 1 — clip bites (good token). $A_{t} = + 2$ , $ρ_{t} = 1.6$ , $ϵ = 0.2$ .

greedy: $1.6 \times 2 = 3.2$
capped: $clip (1.6, 0.8, 1.2) = 1.2$ , so $1.2 \times 2 = 2.4$
$min (3.2, 2.4) = 2.4$ . We refuse the extra $0.8$ ; the gradient here is zero — this is the flat part of the blue curve past $1.2$ .

Example 2 — clip does not rescue (bad token). $A_{t} = - 3$ , $ρ_{t} = 0.5$ .

greedy: $0.5 \times (- 3) = - 1.5$
capped: $clip (0.5, 0.8, 1.2) = 0.8$ , so $0.8 \times (- 3) = - 2.4$
$min (- 1.5, - 2.4) = - 2.4$ . For negative advantage the $min$ picks the more negative value — the pessimistic one — so the model is still allowed (indeed encouraged) to keep dropping this bad token. This is the still-falling pink curve.

PICTURE.

Figure s07 — the same two curves as s06, now with two yellow dots marking the worked examples. The dot on the blue curve sits at $ρ_{t} = 1.6$ , height $2.4$ , labelled "Ex1: min(3.2,2.4)=2.4," on the flat plateau. The dot on the pink curve sits at $ρ_{t} = 0.5$ , height $- 2.4$ , labelled "Ex2: min(-1.5,-2.4)=-2.4," on the still-descending part. Caption: the arithmetic of the two examples is literally the same picture as the clip curves.

Step 8 — The clip is a local guard, so add a global leash

WHAT. The clip only flattens the curve one step at a time, one token at a time. Over many steps and many tokens the policy can still crawl a long way from where it started. We need a second, global restraint: a penalty on drifting away from a frozen reference model (the SFT model), measured by KL divergence.

Recall from Step 2 that the reward model only pays out at the final token $t = T$ . To write that cleanly we use the indicator bracket $[t = T]$ :

The reward actually fed to PPO per token becomes: $r_{t} = RM score, only at the last token r_{ϕ} (full response) [t = T] - per-token KL leash β lo g \frac{π _{θ} ( a _{t} ∣ s _{t} )}{π _{ref} ( a _{t} ∣ s _{t} )}$

Term by term:

$r_{ϕ} (full response) [t = T]$ — the reward model's score, switched on only at the final token $T$ (the whole answer is what gets judged).
$π_{ref}$ — the frozen reference (SFT) model; it never moves.
$lo g \frac{π _{θ}}{π _{ref}}$ — how much more the current model likes this token than the reference did. Big and positive = the model is getting greedy about something the reference was unsure of.
$β$ — how hard the leash pulls.

WHY. The reward model is only trustworthy near the SFT distribution — it was trained on answers from that neighborhood. Wander too far and the RM is easily fooled, which is reward hacking. The KL term keeps the model in-distribution. (DPO later folds this leash directly into the loss and skips the reward model entirely.)

Example 3 — the leash in numbers. RM score $r_{ϕ} = 1.0$ . On a token: $π_{θ} = 0.9$ , $π_{ref} = 0.3$ , so $lo g (0.9/0.3) = lo g 3 \approx 1.10$ . With $β = 0.2$ : penalty $= 0.2 \times 1.10 = 0.22$ . Effective reward $= 1.0 - 0.22 = 0.78$ .

PICTURE.

Figure s08 — a cartoon. A pink dog (the policy) stands between a blue square on the left (the frozen reference/SFT model, "home") and a yellow star on the right (the reward). A dashed white line — the KL leash, labelled $β$ — ties the dog back to home, while a yellow arrow shows the dog straining toward the treat. Caption: the reward pulls the model forward, the KL leash yanks it back before it bolts into nonsense. Clip = short reins on each step; KL = leash to home.

The one-picture summary

Figure s09 — a left-to-right pipeline of five boxes joined by white arrows: (1) reward $r_{ϕ}$ → (2) advantage $A = Q - V$ → (3) ratio $ρ = π / π_{o l d}$ → (4) clip $1 \pm ϵ$ → (5) KL leash to ref. Underneath runs the caption "small safe steps, stays in-distribution." This compresses the whole derivation: turn a score into an above-average signal, track how far the model moved, clip each step, and leash the whole run to the reference.

Reward score → advantage (subtract the critic's sea level) → probability ratio new/old → clip each step → leash the whole thing to the reference. That is PPO for LLMs, start to finish.

Recall Feynman: retell the whole walkthrough

Picture the model writing text as walking down a branching tree, choosing an arrow at each fork, one fork per token, until the answer ends after $T$ tokens. We want it to walk toward answers a judge likes, on average (Step 2) — and since only the finished answer is scored, the whole path earns one number at the end. We can't reach through the judge's opinion with a derivative, so we use a trick: nudge the log-odds of each chosen arrow, scaled by how good the answer was (Step 3). But raw scores are all big and similar, so we subtract the "usual" level $V (s)$ to get a clean "better-or-worse-than-average" signal — the advantage (Step 4). Because generating text is expensive, we reuse one batch by tracking a ratio of new-odds to old-odds at each token (Step 5). That ratio is dangerous: chasing it too hard wrecks the model, so we clip it — flat curve past the band for good tokens, but always free to keep fixing bad ones (Steps 6–7). The clip only guards each little step, so we add a leash back to the original SFT model measured by KL, keeping the model near where the judge is actually reliable (Step 8). Small safe steps, no running off — that is PPO.

Recall Active recall — cover the answers

Why can't we differentiate $J$ directly? ::: Because $J$ averages over discrete token samples the model chooses; we use the log-derivative trick to turn sampling into a differentiable gradient. In LLM-RLHF, when is a nonzero reward paid, and why can we use $γ = 1$ ? ::: Only at the final token $T$ (the reward model scores the whole answer); episodes are finite/short so the discounted sum can't blow up. What are $Q (s, a)$ , $V (s)$ and the advantage $A$ ? ::: $V (s)$ = expected return from a state; $Q (s, a)$ = expected return from taking action $a$ then continuing; $A = Q - V$ = how much better than average that action was. What does subtracting $V (s)$ do to the gradient, and why is it allowed? ::: It leaves the average gradient unchanged (a mean-zero baseline) but cuts variance — the signal becomes "better/worse than average" instead of a big number. On the clip curve, what happens past $ρ_{t} = 1 + ϵ$ for a good token? ::: The curve goes flat → gradient zero → no reward for pushing the ratio higher. Why does the clip alone not keep the model in-distribution? ::: It's a per-step, per-token guard; the policy can still drift far over many steps — the global KL leash to the reference fixes that.