On KL-Regularized Policy Optimization

On KL-Regularized Policy Optimization. By Yifan Zhang, from Princeton University. Published on the sixth of October.

Asynchronous reinforcement learning for large language model agents has a built-in mismatch. One policy is trained on trajectories generated by another. Rollouts come from stale checkpoints, and even at identical parameters the inference engine's probabilities differ from the trainer's. Two standard remedies exist and each has a cost. Clipping importance ratios biases the update. Sampling a group of responses per prompt, as GRPO does, becomes expensive when episodes are long. So the field is left choosing between an update that is biased and an update that is costly, and the cost bites hardest exactly when episodes run long.

The proposal is KL-Regularized Policy Optimization, or KLPO, a framework that anchors the KL regularizer at the sampler, the policy that actually produced the data, rather than somewhere else. The end product is a critic-free update that uses one rollout per prompt, and needs neither a learned normalizer nor a group of responses. Each of those absences answers something in the setting. One rollout per prompt means a long episode is paid for once rather than once per group member. No importance weights means the update does not lean on clipping. The rest of the episode follows how the derivation gets there.

The derivation runs in steps. With the regularizer anchored at the sampler, the regularized improvement step has a closed-form Gibbs solution. That solution implies an optimality condition on a log-ratio, and KLPO fits that condition by least squares on the sampler's own trajectories. Because the sampler's probability enters only through a log-ratio, no importance weights are needed at all.

One obstacle remains: the log-partition function, which is intractable. Profiling out the regression intercept removes it, replacing it with two quantities that can be computed, the signal's mean under the sampler plus a KL divergence from sampler to trainer. For token-level policy mirror descent targets, the paper then shows that the resulting gradient can be computed from terminal returns without a critic. It can be done through sampler-centred scores, or through a single trajectory residual, and this holds even when tool outputs are stochastic.

Further results tighten the picture. Independent Monte Carlo estimates of the KL term are proved to keep these gradients unbiased. Cheaper top-K and binary approximations are offered, with the exact KL gap each one introduces derived rather than assumed. The critic-free gradient result is stated for token-level policy mirror descent targets specifically, which is the scope within which it is claimed.

The framework also turns out to be unifying. SPPO, GPO, REBEL and BPO each arise as special cases of KLPO. The step that makes all of it work is the choice of anchor: put the regularizer where the data came from, and the mismatch between sampler and trainer no longer calls for clipped importance ratios or a group of responses per prompt.