From 50af3c5037832e8b11a3051b9a7828a9960a9248 Mon Sep 17 00:00:00 2001
From: flyingdoubleG <53119724+flyingdoubleG@users.noreply.github.com>
Date: Tue, 20 Oct 2026 21:21:58 +0800
Subject: [PATCH] fix(reward): keep the grouped shape in reward_norm when only
 std is per group

With norm_std_type="group" and norm_mean_type "batch", "running" or None,
reward_norm computes the per-group std on the [n_groups, n_sample] view
but subtracts the mean from the flat [N] rewards, so the division
broadcasts to [n_groups, N] and the RLVR pipeline fails when it writes
the result back to the batch. The agentic agentic_reward_norm supports
this combination.

Compute on the grouped view whenever a group statistic is used and
restore the original shape afterwards; the other combinations are
unchanged.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 roll/utils/functionals.py       |  9 +++++----
 tests/utils/test_functionals.py | 20 ++++++++++++++++++++
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/roll/utils/functionals.py b/roll/utils/functionals.py
index c3964bb82..d6a6f0655 100644
--- a/roll/utils/functionals.py
+++ b/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
 
diff --git a/tests/utils/test_functionals.py b/tests/utils/test_functionals.py
index 143ba6fdc..cc9f00761 100644
--- a/tests/utils/test_functionals.py
+++ b/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()