Fix custom preprocess and postprocess in judger pools

tests/rl/test_judger.py

@@ -140,20 +140,27 @@ def test_native_judger_batch_judge_not_supported(self):
         with self.assertRaisesRegex(NotImplementedError, "does not support batch_judge"):
             asyncio.run(judger.batch_judge([self._make_rollout_state("correctness")]))
 
-    def test_remote_judger_uses_driver_side_preprocess_judger(self):
+    def test_remote_judger_preserves_driver_side_judger_contract(self):
         from xtuner.v1.rl.judger import Judger, RemoteJudger
 
-        class CustomPreprocessJudger(Judger):
+        class CustomJudger(Judger):
             def preprocess(self, rollout_state):
                 return {"custom_value": rollout_state.extra_fields["custom_value"]}
 
+            def postprocess(self, rollout_state, output):
+                rollout_state.reward = {
+                    "score": output["raw_score"],
+                    "postprocessed": True,
+                }
+                return rollout_state
+
         class RemoteMethod:
             def __init__(self):
                 self.payload = None
 
             async def remote(self, payload):
                 self.payload = payload
-                return {"score": payload["custom_value"]}
+                return {"raw_score": payload["custom_value"]}
 
         class FakeActor:
             def __init__(self):
@@ -162,16 +169,93 @@ def __init__(self):
         actor = FakeActor()
         judger = RemoteJudger(
             actor=actor,
-            judger_name="remote_custom_preprocess",
-            preprocess_judger=CustomPreprocessJudger(),
+            judger_name="remote_custom_contract",
+            preprocess_judger=CustomJudger(),
         )
         rollout_state = self._make_rollout_state("correctness")
         rollout_state.extra_fields["custom_value"] = 7
 
         judged_state = asyncio.run(judger.judge(rollout_state))
 
         self.assertEqual(actor.judge_payload.payload, {"custom_value": 7})
-        self.assertEqual(judged_state.reward, {"score": 7})
+        self.assertEqual(judged_state.reward, {"score": 7, "postprocessed": True})
+
+    def test_composed_judger_pool_preserves_branch_contract_and_routing(self):
+        from xtuner.v1.rl.judger import ComposedJudger, Judger, JudgerPool, RemoteJudger
+
+        class CustomJudger(Judger):
+            def preprocess(self, rollout_state):
+                return {
+                    "response": rollout_state.response,
+                    "reward_model": rollout_state.reward_model,
+                    "finish_reason": rollout_state.finish_reason,
+                    "extra_fields": rollout_state.extra_fields,
+                }
+
+            def postprocess(self, rollout_state, output):
+                rollout_state.reward = {
+                    "score": output["raw_score"],
+                    "postprocessed": True,
+                }
+                return rollout_state
+
+        class RemoteMethod:
+            def __init__(self):
+                self.payload = None
+
+            async def remote(self, payload):
+                self.payload = payload
+                return {"raw_score": payload["extra_fields"]["expected_score"]}
+
+        class FakeActor:
+            def __init__(self):
+                self.judge_payload = RemoteMethod()
+
+        branch_judger = CustomJudger()
+        actors = [FakeActor(), FakeActor()]
+        replicas = [
+            RemoteJudger(
+                actor=actor,
+                judger_name="remote_custom_contract",
+                preprocess_judger=branch_judger,
+            )
+            for actor in actors
+        ]
+        judger = ComposedJudger(
+            branches={
+                "biology": JudgerPool(
+                    replicas=replicas,
+                    judger_name="biology",
+                )
+            }
+        )
+        rollout_states = [self._make_rollout_state("biology"), self._make_rollout_state("biology")]
+        for expected_score, rollout_state in enumerate(rollout_states, start=1):
+            rollout_state.finish_reason = "stop"
+            rollout_state.extra_fields = {
+                "expected_score": expected_score,
+                "task_name": "MCC",
+            }
+
+        async def judge_all():
+            return [await judger.judge(state) for state in rollout_states]
+
+        judged_states = asyncio.run(judge_all())
+
+        for expected_score, (actor, state) in enumerate(zip(actors, judged_states), start=1):
+            self.assertEqual(
+                actor.judge_payload.payload,
+                {
+                    "response": state.response,
+                    "reward_model": state.reward_model,
+                    "finish_reason": "stop",
+                    "extra_fields": {
+                        "expected_score": expected_score,
+                        "task_name": "MCC",
+                    },
+                },
+            )
+            self.assertEqual(state.reward, {"score": expected_score, "postprocessed": True})
 
     def test_composed_judger_config(self):
         def merge_fn(original, judged):

xtuner/v1/rl/judger/native.py

@@ -229,11 +229,11 @@ class RemoteJudger(Judger):
 
     ``RemoteJudger`` keeps the same ``Judger`` interface as local judgers, so
     callers still pass ``RolloutState`` to ``judge``. This proxy runs the same
-    ``preprocess`` implementation as the actor-side judger on the driver, then
-    sends only the lightweight payload to ``JudgerActor``. ``JudgerActor`` lives
-    in the Ray worker process and owns the real local judger instance that
-    executes ``judge_payload``. Batch support is determined by that actor-side
-    judger.
+    ``preprocess`` and ``postprocess`` implementations as the actor-side judger
+    on the driver, and sends only the lightweight payload to ``JudgerActor``.
+    ``JudgerActor`` lives in the Ray worker process and owns the real local
+    judger instance that executes ``judge_payload``. Batch support is determined
+    by that actor-side judger.
     """
 
     def __init__(self, actor: RayJudgerProxy, judger_name: str, preprocess_judger: Judger | None = None):
@@ -244,17 +244,28 @@ def __init__(self, actor: RayJudgerProxy, judger_name: str, preprocess_judger: J
     # Preprocess must run on the driver before the Ray call. Otherwise the full
     # RolloutState would be serialized to the actor, and remote branches with
     # custom preprocess logic would lose the payload fields they require.
+    # Postprocess also runs on the driver because it writes the output back to
+    # the original RolloutState retained by the caller.
     def preprocess(self, rollout_state: RolloutState) -> JudgerPayload:
         if self.preprocess_judger is None:
             return super().preprocess(rollout_state)
         return self.preprocess_judger.preprocess(rollout_state)
 
+    def postprocess(self, rollout_state: RolloutState, output: JudgerOutput) -> RolloutState:
+        if self.preprocess_judger is None:
+            return super().postprocess(rollout_state, output)
+        return self.preprocess_judger.postprocess(rollout_state, output)
+
     async def judge_payload(self, payload: JudgerPayloadBatch) -> JudgerOutputBatch:
         return await self.actor.judge_payload.remote(payload)
 
 
 class JudgerPool(Judger):
-    """Round-robin dispatch across replicas of the same judger type."""
+    """Round-robin dispatch across replicas of the same judger type.
+
+    Replicas are homogeneous and therefore share the same payload contract. Preprocessing and postprocessing are
+    delegated to a representative replica, while payload execution is dispatched round-robin across all replicas.
+    """
 
     def __init__(self, replicas: list[Judger], judger_name: str):
         super().__init__(judger_name=judger_name)
@@ -265,6 +276,12 @@ def __init__(self, replicas: list[Judger], judger_name: str):
         self._lock = asyncio.Lock()
         self._worker_loads = dict.fromkeys(range(len(replicas)), 0)
 
+    def preprocess(self, rollout_state: RolloutState) -> JudgerPayload:
+        return self.replicas[0].preprocess(rollout_state)
+
+    def postprocess(self, rollout_state: RolloutState, output: JudgerOutput) -> RolloutState:
+        return self.replicas[0].postprocess(rollout_state, output)
+
     async def _pick_replica(self) -> tuple[int, Judger]:
         async with self._lock:
             replica_idx = self._rr_index % len(self.replicas)