Implementing GRPO Loss with Clipping and KL Penalty in Python

Welcome back! In our previous lesson, we successfully implemented a function to calculate the group-relative advantage, which provides the crucial reward signal for GRPO. We now have the $\hat{A}$ term, representing how much better or worse a given response is compared to its peers.

Today, we will use that advantage to construct the complete GRPO loss function. This lesson is the heart of our implementation module, where we translate the full GRPO objective from mathematical theory into a functional piece of Python code. We will assemble the two main components: the PPO-style clipped surrogate objective for policy updates and the KL divergence penalty for regularization.

By the end of this lesson, you will have a complete function ready to be plugged into a training loop.

The Anatomy of the GRPO Loss

The GRPO objective function is designed to be maximized. However, deep learning frameworks are built to minimize a loss function. Therefore, our implementation will compute the negative of the GRPO objective.

The full objective is composed of two key parts:

1. A clipped policy gradient term: This drives the model to increase the probability of high-advantage responses.
2. A KL divergence penalty: This regularizes the model, preventing it from deviating too far from a trusted reference policy.

Let's look at the formula we are about to implement.

Our task is to translate this into a loss function that a standard optimizer like Adam can minimize. As is common in RL implementations, we will perform calculations in log-space for numerical stability.

1. The Clipped Surrogate Objective

This part of the loss is inherited from Proximal Policy Optimization (PPO) and is designed to take stable policy update steps. It prevents the model from changing too drastically in a single gradient update, which could lead to a collapse in performance.

The term for a single token is:

Where:

• $\hat{A}$ is the group-relative advantage we calculated in the previous lesson.
• $\text{ratio}(\theta )=\frac{{\pi }_{\theta }(\text{output})}{{\pi }_{{\theta }_{old}}(\text{output})}$ is the probability ratio between the current policy ${\pi }_{\theta }$ and the old policy ${\pi }_{{\theta }_{old}}$ that was used to generate the data. In log space, this is computed as $\exp (\log p_{\theta }-\log p_{{\theta }_{old}})$.
• $ϵ$ is a small hyperparameter (e.g., 0.2) that defines the clipping range.

The min function provides a pessimistic bound on the policy update. If the advantage $\hat{A}$ is positive, it caps the update to prevent over-optimism. If $\hat{A}$ is negative, it uses the clipped term, which can allow for a larger (more negative) update to more strongly discourage bad actions.

The following resource provides a superb breakdown of this objective, complete with a step-by-step numerical example.

As you've seen, this objective cleverly balances making progress with maintaining stability. The final policy loss to be minimized is the negative of this term, averaged over all tokens in the batch.

2. The KL Divergence Penalty

The second component of our loss is a regularization term. Its purpose is to keep the current policy ${\pi }_{\theta }$ from straying too far from a trusted reference policy ${\pi }_{ref}$ (usually the initial Supervised Fine-Tuned model). This helps prevent "catastrophic forgetting," where the model unlearns its general language capabilities while optimizing for the new task.

The KL divergence measures the "distance" between two probability distributions. The specific estimator for the KL divergence used in the GRPO paper (and many PPO implementations) is:

where $r=\frac{{\pi }_{ref}(\text{output})}{{\pi }_{\theta }(\text{output})}$. This is an unbiased estimator of the KL divergence with favorable variance properties.

Given your background in statistics, you might find the origin of this estimator interesting. It's a refined version of simpler estimators, designed to be computationally efficient while remaining statistically sound. The following video explains the family of KL estimators that PPO and GRPO draw upon.

This KL penalty is added to the total loss, weighted by a hyperparameter $\beta$.

3. Implementing the Full GRPO Loss

Now we can combine these pieces into a single Python function. Our function will need the following inputs:

• log_probs: Log probabilities of the generated tokens from the current policy ${\pi }_{\theta }$.
• old_log_probs: Log probabilities from the policy that generated the data, ${\pi }_{{\theta }_{old}}$.
• ref_log_probs: Log probabilities from the reference policy ${\pi }_{ref}$.
• advantages: The group-relative advantages we implemented in the last lesson.
• completion_mask: A binary mask to ensure we only compute the loss on the generated tokens, not the prompt.

