feat: Necessary changes for Gym GRPO tutorial

3rdparty/Gym-workspace/setup.py

@@ -42,6 +42,7 @@
     "yappi",
     "ray[default]",
     "psutil",
+    "datasets",
 ]
 
 if src_dir.exists():

3rdparty/Megatron-Bridge-workspace/setup.py

@@ -26,6 +26,7 @@
 bridge_package_name = "megatron.bridge"
 
 CACHED_DEPENDENCIES = [
+    "transformers>=4.57.1",
     "datasets",
     "omegaconf>=2.3.0",
     "tensorboard>=2.19.0",
@@ -37,12 +38,15 @@
     "pyyaml>=6.0.2",
     "tqdm>=4.67.1",
     "hydra-core>1.3,<=1.3.2",
-    "megatron-core[dev,mlm]>=0.15.0a0,<0.16.0",
+    "megatron-core[dev,mlm]>=0.15.0a0,<0.17.0",
     "qwen-vl-utils",
     "transformer-engine[pytorch]>=2.9.0a0,<2.10.0",
     "mamba-ssm",
     "nvidia-resiliency-ext",
     "causal-conv1d",
+    "flash-linear-attention",
+    "timm",
+    "open-clip-torch>=3.2.0",
 ]
 
 # If the bridge source exists, compare cached dependencies with the submodule's pyproject

3rdparty/Megatron-LM-workspace/setup.py

@@ -15,6 +15,8 @@
 
 import os
 import subprocess
+import sys
+import tomllib
 
 import setuptools
 from setuptools import Extension
@@ -35,9 +37,119 @@
 megatron_core_package_name = "megatron.core"
 
 # Path for the C++ extension's source file, relative to setup.py
-# This path is taken from your original setup.py
 megatron_core_cpp_extension_source_file = "megatron/core/datasets/helpers.cpp"
 
