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.