Add KL-penalized advantage adjustment

[arcticfly]

Summary

• Adds a new mechanism that adjusts per-token advantages based on KL divergence from a reference model — tokens where the policy has drifted more get reduced advantages, tokens that drifted less get increased advantages. The adjustment is zero-mean (centered) across tokens.
• New LocalBackend.train() parameters: kl_penalty_coef, kl_penalty_reference_step, and kl_ref_adapter_path
• Fixes a pre-existing bug in preprocessing/inputs.py where warmup config used incorrect field names (lr → learning_rate, kl_coef → kl_penalty_coef)

Usage

To penalize divergence from the base (step-0) checkpoint:

result = await backend.train(
    model,
    trajectory_groups,
    kl_penalty_coef=1.0,
    kl_penalty_reference_step=0,
)

• kl_penalty_coef controls the strength of the penalty (default 1.0, set to 0.0 to disable).
• kl_penalty_reference_step selects which checkpoint to use as the reference. Use 0 for the base checkpoint, or any other saved step. If omitted, the base model (LoRA disabled) is used.

Test plan

• All linting/formatting checks pass (uv run prek run --all-files)
• 5 unit tests for the advantage adjustment formula pass
• Remote sweep of 9 kl_penalty_coef values (0.0001–1.0001) with kl_penalty_reference_step=0 completed successfully on Kubernetes H200 GPUs, all 20 steps each

🤖 Generated with Claude Code

[bradhilton]

LGTM