+# Cached dependencies: default + dev from pyproject.toml
+# VCS dependencies use full "pkg @ git+URL@rev" format matching pyproject.toml [tool.uv.sources]
+CACHED_DEPENDENCIES = [
+    # Default dependencies from pyproject.toml
+    "torch",
+    "numpy<2.0.0",
+    "packaging>=24.2",
+    # Dev dependencies from pyproject.toml
+    "nvidia-modelopt[torch]>=0.33.0a0,<0.34.0; sys_platform != 'darwin'",
+    "transformer-engine[pytorch]>=2.9.0a0,<2.10.0",
+    "nvidia-resiliency-ext>=0.4.0a0,<0.5.0",
+    "tqdm",
+    "einops~=0.8",
+    "tensorstore~=0.1,!=0.1.46,!=0.1.72",
+    "nvtx~=0.2",
+    "multi-storage-client~=0.27",
+    "opentelemetry-api~=1.33.1",
+    "setuptools<80.0.0",
+    "mamba-ssm~=2.2",
+    "causal-conv1d~=1.5",
+    "nv-grouped-gemm~=1.1",
+    "megatron-energon[av_decode]~=6.0",
+    "av<16.0.0",
+    "flashinfer-python",
+    "wget",
+    "onnxscript",
+    "flash-linear-attention~=0.3.2",
+    # VCS dependency - must match pyproject.toml [tool.uv.sources]
+    "emerging_optimizers @ git+https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git@v0.1.0",
+]
+
+
+def build_vcs_dependency(pkg_name: str, source_info: dict) -> str:
+    """Build a PEP 440 VCS dependency string from pyproject.toml [tool.uv.sources] entry."""
+    git_url = source_info.get("git")
+    rev = source_info.get("rev")
+    if not git_url:
+        raise ValueError(f"No git URL found for VCS dependency: {pkg_name}")
+    if not rev:
+        raise ValueError(f"No rev/commit found for VCS dependency: {pkg_name}")
+    return f"{pkg_name} @ git+{git_url}@{rev}"
+
+
+# Read pyproject.toml to validate dependencies
+pyproject_path = os.path.join("Megatron-LM", "pyproject.toml")
+
+if os.path.exists(megatron_core_python_package_source_dir):
+    if not os.path.exists(pyproject_path):
+        raise FileNotFoundError(
+            f"[megatron-core][setup] {pyproject_path} not found; skipping dependency consistency check."
+        )
+
+    with open(pyproject_path, "rb") as f:
+        data = tomllib.load(f)
+
+    # Extract [tool.uv.sources] for VCS dependencies
+    uv_sources = data.get("tool", {}).get("uv", {}).get("sources", {})
+
+    # Combine default dependencies + dev optional-dependencies
+    project = data["project"]
+    default_deps = project.get("dependencies", [])
+    optional_deps = project.get("optional-dependencies", {})
+    dev_deps = optional_deps.get("dev", [])
+
+    submodule_deps = set(str(d).strip() for d in default_deps + dev_deps)
+
+    # Build expected dependencies, converting any in [tool.uv.sources] to full VCS strings
+    submodule_deps_with_vcs = set()
+    for dep in submodule_deps:
+        if dep in uv_sources:
+            # Replace with full VCS string constructed from [tool.uv.sources]
+            vcs_dep = build_vcs_dependency(dep, uv_sources[dep])
+            submodule_deps_with_vcs.add(vcs_dep)
+        else:
+            submodule_deps_with_vcs.add(dep)
+
+    cached_deps_set = set(CACHED_DEPENDENCIES)
+
+    missing_in_cached = submodule_deps_with_vcs - cached_deps_set
+    extra_in_cached = cached_deps_set - submodule_deps_with_vcs
+
+    if missing_in_cached or extra_in_cached:
+        print(
+            "[megatron-core][setup] Dependency mismatch between Megatron-LM-workspace/Megatron-LM/pyproject.toml vs Megatron-LM-workspace/setup.py::CACHED_DEPENDENCIES.",
+            file=sys.stderr,
+        )
+        if missing_in_cached:
+            print(
+                "  - Present in Megatron-LM/pyproject.toml (default+dev) but missing from CACHED_DEPENDENCIES:",
+                file=sys.stderr,
+            )
+            for dep in sorted(missing_in_cached):
+                print(f"    * {dep}", file=sys.stderr)
+        if extra_in_cached:
+            print(
+                "  - Present in CACHED_DEPENDENCIES but not in Megatron-LM/pyproject.toml (default+dev):",
+                file=sys.stderr,
+            )
+            for dep in sorted(extra_in_cached):
+                print(f"    * {dep}", file=sys.stderr)
+        print(
+            "  Please update CACHED_DEPENDENCIES or the submodule pyproject to keep them in sync.",
+            file=sys.stderr,
+        )
+        sys.exit(1)
+    else:
+        print(
+            "[megatron-core][setup] Dependency sets are consistent with the submodule pyproject (default+dev).",
+            file=sys.stderr,
+        )
+
 # Check if the main directory for the megatron.core Python package exists
 if os.path.exists(megatron_core_python_package_source_dir):
     # Add Python package 'megatron.core'
@@ -74,25 +186,5 @@
     ext_modules=final_ext_modules,
     # Add in any packaged data.
     include_package_data=True,
-    install_requires=[
-        # From requirements/pytorch_25.03/requirements.txt
-        "einops",
-        "flask-restful",
-        "nltk",
-        "pytest",
-        "pytest-cov",
-        "pytest_mock",
-        "pytest-random-order",
-        "sentencepiece",
-        "tiktoken",
-        "wrapt",
-        "zarr",
-        "wandb",
-        "tensorstore!=0.1.46,!=0.1.72",
-        "torch",
-        "nvidia-modelopt[torch]>=0.23.2; sys_platform != 'darwin'",
-        # From megatron/core/requirements.txt
-        "torch",  # Repeated with ^ just to make it easy to map back to the original requirements.txt
-        "packaging",
-    ],
+    install_requires=CACHED_DEPENDENCIES,
 )

