add metrics for grouping and advantage stats

Author: jeffreysijuntan

rllm/trainer/tinker/tinker_agent_trainer.py

@@ -201,6 +201,7 @@ async def _fit_agent_async(self):
 
                 # Stream: train on each minibatch as it arrives
                 train_step_start = time.time()
+                all_grouping_metrics = []
                 async for minibatch_episodes in self.generate_agent_episodes(group_size=self.config.training.group_size, minibatch_size=minibatch_size):
                     episodes.extend(minibatch_episodes)
                     minibatch_count += 1
@@ -216,10 +217,11 @@ async def _fit_agent_async(self):
 
                     # Train immediately (streaming), only optimize on last minibatch
                     t_train_start = time.time()
-                    logprobs, datums = await self.trainer.step(minibatch_episodes, learning_rate=learning_rate, beta1=beta1, beta2=beta2, eps=eps, optimizer_step=False)
+                    logprobs, datums, grouping_metrics = await self.trainer.step(minibatch_episodes, learning_rate=learning_rate, beta1=beta1, beta2=beta2, eps=eps, optimizer_step=False)
                     forward_backward_times.append(time.time() - t_train_start)
                     training_logprobs.extend(logprobs)
                     training_datums.extend(datums)
+                    all_grouping_metrics.append(grouping_metrics)
                     logger.info(f"Processed minibatch {minibatch_count}/{num_minibatches} with {len(minibatch_episodes)} episodes")
 
                 optim_step_time = time.time()
@@ -246,6 +248,19 @@ async def _fit_agent_async(self):
                     training_datums=training_datums,  # Pass datums for KL/perplexity metrics
                     training_logprobs=training_logprobs,
                 )
+
+                # Aggregate grouping metrics from all minibatches
+                if all_grouping_metrics:
+                    import numpy as np
+
+                    # Average numeric metrics across minibatches
+                    aggregated_grouping_metrics = {}
+                    for key in all_grouping_metrics[0].keys():
+                        values = [m[key] for m in all_grouping_metrics if key in m]
+                        if values:
+                            aggregated_grouping_metrics[key] = np.mean(values)
+                    metrics.update(aggregated_grouping_metrics)
+
                 tracking_logger.log(data=metrics, step=batch_idx)
                 print_metrics_table(metrics, batch_idx)
 

rllm/trainer/tinker/tinker_data_processor.py

@@ -401,7 +401,7 @@ def process_episodes(
     advantage_computer: TinkerAdvantageComputer,
     trajectory_filter: TinkerTrajectoryFilter,
     algorithm_config,
-) -> list[tinker.Datum]:
+) -> tuple[list[tinker.Datum], dict]:
     """
     Main pipeline to convert Episode objects to training datums.
 
@@ -423,10 +423,14 @@ def process_episodes(
         algorithm_config: Configuration with grouping_level setting
 
     Returns:
-        List of Tinker Datum objects ready for training
+        Tuple of (datums, metrics_dict):
+        - datums: List of Tinker Datum objects ready for training
+        - metrics_dict: Dictionary with grouping and advantage statistics
     """
     from collections import defaultdict
 
+    import numpy as np
+
     grouping_level = algorithm_config.get("grouping_level", "episode")
 
     # Group trajectories based on grouping_level
@@ -469,6 +473,10 @@ def get_task_id(episode):
     # Apply filtering based on configuration
     filtered_groups = trajectory_filter.filter_groups(trajectory_groups)
 
+    # Track metrics
+    all_advantages = []
+    group_sizes = []
+
     training_datums = []
     for group in filtered_groups:
         # Extract rewards for the group (from all trajectories)
@@ -477,13 +485,33 @@ def get_task_id(episode):
         # Compute advantages
         advantages = advantage_computer.compute(group_rewards)
 
+        # Track for metrics
+        all_advantages.extend(advantages)
+        group_sizes.append(len(group.trajectories))
+
         # Create datums for all trajectories in the group
         for trajectory, advantage in zip(group.trajectories, advantages, strict=False):
             # Use trajectory-level building (merges steps when possible)
             new_datums = TinkerDatumBuilder.build_datum_from_trajectory(trajectory, advantage)
             training_datums.extend(new_datums)
 
-    return training_datums
+    # Compute grouping and advantage metrics
+    metrics = {}
+    if filtered_groups:
+        metrics["grouping/num_groups"] = len(filtered_groups)
+        metrics["grouping/num_groups_before_filter"] = len(trajectory_groups)
+        metrics["grouping/avg_group_size"] = np.mean(group_sizes)
+        metrics["grouping/max_group_size"] = np.max(group_sizes)
+        metrics["grouping/min_group_size"] = np.min(group_sizes)
+
+    if all_advantages:
+        metrics["advantage/mean"] = np.mean(all_advantages)
+        metrics["advantage/std"] = np.std(all_advantages)
+        metrics["advantage/max"] = np.max(all_advantages)
+        metrics["advantage/min"] = np.min(all_advantages)
+        metrics["advantage/fraction_zero"] = np.sum(np.abs(all_advantages) < 1e-8) / len(all_advantages)
+
+    return training_datums, metrics
 
 
 def process_trajectory_groups(

rllm/trainer/tinker/tinker_policy_trainer.py

@@ -132,7 +132,7 @@ async def step(
         beta2: float = 0.95,
         eps: float = 1e-8,
         optimizer_step: bool = True,
-    ) -> tuple[list[torch.Tensor], list[tinker.Datum]]:
+    ) -> tuple[list[torch.Tensor], list[tinker.Datum], dict]:
         """
         Complete training step: process episodes and update policy.
 
@@ -147,15 +147,16 @@ async def step(
             optimizer_step: Whether to apply optimizer step after forward-backward
 
         Returns:
-            Tuple of (training_logprobs, training_datums)
+            Tuple of (training_logprobs, training_datums, grouping_metrics)
             - training_logprobs: List of training logprobs for KL computation
             - training_datums: List of datums WITH masks for metrics
+            - grouping_metrics: Dict with grouping and advantage statistics
         """
         if learning_rate is None:
             learning_rate = self.config.training.learning_rate
 
         # Step 1: Process to datums (includes filtering and advantage computation)
-        training_datums = process_episodes(
+        training_datums, grouping_metrics = process_episodes(
             episodes,
             self.advantage_computer,
             self.trajectory_filter,
@@ -193,11 +194,11 @@ async def step(
             training_logprobs = output["logprobs"].to_torch()
             training_logprobs_D.append(training_logprobs)
 
-        # Return both logprobs and datums (with masks for metrics)
-        return training_logprobs_D, training_datums
+        # Return logprobs, datums (with masks for metrics), and grouping metrics
+        return training_logprobs_D, training_datums, grouping_metrics
 
     async def forward_backward_future(self, episodes: list):
-        training_datums = process_episodes(
+        training_datums, grouping_metrics = process_episodes(
             episodes,
             self.advantage_computer,
             self.trajectory_filter,
@@ -211,7 +212,7 @@ async def forward_backward_future(self, episodes: list):
             loss_fn="importance_sampling",
         )
 
-        return fwd_bwd_future
+        return fwd_bwd_future, grouping_metrics
 
     async def optim_step_future(self, learning_rate: float = None, beta1: float = 0.9, beta2: float = 0.95, eps: float = 1e-8):
         if learning_rate is None: