RL Post-Training — Concepts & Diagnostic Guide

Core Concepts

RL post-training optimizes a language model's policy using reward signals. The standard pipeline:

prompt → generate completions → score with reward → compute advantages → policy gradient update

Each stage has distinct failure modes. When a model "shows no improvement," the bug could be anywhere in this pipeline.

Diagnostic Methodology

When RL training produces no improvement, work through these stages in order. Each stage depends on the previous one being correct.

Stage 1: Verify Reward Signal

• Are rewards non-constant? If all rewards are identical, there is no learning signal.
• Do rewards correlate with completion quality? Spot-check decoded completions against their scores.
• Is the reward function being called on the correct text? Check that decoding/stripping preserves the content the reward function needs to evaluate.

Stage 2: Verify Advantage Computation

• Are advantages non-zero when rewards vary? If they collapse to ~0, the policy gradient vanishes.
• Check the magnitude and dtype of every numerical-stability constant in the advantage path (additive epsilons, clipping bounds). Compare each to what the math requires.
• Check the group size G. G ≤ 2 makes std either undefined or extremely noisy.

Stage 3: Verify Log-Probability Computation

• Verify bounds: log-probs of valid tokens must be non-positive.
• Compare your implementation against F.log_softmax on a small deterministic input — a numerical match rules out sign errors, wrong gathering axis, and off-by-one subtraction.
• On a near-one-hot input, confirm the dominant token's log-prob is close to 0 (not close to the min).

Stage 4: Verify Loss Computation

• Is the loss changing across steps? Flat loss suggests zero gradients upstream.
• Log the KL term and the policy-gradient term separately. Either dominating the other is diagnostic.
• Check the fraction of clipped samples. Near-100% clipping means the clip range is starving the signal.

Stage 5: Verify Generation and Decoding

• Are padding and decoder artefacts stripped from the text the reward function sees?
• Does every completion shape your model can emit survive the decoding path with non-empty output where a human would expect non-empty output? Include degenerate cases (no formatting markers, only a prefix, only a suffix) in the round-trip test.
• Print or log a sample of the actual strings handed to the reward function — mismatches between "what the model generated" and "what the reward saw" are often visible at a glance.

Fixing, Not Rewriting

A bug in a branched function is a bug in one branch, not a verdict on the whole function. Before editing:

• Enumerate the input shapes the function handles today and the output each branch produces. The branches exist because callers rely on them.
• Identify which input-output pairs violate the intended contract. Those are the only branches you need to change.
• If your diff collapses a multi-branch function to a one-liner, you have almost certainly broken a contract a different caller depends on. Re-read the call sites before committing.

The same principle applies to epsilon values, sign conventions, and clipping bounds — if a constant looks wrong, replace it with a correct constant; don't remove the surrounding numerical-stability logic.

Key Mathematical Invariants

These invariants should always hold. If any is violated, there is a bug.

Common Pitfall Categories

Detailed pitfall catalog with examples is in references/common-pitfalls.md. Summary:

1. Sign errors in log-space — log-softmax, KL divergence, DPO log-ratio
2. Numerical stability constants out of range — additive epsilons, clip bounds, temperature
3. String processing in decoding — cleanup that blanks shapes the rules didn't anticipate
4. Reward / decoding format mismatch — reward function sees text with the shape it needs removed
5. Reference model drift — reference model not frozen, or gradients flowing through it
6. Gradient flow breakage — detached tensors, in-place ops, .item()/numpy round-trips in the loss path
7. Configuration misuse — SFT-scale LR, KL coefficient dominating, clip range too tight, G too small

Available References