Implementing DeekSeek-R1 GRPO in Apple MLX framework

Table of Contents

• Motivation
• Show me the code: Jupyter notebook
• Peering into the GRPO equation
• Part 1: E[q∼P(Q),⟨oi⟩i=1G∼πθold(O∣q)]\mathbb{E}[q \sim P(Q), \langle o_i\rangle_{i=1}^{G} \sim \pi_{\theta_{old}}(O|q)]E[q∼P(Q),⟨oi​⟩i=1G​∼πθold​​(O∣q)]
• The rollout phase: ⟨oi⟩i=1G\langle o_i\rangle_{i=1}^{G}⟨oi​⟩i=1G​
• Part 2: πθ(oi∣q)πθold(oi∣q)\frac{\pi_\theta(o_i|q)}{\pi_{\theta_{\text{old}}}(o_i|q)}πθold​​(oi​∣q)πθ​(oi​∣q)​
• Part 3: clip(πθ(oi∣q)πθold(oi∣q),1−ϵ,1+ϵ)\text{clip}\left(\frac{\pi_\theta(o_i|q)}{\pi_{\theta_{\text{old}}}(o_i|q)}, 1-\epsilon, 1+\epsilon\right)clip(πθold​​(oi​∣q)πθ​(oi​∣q)​,1−ϵ,1+ϵ)
• Part 4: min⁡(πθ(oi∣q)πθold(oi∣q)Ai,clip(πθ(oi∣q)πθold(oi∣q),1−ϵ,1+ϵ)Ai)\min\left( \frac{\pi_\theta(o_i|q)}{\pi_{\theta_{\text{old}}}(o_i|q)} A_i, \text{clip}\left(\frac{\pi_\theta(o_i|q)}{\pi_{\theta_{\text{old}}}(o_i|q)}, 1-\epsilon, 1+\epsilon\right) A_i \right)min(πθold​​(oi​∣q)πθ​(oi​∣q)​Ai​,clip(πθold​​(oi​∣q)πθ​(oi​∣q)​,1−ϵ,1+ϵ)Ai​)
• Part 5: DKL(πθ∣∣πref))D_{KL}\left(\pi_\theta || \pi_{\text{ref}}) \right)DKL​(πθ​∣∣πref​))
• Part 6: AiA_iAi​
• Bringing it all together
• Concrete example of how to train using GRPO
• Conclusion

Motivation

Group Relative Policy Optimisation (GRPO) was a method developed by DeepSeek to improve "language model reasoning capabilities using pure reinforcement learning (RL)", with the specific goal to "develop reasoning capabilities without any supervised data, focusing on their self-evolution through a pure RL process" (source: DeepSeek-R1: Incentivising Reasoning Capabilitiy in LLMs via Reinforcement Learning). GRPO was the method used to train DeepSeek-R1 released in January 2025, which crashed tech stocks such as Nvidia, and served as the basis of subsequent reasoning models such as the Mistrel Magistral (source: Magistral).

This article seeks to explain the math of GRPO, and how to implement GRPO from scratch to train LLMs using Apple MLX framework. Hence, if you are an Apple silicon user, you are in luck, you can run the Jupyter notebook right on your laptop.

This article is created with the assistance of Google Gemini 2.5 Pro.

Show me the code: Jupyter notebook

For those who are keen to dive right into running the code, you may access it here. If you discover any mistakes or have any improvements to suggest, please feel free to make a pull request! I will look into all requests as soon as I can.

Peering into the GRPO equation

We will dissect is this scary-looking (at least to me!) equation:

This is the theoretical expectation, which theoretically considers every possible prompt, generates a group of responses for each, and then average the improvement we get from the objective function (see below) over all these possibilities.

