fix(grpo): use rank-penalized rewards in the advantage numerator

tests/algorithm/advantage_fn_test.py

@@ -146,6 +146,30 @@ def test_grpo_correct_bias(self):
             places=6,
         )
 
+    def test_grpo_rank_penalty_advantages_use_adjusted_rewards(self):
+        advantage_fn_cls = ADVANTAGE_FN.get("grpo")
+        advantage_fn = advantage_fn_cls(epsilon=1e-6, rank_penalty=0.25)
+        rewards = [1.0, 1.0, 0.0, 0.0]
+        mean_logprobs = [-0.1, -2.0, -0.5, -1.0]  # ranks (ascending): 3, 0, 2, 1
+        exps = [
+            Experience(
+                eid=EID(batch=0, task=0, run=i),
+                tokens=torch.zeros(5),
+                logprobs=torch.full((3,), mean_logprobs[i]),
+                prompt_length=2,
+                reward=rewards[i],
+            )
+            for i in range(4)
+        ]
+        exps, _ = advantage_fn(exps)
+
+        adjusted = torch.tensor([1.0 * (1 - 0.25 * 3 / 4), 1.0, 0.0, 0.0])
+        expected = (adjusted - adjusted.mean()) / (adjusted.std() + 1e-6)
+        advantages = torch.stack([exp.advantages[0] for exp in exps])
+        self.assertTrue(torch.allclose(advantages, expected, atol=1e-5))
+        # the less likely correct response gets the larger advantage
+        self.assertGreater(advantages[1].item(), advantages[0].item())
+
     def test_batch_level_std_grpo(self):
         advantage_fn_cls = ADVANTAGE_FN.get("grpo")
         self.assertIsNotNone(advantage_fn_cls)

trinity/algorithm/advantage_fn/grpo_advantage.py

@@ -134,6 +134,7 @@ def calculate_group_advantage(
     ) -> Tuple[List[Experience], Dict]:
         metrics = {}
         with torch.no_grad():
+            rewards = torch.tensor([exp.reward for exp in exps], dtype=torch.float32)
             if len(exps) == 1:
                 group_reward_mean = torch.tensor(
                     0.0
@@ -143,8 +144,6 @@ def calculate_group_advantage(
                     metrics["skipped_count"] = 1
                     exps.clear()  # Clear experiences if only one experience
             else:
-                rewards = torch.tensor([exp.reward for exp in exps], dtype=torch.float32)
-
                 if self.rank_penalty is not None:
                     # Correct bias by adjusting rewards based on their ranks
                     old_log_probs = torch.tensor(
@@ -162,11 +161,12 @@ def calculate_group_advantage(
                     metrics["skipped_count"] = len(exps)
                     exps.clear()
 
-            for exp in exps:
+            # use the (possibly rank-penalized) rewards the group statistics were computed from
+            for exp, reward in zip(exps, rewards):
                 if self.std_cal_level == "batch" and precomputed_std is not None:
-                    score = (exp.reward - group_reward_mean) / (precomputed_std + self.epsilon)
+                    score = (reward - group_reward_mean) / (precomputed_std + self.epsilon)
                 else:
-                    score = (exp.reward - group_reward_mean) / (group_reward_std + self.epsilon)
+                    score = (reward - group_reward_mean) / (group_reward_std + self.epsilon)
                 exp.advantages = score * exp.action_mask
                 exp.returns = exp.advantages.clone()