Calculating Group Relative Advantage in Python

Welcome to the first lesson of our third module. In the previous module, we fully deconstructed the theory behind the Group Relative Policy Optimization objective function. We saw how it cleverly combines a critic-free advantage signal with the stability mechanisms of PPO. Now, we'll begin translating that theory into practice.

This lesson marks our first step into implementation. Our goal is to write a Python function that computes the group-relative advantage, a core component of the GRPO loss. We will revisit the formula, examine how to structure the data, and walk through a concrete implementation using standard tensor libraries.

From Theory to Code: The Group Relative Advantage

In our theoretical exploration, we defined the group-relative advantage as the standardized reward of a single generation within a group of its peers. This critic-free approach is what distinguishes GRPO from methods like PPO. The formula is a direct translation of this concept.

The formula for the advantage of the $i$-th output in a group is:

where:

• $r_i$ is the reward for the $i$-th output.
• $\bar{r}$ is the mean of all rewards in the group.
• ${\sigma }_r$ is the standard deviation of all rewards in the group.
• $ϵ$ is a small constant (e.g., $10^{-8}$) added for numerical stability to prevent division by zero.

Our task is to implement a function that performs this calculation across a batch of prompts, where for each prompt we have generated a group of multiple responses.

Implementing the Advantage Calculation

Let's imagine our data flow. For a batch of $B$ prompts, we generate $G$ responses for each. This gives us a total of $B\times G$ responses and a corresponding list of $B\times G$ rewards. Our implementation must first group these rewards by their original prompt and then apply the formula to each group.

The following video segments walk through both the concept and the practical code for this process.

As the video demonstrates, the core steps in code (using a library like PyTorch or JAX) are:

1. Reshape Rewards: Convert the flat tensor of rewards with shape (B * G) into a grouped tensor of shape (B, G).
2. Compute Group Statistics: Calculate the mean and standard deviation for each group. This is typically done by applying the functions along dim=1 of the (B, G) tensor, resulting in two tensors of shape (B,).
3. Normalize: Use broadcasting to subtract the mean and divide by the standard deviation. You'll need to expand the (B,) statistics tensors back to match the (B, G) shape.
4. Flatten: Reshape the final advantages tensor back to (B * G) to match the original list of responses.

A Step-by-Step Code Example

Let's solidify this with a concrete example. The article "The Illustrated GRPO" provides an excellent, minimal implementation using PyTorch that we can follow.

Let's trace the logic from the article. Suppose we have batch_size = 2 prompts and num_generations = 4 responses per prompt. Our model generates responses, and we compute rewards for them.

• Prompt 1 (correct answer '5'): generated [5, 6, 7, 5]. Binary rewards: [1, 0, 0, 1].
• Prompt 2 (correct answer '9'): generated [10, 2, 9, 9]. Binary rewards: [0, 0, 1, 1].

Our initial rewards tensor is flat: tensor([1, 0, 0, 1, 0, 0, 1, 1]).

Here is a Python function that encapsulates the logic:

import torch

def compute_group_relative_advantage(
    rewards: torch.Tensor, 
    batch_size: int, 
    num_generations: int,
    epsilon: float = 1e-8
) -> torch.Tensor:
    """
    Computes the group-relative advantage from a batch of rewards.

    Args:
        rewards: A flat tensor of rewards of shape (batch_size * num_generations).
        batch_size: The number of prompts in the batch.
        num_generations: The number of generations per prompt.
        epsilon: A small value to add to the standard deviation for stability.

    Returns:
        A flat tensor of advantages of shape (batch_size * num_generations).
    """



    # 1. Reshape rewards to group them by prompt
    # Shape: (batch_size, num_generations)
    rewards_grouped = rewards.view(batch_size, num_generations)




    # 2. Compute per-group statistics
    # Shape: (batch_size,)
    mean_grouped = rewards_grouped.mean(dim=1)
    std_grouped = rewards_grouped.std(dim=1)




    # 3. Broadcast statistics to match the grouped rewards shape
    # Shape: (batch_size, num_generations)
    # We need to unsqueeze to allow for broadcasting along dim=1
    mean_broadcast = mean_grouped.unsqueeze(1)
    std_broadcast = std_grouped.unsqueeze(1)




    # 4. Compute advantages using the formula
    advantages_grouped = (rewards_grouped - mean_broadcast) / (std_broadcast + epsilon)
    



    # 5. Flatten the result to match the original input's structure
    advantages_flat = advantages_grouped.view(-1)

    return advantages_flat




# --- Example Usage ---
# Our example rewards
rewards_flat = torch.tensor([1, 0, 0, 1, 0, 0, 1, 1], dtype=torch.float32)
B, G = 2, 4




# Grouped for inspection: [[1, 0, 0, 1], [0, 0, 1, 1]]
# Mean per group: [0.5, 0.5]
# Std per group:  [0.5774, 0.5774] (approx)

advantages = compute_group_relative_advantage(rewards_flat, batch_size=B, num_generations=G)
print(advantages)



# tensor([ 0.8660, -0.8660, -0.8660,  0.8660, -0.8660, -0.8660,  0.8660,  0.8660])

Notice the calculation for the first advantage value: $(1.0-0.5)/0.5774\approx 0.866$. This positive advantage will encourage the policy to increase the probability of that correct response. Conversely, the second value, $(0.0-0.5)/0.5774\approx -0.866$, is a negative advantage that will discourage that incorrect response.

The Case of Zero Variance

A key detail is handling the case where all rewards in a group are identical (e.g., all 1.0 or all 0.0). This results in a standard deviation of zero.

• The epsilon term is crucial here to prevent division by zero, ensuring numerical stability.
• The resulting advantage for all items in that group will be 0. This is intuitively correct: if all responses are equally good, no single response has a relative advantage over its peers. The policy receives no gradient signal from this group, which is the desired behavior.

You can find another perspective on the implementation in a different framework (Apple's MLX) in the dev.to article "Implementing DeepSeek-R1 GRPO." Although the framework is different, the logic for calculating the advantage is identical, demonstrating the universality of the concept. You can inspect the code snippets in Part 6: Ai and the concrete example in Step 2: calculate normalised advantage.

Conclusion

In this lesson, we made the critical transition from the theory of GRPO to its practical implementation. We focused on the cornerstone of the algorithm: the group-relative advantage.

The key takeaways are:

• The group-relative advantage is calculated by standardizing rewards within a peer group generated from the same prompt.
• The implementation requires reshaping the rewards tensor to align with the (batch_size, num_generations) structure.
• Standard tensor library functions (mean, std) can be applied along the appropriate dimension to compute group-wise statistics.
• Broadcasting is used to apply the normalization formula element-wise, and a small epsilon is essential for numerical stability.

You have now implemented the function that provides the core $\hat{A}$ term for the GRPO objective. In our next lesson, we will build upon this foundation. We will implement the complete GRPO loss function, integrating the advantages you just learned to compute with the PPO-style clipped policy ratio and the KL divergence penalty.