feat: configurable GDPO per-reward weights and multi-reward NeMo Gym bridge by anjalibshah · Pull Request #2680 · NVIDIA-NeMo/RL

anjalibshah · 2026-06-03T23:15:19Z

Summary

Two GDPO (arXiv:2601.05242) capability additions on top of the existing GDPO support (#2069).

1. Configurable per-reward weights. GDPO's aggregation A = Σ wₙ·Aₙ is now tunable via grpo.adv_estimator.reward_weights (one entry per component, ordered to match reward1, reward2, …). Defaults to equal weights (all 1.0) for backward compatibility; a wrong-length list raises ValueError. Files: nemo_rl/algorithms/advantage_estimator.py, nemo_rl/algorithms/grpo.py, docs/guides/grpo.md.

2. Multi-reward NeMo Gym bridge. run_async_nemo_gym_rollout now surfaces per-component rewards (reward1…rewardN) when a Gym verifier returns reward_components, mirroring the native multi-reward rollout path. Single-reward environments are unaffected (fall back to the scalar reward). Files: nemo_rl/environments/nemo_gym.py (new extract_reward_components), nemo_rl/experience/rollouts.py, examples/nemo_gym/gdpo_multireward.yaml.

Dependency

The bridge consumes the reward_components field added in NVIDIA-NeMo/Gym#1525 — that should land first. Until it does, the bridge code is inert (single-reward fall-back), so this PR is safe to merge independently.

Test plan

Validated in the NeMo-RL container (8×H100):

tests/unit/algorithms/test_grpo.py::test_gdpo_advantage_estimator_reward_weights — PASSED
tests/unit/environments/test_nemo_gym.py::test_extract_reward_components — PASSED

copy-pr-bot · 2026-06-03T23:15:23Z

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.

ZhiyuLi-Nvidia

Generally looks good to me. Only 2 comments here.

ZhiyuLi-Nvidia · 2026-07-15T22:46:14Z

/ok to test cf64525

ZhiyuLi-Nvidia · 2026-07-15T22:46:41Z

Also could you rebase on to main branch after resolving conflicts?

…bridge Add reward_weights to the GDPO advantage estimator so components aggregate as A = sum_n w_n * A_n (defaults to equal weights; wrong length raises). Extend the NeMo Gym bridge to surface per-component rewards (reward1..rewardN) from a verifier that returns reward_components, mirroring the native multi-reward rollout path, so GDPO can train on NeMo Gym environments. Includes unit tests and docs. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Anjali Shah <anjshah@nvidia.com>

…ard env Name the ready-made multi-reward reference env (NVIDIA-NeMo/Gym) so users know exactly what to set env.nemo_gym to for an end-to-end GDPO multi-reward run. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Anjali Shah <anjshah@nvidia.com>

anjalibshah · 2026-07-16T17:13:03Z

Also could you rebase on to main branch after resolving conflicts?

yes done!

anjalibshah · 2026-07-16T17:15:11Z

Generally looks good to me. Only 2 comments here.

Thanks @ZhiyuLi-Nvidia. Resolved conflicts and rebased. Please merge, thanks.

ZhiyuLi-Nvidia · 2026-07-16T17:20:49Z

Found DCO failure.

Could you finally git commit / amend -s to sign the commits?

…ate reward == sum(components) Addresses review on NVIDIA-NeMo#2680: - Stop silently overwriting final_batch["total_reward"] with the component sum. Keep the verifier's scalar `reward`. When reward_components is present the contract is reward == sum(components), so overwriting is a no-op in the correct case and only masks a misconfigured verifier otherwise. Validate the contract instead and raise on a real mismatch. - Expand the sorted() comment to explain it is a canonical ordering (stable reward{n} -> component mapping for the GDPO per-component baseline and reward_weights), not just deduplication. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Anjali Shah <anjshah@nvidia.com>

ZhiyuLi-Nvidia · 2026-07-17T17:32:16Z

/ok to test cb7f6b5

- Move the reward/component-sum consistency check into validate_reward_components_match_scalar() in nemo_gym.py, next to extract_reward_components, so it is unit-testable without the full rollout machinery. run_async_nemo_gym_rollout now calls the helper. - Add test_validate_reward_components_match_scalar covering the matching, float-tolerance, single-reward-skip, and mismatch-raises cases. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Anjali Shah <anjshah@nvidia.com>

anjalibshah · 2026-07-17T17:54:12Z

/ok to test de41394

ZhiyuLi-Nvidia · 2026-07-17T18:08:56Z

/ok to test de41394

yuki-97

thanks @anjalibshah for supporting the two features! left some comments.

…eward2) main's get_gdpo_reward_component_keys() selects batch keys by `k.startswith("reward/")`, but the NeMo Gym bridge emitted positional `reward1, reward2, ...`. Those never match, so GDPO saw zero components and raised "GDPO requires multiple reward components". Emit `reward/<name>` instead, matching the native multi-reward path. - rollouts.py: emit final_batch[f"reward/{name}"] (component name carries identity; get_gdpo_reward_component_keys sorts, so no positional index). - test_grpo.py: reward_weights test now uses reward/correctness, reward/format. - Reword reward_weights ordering from "reward1, reward2, ..." to "alphabetical by component name" across advantage_estimator.py, grpo.py, grpo.md, and gdpo_multireward.yaml; update nemo_gym.py docstring. - Remove dead private-fork link (gym-multireward-scope.md) from the recipe yaml. Addresses review feedback on NVIDIA-NeMo#2680. Signed-off-by: Anjali Shah <anjshah@nvidia.com>

