Fix advantage computation + add advantage to metric logs

tinker_cookbook/rl/data_processing.py

@@ -18,14 +18,25 @@
 
 
 def compute_advantages(trajectory_groups_P: List[TrajectoryGroup]) -> List[torch.Tensor]:
-    """Compute advantages for each trajectory, centered within groups."""
+    """Compute advantages for each trajectory, centered within groups.
+
+    For single-trajectory groups, centers across the entire batch.
+    """
+    # Flatten all rewards
+    all_rewards = torch.cat(
+        [torch.tensor(traj_group.get_total_rewards()) for traj_group in trajectory_groups_P]
+    )
+
+    # Compute baseline per group (or global if group size is 1)
     advantages_P: list[torch.Tensor] = []
 
     for traj_group in trajectory_groups_P:
         rewards_G = torch.tensor(traj_group.get_total_rewards())
-        # Center advantages within the group
-        advantages_G = rewards_G - rewards_G.mean()
-        advantages_P.append(advantages_G)
+        group_size = len(rewards_G)
+
+        # Use group mean if > 1 trajectory, else use batch mean
+        baseline = rewards_G.mean() if group_size > 1 else all_rewards.mean()
+        advantages_P.append(rewards_G - baseline)
 
     return advantages_P
 

tinker_cookbook/rl/train.py

@@ -725,6 +725,17 @@ async def prepare_minibatch(
         advantages_P = compute_advantages(trajectory_groups_P)
         data_D, _metadata_D = assemble_training_data(trajectory_groups_P, advantages_P)
 
+    # Log advantage statistics
+    all_advantages = torch.cat(advantages_P)
+    metrics.update(
+        {
+            "advantages/mean": all_advantages.mean().item(),
+            "advantages/std": all_advantages.std().item(),
+            "advantages/min": all_advantages.min().item(),
+            "advantages/max": all_advantages.max().item(),
+        }
+    )
+
     # Incorporate KL penalty if configured
     if kl_penalty_coef > 0:
         with timed("kl_vs_base", metrics):