[Kingpython123]

Motivation

Closes #621. BAPO is one of the algorithms listed in #470.

BAPO (Balanced Policy Optimization with Adaptive Clipping) targets off-policy / stale-data training (experience reuse, partial rollouts, async mode). With stale data, negative-advantage tokens tend to dominate the policy gradient, and a fixed clip range blocks the updates from low-probability positive tokens. BAPO keeps the PPO-clip surrogate but picks the clipping bounds for every update, which fits Trinity's async and off-policy modes.

Implementation

BAPOPolicyLossFn in trinity/algorithm/policy_loss_fn/bapo_policy_loss.py, following the structure of PPOPolicyLossFn / SAPOPolicyLossFn. It is registered as bapo in POLICY_LOSS_FN, with a bapo algorithm preset in ALGORITHM_TYPE that mirrors the SAPO preset (GRPO advantage, no critic).

Bound search (Algorithm 1). Starting from c_low = a-, c_high = a+, raise c_high by δ1 up to b+, then c_low by δ2 up to b-, until the positive contribution ratio ρ ≥ ρ0. If the target is never met, the widest bounds reached are used. The bounds are explicit importance-ratio values (e.g. 0.6 / 1.2), not 1 ± clip_range epsilons, and the argument names say so:

Argument	Paper	Default
lower_bound_start / lower_bound_max / lower_bound_step	a- / b- / δ2	0.6 / 0.9 / 0.02
upper_bound_start / upper_bound_max / upper_bound_step	a+ / b+ / δ1	1.2 / 3.0 / 0.05
target_positive_ratio	ρ0	0.4
clip_ratio_c	dual-clip, as in PPO	3.0

Defaults follow the paper's Algorithm 1 order and Sec. 5.1 hyperparameters (the reference code uses a different order and a pos/neg ratio target of 1; both are reachable through the arguments except the order).

Definition of ρ (Eq. 8). ρ is the share of the policy gradient that comes from positive-advantage tokens. Each token contributes ratio * |advantage| (the magnitude of its gradient coefficient on log π) when it is not clipped and nothing otherwise: a positive token is clipped when ratio > c_high, a negative token when ratio < c_low or ratio > clip_ratio_c. With this definition raising c_high admits more positive tokens and raising c_low drops low-ratio negative tokens, so both steps of Algorithm 1 move ρ towards the target. Measuring the clipped surrogate values instead would make raising c_low increase the negative share, which defeats the second step. When no negative token contributes, ρ = 1.

Cost. The schedule of candidate (c_low, c_high) pairs is fixed at construction (52 pairs with the defaults). All candidates are evaluated in one vectorized pass over the valid tokens, so the search needs a single host sync per call instead of one per loop iteration.

Scope. The search runs per call, i.e. per micro-batch on each rank (option (a) in the issue), so the policy-loss interface is unchanged. Synchronizing the contribution sums across DP ranks would need dp_group to be forwarded to the policy loss and can be a follow-up.

Metrics. In addition to the PPO metrics (pg_clipfrac, pg_clipfrac_lower, ppo_kl, pg_loss), the loss logs bapo_clip_low, bapo_clip_high and bapo_positive_ratio.

Example and docs

• examples/bapo_gsm8k/ (config + README), derived from the CISPO GSM8K example.
• BAPO added to the algorithm tables in README.md, README_zh.md, docs/sphinx_doc/source{,_zh}/main.md, and to the example lists in example_dataset_perspective.md (en/zh).

Tests

tests/algorithm/bapo_policy_loss_test.py (CPU-only, 11 tests, no verl import needed):

• the search matches a literal transcription of Algorithm 1 (with an independent token-by-token ρ) for 20 random batches × 4 targets;
• no adjustment when ρ0 is already met (target 0, and all-positive advantages);
• the upper bound is raised first; the lower bound is raised once the upper bound saturates (hand-computed case);
• all-negative advantages fall back to the widest bounds (0.9, 3.0);
• with collapsed ranges the loss and metrics equal PPOPolicyLossFn(clip_range_low=0.2, clip_range_high=0.28);
• gradients are finite, an empty mask does not fail, invalid arguments are rejected, and both registries resolve bapo.

python -m pytest tests/algorithm/bapo_policy_loss_test.py -q
11 passed

The example config also loads and passes AlgorithmConfigValidator, and pre-commit run passes on all changed files. I have not run an end-to-end GPU training job.