Support multi-objective (GDPO) rewards for NeMo Gym environments
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
## Summary
A NeMo Gym verifier can already return decoupled reward components — the `example_tool_call_multireward` environment returns `reward_components: {name: score}` plus the summed scalar `reward`. But when such an environment is trained through NeMo-RL, those components are **silently dropped**: NeMo-RL's NeMo Gym rollout path in `nemo_rl/experience/rollouts.py` reads only the scalar `full_result["reward"]` (into `total_reward`) and ignores `reward_components` entirely. This gap should be addressed to support use of multi-objective training with NeMo Gym environments.
Why it happens: GDPO's advantage estimator consumes per-component keys (`reward1, reward2, ...`) from the training batch, and those keys are only produced by NeMo-RL's generic `EnvironmentInterface.step()` path (which returns a 2D `rewards` tensor). The NeMo Gym path bypasses `step()` entirely (`NemoGym.step()` is `NotImplementedError`), so nothing ever converts the verifier's components into `reward1..K`.
## Current behavior (RL v0.6)
- `nemo_rl/experience/rollouts.py` (NeMo Gym path): the final batch sets only `"total_reward": torch.tensor([r["full_result"]["reward"] for r in results])`. No `reward1..K` keys are emitted.
- `nemo_rl/environments/nemo_gym.py`: `NemoGym.step()` raises `NotImplementedError` ("NeMo-Gym will handle the rollouts entirely"); `_postprocess_nemo_gym_to_nemo_rl_result()` carries the verify result through as `full_result`.
- `nemo_rl/algorithms/advantage_estimator.py`: GDPO reads `reward1, reward2, ...` via `get_gdpo_reward_component_keys()` (requires >= 2 components).
- A repo-wide search finds no code that reads a `reward_components` field (only `num_reward_components`, a local variable, and a comment/doc reference).
- On the Gym side, the `example_tool_call_multireward` env's `verify()` returns `reward_components: {name: score}` plus `reward` (the sum), but those components never reach training.
## Proposed change
In the NeMo Gym rollout path, when `full_result` contains a `reward_components` dict (`{component_name: score}`):
- Map it to per-component batch keys `reward1, reward2, ...`. GDPO computes each component's advantage independently and then **sums** them (`advantage_estimator.py`)
- Per-index consistency across samples: `reward1` must denote the same component for every response, because `calculate_baseline_and_std_per_prompt` builds a per-component baseline across all responses to a prompt. So derive the dict→`reward{n}` mapping deterministically (e.g. sort component names) and apply it identically to every sample, and **validate that all rollouts in the batch expose the same component key set** (fail fast on mismatch).
- Keep `total_reward = full_result["reward"]` (the scalar sum) for GRPO / single-reward consumers.
- When `reward_components` is absent, behavior is unchanged (scalar-only; GRPO).
This mirrors the generic `step()` 2D-rewards path, but sourced from the verify-result dict instead of a tensor.
## Contract notes / open questions
- Keep `reward_components` as a **dict** (`dict[str, float]`) on the producer (NeMo Gym) side — the names are explicit and self-documenting. Converting to ordered `reward{n}` keys is the bridge's responsibility, not the verifier's.
- The field name `reward_components` and its `dict[str, float]` shape would become the cross-repo contract between NeMo Gym verifiers and NeMo-RL GDPO. Worth documenting in `docs/guides/environments.md`.
- Companion NeMo Gym issue proposes a shared `MultiRewardVerifyResponse` (defining `reward_components`) once this bridge lands, to formalize the contract on the producer side.
## Acceptance criteria
- A NeMo Gym env whose verifier returns `reward_components` trains under `grpo.adv_estimator.name: gdpo`, with `reward1..K` populated from the components.
- Scalar-only NeMo Gym envs are unaffected (still GRPO via `total_reward`).
Contributor guide
Assessment
This issue has not been assessed yet.