@nnx.jit
def _ppo_step(policy, opt, obs, act, adv, logp_old, mask, epsilon,
entropy_coef, use_clip):
def loss_fn(policy):
logp_all = jax.nn.log_softmax(policy(obs), axis=-1)
logp = jnp.take_along_axis(logp_all, act[:, None], -1).squeeze(-1)
rho = jnp.exp(logp - logp_old)
surr = jnp.where(use_clip, jnp.minimum(
rho * adv, jnp.clip(rho, 1 - epsilon, 1 + epsilon) * adv),
rho * adv)
entropy = -(jnp.exp(logp_all) * logp_all).sum(-1)
loss = -(mask * (surr + entropy_coef * entropy)).sum() / mask.sum()
return loss, (rho, logp, entropy)
(_, (rho, logp, entropy)), grads = nnx.value_and_grad(
loss_fn, has_aux=True)(policy)
opt.update(policy, grads)
n = mask.sum()
return ((mask * (jnp.abs(rho - 1) > epsilon)).sum() / n,
(mask * (logp_old - logp)).sum() / n, (mask * entropy).sum() / n)
def ppo_epochs(ac, batch, adv, logp_old, epsilon, num_epochs,
entropy_coef=0.01, use_clip=True):
"""num_epochs clipped-surrogate passes on one frozen batch; returns
[num_epochs, 3] numpy diagnostics: fraction of ratios outside the
band, approximate KL, mean policy entropy."""
size = 1 << max(6, (len(adv) - 1).bit_length())
mask = jnp.asarray((np.arange(size) < len(adv)).astype(np.float32))
obs, act, adv, logp_old = (_pad(np.asarray(x), size) for x in
(batch.obs, batch.act, adv, logp_old))
step = nnx.cached_partial(_ppo_step, ac.policy, ac.opt_pi)
return np.array([step(obs, act, adv, logp_old, mask, epsilon,
entropy_coef, use_clip)
for _ in range(num_epochs)])Parameter Space versus Policy Space
One parameter, two actions (example due to Joshua Achiam), two updates of the same size \Delta\theta = 2:
\sigma'(0) = 0.25 against \sigma'(6) \approx 0.0025: no learning rate is right in both regions. One oversized step near indifference throws the policy into saturation, scores vanish, and the on-policy data is collected by the broken policy: the run is over. Capping the step in \theta caps the wrong quantity.
TRPO: a Bound, then a Constraint
J(\theta) \geq J(\theta_{\text{old}}) + \bar L(\theta) - \frac{4\gamma A_{\max}}{(1-\gamma)^2} \max_s D_{\text{KL}}\big(\pi_{\theta_{\text{old}}} \Vert \pi_\theta\big)
\bar L: the population surrogate over the discounted occupancy \rho_{\text{old}}; \bar L(\theta_{\text{old}}) = 0 exactly (the sampled \hat L need not vanish). Ascend the lower bound and J ascends monotonically. In practice: sampled \hat L, mean KL \leq \delta_{\text{KL}} over visited states (a proxy, no bound attached), second-order machinery to solve it.
Measuring steps in KL is steepest ascent under the Fisher metric: the natural gradient, :numref:sec_muon’s norm story again.
The Clip
L^{\text{CLIP}} = \frac1n\sum_{i,t} \min\!\big(\rho\hat A,\ \text{clip}(\rho,1-\epsilon,1+\epsilon) \hat A\big)
Once a ratio leaves the band in the paying direction, that sample’s gradient is zero; the pessimistic side stays open. PPO keeps the shape of the guarantee and none of the guarantee.
The Ablation: Eight Seeds, Clip On and Off
Same batches, same twenty passes: half or more of the unclipped seeds die near return 9 (saturation, for real); every clipped seed reaches the ceiling. The insurance pays out on about one ratio check in twenty. Which seeds die reshuffles; the rate is what is stable.
How to Know Your RL Is Broken
entropy: 0.64 over the first five updates, 0.26 over the last five- in these runs, KL and band-exits are front-loaded within a batch; the clip stalls the drift
- across training: entropy decays from about 0.65 to about 0.25 nats; the bonus slows the slide, :numref:
sec_regularizedexplains it
The Batch Goes Stale, Measured
Ratios are importance weights, so the appendix’s effective sample size applies, as a ratio-concentration diagnostic:
after 20 epochs the batch is worth 95% (clipped (PPO)) vs 42% (no clip) of its 136 steps“Reuse for a few epochs, then stop” as a dial: with the clip the weight spectrum stays nearly flat; without it, concentration to half or less. Weights-only: blind to advantages, dependence, and state-distribution staleness.