feat(sc): apply reward penalties to Gym rollouts

examples/configs/grpo_math_1B_megatron_single_controller.yaml

@@ -8,6 +8,15 @@ grpo:
   async_grpo: null
   val_period: 0
 
+# 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}
+
 async_rl:
   sampler:
     name: in_order

examples/configs/ppo_math_1B_megatron_single_controller.yaml

@@ -16,6 +16,15 @@ ppo:
   reward_scaling:
     enabled: false
 
+# 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}
+
 async_rl:
   sampler:
     name: in_order

nemo_rl/algorithms/single_controller_utils/config.py

@@ -35,7 +35,12 @@
     SamplerConfig,
     required_buffer_capacity_for_config,
 )
-from nemo_rl.algorithms.grpo import GRPOConfig, GRPOLoggerConfig
+from nemo_rl.algorithms.grpo import (
+    _REWARD_PENALTY_FLAGS,
+    GRPOConfig,
+    GRPOLoggerConfig,
+    RewardPenaltyConfig,
+)
 from nemo_rl.algorithms.loss import ClippedPGLossConfig
 from nemo_rl.algorithms.loss.loss_functions import MseValueLossConfig
 from nemo_rl.algorithms.opd import OnPolicyDistillationConfig
@@ -586,6 +591,7 @@ class MasterConfig(BaseModel, extra="allow"):
     logger: GRPOLoggerConfig
     cluster: ClusterConfig
     checkpointing: CheckpointingConfig
+    reward_penalties: RewardPenaltyConfig = Field(default_factory=RewardPenaltyConfig)
     data_plane: DataPlaneConfig
     async_rl: AsyncRLConfig
     on_policy_distillation: Optional[OnPolicyDistillationConfig] = None
@@ -933,6 +939,15 @@ def validate_single_controller_config(master_config: MasterConfig) -> None:
     async_config = master_config.async_rl
     algo_cfg = algo_config(master_config)
 
+    reward_penalties_enabled = any(
+        getattr(master_config.reward_penalties, flag) for flag in _REWARD_PENALTY_FLAGS
+    )
+    if reward_penalties_enabled and not master_config.env.get("should_use_nemo_gym"):
+        raise ValueError(
+            "reward_penalties require the NeMo-Gym rollout path "
+            "(env.should_use_nemo_gym=true) on SingleController"
+        )
+
     if algo_cfg.num_prompts_per_step < async_config.min_groups_for_streaming_train:
         raise ValueError(
             f"num_prompts_per_step ({algo_cfg.num_prompts_per_step}) "

nemo_rl/algorithms/single_controller_utils/setup.py

@@ -93,7 +93,11 @@
     RolloutRetryPolicy,
     RolloutTimeouts,
 )
-from nemo_rl.experience.rollouts import should_mask_flagged_samples
+from nemo_rl.experience.rollouts import (
+    get_nemo_gym_thinking_tags,
+    resolve_reward_penalty_config,
+    should_mask_flagged_samples,
+)
 from nemo_rl.models.generation import resolve_generation_class
 from nemo_rl.models.generation.fleet_health import (
     FleetHealthPolicy,
@@ -881,6 +885,11 @@ def setup_single_controller(
         logged by the SC actor).
     """
     validate_single_controller_config(master_config)
+    resolved_reward_penalty_config = resolve_reward_penalty_config(
+        master_config.reward_penalties,
+        tokenizer,
+        thinking_tags=get_nemo_gym_thinking_tags(master_config.env),
+    )
 
     # short names for config sections
     algo_cfg = algo_config(master_config)
@@ -1405,6 +1414,7 @@ def _build_generation_then_trainer(
         generation_config=generation_config,
         use_nemo_gym=use_nemo_gym,
         mask_env_flagged_samples=should_mask_flagged_samples(master_config.env),
+        reward_penalty_config=resolved_reward_penalty_config,
         tq_buffer=tq_buffer,
         timeouts=RolloutTimeouts(
             rollout_s=master_config.async_rl.rollout_failure.nemo_gym.rollout_timeout_s,

nemo_rl/experience/rollout_manager.py

@@ -56,8 +56,10 @@
     _effort_shaping_metrics,
     _find_routed_experts_template,
     _tensorize_by_key,
+    apply_reward_penalties,
     attach_static_multimodal_payload,
     calculate_rewards,
+    compute_reward_penalty_metrics,
 )
 from nemo_rl.models.generation.interfaces import (
     GenerationConfig,
@@ -765,6 +767,7 @@ def __init__(
         max_rollout_turns: int,
         generation_config: GenerationConfig,
         mask_env_flagged_samples: bool = True,
+        reward_penalty_config: Optional[dict[str, Any]] = None,
         # Optional so direct construction does not have to carry the resiliency wiring;
         # RolloutManager always passes both explicitly.
         timeouts: Optional[RolloutTimeouts] = None,
@@ -783,6 +786,7 @@ def __init__(
         self._max_rollout_turns = max_rollout_turns
         self._generation_config = generation_config
         self._mask_env_flagged_samples = mask_env_flagged_samples
+        self._reward_penalty_config = reward_penalty_config
         self._timeouts = timeouts if timeouts is not None else RolloutTimeouts()
         self._max_gym_row_attempts = (
             retry_policy
@@ -1018,10 +1022,11 @@ async def _run_rollouts(
             # All N rollouts share the same input prompt; tensorize one copy.
             prompt_message_log = completed_results[0]["input_message_log"]
             _tensorize_by_key(prompt_message_log, "token_ids")
-            # Convert results to completions.
-            completions = [
-                self._result_to_completion(result) for result in completed_results
-            ]
+            # Apply penalties before Completion captures each result's reward, while
+            # preserving the batch-level counts used by legacy Gym metrics.
+            completions, penalty_counts = self._results_to_completions(
+                completed_results
+            )
 
         # Compute rollout metrics.
         with timer.time(f"{timer_prefix}/compute_metrics"):
@@ -1030,36 +1035,59 @@ async def _run_rollouts(
             )
             # Same helper the batched path uses, so the two cannot drift apart.
             rollout_metrics.update(_effort_shaping_metrics(shaping))
+            rollout_metrics.update(
+                self._compute_reward_penalty_metrics(
+                    penalty_counts, len(completed_results)
+                )
+            )
 
         rollout_metrics.update(env_timing_metrics)
 
         return completions, prompt_message_log, rollout_metrics
 
-    def _result_to_completion(self, result: dict) -> Completion:
-        """Convert one run_rollouts result dict into a Completion."""
-        # Tensorize token fields.
-        _tensorize_by_key(result["message_log"], "token_ids")
-        _tensorize_by_key(
-            [m for m in result["message_log"] if m["role"] == "assistant"],
-            "generation_logprobs",
-        )
-        # Calculate truncation.
-        truncated = (
-            sum(len(m["token_ids"]) for m in result["message_log"]) == self._max_seq_len
-        )
-
-        # Same gate as the batched path: when masking is off, drop the env
-        # mask flag so later batch building never sees it.
-        if not self._mask_env_flagged_samples:
-            (result["full_result"].get("instance_config") or {}).pop(
-                "mask_sample", None
+    def _results_to_completions(
+        self, results: list[dict]
+    ) -> tuple[list[Completion], dict[str, int]]:
+        """Apply configured penalties and convert a Gym result batch."""
+        for result in results:
+            _tensorize_by_key(result["message_log"], "token_ids")
+            _tensorize_by_key(
+                [m for m in result["message_log"] if m["role"] == "assistant"],
+                "generation_logprobs",
             )
 
-        return Completion(
-            message_log=result["message_log"],
-            env_extras=result["full_result"],
-            truncated=truncated,
-            reward=float(result["full_result"]["reward"]),
+            # Same gate as the batched path: when masking is off, drop the env
+            # mask flag so later batch building never sees it.
+            if not self._mask_env_flagged_samples:
+                (result["full_result"].get("instance_config") or {}).pop(
+                    "mask_sample", None
+                )
+
+        penalty_counts = apply_reward_penalties(results, self._reward_penalty_config)
+        completions = []
+        for result in results:
+            truncated = (
+                sum(len(m["token_ids"]) for m in result["message_log"])
+                == self._max_seq_len
+            )
+            completions.append(
+                Completion(
+                    message_log=result["message_log"],
+                    env_extras=result["full_result"],
+                    truncated=truncated,
+                    reward=float(result["full_result"]["reward"]),
+                )
+            )
+        return completions, penalty_counts
+
+    def _compute_reward_penalty_metrics(
+        self, penalty_counts: dict[str, int], num_results: int
+    ) -> dict[str, float]:
+        """Return enabled penalty rates using the legacy Gym metric names."""
+        return compute_reward_penalty_metrics(
+            penalty_counts,
+            num_results,
+            self._reward_penalty_config,
         )
 
     def _compute_rollout_metrics(
@@ -1154,6 +1182,7 @@ def __init__(
         generation_config: Optional[GenerationConfig] = None,
         use_nemo_gym: bool = False,
         mask_env_flagged_samples: bool = True,
+        reward_penalty_config: Optional[dict[str, Any]] = None,
         tq_buffer: Optional[TQReplayBuffer] = None,
         timeouts: Optional[RolloutTimeouts] = None,
         retry_policy: Optional[RolloutRetryPolicy] = None,
@@ -1193,6 +1222,7 @@ def __init__(
             generation_config=generation_config,
             # Only used by AsyncNemoGymRolloutImpl; AsyncRolloutImpl ignores it.
             mask_env_flagged_samples=mask_env_flagged_samples,
+            reward_penalty_config=reward_penalty_config,
             # None means "no deadlines", which is what async_rl's own defaults resolve
             # to; callers that have a config pass the resolved values in.
             timeouts=timeouts if timeouts is not None else RolloutTimeouts(),

nemo_rl/experience/rollouts.py

@@ -76,6 +76,22 @@
 
 TokenizerType = PreTrainedTokenizerBase
 
+_REWARD_PENALTY_METRICS = {
+    "duplicated_reasoning": (
+        "penalize_duplicated_reasoning",
+        "reasoning_equal_to_final_answer_rate",
+    ),
+    "empty_final_answer": (
+        "penalize_empty_final_answer",
+        "empty_final_answer_rate",
+    ),
+    "unwanted_token": ("penalize_unwanted_tokens", "unwanted_token_rate"),
+    "malformed_think_tag": (
+        "penalize_malformed_think_tag",
+        "malformed_think_tag_rate",
+    ),
+}
+
 
 def attach_initial_nemo_gym_image_payloads(
     batch: BatchedDataDict[DatumSpec],
@@ -1875,6 +1891,22 @@ def _get_reward_penalty_config_value(
     return getattr(reward_penalty_config, key, None)
 
 
+def compute_reward_penalty_metrics(
+    penalty_counts: dict[str, int],
+    num_results: int,
+    reward_penalty_config: dict[str, Any] | BaseModel | None,
+) -> dict[str, float]:
+    """Return enabled penalty rates using the legacy NeMo-Gym metric names."""
+    if reward_penalty_config is None or not num_results:
+        return {}
+
+    return {
+        metric_name: penalty_counts[count_key] / num_results
+        for count_key, (flag, metric_name) in _REWARD_PENALTY_METRICS.items()
+        if _get_reward_penalty_config_value(reward_penalty_config, flag)
+    }
+
+
 def _get_reward_penalty_token_id(
     reward_penalty_config: dict[str, Any] | BaseModel,
     key: str,
@@ -2823,26 +2855,13 @@ def _postprocess_single_nemo_gym_group(
 
     rollout_metrics.update(_effort_shaping_metrics(shaping))
 
-    # Penalty metrics — map count keys to (config flag, metric name)
-    _PENALTY_METRICS = {
-        "duplicated_reasoning": (
-            "penalize_duplicated_reasoning",
-            "reasoning_equal_to_final_answer_rate",
-        ),
-        "empty_final_answer": (
-            "penalize_empty_final_answer",
-            "empty_final_answer_rate",
-        ),
-        "unwanted_token": ("penalize_unwanted_tokens", "unwanted_token_rate"),
-        "malformed_think_tag": (
-            "penalize_malformed_think_tag",
-            "malformed_think_tag_rate",
-        ),
-    }
-    if resolved_reward_penalty_config and results:
-        for key, (flag, metric_name) in _PENALTY_METRICS.items():
-            if _get_reward_penalty_config_value(resolved_reward_penalty_config, flag):
-                rollout_metrics[metric_name] = penalty_counts[key] / len(results)
+    rollout_metrics.update(
+        compute_reward_penalty_metrics(
+            penalty_counts,
+            len(results),
+            resolved_reward_penalty_config,
+        )
+    )
 
     # Expose per-component rewards as `reward/<name>` batch keys for multi-reward NeMo
     # Gym environments so GDPO can compute per-component advantages; single-reward envs

tests/unit/experience/test_rollout_generation_failures.py

@@ -537,6 +537,8 @@ def _make_gym_impl(
     impl._stats = stats if stats is not None else RolloutStats()
     # Upstream default; this fixture is about re-dispatch, not sample masking.
     impl._mask_env_flagged_samples = True
+    # Reward penalties are off; direct construction must still satisfy the impl contract.
+    impl._reward_penalty_config = None
     # Effort-level reward shaping is off unless env.nemo_gym.effort_levels is set.
     impl._effort_config = None
     return impl