Dynamically reweight advantage functions based on trajectory certainty to improve policy optimization in foundation models. Addresses advantage reversion and mirror problems by mixing standardized and mean-normalized advantage formulations. Enables more stable gradient signals across high- and low-certainty samples.
Dynamically reweight advantage functions based on trajectory certainty to improve policy optimization in foundation models. Addresses advantage reversion and mirror problems by mixing standardized and mean-normalized advantage formulations. Enables more stable gradient signals across high- and low-certainty samples.
Group Relative Policy Optimization (GRPO) has become the standard approach for aligning foundation models through reinforcement learning, but existing advantage formulations suffer from two critical problems. When calculating advantage estimates, GRPO normalizes by standard deviation, which works poorly when samples have vastly different certainty levels. High-certainty samples (where most trajectories succeed or fail) create numerical instability and receive disproportionate advantage signals, while symmetric reward distributions produce identical normalized advantages despite representing semantically distinct outcomes. These issues result in misleading gradient signals during policy updates, hindering proper sample ranking and skill improvement.
MAPO solves this by recognizing that trajectories exhibit varying levels of certainty during training. Rather than applying a fixed advantage formulation to all samples, MAPO adaptively mixes two advantage signals—a variance-sensitive standardized approach for uncertain trajectories and a stable mean-normalized approach for confident ones. This trajectory-aware design ensures consistent, interpretable gradient signals across all sample types.
Core Concept
MAPO's core insight is that advantage computation should adapt based on trajectory success distribution. The method measures trajectory certainty using the success ratio p (proportion of trajectories that achieve positive rewards), then uses this to dynamically interpolate between two advantage formulations:
Standard deviation-based advantage for uncertain samples (p ≈ 0.5): Emphasizes relative differences in returns, helpful when outcomes vary widely.
Mean-based advantage for certain samples (p → 0 or 1): Uses proportional deviation from the mean, providing stable signals when outcomes cluster around success or failure.
The interpolation weight λ(p) = 1 - 4p(1-p) smoothly transitions between these formulations. At p=0.5 (maximum uncertainty), λ=0 and standardized advantage dominates. At p→0 or p→1 (high certainty), λ→1 and mean-based advantage dominates. This design ensures gradients remain interpretable: low-certainty samples receive larger gradient multipliers (boosting weak signals), while high-certainty samples receive smaller multipliers (dampening noisy signals).
Architecture Overview
Trajectory Certainty Measurement: Calculate success ratio p = N/G (number of successful trajectories divided by total trajectories).
Advantage Percent Deviation (APD): Compute proportional deviation from mean without relying on standard deviation: A_i^APD = (r_i - μ) / μ.
Trajectory Certainty Reweight (TCR): Determine mixing weight λ(p) using the Bernoulli variance function.
Mixed Advantage Construction: Interpolate between standardized (high variance sensitivity) and mean-based (high stability) formulations based on λ(p).
Policy Update: Apply GRPO with mixed advantages using clipped policy gradient objective with KL divergence regularization.
Implementation
Step 1: Compute Reward Statistics and Success Ratio
Before constructing advantages, compute aggregate statistics across the trajectory batch. For a group of G trajectories with rewards r_1...r_G, calculate the mean μ, standard deviation σ, and success ratio p. The success ratio treats binary rewards (r_i ∈ {0,1}) as Bernoulli trials: p estimates the empirical probability that a trajectory achieves the target objective.
import numpy as np
defcompute_reward_statistics(rewards):
"""
Compute mean, std, and success ratio from reward batch.
Args:
rewards: ndarray of shape (G,) containing per-trajectory rewards
Returns:
mu: float, mean reward
sigma: float, standard deviation of rewards
p: float, success ratio (fraction of rewards == 1)
"""
mu = np.mean(rewards)
sigma = np.std(rewards) + 1e-8# Small epsilon for numerical stability
success_count = np.sum(rewards == 1)
p = success_count / len(rewards)
return mu, sigma, p
The epsilon term prevents division by zero when all trajectories have identical rewards (zero variance edge case).
Step 2: Compute Certainty Weighting Function
The trajectory certainty reweight function λ(p) determines how much to trust each advantage formulation. It's based on the variance of a Bernoulli distribution with success probability p. At p=0.5, the Bernoulli variance p(1-p) is maximum (0.25), making trajectories most uncertain. At p→0 or p→1, the variance approaches zero, making outcomes deterministic and certain.
defcompute_certainty_weight(p):
"""
Compute trajectory certainty weight using Bernoulli variance.
Args:
p: float in [0, 1], success ratio
Returns:
lambda_weight: float in [0, 1], mixing weight between formulations
"""
lambda_weight = 1 - 4 * p * (1 - p)
return lambda_weight
This function guarantees λ(p) ∈ [0, 1] with minimum value 0 at p=0.5 and maximum values 1 at p∈{0, 1}. The factor 4 normalizes the Bernoulli variance to the [0, 1] range.
Step 3: Construct Standardized and Mean-Normalized Advantages
Compute both advantage formulations for every trajectory. The standardized advantage A_i is the classic z-score normalization. The advantage percent deviation A_i^APD is a mean-normalized variant that measures proportional deviation without relying on variance.
defconstruct_advantages(rewards, mu, sigma):
"""
Compute both standardized and mean-normalized advantages.
Args:
rewards: ndarray of shape (G,)
mu: float, mean reward
sigma: float, std deviation (with epsilon)
Returns:
adv_std: ndarray of shape (G,), standardized advantages
adv_apd: ndarray of shape (G,), mean-normalized advantages
"""
adv_std = (rewards - mu) / sigma
# Avoid division by zero if mean is near zero
adv_apd = np.where(
np.abs(mu) > 1e-8,
(rewards - mu) / mu,
np.zeros_like(rewards)
)
return adv_std, adv_apd
The mean-normalized form naturally handles cases where μ is very small by zeroing the advantage, preventing extreme values.
Step 4: Mix Advantages Using Trajectory Certainty Weight
Interpolate between the two advantage formulations using the certainty weight λ(p). Low-certainty samples (λ ≈ 0) rely on standardized advantages, while high-certainty samples (λ ≈ 1) rely on mean-normalized advantages. This adaptive mixing ensures stable gradient signals across the full spectrum of sample certainty.
defmix_advantages(adv_std, adv_apd, lambda_weight):
"""
Create mixed advantage by interpolating between two formulations.
Args:
adv_std: ndarray of shape (G,), standardized advantages
adv_apd: ndarray of shape (G,), mean-normalized advantages
lambda_weight: float in [0, 1], mixing coefficient
Returns:
adv_mixed: ndarray of shape (G,), interpolated advantages
"""
adv_mixed = (1 - lambda_weight) * adv_std + lambda_weight * adv_apd
return adv_mixed
When λ=0 (p=0.5): adv_mixed = adv_std (pure standardized). When λ=1 (p→0 or 1): adv_mixed = adv_apd (pure mean-normalized).
Step 5: Integrate into GRPO Policy Update
Apply the mixed advantages within the standard GRPO objective. GRPO uses a clipped policy ratio update similar to PPO, but with a KL divergence penalty to the reference model. The policy is updated to increase log-probability of trajectories with positive mixed advantages and decrease log-probability of negative advantage trajectories.
import torch
import torch.nn.functional as F
defgrpo_policy_loss(
log_prob_new,
log_prob_old,
advantages,
epsilon=0.2,
beta_kl=0.01):
"""
Compute GRPO loss with mixed advantages.
Args:
log_prob_new: tensor of shape (G,), log probability under new policy
log_prob_old: tensor of shape (G,), log probability under old policy
advantages: tensor of shape (G,), mixed advantages
epsilon: float, clipping range for policy ratio
beta_kl: float, KL divergence penalty weight
Returns:
loss: scalar tensor, negative of objective (for minimization)
"""
ratio = torch.exp(log_prob_new - log_prob_old)
# Clipped policy gradient (similar to PPO)
surr_unclipped = ratio * advantages
surr_clipped = torch.clamp(ratio, 1 - epsilon, 1 + epsilon) * advantages
policy_loss = -torch.min(surr_unclipped, surr_clipped).mean()
# KL penalty to reference policy (prevents divergence)
kl_penalty = beta_kl * (log_prob_old - log_prob_new).mean()
total_loss = policy_loss + kl_penalty
return total_loss
This loss encourages the policy to match high-advantage trajectories while regularizing divergence from the reference model.
Step 6: Complete Training Loop
Combine all components into a single training iteration. For each batch of trajectories, compute mixed advantages and apply policy updates.
defmapo_training_step(
trajectories,
log_probs_old,
rewards,
policy_model,
optimizer,
epsilon=0.2,
beta_kl=0.01):
"""
Execute one MAPO training step with trajectory-aware advantages.
Args:
trajectories: list of trajectory objects or tensors
log_probs_old: tensor of shape (G,), old policy log probs
rewards: tensor of shape (G,), reward per trajectory
policy_model: torch.nn.Module, the policy network
optimizer: torch optimizer
epsilon: float, clipping range
beta_kl: float, KL weight
Returns:
loss: float, scalar loss value
"""# Step 1: Compute statistics
mu = rewards.mean()
sigma = rewards.std() + 1e-8
success_count = (rewards == 1).sum().float()
p = success_count / len(rewards)
# Step 2: Compute certainty weight
lambda_weight = 1 - 4 * p * (1 - p)
# Step 3: Construct both advantage types
adv_std = (rewards - mu) / sigma
adv_apd = torch.where(
mu.abs() > 1e-8,
(rewards - mu) / mu,
torch.zeros_like(rewards)
)
# Step 4: Mix advantages
advantages = (1 - lambda_weight) * adv_std + lambda_weight * adv_apd
# Normalize advantages for stability
advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
# Step 5: Forward pass and compute loss
log_probs_new = policy_model(trajectories)
loss = grpo_policy_loss(
log_probs_new,
log_probs_old,
advantages,
epsilon=epsilon,
beta_kl=beta_kl
)
# Step 6: Backprop and update
optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(policy_model.parameters(), max_norm=1.0)
optimizer.step()
return loss.item()
This loop handles the full pipeline: computing statistics, determining certainty, constructing mixed advantages, computing loss, and updating the policy.
Trajectories per update; larger = more stable statistics
learning_rate
5e-5
[1e-5, 1e-4]
Policy model learning rate; tune per foundation model
advantage_norm_eps
1e-8
[1e-10, 1e-6]
Epsilon for advantage normalization; prevents division by zero
When to Use MAPO
Reasoning and coding tasks where trajectories have binary success/failure outcomes (correct answer vs. incorrect).
Imbalanced reward distributions where some queries lead to mostly successes (p → 1) and others mostly failures (p → 0).
Foundation model alignment at scale (7B+ parameters) where stable gradient signals across diverse samples improve convergence.
Long-horizon reasoning where trajectory certainty varies widely across different problem types within a single batch.
When NOT to Use MAPO
Continuous control tasks with dense, smooth reward functions (MAPO designed for discrete binary outcomes).
Small batch sizes (< 32 trajectories) where success ratio p becomes unstable and unreliable.
Uniform reward distributions where all samples have similar certainty; standard GRPO suffices.
Real-time or latency-critical systems where the extra computation of dual advantage formulations is unacceptable.
Common Pitfalls
Forgetting advantage normalization: Always normalize mixed advantages (subtract mean, divide by std + epsilon) before use. Unnormalized advantages can have arbitrary scales, breaking the clipping mechanism and destabilizing training.
Using rewards outside [0, 1]: The method assumes binary success/failure. Scale non-binary rewards to [0, 1] or adapt the success ratio calculation (e.g., use percentiles instead of binary thresholding).
Small batch sizes: Computing p from < 32 samples produces noisy estimates. If batch size is small, use a moving average or episodic estimate of p from historical data.
Ignoring epsilon stability in APD: When μ is very close to zero, the mean-normalized advantage A^APD can explode. Always use np.where or similar guards to zero the advantage when |μ| < 1e-8.
Not adjusting beta_kl: KL penalty strength should scale with batch size and policy learning rate. If loss oscillates, increase beta_kl; if policy drifts from reference, decrease it.
