fix(sc): apply advantage clipping and reward scaling in the advantage stage

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 — `overlong_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.
+- Reward shaping and sample filtering — `overlong_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, in the advantage stage, before advantages are computed.
 - 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

@@ -77,8 +77,10 @@
 )
 from nemo_rl.algorithms.grpo import (
     GRPOSaveState,
+    _clip_grpo_advantages,
     _write_latest_checkpoint_status,
     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
@@ -3170,6 +3172,20 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]:
             tensor_field(data, adv_cfg.sample_mask_field)
         ).float()
 
+        # Before anything reads the rewards, matching grpo.py's ordering. A
+        # no-op unless reward_scaling.enabled, and it is the same helper, so the
+        # source-range clamp and the warning stay identical.
+        # ``getattr``: not every algorithm block defines reward_scaling -- the
+        # OPD test doubles build a bare namespace, and DistillationConfig has
+        # no reward to scale. A missing block is a disabled one, which is the
+        # shipped default anyway.
+        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"]
+
         seq_logprob_error_threshold = self._algo_cfg.seq_logprob_error_threshold
         # Match the legacy path: whenever real policy logprobs are available,
         # report sequence-level generation/training mismatch. A threshold adds
@@ -3276,6 +3292,21 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]:
             self._opd_stat_sumsq += float((valid * valid).sum())
             self._opd_stat_count += int(valid.numel())
 
+        # Clip after logging, not before: grpo.py logs at :3471 and clips at
+        # :3487, grpo_sync.py at :903 then :915. Both report pre-clip
+        # advantages, so clipping first would make SC's advantages/mean|max|min
+        # mean something different from every other driver's.
+        #
+        # A no-op unless advantage_clip_low/high are set, and GRPO-only:
+        # _clip_grpo_advantages has no counterpart in ppo.py.
+        # ``hasattr``: only GRPOConfig carries the two knobs. DistillationConfig
+        # does not, and the OPD test doubles build a bare namespace -- both mean
+        # "no clipping configured". ``not self._is_ppo`` is load-bearing on top
+        # of that, because PPOConfig is ``extra="allow"``: a user who sets
+        # ppo.advantage_clip_low would otherwise get GRPO clipping on a PPO run.
+        if not self._is_ppo and hasattr(self._algo_cfg, "advantage_clip_low"):
+            advantages = _clip_grpo_advantages(advantages, self._algo_cfg)
+
         fields_to_put = {adv_cfg.output_field: advantages}
         if seq_logprob_error_threshold is not None:
             fields_to_put[adv_cfg.sample_mask_field] = sample_mask

nemo_rl/algorithms/single_controller_utils/config.py

@@ -775,14 +775,15 @@ def _validate_algo_settings(master_config: MasterConfig) -> None:
             "with max_num_steps."
         )
 
