feat(grpo): make env-flagged sample masking optional

examples/nemo_gym/grpo_nanov3.yaml

@@ -304,6 +304,8 @@ env:
   should_use_nemo_gym: true
   # true: skip expensive train_data_step*.jsonl; false: write full jsonl.
   should_log_nemo_gym_responses: true
+  # false: ignore env mask_sample flags so the loss trains on every sample.
+  should_mask_flagged_samples: true
   nemo_gym:
     # Port range for Gym HTTP servers, kept below the OS ephemeral range
     # (below the 9000 ephemeral floor on some GB200 nodes) and non-overlapping with NeMo

examples/nemo_gym/grpo_qwen3_30ba3b_thinking_swe1.yaml

@@ -219,6 +219,8 @@ data:
 env:
   should_use_nemo_gym: true
   should_log_nemo_gym_responses: false
+  # false: ignore env mask_sample flags so the loss trains on every sample.
+  should_mask_flagged_samples: true
   nemo_gym:
     skip_venv_if_present: true
     port_range_low: 5000

examples/nemo_gym/grpo_qwen3_30ba3b_thinking_swe2.yaml

@@ -218,6 +218,8 @@ data:
 env:
   should_use_nemo_gym: true
   should_log_nemo_gym_responses: false
+  # false: ignore env mask_sample flags so the loss trains on every sample.
+  should_mask_flagged_samples: true
   nemo_gym:
     skip_venv_if_present: true
     port_range_low: 5000

examples/nemo_gym/grpo_workplace_assistant_nemotron_nano_v2_9b.yaml

@@ -291,6 +291,8 @@ env:
   should_use_nemo_gym: true
   # true: skip expensive train_data_step*.jsonl; false: write jsonl (heavy). Prefer true if disk is tight.
   should_log_nemo_gym_responses: true
+  # false: ignore env mask_sample flags so the loss trains on every sample.
+  should_mask_flagged_samples: true
   nemo_gym:  # This is passed into NeMo-Gym as the initial_global_config_dict
     # Port range for Gym HTTP servers, kept below the OS ephemeral range
     # (below the 9000 ephemeral floor on some GB200 nodes) and non-overlapping with NeMo

examples/nemo_gym/nemotron-3-super/stage1_rlvr.yaml

@@ -280,6 +280,8 @@ env:
   should_use_nemo_gym: true
   # true: skip expensive train_data_step*.jsonl (recommended for large Gym runs); false: write full jsonl.
   should_log_nemo_gym_responses: true
+  # false: ignore env mask_sample flags so the loss trains on every sample.
+  should_mask_flagged_samples: true
   nemo_gym:
     uv_venv_dir: "${oc.env:PERSISTENT_CACHE}/gym_venvs"
     skip_venv_if_present: true

examples/nemo_gym/nemotron-3-super/stage2_swe1.yaml

@@ -282,6 +282,8 @@ data:
 
 env:
   should_use_nemo_gym: true
+  # false: ignore env mask_sample flags so the loss trains on every sample.
+  should_mask_flagged_samples: true
   nemo_gym:
     uv_venv_dir: "${oc.env:PERSISTENT_CACHE}/gym_venvs"  # overridden by super_launch.sh to $GYM_VENV_DIR
     skip_venv_if_present: true

examples/nemo_gym/nemotron-3-super/stage2_swe2.yaml

@@ -278,6 +278,8 @@ sif_dir: "/path/to/sif"
 env:
   should_use_nemo_gym: true
   should_log_nemo_gym_responses: false
+  # false: ignore env mask_sample flags so the loss trains on every sample.
+  should_mask_flagged_samples: true
   nemo_gym:
     uv_venv_dir: "${oc.env:PERSISTENT_CACHE}/gym_venvs"  # overridden by super_launch.sh to $GYM_VENV_DIR
     skip_venv_if_present: true

examples/nemo_gym/nemotron-3-super/stage3_rlhf.yaml

@@ -281,6 +281,8 @@ env:
   should_use_nemo_gym: true
   # true: skip expensive train_data_step*.jsonl (recommended for large Gym runs); false: enable full jsonl.
   should_log_nemo_gym_responses: true
+  # false: ignore env mask_sample flags so the loss trains on every sample.
+  should_mask_flagged_samples: true
   nemo_gym:
     uv_venv_dir: "${oc.env:PERSISTENT_CACHE}/gym_venvs"  # overridden by super_launch.sh to $GYM_VENV_DIR
     skip_venv_if_present: true

nemo_rl/algorithms/async_utils/trajectory_collector.py

@@ -752,6 +752,7 @@ async def _iter_rollout_groups(
             from nemo_rl.experience.rollouts import (
                 get_nemo_gym_thinking_tags,
                 run_async_nemo_gym_rollout,
+                should_mask_flagged_samples,
             )
 
             # NeMo-Gym owns stop criteria. Configuration fills the policy EOS token
@@ -777,6 +778,9 @@ async def _iter_rollout_groups(
                 greedy=False,
                 reward_penalty_config=self.master_config.reward_penalties,
                 thinking_tags=get_nemo_gym_thinking_tags(self.master_config.env),
+                mask_env_flagged_samples=should_mask_flagged_samples(
+                    self.master_config.env
+                ),
             ):
                 task_index = rollout_result.task_index
                 if task_index is None:

nemo_rl/algorithms/grpo.py

@@ -90,6 +90,7 @@
     run_async_multi_turn_rollout,
     run_multi_turn_rollout,
     run_nemo_gym_rollout_sync,
+    should_mask_flagged_samples,
 )
 from nemo_rl.models.generation.interfaces import (
     GenerationConfig,
@@ -2769,6 +2770,9 @@ def grpo_train(
                             effort_config=_get_effort_config(master_config),
                             reward_penalty_config=master_config.reward_penalties,
                             thinking_tags=get_nemo_gym_thinking_tags(master_config.env),
+                            mask_env_flagged_samples=should_mask_flagged_samples(
+                                master_config.env
+                            ),
                         )
                         input_ids = nemo_gym_rollout_result.input_ids
                         repeated_batch = nemo_gym_rollout_result.final_batch
@@ -3614,6 +3618,9 @@ def validate(
                     effort_config=_get_effort_config(master_config),
                     reward_penalty_config=master_config.reward_penalties,
                     thinking_tags=get_nemo_gym_thinking_tags(master_config.env),
+                    mask_env_flagged_samples=should_mask_flagged_samples(
+                        master_config.env
+                    ),
                 )
                 val_batch = nemo_gym_rollout_result.final_batch
                 gen_metrics = nemo_gym_rollout_result.rollout_metrics

nemo_rl/experience/rollout_manager.py

@@ -405,6 +405,7 @@ def __init__(
         max_seq_len: int,
         max_rollout_turns: int,
         generation_config: GenerationConfig,
+        mask_env_flagged_samples: bool = True,
         **kwargs: Any,
     ) -> None:
         self._tokenizer = tokenizer
@@ -413,6 +414,7 @@ def __init__(
         self._max_seq_len = max_seq_len
         self._max_rollout_turns = max_rollout_turns
         self._generation_config = generation_config
+        self._mask_env_flagged_samples = mask_env_flagged_samples
 
         self._validate_init_params()
 
@@ -553,6 +555,13 @@ def _result_to_completion(self, result: dict) -> Completion:
             sum(len(m["token_ids"]) for m in result["message_log"]) == self._max_seq_len
         )
 
+        # Same gate as the batched path: when masking is off, drop the env
+        # mask flag so later batch building never sees it.
+        if not self._mask_env_flagged_samples:
+            (result["full_result"].get("instance_config") or {}).pop(
+                "mask_sample", None
+            )
+
         return Completion(
             message_log=result["message_log"],
             env_extras=result["full_result"],

nemo_rl/experience/rollouts.py

@@ -1558,6 +1558,21 @@ def get_nemo_gym_thinking_tags(env_config: dict[str, Any]) -> list[str]:
     return list(DEFAULT_THINKING_TAGS)
 
 
+def should_mask_flagged_samples(env_config: dict[str, Any]) -> bool:
+    """Read ``env.should_mask_flagged_samples``; absent means True.
+
+    True (the default): env-driven ``mask_sample`` flags are carried in the
+    rollout batch and flagged samples are dropped from the loss.
+
+    Set false when the flags are too coarse to honor: for example, Gym flags
+    rollouts that hit max iterations even when they solve the task, and those
+    are samples worth training on. It also keeps batch composition
+    deterministic for controlled benchmark runs — how many samples get
+    flagged varies run to run.
+    """
+    return env_config.get("should_mask_flagged_samples") is not False
+
+
 def _get_reward_penalty_config_value(
     reward_penalty_config: dict[str, Any] | BaseModel | None,
     key: str,
@@ -1997,6 +2012,7 @@ async def run_async_nemo_gym_rollout(
     effort_config: Optional[EffortLevelsConfig] = None,
     reward_penalty_config: dict[str, Any] | BaseModel | None = None,
     thinking_tags: list[str] | tuple[str, ...] | None = None,
+    mask_env_flagged_samples: bool = True,
     returns_entire_batch: bool = False,
 ) -> AsyncGenerator[NemoGymRolloutResult, None]:
     """Stream complete NeMo-Gym prompt groups in group-completion order.
@@ -2023,6 +2039,8 @@ async def run_async_nemo_gym_rollout(
         effort_config: Optional configuration for effort-based reward shaping.
         reward_penalty_config: Optional reward-penalty configuration.
         thinking_tags: Optional opening and closing tags used by thinking penalties.
+        mask_env_flagged_samples: Whether to carry env-driven ``mask_sample``
+            flags in the rollout batch for loss masking.
         returns_entire_batch: Whether to treat the input as one potentially
             heterogeneous group. This requires ``num_generations`` to equal the
             batch size and is used by synchronous callers.
@@ -2148,6 +2166,7 @@ async def run_async_nemo_gym_rollout(
                         effort_config=effort_config,
                         reward_penalty_config=reward_penalty_config,
                         thinking_tags=thinking_tags,
+                        mask_env_flagged_samples=mask_env_flagged_samples,
                     )
                     if accumulator.is_complete:
                         final_rollout_result = rollout_result
@@ -2184,6 +2203,7 @@ def run_nemo_gym_rollout_sync(
     effort_config: Optional[EffortLevelsConfig] = None,
     reward_penalty_config: dict[str, Any] | BaseModel | None = None,
     thinking_tags: list[str] | tuple[str, ...] | None = None,
+    mask_env_flagged_samples: bool = True,
 ) -> NemoGymRolloutResult:
     """Run and return one complete NeMo-Gym batch synchronously.
 
@@ -2206,6 +2226,8 @@ def run_nemo_gym_rollout_sync(
         effort_config: Optional configuration for effort-based reward shaping.
         reward_penalty_config: Optional reward-penalty configuration.
         thinking_tags: Optional opening and closing tags used by thinking penalties.
+        mask_env_flagged_samples: Whether to carry env-driven ``mask_sample``
+            flags in the rollout batch for loss masking.
 
     Returns:
         The fully postprocessed NeMo-Gym rollout batch in input-row order.
@@ -2235,6 +2257,7 @@ async def _consume_rollout() -> NemoGymRolloutResult:
             effort_config=effort_config,
             reward_penalty_config=reward_penalty_config,
             thinking_tags=thinking_tags,
+            mask_env_flagged_samples=mask_env_flagged_samples,
             returns_entire_batch=True,
         ):
             pass
@@ -2257,6 +2280,7 @@ def _postprocess_single_nemo_gym_group(
     effort_config: Optional[EffortLevelsConfig] = None,
     reward_penalty_config: dict[str, Any] | BaseModel | None = None,
     thinking_tags: list[str] | tuple[str, ...] | None = None,
+    mask_env_flagged_samples: bool = True,
 ) -> NemoGymRolloutResult:
     """Postprocess one complete prompt group from the NeMo-Gym stream."""
     # Length-based reward shaping for low-effort prompts
@@ -2437,11 +2461,12 @@ def _postprocess_single_nemo_gym_group(
             "truncated": torch.tensor(
                 [m["hit_max_tokens"] for m in all_sample_metrics], dtype=torch.bool
             ),
-            # Agent/env-driven mask flag — True means this sample should be masked
-            # from the GRPO gradient (kept for advantage computation).
-            "mask_sample": _extract_mask_sample_flags(results),
         }
     )
+    # Env/agent mask flag: flagged samples are dropped from the loss but still
+    # count for advantages. env.should_mask_flagged_samples=false skips this.
+    if mask_env_flagged_samples:
+        final_batch["mask_sample"] = _extract_mask_sample_flags(results)
 
     if length_rewards_low:
         rollout_metrics["mean_length_reward_low"] = sum(length_rewards_low) / len(