KL-Regularized
Policy Optimization

Asynchronous reinforcement learning (RL) for large language model (LLM) agents trains one policy on trajectories generated by another: rollouts come from stale checkpoints, and the inference engine’s probabilities differ from the trainer’s even at identical parameters. Standard remedies either clip importance ratios, which biases the update, or, as in GRPO, sample a group of responses per prompt, which is costly when episodes are long. We propose KL-Regularized Policy Optimization (KLPO), a framework that anchors the KL regularizer at the sampler. The regularized improvement step then has a closed-form Gibbs solution, and KLPO fits its log-ratio optimality condition by least squares on the sampler’s own trajectories, so the sampler probability enters through a log-ratio and no importance weights are needed. Profiling out the regression intercept replaces the intractable log-partition function with the signal’s sampler mean plus a sampler-to-trainer KL divergence. For token-level policy mirror descent targets, we show that the resulting gradient can be computed from terminal returns without a critic, via sampler-centered scores or a single trajectory residual, even under stochastic tool outputs. We further prove that independent Monte Carlo estimates of the KL term keep these gradients unbiased, derive the exact KL gap of cheaper top-K and binary approximations, and show that SPPO, GPO, REBEL, and BPO arise as special cases of KLPO. The result is a critic-free update that uses one rollout per prompt and requires neither a learned normalizer nor a group of responses.

Historical sampler. Current trainer.

Retain the collection version and actual sampler probabilities. Recompute trainer scores and feedback at each update, with no multiplicative importance weights. The default is KLPO token regression + MC-KL.

From local policy improvement to a critic-free update.

The regression route determines how tokens are weighted. The KL estimator determines how their scores are corrected. Explore all eight implemented combinations.

−∑_u h_u [∇ log p(a_u) − (1/M) ∑_j ∇ log p(v_u,j)]

Here q is the actual sampler and p is the trainer; all probabilities are conditional on the visited prefix. Sum over policy tokens, then average over complete responses. M counts random draws; K counts head tokens.

A CPU example to inspect the update, with a native Molt interface for model training. The toy policy uses variable-length responses, a terminal verifier, historical sampler versions, and repeated learner updates. No model or dataset download is needed.

Requirements: uv, which sets up Python 3.10+ and PyTorch 2.2+ from the lockfile. Code is licensed under Apache 2.0.

# Clone and install
git clone https://github.com/yifanzhang-pro/KLPO.git
cd KLPO
uv sync

# Token regression + MC-KL (default, M=16)
uv run python examples/train_toy.py

# A single auxiliary draw is also valid
uv run python examples/train_toy.py --mc-samples 1

# Verify the loss implementation
uv run pytest -q

On KL-Regularized Policy Optimization
Yifan Zhang · Princeton University · arXiv:2610.08963 · 58 pages
September 18, 2026 · Updated October 5, 2026

The paper includes the complete derivation, proofs, implementation details, and the response-level instance (SKLPO) with its normalization and implicit soft-Q interpretation.

Download paper ↓arXiv ↗LaTeX source ↗GitHub ↗

The PDF here is an unmodified snapshot of paper-source commit 9a778da, which revises arXiv v1 (built from b485b29) with the β → 0⁺ extension of the update.

If you find this work useful, please cite:

@article{zhang2026klpo,
  title         = {On {KL}-Regularized Policy Optimization},
  author        = {Zhang, Yifan},
  journal       = {arXiv preprint arXiv:2610.08963},
  year          = {2026},
  eprint        = {2610.08963},
  archivePrefix = {arXiv},
  primaryClass  = {cs.LG},
  url           = {https://arxiv.org/abs/2610.08963}
}