GRPO: Group-Relative Advantages & Objective - Post-training & Efficiency

GRPO: Group-Relative Advantages & Objective

What you'll build

Two functions: group_advantages, which turns each prompt's group of rewards into normalized advantages, and grpo_objective, the clipped, masked GRPO objective with a k3 KL penalty.

Background

GRPO (Group Relative Policy Optimization, DeepSeekMath, Shao et al. 2024) is PPO without a value network. For each prompt, sample a group of G responses and score each one with a reward model or a verifier (e.g. 1 if the final answer is correct, 0 otherwise). The group mean plays the role of the critic's baseline, so there is no critic to train. Compared with PPO-RLHF Objective (Reinforcement Learning), the KL to the reference model is a separate term in the objective rather than part of the reward, and it uses the k3 estimator, which is never negative (compare KL Divergence Penalty (RLHF)).

Formula

Step 1: group-normalized advantages. For prompt b with rewards r_{b,1..G}:

A_{b,i} = (r_{b,i} - mean_i r_{b,:}) / (std_i r_{b,:} + eps)

• std is the population standard deviation (ddof = 0: np.std's default; PyTorch unbiased=False)
• This is outcome supervision: every token of response i gets the same advantage A_{b,i}

Step 2: per-token clipped surrogate and k3 KL. For token t of response i (log-probs of the sampled token):

ratio = exp(logp - old_logp)
surr  = min(ratio · A, clip(ratio, 1 - ε, 1 + ε) · A)
d     = ref_logp - logp
kl    = exp(d) - d - 1                 k3 estimate of KL(π_θ ‖ π_ref), always ≥ 0
per_token = surr - β · kl

Step 3: average (as in the DeepSeekMath objective):

J = mean over all (b, i) of [ Σ_t mask_t · per_token_t / Σ_t mask_t ]

First average over each response's valid tokens (mask = 1 for real tokens, 0 for padding), then over all B·G responses. 'mean_kl' uses the same averaging applied to kl. The training loss is -J.

Inputs

• rewards: (B, G) floats, one scalar per response
• logp, old_logp, ref_logp: (B, G, T) token log-probs under the current policy, the policy that sampled the data, and the frozen reference model
• mask: (B, G, T) of 0/1, with at least one valid token per response
• epsilon (clip range, default 0.2), beta (KL coefficient, default 0.04), eps (std epsilon, default 1e-8)

Expected return format

group_advantages(rewards, eps=1e-8) returns the (B, G) float array of advantages.

grpo_objective(...) returns a dictionary with these keys:

• 'advantages': the (B, G) array from step 1
• 'objective': J as a Python float (to maximize)
• 'loss': -J as a Python float (to minimize)
• 'mean_kl': the masked, averaged k3 KL as a Python float

Edge cases and conventions

• A group where every response has the same reward gets advantage 0 everywhere: no learning signal, and no division by zero thanks to eps.
• Values at padded positions (mask = 0) must not affect any output.
• On-policy with no drift (logp = old_logp = ref_logp), ratio = 1 and kl = 0, so J is the mean advantage, which is 0.

Constraints

NumPy only, fully vectorized (no Python loops over B, G or T).

Follow-up

• Dr. GRPO argues that dividing by the response length (1/|o_i|) and by the group std biases training toward particular lengths and prompts. What changes if you drop them?
• How large should G be? What happens to the advantage estimates with G = 2?

Input: group_advantages([[1.0, 0.0, 0.0, 1.0]])
Output: [[1.0, -1.0, -1.0, 1.0]]

Input: group_advantages([[0.5, 0.5, 0.5], [1.0, 2.0, 3.0]])
Output: [[0.0, 0.0, 0.0], [-1.224745, 0.0, 1.224745]]

Input: grpo_objective(logp=[[[0.0], [0.0]]], old_logp=[[[-1.0], [0.0]]], ref_logp=[[[0.0], [0.0]]], rewards=[[1.0, 0.0]], mask=[[[1], [1]]], epsilon=0.2, beta=0.0)
Output: {'objective': 0.1, 'loss': -0.1, 'mean_kl': 0.0, ...}

import numpy as np

def group_advantages(rewards, eps=1e-8):
    """
    Group-normalized advantages.

    Args:
        rewards: (B, G) rewards, G sampled responses per prompt
        eps: added to the std

    Returns:
        (B, G) array: (r - group mean) / (group std + eps), population std
    """
    # Your code here
    pass


def grpo_objective(logp, old_logp, ref_logp, rewards, mask,
                   epsilon=0.2, beta=0.04, eps=1e-8):
    """
    GRPO clipped objective with a k3 KL penalty.

    Args:
        logp: (B, G, T) token log-probs under the current policy
        old_logp: (B, G, T) token log-probs under the policy that sampled
        ref_logp: (B, G, T) token log-probs under the frozen reference model
        rewards: (B, G) one scalar reward per response
        mask: (B, G, T) 1 for real tokens, 0 for padding
        epsilon: clipping range
        beta: KL coefficient
        eps: std epsilon for the advantages

    Returns:
        dict with 'advantages' (B, G) array, and 'objective', 'loss' (= -objective)
        and 'mean_kl' as Python floats; token terms are averaged over each
        response's valid tokens, then over all responses
    """
    # Your code here
    pass