Skip to content

feat(sc): apply reward penalties to Gym rollouts - #3741

Merged
yuki-97 merged 3 commits into
mainfrom
amahishi/gym_penalties_sc
Aug 31, 2026
Merged

feat(sc): apply reward penalties to Gym rollouts#3741
yuki-97 merged 3 commits into
mainfrom
amahishi/gym_penalties_sc

Conversation

@macandro96

@macandro96 macandro96 commented Aug 20, 2026

Copy link
Copy Markdown
Contributor

What does this PR do ?

Support reward penalties for gym rollouts.

Issues

List issues that this PR closes (syntax):

Usage

# Reward-zeroing penalties applied to NeMo-Gym rollout results.
reward_penalties:
  penalize_duplicated_reasoning: false
  penalize_empty_final_answer: false
  penalize_unwanted_tokens: false
  penalize_malformed_think_tag: false
  # Optional model/tokenizer-specific IDs. Example:
  # token_ids: {unwanted: [2], think_open: 12, think_close: 13}

Before your PR is "Ready for review"

Pre checks:

  • Make sure you read and followed Contributor guidelines
  • Did you write any new necessary tests?
  • Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
  • Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

  • ...

@macandro96
macandro96 requested review from a team as code owners August 20, 2026 23:35
@copy-pr-bot

copy-pr-bot Bot commented Aug 20, 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.

@macandro96 macandro96 left a comment

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Nice work extending the reward-penalty feature to the SingleController Gym path. Overall the wiring is clean: resolving reward_penalty_config once at setup enables the fail-fast test, and the batched _results_to_completions refactor keeps the singleton hook. Tests for the new SC surface pass locally (16 + 35).

Findings from the parallel review team (bug-finder, rl-expert, test-agent, design-reviewer, devil-advocate):

  • 1 pre-existing test regression that will break CI (inline).
  • 3 DRY / config-conventions suggestions (inline).
  • 1 test-coverage suggestion (inline).
  • Body: the PR description is the unfilled auto-template — please add a one-line summary and link the non-SC parity motivation (existing feature at grpo.py:2960 / :3893 / trajectory_collector.py:936).

Pre-commit passed clean. Local pytest runs:

  • tests/unit/experience/test_rollout_manager.py: 16 passed
  • tests/unit/single_controller/test_single_controller_setup.py: 35 passed
  • tests/unit/experience/test_rollouts.py::test_rollout_manager_consumes_stream_and_restores_input_order: FAILED (see inline)

Generated by Claude Code

Comment thread nemo_rl/experience/rollout_manager.py
Comment thread nemo_rl/algorithms/single_controller_utils/config.py
Comment thread nemo_rl/experience/rollout_manager.py Outdated
Comment thread nemo_rl/algorithms/single_controller_utils/config.py
Comment thread tests/unit/experience/test_rollout_manager.py
@macandro96
macandro96 requested a review from yfw August 21, 2026 23:27

@yuki-97 yuki-97 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.

Traced the reward-penalty wiring for SC end-to-end (RewardPenaltyConfig -> validate_single_controller_config -> resolve_reward_penalty_config -> RolloutManager -> apply_reward_penalties -> compute_reward_penalty_metrics) — logic is correct and the prior self-review threads all genuinely landed. Two minor items below: a dead private wrapper left over from the streaming refactor, and a heads-up that main has since added PPO to SingleController, so this needs a small parity follow-up after rebase.

Comment thread nemo_rl/experience/rollout_manager.py Outdated
Comment thread nemo_rl/algorithms/single_controller_utils/config.py
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
@macandro96
macandro96 force-pushed the amahishi/gym_penalties_sc branch from bd60fb7 to 2fe8a47 Compare August 30, 2026 04:17
@macandro96

Copy link
Copy Markdown
Contributor Author

/ok to test 2fe8a47

@macandro96
macandro96 requested a review from yuki-97 August 30, 2026 04:18
@macandro96 macandro96 added the CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) label Aug 30, 2026

@yuki-97 yuki-97 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.

Re-reviewed at 2fe8a47. Both open threads check out against the current head:

  • Dead _result_to_completion wrapper — deleted, and both mask-gate tests now call _results_to_completions([result])[0][0]. A repo-wide grep at head finds no definition and no caller left.
  • PPO parity after the rebasereward_penalties sits on the new grpo/ppo-Optional MasterConfig, the PPO SC exemplar carries the block, and neither validate_single_controller_config nor resolve_reward_penalty_config reads master_config.grpo, so the PPO SC path genuinely gets the penalties rather than silently dropping them.

One nit below, nothing blocking.

L0_Unit_Tests_Other is red at this head, needs to be fixed before merge.

Comment thread nemo_rl/experience/rollout_manager.py
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
@macandro96

Copy link
Copy Markdown
Contributor Author

/ok to test 5e48ef6

@macandro96
macandro96 requested a review from yuki-97 August 30, 2026 23:34
@yuki-97
yuki-97 enabled auto-merge (squash) August 31, 2026 05:18
@yuki-97
yuki-97 merged commit e27a64e into main Aug 31, 2026
149 of 151 checks passed
@yuki-97
yuki-97 deleted the amahishi/gym_penalties_sc branch August 31, 2026 05:35
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants