support run nemo-gym grpo

Signed-off-by: ruit <ruit@nvidia.com>

Author: RayenTian

examples/nemo_gym/grpo_dapo17k_bytedtsinghua_qwen3_4binstruct_nf.yaml

@@ -211,7 +211,7 @@ policy:
       num_first_layers_in_bf16: 0
       expose_http_server: true
       skip_tokenizer_init: false
-      kv_cache_dtype: null
+      kv_cache_dtype: ${policy.precision}
       http_server_serving_chat_kwargs:
         # This is the tool parser for Qwen 3 4B Instruct. This needs to be changed for other models.
         enable_auto_tools: true
@@ -234,10 +234,21 @@ policy:
         num_nodes: null # Decides number of nodes to be dedicated to generation
 
 data:
-  train_jsonl_fpath: 3rdparty/Gym-workspace/Gym/data/train.jsonl
-  validation_jsonl_fpath: 3rdparty/Gym-workspace/Gym/data/validation.jsonl
+  max_input_seq_length: ${policy.max_total_sequence_length}
   shuffle: true
   num_workers: 0
+  train:
+    dataset_name: NemoGymDataset
+    data_path: 3rdparty/Gym-workspace/Gym/data/train.jsonl
+    repeat: 1
+  validation:
+    dataset_name: NemoGymDataset
+    data_path: 3rdparty/Gym-workspace/Gym/data/validation.jsonl
+  default:
+    env_name: "nemo_gym"
+    prompt_file: null
+    system_prompt_file: null
+    processor: "nemo_gym_data_processor"
 
 env:
   should_use_nemo_gym: true

examples/nemo_gym/run_grpo_nemo_gym.py

@@ -17,14 +17,15 @@
 import os
 import pprint
 from itertools import chain, repeat
-from typing import Optional
+from typing import Dict, Optional
 
 # Increase the W&B single object size warning threshold. Initially 100_000 (100 KB) -> 10_000_000 (10 MB)
 import wandb.util
 
 wandb.util.VALUE_BYTES_LIMIT = 10_000_000
 
 import ray
+from datasets import concatenate_datasets
 from omegaconf import OmegaConf
 from wandb import Table
 
@@ -42,18 +43,19 @@
     setup,
 )
 from nemo_rl.algorithms.utils import get_tokenizer
