GRPO and the OlmoRL Objective

The problem

GRPO (Group Relative Policy Optimization) trains a model with reinforcement learning without a second "value" model. For each prompt it samples a group of $G$ answers, scores each with a reward (for example 1 if a math answer is correct, 0 if not), and uses each answer's reward relative to its group as the advantage of all its tokens. OLMo 3's OlmoRL keeps that idea with a few changes (report, section 4.4.1). You will implement them.

Rewards come as an array of shape $(P \cdot G,)$: $P$ prompts, each with $G$ consecutive answers.

Change 1: advantages without dividing by the spread

The advantage of answer $i$ is its reward minus the mean of its group: $A_i = r_i - \mathrm{mean}(\text{group})$. Original GRPO also divides by the group's standard deviation (plus 1e-6, to avoid dividing by 0). OlmoRL does not: dividing would overweight prompts that are almost always solved or almost never.

Example. $G = 4$.

• $(1, 0, 0, 1)$ — advantages (OlmoRL): $(0.5, -0.5, -0.5, 0.5)$; advantages (divided by the std): $(1, -1, -1, 1)$
• $(1, 1, 1, 0)$ — advantages (OlmoRL): $(0.25, 0.25, 0.25, -0.75)$; advantages (divided by the std): $(0.58, 0.58, 0.58, -1.73)$
• $(1, 0, 0, 0)$ — advantages (OlmoRL): $(0.75, -0.25, -0.25, -0.25)$; advantages (divided by the std): $(1.73, -0.58, -0.58, -0.58)$

A prompt with a small spread (almost all answers the same) gets inflated advantages when divided.

Change 2: drop groups with no signal

If all $G$ rewards in a group are identical, every advantage is 0: the group carries no gradient, and costs compute. nonzero_signal_groups marks the groups to keep: those whose rewards are not all identical.

Example. Rewards $(1, 1, 1, 1)$: all correct, dropped. $(0, 0, 0, 0)$: all wrong, dropped. $(1, 0, 0, 0)$: kept.

Changes 3 and 4: the objective

The per-token term is the one of PPO (Proximal Policy Optimization) (a clipped ratio), but the clip is asymmetric: $\rho$ is limited to $[1 - \epsilon_{\text{low}}, 1 + \epsilon_{\text{high}}]$ with $\epsilon_{\text{low}} = 0.2$ and $\epsilon_{\text{high}} = 0.272$ ("clip-higher"), which lets good tokens move up a little more. With $\rho = e^{\log\pi_{\text{new}} - \log\pi_{\text{old}}}$ and advantage $A$:

And the objective is the sum over all real tokens of all answers, divided by the total number of real tokens: token-level normalization.

Example. Token-level against per-answer. Two answers: the first has 1 token and the second has 3, and say the per-token terms are $1.0$ for the first and $0.2$ for each of the second.

• Token-level: $(1.0 + 0.2 + 0.2 + 0.2)/4 = 0.4$. Every token counts the same.
• Per-answer (mean of each answer, then mean over answers): $(1.0 + 0.2)/2 = 0.6$. The token of the short answer counts three times as much as one of the long answer: a length bias.

Clip-higher: a token with $A > 0$ and $\rho = 1.25$. With the symmetric limit $1.2$ its term is $1.2A$, constant: no gradient. With the limit $1.272$ the term is $1.25A$: not clipped yet.

What to implement

• group_advantages(rewards, G, normalize_std): the advantages; divide by the standard deviation plus 1e-6 only if normalize_std is true (original GRPO). OlmoRL uses normalize_std=False.
• nonzero_signal_groups(rewards, G): a boolean array of shape $(P,)$, true for the groups whose rewards are not all identical.
• olmorl_objective(logp_new, logp_old, adv, mask, eps_low=0.2, eps_high=0.272): token-level arrays of shape $(N, T)$ ($N$ answers of up to $T$ tokens) and adv of shape $(N,)$, the same for every token of an answer. mask marks the real tokens (not padding). Return the token-level objective.

Requirements

• Token-level normalization weighs every real token equally, as in the example.
• No loops.

Hint

adv[:, None] has shape $(N, 1)$ and broadcasts over the $T$ tokens of each answer. For the groups, rewards.reshape(-1, G) puts one group per row.