Skip to content

feat: improve rollout-only trajectory collection #3215

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
joyang-nv wants to merge 1 commit into
base: main
Choose a base branch
from
+435 −37
Open

feat: improve rollout-only trajectory collection #3215

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions examples/nemo_gym/grpo_workplace_assistant_nemotron_nano_v2_9b.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
187 changes: 150 additions & 37 deletions examples/nemo_gym/run_grpo_nemo_gym.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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.
Expand All @@ -13,6 +13,7 @@
# limitations under the License.

import argparse
import json
import os
import pprint
import time
Expand All @@ -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 (
Expand Down Expand Up @@ -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,
Expand All @@ -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)
Expand All @@ -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:
Expand Down Expand Up @@ -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)

Expand All @@ -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:")
Expand All @@ -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,
Expand Down
Loading
Loading