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()