This is the sample-based objective function (the "estimator"). Since we cannot possibly compute the true expectation over all prompts and outputs, we approximate it during training where we take one batch ( a prompt q and its G generated outputs {oi}\{o_i\}{oi​} and we compute the numerical estimate of our objective using the objective function. We then use this numerical estimate to calculate a gradient and update our model's weights, i.e. ( θold\theta_{\text{old}}θold​ becomes $\theta$ ).

I will break down the equations into the following parts:

• Part 1: E[q∼P(Q),⟨oi⟩i=1G∼πθold(O∣q)]\mathbb{E}[q \sim P(Q), \langle o_i\rangle_{i=1}^{G} \sim \pi_{\theta_{old}}(O|q)]E[q∼P(Q),⟨oi​⟩i=1G​∼πθold​​(O∣q)]

• Part 2: πθ(oi∣q)πθold(oi∣q)\frac{\pi_\theta(o_i|q)}{\pi_{\theta_{\text{old}}}(o_i|q)}πθold​​(oi​∣q)πθ​(oi​∣q)​

• Part 3: clip(πθ(oi∣q)πθold(oi∣q),1−ϵ,1+ϵ)\text{clip}\left(\frac{\pi_\theta(o_i|q)}{\pi_{\theta_{\text{old}}}(o_i|q)}, 1-\epsilon, 1+\epsilon\right)clip(πθold​​(oi​∣q)πθ​(oi​∣q)​,1−ϵ,1+ϵ)

• Part 4: min⁡(πθ(oi∣q)πθold(oi∣q)Ai,clip(πθ(oi∣q)πθold(oi∣q),1−ϵ,1+ϵ)Ai)\min\left( \frac{\pi_\theta(o_i|q)}{\pi_{\theta_{\text{old}}}(o_i|q)} A_i, \text{clip}\left(\frac{\pi_\theta(o_i|q)}{\pi_{\theta_{\text{old}}}(o_i|q)}, 1-\epsilon, 1+\epsilon\right) A_i \right)min(πθold​​(oi​∣q)πθ​(oi​∣q)​Ai​,clip(πθold​​(oi​∣q)πθ​(oi​∣q)​,1−ϵ,1+ϵ)Ai​)

• Part 5: DKL(πθ∣∣πref))D_{KL}\left(\pi_\theta || \pi_{\text{ref}}) \right)DKL​(πθ​∣∣πref​))

• Part 6: AiA_iAi​

Part 1

The equation that will be dissected in this part is:

• The tilde (~): It means "distributed according to" or "is sampled from". So:

  • $q\sim P(Q)$ This represents questions (q) sampled from the overall distribution of questions (P(Q)). This is a standard practice where models learn to respond to various prompts during training.
  • ⟨oi⟩i=1G∼πθold(O∣q)\langle o_i\rangle_{i=1}^G \sim \pi_{\theta_{\text{old}}}(O|q)⟨oi​⟩i=1G​∼πθold​​(O∣q) : This means the group of G outputs ⟨oi⟩i=1G\langle o_i\rangle_{i=1}^G⟨oi​⟩i=1G​ is sampled from the policy πθold\pi_{\theta_{\text{old}}}πθold​​ given question q.

• The expectation function i.e.
  
  $\mathbb{E}[\bullet ,\bullet ]$
  
  signifies a joint expectation over multiple random variables. One commonly seen example is
  
  $\mathbb{E}[X\sim P(X),Y\sim P(Y|X)][\text{function}(X,Y)]$
  
  which means taking the expectation over the combined process of first sampling X from P(X), and then sampling Y from P(Y|X). The expectation then applies to the function of both X and Y.

• In probability theory,
  
  EX,Y[f(X,Y)]\mathbb{E}_{X,Y}[f(X,Y)]EX,Y​[f(X,Y)]
  
  can be written as
  
  ∑XP(X)∑YP(Y∣X)f(X,Y)\sum_X P(X) \sum_Y P(Y|X) f(X,Y)∑X​P(X)∑Y​P(Y∣X)f(X,Y)
  
  for discrete variables, or
  
  ∫X∫YP(X,Y)f(X,Y)dXdY\int_X \int_Y P(X,Y) f(X,Y) dX dY∫X​∫Y​P(X,Y)f(X,Y)dXdY
  
  for continuous variables. The comma notation is a shorthand for this sequential or joint sampling process.

• Applying to our equation:

  • q ~ P(Q): First, a question q is randomly chosen from the pool of all possible questions.
  • ⟨oi⟩i=1G∼πθold(O∣q)\langle o_i\rangle_{i=1}^G \sim \pi_{\theta_{\text{old}}}(O|q)⟨oi​⟩i=1G​∼πθold​​(O∣q) : Given that specific question q, a group of G outputs ⟨oi⟩i=1G\langle o_i\rangle_{i=1}^G⟨oi​⟩i=1G​ is then generated by the old policy πθold\pi_{\theta_{\text{old}}}πθold​​ .
  • The expression $[\text{function}(Q,O)]$ that follows (the GRPO objective function) then depends on both q and the generated output ⟨oi⟩i=1G\langle o_i\rangle_{i=1}^G⟨oi​⟩i=1G​ .
  • So, the expectation is taken over the entire data collection process: first randomly pick a question from the bank of available questions, and then generating multiple responses for that question using the πθold\pi_{\theta_{\text{old}}}πθold​​ policy.

In summary:

• The tilde (~) tells you how the data is being generated (which distribution).
• The comma (,) separates independent (or conditionally independent) sampling steps that define the full set of random variables over which the expectation is taken. It implies a joint probability distribution, often constructed sequentially.

Why E[q∼P(Q),⟨oi⟩i=1G∼πθold(O∣q)][function(Q,O)]\mathbb{E}[q \sim P(Q), \langle o_i\rangle_{i=1}^{G} \sim \pi_{\theta_{\text{old}}}(O|q)][\text{function}(Q, O)]E[q∼P(Q),⟨oi​⟩i=1G​∼πθold​​(O∣q)][function(Q,O)] instead of simply E[P(Q),πθold(O∣q)][function(Q,O)]\mathbb{E}[P(Q), \pi_{\theta_{\text{old}}}(O|q)][\text{function}(Q, O)]E[P(Q),πθold​​(O∣q)][function(Q,O)] ?

• The notation $q\sim P(Q)$ explicitly states that q is the random variable being sampled, and P(Q) is its probability distribution.
• Similarly, ⟨oi⟩i=1G∼πθold(O∣q)\langle o_i\rangle_{i=1}^{G} \sim \pi_{\theta_{\text{old}}}(O|q)⟨oi​⟩i=1G​∼πθold​​(O∣q) states that ⟨oi⟩i=1G\langle o_i\rangle_{i=1}^{G}⟨oi​⟩i=1G​ are the random variables (the sampled outputs), and πθold(O∣q)][function(Q,O)]\pi_{\theta_{\text{old}}}(O|q)][\text{function}(Q, O)]πθold​​(O∣q)][function(Q,O)] is the conditional probability distribution from which they are drawn (conditioned on q).

Without q ~ and ⟨oi⟩i=1G\langle o_i\rangle_{i=1}^{G}⟨oi​⟩i=1G​ , the expression E[P(Q),πθold(O∣q)][function(Q,O)]\mathbb{E}[P(Q), \pi_{\theta_{\text{old}}}(O|q)][\text{function}(Q, O)]E[P(Q),πθold​​(O∣q)][function(Q,O)] is ambiguous as it is not clear which variables are being sampled or how are they related to the function $[\text{function}(Q,O)]$ . The P(Q) and πθold(O∣q)\pi_{\theta_{\text{old}}}(O|q)πθold​​(O∣q) are distributions, not actual values that vary and contribute to the average.

Show me the code

To implement the Expectation operator E[...]. we can use a loop where each iteration processes a new, randomly sampled mini-batch of prompts, calculate the loss for that batch, and performs an update. Over many iterations, this process approximates the expected value of the objective over the entire data distribution.

