Reinforcement Learning
Policy Gradient
The objective in reinforcement learning is to optimize a policy $\pi$ by maximizing the expected return $J(\pi)$.
\[\pi ^* = \arg \max \limits _{\pi} J (\pi)\]The return of a policy is defined over all possible trajectories $\tau$, which is a sequence of states and actions $(s_0, a_0, \dots, s_T, a_T)$:
\[J (\pi _{\theta}) = \int _{\tau} P(\tau | \pi) R(\tau) d \tau = \mathbb{E} _{\tau \sim \pi} \left[ R(\tau) \right]\]where
\[P(\tau | \pi) = \rho _0 (s_0) \prod \limits _{t=0} ^{T-1} P (s _{t+1} | s_t, a_t) \pi (a_t | s_t)\] \[R(\tau) = \sum \limits _{t=0} ^{T} \gamma ^t r_t\]where $\gamma \in [0, 1]$ is the discount factor and $r_t$ is the reward received at time $t$.
To maximize the return, we use stochastic gradient ascent to update the policy:
\[\theta _{k+1} = \theta _k + \alpha \nabla _{\theta} J (\pi _{\theta}) | _{\theta _k}\]Here, $\nabla _{\theta} J (\pi _{\theta})$ is known as the policy gradient.
\[\begin{array}{rcl} \nabla _{\theta} J (\pi _{\theta}) & = & \nabla _{\theta} \mathbb{E} _{\tau \sim \pi _{\theta}} \left[ R(\tau) \right] \\ & = & \nabla _{\theta} \int _{\tau} P(\tau | \theta) R(\tau) d \tau \\ & = & \int _{\tau} \nabla _{\theta} P(\tau | \theta) R(\tau) d \tau \\ & = & \int _{\tau} P(\tau | \theta) \nabla _{\theta} \log P(\tau | \theta) \cdot R(\tau) d \tau \\ & = & \mathbb{E} _{\tau \sim \pi _{\theta}} \left[ \nabla _{\theta} \log P(\tau | \theta) \cdot R(\tau) \right] \\ & = & \mathbb{E} _{\tau \sim \pi _{\theta}} \left[ \sum \limits _{t=0} ^T \nabla _{\theta} \log \pi _{\theta} (a_t | s_t) \cdot R(\tau) \right] \\ & \approx & \frac{1}{|\mathcal{D}|} \sum \limits _{\tau \in \mathcal{D}} \sum \limits _{t=0} ^T \nabla _{\theta} \log \pi _{\theta} (a_t | s_t) \cdot R(\tau) \end{array}\]Baseline
The policy gradient estimator is unbiased but often has high variance. To reduce this variance, we introduce a baseline $b(s_t)$ and subtract it from the return:
\[\nabla _{\theta} J(\pi _{\theta}) \approx \frac{1}{|\mathcal{D}|} \sum _{\tau \in \mathcal{D}} \sum _{t=0} ^T \nabla _{\theta} \log \pi _{\theta} (a_t|s_t) \left( R(\tau) - b(s_t) \right)\]Subtracting a baseline does not change the expectation because
\[\begin{array}{rcl} \mathbb{E} _{a \sim \pi _{\theta}} \left[ \nabla_{\theta} \log \pi_{\theta}(a|s) b(s) \right] & = & b(s) \mathbb{E} _{a \sim \pi _{\theta}} \left[ \nabla_{\theta} \log \pi_{\theta}(a|s) \right] \\ & = & b(s) \sum _{a} \pi_{\theta}(a|s) \nabla_{\theta} \log \pi_{\theta}(a|s) \\ & = & b(s) \sum _{a} \nabla_{\theta} \pi_{\theta}(a|s) \\ & = & b(s) \nabla_{\theta} \sum _{a} \pi_{\theta}(a|s) \\ & = & b(s) \nabla_{\theta} 1 \\ & = & 0 \end{array}\]However, it can significantly reduce the variance of the gradient estimator.
Let $\psi$ be the score,
\[\psi = \nabla _{\theta} \log \pi _{\theta} (a|s)\]and $g$ the gradient with the baseline subtracted,
\[g = \psi (R - b)\]So
\[\begin{array}{rcl} b^*(s) & = & \arg \min \limits _{b(s)} \mathrm{Var}(g) \\ & = & \arg \min \limits _{b(s)} \mathbb{E}[g^2] - (\mathbb{E}[g])^2 \\ & = & \arg \min \limits _{b(s)} \mathbb{E}[g^2] \end{array}\] \[\begin{array}{rcl} \frac{d\mathbb{E}[g^2]}{db} & = & \frac{d}{db} \mathbb{E} \left[ \psi ^2 (R - b)^2 \right] \\ & = & \mathbb{E} \left[ 2 \psi ^2 (b - R) \right] \end{array}\]Setting the derivative to zero gives the optimal baseline,
\[b^*(s) = \frac{\mathbb{E} \left[ \psi ^2 R \right]}{\mathbb{E} \left[ \psi ^2 \right]}\]This is actually a weighted average of $R$, where sensitive actions have larger impacts on the optimal baseline. In practice, $b(s)$ is often approximated by the state value function $V^{\pi}(s)$, leading to the advantage function formulation:
\[A^{\pi}(s,a) = R(\tau) - V^{\pi}(s)\]$R(\tau)$ includes rewards from the entire trajectory, but a decision should only consider its impacts on future outcomes, which is known as rewards-to-go or Q-function.
\[Q^{\pi}(s,a) = \sum \limits _{t'=t} ^T r(s_{t'},a_{t'})\]So
\[A^{\pi}(s,a) = Q^{\pi}(s,a) - V^{\pi}(s)\]Then the policy gradient becomes:
\[\nabla _{\theta} J(\pi _{\theta}) \approx \frac{1}{|\mathcal{D}|} \sum _{\tau \in \mathcal{D}} \sum _{t=0} ^T \nabla _{\theta} \log \pi _{\theta} (a_t|s_t) A^{\pi}(s_t, a_t)\]This is the foundation of Advantage Actor-Critic (A2C) and related algorithms.
Score Centering
The identity behind the baseline,
\[\mathbb{E} _{a \sim \pi _{\theta}} \left[ \psi \right] = 0\]held because the actions were drawn from the same $\pi _{\theta}$ the score is taken through. In LLM training that is usually false. Rollouts come from an inference engine and gradients from a training engine, and the two disagree through numerical precision, kernel choice, or stale weights.
Write $\mu$ for the policy that actually produced the actions. Then $\mathbb{E} _{a \sim \mu} \left[ \psi \right]$ is no longer zero, and the covariance identity splits the update at a state in two:
\[\mathbb{E} _{\mu} \left[ R \psi \right] = \mathbb{E} _{\mu} \left[ R \right] \mathbb{E} _{\mu} \left[ \psi \right] + \mathrm{Cov} _{\mu} \left( R, \psi \right)\]Only the covariance knows which action earned which reward. The first term is the drift: it sees $R$ through its mean alone, so it does no credit assignment. Its direction is
\[\mathbb{E} _{\mu} \left[ \psi \right] = \sum _{a} \mu (a|s) \nabla _{\theta} \log \pi _{\theta} (a|s) = \nabla _{\theta} \sum _{a} \mu (a|s) \log \pi _{\theta} (a|s)\]which is a supervised fit of $\pi _{\theta}$ to $\mu$, scaled by the mean reward. Fitting a fixed teacher would simply converge to that teacher. But $\mu$ is a biased copy of $\pi _{\theta}$ that gets refreshed from it, so the bias is learned, copied back, and learned again.
A baseline does not remove the drift. Centering the reward makes the advantage average to zero over a group of rollouts, but the drift lives at each state, and $\mathbb{E} _{\mu} \left[ A \mid s \right]$ is nonzero exactly where the learning signal is: a state that already leads to the right answer has positive expected advantage, and one that already contains a mistake has negative. Centering the reward shrinks the drift without cancelling it.
Since the problem is that $\mathbb{E} _{\mu} \left[ \psi \right] \neq 0$, subtract it:
\[\tilde{\psi} (a) = \psi (a) - \mathbb{E} _{\mu} \left[ \psi \right]\]The subtracted term is an expectation under the same $\mu$, so $\mathbb{E} _{\mu} [ \tilde{\psi} ] = 0$ and
\[\begin{array}{rcl} \mathbb{E} _{\mu} \left[ R \tilde{\psi} \right] & = & \mathbb{E} _{\mu} \left[ R \right] \underbrace{\mathbb{E} _{\mu} \left[ \tilde{\psi} \right]} _{0} + \mathrm{Cov} _{\mu} \left( R, \tilde{\psi} \right) \\ & = & \mathrm{Cov} _{\mu} \left( R, \psi \right) \end{array}\]where the last step uses that shifting $\psi$ by a constant leaves the covariance alone. The drift is cancelled exactly.
This mirrors the baseline. With an exact value function $b(s) = \mathbb{E} _{\mu} \left[ R \mid s \right]$, centering the reward does the same job,
\[\mathbb{E} _{\mu} \left[ \left( R - \mathbb{E} _{\mu} \left[ R \mid s \right] \right) \psi \right] = \mathrm{Cov} _{\mu} \left( R, \psi \right)\]but that needs a critic. Centering the score needs only $\mu (a|s)$, and for a language model the state is the prefix and the actions are the vocabulary, so $\mathbb{E} _{\mu} \left[ \psi \right]$ is one sum over next-token probabilities. On-policy a baseline changes the variance and not the mean. Off-policy, centering the score changes the mean.
It is not a full off-policy correction. On-policy the same identity gives $\mathbb{E} _{\pi _{\theta}} \left[ R \psi \right] = \mathrm{Cov} _{\pi _{\theta}} \left( R, \psi \right)$, so centering recovers the on-policy update except for which distribution the covariance is taken over. Importance sampling attacks that subscript instead, reweighting each action by $\pi _{\theta} (a|s) / \mu (a|s)$ to put the covariance back under $\pi _{\theta}$, at the cost of a ratio that can be arbitrarily large on rare actions. Clipping the ratio bounds the variance but breaks the correction and lets the drift back in, so the two are usually composed: the ratio pulls the sampling distribution toward $\pi _{\theta}$, and centering removes the drift left behind by the clipping. Numerical mismatch keeps $\mu$ close enough that centering alone is enough, while badly stale rollouts need both.
Implementation
Computing $\mathbb{E} _{\mu} \left[ \psi \right]$ exactly means summing over the whole vocabulary, so the sampler would have to log its full next-token distribution at every generated token, which is far too much to store. Write $\pi _v$ and $\mu _v$ for the trainer’s and sampler’s probabilities of token $v$ at this state, and let $H$ be the sampler’s top $k$ tokens. Log only those, and model the tail with the trainer’s own distribution rescaled to match the sampler’s tail mass:
\[\hat{\mu} _v = \begin{cases} \mu _v & v \in H \\ \rho \, \pi _v & v \notin H \end{cases} \qquad \rho = \frac{1 - \sum _{v \in H} \mu _v}{1 - \sum _{v \in H} \pi _v}\]Substituting $\hat{\mu}$ leaves a head term over $k$ tokens and a tail term over everything else. Only the tail still ranges over the full vocabulary, but the same identity holds for the trainer at this state, $\sum _v \pi _v \psi _v = 0$, so its head and tail sums are negatives of each other,
\[\begin{array}{rcl} \mathbb{E} _{\hat{\mu}} \left[ \psi \right] & = & \sum \limits _{v \in H} \mu _v \psi _v + \rho \sum \limits _{v \notin H} \pi _v \psi _v \\ & = & \sum \limits _{v \in H} \mu _v \psi _v - \rho \sum \limits _{v \in H} \pi _v \psi _v \\ & = & \sum \limits _{v \in H} \left( \mu _v - \rho \, \pi _v \right) \psi _v \end{array}\]so the tail is never touched and the correction is a sum over $k$ terms. Since $\psi _v = \nabla _{\theta} \log \pi _v$, the correction is the gradient of a scalar and needs no explicit gradient vectors,
\[L = - R \left( \log \pi _a - \sum _{v \in H} \mathrm{sg} \left[ \mu _v - \rho \, \pi _v \right] \log \pi _v \right)\]where $a$ is the sampled action and $\mathrm{sg}$ is a stop-gradient, applied to the coefficients only.
import jax.numpy as jnp
from jax.lax import stop_gradient
def score_centering_loss(train_logp, samp_logp, topk_ids, sampled_token, advantage):
train_head_logp = train_logp[topk_ids]
tail_mass_ratio = (1 - jnp.exp(samp_logp).sum()) / (1 - jnp.exp(train_head_logp).sum())
head_prob_residual = jnp.exp(samp_logp) - tail_mass_ratio * jnp.exp(train_head_logp)
logp_correction = (stop_gradient(head_prob_residual) * train_head_logp).sum()
return -advantage * (train_logp[sampled_token] - logp_correction)
So the rollout engine only has to return the top-$k$ token ids and their logprobs alongside the sampled token, and $k = 32$ already matches the exact correction.