-    # SC reads none of these on either path, so an enabled one describes shaping
-    # this run does not do. Async GRPO rejects three of them the same way.
+    # An enabled one here describes shaping this run does not do; async GRPO
+    # rejects three of them the same way. An entry leaves this list when the SC
+    # path starts implementing it -- rejecting a knob is only right while
+    # nobody honours it, which is why reward_scaling is no longer here.
     unsupported = [
         name
         for name, enabled in (
             ("overlong_filtering", algo_cfg.overlong_filtering),
             ("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

@@ -331,13 +331,11 @@ def test_accepts_varying_ckpt_structure_with_warmup(self):
         [
             lambda cfg: setattr(cfg, "overlong_filtering", True),
             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=[
             "overlong_filtering",
             "use_dynamic_sampling",
-            "reward_scaling",
             "reward_shaping",
         ],
     )
@@ -740,3 +738,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

@@ -44,6 +44,37 @@
 from nemo_rl.utils.timer import TimeoutChecker, Timer
 
 
+def _grpo_stub(**overrides):
+    """The grpo config fields ``_advantage_stage`` reads, with real defaults.
+
+    Kept in one place: the stage reads a widening set of algorithm-config knobs,
+    and four tests were previously constructing that namespace by hand.
+    """
+    fields = {
+        "seq_logprob_error_threshold": None,
+        "reward_scaling": GRPOConfig.model_fields["reward_scaling"].default_factory(),
+        "advantage_clip_low": None,
+        "advantage_clip_high": None,
+    }
+    fields.update(overrides)
+    return SimpleNamespace(**fields)
+
+
+def _ppo_stub(**overrides):
+    """The ppo config fields ``_advantage_stage`` reads.
+
+    Deliberately without ``advantage_clip_low``/``_high``: real PPOConfig
+    declares neither, so a stub that carries them would hide whether the
+    ``not self._is_ppo`` half of the clip guard does anything.
+    """
+    fields = {
+        "seq_logprob_error_threshold": None,
+        "reward_scaling": GRPOConfig.model_fields["reward_scaling"].default_factory(),
+    }
+    fields.update(overrides)
+    return SimpleNamespace(**fields)
+
+
 class FakeWeightSynchronizer:
     pass
 
@@ -567,7 +598,7 @@ def test_advantage_stage_applies_seq_logprob_error_mask_before_streaming_train(
     ctrl._teacher_logprobs_required = False
     ctrl._is_ppo = False
     ctrl._master_config = SimpleNamespace(
-        grpo=SimpleNamespace(seq_logprob_error_threshold=2.0)
+        grpo=_grpo_stub(seq_logprob_error_threshold=2.0)
     )
     ctrl._algo_cfg = ctrl._master_config.grpo
     ctrl._step_log_dict = {
@@ -636,7 +667,7 @@ def test_advantage_stage_reports_seq_logprob_metrics_without_masking() -> None:
     ctrl._teacher_logprobs_required = False
     ctrl._is_ppo = False
     ctrl._master_config = SimpleNamespace(
-        grpo=SimpleNamespace(seq_logprob_error_threshold=None)
+        grpo=_grpo_stub(seq_logprob_error_threshold=None)
     )
     ctrl._algo_cfg = ctrl._master_config.grpo
     ctrl._step_log_dict = {
@@ -699,7 +730,7 @@ def test_advantage_stage_skips_estimator_when_seq_mask_removes_whole_chunk(
     ctrl._teacher_logprobs_required = False
     ctrl._is_ppo = False
     ctrl._master_config = SimpleNamespace(
-        grpo=SimpleNamespace(seq_logprob_error_threshold=2.0)
+        grpo=_grpo_stub(seq_logprob_error_threshold=2.0)
     )
     ctrl._algo_cfg = ctrl._master_config.grpo
     ctrl._step_log_dict = {
@@ -755,7 +786,7 @@ def test_advantage_stage_skips_preexisting_empty_mask_without_seq_threshold() ->
     ctrl._teacher_logprobs_required = False
     ctrl._is_ppo = False
     ctrl._master_config = SimpleNamespace(
-        grpo=SimpleNamespace(seq_logprob_error_threshold=None)
+        grpo=_grpo_stub(seq_logprob_error_threshold=None)
     )
     ctrl._algo_cfg = ctrl._master_config.grpo
     ctrl._step_log_dict = {
@@ -1739,9 +1770,7 @@ def compute_advantage(self, *, rewards, mask, **kwargs):
     ctrl._reference_logprobs_required = False
     ctrl._teacher_logprobs_required = False
     ctrl._is_ppo = True
-    ctrl._master_config = SimpleNamespace(
-        ppo=SimpleNamespace(seq_logprob_error_threshold=None)
-    )
+    ctrl._master_config = SimpleNamespace(ppo=_ppo_stub())
     ctrl._algo_cfg = ctrl._master_config.ppo
     ctrl._step_log_dict = {
         "rewards": [],
@@ -1769,3 +1798,155 @@ def compute_advantage(self, *, rewards, mask, **kwargs):
     )
     assert "returns" in (result_meta.fields or [])
     assert "advantages" in (result_meta.fields or [])
+
+
+# ---------------------------------------------------------------------------
+# grpo.* knobs that every other driver applies and SC did not.
+# grpo.py:3321 / grpo.py:4883 / grpo_sync.py:915 clip advantages; grpo.py:3024
+# scales rewards. The SC advantage stage is the same point in the pipeline.
+# ---------------------------------------------------------------------------
+
+
+class _RewardRecordingAdvantageEstimator:
+    """Advantage == reward, broadcast, so the written column is readable."""
+
+    def __init__(self) -> None:
+        self.rewards: torch.Tensor | None = None
+
+    def compute_advantage(self, *, rewards, mask, **kwargs):
+        del kwargs
+        self.rewards = rewards.clone()
+        return rewards.unsqueeze(-1).expand_as(mask).clone()
+
+
+def _knob_ctrl(estimator, grpo_stub, *, is_ppo=False):
+    batch, seq = 2, 4
+    columns = {
+        "prompt_ids_for_adv": torch.zeros(batch, seq, dtype=torch.long),
+        "total_reward": torch.tensor([-4.0, 6.0]),
+        "token_mask": torch.ones(batch, seq),
+        "sample_mask": torch.ones(batch),
+    }
+    if is_ppo:
+        # The PPO branch fetches the critic's values column.
+        columns["values"] = torch.zeros(batch, seq)
+    data = TensorDict(columns, batch_size=[batch])
+    data_plane = _AdvantageDataPlane(data)
+    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._master_config = (
+        SimpleNamespace(ppo=grpo_stub) if is_ppo else SimpleNamespace(grpo=grpo_stub)
+    )
+    ctrl._algo_cfg = grpo_stub
+    ctrl._is_ppo = is_ppo
+    ctrl._step_log_dict = {
+        "rewards": [],
+        "masked_advantages": [],
+        "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)],
+        fields=list(data.keys()),
+    )
+    return ctrl, data_plane, meta
+
+
+def test_advantage_clip_bounds_are_applied_before_the_write() -> None:
+    estimator = _RewardRecordingAdvantageEstimator()
+    ctrl, data_plane, meta = _knob_ctrl(
+        estimator, _grpo_stub(advantage_clip_low=-1.0, advantage_clip_high=2.0)
+    )
+
+    asyncio.run(ctrl._advantage_stage(meta))
+
+    written = data_plane.written_fields["advantages"]
+    assert written.min().item() == pytest.approx(-1.0)
+    assert written.max().item() == pytest.approx(2.0)
+
+    # ...but the logged advantages stay pre-clip, because grpo.py:3471 and
+    # grpo_sync.py:903 both log before they clip. Clipping first would leave
+    # SC's advantages/mean|max|min meaning something no other driver's does.
+    logged = torch.cat([t.flatten() for t in ctrl._step_log_dict["masked_advantages"]])
+    assert logged.min().item() == pytest.approx(-4.0)
+    assert logged.max().item() == pytest.approx(6.0)
+
+
+def test_a_ppo_run_ignores_advantage_clip_knobs_set_on_its_own_block() -> None:
+    """``not self._is_ppo`` is the guard, not ``hasattr``.
+
+    PPOConfig is ``extra="allow"``, so a user who sets ``ppo.advantage_clip_low``
+    gets the field -- and ``hasattr`` alone would then run GRPO clipping on a
+    PPO run, which ppo.py never does.
+    """
+
+    class _GaeLikeEstimator:
+        def compute_advantage(self, *, rewards, mask, **kwargs):
+            del kwargs
+            adv = rewards.unsqueeze(-1).expand_as(mask).clone()
+            return adv, adv + 1.0
+
+    ctrl, data_plane, meta = _knob_ctrl(
+        _GaeLikeEstimator(),
+        _ppo_stub(advantage_clip_low=-1.0, advantage_clip_high=2.0),
+        is_ppo=True,
+    )
+
+    asyncio.run(ctrl._advantage_stage(meta))
+
+    written = data_plane.written_fields["advantages"]
+    assert written.min().item() == pytest.approx(-4.0)
+    assert written.max().item() == pytest.approx(6.0)
+
+
+def test_advantages_are_untouched_when_no_clip_bounds_are_set() -> None:
+    estimator = _RewardRecordingAdvantageEstimator()
+    ctrl, data_plane, meta = _knob_ctrl(estimator, _grpo_stub())
+
+    asyncio.run(ctrl._advantage_stage(meta))
+
+    written = data_plane.written_fields["advantages"]
+    assert written.min().item() == pytest.approx(-4.0)
+    assert written.max().item() == pytest.approx(6.0)
+
+
+def test_only_the_configured_clip_bound_applies() -> None:
+    estimator = _RewardRecordingAdvantageEstimator()
+    ctrl, data_plane, meta = _knob_ctrl(estimator, _grpo_stub(advantage_clip_high=2.0))
+
+    asyncio.run(ctrl._advantage_stage(meta))
+
+    written = data_plane.written_fields["advantages"]
+    assert written.max().item() == pytest.approx(2.0)
+    assert written.min().item() == pytest.approx(-4.0), "low bound was not set"
+
+
+def test_reward_scaling_reaches_the_estimator() -> None:
+    scaling = GRPOConfig.model_fields["reward_scaling"].default_factory()
+    scaling.enabled = True
+    scaling.source_min, scaling.source_max = -4.0, 6.0
+    scaling.target_min, scaling.target_max = 0.0, 1.0
+
+    estimator = _RewardRecordingAdvantageEstimator()
+    ctrl, _, meta = _knob_ctrl(estimator, _grpo_stub(reward_scaling=scaling))
+
+    asyncio.run(ctrl._advantage_stage(meta))
+
+    torch.testing.assert_close(estimator.rewards, torch.tensor([0.0, 1.0]))
+
+
+def test_reward_scaling_disabled_leaves_rewards_alone() -> None:
+    estimator = _RewardRecordingAdvantageEstimator()
+    ctrl, _, meta = _knob_ctrl(estimator, _grpo_stub())
+
+    asyncio.run(ctrl._advantage_stage(meta))
+
+    torch.testing.assert_close(estimator.rewards, torch.tensor([-4.0, 6.0]))