Lecture 4: RL Implementation & Practice

Lecture 3 was the math: policy gradient theorem, REINFORCE, PPO, GRPO.

This lecture: how to actually implement, debug, and run RL training for LLMs.

The hardest bugs aren’t math errors — they’re silent implementation mistakes: wrong masking, stale caches, shape mismatches.

Lecture 3 was the math: policy gradient theorem, REINFORCE, PPO, GRPO.

This lecture: how to actually implement, debug, and run RL training for LLMs.

The hardest bugs aren’t math errors — they’re silent implementation mistakes: wrong masking, stale caches, shape mismatches.

A reminder on notation: As in Chapter 6, we use (s, a) from the reinforcement learning literature and (x, y) when prompt-completion notation is more natural. The (s, a) framing reflects the token-level gradient computation; (x, y) reflects the sequence-level reward. Both perspectives appear throughout.

The objective and its gradient:

The gradient says: for each token, compute the direction that makes it more likely (\nabla \log \pi), then scale by how good it was (\Psi_t).

All methods minimize the same family of losses (note the leading minus signs) — they differ in \Psi_t and how updates are bounded:

Where \rho_t = \frac{\pi_\theta(a_t \mid s_t)}{\pi_{\theta_\text{old}}(a_t \mid s_t)} is the importance-sampling ratio. PPO: per-token advantage A_t via GAE. GRPO: per-token ratio \rho_{i,t} but sequence-level advantage \hat{A}_i = \frac{R_i - \mu}{\sigma}.

From lecture 3, the policy gradient derivation showed:

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).mean()
loss.backward()  # autodiff gives ∑ Ψ_t ∇ log π

Every loss function in this lecture is a variation on this pattern.

The fundamental building block: per-token log-probabilities from the policy. Here’s every step explicitly:

# 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.

In practice, this is usually wrapped in a helper. From the book’s code:

def compute_log_probs(model, input_ids, attention_mask):
    logits = model(input_ids=input_ids,
                   attention_mask=attention_mask).logits
    logits = logits[:, :-1, :].to(torch.float32)
    log_probs = F.log_softmax(logits, dim=-1)
    targets = input_ids[:, 1:].unsqueeze(-1)
    return torch.gather(log_probs, dim=-1,
                        index=targets).squeeze(-1)

The shift, gather, and squeeze are the same — just condensed. It returns log-probs for all positions (prompt + completion); masking happens later. At rollout time, call this for the old policy and reference model (cache the results). During training, call it again for the current policy (recomputed each step).

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()

That’s it. Advantages weight the log-probabilities. Positive advantage → increase probability. Negative → decrease. (How we reduce per-token losses to a scalar — .mean() here — turns out to matter more than you’d expect. We return to this in the loss aggregation section.)

Generate K completions per prompt, compute leave-one-out baselines:

# rlhf_reward: (B*K,) flat tensor of rewards
# Prompt-major layout: K sibling completions stay together
rlhf_reward = rlhf_reward.view(-1, rloo_k)  # (B, K)

# Leave-one-out baseline: avg of other K-1 rewards per prompt
baseline = (rlhf_reward.sum(dim=1, keepdim=True) - rlhf_reward) / (rloo_k - 1)

advantages = rlhf_reward - baseline        # (B, K)
advantages = advantages.reshape(-1)        # (B*K,)

The rest follows standard policy gradient — multiply advantages by log-probs.

Note: the data loader must generate K completions per prompt and group them together. The per-prompt baseline helps reduce variance. Standard REINFORCE computes the baseline across all the states in the batch.

The clipped surrogate loss (minimized):

where \rho_t = \frac{\pi_\theta(a_t \mid s_t)}{\pi_{\theta_\text{old}}(a_t \mid s_t)} is the importance-sampling ratio.

Generalized Advantage Estimation (GAE):

GAE gives per-token advantages by propagating Temporal-Difference (TD) errors backward with exponential decay. \lambda = 0 is pure TD (low variance, high bias); \lambda = 1 is Monte Carlo (high variance, low bias).

PPO adds significant complexity over REINFORCE/RLOO:

• Lower variance: per-token credit assignment via GAE and a learned value function
• Token-level credit: each token gets its own advantage signal, not just a sequence-level reward
• Sample reuse: multiple gradient steps per batch via importance sampling

