feat: improve rollout-only trajectory collection

joyang-nv · joyang-nv · commit 6ae9227ab433 · 2026-07-15T16:12:21.000+08:00
Signed-off-by: Jonas Yang &lt;joyang@nvidia.com&gt;
diff --git a/examples/nemo_gym/grpo_workplace_assistant_nemotron_nano_v2_9b.yaml b/examples/nemo_gym/grpo_workplace_assistant_nemotron_nano_v2_9b.yaml
@@ -299,6 +299,7 @@ env:
     port_range_high: 5999
     rollout_max_attempts_to_avoid_lp_nan: 1
     is_trajectory_collection: false  # Set this to true to enable trajectory collection (no training). You may also want to increase `policy.generation.vllm_cfg.gpu_memory_utilization`
+    trajectory_collection_batch_size: null  # Optional positive integer; null collects the validation set in one batch
     config_paths:
     - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml  # Required! And it must be *for_training
     - resources_servers/workplace_assistant/configs/workplace_assistant.yaml
diff --git a/examples/nemo_gym/run_grpo_nemo_gym.py b/examples/nemo_gym/run_grpo_nemo_gym.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+# 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.
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import argparse
+import json
 import os
 import pprint
 import time
@@ -38,7 +39,7 @@
     refit_policy_generation,
     setup,
 )
