Reinforcement Learning for Language Models: From Policy Gradient to PPO

Every model in this series so far learned from a target: the next token in the text, a label, a teacher’s distribution. Some of what we want from a language model has no target to copy. “A helpful answer”, “code that passes the tests”, “a proof that checks” can be scored once the model has written something, but there’s no single right sequence of tokens to train on, and the score isn’t a differentiable function of the model’s weights: it comes from a test runner, a checker, or a person. Reinforcement learning trains on scores like that. This post builds the method that RLHF, DPO, and GRPO and Reinforcement Learning from Verifiable Rewards use to align and train today’s models.

Table of contents

The vocabulary, mapped onto text

Reinforcement learning describes an agent acting in an environment (Sutton and Barto, 2018). For a language model generating a response:

• the state $s_t$ is the prompt plus the tokens generated so far;
• an action $a_t$ is the next token;
• the policy $\pi_\theta(a_t \mid s_t)$ is the language model itself, its next-token distribution;
• an episode is one complete response, ending at the end-of-turn token;
• the reward $r$ usually comes once, at the end: a reward model’s score (see RLHF), or 1 if the answer checks out and 0 if not (see GRPO and Reinforcement Learning from Verifiable Rewards);
• the return $R$ is the sum of rewards over the episode. Episodes are short and end for certain, so future rewards are usually not discounted.

The goal is to maximize the expected return, $J(\theta) = \mathbb{E}_{y \sim \pi_\theta}[R(y)]$, over responses $y$ the model itself samples.

REINFORCE, derived

$R$ can’t be differentiated with respect to $\theta$, but the probability of each response can, and that’s enough. Writing the expectation as a sum over responses,

using $\nabla \pi = \pi \nabla \log \pi$, the log-derivative trick. The last form is an average over samples, so it can be estimated by sampling: generate responses, score them, and take a gradient step on $R(y) \log \pi_\theta(y)$, treating $R$ as a constant. And $\log \pi_\theta(y) = \sum_t \log \pi_\theta(a_t \mid s_t)$ is the same sum of token log-probabilities as the training loss of Cross-Entropy and the Softmax Gradient. This is REINFORCE (Williams, 1992): ordinary maximum-likelihood training on the model’s own samples, each weighted by its reward. A high-reward response is made more likely, token by token.

Variance, and the free baseline

The estimate is unbiased but noisy. Worse, its noise depends on where rewards happen to sit. If every response scores between 10 and 13, every sampled response is pushed up, and only the small differences between the pushes carry information: most of each step is noise that later steps must undo.

Subtracting any baseline $b$ that doesn’t depend on the response fixes that without adding bias, because its contribution averages to zero:

With $b$ the average reward, responses better than average are pushed up and worse ones down. The weight $R(y) - b$ is the advantage: how much better this response was than expected.

Per token, the baseline may depend on the state $s_t$ but not on the token chosen there (the derivation above then still gives zero bias). The standard choice is the expected return from that state, the value $V(s_t)$, which PPO learns with a value head, a second output trained jointly with the policy by regression to predict returns. The advantage of a token is then estimated from the value’s change, $\delta_t = r_t + \gamma V(s_{t+1}) - V(s_t)$, with $V = 0$ after the last token: with $\gamma = 1$, as is usual for text, this is the change in value plus any reward, positive if the token moved the response to a state worth more than before. Generalized advantage estimation (GAE; Schulman et al., 2016) averages these over the following tokens with decaying weights, $\hat{A}_t = \sum_{l \ge 0} (\gamma \lambda)^l \delta_{t+l}$. At $\lambda = 0$ it is the one-step $\delta_t$, trusting the value head for everything after one token (low variance, biased if the value head is wrong). At $\lambda = 1$ the sum telescopes to the actual return minus $V(s_t)$, which is unbiased however wrong the value head is, but high-variance.

Trust regions and PPO

A policy-gradient step is only accurate near the policy that sampled the data. Take too large a step and the policy can collapse, assigning nearly all probability to some degenerate output, after which its own samples no longer teach it anything. TRPO (Schulman et al., 2015) constrains each update so that the average KL divergence from the old policy stays below a fixed bound. PPO (Schulman et al., 2017) gets most of the effect with a simpler objective. For each token, take the ratio of new to old probability, $\rho = \pi_\theta(a_t \mid s_t) / \pi_{\text{old}}(a_t \mid s_t)$, and maximize (the paper also proposes a variant with a KL penalty in place of the clip)

Why a ratio? The samples came from $\pi_\text{old}$, so to estimate the gradient for a different $\pi_\theta$, each sample is reweighted by $\pi_\theta / \pi_\text{old}$ (importance sampling). At $\theta = \theta_\text{old}$, $\nabla_\theta \rho = \rho\, \nabla_\theta \log \pi_\theta = \nabla_\theta \log \pi_\theta$, so the gradient of $\rho \hat{A}$ is exactly the policy gradient of the previous section; the clip only matters once later steps on the same batch move $\rho$ away from 1.

Worked, for one token whose probability has already risen to 1.3 times its old value, $\rho = 1.3$:

• Advantage $+1$ (the token was good): $\min(1.3, 1.2) = 1.2$. The clipped term wins, and it doesn’t depend on $\theta$: zero gradient. The token has been pushed up enough for this batch.
• Advantage $-1$ (the token was bad): $\min(-1.3, -1.2) = -1.3$. The unclipped term wins, and its gradient pulls the probability back down. A step that made a bad token more likely is corrected.

And for $\rho = 0.7$, the mirror image: with advantage $+1$, $\min(0.7, 0.8) = 0.7$, unclipped, so the gradient pushes it back up; with advantage $-1$, $\min(-0.7, -0.8) = -0.8$, clipped, no more pushing down. The objective stops rewarding movement beyond $1 \pm \epsilon$ in the direction the advantage wants, and never stops corrections. That makes it safer to take several gradient steps on each batch of samples, which is what makes PPO more sample-efficient than REINFORCE. It’s a soft limit, not a hard one: other tokens’ gradients still move shared weights, so implementations also watch the KL to the old policy and end a batch’s passes early if it grows too large. The value head’s regression loss is added to the same objective.

An experiment: policy gradients on a tiny policy

The task is easy to state: a policy writes three tokens from a vocabulary of 8 (V), one at a time, and earns one point for each that matches a hidden target (TARGET). Every reward also carries +10, because real rewards are rarely centered at zero, and that offset is what the first learner has to cope with. It is the structure of a language model’s RL problem in miniature: tokens chosen autoregressively, one reward at the end, no gradient through it. Five learners share the same samples, 16 sequences per update: REINFORCE, REINFORCE with the +10 removed, REINFORCE with a baseline, the baseline with four passes per batch, and PPO; the last two spend four optimizer steps on each batch instead of one. The PPO here is simplified: one ratio per sequence rather than per token, and the batch-mean baseline instead of a value head.

def run(method, seed, updates=300, lr=0.1):
    g = torch.Generator().manual_seed(seed)
    logits = torch.zeros(1 + V + V * V, V, requires_grad=True)
    opt = torch.optim.Adam([logits], lr=lr)
    curve = []
    global clipped
    for _ in range(updates):
        seqs = sample(logits.detach(), g)
        R = torch.tensor([10.0 + sum(a == b for a, b in zip(s, TARGET)) for s in seqs])
        curve.append(R.mean().item() - 10)
        adv = {"REINFORCE": R, "REINFORCE, no offset": R - 10}.get(method, R - R.mean())
        old = logprob(logits, seqs).detach()
        for epoch in range(1 if method in ("REINFORCE", "REINFORCE, no offset", "REINFORCE + baseline") else 4):
            lp = logprob(logits, seqs)
            if method == "PPO":
                ratio = (lp - old).exp()
                if epoch > 0:
                    clipped.append(((ratio < 0.8) | (ratio > 1.2)).float().mean().item())
                loss = -torch.min(ratio * adv, ratio.clamp(0.8, 1.2) * adv).mean()
            else:
                loss = -(adv * lp).mean()
            opt.zero_grad(); loss.backward(); opt.step()
    return np.array(curve)

That is one learner’s training loop. The policy is a table of logits, one row per prefix, and the rest of the script samples from it, runs each learner for 10 seeds at two learning rates, and at the end compares one batch’s gradient with the exact one, computed over all 512 possible sequences.

Full script: tiny_policy.py

Adam learning rate 0.1: mean matches per sequence (of 3) over 10 seeds, at update
                            10     50    100    200    300   seeds solved (>= 2.9 at the end)
REINFORCE                 0.37   0.38   0.53   1.07   1.51   0 of 10
REINFORCE, no offset      0.46   1.76   2.48   2.64   2.78   8 of 10
REINFORCE + baseline      0.47   2.28   2.97   2.99   2.99   10 of 10
baseline, 4 passes        1.11   2.63   2.70   2.70   2.70   7 of 10
PPO                       0.79   2.69   2.69   2.80   2.90   9 of 10
PPO's clip was active for 3% of sequences in passes 2 to 4
Adam learning rate 0.03: mean matches per sequence (of 3) over 10 seeds, at update
                            10     50    100    200    300   seeds solved (>= 2.9 at the end)
REINFORCE                 0.36   0.37   0.41   0.49   0.66   0 of 10
REINFORCE, no offset      0.38   0.78   1.50   2.63   2.89   5 of 10
REINFORCE + baseline      0.38   0.86   1.92   2.91   2.95   10 of 10
baseline, 4 passes        0.51   2.69   2.97   2.99   2.99   10 of 10
PPO                       0.48   2.57   2.96   2.99   2.99   10 of 10
PPO's clip was active for 2% of sequences in passes 2 to 4
one batch's gradient vs the exact one, rewards 10 to 13    : cosine 0.02
one batch's gradient vs the exact one, rewards 0 to 3      : cosine 0.40
one batch's gradient vs the exact one, minus the batch mean: cosine 0.47

• REINFORCE barely learns with shifted rewards. With rewards of 10 to 13, no seed out of ten solves the task in 300 updates at either learning rate. Every sample is pushed up, and the useful signal, the difference between 10 and 13, is a small part of each step. Adam rescaling the step doesn’t help: it normalizes the gradient’s size, not its direction. At the start of training one batch’s gradient points almost at random, cosine 0.02 with the exact gradient (all 512 sequences enumerated), against 0.40 without the offset and 0.47 with a baseline.
• The offset is the cause. The counterfactual removes the +10 and changes nothing else: 8 of 10 seeds solve it at the higher learning rate, 5 of 10 at the lower. Rewards of 0 to 3 are still never negative, so the wrong tokens inside a partly right answer are pushed up along with the right ones.
• A baseline fixes it for free. Subtracting the batch’s mean reward solves all 10 seeds at both learning rates. It costs one subtraction. Because the mean includes each sample’s own reward, it is not quite the response-independent baseline derived above: in expectation it scales the gradient by $(N-1)/N$, 15/16 for a batch of 16, which changes its size, not its direction.
• More steps per batch learn faster per sample, and the clip barely matters here. PPO’s four clipped passes over each batch reach 2.6 matches by update 50, where single-step REINFORCE with a baseline has 0.9 at the lower learning rate. But four unclipped passes of REINFORCE with a baseline, the control, reach 2.7: the gain is the extra optimizer steps. In a problem this small the clip binds on 2 to 3% of sequences. At the higher learning rate the unclipped control got stuck on 3 seeds and PPO on 1, a hint that the clip helps when steps are large, which ten seeds can’t settle. Clipping earns its keep on large models and large steps, where repeated passes can carry a policy far from the data that sampled it.
• A policy can lock onto a wrong answer. At the higher learning rate, one PPO seed of ten had not solved the task by the end. A policy that puts nearly all its probability on a wrong token almost never samples anything better, so nothing in its batches pushes it off. That is the exploration problem: a policy only learns about what it samples. Real systems keep the policy from becoming deterministic too early, with an entropy bonus, a KL penalty to the starting model (see RLHF), or simply smaller steps, as the lower learning rate shows here.

None of this needs a language model to see, and all of it carries over: RLHF’s PPO uses a learned value as the baseline, and GRPO (see GRPO and Reinforcement Learning from Verifiable Rewards) uses the same idea as the batch-mean baseline above: the mean reward of several responses to the same prompt, with the difference also divided by those rewards’ standard deviation.

Why PPO for a language model needs four models

PPO as used in RLHF (see RLHF) keeps four networks in memory:

• the policy, being trained;
• the reference policy, a frozen copy of the starting model, for the KL penalty that keeps the policy from drifting too far (see RLHF);
• the reward model, frozen, which scores each response;
• the value model, being trained, usually initialized from the reward model and the same size.

In RLHF the KL penalty is applied per token, as a reward of $-\beta(\log \pi_\theta - \log \pi_\text{ref})$ on each token, with the reward model’s score added at the last one, which is why per-token advantages, and so the value model, matter.

For a 7B policy and a 7B reward model, the two trained ones cost 16 bytes per parameter in mixed-precision AdamW (Initialization and Optimizers), 112 GB each, and the two frozen ones 2 bytes per parameter, 14 GB each: about 250 GB before activations, against 112 GB for supervised fine-tuning of the same model. Generation is also part of every step, and is slow (see Decoding). That cost is what DPO (see DPO) and GRPO (see GRPO and Reinforcement Learning from Verifiable Rewards) attack: DPO drops the reward and value models, GRPO drops the value model.

Where this leads

• RLHF builds RLHF from this: a reward model learned from human comparisons, and PPO with a KL leash.
• DPO removes the RL loop entirely for preference data; GRPO and Reinforcement Learning from Verifiable Rewards removes the value model by scoring each response against others to the same prompt, the setting used for checkable rewards.

Takeaways

• When the score isn’t differentiable, the log-derivative trick turns it into a gradient: $\mathbb{E}[R\, \nabla \log \pi]$, maximum-likelihood training on the model’s own samples, weighted by reward.
• A baseline subtracted from the reward adds no bias and removes the noise an offset in the rewards adds (a batch mean that includes the sample’s own reward only scales the gradient by $(N-1)/N$). On a tiny autoregressive policy, REINFORCE with rewards of 10 to 13 solved 0 of 10 seeds, and with a mean baseline 10 of 10.
• The advantage, reward minus what was expected, is what every method here trains on; GAE estimates it per token from a learned value.
• PPO’s ratio is importance sampling, and its clip stops rewarding movement past $1 \pm \epsilon$ while allowing corrections, so several steps per batch are safer. Here the extra steps, not the clip, made it fast: four unclipped passes did as well. A policy can still lock onto a wrong answer it no longer explores past.
• PPO on a language model holds four models, two of them trained: about 250 GB for a 7B policy and reward model, which is what DPO and GRPO cut.

References

Start with OpenAI’s Spinning Up in Deep RL for the policy gradient derived step by step, then the PPO paper, which is short. Sutton and Barto’s chapter 13 is the textbook treatment.

• Sutton and Barto, Reinforcement Learning: An Introduction (2nd ed., MIT Press, 2018): the textbook treatment, with policy gradients in chapter 13 and the agent and environment vocabulary.
• Williams, Simple Statistical Gradient-Following Algorithms for Connectionist Reinforcement Learning (Machine Learning, 1992): REINFORCE.
• Schulman, Levine, Moritz, Jordan and Abbeel, Trust Region Policy Optimization (ICML, 2015): TRPO, which bounds the KL divergence from the old policy at each update.
• Schulman, Moritz, Levine, Jordan and Abbeel, High-Dimensional Continuous Control Using Generalized Advantage Estimation (ICLR, 2016): GAE, the per-token advantage estimate.
• Schulman, Wolski, Dhariwal, Radford and Klimov, Proximal Policy Optimization Algorithms (2017): PPO’s clipped objective, and its KL-penalty variant.