KLPO

KL-Regularized Policy Optimization for Critic-Free Agentic Reinforcement Learning

KLPO is a critic-free, single-rollout method for off-policy agentic reinforcement learning. This repository implements KLPO token regression + Monte Carlo KL (MC-KL) by default: terminal rewards provide the feedback, and independent auxiliary token draws estimate the sampler-conditioned score correction. It needs no same-prompt response group or learned value/normalizer model.

Authors: Yifan Zhang et al.
Technical report: September 18, 2026 · Revised: September 20, 2026

[Project website] [Paper] [Algorithm reference] [Training guide] [Paper source]

• One complete response per prompt is sufficient. MC-KL uses auxiliary token draws at visited prefixes, not extra response rollouts.
• Token regression + MC-KL is the default. Each token uses its own coefficient R - beta * log(p_action / q_action) and an independently estimated score correction.
• Two regression routes, four KL estimators. Sequence regression is an alternative; TopK-KL, Binary KL, and Full KL are selectable for either route.
• Single-turn and multi-turn Molt integration. Experiment recipes cover R1/Qwen-Math, Python tools, and ALFWorld. The labs-molt KLPO branch preserves sampler records across assistant turns and tool observations; the pinned backend uses synchronous collection.

Figure 1 from the report. Step 7 builds on BPO’s critic-free PMD reformulation; Step 8 uses Score Centering. Independent MC-KL recovers the full-KL gradient in expectation under the stated sampling assumptions. TopK-KL and Binary KL are optional approximations.

Requires Python 3.10+ and PyTorch 2.2+. The toy example runs on CPU without a model or dataset download.

git clone https://github.com/yifanzhang-pro/KLPO.git
cd KLPO
python -m venv .venv
source .venv/bin/activate
pip install -e '.[test]'

python examples/train_toy.py  # Default: token regression + MC-KL, M=128
python examples/train_toy.py --mc-samples 1
python -m pytest -q

The example uses variable-length responses, a terminal verifier, historical sampler versions, and repeated learner updates. M counts independent auxiliary tokens per prefix; token regression supports M ≥ 1. The example is an implementation check, not a benchmark reproduction.

Let q be the actual collection sampler, p the current trainer, and R the terminal reward. At each visited prefix, draw v_j ~ q IID with replacement, independently of the complete rollout. The backward surrogate is:

ell_u = log p(a_u) - log q(a_u)
z_u   = log p(a_u) - mean_j log p(v_j)
loss  = -mean_responses sum_tokens stopgrad(R - beta * ell_u) * z_u

Sum over generated policy tokens and average over complete responses, without length normalization. Keep the recorded sampler probabilities and version fixed during reuse. See the loss API and sampling contracts for differentiable trainer inputs, masks, and adaptive replay requirements.

The regression route selects the feedback coefficient; the KL estimator selects the conditional score correction. All eight combinations are implemented. M is the number of random draws; K is the TopK-KL head size.

MC-KL averages independent sampler log-ratios. Top-K Aggregated KL (TopK-KL) keeps the sampler's K largest probabilities and combines all remaining tokens into one tail bucket. Binary KL groups the sampled action against its complement; Full KL uses the entire vocabulary. The toy vocabulary has four tokens, so --top-k 2 illustrates a nontrivial head/tail split; the training launcher defaults to K=128 when TopK-KL is selected.

Token regression uses per-token feedback; sequence regression uses trajectory feedback. Sequence MC-KL requires M ≥ 2 for leave-one-out residuals. Full-KL and independent-MC population equivalences require the report's assumptions; individual sample gradients can differ, and finite TopK-KL/Binary approximations need not preserve those equivalences. KLPO sequence regression is distinct from the SKLPO response-Gibbs comparison in the paper's appendix.

All recipes default to --route token --kl-estimator mc --mc-samples 128. The pinned backend requires synchronous collection with complete trajectories and one update per batch. The unified examples/run_experiment.py launcher sets this automatically for single-turn and multi-turn settings; pass --sync to the retained scripts/train_molt.py launcher (its asynchronous switches are not supported by the current pin). GPU training requires a compatible Linux/CUDA environment; follow the training guide. The optional predicted-KL budget is disabled in these launchers.

Preview single-turn and multi-turn experiments without installing Molt:

python examples/run_experiment.py \
  --setting deepseek_r1_distill_qwen_1p5b_math_sanity --print-config
python examples/run_experiment.py \
  --setting qwen2p5_7b_instruct_python_tool_10turn --print-config

The experiment guide provides all five launch settings, agent/data requirements, and evaluation protocols. Single-turn settings use Molt's built-in math agent; multi-turn agents, datasets, and checkpoints are supplied separately.

CPU tests cover all eight combinations, independent MC gradient expectations, leave-one-out residuals, M=1 token updates, TopK-KL tail corrections, masked/extreme probabilities, and trajectory/microbatch normalization. To also check the native backend contract:

MOLT_SOURCE_PATH=/path/to/labs-molt python -m pytest -q

This release provides the theory, loss implementation, CPU verification, and native training integration. GPU training and paper-scale benchmark reproduction have not been validated. Example hyperparameters are starting values, not tuned benchmark settings. Multi-turn tests cover native command validation, worker environment propagation, and complete-trajectory masking. The supported synchronous recipes consume each auxiliary record bank once; fixed-record reuse remains an empirical surrogate rather than a fresh conditionally unbiased MC estimate.

@techreport{zhang2026klpo,
  title  = {{KL}-Regularized Policy Optimization for Critic-Free Agentic Reinforcement Learning},
  author = {Zhang, Yifan and others},
  year   = {2026},
  month  = sep,
  url    = {https://yifanzhang-pro.github.io/KLPO/}
}

Download BibTeX.

Code is licensed under Apache 2.0. Training launcher scaffolding is adapted from FlashREINFORCE. The separately installed labs-molt backend retains its own licenses and notices; see NOTICE.