Everything in this chapter so far has been about robots in grids and agents in corridors. Here is why we built it. Take a language model fine-tuned to follow instructions (Chapter 14), give it a prompt, let it write a response, and score the response with a reward model trained on human preferences or with a program that checks the final answer. You want the model to produce higher-scoring responses. That is a reinforcement learning problem, and the model is the policy.

This last section makes that sentence precise. We set up the token-level MDP and the sequence-level bandit and show that they give the same gradient, look at what a single reward at the end of a long response does to credit assignment, and then write down the objective that RLHF, DPO, and GRPO all start from, expected reward minus a KL penalty toward a reference model. We derive its per-token form, run it on a toy language model small enough to solve exactly, find two ways that real implementations get its gradient wrong, and finish with a frank list of what makes RL on language models unusual RL.

Tokens as actions, prefixes as states

Write $x$ for the prompt, drawn from a distribution $\mathcal{D}$ of prompts, and $y = (y_1, \dots, y_T)$ for the response, a sequence of tokens from the vocabulary $\mathcal{V}$ ending with an end-of-sequence token or at a maximum length. (We index response tokens from 1, as in the language-model chapters, so step $t$ produces token $y_t$.) The MDP is

• state $s_t = (x, y_{<t})$, the prompt and everything generated so far,
• action $a_t = y_t \in \mathcal{V}$, the next token,
• policy $\pi_\theta(y_t \mid x, y_{<t})$, the model's next-token distribution, the softmax over the vocabulary of §6.4,
• transition $s_{t+1} = (x, y_{\le t})$, append the token, which is deterministic and known,
• start distribution $p_0 = \mathcal{D}$, the prompts,
• reward $r_t = 0$ for $t < T$ and $r_T = r(x, y)$, a single score for the finished response,

with no discounting, $\gamma = 1$, which is safe because every episode ends by the maximum length. The state is the whole history, so the process is Markov by construction (§15.1). The state space is astronomically large, and nobody enumerates it. The network generalizes across prefixes, exactly as it does in pretraining.

Two views of the same policy

The same model is also a policy over whole responses, since the chain rule of probability gives

In this sequence-level view, the problem is a contextual bandit (§15.1). The context is the prompt, the action is the entire response, and the reward arrives at once. The objective is

and the policy gradient of §15.3, with a baseline $b(x)$ that may depend on the prompt but not on the response, is

In the token-level view, the reward-to-go from any step is $G_t = \sum_{k\ge t}r_k = r(x, y)$, since the only reward is at the end. So REINFORCE with reward-to-go gives every token the same weight, $r(x,y)$, and the two views produce exactly the same estimator. They are two descriptions of one problem.

They differ in how you estimate the advantage, and that is the design choice that separates the main LLM RL algorithms.

• Token-level critic. Learn a value $V_\psi(x, y_{<t})$ (a network with parameters $\psi$) for every prefix and use GAE over tokens (§15.5). Each token gets its own advantage, $\hat A_t \approx r(x,y) - V_\psi(s_t)$ when $\lambda = 1$, so a token that made a good ending more likely gets more credit than one that did not. This is PPO for RLHF, and the critic is a whole second language model with a scalar head (§16.4).
• Sequence-level baseline. Sample several responses to the same prompt and use their average reward as $b(x)$, with every token of a response sharing its advantage $r(x, y) - b(x)$. This is RLOO and GRPO, which drop the critic entirely (§18.4).

Here is a toy language model small enough to check everything exactly. It has a vocabulary of 4 tokens and always writes responses of exactly 3 tokens, so there are only $4^3 = 64$ possible responses. Its policy is a table of logits, one row for each of the $1 + 4 + 16 = 21$ possible prefixes, so it can represent any distribution over responses. The reference model's logits are random. The reward is a "verifier" that pays 1 when the response contains token 0 exactly twice.

The per-token probabilities multiply into a proper distribution over all 64 responses, and the reference model earns the reward about 8% of the time. Computing sequence log-probabilities this way, by gathering each token's log-probability from a stable log-softmax and summing over the real tokens, is the basic operation of every LLM RL and preference method, and it is worth writing once carefully.

import numpy as np

def sequence_logprobs(logits, tokens, mask):
    raise NotImplementedError

Sparse terminal rewards and credit assignment

A reward that arrives only at the end of a 1,000-token response tells you that the response was good or bad, not which tokens made it so. This is the credit assignment problem of §15.1 at its most extreme, and it is why LLM RL needs so many samples. With a sequence-level baseline every token of a good response is pushed up by the same amount, the crucial step in a proof and the filler words around it alike. Over many samples the noise averages out, since tokens that matter are correlated with reward across samples and tokens that do not are not, but the averaging is slow, and the variance grows with the response length (§15.3).

A token-level critic can do better in principle, because $V_\psi(x, y_{<t})$ can notice that the response became much more likely to succeed after a particular step, giving that step a large TD error. In practice, learning such a critic is hard. It must predict the final reward from a partial response, with a single noisy label per response, and an inaccurate critic biases the advantages (§15.4). That difficulty is one reason the critic-free methods of Chapter 18 took over reasoning RL, and why process rewards that score intermediate steps are an active topic (§18.6).

The KL-regularized objective

Maximizing $J(\theta)$ alone is dangerous when the reward comes from a learned model. A reward model is accurate on responses that look like its training data and arbitrary elsewhere, and an optimizer strong enough to find its blind spots will find them. This is reward hacking, and Gao, Schulman, and Hilton (2023) measured how the true quality of a policy first rises and then falls as optimization against a proxy reward model continues (§16.5). The standard remedy is to penalize the policy for drifting from a reference policy $\pi_{\text{ref}}$, usually the supervised fine-tuned model the RL starts from,

with a coefficient $\beta > 0$. In plain words, earn reward, but pay $\beta$ for every nat of divergence from the model you started with.

The penalty does three jobs. It keeps the policy in the region where the reward model can be trusted. It keeps the language fluent and the knowledge intact, since the reference model is a good language model. And it acts like an entropy bonus that keeps the policy from collapsing onto one response. Stiennon et al. (2020) name the first and the last explicitly, and Ziegler et al. (2019) introduced the penalty to keep the policy from drifting into text the reward model was never trained on. Typical values of $\beta$ are small, for example 0.05 in Stiennon et al.'s summarization work and 0.02 in InstructGPT (Ouyang et al. 2022), and Ziegler et al. adapted $\beta$ during training with a controller that targets a chosen KL.

The objective has a closed-form maximizer. Over all policies, the optimum tilts the reference distribution toward high reward,

In plain words, take the reference model's probability of every response, multiply it by $e^{r(x,y)/\beta}$, and renormalize. Two responses that the reference finds equally likely end up in the ratio $e^{\Delta r/\beta}$ if their rewards differ by $\Delta r$, so a small $\beta$ piles the probability onto the best responses and a large one barely moves the reference.

Here is the one-line reason it has this shape. Moving a little probability onto a response $y$ earns its reward $r(x, y)$ and costs $\beta$ times its log-ratio $\log(\pi(y\mid x)/\pi_{\text{ref}}(y\mid x))$, plus a constant, so at the optimum that net gain must be the same for every response (otherwise shifting probability from a worse one to a better one would help). That balance condition reads $\beta\log\frac{\pi^\star(y\mid x)}{\pi_{\text{ref}}(y\mid x)} = r(x,y) - \beta\log Z(x)$, and solving it for $\pi^\star$ gives the formula, with $Z(x)$ the constant that makes the probabilities sum to one.

The result comes with an exact accounting that we will use. For every policy $\pi$ and prompt $x$, the objective at $x$ equals $\beta\log Z(x) - \beta\,D_{\mathrm{KL}}\big(\pi(\cdot\mid x)\,\Vert\,\pi^\star(\cdot\mid x)\big)$. The first term does not depend on $\pi$, and a KL divergence is zero only when its arguments are equal, so $\pi^\star$ is the unique maximizer, the best value is $\beta\log Z(x)$, and every other policy falls short of it by exactly $\beta$ times its KL to $\pi^\star$. That gap is a convenient yardstick, and we will use it to grade training runs below. Both facts are proved step by step, twice, in the derivation of the closed-form optimal policy, which also builds direct preference optimization on the rearranged balance condition, the reward written through the policy as $r = \beta\log(\pi^\star/\pi_{\text{ref}}) + \beta\log Z$.

It is the same exponential tilt that the natural gradient produced in §15.6, now applied to whole responses. The partition function $Z(x)$ sums over every possible response, so nobody can compute $\pi^\star$ for a real model, which is why we optimize toward it with RL instead. For our toy, with only 64 responses, we can compute it exactly, and check the accounting identity numerically.

A large $\beta$ barely moves the policy from the reference. A small one puts nearly all the probability on rewarded responses, at a large KL. And for a policy drawn at random, its objective equals $\beta\log Z$ minus $\beta$ times its KL to $\pi^\star$, to every printed digit. The widget shows the same trade-off with a reward model that is wrong in a particular way, over-rating confident, flattering answers.

The per-token form

The KL between two distributions over whole responses looks impossible to work with, but it decomposes over tokens. Write the log-ratio of sequence probabilities as a sum over tokens, using the chain rule above,

with $s_t = (x, y_{<t})$. Taking the expectation of each term first over $y_t$ given the prefix gives the chain rule for KL divergence, $D_{\mathrm{KL}}(\pi_\theta(\cdot\mid x)\Vert\pi_{\text{ref}}(\cdot\mid x)) = \sum_t\mathbb{E}_{s_t\sim\pi_\theta}[D_{\mathrm{KL}}(\pi_\theta(\cdot\mid s_t)\Vert\pi_{\text{ref}}(\cdot\mid s_t))]$, the sum over positions of the per-token KL at the prefixes the policy actually reaches.

Substitute the first form into the objective. Everything inside the expectation over $y$ becomes a sum over tokens, which is exactly a token-level MDP with the per-token reward

whose sum over the response is $r(x,y) - \beta\log\frac{\pi_\theta(y\mid x)}{\pi_{\text{ref}}(y\mid x)}$, the "full reward" $R(x,y)$ written down by Ziegler et al. and Stiennon et al. In plain words, every token pays a small toll, $\beta$ times how much more likely the policy made it than the reference would have, and the last token also collects the reward-model score. InstructGPT describes exactly this, a per-token KL penalty from the SFT model at each token, with the reward model's score at the end of the episode.

One subtlety makes this legitimate. The per-token reward depends on $\theta$ itself, unlike any reward in an ordinary MDP. Does running a policy gradient on it, with the KL term treated as a fixed reward, give the gradient of $J_\beta$? It does. The gradient of the penalty term is

by the product rule (the first term from differentiating the sampling distribution, the second from differentiating the integrand), and the second term is zero by the expected-score lemma of §15.4. What remains is the score-function gradient with the log-ratio as the reward, which is exactly what a policy gradient method computes when the KL term is placed in the reward and not differentiated. Keep this in mind, because the next part shows what happens when it is differentiated.

import numpy as np

def kl_penalized_rewards(logp, logp_ref, score, mask, beta):
    raise NotImplementedError

Two ways to get the KL gradient wrong

Real implementations estimate the KL from the sampled tokens, with the per-token log-ratio (the $k_1$ estimator of §0.5) or with $k_3 = (\rho' - 1) - \log\rho'$ where $\rho' = \pi_{\text{ref}}/\pi_\theta$ is the reference-to-policy ratio of the sampled token, as GRPO does. Both are unbiased estimates of the KL value. The trouble starts when, instead of putting the estimate in the reward, an implementation adds it to the loss and lets autograd differentiate it. Tang and Munos (2025) pointed out that this is common in open-source RLHF code and in published pseudocode, and that it gives the wrong gradient.

Pitfall one, differentiating the estimate. Differentiating $k_1 = \log\pi_\theta(y) - \log\pi_{\text{ref}}(y)$ at a sampled $y$ gives $\nabla_\theta\log\pi_\theta(y)$, whose expectation is zero by the expected-score lemma. The penalty's gradient is pure noise, and the KL is not regularized at all. Differentiating $k_3$ at the sampled token gives $(1 - \pi_{\text{ref}}/\pi_\theta)\nabla\log\pi_\theta$. The first part has expectation zero again, and the second has expectation $-\sum_a\pi_{\text{ref}}(a)\nabla\log\pi_\theta(a)$, which is the gradient of a different divergence, the KL with its arguments swapped, $D_{\mathrm{KL}}(\pi_{\text{ref}}\Vert\pi_\theta)$. In the language of §0.5, the objective penalizes the reverse, mode-seeking KL, while this gradient follows the forward, mass-covering one. It still pulls the policy toward the reference, so training looks fine, but it optimizes a different objective, and the optimum is no longer the tilted distribution above. Differentiating $k_2 = \frac12(\log\pi_\theta/\pi_{\text{ref}})^2$, by contrast, gives $\log\frac{\pi_\theta}{\pi_{\text{ref}}}\nabla\log\pi_\theta$, whose expectation is the correct gradient of $D_{\mathrm{KL}}(\pi_\theta\Vert\pi_{\text{ref}})$ for a single step, even though $k_2$'s value is a biased estimate of the KL. (Pitfall two is about what goes wrong over a whole sequence.) Exercise 15.7c asks you to verify all three by exact enumeration.

The same algebra settles what happens when each estimate is used as a reward instead, with the score-function gradient $\mathbb{E}[k\,\nabla\log\pi_\theta]$. The product rule gives $\nabla_\theta\mathbb{E}_{\pi_\theta}[k] = \mathbb{E}[k\,\nabla\log\pi_\theta] + \mathbb{E}[\nabla_\theta k]$, so the reward version equals the gradient of the estimator's expected value minus its loss version. Only $k_1$ comes out right both ways of reading. Here is the whole picture for one step, with the token $a\sim\pi_\theta$ and $D = D_{\mathrm{KL}}(\pi_\theta\Vert\pi_{\text{ref}})$.

In plain words, a good estimate of the KL's value does not come with a good estimate of its gradient. RLHF with PPO uses $k_1$ in the reward, which is correct. GRPO uses $k_3$ in the loss, which regularizes through the swapped KL (§18.4).

Pitfall two, ignoring the sequence. Even with a correct per-token gradient, the penalty at token $t$ depends on all the earlier tokens, because they decide which prefix the policy is in. A loss that adds the per-token KL at each position and differentiates only through that position's own log-probability misses the effect of earlier tokens on later penalties. The correct gradient multiplies each token's score by the KL to go, the sum of the log-ratios from that token to the end, which is exactly what you get for free by putting the per-token KL in the reward and using reward-to-go or GAE.

The toy model shows the first pitfall clearly. Below we train it with REINFORCE three ways, all with reward-to-go and a per-position batch-mean baseline. In the first, the KL penalty sits in the reward. In the second, the $k_1$ estimate is added to the loss and differentiated. In the third, there is no penalty at all.

The last column is the yardstick from above, the KL from the trained policy to $\pi^\star$, which is the shortfall in $J_\beta$ divided by $\beta$. With the KL in the reward, REINFORCE ends close to the exact optimum $\pi^\star$ for $\beta = 0.5$, 0.02 nats from it, with a reward of about 0.35 against the optimum's 0.40 and a KL to the reference of 0.30 against 0.38, and it is still converging. With the $k_1$ estimate differentiated in the loss, the result is indistinguishable from having no penalty at all. Both runs push the reward near 1, drift more than 3 nats from the reference, and end about 1.9 nats from $\pi^\star$. The penalty was there in the code and did nothing.

import numpy as np

def kl_gradient_audit(z, z_ref):
    raise NotImplementedError

Why RL on language models is unusual RL

Language models fit the RL framework exactly, but they sit in an odd corner of it, and most of the design choices of Chapters 16 to 18 follow from that corner.

• The dynamics are deterministic and known. The environment is string concatenation. There is nothing to learn about transitions, the agent can "reset" to any prefix it likes, and it can compute the exact probability of any response under itself or the reference. Branching from a shared prefix to compare continuations, tree search over reasoning steps (§18.6), and exact sequence log-probabilities (DPO, §17.2) all depend on this. The whole problem of learning a model of the environment, which model-based reinforcement learning is built around for robots and games, does not arise here.
• The action space is huge, but the prior is strong. Each step chooses among 50,000 to 200,000 tokens, which would make undirected exploration hopeless. But RL starts from a pretrained model that already produces sensible text, so exploration is sampling from that prior, and RL mostly reweights behaviors the model can already produce. Yue et al. (2025) found that RL-trained reasoning models beat their base models at pass@1 while the base models solve more problems at large $k$ in pass@k, which they read as RL sharpening rather than expanding the base model's abilities. That interpretation is debated, but the strong prior is not.
• The horizon is long and the reward is single. Responses run from hundreds to tens of thousands of tokens with one reward at the end, so credit assignment is the hard part, and the choice between a token-level critic and a sequence-level baseline is the main fork between PPO-style RLHF and GRPO-style reasoning RL.
• The reward is learned and exploitable. A game score is the game. A reward model is a proxy, so optimization pressure finds its flaws, which is why the KL penalty, reward-model ensembles, and early stopping exist (§16.5). Verifiable rewards are harder to hack but not impossible (§18.3).
• Episodes are expensive and models are big. One episode is a full generation, so reusing each batch for several PPO epochs pays off, and the policy, the reference, the reward model, and a critic can each be a full-size network, the four-model memory bill of §16.4.
• The data is slightly off-policy for systems reasons. Large systems generate with a separate, optimized inference engine and train with another, and the two compute slightly different log-probabilities for the same tokens. Ratios that should be exactly 1 are not, which is one more reason the importance-sampling machinery of §15.6 matters even in "on-policy" training.

In PyTorch

Here is the per-token RLHF reward computation and a masked PPO loss on a batch of sampled sequences from a tiny GRU language model, with a frozen copy as the reference, the KL penalty in the reward, GAE over tokens with $\gamma = 1$, and a critic head. It is a miniature of the training step §16.4 builds at full scale.

Read more

• Ziegler et al. (2019). Fine-Tuning Language Models from Human Preferences. The KL-penalized reward $r - \beta\log(\pi/\rho)$ with PPO on language models, and the adaptive KL controller.
• Stiennon et al. (2020). Learning to summarize from human feedback. Each token as a step of PPO, $\gamma = 1$, a separate value network, and why the KL term matters.
• Ouyang et al. (2022). Training language models to follow instructions with human feedback. InstructGPT, which frames RLHF as a bandit environment with a per-token KL penalty.
• Tang & Munos (2025). On a few pitfalls in KL divergence gradient estimation for RL. Why differentiating KL estimates in the loss gives the wrong gradient, in tabular and LLM experiments.
• Lambert (2025). Reinforcement Learning from Human Feedback. A book-length treatment of the whole RLHF stack, including the policy gradient and regularization chapters this section prepares you for.
• Yue et al. (2025). Does Reinforcement Learning Really Incentivize Reasoning Capacity in LLMs Beyond the Base Model? The pass@k evidence that RL sharpens rather than extends the base model.
• Gao, Schulman & Hilton (2023). Scaling Laws for Reward Model Overoptimization. ICML 2023. How true quality rises and then falls as a policy is optimized against a proxy reward model.

Watch next

▶

▶