PPO for Language Models › PPO Objective for Language Models

The PPO objective for language models combines the clipped surrogate loss from standard PPO with the KL-shaped reward structure we just described. Understanding how these pieces fit together reveals why PPO works for alignment: it maximizes expected reward while maintaining two separate safeguards against instability. The KL penalty prevents the policy from drifting too far from the reference in absolute terms, while the clipping mechanism prevents any single update step from moving too far from the data collection policy.

Think of the full PPO objective as a three-way balancing act. The reward signal pulls the policy toward higher-scoring behaviors. The KL penalty pulls it back toward the reference model. The clipping mechanism ensures that each individual gradient step stays within a safe neighborhood. These three forces together create a training dynamic that is far more stable than any one of them alone.

Combining all these elements, the PPO objective for language models takes the following form. For a batch of prompt-response pairs, we compute:

$$\mathcal{L}_{\text{PPO}}(\theta) = \mathbb{E}_{(x,y) \sim \pi_{\theta_{\text{old}}}} \left[ \sum_{t=1}^{T} \min\left( \rho_t(\theta) \hat{A}_t, \text{clip}(\rho_t(\theta), 1-\epsilon, 1+\epsilon) \hat{A}_t \right) \right]$$

where:

• $\theta$: the parameters of the language model being optimized
• $\mathbb{E}_{(x,y) \sim \pi_{\theta_{\text{old}}}}$: the expectation over trajectories generated by the policy that collected the data
• $\rho_t(\theta)$: the probability ratio between the current and old policies at step $t$
• $\hat{A}_t$: the estimated advantage at step $t$, showing how much better the chosen action was compared to the average
• $T$: the length of the generated sequence
• $\epsilon$: the clipping parameter (typically 0.1 or 0.2) that defines the trust region

The min operator serves a important purpose in this objective. It takes the conservative lower bound between the unclipped and clipped objectives. This keeps updates cannot be too aggressive. If the advantage is positive, showing that the chosen action was better than expected, clipping limits how much we increase the probability of the action. If the advantage is negative, showing the action was worse than expected, clipping limits how much we decrease the probability. This prevents large updates that could collapse the policy.

The probability ratio $\rho_t(\theta)$ is the core mechanism that allows PPO to reuse trajectory data across multiple gradient steps. Without it, policy gradient methods can only take one gradient step per batch of data, making them extremely sample-inefficient. With the ratio, we can check whether the current policy has drifted too far from the data-collection policy and clip the gradient accordingly.

The ratio is defined as:

$$\rho_t(\theta) = \frac{\pi_\theta(y_t | x, y_{<t})}{\pi_{\theta_{\text{old}}}(y_t | x, y_{<t})}$$

where:

• $\pi_\theta(y_t | x, y_{<t})$: the probability of the token $y_t$ under the current policy (being updated)
• $\pi_{\theta_{\text{old}}}(y_t | x, y_{<t})$: the probability of the same token under the policy that collected the data (frozen for this update step)
• $y_t$: the token generated at step $t$
• $x$: the input prompt
• $y_{<t}$: the sequence of tokens generated prior to step $t$

This ratio measures how the policy has changed since the data was collected. A ratio of 1.0 means the policy assigns exactly the same probability as before. A ratio greater than 1 means the policy now assigns higher probability to this token. A ratio less than 1 means the probability has decreased. By basing updates on this ratio rather than directly on log probabilities, PPO can perform multiple gradient steps on the same batch of data without the policy drifting too far from the data-collection policy.

The key insight here is why we use the ratio rather than log probabilities directly. In standard REINFORCE, we would multiply the log probability by the advantage and take a gradient step. This works for a single step but fails for multiple steps: after updating the policy, the gradient computed from the old trajectories is no longer valid because the old trajectories were generated by a different policy. The ratio corrects for this by weighting each token's contribution by how much the current policy differs from the data-collection policy.

The advantage $\hat{A}_t$ is estimated using Generalized Advantage Estimation (GAE) applied to the KL-shaped rewards. GAE is a technique that balances the bias-variance tradeoff in advantage estimation by combining multi-step returns through exponential weighting:

$$\hat{A}_t = \sum_{l=0}^{T-t} (\gamma \lambda)^l \delta_{t+l}$$

where:

• $\gamma$: the discount factor (usually near 1.0 for RLHF)
• $\lambda$: the GAE smoothing parameter (typically 0.95)
• $T$: the length of the generated sequence
• $\delta_{t+l}$: the temporal difference (TD) error at step $t+l$

This summation calculates an exponentially weighted average of future TD errors. The parameter $\lambda$ controls the bias-variance trade-off: higher values rely more on observed returns (reducing bias but increasing variance), while lower values rely more on value estimates (reducing variance but potentially introducing bias from imperfect value predictions). Setting $\lambda = 1$ recovers Monte Carlo returns, while $\lambda = 0$ gives pure one-step TD learning. The typical value of 0.95 balances these considerations.

The TD error $\delta_t$ measures the surprise at each step: the difference between the observed reward plus the estimated value of the next state, and the estimated value of the current state:

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

where:

• $\tilde{r}_t$: the KL-shaped reward at step $t$
• $\gamma$: the discount factor
• $V(s_t)$: the value function estimate for state $s_t$ (how good is the current context expected to be?)
• $V(s_{t+1})$: the value function estimate for the next state (defined as 0 if $t=T$)

The TD error captures whether the transition was better or worse than expected. If the reward plus next-state value exceeds the current-state value, the transition was unexpectedly good, and the TD error is positive. If the sum falls short, the transition was unexpectedly bad, and the TD error is negative. These signals, accumulated through GAE, produce advantage estimates that indicate which actions led to better outcomes than the value function predicted.