-from nemo_rl.data.datasets import AllTaskProcessedDataset
-from nemo_rl.data.interfaces import DatumSpec
-from nemo_rl.distributed.ray_actor_environment_registry import (
-    get_actor_python_env,
+from nemo_rl.data.datasets import (
+    AllTaskProcessedDataset,
+    load_response_dataset,
+    update_single_dataset_config,
 )
+from nemo_rl.data.interfaces import DatumSpec
 from nemo_rl.distributed.virtual_cluster import init_ray
 from nemo_rl.environments.nemo_gym import (
-    NemoGym,
     NemoGymConfig,
     nemo_gym_example_to_nemo_rl_datum_spec,
     setup_nemo_gym_config,
 )
+from nemo_rl.environments.utils import create_env
 from nemo_rl.experience.rollouts import run_async_nemo_gym_rollout
 from nemo_rl.models.generation import configure_generation_config
 from nemo_rl.utils.config import load_config, parse_hydra_overrides
@@ -109,6 +111,80 @@ def setup_single_nemo_gym_dataset(
     )
 
 
+def setup_data(
+    tokenizer: TokenizerType,
+    data_config: Dict,
+    env_configs: Dict,
+    seed: int,
+) -> tuple[
+    AllTaskProcessedDataset,
+    Optional[AllTaskProcessedDataset],
+    dict[str, EnvironmentInterface],
+    dict[str, EnvironmentInterface],
+]:
+    print("\n▶ Setting up data...")
+    # setup train dataset
+    data_list = []
+    task_data_processors = {}
+
+    if isinstance(data_config["train"], dict):
+        data_config["train"] = [data_config["train"]]
+    for cfg in data_config["train"]:
+        update_single_dataset_config(cfg, data_config["default"])
+        data = load_response_dataset(cfg, seed)
+        data_list.append(data)
+        task_data_processors[data.task_name] = (data.task_spec, data.processor)
+
+    merged_data = concatenate_datasets([data.dataset for data in data_list])
+    dataset = AllTaskProcessedDataset(
+        merged_data,
+        tokenizer,
+        None,
+        task_data_processors,
+        max_seq_length=data_config["max_input_seq_length"],
+    )
+    print(f"  ✓ Training dataset loaded with {len(dataset)} samples.")
+
+    # setup validation dataset
+    val_task_data_processors = {}
+    val_data_list = []
+
+    for data in data_list:
+        if hasattr(data, "val_dataset") and data.val_dataset is not None:
+            val_data_list.append(data.val_dataset)
+            # bind task_name to task_data_processors
+            task_name = data.task_name
+            val_task_data_processors[task_name] = task_data_processors[task_name]
+
+    if data_config["validation"] is not None:
+        if isinstance(data_config["validation"], dict):
+            data_config["validation"] = [data_config["validation"]]
+
+        for cfg in data_config["validation"]:
+            update_single_dataset_config(cfg, data_config["default"])
+            val_data = load_response_dataset(cfg, seed)
+            val_data_list.append(val_data.dataset)
+            # bind task_name to task_data_processors
+            val_task_data_processors[val_data.task_name] = (
+                val_data.task_spec,
+                val_data.processor,
+            )
+
+    val_dataset = None
+    if len(val_data_list) > 0:
+        merged_val_data = concatenate_datasets(val_data_list)
+        val_dataset = AllTaskProcessedDataset(
+            merged_val_data,
+            tokenizer,
+            None,
+            val_task_data_processors,
+            max_seq_length=data_config["max_input_seq_length"],
+        )
+        print(f"  ✓ Validation dataset loaded with {len(val_dataset)} samples.")
+
+    return dataset, val_dataset
+
+
 # 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,
@@ -202,13 +278,11 @@ def main() -> None:
     assert _should_use_nemo_gym(config)
 
     print("\n▶ Setting up data...")
-    train_dataset = setup_single_nemo_gym_dataset(
-        jsonl_fpath=config["data"]["train_jsonl_fpath"],
-        tokenizer=tokenizer,
-    )
-    val_dataset = setup_single_nemo_gym_dataset(
-        jsonl_fpath=config["data"]["validation_jsonl_fpath"],
+    train_dataset, val_dataset = setup_data(
         tokenizer=tokenizer,
+        data_config=config["data"],
+        env_configs=config["env"],
+        seed=config["grpo"]["seed"],
     )
 
     # Validation dataset config setup.
@@ -254,13 +328,7 @@ def main() -> None:
         base_urls=policy_generation.dp_openai_server_base_urls,
         initial_global_config_dict=config["env"]["nemo_gym"],
     )
-    nemo_gym = NemoGym.options(
-        runtime_env={
-            "py_executable": get_actor_python_env(
-                "nemo_rl.environments.nemo_gym.NemoGym"
-            ),
-        }
-    ).remote(nemo_gym_config)
+    nemo_gym = create_env(env_name="nemo_gym", env_config=nemo_gym_config)
     # Blocking wait for NeMo-Gym to spin up
     ray.get(nemo_gym.health_check.remote())
     task_to_env = {"nemo_gym": nemo_gym}

nemo_rl/data/datasets/response_datasets/__init__.py

@@ -24,6 +24,7 @@
 from nemo_rl.data.datasets.response_datasets.deepscaler import DeepScalerDataset
 from nemo_rl.data.datasets.response_datasets.geometry3k import Geometry3KDataset
 from nemo_rl.data.datasets.response_datasets.helpsteer3 import HelpSteer3Dataset
+from nemo_rl.data.datasets.response_datasets.nemogym_dataset import NemoGymDataset
 from nemo_rl.data.datasets.response_datasets.oai_format_dataset import (
     OpenAIFormatDataset,
 )
@@ -87,6 +88,8 @@ def load_response_dataset(data_config: ResponseDatasetConfig, seed: int = 42):
             **data_config,  # pyrefly: ignore[missing-argument]  `data_path` is required for this class
             seed=seed,
         )
+    elif dataset_name == "NemoGymDataset":
+        base_dataset: Any = NemoGymDataset(**data_config)
     else:
         raise ValueError(
             f"Unsupported {dataset_name=}. "
@@ -115,4 +118,5 @@ def load_response_dataset(data_config: ResponseDatasetConfig, seed: int = 42):
     "SquadDataset",
     "Tulu3SftMixtureDataset",
     "HelpSteer3Dataset",
+    "NemoGymDataset",
 ]

nemo_rl/data/datasets/response_datasets/nemogym_dataset.py

@@ -0,0 +1,59 @@
+# Copyright (c) 2025, 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.
+
+from typing import Any, Optional
+
+import torch
+
+from nemo_rl.data.datasets.raw_dataset import RawDataset
+from nemo_rl.data.datasets.utils import load_dataset_from_path
+
+
+class NemoGymDataset(RawDataset):
+    """Simple wrapper around the Nemo Gym dataset."""
+
+    def __init__(self, data_path: Optional[str] = None, **kwargs) -> None:
+        self.task_name = "NemoGymDataset"
+
+        # load from jsonl
+        if data_path is None:
+            # Allow optional at type level for config validation; enforce at runtime for clarity
+            raise ValueError(
+                "NemoGymDataset requires `data_path` in data_config to load examples."
+            )
+        self.dataset = load_dataset_from_path(data_path)
+
+        # format the dataset
+        # HuggingFace Dataset 在 map/写入 Arrow 时不会持久化 torch.Tensor，会把它序列化成 Python 列表。因此下游在取样时读到的是 []（list），触发断言
+        self.dataset = self.dataset.map(
+            self.format_data,
+            with_indices=True,
+        )
+        if "repeat" in kwargs:
+            self.dataset = self.dataset.repeat(kwargs["repeat"])
+
+    def format_data(self, data: dict[str, Any], idx: int) -> dict[str, Any]:
+        return {
+            "message_log": [
+                {"role": "user", "content": "", "token_ids": torch.tensor([])}
+            ],
+            "task_name": self.task_name,
+            "length": 0,
+            "extra_env_info": data,
+            "loss_multiplier": 1.0,  # Fix to 1.0 to backprop on all examples
+            "idx": idx,
+            "stop_strings": None,
+            # Extra vars
+            "token_ids": [],  # Just need this empty key to be compatible with the current NeMo RL GRPO impl
+        }

nemo_rl/data/processors.py

@@ -538,6 +538,26 @@ def multichoice_qa_processor(
     return output
 
 
+def nemo_gym_data_processor(
+    datum_dict: dict[str, Any],
+    *args,
+    **kwargs,
+) -> DatumSpec:
+    """Process a datum dictionary (directly loaded from dataset) into a DatumSpec for Nemo Gym."""
+    # Ensure message_log exists and contains tensor token_ids so downstream padding works
+    if "message_log" not in datum_dict or not datum_dict["message_log"]:
+        datum_dict["message_log"] = [
+            {"role": "user", "content": "", "token_ids": torch.tensor([])}
+        ]
+    else:
+        for msg in datum_dict["message_log"]:
+            if "token_ids" not in msg:
+                msg["token_ids"] = torch.tensor([])
+            elif not isinstance(msg["token_ids"], torch.Tensor):
+                msg["token_ids"] = torch.tensor(msg["token_ids"])
+    return cast(DatumSpec, datum_dict)
+
+
 # Processor registry. Key is the processor name, value is the processor function.
 # Note: We cast the literal dict to Dict[str, TaskDataProcessFnCallable] because
 # type checkers see each concrete function's signature as a distinct callable type.
@@ -554,6 +574,7 @@ def multichoice_qa_processor(
         "multichoice_qa_processor": multichoice_qa_processor,
         "sft_processor": sft_processor,
         "vlm_hf_data_processor": vlm_hf_data_processor,
+        "nemo_gym_data_processor": nemo_gym_data_processor,
     },
 )
 

nemo_rl/environments/utils.py

@@ -46,6 +46,9 @@ class EnvRegistryEntry(TypedDict, total=False):
     "vlm": {
         "actor_class_fqn": "nemo_rl.environments.vlm_environment.VLMEnvironment",
     },
+    "nemo_gym": {
+        "actor_class_fqn": "nemo_rl.environments.nemo_gym.NemoGym",
+    },
 }