Do not fret that the code is long. I will break it down and explain each piece accordingly.

# --- File: MLX LM GRPO.ipynb, Cell: grpo-loop ---
def grpo_train_loop(
    model, model_old, model_ref, tokenizer, optimizer, train_set,
    iters=200, group_size=4, batch_size=2, epsilon=0.2, beta=0.01,
    update_every=10, max_ans_len=4
):
    # Create a grad function for the trainable model
    loss_and_grad_fn = nn.value_and_grad(model, grpo_loss_fn)

    losses = []
    all_rewards = []

    # Start training
    pbar = tqdm.tqdm(range(iters))
    for it in pbar:
        batch_prompts = []
        batch_answers = []

        # 1. Sample a batch of prompts
        indices = np.random.randint(0, len(train_set), batch_size)
        for i in indices:
            # The last word of the output is the ground truth answer (e.g., "ending4")
            prompt_text, answer_text = train_set[i]["output"].rsplit(" ", maxsplit=1)
            full_prompt = [
                {"role": "user", "content": train_set[i]["instruction"]},
                {"role": "assistant", "content": prompt_text}
            ]
            batch_prompts.append(full_prompt)
            batch_answers.append(answer_text)

        # 2. Rollout: Generate G responses for each prompt using the old model
        rollout_sequences = []
        rollout_rewards = []
        rollout_log_probs = []
        rollout_a_toks = []

        for i in range(batch_size):
            prompt_tokens = tokenizer.apply_chat_template(batch_prompts[i], continue_final_message=True)
            group_rewards = []

            for _ in range(group_size):
                # Generate a response
                response = generate(model_old, tokenizer, prompt_tokens, max_tokens=max_ans_len)
                answer_tokens = tokenizer.encode(response, add_special_tokens=False)

                # 3. Get Reward
                reward = 1.0 if batch_answers[i] in response else 0.0
                group_rewards.append(reward)

                # Store data for the optimization step
                full_sequence = mx.array(prompt_tokens + answer_tokens)
                rollout_sequences.append(full_sequence)
                rollout_a_toks.append(mx.array(answer_tokens))

            all_rewards.extend(group_rewards)
            rollout_rewards.append(mx.array(group_rewards))

        # 4. Compute Advantages
        advantages = []
        for rewards in rollout_rewards:
            mean_reward = mx.mean(rewards)
            std_reward = mx.sqrt(mx.var(rewards)) + 1e-8 # Add epsilon for stability
            adv = (rewards - mean_reward) / std_reward
            advantages.append(adv)

        advantages = mx.concatenate(advantages)
        sequences = pad_sequences(rollout_sequences, tokenizer.pad_token_id)
        a_toks = pad_sequences(rollout_a_toks, tokenizer.pad_token_id)

        # Calculate log_probs with the old model for the ratio calculation
        old_log_probs = calculate_log_probs(model_old, sequences, a_toks)

        # 5. Optimization Step
        (loss, policy_reward, kl_div), grads = loss_and_grad_fn(
            model, model_ref, sequences, a_toks, advantages, old_log_probs, beta, epsilon
        )

        optimizer.update(model, grads)
        mx.eval(model.parameters(), optimizer.state)

        losses.append(loss.item())
        pbar.set_description(f"Loss: {np.mean(losses[-10:]):.3f}, Mean Reward: {np.mean(all_rewards[-20:]):.3f}")

        # Sync old model weights
        if (it + 1) % update_every == 0:
            model_old.update(model.parameters())
            print(f"\nIter {it+1}: Synced old model weights.")

    # Final save of adapter weights
    model.save_weights(str(adapter_path / "adapters.safetensors"))
    print("Saved final weights to adapters/adapters.safetensors.")
    return losses, all_rewards