The cost: four models in memory (policy, value, reference, RM), fragile value function initialization, and more hyperparameters to tune.

PPO in the late 2010s and early 2020s was by far the most developed and understood RL algorithm, used widely across RL domains. In tasks other than language models, PPO was far superior to REINFORCE in performance.

for each batch:
  1. Sample prompts from dataset
  2. Generate completions with current policy π_θ
  3. Score with reward model → per-sequence rewards
  4. Compute ref model log-probs → per-token KL penalty
  5. Shape rewards: r_t = r_t - β * KL_t (per token)
  6. Compute returns via backward pass (GAE or MC)
  7. Compute advantages: A_t = returns_t - V(s_t)
  8. For k = 1 to K epochs on this batch:
     - Compute ratio = π_θ(a_t|s_t) / π_old(a_t|s_t)
     - Policy loss: clipped surrogate objective
     - Value loss: MSE on returns
     - total_loss = policy_loss + vf_coef * value_loss
     - Backward + optimizer step
  9. Sync π_old ← π_θ

The reward model gives one scalar R(x, y) at the end of a completion. Per-token rewards are shaped via KL:

Where \text{KL}_t = \log \pi_\theta(a_t \mid s_t) - \log \pi_\text{ref}(a_t \mid s_t).

These per-token rewards feed into GAE, which propagates credit backward to assign per-token advantages.

Before training starts, the rollout phase must compute and store everything the loss needs:

# rewards: (B,) terminal only (KL shaping omitted), values_old: (B, L) from rollout
B, L = completion_mask.shape
advantages = torch.zeros_like(values_old)
next_v = torch.zeros(B, device=values_old.device)
gae = torch.zeros(B, device=values_old.device)

last_idx = completion_mask.long().cumsum(-1).argmax(-1, keepdim=True)
done_mask = (torch.arange(L, device=values_old.device).unsqueeze(0) >= last_idx).float()
rewards_t = torch.zeros_like(values_old).scatter_(-1, index=last_idx, src=rewards)

for t in reversed(range(L)):
    not_done = 1.0 - done_mask[:, t]
    delta = rewards_t[:, t] + gamma * not_done * next_v - values_old[:, t]
    gae = delta + gamma * lam * not_done * gae
    advantages[:, t] = gae
    next_v = values_old[:, t]

advantages = advantages * completion_mask
targets = (advantages + values_old).detach()
advantages = advantages.detach()

Simplified: terminal reward only. For KL-shaped rewards, add -\beta \cdot \text{KL}_t per token before this loop.

The clipped surrogate objective:

# Compute probability ratio
ratio = torch.exp(new_per_token_logps - per_token_logps)  # (B, L), per_token_logps cached from rollout

# Clipped surrogate objective
eps = 0.2  # clip range
pg_losses1 = -advantages * ratio
pg_losses2 = -advantages * torch.clamp(ratio, 1.0 - eps, 1.0 + eps)
pg_loss = torch.max(pg_losses1, pg_losses2)

torch.max selects the more pessimistic (conservative) gradient. Because we minimize a negative loss, this prevents over-committing to any single update.

PPO (and GRPO) optionally reuse each rollout batch for multiple gradient steps. Clipping activates whenever \pi_\theta has drifted from \pi_{\theta_\text{old}} — two mechanisms cause this:

PPO (and GRPO) optionally reuse each rollout batch for multiple gradient steps. Clipping activates whenever \pi_\theta has drifted from \pi_{\theta_\text{old}} — two mechanisms cause this:

Minibatching: split the rollout batch into smaller minibatches to allow a larger, total batch size fit on a certain GPU setup. After updating on the first minibatch, \pi_\theta has changed — so later minibatches in the same epoch already see \rho_t \neq 1. Clipping can activate even with K = 1 epoch.

PPO (and GRPO) optionally reuse each rollout batch for multiple gradient steps. Clipping activates whenever \pi_\theta has drifted from \pi_{\theta_\text{old}} — two mechanisms cause this:

Minibatching: split the rollout batch into smaller minibatches to allow a larger, total batch size fit on a certain GPU setup. After updating on the first minibatch, \pi_\theta has changed — so later minibatches in the same epoch already see \rho_t \neq 1. Clipping can activate even with K = 1 epoch.

Multiple epochs: loop over the full batch K times to learn more from a given rollout (which can be expensive). Each pass sees a more-updated \pi_\theta, making ratios drift further. Typical K = 2–4; beyond \sim6 the policy is too far off-policy.

With K = 1 and no minibatching: \pi_\theta = \pi_{\theta_\text{old}}, ratios are always 1, and clipping never activates — PPO reduces to vanilla policy gradient with GAE.

Both PPO and GRPO use this same sample-reuse structure.

The critic V_\phi learns to predict returns — the total discounted reward from each token onward. GAE gave us advantages by combining actual rewards with the critic’s own predictions:

\hat{G}_t is a better estimate of the true return than V_\phi currently produces. We use it as the regression target to improve the critic:

targets = (advantages + values_old).detach()

The .detach() is critical — targets are fixed from the rollout, not something we backpropagate through.

The value function has its own clipping — same idea as the policy clip, but easy to overlook. It prevents the critic from jumping too far from its rollout-time predictions in a single update:

# Current critic predictions
v_pred = value_net(completions)  # (B, L)

# old_values: critic predictions from rollout time (before any training updates, gradient detached)
# Clamp new predictions to stay within eps of the rollout values
v_clip = torch.clamp(v_pred, old_values - eps, old_values + eps)
vf_unclipped = 0.5 * (v_pred - targets) ** 2
vf_clipped   = 0.5 * (v_clip - targets) ** 2
vf_loss = torch.max(vf_unclipped, vf_clipped)

torch.max picks the worse (more conservative) loss — if the unclipped prediction is already close to the target, the clipped version won’t interfere.

Combined loss: policy + value (KL enters via reward shaping, not as a separate loss term):

per_token_loss = pg_loss + vf_coef * vf_loss  # (B, L)

# Apply completion mask and aggregate
loss = ((per_token_loss * completion_mask).sum(dim=1) /
         completion_mask.sum(dim=1)).mean()

The vf_coef (typically 0.5–1.0) balances the two objectives.

Normalize advantages to zero mean, unit variance within the batch:

valid_adv = advantages[completion_mask.bool()]
advantages = ((advantages - valid_adv.mean()) /
              (valid_adv.std() + 1e-8)) * completion_mask

Why: stabilizes gradient magnitudes across batches. Without whitening, batches with uniformly high or low rewards can produce outsized gradients.

The value function V_\phi needs to produce reasonable estimates from the start:

• Initialize from RM backbone (InstructGPT convention): value predictions start near actual rewards
• Initialize from SFT model + random head: cheaper but early training unstable
• Cold-start issues: if initial value estimates are bad, GAE advantages are noisy → early training can be chaotic

Tülu 3 (Lambert et al., 2024) initializes from the reward model.

Detail: Many RL for LLM setups do “value function warmup” where they take training steps over data with measured rewards to help the value function initialize, so it is stable before taking policy steps.

Illustrative ranges from common LLM RLHF setups — not universal defaults:

In supervised learning, gradient noise is moderate — critical batch sizes are in the thousands. In RL, the gradient noise scale is orders of magnitude higher because gradients come from Monte Carlo rollouts, not labeled data.

• OpenAI Dota 2: critical batch size was in the millions of transitions (McCandlish et al., 2018)
• PPO specifically is not batch-size-invariant — clipping couples batch size to effective step size, so you can’t just compensate with learning rate (Hilton et al., 2021)
• PPO plateaus are often caused by noisy loss estimates that only resolve with more samples (Beukman et al., 2026)

Large batches reduce gradient variance proportional to 1/N. In RLHF, this is one of the cheapest ways to stabilize training — more effective than most hyperparameter tuning.

What to check during training:

• Clip fraction: percentage of tokens clipped — 0% means clipping never activates, >50% means policy is changing too fast
• Advantage statistics: mean should be near 0 (if whitened), std should be stable
• Value loss: should decrease — if not, value function isn’t learning
• KL divergence: should increase gradually, not explode

What to check during training:

