From 28e306731fa0f04f94dc7056d2427c368fd63550 Mon Sep 17 00:00:00 2001
From: Alpamys <vpn.alpamys@gmail.com>
Date: Sun, 5 Jul 2026 16:56:20 +0500
Subject: [PATCH] feat(prm): wire PRM reward into GRPO setup via
 _select_reward_fn (v0.71.30)

---
 src/soup_cli/trainer/grpo.py | 26 +++++++++++++++++----
 tests/test_v07130.py         | 44 ++++++++++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+), 5 deletions(-)

diff --git a/src/soup_cli/trainer/grpo.py b/src/soup_cli/trainer/grpo.py
index 751a984..67dbc1f 100644
--- a/src/soup_cli/trainer/grpo.py
+++ b/src/soup_cli/trainer/grpo.py
@@ -134,6 +134,24 @@ def _read_attr(obj: Any, name: str) -> Any:
         return obj.get(name)
     return getattr(obj, name, None)
 
+
+def _select_reward_fn(tcfg: Any, device: str, trust_remote_code: bool) -> Any:
+    """Choose the GRPO reward function (v0.71.30).
+
+    When ``tcfg.prm_reward`` is set, a trained Soup PRM scores each completion's
+    steps and REPLACES the configured ``reward_fn`` (process-supervision). The
+    returned callable rides the existing shaping + ``wrap_reward_funcs`` seam in
+    :meth:`GRPOTrainerWrapper.setup` unchanged, so the v0.71.26 reward-hack
+    mitigation controller observes the PRM reward for free.
+    """
+    if tcfg.prm_reward is not None:
+        from soup_cli.utils.prm_reward import build_prm_reward_fn
+
+        return build_prm_reward_fn(tcfg, device, trust_remote_code)
+    from soup_cli.trainer.rewards import load_reward_fn
+
+    return load_reward_fn(tcfg.reward_fn, verifiable_domain=tcfg.verifiable_domain)
+
 class GRPOTrainerWrapper:
     """High-level wrapper for GRPO training from SoupConfig.
 
@@ -222,11 +240,9 @@ class GRPOTrainerWrapper:
         use_unsloth = cfg.backend == "unsloth"
 
         # --- Load reward function ---
-        from soup_cli.trainer.rewards import load_reward_fn
-
-        reward_fn = load_reward_fn(
-            tcfg.reward_fn, verifiable_domain=tcfg.verifiable_domain
-        )
+        # v0.71.30 — when tcfg.prm_reward is set, a trained Soup PRM replaces
+        # the configured reward (process-supervision); otherwise load reward_fn.
+        reward_fn = _select_reward_fn(tcfg, self.device, self._trust_remote_code)
 
         # v0.71.11 #235/#240 — when the reward-hack or echo-trap detector is
         # enabled, wrap the reward function(s) with a capture shim so the
diff --git a/tests/test_v07130.py b/tests/test_v07130.py
index 6515742..6ad3c38 100644
--- a/tests/test_v07130.py
+++ b/tests/test_v07130.py
@@ -349,6 +349,50 @@ class TestBuildPrmRewardFn:
             mod.build_prm_reward_fn(_T(), device="cpu", trust_remote_code=False)
 
 
+# ---------------------------------------------------------------------------
+# Task 4 — GRPO wiring
+# ---------------------------------------------------------------------------
+class TestGrpoPrmWiring:
+    def test_prm_reward_selected(self, monkeypatch):
+        import soup_cli.trainer.grpo as grpo
+
+        sentinel = object()
+        monkeypatch.setattr(
+            "soup_cli.utils.prm_reward.build_prm_reward_fn",
+            lambda tcfg, device, trust_remote_code: sentinel,
+        )
+
+        class _T:
+            prm_reward = "./prm"
+            prm_aggregate = "min"
+            reward_fn = "accuracy"
+            verifiable_domain = None
+
+        out = grpo._select_reward_fn(_T(), "cpu", False)
+        assert out is sentinel
+
+    def test_standard_reward_selected(self, monkeypatch):
+        import soup_cli.trainer.grpo as grpo
+
+        captured = {}
+
+        def _fake_load(spec, verifiable_domain=None):
+            captured["spec"] = spec
+            return "REWARD_FN"
+
+        monkeypatch.setattr("soup_cli.trainer.rewards.load_reward_fn", _fake_load)
+
+        class _T:
+            prm_reward = None
+            prm_aggregate = "min"
+            reward_fn = "accuracy"
+            verifiable_domain = None
+
+        out = grpo._select_reward_fn(_T(), "cpu", False)
+        assert out == "REWARD_FN"
+        assert captured["spec"] == "accuracy"
+
+
 class TestNoTopLevelTorch:
     def test_prm_reward_has_no_top_level_torch(self):
         import soup_cli.utils.prm_reward as mod