feat: reward profiling

nemo_gym/profile.py

@@ -0,0 +1,73 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+import json
+from collections import defaultdict
+from pathlib import Path
+from typing import Optional
+
+from pydantic import Field
+
+from nemo_gym.config_types import BaseNeMoGymCLIConfig
+from nemo_gym.global_config import get_global_config_dict
+
+
+class ProfileConfig(BaseNeMoGymCLIConfig):
+    input_jsonl_fpath: str = Field(description="Original task dataset.")
+    rollouts_jsonl_fpath: str = Field(description="Rollouts file from ng_collect_rollouts with num_repeats.")
+    output_jsonl_fpath: str = Field(description="Output file for profiled dataset.")
+    pass_threshold: Optional[float] = Field(
+        default=None, description="Reward threshold for pass_rate. If None, pass_rate not computed."
+    )
+
+
+def profile():
+    config = ProfileConfig.model_validate(get_global_config_dict())
+
+    with open(config.input_jsonl_fpath) as f:
+        tasks = [json.loads(line) for line in f]
+
+    with open(config.rollouts_jsonl_fpath) as f:
+        rollouts = [json.loads(line) for line in f]
+
+    grouped = defaultdict(list)
+    for rollout in rollouts:
+        task_idx = rollout.get("_task_index")
+        if task_idx is not None:
+            grouped[task_idx].append(rollout)
+
+    Path(config.output_jsonl_fpath).parent.mkdir(exist_ok=True, parents=True)
+    with open(config.output_jsonl_fpath, "w") as f:
+        for task_idx, task_rollouts in sorted(grouped.items()):
+            if task_idx >= len(tasks):
+                continue
+
+            rewards = [r.get("reward", 0.0) for r in task_rollouts]
+            profiled_task = {**tasks[task_idx]}
+
+            avg = sum(rewards) / len(rewards)
+            profiled_task["avg_reward"] = avg
+            profiled_task["std_reward"] = (sum((r - avg) ** 2 for r in rewards) / len(rewards)) ** 0.5
+            profiled_task["min_reward"] = min(rewards)
+            profiled_task["max_reward"] = max(rewards)
+            profiled_task["total_samples"] = len(rewards)
+
+            if config.pass_threshold is not None:
+                passed = sum(1 for r in rewards if r >= config.pass_threshold)
+                profiled_task["pass_rate"] = passed / len(rewards)
+                profiled_task["pass_rate_total"] = len(rewards)
+                profiled_task["pass_rate_passed"] = passed
+                profiled_task["pass_threshold"] = config.pass_threshold
+
+            f.write(json.dumps(profiled_task) + "\n")

nemo_gym/rollout_collection.py

@@ -17,7 +17,7 @@
 from asyncio import Future, Semaphore
 from collections import Counter
 from contextlib import nullcontext
-from itertools import chain, repeat
+from itertools import repeat
 from pathlib import Path
 from typing import Any, Dict, Iterator, List, Optional, Tuple
 
@@ -90,7 +90,11 @@ async def run_from_config(self, config: RolloutCollectionConfig):
 
         if config.num_repeats:
             previous_length = len(rows)
-            rows = list(chain.from_iterable(repeat(row, config.num_repeats) for row in rows))
+            expanded = []
+            for task_idx, row in enumerate(rows):
+                for _ in range(config.num_repeats):
+                    expanded.append({**row, "_task_index": task_idx})
+            rows = expanded
             print(f"Repeating rows (in a pattern of abc to aabbcc) from {previous_length} to {len(rows)}!")
 
         semaphore = nullcontext()
@@ -128,8 +132,12 @@ async def _post_coroutine(row: dict) -> None:
                     response = await server_client.post(server_name=agent_name, url_path="/run", json=row)
                     await raise_for_status(response)
                     result = await get_response_json(response)
+                    if "_task_index" in row:
+                        result["_task_index"] = row["_task_index"]
                     f.write(json.dumps(result) + "\n")
-                    metrics.update({k: v for k, v in result.items() if isinstance(v, (int, float))})
+                    metrics.update(
+                        {k: v for k, v in result.items() if isinstance(v, (int, float)) and not k.startswith("_")}
+                    )
 
             await tqdm.gather(*map(_post_coroutine, rows), desc="Collecting rollouts", miniters=tqdm_miniters)
 

pyproject.toml

@@ -76,18 +76,15 @@ dependencies = [
     #
     # By design, most (if not all) dependencies are unfrozen here to be easier to consume. The core pieces we need are server infra like FastAPI, etc.
     ########################################
-
     # OpenAI: We leverage OpenAI Responses, Chat Completions, and Completions schemas for Nemo Gym abstractions. It may also be used to directly query endpoints.
     # We specifically upper bound this OpenAI dependency since the version bumps so frequently.
     # Updated Wed Oct 29, 2025 with openai==2.6.1
     # License: Apache 2.0 https://github.com/openai/openai-python/blob/a8258744cbecf51321587fc870e8920bd2c07809/LICENSE
     "openai<=2.6.1",
-
     # tqdm: Used for progress tracking on batch operations.
     # Updated Fri Jul 25, 2025 with tqdm==4.67.1
     # License: MIT https://github.com/tqdm/tqdm/blob/0ed5d7f18fa3153834cbac0aa57e8092b217cc16/LICENCE
     "tqdm",
-
     # Pydantic: Used for typing and import/export.
     # Updated Fri Jul 25, 2025 with pydantic==2.11.7 and pydantic_core==2.33.2
     # pydantic license: MIT https://github.com/pydantic/pydantic/blob/1c79f0e4d3fbdb8b93e837175d7098e016117237/LICENSE
@@ -96,64 +93,52 @@ dependencies = [
     "pydantic",
     "pydantic_core",
     "devtools",
-
     # FastAPI: Used for server infrastructure
     # Updated Fri Jul 25, 2025 with fastapi==0.116.1
     # License: MIT https://github.com/fastapi/fastapi/blob/6df50d40fe195adc026af169d6ebf298a1c183a5/LICENSE
     "fastapi",
-
     # Uvicorn: Used to serve FastAPI apps
     # Updated Mon Jul 28, 2025 with uvicorn==0.35.0
     # License: BSD 3-Clause https://github.com/encode/uvicorn/blob/c1144fd4f130388cffc05ee17b08747ce8c1be11/LICENSE.md
     "uvicorn",
-
     # UVLoop: a faster async event loop than Python's native asyncio. Used automatically by Uvicorn as an async loop backend.
     # Updated Fri Aug 01, 2025 with uvloop==0.21.0
     # License: Apache 2.0 and MIT https://github.com/MagicStack/uvloop/blob/96b7ed31afaf02800d779a395591da6a2c8c50e1/LICENSE-APACHE https://github.com/MagicStack/uvloop/blob/96b7ed31afaf02800d779a395591da6a2c8c50e1/LICENSE-MIT
     "uvloop",
-
     # Hydra and OmegaConf: CLI Configuration utilities
     # Updated Tue Jul 29, 2025 with hydra-core==1.3.2 and omegaconf==2.3.0
     # hydra-core license: MIT https://github.com/facebookresearch/hydra/blob/737fc3349ef3a4031035645f7e8c80be66a57042/LICENSE
     # omegaconf license: BSD 3-Clause https://github.com/omry/omegaconf/blob/117f7de07285e4d1324b9229eaf873de15279457/LICENSE
     "hydra-core",
     "omegaconf",
-
     # Gradio: For simple frontend interfaces for viewing data
     # Updated Sun Aug 03, 2025 with gradio==5.16.0
     # License: Apache 2.0 https://github.com/gradio-app/gradio/blob/2b4432edea8a62659e180e24eedd2bddbed08e77/LICENSE
     "gradio",
-
     # MLFlow: used for interacting with the Gitlab model registry
     # Updated Tue Aug 05, 2025 with mlflow==3.2.0
     # License: Apache 2.0 https://github.com/mlflow/mlflow/blob/1510ed1bc92d3a4258973005d64f64a43136e251/LICENSE.txt
     "mlflow",
-
     # aiohttp: async http backend
     # Updated Sun Sep 21, 2025 with aiohttp==3.12.15
     # License: Apache 2.0 https://github.com/aio-libs/aiohttp/blob/9a2f146a12e3525b43e96723ef41584bf9cf784e/LICENSE.txt
     "aiohttp",
-
     # yappi: profiling tool
     # Updated Mon Sep 22, 2025 with yappi==1.6.10
     # License: MIT https://github.com/sumerc/yappi/blob/1d3f7501701e1f050b6dcd6a86fd36aec08185c7/LICENSE
     "yappi",
-
     # Ray: Used for distributed processing
     # Updated Fri Oct 18, 2025 with ray[default]==2.46.0
     # License: Apache 2.0 https://github.com/ray-project/ray/blob/master/LICENSE
     "ray[default]",
-
     # psutil: Cross-platform process and system utilities
     # Updated: Fri Nov 07, 2025 with psutil==6.1.1
     # License: BSD 3-Clause https://github.com/giampaolo/psutil/blob/master/LICENSE
     "psutil",
-
     # HuggingFace datasets: for loading and converting parquet datasets
     # Updated Thu Dec 04, 2025 with datasets==4.4.1
     # License: Apache 2.0 https://github.com/huggingface/datasets/blob/main/LICENSE
     "datasets",
-
     # orjson: JSON library for faster serialization/deserialization
     # Updated: Thu Jan 08, 2026 with orjson==3.11.3
     # License: Apache 2.0 https://github.com/ijl/orjson/blob/fb3eb1f729c7e7b019f780af5695722c99c7c695/LICENSE-APACHE
@@ -255,6 +240,9 @@ ng_init_resources_server = "nemo_gym.cli:init_resources_server"
 nemo_gym_collect_rollouts = "nemo_gym.rollout_collection:collect_rollouts"
 ng_collect_rollouts = "nemo_gym.rollout_collection:collect_rollouts"
 
+# Reward profiling
+ng_profile = "nemo_gym.profile:profile"
+
 # Dataset management
 nemo_gym_upload_dataset_to_gitlab = "nemo_gym.gitlab_utils:upload_jsonl_dataset_cli"
 ng_upload_dataset_to_gitlab = "nemo_gym.gitlab_utils:upload_jsonl_dataset_cli"

tests/unit_tests/test_profile.py

@@ -0,0 +1,24 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+from nemo_gym.profile import ProfileConfig
+
+
+class TestProfile:
+    def test_sanity(self) -> None:
+        ProfileConfig(
+            input_jsonl_fpath="",
+            rollouts_jsonl_fpath="",
+            output_jsonl_fpath="",
+        )