A Complete Breakdown of PPO: From Policy Gradients to the Full LLM RLHF Training Pipeline

PPO (Proximal Policy Optimization) is the reinforcement learning algorithm most commonly used in the RLHF (Reinforcement Learning from Human Feedback) pipeline: the early alignment training of ChatGPT and Claude both used it to turn a model that already knows how to talk into one that “talks in a way people prefer.” This article walks through PPO in the order “why it’s needed → how it’s derived → how it fits into LLM training,” focusing on its concrete form in the LLM RLHF setting rather than the classic Atari/MuJoCo setting.

1. The RLHF Big Picture and MDP Formulation

A typical RLHF pipeline has three steps: SFT — using human-annotated, high-quality question-answer pairs to supervise-finetune the pretrained model, producing a base policy that “understands instructions”; training a Reward Model (RM) — having annotators rank multiple model outputs for the same prompt, training a scoring model $r_\phi(x, y)$ that takes prompt $x$ and response $y$ as input and outputs a scalar score representing “how much a human would like this response”; and the PPO stage — treating the SFT model as the initial policy and using the RM’s score as the reinforcement-learning reward signal to keep optimizing this policy. PPO only handles the third step, but how well it works depends heavily on the quality of the first two — if the RM learns inaccurately, PPO will faithfully optimize the policy straight into the RM’s blind spots (this is reward hacking, discussed later).

To handle language generation within the RL framework, the generation process first needs to be translated into the language of an MDP (Markov Decision Process): the state $s_t$ is the prompt plus the tokens generated so far, i.e., $s_t = (x, y_{<t})$; the action $a_t$ is the next token to generate, with the action space being the entire vocabulary (usually tens of thousands to over a hundred thousand tokens); the policy $\pi_\theta$ is the LLM itself, with $\pi_\theta(a_t \mid s_t)$ being the model’s probability distribution over the next token given the current context; an episode is one complete generation, from the prompt until EOS or truncation. The reward is where things differ most from classic RL: the RM usually only scores the complete sequence once generation finishes — no intermediate token gets an immediate reward from an “environment,” so the raw reward is sequence-level and sparse, landing only on the final token. This structure — “each token is one decision step, but the reward only appears at the end” — determines how GAE and the KL penalty need to be designed later on.

2. From Policy Gradients to the Clipped Objective

The most direct approach in reinforcement learning is the policy gradient: take the gradient of the expected return $J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta}[R(\tau)]$ directly. REINFORCE gives the estimator:

$$\nabla_\theta J(\theta) = \mathbb{E}_t\big[\nabla_\theta \log \pi_\theta(a_t \mid s_t) \cdot R_t\big]$$

The intuition is simple: if a sampled trajectory’s return $R_t$ is positive, raise the probability of the action that was sampled; if it’s negative, lower it. The problem is that $R_t$ comes from Monte Carlo sampling and has high variance, which gets amplified further by long sequences (LLM generation often runs to hundreds of tokens) and a huge action space (tens of thousands of vocabulary tokens). The standard fix is to introduce a baseline (usually a state-value function $V(s_t)$) and replace the raw return with the advantage function $A_t = R_t - V(s_t)$ — measuring “how much better this action is than the average at this state,” with mean zero and noticeably lower variance:

$$\nabla_\theta J(\theta) = \mathbb{E}_t\big[\nabla_\theta \log \pi_\theta(a_t \mid s_t) \cdot A_t\big]$$

Policy gradients have another, thornier problem: the size of an update is hard to control. If the learning rate is even slightly too large, a single gradient update can drastically shift the policy’s output distribution — for an LLM, this can mean the model starts producing gibberish within a few steps, or collapses into a degenerate mode that fools the RM while producing poor-quality language. Once this shift happens it’s hard to recover from, since the next batch of rollouts is already sampled from the damaged policy. TRPO’s (Trust Region Policy Optimization) approach is to explicitly constrain the KL divergence between the old and new policy, keeping each update inside a “trust region” — but this requires solving a constrained second-order optimization problem, which is complex and expensive to implement. PPO’s starting point is: can a simpler first-order method achieve a similar effect?

PPO’s answer is to use a probability ratio to measure the difference between the old and new policy: $r_t(\theta) = \pi_\theta(a_t \mid s_t) / \pi_{\theta_{\text{old}}}(a_t \mid s_t)$. If $r_t(\theta) A_t$ is optimized directly, when $A_t$ is large the optimization process will aggressively push $r_t(\theta)$ up, causing the policy to step outside the trust region in one move. PPO’s clipped objective squeezes this ratio into a fixed interval:

$$L^{\text{CLIP}}(\theta) = \mathbb{E}_t\Big[\min\big(r_t(\theta)\, A_t,\ \text{clip}(r_t(\theta),\, 1-\epsilon,\, 1+\epsilon)\, A_t\big)\Big]$$

$\epsilon$ is usually taken between 0.1 and 0.2. When $A_t > 0$ (this action is better than average), once $r_t(\theta)$ exceeds $1+\epsilon$, the clipped term stops growing, so the gradient no longer has any incentive to keep pushing $r_t$ higher; the symmetric case holds when $A_t < 0$ — once $r_t(\theta)$ drops below $1-\epsilon$, the objective no longer improves by lowering $r_t$ further. Taking the $\min$ rather than just using the clipped value ensures this objective is always a pessimistic lower bound on the true objective — regardless of the sign of $A_t$, what actually gets optimized is always the more conservative estimate, preventing the model from making an overconfident, irreversible update based on a single sample’s advantage estimate. This is where the name “proximal” comes from: a cheap first-order method achieving an effect similar to TRPO’s “each update’s magnitude is bounded.”

3. GAE and the Full Loss Function

The $A_t$ in the clipped objective needs a value function $V_\phi(s_t)$ to estimate it. PPO is usually paired with GAE (Generalized Advantage Estimation) to compute the advantage:

$$\delta_t = r_t + \gamma V_\phi(s_{t+1}) - V_\phi(s_t), \qquad A_t^{\text{GAE}(\gamma,\lambda)} = \sum_{l=0}^{\infty} (\gamma\lambda)^l\, \delta_{t+l}$$

$\delta_t$ is the single-step TD error, and $\lambda$ controls “how far to look ahead”: at $\lambda=0$, GAE reduces to using only the single-step TD error (low variance, high bias); at $\lambda=1$, it reduces to the Monte Carlo advantage estimate (high variance, unbiased), with common values trading off between 0.9 and 0.97. In the LLM RLHF setting, since the raw reward only lands on the final token of the sequence, $V_\phi(s_t)$ needs to learn to “predict the RM score that will ultimately be received from the current point of generation to the end” — which is exactly why a separate Critic (value model) needs to be trained, usually initialized from the SFT model or the RM, with its output dimension changed to one scalar per token position.

In practice, PPO’s total loss is a combination of three terms:

$$L(\theta, \phi) = L^{\text{CLIP}}(\theta) - c_1 \, L^{\text{VF}}(\phi) + c_2 \, S[\pi_\theta](<s_t>)$$

$L^{\text{CLIP}}(\theta)$ is the clipped policy objective derived above, training the Actor; $L^{\text{VF}}(\phi) = (V_\phi(s_t) - V_t^{\text{target}})^2$ is the Critic’s regression loss (with $V_t^{\text{target}} = A_t^{\text{GAE}} + V_\phi(s_t)$); $S[\pi_\theta](<s_t>)$ is the policy’s entropy, with an entropy bonus added to encourage the policy to retain some randomness and prevent premature collapse to deterministic output. $c_1$ and $c_2$ control the weight of the value loss and the entropy bonus, respectively.

4. PPO in LLM RLHF: Four Models and the KL Penalty

In a classic RL setting (such as MuJoCo), PPO only needs two networks, an Actor and a Critic. The PPO stage of LLM RLHF has to keep four models in GPU memory at once, which is its biggest engineering challenge compared to classic PPO:

Model	Trained?	Role
Policy / Actor	Trained	The LLM currently being optimized, initialized from the SFT model, responsible for generating responses
Reference Model	Frozen	A read-only copy of the SFT model, used only to compute the KL penalty, keeping the Policy from drifting too far
Reward Model	Frozen	The separately trained scoring model, giving a scalar score only at the end of the sequence
Critic / Value Model	Trained	Predicts the value $V_\phi(s_t)$ at each token position, supplying the advantage estimate for GAE

