RLHF with PPO

RLHF with PPO is the classic recipe for teaching a language model to produce responses people prefer: sample responses from the current model, score them with a reward model trained on human comparisons, and update the model with Proximal Policy Optimization while a KL penalty keeps it close to where it started. It is the third stage of the pipeline that produced InstructGPT: SFT first, a reward model second, PPO third. Most of what came after, from DPO to RLOO, ReMax and GRPO, is defined by which parts of this recipe it removes. Knowing what each part does tells you which simplification you can afford.

Where it came from#

PPO itself is a general RL algorithm, published by Schulman, Wolski, Dhariwal, Radford and Klimov in 2017 for robotics and Atari. Its contribution was a cheap way to take several gradient steps on the same batch of experience without the policy running away: clip the probability ratio between new and old policy so the objective stops rewarding changes beyond a small trust region.

The language-model application came from OpenAI's alignment team. Ziegler et al. (2019) fine-tuned GPT-2 with PPO against a reward model learned from human comparisons, and introduced the KL-shaped reward that every RLHF implementation still uses. Stiennon et al. (2020) scaled it to summarisation and showed policies trained this way beat larger supervised models in human evaluation. Ouyang et al. (2022) applied it to instruction following, where a 1.3B InstructGPT model was preferred by labelers over the 175B GPT-3 base. Anthropic's helpful-and-harmless assistant (Bai et al., 2022) and Llama 2 (Touvron et al., 2023) followed with their own PPO-based pipelines. The problem it solves is that "good response" has no closed-form definition. Supervised learning can only imitate demonstrations; RL against a learned reward lets the model explore and be graded on outputs nobody wrote.

The objective#

The KL-regularised RL problem

Let ${\pi }_{\theta }$ be the policy (the model being trained), ${\pi }_{\text{ref}}$ the frozen SFT model, $r_{\varphi }(x,y)$ the reward model's score of response $y$ to prompt $x$, and $\beta >0$ the KL coefficient. RLHF maximises

The KL term does two jobs. It stops the policy from drifting into text the reward model has never seen, where its scores are unreliable, and it preserves the fluency and knowledge of the SFT model. This objective has a known optimum, which DPO exploits and which we use to check our toy run below:

PPO does not compute ${\pi }^{\ast }$ (the sum over all responses is intractable); it climbs toward it by sampling.

Responses as token-level episodes

PPO treats generation as an episode. The state $s_t$ is the prompt plus the first $t$ tokens, the action $a_t$ is the next token, and the response ends at the end-of-sequence token at step $T$. The KL penalty is spread over tokens using the per-token log-ratio, and the reward model's score arrives only at the last token:

Summed over the episode, the KL part is a single-sample estimate of the sequence-level KL. This shaping is what Ziegler et al. introduced, and it is why PPO for RLHF needs a value function: a reward that appears only at the end has to be credited back to the tokens that earned it.

Advantages with GAE

A value model $V_{\psi }(s_t)$ predicts the expected remaining reward from each position. The temporal-difference error and the generalised advantage estimate (Schulman et al., 2015) are

with $V_{\psi }(s_{T+1})=0$. $\gamma$ is the discount (1.0 in almost every RLHF setup, since responses are short and finite) and $\lambda \in [0,1]$ trades bias for variance: $\lambda =1$ gives the Monte Carlo return minus the baseline, $\lambda =0$ trusts the value model completely. The value target is the return ${\hat{R}}_t={\hat{A}}_t+V_{\psi }(s_t)$.

The clipped surrogate

With ${\rho }_t(\theta )={\pi }_{\theta }(a_t\mid s_t)/{\pi }_{{\theta }_{\text{old}}}(a_t\mid s_t)$ the probability ratio against the policy that generated the data, PPO maximises

and the value model minimises a (usually also clipped) squared error to ${\hat{R}}_t$. The minimum makes the objective pessimistic. When an action has positive advantage and its probability has already risen more than $\epsilon$, the clipped term wins and its gradient is zero. When the change goes the wrong way, the unclipped term wins and the gradient pushes back at full strength. The policy gradient in the unclipped region is the familiar ${\hat{A}}_t{\rho }_t{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)$.

How the siblings simplify it

Worked example by hand#

Take a four-token response with $\beta =0.1$, $\gamma =1$, $\lambda =0.95$. The per-token log-ratios against the reference are 0.10, -0.05, 0.20 and 0.00, the reward model scores the full response at 1.0, and the value model predicts 0.5, 0.6, 0.4, 0.8. The shaped rewards are -0.010, +0.005, -0.020 and 1.000 (the RM score lands on the EOS token). TD errors work backwards from the end: ${\delta }_3=1.0+0-0.8=0.200$; ${\delta }_2=-0.02+0.8-0.4=0.380$; ${\delta }_1=0.005+0.4-0.6=-0.195$; ${\delta }_0=-0.01+0.6-0.5=0.090$. Accumulating with $\gamma \lambda =0.95$ gives advantages 0.419, 0.347, 0.570, 0.200. Token 2 gets the largest credit because the value model under-predicted what followed it; token 1 gets less than its neighbours because its negative TD error says the value model was over-optimistic there.

Now the clip, with $\epsilon =0.2$ and $\hat{A}=2$. If an update has pushed the token's ratio to 1.3, the unclipped term is 2.6, the clipped term 2.4, and the minimum is 2.4: a constant in $\theta$, so no gradient. That token has moved enough. If the ratio fell to 0.7 instead, the unclipped term 1.4 is smaller than the clipped 1.6, so the gradient flows and pulls the probability back up. With $\hat{A}=-2$ the cases mirror: a ratio of 1.3 keeps its gradient (objective -2.6), a ratio of 0.7 is clipped (objective -1.6, no gradient). Clipping only ever removes the incentive to move further in the direction the advantage already favours.

The algorithm#

inputs: SFT model π_sft, reward model r_φ, prompt set D
π_θ ← π_sft;  π_ref ← π_sft (frozen);  V_ψ ← r_φ backbone + new scalar head
repeat for N iterations:
    # rollout phase (no gradients)
    sample B prompts x from D
    y ~ π_θ(·|x) with temperature 1.0, up to L_max tokens
    logp_old  = log π_θ(y_t | s_t)          per token
    logp_ref  = log π_ref(y_t | s_t)        per token
    values    = V_ψ(s_t)                    per token
    score     = r_φ(x, y)                   one per sequence (penalise if no EOS)
    r_t = -β (logp_old - logp_ref);  r_T += score
    A_t, R_t = GAE(r_t, values, γ=1, λ=0.95);  whiten A over the batch
    # update phase
    for epoch in 1..E (E = 1 to 4):
        for minibatch in shuffle(rollouts):
            ρ = exp(log π_θ(y_t|s_t) - logp_old)
            L_pi = -mean(min(ρ A, clip(ρ, 1-ε, 1+ε) A))
            L_v  = 0.5 mean(max((V_ψ - R)^2, (clip(V_ψ, V_old ± ε_v) - R)^2))
            step on L_pi + c_v L_v
    optionally adapt β toward a KL target

Data. PPO needs prompts only; responses are generated online. Typical prompt sets run from 10k to a few hundred thousand, drawn from the same distribution the reward model was trained on. InstructGPT's PPO set was on the order of tens of thousands of prompts. A run of 1,000 iterations at 512 prompts each sees about half a million rollouts, so prompts are reused across epochs of the outer loop.

Implementation details that matter. Huang et al. (2024) reproduced the OpenAI summarisation results and catalogued the details that decide whether PPO works: score with the reward model at the EOS token and give a fixed penalty to responses that never emit EOS; normalise the reward model's outputs before RL so their scale is known; initialise the value model from the reward model; whiten advantages per batch; and use separate networks for policy and value. Zheng et al. (2023) reached similar conclusions in "Secrets of RLHF", singling out the token-level KL penalty and reward normalisation as the constraints that keep training stable.

Hyperparameters that matter#

Ziegler et al. used an adaptive controller instead of a fixed $\beta$: after each batch, compute the error $e=\mathrm{clip}((\mathrm{KL}-{\mathrm{KL}}_{\text{target}})/{\mathrm{KL}}_{\text{target}},-0.2,0.2)$ and set $\beta \leftarrow \beta (1+K_{\beta }e)$ with a small gain $K_{\beta }$ such as 0.1. This holds KL roughly constant, which is often easier to reason about than a fixed coefficient whose effect depends on the reward model's scale. Tune in this order: fix $\beta$ and the learning rate on a short run while watching KL and clip fraction, then lengthen the run, then adjust batch size if curves are noisy.

Compute and memory#

PPO keeps four networks: the policy and value model are trained, the reference and reward model are inference-only. For full fine-tuning with Adam in mixed precision, a trained model costs about 16 bytes per parameter and a frozen bf16 model 2 bytes. With a policy of $P$ parameters and RM and value models of $P_r$ parameters each:

For a 7B policy with a 7B value and reward model that is about $36\times 7=252$ GB of model state before activations, which is why PPO at 7B usually means a multi-GPU node with ZeRO-3 sharding, and at 70B means multiple nodes. LoRA on the policy and value model cuts the trained-model term sharply, and some pipelines share a backbone between policy and value with two heads at the cost of interference between the losses.

Time is dominated by generation. Autoregressive sampling of long responses runs at a small fraction of training throughput, so in many setups the rollout phase takes the majority of wall-clock time. Modern stacks put generation on a dedicated inference engine (for example vLLM) and sync weights after each update. Relative to SFT on the same number of response tokens, a PPO run typically costs 5× to 20× more compute; relative to DPO on a comparable number of pairs, several times more, plus the reward model you had to train first.

A toy run: PPO against its closed-form target#

Because the KL-regularised objective has a known optimum, we can check what PPO actually converges to. We built a single-step "language model" over 50 possible responses with reference logits drawn from $\mathcal{N}(0,1)$ (reference entropy 3.55 nats), a true reward per response drawn from $\mathcal{N}(0,1)$, and a proxy reward model equal to the truth plus Gaussian error whose standard deviation grows from 0.25 on common responses to 1.75 on responses the reference rarely produces, mimicking an RM that is unreliable off-distribution (correlation with truth 0.80). PPO used the shaped reward $r-\beta \log (\pi /{\pi }_{\text{ref}})$, a moving-average baseline, whitened advantages, clip 0.2, 4 epochs per batch of 64 samples, learning rate 0.3 on the logits, 400 iterations, seed 1.

For $\beta \ge 0.2$, PPO lands within about 0.13 nats of the analytic optimum and within 0.08 of its proxy reward: the algorithm does solve the objective it is given. Two things in the table are worth more than that confirmation.

First, look at the gap between proxy and true reward. At $\beta =2$ it is 0.25; at $\beta =0.2$ it is 1.50. Going from 1.7 to 4.4 nats of KL buys 0.85 units of proxy reward and only 0.13 units of true reward. Most of the extra optimisation is spent exploiting the reward model's errors on rare responses. The chart shows the same thing along the whole frontier: proxy reward keeps climbing with KL while true reward flattens.

Second, at $\beta \le 0.1$ PPO does not reach the optimum. Within 80 iterations it collapsed onto a single response with proxy reward 2.2 and stayed there for the remaining 320; the optimum puts its mass on a different, rarer response with proxy reward 3.2. Once the policy is nearly deterministic, it stops sampling alternatives and has no way to discover them. That is entropy collapse, and in this toy the collapsed policy happened to have higher true reward by luck of the draw, which is a reminder that a run's final number can mislead in either direction. With a weak KL penalty, the policy you get depends on what it found first.

Failure modes and how to detect them#

Reward hacking. RM score rises smoothly while held-out win rate (human or LLM judge) stalls or falls. Responses grow longer, adopt a fixed template, or repeat flattering phrases. Plot judge win rate against KL, not steps, and stop where it flattens. Raise $\beta$, or retrain the RM on samples from the current policy.

KL blow-up. KL jumps by several nats in a few iterations, often with the clip fraction (share of tokens where clipping is active) above 0.3 and the approximate KL between old and new policy per update above 0.05. Cause: learning rate too high, too many PPO epochs, or a sampling mismatch (temperature, top-p or a different tokenizer path between generator and trainer, which makes $\rho$ wrong from the first step). Check that $\rho$ averages 1.0 on the first minibatch of each iteration; if it doesn't, generator and trainer disagree.

Negative KL estimates. The per-token log-ratio estimator can go negative in expectation if responses are truncated or sampled with top-p, which the policy then exploits by producing tokens with low policy probability. Sample with pure temperature 1.0 during training.

Value model lag. Explained variance of the value predictions stays near zero or negative, and advantages look like noise. Increase the value learning rate, initialise the critic from the RM, or warm it up for a few iterations before updating the policy.

Entropy collapse. Per-token entropy drops steadily and samples for the same prompt become near-identical. Exploration stops and the run plateaus, as in the toy above. Watch entropy and the number of distinct responses per prompt; raise $\beta$ or lower the learning rate.

Length exploitation and truncation. Mean response length climbs until it hits the maximum and many responses end without EOS. Apply the missing-EOS penalty and track the fraction of truncated responses.

Capability regression. Benchmarks the SFT model was good at (math, code, factual QA) drop. InstructGPT's PPO-ptx variant mixes a pretraining language-modelling loss into the PPO update to counter this; a cheaper check is a fixed regression suite run every few hundred steps.

When to use PPO, and what to use instead#

Minimal implementation sketch#

The core of a PPO update in PyTorch, given rollouts already generated and scored. TRL and OpenRLHF provide complete trainers with generation, sharding and logging; the loss below is what they compute.

import torch

def gae(rewards, values, mask, gamma=1.0, lam=0.95):
    # response-only tensors [B, T], right-padded; mask is 1 on real tokens
    B, T = rewards.shape
    adv = torch.zeros_like(rewards); last = torch.zeros(B, device=rewards.device)
    for t in reversed(range(T)):
        alive = mask[:, t + 1] if t + 1 < T else torch.zeros(B, device=rewards.device)
        next_v = values[:, t + 1] * alive if t + 1 < T else alive
        delta = rewards[:, t] + gamma * next_v - values[:, t]
        last = (delta + gamma * lam * alive * last) * mask[:, t]
        adv[:, t] = last
    returns = adv + values
    m = mask.bool()
    adv = (adv - adv[m].mean()) / (adv[m].std() + 1e-8)        # whiten
    return adv * mask, returns

def shaped_rewards(logp_old, logp_ref, score, mask, beta=0.05):
    r = -beta * (logp_old - logp_ref) * mask
    last = mask.sum(1).long() - 1                               # index of final response token
    r[torch.arange(r.size(0)), last] += score
    return r

def ppo_loss(logp_new, logp_old, adv, v_new, v_old, returns, mask, eps=0.2, eps_v=0.2):
    ratio = torch.exp(logp_new - logp_old)
    pg = -torch.min(ratio * adv, torch.clamp(ratio, 1 - eps, 1 + eps) * adv)
    v_clip = v_old + torch.clamp(v_new - v_old, -eps_v, eps_v)
    vf = 0.5 * torch.max((v_new - returns) ** 2, (v_clip - returns) ** 2)
    n = mask.sum()
    stats = {"clipfrac": (((ratio - 1).abs() > eps).float() * mask).sum() / n,
             "approx_kl": (0.5 * (logp_new - logp_old) ** 2 * mask).sum() / n}
    return (pg * mask).sum() / n, (vf * mask).sum() / n, stats

Log at least: mean RM score, KL to reference per sequence, clip fraction, approximate KL per update, value explained variance, entropy, mean response length and the truncation rate. Those eight numbers diagnose nearly every failure in the previous section. In POFE these are the default panels on an RL run, next to a judge win-rate curve plotted against KL.

Further reading#

Schulman et al. (2017) for PPO and Schulman et al. (2015) for GAE. Ziegler et al. (2019) for the KL-shaped reward and adaptive KL controller. Stiennon et al. (2020) and Ouyang et al. (2022) for the full RLHF pipeline at scale. Huang et al. (2024) and Zheng et al. (2023) for the implementation details that separate working runs from broken ones. Ahmadian et al. (2024) for the argument that PPO's machinery is more than LLM fine-tuning needs, which leads to RLOO.

Frequently asked questions#

References#

1. John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, Oleg Klimov (2017). Proximal Policy Optimization Algorithms. arxiv.org/abs/1707.06347
2. John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, Pieter Abbeel (2015). High-Dimensional Continuous Control Using Generalized Advantage Estimation. arxiv.org/abs/1506.02438
3. Daniel M. Ziegler, Nisan Stiennon, Jeffrey Wu, Tom B. Brown, et al. (2019). Fine-Tuning Language Models from Human Preferences. arxiv.org/abs/1909.08593
4. Nisan Stiennon, Long Ouyang, Jeff Wu, Daniel M. Ziegler, et al. (2020). Learning to summarize from human feedback. arxiv.org/abs/2009.01325
5. Long Ouyang, Jeff Wu, Xu Jiang, Diogo Almeida, et al. (2022). Training language models to follow instructions with human feedback. arxiv.org/abs/2203.02155
6. Yuntao Bai, Andy Jones, Kamal Ndousse, Amanda Askell, et al. (2022). Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedback. arxiv.org/abs/2204.05862
7. Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, et al. (2023). Llama 2: Open Foundation and Fine-Tuned Chat Models. arxiv.org/abs/2307.09288
8. Shengyi Huang, Michael Noukhovitch, Arian Hosseini, Kashif Rasul, Weixun Wang, Lewis Tunstall (2024). The N+ Implementation Details of RLHF with PPO: A Case Study on TL;DR Summarization. arxiv.org/abs/2403.17031
9. Rui Zheng, Shihan Dou, Songyang Gao, Yuan Hua, et al. (2023). Secrets of RLHF in Large Language Models Part I: PPO. arxiv.org/abs/2307.04964
10. Arash Ahmadian, Chris Cremer, Matthias Gallé, Marzieh Fadaee, et al. (2024). Back to Basics: Revisiting REINFORCE Style Optimization for Learning from Human Feedback in LLMs. arxiv.org/abs/2402.14740

Related reading