From b8cc8442a58d270a656dbafded67312763516a9e Mon Sep 17 00:00:00 2001 From: Kingpython123 Date: Sat, 31 Oct 2026 11:02:17 -0700 Subject: [PATCH] Add BAPO policy loss with adaptive clipping bounds (#621) Add BAPOPolicyLossFn (arXiv:2510.18927), which keeps the PPO-clip surrogate but searches the clipping bounds for every call. Starting from [a-, a+], the upper bound is raised by delta_1 up to b+ and then the lower bound by delta_2 up to b-, until positive-advantage tokens contribute at least rho_0 of the policy gradient (Eq. 8, Algorithm 1). Bounds are explicit importance-ratio values. All candidate bounds are evaluated in one vectorized pass, so the search costs a single host sync. - Register the loss as bapo and add a bapo algorithm preset (GRPO advantage, no critic) - Log the selected bounds and the positive contribution ratio - Add unit tests against a literal transcription of Algorithm 1, edge cases and equivalence to PPO with fixed bounds - Add a GSM8K example and list BAPO in the README and docs (en/zh) --- README.md | 1 + README_zh.md | 1 + docs/sphinx_doc/source/main.md | 1 + .../tutorial/example_dataset_perspective.md | 1 + docs/sphinx_doc/source_zh/main.md | 1 + .../tutorial/example_dataset_perspective.md | 1 + examples/bapo_gsm8k/README.md | 9 + examples/bapo_gsm8k/gsm8k.yaml | 76 ++++++ tests/algorithm/bapo_policy_loss_test.py | 196 +++++++++++++++ trinity/algorithm/__init__.py | 1 + trinity/algorithm/algorithm.py | 27 ++ trinity/algorithm/policy_loss_fn/__init__.py | 1 + .../policy_loss_fn/bapo_policy_loss.py | 236 ++++++++++++++++++ 13 files changed, 552 insertions(+) create mode 100644 examples/bapo_gsm8k/README.md create mode 100644 examples/bapo_gsm8k/gsm8k.yaml create mode 100644 tests/algorithm/bapo_policy_loss_test.py create mode 100644 trinity/algorithm/policy_loss_fn/bapo_policy_loss.py diff --git a/README.md b/README.md index ae6cf35dd1c..e2e75dc04e1 100644 --- a/README.md +++ b/README.md @@ -132,6 +132,7 @@ Trinity-RFT provides functionalities for users with different backgrounds and ob | AsymRE [[Paper](https://arxiv.org/pdf/2506.20520)] | [[GSM8K Example](https://github.com/agentscope-ai/Trinity-RFT/tree/main/examples/asymre_gsm8k)] | [[Code](https://github.com/agentscope-ai/Trinity-RFT/tree/main/trinity/algorithm/advantage_fn/asymre_advantage.py)] | `algorithm_type: asymre` | | CISPO [[Paper](https://arxiv.org/pdf/2506.13585)] | - | [[Code](https://github.com/agentscope-ai/Trinity-RFT/tree/main/trinity/algorithm/policy_loss_fn/cispo_policy_loss.py)] | `algorithm_type: cispo` | | SAPO [[Paper](https://arxiv.org/pdf/2511.20347)] | - | [[Code](https://github.com/agentscope-ai/Trinity-RFT/tree/main/trinity/algorithm/policy_loss_fn/sapo_policy_loss.py)] | `algorithm_type: sapo` | +| BAPO [[Paper](https://arxiv.org/pdf/2510.18927)] | [[GSM8K Example](https://github.com/agentscope-ai/Trinity-RFT/tree/main/examples/bapo_gsm8k)] | [[Code](https://github.com/agentscope-ai/Trinity-RFT/tree/main/trinity/algorithm/policy_loss_fn/bapo_policy_loss.py)] | `algorithm_type: bapo` | | On-Policy Distillation [[Blog](https://thinkingmachines.ai/blog/on-policy-distillation/)] [[Paper](https://arxiv.org/pdf/2306.13649)] | [[GSM8K Example](https://github.com/agentscope-ai/Trinity-RFT/tree/main/examples/opd_gsm8k)] | [[Code](https://github.com/agentscope-ai/Trinity-RFT/tree/main/trinity/common/workflows/on_policy_distill_workflow.py)] | `algorithm_type: on_policy_distill` | | JSD (Jensen-Shannon Divergence) | [[GSM8K Example](https://github.com/agentscope-ai/Trinity-RFT/tree/main/examples/opd_gsm8k/opd_gsm8k_jsd.yaml)] | [[Code](https://github.com/agentscope-ai/Trinity-RFT/tree/main/trinity/algorithm/advantage_fn/jsd_advantage.py)] | `algorithm_type: jsd` | diff --git a/README_zh.md b/README_zh.md index 1442f841c51..5f9145c0608 100644 --- a/README_zh.md +++ b/README_zh.md @@ -152,6 +152,7 @@ Trinity-RFT 面向不同背景和目标的用户提供相应功能： | AsymRE [[论文](https://arxiv.org/pdf/2506.20520)] | [[GSM8K 示例](https://github.com/agentscope-ai/Trinity-RFT/tree/main/examples/asymre_gsm8k)] | [[代码](https://github.com/agentscope-ai/Trinity-RFT/tree/main/trinity/algorithm/advantage_fn/asymre_advantage.py)] | `algorithm_type: asymre` | | CISPO [[论文](https://arxiv.org/pdf/2506.13585)] | - | [[代码](https://github.com/agentscope-ai/Trinity-RFT/tree/main/trinity/algorithm/policy_loss_fn/cispo_policy_loss.py)] | `algorithm_type: cispo` | | SAPO [[论文](https://arxiv.org/pdf/2511.20347)] | - | [[代码](https://github.com/agentscope-ai/Trinity-RFT/tree/main/trinity/algorithm/policy_loss_fn/sapo_policy_loss.py)] | `algorithm_type: sapo` | +| BAPO [[论文](https://arxiv.org/pdf/2510.18927)] | [[GSM8K 例子](https://github.com/agentscope-ai/Trinity-RFT/tree/main/examples/bapo_gsm8k)] | [[代码](https://github.com/agentscope-ai/Trinity-RFT/tree/main/trinity/algorithm/policy_loss_fn/bapo_policy_loss.py)] | `algorithm_type: bapo` | | On-Policy Distillation [[博客](https://thinkingmachines.ai/blog/on-policy-distillation/)] [[论文](https://arxiv.org/pdf/2306.13649)] | [[GSM8K 示例](https://github.com/agentscope-ai/Trinity-RFT/tree/main/examples/opd_gsm8k)] | [[代码](https://github.com/agentscope-ai/Trinity-RFT/tree/main/trinity/common/workflows/on_policy_distill_workflow.py)] | `algorithm_type: on_policy_distill` | | JSD（Jensen-Shannon 散度） | [[GSM8K 示例](https://github.com/agentscope-ai/Trinity-RFT/tree/main/examples/opd_gsm8k/opd_gsm8k_jsd.yaml)] | [[代码](https://github.com/agentscope-ai/Trinity-RFT/tree/main/trinity/algorithm/advantage_fn/jsd_advantage.py)] | `algorithm_type: jsd` | diff --git a/docs/sphinx_doc/source/main.md b/docs/sphinx_doc/source/main.md index 066fbbc0f42..bcadf1c4352 100644 --- a/docs/sphinx_doc/source/main.md +++ b/docs/sphinx_doc/source/main.md @@ -88,6 +88,7 @@ We list some algorithms supported by Trinity-RFT in the following table. For mor | AsymRE [[Paper](https://arxiv.org/pdf/2506.20520)] | [[GSM8K Example](https://github.com/agentscope-ai/Trinity-RFT/tree/main/examples/asymre_gsm8k)] | [[Code](https://github.com/agentscope-ai/Trinity-RFT/tree/main/trinity/algorithm/advantage_fn/asymre_advantage.py)] | `algorithm_type: asymre` | | CISPO [[Paper](https://arxiv.org/pdf/2506.13585)] | - | [[Code](https://github.com/agentscope-ai/Trinity-RFT/tree/main/trinity/algorithm/policy_loss_fn/cispo_policy_loss.py)] | `algorithm_type: cispo` | | SAPO [[Paper](https://arxiv.org/pdf/2511.20347)] | - | [[Code](https://github.com/agentscope-ai/Trinity-RFT/tree/main/trinity/algorithm/policy_loss_fn/sapo_policy_loss.py)] | `algorithm_type: sapo` | +| BAPO [[Paper](https://arxiv.org/pdf/2510.18927)] | [[GSM8K Example](https://github.com/agentscope-ai/Trinity-RFT/tree/main/examples/bapo_gsm8k)] | [[Code](https://github.com/agentscope-ai/Trinity-RFT/tree/main/trinity/algorithm/policy_loss_fn/bapo_policy_loss.py)] | `algorithm_type: bapo` | | On-Policy Distillation [[Blog](https://thinkingmachines.ai/blog/on-policy-distillation/)] [[Paper](https://arxiv.org/pdf/2306.13649)] | [[GSM8K Example](https://github.com/agentscope-ai/Trinity-RFT/tree/main/examples/opd_gsm8k)] | [[Code](https://github.com/agentscope-ai/Trinity-RFT/tree/main/trinity/common/workflows/on_policy_distill_workflow.py)] | `algorithm_type: on_policy_distill` | | JSD (Jensen-Shannon Divergence) | [[GSM8K Example](https://github.com/agentscope-ai/Trinity-RFT/tree/main/examples/opd_gsm8k/opd_gsm8k_jsd.yaml)] | [[Code](https://github.com/agentscope-ai/Trinity-RFT/tree/main/trinity/algorithm/advantage_fn/jsd_advantage.py)] | `algorithm_type: jsd` | diff --git a/docs/sphinx_doc/source/tutorial/example_dataset_perspective.md b/docs/sphinx_doc/source/tutorial/example_dataset_perspective.md index 786fa230479..ff064b2b455 100644 --- a/docs/sphinx_doc/source/tutorial/example_dataset_perspective.md +++ b/docs/sphinx_doc/source/tutorial/example_dataset_perspective.md @@ -11,6 +11,7 @@ This guide provides an example list from the dataset perspective, where you can | | Multi-Step GRPO | AgentScope ReAct agent training | [example](https://github.com/agentscope-ai/Trinity-RFT/tree/main/examples/agentscope_react), [doc](https://agentscope-ai.github.io/Trinity-RFT/en/main/tutorial/example_react.html) | | | AsymRE | Regular RFT | [example](https://github.com/agentscope-ai/Trinity-RFT/tree/main/examples/asymre_gsm8k) | | | CISPO | Regular RFT | [example](https://github.com/agentscope-ai/Trinity-RFT/tree/main/examples/cispo_gsm8k) | +| | BAPO | Regular RFT with adaptive clipping | [example](https://github.com/agentscope-ai/Trinity-RFT/tree/main/examples/bapo_gsm8k) | | | GRPO | Training with prioritized tasks | [example](https://github.com/agentscope-ai/Trinity-RFT/tree/main/examples/grpo_gsm8k_task_pipeline), [doc](https://agentscope-ai.github.io/Trinity-RFT/en/main/tutorial/example_data_functionalities.html#example-data-processor-for-task-pipeline) | | | GRPO | Training with reward reshaping on experiences | [example](https://github.com/agentscope-ai/Trinity-RFT/tree/main/examples/grpo_gsm8k_experience_pipeline), [doc](https://agentscope-ai.github.io/Trinity-RFT/en/main/tutorial/example_data_functionalities.html#example-data-processor-for-experience-pipeline) | | | GRPO | Training with RULER (Relative Universal LLM-Elicited Rewards) | [example](https://github.com/agentscope-ai/Trinity-RFT/tree/main/examples/grpo_gsm8k_ruler) | diff --git a/docs/sphinx_doc/source_zh/main.md b/docs/sphinx_doc/source_zh/main.md index 9ed175386c6..f0c0deb25bc 100644 --- a/docs/sphinx_doc/source_zh/main.md +++ b/docs/sphinx_doc/source_zh/main.md @@ -83,6 +83,7 @@ Trinity-RFT 面向不同背景和目标的用户提供相应功能： | AsymRE [[论文](https://arxiv.org/pdf/2506.20520)] | [[GSM8K 示例](https://github.com/agentscope-ai/Trinity-RFT/tree/main/examples/asymre_gsm8k)] | [[代码](https://github.com/agentscope-ai/Trinity-RFT/tree/main/trinity/algorithm/advantage_fn/asymre_advantage.py)] | `algorithm_type: asymre` | | CISPO [[论文](https://arxiv.org/pdf/2506.13585)] | - | [[代码](https://github.com/agentscope-ai/Trinity-RFT/tree/main/trinity/algorithm/policy_loss_fn/cispo_policy_loss.py)] | `algorithm_type: cispo` | | SAPO [[论文](https://arxiv.org/pdf/2511.20347)] | - | [[代码](https://github.com/agentscope-ai/Trinity-RFT/tree/main/trinity/algorithm/policy_loss_fn/sapo_policy_loss.py)] | `algorithm_type: sapo` | +| BAPO [[论文](https://arxiv.org/pdf/2510.18927)] | [[GSM8K 例子](https://github.com/agentscope-ai/Trinity-RFT/tree/main/examples/bapo_gsm8k)] | [[代码](https://github.com/agentscope-ai/Trinity-RFT/tree/main/trinity/algorithm/policy_loss_fn/bapo_policy_loss.py)] | `algorithm_type: bapo` | | On-Policy Distillation [[博客](https://thinkingmachines.ai/blog/on-policy-distillation/)] [[论文](https://arxiv.org/pdf/2306.13649)] | [[GSM8K 示例](https://github.com/agentscope-ai/Trinity-RFT/tree/main/examples/opd_gsm8k)] | [[代码](https://github.com/agentscope-ai/Trinity-RFT/tree/main/trinity/common/workflows/on_policy_distill_workflow.py)] | `algorithm_type: on_policy_distill` | | JSD（Jensen-Shannon 散度） | [[GSM8K 示例](https://github.com/agentscope-ai/Trinity-RFT/tree/main/examples/opd_gsm8k/opd_gsm8k_jsd.yaml)] | [[代码](https://github.com/agentscope-ai/Trinity-RFT/tree/main/trinity/algorithm/advantage_fn/jsd_advantage.py)] | `algorithm_type: jsd` | diff --git a/docs/sphinx_doc/source_zh/tutorial/example_dataset_perspective.md b/docs/sphinx_doc/source_zh/tutorial/example_dataset_perspective.md index 07a42551c59..07090c1e19b 100644 --- a/docs/sphinx_doc/source_zh/tutorial/example_dataset_perspective.md +++ b/docs/sphinx_doc/source_zh/tutorial/example_dataset_perspective.md @@ -11,6 +11,7 @@ | | Multi-Step GRPO | AgentScope ReAct 智能体训练 | [样例位置](https://github.com/agentscope-ai/Trinity-RFT/tree/main/examples/agentscope_react), [相关文档](https://agentscope-ai.github.io/Trinity-RFT/zh/main/tutorial/example_react.html) | | | AsymRE | 常规 RFT | [样例位置](https://github.com/agentscope-ai/Trinity-RFT/tree/main/examples/asymre_gsm8k) | | | CISPO | 常规 RFT | [样例位置](https://github.com/agentscope-ai/Trinity-RFT/tree/main/examples/cispo_gsm8k) | +| | BAPO | 自适应裁剪的常规 RFT | [样例位置](https://github.com/agentscope-ai/Trinity-RFT/tree/main/examples/bapo_gsm8k) | | | GRPO | 使用优先级任务进行训练 | [样例位置](https://github.com/agentscope-ai/Trinity-RFT/tree/main/examples/grpo_gsm8k_task_pipeline), [相关文档](https://agentscope-ai.github.io/Trinity-RFT/en/main/tutorial/example_data_functionalities.html#example-data-processor-for-task-pipeline) | | | GRPO | 在经验上进行奖励重塑的训练 | [样例位置](https://github.com/agentscope-ai/Trinity-RFT/tree/main/examples/grpo_gsm8k_experience_pipeline), [相关文档](https://agentscope-ai.github.io/Trinity-RFT/zh/main/tutorial/example_data_functionalities.html#example-data-processor-for-experience-pipeline) | | | GRPO | 使用 RULER (Relative Universal LLM-Elicited Rewards) 进行训练 | [样例位置](https://github.com/agentscope-ai/Trinity-RFT/tree/main/examples/grpo_gsm8k_ruler) | diff --git a/examples/bapo_gsm8k/README.md b/examples/bapo_gsm8k/README.md new file mode 100644 index 00000000000..5a919284bce --- /dev/null +++ b/examples/bapo_gsm8k/README.md @@ -0,0 +1,9 @@ +# BAPO on GSM8K dataset + +This example shows the usage of [BAPO](https://arxiv.org/abs/2510.18927) (Balanced Policy Optimization with Adaptive Clipping) on the GSM8K dataset. + +BAPO keeps the PPO-clip surrogate but chooses the clipping bounds for every update. Starting from `[lower_bound_start, upper_bound_start]`, it first raises the upper bound by `upper_bound_step` (up to `upper_bound_max`) and then the lower bound by `lower_bound_step` (up to `lower_bound_max`), until tokens with positive advantages contribute at least `target_positive_ratio` of the policy gradient. The bounds are explicit importance-ratio values, not the `1 ± clip_range` epsilons used by PPO. The default values follow Sec. 5.1 of the paper. + +The selected bounds and the resulting positive contribution ratio are logged as `actor/bapo_clip_low`, `actor/bapo_clip_high` and `actor/bapo_positive_ratio`. + +The config file is located in [`gsm8k.yaml`](gsm8k.yaml). diff --git a/examples/bapo_gsm8k/gsm8k.yaml b/examples/bapo_gsm8k/gsm8k.yaml new file mode 100644 index 00000000000..7fe378e079c --- /dev/null +++ b/examples/bapo_gsm8k/gsm8k.yaml @@ -0,0 +1,76 @@ +project: "Trinity-RFT-gsm8k" +name: "qwen2.5-1.5B-gsm8k-bapo" +checkpoint_root_dir: ${oc.env:TRINITY_CHECKPOINT_ROOT_DIR,./checkpoints} +algorithm: + algorithm_type: bapo + repeat_times: 8 + policy_loss_fn_args: + # clipping bounds are explicit importance-ratio values (Sec. 5.1 of the paper) + lower_bound_start: 0.6 # a- + lower_bound_max: 0.9 # b- + lower_bound_step: 0.02 # delta_2 + upper_bound_start: 1.2 # a+ + upper_bound_max: 3.0 # b+ + upper_bound_step: 0.05 # delta_1 + target_positive_ratio: 0.4 # rho_0 + clip_ratio_c: 3.0 + loss_agg_mode: token-mean + optimizer: + lr: 1e-6 +model: + model_path: ${oc.env:TRINITY_MODEL_PATH,Qwen/Qwen2.5-1.5B-Instruct} + max_response_tokens: 1024 + max_model_len: 2048 +cluster: + node_num: 1 + gpu_per_node: 8 +buffer: + total_epochs: 1 + batch_size: 96 + explorer_input: + taskset: + name: gsm8k + storage_type: file + path: ${oc.env:TRINITY_TASKSET_PATH,openai/gsm8k} + subset_name: 'main' + split: 'train' + format: + prompt_key: 'question' + response_key: 'answer' + rollout_args: + temperature: 1.0 + eval_tasksets: + - name: gsm8k-eval + storage_type: file + path: ${oc.env:TRINITY_TASKSET_PATH,openai/gsm8k} + subset_name: 'main' + split: 'test' + format: + prompt_key: 'question' + response_key: 'answer' + default_workflow_type: 'math_workflow' + trainer_input: + experience_buffer: + name: gsm8k_buffer + storage_type: queue + path: 'sqlite:///gsm8k.db' +explorer: + eval_interval: 50 + runner_per_model: 8 + rollout_model: + engine_num: 2 + tensor_parallel_size: 1 + enable_prefix_caching: false + enforce_eager: true + dtype: bfloat16 + seed: 42 +synchronizer: + sync_method: 'nccl' + sync_interval: 4 + sync_timeout: 1200 +trainer: + save_interval: 100 + grad_clip: 1.0 + use_dynamic_bsz: true + max_token_len_per_gpu: 16384 + ulysses_sequence_parallel_size: 1 diff --git a/tests/algorithm/bapo_policy_loss_test.py b/tests/algorithm/bapo_policy_loss_test.py new file mode 100644 index 00000000000..e1b9376ca1c --- /dev/null +++ b/tests/algorithm/bapo_policy_loss_test.py @@ -0,0 +1,196 @@ +# -*- coding: utf-8 -*- +"""Tests for the BAPO policy loss function.""" + +import unittest + +import torch + +from trinity.algorithm import ALGORITHM_TYPE +from trinity.algorithm.policy_loss_fn import POLICY_LOSS_FN +from trinity.algorithm.policy_loss_fn.bapo_policy_loss import BAPOPolicyLossFn + + +def reference_positive_ratio(ratio, advantages, mask, c_low, c_high, clip_ratio_c): + """Token-by-token computation of the positive contribution ratio.""" + pos, neg = 0.0, 0.0 + for r, a, m in zip(ratio.flatten().tolist(), advantages.flatten().tolist(), mask.flatten()): + if not m: + continue + if a > 0 and r <= c_high: + pos += r * a + elif a < 0 and c_low <= r <= clip_ratio_c: + neg += r * -a + if neg <= 0: + return 1.0 + return pos / (pos + neg) + + +def reference_algorithm_1(ratio, advantages, mask, args): + """Literal transcription of the bound search in Algorithm 1 of the BAPO paper.""" + c_low, c_high = args["lower_bound_start"], args["upper_bound_start"] + + def rho(): + return reference_positive_ratio( + ratio, advantages, mask, c_low, c_high, args["clip_ratio_c"] + ) + + tol = 1e-6 # float tolerance for the "<=" comparisons + while ( + rho() < args["target_positive_ratio"] + and c_low + args["lower_bound_step"] <= args["lower_bound_max"] + tol + ): + if c_high + args["upper_bound_step"] <= args["upper_bound_max"] + tol: + c_high += args["upper_bound_step"] + else: + c_low += args["lower_bound_step"] + return c_low, c_high, rho() + + +class BAPOPolicyLossTest(unittest.TestCase): + def setUp(self): + torch.manual_seed(42) + shape = (6, 24) + self.logprob = 2 * torch.rand(shape) - 1 + self.old_logprob = 2 * torch.rand(shape) - 1 + self.mask = torch.rand(shape) > 0.3 + self.advantages = 2 * torch.rand(shape) - 1 + self.args = BAPOPolicyLossFn.default_args() + + def _call(self, fn, advantages=None): + return fn( + log_prob=self.logprob, + old_log_probs=self.old_logprob, + response_mask=self.mask, + advantages=self.advantages if advantages is None else advantages, + ) + + def _ratio(self): + return torch.exp(self.logprob - self.old_logprob) + + def test_registered(self): + self.assertIs(POLICY_LOSS_FN.get("bapo"), BAPOPolicyLossFn) + algorithm = ALGORITHM_TYPE.get("bapo") + self.assertEqual(algorithm.default_config()["policy_loss_fn"], "bapo") + self.assertEqual(algorithm.default_config()["advantage_fn"], "grpo") + self.assertFalse(algorithm.use_critic) + + def test_bound_schedule_follows_algorithm_1(self): + fn = BAPOPolicyLossFn(**self.args) + schedule = fn.bound_schedule + # 1 initial pair + 36 upper steps (1.2 -> 3.0) + 15 lower steps (0.6 -> 0.9) + self.assertEqual(len(schedule), 52) + self.assertEqual(schedule[0], (0.6, 1.2)) + self.assertEqual(schedule[1], (0.6, 1.25)) + self.assertEqual(schedule[36], (0.6, 3.0)) + self.assertEqual(schedule[37], (0.62, 3.0)) + self.assertEqual(schedule[-1], (0.9, 3.0)) + + def test_matches_literal_algorithm_1(self): + for seed in range(20): + torch.manual_seed(seed) + self.logprob = 2 * torch.rand(self.logprob.shape) - 1 + self.old_logprob = 2 * torch.rand(self.logprob.shape) - 1 + self.advantages = 2 * torch.rand(self.logprob.shape) - 1.2 # mostly negative + for target in (0.2, 0.4, 0.5, 0.7): + args = dict(self.args, target_positive_ratio=target) + fn = BAPOPolicyLossFn(**args) + c_low, c_high, rho = fn.select_bounds(self._ratio(), self.advantages, self.mask) + ref_low, ref_high, ref_rho = reference_algorithm_1( + self._ratio(), self.advantages, self.mask, args + ) + self.assertAlmostEqual(c_low, ref_low, places=6) + self.assertAlmostEqual(c_high, ref_high, places=6) + self.assertAlmostEqual(rho, ref_rho, places=5) + + def test_no_adjustment_when_target_already_met(self): + fn = BAPOPolicyLossFn(**dict(self.args, target_positive_ratio=0.0)) + _, metrics = self._call(fn) + self.assertEqual(metrics["bapo_clip_low"], 0.6) + self.assertEqual(metrics["bapo_clip_high"], 1.2) + + # all advantages positive: rho is 1 for the initial bounds + fn = BAPOPolicyLossFn(**self.args) + _, metrics = self._call(fn, advantages=self.advantages.abs() + 0.1) + self.assertEqual(metrics["bapo_clip_low"], 0.6) + self.assertEqual(metrics["bapo_clip_high"], 1.2) + self.assertEqual(metrics["bapo_positive_ratio"], 1.0) + + def test_upper_bound_raised_first(self): + # One positive token with ratio 2.0 (clipped at 1.2) and one negative token. + ratio = torch.tensor([[2.0, 1.0]]) + advantages = torch.tensor([[1.0, -1.0]]) + mask = torch.ones_like(ratio, dtype=torch.bool) + fn = BAPOPolicyLossFn(**self.args) + c_low, c_high, rho = fn.select_bounds(ratio, advantages, mask) + self.assertEqual(c_low, 0.6) + self.assertAlmostEqual(c_high, 2.0) + self.assertAlmostEqual(rho, 2.0 / 3.0, places=6) + + def test_lower_bound_raised_after_upper_saturates(self): + # Positive token is never clipped; negative tokens with small ratios dominate. + ratio = torch.tensor([[1.0, 0.65, 0.7, 0.95]]) + advantages = torch.tensor([[0.5, -1.0, -1.0, -0.2]]) + mask = torch.ones_like(ratio, dtype=torch.bool) + fn = BAPOPolicyLossFn(**self.args) + c_low, c_high, rho = fn.select_bounds(ratio, advantages, mask) + # 0.5 / (0.5 + 0.65 + 0.7 + 0.19) < 0.4: raising c_high cannot help, so c_low + # rises past 0.65 (rho = 0.5 / 1.39 < 0.4) and then past 0.7 (rho = 0.5 / 0.69). + self.assertAlmostEqual(c_high, 3.0) + self.assertAlmostEqual(c_low, 0.72) + self.assertAlmostEqual(rho, 0.5 / 0.69, places=5) + + def test_target_unreachable_uses_widest_bounds(self): + fn = BAPOPolicyLossFn(**self.args) + _, metrics = self._call(fn, advantages=-(self.advantages.abs() + 0.1)) + self.assertAlmostEqual(metrics["bapo_clip_low"], 0.9) + self.assertAlmostEqual(metrics["bapo_clip_high"], 3.0) + self.assertEqual(metrics["bapo_positive_ratio"], 0.0) + + def test_equivalent_to_ppo_with_fixed_bounds(self): + bapo = BAPOPolicyLossFn( + lower_bound_start=0.8, + lower_bound_max=0.8, + upper_bound_start=1.28, + upper_bound_max=1.28, + target_positive_ratio=1.0, + ) + ppo = POLICY_LOSS_FN.get("ppo")(clip_range_low=0.2, clip_range_high=0.28) + bapo_loss, bapo_metrics = self._call(bapo) + ppo_loss, ppo_metrics = self._call(ppo) + self.assertTrue(torch.allclose(bapo_loss, ppo_loss)) + for key in ("pg_clipfrac", "pg_clipfrac_lower", "ppo_kl", "pg_loss"): + self.assertAlmostEqual(bapo_metrics[key], ppo_metrics[key], places=6) + + def test_loss_is_differentiable(self): + self.logprob.requires_grad_(True) + fn = BAPOPolicyLossFn(**self.args) + loss, _ = self._call(fn) + loss.backward() + self.assertIsNotNone(self.logprob.grad) + self.assertTrue(torch.isfinite(self.logprob.grad).all()) + + def test_empty_mask(self): + fn = BAPOPolicyLossFn(**self.args) + loss, metrics = fn( + log_prob=self.logprob, + old_log_probs=self.old_logprob, + response_mask=torch.zeros_like(self.mask), + advantages=self.advantages, + ) + self.assertTrue(torch.isfinite(loss)) + self.assertEqual(metrics["bapo_clip_low"], 0.6) + self.assertEqual(metrics["bapo_clip_high"], 1.2) + + def test_invalid_args(self): + with self.assertRaises(AssertionError): + BAPOPolicyLossFn(lower_bound_start=0.9, lower_bound_max=0.6) + with self.assertRaises(AssertionError): + BAPOPolicyLossFn(lower_bound_max=1.3, upper_bound_start=1.2) + with self.assertRaises(AssertionError): + BAPOPolicyLossFn(upper_bound_step=-0.1) + with self.assertRaises(AssertionError): + BAPOPolicyLossFn(target_positive_ratio=1.5) + + +if __name__ == "__main__": + unittest.main() diff --git a/trinity/algorithm/__init__.py b/trinity/algorithm/__init__.py index 787ac233e92..75cc95cdd75 100644 --- a/trinity/algorithm/__init__.py +++ b/trinity/algorithm/__init__.py @@ -23,6 +23,7 @@ "cispo": "trinity.algorithm.algorithm.CISPOAlgorithm", "gspo": "trinity.algorithm.algorithm.GSPOAlgorithm", "sapo": "trinity.algorithm.algorithm.SAPOAlgorithm", + "bapo": "trinity.algorithm.algorithm.BAPOAlgorithm", "mix": "trinity.algorithm.algorithm.MIXAlgorithm", "mix_chord": "trinity.algorithm.algorithm.MIXCHORDAlgorithm", "raft": "trinity.algorithm.algorithm.RAFTAlgorithm", diff --git a/trinity/algorithm/algorithm.py b/trinity/algorithm/algorithm.py index 2cdf4cc08e1..a5372f08af8 100644 --- a/trinity/algorithm/algorithm.py +++ b/trinity/algorithm/algorithm.py @@ -396,6 +396,33 @@ def default_config(cls) -> Dict: } +class BAPOAlgorithm(AlgorithmType): + """BAPO (Balanced Policy Optimization with Adaptive Clipping) algorithm. + + BAPO keeps the PPO-clip surrogate but adapts the clipping bounds for every update, + so that positive-advantage tokens keep a target share of the policy gradient. + See https://arxiv.org/abs/2510.18927 + """ + + use_critic: bool = False + use_reference: bool = True + compute_advantage_in_trainer: bool = False + can_balance_batch: bool = True + schema: str = "experience" + + @classmethod + def default_config(cls) -> Dict: + return { + "repeat_times": 2, + "advantage_fn": "grpo", + "sample_strategy": "default", + "policy_loss_fn": "bapo", + "kl_penalty_fn": "none", + "kl_loss_fn": "k2", + "entropy_loss_fn": "default", + } + + class MIXAlgorithm(AlgorithmType): """MIX algorithm.""" diff --git a/trinity/algorithm/policy_loss_fn/__init__.py b/trinity/algorithm/policy_loss_fn/__init__.py index 6d1b0a2465c..cb4e7d990dc 100644 --- a/trinity/algorithm/policy_loss_fn/__init__.py +++ b/trinity/algorithm/policy_loss_fn/__init__.py @@ -18,6 +18,7 @@ "sppo": "trinity.algorithm.policy_loss_fn.sppo_loss_fn.sPPOPolicyLossFn", "rec": "trinity.algorithm.policy_loss_fn.rec_policy_loss.RECPolicyLossFn", "sapo": "trinity.algorithm.policy_loss_fn.sapo_policy_loss.SAPOPolicyLossFn", + "bapo": "trinity.algorithm.policy_loss_fn.bapo_policy_loss.BAPOPolicyLossFn", "importance_sampling": "trinity.algorithm.policy_loss_fn.importance_sampling_policy_loss.ImportanceSamplingLossFn", }, ) diff --git a/trinity/algorithm/policy_loss_fn/bapo_policy_loss.py b/trinity/algorithm/policy_loss_fn/bapo_policy_loss.py new file mode 100644 index 00000000000..b7ed6cfec2e --- /dev/null +++ b/trinity/algorithm/policy_loss_fn/bapo_policy_loss.py @@ -0,0 +1,236 @@ +"""BAPO policy loss function. + +Balanced Policy Optimization with Adaptive Clipping (BAPO) keeps the PPO-clip +surrogate but chooses the clipping bounds for every update, so that tokens with +positive advantages contribute at least a target share of the policy gradient. + +Refer to the BAPO paper for details. https://arxiv.org/abs/2510.18927 +""" + +import math +from typing import Dict, List, Tuple + +import torch + +from trinity.algorithm.policy_loss_fn.policy_loss_fn import PolicyLossFn +from trinity.algorithm.utils import aggregate_loss, masked_mean + + +class BAPOPolicyLossFn(PolicyLossFn): + """BAPO policy loss (arXiv:2510.18927, Eq. 8 and Algorithm 1). + + The clipping bounds are **explicit importance-ratio values** (e.g. ``0.6`` and + ``1.2``), not the ``1 ± clip_range`` epsilons used by ``PPOPolicyLossFn``. + + For every call, the bounds start at ``c_low = lower_bound_start`` (a-) and + ``c_high = upper_bound_start`` (a+). While the positive contribution ratio + ``rho`` is below ``target_positive_ratio`` (rho_0), ``c_high`` is raised by + ``upper_bound_step`` (delta_1) up to ``upper_bound_max`` (b+), and then ``c_low`` + is raised by ``lower_bound_step`` (delta_2) up to ``lower_bound_max`` (b-). If the + target is never met, the widest bounds reached are used. This matches the loop in + Algorithm 1 whenever ``lower_bound_start + lower_bound_step <= lower_bound_max`` + (true for the paper's settings); when the lower range is empty, the upper bound is + still searched instead of skipping the search altogether. + + ``rho`` is the share of the policy gradient that comes from positive-advantage + tokens. A token contributes ``ratio * |advantage|`` (the magnitude of its + gradient coefficient on ``log pi``) when it is not clipped, and nothing otherwise: + + - a positive-advantage token is clipped when ``ratio > c_high``; + - a negative-advantage token is clipped when ``ratio < c_low`` or, with the + dual-clip, when ``ratio > clip_ratio_c``. + + Raising ``c_high`` therefore lets more positive tokens contribute, and raising + ``c_low`` removes low-ratio negative tokens, so both steps increase ``rho``. + When no negative-advantage token contributes, ``rho`` is taken as ``1``. + + The bounds are searched independently for every call, i.e. per micro-batch on + each data-parallel rank. + """ + + def __init__( + self, + backend: str = "verl", + lower_bound_start: float = 0.6, + lower_bound_max: float = 0.9, + lower_bound_step: float = 0.02, + upper_bound_start: float = 1.2, + upper_bound_max: float = 3.0, + upper_bound_step: float = 0.05, + target_positive_ratio: float = 0.4, + clip_ratio_c: float = 3.0, + loss_agg_mode: str = "token-mean", + ) -> None: + """Initialize the BAPO policy loss function. + + Args: + backend: The training framework/backend to use (e.g., "verl"). + lower_bound_start: Initial lower clipping bound ``a-``. + lower_bound_max: Largest allowed lower clipping bound ``b-``. + lower_bound_step: Step ``delta_2`` used to raise the lower bound. + upper_bound_start: Initial upper clipping bound ``a+``. + upper_bound_max: Largest allowed upper clipping bound ``b+``. + upper_bound_step: Step ``delta_1`` used to raise the upper bound. + target_positive_ratio: Target share ``rho_0`` of the policy gradient + contributed by positive-advantage tokens, in ``[0, 1]``. + clip_ratio_c: Dual-clip constant for negative-advantage tokens, as in PPO. + loss_agg_mode: Mode for aggregating the loss across tokens. + """ + super().__init__(backend=backend) + assert 0.0 < lower_bound_start <= lower_bound_max <= upper_bound_start <= upper_bound_max, ( + "BAPO bounds must satisfy 0 < lower_bound_start <= lower_bound_max " + "<= upper_bound_start <= upper_bound_max, got " + f"[{lower_bound_start}, {lower_bound_max}] and [{upper_bound_start}, {upper_bound_max}]." + ) + assert lower_bound_step >= 0.0, "lower_bound_step must be non-negative." + assert upper_bound_step >= 0.0, "upper_bound_step must be non-negative." + assert 0.0 <= target_positive_ratio <= 1.0, "target_positive_ratio must be in [0, 1]." + assert clip_ratio_c > 1.0, "clip_ratio_c must be greater than 1.0." + self.lower_bound_start = lower_bound_start + self.lower_bound_max = lower_bound_max + self.lower_bound_step = lower_bound_step + self.upper_bound_start = upper_bound_start + self.upper_bound_max = upper_bound_max + self.upper_bound_step = upper_bound_step + self.target_positive_ratio = target_positive_ratio + self.clip_ratio_c = clip_ratio_c + self.loss_agg_mode = loss_agg_mode + self.bound_schedule = self._build_bound_schedule() + + @staticmethod + def _num_steps(start: float, end: float, step: float) -> int: + """Number of whole steps of size ``step`` that fit between ``start`` and ``end``.""" + if step <= 0.0: + return 0 + return int(math.floor((end - start) / step + 1e-6)) + + def _build_bound_schedule(self) -> List[Tuple[float, float]]: + """Candidate ``(c_low, c_high)`` pairs in the order Algorithm 1 visits them.""" + num_high = self._num_steps( + self.upper_bound_start, self.upper_bound_max, self.upper_bound_step + ) + num_low = self._num_steps( + self.lower_bound_start, self.lower_bound_max, self.lower_bound_step + ) + schedule = [ + (self.lower_bound_start, round(self.upper_bound_start + k * self.upper_bound_step, 8)) + for k in range(num_high + 1) + ] + final_high = schedule[-1][1] + schedule.extend( + (round(self.lower_bound_start + k * self.lower_bound_step, 8), final_high) + for k in range(1, num_low + 1) + ) + return schedule + + def positive_contribution_ratio( + self, + ratio: torch.Tensor, + advantages: torch.Tensor, + action_mask: torch.Tensor, + clip_low: torch.Tensor, + clip_high: torch.Tensor, + ) -> torch.Tensor: + """Compute ``rho`` for one or more candidate bounds. + + Args: + ratio: Importance ratio ``pi_theta / pi_old``, shape ``[batch, seq_len]``. + advantages: Advantages, shape ``[batch, seq_len]``. + action_mask: Mask of valid tokens, shape ``[batch, seq_len]``. + clip_low: Lower bounds, shape ``[num_candidates]``. + clip_high: Upper bounds, shape ``[num_candidates]``. + + Returns: + torch.Tensor: ``rho`` for every candidate, shape ``[num_candidates]``. + """ + valid = action_mask.bool() + r = ratio.detach()[valid].unsqueeze(0) # [1, num_tokens] + adv = advantages.detach()[valid].unsqueeze(0) + weight = r * adv.abs() + pos_active = (adv > 0) & (r <= clip_high.unsqueeze(1)) + neg_active = (adv < 0) & (r >= clip_low.unsqueeze(1)) & (r <= self.clip_ratio_c) + pos_contrib = (weight * pos_active).sum(dim=-1) + neg_contrib = (weight * neg_active).sum(dim=-1) + total = (pos_contrib + neg_contrib).clamp(min=torch.finfo(weight.dtype).tiny) + return torch.where(neg_contrib > 0, pos_contrib / total, torch.ones_like(total)) + + def select_bounds( + self, + ratio: torch.Tensor, + advantages: torch.Tensor, + action_mask: torch.Tensor, + ) -> Tuple[float, float, float]: + """Run the bound search of Algorithm 1. + + Returns: + Tuple[float, float, float]: The selected ``c_low``, ``c_high`` and the + resulting positive contribution ratio ``rho``. + """ + device = ratio.device + dtype = torch.float32 if ratio.dtype in (torch.float16, torch.bfloat16) else ratio.dtype + clip_low = torch.tensor([b[0] for b in self.bound_schedule], device=device, dtype=dtype) + clip_high = torch.tensor([b[1] for b in self.bound_schedule], device=device, dtype=dtype) + rho = self.positive_contribution_ratio( + ratio.to(dtype), advantages.to(dtype), action_mask, clip_low, clip_high + ) + met = rho >= self.target_positive_ratio + # Index of the first candidate that meets the target. A sentinel appended at the + # end makes argmax fall back to the last (widest) candidate when none does. + num_candidates = len(self.bound_schedule) + index_t = torch.cat([met, met.new_ones(1)]).int().argmax().clamp(max=num_candidates - 1) + index, rho_value = torch.stack([index_t.to(rho.dtype), rho[index_t]]).tolist() + c_low, c_high = self.bound_schedule[int(index)] + return c_low, c_high, rho_value + + def __call__( # type: ignore + self, + logprob: torch.Tensor, + old_logprob: torch.Tensor, + action_mask: torch.Tensor, + advantages: torch.Tensor, + **kwargs, + ) -> Tuple[torch.Tensor, Dict]: + negative_approx_kl = torch.clamp(logprob - old_logprob, min=-20.0, max=20.0) + ratio = torch.exp(negative_approx_kl) + ppo_kl = masked_mean(-negative_approx_kl, action_mask) + + clip_low, clip_high, rho = self.select_bounds(ratio, advantages, action_mask) + + pg_losses1 = -advantages * ratio + pg_losses2 = -advantages * torch.clamp(ratio, clip_low, clip_high) + clip_pg_losses1 = torch.maximum(pg_losses1, pg_losses2) + pg_clip_frac = masked_mean(torch.gt(pg_losses2, pg_losses1).float(), action_mask) + + pg_losses3 = -advantages * self.clip_ratio_c + clip_pg_losses2 = torch.min(pg_losses3, clip_pg_losses1) + pg_clipfrac_lower = masked_mean( + torch.gt(clip_pg_losses1, pg_losses3) * (advantages < 0).float(), action_mask + ) + pg_losses = torch.where(advantages < 0, clip_pg_losses2, clip_pg_losses1) + + pg_loss = aggregate_loss(pg_losses, action_mask, loss_agg_mode=self.loss_agg_mode) + metrics = { + "pg_clipfrac": pg_clip_frac.detach().item(), + "pg_clipfrac_lower": pg_clipfrac_lower.detach().item(), + "ppo_kl": ppo_kl.detach().item(), + "pg_loss": pg_loss.detach().item(), + "bapo_clip_low": clip_low, + "bapo_clip_high": clip_high, + "bapo_positive_ratio": rho, + } + return pg_loss, metrics + + @classmethod + def default_args(cls) -> Dict: + """Default arguments, following Sec. 5.1 of the BAPO paper.""" + return { + "lower_bound_start": 0.6, + "lower_bound_max": 0.9, + "lower_bound_step": 0.02, + "upper_bound_start": 1.2, + "upper_bound_max": 3.0, + "upper_bound_step": 0.05, + "target_positive_ratio": 0.4, + "clip_ratio_c": 3.0, + "loss_agg_mode": "token-mean", + }