fix(sc): address reward penalty review feedback

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>

Author: macandro96

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/experience/rollout_manager.py

@@ -56,8 +56,8 @@
     _effort_shaping_metrics,
     _find_routed_experts_template,
     _tensorize_by_key,
-    attach_static_multimodal_payload,
     apply_reward_penalties,
+    attach_static_multimodal_payload,
     calculate_rewards,
     compute_reward_penalty_metrics,
 )
@@ -1085,7 +1085,9 @@ def _compute_reward_penalty_metrics(
     ) -> 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
+            penalty_counts,
+            num_results,
+            self._reward_penalty_config,
         )
 
     def _compute_rollout_metrics(

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_manager.py

@@ -612,12 +612,16 @@ def _mask_gate_result():
 
 
 def test_result_to_completion_keeps_mask_flag_when_gate_on():
-    completion = _nemo_gym_impl(True)._results_to_completions([_mask_gate_result()])[0][0]
+    completion = _nemo_gym_impl(True)._results_to_completions([_mask_gate_result()])[0][
+        0
+    ]
     assert completion.env_extras["instance_config"]["mask_sample"] is True
 
 
 def test_result_to_completion_drops_mask_flag_when_gate_off():
-    completion = _nemo_gym_impl(False)._results_to_completions([_mask_gate_result()])[0][0]
+    completion = _nemo_gym_impl(False)._results_to_completions([_mask_gate_result()])[
+        0
+    ][0]
     assert "mask_sample" not in completion.env_extras["instance_config"]
     assert completion.env_extras["instance_config"]["other_key"] == "kept"
 
@@ -714,6 +718,22 @@ def test_nemo_gym_reward_penalties_match_legacy_rewards_counts_and_metrics(
     assert impl._compute_reward_penalty_metrics(penalty_counts, 1) == {metric_name: 1.0}
 
 
+def test_nemo_gym_reward_penalty_metrics_compute_fractional_rate():
+    impl = _nemo_gym_impl(True, {"penalize_empty_final_answer": True})
+
+    metrics = impl._compute_reward_penalty_metrics(
+        {
+            "duplicated_reasoning": 0,
+            "empty_final_answer": 1,
+            "unwanted_token": 0,
+            "malformed_think_tag": 0,
+        },
+        3,
+    )
+
+    assert metrics == {"empty_final_answer_rate": 1 / 3}
+
+
 # ---------------------------------------------------------------------------
 # Tests for AsyncRolloutManager (native async path)
 # ---------------------------------------------------------------------------

tests/unit/experience/test_rollouts.py

@@ -2210,7 +2210,11 @@ def remote(self, inputs, timer_prefix):
     }
     manager._tokenizer = None
     manager._effort_config = None
-    manager._result_to_completion = lambda result: result["value"]
+    manager._results_to_completions = lambda results: (
+        [result["value"] for result in results],
+        {},
+    )
+    manager._compute_reward_penalty_metrics = lambda counts, num_results: {}
     manager._compute_rollout_metrics = lambda completions, agent: {
         "completion_count": len(completions),
         "agent": agent,