[RL] Add score-centered policy gradients

[nanjiangwill]

Summary

Implement score centering from Score Centering Stabilizes Off-policy Reinforcement Learning end to end for Miles.

• Add the paper-exact modeled-tail estimator using the sampler top-k head (k=128 by default), raw fp32 behavior log-probabilities, the sampled-token behavior score, and the detached additive correction.
• Add an exact finite-support estimator that composes with the bounded sampling-support replay already present on this target branch.
• Support score centering alone, truncated importance sampling, and the built-in IcePop masked-importance-sampling bounds.
• Carry typed distribution metadata through direct rollout, session v1/v2, multi-turn/tool calls, sample codecs, Ray transport, FSDP, Megatron TP, and context parallelism.
• Treat environment/tool tokens as forced singleton distributions, giving them exactly zero policy gradient.
• Fail closed when request semantics or response metadata cannot represent the behavior policy faithfully.

Estimator semantics

Modeled tail

The implementation matches the authors’ released PyTorch reference:

• Keep the top-k behavior probabilities as a subdistribution; do not renormalize the retained head.
• Preserve the sampled-token behavior log-probability even when the sampled token is outside the head.
• Compute trainer probabilities in fp32 over the real, unpadded vocabulary.
• Model the omitted tail using the trainer distribution rescaled to the behavior tail mass, with eps=1e-6.
• Detach the residual centering coefficient and the sampled-token importance weight exactly as in the reference loss.

Without replay, request validation requires an unconstrained behavior distribution and rejects value-changing logit transforms whose full behavior distribution is unavailable.

Exact replay support

When bounded top-p/top-k replay is enabled, SGLang supplies the complete realized support and normalized behavior probability for every support token. Miles normalizes current actor logits on that same fixed support and evaluates the centering expectation exactly—there is no modeled tail. Support filters and constraints are therefore orthogonal to score centering in this mode.

Training integration

• Uses a sparse, autograd-aware tensor-parallel log-probability primitive, avoiding materialization of a full-vocabulary log-softmax solely for the correction.
• Excludes padded vocabulary columns from normalization.
• Preserves duplicate sparse queries and accumulates their gradients correctly.
• Applies score-centered REINFORCE directly rather than passing the surrogate through PPO clipping.
• Preserves optional entropy and reference-KL objectives where their actor-score semantics remain valid.
• Reports sampler/train head mass, centered importance weight, ESS, rollout/train KL, and log-probability mismatch metrics.

SGLang dependency

Requires the generic response metadata proposed in sgl-project/sglang#40932: raw rectangular output top-logprobs and normalized probabilities aligned with realized sampling-support IDs. The SGLang API contains no Miles- or score-centering-specific naming.

Validation

• 666 affected CPU tests passed on modal commit a59ed7bc6.
• Includes a real two-process Gloo tensor-parallel forward/gradient test covering duplicate token queries, exact-support normalization, modeled full-vocabulary normalization, and padded vocabulary exclusion.
• Formula tests compare modeled-tail score centering and exact-support score centering against dense reference implementations for score centering alone, TIS, and MIS.
• Covers request validation, metadata corruption, aborts, sample trimming/merging, forced tokens, session codecs, Ray transport, FSDP/Megatron batch keys, and argument validation.
• Black, isort, Ruff, compilation, and git diff --check pass under the repository-pinned tool versions.
• One test newly present on the target branch (test_external_endpoint_owns_sampling_support_forwarding) was deselected because it independently fails on the clean target commit: it creates an external endpoint without setting the required --rollout-num-gpus 0.
• GPU/E2E experiments were intentionally not run.