[flyingdoubleG]

Description

Background. SAPO (arXiv:2511.20347) maximizes E[f(r) * A] with the soft gate f(r) = sigmoid(tau (r - 1)) * 4 / tau, r = pi_theta / pi_old. Differentiating gives (paper Eq. 7-8)

grad J = E[ w * r * A * grad log pi ],   w = 4 p (1 - p),   p = sigmoid(tau (r - 1))

so the per-token weight w * r is 1 at r = 1 (plain policy gradient) and decays as r moves away from 1. That decay is the soft trust region.

Problem. SAPOPolicyLossFn computes

sapo_loss = -advantages * soft_gate.detach() * logprob

whose gradient weight is f(r) itself, not f'(r) * r. With the default tau_pos = 1:

ratio r	weight now: f(r)	SAPO: 4p(1-p)·r
0.5	1.510	0.470
1.0	2.000	1.000
2.0	2.924	1.573
5.0	3.928	0.353

So on-policy tokens get 2 / tau times the policy gradient, and the weight increases with r (towards 4 / tau) instead of decaying: off-policy tokens are weighted up rather than down. (Computed with autograd on soft_gate_function from this file.)

Change. sapo_loss = -advantages * soft_gate, with the gate differentiated through ratio, which is how verl (compute_policy_loss_sapo) and ms-swift implement it. The logged gate / ratio metrics are unchanged; the logged loss value is now -mean(f(r) A) instead of -mean(f(r) A log pi).

Tests. tests/algorithm/policy_loss_test.py: updated the loss snapshot in test_sapo_policy_loss (other metrics unchanged) and added test_sapo_policy_loss_gradient, which checks d loss / d logprob against Eq. 7-8 for the default temperatures. Both fail before the change; the file passes after (9 tests, CPU).

Note, not changed here: the paper averages over tokens within each sequence and then over sequences (1/G Σ 1/|y_i| Σ), i.e. seq-mean-token-mean, and verl forces that mode for SAPO; this class defaults to token-mean. Happy to change the default in a follow-up if you prefer.

Checklist

• Code has passed all tests
• Docstrings have been added/updated in Google Style
• Documentation has been updated
• Code is ready for review

🤖 Generated with Claude Code