The following resources show practical implementations in both the TRL library and a standalone script, which we will use as a guide.

Let's write our function, bringing everything together.

import torch

def compute_grpo_loss(
    log_probs: torch.Tensor,
    old_log_probs: torch.Tensor,
    ref_log_probs: torch.Tensor,
    advantages: torch.Tensor,
    completion_mask: torch.Tensor,
    epsilon: float,
    beta: float,
) -> torch.Tensor:
    """
    Computes the complete GRPO loss.

    Args:
        log_probs: Per-token log probabilities from the current policy (B*G, L).
        old_log_probs: Per-token log probabilities from the old policy (B*G, L).
        ref_log_probs: Per-token log probabilities from the reference policy (B*G, L).
        advantages: Per-sequence advantages (B*G,).
        completion_mask: Mask to select completion tokens (B*G, L).
        epsilon: Clipping hyperparameter for the surrogate objective.
        beta: Weight for the KL divergence penalty.

    Returns:
        The final GRPO loss as a scalar tensor.
    """



    # 1. Clipped Surrogate Objective calculation
    



    # Reshape advantages to be per-token
    # from (B*G,) to (B*G, 1) to allow broadcasting across sequence length L
    advantages_per_token = advantages.unsqueeze(-1)
    



    # Calculate the probability ratio in log-space, then exponentiate
    # ratio shape: (B*G, L)
    ratio = torch.exp(log_probs - old_log_probs)
    



    # Unclipped policy objective
    policy_loss_unclipped = ratio * advantages_per_token
    



    # Clipped policy objective
    clipped_ratio = torch.clamp(ratio, 1.0 - epsilon, 1.0 + epsilon)
    policy_loss_clipped = clipped_ratio * advantages_per_token
    



    # PPO's pessimistic objective: take the minimum
    # Because we MINIMIZE a loss, but the objective is MAXIMIZED, we negate it.
    surrogate_loss = -torch.min(policy_loss_unclipped, policy_loss_clipped)
    



    # 2. KL Divergence Penalty calculation
    



    # Use the r - log(r) - 1 approximation for KL divergence
    # r_kl = pi_ref / pi_theta
    log_r_kl = ref_log_probs - log_probs
    r_kl = torch.exp(log_r_kl)
    kl_penalty = r_kl - log_r_kl - 1
    



    # 3. Combine and apply mask
    



    # Combine the surrogate loss and the KL penalty
    per_token_loss = surrogate_loss + beta * kl_penalty
    



    # Apply the completion mask to only calculate loss over generated tokens
    # and average it to get a single scalar value.
    masked_loss = per_token_loss * completion_mask
    final_loss = masked_loss.sum() / completion_mask.sum()

    return final_loss

Let's break down the implementation:

1. Surrogate Objective: We compute the ratio from log probabilities, calculate both the unclipped and clipped policy objectives (policy_loss_unclipped, policy_loss_clipped), and take the negative of their minimum. The advantages tensor is unsqueezed to enable broadcasting across the token sequence dimension.
2. KL Penalty: We implement the r - log(r) - 1 formula, where r is the ratio between the reference and current policies.
3. Final Loss: We combine the two terms, apply the completion_mask to zero out the loss for prompt tokens, and then compute the mean loss over only the valid (completion) tokens.

Conclusion

In this lesson, we have successfully constructed the complete GRPO loss function in Python. We have translated the core theoretical concepts of PPO's stable policy updates and KL regularization into a practical, usable function.

Here are the key takeaways:

• The GRPO loss consists of a clipped surrogate objective to drive policy improvement and a KL penalty for regularization.
• The surrogate objective uses a probability ratio and clipping to ensure stable updates. The entire term is negated to frame it as a loss to be minimized.
• The KL penalty keeps the learned policy from diverging too far from a trusted reference model, using the computationally efficient r - log(r) - 1 estimator.
• A completion mask is a crucial practical detail, ensuring that the loss is only computed on the tokens generated by the model, not the input prompt.

You now have the two most important functions of our project: compute_group_relative_advantage and compute_grpo_loss. In our final lesson, we will put it all together. We will assemble a complete training loop that uses these functions to fine-tune a language model on a dummy dataset, taking our GRPO implementation from a set of functions to a running process.