From 3806e0e1f5ec9732d299e8ffa54e46e6d5ce4e8f Mon Sep 17 00:00:00 2001 From: David-Wu1119 <133224895+David-Wu1119@users.noreply.github.com> Date: Wed, 11 Nov 2026 00:48:51 -0500 Subject: [PATCH] fix(grpo): use rank-penalized rewards in the advantage numerator With `rank_penalty` set, GRPOGroupedAdvantage rescales each reward by its likelihood rank (the unlikeliness reward of arXiv:2506.02355) and computes the group mean and std from the rescaled rewards, but the advantage itself was still computed from the raw `exp.reward`. With binary rewards every correct response then received the same advantage regardless of its rank, so the option shifted the baseline but never favoured unlikely correct responses, which is its whole purpose. Compute the advantage from the same (possibly rank-penalized) rewards that the group statistics use. Behaviour without `rank_penalty` is unchanged. Co-Authored-By: Claude Opus 5.5 --- tests/algorithm/advantage_fn_test.py | 24 +++++++++++++++++++ .../algorithm/advantage_fn/grpo_advantage.py | 10 ++++---- 2 files changed, 29 insertions(+), 5 deletions(-) diff --git a/tests/algorithm/advantage_fn_test.py b/tests/algorithm/advantage_fn_test.py index b9f8dab858a..a65dcc89d2b 100644 --- a/tests/algorithm/advantage_fn_test.py +++ b/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) diff --git a/trinity/algorithm/advantage_fn/grpo_advantage.py b/trinity/algorithm/advantage_fn/grpo_advantage.py index 64dca61a175..f974a5f992b 100644 --- a/trinity/algorithm/advantage_fn/grpo_advantage.py +++ b/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()