docs/design-docs/fsdp2-parallel-plan.md

@@ -20,11 +20,11 @@ The Hugging Face tensor parallel plan is the default. It's available for most mo
 
 ## Custom Parallel Plan Example
 
-A custom parallel plan should be defined in a separate file, such as the example provided in `examples/custom_parallel.py`.
+A custom parallel plan should be defined in a separate file, such as the example provided in `examples/custom_parallel/custom_parallel.py`.
 
 To implement the custom parallel plan, either update the value of `custom_parallel_plan` in the `yaml` file directly, or pass the override via the command line. For example:
 
 ```bash
 uv run examples/run_grpo_math.py \
-    policy.dtensor_cfg.custom_parallel_plan=examples.custom_parallel.custom_parallel_plan
+    policy.dtensor_cfg.custom_parallel_plan=examples.custom_parallel.custom_parallel.custom_parallel_plan
 ```

docs/guides/environments.md

@@ -1,6 +1,6 @@
 # Environments for GRPO Training
 
-GRPO supports several examples of environments for different tasks. Each environment provides a standardized interface for reward computation and evaluation.
+GRPO includes multiple environments, each offering a standard interface for reward computation and evaluation.
 
 ## Math Environment
 
@@ -40,9 +40,51 @@ code_env = CodeEnvironment.remote(env_config)
 
 ### Configuration
 - `num_workers`: Number of parallel workers for code execution
-- `terminate_on_evaluation`: Whether to terminate after code execution (True for single-turn, False for multi-turn)
+- `terminate_on_evaluation`: Whether to terminate after code execution (True for single-turn, False for multi-turn).
 