The rollout phase: ⟨oi⟩i=1G\langle o_i\rangle_{i=1}^{G}⟨oi​⟩i=1G​

For each prompt q in our batch, we need to generate a group G possible outputs ( ⟨oi⟩i=1G\langle o_i\rangle_{i=1}^{G}⟨oi​⟩i=1G​ ) using the fixed, old policy πθold\pi_{\theta_{\text{old}}}πθold​​ . This is the data collection or "rollout" phase.

The code is implemented with a nested for loop within grpo_train_loop that calls generate. In the code, model_old is πθold\pi_{\theta_{\text{old}}}πθold​​ . The outer loop iterates through prompts in the batch, and the inner loop runs group_size (G) times to generate each output oio_ioi​ for that prompt.

# --- File: MLX LM GRPO.ipynb, Cell: grpo-loop ---
# Rollout: Generate G responses for each prompt using the old model
...
for i in range(batch_size):
            prompt_tokens = tokenizer.apply_chat_template(batch_prompts[i], continue_final_message=True)
            group_rewards = []

            for _ in range(group_size):
                # Generate a response
                response = generate(model_old, tokenizer, prompt_tokens, max_tokens=max_ans_len)
                answer_tokens = tokenizer.encode(response, add_special_tokens=False)

                # 3. Get Reward
                reward = 1.0 if batch_answers[i] in response else 0.0
                group_rewards.append(reward)

                # Store data for the optimization step
                full_sequence = mx.array(prompt_tokens + answer_tokens)
                rollout_sequences.append(full_sequence)
                rollout_a_toks.append(mx.array(answer_tokens))

            all_rewards.extend(group_rewards)
            rollout_rewards.append(mx.array(group_rewards))

Part 2: πθ(oi∣q)πθold(oi∣q)\frac{\pi_\theta(o_i|q)}{\pi_{\theta_{\text{old}}}(o_i|q)}πθold​​(oi​∣q)πθ​(oi​∣q)​

The concept behind this equation is importance sampling:

• πθold(oi∣q)\pi_{\theta_{\text{old}}}(o_i|q)πθold​​(oi​∣q) (The Denominator): This represents the probability of generating a specific output oio_ioi​ (a complete reasoning trace and answer) given the input question $q$ , using the old policy. The "old" policy is the version of the model that was used to generate the batch of data for the current training step. It is frozen during this step.

• πθ(oi∣q)\pi_\theta(o_i|q)πθ​(oi​∣q) (The Numerator): This represents the probability of generating that exact same output oio_ioi​ given the question $q$ , but using the current policy πθ\pi_\thetaπθ​ . This is the policy we are actively training and updating in this step.

• The Ratio: The ratio πθ(oi∣q)πθold(oi∣q)\frac{\pi_\theta(o_i|q)}{\pi_{\theta_{\text{old}}}(o_i|q)}πθold​​(oi​∣q)πθ​(oi​∣q)​ is the likelihood ratio or importance weight.

  • If $r(\theta )>1$ , the new policy πθ\pi_\thetaπθ​ is more likely to generate output oio_ioi​ than the old one was.
  • If $r(\theta )<1$ , the new policy πθ\pi_\thetaπθ​ is less likely to generate output oio_ioi​ .
  • If $r(\theta )=1$ , the policy has not changed with respect to this specific output.

What is its purpose? (Off-Policy Learning)

The primary purpose of this ratio is to enable off-policy learning.

In reinforcement learning, the ideal way to evaluate a policy is to use it to generate actions and see what rewards you get. However, generating new outputs ( oio_ioi​ ) from the model for every single gradient update is computationally very expensive.

Off-policy methods solve this. We can:

1. Generate a large batch of experiences (the outputs oio_ioi​ ) using the πθold\pi_{\theta_{\text{old}}}πθold​​ policy.
2. Then, perform several steps of optimization on our πθ\pi_\thetaπθ​ policy using that same batch of data.

