fix(sc): apply reward scaling before advantage estimation

docs/guides/single-controller.md

@@ -241,6 +241,6 @@ The SC path is still under active development. Feature gaps are tracked in [issu
 - Generation backend: vLLM and Megatron generation are supported; SGLang and TRT-LLM have not been tested on SC.
 - Validation is not yet supported (setup raises on `val_period > 0`, `val_at_start`, or `val_at_end`); checkpointing is.
 - (PPO) Rollout drop budgets — `async_rl.rollout_failure.max_skipped_prompts` and `max_consecutive_dropped_prompts` must both be `0`. A drop shortens the step, and the critic shards it against the configured `value.train_global_batch_size` rather than its actual size, so setup rejects a non-zero budget. The resiliency layer stays available on GRPO.
-- Reward shaping and sample filtering — `reward_shaping`, `reward_scaling`, and `use_dynamic_sampling` are implemented on neither algorithm block, so setup rejects them rather than silently skipping the shaping. Environment-flagged sample masking and `overlong_filtering` are supported; truncated completions are excluded from the loss through `sample_mask`, and a step in which every completion is filtered is rejected rather than skipped.
+- Reward shaping and sample filtering — `reward_shaping` and `use_dynamic_sampling` are implemented on neither algorithm block, so setup rejects them rather than silently skipping the shaping. `reward_scaling` is applied before advantages are computed. Environment-flagged sample masking and `overlong_filtering` are supported; truncated completions are excluded from the loss through `sample_mask`, and a step in which every completion is filtered is rejected rather than skipped.
 - The `windowed` sampler has no `over_sampling_ratio` cap — over-produced groups aged past the window are evicted, wasting rollout compute.
 - The drain gate in refit is not yet supported.

nemo_rl/algorithms/single_controller.py

@@ -83,6 +83,7 @@
     _write_latest_checkpoint_status,
     aggregate_rollout_metrics,
     compute_and_apply_seq_logprob_error_masking,
+    scale_rewards,
 )
 from nemo_rl.algorithms.metric_utils import SetupTimingMetrics
 from nemo_rl.algorithms.ppo import _compute_critic_metrics
@@ -3254,6 +3255,16 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]:
             tensor_field(data, adv_cfg.truncated_field)
         ).bool()
 
+        # Match the legacy drivers: scale rewards before the estimator reads
+        # them. Algorithm blocks without reward scaling keep their existing
+        # behavior, and a disabled configuration is a no-op in the shared helper.
+        reward_scaling_cfg = getattr(self._algo_cfg, "reward_scaling", None)
+        if reward_scaling_cfg is not None:
+            rewards = scale_rewards(
+                BatchedDataDict({"total_reward": rewards}),
+                reward_scaling_cfg,
+            )["total_reward"]
+
         num_mask_sample_filtered = int(mask_sample.sum().item())
         self._step_log_dict["num_mask_sample_filtered"].append(num_mask_sample_filtered)
         final_sample_mask = sample_mask * (~mask_sample).to(sample_mask.dtype)

nemo_rl/algorithms/single_controller_utils/config.py

@@ -792,13 +792,12 @@ def _validate_algo_settings(master_config: MasterConfig) -> None:
         )
 
     # An enabled one here describes shaping this run does not do. An entry leaves
-    # this list once the SC path implements it; overlong_filtering is applied in
-    # the advantage stage from the raw completion flags in the TransferQueue.
+    # this list once the SC path implements it; overlong_filtering and
+    # reward_scaling are applied in the advantage stage.
     unsupported = [
         name
         for name, enabled in (
             ("use_dynamic_sampling", algo_cfg.use_dynamic_sampling),
-            ("reward_scaling", algo_cfg.reward_scaling.enabled),
             ("reward_shaping", algo_cfg.reward_shaping.enabled),
         )
         if enabled

tests/unit/single_controller/test_ppo_setup.py

@@ -329,12 +329,10 @@ def test_accepts_varying_ckpt_structure_with_warmup(self):
         "enable",
         [
             lambda cfg: setattr(cfg, "use_dynamic_sampling", True),
-            lambda cfg: setattr(cfg.reward_scaling, "enabled", True),
             lambda cfg: setattr(cfg.reward_shaping, "enabled", True),
         ],
         ids=[
             "use_dynamic_sampling",
-            "reward_scaling",
             "reward_shaping",
         ],
     )
