Understanding PPO for Language Models

A rollout is one prompt plus a completion sampled from the language model. Proximal Policy Optimization (PPO) scores that completed response, assigns each sampled token a training signal, and optimizes a policy and a critic: the policy changes token probabilities, while the critic predicts expected future reward.

1.The whole loop, before the equations

PPO alternates between collecting experience and learning from that fixed experience. For a language model, one experience is a prompt plus a sampled completion.

The critic and policy solve different problems:

The policy does not directly maximize the critic's output. It increases the probability of sampled tokens with positive estimated advantage and decreases the probability of tokens with negative estimated advantage. Observed rewards ground that estimate, while intermediate critic predictions may also shape it.

2.Language generation is a trajectory

Start with a prompt. Before generating token \(t\), the model has a context \(x_t\): the prompt plus all response tokens already generated. The next token \(y_t\) is sampled from the policy:

\(\pi\) means policy; for a language model it is the next-token sampling distribution. \(\theta_{\mathrm{old}}\) denotes the frozen policy parameters used to collect this rollout; later, \(\theta\) will denote the current policy parameters being optimized. The critic has its own parameters \(\phi\). After the token is chosen, the environment produces a reward \(r_t\). The response ends after \(T\) optimized tokens.

The return \(G_t\) is the sum of rewards from position \(t\) onward. The discount factor \(\gamma\in[0,1]\) controls how much less a later reward counts:

The true state value is an expectation over all possible future tokens and environment outcomes under the rollout policy:

The critic \(V_\phi(x_t)\) is a learned approximation to that expectation. A forward pass gives a definite scalar, but the quantity remains an estimate—not an observed fact.

3.The rollout happens first

The credit calculation will use the critic's value at the next token position and evidence from later positions. It does not predict an unknown future during generation. PPO first samples the complete response and stores it. Only then does optimization begin.

4.One terminal-reward math rollout

Take the prompt “Solve \(2x+3=11\)”. To keep the arithmetic visible, use four toy token positions whose labels summarize the generated fragments. A real tokenizer may split the same text into more positions; PPO applies the identical calculation to every actual token.

The answer is correct, so the verifier emits \(+1\) at the final position and zero task reward before it. Let \(\gamma=0.9\). During the rollout, the critic made these predictions:

Each row pairs the sampled fragment \(y_t\) with the value of its pre-token context \(x_t\). Thus \(V(x_2)=0.35\) sees the first fragment, but not the second. The values rise as the sampled solution becomes more promising, but a rise is not itself the learning target. We need a consistent way to compare each stored prediction with what was learned one step later.

5.The TD residual: one-step new evidence

After token \(y_t\) creates the next context \(x_{t+1}\), compare the old prediction with the immediate reward plus the discounted next prediction. Let \(d_t=1\) when this transition truly terminates the episode and \(d_t=0\) otherwise. Using the next critic prediction in a target is called bootstrapping:

\(\delta_t\) is the temporal-difference (TD) residual. Positive means this realized transition supplied better evidence than the critic expected; negative means worse. It is computed for the sampled transition in this rollout, not for every token the model could have chosen. The notation \(V_{\mathrm{old}}\) emphasizes that these are frozen rollout-time predictions.

At a true terminal boundary the multiplier \(1-d_t\) removes the bootstrap—equivalently, set \(V_{\mathrm{old}}(x_{T+1})=0\), because no return remains. A response cut off only by a length or rollout boundary is a truncation, not necessarily a terminal state; then \(d_t=0\) and the trainer bootstraps from \(V_{\mathrm{old}}(x_{T+1})\). For the truly terminal four-position example:

6.GAE: let later evidence reach earlier tokens

A one-step residual is local. Generalized Advantage Estimation (GAE) combines it with later residuals. The result \(\widehat A_t\) estimates how much better or worse the sampled continuation was than the stored baseline. A parameter \(\lambda\in[0,1]\) controls how much later evidence enters. The implementation is a reverse recursion:

For the math rollout, choose \(\lambda=0.8\), so \(\gamma\lambda=0.72\). Start at the end: \(\widehat A_4=\delta_4=0.200\). Then \(\widehat A_3=0.170+0.72(0.200)=0.314\). Continuing backward gives \(\widehat A_2=0.145+0.72(0.314)=0.371\) and \(\widehat A_1=0.115+0.72(0.371)=0.382\). The early positions are positive because later transitions exceeded the critic's stored predictions—not because the final reward was copied evenly to every token.

Within one uninterrupted episode, expanding the recursion shows exactly where a residual five positions later enters the current advantage:

The sum looks forward from \(t\) to the end. GAE is computed backward only as an efficient implementation. The current residual has weight 1; a residual \(\ell\) positions later has weight \((\gamma\lambda)^\ell\), which shrinks with distance when \(\gamma\lambda<1\).

Within an episode, \(\delta_k\) enters \(\widehat A_t\) for \(t\le k\) with coefficient \((\gamma\lambda)^{k-t}\); when \(\gamma\lambda=0\), it enters only at \(t=k\). That is the backward credit path. There is no gradient through the sampled token sequence: during the policy update, the computed advantages are fixed numbers.

RolesWhat \(\gamma\) and \(\lambda\) each do

Both appear in the propagation weight \((\gamma\lambda)^\ell\), but they are not interchangeable: \(\gamma\) belongs to the task's return definition. Lower \(\lambda\) leans more on nearby critic predictions; higher \(\lambda\) leans more on the sampled continuation.

7.The return target trains the critic

GAE produces an advantage relative to the stored baseline. Add that baseline back to form the critic's regression target:

The subscript “old” matters in implementation. Values are computed during the rollout, stored, and detached. \(\widehat R_t\) is then a fixed label while the current critic \(V_\phi\) is optimized. Let \(N\) be the number of valid response-token positions in the minibatch; \(\operatorname{stopgrad}\) means treat its argument as a constant during backpropagation; the superscript \(\mathrm{VF}\) labels the value-function loss:

The expectation symbol often used in papers is implemented as a mean over valid response-token positions in the minibatch. We minimize this mean-squared error.

The target has an especially useful recursive form:

Use the boundary \(\widehat R_{T+1}=V_{\mathrm{old}}(x_{T+1})\). At true termination that value is zero; at truncation it supplies the continuation bootstrap. At \(\lambda=0\), the target is the one-step TD target. At \(\lambda=1\), it is the sampled discounted return for a terminal rollout, or that return plus the discounted boundary bootstrap for a truncation. Values between 0 and 1 mix shorter critic bootstraps with longer sampled evidence.

8.The advantage trains the language model

Now freeze the advantages. For each recorded pair \((x_t,y_t)\), the basic policy-gradient term is:

• If \(\widehat A_t>0\), gradient ascent increases the probability of the sampled token in that context.
• If \(\widehat A_t<0\), it decreases that probability.
• The magnitude controls the strength of that sample's contribution, often after batchwise advantage normalization.

A downstream residual affects an earlier token because it is already included in that token's scalar \(\widehat A_t\). Standard PPO does not backpropagate through \(y_t\rightarrow x_{t+1}\), through the reward, or through the critic while updating the policy:

advantages = advantages.detach()
policy_loss = -(logprob_of_sampled_token * advantages).mean()

This separation is essential. The sampled actions and future contexts are data. The current token log-probability is the differentiable path into the policy parameters.

9.PPO: limit the incentive when reusing a rollout

At the start of a synchronous PPO update, \(\pi_\theta\) matches the rollout policy. PPO may then take one or more minibatch updates—sometimes across several epochs—on the same rollout. After the first optimizer step, the current policy has changed, so it forms the probability ratio:

The denominator is a frozen log-probability stored with the rollout. The numerator comes from the model being optimized. Therefore the actor gradient enters through the current log-probability—and thus through \(\rho_t\)—while \(\widehat A_t\), \(x_t\), \(y_t\), and the denominator stay fixed.

For one concrete update, suppose a sampled token had probability \(0.20\) under the rollout policy and \(0.26\) now. Then \(\rho=0.26/0.20=1.30\). If its advantage is \(+1\) and the clip threshold is \(\epsilon=0.20\), PPO compares the raw term \(1.30\) with the clipped term \(1.20\), then selects the smaller \(1.20\). This sample gains nothing from pushing the ratio higher.

In general, \(\epsilon\in(0,1)\) sets the clip band \([1-\epsilon,1+\epsilon]\). Without clipping, repeatedly maximizing \(\rho_t\widehat A_t\) can move the policy too far from the rollout policy, so PPO uses:

Theoretical notation maximizes \(J^{\mathrm{CLIP}}\). PyTorch optimizers minimize, so code uses \(L^{\mathrm{policy}}=-J^{\mathrm{CLIP}}\). The inner min is part of the pessimistic objective; it does not mean the outer training goal is minimization.

PiecewiseWhat clipping does for each sign

For a positive advantage, the sample stops benefiting once its probability ratio exceeds \(1+\epsilon\). For a negative advantage, it stops benefiting once the ratio falls below \(1-\epsilon\). The objective still follows harmful movement in the opposite direction; clipping removes incentive only where the sample would otherwise improve too much.

10.PPO update ratio versus reference-model penalty

Reinforcement learning from human feedback (RLHF) implementations commonly keep a frozen reference policy initialized from the supervised fine-tuning (SFT) policy. This introduces a second comparison, which must not be confused with PPO clipping. The reference term estimates Kullback–Leibler (KL) divergence, a measure of distribution shift.

Let \(\beta\ge0\) control the strength of the reference penalty, and let \(R_{\mathrm{score}}\) be the scalar reward-model or verifier score for the completed response. A common token reward is:

The scalar reward-model or verifier score is added at the final token; the sampled log-ratio penalty is dense across the response. Individual sampled log-ratios can have either sign—their expectation under the rollout policy is the nonnegative KL divergence. Process reward models or environment feedback may add other intermediate rewards.

11.The training loop in implementation order

This compact skeleton puts the core quantities in execution order. Production code typically handles response masks and may also use advantage normalization, value clipping, entropy terms, and separate actor/critic optimizers. In the combined-loss skeleton below, value_coef weights the value loss relative to the policy loss:

# 1. Freeze a behavior snapshot and collect complete responses.
old_policy = snapshot(policy)
x, y, old_logprob, old_value, reward, terminated = rollout(old_policy, critic)

# 2. Post-rollout: reverse scan through valid response tokens.
delta[t] = reward[t] + gamma * (1-terminated[t]) * old_value[t + 1] - old_value[t]
adv[t]   = delta[t] + gamma * gae_lambda * (1-terminated[t]) * adv[t + 1]
target[t] = adv[t] + old_value[t]

# 3. Reuse this fixed batch for the configured minibatch updates.
ratio = exp(policy.logprob(x, y) - old_logprob)
policy_objective = min(ratio * stopgrad(adv),
                       clip(ratio, 1-eps, 1+eps) * stopgrad(adv))
value_loss = (critic.value(x) - stopgrad(target)) ** 2

minimize(-mean(policy_objective) + value_coef * mean(value_loss))

The policy and critic may be separate language models, or they may share a transformer trunk with a token-logit head and a scalar value head. Detaching targets prevents gradients through target construction. If the heads share a trunk, however, both losses still update those shared parameters unless the implementation explicitly routes or blocks one gradient.

After the configured minibatch updates, discard the rollout buffer, refresh \(\pi_{\mathrm{old}}\) from the updated policy, and collect new on-policy data. In asynchronous systems, policy staleness makes this synchronization boundary especially important.

12.What modern run settings look like

There is no single “modern PPO configuration.” Sequence length, reward density, model size, rollout parallelism, and critic quality change the useful settings. These are named examples, not defaults to copy blindly:

The OpenRLHF learning rates and batch sizes above are overrides from its published Llama-3 PPO example script, not library defaults. CompactionRL's value-to-policy entry is an update ratio, not a number of PPO epochs.

Three patterns are more transferable than the numbers:

1. The critic often needs extra care. It may use a higher learning rate, warmup, extra updates, value clipping, or a separate \(\lambda\) for return targets.
2. Outcome-based LLM tasks often set \(\gamma=1\). When correctness is the only terminal reward, \(\gamma=1\) avoids implicitly discounting correct outcomes that arrive after more tokens; length costs can be added explicitly.
3. Normalize and monitor. Advantage whitening stabilizes gradient scale. Useful diagnostics include reward, reference KL, approximate old/new KL, clip fraction, value loss, value explained variance, response length, and entropy.

13.Limits, alternatives, and the compact mental model

1. The critic is expensive. A separate value model can approach the policy's memory and compute footprint. Group Relative Policy Optimization (GRPO), REINFORCE Leave-One-Out (RLOO), and REINFORCE++ avoid a learned critic, using group or leave-one-out baselines, or globally normalized sampled returns, instead.
2. GAE assigns statistical credit, not causal proof. An earlier sampled token can receive signal from downstream residuals, even though the rollout does not reveal the counterfactual outcome of choosing another token.
3. Clipping is a heuristic trust-region proxy. It improves stability but does not impose a hard KL constraint or prevent every destructive update.
4. The objective is only as good as the reward. Reward-model errors, verifier loopholes, and overly strong KL or length shaping can produce the wrong behavior efficiently.
5. Stale rollouts increase off-policy mismatch. Too many epochs or asynchronous lag can produce extreme ratios and push training away from the near-on-policy regime PPO expects.

References

• Schulman et al., Proximal Policy Optimization Algorithms (2017) — clipped surrogate objective.
• Schulman et al., High-Dimensional Continuous Control Using Generalized Advantage Estimation (2015) — TD residuals and GAE.
• Ouyang et al., Training Language Models to Follow Instructions with Human Feedback (2022) — InstructGPT's PPO + reference-KL setup.
• NVIDIA, An In-Depth Walkthrough of PPO in NeMo RL — current LLM PPO implementation and configuration.
• OpenRLHF, train_ppo_ray.py — current open-source PPO arguments and defaults.
• Li et al., CompactionRL (2026) — long-horizon PPO with value pretraining and length-adaptive GAE.
• Huang et al., The N Implementation Details of RLHF with PPO — value clipping, whitening, and reproducibility details.
• OpenAI, Spinning Up: PPO — concise treatment of the piecewise clipping behavior.

How to cite this post

Dong, S. (2026). Understanding PPO for Language Models.
https://simondong1.github.io/ppo.html

@misc{dong2026understandingppo,
  author = {Dong, Simon},
  title  = {Understanding PPO for Language Models},
  year   = {2026},
  month  = {July},
  url    = {https://simondong1.github.io/ppo.html},
  note   = {Technical blog post, RL for LLMs series, Part 1}
}