Skip to content

feat: multi-reward tool-call environment and reward_components for GDPO - #1525

Merged
cmunley1 merged 9 commits into
NVIDIA-NeMo:mainfrom
anjalibshah:gdpo-multireward-env-upstream
Jun 12, 2026
Merged

feat: multi-reward tool-call environment and reward_components for GDPO#1525
cmunley1 merged 9 commits into
NVIDIA-NeMo:mainfrom
anjalibshah:gdpo-multireward-env-upstream

Conversation

@anjalibshah

Copy link
Copy Markdown
Contributor

Summary

Adds multi-reward support for GDPO (arXiv:2601.05242) on the NeMo Gym side.

  • nemo_gym/base_resources_server.py — optional reward_components: dict[str, float] | None = None on BaseVerifyResponse. Backward-compatible (defaults to None); single-reward environments are unaffected.
  • resources_servers/tool_call_multireward/ — a new environment that asks the model to call get_weather and scores the response on three independent {0,1} components — correctness, schema_valid, format — returned in reward_components (with reward set to their sum). Includes dataset generator, example data, config, README, and tests.

Why

GRPO sums reward components and normalizes the sum, so responses with the same total but different composition collapse to an identical advantage. GDPO normalizes each component independently and keeps them distinct. This environment exposes exactly that structure (e.g. {0,1,1} and {1,1,0} both total 2.0). Downstream, NeMo-RL consumes reward_components as reward1…rewardN for the GDPO advantage estimator.

Test plan

resources_servers/tool_call_multireward/tests/test_app.py — 7 scoring cases including the same-total/different-composition case. Verified green (8/8) via the NeMo Gym test runner.

Add an optional reward_components field to BaseVerifyResponse so environments can
return decoupled per-component scores for multi-reward RL (GDPO,
https://arxiv.org/abs/2601.05242). Add the tool_call_multireward environment, which
grades a single get_weather tool call on three independent components (correctness,
schema_valid, format) and returns them in reward_components with reward set to their
sum. Includes dataset generator, example data, config, README, and unit tests.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Anjali Shah <anjshah@nvidia.com>
@copy-pr-bot

copy-pr-bot Bot commented Jun 3, 2026

Copy link
Copy Markdown

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.

anjalibshah and others added 2 commits June 3, 2026 17:35
…tting

Move reward_components off BaseVerifyResponse onto ToolCallMultiRewardVerifyResponse
so other environments' verify responses are unchanged (the base-class field was
serializing reward_components: null into every response and breaking other servers'
exact-match tests). Regenerate example.jsonl via create_examples.py and apply ruff
format/lint. NeMo-RL's bridge still reads reward_components from the verify result.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Anjali Shah <anjshah@nvidia.com>
Add example_metrics.json and example_rollouts.jsonl (5 rollouts) generated by
running the 5 example prompts through the agent (Qwen2.5-7B-Instruct via a local
vLLM server), plus a data/.gitignore matching the other example environments. This
satisfies ng_test's example-data validation (ng_test +should_validate_data=true).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Anjali Shah <anjshah@nvidia.com>
@anjalibshah

Copy link
Copy Markdown
Contributor Author

Note on the example-validation artifacts

data/example_metrics.json and data/example_rollouts.jsonl (5 rollouts) were generated by running the 5 example prompts through the agent against a local vLLM server (the canonical NeMo Gym serving path), not an external endpoint:

  • Model: Qwen/Qwen2.5-7B-Instruct served via vllm serve with --enable-auto-tool-choice --tool-call-parser hermes (exposes /v1/responses + tool calling).
  • ng_run with responses_api_models/openai_model + this env's config pointed at the local server, then ng_collect_rollouts over data/example.jsonl.
  • example_metrics.json was regenerated from the same model (ng_prepare_data +mode=example_validation) so metrics and rollouts are consistent.

Validated locally:

ng_test +entrypoint=resources_servers/tool_call_multireward +should_validate_data=true
# 8 passed; "The data for resources_servers/tool_call_multireward has been successfully validated!"

Reviewers/maintainers can regenerate these against the standard serving stack if a different reference model is preferred — the env scoring logic is model-agnostic.

Regenerated via scripts/update_env_list.py (the update-readme-table pre-commit hook)
so the new multi-reward environment appears in the environment list.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Anjali Shah <anjshah@nvidia.com>
Comment thread resources_servers/tool_call_multireward/configs/tool_call_multireward.yaml Outdated
anjalibshah and others added 2 commits June 4, 2026 18:19
Per review feedback, prefix the environment with example_ to match the repo
convention for illustrative envs (example_single_tool_call, example_multi_step, ...),
since it's a reference/test-bed for the reward_components convention rather than a
training/eval dataset. Renames the directory, config + server/agent keys, dataset
paths, the agent reference inside the example data artifacts, and regenerates the
README environment catalog (now listed under examples).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Anjali Shah <anjshah@nvidia.com>
Add the value field (per review) for catalog consistency with other resources servers.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Anjali Shah <anjshah@nvidia.com>

@cmunley1 cmunley1 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

looks good!

@cmunley1

cmunley1 commented Jun 9, 2026

Copy link
Copy Markdown
Contributor

Is there any docs how to use it with NeMo RL or is that a future thing?

@cmunley1

Copy link
Copy Markdown
Contributor

/ok to test 1146e5d

@cmunley1

Copy link
Copy Markdown
Contributor

/ok to test 603a0a7

@cmunley1
cmunley1 merged commit 0825b44 into NVIDIA-NeMo:main Jun 12, 2026
26 of 29 checks passed
wprazuch pushed a commit that referenced this pull request Jun 22, 2026
…PO (#1525)

## Summary
Adds multi-reward support for GDPO
([arXiv:2601.05242](https://arxiv.org/abs/2601.05242)) on the NeMo Gym
side.

- **`nemo_gym/base_resources_server.py`** — optional `reward_components:
dict[str, float] | None = None` on `BaseVerifyResponse`.
Backward-compatible (defaults to `None`); single-reward environments are
unaffected.
- **`resources_servers/tool_call_multireward/`** — a new environment
that asks the model to call `get_weather` and scores the response on
three independent `{0,1}` components — `correctness`, `schema_valid`,
`format` — returned in `reward_components` (with `reward` set to their
sum). Includes dataset generator, example data, config, README, and
tests.

## Why
GRPO sums reward components and normalizes the sum, so responses with
the same total but different composition collapse to an identical
advantage. GDPO normalizes each component independently and keeps them
distinct. This environment exposes exactly that structure (e.g.
`{0,1,1}` and `{1,1,0}` both total `2.0`). Downstream, NeMo-RL consumes
`reward_components` as `reward1…rewardN` for the GDPO advantage
estimator.

## Test plan
`resources_servers/tool_call_multireward/tests/test_app.py` — 7 scoring
cases including the same-total/different-composition case. Verified
green (8/8) via the NeMo Gym test runner.

---------

Signed-off-by: Anjali Shah <anjshah@nvidia.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: Christian Munley <cmunley@nvidia.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants