5.6 Policy Gradients & PPO for Language Models¶

In The RLHF Pipeline & Reward Modeling we built two of the three pieces of Reinforcement Learning from Human Feedback (RLHF): a supervised fine-tuned (SFT) policy that can follow instructions, and a reward model (RM) that maps a prompt-response pair to a scalar “how good is this.” This chapter is about the third piece — the part that actually changes the policy’s weights using the reward. That part is reinforcement learning, and for years the default algorithm was Proximal Policy Optimization (PPO). It is the optimizer behind InstructGPT and the original ChatGPT, and even though critic-free methods like GRPO (GRPO, RLOO & Critic-Free RL) and offline methods like DPO (Direct Preference Optimization & Its Variants) have eaten much of its market share, PPO remains the canonical, most-general, and most-instructive RL algorithm for language models. Everything that came after is best understood as “PPO, but we deleted/changed this part.”

The intellectual problem PPO solves is deceptively hard. Supervised fine-tuning has a differentiable target: for each token there is a “correct” next token, and we minimize cross-entropy. But “produce a response a human would prefer” gives you a single scalar score at the end of a whole generated sequence, with no per-token target and no gradient through the sampling step. We cannot backpropagate through “sample a token from a categorical distribution.” Reinforcement learning is precisely the toolkit for optimizing an expected reward when the only thing you can do is sample actions, observe a scalar, and nudge the action distribution. This chapter develops that toolkit from first principles — the score-function (REINFORCE) estimator, baselines and the advantage, Generalized Advantage Estimation (GAE) — and then assembles the full PPO-for-LLM loss and training loop, with runnable code and a worked numerical example. We end with an honest accounting of why PPO is finicky, which is the reason the rest of Part V exists.

This chapter assumes you are comfortable with the transformer forward pass (Building a GPT From Scratch), autodiff (Automatic Differentiation & PyTorch Internals), and the probability background in Probability, Statistics & Information Theory.

Generation as a Markov decision process¶

To run RL on a language model we first have to see generation as a sequential decision problem. The mapping is exact and worth internalizing because every term in the PPO loss corresponds to one of these objects.

