From policy gradients to train–rollout consistency

Return to the experiment lab. Read the systems companion for precision, determinism, asynchronous execution, PD, and HiCache.

The equations below are a tutorial derivation of the objectives implemented in this repository. They distinguish identities, biased surrogates, and engineering approximations. Paper links identify the original work; the implementation links explain the particular variant exported by the lab.

A common language

Let \(x\) be a prompt, \(y=(a_1,\ldots,a_T)\) a response, and \(h_t=(x,a_{<t})\) a prefix. Let \(R(x,y)\) be a reward independent of the parameters during differentiation. Use three different policies:

Symbol	Meaning	Where it comes from
\(q(a\mid h)\)	Behavior distribution that actually sampled the token	SGLang, including quantization, sampling transforms, and potentially older weights
\(p_0(a\mid h)\)	Frozen trainer before the optimizer update	Megatron’s recomputed old log-probs
\(p_\theta(a\mid h)\)	Current trainable policy	Megatron’s differentiable forward
\(\operatorname{sg}(z)\)	Stop-gradient: value \(z\), derivative zero	Detaching weights, rewards, and correction coefficients

At a fixed prefix define the score \(s_a=\nabla_\theta\log p_\theta(a\mid h)\). Assuming normalized differentiable probabilities and legitimate interchange of summation and differentiation:

\[ \sum_a p_\theta(a\mid h)s_a =\sum_a\nabla_\theta p_\theta(a\mid h) =\nabla_\theta 1=0. \]

Starting with \(J(\theta)=\sum_y p_\theta(y\mid x)R(x,y)\), apply \(\nabla p=p\nabla\log p\) and autoregressive factorization:

\[ \nabla J =\mathbb E_{y\sim p_\theta}\left[R\nabla\log p_\theta(y\mid x)\right] =\mathbb E_{y\sim p_\theta}\left[R\sum_{t=1}^{T}s_{a_t}\right]. \]

An action-independent baseline \(b(h_t)\) has zero expectation against the score under \(p_\theta\), so it can reduce variance without changing that expectation. Sampling from \(q\) generally breaks this cancellation. This is why identical checkpoint names do not prove an on-policy update.

GRPO

Original work: DeepSeekMath; clipping objective: PPO.

For \(G\) responses to one prompt, form group statistics and a response advantage:

\[ \bar R=\frac1G\sum_{i=1}^G R_i,\quad \sigma_R=\sqrt{\frac1G\sum_i(R_i-\bar R)^2},\quad A_i=\frac{R_i-\bar R}{\sigma_R+\epsilon}. \]

The response advantage is shared across its response tokens. The exact standard-deviation convention and numerical epsilon are implementation details; disabling standard-deviation normalization gives \(A_i=R_i-\bar R\). Group baselines include the sampled response itself, so the usual unbiased-baseline argument does not apply literally. For independent responses, mean centering alone scales the expected on-policy gradient by \(1-1/G\); standard-deviation normalization adds a random scale.

Define \(r_{i,t}(\theta)=p_\theta(a_{i,t}\mid h_{i,t})/p_0(a_{i,t}\mid h_{i,t})\). A clipped surrogate to maximize, omitting optional KL terms, is

\[ J_{\rm clip}=\frac1G\sum_i\frac1{T_i}\sum_t \min\left(r_{i,t} A_i, \operatorname{clip}(r_{i,t},1-\epsilon_-,1+\epsilon_+)A_i\right). \]

For positive \(A_i\), the upper clip limits incentive to increase token probability. For negative \(A_i\), the lower clip limits incentive to decrease it. At \(\theta=\theta_0\), inside the clip, \(\nabla r=r\nabla\log p_\theta\) reduces to the policy-gradient direction. This is a local surrogate, not a general identity for arbitrarily stale data.

The lab emits --advantage-estimator grpo, eight samples per prompt, and PPO clipping unless SC is enabled. --calculate-per-token-loss changes the outer reduction to total loss divided by total response-token count; it is not the same weighting as averaging each response’s mean.

TIS

Original work: On the Rollout-Training Mismatch in Modern RL Systems, with the authors’ technical article. The following algebra also describes the repository’s built-in callback.

At a fixed prefix, provided \(q(a)>0\) wherever \(p_0(a)f(a)\ne0\), change of measure gives

\[ \mathbb E_{a\sim p_0}[f(a)] =\sum_a q(a)\frac{p_0(a)}{q(a)}f(a) =\mathbb E_{a\sim q}[\rho_a f(a)],\qquad \rho_a=\exp(\log p_0(a)-\log q(a)). \]

The PPO policy-update ratio \(r=p_\theta/p_0\) and behavior-correction ratio \(\rho=p_0/q\) serve different purposes. Before clipping, their product is \(p_\theta/q\). slime multiplies the PPO loss by detached \(w(\rho)=\operatorname{clip}(\rho,L,U)\); the lab uses \(L=0,U=2\).

Derive its bias by subtracting the exact expression:

\[ \mathbb E_q[w(\rho)f]-\mathbb E_{p_0}[f] =\mathbb E_q[(w(\rho)-\rho)f]. \]

For \(L=0\), this becomes \(-\mathbb E_q[(\rho-U)f\,\mathbf1_{\rho>U}]\). Capping weights limits outliers but does not give unbiased importance sampling. With nonzero \(L\), small weights are also increased.

This fixed-prefix identity does not correct the distribution of prefixes. Full-sequence change of measure would use \(p_0(y)/q(y)=\prod_t\rho_t\), often with high variance. Token-wise TIS is a practical surrogate, not an exact replacement for that product.

Implementation: vanilla_tis_function, importance_weights. REINFORCE, including SC, uses detached current trainer probabilities in the correction numerator instead of \(p_0\).

ICE-POP

Original work: Every Step Evolves: Scaling Reinforcement Learning for Trillion-Scale Thinking Model.

The lab selects slime’s built-in icepop_function. Its behavior-correction rule is

\[ w_{\rm ice}(\rho)=\rho\,\mathbf1_{L\le\rho\le U},\qquad L=0.5,\quad U=2. \]

Inside the interval this equals importance sampling. Outside it, the contribution is zero. Its bias relative to the same fixed-prefix target is therefore

\[ \mathbb E_q[w_{\rm ice}(\rho)f]-\mathbb E_{p_0}[f] =-\mathbb E_q[\rho f\,\mathbf1_{\rho<L\;\lor\;\rho>U}]. \]

A token with \(\rho=4\) receives weight 2 under the lab’s TIS rule and weight 0 under IcePop. These are distinct estimators. The callback sets weights to zero; it does not renormalize the loss by the number of retained tokens. PPO clipping can still act on \(r\) independently. The paper’s complete training recipe contains other components; selecting this callback does not reproduce the whole paper.

Exported settings: --use-tis --custom-tis-function-path slime.backends.megatron_utils.loss.icepop_function --tis-clip-low 0.5 --tis-clip 2. Observe tis, tis_clipfrac, reward, and effective signal after masking.

R3

Original work: Stabilizing MoE Reinforcement Learning by Aligning Training and Inference Routers.

Write a MoE layer as \(F(x)=\sum_{i\in I(x)}g_i(x)E_i(x)\), with discrete expert set \(I=\operatorname{TopK}(u(x))\). A perturbation near a top-k boundary can swap an expert even when its magnitude is tiny. Decompose the output difference by adding and subtracting the same intermediate term:

\[ F_{I_q}(x_q)-F_{I_p}(x_p) =\underbrace{F_{I_q}(x_q)-F_{I_q}(x_p)}_{\text{same-route numerical/input difference}} +\underbrace{F_{I_q}(x_p)-F_{I_p}(x_p)}_{\text{expert-selection difference}}. \]

Replay sets the trainer route to recorded \(I_q\), eliminating the second term at this layer for this comparison. It does not eliminate the first. For a softmax-gated example,

\[ g_i=\frac{\mathbf1_{i\in I_q}\exp u_i}{\sum_{j\in I_q}\exp u_j},\qquad \frac{\partial g_i}{\partial u_j}=g_i(\mathbf1_{i=j}-g_j),\quad i,j\in I_q. \]

Replay expert identities, not detached expert outputs. Gradients can still flow through selected gates and experts. GLM uses sigmoid-based routing and its own normalization/scaling, so this softmax example is illustrative; keep the architecture’s actual gate rule.

The lab emits --use-rollout-routing-replay; slime-managed SGLang requests routed experts automatically. External engines must be launched with --enable-return-routed-experts. Storage grows with tokens × MoE layers × selected experts. R3 is independent of TIS/SC and is not meaningful for a dense model. Implementation: sglang_engine.py and the replay code in megatron_utils.

SC

