8. GRPO-Style Group-Relative Advantage Normalization

Problem

GRPO (Group Relative Policy Optimization) replaces the separate value/critic network that PPO-style RL normally needs with something much cheaper: sample a group of G completions for the same prompt, score each with a reward model, and turn those raw rewards directly into advantages by normalizing within the group — subtract the group's mean reward and divide by the group's standard deviation.

Every completion's advantage is then just "how much better or worse was this response than the other responses to the same prompt," with no learned baseline required. The tricky part is doing this without producing NaNs when a group happens to have zero variance (e.g. a group of size 1, or every sampled completion getting the identical reward).

Source: src/8_grpo_advantage_normalization.py

def grpo_group_advantages(rewards: torch.Tensor, eps: float = 1e-6) -> torch.Tensor: ...

>>> rewards = torch.tensor([[1.0, 2.0, 3.0]])
>>> grpo_group_advantages(rewards)
tensor([[-1.2247,  0.0000,  1.2247]])

Step-by-Step Approach

1. Compute each group's mean reward along the group axis (the last dimension): rewards.mean(dim=-1, keepdim=True).
2. Compute each group's standard deviation along the same axis with unbiased=False (population variance, divide by G), not the default unbiased=True (divide by G-1), which is undefined (0/0 = NaN) whenever a group has only 1 member.
3. Add a small eps to the standard deviation before dividing, so even a genuinely zero-variance group (every reward in the group identical) produces 0 / eps = 0 instead of 0 / 0 = NaN.
4. Subtract the mean and divide by (std + eps), relying on broadcasting since mean and std both have shape (num_prompts, 1) against rewards's (num_prompts, group_size).
5. Sanity check: a group of all-identical rewards should produce all-zero advantages, and a group of size 1 should always produce advantage 0 (there's nothing to compare it against) — verify neither case ever produces NaN.

The key insight is that GRPO's efficiency claim ("no critic network needed") only holds up if the normalization itself is numerically bulletproof — a critic-free RL algorithm that occasionally emits NaN advantages because of a same-reward or singleton group is not actually simpler in practice, just simpler on paper.

Reference solution

import torch


def grpo_group_advantages(rewards: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
    mean = rewards.mean(dim=-1, keepdim=True)  # (num_prompts, 1)
    # unbiased=False (population std, divide by G not G-1): the default
    # unbiased=True is NaN for group_size == 1 (0/0), which a real GRPO
    # rollout can hit whenever a prompt only gets one sample
    std = rewards.std(dim=-1, keepdim=True, unbiased=False)  # (num_prompts, 1)
    return (rewards - mean) / (std + eps)  # broadcast over group_size

Key Functions & Tricks

• tensor.mean(dim=-1, keepdim=True) — per-group mean, keeping a broadcastable singleton dimension instead of squeezing it away.
• tensor.std(dim=-1, keepdim=True, unbiased=False) — population standard deviation; the unbiased=False flag is the critical numerical-stability choice here, avoiding the N-1 denominator's 0/0 failure mode at group size 1.
• std + eps before dividing — the standard epsilon-stabilization pattern seen throughout normalization layers (LayerNorm, BatchNorm), applied here to a reward-normalization context instead.
• Broadcasting (rewards - mean) / (std + eps) — shape (num_prompts, 1) statistics broadcast automatically against (num_prompts, group_size) rewards.

How to Recognize This Pattern