fix: fix rollout metrics fail in unit test (#3173)

Signed-off-by: Yuki Huang <yukih@nvidia.com>

Author: yuki-97

nemo_rl/experience/metric_utils.py

@@ -0,0 +1,53 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Shared aggregation helpers for rollout metrics."""
+
+import math
+import statistics
+from collections.abc import Sequence
+
+from wandb import Histogram
+
+
+def calculate_single_metric(
+    values: Sequence[float | int], batch_size: int, key_name: str
+) -> dict:
+    """Compute summary statistics for a metric as slash-prefixed keys.
+
+    Args:
+        values: Per-sample metric values to aggregate.
+        batch_size: Denominator for the mean (sum(values) / batch_size, not len(values)); stddev still uses len(values).
+        key_name: Prefix for the returned metric keys (e.g. "total_reward").
+
+    Returns:
+        Dict mapping "{key_name}/{stat}" to its value for stat in mean, max, min, median, stddev (nan for a single value), and histogram (a wandb.Histogram).
+    """
+    return {
+        f"{key_name}/mean": sum(values) / batch_size,
+        f"{key_name}/max": max(values),
+        f"{key_name}/min": min(values),
+        f"{key_name}/median": statistics.median(values),
+        f"{key_name}/stddev": statistics.stdev(values) if len(values) > 1 else math.nan,
+        f"{key_name}/histogram": Histogram(values),
+    }
+
+
+def pct(values: Sequence[float | int], p: float) -> float:
+    """Percentile helper for buffer starvation diagnostics."""
+    if not values:
+        return 0.0
+    sorted_v = sorted(values)
+    idx = min(int(len(sorted_v) * p / 100), len(sorted_v) - 1)
+    return float(sorted_v[idx])

nemo_rl/experience/rollout_manager.py

@@ -25,11 +25,8 @@
 from nemo_rl.distributed.batched_data_dict import BatchedDataDict
 from nemo_rl.environments.interfaces import EnvironmentInterface
 from nemo_rl.experience.interfaces import Completion, PromptGroupRecord
-from nemo_rl.experience.rollouts import (
-    _calculate_single_metric,
-    _tensorize_by_key,
-    calculate_rewards,
-)
+from nemo_rl.experience.metric_utils import calculate_single_metric, pct
+from nemo_rl.experience.rollouts import _tensorize_by_key, calculate_rewards
 from nemo_rl.models.generation.interfaces import (
     GenerationConfig,
     GenerationDatumSpec,
@@ -337,17 +334,29 @@ def _aggregate_rollout_metrics(
         terminated = [m["terminated"] for m in all_sample_metrics]
         max_turns_reached = [m["max_turns_reached"] for m in all_sample_metrics]
 
+        # max_gen_tokens_per_turn: Diagnostic for long single generations
+        max_gen_tokens_per_turn = [
+            max(m["turn_gen_tokens"]) if m["turn_gen_tokens"] else 0
+            for m in all_sample_metrics
+        ]
+
         # Aggregate metrics across all samples.
         n = len(all_sample_metrics)
         rollout_metrics: dict[str, Any] = {
-            **_calculate_single_metric(total_reward, n, "total_reward"),
+            **calculate_single_metric(total_reward, n, "total_reward"),
             # turn metrics
             "total_turns": sum(turn_count),
-            **_calculate_single_metric(turn_count, n, "turns_per_sample"),
+            **calculate_single_metric(turn_count, n, "turns_per_sample"),
+            "turns_per_sample/p95": pct(turn_count, 95),
+            "turns_per_sample/p99": pct(turn_count, 99),
             # token metrics
-            **_calculate_single_metric(total_tokens, n, "total_tokens_per_sample"),
-            **_calculate_single_metric(assistant_tokens, n, "gen_tokens_per_sample"),
-            **_calculate_single_metric(env_tokens, n, "env_tokens_per_sample"),
+            **calculate_single_metric(total_tokens, n, "total_tokens_per_sample"),
+            **calculate_single_metric(assistant_tokens, n, "gen_tokens_per_sample"),
+            **calculate_single_metric(env_tokens, n, "env_tokens_per_sample"),
+            # max_gen_tokens_per_turn: Diagnostic for long single generations
+            "max_gen_tokens_per_turn/max": max(max_gen_tokens_per_turn),
+            "max_gen_tokens_per_turn/mean": sum(max_gen_tokens_per_turn) / n,
+            "max_gen_tokens_per_turn/p95": pct(max_gen_tokens_per_turn, 95),
             # truncated metrics
             "truncation_rate": sum(truncated) / n,
             "natural_termination_rate": sum(terminated) / n,
@@ -362,7 +371,7 @@ def _aggregate_rollout_metrics(
             rollout_metrics["per_worker_token_counts"] = per_worker_token_counts
 
         # Per-turn token histograms (flat across all turns, distinct from the
-        # per-sample histograms emitted via _calculate_single_metric above).
+        # per-sample histograms emitted via calculate_single_metric above).
         rollout_metrics["histogram/gen_tokens_length"] = [
             t for m in all_sample_metrics for t in m["turn_gen_tokens"]
         ]
@@ -542,18 +551,36 @@ def _compute_rollout_metrics(
             sum(len(m["token_ids"]) for m in c.message_log if m["role"] == "assistant")
             for c in completions
         ]
+        # max_gen_tokens_per_turn: Diagnostic for long single generations
+        max_gen_tokens_per_turn = [
+            max(
+                (
+                    len(m["token_ids"])
+                    for m in c.message_log
+                    if m["role"] == "assistant"
+                ),
+                default=0,
+            )
+            for c in completions
+        ]
         # truncated metrics
         truncated = [c.truncated for c in completions]
 
         # Aggregate metrics across all samples.
         n = len(completions)
         rollout_metrics: dict[str, Any] = {
-            **_calculate_single_metric(total_reward, n, "total_reward"),
+            **calculate_single_metric(total_reward, n, "total_reward"),
             # turn metrics
-            **_calculate_single_metric(turn_count, n, "turns_per_sample"),
+            **calculate_single_metric(turn_count, n, "turns_per_sample"),
+            "turns_per_sample/p95": pct(turn_count, 95),
+            "turns_per_sample/p99": pct(turn_count, 99),
             # token metrics
-            **_calculate_single_metric(total_tokens, n, "total_tokens_per_sample"),
-            **_calculate_single_metric(assistant_tokens, n, "gen_tokens_per_sample"),
+            **calculate_single_metric(total_tokens, n, "total_tokens_per_sample"),
+            **calculate_single_metric(assistant_tokens, n, "gen_tokens_per_sample"),
+            **calculate_single_metric(
+                max_gen_tokens_per_turn, n, "max_gen_tokens_per_turn"
+            ),
+            "max_gen_tokens_per_turn/p95": pct(max_gen_tokens_per_turn, 95),
             # truncated metrics
             "natural_termination_rate": sum(not t for t in truncated) / n,
             "truncation_rate": sum(truncated) / n,
@@ -569,7 +596,7 @@ def _compute_rollout_metrics(
             ]
             if values:
                 rollout_metrics.update(
-                    _calculate_single_metric(values, n, f"{agent_name}/{key}")
+                    calculate_single_metric(values, n, f"{agent_name}/{key}")
                 )
         rollout_metrics[f"{agent_name}/full_result"] = Table(
             data=[[json.dumps(r, separators=(",", ":"))] for r in agent_extras],

nemo_rl/experience/rollouts.py

@@ -18,19 +18,17 @@
 import asyncio
 import copy
 import json
-import math
 import statistics
 import warnings
 from collections import defaultdict
-from collections.abc import Sequence
 from dataclasses import dataclass
 from typing import Any, Optional
 
 import ray
 import torch
 from pydantic import BaseModel
 from transformers import PreTrainedTokenizerBase
-from wandb import Histogram, Table
+from wandb import Table
 
 from nemo_rl.algorithms.utils import get_gdpo_reward_component_keys
 from nemo_rl.data.interfaces import (
@@ -48,6 +46,7 @@
     EnvironmentReturn,
 )
 from nemo_rl.environments.nemo_gym import DEFAULT_THINKING_TAGS
+from nemo_rl.experience.metric_utils import calculate_single_metric, pct
 from nemo_rl.models.generation.interfaces import (
     GenerationConfig,
     GenerationDatumSpec,
@@ -1263,14 +1262,6 @@ async def run_single_sample_with_error_handling(i, sample_state):
             if key not in final_batch:
                 final_batch[key] = input_batch[key]
 
-        # Helper for percentile (buffer starvation diagnostics)
-        def _pct(values: Sequence[float | int], p: float) -> float:
-            if not values:
-                return 0.0
-            sorted_v = sorted(values)
-            idx = min(int(len(sorted_v) * p / 100), len(sorted_v) - 1)
-            return float(sorted_v[idx])
-
         turn_counts = [m["turn_count"] for m in all_sample_metrics]
         max_gen_tokens_per_turn_values = [
             m["max_gen_tokens_per_turn"] for m in all_sample_metrics
@@ -1282,8 +1273,8 @@ def _pct(values: Sequence[float | int], p: float) -> float:
             "total_turns": sum(turn_counts),
             "avg_turns_per_sample": sum(turn_counts) / batch_size,
             "max_turns_per_sample": max(turn_counts),
-            "turns_per_sample/p95": _pct(turn_counts, 95),
-            "turns_per_sample/p99": _pct(turn_counts, 99),
+            "turns_per_sample/p95": pct(turn_counts, 95),
+            "turns_per_sample/p99": pct(turn_counts, 99),
             "natural_termination_rate": sum(m["terminated"] for m in all_sample_metrics)
             / batch_size,
             "truncation_rate": sum(m["truncated"] for m in all_sample_metrics)
@@ -1312,7 +1303,7 @@ def _pct(values: Sequence[float | int], p: float) -> float:
             "max_gen_tokens_per_turn/max": max(max_gen_tokens_per_turn_values),
             "max_gen_tokens_per_turn/mean": sum(max_gen_tokens_per_turn_values)
             / batch_size,
-            "max_gen_tokens_per_turn/p95": _pct(max_gen_tokens_per_turn_values, 95),
+            "max_gen_tokens_per_turn/p95": pct(max_gen_tokens_per_turn_values, 95),
             # Reward metrics
             "mean_total_reward": sum(m["total_reward"] for m in all_sample_metrics)
             / batch_size,
@@ -1359,19 +1350,6 @@ class AsyncNemoGymRolloutResult:
     rollout_metrics: dict[str, Any]
 
 
-def _calculate_single_metric(
-    values: Sequence[float | int], batch_size: int, key_name: str
-) -> dict:
-    return {
-        f"{key_name}/mean": sum(values) / batch_size,
-        f"{key_name}/max": max(values),
-        f"{key_name}/min": min(values),
-        f"{key_name}/median": statistics.median(values),
-        f"{key_name}/stddev": statistics.stdev(values) if len(values) > 1 else math.nan,
-        f"{key_name}/histogram": Histogram(values),
-    }
-
-
 def get_nemo_gym_thinking_tags(env_config: dict[str, Any]) -> list[str]:
     """Return thinking tags used by the Gym-side detector."""
     nemo_gym_config = env_config.get("nemo_gym")
@@ -1918,39 +1896,32 @@ def run_async_nemo_gym_rollout(
             m["max_gen_tokens_per_turn"] for m in all_sample_metrics
         ]
 
-        def _pct(values: Sequence[float | int], p: float) -> float:
-            if not values:
-                return 0.0
-            sorted_v = sorted(values)
-            idx = min(int(len(sorted_v) * p / 100), len(sorted_v) - 1)
-            return float(sorted_v[idx])
-
         rollout_metrics = {
             **rollout_loop_timing_metrics,
-            **_calculate_single_metric(
+            **calculate_single_metric(
                 turn_counts,
                 batch_size,
                 "turns_per_sample",
             ),
-            "turns_per_sample/p95": _pct(turn_counts, 95),
-            "turns_per_sample/p99": _pct(turn_counts, 99),
-            **_calculate_single_metric(
+            "turns_per_sample/p95": pct(turn_counts, 95),
+            "turns_per_sample/p99": pct(turn_counts, 99),
+            **calculate_single_metric(
                 [m["total_tokens"] for m in all_sample_metrics],
                 batch_size,
                 "total_tokens_per_sample",
             ),
-            **_calculate_single_metric(
+            **calculate_single_metric(
                 [m["assistant_tokens"] for m in all_sample_metrics],
                 batch_size,
                 "gen_tokens_per_sample",
             ),
-            **_calculate_single_metric(
+            **calculate_single_metric(
                 max_gen_tokens_per_turn_values,
                 batch_size,
                 "max_gen_tokens_per_turn",
             ),
-            "max_gen_tokens_per_turn/p95": _pct(max_gen_tokens_per_turn_values, 95),
-            **_calculate_single_metric(
+            "max_gen_tokens_per_turn/p95": pct(max_gen_tokens_per_turn_values, 95),
+            **calculate_single_metric(
                 [m["total_reward"] for m in all_sample_metrics],
                 batch_size,
                 "total_reward",
@@ -1989,7 +1960,7 @@ def _pct(values: Sequence[float | int], p: float) -> float:
                 ]
                 if values:
                     per_agent_metrics.update(
-                        _calculate_single_metric(
+                        calculate_single_metric(
                             values, len(agent_results), f"{agent_name}/{key}"
                         )
                     )

pyrefly.toml

@@ -133,6 +133,7 @@ project-includes = [
   "nemo_rl/evals/answer_parsing.py",
   "nemo_rl/experience/__init__.py",
   "nemo_rl/experience/interfaces.py",
+  "nemo_rl/experience/metric_utils.py",
   "nemo_rl/experience/rollout_manager.py",
   "nemo_rl/experience/rollouts.py",
   "nemo_rl/modelopt/__init__.py",

tests/unit/excluded_unit_tests.sh

@@ -209,10 +209,18 @@ EXCLUDED_UNIT_TESTS=(
     --deselect=tests/unit/models/dtensor/test_parallelize.py::test_parallelize_plan_keys
 
     ###########################################################################
-    # EXPERIENCE — rollout tests all require vLLM setup (~67-91s setup each)
-    ###########################################################################
-
-    --ignore=unit/experience/test_rollouts.py
+    # EXPERIENCE — rollout tests need vLLM setup (~67-91s each).
+    # Keep pure Python tests, plus the 3 matches_original tests that guard correctness.
+    ###########################################################################
+
+    --deselect=tests/unit/experience/test_rollouts.py::test_run_multi_step_calculator_vllm_sync
+    --deselect=tests/unit/experience/test_rollouts.py::test_run_multi_step_calculator_vllm_async
+    --deselect=tests/unit/experience/test_rollouts.py::test_max_seqlen_respected_sync
+    --deselect=tests/unit/experience/test_rollouts.py::test_max_seqlen_respected_async
+    --deselect=tests/unit/experience/test_rollouts.py::test_run_sliding_puzzle_vllm
+    --deselect=tests/unit/experience/test_rollouts.py::test_async_rollout_manager
+    --deselect=tests/unit/experience/test_rollouts.py::test_async_rollout_manager_truncation
+    --deselect=tests/unit/experience/test_rollouts.py::test_async_nemo_gym_rollout_manager
 
     ###########################################################################
     # ENVIRONMENTS