10.2 PPO for Language-Model Alignment

Proximal Policy Optimization (PPO) is an online reinforcement-learning method used in classic RLHF systems [1] [2]. In language-model PPO, the policy generates responses, a reward model scores them, a value head estimates return, and a frozen reference policy constrains drift. Leaving out any of these pieces turns the example into another algorithm.

The Four-Model View

A practical RLHF PPO run distinguishes:

policy: initialized from an SFT checkpoint and updated;
reference policy: frozen snapshot used for token-level KL control;
reward model (RM): converts a prompt/response into terminal or shaped reward;
value head/critic: predicts return for each generated response position.

Some implementations share a backbone for policy and value or offload the reference/RM. Record revisions, tokenizer/chat templates, precision, quantization, and which parameters are shared. The evaluated artifact includes decoding, tools, and reward preprocessing—not weights alone.

Rollout and Reward Shaping

Render prompts with the tokenizer-native chat template and add_generation_prompt=True. Generate variable-length responses, preserve EOS, and build a response mask that is one only for generated tokens before padding. Prompt, padding, and tokens after terminal EOS do not receive policy/value loss.

For response token $t$ , a common shaped reward is

r_t = -\beta\left(\log\pi_{\text{old}}(a_t\mid s_t)- \log\pi_{\text{ref}}(a_t\mid s_t)\right)

plus the RM score at the final valid response token. This per-token KL sample penalty is not identical to every analytic KL estimator; document the implementation. Log raw RM reward, KL penalty, and final shaped reward separately so reward hacking is visible.

The RM must use its own pinned tokenizer/template and calibrated domain. A scalar reward can carry length, position, and style bias. Evaluate reward distributions and human agreement by slice before trusting optimization.

Returns, Value Head, and GAE

With value prediction $V_t$ , temporal-difference residual

\delta_t = r_t + \gamma V_{t+1} - V_t,

and Generalized Advantage Estimation (GAE)

A_t = \delta_t + \gamma\lambda(1-d_t)A_{t+1},

where $d_t$ marks terminal/padded positions. Compute the recursion only over the response mask. The return target is $R_t=A_t+V_t$ . Many systems apply masked advantage whitening across the rollout batch; distributed implementations must use consistent global statistics and must not include padding zeros.

Train the value head with a masked value loss, often with value clipping, and the policy with the PPO clipped surrogate:

L^{\text{clip}}_t=\min\left(r_t(\theta)A_t, \operatorname{clip}(r_t(\theta),1-\epsilon,1+\epsilon)A_t\right),

where now $r_t(\theta)=\exp(\log\pi_\theta-\log\pi_{\text{old}})$ is the probability ratio, not the reward above.

Advantage (\hat{A}_t): 1.0

Positive Advantage: The action was better than expected. Objective is clipped when r_t > 1.2 to prevent over-updating.

Unclipped Value

Final Objective

Canonical Update Skeleton

This pseudocode states the tensor semantics; it is deliberately not labeled runnable because model wrappers, generation backends, and distributed rollout stores differ.

rollout = generate_with_old_policy(prompts)
response_mask = rollout.response_mask  # [batch, response_tokens]

with no_grad():
    old_logp = completion_logp(old_policy, rollout, response_mask)
    ref_logp = completion_logp(reference, rollout, response_mask)
    rm_score = reward_model_score(rollout)
    old_values = value_head(rollout)

token_rewards = -kl_beta * (old_logp - ref_logp)
token_rewards[rollout.last_valid_index] += rm_score
advantages, returns = masked_gae(
    token_rewards, old_values, response_mask, gamma, gae_lambda
)
advantages = masked_global_whiten(advantages, response_mask)

for _ in range(ppo_epochs):
    new_logp = completion_logp(policy, rollout, response_mask)
    new_values = value_head(rollout)
    ratio = (new_logp - old_logp).exp()
    policy_loss = masked_clipped_surrogate(ratio, advantages, response_mask)
    value_loss = masked_clipped_value_loss(new_values, old_values, returns, response_mask)
    entropy = masked_entropy(policy, rollout, response_mask)
    optimize(policy_loss + value_coef * value_loss - entropy_coef * entropy)

Store the sampling policy revision and old log probabilities with every rollout. Too many epochs or delayed consumers increase policy lag: data becomes stale relative to the current policy and the importance ratio becomes unreliable. Set maximum rollout age and discard or regenerate stale batches.

Stability and Systems Contract

Pin the SFT policy, reference, RM, value initialization, tokenizer/templates, dataset manifest, generation settings, reward normalization, KL controller, $\gamma$ , $\lambda$ , clipping coefficients, batch in response tokens, seeds, and distributed topology. Save all four model identities plus optimizer, scheduler, scaler/FP8 state, RNG, rollout cursor/store state, and global tokens.

Monitor:

raw RM reward, non-score reward, per-token KL, shaped reward;
policy/value loss, value explained variance, advantage mean/std, entropy;
clip fraction, ratio distribution, response length/EOS, truncation and padding;
nonfinite/overflow, gradient norms, rollout/train throughput, queue age and policy lag;
held-out human preference, capability retention, safety, format/tool, and cost slices.

Abort on persistent nonfinite values, value divergence, KL/entropy/length outside bounds, excessive clip fraction or stale rollouts, reward increase without held-out improvement, or critical safety regression. Keep the frozen pre-PPO bundle as last-known-good.

Emerging Variants

Research systems propose decoupled rollouts, alternative trust regions, replay, and hybrid-policy methods such as Outer-PPO or HP3O. These are emerging design points, not a replacement for understanding the canonical loop. Compare them only with matched models, reward models, rollout budgets, data, and evaluation; describe maturity and implementation availability explicitly.

Quizzes

Quiz 1: Why does LLM PPO need a value head?

The value head provides a baseline and return estimate for each valid response position. GAE uses it to reduce variance and assign credit across the generated trajectory.

Quiz 2: Why is the response mask required even when padding uses an EOS token?

EOS and padding IDs can coincide or appear at different lengths. The explicit mask prevents prompt, padding, and post-terminal positions from entering KL, advantage, policy, value, and entropy reductions.

Quiz 3: What signals reward hacking during PPO?

Raw RM reward rises while held-out human preference, correctness, or safety does not; length/style shifts or KL grows; or reward gains concentrate in biased slices. Separate reward components and fixed evaluations expose it.

Quiz 4: Why is policy lag dangerous?

Rollouts generated by an old policy become increasingly off-policy. Importance ratios spread, clipping dominates, and updates use stale behavior. Limit rollout age and regenerate data rather than silently reusing it.

References

Schulman, J., et al. (2017). Proximal Policy Optimization Algorithms. arXiv:1707.06347.
Ouyang, L., et al. (2022). Training language models to follow instructions with human feedback. arXiv:2203.02155.
Zheng, R., et al. (2023). Secrets of RLHF in Large Language Models Part I: PPO. arXiv:2307.04964.