feat: support R3 async RL without TQ

nemo_rl/algorithms/grpo.py

@@ -1564,6 +1564,36 @@ def _should_use_async_rollouts(master_config: MasterConfig) -> bool:
     return False
 
 
+def _preserve_router_replay_routed_experts(
+    target: BatchedDataDict,
+    flat_messages: BatchedDataDict,
+    policy_config: PolicyConfig,
+) -> None:
+    """Carry rollout-recorded routes into policy worker inputs when R3 is enabled."""
+    if router_replay_enabled(policy_config) and "routed_experts" in flat_messages:
+        target["routed_experts"] = flat_messages["routed_experts"]
+
+
+def _build_async_grpo_train_data(
+    flat_messages: BatchedDataDict,
+    input_lengths: torch.Tensor,
+    repeated_batch: BatchedDataDict,
+    policy_config: PolicyConfig,
+) -> BatchedDataDict[ClippedPGLossDataDict]:
+    """Build the async no-TQ policy train batch from flattened rollout messages."""
+    train_data = BatchedDataDict[ClippedPGLossDataDict](
+        {
+            "input_ids": flat_messages["token_ids"],
+            "input_lengths": input_lengths,
+            "generation_logprobs": flat_messages["generation_logprobs"],
+            "token_mask": flat_messages["token_loss_mask"],
+            "sample_mask": repeated_batch["loss_multiplier"],
+        }
+    )
+    _preserve_router_replay_routed_experts(train_data, flat_messages, policy_config)
+    return train_data
+
+
 def _should_use_nemo_gym(master_config: MasterConfig) -> bool:
     """Determine if NeMo-Gym should be used for rollouts and validation based on the configuration."""
     env_config = master_config.env
@@ -2383,11 +2413,9 @@ def grpo_train(
                     # but the train_data whitelist above drops it. Copy it back so
                     # the Megatron worker's train-stage router-replay guard finds
                     # it. Mirrors the TQ producer (sync_rollout_actor.py).
-                    if (
-                        router_replay_enabled(master_config.policy)
-                        and "routed_experts" in flat_messages
-                    ):
-                        train_data["routed_experts"] = flat_messages["routed_experts"]
+                    _preserve_router_replay_routed_experts(
+                        train_data, flat_messages, master_config.policy
+                    )
                     train_data.to("cpu")
 
                     metrics_logging_data["content"] = flat_messages["content"]
@@ -2438,11 +2466,9 @@ def grpo_train(
                     # intentionally ignores routed_experts (require_router_replay
                     # =False short-circuits before the field is read), so a
                     # present-but-unused field here is safe.
-                    if (
-                        router_replay_enabled(master_config.policy)
-                        and "routed_experts" in flat_messages
-                    ):
-                        logprob_data["routed_experts"] = flat_messages["routed_experts"]
+                    _preserve_router_replay_routed_experts(
+                        logprob_data, flat_messages, master_config.policy
+                    )
 
                     if not skip_prev_logprobs:
                         train_data["prev_logprobs"] = policy.get_logprobs(
@@ -3185,6 +3211,14 @@ def async_grpo_train(
     assert master_config.loss_fn.use_importance_sampling_correction, (
         "Importance sampling correction must be enabled for async GRPO for good convergence due to off-policy samples!"
     )
+    if router_replay_enabled(master_config.policy) and (
+        master_config.data_plane or {}
+    ).get("enabled", False):
+        raise NotImplementedError(
+            "policy.router_replay.enabled=true with async GRPO is currently "
+            "supported only when data_plane.enabled=false. Async + TQ support "
+            "has not been merged yet."
+        )
 
     if master_config.grpo["async_grpo"]["max_trajectory_age_steps"] > 1:
         if not master_config.grpo["async_grpo"].get("in_flight_weight_updates", False):
@@ -3600,16 +3634,12 @@ def async_grpo_train(
                         ],
                     )
 
-                    # Create training data
-                    # Note: advantages will be computed and added after logprobs are available
-                    train_data = BatchedDataDict[ClippedPGLossDataDict](
-                        {
-                            "input_ids": flat_messages["token_ids"],
-                            "input_lengths": input_lengths,
-                            "generation_logprobs": flat_messages["generation_logprobs"],
-                            "token_mask": flat_messages["token_loss_mask"],
-                            "sample_mask": repeated_batch["loss_multiplier"],
-                        }
+                    # Create training data. Advantages are added after logprobs.
+                    train_data = _build_async_grpo_train_data(
+                        flat_messages,
+                        input_lengths,
+                        repeated_batch,
+                        master_config.policy,
                     )
                     train_data.to("cpu")
 

tests/unit/algorithms/test_grpo_router_replay_async.py

@@ -0,0 +1,129 @@
+# 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.
+# 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 types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+from nemo_rl.algorithms.grpo import (
+    MasterConfig,
+    _build_async_grpo_train_data,
+    _default_grpo_save_state,
+    async_grpo_train,
+)
+from nemo_rl.distributed.batched_data_dict import BatchedDataDict
+
+
+@pytest.fixture(scope="session", autouse=True)
+def init_ray_cluster():
+    yield
+
+
+@pytest.fixture(scope="session", autouse=True)
+def ray_gpu_monitor():
+    class _NoopGpuMonitor:
+        def _collect_metrics(self):
+            return {}
+
+        def stop(self):
+            pass
+
+    yield _NoopGpuMonitor()
+
+
+@pytest.fixture(scope="session", autouse=True)
+def session_data(_unit_test_data):
+    yield _unit_test_data
+
+
+def _make_async_master_config(data_plane=None) -> MasterConfig:
+    return MasterConfig.model_construct(
+        **{
+            "policy": {
+                "router_replay": {"enabled": True},
+                "generation": {
+                    "backend": "vllm",
+                    "vllm_cfg": {"async_engine": True},
+                },
+            },
+            "loss_fn": SimpleNamespace(use_importance_sampling_correction=True),
+            "data_plane": data_plane,
+        }
+    )
+
+
+# Keep this focused on async no-TQ batch construction instead of full Ray orchestration.
+@pytest.mark.parametrize(
+    ("policy_config", "expect_routed_experts"),
+    [
+        ({"router_replay": {"enabled": True}}, True),
+        ({"router_replay": {"enabled": False}}, False),
+    ],
+)
+def test_build_async_grpo_train_data_preserves_routed_experts_for_r3(
+    policy_config, expect_routed_experts
+):
+    routes = torch.arange(1 * 3 * 2 * 4, dtype=torch.int32).reshape(1, 3, 2, 4)
+    flat_messages = BatchedDataDict(
+        {
+            "token_ids": torch.tensor([[1, 2, 3]]),
+            "generation_logprobs": torch.zeros(1, 3),
+            "token_loss_mask": torch.tensor([[0, 1, 1]]),
+            "routed_experts": routes,
+        }
+    )
+    input_lengths = torch.tensor([3])
+    repeated_batch = BatchedDataDict({"loss_multiplier": torch.tensor([1.0])})
+
+    train_data = _build_async_grpo_train_data(
+        flat_messages,
+        input_lengths,
+        repeated_batch,
+        policy_config,
+    )
+
+    assert torch.equal(train_data["input_ids"], flat_messages["token_ids"])
+    assert torch.equal(train_data["input_lengths"], input_lengths)
+    assert torch.equal(
+        train_data["generation_logprobs"], flat_messages["generation_logprobs"]
+    )
+    assert torch.equal(train_data["token_mask"], flat_messages["token_loss_mask"])
+    assert torch.equal(train_data["sample_mask"], repeated_batch["loss_multiplier"])
+
+    if expect_routed_experts:
+        assert torch.equal(train_data["routed_experts"], routes)
+    else:
+        assert "routed_experts" not in train_data
+
+
+def test_async_grpo_r3_rejects_data_plane_until_async_tq_exists():
+    master_config = _make_async_master_config(data_plane={"enabled": True})
+
+    with pytest.raises(NotImplementedError, match="data_plane.enabled=false"):
+        async_grpo_train(
+            MagicMock(),
+            MagicMock(),
+            MagicMock(),
+            None,
+            MagicMock(),
+            MagicMock(),
+            {"math": MagicMock()},
+            None,
+            MagicMock(),
+            MagicMock(),
+            _default_grpo_save_state(),
+            master_config,
+        )