-from nemo_rl.algorithms.utils import get_tokenizer
+from nemo_rl.algorithms.utils import get_tokenizer, log_generation_metrics_to_wandb
 from nemo_rl.data.utils import setup_response_data
 from nemo_rl.distributed.virtual_cluster import init_ray
 from nemo_rl.environments.nemo_gym import (
@@ -68,6 +69,34 @@ def parse_args() -> tuple[argparse.Namespace, list[str]]:
     return args, overrides
 
 
+def _pop_trajectory_collection_settings(
+    nemo_gym_config: dict[str, object],
+) -> tuple[bool, int | None]:
+    """Remove and validate NeMo-RL trajectory-collection settings."""
+    is_trajectory_collection = bool(
+        nemo_gym_config.pop("is_trajectory_collection", False)
+    )
+    batch_size = nemo_gym_config.pop("trajectory_collection_batch_size", None)
+    if batch_size is None:
+        return is_trajectory_collection, None
+
+    if not is_trajectory_collection:
+        raise ValueError(
+            "env.nemo_gym.trajectory_collection_batch_size requires "
+            "env.nemo_gym.is_trajectory_collection=true"
+        )
+    if (
+        isinstance(batch_size, bool)
+        or not isinstance(batch_size, int)
+        or batch_size <= 0
+    ):
+        raise ValueError(
+            "env.nemo_gym.trajectory_collection_batch_size must be a positive integer"
+        )
+
+    return is_trajectory_collection, batch_size
+
+
 # These types are directly imported from grpo_train since if something about the architecture changes we want to immediately fail.
 def collect_trajectories(
     policy: ColocatablePolicyInterface,
@@ -78,7 +107,13 @@ def collect_trajectories(
     logger: Logger,
     master_config: MasterConfig,
 ) -> None:
-    """Run trajectory collection."""
+    """Run trajectory collection and persist every completed batch."""
+    expected_trajectories = master_config.grpo["max_val_samples"]
+    if expected_trajectories is None or expected_trajectories <= 0:
+        raise ValueError(
+            "Trajectory collection requires a non-empty validation dataset"
+        )
+
     # common config/state items
     colocated_inference = master_config.policy["generation"]["colocated"]["enabled"]
     refit_policy_generation(policy, policy_generation, colocated_inference)
@@ -87,33 +122,108 @@ def collect_trajectories(
 
     print("\n🔍 Running trajectory collection...", flush=True)
     generation_config = master_config.policy["generation"]
-    for val_batch in val_dataloader:
-        nemo_gym_rollout_result = run_async_nemo_gym_rollout(
-            policy_generation=policy_generation,
-            input_batch=val_batch,
-            tokenizer=tokenizer,
-            task_to_env=val_task_to_env,
-            max_seq_len=master_config.policy["max_total_sequence_length"],
-            generation_config=generation_config,
-            max_rollout_turns=None,
-            greedy=False,
-        )
-
-        rows_to_log: list[str] = []
-        for key, value in nemo_gym_rollout_result.rollout_metrics.items():
-            if "full_result" not in key:
-                continue
-
-            value: Table
-            data: list[list[str]] = value.data  # (n, 1)
-            rows_to_log.extend(v[0] for v in data)
-
-        logger.log_string_list_as_jsonl(rows_to_log, log_filename)
+    vllm_config = generation_config.get("vllm_cfg", {})
+    should_log_generation_metrics = (
+        vllm_config.get("enable_vllm_metrics_logger", False)
+        and vllm_config.get("async_engine", False)
+        and master_config.logger["wandb_enabled"]
+    )
+    collected_trajectories = 0
+    total_reward = 0.0
+
+    try:
+        for batch_idx, val_batch in enumerate(val_dataloader):
+            batch_step = batch_idx + 1
+            if should_log_generation_metrics:
+                policy_generation.clear_logger_metrics()
+
+            nemo_gym_rollout_result = run_async_nemo_gym_rollout(
+                policy_generation=policy_generation,
+                input_batch=val_batch,
+                tokenizer=tokenizer,
+                task_to_env=val_task_to_env,
+                max_seq_len=master_config.policy["max_total_sequence_length"],
+                generation_config=generation_config,
+                max_rollout_turns=None,
+                greedy=False,
+            )
+            if should_log_generation_metrics:
+                generation_logger_metrics = policy_generation.get_logger_metrics()
+
+            rows_to_log: list[str] = []
+            for key, value in nemo_gym_rollout_result.rollout_metrics.items():
+                if "full_result" not in key:
+                    continue
+
+                value: Table
+                data: list[list[str]] = value.data  # (n, 1)
+                rows_to_log.extend(v[0] for v in data)
+
+            if not rows_to_log:
+                raise RuntimeError(
+                    f"Trajectory batch {batch_idx} did not contain any full Gym results"
+                )
+
+            attributed_rows: list[str] = []
+            batch_size = len(rows_to_log)
+            batch_reward = 0.0
+            for batch_position, serialized_result in enumerate(rows_to_log):
+                result = json.loads(serialized_result)
+                result["trajectory_collection_batch_index"] = batch_idx
+                result["trajectory_collection_batch_position"] = batch_position
+                result["trajectory_collection_batch_size"] = batch_size
+                batch_reward += float(result["reward"])
+                attributed_rows.append(json.dumps(result, separators=(",", ":")))
+
+            # Append after every completed batch so earlier trajectories survive a later
+            # batch or worker failure during a long collection run.
+            logger.log_string_list_as_jsonl(attributed_rows, log_filename)
+            collected_trajectories += batch_size
+            total_reward += batch_reward
+
+            batch_rollout_metrics = {
+                key: value
+                for key, value in nemo_gym_rollout_result.rollout_metrics.items()
+                if "full_result" not in key
+            }
+            # Match the training prefix so rollout-only and GRPO runs expose the same
+            # timing and rollout metric names for direct comparison.
+            logger.log_metrics(batch_rollout_metrics, batch_step, prefix="train")
+            if should_log_generation_metrics:
+                log_generation_metrics_to_wandb(
+                    generation_logger_metrics,
+                    batch_step,
+                    vllm_config["vllm_metrics_logger_interval"],
+                    logger,
+                )
+            logger.log_metrics(
+                {
+                    "mean_reward": total_reward / collected_trajectories,
+                    "num_trajectories": collected_trajectories,
+                },
+                batch_step,
+                prefix="trajectory_collection",
+                step_finished=True,
+            )
+            print(
+                f"Collected {collected_trajectories}/{expected_trajectories} "
+                f"trajectories after batch {batch_idx + 1}",
+                flush=True,
+            )
+    finally:
+        policy_generation.finish_generation()
 
-        # TODO: eventually as trajectory collection use cases exceed 4 hours, we can leverage the dataloader save functionality to resume
-        # And also leverage the TimeoutChecker functionality as well
+    if collected_trajectories != expected_trajectories:
+        raise RuntimeError(
+            "Trajectory collection was incomplete: "
+            f"expected {expected_trajectories}, got {collected_trajectories}"
+        )
 
-    policy_generation.finish_generation()
+    print(
+        f"Trajectory collection complete: {collected_trajectories} trajectories, "
+        f"mean reward {total_reward / collected_trajectories:.6f}",
+        flush=True,
+    )
 
 
 def main() -> None:
@@ -163,6 +273,12 @@ def main() -> None:
         # NeMo-Gym specific config setup.
         setup_nemo_gym_config(config, tokenizer)
 
+        # These are NeMo-RL control-flow settings, not NeMo-Gym global config.
+        (
+            is_trajectory_collection,
+            trajectory_collection_batch_size,
+        ) = _pop_trajectory_collection_settings(config.env["nemo_gym"])
+
     # We assert here since this is right after the final config has been materialized.
     assert _should_use_nemo_gym(config)
 
@@ -184,11 +300,15 @@ def main() -> None:
         )
 
     if val_dataset is not None:
+        val_batch_size = len(val_dataset)
+        if trajectory_collection_batch_size is not None:
+            val_batch_size = min(trajectory_collection_batch_size, len(val_dataset))
         print(
-            f"Setting `grpo.max_val_samples` and `grpo.val_batch_size` to the length of the validation dataset, which is {len(val_dataset)}"
+            f"Setting `grpo.max_val_samples` to {len(val_dataset)} and "
+            f"`grpo.val_batch_size` to {val_batch_size}"
         )
         config.grpo["max_val_samples"] = len(val_dataset)
-        config.grpo["val_batch_size"] = config.grpo["max_val_samples"]
+        config.grpo["val_batch_size"] = val_batch_size
 
     # Print config
     print("Final config:")
@@ -197,13 +317,6 @@ def main() -> None:
     with rl_init_timer.time("ray_connect"):
         init_ray()
 
-    # `is_trajectory_collection` is a NeMo-RL-side control-flow knob; pop it
-    # before setup() so it is not forwarded into NeMo-Gym's global config (the
-    # gym actor is now created inside setup()).
-    is_trajectory_collection = (
-        config.env["nemo_gym"].pop("is_trajectory_collection", False) or False
-    )
-
     with rl_init_timer.time("setup"):
         (
             policy,
diff --git a/tests/unit/test_run_grpo_nemo_gym.py b/tests/unit/test_run_grpo_nemo_gym.py
