Debug GRPO / RL Training Code

The Round

What You Walk Out With

• All three bugs identified and fixed
• Code that runs without NaN errors and produces a non-trivial training signal
• Verbal explanation of each bug's root cause + fix
• Answers to RL-theory follow-ups about ratio clipping and on-policy/off-policy drift
• A diagnostic — typically `print(ratio.mean())` — that surfaces the on-policy follow-up cleanly

Problem Statement

Prerequisites

Canonical Solution

Process Playbook

Bug Catalogue

# Before:
action = torch.multinomial(logits, 1)

# After:
probs = F.softmax(logits, dim=-1)
action = torch.multinomial(probs, 1)

# Cleaner:
action = torch.distributions.Categorical(logits=logits).sample()

# Before:
advantage = (returns - returns.mean()) / returns.std()

# After:
advantage = (returns - returns.mean()) / (returns.std(correction=0) + 1e-8)

# Before:
ratio = new_logprob - old_logprob   # WRONG — this is the log-ratio, not the ratio

# After:
ratio = torch.exp(new_logprob - old_logprob)
# Sanity check: on the first mini-epoch step π_new == π_old, so ratio == 1 exactly:
# assert torch.allclose(ratio, torch.ones_like(ratio)) when step == 0

Follow-up Arc

Alternative Approaches

Observed Variants

• Base: debug NaN errors in GRPO code (softmax + epsilon fixes)
• With three explicit bugs to find
• With deep follow-up on ratio clipping theory
• With follow-up: why ratio is not strictly 1 at runtime (not strictly on-policy)

Pitfalls

• Forgetting to apply softmax before multinomial sampling — `torch.multinomial` raises a RuntimeError on any negative logit, and silently samples the WRONG distribution if logits are all non-negative (it does NOT NaN-explode). Easy to miss if your test env happens to produce non-negative logits early.
• Computing advantage as `(returns - baseline) / std` without an epsilon — division by zero on early rollouts where all returns are similar (sparse-reward envs, early training). Note: in GRPO, the group must have size > 1 — PyTorch's default `unbiased=True` std of a size-1 group is itself NaN, which epsilon won't save; use `correction=0` or guarantee group size > 1.
• Computing the importance-sampling ratio as a log-difference instead of `exp(log-difference)` — the model still 'trains' but the surrogate objective is wrong, the clip becomes meaningless, and updates are biased. The interviewer's first follow-up question is specifically about this: 'if ratio = model_logprob − old_logprob, could the model still train? Why not?'
• Assuming the script is on-policy and never printing the ratio — misses the entire mini-epoch follow-up. The whole point of the third follow-up is that 'on-policy' GRPO with multiple gradient steps per rollout batch ISN'T strictly on-policy after step 1.
• Spending too much time reading top-to-bottom before running — the first run reveals 2/3 bugs immediately. Read for ~5 minutes for context, then run.
• Answering 'why do we clip?' with only 'to prevent large updates' — the great answer goes further: clipping is a trust-region HEURISTIC that removes the policy-improvement incentive when the per-action ratio leaves `[1−ε, 1+ε]`. It does NOT analytically bound KL(new || old) — that's why PPO has a separate adaptive-KL-penalty variant and why GRPO/DeepSeek-Math add an explicit KL term to a reference policy on top of the clip. Whether the gradient drops to zero depends on the advantage sign and which branch of the `min()` is active per sample, not just on whether the ratio is in the clip band — so 'clipping zeros the gradient' is too strong a generalization.

What Each Level Looks Like

Insider Notes