Log individual trajetory rewards, rather than average across microbatch

PiperOrigin-RevId: 816502498

Author: The tunix Authors

tests/rl/grpo/grpo_learner_test.py

@@ -325,7 +325,7 @@ def wrapper(*args, **kwargs):
         else ('rewards/' + reward_fns.__name__,)
     )
     for metric_name in [
-        'rewards/overall',
+        'rewards/sum',
         *rewards_metrics,
         'completions/mean_length',
         'completions/max_length',
@@ -856,7 +856,7 @@ def test_trajectory_ids(self):
     def my_reward_fn(trajectories, prompts, **kwargs):
       for t_id, prompt in zip(kwargs['trajectory_ids'], prompts):
         trajectories[kwargs['mode']][t_id] = prompt
-      return 1.0
+      return [1.0] * len(prompts)
 
     vocab = tc.MockVocab()
     model = tc.ToyTransformer(rngs=nnx.Rngs(0), vocab_size=vocab.GetPieceSize())

tunix/rl/rl_learner.py

@@ -185,35 +185,38 @@ def _compute_rewards(
             f"Content of r: {r}"
         )
       rewards[:, i] = np.array(r)
+      for reward in r:
+        self.rl_cluster.buffer_metrics(
+            {
+                f"rewards/{reward_fn.__name__}": (
+                    reward,
+                    np.mean,
+                ),
+            },
+            mode=mode,
+        )
+
+    rewards = np.nansum(rewards, axis=1)
+    for trajectory_idx in range(len(prompts)):
+      trajectory_rewards = rewards[trajectory_idx]
       self.rl_cluster.buffer_metrics(
           {
-              f"rewards/{reward_fn.__name__}": (
-                  np.mean(r),
+              "rewards/sum": (
+                  np.sum(trajectory_rewards),
                   np.mean,
               ),
           },
           mode=mode,
       )
-
-    rewards = np.nansum(rewards, axis=1)
-    self.rl_cluster.buffer_metrics(
-        {
-            "rewards/overall": (
-                np.mean(rewards),
-                np.mean,
-            ),
-        },
-        mode=mode,
-    )
-    self.rl_cluster.buffer_metrics(
-        {
-            "rewards/min": (
-                np.min(rewards),
-                np.min,
-            ),
-        },
-        mode=mode,
-    )
+      self.rl_cluster.buffer_metrics(
+          {
+              "rewards/min": (
+                  np.min(trajectory_rewards),
+                  np.min,
+              ),
+          },
+          mode=mode,
+      )
     for p, c in zip(prompts, completions):
       self.rl_cluster.buffer_metrics(
           {