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

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

Author: anjalibshah

docs/guides/grpo.md

@@ -650,6 +650,35 @@ grpo:
 ```
 The number of weights must equal the number of reward components, or a `ValueError` is raised.
 
+#### Comparing GDPO against a GRPO baseline
+
+The clearest way to see GDPO's effect is a matched comparison on the same multi-reward environment, changing only the advantage estimator. `examples/configs/gdpo_math_1B.yaml` runs GDPO; the GRPO control is the same config with `adv_estimator.name=grpo` (plus a separate checkpoint dir and logger run name so the two runs don't clobber each other):
+
+```
+# GDPO
+uv run examples/run_grpo.py --config examples/configs/gdpo_math_1B.yaml
+# GRPO control on the SAME multi-reward env — only the advantage estimator changes
+uv run examples/run_grpo.py --config examples/configs/gdpo_math_1B.yaml \
+    grpo.adv_estimator.name=grpo \
+    checkpointing.checkpoint_dir=results/grpo_math_1B_control \
+    logger.wandb.name=grpo-math-1b-control
+```
+
+For intuition on why this matters, see the advantage-collapse example in the [GDPO paper](https://arxiv.org/abs/2601.05242): two responses with the same total reward but different composition (e.g. `(1, 0)` vs. `(0, 1)`) receive an identical advantage under GRPO but distinct advantages under GDPO.
+
+#### What to measure
+
+The headline metric is per-reward convergence, not just aggregate reward. NeMo-RL does not log per-component rewards out of the box (only `total_reward` aggregates are logged), but the components are available as `reward/<name>` keys (e.g. `reward/correctness`) in the rollout batch — add your own logging for them and compare the GDPO and GRPO curves:
+
+- Under GRPO, expect components to move together, or low-variance components to stall as their signal is swamped by the dominant term in the sum.
+- Under GDPO, expect each component to improve on its own schedule.
+- Watch the per-prompt advantage spread: if distinct reward combinations produce near-identical advantages under GRPO, that is the collapse GDPO is meant to fix.
+
+#### Practical notes
+
+- **Reward scaling applies per component.** `reward_scaling` rescales each `reward/<name>` component as well as `total_reward`, so keep components on comparable scales or rely on GDPO's per-component normalization.
+- **Final batch normalization always runs.** In `GDPOAdvantageEstimator`, the final per-batch normalization applies regardless of `normalize_rewards` (which only gates the per-component std division). Account for this when reasoning about advantage magnitudes.
+
 ## LoRA Configuration
 
 GRPO supports LoRA on both the DTensor and Megatron backends. To enable LoRA on the default DTensor backend:

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

@@ -0,0 +1,32 @@
+defaults: ../../../nemo_gym/gdpo_example_tool_call_multireward.yaml
+grpo:
+  num_prompts_per_step: 8
+  num_generations_per_prompt: 8
+policy:
+  model_name: Qwen/Qwen2.5-1.5B-Instruct
+  max_total_sequence_length: 2048
+  dtensor_cfg:
+    enabled: true
+    tensor_parallel_size: 1
+  megatron_cfg:
+    _override_: true
+    enabled: false
+  optimizer:
+    name: torch.optim.AdamW
+    kwargs:
+      lr: 5.0e-06
+      weight_decay: 0.01
+      betas:
+      - 0.9
+      - 0.999
+      eps: 1.0e-08
+  generation:
+    vllm_cfg:
+      http_server_serving_chat_kwargs:
+        tool_parser: hermes
+checkpointing:
+  checkpoint_dir: results/gdpo-qwen2.5-1.5b-1n8g-gym-multireward
+logger:
+  wandb:
+    project: nemo-rl
+    name: gdpo-qwen2.5-1.5b-1n8g-gym-multireward

examples/configs/recipes/llm/gdpo-qwen2.5-1.5b-1n8g-gym-multireward.yaml

@@ -0,0 +1,43 @@
+# GDPO over the `example_tool_call_multireward` NeMo Gym environment.
+#
+# The env asks the model for a single `get_weather` tool call and its verifier returns
+# three decoupled reward components (correctness / format / schema_valid) via
+# `reward_components`. The Gym->NeMo-RL bridge (extract_reward_components /
+# build_reward_component_columns in nemo_rl/environments/nemo_gym.py) surfaces them as
+# `reward/<name>` batch keys, which GDPO's advantage estimator consumes. This is the
+# multi-reward gym path exercised by the nightly recipe under
+# examples/configs/recipes/llm/gdpo-qwen2.5-1.5b-1n8g-gym-multireward.yaml.
+#
+# Inherits the workplace-assistant gym recipe (policy / vLLM / gym plumbing) and swaps in
+# the multi-reward env, its data, and the GDPO estimator.
+defaults: grpo_workplace_assistant_nemotron_nano_v2_9b.yaml
+
+grpo:
+  adv_estimator:
+    name: "gdpo"
+    normalize_rewards: true
+    use_leave_one_out_baseline: false
+    # One weight per component, ordered alphabetically by name:
+    # reward/correctness, reward/format, reward/schema_valid.
+    reward_weights: [1.0, 0.5, 0.25]
+
+data:
+  # example_tool_call_multireward ships its data (data/example.jsonl); no HF download /
+  # ng_prepare_data step is needed. Same file for train/validation (smoke-scale env).
+  train:
+    data_path: 3rdparty/Gym-workspace/Gym/resources_servers/example_tool_call_multireward/data/example.jsonl
+  validation:
+    data_path: 3rdparty/Gym-workspace/Gym/resources_servers/example_tool_call_multireward/data/example.jsonl
+
+env:
+  nemo_gym:
+    # Replaces the workplace_assistant env list. vllm_model_for_training is required and
+    # must come first; then the multi-reward resources server + its agent.
+    config_paths:
+    - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml
+    - resources_servers/example_tool_call_multireward/configs/example_tool_call_multireward.yaml
+
+logger:
+  wandb:
+    project: "gdpo-dev"
+    name: "gdpo-example-tool-call-multireward"

examples/nemo_gym/gdpo_example_tool_call_multireward.yaml

@@ -525,6 +525,32 @@ def extract_reward_components(nemo_gym_result: dict) -> Dict[str, float] | None:
     return {str(name): float(score) for name, score in components.items()}
 
 
+def build_reward_component_columns(
+    component_dicts: List[Dict[str, float] | None],
+) -> Dict[str, torch.Tensor]:
+    """Build ``reward/<name>`` batch columns from per-sample reward-component dicts.
+
+    Takes the union of component names across the batch in sorted (deterministic) order
+    and, for each, emits a ``reward/<name>`` tensor with one entry per sample. A
+    component absent on a given sample is filled with ``0.0`` so every column covers all
+    samples (the per-prompt baseline requires each component present for all responses).
+
+    Keys are prefixed ``reward/`` so they are exactly what
+    ``nemo_rl.algorithms.utils.get_gdpo_reward_component_keys`` selects (it matches
+    ``startswith("reward/")`` and sorts by name); the name carries the component identity,
+    so no positional index is needed. Returns an empty dict when no sample has components.
+    """
+    component_names = sorted(
+        {name for c in component_dicts if c is not None for name in c}
+    )
+    return {
+        f"reward/{name}": torch.tensor(
+            [c[name] if c is not None and name in c else 0.0 for c in component_dicts]
+        )
+        for name in component_names
+    }
+
+
 def validate_reward_components_match_scalar(nemo_gym_results: List[dict]) -> None:
     """Assert each multi-reward result sets ``reward == sum(reward_components)``.
 

nemo_rl/environments/nemo_gym.py

@@ -2471,30 +2471,17 @@ def _postprocess_single_nemo_gym_group(
     # Gym environments so GDPO can compute per-component advantages; single-reward envs
     # are unaffected. Mirrors the native rollout path's reward-component handling above.
     from nemo_rl.environments.nemo_gym import (
+        build_reward_component_columns,
         extract_reward_components,
         validate_reward_components_match_scalar,
     )
 
     component_dicts = [extract_reward_components(r["full_result"]) for r in results]
     if any(c is not None for c in component_dicts):
-        # Emit each component under a `reward/<name>` key, matching the native
-        # multi-reward path and what get_gdpo_reward_component_keys() consumes (it selects
-        # keys starting with "reward/" and sorts them by name). The name carries the
-        # component identity, so ordering is handled downstream by that sort — no
-        # positional index needed. Take the union of names across the batch and default a
-        # component absent on a given sample to 0.0, so every sample carries the same key
-        # set (the per-prompt baseline requires each reward/<name> present for all
-        # responses to a prompt).
-        component_names = sorted(
-            {name for c in component_dicts if c is not None for name in c}
-        )
-        for name in component_names:
-            final_batch[f"reward/{name}"] = torch.tensor(
-                [
-                    c[name] if c is not None and name in c else 0.0
-                    for c in component_dicts
-                ]
-            )
+        # Emit each component under a `reward/<name>` key (see
+        # build_reward_component_columns): matches the native multi-reward path and what
+        # get_gdpo_reward_component_keys() consumes.
+        final_batch.update(build_reward_component_columns(component_dicts))
         # Leave total_reward as the verifier's scalar `reward` (set above); do not
         # silently overwrite it. When a verifier emits reward_components, the contract is
         # reward == sum(components), so overwriting would be a no-op in the correct case

nemo_rl/experience/rollouts.py

@@ -0,0 +1,62 @@
+#!/bin/bash
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd)
+source $SCRIPT_DIR/common.env
+
+# ===== BEGIN CONFIG =====
+NUM_NODES=1
+GPUS_PER_NODE=8
+STEPS_PER_RUN=10
+MAX_STEPS=10
+NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN ))  # Round up
+# Small model, but nemo_gym server startup + vLLM warmup adds overhead on top of the
+# 10 steps; 60 min leaves margin for teardown + metric dump.
+NUM_MINUTES=60
+# ===== END CONFIG =====
+
+exit_if_max_steps_reached
+
+cd $PROJECT_ROOT
+
+# example_tool_call_multireward ships its data (data/example.jsonl) — no HF download or
+# ng_prepare_data needed. Regenerate it so the run never trains on stale example data.
+( cd 3rdparty/Gym-workspace/Gym && uv run python resources_servers/example_tool_call_multireward/create_examples.py )
+
+# Prepare the training data:
+#  - attach `agent_ref` to each row (points rollouts at the env's agent),
+#  - replicate so there are enough prompts to run MAX_STEPS (the shipped example set is
+#    just a handful of rows; one epoch would otherwise end before a single step).
+DATA_DIR=$EXP_DIR/data
+mkdir -p $DATA_DIR
+ENV_DATA=3rdparty/Gym-workspace/Gym/resources_servers/example_tool_call_multireward/data/example.jsonl
+VALIDATION_PATH=$DATA_DIR/example_tool_call_multireward_validation.jsonl
+TRAIN_PATH=$DATA_DIR/example_tool_call_multireward_train.jsonl
+jq -c '. + {agent_ref: {name: "example_tool_call_multireward_simple_agent"}}' $ENV_DATA > $VALIDATION_PATH
+for _ in $(seq 1 60); do cat $VALIDATION_PATH; done > $TRAIN_PATH
+
+# Run the experiment via the gym entrypoint
+uv run examples/nemo_gym/run_grpo_nemo_gym.py \
+    --config $CONFIG_PATH \
+    grpo.max_num_steps=$MAX_STEPS \
+    logger.log_dir=$LOG_DIR \
+    logger.wandb_enabled=True \
+    logger.wandb.project=nemo-rl \
+    logger.wandb.name=$EXP_NAME \
+    logger.monitor_gpus=True \
+    logger.tensorboard_enabled=True \
+    checkpointing.enabled=True \
+    checkpointing.checkpoint_dir=$CKPT_DIR \
+    data.train.data_path=$TRAIN_PATH \
+    data.validation.data_path=$VALIDATION_PATH \
+    $@ \
+    2>&1 | tee $RUN_LOG
+
+uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS
+
+# Smoke-level threshold; tighten after observing real runs.
+if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then
+    uv run tests/check_metrics.py $JSON_METRICS \
+        'max(data["train/reward"]) > 0.0'
+
+    # Clean up checkpoint directory after successful run to save space.
+    rm -rf "$CKPT_DIR"
+fi

tests/test_suites/llm/gdpo-qwen2.5-1.5b-1n8g-fsdp2-reward-weights.sh

@@ -275,7 +275,9 @@ tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-1n8g-megatron-valuetp2sp-pp2cp2-pac
 # GDPO #
 ########
 
-# DTensor GDPO with per-reward weights on the native math_multi_reward env
-# (correctness / format / integer), Qwen2.5-1.5B, GSM8K — guards the reward_weights
-# aggregation path in GDPOAdvantageEstimator.
-tests/test_suites/llm/gdpo-qwen2.5-1.5b-1n8g-fsdp2-reward-weights.sh
+# GDPO over the example_tool_call_multireward NeMo Gym env (correctness / format /
+# schema_valid), Qwen2.5-1.5B — exercises the multi-reward gym rollout path end to end:
+# reward_components -> reward/<name> bridge -> GDPO with per-reward weights. Replaces the
+# native-rollout reward-weights nightly (that path is already covered on main; the
+# reward_weights math itself is covered by unit tests).
+tests/test_suites/llm/gdpo-qwen2.5-1.5b-1n8g-gym-multireward.sh

tests/test_suites/llm/gdpo-qwen2.5-1.5b-1n8g-gym-multireward.sh

@@ -29,6 +29,7 @@
 from nemo_rl.environments.nemo_gym import (
     NemoGym,
     NemoGymConfig,
+    build_reward_component_columns,
     extract_reward_components,
     setup_nemo_gym_config,
     validate_reward_components_match_scalar,
@@ -62,6 +63,40 @@ def test_extract_reward_components():
     assert all(isinstance(v, float) for v in components.values())
 
 
+def test_build_reward_component_columns():
+    """The bridge emission helper: reward/<name> keys, 0.0-fill, deterministic order.
+
+    Guards the producer->consumer contract end to end — the keys built here must be
+    exactly what get_gdpo_reward_component_keys() selects (this is what would have caught
+    the earlier reward1/reward2 vs reward/<name> mismatch).
+    """
+    from nemo_rl.algorithms.utils import get_gdpo_reward_component_keys
+
+    # Keys are reward/<name>; one entry per sample; values preserved.
+    cols = build_reward_component_columns(
+        [
+            {"correctness": 1.0, "format": 0.0},
+            {"correctness": 0.0, "format": 1.0},
+        ]
+    )
+    assert set(cols) == {"reward/correctness", "reward/format"}
+    assert torch.equal(cols["reward/correctness"], torch.tensor([1.0, 0.0]))
+    assert torch.equal(cols["reward/format"], torch.tensor([0.0, 1.0]))
+
+    # Union across the batch, deterministic (sorted) order, 0.0-fill for missing
+    # components (and None samples).
+    cols = build_reward_component_columns([{"b": 2.0}, {"a": 1.0, "b": 3.0}, None])
+    assert list(cols.keys()) == ["reward/a", "reward/b"]
+    assert torch.equal(cols["reward/a"], torch.tensor([0.0, 1.0, 0.0]))
+    assert torch.equal(cols["reward/b"], torch.tensor([2.0, 3.0, 0.0]))
+
+    # The emitted keys are exactly what GDPO's consumer selects.
+    assert get_gdpo_reward_component_keys(cols) == ["reward/a", "reward/b"]
+
+    # No components anywhere -> no columns (single-reward path is untouched).
+    assert build_reward_component_columns([None, None]) == {}
+
+
 def test_validate_reward_components_match_scalar():
     """Multi-reward verifiers must set reward == sum(reward_components); mismatch raises."""
     # Contract satisfied: reward equals the component sum -> no error.