Fix training metrics before and after processing (#145)

agentlightning/verl/trainer.py

@@ -19,7 +19,7 @@ from verl import DataProto
 from verl.protocol import pad_dataproto_to_divisor, unpad_dataproto
 from verl.trainer.ppo.core_algos import agg_loss
 from verl.trainer.ppo.metric_utils import (
-    compute_data_metrics,
+    _compute_response_info,
     compute_throughout_metrics,
     compute_timing_metrics,
 )
@@ -53,6 +53,108 @@ def _timer(name: str, timing_raw: Dict[str, float]):
     timing_raw[name] += timer.last
 
 
+# This function is adapted from verl.
+# We introduce a new parameter `suffix` to distinguish between metrics computed
+# before and after AgentLightning’s post-processing.
+# - "Before" refers to raw reward and advantage values.
+# - "After" refers to values computed following post-processing, which involves:
+#     (1) Dropping prompts that exceed the maximum allowed length.
+#     (2) Adjusting the batch size to be a multiple of the mini PPO size.
+# Different suffixes are used to label these two stages accordingly.
+def compute_data_metrics(batch: DataProto, use_critic: bool = True, suffix: str = "") -> Dict[str, Any]:
+    """
+    Computes various metrics from a batch of data for PPO training.
+
+    This function calculates metrics related to scores, rewards, advantages, returns, values,
+    and sequence lengths from a batch of data. It provides statistical information (mean, max, min)
+    for each metric category.
+
+    Args:
+        batch: A DataProto object containing batch data with token-level scores, rewards, advantages, etc.
+        use_critic: Whether to include critic-specific metrics. Defaults to True.
+
+    Returns:
+        A dictionary of metrics including:
+            - critic/score/mean, max, min: Statistics about sequence scores
+            - critic/rewards/mean, max, min: Statistics about sequence rewards
+            - critic/advantages/mean, max, min: Statistics about advantages
+            - critic/returns/mean, max, min: Statistics about returns
+            - critic/values/mean, max, min: Statistics about critic values (if use_critic=True)
+            - critic/vf_explained_var: Explained variance of the value function (if use_critic=True)
+            - response_length/mean, max, min, clip_ratio: Statistics about response lengths
+            - prompt_length/mean, max, min, clip_ratio: Statistics about prompt lengths
+    """
+    sequence_score = batch.batch["token_level_scores"].sum(-1)
+    sequence_reward = batch.batch["token_level_rewards"].sum(-1)
+
+    advantages = batch.batch["advantages"]
+    returns = batch.batch["returns"]
+
+    max_response_length = batch.batch["responses"].shape[-1]
+
+    prompt_mask = batch.batch["attention_mask"][:, :-max_response_length].bool()
+    response_mask = batch.batch["attention_mask"][:, -max_response_length:].bool()
+
+    max_prompt_length = prompt_mask.size(-1)
+
+    response_info = _compute_response_info(batch)
+    prompt_length = response_info["prompt_length"]
+    response_length = response_info["response_length"]
+
+    valid_adv = torch.masked_select(advantages, response_mask)
+    valid_returns = torch.masked_select(returns, response_mask)
+
+    if use_critic:
+        values = batch.batch["values"]
+        valid_values = torch.masked_select(values, response_mask)
+        return_diff_var = torch.var(valid_returns - valid_values)
+        return_var = torch.var(valid_returns)
+
+    metrics = {
+        # score
+        "critic/score/mean" + suffix: torch.mean(sequence_score).detach().item(),
+        "critic/score/max" + suffix: torch.max(sequence_score).detach().item(),
+        "critic/score/min" + suffix: torch.min(sequence_score).detach().item(),
+        # reward
+        "critic/rewards/mean" + suffix: torch.mean(sequence_reward).detach().item(),
+        "critic/rewards/max" + suffix: torch.max(sequence_reward).detach().item(),
+        "critic/rewards/min" + suffix: torch.min(sequence_reward).detach().item(),
+        # adv
+        "critic/advantages/mean" + suffix: torch.mean(valid_adv).detach().item(),
+        "critic/advantages/max" + suffix: torch.max(valid_adv).detach().item(),
+        "critic/advantages/min" + suffix: torch.min(valid_adv).detach().item(),
+        # returns
+        "critic/returns/mean" + suffix: torch.mean(valid_returns).detach().item(),
+        "critic/returns/max" + suffix: torch.max(valid_returns).detach().item(),
+        "critic/returns/min" + suffix: torch.min(valid_returns).detach().item(),
+        **(
+            {
+                # values
+                "critic/values/mean" + suffix: torch.mean(valid_values).detach().item(),
+                "critic/values/max" + suffix: torch.max(valid_values).detach().item(),
+                "critic/values/min" + suffix: torch.min(valid_values).detach().item(),
+                # vf explained var
+                "critic/vf_explained_var" + suffix: (1.0 - return_diff_var / (return_var + 1e-5)).detach().item(),
+            }
+            if use_critic
+            else {}
+        ),
+        # response length
+        "response_length/mean" + suffix: torch.mean(response_length).detach().item(),
+        "response_length/max" + suffix: torch.max(response_length).detach().item(),
+        "response_length/min" + suffix: torch.min(response_length).detach().item(),
+        "response_length/clip_ratio"
+        + suffix: torch.mean(torch.eq(response_length, max_response_length).float()).detach().item(),
+        # prompt length
+        "prompt_length/mean" + suffix: torch.mean(prompt_length).detach().item(),
+        "prompt_length/max" + suffix: torch.max(prompt_length).detach().item(),
+        "prompt_length/min" + suffix: torch.min(prompt_length).detach().item(),
+        "prompt_length/clip_ratio"
+        + suffix: torch.mean(torch.eq(prompt_length, max_prompt_length).float()).detach().item(),
+    }
+    return metrics
+
+
 class AgentLightningTrainer(RayPPOTrainer):
     """
     Specialized PPO trainer for agent-based reinforcement learning.
@@ -215,6 +317,9 @@ class AgentLightningTrainer(RayPPOTrainer):
                     config=self.config.algorithm,
                 )
 
+            # Calculate the metrics before processing. Refer to the comments of function `compute_data_metrics` for details.
+            metrics.update(compute_data_metrics(batch=batch, use_critic=self.use_critic, suffix="_before_processing"))
+
             # after advantages are assinged, we begin to drop (1) long prompt (2) floor to ppo minisize
             keep_indices = (~batch.batch["is_drop_mask"]).nonzero(as_tuple=True)[0]
             metrics["training/n_triplets_prompt_too_long"] = (
@@ -274,7 +379,7 @@ class AgentLightningTrainer(RayPPOTrainer):
                     )
 
         # compute training metrics
-        metrics.update(compute_data_metrics(batch=batch, use_critic=self.use_critic))
+        metrics.update(compute_data_metrics(batch=batch, use_critic=self.use_critic, suffix="_after_processing"))
         metrics.update(compute_timing_metrics(batch=batch, timing_raw=timing_raw))
         # TODO: implement actual tflpo and theoretical tflpo
         n_gpus = self.resource_pool_manager.get_n_gpus()