[reward, trainer] feat: add managed multi-reward deployments

docs/examples/config.md

@@ -334,7 +334,8 @@ Diffusion recipes compose `reward@reward: reward` (`verl_omni/trainer/config/rew
 
 - `reward.num_workers`: Parallel reward-manager workers.
 - `reward.custom_reward_function.path` / `name`: Single custom score function.
-- `reward.reward_functions`: Multi-reward dict (`{name: {path, name, weight}}`); mutually exclusive with `custom_reward_function`.
+- `reward.reward_functions`: Multi-reward dict (`{name: {path, name, weight}}`). A term uses a same-name entry in `reward.models` automatically; set `model` only when the names differ.
+- `reward.models`: Optional named `engine` and `native` reward models. See {doc}`../algo/named_reward_models` for lifecycle, pool, placement, and extension details.
 - `reward.aggregation`: Multi-reward aggregation (`weighted_sum` only).
 - `reward.reward_manager`: Defaults to `VisualRewardManager` from `pkg://verl_omni.reward_loop.reward_manager`.
 - `reward.reward_model.*`: Optional model-based RM (resource pool, rollout engine knobs). See {doc}`../algo/async_reward` and {doc}`../start/http_scorer`.

docs/index.md

@@ -46,6 +46,7 @@ examples/config.md
 :caption: Advanced Features
 
 algo/async_reward.md
+algo/named_reward_models.md
 algo/rollout_correction.md
 algo/separate_async_omni.md
 start/rollout_batching.md

examples/flowgrpo_trainer/qwen_image_edit/README.md

@@ -98,29 +98,23 @@ bash examples/flowgrpo_trainer/qwen_image_edit/run_qwen_image_edit_lora.sh \
 
 ### Ascend NPU
 
-The NPU recipe uses the synchronous V1 diffusion trainer and assigns each
-local PickScore reward worker to an accelerator bundle from the actor/rollout
-resource pool:
+The NPU recipe uses the synchronous V1 diffusion trainer and a named native
+PickScore model. Its `placement.devices` list assigns one full PickScore
+model worker to each selected parent-pool bundle:
 
 ```bash
 WORKSPACE=$PWD \
 NUM_GPUS_ACTOR_ROLLOUT_REWARD=16 \
 bash examples/flowgrpo_trainer/qwen_image_edit/run_qwen_image_edit_lora_v1_npu.sh
 ```
 
-The launcher enables accelerator placement for the local reward workers:
-
-```text
-reward.custom_reward_function.use_accelerator=True
-```
-
 Its checked-in resource layout uses 16 NPUs, rollout tensor parallelism 4, and
-16 local reward workers, creating four rollout replicas. Each completed sample
-selects from all configured reward workers at random; rollout tensor
-parallelism does not partition or bind the reward worker set. Adjust
-`ROLLOUT_TP` for the rollout topology and `REWARD_WORKERS` for reward
-parallelism when adapting the recipe. The model, reward function, and parquet
-inputs use the same environment overrides listed below.
+16 native PickScore workers, creating four rollout replicas. The default
+native placement is `[0, ..., 15]`; these are native-subpool bundle indices,
+not physical NPU IDs or tensor-parallel ranks. Adjust `ROLLOUT_TP` for rollout
+topology and set `NATIVE_REWARD_DEVICES` explicitly when changing the number
+of native PickScore replicas. The model and parquet inputs use the same
+environment overrides listed below.
 
 By default, the launcher reads:
 
@@ -144,9 +138,20 @@ Set `TRAIN_FILES` and `VAL_FILES` to use different parquet files.
 | `REWARD_WORKERS` | `4` | Asynchronous reward worker count. |
 | `IMAGE_RESOLUTION` | `512` | Square target output resolution. |
 | `MAX_PROMPT_LENGTH` | `8192` | Token and prompt-embedding length limit. |
-| `REWARD_FUNCTION_PATH` | `pkg://verl_omni.utils.reward_score.pickscore_reward` | Reward module import path. |
+| `PICKSCORE_MODEL_PATH` | `yuvalkirstain/PickScore_v1` | PickScore checkpoint for the native reward model. |
+| `NATIVE_REWARD_DEVICES` | `[0,1,2,3]` (CUDA) / `[0,...,15]` (NPU) | Native-subpool bundle indices; one full PickScore instance per entry. |
+| `REWARD_OFFLOAD` | `true` | `true` wakes/sleeps around scoring; `false` keeps the reward model resident. The meaning is identical for engine and native models. |
+
+The launcher configures `reward.models.pickscore.backend=native`; the same-name
+`reward.reward_functions.pickscore` entry binds automatically. Native workers
+preserve the PickScore model's local inference queue; they are not TP shards.
+The launcher explicitly configures the PickScore model class through
+`executor.model`; the generic native backend does not select model-specific
+behavior.
 
-The launcher selects `compute_score_pickscore` from the reward module.
+See [Named Reward Models](../../../docs/algo/named_reward_models.md)
+for native-only and engine-only configurations, mixed-model resource placement,
+lifecycle settings, and the custom-model extension contract.
 
 Additional Hydra overrides can be appended to the command:
 
@@ -160,7 +165,8 @@ bash examples/flowgrpo_trainer/qwen_image_edit/run_qwen_image_edit_lora.sh \
 
 Keep `actor_rollout_ref.rollout.n` greater than one for group-relative
 advantages. When reducing GPU count, also reduce the training batch size,
-rollout count, micro-batch sizes, and reward worker count to fit memory.
+rollout count, micro-batch sizes, and the explicit native placement list to
+fit memory.
 
 The launcher enables console, TensorBoard, and W&B logging by default. The
 first command overrides this with `trainer.logger=console` so it can run

examples/flowgrpo_trainer/qwen_image_edit/run_qwen_image_edit_lora_v1_npu.sh

@@ -6,11 +6,14 @@ set -x
 export VLLM_ASCEND_ENABLE_NZ=0
 export VERL_DATAPROTO_SERIALIZATION_METHOD=numpy
 model_name=${MODEL_PATH:-Qwen/Qwen-Image-Edit-2511}
-reward_function_path=${REWARD_FUNCTION_PATH:-pkg://verl_omni.utils.reward_score.pickscore_reward}
+pickscore_model_path=${PICKSCORE_MODEL_PATH:?PICKSCORE_MODEL_PATH must be set}
 
 NUM_GPUS_ACTOR_ROLLOUT_REWARD=${NUM_GPUS_ACTOR_ROLLOUT_REWARD:-16}
 ROLLOUT_TP=${ROLLOUT_TP:-4}
-REWARD_WORKERS=${REWARD_WORKERS:-16}
+# Parent-pool bundle indices. Each entry loads one complete PickScore model;
+# these are not physical NPU IDs and are not tensor-parallel ranks.
+NATIVE_REWARD_DEVICES=${NATIVE_REWARD_DEVICES:-"[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15]"}
+REWARD_OFFLOAD=${REWARD_OFFLOAD:-true}
 IMAGE_RESOLUTION=${IMAGE_RESOLUTION:-512}
 MAX_PROMPT_LENGTH=${MAX_PROMPT_LENGTH:-1024}
 
@@ -60,7 +63,6 @@ python3 -m verl_omni.trainer.main_diffusion_v1 \
     actor_rollout_ref.rollout.n=8 \
     actor_rollout_ref.model.attn_backend='_native_npu' \
     actor_rollout_ref.rollout.rollout_attn_backend=TORCH_SDPA \
-    reward.custom_reward_function.use_accelerator=True \
     trainer.device=npu \
     trainer.use_v1=True \
     actor_rollout_ref.rollout.cudagraph_capture_sizes="[1,2,4,8,16,32,64,128,256,384,512,640,768,896,1024]" \
@@ -83,10 +85,15 @@ python3 -m verl_omni.trainer.main_diffusion_v1 \
     actor_rollout_ref.rollout.val_kwargs.pipeline.num_inference_steps=40 \
     actor_rollout_ref.rollout.val_kwargs.algo.noise_level=0.0 \
     actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \
-    reward.num_workers=$REWARD_WORKERS \
     reward.reward_model.enable=False \
-    reward.custom_reward_function.path=$reward_function_path \
-    reward.custom_reward_function.name=compute_score_pickscore \
+    reward.reward_manager.name=MultiVisualRewardManager \
+    +reward.models.pickscore.backend=native \
+    +reward.models.pickscore.offload=$REWARD_OFFLOAD \
+    +reward.models.pickscore.model_path=$pickscore_model_path \
+    +reward.models.pickscore.placement.devices="$NATIVE_REWARD_DEVICES" \
+    +reward.models.pickscore.executor.model=verl_omni.utils.reward_score.pickscore_reward:PickScoreNativeModel \
+    +reward.reward_functions.pickscore.path=pkg://verl_omni.utils.reward_score.pickscore_reward \
+    +reward.reward_functions.pickscore.name=compute_score_pickscore_native \
     trainer.logger='["console", "tensorboard"]' \
     trainer.project_name=flow_grpo \
     trainer.experiment_name=qwen_image_edit_lora_pickscore \

tests/reward_loop/test_multi_reward_manager_on_cpu.py

@@ -67,6 +67,32 @@ async def reward_asserts_float_latent_contract(data_source, solution_image, grou
     return float(solution_image[0, 0, 0])
 
 
+async def reward_uses_named_engine_router(reward_router_address, model_name):
+    assert reward_router_address == "engine-router"
+    assert model_name == "ocr-model"
+    return {"score": 0.6, "backend": "engine-function"}
+
+
+async def reward_uses_native_model(reward_model, ground_truth, solution_image):
+    output = await reward_model.infer(prompt=ground_truth, image=solution_image)
+    return {"score": output["value"], "backend": "native-function"}
+
+
+class _NativeModelExecutor:
+    def reward_kwargs(self):
+        return {"reward_model": self}
+
+    async def infer(self, prompt, image):
+        assert prompt == "hello"
+        assert image.dtype == torch.uint8
+        return {"value": 0.75}
+
+
+class _EngineRouterClient:
+    def reward_kwargs(self):
+        return {"reward_router_address": "engine-router", "model_name": "ocr-model"}
+
+
 # ---------------------------------------------------------------------------
 # Helpers
 # ---------------------------------------------------------------------------
@@ -253,20 +279,36 @@ def test_weighted_aggregation(self):
         assert "reward/dict_result/score" not in result["reward_extra_info"]
         assert result["reward_extra_info"]["reward/combined"] == pytest.approx(2.0)
 
-    def test_exception_contributes_zero(self):
-        """A failing sub-reward contributes 0 without breaking others."""
+    def test_required_exception_fails_fast(self):
+        """A failing required sub-reward aborts reward computation."""
         reward_fns = {
             "good": {"path": DUMMY_REWARDS_PATH, "name": "reward_fixed_score", "weight": 1.0},
-            "bad": {"path": DUMMY_REWARDS_PATH, "name": "reward_raises", "weight": 1.0},
+            "bad": {
+                "path": DUMMY_REWARDS_PATH,
+                "name": "reward_raises",
+                "weight": 1.0,
+                "required": True,
+            },
         }
         manager = _build_manager(reward_fns)
         data = _make_single_data()
 
-        result = manager.loop.run_until_complete(manager.run_single(data))
+        with pytest.raises(RuntimeError, match="Required sub-reward 'bad' failed: intentional failure"):
+            manager.loop.run_until_complete(manager.run_single(data))
 
-        # combined = 1.0 * 0.5 + 1.0 * 0.0 = 0.5
-        assert result["reward_score"] == pytest.approx(0.5)
+    def test_optional_exception_contributes_zero(self):
+        """A failing optional sub-reward records the error and contributes zero."""
+        reward_fns = {
+            "good": {"path": DUMMY_REWARDS_PATH, "name": "reward_fixed_score", "weight": 2.0},
+            "bad": {"path": DUMMY_REWARDS_PATH, "name": "reward_raises", "weight": 3.0},
+        }
+        manager = _build_manager(reward_fns)
+
+        result = manager.loop.run_until_complete(manager.run_single(_make_single_data()))
+
+        assert result["reward_score"] == pytest.approx(1.0)
         assert result["reward_extra_info"]["reward/bad"] == pytest.approx(0.0)
+        assert result["reward_extra_info"]["reward/bad/errors"] == 1
 
     def test_async_reward_function(self):
         """Async reward functions are awaited correctly."""
@@ -280,6 +322,72 @@ def test_async_reward_function(self):
 
         assert result["reward_score"] == pytest.approx(0.8)
 
+    def test_mixes_rule_engine_and_native_models(self):
+        manager = _build_manager(
+            {
+                "rule": {"path": DUMMY_REWARDS_PATH, "name": "reward_fixed_score", "weight": 1.0},
+                "engine": {
+                    "model": "ocr_engine",
+                    "path": DUMMY_REWARDS_PATH,
+                    "name": "reward_uses_named_engine_router",
+                    "weight": 2.0,
+                },
+                "native": {
+                    "model": "native_pickscore",
+                    "path": DUMMY_REWARDS_PATH,
+                    "name": "reward_uses_native_model",
+                    "weight": 1.0,
+                },
+            }
+        )
+        manager.set_reward_executors(
+            {"ocr_engine": _EngineRouterClient()}, {"native_pickscore": _NativeModelExecutor()}
+        )
+
+        result = manager.loop.run_until_complete(manager.run_single(_make_single_data()))
+
+        assert result["reward_score"] == pytest.approx(2.45)
+        assert result["reward_extra_info"]["reward/rule"] == pytest.approx(0.5)
+        assert result["reward_extra_info"]["reward/engine"] == pytest.approx(0.6)
+        assert result["reward_extra_info"]["reward/native"] == pytest.approx(0.75)
+        assert result["reward_extra_info"]["reward/engine/backend"] == "engine-function"
+        assert result["reward_extra_info"]["reward/native/backend"] == "native-function"
+
+    def test_native_reward_function_receives_only_inference_handle(self):
+        manager = _build_manager(
+            {
+                "native": {
+                    "model": "native_model",
+                    "path": DUMMY_REWARDS_PATH,
+                    "name": "reward_uses_native_model",
+                }
+            }
+        )
+        manager.set_reward_executors(None, {"native_model": _NativeModelExecutor()})
+
+        result = manager.loop.run_until_complete(manager.run_single(_make_single_data()))
+
+        assert result["reward_score"] == pytest.approx(0.75)
+        assert result["reward_extra_info"]["reward/native/backend"] == "native-function"
+
+    def test_engine_reward_function_uses_its_named_router(self):
+        manager = _build_manager(
+            {
+                "ocr": {
+                    "model": "ocr_engine",
+                    "path": DUMMY_REWARDS_PATH,
+                    "name": "reward_uses_named_engine_router",
+                    "weight": 1.0,
+                },
+            }
+        )
+        manager.set_reward_executors({"ocr_engine": _EngineRouterClient()}, None)
+
+        result = manager.loop.run_until_complete(manager.run_single(_make_single_data()))
+
+        assert result["reward_score"] == pytest.approx(0.6)
+        assert result["reward_extra_info"]["reward/ocr/backend"] == "engine-function"
+
     def test_jpeg_reward_via_file_path(self):
         """JPEG reward loaded via file path must not fail on relative imports."""
         reward_fns = {
@@ -330,3 +438,34 @@ def test_loads_sub_rewards(self):
         assert len(manager._sub_rewards) == 1
         assert manager._sub_rewards[0]["key"] == "a"
         assert manager._sub_rewards[0]["weight"] == 0.5
+        assert manager._sub_rewards[0]["required"] is False
+
+    @pytest.mark.parametrize(
+        ("configured", "expected"),
+        [(True, True), (False, False), ("true", True), ("false", False)],
+    )
+    def test_parses_required(self, configured, expected):
+        manager = _build_manager(
+            {
+                "a": {
+                    "path": DUMMY_REWARDS_PATH,
+                    "name": "reward_fixed_score",
+                    "required": configured,
+                }
+            }
+        )
+
+        assert manager._sub_rewards[0]["required"] is expected
+
+    @pytest.mark.parametrize("configured", ["yes", 1, None])
+    def test_rejects_invalid_required(self, configured):
+        with pytest.raises((TypeError, ValueError), match="required"):
+            _build_manager(
+                {
+                    "a": {
+                        "path": DUMMY_REWARDS_PATH,
+                        "name": "reward_fixed_score",
+                        "required": configured,
+                    }
+                }
+            )