feat(grpo): add wall-clock time-efficiency reward for NeMo-Gym rollouts

[yfw]

What does this PR do ?

Adds an opt-in wall-clock time-efficiency reward (grpo.time_efficiency) for NeMo-Gym agentic (SWE) rollouts, ported from the internal sdd/swe-opencode-superv35 time-efficiency experiments.

reward_i = reward_i
           + lambda_call_bonus * min(calls_i / call_bonus_ref, 1)   # 0 by default
           - lambda_time * (openhands_run_time_i / 60)

openhands_run_time is the wall time of the agent container from launch to exit, reported by the Gym swe_agents server (it includes apptainer spin-up and excludes final evaluation and Ray queueing). With the default lambda_time = 1/60 a 60-minute rollout costs exactly 1.0. apply_to selects whether failures are charged too ("all", the original semantics) or only resolved rollouts whose reward survived the reward-zeroing penalties ("correct"); lambda_call_bonus > 0 (with call_bonus_ref) adds the per-tool-call arm from the internal experiments: a saturating bonus for well-formed function_call items that pays for exploring, following the same apply_to gate (use it with "correct", since paying it on failures rewards junk calls). floor optionally clamps the adjusted reward. time_efficiency/{minutes,deduction}/{mean,max}, calls/mean, bonus/mean, seconds_per_call, bonus_saturated_frac, floored_frac and group_has_signal are logged next to the reward metrics.

• nemo_rl/utils/time_efficiency.py: TimeEfficiencyConfig (pydantic BaseModel) and apply_time_efficiency_reward, mirroring how length_penalty lives in utils.
• Threaded explicitly through run_nemo_gym_rollout_sync, run_async_nemo_gym_rollout and _postprocess_single_nemo_gym_group; runs after the other reward shapers (effort shaping, reward-zeroing penalties, length penalties), which all assume a binary env reward, so the continuous deduction composes with them.
• Wired at the training rollout call sites (grpo_train, AsyncTrajectoryCollector, SyncRolloutActor). Validation keeps the raw env reward on both the grpo.validate() path and the single-controller validate_sync path (rollout_to_tq(..., is_validation=True)); the source branch also applied it during validation.
• Rollouts the env flags mask_sample (e.g. swe_agents timeouts) are not charged while env.should_mask_flagged_samples is on: they are dropped from the loss, so charging them could only drag their siblings' group baseline down. With the flag off they train and are charged.
• Enabling it without the NeMo-Gym path raises at setup (mirrors the reward_penalties guard), since only Gym agents emit openhands_run_time.
• GRPOConfig.time_efficiency field; documented in the grpo_math_1B exemplar and reference configs. Disabled by default, so existing runs are unaffected.

The SWE e2e config that uses this lives in the post-training pipeline repo rather than here.

Issues

N/A

Usage

uv run examples/nemo_gym/run_grpo_nemo_gym.py --config <swe-e2e-config.yaml> \
  ++grpo.time_efficiency.enabled=true \
  ++grpo.time_efficiency.lambda_time=0.016666667 \
  ++grpo.time_efficiency.apply_to=all \
  ++grpo.advantage_clip_low=-3 ++grpo.advantage_clip_high=3```

The advantage clip is strongly recommended: with `normalize_rewards` the small continuous term is divided by the group's reward spread and can explode.

Per-tool-call arm (the internal `per_tool_time_efficiency` defaults; `call_bonus_ref` was calibrated on the super-v3.5 base checkpoint and should be re-read from `time_efficiency/calls/mean` at step 1 for another checkpoint):

```bash
  ++grpo.time_efficiency.enabled=true ++grpo.time_efficiency.apply_to=correct \
  ++grpo.time_efficiency.lambda_call_bonus=0.10 ++grpo.time_efficiency.call_bonus_ref=56.0 \
  ++grpo.time_efficiency.floor=0.05

The advantage clip is strongly recommended: with normalize_rewards the small continuous term is divided by the group's reward spread and can explode.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests? tests/unit/utils/test_time_efficiency.py (22 tests), tests/unit/experience/test_time_efficiency_rollouts.py (ordering vs reward-zeroing penalties under both apply_to modes), and setup-guard tests in tests/unit/algorithms/test_grpo.py
• Did you run the unit tests and functional tests locally? Only the pure unit tests, ruff (v0.9.9), and the exemplar/reference config check ran locally (no Python 3.13.14/torch env available on the login node). The torch-dependent test and the rest of the suite need CI.
• Did you add or update any necessary documentation? The new block is documented in the exemplar YAML and on the BaseModel; no docs page changes.

Additional Information

Differences from the internal implementation (sdd/swe-opencode-superv35, grpo.time_efficiency in _postprocess_single_group). The reward formula, lambda_time, floor, and the handling of missing or malformed openhands_run_time are unchanged. What differs:

1. Loss-masked rollouts are not charged while env.should_mask_flagged_samples is on, in both apply_to modes. This matters for "all": internally a swe_agents timeout (flagged mask_sample, always the group's longest rollout) was charged the full deduction even though it is dropped from the loss, which pulled the group ba
   seline down for its siblings; here it keeps its raw 0.0. Under "correct" a timeout is unresolved and was never charged, so the only change is the rare resolved rollout that Gym flags for max iterations or context window, which is now exempt too. With the flag off, flagged rows train and are charged as before. For the SWE config
   as run, this is the one change that alters training rewards, and only in the "all" arm.
2. Ordering. The deduction runs after effort shaping, the reward-zeroing penalties and the length penalties instead of before them, because those assume a binary env reward. Internally a rollout zeroed by a penalty was reset to exactly 0.0 after being charged; here it is charged after zeroing. No effect on the SWE config, which enables none of those shapers.
3. "correct" gate. Requires resolved and a positive post-penalty reward; internally resolved alone. Identical for swe_agents unless a reward-zeroing penalty is enabled.
4. Validation is not charged. Internally the deduction also applied during validation, so val:total_reward/mean was accuracy minus the mean deduction. Here validation reports the raw env reward on both the grpo.validate() and single-controller validate_sync paths, so validation curves are not directly comparable to the internal runs.
5. Metric names: time_efficiency/{minutes,deduction}/{mean,max} (was minutes_mean, minutes_max, deduction_mean, deduction_max), so the async aggregator takes true maxima across groups. group_has_signal is keyed on the deductions of trainable rows rather than on raw wall times, so it reads 0 when nothing was actually charged.
6. Config and wiring. grpo.time_efficiency is a typed TimeEfficiencyConfig (an invalid apply_to fails at load) threaded explicitly through the rollout functions, instead of an untyped dict read from master_config inside the postprocess. Enabling it without the NeMo-Gym path raises at setup, mirroring the reward_penalties guard.
7. Docs. openhands_run_time is documented as spanning the agent container from launch to exit (it includes apptainer spin-up); the internal comments described it as the agent loop only.
8. Per-tool-call arm. The internal per_tool_time_efficiency block (call bonus plus time deduction) is folded into this same block as lambda_call_bonus / call_bonus_ref, off by default. call_bonus_ref is required when the bonus is on instead of defaulting to 56, and the bonus follows the same apply_to and mask_sample gates and the same last-in-the-postprocess ordering as the deduction. Metrics for it: calls/mean, bonus/mean, seconds_per_call, bonus_saturated_frac, floored_frac.

Net effect for comparing against the completed internal runs: training rewards differ only through item 1, and validation metrics through item 4.

• Credit to Mengru Wang, who wrote the original implementation and ran the experiments on the internal branch.
• Experiment notes from those runs: apply_to: all charged failures for elapsed time and taught the policy to quit early, costing accuracy on long problems; apply_to: correct is the follow-up arm.

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

[yfw]

Reviewed at 0ed0001 against super-v3.5-posttraining.

Merge safety. Nothing in this PR changes behaviour while time_efficiency.enabled: false, which is the default and the setting of every shipped recipe. apply_time_efficiency_reward returns {} before touching results; all three new parameters default to None and every caller of the three widened signatures is keyword-only, so ppo.py, distillation.py and grpo.validate() are unaffected; a resolved-config scan across both examples/configs/** and examples/nemo_gym/** found no recipe that sets the block; and the saved config.yaml is write-only, so a checkpoint predating this field resumes on the pydantic default.

Verification run against this head. The 20 new tests pass. tests/unit/test_config_v2.py passes (14), and removing the reference-config block makes it fail with missing=['grpo.time_efficiency'] — the pairing is load-bearing and present. tests/unit/utils + tests/unit/experience + tests/unit/algorithms + test_config_v2.py: 1363 passed; the 8 non-passing tests are all tests/unit/utils/test_native_checkpoint.py (NVIDIA driver ... too old (found version 12080) while building nv-grouped-gemm), pre-existing on the base branch and with no import path to anything this PR touches. A 22-mutation sweep over nemo_rl/utils/time_efficiency.py killed 21, including replacing the whole function body with return {} — the new tests have real teeth. The 2 pyrefly errors and the 3 files ruff reformats in this tree are pre-existing on the base branch (git blame dates them to Sept 2 and Sept 4), not from this PR.

Both CI ratchets are satisfied unprompted: nemo_rl/utils/time_efficiency.py is in pyrefly.toml project-includes, and the exemplar/reference-config edit is paired. The config shape also follows the right internal analog — a typed BaseModel sub-block like reward_scaling/reward_shaping — rather than length_penalty's untyped dumped-dict shape that the PR body cites; extra="allow" matches all five sibling reward configs and the config-conventions rule.

One inconsistency not posted inline because no shipped config reaches it: on the single-controller path, SyncRolloutActor.rollout_to_tq passes the block for both training and validate_sync, so the new comment at grpo.py:357-358 ("validation reports the raw env reward") holds for the legacy validate() but not for grpo_train_sync — the same way effort_levels and reward_penalties already behave there. Worth a word in that comment.

Three comments inline: one reachable reward inversion with a two-part fix, one factual correction to the module docstring, and one metric that misreports on the apply_to: "correct" arm.

Generated by Claude Code