From 512fb935190a45838e21d9c0fc32694c2752fb2b Mon Sep 17 00:00:00 2001
From: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
Date: Mon, 24 Aug 2026 00:40:00 -0400
Subject: [PATCH 1/8] fix(sc): apply advantage clipping and reward scaling in
 the advantage stage

grpo.advantage_clip_low/high and grpo.reward_scaling parse on the
SingleController path and then do nothing: grep for either in
single_controller.py or single_controller_utils/ returns no hits.

Every other driver applies them. _clip_grpo_advantages is called from
grpo.py:3321, grpo.py:4883 and grpo_sync.py:915; scale_rewards from
grpo.py:3024. grpo_sync.py:915 is the closest analog to SC's
_advantage_stage -- it clips immediately before its data-plane write, which is
the same point in the pipeline -- so this calls the same two helpers rather
than reimplementing either, and the source-range clamp and its warning stay
identical to the batched path.

Both are no-ops at their defaults: advantage_clip_* default to None and
reward_scaling.enabled to False, so a run that does not set them is unchanged.

Two knobs are deliberately left out because SC does not fetch what they need.
grpo.reward_shaping reads batch["message_log"] (reward_functions.py:145) and
optionally response_token_lengths; the message-level advantage penalties behind
grpo.invalid_tool_call_advantage / malformed_thinking_advantage need the message
logs too. _advantage_stage pulls only _advantage_input_fields() from the data
plane, so both want plumbing beyond this change.

Five tests. The three that assert a knob does something fail on main; the two
that assert the disabled default changes nothing pass either way, which is the
point of having them. The grpo config namespace those tests build by hand moved
into a shared _grpo_stub so the next field the stage reads does not break four
unrelated tests again.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
---
 nemo_rl/algorithms/single_controller.py       |  13 ++
 .../test_single_controller.py                 | 136 +++++++++++++++++-
 2 files changed, 145 insertions(+), 4 deletions(-)

diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py
index 41d64add4b..0548ffd15e 100644
--- a/nemo_rl/algorithms/single_controller.py
+++ b/nemo_rl/algorithms/single_controller.py
@@ -49,8 +49,10 @@
 from nemo_rl.algorithms.async_utils.staleness_sampler import create_sampler
 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.single_controller_utils.config import (
@@ -1705,6 +1707,14 @@ 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 grpo.reward_scaling.enabled, and it is the same helper,
+        # so the source-range clamp and the warning stay identical.
+        rewards = scale_rewards(
+            BatchedDataDict({"total_reward": rewards}),
+            self._master_config.grpo.reward_scaling,
+        )["total_reward"]
+
         seq_logprob_error_threshold = (
             self._master_config.grpo.seq_logprob_error_threshold
         )
@@ -1784,6 +1794,9 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]:
             )
         else:
             advantages = torch.zeros_like(mask)
+        # grpo_sync.py:915 does the same, immediately before its data-plane
+        # write. A no-op unless grpo.advantage_clip_low/high are set.
+        advantages = _clip_grpo_advantages(advantages, self._master_config.grpo)
         response_advantages = torch.masked_select(advantages, mask.bool())
         self._step_log_dict["rewards"].append(rewards.detach().cpu())
         self._step_log_dict["masked_advantages"].append(
diff --git a/tests/unit/single_controller/test_single_controller.py b/tests/unit/single_controller/test_single_controller.py
index 3573ba901a..15e380c670 100644
--- a/tests/unit/single_controller/test_single_controller.py
+++ b/tests/unit/single_controller/test_single_controller.py
@@ -303,6 +303,22 @@ def put_samples(self, *, fields, **kwargs) -> None:
         self.written_fields = fields
 
 
+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 ``master_config.grpo``
+    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)
+
+
 class _MaskRecordingAdvantageEstimator:
     def __init__(self) -> None:
         self.mask: torch.Tensor | None = None
@@ -345,7 +361,7 @@ def test_advantage_stage_applies_seq_logprob_error_mask_before_streaming_train(
     ctrl._policy_logprobs_required = True
     ctrl._reference_logprobs_required = False
     ctrl._master_config = SimpleNamespace(
-        grpo=SimpleNamespace(seq_logprob_error_threshold=2.0)
+        grpo=_grpo_stub(seq_logprob_error_threshold=2.0)
     )
     ctrl._step_log_dict = {
         "rewards": [],
@@ -411,7 +427,7 @@ def test_advantage_stage_reports_seq_logprob_metrics_without_masking() -> None:
     ctrl._policy_logprobs_required = True
     ctrl._reference_logprobs_required = False
     ctrl._master_config = SimpleNamespace(
-        grpo=SimpleNamespace(seq_logprob_error_threshold=None)
+        grpo=_grpo_stub(seq_logprob_error_threshold=None)
     )
     ctrl._step_log_dict = {
         "rewards": [],
@@ -471,7 +487,7 @@ def test_advantage_stage_skips_estimator_when_seq_mask_removes_whole_chunk(
     ctrl._policy_logprobs_required = True
     ctrl._reference_logprobs_required = False
     ctrl._master_config = SimpleNamespace(
-        grpo=SimpleNamespace(seq_logprob_error_threshold=2.0)
+        grpo=_grpo_stub(seq_logprob_error_threshold=2.0)
     )
     ctrl._step_log_dict = {
         "rewards": [],
@@ -524,7 +540,7 @@ def test_advantage_stage_skips_preexisting_empty_mask_without_seq_threshold() ->
     ctrl._policy_logprobs_required = False
     ctrl._reference_logprobs_required = False
     ctrl._master_config = SimpleNamespace(
-        grpo=SimpleNamespace(seq_logprob_error_threshold=None)
+        grpo=_grpo_stub(seq_logprob_error_threshold=None)
     )
     ctrl._step_log_dict = {
         "rewards": [],
@@ -981,3 +997,115 @@ def test_train_pump_keeps_train_buffers_once_the_step_is_open(monkeypatch) -> No
 
     assert ctrl._train_steps == 1
     assert trainer.keep_train_buffers_calls == [False, True]
+
+
+# ---------------------------------------------------------------------------
+# 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):
+    batch, seq = 2, 4
+    data = TensorDict(
+        {
+            "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),
+        },
+        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._master_config = SimpleNamespace(grpo=grpo_stub)
+    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)
+
+
+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]))

From cbd29ffc777c6b8f6b687ee22d7a7e2527738a3f Mon Sep 17 00:00:00 2001
From: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
Date: Wed, 26 Aug 2026 18:15:57 -0400
Subject: [PATCH 2/8] style(sc): make the unsupported-list comment identical
 across the parity PRs

#3786 and #3787 each rewrote this comment to name the knob they removed, so
they conflicted on nothing but the prose -- the tuple entries merged fine.
Stating the rule once, without naming a knob, lets the two auto-merge in
either order.

Signed-off-by: Tianyi Zhang <zhangtianyi975@gmail.com>
Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
---
 nemo_rl/algorithms/single_controller_utils/config.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py
index 9ac73f5379..06ae99c434 100644
--- a/nemo_rl/algorithms/single_controller_utils/config.py
+++ b/nemo_rl/algorithms/single_controller_utils/config.py
@@ -728,10 +728,9 @@ def _validate_algo_settings(master_config: MasterConfig) -> None:
     """
     algo_cfg = algo_config(master_config)
 
-    # An enabled one here describes shaping this run does not do. reward_scaling
-    # is off this list because _advantage_stage now applies it through the same
-    # scale_rewards helper grpo.py uses -- rejecting a knob is only right while
-    # nobody implements it.
+    # An enabled one here describes shaping this run does not do. An entry
+    # leaves this list when the SC path starts implementing it -- rejecting a
+    # knob is only right while nobody honours it.
     unsupported = [
         name
         for name, enabled in (

From add1837f3af1dcc6949778f5b3e06b272ee7c634 Mon Sep 17 00:00:00 2001
From: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
Date: Wed, 26 Aug 2026 18:17:11 -0400
Subject: [PATCH 3/8] test(sc): move the reward_scaling case off #3786's
 insertion point

Both parity PRs added a test right before the same anchor, so they conflicted
on nothing but placement. Moving this one to the end of the file lets the two
auto-merge in either order.

Signed-off-by: Tianyi Zhang <zhangtianyi975@gmail.com>
Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
---
 .../unit/single_controller/test_ppo_setup.py  | 24 +++++++++----------
 1 file changed, 12 insertions(+), 12 deletions(-)

diff --git a/tests/unit/single_controller/test_ppo_setup.py b/tests/unit/single_controller/test_ppo_setup.py
index ce11cff668..8c424ce260 100644
--- a/tests/unit/single_controller/test_ppo_setup.py
+++ b/tests/unit/single_controller/test_ppo_setup.py
@@ -343,18 +343,6 @@ def test_rejects_shaping_the_sc_path_does_not_implement(self, enable):
         ):
             validate_single_controller_config(mc)
 
-    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)
-
     def test_rejects_shaping_on_a_grpo_run_too(self):
         mc = _make_master_config()
         mc.grpo.overlong_filtering = True
@@ -637,3 +625,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)

From f768ea7f434be3c27581689699d010974bd6fca4 Mon Sep 17 00:00:00 2001
From: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
Date: Wed, 26 Aug 2026 22:19:11 -0400
Subject: [PATCH 4/8] fix(sc): read reward_scaling defensively in the advantage
 stage

Not every algorithm block defines it. The OPD test doubles main added in
#3768 build a bare namespace, and DistillationConfig has no reward to scale
at all, so a hard attribute read turns this into an AttributeError on paths
that never asked for scaling.

A missing block is a disabled one, which is the shipped default anyway.

Found by merging all twenty of my open PRs together: this passed alone and
failed against main's MOPD tests.

Signed-off-by: Tianyi Zhang <zhangtianyi975@gmail.com>
Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
---
 nemo_rl/algorithms/single_controller.py | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py
index d854deb1e7..f980c4a4f8 100644
--- a/nemo_rl/algorithms/single_controller.py
+++ b/nemo_rl/algorithms/single_controller.py
@@ -1931,10 +1931,16 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]:
         # 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.
-        rewards = scale_rewards(
-            BatchedDataDict({"total_reward": rewards}),
-            self._algo_cfg.reward_scaling,
-        )["total_reward"]
+        # ``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,

From 6ae01300e36d171d4ff861a4344d2a70db59f746 Mon Sep 17 00:00:00 2001
From: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
Date: Wed, 26 Aug 2026 22:38:06 -0400
Subject: [PATCH 5/8] fix(sc): gate the advantage clip on the knob existing,
 not just on not-PPO

Same class as the reward_scaling read one commit up. Only GRPOConfig carries
advantage_clip_low/high; DistillationConfig does not, and the OPD test doubles
main added in #3768 build a bare namespace. "Not PPO" is not the same as
"has the knob", so the not-is_ppo guard alone turns this into an
AttributeError on paths that never configured clipping.

Signed-off-by: Tianyi Zhang <zhangtianyi975@gmail.com>
Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
---
 nemo_rl/algorithms/single_controller.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py
index f980c4a4f8..40aad2eb46 100644
--- a/nemo_rl/algorithms/single_controller.py
+++ b/nemo_rl/algorithms/single_controller.py
@@ -2035,7 +2035,10 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]:
         # write. A no-op unless advantage_clip_low/high are set, and GRPO-only:
         # _clip_grpo_advantages has no counterpart in ppo.py, and PPOConfig
         # carries neither knob, so a PPO run must not reach for them.
-        if not self._is_ppo:
+        # ``hasattr``: only GRPOConfig carries the two knobs. PPOConfig does
+        # not, DistillationConfig does not, and the OPD test doubles build a
+        # bare namespace -- all of which mean "no clipping configured".
+        if not self._is_ppo and hasattr(self._algo_cfg, "advantage_clip_low"):
             advantages = _clip_grpo_advantages(advantages, self._algo_cfg)
         response_advantages = torch.masked_select(advantages, mask.bool())
         self._step_log_dict["rewards"].append(rewards.detach().cpu())

From de259f4850c73fea37c9cdf6ff6f59394e3128d5 Mon Sep 17 00:00:00 2001
From: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
Date: Wed, 26 Aug 2026 22:42:33 -0400
Subject: [PATCH 6/8] test(sc): stub _teacher_logprobs_required, which main's
 advantage stage reads

Nine controller stubs in this file predate that attribute, so they pass here
and fail once merged with current main -- #3768 made _advantage_stage read it
unconditionally. Same shape as the _rollout_manager stub gap in #3783.

Signed-off-by: Tianyi Zhang <zhangtianyi975@gmail.com>
Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
---
 .../single_controller/test_single_controller_actor.py    | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/tests/unit/single_controller/test_single_controller_actor.py b/tests/unit/single_controller/test_single_controller_actor.py
index 63fd8e351a..d642579441 100644
--- a/tests/unit/single_controller/test_single_controller_actor.py
+++ b/tests/unit/single_controller/test_single_controller_actor.py
@@ -590,6 +590,7 @@ def test_advantage_stage_applies_seq_logprob_error_mask_before_streaming_train(
     ctrl._advantage_estimator = estimator
     ctrl._policy_logprobs_required = True
     ctrl._reference_logprobs_required = False
+    ctrl._teacher_logprobs_required = False
     ctrl._is_ppo = False
     ctrl._master_config = SimpleNamespace(
         grpo=_grpo_stub(seq_logprob_error_threshold=2.0)
@@ -658,6 +659,7 @@ def test_advantage_stage_reports_seq_logprob_metrics_without_masking() -> None:
     ctrl._advantage_estimator = estimator
     ctrl._policy_logprobs_required = True
     ctrl._reference_logprobs_required = False
+    ctrl._teacher_logprobs_required = False
     ctrl._is_ppo = False
     ctrl._master_config = SimpleNamespace(
         grpo=_grpo_stub(seq_logprob_error_threshold=None)
@@ -720,6 +722,7 @@ def test_advantage_stage_skips_estimator_when_seq_mask_removes_whole_chunk(
     ctrl._advantage_estimator = estimator
     ctrl._policy_logprobs_required = True
     ctrl._reference_logprobs_required = False
+    ctrl._teacher_logprobs_required = False
     ctrl._is_ppo = False
     ctrl._master_config = SimpleNamespace(
         grpo=_grpo_stub(seq_logprob_error_threshold=2.0)
@@ -775,6 +778,7 @@ def test_advantage_stage_skips_preexisting_empty_mask_without_seq_threshold() ->
     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=_grpo_stub(seq_logprob_error_threshold=None)
@@ -990,6 +994,7 @@ def _train_pump_controller(*, sampler) -> object:
     ctrl._advantage_cfg = AdvantageConfig()
     ctrl._policy_logprobs_required = False
     ctrl._reference_logprobs_required = False
+    ctrl._teacher_logprobs_required = False
     ctrl._advantage_estimator = None
     ctrl._partition_id = "rollout_data"
     ctrl._sampler = sampler
@@ -1310,6 +1315,7 @@ def test_train_pump_does_not_offload_the_policy_on_a_grpo_run(monkeypatch) -> No
     ctrl = _train_pump_controller(sampler=_ChunkedSampler(meta, chunks=2))
     ctrl._policy_logprobs_required = False
     ctrl._reference_logprobs_required = False
+    ctrl._teacher_logprobs_required = False
     ctrl._trainer = _OrderRecordingTrainer(calls)
     ctrl._sync_weights = AsyncMock(return_value=1)
     ctrl._logger = MagicMock()
@@ -1449,6 +1455,7 @@ def test_train_pump_parks_the_policy_when_neither_logprob_is_needed(
     )
     ctrl._policy_logprobs_required = False
     ctrl._reference_logprobs_required = False
+    ctrl._teacher_logprobs_required = False
     ctrl._trainer = _OrderRecordingTrainer(calls)
     ctrl._advantage_stage = AsyncMock(return_value=(meta, True))
     monkeypatch.setattr(single_controller.ray, "cluster_resources", lambda: {})
@@ -1647,6 +1654,7 @@ def compute_advantage(self, *, rewards, mask, **kwargs):
     ctrl._advantage_estimator = estimator
     ctrl._policy_logprobs_required = False
     ctrl._reference_logprobs_required = False
+    ctrl._teacher_logprobs_required = False
     ctrl._is_ppo = True
     ctrl._master_config = SimpleNamespace(
         ppo=_grpo_stub(seq_logprob_error_threshold=None)
@@ -1718,6 +1726,7 @@ def _knob_ctrl(estimator, grpo_stub):
     ctrl._advantage_estimator = estimator
     ctrl._policy_logprobs_required = False
     ctrl._reference_logprobs_required = False
+    ctrl._teacher_logprobs_required = False
     ctrl._master_config = SimpleNamespace(grpo=grpo_stub)
     ctrl._algo_cfg = grpo_stub
     ctrl._is_ppo = False

From 65352952e7a7e08654c638b336a63ab438fca624 Mon Sep 17 00:00:00 2001
From: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
Date: Wed, 26 Aug 2026 22:48:31 -0400
Subject: [PATCH 7/8] test(sc): let the estimator double return whichever shape
 the checkout has

#3512 replaces the bare-tensor return with AdvantageResult. The double now
resolves that class through the module instead of importing it, so it returns
a tensor on this branch and an AdvantageResult once #3512 lands, and this file
does not depend on a name that exists only there.

An earlier version of this commit also added _teacher_logprobs_required to
nine controller stubs. main already sets it in seven places and the suite is
green without the other two, so that was duplicate work that collided with
#3512 on the same lines -- dropped.

Signed-off-by: Tianyi Zhang <zhangtianyi975@gmail.com>
Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
---
 .../test_single_controller_actor.py           | 20 ++++++++++---------
 1 file changed, 11 insertions(+), 9 deletions(-)

diff --git a/tests/unit/single_controller/test_single_controller_actor.py b/tests/unit/single_controller/test_single_controller_actor.py
index d642579441..f5f3d29b51 100644
--- a/tests/unit/single_controller/test_single_controller_actor.py
+++ b/tests/unit/single_controller/test_single_controller_actor.py
@@ -15,6 +15,7 @@
 """Tests for SingleController initialization and pump lifecycle."""
 
 import asyncio
+import importlib
 import math
 from types import SimpleNamespace
 from unittest.mock import AsyncMock, MagicMock
@@ -590,7 +591,6 @@ def test_advantage_stage_applies_seq_logprob_error_mask_before_streaming_train(
     ctrl._advantage_estimator = estimator
     ctrl._policy_logprobs_required = True
     ctrl._reference_logprobs_required = False
-    ctrl._teacher_logprobs_required = False
     ctrl._is_ppo = False
     ctrl._master_config = SimpleNamespace(
         grpo=_grpo_stub(seq_logprob_error_threshold=2.0)
@@ -659,7 +659,6 @@ def test_advantage_stage_reports_seq_logprob_metrics_without_masking() -> None:
     ctrl._advantage_estimator = estimator
     ctrl._policy_logprobs_required = True
     ctrl._reference_logprobs_required = False
-    ctrl._teacher_logprobs_required = False
     ctrl._is_ppo = False
     ctrl._master_config = SimpleNamespace(
         grpo=_grpo_stub(seq_logprob_error_threshold=None)
@@ -722,7 +721,6 @@ def test_advantage_stage_skips_estimator_when_seq_mask_removes_whole_chunk(
     ctrl._advantage_estimator = estimator
     ctrl._policy_logprobs_required = True
     ctrl._reference_logprobs_required = False
-    ctrl._teacher_logprobs_required = False
     ctrl._is_ppo = False
     ctrl._master_config = SimpleNamespace(
         grpo=_grpo_stub(seq_logprob_error_threshold=2.0)
@@ -778,7 +776,6 @@ def test_advantage_stage_skips_preexisting_empty_mask_without_seq_threshold() ->
     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=_grpo_stub(seq_logprob_error_threshold=None)
@@ -994,7 +991,6 @@ def _train_pump_controller(*, sampler) -> object:
     ctrl._advantage_cfg = AdvantageConfig()
     ctrl._policy_logprobs_required = False
     ctrl._reference_logprobs_required = False
-    ctrl._teacher_logprobs_required = False
     ctrl._advantage_estimator = None
     ctrl._partition_id = "rollout_data"
     ctrl._sampler = sampler
@@ -1315,7 +1311,6 @@ def test_train_pump_does_not_offload_the_policy_on_a_grpo_run(monkeypatch) -> No
     ctrl = _train_pump_controller(sampler=_ChunkedSampler(meta, chunks=2))
     ctrl._policy_logprobs_required = False
     ctrl._reference_logprobs_required = False
-    ctrl._teacher_logprobs_required = False
     ctrl._trainer = _OrderRecordingTrainer(calls)
     ctrl._sync_weights = AsyncMock(return_value=1)
     ctrl._logger = MagicMock()
@@ -1455,7 +1450,6 @@ def test_train_pump_parks_the_policy_when_neither_logprob_is_needed(
     )
     ctrl._policy_logprobs_required = False
     ctrl._reference_logprobs_required = False
-    ctrl._teacher_logprobs_required = False
     ctrl._trainer = _OrderRecordingTrainer(calls)
     ctrl._advantage_stage = AsyncMock(return_value=(meta, True))
     monkeypatch.setattr(single_controller.ray, "cluster_resources", lambda: {})
@@ -1654,7 +1648,6 @@ def compute_advantage(self, *, rewards, mask, **kwargs):
     ctrl._advantage_estimator = estimator
     ctrl._policy_logprobs_required = False
     ctrl._reference_logprobs_required = False
-    ctrl._teacher_logprobs_required = False
     ctrl._is_ppo = True
     ctrl._master_config = SimpleNamespace(
         ppo=_grpo_stub(seq_logprob_error_threshold=None)
@@ -1704,7 +1697,16 @@ def __init__(self) -> None:
     def compute_advantage(self, *, rewards, mask, **kwargs):
         del kwargs
         self.rewards = rewards.clone()
-        return rewards.unsqueeze(-1).expand_as(mask).clone()
+        adv = rewards.unsqueeze(-1).expand_as(mask).clone()
+        # Bare tensor today, AdvantageResult once #3512 lands. Resolved through
+        # the module so this file does not import a name that exists only on
+        # that branch.
+        result_cls = getattr(
+            importlib.import_module("nemo_rl.algorithms.advantage_estimator"),
+            "AdvantageResult",
+            None,
+        )
+        return adv if result_cls is None else result_cls(advantages=adv)
 
 
 def _knob_ctrl(estimator, grpo_stub):

From e6ddf3d9134abd28338f45667b02387b4963b531 Mon Sep 17 00:00:00 2001
From: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
Date: Fri, 28 Aug 2026 13:29:00 -0400
Subject: [PATCH 8/8] fix(sc): log advantages pre-clip, and pin the guard that
 keeps PPO out

Two gaps in this PR's own parity claim.

grpo.py logs advantages at :3471 and clips at :3487; grpo_sync.py at :903
then :915. Both report pre-clip. This PR clipped first, so SC's
advantages/mean|max|min would have meant something no other driver's
does. Clip after the log; the data-plane write is unchanged.

The clip guard's 'not self._is_ppo' half was untested, because the PPO
test handed the ppo block a GRPO-shaped stub carrying advantage_clip_low
-- a field real PPOConfig does not declare. Deleting the clause left
every test green. PPOConfig is extra="allow", so the clause is exactly
what stops a user-set ppo.advantage_clip_low from running GRPO clipping
on a PPO run; there is a test for that now, and the PPO stub matches the
real config.

Also drop the importlib lookup for AdvantageResult: it does not exist on
main, #3512 is unmerged, and if that branch ever returned the wrapper the
production code at single_controller.py:2028 would break anyway.

The guide said setup rejects reward_scaling. It does not any more.

Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
---
 docs/guides/single-controller.md              |  2 +-
 nemo_rl/algorithms/single_controller.py       | 24 +++--
 .../test_single_controller_actor.py           | 92 +++++++++++++------
 3 files changed, 82 insertions(+), 36 deletions(-)

diff --git a/docs/guides/single-controller.md b/docs/guides/single-controller.md
index ace04ec811..6d2e2411bc 100644
--- a/docs/guides/single-controller.md
+++ b/docs/guides/single-controller.md
@@ -179,6 +179,6 @@ The SC path is still under active development. Feature gaps are tracked in [issu
 - Generation backend: only vLLM is supported and validated; Megatron generation, 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`).
 - (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.
diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py
index 40aad2eb46..8d2dd2dbd5 100644
--- a/nemo_rl/algorithms/single_controller.py
+++ b/nemo_rl/algorithms/single_controller.py
@@ -2031,21 +2031,27 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]:
             if self._is_ppo:
                 returns = torch.zeros_like(mask)
 
-        # grpo_sync.py:915 does the same, immediately before its data-plane
-        # write. A no-op unless advantage_clip_low/high are set, and GRPO-only:
-        # _clip_grpo_advantages has no counterpart in ppo.py, and PPOConfig
-        # carries neither knob, so a PPO run must not reach for them.
-        # ``hasattr``: only GRPOConfig carries the two knobs. PPOConfig does
-        # not, DistillationConfig does not, and the OPD test doubles build a
-        # bare namespace -- all of which mean "no clipping configured".
-        if not self._is_ppo and hasattr(self._algo_cfg, "advantage_clip_low"):
-            advantages = _clip_grpo_advantages(advantages, self._algo_cfg)
         response_advantages = torch.masked_select(advantages, mask.bool())
         self._step_log_dict["rewards"].append(rewards.detach().cpu())
         self._step_log_dict["masked_advantages"].append(
             response_advantages.detach().cpu()
         )
 
+        # 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
diff --git a/tests/unit/single_controller/test_single_controller_actor.py b/tests/unit/single_controller/test_single_controller_actor.py
index f5f3d29b51..5a404c5576 100644
--- a/tests/unit/single_controller/test_single_controller_actor.py
+++ b/tests/unit/single_controller/test_single_controller_actor.py
@@ -15,7 +15,6 @@
 """Tests for SingleController initialization and pump lifecycle."""
 
 import asyncio
-import importlib
 import math
 from types import SimpleNamespace
 from unittest.mock import AsyncMock, MagicMock
@@ -57,6 +56,21 @@ def _grpo_stub(**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
 
@@ -1649,9 +1663,7 @@ def compute_advantage(self, *, rewards, mask, **kwargs):
     ctrl._policy_logprobs_required = False
     ctrl._reference_logprobs_required = False
     ctrl._is_ppo = True
-    ctrl._master_config = SimpleNamespace(
-        ppo=_grpo_stub(seq_logprob_error_threshold=None)
-    )
+    ctrl._master_config = SimpleNamespace(ppo=_ppo_stub())
     ctrl._algo_cfg = ctrl._master_config.ppo
     ctrl._step_log_dict = {
         "rewards": [],
@@ -1697,29 +1709,21 @@ def __init__(self) -> None:
     def compute_advantage(self, *, rewards, mask, **kwargs):
         del kwargs
         self.rewards = rewards.clone()
-        adv = rewards.unsqueeze(-1).expand_as(mask).clone()
-        # Bare tensor today, AdvantageResult once #3512 lands. Resolved through
-        # the module so this file does not import a name that exists only on
-        # that branch.
-        result_cls = getattr(
-            importlib.import_module("nemo_rl.algorithms.advantage_estimator"),
-            "AdvantageResult",
-            None,
-        )
-        return adv if result_cls is None else result_cls(advantages=adv)
+        return rewards.unsqueeze(-1).expand_as(mask).clone()
 
 
-def _knob_ctrl(estimator, grpo_stub):
+def _knob_ctrl(estimator, grpo_stub, *, is_ppo=False):
     batch, seq = 2, 4
-    data = TensorDict(
-        {
-            "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),
-        },
-        batch_size=[batch],
-    )
+    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)
@@ -1729,9 +1733,11 @@ def _knob_ctrl(estimator, grpo_stub):
     ctrl._policy_logprobs_required = False
     ctrl._reference_logprobs_required = False
     ctrl._teacher_logprobs_required = False
-    ctrl._master_config = SimpleNamespace(grpo=grpo_stub)
+    ctrl._master_config = (
+        SimpleNamespace(ppo=grpo_stub) if is_ppo else SimpleNamespace(grpo=grpo_stub)
+    )
     ctrl._algo_cfg = grpo_stub
-    ctrl._is_ppo = False
+    ctrl._is_ppo = is_ppo
     ctrl._step_log_dict = {
         "rewards": [],
         "masked_advantages": [],
@@ -1759,6 +1765,40 @@ def test_advantage_clip_bounds_are_applied_before_the_write() -> None:
     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()