PPO for Language Models › Worked Example: Tracing a Complete PPO Update

To solidify the abstract mechanics, let us trace through a complete PPO update for a short response, computing every quantity step by step. We will use a simplified three-token response to keep the numbers manageable.

Setup: Consider a prompt "What is 2+2?" and a response with three tokens: "The", "answer", "is" (we omit the actual number for brevity). The reward model scores this as 0.7 (somewhat helpful but incomplete). The reference model and current policy have slightly different log probabilities for each token.

Step 1: Assign the sparse reward. The terminal reward of 0.7 goes to the final token, while earlier tokens receive zero:

$$r_1 = 0, \quad r_2 = 0, \quad r_3 = 0.7$$

Step 2: Compute the per-token KL divergence. Suppose the log probabilities are:

$$\log \pi_\theta(\text{``The''}) = -0.9, \quad \log \pi_{\text{ref}}(\text{``The''}) = -1.2$$ $$\log \pi_\theta(\text{``answer''}) = -1.4, \quad \log \pi_{\text{ref}}(\text{``answer''}) = -0.8$$ $$\log \pi_\theta(\text{``is''}) = -0.6, \quad \log \pi_{\text{ref}}(\text{``is''}) = -0.7$$

The KL contribution at each step is the log ratio:

$$\text{KL}_1 = -0.9 - (-1.2) = +0.3$$ $$\text{KL}_2 = -1.4 - (-0.8) = -0.6$$ $$\text{KL}_3 = -0.6 - (-0.7) = +0.1$$

Step 3: Apply KL shaping with $\beta = 0.2$. The shaped rewards are:

$$\tilde{r}_1 = 0 - 0.2 \times 0.3 = -0.06$$ $$\tilde{r}_2 = 0 - 0.2 \times (-0.6) = +0.12$$ $$\tilde{r}_3 = 0.7 - 0.2 \times 0.1 = +0.68$$

Token 1 receives a small penalty because the policy moved toward higher probability than the reference. Token 2 receives a small bonus because the policy moved toward lower probability than the reference. Token 3 receives most of the terminal reward, slightly discounted by the KL penalty.

Step 4: Compute TD errors and GAE advantages. Suppose the value function estimates are:

$$V(s_1) = 0.3, \quad V(s_2) = 0.5, \quad V(s_3) = 0.6, \quad V(s_4) = 0 \text{ (terminal)}$$

With $\gamma = 1.0$, the TD errors are:

$$\delta_1 = \tilde{r}_1 + V(s_2) - V(s_1) = -0.06 + 0.5 - 0.3 = 0.14$$ $$\delta_2 = \tilde{r}_2 + V(s_3) - V(s_2) = 0.12 + 0.6 - 0.5 = 0.22$$ $$\delta_3 = \tilde{r}_3 + V(s_4) - V(s_3) = 0.68 + 0 - 0.6 = 0.08$$

With $\lambda = 0.95$, the GAE advantages work backward:

$$\hat{A}_3 = \delta_3 = 0.08$$ $$\hat{A}_2 = \delta_2 + 0.95 \times \hat{A}_3 = 0.22 + 0.076 = 0.296$$ $$\hat{A}_1 = \delta_1 + 0.95 \times \hat{A}_2 = 0.14 + 0.281 = 0.421$$

All three advantages are positive, meaning all three tokens led to outcomes better than what the value function predicted. This makes sense: the response received a reward of 0.7 while the value function only predicted around 0.3 to 0.6 for the intermediate states.

Step 5: Normalize advantages. We subtract the mean and divide by the standard deviation:

$$\mu = (0.421 + 0.296 + 0.08) / 3 \approx 0.266$$ $$\sigma = \sqrt{((0.421-0.266)^2 + (0.296-0.266)^2 + (0.08-0.266)^2)/3} \approx 0.141$$ $$\hat{A}_1^{\text{norm}} \approx 1.10, \quad \hat{A}_2^{\text{norm}} \approx 0.21, \quad \hat{A}_3^{\text{norm}} \approx -1.31$$

After normalization, tokens 1 and 2 have positive advantages (they were above average), while token 3 has a negative advantage (it was below average even though the reward was positive, because the value function had already predicted a high value for that state).

Step 6: Compute probability ratios. Suppose during the gradient update epoch, the policy has shifted slightly:

$$\log \pi_{\theta_{\text{new}}}(\text{``The''}) = -0.85, \quad \log \pi_{\theta_{\text{old}}}(\text{``The''}) = -0.9$$

The ratio for token 1 is $\exp(-0.85 - (-0.9)) = \exp(0.05) \approx 1.05$. The advantage is positive, so the policy is trying to increase this token's probability. The ratio of 1.05 is within the trust region $[0.8, 1.2]$, so clipping does not activate.

Step 7: Apply the clipped objective. For token 1 with advantage $\hat{A}_1^{\text{norm}} = 1.10$:

$$\text{unclipped} = \rho_1 \times \hat{A}_1^{\text{norm}} = 1.05 \times 1.10 = 1.155$$ $$\text{clipped} = \text{clip}(1.05, 0.8, 1.2) \times 1.10 = 1.05 \times 1.10 = 1.155$$

Both are equal here because the ratio is within the trust region. The gradient flows through the unclipped objective. For a token where the ratio had exceeded 1.2, clipping would have capped the objective at $1.2 \times \hat{A}$, preventing the gradient from pushing the probability any higher.

This walkthrough shows the complete flow: sparse reward becomes shaped reward through KL penalty, shaped reward becomes advantages through GAE, and advantages become policy gradients through the clipped surrogate objective. Each step has a clear purpose, and together they produce a training signal that is both informative and stable.