PPO Clipped Objective

The problem

PPO (Proximal Policy Optimization) improves a policy using samples collected by an older copy of it. Taking several gradient steps on the same samples is efficient but dangerous: the policy could move so far that the samples no longer describe it. PPO limits how far each token's probability can move, by clipping the update.

For every sampled token, with $\rho = e^{\log \pi_{\text{new}} - \log \pi_{\text{old}}}$ (how much more, or less, likely the new policy makes it than the old one) and advantage $A$ (how much better than expected the token was, positive or negative):

where $\epsilon$ (epsilon) is 0.2, so $\tilde\rho$ is $\rho$ forced into $[0.8, 1.2]$. $J$ is maximized: a token with positive advantage wants its probability to go up, one with negative advantage wants it down, but the clip stops the update once $\rho$ has gone past the limit.

Example. $\epsilon = 0.2$. Four tokens, all with $\log \pi_{\text{old}} = 0$:

• $\rho = 1.5$, $A = +1$: $\rho A = 1.5$, $\tilde\rho A = 1.2$, minimum $1.2$: gradient 0 (clipped).
• $\rho = 0.5$, $A = +1$: $\rho A = 0.5$, $\tilde\rho A = 0.8$, minimum $0.5$: gradient $\rho A = 0.5$.
• $\rho = 0.5$, $A = -1$: $\rho A = -0.5$, $\tilde\rho A = -0.8$, minimum $-0.8$: gradient 0 (clipped).
• $\rho = 1.5$, $A = -1$: $\rho A = -1.5$, $\tilde\rho A = -1.2$, minimum $-1.5$: gradient $\rho A = -1.5$.

Objective: $(1.2 + 0.5 - 0.8 - 1.5)/4 = -0.15$. Gradient: $(0, 0.5, 0, -1.5) / 4 = (0, 0.125, 0, -0.375)$. Clip fraction: 2 of 4 tokens are clipped, $0.5$.

Where the gradient is zero

Where the unclipped term $\rho A$ is the minimum, its derivative with respect to $\log \pi_{\text{new}}$ is $\rho A$ (since $d\rho / d\log\pi = \rho$). Where the clipped term is the strict minimum, it does not change when the policy moves: derivative 0. That happens when:

• $A > 0$ and $\rho > 1 + \epsilon$: a good token already made more likely than allowed;
• $A < 0$ and $\rho < 1 - \epsilon$: a bad token already made less likely than allowed.

The other two cases keep their gradient: a mistake in the wrong direction (the second and fourth cases) is always corrected.

What to implement

For arrays of shape $(N,)$:

• ppo_objective(logp_new, logp_old, adv, eps=0.2): $J$.
• ppo_grad(logp_new, logp_old, adv, eps=0.2): $\partial J / \partial \log \pi_{\text{new}}$, shape $(N,)$, including the $1/N$ of the mean.
• clip_fraction(logp_new, logp_old, adv, eps=0.2): the fraction of tokens whose term is clipped, that is, where the gradient is 0. Training logs track it: if it is large, the updates are being cut a lot.

Requirements

• Work from log-probabilities: compute $\rho$ with np.exp of the difference.
• No loops.