• Clip fraction: percentage of tokens clipped — 0% means clipping never activates, >50% means policy is changing too fast
• Advantage statistics: mean should be near 0 (if whitened), std should be stable
• Value loss: should decrease — if not, value function isn’t learning
• KL divergence: should increase gradually, not explode

Common silent bugs — training runs but learns the wrong thing:

• Prompt tokens included in policy loss (inflates loss, wastes gradient on unchangeable tokens)
• Reward scattered to wrong position (e.g., [:, -1] instead of last generated token)
• Stale log-probs used as “current” (ratio stuck near 1, clipping never activates)
• Zero-std GRPO groups (all completions get same reward → NaN advantages)
• Value or policy loss computed outside completion mask (trains on padding)

For a 7B model with fp16:

~56 GB just for model weights before optimizer states, activations, or gradients. This is why PPO often requires model parallelism or offloading. Some implementations reduce this by sharing backbones (e.g., a value head on the policy network).

For each prompt, sample G completions and compute group-normalized advantages:

Then apply the same clipped loss as PPO (minimized) — per-token ratios but sequence-level advantages — plus an optional KL penalty (more in the reasoning lecture):

No value function, no GAE — advantages come entirely from comparing siblings within a group.

# Generate G completions per prompt, then compute group advantages
mean_r = rewards.view(-1, G).mean(dim=1)
std_r = rewards.view(-1, G).std(dim=1)
mean_r = mean_r.repeat_interleave(G)
std_r = std_r.repeat_interleave(G)
advantages = ((rewards - mean_r) / (std_r + 1e-4)).unsqueeze(1)

# Importance sampling ratio
ratio = torch.exp(new_logps - old_logps)   # (B*G, L)

# Clipped surrogate (same as PPO)
pg_losses1 = -advantages * ratio
pg_losses2 = -advantages * torch.clamp(ratio, 1 - eps, 1 + eps)
pg_loss = torch.max(pg_losses1, pg_losses2)

# KL penalty in loss (not in reward)
per_token_loss = pg_loss + beta * per_token_kl

loss = ((per_token_loss * mask).sum(dim=1) / mask.sum(dim=1)).mean()

What GRPO removes relative to PPO:

Significantly less code and ~1 fewer model copy in memory.

Ideal / Theory (on-policy): generate → update → generate → update

Each batch of completions is scored and used for a short update window (one or a few epochs), then discarded.

Reality (async): generation and training overlap on different GPU groups for better throughput. The model used for generation may be 1–N steps behind the training model.

Tradeoff: perfect on-policy is slow (GPUs idle during generation or training). Slight staleness is usually fine.

• A video I recorded looking at codebases implementing GRPO, DAPO, Dr. GRPO, and other papers.
• ~24min in, talk on scaling RL for Olmo 3.
• Finbarr Timbers’s blog post on making RL fast.

• TRL (Werra et al., 2020) — Hugging Face ecosystem, PPO/GRPO/DPO. Best starting point for getting started
• Open Instruct (Ivison et al., 2024) — Allen AI, multi-algorithm. Best for research and reproduction
• veRL — Very popular in the RLVR era.
• OpenRLHF — Started for RLHF work, popular with RLVR, etc. too.

• Log-space arithmetic: always use log_softmax + gather, never softmax then log — softmax squashes small probabilities to tiny floats, then log amplifies the precision loss
• Masking padding tokens: completion_mask must be 1 only for completion tokens — exclude prompt tokens, post-EOS padding, and the EOS token itself (or include EOS consistently). Multiply losses by this mask before aggregation

• Log-space arithmetic: always use log_softmax + gather, never softmax then log — softmax squashes small probabilities to tiny floats, then log amplifies the precision loss
• Masking padding tokens: completion_mask must be 1 only for completion tokens — exclude prompt tokens, post-EOS padding, and the EOS token itself (or include EOS consistently). Multiply losses by this mask before aggregation
• Stop-token handling: EOS tokens need consistent treatment — include in loss or not, but be consistent. Mishandling is a common silent error
• Sequence length handling: variable-length completions need careful normalization (next section)

• Log-space arithmetic: always use log_softmax + gather, never softmax then log — softmax squashes small probabilities to tiny floats, then log amplifies the precision loss
• Masking padding tokens: completion_mask must be 1 only for completion tokens — exclude prompt tokens, post-EOS padding, and the EOS token itself (or include EOS consistently). Multiply losses by this mask before aggregation
• Stop-token handling: EOS tokens need consistent treatment — include in loss or not, but be consistent. Mishandling is a common silent error
• Sequence length handling: variable-length completions need careful normalization (next section)
• Detaching baselines: advantages.detach() — don’t backpropagate through the advantage computation
• Division guards: anywhere you divide by mask.sum() or sequence lengths, use .clamp_min(1) or + eps to avoid NaN from empty completions (e.g. immediate EOS)

Every algorithm we’ve seen computes the same core gradient: \nabla \log \pi \cdot A. But how you aggregate per-token losses into a scalar changes training dynamics more than you’d expect.

The next section covers three strategies — per-sequence, per-token, and fixed-length normalization — and why the choice matters in practice.

Same algorithm, different aggregation → different training dynamics.

This is often measured by how the sequence length changes through training, especially in RLVR.

Each sequence contributes equally to the batch loss, regardless of length:

# Strategy 1: Per-sequence normalization
loss = ((per_token_loss * completion_mask).sum(dim=1) /
         completion_mask.sum(dim=1)).mean()

Standard in GRPO and some PPO implementations.

Each sequence gets equal weight → per-token gradients are inversely proportional to sequence length:

• Short sequence (5 tokens): each token gets gradient \propto 1/5 = 0.20
• Long sequence (10 tokens): each token gets gradient \propto 1/10 = 0.10

Short sequences have larger per-token gradients. This can bias the model away from lengthy responses.

Each token contributes equally across the entire batch:

# Strategy 2: Per-token normalization
loss = (per_token_loss * completion_mask).sum() / completion_mask.sum()

Used in DAPO (Yu & others, 2025). Longer sequences contribute proportionally more gradient.

All tokens get equal gradient magnitude. Longer sequences contribute more to the total gradient because they have more tokens.

Can bias toward verbose completions — the model gets more gradient signal from longer answers, which may encourage longer generations.

Normalize by a constant L_\text{max} (max generation length):

# Strategy 3: Fixed-length normalization
loss = ((per_token_loss * completion_mask).sum(dim=1) /
         L_max).mean()

From Dr. GRPO (Liu et al., 2025). Equalizes per-token scale while letting longer sequences contribute more total gradient (more active tokens in the sum).

seq_1_losses = [1, 1, 1, 1, 10]           # 5 tokens, mean = 2.8
seq_2_losses = [1, 1, 1, 1, 1, 1, 1, 1, 1, 10]  # 10 tokens, mean = 1.9

Per-sequence gives short sequences bigger per-token gradients. Per-token and fixed-length equalize.

Illustrative ranges — thresholds vary by model size, task, and algorithm:

Also monitor: eval scores on held-out benchmarks, and read sample outputs for coherence.

When RL runs go wrong, it is often due to RL latching onto spurious signals rather than the intended reward. These topics deserve their own lecture:

• KL penalties: forward vs. reverse KL, math and code
• Goodhart’s Law: when the proxy objective diverges from true quality
• Reward hacking: sycophancy, verbosity, over-refusal
• Implicit regularization: “SFT Memorizes, RL Generalizes”, “RL’s Razor”

Covered in depth in a future lecture on Chapters 14 & 15.

From math to running code — the implementation details that matter:

1. Core code patterns — log-probs, masking, the REINFORCE/RLOO/PPO/GRPO loss functions
2. Advantage estimation — per-prompt baselines (RLOO, GRPO z-score) vs per-token credit (GAE)
3. Loss aggregation — per-sequence vs per-token vs fixed-length normalization changes training dynamics
4. PPO’s complexity budget — value function targets, critic clipping, double regularization
5. Async training — actor/learner splits, staleness budgets, distributed infrastructure
6. Practical engineering — batch size, monitoring, debugging, training recipes

1. Introduction & training overview — Chapters 1–3
2. Instruction tuning, reward models, rejection sampling — Chapters 4, 5, 9
3. RL theory — Chapter 6 (Part 1)
4. RL implementation & practice — Chapter 6 (Part 2)
5. Reasoning — Chapter 7
6. Direct alignment algorithms — Chapter 8
7. …