balanced-policy-optimization-rl

name	balanced-policy-optimization-rl
title	BAPO: Stabilizing Off-Policy RL for LLMs via Balanced Policy Optimization with Adaptive Clipping
version	0.0.2
engine	skillxiv-v0.0.2-claude-opus-4.6
license	MIT
url	https://arxiv.org/abs/2510.18927
keywords	["off-policy RL","LLM training","policy optimization","gradient stability","entropy"]
description	Stabilize off-policy RL for LLMs using adaptive clipping that dynamically rebalances positive/negative gradients and preserves entropy, improving mathematical reasoning performance vs standard PPO.

Technique: Balanced Policy Optimization — Adaptive Clipping for Stable LLM RL

Off-policy reinforcement learning for LLMs faces two fundamental instability problems: positive-advantage samples get drowned out by negative-advantage samples during gradient updates, and fixed clipping (PPO-style) systematically suppresses entropy-increasing updates, leading to premature convergence and over-exploitation.

BAPO solves both problems through adaptive clipping: instead of fixed clip ratios, dynamically adjust clipping bounds to balance gradients across positive and negative samples while explicitly preserving entropy. This enables more stable training and better final performance on complex reasoning tasks.

Core Concept

BAPO operates on three principles:

• Gradient Balancing: Detect when negative samples dominate, adjust clipping to balance contributions
• Entropy Preservation: Monitor when clipping suppresses entropy-increasing moves, override clipping for those updates
• Adaptive Thresholds: Learn task-specific clipping ranges during training rather than using fixed values
• Theoretical Grounding: Formalize as an "Entropy-Clip Rule" showing PPO's implicit entropy suppression

The result is more stable training curves and stronger final performance on mathematical reasoning (+state-of-the-art on AIME).

Architecture Overview

• Policy Network: LLM backbone producing action probabilities
• Advantage Estimator: Compute advantage from rewards or returns
• Adaptive Clipper: Monitor gradient distributions, adjust clip bounds dynamically
• Entropy Monitor: Track KL divergence, detect when clipping suppresses exploration
• Gradient Aggregator: Combine clipped and unclipped gradients intelligently
• Training Loop: Standard RL trajectory collection + BAPO-modified policy updates

Implementation Steps

The key innovation is replacing fixed PPO clipping with dynamic bounds that balance gradients. This example shows the core algorithm.

import torch
import torch.nn as nn
from typing import Tuple

class BalancedPolicyOptimizer:
    

     ():
        .model = model
        .clip_ratio = initial_clip_ratio
        .entropy_weight = entropy_weight
        .gradient_balance_target = gradient_balance_target
        .optimizer = torch.optim.Adam(model.parameters(), lr=)

     () -> torch.Tensor:
        
        advantages = torch.zeros_like(rewards)
        gae = 

         t  (((rewards))):
             t == (rewards) - :
                next_value = 
            :
                next_value = values[:, t + ]

            delta = rewards[:, t] + gamma * next_value - values[:, t]
            gae = delta + gamma * gae_lambda * gae
            advantages[:, t] = gae

         advantages

     () -> [torch.Tensor, ]:
        
        
        ratio = torch.exp(new_log_probs - old_log_probs)

        
        positive_mask = advantages > 
        negative_mask = advantages <= 

        
        clipped_ratio = torch.clamp(ratio,  - .clip_ratio,  + .clip_ratio)

        
        positive_loss = torch.where(
            positive_mask,
            torch.(ratio * advantages, clipped_ratio * advantages),
            torch.zeros_like(advantages)
        )

        negative_loss = torch.where(
            negative_mask,
            torch.(ratio * advantages, clipped_ratio * advantages),
            torch.zeros_like(advantages)
        )

        
        positive_grad = positive_loss.()
        negative_grad = -negative_loss.()  

        gradient_ratio = torch.(negative_grad) / (torch.(positive_grad) + )

        
         gradient_ratio > .gradient_balance_target:
            
            adaptive_clip_negative = .clip_ratio * ( +  * gradient_ratio)
            adaptive_ratio_neg = torch.clamp(
                ratio,
                 - adaptive_clip_negative,
                 + adaptive_clip_negative
            )
            negative_loss = torch.where(
                negative_mask,
                torch.(ratio * advantages, adaptive_ratio_neg * advantages),
                torch.zeros_like(advantages)
            )

        
        policy_loss = -(positive_loss + negative_loss).mean()

        
        entropy_loss = -.entropy_weight * entropy.mean()

        total_loss = policy_loss + entropy_loss

         total_loss, {
            : policy_loss.item(),
            : entropy_loss.item(),
            : gradient_ratio.item(),
            : .clip_ratio
        }

     ():
        
         gradient_ratio > .gradient_balance_target:
            
            .clip_ratio = (.clip_ratio * , )
         gradient_ratio < .gradient_balance_target * :
            
            .clip_ratio = (.clip_ratio * , )


 () -> :
    
    
    advantages = optimizer.compute_advantages(rewards, values)

    
    advantages = (advantages - advantages.mean()) / (advantages.std() + )

    
    loss, metrics = optimizer.compute_policy_loss_with_adaptive_clipping(
        old_log_probs,
        new_log_probs,
        advantages,
        entropy
    )

    
    optimizer.optimizer.zero_grad()
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), )
    optimizer.optimizer.step()

    
    optimizer.update_adaptive_clip_ratio(metrics[])

     metrics