Support DROID policy server for Cosmos3 OpenPI

recipes/cosmos3/Cosmos3-Nano.md

@@ -47,6 +47,10 @@ mode is selected per request:
   **`nvidia/Cosmos3-Nano-Policy-DROID`** is served the same way
   (`domain_name=droid_lerobot`).
 
+- **DROID OpenPI policy server** — serve `nvidia/Cosmos3-Nano-Policy-DROID` and
+  connect an OpenPI-compatible websocket client to `/v1/realtime/robot/openpi`.
+  This path returns action chunks directly instead of an mp4.
+
   Action requests can use `input_reference` or `video_reference` for video input.
   `policy` and `forward_dynamics` can also use an image reference; `inverse_dynamics`
   requires a video reference.
@@ -245,6 +249,39 @@ VIDEO_ID=$(curl -sS -X POST http://localhost:8000/v1/videos \
 # poll until status == completed, then:
 curl -sS "http://localhost:8000/v1/videos/$VIDEO_ID" | jq '.action | {shape, dtype, raw_action_dim, domain_id}'
 curl -sS -L "http://localhost:8000/v1/videos/$VIDEO_ID/content" -o cosmos3_inverse_dynamics.mp4
+
+# DROID OpenPI policy server (websocket action serving).
+# Requires cosmos_framework on PYTHONPATH because the pipeline reuses the
+# reference RoboLab action transforms. If your checkpoint config already
+# includes policy_server_config, omit the stage_overrides file and flag.
+cat > cosmos3_droid_openpi_stage_overrides.json <<'JSON'
+{
+  "0": {
+    "model_config": {
+      "policy_server_config": {
+        "image_resolution": [540, 640],
+        "n_external_cameras": 2,
+        "needs_wrist_camera": true,
+        "needs_stereo_camera": false,
+        "needs_session_id": true,
+        "action_space": "joint_position"
+      }
+    }
+  }
+}
+JSON
+
+vllm serve nvidia/Cosmos3-Nano-Policy-DROID \
+  --omni \
+  --host 0.0.0.0 --port 8000 \
+  --model-class-name Cosmos3OmniDiffusersPipeline \
+  --no-guardrails \
+  --stage-overrides "$(cat cosmos3_droid_openpi_stage_overrides.json)"
+
+# Point an OpenPI websocket client at:
+#   ws://localhost:8000/v1/realtime/robot/openpi
+# The first server message is policy_server_config. Each infer request sends a
+# msgpack-numpy observation dict and receives a writable float32 action array.
 ```
 
 #### Notes
@@ -272,6 +309,13 @@ curl -sS -L "http://localhost:8000/v1/videos/$VIDEO_ID/content" -o cosmos3_inver
   For V2V, `condition_frame_indexes_vision` selects the clean conditioned latent
   frame indexes (default `[0, 1]`), and `condition_video_keep` selects whether the
   API decodes the first or last needed reference frames (`"first"` by default).
+- **DROID OpenPI observations:** include a string `prompt`, either
+  `observation/image` or the three-view DROID camera keys
+  (`observation/wrist_image_left`, `observation/exterior_image_1_left`,
+  `observation/exterior_image_2_left`), plus `observation/gripper_position` and
+  `observation/joint_position`. Optional extra params include `history_length`,
+  `conditioning_fps`, `action_chunk_size`, `raw_action_dim`, `deterministic_seed`,
+  and `session_id`.
 - **Known limitations:**
   - Guardrails-on requires `cosmos-guardrail` **and** access to the gated
     `nvidia/Cosmos-1.0-Guardrail` repo (accept license + `HF_TOKEN`); otherwise

requirements/common.txt

@@ -22,3 +22,4 @@ prettytable>=3.8.0
 aenum==3.1.16
 pyzmq>=25.0.0
 janus>=1.0.0
+msgpack>=1.0.0

tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py

@@ -5,6 +5,7 @@
 
 import sys
 import types
+from dataclasses import dataclass
 from types import SimpleNamespace
 from typing import Any
 
@@ -230,6 +231,8 @@ def test_pipeline_registered_and_exported() -> None:
     from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline
     from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
     from vllm_omni.diffusion.registry import (
+        _DIFFUSION_ACTION_POST_PROCESS_FUNCS,
+        _DIFFUSION_IR_OP_PRIORITY_FUNCS,
         _DIFFUSION_MODELS,
         _DIFFUSION_POST_PROCESS_FUNCS,
         _DIFFUSION_PRE_PROCESS_FUNCS,
@@ -245,6 +248,10 @@ def test_pipeline_registered_and_exported() -> None:
     )
     assert _DIFFUSION_PRE_PROCESS_FUNCS["Cosmos3OmniDiffusersPipeline"] == "get_cosmos3_pre_process_func"
     assert _DIFFUSION_POST_PROCESS_FUNCS["Cosmos3OmniDiffusersPipeline"] == "get_cosmos3_post_process_func"
+    assert (
+        _DIFFUSION_ACTION_POST_PROCESS_FUNCS["Cosmos3OmniDiffusersPipeline"] == "get_cosmos3_action_post_process_func"
+    )
+    assert _DIFFUSION_IR_OP_PRIORITY_FUNCS["Cosmos3OmniDiffusersPipeline"] == "get_cosmos3_ir_op_priority_func"
     assert "Cosmos3OmniDiffusersPipeline" in CUSTOM_DIT_ENABLERS
     assert "Cosmos3OmniDiffusersPipeline" in cosmos3.__all__
 
@@ -413,6 +420,73 @@ def test_postprocess_handles_image_video_audio_and_validation() -> None:
         func({"image": video, "video": video})
 
 
+def test_action_postprocess_handles_robolab_policy_outputs() -> None:
+    from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import (
+        RoboLabPolicyInputs,
+        get_cosmos3_action_post_process_func,
+        make_robolab_action_postprocess_inputs,
+    )
+
+    func = get_cosmos3_action_post_process_func(SimpleNamespace())
+    inputs = RoboLabPolicyInputs(
+        prompt="Pick the cube.",
+        video_tensor=torch.zeros(1, 3, 3, 16, 16),
+        action_tensor=torch.zeros(2, 2),
+        action_condition_indexes=[0],
+        action_start_frame_offset=1,
+        raw_action_dim=2,
+        domain_id=7,
+        fps=15.0,
+        height=16,
+        width=16,
+        image_size=None,
+        num_frames=3,
+        num_inference_steps=4,
+        guidance_scale=3.0,
+        flow_shift=5.0,
+        seed=11,
+        history_length=1,
+        action_space="joint_pos",
+        observation={},
+    )
+
+    action = torch.tensor([[[0.0, 0.25], [1.0, 0.75]]])
+    custom_output = {"robolab_action_postprocess": make_robolab_action_postprocess_inputs(inputs)}
+    processed = func(action, custom_output=custom_output)
+
+    assert processed.shape == (1, 2)
+    assert processed.dtype == torch.zeros((), dtype=torch.float32).numpy().dtype
+    torch.testing.assert_close(torch.from_numpy(processed), torch.tensor([[1.0, 0.25]]))
+    assert "robolab_action_postprocess" not in custom_output
+
+
+def test_ir_op_priority_hook_preserves_platform_fields(monkeypatch: pytest.MonkeyPatch) -> None:
+    from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import get_cosmos3_ir_op_priority_func
+
+    @dataclass
+    class FakeIrOpPriorityConfig:
+        rms_norm: list[str]
+        fused_add_rms_norm: list[str]
+        custom_op: list[str]
+
+    fake_kernel = types.ModuleType("vllm.config.kernel")
+    fake_kernel.IrOpPriorityConfig = FakeIrOpPriorityConfig
+    monkeypatch.setitem(sys.modules, fake_kernel.__name__, fake_kernel)
+
+    func = get_cosmos3_ir_op_priority_func(SimpleNamespace())
+    default_priority = FakeIrOpPriorityConfig(
+        rms_norm=["vllm_c", "native"],
+        fused_add_rms_norm=["vllm_c", "native"],
+        custom_op=["platform_kernel", "native"],
+    )
+
+    merged = func(default_priority, vllm_config=SimpleNamespace())
+
+    assert merged.rms_norm == ["native"]
+    assert merged.fused_add_rms_norm == ["native"]
+    assert merged.custom_op == ["platform_kernel", "native"]
+
+
 def test_prompt_formatting_and_checkpoint_key_remap(make_cosmos3_pipeline) -> None:
     from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline
 
@@ -762,6 +836,95 @@ def test_forward_i2v_sound_and_action_routes(self, make_cosmos3_pipeline) -> Non
         )
         assert captured["diffuse_calls"][-1]["shared_kwargs"]["action_domain_ids"].tolist() == [7]
         assert output.custom_output["action"].shape == (1, 2, 2)
+        assert "action_only_output" not in output.custom_output
+
+    def test_forward_dispatches_robolab_policy_flow(
+        self,
+        make_cosmos3_pipeline,
+        monkeypatch: pytest.MonkeyPatch,
+    ) -> None:
+        from vllm_omni.diffusion.models.cosmos3 import pipeline_cosmos3
+
+        pipeline = make_cosmos3_pipeline()
+        pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, action_gen=True, action_dim=4)
+        captured = self._install_forward_stubs(pipeline)
+        video_latents = torch.zeros(1, 2, 1, 1, 1)
+        velocity_mask = torch.ones(1, 1, 1, 1, 1)
+        condition_latents = torch.zeros_like(video_latents)
+
+        inputs = pipeline_cosmos3.RoboLabPolicyInputs(
+            prompt="Pick the cube.",
+            video_tensor=torch.zeros(1, 3, 3, 16, 16),
+            action_tensor=torch.zeros(2, 2),
+            action_condition_indexes=[0],
+            action_start_frame_offset=1,
+            raw_action_dim=2,
+            domain_id=7,
+            fps=15.0,
+            height=16,
+            width=16,
+            image_size=None,
+            num_frames=3,
+            num_inference_steps=4,
+            guidance_scale=3.0,
+            flow_shift=5.0,
+            seed=11,
+            history_length=1,
+            action_space="joint_pos",
+            observation={},
+        )
+
+        def fake_prepare_action_latents(**kwargs):
+            captured["prepare_action"] = kwargs
+            action_chunk_size = kwargs["action_chunk_size"]
+            raw_action_dim = int(kwargs["raw_action_dim"])
+            return (
+                torch.zeros(1, action_chunk_size, 4),
+                torch.ones(1, action_chunk_size, 1),
+                torch.zeros(1, action_chunk_size, 4),
+                raw_action_dim,
+            )
+
+        def fake_prepare_action_video(*args, **kwargs):
+            captured["prepare_action_video"] = {"args": args, "kwargs": kwargs}
+            return video_latents, velocity_mask, condition_latents
+
+        monkeypatch.setattr(
+            pipeline_cosmos3,
+            "build_robolab_unipc_scheduler",
+            lambda num_steps, shift, device: StubScheduler(list(range(num_steps, 0, -1)), flow_shift=shift),
+        )
+        pipeline._build_robolab_policy_inputs = lambda sp, prompt_data, request_id=None: inputs
+        pipeline._prepare_action_latents = fake_prepare_action_latents
+        pipeline._prepare_latents_action_video = fake_prepare_action_video
+        pipeline._decode_latents = lambda latents: (_ for _ in ()).throw(
+            AssertionError("RoboLab should not decode video")
+        )
+
+        output = pipeline.forward(SimpleNamespace(prompts=["ignored"], sampling_params=make_sampling_params()))
+
+        assert captured["format"] == {
+            "prompt": "Pick the cube.",
+            "negative_prompt": "",
+            "num_frames": 3,
+            "frame_rate": 15.0,
+            "height": 16,
+            "width": 16,
+            "is_t2i": False,
+        }
+        assert "flow_shifts" not in captured
+        assert pipeline.scheduler.set_timesteps_calls == []
+        assert captured["prepare_action"]["clean_action"] is inputs.action_tensor
+        assert captured["prepare_action"]["condition_indexes"] == [0]
+        assert captured["prepare_action_video"]["kwargs"] == {"image_size": None}
+        assert captured["diffuse_calls"][-1]["shared_kwargs"]["action_domain_ids"].tolist() == [7]
+        assert captured["diffuse_calls"][-1]["timesteps"].tolist() == [4, 3, 2, 1]
+        assert output.output == {}
+        assert output.custom_output["action_only_output"] is True
+        assert output.custom_output["action"].shape == (1, 2, 2)
+        assert "actions" not in output.custom_output
+        assert "robolab_action_postprocess" in output.custom_output
+        assert "robolab_policy_inputs" not in output.custom_output
 
     @pytest.mark.parametrize(
         ("prompt", "sampling_params", "message"),

tests/diffusion/test_diffusion_plugin_hooks.py

@@ -10,11 +10,14 @@
 - Worker integration: model runner resolved via platform hook
 """
 
-from unittest.mock import patch
+from types import SimpleNamespace
+from unittest.mock import Mock, patch
 
 import pytest
 
 from vllm_omni.diffusion.registry import (
+    _DIFFUSION_ACTION_POST_PROCESS_FUNCS,
+    _DIFFUSION_IR_OP_PRIORITY_FUNCS,
     _DIFFUSION_MODELS,
     _DIFFUSION_POST_PROCESS_FUNCS,
     _DIFFUSION_PRE_PROCESS_FUNCS,
@@ -59,13 +62,19 @@ def cleanup_registry(self):
         original_models = _DIFFUSION_MODELS.copy()
         original_pre = _DIFFUSION_PRE_PROCESS_FUNCS.copy()
         original_post = _DIFFUSION_POST_PROCESS_FUNCS.copy()
+        original_action_post = _DIFFUSION_ACTION_POST_PROCESS_FUNCS.copy()
+        original_ir_op_priority = _DIFFUSION_IR_OP_PRIORITY_FUNCS.copy()
         yield
         _DIFFUSION_MODELS.clear()
         _DIFFUSION_MODELS.update(original_models)
         _DIFFUSION_PRE_PROCESS_FUNCS.clear()
         _DIFFUSION_PRE_PROCESS_FUNCS.update(original_pre)
         _DIFFUSION_POST_PROCESS_FUNCS.clear()
         _DIFFUSION_POST_PROCESS_FUNCS.update(original_post)
+        _DIFFUSION_ACTION_POST_PROCESS_FUNCS.clear()
+        _DIFFUSION_ACTION_POST_PROCESS_FUNCS.update(original_action_post)
+        _DIFFUSION_IR_OP_PRIORITY_FUNCS.clear()
+        _DIFFUSION_IR_OP_PRIORITY_FUNCS.update(original_ir_op_priority)
 
     def test_register_new_model(self):
         """Test registering a new diffusion model with pre/post process functions."""
@@ -75,6 +84,8 @@ def test_register_new_model(self):
             class_name="TestPipeline",
             pre_process_func_name="test_pre_process",
             post_process_func_name="test_post_process",
+            action_post_process_func_name="test_action_post_process",
+            ir_op_priority_func_name="test_ir_op_priority",
         )
         assert "TestPipeline" in _DIFFUSION_MODELS
         assert _DIFFUSION_MODELS["TestPipeline"] == (
@@ -84,6 +95,8 @@ def test_register_new_model(self):
         )
         assert _DIFFUSION_PRE_PROCESS_FUNCS["TestPipeline"] == "test_pre_process"
         assert _DIFFUSION_POST_PROCESS_FUNCS["TestPipeline"] == "test_post_process"
+        assert _DIFFUSION_ACTION_POST_PROCESS_FUNCS["TestPipeline"] == "test_action_post_process"
+        assert _DIFFUSION_IR_OP_PRIORITY_FUNCS["TestPipeline"] == "test_ir_op_priority"
 
 
 class TestWorkerUsesHook:
@@ -108,3 +121,22 @@ def test_model_runner_resolved_via_platform(self, mock_platform, mock_resolve):
         assert worker.model_runner is mock_runner_instance
         mock_platform.get_diffusion_model_runner_cls.assert_called_once()
         mock_resolve.assert_called_once_with("custom.path")
+
+    @patch("vllm_omni.diffusion.worker.diffusion_worker.get_diffusion_ir_op_priority_func")
+    @patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
+    def test_ir_op_priority_hook_receives_platform_default(self, mock_platform, mock_get_hook):
+        """Test model IR priority hook merges from the platform default."""
+        from vllm_omni.diffusion.worker.diffusion_worker import _resolve_ir_op_priority
+
+        od_config = SimpleNamespace(model_class_name="TestPipeline")
+        vllm_config = SimpleNamespace()
+        default_priority = object()
+        merged_priority = object()
+        hook = Mock(return_value=merged_priority)
+        mock_platform.get_default_ir_op_priority.return_value = default_priority
+        mock_get_hook.return_value = hook
+
+        assert _resolve_ir_op_priority(od_config, vllm_config) is merged_priority
+        mock_platform.get_default_ir_op_priority.assert_called_once_with(vllm_config)
+        mock_get_hook.assert_called_once_with(od_config)
+        hook.assert_called_once_with(default_priority, vllm_config=vllm_config)