A Markov decision process (MDP) is a tuple \((\mathcal{S}, \mathcal{A}, P, r, \gamma)\): states, actions, a transition kernel \(P(s'\mid s,a)\), a reward \(r(s,a)\), and a discount \(\gamma\). For autoregressive text generation, the correspondence is:

• State \(s_t\) = the prompt plus all tokens generated so far: \(s_t = (q, o_{<t})\) where \(q\) is the prompt and \(o_{<t} = (o_1, \dots, o_{t-1})\). The state is the full context window the model conditions on.
• Action \(a_t = o_t\) = the next token emitted, drawn from the vocabulary \(\mathcal{V}\). So the action space is huge but discrete: \(|\mathcal{A}| = |\mathcal{V}|\), on the order of \(10^5\).
• Policy \(\pi_\theta(a_t \mid s_t) = \pi_\theta(o_t \mid q, o_{<t})\) = exactly the model’s next-token softmax distribution. The language model is the policy. This is the single most important identification in the chapter.
• Transition \(P(s_{t+1}\mid s_t, a_t)\) is deterministic: appending token \(o_t\) to \(s_t\) gives \(s_{t+1} = (q, o_{\le t})\) with probability 1. There is no environment stochasticity; all randomness is in the policy’s own sampling. This is a special, friendly property of the text MDP.
• Reward. In RLHF the reward is terminal and sparse: zero for every token until the sequence ends (at the EOS token or length cap), and then a single scalar \(R(q, o)\) from the reward model. There is no per-token reward signal. (Contrast this with RL with verifiable rewards, where the terminal reward is a programmatic checker; see RL with Verifiable Rewards & The Reasoning Recipe.)
• Episode / trajectory \(\tau = (s_0, a_0, s_1, a_1, \dots)\) = one full generated response. The “horizon” is the response length \(T = |o|\).

The objective we want to maximize is the expected terminal reward of responses sampled from our policy, over the distribution of prompts \(\mathcal{D}\):

The whole game is: nudge \(\theta\) so that responses the reward model likes become more probable. We almost always set the discount \(\gamma = 1\) for LLMs (we care about the whole response equally; there is no infinite horizon to tame), with one important caveat about GAE we will return to.

REINFORCE: the score-function gradient¶

Deriving the policy gradient¶

We want \(\nabla_\theta J(\theta)\), but \(J\) is an expectation over a distribution that depends on \(\theta\) — we cannot just push the gradient inside, because the thing we average over is itself changing. The score-function estimator (a.k.a. the likelihood-ratio trick, a.k.a. REINFORCE) resolves this with one identity. Write the expectation as a sum (or integral) over trajectories, where \(p_\theta(\tau)\) is the probability of trajectory \(\tau\) under the policy:

Now use the log-derivative identity \(\nabla_\theta p_\theta(\tau) = p_\theta(\tau)\,\nabla_\theta \log p_\theta(\tau)\) (which is just the chain rule on \(\log\)). Substituting turns the sum back into an expectation we can estimate by sampling:

This is the crucial move: the gradient of an expectation became an expectation of a gradient, and the troublesome \(\nabla R\) never appears — we do not need the reward to be differentiable. Now expand \(\log p_\theta(\tau)\). The trajectory probability factorizes; the transition terms \(P(s_{t+1}\mid s_t,a_t)\) do not depend on \(\theta\) (the environment is fixed), so their gradient is zero and they drop out:

So \(\nabla_\theta \log p_\theta(\tau) = \sum_t \nabla_\theta \log \pi_\theta(a_t\mid s_t)\), and we arrive at the REINFORCE estimator:

In words: sample a response, then push up the log-probability of every token in it, scaled by the whole-trajectory reward. Good responses (high \(R\)) make all their tokens more likely; bad responses make all theirs less likely. It is almost embarrassingly simple, and it is unbiased.

import torch
import torch.nn.functional as F

def reinforce_loss(token_logprobs, response_mask, rewards):
    """
    Vanilla REINFORCE for a batch of sampled responses.

    token_logprobs : (B, T) log pi_theta(o_t | q, o_<t) for the SAMPLED tokens.
    response_mask  : (B, T) 1.0 on generated tokens, 0.0 on prompt/padding.
    rewards        : (B,)   scalar terminal reward per response.

    We MAXIMIZE E[R * sum_t logprob], so we MINIMIZE its negative.
    rewards are detached: they are weights, not differentiable targets.
    """
    seq_logprob = (token_logprobs * response_mask).sum(dim=-1)   # (B,) sum over tokens
    loss = -(rewards.detach() * seq_logprob).mean()
    return loss

Why vanilla REINFORCE is unusable as-is¶

REINFORCE is unbiased but has catastrophic variance, for two reasons that we fix in the next two sections.

Problem 1 — the common-mode offset. Suppose your reward model outputs scores in, say, \([0, 10]\) (all positive). Then every sampled response has positive reward, so every response’s tokens get pushed up — the gradient says “make everything more likely,” which is incoherent (probabilities sum to 1). The only useful information is the relative differences between responses, and that signal is a tiny ripple on top of a huge positive DC offset. The estimator is unbiased, so in infinite samples the offset cancels (because \(\sum_o \pi_\theta(o)=1\) makes a constant baseline contribute zero in expectation), but with a finite batch the variance is enormous. The fix is a baseline (next section).

Problem 2 — credit assignment. Multiplying every token’s gradient by the whole-trajectory reward means a token early in the sequence is credited for rewards earned by tokens after it (fine, it caused them) but also blamed/credited symmetrically with no notion of which tokens actually mattered. For LLMs with a single terminal reward this is less acute than in dense-reward RL, but the value function and advantage still let us assign credit far more sharply. The fix is the actor-critic advantage and GAE (two sections down).

Baselines, the value function, and the advantage¶

The baseline trick¶

We can subtract from the reward any quantity \(b(s_t)\) that does not depend on the action \(a_t\) without biasing the gradient. The proof is one line: the expected contribution of the baseline term is

using the same log-derivative identity in reverse (\(\sum_a \pi \nabla \log\pi = \sum_a \nabla\pi = \nabla\sum_a\pi = \nabla 1 = 0\)). So baselines change variance but not the mean — they are a free lunch. The variance-minimizing baseline is (approximately) the expected reward from state \(s_t\), which is exactly the value function.

The value function and the advantage¶

Define the state-value function as the expected return (sum of future rewards) from state \(s_t\) under the current policy:

For RLHF with terminal reward only, \(V^\pi(s_t)\) is “the reward model score we expect to earn if we continue generating from this partial response with the current policy.” Define the action-value \(Q^\pi(s_t, a_t)\) analogously but committing to action \(a_t\) first. The advantage is their difference:

The advantage answers the precise question we care about: was taking token \(a_t\) in state \(s_t\) better or worse than the policy’s average behavior there? Positive advantage → make that token more likely; negative → less likely. Using \(V\) as the baseline, the policy gradient becomes the actor-critic form:

This is the engine of PPO. The actor is the policy \(\pi_\theta\); the critic is a learned estimate \(V_\phi \approx V^\pi\). In an LLM, the critic is typically a separate value head on top of a transformer (often a copy of the policy or the reward model with a scalar output per token) — a second model nearly as expensive as the policy itself. That expense is exactly what GRPO later deletes.

We will estimate \(V_\phi\) by regression (the critic loss) and combine it with the observed return to estimate \(A_t\). The cleanest, lowest-variance way to do that combination is GAE.

Generalized Advantage Estimation (GAE)¶

The bias-variance dial¶

We need an estimate \(\hat A_t\) of the advantage from sampled data plus the critic \(V_\phi\). There is a spectrum of estimators trading bias against variance, indexed by how many real reward steps you use before falling back on the critic’s guess.

Define the TD residual (temporal-difference error) at step \(t\):

This is “the reward you actually got plus the critic’s value of where you landed, minus the critic’s prediction for where you were” — a one-step correction to the value estimate. The two extremes:

• One-step TD (\(\hat A_t = \delta_t\)): low variance (uses only one real reward), but biased whenever \(V_\phi\) is wrong.
• Monte Carlo (\(\hat A_t = \sum_{k\ge t}\gamma^{k-t} r_k - V_\phi(s_t)\)): unbiased (uses the actual full return), but high variance.

GAE (Schulman et al., 2016) interpolates with an exponential weighting controlled by \(\lambda \in [0,1]\):

\(\lambda = 0\) recovers one-step TD (low variance, high bias); \(\lambda = 1\) recovers the Monte Carlo advantage (high variance, unbiased — the \(V_\phi\) terms telescope away). In practice \(\lambda \approx 0.95\) and \(\gamma \approx 1.0\) for LLMs. The advantage and the critic’s regression target (the return) are computed together:

GAE is computed by a single backward recursion over the sequence — start at the last token and accumulate:

import torch

def compute_gae(rewards, values, response_mask, gamma=1.0, lam=0.95):
    """
    Generalized Advantage Estimation, computed per token by a backward pass.

    rewards       : (B, T) per-token reward. For RLHF this is mostly zeros,
                    with the reward-model score (and any KL penalty) placed
                    at the last response token of each sequence.
    values        : (B, T) critic estimates V_phi(s_t) for each token.
    response_mask : (B, T) 1.0 on response tokens, 0.0 elsewhere.
    Returns:
      advantages (B, T), returns (B, T) = advantages + values.
    """
    B, T = rewards.shape
    advantages = torch.zeros_like(rewards)
    last_adv = torch.zeros(B, device=rewards.device)
    # V(s_{T}) for the bootstrap of the final step is 0 (episode ends at EOS).
    next_value = torch.zeros(B, device=rewards.device)

    for t in reversed(range(T)):
        mask_t = response_mask[:, t]
        # delta_t = r_t + gamma * V(s_{t+1}) - V(s_t)
        delta = rewards[:, t] + gamma * next_value - values[:, t]
        last_adv = delta + gamma * lam * last_adv
        # zero out advantage/last_adv outside the response so padding can't leak
        advantages[:, t] = last_adv * mask_t
        last_adv = last_adv * mask_t
        next_value = values[:, t] * mask_t  # carry V(s_t) leftward as V(s_{t+1})

    returns = advantages + values
    # Whitening (mean 0, std 1) over response tokens stabilizes the scale.
    flat = advantages[response_mask.bool()]
    advantages = (advantages - flat.mean()) / (flat.std() + 1e-8)
    advantages = advantages * response_mask
    return advantages, returns

The PPO objective¶

From policy gradient to a surrogate we can reuse¶

REINFORCE/actor-critic is an on-policy estimator: the gradient is only valid for data generated by the current \(\theta\). But generating rollouts is the expensive part (it’s autoregressive decoding on a server; see The Generation–Training Loop & Rollout Engines), so we badly want to take several gradient steps on each batch of rollouts. The moment we update \(\theta\) once, the data is off-policy and the naive gradient is biased.

The principled fix is importance sampling. If data was generated by a behavior policy \(\pi_{\theta_{\text{old}}}\) but we want the gradient for the current \(\pi_\theta\), we reweight by the ratio:

The corresponding surrogate objective \(L^{\text{IS}}(\theta) = \mathbb{E}[ r_t(\theta)\,\hat A_t ]\) has the right gradient at \(\theta = \theta_{\text{old}}\). But importance sampling is dangerous: if \(\pi_\theta\) moves far from \(\pi_{\theta_{\text{old}}}\), the ratio \(r_t\) can blow up (a token that was rare under old policy becomes likely under new), the estimator variance explodes, and a single huge ratio times a large advantage can destroy the policy in one step. The whole point of PPO is to bound how far the update can exploit this ratio.

The clipped surrogate¶

PPO’s predecessor, TRPO (Trust Region Policy Optimization), enforced a hard KL constraint between old and new policy via a second-order optimization — correct but heavy. PPO (Schulman et al., 2017) achieves a similar trust-region effect with a stunningly simple first-order trick: clip the ratio. The PPO clipped objective per token is

and we maximize \(\mathbb{E}_t[L^{\text{CLIP}}_t]\). Typical \(\epsilon = 0.2\). Let us read this carefully, because the min and the sign of \(\hat A\) interact in a way that trips people up.

• When \(\hat A_t > 0\) (good token, we want \(r_t \uparrow\)): the unclipped term \(r_t \hat A\) grows with \(r_t\), but the clipped term caps \(r_t\) at \(1+\epsilon\). The min picks the smaller, so once \(r_t > 1+\epsilon\) the objective is flat at \((1+\epsilon)\hat A\) — its gradient is zero. Translation: you get full credit for boosting a good token up to \(1+\epsilon\), but no extra reward for boosting it further. No incentive to take a huge step.
• When \(\hat A_t < 0\) (bad token, we want \(r_t \downarrow\)): the unclipped term \(r_t\hat A\) becomes more negative as \(r_t\) grows and less negative as \(r_t\) shrinks. The clipped term floors \(r_t\) at \(1-\epsilon\). The min of two negatives picks the more negative one… which means the clip engages on the lower side: once \(r_t < 1-\epsilon\), the objective is flat at \((1-\epsilon)\hat A\), zero gradient. Translation: you get full credit for suppressing a bad token down to \(1-\epsilon\), but no extra reward for crushing it further.

The asymmetry is the key insight: the min makes the clip into a one-sided trust region that only ever removes the incentive to move too far in the rewarding direction. It never clips a step that is correcting an overshoot back toward the old policy. The objective is a pessimistic (lower) bound on the unclipped surrogate.

def ppo_policy_loss(new_logprobs, old_logprobs, advantages, mask, clip_eps=0.2):
    """
    PPO clipped policy (actor) loss for one minibatch.

    new_logprobs : (B, T) log pi_theta(o_t|s_t)  -- requires grad
    old_logprobs : (B, T) log pi_old(o_t|s_t)    -- detached, from rollout time
    advantages   : (B, T) detached GAE advantages (already whitened)
    mask         : (B, T) response mask
    """
    log_ratio = new_logprobs - old_logprobs            # log(π_θ / π_old)
    ratio = log_ratio.exp()                            # the importance ratio r_t

    unclipped = ratio * advantages
    clipped   = torch.clamp(ratio, 1 - clip_eps, 1 + clip_eps) * advantages
    # maximize min(...)  ==  minimize -min(...)
    per_token_loss = -torch.min(unclipped, clipped)

    loss = (per_token_loss * mask).sum() / mask.sum().clamp(min=1.0)

    # Diagnostics every practitioner watches:
    with torch.no_grad():
        clipfrac = ((ratio - 1.0).abs() > clip_eps)[mask.bool()].float().mean()
        approx_kl = ((ratio - 1) - log_ratio)[mask.bool()].mean()  # k3 KL(old||new)
    return loss, clipfrac, approx_kl

The value (critic) loss¶

The critic \(V_\phi\) is trained by regression toward the GAE returns \(\hat R_t\). PPO uses a clipped value loss too — to stop the critic from moving too far per update, mirroring the policy clip:

The max makes this a pessimistic (larger) loss, so the critic is penalized for both under- and over-shooting beyond the clip window. Many modern implementations find the value-clip optional or even slightly harmful and use plain MSE; we include it because the canonical PPO-RLHF recipe (and TRL) uses it.

def ppo_value_loss(values, old_values, returns, mask, clip_eps_v=0.2):
    """Clipped value-function regression loss (critic)."""
    v_clipped = old_values + torch.clamp(values - old_values, -clip_eps_v, clip_eps_v)
    loss_unclipped = (values   - returns) ** 2
    loss_clipped   = (v_clipped - returns) ** 2
    per_token = 0.5 * torch.max(loss_unclipped, loss_clipped)   # pessimistic
    return (per_token * mask).sum() / mask.sum().clamp(min=1.0)

The KL penalty to the reference policy¶

Why we need it¶

Maximize reward-model score alone and the policy will over-optimize: it finds adversarial responses that the (imperfect, finite-data) reward model scores highly but humans hate — repetitive sycophancy, weird formatting, exploitation of RM blind spots. This is reward hacking (Reward Hacking, Over-Optimization & Alignment Failures), and it is the dominant failure mode of unconstrained RLHF. The standard guardrail is a KL penalty that keeps the policy close to the frozen reference (the SFT model \(\pi_{\text{ref}}\)): drift too far and you pay.

We penalize the per-token KL divergence \(\mathbb{D}_{\text{KL}}[\pi_\theta(\cdot\mid s_t)\,\|\,\pi_{\text{ref}}(\cdot\mid s_t)]\). The full RLHF objective is therefore the KL-regularized expected reward:

This regularized objective is exactly what DPO solves in closed form (Direct Preference Optimization & Its Variants) — DPO and PPO optimize the same target by different means. The coefficient \(\beta\) is the single most consequential hyperparameter in PPO-RLHF: too small and you reward-hack; too large and the policy barely moves and you get no improvement.

Folding KL into the per-token reward¶

The dominant implementation trick (from InstructGPT) is to not add KL as a separate loss term but to bake it into the per-token reward before computing GAE. At each response token we set

a per-token KL penalty, and we add the reward-model score \(R(q,o)\) only at the final token:

This is elegant: GAE then propagates the KL penalty through the value function and advantage automatically, so the policy is shaped at every token by “stay near the reference” while being pulled at the end toward “score well.” Note \(\log\pi_\theta - \log\pi_{\text{ref}}\) for the sampled token is the k1 single-sample KL estimator; it is unbiased but can be negative for individual tokens. Some recipes instead use Schulman’s always-positive k3 estimator \(\rho - \log\rho - 1\) (with \(\rho = \pi_{\text{ref}}/\pi_\theta\)) for lower variance; see the discussion in GRPO, RLOO & Critic-Free RL.

def make_token_rewards(policy_logprobs, ref_logprobs, rm_scores, mask, kl_beta=0.05):
    """
    Build the per-token reward tensor fed to GAE: a per-token KL penalty
    everywhere, plus the scalar reward-model score at the last response token.

    policy_logprobs, ref_logprobs : (B, T) log-probs of the SAMPLED tokens.
    rm_scores : (B,) one scalar per response from the reward model.
    """
    # k1 per-token KL penalty: -beta * (logπ_θ - logπ_ref)
    kl = policy_logprobs - ref_logprobs                  # (B, T)
    token_rewards = -kl_beta * kl * mask                 # KL on every response token

    # Index of the last response token per row, and drop the RM score there.
    # Found positionally (not via mask.sum()-1): after the next-token shift,
    # `mask` still carries the prompt's leading zeros, so "count of response
    # tokens minus one" is NOT the same column as "the last column where
    # mask==1" unless the prompt has length 1. Locating the rightmost 1
    # directly keeps this correct for any prompt length.
    positions = torch.arange(mask.size(1), device=mask.device).unsqueeze(0).expand_as(mask)
    last_idx = (positions * mask).amax(dim=1).long()      # (B,)
    rows = torch.arange(token_rewards.size(0), device=mask.device)
    token_rewards[rows, last_idx] += rm_scores           # terminal RM reward
    return token_rewards

The full PPO-for-LLM loop¶

We now assemble everything. The total per-minibatch loss combines the three pieces — clipped policy loss, value loss, and an entropy bonus \(S[\pi_\theta]\) that discourages premature collapse to a deterministic policy (it keeps exploration alive):

(The KL-to-reference is already inside \(\hat A_t\) via the per-token reward; the entropy term here is a different, intra-policy regularizer.) Typical coefficients: \(c_v \approx 0.1\)–\(1.0\), \(c_e \approx 0\)–\(0.01\).

The training loop has a characteristic two-phase rhythm — a rollout phase (generate, score, compute advantages) and an optimization phase (several epochs of minibatch gradient steps on the frozen rollouts), which is the structure every RL-infra system in Part VI is built around:

Here is a compact but complete PPO step that ties the helper functions together. It is written for clarity over speed; a production system (TRL’s PPOTrainer, OpenRLHF, veRL) disaggregates generation onto an inference engine and shards the four models, but the math is identical. The one helper left to you is generate_batch: ordinary batched temperature sampling from the policy (Sampling Strategies & Decoding Algorithms) that returns the prompt-plus-response input_ids and a mask that is 1 exactly on generated tokens — sample at temperature 1.0 with no top-\(k\)/top-\(p\) truncation, because the importance ratio and the KL are only valid if the tokens really were drawn from \(\pi_{\text{old}}\) and not from a truncated version of it.

import torch
import torch.nn.functional as F

# Assume: policy (with value head), ref_model, reward_model, tokenizer, optimizer.
# policy(input_ids) returns logits AND a per-token scalar value (value head).

PPO_EPOCHS   = 4
MINIBATCHES  = 4
CLIP_EPS     = 0.2
KL_BETA      = 0.05
VF_COEF      = 0.5
ENT_COEF     = 0.0
TARGET_KL    = 0.02   # early-stop the epoch loop if approx_kl exceeds this

def token_logprobs_and_values(model, input_ids):
    """Per-token log-prob of the realized next token, plus value-head output."""
    out = model(input_ids)                                   # logits (B,T,V), values (B,T)
    logits, values = out.logits[:, :-1], out.values[:, :-1]
    logp = F.log_softmax(logits.float(), dim=-1)
    targets = input_ids[:, 1:].unsqueeze(-1)
    token_lp = logp.gather(-1, targets).squeeze(-1)          # (B, T-1)
    return token_lp, values

@torch.no_grad()
def ppo_rollout(prompts):
    # 1-2: generate responses with the current (behavior) policy.
    input_ids, resp_mask = generate_batch(policy, prompts)   # your decode fn
    # 3: reward-model scores (one scalar per sequence).
    rm_scores = reward_model.score(input_ids, resp_mask)     # (B,)
    # 4: cache behavior log-probs, reference log-probs, and values.
    old_lp,  old_values = token_logprobs_and_values(policy,    input_ids)
    ref_lp,  _          = token_logprobs_and_values(ref_model, input_ids)
    mask = resp_mask[:, 1:]                                   # align with shifted targets
    # 5: per-token rewards = KL penalty + terminal RM score.
    token_rewards = make_token_rewards(old_lp, ref_lp, rm_scores, mask, KL_BETA)
    # 6: GAE advantages and returns.
    advantages, returns = compute_gae(token_rewards, old_values, mask)
    return dict(input_ids=input_ids, mask=mask, old_lp=old_lp,
                old_values=old_values, advantages=advantages, returns=returns,
                mean_reward=rm_scores.mean().item())

def ppo_update(buf):
    B = buf["input_ids"].size(0)
    idx = torch.randperm(B)
    mb_size = B // MINIBATCHES
    for _ in range(PPO_EPOCHS):
        for m in range(MINIBATCHES):
            sl = idx[m*mb_size:(m+1)*mb_size]
            new_lp, new_values = token_logprobs_and_values(policy, buf["input_ids"][sl])

            p_loss, clipfrac, approx_kl = ppo_policy_loss(
                new_lp, buf["old_lp"][sl], buf["advantages"][sl], buf["mask"][sl], CLIP_EPS)
            v_loss = ppo_value_loss(
                new_values, buf["old_values"][sl], buf["returns"][sl], buf["mask"][sl])
            entropy = policy_entropy(policy, buf["input_ids"][sl], buf["mask"][sl])

            loss = p_loss + VF_COEF * v_loss - ENT_COEF * entropy
            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(policy.parameters(), 1.0)
            optimizer.step()
        if approx_kl > 1.5 * TARGET_KL:        # KL early-stopping safeguard
            break

# Driver:
# for it in range(NUM_ITERS):
#     buf = ppo_rollout(sample_prompts(batch_size))
#     ppo_update(buf)

The same loop in a real library¶

You write the loop above once, to understand it; in production you use a library that has already fixed the fifty details (padding, EOS handling, sharding, generation offload). In HuggingFace TRL the whole two-phase rhythm is one object, and every constructor argument is a symbol derived above — which is the fastest way to check your mental model:

# pip install "trl>=0.12" transformers accelerate datasets
from datasets import load_dataset
from transformers import (AutoModelForCausalLM, AutoModelForSequenceClassification,
                          AutoTokenizer)
from trl import PPOConfig, PPOTrainer

BASE    = "Qwen/Qwen2.5-0.5B-Instruct"
RM_PATH = "path/to/your-reward-model"   # the RM you trained in chapter 5.5

tok       = AutoTokenizer.from_pretrained(BASE)
policy    = AutoModelForCausalLM.from_pretrained(BASE)   # actor, trains
ref_model = AutoModelForCausalLM.from_pretrained(BASE)   # frozen π_ref for the KL leash
# RM and critic are both scalar-head models; the critic trains, the RM is frozen.
# The critic is conventionally *initialized from the RM* — it already knows what
# "a good continuation" looks like, so V_phi starts far from random.
reward_model = AutoModelForSequenceClassification.from_pretrained(RM_PATH, num_labels=1)
value_model  = AutoModelForSequenceClassification.from_pretrained(RM_PATH, num_labels=1)

cfg = PPOConfig(
    output_dir="ppo-out",
    learning_rate=1e-6,
    per_device_train_batch_size=8,
    num_ppo_epochs=4,          # PPO_EPOCHS above
    num_mini_batches=4,        # MINIBATCHES above
    total_episodes=20_000,     # total rollouts over the run
    kl_coef=0.05,              # β in the per-token KL reward
    cliprange=0.2,             # ε in L^CLIP
    cliprange_value=0.2,       # ε_v in the clipped value loss
    vf_coef=0.1,               # c_v
    gamma=1.0, lam=0.95,       # GAE dial
    whiten_rewards=False,      # advantages are whitened regardless
    missing_eos_penalty=1.0,   # penalize samples that never terminated
)

# The dataset must supply tokenized prompts in an "input_ids" column.
ds = load_dataset("trl-lib/tldr", split="train")

trainer = PPOTrainer(args=cfg, processing_class=tok, model=policy, ref_model=ref_model,
                     reward_model=reward_model, value_model=value_model,
                     train_dataset=ds)
trainer.train()

TRL rewrote PPOTrainer around the standard Trainer interface in v0.12 (the older AutoModelForCausalLMWithValueHead + trainer.step(...) style predates it), and config field names have shifted across releases — always read PPOConfig’s docstring for the version you install (TRL: HuggingFace’s RL Library). At larger scale the same objects appear in veRL (whose verl/trainer/ppo/core_algos.py contains a GAE routine that is line-for-line the compute_gae above, plus the clipped policy/value losses) and OpenRLHF (Ray actors for policy/critic/RM/reference with vLLM rollouts) — see veRL: HybridFlow & The Single-Controller Architecture and OpenRLHF, NeMo-Aligner & Ray-Based Systems. Reading one of those core_algos-style files after this chapter is a genuinely short trip: you will recognize every function.

Why PPO is finicky¶

PPO works, and InstructGPT proved it scales. But practitioners universally describe it as brittle and operationally heavy, and understanding why is exactly what motivates the rest of Part V and Part VI. The difficulties stack:

1. Four models in memory at once. A PPO step needs the policy (trains), the critic/value net (trains, often as large as the policy), the reward model (frozen, large), and the reference model (frozen). On a 7B policy that is roughly \(4\times\) the parameter footprint plus the policy’s and critic’s optimizer states (Adam keeps two moments per trainable parameter; see Optimizers). Fitting this requires sharding (FSDP/ZeRO; see Distributed Training I) and careful memory choreography — and you still have to generate with the policy, which competes for the same GPUs (Colocated vs Disaggregated RL & Weight Synchronization).

2. The critic is hard to train and a major instability source. A value head regressing sparse terminal rewards through a 7B transformer is itself a finicky learning problem. If \(V_\phi\) is biased, your advantages are biased, the policy chases a wrong gradient, and the whole thing diverges quietly — reward looks fine, then collapses. Many practitioners spend more time debugging the critic than the policy. This single pain point is the entire reason GRPO and RLOO delete the critic (GRPO, RLOO & Critic-Free RL). If you do keep a critic, the standard hardening is: initialize it from the reward model (not from random, and not by bolting an untrained scalar head onto the policy), zero-initialize the final scalar layer so \(V_\phi \approx 0\) at step 0 rather than emitting arbitrary large values, run a critic-only warmup of a few dozen steps on frozen rollouts before letting the policy move at all, and give the critic its own (usually larger) learning rate. For long chain-of-thought responses a length-adaptive GAE \(\lambda\) helps too, since a \(\lambda\) tuned at 500 tokens over-smooths at 8,000 — the VAPO line of work is essentially a catalogue of these critic fixes (Scaling RL: Throughput, Load Balancing & The Latest Tricks, Advantage Estimation, KL Control & Stability Tricks).

3. A thicket of coupled hyperparameters. \(\beta\) (KL), \(\epsilon\) (clip), \(\lambda\) and \(\gamma\) (GAE), \(c_v\) and \(c_e\), learning rates for policy and critic, number of PPO epochs, minibatch size, the reward-normalization scheme, the target KL for adaptive control. These interact non-linearly: a too-large LR with too-small \(\beta\) reward-hacks; a too-large \(\beta\) stalls; too many PPO epochs makes the data badly off-policy and the clip can’t save you. The viable region is narrow and problem-dependent, and there is no clean loss curve telling you you’re in it — you must watch reward, KL, clipfrac, value loss, and entropy together.

4. Reward over-optimization is always lurking. Even with the KL leash, push long enough and the policy finds RM blind spots; reward climbs while true quality falls (the “Goodharting” curve, often visible as reward up / human-eval down). The KL coefficient is a blunt instrument against a sharp problem. See Reward Hacking, Over-Optimization & Alignment Failures.

5. Generation–training skew. In production you generate on a fast inference engine (vLLM/SGLang) and train on a different stack; numerical differences (kernels, precision, KV-cache layout) mean the log-probs you think the behavior policy had can differ slightly from what the trainer recomputes — biasing the importance ratio. Keeping the sampler and trainer log-probs consistent is a real engineering tax (The Generation–Training Loop & Rollout Engines).

6. Sample inefficiency and cost. Each iteration regenerates fresh rollouts (you can only reuse them for a few epochs before the clip stops being valid). Generation dominates wall-clock. This is why DPO — which needs no online generation — is so attractive when you already have preference pairs, and why so much RL-infra work targets rollout throughput (Scaling RL: Throughput, Load Balancing & The Latest Tricks).

The honest summary: PPO is the right general algorithm and the most educational one — every term has a clean theoretical justification — but it asks a lot of the engineer. The field’s trajectory since 2023 has been a steady search for things that keep PPO’s good behavior (trust-region stability, online exploration) while removing its pain (the critic, the hyperparameter thicket, the offline-impossibility). DPO removed the RL loop entirely; GRPO/RLOO removed the critic; RLVR removed the reward model in favor of verifiers. You cannot understand why any of those exist without first understanding the PPO they are simplifying — which is this chapter.

This is also the calculus we make explicitly in the capstone: Stack-100M’s RL stage uses critic-free GRPO on a verifiable task, not PPO, precisely because a second ~100M network plus its optimizer state buys almost nothing when the reward is terminal and a group baseline is available (Post-Training: SFT, DPO, and Narrow RLVR (GRPO) That Works at 100M). But essentially all of the bookkeeping in this chapter carries over verbatim to that run — gathering per-token log-probs of the sampled tokens, the response mask and its off-by-one shift, the reference-model forward pass and the per-token KL, advantage whitening, the clipped ratio, clipfrac/KL/entropy diagnostics. Write these helpers once, here, and the capstone’s RL loop is a small edit away (replace GAE with a group mean, delete the value model).

Further reading¶

• Sutton & Barto, Reinforcement Learning: An Introduction (2^nd ed.) — the canonical text for MDPs, policy gradients, value functions, and the actor-critic framework.
• Williams, Simple Statistical Gradient-Following Algorithms for Connectionist Reinforcement Learning (1992) — the original REINFORCE / score-function estimator.
• Schulman, Levine, Moritz, Jordan, Abbeel, Trust Region Policy Optimization (2015) — the trust-region predecessor PPO approximates.
• Schulman, Moritz, Levine, Jordan, Abbeel, High-Dimensional Continuous Control Using Generalized Advantage Estimation (2016) — GAE.
• Schulman, Wolski, Dhariwal, Radford, Klimov, Proximal Policy Optimization Algorithms (2017) — the clipped objective itself.
• Ouyang, Wu, Jiang, et al. (OpenAI), Training Language Models to Follow Instructions with Human Feedback (InstructGPT, 2022) — the canonical PPO-for-RLHF recipe, per-token KL reward, and adaptive KL control.
• Ziegler, Stiennon, Wu, et al., Fine-Tuning Language Models from Human Preferences (2019) — the earlier blueprint for KL-regularized RLHF.
• Huang et al., The N Implementation Details of RLHF with PPO / the CleanRL and TRL PPO implementations — concrete, debuggable reference code; see TRL: HuggingFace’s RL Library.
• John Schulman, Approximating KL Divergence (blog note) — the k1/k2/k3 KL estimators used for the penalty.

Exercises¶

1. In the text MDP the transition kernel \(P(s_{t+1}\mid s_t, a_t)\) is deterministic, unlike robotics or Atari where the environment injects its own randomness. (a) Explain concretely, in terms of tokens and context windows, why appending a token is a deterministic transition. (b) Name two distinct mechanical consequences of this determinism that make RL for LLMs easier than classic RL, and tie each to a specific quantity that appears in the PPO loss.

2. A single-token bandit (“choose one of three responses”) has current policy probabilities \(\pi = (0.2, 0.5, 0.3)\) for responses \((A, B, C)\) and reward-model scores \(R = (8.0, 7.0, 9.0)\). (a) Compute the value-function baseline \(b = V^\pi = \mathbb{E}_{o\sim\pi}[R]\). (b) Compute the advantage \(A_i = R_i - b\) for each response and state which tokens get pushed up vs. down. © The chapter says a constant baseline does not bias the gradient. Verify this for this bandit by showing that \(\sum_i \pi_i (R_i - b)\,\nabla_\theta\log\pi_i\) has the same expectation as \(\sum_i \pi_i R_i \nabla_\theta\log\pi_i\), i.e. that the extra term \(\sum_i \pi_i\, b\, \nabla_\theta \log \pi_i\) vanishes.

3. Take one response with three real tokens (response_mask = [1, 1, 1]), gamma = 1.0, lam = 0.5. The terminal reward-model score \(R = 6\) sits on the last token: rewards = [0, 0, 6]. The critic outputs values = [2, 4, 1]. Using the terminal bootstrap \(V(s_T) = 0\), compute (a) the TD residuals \(\delta_t\), (b) the raw GAE advantages \(\hat A_t\) via the backward recursion, and © the critic’s regression targets \(\hat R_t = \hat A_t + V_\phi(s_t)\) (returns, before whitening).

4. With \(\epsilon = 0.2\), evaluate the per-token clipped objective \(L^{\text{CLIP}}_t = \min\!\big(r_t\hat A_t,\ \operatorname{clip}(r_t,\,0.8,\,1.2)\,\hat A_t\big)\) and state whether the clip engages (i.e. whether the gradient through this token is zero) for each token. (a) \(\hat A_t = +0.5\), \(r_t = 1.35\). (b) \(\hat A_t = -0.5\), \(r_t = 0.70\). © \(\hat A_t = +0.5\), \(r_t = 1.10\). Then (d) explain in one sentence why the clip in (a) is not a bug even though we wanted to push this good token up.

5. The chapter’s make_token_rewards uses the k1 KL estimator \(\log\pi_\theta - \log\pi_{\text{ref}}\), which is unbiased but can be negative for an individual token (giving that token a positive reward for sitting below the reference). Reimplement it using Schulman’s always-non-negative k3 estimator \(\rho - \log\rho - 1\) with \(\rho = \pi_{\text{ref}}/\pi_\theta\), keeping the terminal-RM-score placement identical. Then compute the k3 per-token KL penalty for the chapter’s interior token (\(\pi_\theta = 0.14\), \(\pi_{\text{ref}} = 0.08\), \(\beta = 0.05\)) and compare its sign and magnitude to the k1 value of \(-0.028\).

def make_token_rewards_k3(policy_logprobs, ref_logprobs, rm_scores, mask, kl_beta=0.05):
    """Same as make_token_rewards but with Schulman's k3 (always >= 0) KL estimator."""
    log_rho = ref_logprobs - policy_logprobs          # log(pi_ref / pi_theta)
    rho = log_rho.exp()
    kl_k3 = rho - log_rho - 1.0                        # k3, elementwise, always >= 0
    token_rewards = -kl_beta * kl_k3 * mask            # KL penalty on every response token

    # Terminal RM score on the last unmasked token (found positionally, as before).
    positions = torch.arange(mask.size(1), device=mask.device).unsqueeze(0).expand_as(mask)
    last_idx = (positions * mask).amax(dim=1).long()   # (B,)
    rows = torch.arange(token_rewards.size(0), device=mask.device)
    token_rewards[rows, last_idx] += rm_scores
    return token_rewards

6. The full PPO loss uses an entropy bonus, and the training loop calls entropy = policy_entropy(policy, buf["input_ids"][sl], buf["mask"][sl]), but policy_entropy is never defined in the chapter. Implement it: return the mask-averaged per-token entropy \(S = -\sum_{v}\pi_\theta(v\mid s_t)\log\pi_\theta(v\mid s_t)\) of the policy’s next-token distribution, consistent with the chapter’s shift/alignment convention (logits [:, :-1] line up with the mask derived from resp_mask[:, 1:]). Explain in one line why maximizing this entropy (adding -ENT_COEF * entropy to the loss) helps.

import torch
import torch.nn.functional as F

def policy_entropy(model, input_ids, mask):
    """
    Mean per-token entropy of the policy's next-token distribution over the
    response tokens. Shapes match the chapter: logits[:, :-1] aligns with the
    mask taken from resp_mask[:, 1:].
    """
    out = model(input_ids)                       # logits (B, T, V)
    logits = out.logits[:, :-1]                  # (B, T-1, V), drop last position
    logp = F.log_softmax(logits.float(), dim=-1) # (B, T-1, V)
    p = logp.exp()
    token_entropy = -(p * logp).sum(dim=-1)      # (B, T-1)  S = -sum_v p log p
    return (token_entropy * mask).sum() / mask.sum().clamp(min=1.0)