Study Plan

Normalize GRPO Group Advantages

Language ModelsReinforcement Learning

Medium

Normalize rollout rewards independently within each prompt group. For rollout i in group g, define

A_i=\frac{r_i-\mu_g}{\sigma_g+\epsilon}.

Here, $r_i$ is the reward, $\mu_g$ is its group's mean, $\sigma_g$ is the population standard deviation, and $\epsilon$ is the supplied positive stabilizer. A group with identical rewards has zero advantages and is inactive; groups with nonzero standard deviation are active. Preserve rollout order, and order group-level results by group ID. Empty inputs produce empty tensors and an active count of zero.

Return a Python dictionary with advantages, a tensor matching the rewards' shape, dtype and device; active_group_mask, a Boolean tensor of shape (number of groups,) on that device; and active_group_count, a Python integer.

Loading visualization...

Examples

Input: rewards = [1,3], group_ids = [0,0], epsilon = 0.00000001

Output: {"advantages":[-1,1],"active_group_mask":[true],"active_group_count":1}

Explanation: The group has mean 2 and population standard deviation 1.

Input: rewards = [1,3,4,4], group_ids = [0,0,1,1], epsilon = 0.00000001

Output: {"advantages":[-1,1,0,0],"active_group_mask":[true,false],"active_group_count":1}

Input: rewards = [4,4], group_ids = [0,0], epsilon = 0.00000001

Output: {"advantages":[0,0],"active_group_mask":[false],"active_group_count":0}

Hint 1

Boolean indexing selects the rollouts belonging to one group.

Hint 2

torch.var supports population variance without a sample correction.

Hint 3

torch.all checks whether every selected reward equals the first.

Requirements

Use PyTorch and preserve reward dtype and device.
A constant-reward group must produce exact zero advantages and remain inactive.
Do not modify input tensors.

Constraints

rewards is a finite one-dimensional floating PyTorch tensor; editor inputs use float32.
group_ids is an aligned int64 tensor on the same device, containing contiguous IDs from zero.
epsilon is positive and finite.
Empty aligned tensors and noncontiguous inputs are valid.

Try Similar Problems

Compare Sequence and Token Length NormalizationMedium Compute a Clipped GRPO SurrogateHard Canonicalize Pairwise Preference DataMedium Implement the DPO LossMedium Score a KL-Regularized Policy BatchMedium Train a BPE Tokenizer From Scratch in PythonMedium Compute Temperature-Based Data MixturesMedium Apply Ranked BPE MergesMedium Value IterationMedium Expected SARSAMedium Bradley-Terry Reward Model LossMedium RLOO (REINFORCE Leave-One-Out)Medium