[FEATURE] Add BAPO (Balanced Policy Optimization with Adaptive Clipping) policy loss

[WanLanglin]

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 BAPOPolicyLossFn in trinity/algorithm/policy_loss_fn/bapo_policy_loss.py, following the structure of PPOPolicyLossFn and SAPOPolicyLossFn, registered in POLICY_LOSS_FN with an algorithm preset (GRPO advantage, no critic).
• The bounds are explicit importance-ratio values (e.g. 0.6 / 1.2), not the 1 ± clip_range epsilons that PPOPolicyLossFn uses. The argument names will make this explicit.
• Metrics: the selected clip_low, clip_high and 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_group to 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!