From 55126e3343b282b3d400707ce02f71a3d80476c9 Mon Sep 17 00:00:00 2001
From: Yuki Huang <yukih@nvidia.com>
Date: Sat, 18 Jul 2026 03:00:07 -0700
Subject: [PATCH 01/13] feat(sc): wire NeMo-Gym integration into setup

Signed-off-by: Yuki Huang <yukih@nvidia.com>
---
 .../single_controller_utils/setup.py          | 44 +++++++++++++++----
 1 file changed, 35 insertions(+), 9 deletions(-)

diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py
index e9db9d6690..4ae9058a8a 100644
--- a/nemo_rl/algorithms/single_controller_utils/setup.py
+++ b/nemo_rl/algorithms/single_controller_utils/setup.py
@@ -49,11 +49,13 @@
 from nemo_rl.data_plane import DataPlaneClient, build_data_plane_client
 from nemo_rl.distributed.virtual_cluster import RayVirtualCluster
 from nemo_rl.environments.interfaces import EnvironmentInterface
+from nemo_rl.environments.nemo_gym import spinup_nemo_gym_actor
 from nemo_rl.experience.rollout_manager import RolloutManager
 from nemo_rl.models.generation.sglang.config import SGLangConfig
 from nemo_rl.models.generation.sglang.sglang_generation import SGLangGeneration
 from nemo_rl.models.generation.vllm import VllmGeneration
 from nemo_rl.models.generation.vllm.config import VllmConfig
+from nemo_rl.models.megatron.router_replay import router_replay_enabled
 from nemo_rl.models.policy.tq_policy import TQPolicy
 from nemo_rl.weight_sync import WeightSynchronizer, create_weight_synchronizer
 
@@ -318,16 +320,22 @@ def setup_single_controller(
     # Setup Dataset & Environments
     # ==========================
     # TODO: add validate dataset wiring.
-    if _should_use_nemo_gym(cast(GrpoMasterConfig, master_config)):
-        raise NotImplementedError(
-            "NeMo-Gym integration for SingleController is not supported yet; "
-            "it will land in https://github.com/NVIDIA-NeMo/RL/pull/3267."
+    use_nemo_gym = _should_use_nemo_gym(cast(GrpoMasterConfig, master_config))
+    if use_nemo_gym:
+        # NeMo-Gym creates the env actor outside setup_response_data; we wire
+        # it in after generation is up (it needs the OpenAI server URLs).
+        response_data = setup_response_data(
+            tokenizer, data_config, env_configs=None
         )
-    response_data = setup_response_data(
-        tokenizer, data_config, env_configs=master_config.env
-    )
-    assert len(response_data) == 4
-    dataset, _val_dataset, env_handles, _val_env_handles = response_data
+        assert len(response_data) == 2
+        dataset, _val_dataset = response_data
+        env_handles: dict[str, EnvironmentInterface] = {}
+    else:
+        response_data = setup_response_data(
+            tokenizer, data_config, env_configs=master_config.env
+        )
+        assert len(response_data) == 4
+        dataset, _val_dataset, env_handles, _val_env_handles = response_data
     dataloader = StatefulDataLoader(
         dataset,
         batch_size=grpo_config["num_prompts_per_step"],
@@ -363,6 +371,23 @@ def setup_single_controller(
             generation = gen_future.result()
             policy = policy_future.result()
 
+    # ==========================
+    # NeMo-Gym actor (after generation is up so OpenAI URLs are available)
+    # ==========================
+    if use_nemo_gym:
+        if generation_config["backend"] != "vllm":
+            raise NotImplementedError(
+                "SC NeMo-Gym integration currently supports the vllm backend "
+                f"only; got {generation_config['backend']!r}"
+            )
+        enable_router_replay = router_replay_enabled(master_config.policy)
+        env_handles["nemo_gym"] = spinup_nemo_gym_actor(
+            env_configs=master_config.env,
+            base_urls=generation.dp_openai_server_base_urls,
+            model_name=generation_config["model_name"],
+            enable_router_replay=enable_router_replay,
+        )
+
     # ==========================
     # Setup Data Plane Client & Weight Sync
     # ==========================
@@ -403,6 +428,7 @@ def setup_single_controller(
         max_rollout_turns=grpo_config.get("max_rollout_turns"),
         policy_generation=generation,
         generation_config=generation_config,
+        use_nemo_gym=use_nemo_gym,
         tq_buffer=tq_buffer,
     )
 

From 46dda58d137854d42c43665ab2cb7690fb3232e8 Mon Sep 17 00:00:00 2001
From: Yuki Huang <yukih@nvidia.com>
Date: Wed, 15 Jul 2026 09:04:53 -0700
Subject: [PATCH 02/13] feat(sc): NeMo-Gym path

Signed-off-by: Yuki Huang <yukih@nvidia.com>
---
 .../nemo_gym/run_distillation_nemo_gym.py     |   4 +-
 nemo_rl/environments/nemo_gym.py              |  68 ++++++++++-
 .../grpo_async_gym_single_controller.sh       | 112 ++++++++++++++++++
 3 files changed, 180 insertions(+), 4 deletions(-)
 create mode 100755 tests/functional/grpo_async_gym_single_controller.sh

diff --git a/examples/nemo_gym/run_distillation_nemo_gym.py b/examples/nemo_gym/run_distillation_nemo_gym.py
index 56fe23eb6b..ca2fad0d95 100644
--- a/examples/nemo_gym/run_distillation_nemo_gym.py
+++ b/examples/nemo_gym/run_distillation_nemo_gym.py
@@ -32,9 +32,7 @@
 from nemo_rl.algorithms.utils import get_tokenizer
 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 (
-    setup_nemo_gym_config,
-)
+from nemo_rl.environments.nemo_gym import setup_nemo_gym_config
 from nemo_rl.models.generation import configure_generation_config
 from nemo_rl.utils.config import (
     load_config,
diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py
index 513cb00b18..107ca98bd9 100644
--- a/nemo_rl/environments/nemo_gym.py
+++ b/nemo_rl/environments/nemo_gym.py
@@ -18,12 +18,14 @@
 from collections import Counter
 from collections.abc import AsyncGenerator
 from pathlib import Path
-from typing import Any, Dict, List, NotRequired, TypedDict
+from typing import Any, Dict, List, NotRequired, Optional, TypedDict
 
 import ray
 import torch
+from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
 from transformers import PreTrainedTokenizerBase
 
+from nemo_rl.distributed.ray_actor_environment_registry import get_actor_python_env
 from nemo_rl.distributed.virtual_cluster import (
     DEFAULT_GYM_PORT_RANGE_HIGH,
     DEFAULT_GYM_PORT_RANGE_LOW,
@@ -32,6 +34,7 @@
 )
 from nemo_rl.environments.interfaces import EnvironmentInterface
 from nemo_rl.utils.timer import Timer
+from nemo_rl.utils.venvs import create_local_venv_on_each_node
 
 # Kept local (not imported from models.generation) so the gym actor stays free of
 # generation-module imports. Must cover every name resolve_routed_experts_dtype
@@ -600,3 +603,66 @@ def setup_nemo_gym_config(config, tokenizer) -> None:
     # Stop strings or token ids are not supported
     generation_config["stop_strings"] = None
     generation_config["stop_token_ids"] = None
+
+
+def spinup_nemo_gym_actor(
+    env_configs: dict[str, Any],
+    base_urls: list[Optional[str]],
+    model_name: str,
+    enable_router_replay: bool,
+) -> Any:
+    """Spin up the NeMo-Gym actor against the given generation server URLs.
+
+    When ``env_configs["nemo_gym"]["num_gpu_nodes"] > 0``, the actor is
+    scheduled with soft NodeAffinity to the current Ray node so its colocated
+    GPU resources land where the caller expects.
+    """
+    nemo_gym_dict = dict(env_configs["nemo_gym"])
+
+    # NeMo-RL-side detection knobs are top-level NemoGymConfig fields
+    # (where the detector reads them), not part of Gym's global config.
+    invalid_tool_call_patterns = nemo_gym_dict.pop("invalid_tool_call_patterns", None)
+    thinking_tags = nemo_gym_dict.pop("thinking_tags", None)
+
+    # Pass prebuilt cache + venv dirs through the global config so the gym reuses
+    # image-baked venvs instead of rebuilding them.
+    uv_cache_dir = get_nemo_gym_uv_cache_dir()
+    if uv_cache_dir is not None:
+        nemo_gym_dict.setdefault("uv_cache_dir", uv_cache_dir)
+    uv_venv_dir = get_nemo_gym_venv_dir()
+    if uv_venv_dir is not None:
+        nemo_gym_dict.setdefault("uv_venv_dir", uv_venv_dir)
+
+    nemo_gym_cfg = NemoGymConfig(
+        model_name=model_name,
+        base_urls=base_urls,
+        invalid_tool_call_patterns=invalid_tool_call_patterns,
+        thinking_tags=thinking_tags,
+        require_routed_experts=enable_router_replay,
+        initial_global_config_dict=nemo_gym_dict,
+    )
+
+    nemo_gym_py_exec = get_actor_python_env("nemo_rl.environments.nemo_gym.NemoGym")
+    if nemo_gym_py_exec.startswith("uv"):
+        nemo_gym_py_exec = create_local_venv_on_each_node(
+            nemo_gym_py_exec, "nemo_rl.environments.nemo_gym.NemoGym"
+        )
+
+    nemo_gym_opts: dict[str, Any] = {}
+    if nemo_gym_dict.get("num_gpu_nodes", 0):
+        nemo_gym_opts["scheduling_strategy"] = NodeAffinitySchedulingStrategy(
+            node_id=ray.get_runtime_context().get_node_id(),
+            soft=True,
+        )
+    nemo_gym_opts["runtime_env"] = {
+        "py_executable": nemo_gym_py_exec,
+        "env_vars": {
+            **os.environ,
+            "VIRTUAL_ENV": nemo_gym_py_exec,
+            "UV_PROJECT_ENVIRONMENT": nemo_gym_py_exec,
+        },
+    }
+
+    actor = NemoGym.options(**nemo_gym_opts).remote(nemo_gym_cfg)
+    ray.get(actor._spinup.remote())
+    return actor
diff --git a/tests/functional/grpo_async_gym_single_controller.sh b/tests/functional/grpo_async_gym_single_controller.sh
new file mode 100755
index 0000000000..f4f283bfff
--- /dev/null
+++ b/tests/functional/grpo_async_gym_single_controller.sh
@@ -0,0 +1,112 @@
+#!/bin/bash
+# SingleController + NeMo-Gym e2e smoke. Mirrors grpo_async_gym.sh but
+# routes everything through the SC path (TransferQueue data plane +
+# SingleControllerActor) instead of async_grpo_train.
+
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd)
+PROJECT_ROOT=$(realpath $SCRIPT_DIR/../..)
+# Mark the current repo as safe, since wandb fetches metadata about the repo
+git config --global --add safe.directory $PROJECT_ROOT
+
+set -eou pipefail
+
+EXP_NAME=$(basename $0 .sh)
+EXP_DIR=$SCRIPT_DIR/$EXP_NAME
+LOG_DIR=$EXP_DIR/logs
+JSON_METRICS=$EXP_DIR/metrics.json
+RUN_LOG=$EXP_DIR/run.log
+CHECKPOINT_DIR=$EXP_DIR/checkpoints
+DATA_DIR=$EXP_DIR/data
+export PYTHONPATH=${PROJECT_ROOT}:${PYTHONPATH:-}
+
+rm -rf $EXP_DIR $LOG_DIR
+mkdir -p $EXP_DIR $LOG_DIR $CHECKPOINT_DIR $DATA_DIR
+
+# clean up checkpoint directory on exit
+trap "rm -rf $CHECKPOINT_DIR" EXIT
+
+cd $PROJECT_ROOT
+
+# Follow nemo-gym instructions here to get this data:
+# https://docs.nvidia.com/nemo/gym/0.1.0/tutorials/nemo-rl-grpo/setup.html#training-nemo-rl-grpo-setup
+cd 3rdparty/Gym-workspace/Gym
+
+# We need HF_TOKEN to download the data from huggingface
+if [[ ! -f env.yaml ]]; then
+    if [[ -z "${HF_TOKEN:-}" ]]; then
+        echo "[ERROR] HF_TOKEN is not set"
+        exit 1
+    fi
+    echo "hf_token: $HF_TOKEN" >> env.yaml
+fi
+
+uv run ng_prepare_data "+config_paths=[resources_servers/workplace_assistant/configs/workplace_assistant.yaml]" \
+    +output_dirpath=data/workplace_assistant \
+    +mode=train_preparation \
+    +should_download=true \
+    +data_source=huggingface
+cd -
+
+# This trimming of the workplace assistant dataset is necessary b/c with all the tools the first prompt is >4000 tokens
+# which will cause vllm to return nothing on the first prompt and crash RL. Since we want to keep this test short to
+# smoke test, we trim all but the first tool
+TRAIN_PATH=$DATA_DIR/workplace_assistant_train.jsonl
+VALIDATION_PATH=$DATA_DIR/workplace_assistant_validation.jsonl
+jq -c '.responses_create_params.tools |= (.[0:1])' 3rdparty/Gym-workspace/Gym/data/workplace_assistant/train.jsonl > $TRAIN_PATH
+jq -c '.responses_create_params.tools |= (.[0:1])' 3rdparty/Gym-workspace/Gym/data/workplace_assistant/validation.jsonl > $VALIDATION_PATH
+
+uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \
+    $PROJECT_ROOT/examples/run_grpo_single_controller.py \
+    --config $PROJECT_ROOT/examples/nemo_gym/grpo_qwen3_30ba3b_instruct.yaml \
+    policy.model_name=Qwen/Qwen3-0.6B \
+    policy.dtensor_cfg.enabled=false \
+    policy.megatron_cfg.enabled=true \
+    policy.megatron_cfg.tensor_model_parallel_size=1 \
+    policy.megatron_cfg.pipeline_model_parallel_size=1 \
+    policy.megatron_cfg.expert_model_parallel_size=1 \
+    policy.megatron_cfg.context_parallel_size=1 \
+    policy.megatron_cfg.sequence_parallel=false \
+    policy.generation.vllm_cfg.tensor_parallel_size=1 \
+    policy.generation.vllm_cfg.async_engine=true \
+    policy.max_total_sequence_length=512 \
+    policy.generation.colocated.enabled=false \
+    policy.generation.colocated.resources.num_nodes=1 \
+    policy.generation.colocated.resources.gpus_per_node=1 \
+    grpo.num_prompts_per_step=4 \
+    grpo.num_generations_per_prompt=2 \
+    grpo.max_num_steps=10 \
+    grpo.val_period=-1 \
+    policy.train_global_batch_size=8 \
+    policy.train_micro_batch_size=1 \
+    cluster.gpus_per_node=2 \
+    loss_fn.reference_policy_kl_penalty=0.01 \
+    loss_fn.use_importance_sampling_correction=true \
+    logger.tensorboard_enabled=true \
+    logger.log_dir=$LOG_DIR \
+    logger.wandb_enabled=false \
+    logger.monitor_gpus=true \
+    checkpointing.enabled=false \
+    data.train.data_path=$TRAIN_PATH \
+    data.validation.data_path=$VALIDATION_PATH \
+    ++data_plane.enabled=true \
+    ++data_plane.impl=transfer_queue \
+    ++data_plane.backend=simple \
+    ++data_plane.storage_capacity=1000000 \
+    ++data_plane.num_storage_units=2 \
+    ++data_plane.claim_meta_poll_interval_s=0.5 \
+    ++data_plane.global_segment_size=549755813888 \
+    ++data_plane.local_buffer_size=68719476736 \
+    ++async_rl.batch_selection_strategy=strict_on_policy \
+    ++async_rl.max_weight_staleness_versions=0 \
+    ++async_rl.min_prompt_groups_per_batch=4 \
+    ++async_rl.max_inflight_prompts=4 \
+    ++async_rl.max_buffered_rollouts=4 \
+    $@ \
+    2>&1 | tee $RUN_LOG
+
+uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS
+
+# Observed to be between 0.8-1.3
+uv run tests/check_metrics.py $JSON_METRICS \
+    'median(data["train/gen_kl_error"]) < 1.3' \
+    'max(data["train/reward"]) > 0'

From 56628ae26e7fa600180b29ef35c3c8432525f2f4 Mon Sep 17 00:00:00 2001
From: Yuki Huang <yukih@nvidia.com>
Date: Sat, 18 Jul 2026 08:23:57 -0700
Subject: [PATCH 03/13] Revert "refactor(sc): drop dead nemo-gym setup call
 from entrypoint"

This reverts commit 0eeab0802a27eac1e781f927f76cf1757f25babe.

Signed-off-by: Yuki Huang <yukih@nvidia.com>
---
 examples/run_grpo_single_controller.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/examples/run_grpo_single_controller.py b/examples/run_grpo_single_controller.py
index 3b776ff773..544124c9fa 100644
--- a/examples/run_grpo_single_controller.py
+++ b/examples/run_grpo_single_controller.py
@@ -34,6 +34,7 @@
 )
 from nemo_rl.algorithms.utils import get_tokenizer
 from nemo_rl.distributed.virtual_cluster import init_ray
+from nemo_rl.environments.nemo_gym import setup_nemo_gym_config
 from nemo_rl.models.generation import configure_generation_config
 from nemo_rl.utils.config import (
     load_config,
@@ -117,6 +118,10 @@ def main() -> None:
         trains_mtp=trains_mtp,
     )
 
+    # NeMo-Gym specific config setup.
+    if bool(config.env.get("should_use_nemo_gym")):
+        setup_nemo_gym_config(config, tokenizer)
+
     actor_args = setup_single_controller(config, tokenizer)
 
     print("🚀 Launching SingleControllerActor")

From 2881dc31e2e149efe55d94c91c4fb713b5ddb85c Mon Sep 17 00:00:00 2001
From: Yuki Huang <yukih@nvidia.com>
Date: Sat, 18 Jul 2026 08:24:51 -0700
Subject: [PATCH 04/13] remove test_nemo_gym_not_supported

Signed-off-by: Yuki Huang <yukih@nvidia.com>
---
 .../single_controller/test_single_controller_setup.py    | 9 ---------
 1 file changed, 9 deletions(-)

diff --git a/tests/unit/single_controller/test_single_controller_setup.py b/tests/unit/single_controller/test_single_controller_setup.py
index 3d6ba75aa5..52fd2e65a1 100644
--- a/tests/unit/single_controller/test_single_controller_setup.py
+++ b/tests/unit/single_controller/test_single_controller_setup.py
@@ -282,15 +282,6 @@ def test_returns_actor_args(self, patched_factories):
         assert actor_args.partition_id == "rollout_data"
         assert actor_args.tq_buffer._partition_id == "rollout_data"
 
-    def test_nemo_gym_not_supported(self):
-        """SC path trips the nemo-gym guard until PR #3267 lands."""
-        mc = _make_master_config()
-        with (
-            patch.object(sc_setup_mod, "_should_use_nemo_gym", return_value=True),
-            pytest.raises(NotImplementedError, match="NeMo-Gym"),
-        ):
-            setup_single_controller(mc, MagicMock(pad_token_id=0))
-
     def test_env_handles_sourced_from_setup_response_data(self, patched_factories):
         """setup_response_data receives master_config.env and supplies env handles."""
         math_env_cfg = {"some": "value"}

From 24be481e178f7828c48779fe8bbf3922196ca7f3 Mon Sep 17 00:00:00 2001
From: Yuki Huang <yukih@nvidia.com>
Date: Sat, 18 Jul 2026 08:55:47 -0700
Subject: [PATCH 05/13] fix(sc): correct async_rl overrides in nemo-gym smoke
 test

Signed-off-by: Yuki Huang <yukih@nvidia.com>
---
 tests/functional/grpo_async_gym_single_controller.sh | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/functional/grpo_async_gym_single_controller.sh b/tests/functional/grpo_async_gym_single_controller.sh
index f4f283bfff..7e3cf488d7 100755
--- a/tests/functional/grpo_async_gym_single_controller.sh
+++ b/tests/functional/grpo_async_gym_single_controller.sh
@@ -76,6 +76,7 @@ uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJE
     grpo.num_generations_per_prompt=2 \
     grpo.max_num_steps=10 \
     grpo.val_period=-1 \
+    grpo.val_at_start=false \
     policy.train_global_batch_size=8 \
     policy.train_micro_batch_size=1 \
     cluster.gpus_per_node=2 \
@@ -96,9 +97,8 @@ uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJE
     ++data_plane.claim_meta_poll_interval_s=0.5 \
     ++data_plane.global_segment_size=549755813888 \
     ++data_plane.local_buffer_size=68719476736 \
-    ++async_rl.batch_selection_strategy=strict_on_policy \
     ++async_rl.max_weight_staleness_versions=0 \
-    ++async_rl.min_prompt_groups_per_batch=4 \
+    ++async_rl.min_groups_for_streaming_train=4 \
     ++async_rl.max_inflight_prompts=4 \
     ++async_rl.max_buffered_rollouts=4 \
     $@ \

From 2bfb939d2a73242d0f7680cdbd64ac5c9ecd38dc Mon Sep 17 00:00:00 2001
From: Yuki Huang <yukih@nvidia.com>
Date: Sat, 18 Jul 2026 09:03:31 -0700
Subject: [PATCH 06/13] fix(sc): configure vLLM for router_replay before
 generation build

Signed-off-by: Yuki Huang <yukih@nvidia.com>
---
 nemo_rl/algorithms/single_controller_utils/setup.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py
index 4ae9058a8a..7cac4d5df2 100644
--- a/nemo_rl/algorithms/single_controller_utils/setup.py
+++ b/nemo_rl/algorithms/single_controller_utils/setup.py
@@ -55,7 +55,10 @@
 from nemo_rl.models.generation.sglang.sglang_generation import SGLangGeneration
 from nemo_rl.models.generation.vllm import VllmGeneration
 from nemo_rl.models.generation.vllm.config import VllmConfig
-from nemo_rl.models.megatron.router_replay import router_replay_enabled
+from nemo_rl.models.megatron.router_replay import (
+    configure_vllm_for_router_replay,
+    router_replay_enabled,
+)
 from nemo_rl.models.policy.tq_policy import TQPolicy
 from nemo_rl.weight_sync import WeightSynchronizer, create_weight_synchronizer
 
@@ -169,6 +172,7 @@ def _build_generation(
         vllm_config.setdefault("vllm_kwargs", {})["hf_overrides"] = (
             master_config.policy.get("hf_config_overrides", {})
         )
+        configure_vllm_for_router_replay(master_config.policy)
         gen = VllmGeneration(cluster=inference_cluster, config=vllm_config)
     elif backend == "sglang":
         sglang_config = cast(SGLangConfig, generation_config)

From 082667f81639dd019983ac86f7ae58e77447bb97 Mon Sep 17 00:00:00 2001
From: Yuki Huang <yukih@nvidia.com>
Date: Sat, 18 Jul 2026 09:13:21 -0700
Subject: [PATCH 07/13] test(sc): cover nemo-gym env handle wiring and non-vllm
 backend guard

Signed-off-by: Yuki Huang <yukih@nvidia.com>
---
 .../test_single_controller_setup.py           | 42 +++++++++++++++++++
 1 file changed, 42 insertions(+)

diff --git a/tests/unit/single_controller/test_single_controller_setup.py b/tests/unit/single_controller/test_single_controller_setup.py
index 52fd2e65a1..c295757738 100644
--- a/tests/unit/single_controller/test_single_controller_setup.py
+++ b/tests/unit/single_controller/test_single_controller_setup.py
@@ -391,3 +391,45 @@ def test_megatron_train_iters_not_set_when_disabled(self, patched_factories):
         setup_single_controller(mc, MagicMock(pad_token_id=0))
 
         assert "train_iters" not in mc.policy.get("megatron_cfg", {})
+
+    def test_nemo_gym_wires_env_handle(self, patched_factories):
+        """When _should_use_nemo_gym is True the nemo-gym actor is spun up and stored."""
+        mc = _make_master_config(colocated=True, backend="vllm")
+        mc.policy["generation"]["model_name"] = "test-model"
+        mc.policy["generation"]["stop_strings"] = None
+        mc.policy["generation"]["stop_token_ids"] = None
+        mc.policy["generation"]["top_k"] = None
+        patched_factories["setup_response_data"].return_value = (
+            list(range(8)),
+            None,
+        )
+        fake_gym_actor = MagicMock(name="nemo_gym_actor")
+
+        with (
+            patch.object(sc_setup_mod, "_should_use_nemo_gym", return_value=True),
+            patch.object(
+                sc_setup_mod, "spinup_nemo_gym_actor", return_value=fake_gym_actor
+            ) as mock_spinup,
+            patch.object(sc_setup_mod, "router_replay_enabled", return_value=False),
+        ):
+            actor_args = setup_single_controller(mc, MagicMock(pad_token_id=0))
+
+        mock_spinup.assert_called_once()
+        assert actor_args.env_handles["nemo_gym"] is fake_gym_actor
+
+    @pytest.mark.parametrize("backend", ["sglang", "megatron"])
+    def test_nemo_gym_rejects_non_vllm_backend(self, patched_factories, backend):
+        """SC nemo-gym wiring only supports vLLM; every other backend must raise."""
+        mc = _make_master_config(colocated=True, backend=backend)
+        patched_factories["setup_response_data"].return_value = (
+            list(range(8)),
+            None,
+        )
+
+        with (
+            patch.object(sc_setup_mod, "_should_use_nemo_gym", return_value=True),
+            patch.object(sc_setup_mod, "spinup_nemo_gym_actor") as mock_spinup,
+            pytest.raises(NotImplementedError, match="vllm"),
+        ):
+            setup_single_controller(mc, MagicMock(pad_token_id=0))
+        mock_spinup.assert_not_called()

From a3e32cd5267e3223ad9ec1fc79b2f083610407df Mon Sep 17 00:00:00 2001
From: Yuki Huang <yukih@nvidia.com>
Date: Sat, 18 Jul 2026 21:44:16 -0700
Subject: [PATCH 08/13] fix rebase

Signed-off-by: Yuki Huang <yukih@nvidia.com>
---
 tests/functional/grpo_async_gym_single_controller.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/functional/grpo_async_gym_single_controller.sh b/tests/functional/grpo_async_gym_single_controller.sh
index 7e3cf488d7..f55a2b2316 100755
--- a/tests/functional/grpo_async_gym_single_controller.sh
+++ b/tests/functional/grpo_async_gym_single_controller.sh
@@ -97,7 +97,7 @@ uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJE
     ++data_plane.claim_meta_poll_interval_s=0.5 \
     ++data_plane.global_segment_size=549755813888 \
     ++data_plane.local_buffer_size=68719476736 \
-    ++async_rl.max_weight_staleness_versions=0 \
+    ++async_rl.max_staleness_versions=0 \
     ++async_rl.min_groups_for_streaming_train=4 \
     ++async_rl.max_inflight_prompts=4 \
     ++async_rl.max_buffered_rollouts=4 \

From 7cc931bb6852255e25b88cf5fd05505e76da7bd9 Mon Sep 17 00:00:00 2001
From: ruit <ruit@nvidia.com>
Date: Thu, 23 Jul 2026 22:50:28 -0700
Subject: [PATCH 09/13] fix(sc): configure sampler in NeMo-Gym smoke test

Signed-off-by: ruit <ruit@nvidia.com>
---
 nemo_rl/algorithms/single_controller_utils/setup.py  | 4 +---
 tests/functional/grpo_async_gym_single_controller.sh | 3 ++-
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py
index 7cac4d5df2..0ab895cb0c 100644
--- a/nemo_rl/algorithms/single_controller_utils/setup.py
+++ b/nemo_rl/algorithms/single_controller_utils/setup.py
@@ -328,9 +328,7 @@ def setup_single_controller(
     if use_nemo_gym:
         # NeMo-Gym creates the env actor outside setup_response_data; we wire
         # it in after generation is up (it needs the OpenAI server URLs).
-        response_data = setup_response_data(
-            tokenizer, data_config, env_configs=None
-        )
+        response_data = setup_response_data(tokenizer, data_config, env_configs=None)
         assert len(response_data) == 2
         dataset, _val_dataset = response_data
         env_handles: dict[str, EnvironmentInterface] = {}
diff --git a/tests/functional/grpo_async_gym_single_controller.sh b/tests/functional/grpo_async_gym_single_controller.sh
index f55a2b2316..0178961739 100755
--- a/tests/functional/grpo_async_gym_single_controller.sh
+++ b/tests/functional/grpo_async_gym_single_controller.sh
@@ -97,7 +97,8 @@ uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJE
     ++data_plane.claim_meta_poll_interval_s=0.5 \
     ++data_plane.global_segment_size=549755813888 \
     ++data_plane.local_buffer_size=68719476736 \
-    ++async_rl.max_staleness_versions=0 \
+    ++async_rl.sampler.name=in_order \
+    ++async_rl.sampler.max_lookahead_versions=0 \
     ++async_rl.min_groups_for_streaming_train=4 \
     ++async_rl.max_inflight_prompts=4 \
     ++async_rl.max_buffered_rollouts=4 \

From 0063fcabe0b9f8f1f33a0be8f56227fed894987e Mon Sep 17 00:00:00 2001
From: ruit <ruit@nvidia.com>
Date: Fri, 24 Jul 2026 09:12:52 -0700
Subject: [PATCH 10/13] fix(sc): address NeMo-Gym review feedback

Signed-off-by: ruit <ruit@nvidia.com>
---
 nemo_rl/algorithms/single_controller_utils/setup.py | 10 +++++-----
 nemo_rl/environments/nemo_gym.py                    | 13 +++++++++++++
 .../L1_Functional_Tests_SingleController.sh         |  1 +
 .../test_single_controller_setup.py                 |  9 ++++++++-
 4 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py
index 0ab895cb0c..bbac9a28c2 100644
--- a/nemo_rl/algorithms/single_controller_utils/setup.py
+++ b/nemo_rl/algorithms/single_controller_utils/setup.py
@@ -325,6 +325,11 @@ def setup_single_controller(
     # ==========================
     # TODO: add validate dataset wiring.
     use_nemo_gym = _should_use_nemo_gym(cast(GrpoMasterConfig, master_config))
+    if use_nemo_gym and generation_config["backend"] != "vllm":
+        raise NotImplementedError(
+            "SC NeMo-Gym integration currently supports the vllm backend "
+            f"only; got {generation_config['backend']!r}"
+        )
     if use_nemo_gym:
         # NeMo-Gym creates the env actor outside setup_response_data; we wire
         # it in after generation is up (it needs the OpenAI server URLs).
@@ -377,11 +382,6 @@ def setup_single_controller(
     # NeMo-Gym actor (after generation is up so OpenAI URLs are available)
     # ==========================
     if use_nemo_gym:
-        if generation_config["backend"] != "vllm":
-            raise NotImplementedError(
-                "SC NeMo-Gym integration currently supports the vllm backend "
-                f"only; got {generation_config['backend']!r}"
-            )
         enable_router_replay = router_replay_enabled(master_config.policy)
         env_handles["nemo_gym"] = spinup_nemo_gym_actor(
             env_configs=master_config.env,
diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py
index 107ca98bd9..37f4f13150 100644
--- a/nemo_rl/environments/nemo_gym.py
+++ b/nemo_rl/environments/nemo_gym.py
@@ -616,6 +616,19 @@ def spinup_nemo_gym_actor(
     When ``env_configs["nemo_gym"]["num_gpu_nodes"] > 0``, the actor is
     scheduled with soft NodeAffinity to the current Ray node so its colocated
     GPU resources land where the caller expects.
+
+    Args:
+        env_configs: The ``master_config.env`` mapping; ``env_configs["nemo_gym"]``
+            supplies the Gym global config plus NeMo-RL detection knobs
+            (``invalid_tool_call_patterns``, ``thinking_tags``, ``num_gpu_nodes``).
+        base_urls: Per-DP-rank OpenAI-compatible server base URLs from the
+            generation backend.
+        model_name: Served model name the Gym rollouts should target.
+        enable_router_replay: Sets ``require_routed_experts`` on the
+            ``NemoGymConfig``.
+
+    Returns:
+        The spun-up ``NemoGym`` Ray actor handle (``_spinup`` already awaited).
     """
     nemo_gym_dict = dict(env_configs["nemo_gym"])
 
diff --git a/tests/functional/L1_Functional_Tests_SingleController.sh b/tests/functional/L1_Functional_Tests_SingleController.sh
index ea818aaf8c..2a6fad63ee 100755
--- a/tests/functional/L1_Functional_Tests_SingleController.sh
+++ b/tests/functional/L1_Functional_Tests_SingleController.sh
@@ -35,6 +35,7 @@ run_test() {
 }
 
 run_test fast uv run --no-sync bash ./tests/functional/grpo_dp_single_controller.sh
+run_test fast uv run --no-sync bash ./tests/functional/grpo_async_gym_single_controller.sh
 
 cd ${PROJECT_ROOT}/tests
 if compgen -G ".coverage*" > /dev/null; then
diff --git a/tests/unit/single_controller/test_single_controller_setup.py b/tests/unit/single_controller/test_single_controller_setup.py
index c295757738..4defb58607 100644
--- a/tests/unit/single_controller/test_single_controller_setup.py
+++ b/tests/unit/single_controller/test_single_controller_setup.py
@@ -414,7 +414,14 @@ def test_nemo_gym_wires_env_handle(self, patched_factories):
         ):
             actor_args = setup_single_controller(mc, MagicMock(pad_token_id=0))
 
-        mock_spinup.assert_called_once()
+        mock_spinup.assert_called_once_with(
+            env_configs=mc.env,
+            base_urls=patched_factories[
+                "_build_generation"
+            ].return_value.dp_openai_server_base_urls,
+            model_name="test-model",
+            enable_router_replay=False,
+        )
         assert actor_args.env_handles["nemo_gym"] is fake_gym_actor
 
     @pytest.mark.parametrize("backend", ["sglang", "megatron"])

From 1b9af5d048716d39976b35747fbb3e0772a08cdd Mon Sep 17 00:00:00 2001
From: ruit <ruit@nvidia.com>
Date: Fri, 24 Jul 2026 09:37:23 -0700
Subject: [PATCH 11/13] docs(sc): track deferred NeMo-Gym initialization

Signed-off-by: ruit <ruit@nvidia.com>
---
 nemo_rl/algorithms/single_controller_utils/setup.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py
index bbac9a28c2..821c50cfd9 100644
--- a/nemo_rl/algorithms/single_controller_utils/setup.py
+++ b/nemo_rl/algorithms/single_controller_utils/setup.py
@@ -382,6 +382,8 @@ def setup_single_controller(
     # NeMo-Gym actor (after generation is up so OpenAI URLs are available)
     # ==========================
     if use_nemo_gym:
+        # TODO(#2625): Mirror GRPO's deferred vLLM load so NeMo-Gym spinup
+        # overlaps model loading instead of running serially afterward.
         enable_router_replay = router_replay_enabled(master_config.policy)
         env_handles["nemo_gym"] = spinup_nemo_gym_actor(
             env_configs=master_config.env,

From 4b4397ac68e66c4fd68b718d35c47f4fced8b660 Mon Sep 17 00:00:00 2001
From: Terry Kong <terryk@nvidia.com>
Date: Fri, 24 Jul 2026 13:19:35 -0700
Subject: [PATCH 12/13] test(sc): add env to run_grpo_single_controller
 main_context fixture

The entrypoint restored the NeMo-Gym setup call, which reads
config.env.get("should_use_nemo_gym") in main(). The main_context fixture's
SimpleNamespace config had no `env` attribute, so
test_cleanup_is_best_effort_and_preserves_run_error (and the other main()
tests) raised AttributeError: 'types.SimpleNamespace' object has no
attribute 'env'. Add env={} so the gym branch is skipped in these tests.

Signed-off-by: Terry Kong <terryk@nvidia.com>
---
 tests/unit/single_controller/test_run_grpo_single_controller.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/unit/single_controller/test_run_grpo_single_controller.py b/tests/unit/single_controller/test_run_grpo_single_controller.py
index c1331a66f4..3a22bd44c8 100644
--- a/tests/unit/single_controller/test_run_grpo_single_controller.py
+++ b/tests/unit/single_controller/test_run_grpo_single_controller.py
@@ -31,6 +31,7 @@ def main_context(monkeypatch: pytest.MonkeyPatch) -> SimpleNamespace:
             "draft": {"enabled": False},
             "megatron_cfg": {"mtp_num_layers": 2},
         },
+        env={},
         data_plane={"enabled": True},
         logger={"log_dir": "/tmp/logs"},
         checkpointing={"enabled": False},

From bad63f0494b348275abbd77dc93dddc70f9c22b1 Mon Sep 17 00:00:00 2001
From: Terry Kong <terryk@nvidia.com>
Date: Fri, 24 Jul 2026 13:34:59 -0700
Subject: [PATCH 13/13] fix(sc): compute reference logprobs when KL penalty is
 set in nemo-gym smoke

grpo_async_gym_single_controller.sh overrode
loss_fn.reference_policy_kl_penalty=0.01 but inherited
grpo.skip_reference_policy_logprobs_calculation=true from the
workplace-assistant recipe base, so the loss demanded
reference_policy_logprobs that the SC train pump was told to skip ->
KeyError: 'reference_policy_logprobs' mid-training. Set
skip_reference_policy_logprobs_calculation=false in the smoke so the
reference-logprob path is exercised end-to-end on the gym route (the
non-gym grpo_dp smoke already trains with the KL penalty).

Also harden validate_single_controller_config: a non-zero
reference_policy_kl_penalty combined with
skip_reference_policy_logprobs_calculation=true now raises at setup with
a clear message, instead of failing deep in the loss with an opaque
KeyError after the whole cluster/model build.

Signed-off-by: Terry Kong <terryk@nvidia.com>
---
 .../single_controller_utils/config.py         | 20 +++++++++++++++++++
 .../grpo_async_gym_single_controller.sh       |  1 +
 2 files changed, 21 insertions(+)

diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py
index b9234ae739..d83f07ecf7 100644
--- a/nemo_rl/algorithms/single_controller_utils/config.py
+++ b/nemo_rl/algorithms/single_controller_utils/config.py
@@ -118,6 +118,26 @@ def validate_single_controller_config(master_config: MasterConfig) -> None:
         sampler_name=async_config.sampler.name,
     )
 
+    # A non-zero reference-policy KL penalty makes the loss read
+    # ``reference_policy_logprobs``, but the SC train pump only computes them
+    # when ``skip_reference_policy_logprobs_calculation`` is false (see
+    # SingleControllerActor._reference_logprobs_required). Catch the
+    # inconsistent pair at setup instead of a mid-training KeyError.
+    reference_policy_kl_penalty = getattr(
+        master_config.loss_fn, "reference_policy_kl_penalty", 0
+    )
+    if reference_policy_kl_penalty and master_config.grpo.get(
+        "skip_reference_policy_logprobs_calculation"
+    ):
+        raise ValueError(
+            "loss_fn.reference_policy_kl_penalty="
+            f"{reference_policy_kl_penalty} requires reference_policy_logprobs, "
+            "but grpo.skip_reference_policy_logprobs_calculation=true skips "
+            "computing them on the SingleController path. Set "
+            "grpo.skip_reference_policy_logprobs_calculation=false, or set "
+            "loss_fn.reference_policy_kl_penalty=0."
+        )
+
 
 # ── Internal SingleController configs ────────────────────────────────────
 
diff --git a/tests/functional/grpo_async_gym_single_controller.sh b/tests/functional/grpo_async_gym_single_controller.sh
index 0178961739..54449013fd 100755
--- a/tests/functional/grpo_async_gym_single_controller.sh
+++ b/tests/functional/grpo_async_gym_single_controller.sh
@@ -81,6 +81,7 @@ uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJE
     policy.train_micro_batch_size=1 \
     cluster.gpus_per_node=2 \
     loss_fn.reference_policy_kl_penalty=0.01 \
+    grpo.skip_reference_policy_logprobs_calculation=false \
     loss_fn.use_importance_sampling_correction=true \
     logger.tensorboard_enabled=true \
     logger.log_dir=$LOG_DIR \