feat(nemo-gym): train on rollouts from external agent harnesses

nemo_rl/algorithms/grpo.py

@@ -2305,6 +2305,60 @@ def _apply_mask_sample_filter(repeated_batch: BatchedDataDict[DatumSpec]) -> int
     return num_masked
 
 
+def _prompt_grouping_ids(
+    prompt_token_ids: torch.Tensor,
+    *,
+    num_prompts: int,
+    num_generations_per_prompt: int,
+    use_nemo_gym: bool,
+    group_offset: int = 0,
+) -> torch.Tensor:
+    """Return identities used to group generations from the same prompt."""
+    if not use_nemo_gym:
+        return prompt_token_ids
+
+    return (
+        torch.arange(group_offset, group_offset + num_prompts)
+        .repeat_interleave(num_generations_per_prompt)
+        .unsqueeze(1)
+    )
+
+
+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
+    should_use_nemo_gym = bool(env_config.get("should_use_nemo_gym"))
+    if not should_use_nemo_gym:
+        return should_use_nemo_gym
+
+    # Validate the setup for training with NeMo-Gym.
+    generation_config = master_config.policy["generation"]
+    assert should_use_async_rollouts(generation_config), (
+        "❌ Error: In order to use NeMo-Gym, you must use a generation backend with `async_engine: true`!"
+    )
+
+    # We piggyback off of `should_use_async_rollouts` to guarantee the existence of these configs.
+    if generation_config["backend"] == "vllm":
+        should_expose_http_server = generation_config["vllm_cfg"].get(
+            "expose_http_server"
+        )
+    elif generation_config["backend"] == "megatron":
+        should_expose_http_server = generation_config["mcore_generation_config"].get(
+            "expose_http_server"
+        )
+    elif generation_config["backend"] == "trtllm":
+        should_expose_http_server = generation_config["trtllm_cfg"].get(
+            "expose_http_server"
+        )
+    else:
+        should_expose_http_server = False
+    assert should_expose_http_server, (
+        "In order to use NeMo-Gym, you must expose the generation server via `expose_http_server: true`!"
+    )
+
+    return should_use_nemo_gym
+
+
 def _should_log_nemo_gym_responses(master_config: MasterConfig) -> bool:
     """Whether NeMo Gym is responsible for full response logging.
 
@@ -2937,6 +2991,7 @@ def grpo_train(
         batch_cache: BatchedDataDict[DatumSpec] = None
         # This is the number of batches we processed so far at each step to generate responses whose std is non-zero. Maximum threshold is set by dynamic_sampling_max_gen_batches. Used in the case of dynamic sampling.
         dynamic_sampling_num_gen_batches = 0
+        prompt_grouping_id_offset = 0
 
         # Run grpo/dapo training loop (single-turn)
         for batch in wrapped_dataloader:
@@ -2995,7 +3050,14 @@ def grpo_train(
                         repeated_batch["message_log"],
                         pad_value_dict={"token_ids": tokenizer.pad_token_id},
                     )
-                    input_ids = batched_flat["token_ids"]
+                    prompt_grouping_ids = _prompt_grouping_ids(
+                        batched_flat["token_ids"],
+                        num_prompts=batch.size,
+                        num_generations_per_prompt=master_config.grpo.num_generations_per_prompt,
+                        use_nemo_gym=_should_use_nemo_gym(master_config),
+                        group_offset=prompt_grouping_id_offset,
+                    )
+                    prompt_grouping_id_offset += batch.size
 
                 # Generate responses - this updates the LLMMessageLogType in repeated_batch
                 memory_tracker.snapshot_start_of_stage("Generation", dir())
@@ -3101,7 +3163,8 @@ def grpo_train(
                                 master_config.grpo.debug_payload_metrics
                             ),
                         )
-                        input_ids = nemo_gym_rollout_result.input_ids
+                        # Keep the repeated dataset prompt IDs for GRPO grouping.
+                        # Captured first-call prompts may contain harness-specific paths.
                         repeated_batch = nemo_gym_rollout_result.final_batch
                         rollout_metrics = nemo_gym_rollout_result.rollout_metrics
                         del nemo_gym_rollout_result
@@ -3186,7 +3249,7 @@ def grpo_train(
                         # Just fix the device id for now
                         device_id = 0
                         baseline, std = calculate_baseline_and_std_per_prompt(
-                            input_ids.cuda(device_id),
+                            prompt_grouping_ids.cuda(device_id),
                             rewards.cuda(device_id),
                             torch.ones_like(rewards).cuda(device_id),
                             leave_one_out_baseline=master_config.grpo.use_leave_one_out_baseline,
@@ -3200,13 +3263,16 @@ def grpo_train(
                         std = std.cpu()
                     else:
                         baseline, std = calculate_baseline_and_std_per_prompt(
-                            input_ids,
+                            prompt_grouping_ids,
                             rewards,
                             torch.ones_like(rewards),
                             leave_one_out_baseline=master_config.grpo.use_leave_one_out_baseline,
                             std_rewards=std_rewards,
                         )
 
+                    # Keep source prompt identity aligned with rows selected by dynamic sampling.
+                    repeated_batch["prompt_grouping_ids"] = prompt_grouping_ids
+
                     # Apply dynamic sampling to filter prompts with non-zero std (DAPO algorithm)
                     repeated_batch, is_batch_complete, batch_cache, ds_metrics = (
                         dynamic_sampling(
@@ -3243,24 +3309,12 @@ def grpo_train(
                     # Save baseline for logging (before deletion)
                     baseline_for_log = baseline.clone()
 
-                    # Must precede prompt extraction: it reuses the same message
-                    # dicts, so this also protects the prompt flatten below.
+                    # Must precede the training flatten because it reuses the same message dicts.
                     backfill_missing_routed_experts(repeated_batch["message_log"])
 
-                    # Extract original prompt messages using the length field
-                    # This correctly handles multi-turn prompts that contain assistant messages
-                    initial_prompt_message_logs = extract_initial_prompt_messages(
-                        repeated_batch["message_log"],
-                        repeated_batch["length"],
-                    )
-                    prompt_batched_flat, _ = batched_message_log_to_flat_message(
-                        initial_prompt_message_logs,
-                        pad_value_dict={"token_ids": tokenizer.pad_token_id},
-                    )
-                    prompt_ids_for_adv = prompt_batched_flat["token_ids"]
-                    del initial_prompt_message_logs
-                    del prompt_batched_flat
-                    del input_ids
+                    prompt_ids_for_adv = repeated_batch["prompt_grouping_ids"]
+                    del repeated_batch["prompt_grouping_ids"]
+                    del prompt_grouping_ids
                     del baseline
                     del std
 
@@ -3892,6 +3946,7 @@ def grpo_train(
             # Reset the batch and set dynamic_sampling_num_gen_batches to 0
             batch_cache = None
             dynamic_sampling_num_gen_batches = 0
+            prompt_grouping_id_offset = 0
 
             # Clear mem
             memory_tracker.snapshot_start_of_stage("After CPU memory clear", dir())

nemo_rl/algorithms/ppo.py

@@ -1392,7 +1392,6 @@ def ppo_train(
                             max_rollout_turns=None,
                             greedy=False,
                         )
-                        input_ids = nemo_gym_rollout_result.input_ids
                         repeated_batch = nemo_gym_rollout_result.final_batch
                         rollout_metrics = nemo_gym_rollout_result.rollout_metrics
                         del nemo_gym_rollout_result