The importance sampling ratio is the mathematical "correction factor" that allows us to estimate how good an action is under the new policy ( πθ\pi_\thetaπθ​ ) using data that was collected by the old policy ( πθold\pi_{\theta_{\text{old}}}πθold​​ ).

For the ease of discussion, I will equate:

The objective function multiplies this ratio by the advantage AiA_iAi​ (how good the output oio_ioi​ was). So, the update logic is:

• If oio_ioi​ was a good output ( Ai>0A_i > 0Ai​>0 ), we want to increase its probability. Maximising r(θ)×Air(\theta) \times A_ir(θ)×Ai​ will push $r(\theta )$ to be greater than 1, which in turn pushes πθ(oi∣q)\pi_\theta(o_i|q)πθ​(oi​∣q) to increase.
• If oio_ioi​ was a bad output ( Ai<0A_i < 0Ai​<0 ), we want to decrease its probability. Maximising r(θ)×Air(\theta) \times A_ir(θ)×Ai​ (a negative number) will push $r(\theta )$ to be less than 1, making the overall term less negative and thus decreasing πθ(oi∣q)\pi_\theta(o_i|q)πθ​(oi​∣q) .

Why is it designed that way? (Stability and PPO)

While the ratio allows for efficient learning, it is also a source of instability. If the new policy πθ\pi_\thetaπθ​ becomes very different from the old one πθold\pi_{\theta_{\text{old}}}πθold​​ , the ratio $r(\theta )$ can become extremely large or close to zero. A very large ratio would lead to a massive, noisy gradient update, potentially destroying all the learning the model has already done.

This is the problem that Proximal Policy Optimization (PPO), from which GRPO's objective is derived, was designed to solve. The design in Equation 1 is a direct implementation of the PPO-Clip objective (source: Proximal Policy Optimization).

The goal is to keep the new policy "proximal" (i.e., close) to the old policy. This creates a "trust region" where we can be confident the update is beneficial.

Part 3: clip(πθ(oi∣q)πθold(oi∣q),1−ϵ,1+ϵ)\text{clip}\left(\frac{\pi_\theta(o_i|q)}{\pi_{\theta_{\text{old}}}(o_i|q)}, 1-\epsilon, 1+\epsilon\right)clip(πθold​​(oi​∣q)πθ​(oi​∣q)​,1−ϵ,1+ϵ)

This is the core of the PPO algorithm, which encourages making the new policy πθ\pi_\thetaπθ​ more likely to produce high-advantage outputs, but "clips" the update to prevent it from changing to drastically and destabilsing training.

Example of clipping

Take for instance we set $ϵ$ to be 0.2. We then get the following clipping equation:

Examples of how the clipping equation works is below:

• $\text{clip}(1.15,0.8,1.2)=1.15$ because the value is within the range.
• $\text{clip}(1.5,0.8,1.2)=1.2$ because the value is beyond the range and is clipped down to the maximum value of 1.2.

We can see that with clipping, when the optimiser gets "greedy" and suggests a huge change, the model is still encouraged to make the output more likely, but is prevented from making a dangerously large jump in the policy.

Code Implementation

This logic is encapsulated within the grpo_loss_fn.

