FROM 72e87b1bcff34bc1c2a67f352202d3a2cdbb3b84 async+gym (but missing genrm change)

Signed-off-by: Jiaqi Zeng <jiaqiz@nvidia.com>
Signed-off-by: Terry Kong <terryk@nvidia.com>

Author: HeyyyyyyG

examples/nemo_gym/run_grpo_nemo_gym.py

@@ -231,7 +231,55 @@ def main() -> None:
             logger=logger,
             master_config=master_config,
         )
+    # Check if async mode is enabled
+    elif "async_grpo" in config["grpo"] and config["grpo"]["async_grpo"]["enabled"]:
+        # Async GRPO does not support dynamic sampling, reward scaling, or reward shaping (DAPO features)
+        unsupported_features = [
+            "use_dynamic_sampling",
+            "reward_scaling",
+            "reward_shaping",
+        ]
+
+        for feature in unsupported_features:
+            if feature not in config["grpo"]:
+                continue
+
+            if feature == "use_dynamic_sampling":
+                if config["grpo"][feature]:
+                    raise NotImplementedError(
+                        f"{feature} is not supported with async GRPO"
+                    )
+            else:
+                if config["grpo"][feature]["enabled"]:
+                    raise NotImplementedError(
+                        f"{feature} is not supported with async GRPO"
+                    )
+
+        from nemo_rl.algorithms.grpo import async_grpo_train
+
+        print("🚀 Running async GRPO training")
+
+        async_config = config["grpo"]["async_grpo"]
+        # Run async GRPO training
+        async_grpo_train(
+            policy=policy,
+            policy_generation=policy_generation,
+            dataloader=dataloader,
+            val_dataloader=val_dataloader,
+            tokenizer=tokenizer,
+            loss_fn=loss_fn,
+            task_to_env=task_to_env,
+            val_task_to_env=val_task_to_env,
+            logger=logger,
+            checkpointer=checkpointer,
+            grpo_save_state=grpo_state,
+            master_config=master_config,
+            max_trajectory_age_steps=async_config["max_trajectory_age_steps"],
+        )
     else:
+        print("🚀 Running synchronous GRPO training")
+
+        # Run standard GRPO training
         grpo_train(
             policy,
             policy_generation,

nemo_rl/algorithms/async_utils.py

@@ -642,17 +642,51 @@ def _run_prompt_group_worker(
         prompt_idx: int,
     ) -> None:
         try:
+            # Import here to avoid circular dependency
+            from nemo_rl.algorithms.grpo import _should_use_nemo_gym
+            from nemo_rl.experience.rollouts import run_async_nemo_gym_rollout
+
             # Run rollout for this prompt group
             # Async engine supports concurrent generation; avoid locking
-            final_batch, rollout_metrics = run_async_multi_turn_rollout(
-                policy_generation=self.policy_generation,
-                input_batch=repeated_batch,
-                tokenizer=self.tokenizer,
-                task_to_env=self.task_to_env,
-                max_seq_len=self.master_config["policy"]["max_total_sequence_length"],
-                max_rollout_turns=self.master_config["grpo"]["max_rollout_turns"],
-                greedy=False,
-            )
+            # Check if we should use nemo_gym (similar to synchronous GRPO)
+            if _should_use_nemo_gym(self.master_config):
+                generation_config = self.master_config["policy"]["generation"]
+                env_cfg = self.master_config.get("env") or {}
+                nemo_gym_rollout_result = run_async_nemo_gym_rollout(
+                    policy_generation=self.policy_generation,
+                    input_batch=repeated_batch,
+                    tokenizer=self.tokenizer,
+                    task_to_env=self.task_to_env,
+                    max_seq_len=None,
+                    generation_config=generation_config,
+                    max_rollout_turns=None,
+                    greedy=False,
+                    # GenRM compare config
+                    use_genrm_compare=env_cfg.get("use_genrm_compare", False),
+                    num_generations_per_prompt=self.master_config["grpo"][
+                        "num_generations_per_prompt"
+                    ],
+                    genrm_compare_server_name=env_cfg.get(
+                        "genrm_compare_server_name", "genrm_compare"
+                    ),
+                    genrm_agent_names=env_cfg.get(
+                        "genrm_agent_names", ["genrm_simple_agent"]
+                    ),
+                )
+                final_batch = nemo_gym_rollout_result.final_batch
+                rollout_metrics = nemo_gym_rollout_result.rollout_metrics
+            else:
+                final_batch, rollout_metrics = run_async_multi_turn_rollout(
+                    policy_generation=self.policy_generation,
+                    input_batch=repeated_batch,
+                    tokenizer=self.tokenizer,
+                    task_to_env=self.task_to_env,
+                    max_seq_len=self.master_config["policy"][
+                        "max_total_sequence_length"
+                    ],
+                    max_rollout_turns=self.master_config["grpo"]["max_rollout_turns"],
+                    greedy=False,
+                )
 
             # Move to CPU and push to buffer (avoid blocking on GC/push)
             final_batch_cpu = final_batch.to("cpu")