-We’re tracking an end-to-end example of this environment in [#858](https://github.com/NVIDIA-NeMo/RL/issues/858). Add a 👍 to show your interest.
+We are tracking an end-to-end example of this environment in [#858](https://github.com/NVIDIA-NeMo/RL/issues/858). Add a 👍 to show your interest.
+
+## Code Jaccard Environment
+
+The Code Jaccard Environment evaluates code (or text) responses by measuring Jaccard-based similarity against ground-truth answers. This is a lightweight, text-similarity reward useful when an execution sandbox is unnecessary or unavailable.
+
+### How It Works
+- Extracts the assistant’s response text from each conversation.
+- Computes a Jaccard similarity score between the response and ground truth:
+  - Tokenizes both texts by whitespace, computes intersection/union, then applies a length ratio penalty.
+  - Scores are in [0, 1]. Observations label responses as “aligned/misaligned” using a 0.5 threshold.
+- Returns:
+  - observations: Environment feedback strings.
+  - rewards: Tensor of similarity scores.
+  - terminateds: All ones (single-step episodes).
+  - answers: The response text when requested (optional).
+
+### Usage
+```python
+from nemo_rl.environments.code_jaccard_environment import CodeJaccardEnvironment
+
+env_config = {
+    "num_workers": 2,
+    # Optional default stop strings (unused in scoring but available for consistency)
+    "stop_strings": None,
+}
+
+code_jaccard_env = CodeJaccardEnvironment.remote(env_config)
+```
+
+### Configuration
+- `num_workers` (int): Number of parallel verification workers.
+- `stop_strings` (list[str] | None): Optional default stop strings (propagated downstream; not required for scoring).
+
+### Sample GRPO Config
+```yaml
+env:
+  code_jaccard:
+    num_workers: 2
+    stop_strings: null
+data:
+  env_name: code_jaccard
+```
 
 ## Reward Model Environment
 
@@ -72,9 +114,9 @@ reward_env = RewardModelEnvironment.remote(env_config)
 
 In GRPO training, resources are allocated across three main components:
 
-- **Policy Actor**: The trained model
+- **Policy Actor**: The trained model.
 - **Generation Actor**: Used for generating responses during rollouts (can be colocated with policy or on separate nodes/GPUs).
-- **Reward Model Environment Actor**: Evaluates generated responses and computes rewards
+- **Reward Model Environment Actor**: Evaluates generated responses and computes rewards.
 
 The resource allocation logic works as follows:
 
@@ -86,10 +128,10 @@ The resource allocation logic works as follows:
     2. Policy and generation non-colocated: 8 GPUs total = 2 for policy + 2 for generation + 4 for reward model
 
 #### Multi-Node Setup (`num_nodes > 1`)
-- Policy training, generation, and reward model environment can be distributed across different nodes
-- Reward model gets dedicated resources as specified in `env.reward_model.resources`
-- Generation gets dedicated resources as specified in `policy.generation.colocated.resources`
-- Remaining nodes are allocated to policy training
+- Policy training, generation, and reward model environment can be distributed across different nodes.
+- Reward model gets dedicated resources as specified in `env.reward_model.resources`.
+- Generation gets dedicated resources as specified in `policy.generation.colocated.resources`.
+- Remaining nodes are allocated to policy training.
 
 In the future, the resource control part will be refactored to enable fine-grained resource configuration for each actor. For detailed resource management and optimization strategies, see [#1100](https://github.com/NVIDIA-NeMo/RL/issues/1100).
 
@@ -99,4 +141,86 @@ See [examples/run_grpo_rm.py](../../examples/run_grpo_rm.py) for a complete exam
 
 ### Configuration Examples
 
-See [examples/configs/grpo_rm_1B.yaml](../../examples/configs/grpo_rm_1B.yaml) for a complete configuration example.
+See [examples/configs/grpo_rm_1B.yaml](../../examples/configs/grpo_rm_1B.yaml) for a complete configuration example.
+
+
+## Registering Custom Environments
+
+NeMo RL provides a flexible environment registration mechanism that allows you to add custom environments without modifying the source code.
+
+### Using the `register_env` Interface
+
+You can use the `register_env` function to dynamically register new environments without modifying NeMo RL's internal code.
+
+**Function Signature**
+
+```python
+from nemo_rl.environments.utils import register_env
+
+register_env(env_name: str, actor_class_fqn: str) -> None
+```
+
+**Parameters:**
+
+- `env_name`: Unique identifier name for the environment (string)
+- `actor_class_fqn`: Fully Qualified Name of the environment Actor class, in the format `'module.path.ClassName'`
+
+### Example: Registering a Custom Environment
+
+Suppose you've created a custom reinforcement learning environment for code generation tasks:
+
+**1. Create Your Custom Environment Actor Class**
+
+```python
+# File: my_custom_envs/code_gen_env.py
+import ray
+from nemo_rl.environments.interfaces import EnvironmentInterface
+
+@ray.remote
+class CodeGenEnvironmentActor(EnvironmentInterface):
+    """Custom code generation environment."""
+
+    def __init__(self, config):
+        self.config = config
+        # Initialize your environment
+
+    async def reset(self):
+        # Reset environment logic
+        return initial_state
+
+    async def step(self, action):
+        # Execute action, return reward, etc.
+        return observation, reward, done, info
+
+    # Implement other required interface methods...
+```
+
+**2. Register the Environment in Your Training Script**
+
+```python
+# File: train.py
+from nemo_rl.environments.utils import register_env
+
+# Register your custom environment
+register_env(
+    env_name="code_gen",
+    actor_class_fqn="my_custom_envs.code_gen_env.CodeGenEnvironmentActor"
+)
+
+# Now you can use "code_gen" in your config
+# Training code...
+```
+
+**3. Use the Registered Environment in Your Config**
+
+```yaml
+# config.yaml
+env:
+  code_gen:
+    num_workers: 2
+    max_code_length: 512
+    test_cases_per_problem: 5
+
+data:
+  env_name: code_gen  # Use your registered environment name
+```