The probability ratio ri(θ)r_i(\theta)ri​(θ) is calculated in log-space for
numerical stability: ri(θ)=exp(log probθ−log probθoldr_i(\theta) = exp(\text{log prob}\theta - \text{log prob}{\theta_{\text{old}}}ri​(θ)=exp(log probθ−log probθold​ ).

# --- File: MLX LM GRPO.ipynb, Cell: grpo-helpers ---
def grpo_loss_fn(model, model_ref, sequences, a_toks, advantages, old_log_probs, beta, epsilon):
    """The GRPO loss function."""
    # Get log probs from the trainable model (π_θ)
    log_probs = calculate_log_probs(model, sequences, a_toks)

    # Get log probs from the reference model (π_ref) for KL penalty
    log_probs_ref = calculate_log_probs(model_ref, sequences, a_toks)

    # PPO-clip objective
    ratio = mx.exp(log_probs - old_log_probs)
    clipped_ratio = mx.clip(ratio, 1.0 - epsilon, 1.0 + epsilon)
    policy_reward = mx.minimum(ratio * advantages, clipped_ratio * advantages)

    # KL penalty
    # Step 1: Calculate log(r) where r = π_ref / π_θ
    # log(r) = log(π_ref) - log(π_θ)
    log_ratio_for_kl = log_probs_ref - log_probs

    # Step 2: Calculate r itself by exponentiating log(r)
    # r = exp(log(r))
    ratio_for_kl = mx.exp(log_ratio_for_kl)

    # Step 3: Apply the paper's full formula: r - log(r) - 1
    kl_div = ratio_for_kl - log_ratio_for_kl - 1

    # The objective is to maximize this, so we return the negative for minimization
    loss = -mx.mean(policy_reward - beta * kl_div)
    return loss, mx.mean(policy_reward), mx.mean(kl_div)

The helper function calculate_log_probs is responsible for computing log P(o_i | q) for a given policy.

def calculate_log_probs(model, sequences, a_toks):
    """Calculates the log probabilities of the generated answer tokens."""
    # Pass the full sequence (prompt + answer) to the model
    logits = model(sequences)

    # Convert to log probabilities
    log_probs_full = nn.log_softmax(logits, axis=-1)

    ## Find the actual positions where answer tokens should be extracted
    # This assumes a_toks contains the actual token IDs that were generated
    batch_size, seq_len = sequences.shape
    _, ans_len = a_toks.shape

    # Calculate the starting position for answer tokens (assuming they're at the end)
    start_pos = seq_len - ans_len

    # Extract log probabilities for the answer portion of the sequence
    answer_log_probs = log_probs_full[:, start_pos:start_pos+ans_len, :]

    # Create indices for gathering - ensure proper shape alignment
    indices = a_toks[:, :, None]

    # Extract log probabilities for the actual answer tokens
    selected_log_probs = mx.take_along_axis(answer_log_probs, indices, axis=-1).squeeze(-1)

    # Sum log probabilities across the answer sequence
    return mx.sum(selected_log_probs, axis=-1)

Part 4: min⁡(πθ(oi∣q)πθold(oi∣q)Ai,clip(πθ(oi∣q)πθold(oi∣q),1−ϵ,1+ϵ)Ai)\min\left( \frac{\pi_\theta(o_i|q)}{\pi_{\theta_{\text{old}}}(o_i|q)} A_i, \text{clip}\left(\frac{\pi_\theta(o_i|q)}{\pi_{\theta_{\text{old}}}(o_i|q)}, 1-\epsilon, 1+\epsilon\right) A_i \right)min(πθold​​(oi​∣q)πθ​(oi​∣q)​Ai​,clip(πθold​​(oi​∣q)πθ​(oi​∣q)​,1−ϵ,1+ϵ)Ai​)

The purpose of this min function is to act as a floor, i.e.:

• When an output is good (positive AiA_iAi​ ), it prevents the update from becoming too rewarding.
• When an output is bad (negative AiA_iAi​ ), it acts as a floor on the penalty.

In both cases, the min function prevents the model from making a large policy change too quickly.

Example

Using back the same clipping function in Part 3, where we set $ϵ=0.2$ , $r(\theta )=1.5$ , and we now assume we have an advantage value A1=10A_1 = 10A1​=10 . The equation we thus get is:

Part 5: DKL(πθ∣∣πref))D_{KL}\left(\pi_\theta || \pi_{\text{ref}}) \right)DKL​(πθ​∣∣πref​))