@@ -774,3 +772,15 @@ def test_worker_group_slots(
         else:
             # The critic never lands on the inference cluster.
             assert inference.kwargs["max_colocated_worker_groups"] == 1
+
+    def test_reward_scaling_is_implemented_not_rejected(self):
+        """It is off the unsupported list because _advantage_stage applies it.
+
+        The list exists so an enabled knob cannot silently do nothing. Once the
+        stage calls the same ``scale_rewards`` helper grpo.py uses, rejecting it
+        would refuse a run the path now handles.
+        """
+        mc = _ppo_master_config()
+        mc.ppo.reward_scaling.enabled = True
+
+        validate_single_controller_config(mc)

tests/unit/single_controller/test_single_controller_actor.py

@@ -849,6 +849,73 @@ def test_advantage_stage_clips_training_values_and_metrics() -> None:
     assert logged.max().item() == pytest.approx(2.0)
 
 
+def test_advantage_stage_scales_rewards_before_estimation() -> None:
+    batch_size, sequence_length = 2, 4
+    data = TensorDict(
+        {
+            "prompt_ids_for_adv": torch.zeros(
+                batch_size, sequence_length, dtype=torch.long
+            ),
+            "total_reward": torch.tensor([-4.0, 6.0]),
+            "token_mask": torch.ones(batch_size, sequence_length),
+            "sample_mask": torch.ones(batch_size),
+            "mask_sample": torch.zeros(batch_size, dtype=torch.bool),
+            "truncated": torch.zeros(batch_size, dtype=torch.bool),
+        },
+        batch_size=[batch_size],
+    )
+    data_plane = _AdvantageDataPlane(data)
+    estimator = _MaskRecordingAdvantageEstimator()
+
+    controller_cls = SingleControllerActor.__ray_metadata__.modified_class
+    ctrl = object.__new__(controller_cls)
+    ctrl._dp_client = data_plane
+    ctrl._advantage_cfg = AdvantageConfig()
+    ctrl._advantage_estimator = estimator
+    ctrl._policy_logprobs_required = False
+    ctrl._reference_logprobs_required = False
+    ctrl._teacher_logprobs_required = False
+    ctrl._is_ppo = False
+    ctrl._master_config = SimpleNamespace(
+        grpo=GRPOConfig(
+            seq_logprob_error_threshold=None,
+            reward_scaling={
+                "enabled": True,
+                "source_min": -4.0,
+                "source_max": 6.0,
+                "target_min": 0.0,
+                "target_max": 1.0,
+            },
+        )
+    )
+    ctrl._algo_cfg = ctrl._master_config.grpo
+    ctrl._message_level_advantage_penalties_enabled = False
+    ctrl._step_log_dict = {
+        "rewards": [],
+        "masked_advantages": [],
+        "num_mask_sample_filtered": [],
+        "sequence_lengths": [],
+        "seq_logprob_error_metrics": [],
+    }
+    meta = KVBatchMeta(
+        partition_id="rollout_data",
+        task_name="train",
+        sample_ids=[f"sample-{i}" for i in range(batch_size)],
+        fields=list(data.keys()),
+    )
+
+    asyncio.run(ctrl._advantage_stage(meta))
+
+    assert data_plane.written_fields is not None
+    torch.testing.assert_close(
+        data_plane.written_fields["advantages"],
+        torch.tensor([[0.0] * sequence_length, [1.0] * sequence_length]),
+    )
+    torch.testing.assert_close(
+        ctrl._step_log_dict["rewards"][0], torch.tensor([0.0, 1.0])
+    )
+
+
 def test_advantage_stage_skips_estimator_when_seq_mask_removes_whole_chunk(
     capsys: pytest.CaptureFixture[str],
 ) -> None: