From Classical RL to Frontier LLMs

RL compute relative to pretraining compute for the same model, in GPU-hours:

How RL is approached with LLMs and what the challenges of RL are for frontier models.

Some examples we include:

1. Policy gradients in LLM code. How LLM log-probs, logits, and a surrogate loss turn REINFORCE into a few lines of PyTorch.
2. Asynchronous RL systems. How RL training is usually done for large LLMs.
3. Training-inference mismatch. A common challenge that many of the latest RL algorithms address.

p_v and q_v are shorthand for the probability of token v at one position, given x and the prefix a_{<t}. CS 4789 uses the same x and a.

Math RL needed prompts and answer checkers. Agentic RL needs environments: containers, real file systems, services, verification tests (Ivison et al., 2026). Building worlds for LLMs is hard.

The objective and its gradient:

The gradient says: for each sampled token, compute the direction in parameter space that makes it more likely (\nabla \log \pi), then scale by how much better than expected it was (\Psi_t).

The policy gradient derivation gives:

So we write a loss whose negative gradient is the policy gradient, treating each advantage A_i as a constant across the trajectory:

Gradient descent on L is gradient ascent on J.

In code, we compute the log-probs and let autodiff handle the gradient:

seq_log_probs = (token_log_probs * completion_mask).sum(dim=-1)
loss = -(seq_log_probs * advantages.detach()).mean()
loss.backward()  # fills each parameter's .grad with ∇L
optimizer.step()  # θ ← θ - η ∇L

loss.backward() runs backpropagation: it computes \partial L / \partial \theta for every parameter.

optimizer.step() applies the update.

The loss values themselves are often weird in RL (e.g. zero-centered distributions, rather than nice decreasing accuracy metrics of pretraining)

Most RL algorithms used today are variations on this computation.

The fundamental building block: per-token log-probabilities from the policy.

# Forward pass through the model
logits = model(input_ids).logits              # (B, L, vocab_size)

# Autoregressive shift: logit at position t predicts token at t+1
logits = logits[:, :-1, :]                    # (B, L-1, vocab_size)
labels = input_ids[:, 1:]                     # (B, L-1)
completion_mask = completion_mask[:, 1:]       # (B, L-1)

# Per-token log-probs
log_probs = logits.log_softmax(dim=-1)
token_log_probs = log_probs.gather(dim=-1, index=labels.unsqueeze(-1)).squeeze(-1) # (B, L-1)

# Per-sequence log-prob: sum over completion tokens
seq_log_probs = (token_log_probs * completion_mask).sum(dim=-1)

Everything in log-space to avoid numerical underflow from multiplying many small probabilities.

The simplest policy gradient loss:

# rewards: (B,) — one reward per sequence
# seq_log_probs: (B,) — sum of log-probs over completion tokens

# Baseline: average reward in the batch
baseline = rewards.mean()
advantages = rewards - baseline

# REINFORCE loss (negative because we minimize)
loss = -(advantages * seq_log_probs).mean()

TPositive advantage → increase probability. Negative → decrease.

(Aside: How we reduce per-token losses to a scalar — .mean() here — turns out to matter more than you’d expect. More in the book and lecture video.)

Policy-gradient theory is built on the assumption of on-policy data – from the current model weights in the training framework of the LLM.

As models get bigger, this is hard because: sampling is expensive, sampling is dominated by long-trail generations, most people use separate systems for inference v training, moving large model weights around is slow, inference weights can be further quantized for speed relative to training weights, multi-cluster deployments, etc.

# 1. Sample completions on the inference engine (pseudocode)
input_ids, completion_mask = sampler.generate(prompts, n=G)

# 2. Score them, then center within each group
rewards = verifier(prompts, input_ids)                  # (B*G,)
advantages = group_center(rewards, G)                   # (B*G,)

# 3. Recompute log-probs on the trainer, this time with gradients
logits = model(input_ids).logits[:, :-1, :]
token_logprobs = logits.log_softmax(-1).gather(
    -1, input_ids[:, 1:].unsqueeze(-1)).squeeze(-1)
seq_logprobs = (token_logprobs * completion_mask[:, 1:]).sum(-1)

# 4. Update, then sync the new weights back to the sampler
(-(advantages * seq_logprobs).mean()).backward()
optimizer.step()
sampler.load_weights(model)

Steps 1 and 3 ask two different systems about the same tokens. Do they agree?

At one position, with prompt x and prefix a_{<t}, each system defines a next-token distribution:

The gradient we want samples from the learner:

Reweighting by the sequence ratio recovers the gradient exactly (Marek & Ryabinin, 2026):

In practice, LLM RL systems apply TIS as a per-token correction weight on the policy-gradient loss:

# Shape: (B*G, L)
C = 2.0  # TIS cap

logratio = learner_logprobs - sampler_logprobs
logratio = logratio.clamp(-10.0, 10.0)              # numerical safety
tis_weight = torch.exp(logratio).clamp(max=C)        # one-sided truncation

# Use as a fixed correction weight on the per-token PG loss
per_token_pg_loss = per_token_pg_loss * tis_weight.detach()

TIS for the learner–sampler ratio has been adopted across major open-source RL frameworks (VeRL, TRL, OpenRLHF, SkyRL, OAT, and Open Instruct, which uses C = 2).

The policy gradient for one completion:

At one position, the expected score under the learner, written as a sum:

Tokens now come from q, but the score is still computed through p. Define the expected score under the sampler:

Start from the estimated gradient, and move the sum outside the expectation:

Use the covariance identity, \mathbb{E}[XY] = \mathbb{E}[X] \, \mathbb{E}[Y] + \text{Cov}(X, Y), starting on-policy, the expectation of the gradient:

The drift, \mathbb{E}_q[A] \, \bar{\boldsymbol{\xi}}, depends on the advantages only through their mean, so it moves the weights the same way regardless of which token was sampled.

We don’t want that: in practice, it can pull the learner’s weights toward a stale and/or quantized sampler, causing potentially compounding instabilities.

Baselines center the reward. Score centering centers the per-token log-prob gradients (called “score” in statistics), subtracting its expected value under the sampler:

With the centered score, the policy-gradient estimate becomes:

Apply the covariance identity to the inner, per-token term:

• Score centering removes the drift, the pull toward the sampler that carries no learning signal.
• It does not make training on-policy: the signal is still \text{Cov}_q(A, \boldsymbol{\xi}_{a_t}), measured under the sampler, not the learner.
• Importance sampling (e.g., TIS) reweights samples toward the learner, but clipping the ratios lets some drift back in.
• They fix different problems and are composable. The score-centering paper recommends using both!

(Supposedly) Many labs are supposedly using StableReinforce, the loss Tinker added, rather than GRPO. This is unconfirmed, but the two are related.

The group advantages are nearly the same. GRPO normalizes each reward by the group mean and standard deviation:

Tinker’s docs describe the behavior, but not the main loss:

“In our internal post-training, this loss function has led to extremely stable training, even under highly off-policy async RL and low precision numerics.”

The plot shows gradients with respect to the logits z. For a softmax, the gradient of \log p_v is the one-hot vector e_v minus the probabilities:

Ripple-65B Preview, a 65B open-weight model from Ricursive (built on a lot of Olmo 3 infra!), names score centering in its post-training:

“Our mixed RL recipe uses truncated importance sampling, score centering, and linear length penalty with high-quality long-horizon synthetic environments.”

The ratio between the two systems for a sampled token, and its log:

Tokens are sampled from q, so take the expectation under q:

Flip the sign to recognize the KL divergence:

If the systems agree, every \delta_t is zero. Otherwise, log-ratios are negative on average.

The ratio for a whole completion is the product of the token ratios:

Take the expectation under q, and apply the previous slide at each position:

This sum is the KL between the two sequence distributions, \text{KL}\big(q(a \mid x) \,\Vert\, p_\theta(a \mid x)\big), so it grows with length.

For illustration, a small per-token KL of 0.001 over 10{,}000 tokens gives a sequence KL of 10 nats. A typical weight is then near e^{-10} \approx 5 \times 10^{-5}.

In statistics, the score is the gradient of the log-likelihood with respect to the parameters. For one token v at one position:

• It is a vector the size of the parameters: the direction that makes token v more likely.
• The policy gradient is an advantage-weighted sum of the scores of the sampled tokens:

Distilling toward a fixed teacher is harmless. It converges to the best fit of that teacher.

In online RL, the teacher moves with the student:

• The sampler is a biased copy of the learner, e.g., quantized or stale.
• Each step pulls the learner toward the sampler’s errors.
• The new weights are synced back to the sampler, carrying those errors with them.
• The next sampler is biased again, so the error compounds instead of converging.
• The bias does not average out with more samples, since \bar{\boldsymbol{\xi}} is fixed at each prefix. It needs a direct correction.

Not fully. Group centering makes advantages sum to zero over the rollouts of a prompt:

But drift arises at each prefix, and \mathbb{E}_q[A \mid x, a_{<t}] is generally not zero there:

• A prefix likely to reach a correct answer has positive expected advantage.
• A prefix that already contains a mistake has negative expected advantage.

Group centering shrinks drift relative to all-positive rewards, and it was the most stable reward setting in the paper’s online experiments. It does not eliminate drift.

We want a scalar whose gradient is A \, (\boldsymbol{\xi}_{a_t} - \bar{\boldsymbol{\xi}}).

The first piece is already a gradient, \boldsymbol{\xi}_{a_t} = \nabla_\theta \log p_{a_t}. For the second:

The sampler’s probabilities are logged values, so we mark them as constants with a stop-gradient and pull the gradient out:

Combine the two pieces, and multiply by -A:

The correction is one extra term in the REINFORCE loss from earlier:

# Learner: log-probs over the whole vocabulary, with gradients
learner_vocab_logprobs = logits.log_softmax(-1)              # (B*G, L, V)
learner_logprobs = learner_vocab_logprobs.gather(
    -1, tokens.unsqueeze(-1)).squeeze(-1)                    # (B*G, L)

# Sampler: logged at generation time, so it has no gradient path
sampler_probs = sampler_vocab_logprobs.exp()                 # (B*G, L, V)

# The gradient of `centering` is the expected score, s-bar
centering = (sampler_probs * learner_vocab_logprobs).sum(-1) # (B*G, L)

# advantages: (B*G, 1), broadcast to every token
per_token_loss = -advantages * (learner_logprobs - centering)
loss = ((per_token_loss * mask).sum(1) / mask.sum(1)).mean()

No ratios, no masks, and no clipping. On-policy, the gradient of centering is zero.

Computing \bar{\boldsymbol{\xi}} exactly needs the sampler’s next-token distribution at every position:

• Qwen3’s vocabulary has 152K tokens, so one position in fp32 takes about 608KB.
• A batch of 1024 rollouts with 32K tokens each would take about 20TB.

Instead, log only the sampler’s head H, its k most likely tokens at that position, with k = 128. Model the tail with the learner, rescaled to carry the sampler’s missing mass:

Substitute \hat{q} into the expected score, and split the sum into head and tail:

The zero-mean identity splits the same way, which rewrites the tail in terms of the head:

Substitute, and collect the terms. Only the k head tokens remain:

The paper reports that k = 128 and k = 32 both perform on par with the full vocabulary, at under 1% wall-clock overhead for k = 128.

# Logged during generation: sampler_topk_logprobs, sampler_topk_ids (B*G, L, k)
learner_logprobs = learner_vocab_logprobs.gather(-1, tokens[..., None])[..., 0]
head_logprobs = learner_vocab_logprobs.gather(-1, sampler_topk_ids)  # (B*G, L, k)
learner_head_probs = head_logprobs.exp()
sampler_head_probs = sampler_topk_logprobs.exp()

# rho: the ratio of the sampler's tail mass to the learner's tail mass
sampler_tail = (1 - sampler_head_probs.sum(-1)).clamp_min(1e-6)
learner_tail = (1 - learner_head_probs.sum(-1)).clamp_min(1e-6)
rho = sampler_tail / learner_tail                                    # (B*G, L)

# detach() is required: coef contains the learner's own probabilities
coef = (sampler_head_probs - rho[..., None] * learner_head_probs).detach()
centering = (coef * head_logprobs).sum(-1)                           # (B*G, L)

per_token_loss = -advantages * (learner_logprobs - centering)