[Linbo-Liu]

Summary

Adds score function centering (SFC) to the slime backend, following Tinker's StableReinforce loss. It keeps policy-gradient training stable when rollouts come from an older version of the model, as in asynchronous training. To support it, the rollout gateway can ask SGLang for K extra random draws from the sampling distribution at every generated token.

At each generated token a with draws v_1 … v_K from the same sampler distribution:

c = z_θ(a) − mean_k z_θ(v_k)
L = −A · c  −  β_e · mean_k log p_θ(v_k),    β_e = β₀ / √age  (0 for fresh tokens)

age is the number of optimizer steps since the weights that sampled the token were pushed to SGLang. The centered score needs only K + 1 logit gathers per position; the loss has no importance ratio and no entropy term.

Usage

cd src/agentcore_rl_toolkit/backends/slime/examples/math_agent
SLIME_DIR=... MODEL_DIR=... TRAIN_DATA_PATH=... MODEL_TYPE=... bash train_sfc.sh

train_sfc.sh is train.sh with SFC enabled and asynchronous training: it runs slime's train_async.py on separate trainer and rollout GPUs (TRAINER_GPUS, ROLLOUT_GPUS, 4 each by default), so SGLang generates the next rollout while Megatron trains on the current one. It launches through python3 -m agentcore_rl_toolkit.backends.slime.train ${SLIME_DIR}/train_async.py, which registers the SFC flags with slime's argument parser, and passes:

  --loss-type custom_loss \
  --custom-loss-function-path agentcore_rl_toolkit.backends.slime.integration.score_centering.score_centering_loss \
  --custom-megatron-init-path agentcore_rl_toolkit.backends.slime.integration.score_centering.install \
  --sfc-num-auxiliary-token-samples 8 \
  --sfc-kl-beta 0.3 \
  --sfc-loss-agg-mode token-mean \

It also exports AGENTCORE_SGLANG_TOKEN_SAMPLES=1 to the SGLang engines. Sampling must be untruncated (--rollout-top-p 1.0, --rollout-top-k -1), and --entropy-coef must be 0; both are checked at startup.

Changes

Rollout gateway

• SamplingBackend.generate takes an optional num_auxiliary_token_samples; sessions request it with create_session(sid, num_auxiliary_token_samples=K), and backends without the parameter are rejected at session creation. Backends return the draws as AuxiliaryTokenSamples on TurnRecord, and they reach TraceRecord.auxiliary_token_samples on trained tokens.
• New SGLang plugin sglang_token_samples, registered as a sglang.srt.plugins entry point. With AGENTCORE_SGLANG_TOKEN_SAMPLES=1, a top_logprobs_num=K request returns the sampled token followed by K random draws instead of the top-K tokens. The top-logprobs path is reused as transport because it already carries per-position token ids end to end; while the variable is set, it changes the meaning of top_logprobs_num for every request on that server.
• SglangHttpBackend requests the draws, records SGLang's weight version, and fails the turn if a position doesn't carry exactly the sampled token plus K draws (plugin missing or truncated sampling).

slime backend

• integration/score_centering.py: the loss, and a --custom-megatron-init-path hook. Draws travel to the train actors in Sample.multimodal_train_inputs, the per-sample dict slime forwards unchanged; the hook takes them out of the model's forward kwargs, counts optimizer steps, and records the step of each weight push for token ages.
• integration/rollout.py: requests draws for training rollouts (not evaluation) and attaches them to each sample.
• train.py: entrypoint that runs slime's train.py or train_async.py with the --sfc-* flags registered (parse_args(add_custom_arguments=...)); slime's own entrypoints drop flags they don't define.
• examples/math_agent/train_sfc.sh: GSM8K recipe with SFC enabled.

Results

GSM8K with Qwen3-0.6B, rollouts on AgentCore Runtime, train_async.py with 2 trainer and 2 SGLang GPUs, 64 prompts × 4 samples per rollout, 16 optimizer steps per rollout, lr 1e-6, token-sum aggregation, β₀ = 0.3, K = 8. Rollouts are generated by weights 4 versions behind the trainer (64–79 optimizer steps): slime's async loop gives one version of lag, and an experiment-only hook pushes weight snapshots three versions old. The baseline is ratio-free REINFORCE (−A · log π, same aggregation), matching verl's bypass-mode loss_type=reinforce with rollout_is=null. One seed per run; runs are still in progress.

Run	Validation reward	Train reward	Drift from generator (ppo_kl)
REINFORCE	0.660 at rollout 9 (rollout 19 eval in progress)	peaked at 0.79 (rollout 11), 0.00 by rollout 19	0.8–59.6 over the last 64 steps
SFC + KL	0.253 → 0.679 → 0.713 at rollouts 9, 19, 29	0.66–0.85 over rollouts 35–39	0.002–0.020

As in the verl experiments, REINFORCE collapses under stale rollouts, with responses running to the length limit, while SFC + KL trains steadily. SFC starts slower here because the lag is enforced from step 0: for the first 64 steps the sampler is the initial policy, and the KL anchors the learner to it.

Testing

REINFORCE with async training and staleness=4 collapsed, while SFC is stable.