Original work: Score Centering Stabilizes Off-policy Reinforcement Learning. The derivation here follows the algebra of slime’s score_centering_correction and its loss implementation.

From an unwanted mean to a centered score

Fix a prefix and let \(\mu=\sum_a q_a s_a\). In general \(\mu\ne0\), even with the same checkpoint on both sides. For a detached advantage \(A\), covariance decomposition gives

\[ \mathbb E_q[A s]=\mathbb E_q[A]\mu+\operatorname{Cov}_q(A,s). \]

Use \(\tilde s_a=s_a-\mu\). Then \(\mathbb E_q[\tilde s]=0\), and

\[ \mathbb E_q[A\tilde s]=\operatorname{Cov}_q(A,s). \]

The prefix-conditional mean term vanishes. The covariance is still under \(q\), so this does not establish an unbiased gradient under \(p_\theta\). Group-centered rewards across full responses do not ensure \(\mathbb E[A\mid h]=0\) at every prefix.

Compose with a correction weight

Let \(w_a=f(p_a/q_a)\), with \(f\) either identity, TIS, masking, or 1. Define \(\mu_w=\sum_a q_aw_as_a\) and center the weighted score:

\[ \tilde s^w_a=w_as_a-\mu_w,\qquad \mathbb E_q[\tilde s^w]=0. \]

An implementable minimization loss with exactly this local gradient direction is

\[ \ell=-\operatorname{sg}(A)\left[ \operatorname{sg}(w_a)\log p_a -\sum_v\operatorname{sg}(q_vw_v)\log p_v\right]. \]

Differentiate only \(\log p\): \(-\nabla\ell=A(w_as_a-\mu_w)\). Differentiating through \(w\) would add unwanted derivatives. With exact \(w=p/q\) and full support, \(\mu_w=\sum p s=0\); clipped or masked weights generally make it nonzero.

Why top-k does not require storing the full vocabulary

slime records the sampler’s head \(H\) (128 tokens in the lab). Its probabilities retain full-vocabulary mass, not a head-only renormalization. Set

\[\begin{split} P_H=\sum_{v\in H}p_v,\quad Q_H=\sum_{v\in H}q_v,\quad \alpha=\frac{1-Q_H}{1-P_H},\quad \hat q_v=\begin{cases}q_v&v\in H\\\alpha p_v&v\notin H.\end{cases} \end{split}\]

This approximation sums to one. Since \(\sum_v p_vs_v=0\), the tail score can be replaced algebraically:

\[ \mu_{\hat q}=\sum_{H}q_vs_v+\alpha\sum_{v\notin H}p_vs_v =\sum_H(q_v-\alpha p_v)s_v. \]

In the tail \(p_v/\hat q_v=1/\alpha\), so with \(w=f(p/q)\):

\[ \mu_{\hat q,w}=\sum_{v\in H} \left[q_v f(p_v/q_v)-\alpha f(1/\alpha)p_v\right]s_v. \]

This is the head-only residual used in the implementation. Full-distribution centering is exact; centering with the reconstructed tail is an approximation, and small tail masses need numerical protection. A minimal logit check uses \(p=(0.6,0.4)\), \(q=(0.5,0.5)\) and \(s_a=e_a-p\): the uncentered mean is \(q-p=(-0.1,0.1)\); subtracting this mean makes the expected centered score zero.

Match the sampler’s actual distribution

The lab fixes temperature 1, top-p 1, top-k -1 and selects --pg-loss-type reinforce, --disable-grpo-std-normalization, --calculate-per-token-loss. SC cannot be substituted into GSPO/CISPO/PPO clipping. Streaming, constrained decoding, sampling penalties, and per-request temperature changes are unsupported here.

slime also supports recorded top-p support \(S\): normalize trainer probabilities on \(S\) and use all recorded \(q_v\) on that support. Then \(\mu_w=\sum_{v\in S}q_vw_v\nabla\log p^S_v\) needs no tail approximation. This is exact relative to the recorded support, not to an untruncated policy; storage can greatly exceed top-128. See sampling and compatibility details.

What to measure

Hold prompt data, seed, response budget, and evaluation fixed. Begin with the short baseline, change one mechanism at a time, and record reward alongside train_rollout_logprob_abs_diff, tis_clipfrac when applicable, and async staleness/mean / staleness/max. Smaller mismatch alone does not prove better learning. Save a fixed batch for debug replay to distinguish optimizer behavior from sampling variation.