ZhiyuLi-Nvidia · 2026-07-20T21:56:09Z

/ok to test fce0393

Guards the reward_weights aggregation path added to GDPOAdvantageEstimator (per review on #2680). Adds a DTensor GDPO recipe over the native math_multi_reward env (correctness / format / integer components), Qwen2.5-1.5B on GSM8K, 1n8g, with reward_weights=[1.0, 0.5, 0.25]. Follows the nightly pattern from #1866: - examples/configs/recipes/llm/gdpo-qwen2.5-1.5b-1n8g-fsdp2-reward-weights.yaml - tests/test_suites/llm/gdpo-qwen2.5-1.5b-1n8g-fsdp2-reward-weights.sh - register in tests/test_suites/nightly.txt - add "gdpo" -> gdpo_math_1B.yaml to ALGO_MAPPING_TO_BASE_YAML so the recipe name-prefix check passes. Signed-off-by: Anjali Shah <anjshah@nvidia.com>

Guards the reward_weights aggregation path added to GDPOAdvantageEstimator (per review on NVIDIA-NeMo#2680). Adds a DTensor GDPO recipe over the native math_multi_reward env (correctness / format / integer components), Qwen2.5-1.5B on GSM8K, 1n8g, with reward_weights=[1.0, 0.5, 0.25]. Follows the nightly pattern from NVIDIA-NeMo#1866: - examples/configs/recipes/llm/gdpo-qwen2.5-1.5b-1n8g-fsdp2-reward-weights.yaml - tests/test_suites/llm/gdpo-qwen2.5-1.5b-1n8g-fsdp2-reward-weights.sh - register in tests/test_suites/nightly.txt - add "gdpo" -> gdpo_math_1B.yaml to ALGO_MAPPING_TO_BASE_YAML so the recipe name-prefix check passes. Signed-off-by: Anjali Shah <anjshah@nvidia.com>

ZhiyuLi-Nvidia · 2026-07-20T22:09:14Z

/ok to test 9143bf2

terrykong

Thanks for the quick turnaround on the review feedback — both follow-up commits were verified locally at head fce0393c:

d00b431 fully resolves the reward/<name> key-scheme issue: the previously-failing test_gdpo_advantage_estimator_reward_weights now passes, a repo-wide search finds no remaining reward1/reward2 references, all ordering docs/comments/error messages are reworded consistently, and the dead skill link is gone. Nice, thorough fix.
fce0393 adds a well-plumbed native-path nightly (recipe resolution, naming registration, and alphabetical weight→component mapping all verified). Two follow-ups inline: the recipe yaml currently fails the minimize-check pre-commit hook (CI blocker, one-command fix), and the Gym bridge block itself still has no test coverage since the nightly runs the native math_multi_reward path.

This review was produced by a 5-agent team (RL expert, test agent, bug finder, comment reviewer, devil's advocate); all findings were independently verified, and all other pre-commit hooks (ruff, ruff-format, pyrefly) pass at head.

Relationship to #2681 (same author): no code dependency, but the demo script there feeds reward1/reward2 keys into GDPOAdvantageEstimator, so with the key scheme now settled as reward/<name> it will need the matching update.

Generated by Claude Code

terrykong · 2026-07-21T07:00:55Z

merged. @anjalibshah agreed offline to address the reviews in #2681 regarding:

feat: configurable GDPO per-reward weights and multi-reward NeMo Gym bridge #2680 (comment)
feat: configurable GDPO per-reward weights and multi-reward NeMo Gym bridge #2680 (comment)

…-experiment skill Per review on #2681: experiment-specific content doesn't fit the repo's skill convention (existing skills are general infra/process workflows). Remove the skills/nemo-rl-gdpo-experiment/ directory and fold the genuinely additive pieces into the existing GDPO section of docs/guides/grpo.md: - matched GDPO-vs-GRPO comparison workflow (points at the examples/ control config and CPU collapse demo, both unchanged) - what to measure (per-reward convergence, per-prompt advantage spread) - practical notes: reward_scaling applies per component; final batch normalization always runs regardless of normalize_rewards Deliberately excludes reward_weights guidance to avoid overlap with #2680, which documents that option. Env-authoring contract is left in docs/guides/environments.md where it already lives. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Anjali Shah <anjshah@nvidia.com>

`spinup_server` is not a field Gym has ever read. It is absent from nemo_gym/server_utils.py at both the pin it was written against (d67ad6611, via #2680) and the current pin (c3bac9631), and nothing on the config MRO overrides pydantic's default extra="ignore" -- the only model_config in the chain is ConfigDict(arbitrary_types_allowed=True) at server_utils.py:295 -- so the key was silently discarded on every load. What actually decides whether a server is launched is the presence of an entrypoint (cli/env.py:407, unchanged since d67ad6611:188). These thin clients carry one, so they do spin up; the dry run simply returns early at server_utils.py:698 before serving. The conclusion in the previous comment was right -- validation is unconditional -- but the reason given described a mechanism that does not exist, which made a dead key read as load-bearing. Drop the six occurrences and reword the comment to the real mechanism. Upstream Gym's own resources_servers/jailbreak_detection/configs/safety_judge_model.yaml defines the same server without the key, and it has no reader anywhere in this repo. No behaviour change: tests/unit/environments/test_nemo_gym_prefetch_configs.py still reports 7 passed, 2 skipped against the edited configs, which is the expected result if the key was genuinely inert. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

anjalibshah requested review from a team as code owners June 3, 2026 23:15

github-actions Bot added the Documentation Improvements or additions to documentation label Jun 3, 2026

anjalibshah mentioned this pull request Jun 3, 2026

feat: configurable GDPO per-reward weights and multi-reward NeMo Gym bridge anjalibshah/RL#2

Closed

This was referenced Jun 23, 2026

Support multi-objective (GDPO) rewards for NeMo Gym environments #2892

Open

feat: shared MultiRewardVerifyResponse (reward_components) as a trainer-agnostic multi-reward contract NVIDIA-NeMo/Gym#1664

Closed

ZhiyuLi-Nvidia reviewed Jul 15, 2026

View reviewed changes

Comment thread nemo_rl/experience/rollouts.py Outdated

ZhiyuLi-Nvidia reviewed Jul 15, 2026

View reviewed changes

Comment thread nemo_rl/experience/rollouts.py Outdated

ZhiyuLi-Nvidia reviewed Jul 15, 2026

View reviewed changes

ZhiyuLi-Nvidia added the CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) label Jul 15, 2026

copy-pr-bot Bot temporarily deployed to public July 15, 2026 22:46 Inactive

copy-pr-bot Bot deployed to test July 15, 2026 22:49 Active

copy-pr-bot Bot temporarily deployed to public July 15, 2026 22:50 Inactive

ZhiyuLi-Nvidia added the CI:docs Run doctest label Jul 15, 2026

anjalibshah and others added 2 commits July 16, 2026 10:09

anjalibshah force-pushed the gdpo-weights-and-gym-multireward-upstream branch from 40fc0fa to fe17cee Compare July 16, 2026 17:11

ZhiyuLi-Nvidia previously approved these changes Jul 16, 2026

View reviewed changes

anjalibshah mentioned this pull request Jul 16, 2026

feat: GDPO vs GRPO multi-reward example and collapse demo #2681

Merged

anjalibshah force-pushed the gdpo-weights-and-gym-multireward-upstream branch from fe17cee to cb7f6b5 Compare July 16, 2026 17:26

anjalibshah dismissed ZhiyuLi-Nvidia’s stale review via de41394 July 17, 2026 17:52

anjalibshah force-pushed the gdpo-weights-and-gym-multireward-upstream branch from cb7f6b5 to de41394 Compare July 17, 2026 17:52

yuki-97 removed the CI:docs Run doctest label Jul 18, 2026

yuki-97 reviewed Jul 18, 2026

View reviewed changes

Comment thread nemo_rl/experience/rollouts.py Outdated

Comment thread examples/nemo_gym/gdpo_multireward.yaml Outdated

Comment thread docs/guides/grpo.md

anjalibshah force-pushed the gdpo-weights-and-gym-multireward-upstream branch from fce0393 to 9143bf2 Compare July 20, 2026 22:06

terrykong reviewed Jul 20, 2026

View reviewed changes

Comment thread tests/test_suites/nightly.txt

Comment thread examples/configs/recipes/llm/gdpo-qwen2.5-1.5b-1n8g-fsdp2-reward-weights.yaml

Comment thread nemo_rl/experience/rollouts.py

terrykong enabled auto-merge (squash) July 21, 2026 07:00

terrykong approved these changes Jul 21, 2026

View reviewed changes

terrykong merged commit 9f701f0 into NVIDIA-NeMo:main Jul 21, 2026
78 of 79 checks passed

yfw mentioned this pull request Aug 4, 2026

feat: Multimodal nemo gym compatible grpo pipeline #3414

Merged

4 tasks

Uh oh!

Conversation

anjalibshah commented Jun 3, 2026

Summary

Dependency

Test plan

Uh oh!

copy-pr-bot Bot commented Jun 3, 2026

Uh oh!

Uh oh!

Uh oh!

ZhiyuLi-Nvidia left a comment

Choose a reason for hiding this comment

Uh oh!

ZhiyuLi-Nvidia commented Jul 15, 2026

Uh oh!

ZhiyuLi-Nvidia commented Jul 15, 2026

Uh oh!

anjalibshah commented Jul 16, 2026

Uh oh!

anjalibshah commented Jul 16, 2026

Uh oh!

ZhiyuLi-Nvidia commented Jul 16, 2026

Uh oh!

ZhiyuLi-Nvidia commented Jul 17, 2026

Uh oh!

anjalibshah commented Jul 17, 2026

Uh oh!

ZhiyuLi-Nvidia commented Jul 17, 2026

Uh oh!

yuki-97 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ZhiyuLi-Nvidia commented Jul 20, 2026

Uh oh!

ZhiyuLi-Nvidia commented Jul 20, 2026

Uh oh!

terrykong left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

terrykong commented Jul 21, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants