05 · RL on Token Sequences

What you're learning

How an autoregressive language model is a Markov decision process, so that everything you built in Tier 1 — policy gradients, actor-critic, GAE, PPO — drops onto language generation with no new algorithm, only a new environment. By the end you should be able to write down the token-level MDP precisely, derive how a single terminal reward reaches the first token, explain exactly why we regularize toward a frozen reference, and map every symbol in the PPO loop onto a concrete LLM object.

This is the first lesson of Tier 2 (the bridge to language). Tier 1 gave you the optimizer; this lesson gives you the MDP it runs on. Lesson 04 (04-ppo) already foreshadowed this mapping in its §7 — here we make it rigorous and mechanistic.

---

1. Learning map

graph TD
    A["Tier 1 MDP<br/>(s, a, transition, reward)"] --> B["Recast: what is a<br/>'state' for an LLM?"]
    B --> C["Token-level MDP<br/>s_t = (x, y_&lt;t), a_t = y_t"]
    C --> D["Deterministic transition:<br/>append the token"]
    C --> E["Sequence prob factorizes:<br/>π(y|x)=Π π(y_t|x,y_&lt;t)"]
    C --> F["Reward is terminal &amp; sparse:<br/>RM/verifier at EOS"]
    F --> G["Credit assignment:<br/>how does r_T reach y_1?"]
    F --> H["KL-to-reference shaping<br/>r_t = RM·1[t=T] − β·KL_ref"]
    G --> I["Huge |V|, long T<br/>→ variance, need baselines"]
    H --> I
    D --> J["This is just PPO<br/>on this MDP"]
    E --> J
    I --> J
    J --> K["Four models in memory:<br/>policy, ref, RM, critic"]
    K --> L["Rollout = generation<br/>(dominates cost)"]
    L --> M["Next: 06 where RM comes from;<br/>later: drop critic (GRPO),<br/>drop RL (DPO)"]

    style C fill:#44a,color:#fff
    style J fill:#4a4,color:#fff
    style H fill:#a64,color:#fff

Prerequisites (assumed from Tier 1): the MDP formalism and returns (01-mdps-and-the-rl-objective), the policy-gradient theorem and REINFORCE (02-policy-gradients), actor-critic with GAE (03-actor-critic-and-gae), and PPO (04-ppo).

---

2. Why this matters: RL “for free” on language

The remarkable fact that makes RLHF possible: you do not need a new RL algorithm to fine-tune an LLM with rewards. PPO was designed for robots and Atari, where the environment is a physics simulator with stochastic dynamics.¹ Language has none of that — no simulator, no physics, no exogenous randomness in the transition. And yet, if you squint at decoding the right way, generating a response is exactly an episode in a (degenerate but perfectly valid) MDP.

The payoff of taking this seriously is threefold:

1. Reuse. Every theorem and every trick from Tier 1 — the policy-gradient theorem, advantage baselines, GAE’s bias–variance dial, PPO’s clipped trust region — applies verbatim once you identify the state, action, transition, and reward.
2. Diagnosis. The pathologies of RLHF (reward hacking, mode collapse, entropy craters, KL blowups) are not mysterious LLM phenomena. They are the standard RL failure modes viewed through a language lens, and the fixes are the standard RL fixes.
3. A ladder for what’s next. Once you see RLHF as “PPO on this specific MDP,” the later methods become legible as deliberate deletions: GRPO deletes the critic; DPO deletes the whole online RL loop. You cannot appreciate a deletion until you hold the full thing.

The trap this lesson exists to prevent

Beginners treat “RLHF” as one monolithic new technique. It is not. It is (a) a way of casting text generation as an MDP (this lesson), (b) a reward signal (lesson 06), and (c) an optimizer you already know (PPO, lesson 04). Keeping these three layers separate is the single most clarifying move in the whole subject.

---

3. The token-level MDP

Fix a pretrained/SFT autoregressive LM with parameters $\theta$ over a vocabulary $V$. A prompt $x$ is given. The model generates a response $y=(y_1,y_2,\dots ,y_T)$ one token at a time. We claim this generation process is an MDP $(\mathcal{S},\mathcal{A},P,R,{\rho }_0)$. Let’s fill in each element.

State. The state at step $t$ is the entire context the model conditions on:
$s_t=(x,y_{<t})=(x,y_1,\dots ,y_{t-1}).$
The initial state $s_1=(x)$ is drawn from the prompt distribution ${\rho }_0=\mathcal{D}$ (the dataset of prompts). Note the state grows — this is a variable-length state space, but it is genuinely Markov: the next-token distribution depends only on $s_t$, and $s_t$ literally contains the full history, so the Markov property holds trivially (indeed by construction, since the LM is a function of the whole prefix).

Action. The action is the next token drawn from the vocabulary:
$a_t=y_t\in \mathcal{A}=V,|V|\sim 32\text{k}\text{–}256\text{k}.$²

Policy. The policy is the autoregressive LM itself — its next-token softmax:
${\pi }_{\theta }(a_t\mid s_t)={\pi }_{\theta }(y_t\mid x,y_{<t}).$
No architectural change is needed. The thing that samples your chatbot’s next token is a stochastic policy over $\mathcal{A}=V$.

Transition. Here is the degeneracy that makes language special: the transition is deterministic — it just appends the chosen token.
$P(s_{t+1}\mid s_t,a_t)=1[s_{t+1}=(x,y_{<t},a_t)].$
There is no environment stochasticity, no dice roll after you act. All randomness in an episode comes from the policy’s sampling, not from $P$. (Contrast a robot: you command a torque, physics decides where you land.)

Termination. The episode ends when the policy emits the end-of-sequence token, $y_T=\text{EOS}$, or when a hard length cap $T_{\max }$ is hit. So the horizon is finite but variable and policy-controlled — the policy decides when to stop by choosing to emit EOS.

Reward. Deferred to §5, because it is the subtle part. For now: typically a single scalar delivered at the terminal step.

flowchart LR
    subgraph Episode["One generation episode = one trajectory"]
        direction LR
        S1["s_1 = (x)"] -->|"a_1 = y_1 ~ π_θ(·|x)"| S2["s_2 = (x, y_1)"]
        S2 -->|"a_2 = y_2 ~ π_θ(·|x,y_1)"| S3["s_3 = (x, y_1 y_2)"]
        S3 -->|"..."| ST["s_T = (x, y_<T)"]
        ST -->|"a_T = EOS"| TERM["terminal"]
    end
    TERM -->|"full response y_1:T"| RM["Reward model / verifier<br/>r_ψ(x, y) → scalar"]
    RM -->|"terminal reward r_T"| CREDIT["distributed over tokens<br/>via return + GAE"]

    style S1 fill:#334,color:#fff
    style RM fill:#a64,color:#fff
    style CREDIT fill:#343,color:#fff

Why the deterministic transition is a gift, not a technicality

Because $P$ is deterministic, the only source of variance in a return estimate is the policy’s own sampling and the reward. There is no environment noise to average out. This is why on-policy methods behave relatively benignly here, and why value estimation is “easy” in one sense — the value of a state is a pure function of the policy’s future token choices, with no aleatoric transition noise mixed in. It also means the discount $\gamma$ is usually set to $1$ (or very near it): there is no physical time-preference, and undiscounted return over a bounded response is well-defined.

Quiz 1: In this MDP, where does all the stochasticity of a trajectory come from, and what would happen to the "RL" if you decoded greedily (argmax) at every step?

Answer

All trajectory stochasticity comes from the policy’s sampling of $a_t\sim {\pi }_{\theta }(\cdot \mid s_t)$; the transition $P$ is deterministic (append), and there is no exogenous environment noise. If you decoded greedily, every episode from a given prompt would be identical — a single deterministic trajectory. Policy-gradient RL would then be estimating ${\nabla }_{\theta }{\mathbb{E}}_{y\sim {\pi }_{\theta }}[R]$ from a degenerate distribution with essentially zero exploration: you’d see one $y$ per prompt, the score-function estimator would have no diversity to learn a gradient direction from within a prompt, and you could not discover higher-reward continuations you don’t already rank first. This is exactly why rollouts are done with temperature/nucleus sampling — RL needs the policy’s stochasticity as its exploration mechanism.

---

4. Sequence probability and log-prob

Because generation is autoregressive, the probability the policy assigns to a whole response factorizes by the chain rule:
${\pi }_{\theta }(y\mid x)={\prod }_{t=1}^T{\pi }_{\theta }(y_t\mid x,y_{<t}).$

Taking logs turns the product into the quantity we actually manipulate — the sequence log-probability is a sum of per-token log-probs:
$\log {\pi }_{\theta }(y\mid x)={\sum }_{t=1}^T\log {\pi }_{\theta }(y_t\mid x,y_{<t}).$

This is not a cosmetic identity; it is the hinge that makes token-level RL work. Three consequences you will use constantly:

1. Per-token decomposability. A gradient on the sequence log-prob is a sum of gradients on token log-probs, each attached to its own state $s_t$. That is precisely the shape the policy-gradient theorem wants — a sum over timesteps of ${\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)$ weighted by an advantage.
2. A “trajectory-level” view and a “token-level” view are the same object. You can think of the policy as choosing a whole sequence (bandit view: one action = one response) or as choosing tokens step by step (MDP view). The factorization above is what lets you move between them. The token-level view is strictly more powerful because it exposes intermediate states for value estimation and credit assignment.
3. The PPO ratio is per token. The importance ratio from lesson 04 becomes $r_t(\theta )=\frac{{\pi }_{\theta }(y_t\mid x,y_{<t})}{{\pi }_{{\theta }_{\text{old}}}(y_t\mid x,y_{<t})}$, one ratio per generated token — computed cheaply from stored old log-probs.

Log-probs are the currency of the whole pipeline

Almost every quantity you compute in RLHF is a per-token log-prob or a difference of them: the PPO ratio needs $\log {\pi }_{\theta }-\log {\pi }_{{\theta }_{\text{old}}}$; the reference-KL penalty (§5) needs $\log {\pi }_{\theta }-\log {\pi }_{\text{ref}}$; DPO’s loss (later) is built entirely from summed log-prob differences. Get comfortable thinking in nats-per-token.

---

5. Reward structure: sparse, terminal, and shaped by KL

5.1 The raw reward is a single terminal scalar

In the canonical RLHF setup the reward comes from a learned reward model $r_{\psi }(x,y)$ (lesson 06 — where it comes from) or, in verifiable domains (math, code, unit tests), from a verifier returning e.g. $\{0,1\}$ for correct/incorrect. Either way, the score is defined on the whole response and is only available once the response is complete:
$R(s_t,a_t)=\{\begin{array}{ll}{r}_{\psi }(x,y_{1:T})& t=T(\text{terminal})\\ 0& t<T.\end{array}$

This is a sparse, trajectory-level reward. Every intermediate token earns $0$; the environment says nothing about whether token $y_5$ was good until the entire response is scored at the end. Contrast this with dense per-token rewards (imaginable: a reward at every token, e.g. a per-token quality signal), which are far easier for credit assignment but which we usually don’t have — human preferences and verifiers are naturally sequence-level.

5.2 How a terminal reward reaches early tokens

If only $y_T$ gets a reward, how does the first token ever learn? Through the return and the advantage, exactly as in Tier 1. With discount $\gamma$ (typically $1$ for LLMs), the return-to-go from step $t$ is
$G_t={\sum }_{k=t}^T{\gamma }^{k-t}R(s_k,a_k)={\gamma }^{T-t}{r}_{\psi }(x,y_{1:T}),$
so with $\gamma =1$ every token in the response shares the same return $G_t=r_{\psi }(x,y)$. The score-function gradient
${\nabla }_{\theta }J=\mathbb{E}[{\sum }_{t=1}^T{\nabla }_{\theta }\log {\pi }_{\theta }(y_t\mid x,y_{<t}){\hat{A}}_t]$
then pushes up the log-prob of every token in a high-reward response and down every token in a low-reward one. That is the crude REINFORCE-style credit path: the terminal signal is broadcast to all tokens through the return.

Broadcasting a single scalar to all tokens is high-variance and assigns credit bluntly (a great response with one bad token still reinforces the bad token). The critic + GAE³ sharpen this. A learned value head $V_{\varphi }(s_t)$ predicts the expected terminal reward given the prefix so far, and the TD residual localizes credit:
${\delta }_t=R(s_t,a_t)+\gamma V_{\varphi }(s_{t+1})-V_{\varphi }(s_t),{\hat{A}}_t^{\text{GAE}(\gamma ,\lambda )}={\sum }_{l=0}^{T-t}(\gamma \lambda {)}^l{\delta }_{t+l}.$
For $t<T$ the immediate reward is $0$, so ${\delta }_t=\gamma V_{\varphi }(s_{t+1})-V_{\varphi }(s_t)$ is purely a bootstrapped signal: it measures whether emitting $y_t$ raised or lowered the critic’s estimate of the eventual reward. In effect the critic converts the sparse terminal reward into a dense per-token advantage, letting a token be credited for how much it improved the prospects of the response rather than for the final score alone. The $\lambda$ knob trades the bias of trusting the critic against the variance of trusting the raw terminal return — exactly the bias–variance dial from lesson 03, now operating over a token sequence.

Quiz 2: With a purely terminal reward and $\gamma =1$, explain two distinct mechanisms by which the first token $y_1$ of a long response receives a learning signal — one without a critic and one with a critic — and why the second is usually preferred.

Answer

Without a critic (REINFORCE / return broadcast): with $\gamma =1$, the return-to-go for every token equals the single terminal reward $G_t=r_{\psi }(x,y)$. So $y_1$‘s gradient is ${\nabla }_{\theta }\log {\pi }_{\theta }(y_1\mid x)\cdot {\hat{A}}_1$ where ${\hat{A}}_1$ is built from that shared return (often minus a baseline). The whole response’s success/failure is attributed uniformly to $y_1$ along with every other token. This works but is high variance (one noisy terminal scalar drives thousands of token updates) and imprecise (a single lucky/unlucky token gets the same credit as the pivotal one).
With a critic (GAE): the value head $V_{\varphi }(s_1)$ predicts expected terminal reward from the prompt alone; $V_{\varphi }(s_2)$ predicts it after seeing $y_1$. The TD residual ${\delta }_1=\gamma V_{\varphi }(s_2)-V_{\varphi }(s_1)$ (since $R=0$ at $t=1$) isolates how much choosing $y_1$ changed the expected outcome. GAE accumulates these residuals, giving $y_1$ an advantage that reflects its specific contribution, with variance reduced by the learned baseline.
Why the critic version is preferred: the reward is extremely sparse (one scalar per possibly-thousand-token episode) and the action space is huge, so the raw return estimator is very high variance. A good baseline/critic is the standard variance-reduction tool, and here it is close to essential — which is precisely why later methods that drop the critic (GRPO) must replace it with another strong baseline (a group mean) rather than nothing.

5.3 The KL-to-reference term: shaping the reward to stay on-distribution

If you optimized the raw RM reward alone, PPO would exploit every imperfection in $r_{\psi }$ — it is a learned, imperfect proxy for human preference, and its errors are largest exactly in regions of output space the base model rarely visits. An unconstrained optimizer marches straight into those blind spots and produces text that scores high but is degenerate: repetition, sycophancy, adversarial gibberish, mode collapse onto one high-reward template. This is reward hacking / over-optimization of the proxy (Goodhart’s law made concrete).⁴

The fix — introduced by Ziegler et al. (2019)⁵ and standard since⁶ — is to add a per-token KL penalty to a frozen reference policy ${\pi }_{\text{ref}}$ (the SFT model, held fixed for the whole run) directly into the reward. The effective per-token reward becomes:
$\boxed{r_t=\underset{\text{terminal RM/verifier reward}}{\underbrace{r_{\psi }(x,y_{1:T})1[t=T]}}-\underset{\text{per-token reward-shaping KL penalty}}{\underbrace{\beta (\log {\pi }_{\theta }(y_t\mid x,y_{<t})-\log {\pi }_{\text{ref}}(y_t\mid x,y_{<t}))}}}$

The bracketed difference is a single-sample estimate of the pointwise KL $\mathrm{KL}({\pi }_{\theta }(\cdot \mid s_t)\Vert {\pi }_{\text{ref}}(\cdot \mid s_t))$ at position $t$ (its expectation under $y_t\sim {\pi }_{\theta }$ is the KL). Summing over the response, the objective PPO actually maximizes is
${\mathbb{E}}_{x\sim \mathcal{D},y\sim {\pi }_{\theta }}[r_{\psi }(x,y)-\beta \mathrm{KL}({\pi }_{\theta }(\cdot \mid x)\Vert {\pi }_{\text{ref}}(\cdot \mid x)\left)\right].$

Why this specific form does the job:

• Keeps the policy on-distribution. The penalty makes drifting into regions ${\pi }_{\text{ref}}$ never explores costly, so the policy trades a little reward for staying where $r_{\psi }$ was trained and is therefore trustworthy. The leash length is $\beta$.
• Prevents reward hacking. The degenerate high-RM outputs live precisely in ${\pi }_{\text{ref}}$‘s low-probability tail; the KL term taxes exactly that movement.
• Prevents mode collapse and preserves fluency/capabilities. Because the anchor is frozen, the model cannot slowly walk away from fluent language one update at a time — every deviation is re-measured against the same fixed ${\pi }_{\text{ref}}$. Since KL is minimized by matching ${\pi }_{\text{ref}}$, the base model’s fluency, world knowledge, and instruction-following are actively retained rather than being overwritten by reward chasing.

Reward shaping, not a hard constraint

Folding the KL into the reward is a form of potential-free reward shaping: you’re modifying the per-step reward the RL algorithm sees. This is distinct from enforcing a constraint on the optimizer. Some pipelines instead (or additionally) add the KL to the loss as an explicit penalty term rather than the reward — numerically similar, and both are the “reference KL.” What matters is that the anchor ${\pi }_{\text{ref}}$ is frozen and the penalty is applied at every token.

The two-different-KLs confusion (read this twice)

There are two KL divergences in RLHF-PPO and conflating them is the most common conceptual error in the whole subject:

1. Reward-shaping / reference KL — $\mathrm{KL}({\pi }_{\theta }\Vert {\pi }_{\text{ref}})$ against the frozen SFT model, added explicitly into the reward with coefficient $\beta$. It bounds total drift over the entire run and defends against reward hacking / fluency loss. The anchor never moves.
2. Optimizer / trust-region KL — $\mathrm{KL}({\pi }_{\theta }\Vert {\pi }_{{\theta }_{\text{old}}})$ against the policy that generated the current batch, enforced implicitly by the PPO clip (usually only monitored, not added to the reward). It bounds the size of each optimization step and defends against the policy-gradient step being too large (failure mode 1 from lesson 04). The anchor ${\theta }_{\text{old}}$ is refreshed every iteration.
   Same divergence functional, completely different jobs, different anchors, different horizons, different mechanism (reward vs. clip). If someone says “the KL” without qualification, ask which one.

---

6. Consequences of a huge action space and long episodes

Two numbers dominate the character of this MDP:

• Action space $|V|\sim 32\text{k}$–$256\text{k}$. Each step is a categorical choice over the entire vocabulary — orders of magnitude larger than the handful of discrete actions in a game, or a low-dimensional continuous action in control.
• Horizon $T$ up to thousands of tokens. A single reasoning trace or long answer can be $10^3$–$10^4$ steps, and the only reward arrives at the very end.

What this implies, mechanically:

1. Enormous return variance. The number of possible trajectories is $\sim |V{|}^T$ — astronomically large. A Monte-Carlo return estimate over so many branching choices, credited by one terminal scalar, is extremely noisy. Variance reduction is not a nicety; it is the difference between learning and not.
2. Long-horizon credit assignment. With reward only at $T$, attributing the terminal score to the specific tokens that earned it (the pivotal reasoning step, the correct final digit) across thousands of intervening tokens is genuinely hard. This is where GAE’s $\lambda$ and a well-fit critic earn their keep — and also where sparse-reward RL is fundamentally fragile.
3. The appeal of strong baselines. Because variance is the binding constraint, the quality of the baseline subtracted from the return is decisive. A learned value head (PPO), a group-relative mean (GRPO), or an implicit reference (DPO) — each later method is, at heart, a different answer to “what baseline do we subtract to tame the variance of this huge-action, long-horizon, sparse-reward MDP?”
4. Clipping and small steps matter more. With such large per-step distributions, a modest parameter move can swing many token probabilities; the PPO trust region (small $ϵ$, small LR $\sim 10^{-6}$, few epochs) is what keeps the policy from lurching.

Why entropy and temperature are load-bearing here

A $|V|$-way categorical policy can collapse onto a few tokens frighteningly fast under reward pressure — that is mode collapse. The entropy bonus (from PPO’s objective) and the sampling temperature during rollouts are the exploration controls that keep the policy from prematurely committing across a vocabulary of hundreds of thousands of options. In language, “entropy collapse” shows up as repetitive, templated, low-diversity generations.

---

7. The optimization view: this is just PPO on this MDP

Nothing below is new algorithm — it is lesson 04 with the symbols rebound to LLM objects. That is the whole point of Tier 2’s first lesson.

PPO / Tier-1 concept	Token-sequence instantiation
Policy ${\pi }_{\theta }$	The LLM being fine-tuned (init from the SFT model)
State $s_t$	Prompt + tokens so far, $(x,y_{<t})$
Action $a_t$	Next token $y_t\in V$
Transition $P$	Deterministic append (no environment noise)
Trajectory / episode	One full generated response $y_{1:T}$ (ends at EOS / $T_{\max }$)
Reward $R(s_t,a_t)$	RM/verifier score at $t=T$, plus per-token $-\beta {\mathrm{KL}}_{\text{ref}}$
Value $V_{\varphi }(s_t)$	A value head on the LLM predicting return-to-go per token
Advantage ${\hat{A}}_t$	GAE over the per-token rewards $r_t$
Importance ratio $r_t(\theta )$	${\pi }_{\theta }(y_t\mid x,y_{<t})/{\pi }_{{\theta }_{\text{old}}}(y_t\mid x,y_{<t})$, per token
Clipped objective $L^{\text{CLIP}}$	Same formula, summed over generated tokens of the response
Rollout collection	Generation — sampling responses from ${\pi }_{{\theta }_{\text{old}}}$

The per-token clipped objective is identical to lesson 04, just summed over the response’s tokens:
$L^{\text{CLIP}}(\theta )={\mathbb{E}}_{x,y\sim {\pi }_{{\theta }_{\text{old}}}}\left[\frac{1}{T}{\sum }_{t=1}^T\min (r_t(\theta ){\hat{A}}_t,\mathrm{clip}(r_t(\theta ),1-ϵ,1+ϵ\left){\hat{A}}_t\right)\right],$
with ${\hat{A}}_t$ the GAE advantage computed from the shaped per-token rewards $r_t$ of §5, and the value head trained by the usual (often clipped) squared-error regression toward the returns ${\hat{R}}_t={\hat{A}}_t+V_{{\varphi }_{\text{old}}}(s_t)$. The full RLHF-PPO loss is once again $L^{\text{CLIP}}-c_1{L}^{\text{VF}}+c_2\mathcal{H}[{\pi }_{\theta }]$.

flowchart TD
    A["Prompts x ~ D"] --> B["<b>Rollout = generation</b><br/>sample y ~ π_θ_old(·|x)<br/>store per-token logπ_θ_old, V(s_t)"]
    B --> C["<b>Score</b><br/>RM r_ψ(x,y) at EOS<br/>+ compute logπ_ref per token"]
    C --> D["<b>Shape reward</b><br/>r_t = r_ψ·1[t=T] − β(logπ_θ − logπ_ref)"]
    D --> E["<b>GAE over tokens</b><br/>δ_t = r_t + γV(s_t+1) − V(s_t)<br/>Â_t = Σ(γλ)^l δ_t+l ; R_t = Â_t + V(s_t)"]
    E --> F["<b>K epochs PPO-SGD</b><br/>L^CLIP − c1 L^VF + c2 H<br/>ratio r_t = π_θ/π_θ_old per token"]
    F --> G{"monitor KL(π_θ‖π_θ_old)<br/>and KL(π_θ‖π_ref)"}
    G -->|"refresh θ_old ← θ"| A

    style B fill:#334,color:#fff
    style D fill:#a64,color:#fff
    style E fill:#343,color:#fff
    style F fill:#433,color:#fff

What is genuinely different from control-task PPO

The algorithm is identical, but the operating point is not. Compared to Atari/MuJoCo PPO, RLHF uses a much smaller learning rate ($\sim 10^{-6}$), fewer epochs ($K=1$–$4$), often smaller $ϵ$, and — critically — the dominant hyperparameter becomes $\beta$, the reference-KL coefficient, which has no analogue in a game where there is no “reference policy to stay near.” Also, the critic is frequently a separate full-size model (not a shared trunk) to keep the value updates from corrupting the policy’s representations. These are tuning differences on the same machine.

---

8. Practical framing: rollouts are generation, and four models sit in memory

Casting RLHF as an MDP has a very concrete systems consequence. In classic RL, “rollout” means stepping a cheap simulator. Here, rollout = autoregressive generation — the single most expensive operation in the loop.

The four models often resident at once:

1. Policy ${\pi }_{\theta }$ — trainable; the LLM being optimized.
2. Critic / value head $V_{\varphi }$ — trainable; predicts per-token return-to-go. Often a second full-size model (memory $\approx$ another policy).
3. Reference ${\pi }_{\text{ref}}$ — frozen; the SFT model, used only for forward passes to get $\log {\pi }_{\text{ref}}$ for the KL penalty.
4. Reward model $r_{\psi }$ — frozen; scores completed responses. (In verifiable domains this is replaced by a cheap verifier and can drop out.)

That is roughly $2\times$ the policy’s memory just for the critic and reference, plus a separate RM — a major reason RLHF is operationally heavy and a direct target of the methods to come.

Why generation dominates cost. Each PPO iteration must first sample full responses token-by-token from ${\pi }_{{\theta }_{\text{old}}}$. Autoregressive decoding is inherently sequential — $T$ forward passes per response, memory-bandwidth-bound, no way to parallelize across the time dimension within one sequence. The subsequent gradient update is a handful of parallelizable forward+backward passes over already-generated sequences. So in wall-clock terms the loop is typically generation-bound, which is why RLHF infrastructure invests heavily in fast inference (batched decoding, KV-cache, paged attention, dedicated inference engines like vLLM feeding the trainer).⁷ The forward passes to compute $\log {\pi }_{\text{ref}}$ and $r_{\psi }(x,y)$ add to this but are non-sequential (whole sequence at once).

This memory/compute bill is the "why" of the next lessons

Hold this picture: four models, generation-bound, $\beta$-sensitive. GRPO⁸ attacks it by deleting the critic — replacing $V_{\varphi }$ with the mean reward over a group of sampled responses per prompt (a Monte-Carlo baseline), removing an entire trainable full-size model. DPO⁹ attacks it harder by deleting the online RL loop and the reward model — showing the (reward + reference-KL) objective has a closed-form optimum you can fit directly on preference pairs, leaving essentially just the policy and a frozen reference. Every deletion is legible only because you now hold the full four-model picture.

---

9. Quizzes

Quiz 3: Why regularize toward a frozen reference rather than, say, just adding an entropy bonus, or using the PPO clip's trust region to keep the policy well-behaved?

Answer

The three mechanisms target different failure modes and are not substitutes:

• Entropy bonus keeps the policy stochastic (prevents premature determinism), but does nothing to keep outputs fluent or on the base-model distribution. A high-entropy policy can still be high-entropy garbage. It has no notion of “the language the base model produces.”
• PPO clip / trust-region KL bounds each step relative to ${\pi }_{{\theta }_{\text{old}}}$, which moves every iteration. Over many iterations the policy can walk arbitrarily far from the original SFT model — each step small, cumulative drift unbounded. The clip cannot prevent slow, monotonic drift into reward-model blind spots because its anchor keeps advancing with the policy.
• Reference KL to a frozen ${\pi }_{\text{ref}}$ is the only one whose anchor never moves. It bounds total deviation from the SFT distribution over the whole run, making it the mechanism that (a) taxes movement into $r_{\psi }$‘s untrustworthy tail (anti-reward-hacking), (b) preserves fluency/capabilities by keeping the policy near a known-good distribution, and (c) prevents mode collapse by penalizing peaky deviations from ${\pi }_{\text{ref}}$. Because the reward model is a learned proxy, an anchor to real well-formed language (the SFT model) is exactly what constrains exploitation of that proxy. Entropy and the clip cannot provide this because neither references the base distribution.

Quiz 4: A response is 500 tokens; the verifier returns reward $1.0$ at EOS and the model used $\gamma =1.0$, $\lambda =1.0$, and (hypothetically) a perfectly zero value function $V_{\varphi }\equiv 0$. Write the advantage ${\hat{A}}_t$ for an arbitrary token $t$ (ignore the KL term), and explain what pathology this specific setting reproduces.

Answer

With $V_{\varphi }\equiv 0$, the TD residuals are ${\delta }_t=r_t+\gamma V(s_{t+1})-V(s_t)=r_t$. With $\gamma =\lambda =1$, GAE reduces to the plain return-to-go: ${\hat{A}}_t={\sum }_{l\ge 0}{\delta }_{t+l}={\sum }_{k=t}^T{r}_k=1.0$ for every token (the only nonzero reward is the terminal $1.0$, and undiscounted it reaches all positions). So every one of the 500 tokens receives the identical advantage ${\hat{A}}_t=1$. This is exactly REINFORCE with no baseline: the terminal scalar is broadcast uniformly, giving maximal-variance, minimally-localized credit — a correct token and a filler token in the same response get identical reinforcement. It reproduces the high-variance, poor-credit-assignment regime that motivates having any baseline at all. Setting $\lambda <1$ with a trained critic (or, in GRPO, subtracting a group-mean baseline) is precisely the fix. Note also that in this setting there is no way to prefer a shorter correct response over a longer one — a hint at the length-bias problems that plague sparse-reward LLM RL.

---

10. Practice problems

Problem 1: Trace one token's per-token reward end to end

Setup. A math prompt $x$; the policy samples a 4-token answer $y=(y_1,y_2,y_3,y_4)$ with $y_4=\text{EOS}$. A verifier gives $r_{\psi }(x,y)=1$ (correct). Use $\beta =0.1$, $\gamma =1$, and suppose the per-token log-prob differences $\log {\pi }_{\theta }(y_t\mid \cdot )-\log {\pi }_{\text{ref}}(y_t\mid \cdot )$ are $[0.5,-0.2,1.0,0.0]$ for $t=\mathrm{1..4}$. Write the shaped per-token rewards $r_t$, and then the return-to-go $G_t$ for each token (no critic, so use the raw return).

Worked solution.
The shaped reward is $r_t=r_{\psi }1[t=4]-\beta (\log {\pi }_{\theta }-\log {\pi }_{\text{ref}})$.

• $r_1=0-0.1(0.5)=-0.05$
• $r_2=0-0.1(-0.2)=+0.02$
• $r_3=0-0.1(1.0)=-0.10$
• $r_4=1-0.1(0.0)=+1.00$

Return-to-go with $\gamma =1$ is $G_t={\sum }_{k=t}^4{r}_k$:

• $G_4=1.00$
• $G_3=-0.10+1.00=0.90$
• $G_2=0.02-0.10+1.00=0.92$
• $G_1=-0.05+0.02-0.10+1.00=0.87$

Reading it: the terminal $+1$ dominates and propagates to all tokens (that’s the credit path), but the KL shaping modulates each token: token 3 deviated most from the reference ($\Delta \log \pi =1.0$) and so carries the largest penalty, gently discouraging that particular high-drift choice even though the overall answer was correct. This is how the reference KL trims risky-but-not-yet-punished deviations token by token, rather than waiting for the terminal signal.

Problem 2: Diagnose a training failure from symptoms

Setup. Mid-RLHF-run you observe: mean RM score keeps climbing, but $\mathrm{KL}({\pi }_{\theta }\Vert {\pi }_{\text{ref}})$ has blown past 40 nats/response, generations have become repetitive and oddly phrased, and human raters dislike the outputs despite the high RM score. What is happening mechanistically, which knob is mis-set, and what do you do?

Worked solution.
This is textbook reward hacking / over-optimization of the reward model. The RM $r_{\psi }$ is an imperfect learned proxy; its errors are largest in the low-probability tail of ${\pi }_{\text{ref}}$. The rising RM score with exploding reference-KL means the policy has walked far off the SFT distribution into exactly that tail, finding inputs that fool $r_{\psi }$ into high scores — hence high RM but low actual human preference (the proxy and the true objective have diverged: Goodhart). The repetitive/odd phrasing is the mode-collapse signature of a $|V|$-way policy peaking onto a few high-RM templates.
Mis-set knob: $\beta$ (the reference-KL coefficient) is too small — the leash is too long, so drift into the proxy’s blind spots is nearly free. (A too-aggressive learning rate / too many epochs compounds it.)
Fixes, in order: (1) increase $\beta$ (or switch on an adaptive-$\beta$ controller targeting a KL of a few nats/response) to re-anchor to ${\pi }_{\text{ref}}$; (2) if using KL early-stopping, tighten it; (3) longer-term, the real cause is an over-optimizable RM — improve/retrain it or add RM ensembling (lesson 06). Note which KL you are turning: this is the reference KL ($\beta$ in the reward), not the PPO clip $ϵ$ — tightening $ϵ$ would slow steps but would not stop cumulative drift toward the frozen reference, because $ϵ$‘s anchor ${\pi }_{{\theta }_{\text{old}}}$ moves with the policy.

---

11. Reading order

Work through these in order:

1. The template of the whole pipeline: Training language models to follow instructions with human feedback (InstructGPT) — Ouyang et al., 2022 (arXiv:2203.02155) — §3.5 and the appendix give the exact RM-score-minus-KL reward, the value-head/PPO setup, and the four-model loop. This lesson is that section, formalized.
2. The origin of the reference-KL idea: Fine-Tuning Language Models from Human Preferences — Ziegler et al., 2019 (arXiv:1909.08593) — the first clean statement of “reward = preference-model score − $\beta \cdot$KL-to-reference,” including the over-optimization argument. Read §2–4.
3. The illustrated mental model: Illustrating Reinforcement Learning from Human Feedback (HuggingFace blog) — the canonical diagram of the three-phase RLHF pipeline; use it to lock in the four-model picture from §8.
4. From theory to a real run: HuggingFace TRL — PPOTrainer docs — the concrete four-model config (policy, ref, reward, value), the KL controller, and the knobs from §7–§8. The best bridge to actually running one.
5. The definitive synthesis: Nathan Lambert — The RLHF Book (rlhfbook.com) — chapters on the RL formulation of RLHF, KL regularization, and reward modeling; the most complete modern treatment and an excellent map of the whole subfield.

---

12. What’s next

Two threads, in order:

1. 06-rlhf-pipeline — this lesson took the reward $r_{\psi }$ as given. Lesson 06 opens the black box: where the reward model comes from — the Bradley–Terry preference loss on human comparisons, the SFT → RM → PPO three-stage pipeline, and the failure modes of the RM itself (calibration, over-optimization, ensembling). It is the missing (b) in the “(a) MDP + (b) reward + (c) optimizer” decomposition.
2. Circle back to 04-ppo — reread §7 now that the token-level MDP is rigorous; the value-head/GAE-over-tokens plumbing and the two-KLs distinction will read as concrete rather than foreshadowed.

Then the two deletions from this recipe, both of which only make sense once you hold the full four-model picture from §8: 09-rl-for-reasoning drops the critic (group-mean baseline instead of a value head), and 07-dpo-and-rl-free-preference-optimization drops the online RL loop and the reward model entirely (closed-form optimum of the reward + reference-KL objective, fit directly on preference pairs).

---

References

Topic hub: rl-for-llms | Reference pages: ppo, kl-regularization-rlhf | Filed: 2026-09-02

Footnotes

1. Schulman, Wolski, Dhariwal, Radford, Klimov (2017), “Proximal Policy Optimization Algorithms,” arXiv:1707.06347. https://arxiv.org/abs/1707.06347 — origin of PPO (evaluated on MuJoCo continuous control and Atari). ↩

2. Representative modern LLM vocabulary sizes: GPT-2 uses 50,257 BPE tokens (Radford et al. 2019, “Language Models are Unsupervised Multitask Learners”); Llama 3 uses a 128k-token tokenizer (arXiv:2407.21783, https://arxiv.org/abs/2407.21783); some recent tokenizers (e.g. Gemma) reach ~256k. The 32k–256k range is representative, not a hard bound. [established] ↩

3. Schulman, Moritz, Levine, Jordan, Abbeel (2015/2016), “High-Dimensional Continuous Control Using Generalized Advantage Estimation,” arXiv:1506.02438. https://arxiv.org/abs/1506.02438 — origin of GAE and the $\lambda$ bias–variance dial. ↩

4. Gao, Schulman, Hilton (2022), “Scaling Laws for Reward Model Overoptimization,” arXiv:2210.10760. https://arxiv.org/abs/2210.10760 — quantifies reward hacking / Goodhart over-optimization of a learned proxy RM as a function of KL from the reference. ↩

5. Ziegler, Stiennon, Wu, Brown, Radford, Amodei, Christiano, Irving (2019), “Fine-Tuning Language Models from Human Preferences,” arXiv:1909.08593. https://arxiv.org/abs/1909.08593 — first clean statement of the RM-score-minus-$\beta$·KL-to-reference reward for LM fine-tuning, and the over-optimization argument. ↩

6. Ouyang et al. (2022), “Training language models to follow instructions with human feedback” (InstructGPT), arXiv:2203.02155. https://arxiv.org/abs/2203.02155 — the standard SFT → RM → PPO recipe with the per-token KL penalty and four-model loop. ↩

7. Kwon, Li, Zhuang, Sheng, Zheng, Yu, Gonzalez, Zhang, Stoica (2023), “Efficient Memory Management for Large Language Model Serving with PagedAttention” (vLLM), arXiv:2309.06180. https://arxiv.org/abs/2309.06180 — PagedAttention / vLLM inference engine. ↩

8. Shao et al. (2024), “DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models,” arXiv:2402.03300. https://arxiv.org/abs/2402.03300 — introduces Group Relative Policy Optimization (GRPO), which drops the value critic in favor of a group-mean baseline. ↩

9. Rafailov, Sharma, Mitchell, Ermon, Manning, Finn (2023), “Direct Preference Optimization: Your Language Model is Secretly a Reward Model,” arXiv:2305.18290. https://arxiv.org/abs/2305.18290 — closed-form optimum of the reward + reference-KL objective, fit directly on preference pairs. ↩