From 4175f3ea7b81930818eebaf23edcbb3b5cf7d713 Mon Sep 17 00:00:00 2001 From: flyingdoubleG <53119724+flyingdoubleG@users.noreply.github.com> Date: Tue, 20 Oct 2026 21:31:40 +0800 Subject: [PATCH] fix(sapo): differentiate the soft gate through the ratio SAPO maximizes E[f(r) * A] with f(r) = sigmoid(tau (r - 1)) * 4 / tau; its gradient is f'(r) * r * A * grad(log pi), i.e. the weight w * r with w = 4 p (1 - p), p = sigmoid(tau (r - 1)) (arXiv:2511.20347, Eq. 7-8), which is 1 at r = 1 and decays as r moves away from 1. The loss was written as -A * sg(f(r)) * log pi, whose gradient weight is f(r) itself: 2 / tau at r = 1 (twice the policy gradient for tau = 1) and increasing with r towards 4 / tau, so off-policy tokens were weighted up instead of down. Use -A * f(r) with the gate differentiated through r, as verl and ms-swift do. Update the loss snapshot in the existing test and add a gradient test against Eq. 7-8. Co-Authored-By: Claude Opus 5.5 --- tests/algorithm/policy_loss_test.py | 20 ++++++++++++++++++- .../policy_loss_fn/sapo_policy_loss.py | 6 +++--- 2 files changed, 22 insertions(+), 4 deletions(-) diff --git a/tests/algorithm/policy_loss_test.py b/tests/algorithm/policy_loss_test.py index bac8e69322f..7bba4670e99 100644 --- a/tests/algorithm/policy_loss_test.py +++ b/tests/algorithm/policy_loss_test.py @@ -148,7 +148,7 @@ def test_sapo_policy_loss(self): policy_loss_fn_args = policy_loss_fn_cls.default_args() policy_loss_fn = policy_loss_fn_cls(**policy_loss_fn_args) loss, metrics = policy_loss_fn(log_prob=self.logprob, **self.input_data.batch) - sapo_loss = torch.tensor(-0.05128994956612587) + sapo_loss = torch.tensor(0.24060271680355072) ppo_kl = torch.tensor(-0.21663446724414825) avg_soft_gate = torch.tensor(2.3191137313842773) avg_ratio = torch.tensor(1.630766749382019) @@ -159,3 +159,21 @@ def test_sapo_policy_loss(self): self.assertTrue(torch.allclose(torch.tensor(metrics["avg_soft_gate"]), avg_soft_gate)) self.assertTrue(torch.allclose(torch.tensor(metrics["avg_ratio"]), avg_ratio)) self.assertTrue(torch.allclose(torch.tensor(metrics["pos_adv_frac"]), pos_adv_frac)) + + def test_sapo_policy_loss_gradient(self): + # SAPO (arXiv:2511.20347, Eq. 7-8): d/dlogpi of -f(r) A is -w r A with w = 4 p (1 - p), + # p = sigmoid(tau (r - 1)); the weight is 1 at r = 1 and decays away from it. + policy_loss_fn_cls = POLICY_LOSS_FN.get("sapo") + policy_loss_fn = policy_loss_fn_cls(**policy_loss_fn_cls.default_args()) + batch = self.input_data.batch + logprob = self.logprob.clone().requires_grad_(True) + loss, _ = policy_loss_fn(log_prob=logprob, **batch) + loss.backward() + + mask = batch["response_mask"].float() + advantages = batch["advantages"] + ratio = torch.exp(self.logprob - batch["old_log_probs"]) + tau = torch.where(advantages > 0, torch.tensor(1.0), torch.tensor(1.05)) + p = torch.sigmoid(tau * (ratio - 1)) + expected = -4 * p * (1 - p) * ratio * advantages * mask / mask.sum() + self.assertTrue(torch.allclose(logprob.grad, expected, atol=1e-6)) diff --git a/trinity/algorithm/policy_loss_fn/sapo_policy_loss.py b/trinity/algorithm/policy_loss_fn/sapo_policy_loss.py index 285a24e3385..ffde609fcec 100644 --- a/trinity/algorithm/policy_loss_fn/sapo_policy_loss.py +++ b/trinity/algorithm/policy_loss_fn/sapo_policy_loss.py @@ -112,9 +112,9 @@ def __call__( # type: ignore soft_gate = self.soft_gate_function(ratio, advantages) # SAPO loss: -E[f_{i,t}(r_{i,t}) * Â_i] - # We multiply by logprob to get the policy gradient - # The gradient of log π_θ gives us the policy gradient direction - sapo_loss = -advantages * soft_gate.detach() * logprob + # Differentiate f(r) through the ratio: grad = f'(r) * r * A * grad(log pi), i.e. the weight + # 4 p (1 - p) * r of the paper (Eq. 7-8), which peaks at r = 1 and decays away from it. + sapo_loss = -advantages * soft_gate # Aggregate loss across tokens loss = aggregate_loss(sapo_loss, action_mask, loss_agg_mode=self.loss_agg_mode)