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

[macandro96]

What does this PR do ?

Support #3163 for single controller path

Issues

List issues that this PR closes (syntax):

Usage

• You can potentially add a usage example below

# Add a code snippet demonstrating how to use this

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

• ...

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[tianyi-zhang-02]

Heads-up: I had #3786 open doing an overlapping thing, and yours came first and covers more — I've said so there and will cut mine down to the non-overlapping remainder rather than compete with it.

Two things from reading yours that might be useful:

The advantage_mask split. I checked whether keeping it over the full prompt group is load-bearing for GRPO, since my version narrowed sample_mask before the advantage stage. It isn't a value difference: GRPOAdvantageEstimator.compute_advantage passes torch.ones_like(rewards) as the valid-mask into calculate_baseline_and_std_per_prompt, and uses mask only for advantages.expand(mask.shape). So both arrangements give identical advantages. Yours is still the safer one — it stops depending on that internal detail — just noting it in case a reviewer asks whether the split is required.

One gap that survives this PR. The prompt-level DatumSpec.loss_multiplier is still dropped on the SC path. It's set to 0.0 when the prompt itself exceeded the length budget before generation ran, and record_to_train_batch never reads it — git grep loss_multiplier nemo_rl/experience/payload.py is empty on main, and this PR doesn't add it. So an over-long prompt still trains at full weight after this lands. Completion-independent, so it doesn't fit naturally in your mask_sample/truncated shape — happy to send it separately on top of yours. :)

[tianyi-zhang-02]

Correction to what I wrote above — I said #3766's coverage is a superset of this PR's. It isn't, and the difference matters in your favour, @macandro96.

I checked #3766's config.py hunk properly afterwards. It adds two AdvantageConfig field names and nothing else:

+    mask_sample_field: str = "mask_sample"
+    truncated_field: str = "truncated"

("overlong_filtering", algo_cfg.overlong_filtering) is still in the unsupported list at single_controller_utils/config.py:739 on main, and #3766 does not remove it. So on that branch a config with grpo.overlong_filtering: true raises NotImplementedError at setup and its truncation-filter path is never reached. Its env-flagged mask_sample path works; the truncated half is unreachable until something takes that entry off the list, which is what this PR does.

So the accurate split is:

	env mask_sample	truncated / overlong_filtering	num_mask_sample_filtered	prompt-level loss_multiplier
#3766	yes	implemented but gated off	yes	no
#3786 (this)	yes	yes, and unblocks the knob	no	yes

Which doesn't change my conclusion — #3766 came first, its column-based shape is the better one, and I'd rather not have two mechanisms for one rule. But the right outcome isn't "close this"; it's that #3766 needs the unsupported-list removal to make its own truncation filter live, and the prompt-level loss_multiplier is still dropped by both.

Happy to send those two as a small PR on top of #3766 once it lands, and close this. Or, if it's easier, take the two lines from here directly — no attribution needed. :)

[yfw]

Team review of the mask-composition change: 6 agents (RL expert, bug finder, test agent, design reviewer, comment reviewer, devil's advocate), all findings adversarially challenged before posting. Thanks for this — the env mask_sample gap on the SingleController path is a real parity hole and the column-based shape is the right way to close it.

Please rebase first. This branch is currently CONFLICTING and 35 commits behind main. Three of the six comments below are only about what happens after that rebase, and two of them describe changes that no merge conflict and no CI run will point at — I simulated the merge to check.

The description is the unfilled template; a sentence on which recipe you ran, and on whether the advantage-mask split is intentional, would resolve most of what's below.

Deliberately not asking for convergence curves: the one open question (the advantage-mask split) is semantic rather than empirical — a curve wouldn't tell us whether normalizing over dropped rows is intended. Worth answering yes/no first.

Verified clean and worth noting so nobody re-raises them: pre-commit run --all-files passes every hook; the mask_sample key path matches Gym and the legacy reader exactly; env.should_mask_flagged_samples=false is already honoured upstream at rollout_manager.py:1021 so no second config knob was needed; adding both fields to PROMOTE_1D_FIELDS was mandatory, not optional (the Mooncake adapter rejects undeclared dense-1D tensors) and the round-trip test pins it; num_mask_sample_filtered matches the legacy metric name and emission semantics; and AdvantageConfig is correctly a @dataclass of internal column names, so no exemplar-YAML update is owed.

Generated by Claude Code

[macandro96]

/ok to test 48121bb

[macandro96]

/ok to test 5c540c3

[macandro96]

/ok to test 7482514