fix(grpo): support environment sample masking in single controller

docs/guides/single-controller.md

@@ -241,6 +241,6 @@ The SC path is still under active development. Feature gaps are tracked in [issu
 - Generation backend: vLLM and Megatron generation are supported; SGLang and TRT-LLM have not been tested on SC.
 - Validation is not yet supported (setup raises on `val_period > 0`, `val_at_start`, or `val_at_end`); checkpointing is.
 - (PPO) Rollout drop budgets — `async_rl.rollout_failure.max_skipped_prompts` and `max_consecutive_dropped_prompts` must both be `0`. A drop shortens the step, and the critic shards it against the configured `value.train_global_batch_size` rather than its actual size, so setup rejects a non-zero budget. The resiliency layer stays available on GRPO.
-- Reward shaping and sample filtering — `overlong_filtering`, `reward_shaping`, `reward_scaling`, and `use_dynamic_sampling` are implemented on neither algorithm block, so setup rejects them rather than silently skipping the shaping.
+- Reward shaping — `reward_shaping`, `reward_scaling`, and `use_dynamic_sampling` are implemented on neither algorithm block, so setup rejects them rather than silently skipping the shaping. Environment-flagged sample masking and `overlong_filtering` are supported.
 - The `windowed` sampler has no `over_sampling_ratio` cap — over-produced groups aged past the window are evicted, wasting rollout compute.
 - The drain gate in refit is not yet supported.

nemo_rl/algorithms/single_controller.py

@@ -434,6 +434,7 @@ def __init__(
         self._step_log_dict: dict[str, list] = {
             "rewards": [],
             "masked_advantages": [],
+            "num_mask_sample_filtered": [],
             "sequence_lengths": [],
             "seq_logprob_error_metrics": [],
             **{key: [] for key in VIOLATION_TAG_KEYS},
@@ -1856,8 +1857,9 @@ async def _train_pump(self) -> None:
                     if self._is_ppo and not has_valid_training_tokens:
                         raise RuntimeError(
                             "SingleController has no valid response tokens after "
-                            "filtering. Check ppo.seq_logprob_error_threshold to "
-                            "avoid an optimizer step with an empty batch."
+                            "filtering. Check seq_logprob_error_threshold, "
+                            "overlong_filtering, and environment mask_sample flags "
+                            "to avoid an optimizer step with an empty batch."
                         )
 
                     # ---- 3. Train the model -- train_microbatches_from_meta ----
@@ -1988,8 +1990,9 @@ async def _train_pump(self) -> None:
                     if not step_open:
                         raise RuntimeError(
                             "SingleController has no valid response tokens after "
-                            "filtering. Check grpo.seq_logprob_error_threshold to "
-                            "avoid an optimizer step with an empty batch."
+                            "filtering. Check seq_logprob_error_threshold, "
+                            "overlong_filtering, and environment mask_sample flags "
+                            "to avoid an optimizer step with an empty batch."
                         )
 
                     with self._timer.time("policy_training"):
@@ -3203,6 +3206,18 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]:
         sample_mask = squeeze_trailing_unit_dim(
             tensor_field(data, adv_cfg.sample_mask_field)
         ).float()
+        mask_sample = squeeze_trailing_unit_dim(
+            tensor_field(data, adv_cfg.mask_sample_field)
+        ).bool()
+        truncated = squeeze_trailing_unit_dim(
+            tensor_field(data, adv_cfg.truncated_field)
+        ).bool()
+
+        num_mask_sample_filtered = int(mask_sample.sum().item())
+        self._step_log_dict["num_mask_sample_filtered"].append(num_mask_sample_filtered)
+        final_sample_mask = sample_mask * (~mask_sample).to(sample_mask.dtype)
+        if self._algo_cfg.overlong_filtering:
+            final_sample_mask = final_sample_mask * (~truncated).to(sample_mask.dtype)
 
         seq_logprob_error_threshold = self._algo_cfg.seq_logprob_error_threshold
         # Match the legacy path: whenever real policy logprobs are available,
@@ -3212,7 +3227,7 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]:
             masking_data = BatchedDataDict(
                 {
                     "token_mask": token_mask,
-                    "sample_mask": sample_mask,
+                    "sample_mask": final_sample_mask,
                     "prev_logprobs": tensor_field(
                         data,
                         adv_cfg.policy_logprobs_field,
@@ -3224,7 +3239,7 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]:
                 }
             )
             num_valid_seqs_before = float(
-                ((token_mask[:, 1:] * sample_mask.unsqueeze(-1)).sum(dim=-1) > 0)
+                ((token_mask[:, 1:] * final_sample_mask.unsqueeze(-1)).sum(dim=-1) > 0)
                 .sum()
                 .item()
             )
@@ -3233,9 +3248,9 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]:
                 rewards=rewards,
                 seq_logprob_error_threshold=seq_logprob_error_threshold,
             )
-            sample_mask = masking_data["sample_mask"]
+            final_sample_mask = masking_data["sample_mask"]
             num_valid_seqs_after = float(
-                ((token_mask[:, 1:] * sample_mask.unsqueeze(-1)).sum(dim=-1) > 0)
+                ((token_mask[:, 1:] * final_sample_mask.unsqueeze(-1)).sum(dim=-1) > 0)
                 .sum()
                 .item()
             )
@@ -3246,7 +3261,7 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]:
             seq_error_metrics["_num_valid_seqs_after"] = num_valid_seqs_after
             self._step_log_dict["seq_logprob_error_metrics"].append(seq_error_metrics)
 
-        mask = token_mask * sample_mask.unsqueeze(-1)
+        mask = token_mask * final_sample_mask.unsqueeze(-1)
 
         repeated_batch: dict[str, torch.Tensor] = {
             "total_reward": rewards,
@@ -3332,8 +3347,8 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]:
             self._opd_stat_count += int(valid.numel())
 
         fields_to_put = {adv_cfg.output_field: advantages}
-        if seq_logprob_error_threshold is not None:
-            fields_to_put[adv_cfg.sample_mask_field] = sample_mask
+        if not torch.equal(final_sample_mask, sample_mask):
+            fields_to_put[adv_cfg.sample_mask_field] = final_sample_mask
         new_fields = [adv_cfg.output_field]
         if returns is not None:
             fields_to_put[adv_cfg.returns_field] = returns
@@ -3360,6 +3375,8 @@ def _advantage_input_fields(self) -> list[str]:
             adv_cfg.token_mask_field,
             adv_cfg.sample_mask_field,
             *adv_cfg.repeated_batch_fields,