All four models tend to be roughly similar in parameter count (the Critic and Reward Model sometimes share a backbone, differing only in the output head), meaning the GPU memory and compute cost of the PPO stage is roughly several times that of pure inference — this is also the direct motivation behind later variants like GRPO wanting to cut the Critic.

Relying only on the RM’s score as the reward makes it very easy for the policy to optimize toward the RM’s blind spots instead of genuinely getting better. The standard RLHF approach is to add the KL divergence between the Policy and the Reference Model as a penalty term in the reward:

$$r_t = \underbrace{\mathbb{1}[t = T] \cdot r_\phi(x, y)}_{\text{given only at the final token}} \;-\; \beta \log\frac{\pi_\theta(a_t \mid s_t)}{\pi_{\text{ref}}(a_t \mid s_t)}$$

This KL penalty is computed per token and added to the reward at every single step, while the RM’s score only appears at the final token — GAE automatically propagates the trailing RM score backward to earlier tokens through temporal-difference credit assignment, while each step’s own KL deviation is also counted into the reward immediately; the two don’t conflict with each other. The value of $\beta$ is very sensitive: too small and it can’t constrain the policy, making reward hacking likely; too large and the policy barely dares to deviate from the SFT model, failing to learn the preferences the RM is meant to capture. Work such as InstructGPT uses adaptive KL control: continuously monitoring the deviation between the actual KL value and a target KL, and dynamically adjusting $\beta$ with a simple proportional controller, rather than fixing it as a constant.

5. Practical Pitfalls and Variants Beyond PPO

A few common pitfalls in engineering practice: Reward hacking — the policy finds a loophole in the RM’s scoring function instead of genuinely improving quality, commonly showing up as unusually verbose responses (the RM favors long responses) or piling on phrasing the RM happens to like; mitigations include the KL penalty, RM ensembling, and periodically retraining the RM on fresh data. Too many PPO epochs is a problem — repeating too many gradient updates on the same batch of rollout data causes the policy to overfit to that batch and the KL to spiral out of control quickly; LLM RLHF usually does only 1–2 epochs, far more conservative than in classic RL settings. Value Model training is unstable — it’s genuinely hard for the Critic to learn to predict the final RM score from scratch, and noisy value estimates early in training pollute the advantage estimate, so some practices warm up the Critic separately first. Reward normalization — the RM’s score scale can drift between different batches, so running mean/std normalization (whitening) is commonly used. Sampling parameters — the temperature/top-p used during the rollout stage directly determines the distribution of actions explored, which in turn affects the variance of the advantage estimate — an easily overlooked hyperparameter with a large impact.

PPO’s main cost is that it’s “heavy”: it needs a separately trained RM, a maintained Critic, and four models running at once. Variants over the past couple of years have each cut away part of this complexity: DPO (Direct Preference Optimization) skips the explicit RM and the PPO sample-update loop entirely, constructing a contrastive loss function directly from human preference pairs $(y_w, y_l)$ and optimizing the policy via supervised learning — at the cost of assuming a Bradley-Terry preference model and losing the flexibility of online exploration. GRPO (Group Relative Policy Optimization, proposed in DeepSeekMath) keeps PPO’s clipped objective and online rollouts, but drops the Critic — it samples a group of outputs for the same prompt and normalizes each output’s advantage using the group’s reward mean and variance, directly eliminating a model of the same scale as the Policy, making it a popular simplification direction in current LLM reasoning-capability training. Both trade off “PPO’s generality” against “fewer models, a simpler pipeline” — PPO is still the most mature and general baseline, but if the preference data is sufficiently well-structured, or an approximation via within-group relative rewards is acceptable, the lighter-weight variants often have a better cost-to-benefit ratio.

At bottom, what PPO does in LLM RLHF is constrain “improving the model using the RM’s scores” within a framework where “no single update can drift too far from the reference model”: the clipped surrogate objective bounds the magnitude of a single update, GAE reasonably distributes the sparse, sequence-level reward across every token, and the KL penalty adds another line of defense at the reward level. Understanding what problem each of these three pieces solves matters more than memorizing the final objective function’s formula.