Repository navigation
Description
Feature Request
Following the algorithm item in #470, which lists BAPO as an example, I'd like to add BAPO (Balanced Policy Optimization with Adaptive Clipping) as a policy loss.
- Paper: https://arxiv.org/abs/2510.18927
- Authors' verl-based reference: https://github.com/WooooDyy/BAPO (
recipe/bapo/policy_loss.py)
BAPO keeps the PPO-clip surrogate but picks the clipping bounds for every update. Starting from c_low = a- and c_high = a+, it raises c_high by δ1 (up to b+) and then c_low by δ2 (up to b-) until positive-advantage tokens contribute at least ρ0 of the policy-gradient loss (Eq. 8 and Algorithm 1 in the paper).
Purpose of the Feature
With stale data (experience reuse over several epochs, partial rollouts, async training), the paper finds that negative-advantage tokens dominate the gradient and that fixed clipping blocks the entropy-increasing updates from low-probability positive tokens. BAPO addresses both without hand-tuning the clip range, which fits Trinity's async and off-policy modes.
Expected Implementation (Optional)
- A
BAPOPolicyLossFnintrinity/algorithm/policy_loss_fn/bapo_policy_loss.py, following the structure ofPPOPolicyLossFnandSAPOPolicyLossFn, registered inPOLICY_LOSS_FNwith an algorithm preset (GRPO advantage, no critic). - The bounds are explicit importance-ratio values (e.g. 0.6 / 1.2), not the
1 ± clip_rangeepsilons thatPPOPolicyLossFnuses. The argument names will make this explicit. - Metrics: the selected
clip_low,clip_highand the positive-contribution ratio. - Unit tests: the bound search checked against a literal transcription of Algorithm 1, edge cases (all-positive or all-negative advantages, both bounds saturating), and equivalence to PPO clipping when the movable ranges are collapsed to a point.
- An example config on GSM8K.
Questions before I start
1. The reference code differs from the paper. Which should be the default?
| Paper (Algorithm 1, Sec. 5 settings) | Reference code (bapo_trainer.yaml) |
|
|---|---|---|
| Search order | raise c_high first, then c_low |
raise c_low first, then c_high |
| Target | positive share of the loss, ρ0 = 0.4 |
positive / negative ratio = 1 (i.e. share 0.5) |
| Lower bound range / step | 0.6 → 0.9, δ2 = 0.02 |
0.6 → 0.8, step 0.05 |
| Upper bound range / step | 1.2 → 3.0, δ1 = 0.05 |
1.2 → 2.0, step 0.05 |
My plan is to follow the paper's Algorithm 1 and hyperparameters by default and expose the ranges, steps and ρ0 as arguments.
2. Where should the bound search be synchronized?
As far as I can tell, TrinityPolicyLoss calls the policy loss once per micro-batch and does not forward dp_group to policy_loss_fn. Without synchronization, each rank and micro-batch picks its own bounds, so the all-reduced gradient mixes objectives. The reference code avoids this by all-reducing the positive and negative sums. Options:
- (a) Search per micro-batch on each rank. Simplest, and needs no interface change.
- (b) Forward
dp_groupto the policy loss and all-reduce the contribution sums, so all ranks agree (still per micro-batch).
Would you prefer (a) for a first version, or is changing the loss interface for (b) acceptable?
3. Should this land in trinity/plugins/ first, or go directly into trinity/algorithm/?
Additional Information
I only have a single 16 GB GPU, so I plan to validate with unit tests and a short small-model run (e.g. Qwen3-0.6B on GSM8K). I can't reproduce the paper-scale results.
Are You Willing to Submit a PR?
- Yes, I am willing to submit a PR!
Activity
Kingpython123 commentedon Oct 31, 2026
Hi, since there has been no PR for this in a while, I'm going to open one. Plan, following the proposal above:
BAPOPolicyLossFnintrinity/algorithm/policy_loss_fn/bapo_policy_loss.py, registered asbapo, plus abapoalgorithm preset (GRPO advantage, no critic).c_highby δ1 up to b+, thenc_lowby δ2 up to b-) and Sec. 5 hyperparameters (ρ0 = 0.4, [0.6, 0.9], [1.2, 3.0], δ1 = 0.05, δ2 = 0.02), all configurable.@WanLanglin if you already have a branch for this, let me know and I'll step back.