The LLM as an MDP#

Every method in this chapter — RLHF, DPO, GRPO — is classical machinery from Sections 2 and 3 applied to one specific MDP. Before touching any algorithm, we pin that MDP down, because its peculiarities explain why each method looks the way it does: why advantage estimation is painful, why the bandit view is tempting, and why GRPO can afford to delete the critic.

The mapping#

An LLM is a policy:

RL conceptLLM instantiation
State $s_t$Prompt $x$ + tokens generated so far, $(x, y_1, \dots, y_t)$
Action $a_t$The next token $y_{t+1}$ from the vocabulary $\mathcal{V}$
Transition $P(s_{t+1} \mid s_t, a_t)$Deterministic concatenation: $s_{t+1} = s_t \circ a_t$
Policy $\pi_\theta(a \mid s)$The language model’s next-token distribution
Reward $r$Preference score, verifier output, task success — terminal only
EpisodeOne complete response, ending at <eos> or the length cap
Discount $\gamma$$1$ in essentially all implementations

Formally: a finite-horizon MDP with horizon $T_{\max}$, discrete action space $|\mathcal{A}| = |\mathcal{V}| \approx 10^5$, known deterministic dynamics, and a reward that is zero everywhere except at the final step:

$$ r(s_t, a_t) = \begin{cases} R(x, y) & \text{if } a_t = \langle\texttt{eos}\rangle \text{ (or } t = T_{\max}\text{)} \\ 0 & \text{otherwise.} \end{cases} $$

Generation is a rollout. Sampling temperature is exploration. A conversation is an episode.

Key idea: nothing new needs to be invented to apply RL to language models. What is new is the shape of this particular MDP — and every design choice in RLHF, DPO, and GRPO is a reaction to that shape.

Four properties that shape everything downstream#

1. Transitions are deterministic and known. The “environment” of single-turn generation is just string concatenation. There is no dynamics model to learn — the model-based/model-free distinction from Section 3 collapses, and all randomness in a rollout comes from the policy itself. (Hold this thought: it is exactly what breaks again in Section 4.5, where tool calls and users make transitions stochastic.)

2. Reward is terminal and sparse. Nothing arrives until the full response exists. A 2,000-token chain-of-thought receives a single scalar at the end, and credit assignment — which tokens deserve it? — becomes the central difficulty. Value-based bootstrapping (Section 3) was designed for exactly this, which is why PPO-based RLHF drags a per-token critic along; the pain of training that critic is why later methods try to avoid it.

3. The action space is enormous, but the prior is superb. $10^5$ discrete actions per step would be hopeless from scratch — but we never start from scratch. The pretrained model already concentrates probability mass on sensible continuations, so exploration is warm-started. This is also why every method in this chapter regularizes toward a reference model: the prior is too valuable to wander away from.

4. The reward is not part of nature. In the gridworld, reward was a fixed property of the environment. Here it is a learned model of human preferences, a human directly, or a programmatic verifier — an approximation of what we actually want. Optimizing hard against an approximation is where reward hacking enters, and the chapter’s storyline (RLHF → DPO → RLVR) can be read as successive attempts to make the reward signal harder to game.

The one primitive every method shares#

Because transitions are deterministic, the probability of a whole response factorizes over tokens:

$$ \log \pi_\theta(y \mid x) \;=\; \sum_{t=1}^{|y|} \log \pi_\theta\!\left(y_t \mid x, y_{\lt t}\right), $$

and consequently the score function — the object every policy-gradient estimator in this chapter is built from — decomposes the same way:

$$ \nabla_\theta \log \pi_\theta(y \mid x) \;=\; \sum_{t=1}^{|y|} \nabla_\theta \log \pi_\theta\!\left(y_t \mid x, y_{\lt t}\right). $$

In code, this is one log_softmax, one gather, and one masked sum — the workhorse inside the RLHF, DPO, and GRPO losses of the next three pages:

import torch
torch.manual_seed(0)

# Toy setup: batch of 2 "responses", vocab of 8 tokens, length 5.
# logits[b, t] are the model's next-token logits *before* emitting response[b, t].
B, T, V = 2, 5, 8
logits   = torch.randn(B, T, V)            # from model(prompt + response[:, :-1])
response = torch.randint(0, V, (B, T))     # the tokens actually sampled
mask     = torch.tensor([[1,1,1,1,1],      # response 1: 5 tokens
                         [1,1,1,0,0.]])    # response 2: 3 tokens + padding

# Per-token log-probs: log pi(y_t | x, y_<t)
logps = torch.log_softmax(logits, dim=-1)              # (B, T, V)
token_logps = logps.gather(-1, response.unsqueeze(-1)).squeeze(-1)  # (B, T)

# Sequence log-prob: log pi(y | x) = sum_t log pi(y_t | x, y_<t)
seq_logps = (token_logps * mask).sum(-1)               # (B,)

print("token log-probs:\n", token_logps.round(decimals=3))
print("sequence log-probs:", seq_logps.round(decimals=3))
token log-probs:
 tensor([[-2.2320, -2.8010, -3.9580, -3.3960, -2.2290],
        [-4.4220, -1.5100, -3.0590, -1.9300, -2.2810]])
sequence log-probs: tensor([-14.6160,  -8.9910])

The masking detail is not cosmetic: responses in a batch have different lengths, and whether you sum or average the masked token log-probs — and what you normalize by — turns out to be a genuine design decision with measurable consequences. We will meet it again as a bias in GRPO (Section 4.4).

Deep dive: token-level MDP vs. sequence-level bandit

There are two legitimate readings of the setup, and the literature switches between them freely — often without saying so.

Token-level MDP. Actions are tokens; the horizon is $|y|$; the reward is terminal. This is the reading in the table above.

Sequence-level contextual bandit. The whole response $y$ is a single action taken in "state" $x$; the horizon is $1$. No transitions, no bootstrapping — a bandit from Section 2, just with an astronomically large arm set.

When are they the same? For pure Monte-Carlo policy gradients they are identical. With terminal reward, $\gamma = 1$, and a prompt-dependent baseline $b(x)$, the token-level REINFORCE estimator is

$$ \hat{g} \;=\; \big(R(x, y) - b(x)\big) \sum_{t=1}^{|y|} \nabla_\theta \log \pi_\theta\!\left(y_t \mid x, y_{\lt t}\right) \;=\; \big(R(x, y) - b(x)\big)\, \nabla_\theta \log \pi_\theta(y \mid x), $$

which is exactly the bandit-level REINFORCE estimator, by the factorization above. The same reward multiplies every token's score; the two views produce the same gradient sample-for-sample. GRPO's estimator (Section 4.4) lives precisely in this regime, which is why the bandit reading is the honest one for it.

When do they diverge? The moment anything acts per-token:

  • Critics and GAE. A learned value function $V(s_t)$ over partial sequences assigns different advantages to different tokens — this is meaningful only in the token MDP, and it is what PPO-based RLHF does (Section 4.2).
  • Per-token KL penalties. The standard RLHF implementation adds $-\beta \log\frac{\pi_\theta(y_t \mid \cdot)}{\pi_{\text{ref}}(y_t \mid \cdot)}$ to the reward at every step, giving dense per-token signal — a shaping term that has no bandit analogue.
  • Per-token clipping. PPO clips the importance ratio at each token independently; clipping the single sequence-level ratio would behave very differently (sequence-level ratios are products of hundreds of token ratios and explode or vanish).
  • Discounting. Any $\gamma < 1$ weights early tokens more than late ones and breaks the equivalence. Implementations use $\gamma = 1$ — the same "what implementations actually do" territory as the dropped $\gamma^t$ prefactor in the mathematics companion (Corollary 2).

Rule of thumb: Monte-Carlo methods (REINFORCE-style, GRPO) may be analyzed as bandits; anything with a critic, bootstrapping, or per-token shaping must be analyzed as an MDP. Keeping the two readings straight will save you from several published confusions.

Where this leaves us#

The shape of the MDP now predicts the chapter:

  • Terminal sparse reward + long sequences → advantage estimation is the hard part → RLHF trains a per-token critic and pays for it (Section 4.2).
  • The bandit equivalence above → maybe the critic was never necessary → DPO deletes the RL loop entirely (Section 4.3), and GRPO replaces the critic with a group-mean baseline (Section 4.4).
  • Deterministic known transitions are a luxury of single-turn generation → tools, users, and multi-turn interaction take it away, and the full RL problem of Section 2 returns (Section 4.5).

One connection back to the fundamentals worth making explicit: with $\gamma = 1$ and a hard horizon, this is a finite-horizon MDP — the same setting as the backward-induction framing on the DTR / batch-RL bridge page. The theory we built there transfers unchanged; only the scale is new.

Backward 4. RL for LLMs, Reasoning & Agents 4.2 RLHF: Reward Models + PPO Forward