[David-Wu1119]

Description

Bug. GRPOGroupedAdvantage with rank_penalty set (advantage_fn_args: {rank_penalty: ...}) applies the unlikeliness reward of arXiv:2506.02355 (Sec. 4.1): each reward is rescaled by 1 - rank_penalty * rank / G, where higher-likelihood responses get larger ranks. The group mean and std are computed from the rescaled rewards, but the advantage was still computed from the raw exp.reward:

rewards = rewards * (1 - group_ranks * self.rank_penalty)
group_reward_mean = torch.mean(rewards)
group_reward_std = torch.std(rewards)
...
score = (exp.reward - group_reward_mean) / (group_reward_std + self.epsilon)

With binary rewards, every correct response in a group therefore gets the same advantage whatever its rank. The option only shifts the baseline and never favours the less likely correct responses, which is what the paper's modified reward r_i is for (the advantage is computed from r_i).

Reproduce. One group, rewards [1, 1, 0, 0], mean log-probs [-0.1, -2.0, -0.5, -1.0], rank_penalty=0.25: on main both correct responses get the same advantage; the unlikely one (-2.0) should get the larger one.

Fix. Compute the advantage from the same (possibly rank-penalized) rewards the group statistics use. Without rank_penalty the result is unchanged.

Tests. New tests/algorithm/advantage_fn_test.py::TestGroupedAdvantageFn::test_grpo_rank_penalty_advantages_use_adjusted_rewards checks the advantages against (r_i - mean(r)) / (std(r) + eps) on the adjusted rewards. It fails on main and passes with this change. tests/algorithm and tests/common/rewards pass (108 tests, CPU). pre-commit passes on the changed files. No GPU or model-weight tests were run.

Checklist

Please check the following items before code is ready to be reviewed.

• Code has passed all tests (only the CPU tests in tests/algorithm and tests/common/rewards were run; GPU and Ray tests were not)
• Docstrings have been added/updated in Google Style
• Documentation has been updated
• Code is ready for review

🤖 Generated with Claude Code

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

An organization owner can choose the repository organization as the billing source for reviews requested by members with a paid Copilot license. Learn about Copilot code review billing. Review the Copilot code review billing settings.