+            adv_cfg.mask_sample_field,
+            adv_cfg.truncated_field,
         ]
         if self._message_level_advantage_penalties_enabled:
             fields.extend(

nemo_rl/algorithms/single_controller_utils/config.py

@@ -791,7 +791,6 @@ def _validate_algo_settings(master_config: MasterConfig) -> None:
     unsupported = [
         name
         for name, enabled in (
-            ("overlong_filtering", algo_cfg.overlong_filtering),
             ("use_dynamic_sampling", algo_cfg.use_dynamic_sampling),
             ("reward_scaling", algo_cfg.reward_scaling.enabled),
             ("reward_shaping", algo_cfg.reward_shaping.enabled),
@@ -1156,6 +1155,8 @@ class AdvantageConfig:
     sample_mask_field: str = "sample_mask"
     invalid_tool_call_mask_field: str = INVALID_TOOL_CALL_MASK
     malformed_thinking_mask_field: str = MALFORMED_THINKING_MASK
+    mask_sample_field: str = "mask_sample"
+    truncated_field: str = "truncated"
     repeated_batch_fields: list[str] = field(default_factory=list)
     policy_logprobs_field: str = "prev_logprobs"
     generation_logprobs_field: str = "generation_logprobs"

nemo_rl/algorithms/single_controller_utils/utils.py

@@ -96,6 +96,7 @@ def reduce_advantage_pump_metrics(
     sequence_lengths: list[int],
     *,
     seq_logprob_error_metrics: list[dict[str, float]] | None = None,
+    num_mask_sample_filtered: list[int] | None = None,
     num_invalid_tool_calls: list[int] | None = None,
     num_malformed_thinking: list[int] | None = None,
     num_assistant_messages: list[int] | None = None,
@@ -108,13 +109,16 @@ def reduce_advantage_pump_metrics(
         sequence_lengths: All input_lengths trained on this step.
         seq_logprob_error_metrics: Sequence-error metrics and their aggregation
             counts, one record per streaming chunk.
+        num_mask_sample_filtered: Environment-flagged sample counts, one per
+            streaming chunk.
         num_invalid_tool_calls: Per-sample invalid tool-call counts.
         num_malformed_thinking: Per-sample malformed-thinking counts.
         num_assistant_messages: Per-sample assistant message counts (rate denominator).
 
     Returns:
         Step-level reward, advantage, token-count, optional sequence
-        log-probability error metrics, and per-sample violation counts.
+        log-probability error metrics, the num_mask_sample_filtered count, and
+        per-sample violation counts.
 
     """
     out: dict[str, float] = {}
@@ -132,6 +136,8 @@ def reduce_advantage_pump_metrics(
             out["advantages/min"] = 0.0
     if sequence_lengths:
         out["total_num_tokens"] = float(sum(sequence_lengths))
+    if num_mask_sample_filtered is not None:
+        out["num_mask_sample_filtered"] = float(sum(num_mask_sample_filtered))
     if seq_logprob_error_metrics:
         out.update(_reduce_seq_logprob_error_metrics(seq_logprob_error_metrics))
     n_asst = sum(num_assistant_messages or [])

nemo_rl/data_plane/schema.py

@@ -28,6 +28,8 @@
 INPUT_IDS = "input_ids"
 INPUT_LENGTHS = "input_lengths"
 SAMPLE_MASK = "sample_mask"
+MASK_SAMPLE = "mask_sample"
+TRUNCATED = "truncated"
 META_IDX = "meta_idx"
 
 # Token-aligned message-violation fields consumed by SingleController advantages.
@@ -55,6 +57,8 @@
 # TransferQueue's lazy field-name registration race.
 SC_ROLLOUT_SCHEMA_FIELDS = (
     *DP_TRAIN_FIELDS,
+    MASK_SAMPLE,
+    TRUNCATED,
     "prompt_ids_for_adv",
     "total_reward",
     "values",
@@ -117,8 +121,10 @@
 PROMOTE_1D_FIELDS: frozenset[str] = frozenset(
     {
         INPUT_LENGTHS,
+        MASK_SAMPLE,
         "total_reward",
         SAMPLE_MASK,
+        TRUNCATED,
     }
 )
 

nemo_rl/experience/payload.py

@@ -27,7 +27,9 @@
 from nemo_rl.data_plane.schema import (
     INVALID_TOOL_CALL_MASK,
     MALFORMED_THINKING_MASK,
+    MASK_SAMPLE,
     ROUTED_EXPERTS_FIELD,
+    TRUNCATED,
 )
 from nemo_rl.distributed.batched_data_dict import BatchedDataDict
 from nemo_rl.experience.interfaces import PromptGroupRecord
@@ -102,8 +104,9 @@ def record_to_train_batch(
             flags for configured advantage penalties.
 
     Returns:
-        BatchedDataDict with input_ids, input_lengths, generation_logprobs, token_mask,
-        sample_mask, prompt_ids_for_adv, total_reward, violation counts, and optional
+        BatchedDataDict with input_ids, input_lengths, generation_logprobs,
+        token_mask, an all-ones sample_mask, the raw mask_sample and truncated
+        flags, prompt_ids_for_adv, total_reward, violation counts, and optional
         routed experts and message-violation masks.
     """
     # Lazy imports: grpo and llm_message_utils transitively pull
@@ -113,7 +116,10 @@ def record_to_train_batch(
         extract_initial_prompt_messages,
     )
     from nemo_rl.data.llm_message_utils import batched_message_log_to_flat_message
-    from nemo_rl.experience.rollouts import backfill_missing_routed_experts
+    from nemo_rl.experience.rollouts import (
+        _mask_sample_flags,
+        backfill_missing_routed_experts,
+    )
 
     completions = record.completions
     n = len(completions)
@@ -146,6 +152,8 @@ def record_to_train_batch(
     total_reward = torch.tensor(
         [float(c.reward) for c in completions], dtype=torch.float32
     )
+    mask_sample = _mask_sample_flags(c.env_extras for c in completions)
+    truncated = torch.tensor([c.truncated for c in completions], dtype=torch.bool)
     sample_mask = torch.ones(n, dtype=torch.float32)
 
     train_data: dict[str, Any] = {
@@ -155,6 +163,8 @@ def record_to_train_batch(
         "token_mask": flat["token_loss_mask"],
         "sample_mask": sample_mask,
         "prompt_ids_for_adv": prompt_flat["token_ids"],
+        MASK_SAMPLE: mask_sample,
+        TRUNCATED: truncated,
         "total_reward": total_reward,
         _VIOLATION_COUNTS_KEY: violation_counts,
     }

nemo_rl/experience/rollouts.py

@@ -21,7 +21,7 @@
 import statistics
 import warnings
 from collections import defaultdict
-from collections.abc import AsyncGenerator, Mapping, Sequence
+from collections.abc import AsyncGenerator, Iterable, Mapping, Sequence
 from dataclasses import dataclass
 from typing import Any, Optional
 
@@ -49,6 +49,7 @@
     attach_image_model_inputs_to_message,
     extract_input_images_from_responses_messages,
 )
+from nemo_rl.data_plane.schema import MASK_SAMPLE
 from nemo_rl.distributed.batched_data_dict import BatchedDataDict
 from nemo_rl.environments.interfaces import (
     EnvironmentInterface,
@@ -279,16 +280,12 @@ def _add_r3_fallback_metrics(
     )
 
 
-def _extract_mask_sample_flags(results: list[dict[str, Any]]) -> torch.Tensor:
+def _mask_sample_flags(extras: Iterable[dict[str, Any] | None]) -> torch.Tensor:
     """Return True for samples the environment asks GRPO to mask from loss."""
     return torch.tensor(
         [
-            bool(
-                (result["full_result"].get("instance_config") or {}).get(
-                    "mask_sample", False
-                )
-            )
-            for result in results
+            bool(((extra or {}).get("instance_config") or {}).get(MASK_SAMPLE, False))
+            for extra in extras
         ],
         dtype=torch.bool,
     )
@@ -2848,10 +2845,12 @@ def _postprocess_single_nemo_gym_group(
             ),
         }
     )
-    # Env/agent mask flag: flagged samples are dropped from the loss but still
-    # count for advantages. env.should_mask_flagged_samples=false skips this.
+    # Carry the raw env/agent flag downstream; the advantage stage composes it
+    # into sample_mask. env.should_mask_flagged_samples=false skips this.
     if mask_env_flagged_samples:
-        final_batch["mask_sample"] = _extract_mask_sample_flags(results)
+        final_batch[MASK_SAMPLE] = _mask_sample_flags(
+            result["full_result"] for result in results
+        )
 
     rollout_metrics.update(_effort_shaping_metrics(shaping))
 

tests/unit/data_plane/test_codec_mooncake.py

@@ -75,6 +75,28 @@ def test_promote_1d_roundtrip_via_from_wire() -> None:
     assert torch.equal(back["input_lengths"], original)
 
 
+@pytest.mark.parametrize("field_name", ["mask_sample", "truncated"])
+def test_raw_sample_filter_fields_roundtrip_as_dense_1d(field_name: str) -> None:
+    """Raw loss-filter fields use the Mooncake scalar wire workaround."""
+    from tensordict import TensorDict
+
+    from nemo_rl.data_plane.adapters.transfer_queue import (
+        _from_wire,
+        _promote_1d_leaves,
+    )
+
+    n = 4
+    original = torch.tensor([False, True, False, True])
+    td = TensorDict({field_name: original}, batch_size=[n])
+
+    wire = _promote_1d_leaves(td)
+    assert wire[field_name].shape == (n, 1)
+
+    back = _from_wire(wire)
+    assert back[field_name].shape == (n,)
+    assert torch.equal(back[field_name], original)
+
+
 def test_from_wire_densifies_uniform_nested_rows() -> None:
     """TQ v0.1.9's uniform nested reads are restored to dense tensors."""
     from tensordict import TensorDict