From 650077faaaa0c4efa1bc432cbd73bd2256da4d51 Mon Sep 17 00:00:00 2001
From: Linxiushen <1787979356@qq.com>
Date: Mon, 21 Sep 2026 21:13:48 +0800
Subject: [PATCH] fix(rewards): make DAPO zero-penalty region inclusive so
 `cache_length=0` no longer divides by zero

MathDAPORewardFn.compute_overlong_penalty() returned 0 only while
response_len < max_response_length - cache_length, and otherwise divided
by cache_length. With cache_length=0 (no soft window, hard truncation
only) a response whose length equals max_response_length, which is the
usual length of a truncated rollout, hit the division and raised
ZeroDivisionError; the workflow runner catches it, so that rollout's
experience was silently dropped with an error log.

Use <= for the zero-penalty region. This matches the closed interval in
the DAPO paper's overlong shaping definition, and for cache_length > 0 the
value at the boundary was already 0, so results are unchanged there
(738-cell differential: only the 9 cache_length=0 cells change, from an
exception to 0.0).

Adds tests for the piecewise penalty with a soft window and for
cache_length=0.
---
 tests/algorithm/test_dapo_algorithm.py | 26 ++++++++++++++++++++++++++
 trinity/common/rewards/dapo_reward.py  |  2 +-
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/tests/algorithm/test_dapo_algorithm.py b/tests/algorithm/test_dapo_algorithm.py
index e4c143f55c2..74260b57723 100644
--- a/tests/algorithm/test_dapo_algorithm.py
+++ b/tests/algorithm/test_dapo_algorithm.py
@@ -108,6 +108,32 @@ def test_symmetric_accuracy(self, mock_compute_score):
         self.assertEqual(good["accuracy"], 1.0)
         self.assertEqual(bad["accuracy"], -1.0)
 
+    def test_overlong_penalty_piecewise(self):
+        fn = MathDAPORewardFn(
+            enable_overlong_penalty=True,
+            penalty_factor=1.0,
+            max_response_length=100,
+            cache_length=20,
+        )
+        # DAPO paper (Sec. 2.4): zero until `max_response_length - cache_length`,
+        # then linear, then -penalty_factor beyond `max_response_length`.
+        self.assertEqual(fn.compute_overlong_penalty(torch.zeros(79)), 0.0)
+        self.assertEqual(fn.compute_overlong_penalty(torch.zeros(80)), 0.0)
+        self.assertAlmostEqual(fn.compute_overlong_penalty(torch.zeros(90)), -0.5)
+        self.assertAlmostEqual(fn.compute_overlong_penalty(torch.zeros(100)), -1.0)
+        self.assertEqual(fn.compute_overlong_penalty(torch.zeros(101)), -1.0)
+
+    def test_overlong_penalty_without_soft_window(self):
+        fn = MathDAPORewardFn(
+            enable_overlong_penalty=True,
+            penalty_factor=1.0,
+            max_response_length=100,
+            cache_length=0,
+        )
+        self.assertEqual(fn.compute_overlong_penalty(torch.zeros(99)), 0.0)
+        self.assertEqual(fn.compute_overlong_penalty(torch.zeros(100)), 0.0)
+        self.assertEqual(fn.compute_overlong_penalty(torch.zeros(101)), -1.0)
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/trinity/common/rewards/dapo_reward.py b/trinity/common/rewards/dapo_reward.py
index dcc194e683a..d715b014f38 100644
--- a/trinity/common/rewards/dapo_reward.py
+++ b/trinity/common/rewards/dapo_reward.py
@@ -89,7 +89,7 @@ def compute_overlong_penalty(self, response_token):
         response_len = len(response_token)
         expected_len = self.max_response_length - self.cache_length
 
-        if response_len < expected_len:
+        if response_len <= expected_len:
             return 0.0
         elif response_len > self.max_response_length:
             return -self.penalty_factor