Proximal Policy Optimization Algorithms

The step-size problem

A policy gradient method improves a stochastic policy $\pi_\theta$ by raising the probability of actions that beat expectation, following the estimator $\hat g = \mathbb{E}_t[\nabla_\theta \log \pi_\theta(a_t \mid s_t)\,\hat A_t]$. The trouble is that this gradient is a direction, not a distance: it tells you where to move in parameter space but says nothing about how far is safe.

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

A step that is too small wastes samples and learning crawls; a step that is too large pushes the policy into a region the collected trajectories no longer describe, and performance collapses with no data on hand to undo the damage. TRPO addresses this by maximizing the surrogate subject to a hard constraint $\mathbb{E}_t[\mathrm{KL}(\pi_{\theta_\text{old}} \,\|\, \pi_\theta)] \le \delta$, which keeps each update inside a trust region. That works, but it requires solving a constrained second-order problem every step — a conjugate-gradient inner loop against the Fisher-vector product $F v = \nabla^2_\theta \mathrm{KL} \cdot v$ — which is awkward to implement and incompatible with architectures that share parameters or add noise. PPO keeps the trust-region intuition but throws out the machinery, using only first-order updates.

The surrogate objective

PPO is built around the probability ratio between the new and old policies, which is exactly $1$ before any gradient step is taken because the two policies are identical at that moment.

$$r_t(\theta) = \frac{\pi_\theta(a_t \mid s_t)}{\pi_{\theta_\text{old}}(a_t \mid s_t)}$$

The "conservative policy iteration" surrogate weights each advantage by this ratio, giving $L^{\text{CPI}}(\theta) = \mathbb{E}_t[r_t(\theta)\,\hat A_t]$. Maximizing it directly is the naive move and it fails for the same reason vanilla gradients do: nothing stops the optimizer from driving $r_t$ to enormous values when $\hat A_t > 0$, sending $\pi_\theta$ far outside the region where the old samples — and the advantage estimates computed from them — remain valid.

$$L^{\text{CPI}}(\theta) = \mathbb{E}_t\big[r_t(\theta)\,\hat A_t\big]$$

Clipping the ratio

PPO's central trick is to clip the ratio into $[1-\epsilon,\,1+\epsilon]$ and take the minimum of the clipped and unclipped terms, so the objective becomes a pessimistic lower bound on $L^{\text{CPI}}$.

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

The $\min$ behaves differently depending on the sign of the advantage, and it is worth tracing both cases. When $\hat A_t > 0$ the action was better than expected and we want $r_t$ to rise; the term $\operatorname{clip}(r_t, 1-\epsilon, 1+\epsilon)\hat A_t$ caps the reward at $r_t = 1+\epsilon$, so once the ratio crosses that ceiling the objective flattens and its gradient vanishes — there is no incentive to make an already-favored action drastically more likely. When $\hat A_t < 0$ the action was worse than expected and we want $r_t$ to fall; here the clip floors the contribution at $r_t = 1-\epsilon$, so pushing the probability down past that point stops helping. In both cases the $\min$ ensures the bound only ever removes incentive to move far from $r_t = 1$; it never lets a favorable clip make the objective look better than the honest unclipped term. One asymmetry to keep in mind: when $\hat A_t < 0$ and the ratio has already blown up past $1+\epsilon$ (the policy moved the wrong way), the unclipped term is the smaller one, so the gradient is not cut off and the update can still pull the mistake back.

The full loss

In practice the policy and value function share a network, and exploration needs encouragement, so the quantity actually optimized combines three terms.

$$L_t(\theta) = \mathbb{E}_t\big[L_t^{\text{CLIP}}(\theta) - c_1\, L_t^{\text{VF}}(\theta) + c_2\, H[\pi_\theta](s_t)\big]$$

The first term is the clipped surrogate that drives policy improvement. The second is a value-function regression loss, $L_t^{\text{VF}} = (V_\theta(s_t) - V_t^{\text{targ}})^2$, which trains the critic whose predictions feed the advantage estimates; $c_1$ trades policy against value learning when they share weights. The third is an entropy bonus $H[\pi_\theta]$ that keeps the action distribution from collapsing prematurely, with $c_2$ controlling how much exploration is rewarded.

Token-level PPO for language models

The reason PPO matters far outside robotics is that RLHF casts text generation as exactly this MDP. The state $s_t$ is the prompt plus everything generated so far, the action $a_t$ is the next token, and the policy is the language model itself: $\pi_\theta(a_t \mid s_t)$ is the softmax probability the model assigns to token $t$ given the context. The ratio is therefore computed per token, comparing the model being trained against a frozen snapshot $\pi_{\theta_\text{old}}$ taken at the start of the current PPO iteration.

$$r_t(\theta) = \frac{\pi_\theta(\text{tok}_t \mid \text{tok}_{<t}, \text{prompt})}{\pi_{\theta_\text{old}}(\text{tok}_t \mid \text{tok}_{<t}, \text{prompt})}$$

The efficiency that makes this tractable comes from teacher forcing: a single forward pass over the full completion, with the realized tokens as labels, yields $\log \pi_\theta(\text{tok}_t \mid \cdot)$ at every position at once, so all $L$ ratios for an $L$-token response fall out of one pass rather than $L$ sequential decodes. RLHF then adds a per-token penalty for drifting away from a separate reference policy $\pi_\text{ref}$ — the frozen SFT model — pulling the effective per-token reward toward

$$\tilde r_t = r^{\text{RM}}_t - \beta\, \log \frac{\pi_\theta(\text{tok}_t \mid s_t)}{\pi_\text{ref}(\text{tok}_t \mid s_t)}$$

where $r^{\text{RM}}_t$ is nonzero only at the final token (the scalar reward-model score). This KL term, distinct from PPO's own clipping, is what stops reward hacking: without it the policy would chase whatever degenerate text maximizes the imperfect reward model, and the $\beta$-weighted anchor to $\pi_\text{ref}$ keeps generations fluent and on-distribution.

Advantages via GAE

The advantage $\hat A_t$ that scales every ratio is itself estimated, and PPO almost always uses generalized advantage estimation. The one-step TD residual measures how much better a transition turned out than the critic predicted.

$$\delta_t = r_t + \gamma V(s_{t+1}) - V(s_t)$$

GAE forms an exponentially weighted sum of these residuals, where $\lambda \in [0,1]$ interpolates between the low-variance, high-bias one-step estimate ($\lambda = 0$, just $\delta_t$) and the high-variance, low-bias Monte Carlo return ($\lambda = 1$).

$$\hat A_t^{\text{GAE}(\gamma,\lambda)} = \sum_{l=0}^{\infty} (\gamma\lambda)^l\, \delta_{t+l}$$

In the language-model setting $\gamma$ is typically set near $1$ since the sequence is short and the reward is terminal, and the value head $V(s_t)$ — predicting expected return from each partial generation — is the critic trained by the $L^{\text{VF}}$ term above.