Around the Sampler: Tangent, Remainder and KL

The figures step through like slides. Use → or ▶, and step back with ←.

1. One off-policy gap

Every token is drawn by one policy and trained through another.

The levers that buy throughput (fast kernels, asynchrony, larger batches) open a gap between $\ps$, which drew the token, and $\pt$, which takes the gradient. The gap has two sources: engine mismatch, which is present even at identical weights, and staleness. Part 2 models both. The gap sets the bias of the update, the batch sets its variance, and the step size multiplies both:

This post is about the bias: what the update should be when the data came from $\ps$. From here on, $\mathbb{E}$ and $\mathrm{Cov}$ are over rollouts drawn by the sampler.

2. Gain = tangent + remainder

Measured from the sampler, the gain of an update is exactly a tangent plus a remainder.

With $\ell_t = \log\frac{\pt(y_t)}{\ps(y_t)}$ and $\ell = \sum_t \ell_t$, reweighting whole rollouts gives the gain exactly:

At each prefix, the mean of $\ell_t$ under the sampler is $-\mathrm{KL}_t$. So $\ell$ splits into a centered first-order sum minus a second-order sum of token KLs. The first-order part gives the tangent:

This is the covariance loss. Computing its value needs the ratio, but its gradient needs only the trainer. Whatever the tangent misses is the remainder, and that remainder is the bias:

The bound, in three facts.

1. Both weights are non-negative.
2. Both have the same mean, $\mathbb{E}[e^{\ell}-1-\ell] = \mathbb{E}[\sum_t \mathrm{KL}_t] = \mathrm{KL}_{\text{seq}}$.
3. The reward lies in $[0, R_{\max}]$.

So each expectation lies in $[0, R_{\max}\mathrm{KL}_{\text{seq}}]$, and therefore

The bound holds at any distance from the sampler and has no horizon factor. Per-position trust-region bounds are governed by the worst token instead. For a 4,096-token rollout with worst-token KL $10^{-4}$, sequence KL 0.01 and reward scale 1:

For very long rollouts the sequence KL can exceed one, and then no bound that ignores the reward certifies anything. I use the bound for the shape of a good step, not as a certificate.

3. The best step

The lower edge of the band has a closed-form top, and it is never worse than the sampler.

The last frame of the figure above shows the lower bound $\text{gain} \ge \text{tangent} - R_{\max}\mathrm{KL}_{\text{seq}}$. It is concave in the trainer’s logits and zero at the sampler. Its maximizer reweights the sampler at every prefix by one plus the advantage $A$ over $R_{\max}$:

Since the bound is zero at the sampler, this step is never worse than the sampler. Among the changes a network can make, the natural-gradient step comes closest to it. A KL budget with multiplier $\lambda$ replaces $R_{\max}$ by $R_{\max}+\lambda$: a smaller step in the same direction.

4. Score centering

The gradient of the tangent raises the drawn token and lowers every token by how often the sampler draws it.

Averaged over what the sampler draws, the raises and lowerings cancel, so a constant reward moves nothing. The plain score averages to zero under the trainer, not under the sampler, so without centering even a constant reward pushes the policy. That is the off-policy bias in its simplest form. Score centering has an independent public treatment in Marek & Ryabinin (2026); here it is simply the gradient of the tangent.

5. Per-token importance sampling: same tangent, more variance

PPO and GRPO multiply by a ratio where the covariance loss subtracts a mean, and multiplication costs variance.

At the sampler it has the same tangent, so to first order it takes the same step. But a token that the sampler rarely draws and the trainer now prefers gets a large ratio, and the estimate jumps whenever that token appears. The $\chi^2$ divergence, which is dominated by such tokens, can be far above the KL, and the KL is all the bound controls. Clipping tames the jumps, but gives up first-order correctness on exactly those tokens.

6. One knob: how hard to pull toward the sampler

Every method is the tangent plus some multiple $\beta$ of a pull toward the sampler.

• Pull. Add $\beta\,\mathrm{KL}_{\text{seq}}$ to the loss. Its gradient is the mean score, which points toward the sampler and ignores the reward. The sequence KL is a divergence between the state distributions of the two policies (the state is the prefix), as in Divergence-Augmented Policy Optimization (NeurIPS 2019).
• Reject. Drop a rollout whose sequence KL is too large, using $M(y) = \mathbf{1}[\sum_t \mathrm{KL}_t \le \tau]$. This is Trust Region Masking, which ships in veRL. Inside the mask the estimator is still first-order correct. The threshold must sit above the engine-mismatch floor.

$\beta$ is a stability–plasticity knob: it sets how much the learner trusts where its data came from. How $\beta$, the step size and the queue should move as the run grows is the subject of Part 2.

References

• J. Liu, Y. Li, et al. When Speed Kills Stability: Demystifying RL Collapse from the Training-Inference Mismatch. Blog, 2025.
• Y. Li, J. Liu, et al. Trust Region Masking for Long-Horizon LLM Reinforcement Learning. ICML 2026.
• Q. Wang, Y. Li, J. Xiong, T. Zhang. Divergence-Augmented Policy Optimization. NeurIPS 2019.
• M. Marek, M. Ryabinin. Score Centering Stabilizes Off-policy Reinforcement Learning. arXiv, 2026.
• J. Schulman, S. Levine, P. Abbeel, M. Jordan, P. Moritz. Trust Region Policy Optimization. ICML 2015.

Citation

@online{li-2026-around-the-sampler,
  title  = {Around the Sampler: Tangent, Remainder and KL},
  author = {Li, Yingru},
  year   = {2026},
  month  = oct,
  url    = {https://richardli.xyz/post/offpolicy-around-the-sampler/}
}