fix(reward): keep the grouped shape in reward_norm when only std is per group

roll/utils/functionals.py

@@ -816,16 +816,17 @@ def reward_norm(
         reward_std = torch.std(reshape_reward, dim=-1, keepdim=True)
     elif norm_std_type == "running":
         reward_std = running.std
-    # 选择基础奖励值
-    rewards = reshape_reward if norm_mean_type == "group" else response_level_rewards
+    # 选择基础奖励值: 只要用到 group 统计量 (mean 或 std)，就在分组形状上计算，保证 [*, n_groups, 1] 的统计量正确广播
+    grouped = group_mode and n_sample > 0
+    rewards = reshape_reward if grouped else response_level_rewards
     # 标准化奖励
     if norm_std_type is not None:
         normalized_rewards = (rewards - reward_mean) / (reward_std + 1e-6)
     else:
         normalized_rewards = (rewards - reward_mean)
 
-    # 如果是对 group mean 归一化，需要恢复原始形状
-    if norm_mean_type == "group":
+    # 在分组形状上计算时，需要恢复原始形状
+    if grouped:
         normalized_rewards = normalized_rewards.reshape(*response_level_rewards.size())
     return normalized_rewards
 

tests/utils/test_functionals.py

@@ -8,6 +8,7 @@
     compute_approx_kl,
     divide_by_chunk_size,
     pad_to_length,
+    reward_norm,
     traverse_obj,
 )
 
@@ -76,5 +77,24 @@ def test_compute_approx_kl_k3_clamps_before_exp():
     assert torch.isfinite(result).all()
 
 
+@pytest.mark.parametrize(
+    ("norm_mean_type", "norm_std_type"),
+    [("group", "group"), ("batch", "group"), (None, "group"), ("group", "batch"), ("batch", "batch"), (None, None)],
+)
+def test_reward_norm_combinations(norm_mean_type, norm_std_type):
+    rewards = torch.tensor([1.0, 0.0, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0])  # 2 prompts x 4 samples
+    groups = rewards.reshape(2, 4)
+    means = {"group": groups.mean(-1, keepdim=True), "batch": rewards.mean(), None: 0.0}
+    stds = {"group": groups.std(-1, keepdim=True), "batch": rewards.std()}
+    expected = groups - means[norm_mean_type]
+    if norm_std_type is not None:
+        expected = expected / (stds[norm_std_type] + 1e-6)
+
+    result = reward_norm(rewards, n_sample=4, norm_mean_type=norm_mean_type, norm_std_type=norm_std_type)
+
+    assert result.shape == rewards.shape
+    torch.testing.assert_close(result, expected.reshape(-1))
+
+
 if __name__ == "__main__":
     pytest.main()