This function acts as a regulariser, penalising the policy πθ\pi_{\theta}πθ​ for deviating from a reference policy πref\pi_{\text{ref}}πref​ . (Note: When we say policy, we are actually referring to the LLM as weights, so $\pi \theta$ refers to the LLM with updated weights, while πref\pi_{\text{ref}}πref​ usually refers to the original stock LLM.)_

The paper defines the term as:

What is key is to recognise that this equation is not the standard Kulback-Leibler (KL) divergence, but a more specific, per-sample approximation chosen for being computationally cheaper. This compares to the standard KL divergence which can be expressed as:

Compared to the standard definition, there are two main differences:

1. The standard definition is an expectation over the entire distribution πθ\pi_{\theta}πθ​ , whereas the paper's variant is an expression for a single sample oio_ioi​ .
2. The functional form of the paper's variant is different from the term inside the standard expectation.

Deconstruction and Analysis of the variant of KL divergence

For simplicity, let the probability ratio be r:

Then Equation (2) defines a function $f(r)$ :

This function is evaluated for a single sample oio_ioi​ , which itself was drawn from the old policy, πθold\pi_{\theta_{old}}πθold​​ .

In order for $f(r)$ to be a valid divergence measure, it must satisfy two properties:

1. Non-negativity: $D(p||q)>=0$ .
2. Identity of Indiscernibles: $D(p||q)=0$ if and only if $p=q$ .

Proof that $f(r)$ satisfies the two properties is in Appendix A.

Show me the code

This term is a regularizer. It penalizes the objective if the trainable policy π_θ strays too far from the original, trusted reference policy π_ref, helping to prevent catastrophic forgetting.

Code Implementation: Also within grpo_loss_fn.

# --- File: MLX LM GRPO.ipynb, Cell: grpo-helpers ---
def grpo_loss_fn(...):
    ...
    # Get log probs from the reference model (π_ref)
    log_probs_ref = calculate_log_probs(model_ref, ...)

    # KL penalty
    # Step 1: Calculate log(r) where r = π_ref / π_θ
    # log(r) = log(π_ref) - log(π_θ)
    log_ratio_for_kl = log_probs_ref - log_probs

    # Step 2: Calculate r itself by exponentiating log(r)
    # r = exp(log(r))
    ratio_for_kl = mx.exp(log_ratio_for_kl)

    # Step 3: Apply the paper's full formula: r - log(r) - 1
    kl_div = ratio_for_kl - log_ratio_for_kl - 1
    ...

Part 6: AiA_iAi​

The Advantage Function AiA_iAi​ is a central component in modern policy gradient methods. Intuitively, the advantage tells us not just if an action was "good" (positive reward), but if it was "better than average". It is designed to reduce sensitivity to reward scaling, and stabilises training by preventing outlier rewards. A more technical discussion about the advantage function is available in Appendix B.

Given that the advantage AiA_iAi​ tells us how much better or worse a specific output oio_ioi​ was compared. to the average of its group, this requires two steps:

1. Calculate the raw reward rir_iri​ ; and
2. Normalising the output

In short, the mathematical equation is:

We can implement the code as such:

1. Reward calculation rir_iri​ as a simple rule-based reward.

# --- File: MLX LM GRPO.ipynb, Cell: grpo-loop ---
# Get Reward (r_i)
reward = 1.0 if batch_answers[i] in response else 0.0
group_rewards.append(reward)

1. Normalisating to get Advantage AiA_iAi​

# --- File: MLX LM GRPO.ipynb, Cell: grpo-loop ---
# Compute Advantages
advantages = []
for rewards in rollout_rewards:
    mean_reward = mx.mean(rewards)
    std_reward = mx.sqrt(mx.var(rewards)) + 1e-8 # Epsilon for stability
    # This line directly implements the advantage formula
    adv = (rewards - mean_reward) / std_reward 
    advantages.append(adv)

advantages = mx.concatenate(advantages)

Bringing it all together

Finally, we combine all the pieces, average over the batch, and negate the result to create a loss that can be minimised by the optimiser.

Mathematical Component: