feat(grpo): train all Gym traces per logical rollout

[esarafian]

What does this PR do ?

Add opt-in synchronous Megatron GRPO training for all independently conditioned Gym traces in a logical rollout, including branches and compacted contexts.

Compute GRPO advantages once per original task's logical rollouts, then expand to physical training rows. Preserve captured token IDs, behavior logprobs, masks, and unique sampled-span ownership. The objective averages unique eligible sampled tokens; physical rows share one optimizer update, while the learning-rate scheduler advances by logical rollout count. Zero-loss padding accommodates unequal trace counts without dropping trailing traces. Overlong traces are masked individually, and invalid rollouts are excluded from the logical baseline. The Gym actor assigns capture identity before synchronous dispatch, preserves supplied IDs, and uses Gym's canonical correlation rules for finalization. Zero-call captures become masked inert rows; missing or malformed correlation remains an error. The existing Gym path remains the default.

Issues

No automatic issue closure. This implements a bounded synchronous multi-trace path; asynchronous/TQ/staged delivery and process-level credit assignment remain outside this change.

Usage

Apply these overrides to a working synchronous Gym Megatron recipe:

grpo:
  gym_multi_trace: true
  num_prompts_per_step: 2
  num_generations_per_prompt: 4
  async_grpo:
    enabled: false
data_plane: null
policy:
  train_global_batch_size: 8  # logical rollouts, before trace expansion
  train_micro_batch_size: 1
  megatron_cfg:
    enabled: true
loss_fn:
  token_level_loss: true
  sequence_level_importance_ratios: false
env:
  should_use_nemo_gym: true
  should_mask_flagged_samples: true
  nemo_gym:
    token_id_capture:
      enabled: true
      all_agents: true
      delivery: all_traces
      builder: per_request  # or prefix_merging
      dir: /shared/gym-capture

Merge dependency: NVIDIA-NeMo/Gym#3212. The RL change should merge after that companion PR. The Gym submodule is pinned to its implementation at 407fcf9f, which supplies the training_traces v1 envelope. uv.lock reflects that revision's dependency requirements. See the multi-trace guide for identity, masking, objective, and supported configuration details.

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? Selected unit checks and real two-step Pi + ReasoningGym / Qwen3-1.7B training runs passed for both builders; exact scope is below.
• Did you add or update any necessary documentation? The feature guide and configuration exemplar/reference were updated.

Additional Information

Validation on the matching NeMo RL nightly environment:

• 21 helper tests passed through normal repository pytest with its standard fixtures.
• The final selected suite passed 98 tests: trace custody/grouping/padding and loss/gradient equivalence, capability guards, Gym capture finalization through the RL actor for both builders, actual rollout dispatch across batches with reused task indices, supplied IDs/attempts, empty-capture masking, existing Gym utilities, and configuration migration. This supplemental run used the baked Gym actor interpreter plus driver test dependencies, with --noconftest -o addopts="" to avoid starting or stopping Ray during the separate runtime experiment. It is not a full CI shard result.
• Ruff formatting/checks and full Pyrefly passed (0 errors shown; existing configured ignores retained). git diff --check and uv lock --check passed.
• Actual Pi + ReasoningGym / Qwen3-1.7B training passed for both per_request and prefix_merging on RL a425d8b3c0bb991caeb40cf31ad6dd995d373df6 with Gym 407fcf9f55e31bcd62290e6a3050f01ef8dbba3f. Each run completed two updates, saved a checkpoint, advanced Adam to step 2 and the scheduler by 16 logical rollouts, and changed sampled attention weights.
• Each run contained 16 logical rollouts and 32 physical traces: a bash tool-call response and a subsequent assistant answer per rollout. A raw-capture-to-training-row audit matched all owned token IDs and behavior logprobs, effective masks, and the worker loss denominator: 9,232 tokens for per_request and 9,801 for prefix_merging, with no excluded, invalid, overlong, or padding rows. All had positive importance-sampling weights. The tool-call portion contributed 672 eligible tokens in each run; assistant answers contributed 8,560 and 9,129 respectively.
• The first update had zero advantages. In the second, 2,345 per_request tokens and 2,320 prefix_merging tokens had nonzero advantages, including 168 tool-call tokens in each. Token inclusion and positive weights do not assert nonzero gradient for every token; PPO clipping and zero advantages remain distinct.
• Pi's prompt normalization yielded two independent roots per rollout in these live runs, so both builders retained two traces. Actual prefix joining and equivalent split/merged loss gradients are covered separately by tests; no live row-count reduction or model-quality improvement is claimed.
• The final Gym dependency pin was revalidated with the same 98 selected tests and uv lock --check; no lockfile changes were required.

Setup explicitly rejects unsupported asynchronous/TQ/staged capture, non-Megatron or multimodal policies, alternative advantage estimators, sequence-level loss/ratios, active seq-mask-tis, and legacy filtering/penalty options that would change the logical multi-trace semantics. Outcome reward broadcast is not process credit assignment or session-normalized weighting.

Public CI status: all displayed checks on the final head passed, but the NVIDIA main test workflow remains behind maintainer approval and the required CI:L2 label. Our label request was denied by repository permissions. The official Claude review workflow was triggered and failed before reviewing because its app-token exchange requires repository write access; no review findings were produced. Local tests and independent review do not substitute for these pending official checks.

[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.

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: 98d5449 (PR #4078 from esarafian/polar-rollout-training)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: accc9cd (PR #4078 from esarafian/polar-rollout-training)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: bff7e4a (PR #4078 from esarafian/polar-rollout-training)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

[esarafian]

/ok to test a425d8b

[esarafian]

/claude review

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: a425d8b (PR #4078 from esarafian/polar-rollout-training)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