[Diffusion] DreamZero world model integration with CFG parallel + OpenPI serving (#2162)

Signed-off-by: Meng <meng_chen99@163.com>
Signed-off-by: Yangshen Deng <yangshen.d@outlook.com>
Signed-off-by: Suli Wang <50828270+wsl2000@users.noreply.github.com>
Co-authored-by: Meng <meng_chen99@163.com>
Co-authored-by: Suli Wang <50828270+wsl2000@users.noreply.github.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>

tests/diffusion/test_diffusion_scheduler.py

@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-import asyncio
 import queue
 import threading
 from types import SimpleNamespace
@@ -447,11 +446,7 @@ class TestDiffusionEngine:
     @pytest.mark.asyncio
     async def test_step_raises_aborted_error(self, mocker: MockerFixture) -> None:
         engine = DiffusionEngine.__new__(DiffusionEngine)
-        engine._closed = False
-        engine._loop_started = True
-        engine._init_lock = asyncio.Lock()
-        engine.main_loop = asyncio.get_running_loop()
-        engine.stop_event = threading.Event()
+        engine._check_and_start_background_loop = mocker.AsyncMock()
         engine.pre_process_func = None
         engine.async_add_req_and_wait_for_response = mocker.AsyncMock(
             return_value=DiffusionOutput(aborted=True, abort_message="Request req-abort aborted.")
@@ -526,6 +521,52 @@ class TestDiffusionEngine:
         with pytest.raises(RuntimeError, match="Dummy run failed: boom"):
             engine._dummy_run()
 
+    @pytest.mark.asyncio
+    async def test_step_multi_request_reuses_multimodal_slice_logic(self, mocker: MockerFixture) -> None:
+        engine = DiffusionEngine.__new__(DiffusionEngine)
+        engine.od_config = SimpleNamespace(
+            model_class_name="mock_model",
+            enable_cpu_offload=False,
+        )
+        engine.pre_process_func = None
+        engine.post_process_func = None
+        engine._check_and_start_background_loop = mocker.AsyncMock()
+        engine.async_add_req_and_wait_for_response = mocker.AsyncMock(
+            return_value=DiffusionOutput(
+                output={
+                    "video": ["frame-0", "frame-1"],
+                    "audio": ["audio-0", "audio-1"],
+                    "actions": torch.tensor([[1.0, 2.0], [3.0, 4.0]]),
+                }
+            )
+        )
+
+        request = OmniDiffusionRequest(
+            prompts=["prompt-0", "prompt-1"],
+            sampling_params=OmniDiffusionSamplingParams(
+                num_inference_steps=1,
+                num_outputs_per_prompt=1,
+            ),
+            request_id="req-batch",
+        )
+
+        mocker.patch("vllm_omni.diffusion.diffusion_engine.supports_audio_output", return_value=False)
+        outputs = await engine.step(request)
+
+        assert len(outputs) == 2
+        assert outputs[0].images == ["frame-0"]
+        assert outputs[1].images == ["frame-1"]
+        assert outputs[0].multimodal_output["audio"] == "audio-0"
+        assert outputs[1].multimodal_output["audio"] == "audio-1"
+        torch.testing.assert_close(
+            outputs[0].multimodal_output["actions"],
+            torch.tensor([1.0, 2.0]),
+        )
+        torch.testing.assert_close(
+            outputs[1].multimodal_output["actions"],
+            torch.tensor([3.0, 4.0]),
+        )
+
 
 class TestStepScheduler:
     def setup_method(self) -> None:

tests/dreamzero/openpi_client_helper.py

@@ -0,0 +1,307 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import uuid
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+try:
+    import cv2
+except ImportError:  # pragma: no cover - optional e2e dependency
+    cv2 = None
+
+try:
+    import websockets.sync.client as websockets_client
+except ImportError:  # pragma: no cover - optional e2e dependency
+    websockets_client = None
+
+try:
+    from openpi_client import msgpack_numpy
+except ImportError:  # pragma: no cover - optional e2e dependency
+    msgpack_numpy = None
+
+PING_INTERVAL_SECS = 300
+PING_TIMEOUT_SECS = 3600
+DEFAULT_HOST = "127.0.0.1"
+DEFAULT_PORT = 8000
+DEFAULT_PATH = "/v1/realtime/robot/openpi"
+DEFAULT_PROMPT = "Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan"
+ACTION_HORIZON = 24
+DEFAULT_ACTION_DIM = 8
+RELATIVE_OFFSETS = [-23, -16, -8, 0]
+CAMERA_FILES = {
+    "observation/exterior_image_0_left": "exterior_image_1_left.mp4",
+    "observation/exterior_image_1_left": "exterior_image_2_left.mp4",
+    "observation/wrist_image_left": "wrist_image_left.mp4",
+}
+
+
+def _decode_action_response(response: bytes | str) -> np.ndarray:
+    if isinstance(response, str):
+        raise RuntimeError(f"Inference failed: {response}")
+    decoded = msgpack_numpy.unpackb(response)
+    if isinstance(decoded, dict) and decoded.get("type") == "error":
+        message = decoded.get("message", decoded)
+        raise RuntimeError(f"Inference failed: {message}")
+    return np.asarray(decoded, dtype=np.float32)
+
+
+def require_dependencies() -> None:
+    missing = []
+    if cv2 is None:
+        missing.append("opencv-python")
+    if websockets_client is None:
+        missing.append("websockets")
+    if msgpack_numpy is None:
+        missing.append("openpi-client")
+    if missing:
+        raise ModuleNotFoundError(f"DreamZero OpenPI test dependencies are missing: {', '.join(missing)}")
+
+
+@dataclass(frozen=True)
+class DreamZeroServerMetadata:
+    image_resolution: tuple[int, int]
+    n_external_cameras: int
+    needs_wrist_camera: bool
+    needs_stereo_camera: bool
+    needs_session_id: bool
+    action_space: str
+
+    @classmethod
+    def from_dict(cls, payload: dict[str, Any]) -> DreamZeroServerMetadata:
+        required_keys = (
+            "image_resolution",
+            "n_external_cameras",
+            "needs_wrist_camera",
+            "needs_stereo_camera",
+            "needs_session_id",
+            "action_space",
+        )
+        missing_keys = [key for key in required_keys if key not in payload]
+        if missing_keys:
+            raise ValueError(f"Missing DreamZero metadata keys: {missing_keys}")
+
+        image_resolution = payload["image_resolution"]
+        if not isinstance(image_resolution, (list, tuple)) or len(image_resolution) != 2:
+            raise ValueError(f"Invalid image_resolution: {image_resolution!r}")
+
+        return cls(
+            image_resolution=(int(image_resolution[0]), int(image_resolution[1])),
+            n_external_cameras=int(payload["n_external_cameras"]),
+            needs_wrist_camera=bool(payload["needs_wrist_camera"]),
+            needs_stereo_camera=bool(payload["needs_stereo_camera"]),
+            needs_session_id=bool(payload["needs_session_id"]),
+            action_space=str(payload["action_space"]),
+        )
+
+
+class OpenPIWebsocketClient:
+    def __init__(
+        self,
+        *,
+        host: str = DEFAULT_HOST,
+        port: int = DEFAULT_PORT,
+        path: str = DEFAULT_PATH,
+    ) -> None:
+        require_dependencies()
+        self._uri = f"ws://{host}:{port}{path}"
+        self._packer = msgpack_numpy.Packer()
+        self._ws, self._server_metadata = self._connect()
+
+    def _connect(self):
+        conn = websockets_client.connect(
+            self._uri,
+            compression=None,
+            max_size=None,
+            ping_interval=PING_INTERVAL_SECS,
+            ping_timeout=PING_TIMEOUT_SECS,
+        )
+        metadata = msgpack_numpy.unpackb(conn.recv())
+        if not isinstance(metadata, dict):
+            raise TypeError(f"Expected dict metadata from server, got {type(metadata)!r}")
+        return conn, metadata
+
+    def get_server_metadata(self) -> dict[str, Any]:
+        return dict(self._server_metadata)
+
+    def infer(self, obs: dict[str, Any]) -> np.ndarray:
+        payload = dict(obs)
+        payload["endpoint"] = "infer"
+        self._ws.send(self._packer.pack(payload))
+        response = self._ws.recv()
+        return _decode_action_response(response)
+
+    def reset(self, reset_info: dict[str, Any] | None = None) -> str:
+        payload = dict(reset_info or {})
+        payload["endpoint"] = "reset"
+        self._ws.send(self._packer.pack(payload))
+        response = self._ws.recv()
+        if isinstance(response, str):
+            return response
+        decoded = msgpack_numpy.unpackb(response)
+        if not isinstance(decoded, dict) or decoded.get("status") != "reset successful":
+            raise RuntimeError(f"Unexpected reset response: {decoded!r}")
+        return str(decoded["status"])
+
+    def close(self) -> None:
+        self._ws.close()
+
+
+def load_all_frames(video_path: Path) -> np.ndarray:
+    require_dependencies()
+    cap = cv2.VideoCapture(str(video_path))
+    frames = []
+    while True:
+        ok, frame = cap.read()
+        if not ok:
+            break
+        frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+    cap.release()
+    if not frames:
+        raise RuntimeError(f"No frames loaded from {video_path}")
+    return np.stack(frames, axis=0)
+
+
+def load_camera_frames(video_dir: Path) -> dict[str, np.ndarray]:
+    camera_frames: dict[str, np.ndarray] = {}
+    for camera_key, file_name in CAMERA_FILES.items():
+        video_path = video_dir / file_name
+        if not video_path.exists():
+            raise FileNotFoundError(f"Missing DreamZero test asset: {video_path}")
+        camera_frames[camera_key] = load_all_frames(video_path)
+    return camera_frames
+
+
+def build_frame_schedule(total_frames: int, num_chunks: int) -> list[list[int]]:
+    chunks: list[list[int]] = []
+    current_frame = 23
+    for _ in range(num_chunks):
+        indices = [max(current_frame + offset, 0) for offset in RELATIVE_OFFSETS]
+        if indices[-1] >= total_frames:
+            break
+        chunks.append(indices)
+        current_frame += ACTION_HORIZON
+    return chunks
+
+
+def make_obs_from_video(
+    camera_frames: dict[str, np.ndarray],
+    frame_indices: list[int],
+    *,
+    prompt: str,
+    session_id: str,
+) -> dict[str, Any]:
+    obs: dict[str, Any] = {}
+    for camera_key, all_frames in camera_frames.items():
+        selected = all_frames[frame_indices]
+        obs[camera_key] = selected[0] if len(frame_indices) == 1 else selected
+
+    obs["observation/joint_position"] = np.zeros(7, dtype=np.float32)
+    obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32)
+    obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32)
+    obs["prompt"] = prompt
+    obs["session_id"] = session_id
+    return obs
+
+
+def build_demo_observations(
+    camera_frames: dict[str, np.ndarray],
+    *,
+    prompt: str,
+    session_id: str,
+    num_chunks: int = 2,
+) -> list[dict[str, Any]]:
+    if num_chunks < 1:
+        raise ValueError("num_chunks must be at least 1")
+
+    total_frames = min(frames.shape[0] for frames in camera_frames.values())
+    observations = [
+        make_obs_from_video(
+            camera_frames,
+            [0],
+            prompt=prompt,
+            session_id=session_id,
+        )
+    ]
+    for indices in build_frame_schedule(total_frames, num_chunks - 1):
+        observations.append(
+            make_obs_from_video(
+                camera_frames,
+                indices,
+                prompt=prompt,
+                session_id=session_id,
+            )
+        )
+    return observations
+
+
+def validate_session_result(
+    result: dict[str, Any],
+    *,
+    expected_action_horizon: int = ACTION_HORIZON,
+    expected_action_dim: int = DEFAULT_ACTION_DIM,
+) -> None:
+    metadata = DreamZeroServerMetadata.from_dict(result["metadata"])
+    if metadata.image_resolution != (180, 320):
+        raise AssertionError(f"Unexpected image_resolution: {metadata.image_resolution}")
+    if metadata.n_external_cameras != 2:
+        raise AssertionError(f"Unexpected n_external_cameras: {metadata.n_external_cameras}")
+    if not metadata.needs_wrist_camera:
+        raise AssertionError("DreamZero test expects wrist camera metadata")
+    if metadata.action_space != "joint_position":
+        raise AssertionError(f"Unexpected action_space: {metadata.action_space}")
+
+    actions = result["actions"]
+    if len(actions) != 3:
+        raise AssertionError(f"Expected 3 action tensors, got {len(actions)}")
+    for index, action in enumerate(actions):
+        if action.shape != (expected_action_horizon, expected_action_dim):
+            raise AssertionError(
+                f"Action {index} shape mismatch: expected "
+                f"{(expected_action_horizon, expected_action_dim)}, got {action.shape}"
+            )
+        if not np.isfinite(action).all():
+            raise AssertionError(f"Action {index} contains non-finite values")
+
+    if result["reset_status"] != "reset successful":
+        raise AssertionError(f"Unexpected reset status: {result['reset_status']!r}")
+
+
+def run_policy_session(
+    *,
+    host: str = DEFAULT_HOST,
+    port: int = DEFAULT_PORT,
+    path: str = DEFAULT_PATH,
+    video_dir: Path,
+    prompt: str = DEFAULT_PROMPT,
+    session_id: str | None = None,
+    num_chunks: int = 2,
+) -> dict[str, Any]:
+    session_id = session_id or str(uuid.uuid4())
+    camera_frames = load_camera_frames(video_dir)
+    observations = build_demo_observations(
+        camera_frames,
+        prompt=prompt,
+        session_id=session_id,
+        num_chunks=num_chunks,
+    )
+
+    client = OpenPIWebsocketClient(host=host, port=port, path=path)
+    try:
+        metadata = client.get_server_metadata()
+        actions = [client.infer(obs) for obs in observations]
+        reset_status = client.reset({})
+        actions.append(client.infer(observations[0]))
+        return {
+            "metadata": metadata,
+            "actions": actions,
+            "reset_status": reset_status,
+            "session_id": session_id,
+        }
+    finally:
+        client.close()

tests/dreamzero/test_dreamzero_utils.py

@@ -0,0 +1,24 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+from vllm_omni.diffusion.models.dreamzero.utils import (
+    DEFAULT_CFG_SCALE,
+    DEFAULT_EMBODIMENT_NAME_TO_ID,
+    DEFAULT_NEGATIVE_PROMPT,
+    DEFAULT_NUM_INFERENCE_STEPS,
+    DEFAULT_SEED,
+    DEFAULT_SIGMA_SHIFT,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_dreamzero_default_constants_match_source_baseline():
+    assert DEFAULT_NUM_INFERENCE_STEPS == 16
+    assert DEFAULT_CFG_SCALE == 5.0
+    assert DEFAULT_SIGMA_SHIFT == 5.0
+    assert DEFAULT_SEED == 1140
+    assert "worst quality" in DEFAULT_NEGATIVE_PROMPT
+    assert DEFAULT_EMBODIMENT_NAME_TO_ID["oxe_droid"] == 17

tests/dreamzero/test_openpi_client_helper.py

@@ -0,0 +1,46 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import numpy as np
+import pytest
+
+from tests.dreamzero import openpi_client_helper
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class FakeMsgpackNumpy:
+    payloads = {}
+
+    @staticmethod
+    def packb(obj):
+        key = str(len(FakeMsgpackNumpy.payloads)).encode()
+        FakeMsgpackNumpy.payloads[key] = obj
+        return key
+
+    @staticmethod
+    def unpackb(data):
+        return FakeMsgpackNumpy.payloads[data]
+
+
+def test_decode_action_response_surfaces_structured_error(monkeypatch):
+    monkeypatch.setattr(openpi_client_helper, "msgpack_numpy", FakeMsgpackNumpy)
+    payload = FakeMsgpackNumpy.packb(
+        {
+            "type": "error",
+            "message": "Internal inference error",
+        }
+    )
+
+    with pytest.raises(RuntimeError, match="Internal inference error"):
+        openpi_client_helper._decode_action_response(payload)
+
+
+def test_decode_action_response_converts_action_payload_to_float32(monkeypatch):
+    monkeypatch.setattr(openpi_client_helper, "msgpack_numpy", FakeMsgpackNumpy)
+    payload = FakeMsgpackNumpy.packb(np.asarray([[1.0, 2.0]], dtype=np.float64))
+
+    actions = openpi_client_helper._decode_action_response(payload)
+
+    assert actions.dtype == np.float32
+    np.testing.assert_array_equal(actions, np.asarray([[1.0, 2.0]], dtype=np.float32))

tests/dreamzero/test_pipeline_state.py

@@ -0,0 +1,59 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections import OrderedDict
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
+from vllm_omni.diffusion.models.dreamzero.state_dreamzero import DreamZeroState
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _empty_pipeline() -> DreamZeroPipeline:
+    pipeline = DreamZeroPipeline.__new__(DreamZeroPipeline)
+    pipeline._states = OrderedDict()
+    pipeline._max_session_states = 2
+    return pipeline
+
+
+def test_dreamzero_pipeline_state_is_session_keyed() -> None:
+    pipeline = _empty_pipeline()
+
+    session_a = pipeline._get_or_create_state("session-a")
+    session_b = pipeline._get_or_create_state("session-b")
+    session_a.call_count = 7
+    session_b.call_count = 3
+
+    assert pipeline._get_or_create_state("session-a") is session_a
+    assert pipeline._get_or_create_state("session-b") is session_b
+    assert session_a.call_count == 7
+    assert session_b.call_count == 3
+
+
+def test_dreamzero_pipeline_state_lru_caps_retained_sessions() -> None:
+    pipeline = _empty_pipeline()
+
+    session_a = pipeline._get_or_create_state("session-a")
+    pipeline._get_or_create_state("session-b")
+    assert pipeline._get_or_create_state("session-a") is session_a
+
+    pipeline._get_or_create_state("session-c")
+
+    assert list(pipeline._states) == ["session-a", "session-c"]
+    assert "session-b" not in pipeline._states
+
+
+def test_dreamzero_state_cache_access_requires_initialization() -> None:
+    state = DreamZeroState()
+
+    with pytest.raises(RuntimeError, match="KV caches not initialized"):
+        state.get_kv_caches()
+
+    with pytest.raises(RuntimeError, match="Cross-attn caches not initialized"):
+        state.get_crossattn_caches()
+
+    with pytest.raises(RuntimeError, match="create_kv_caches first"):
+        state.update_kv_cache(0, torch.empty(0))

tests/dreamzero/upstream/openpi_test_client_ar.py

@@ -0,0 +1,381 @@
+#!/usr/bin/env python3
+"""
+Copied from the DreamZero repository's `test_client_AR.py`:
+https://github.com/dreamzero0/dreamzero/blob/main/test_client_AR.py
+
+Kept here for end-to-end compatibility / parity testing against the vLLM
+OpenPI server.
+
+Test client for AR_droid policy server using roboarena interface.
+
+Sends real video frames from debug_image/ directory instead of zero dummy images.
+
+Frame schedule (matching debug_inference.py):
+  - Step 0 (initial): send frame [0]             (1 frame, H W 3)
+  - Step 1: send frames [0, 7, 15, 23]           (4 frames, 4 H W 3)
+  - Step 2: send frames [24, 31, 39, 47]         (4 frames)
+  - Step 3: send frames [48, 55, 63, 71]         (4 frames)
+  - ...
+
+Expected server configuration:
+    - image_resolution: (180, 320)
+    - n_external_cameras: 2
+    - needs_wrist_camera: True
+    - action_space: "joint_position"
+
+Usage:
+    # Start server with roboarena interface:
+    torchrun --nproc_per_node=8 socket_test_optimized_AR.py --port 8000
+
+    # Run against the original DreamZero websocket server (default path):
+    python test_client_AR.py --host <server_host> --port 8000
+
+    # Run against the vLLM OpenPI server:
+    python test_client_AR.py --host <server_host> --port 8000 --path /v1/realtime/robot/openpi
+
+    # Use zero images instead of real video (old behavior):
+    python test_client_AR.py --host <server_host> --port 8000 --use-zero-images
+"""
+
+import argparse
+import logging
+import os
+import sys
+import time
+import uuid
+from pathlib import Path
+
+import cv2
+import numpy as np
+from openpi_client import msgpack_numpy
+
+_DREAMZERO_REPO_ENV = os.environ.get("DREAMZERO_REPO")
+DREAMZERO_REPO = Path(_DREAMZERO_REPO_ENV).expanduser() if _DREAMZERO_REPO_ENV else None
+
+
+def _import_upstream_policy_modules():
+    if DREAMZERO_REPO is None:
+        raise ImportError("Set DREAMZERO_REPO to an upstream DreamZero checkout before using this helper.")
+    if DREAMZERO_REPO.exists() and str(DREAMZERO_REPO) not in sys.path:
+        sys.path.insert(0, str(DREAMZERO_REPO))
+
+    import eval_utils.policy_server as policy_server
+    from eval_utils.policy_client import WebsocketClientPolicy
+
+    return policy_server, WebsocketClientPolicy
+
+
+policy_server, WebsocketClientPolicy = _import_upstream_policy_modules()
+
+VIDEO_DIR = os.environ.get(
+    "DREAMZERO_VIDEO_DIR",
+    str(DREAMZERO_REPO / "debug_image") if DREAMZERO_REPO is not None else "debug_image",
+)
+
+# roboarena key -> video filename
+CAMERA_FILES = {
+    "observation/exterior_image_0_left": "exterior_image_1_left.mp4",
+    "observation/exterior_image_1_left": "exterior_image_2_left.mp4",
+    "observation/wrist_image_left": "wrist_image_left.mp4",
+}
+
+# Frame schedule constants (matching debug_inference.py)
+RELATIVE_OFFSETS = [-23, -16, -8, 0]
+ACTION_HORIZON = 24
+DEFAULT_WEBSOCKET_PATH = ""
+
+
+class OpenPIWebsocketClientPolicy(WebsocketClientPolicy):
+    """DreamZero websocket client with a configurable path suffix.
+
+    The original DreamZero client connects to ``ws://host:port``.
+    vLLM serves the compatible robot policy endpoint at
+    ``/v1/realtime/robot/openpi`` when ``path`` is set accordingly.
+    """
+
+    def __init__(
+        self,
+        host: str = "0.0.0.0",
+        port: int = 8000,
+        path: str = DEFAULT_WEBSOCKET_PATH,
+    ) -> None:
+        self._uri = f"ws://{host}:{port}{path}"
+        self._packer = msgpack_numpy.Packer()
+        self._ws, self._server_metadata = self._wait_for_server()
+
+
+def load_all_frames(video_path: str) -> np.ndarray:
+    """Load all frames from a video file. Returns (N, H, W, 3) uint8 array (RGB)."""
+    cap = cv2.VideoCapture(video_path)
+    frames = []
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+        frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+    cap.release()
+    if not frames:
+        raise RuntimeError(f"No frames loaded from {video_path}")
+    return np.stack(frames, axis=0)
+
+
+def load_camera_frames() -> dict[str, np.ndarray]:
+    """Load all video frames for each camera from the debug_image/ directory.
+
+    Returns:
+        Dict mapping roboarena camera keys to (N, H, W, 3) uint8 arrays.
+    """
+    camera_frames: dict[str, np.ndarray] = {}
+    for cam_key, fname in CAMERA_FILES.items():
+        path = os.path.join(VIDEO_DIR, fname)
+        camera_frames[cam_key] = load_all_frames(path)
+        logging.info(f"Loaded {cam_key}: {camera_frames[cam_key].shape}")
+    return camera_frames
+
+
+def build_frame_schedule(total_frames: int, num_chunks: int) -> list[list[int]]:
+    """Build the frame index schedule for multi-frame chunks.
+
+    Returns a list of frame-index lists. Each inner list has 4 indices.
+    """
+    chunks: list[list[int]] = []
+    current_frame = 23  # first anchor frame
+    for _ in range(num_chunks):
+        indices = [max(current_frame + off, 0) for off in RELATIVE_OFFSETS]
+        if indices[-1] >= total_frames:
+            logging.info(f"Frame {indices[-1]} >= {total_frames}, stopping at {len(chunks)} chunks")
+            break
+        chunks.append(indices)
+        current_frame += ACTION_HORIZON
+    return chunks
+
+
+def _make_obs_from_video(
+    camera_frames: dict[str, np.ndarray],
+    frame_indices: list[int],
+    prompt: str,
+    session_id: str,
+) -> dict:
+    """Build an observation dict from real video frames.
+
+    For 1 frame: each image key is (H, W, 3).
+    For 4 frames: each image key is (4, H, W, 3).
+    """
+    obs: dict = {}
+    for cam_key, all_frames in camera_frames.items():
+        selected = all_frames[frame_indices]  # (T, H, W, 3)
+        if len(frame_indices) == 1:
+            selected = selected[0]  # (H, W, 3)
+        obs[cam_key] = selected
+
+    obs["observation/joint_position"] = np.zeros(7, dtype=np.float32)
+    obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32)
+    obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32)
+    obs["prompt"] = prompt
+    obs["session_id"] = session_id
+    return obs
+
+
+def _make_zero_observation(
+    server_config: policy_server.PolicyServerConfig,
+    prompt: str = "pick up the object",
+    session_id: str | None = None,
+) -> dict:
+    """Create a dummy observation matching AR_droid expectations.
+
+    AR_droid expects:
+        - 2 external cameras (exterior_image_0_left, exterior_image_1_left)
+        - 1 wrist camera (wrist_image_left)
+        - Image resolution: 180x320 (H x W)
+        - joint_position: 7 DoF
+        - gripper_position: 1 DoF
+    """
+    obs = {}
+
+    # Determine image resolution
+    if server_config.image_resolution is not None:
+        h, w = server_config.image_resolution
+    else:
+        # Default for AR_droid
+        h, w = 180, 320
+
+    # External cameras (0-indexed in roboarena)
+    for i in range(server_config.n_external_cameras):
+        obs[f"observation/exterior_image_{i}_left"] = np.zeros((h, w, 3), dtype=np.uint8)
+        if server_config.needs_stereo_camera:
+            obs[f"observation/exterior_image_{i}_right"] = np.zeros((h, w, 3), dtype=np.uint8)
+
+    # Wrist camera
+    if server_config.needs_wrist_camera:
+        obs["observation/wrist_image_left"] = np.zeros((h, w, 3), dtype=np.uint8)
+        if server_config.needs_stereo_camera:
+            obs["observation/wrist_image_right"] = np.zeros((h, w, 3), dtype=np.uint8)
+
+    # Session ID - should be passed in to ensure consistency within a session
+    if server_config.needs_session_id:
+        import uuid
+
+        # Generate unique session ID if not provided
+        obs["session_id"] = session_id if session_id else str(uuid.uuid4())
+
+    # State observations (AR_droid: 7 DoF arm + 1 gripper)
+    obs["observation/joint_position"] = np.zeros(7, dtype=np.float32)
+    obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32)
+    obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32)
+
+    # Language prompt
+    obs["prompt"] = prompt
+
+    return obs
+
+
+def test_ar_droid_policy_server(
+    host: str = "localhost",
+    port: int = 8000,
+    path: str = DEFAULT_WEBSOCKET_PATH,
+    num_chunks: int = 15,
+    prompt: str = "Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan",
+    use_zero_images: bool = False,
+):
+    """Test the AR_droid policy server with roboarena interface.
+
+    When use_zero_images is False (default), loads real video frames from
+    debug_image/ and follows the frame schedule from debug_inference.py.
+    """
+    logging.info(f"Connecting to AR_droid server at ws://{host}:{port}{path} ...")
+
+    client = OpenPIWebsocketClientPolicy(host=host, port=port, path=path)
+
+    # Validate server metadata
+    metadata = client.get_server_metadata()
+    logging.info(f"Server metadata: {metadata}")
+    assert isinstance(metadata, dict), "Metadata should be a dict"
+
+    try:
+        server_config = policy_server.PolicyServerConfig(**metadata)
+    except Exception as e:
+        logging.error(f"Error parsing metadata: {e}")
+        raise e
+
+    # Validate expected AR_droid configuration
+    logging.info(f"Server config: {server_config}")
+    assert server_config.n_external_cameras == 2, f"Expected 2 external cameras, got {server_config.n_external_cameras}"
+    assert server_config.needs_wrist_camera, "Expected wrist camera to be enabled"
+    assert server_config.action_space == "joint_position", (
+        f"Expected joint_position action space, got {server_config.action_space}"
+    )
+
+    logging.info("Server configuration validated for AR_droid")
+
+    # Generate unique session ID for this test run
+    session_id = str(uuid.uuid4())
+    logging.info(f"Session ID: {session_id}")
+
+    # ── Zero-image fallback mode ──────────────────────────────────────
+    if use_zero_images:
+        logging.info("Using ZERO dummy images (legacy mode)")
+        for i in range(num_chunks):
+            obs = _make_zero_observation(server_config, prompt=prompt, session_id=session_id)
+            logging.info(f"Inference {i + 1}/{num_chunks}: prompt='{prompt}'")
+            t0 = time.time()
+            actions = client.infer(obs)
+            dt = time.time() - t0
+            _log_action(actions, dt)
+
+        logging.info("Sending reset...")
+        client.reset({})
+        logging.info("Done (zero-image mode).")
+        return
+
+    # ── Real video frame mode ─────────────────────────────────────────
+    logging.info("Loading real video frames from debug_image/ directory")
+    camera_frames = load_camera_frames()
+
+    total_frames = min(v.shape[0] for v in camera_frames.values())
+    logging.info(f"Total frames available: {total_frames}")
+
+    # Build frame schedule
+    chunks = build_frame_schedule(total_frames, num_chunks)
+
+    logging.info("Frame schedule:")
+    logging.info("  Initial: [0]")
+    for i, indices in enumerate(chunks):
+        logging.info(f"  Chunk {i}: {indices}")
+
+    # Step 0: initial single frame
+    logging.info("=== Initial: frame [0] ===")
+    obs = _make_obs_from_video(camera_frames, [0], prompt, session_id)
+    t0 = time.time()
+    actions = client.infer(obs)
+    dt = time.time() - t0
+    _log_action(actions, dt)
+
+    # Subsequent chunks: send 4 frames at a time
+    for chunk_idx, frame_indices in enumerate(chunks):
+        logging.info(f"=== Chunk {chunk_idx}: frames {frame_indices} ===")
+        obs = _make_obs_from_video(camera_frames, frame_indices, prompt, session_id)
+        t0 = time.time()
+        actions = client.infer(obs)
+        dt = time.time() - t0
+        _log_action(actions, dt)
+
+    # Reset triggers video save on the server
+    logging.info("Sending reset to save video...")
+    client.reset({})
+
+    logging.info("Done.")
+
+
+def _log_action(actions: np.ndarray, dt: float) -> None:
+    """Pretty-print action shape, range, and timing."""
+    assert isinstance(actions, np.ndarray), f"Expected numpy array, got {type(actions)}"
+    assert actions.ndim == 2, f"Expected 2D array, got shape {actions.shape}"
+    assert actions.shape[-1] == 8, f"Expected 8 action dims (7 joints + 1 gripper), got {actions.shape[-1]}"
+    logging.info(f"  Action shape: {actions.shape}, range: [{actions.min():.4f}, {actions.max():.4f}], time: {dt:.2f}s")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Test AR_droid policy server with real video frames from debug_image/")
+    parser.add_argument("--host", default="localhost", help="Server hostname")
+    parser.add_argument("--port", type=int, default=8000, help="Server port")
+    parser.add_argument(
+        "--path",
+        default=DEFAULT_WEBSOCKET_PATH,
+        help="WebSocket path suffix (default: empty string for the original DreamZero server)",
+    )
+    parser.add_argument(
+        "--num-chunks",
+        type=int,
+        default=15,
+        help="Number of 4-frame chunks to send after the initial frame (default: 15)",
+    )
+    parser.add_argument(
+        "--prompt",
+        default="Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan",
+        help="Language prompt for the policy",
+    )
+    parser.add_argument(
+        "--use-zero-images",
+        action="store_true",
+        help="Use zero dummy images instead of real video frames (legacy mode)",
+    )
+
+    args = parser.parse_args()
+
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s [%(levelname)s] %(message)s",
+    )
+
+    test_ar_droid_policy_server(
+        host=args.host,
+        port=args.port,
+        path=args.path,
+        num_chunks=args.num_chunks,
+        prompt=args.prompt,
+        use_zero_images=args.use_zero_images,
+    )
+
+
+if __name__ == "__main__":
+    main()

tests/dreamzero/upstream/test_openpi_e2e_source_parity.py

@@ -0,0 +1,358 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Formal OpenPI end-to-end parity: upstream DreamZero server vs `vllm serve`.
+
+This test uses DreamZero's own client-side observation builders from
+`${DREAMZERO_REPO}/test_client_AR.py`, and client-side websocket protocol from
+`${DREAMZERO_REPO}/eval_utils/policy_client.py`.
+
+The only client-side adaptation for vLLM is the websocket path:
+DreamZero's upstream server serves at `/`, while vLLM serves OpenPI at
+`/v1/realtime/robot/openpi`.
+
+Current scope for this test:
+- default two-GPU run (`nproc_per_node=2` on upstream, `--cfg-parallel-size 2` on `vllm serve`)
+- non-`torch.compile` (upstream launched through
+  `upstream_socket_server_no_compile.py`, vLLM with `--enforce-eager`)
+- non-DiT-cache / non-skip-schedule (`NUM_DIT_STEPS=16`)
+
+Serving contract locked by this test:
+- upstream DreamZero still boots from the local checkpoint directory
+- vLLM boots from the official DreamZero HF repo name (`GEAR-Dreams/DreamZero-DROID`)
+  rather than a prepared local bundle path
+"""
+
+from __future__ import annotations
+
+import os
+import shutil
+import subprocess
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+
+from tests.helpers.runtime import get_open_port
+
+msgpack_numpy = pytest.importorskip("openpi_client.msgpack_numpy")
+
+_DREAMZERO_REPO_ENV = os.environ.get("DREAMZERO_REPO")
+DREAMZERO_REPO = Path(_DREAMZERO_REPO_ENV).expanduser() if _DREAMZERO_REPO_ENV else None
+if DREAMZERO_REPO is not None and str(DREAMZERO_REPO) not in sys.path:
+    sys.path.insert(0, str(DREAMZERO_REPO))
+
+try:
+    import test_client_AR as dreamzero_client
+    from eval_utils.policy_client import WebsocketClientPolicy
+except Exception:  # pragma: no cover - guarded by pytest skip below
+    dreamzero_client = None
+    WebsocketClientPolicy = None
+_BaseWebsocketClientPolicy = WebsocketClientPolicy if WebsocketClientPolicy is not None else object
+
+CHECKPOINT_DIR = DREAMZERO_REPO / "checkpoints" / "dreamzero" if DREAMZERO_REPO is not None else None
+VLLM_MODEL = os.environ.get("VLLM_DREAMZERO_MODEL", "GEAR-Dreams/DreamZero-DROID")
+SERVICE_READY_TIMEOUT_S = int(os.environ.get("OPENPI_SERVICE_READY_TIMEOUT_S", "900"))
+PROMPT = "Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan"
+SESSION_ID = "openpi-e2e-parity-session"
+
+pytestmark = [
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="GPU required"),
+    pytest.mark.skipif(
+        dreamzero_client is None or WebsocketClientPolicy is None,
+        reason="DreamZero client modules are required on PYTHONPATH",
+    ),
+    pytest.mark.skipif(
+        DREAMZERO_REPO is None or not DREAMZERO_REPO.exists(),
+        reason="DreamZero source repo is required at DREAMZERO_REPO",
+    ),
+    pytest.mark.skipif(
+        CHECKPOINT_DIR is None or not CHECKPOINT_DIR.exists(), reason="DreamZero local checkpoint is required"
+    ),
+]
+
+
+class OpenPIWebsocketClientPolicy(_BaseWebsocketClientPolicy):
+    """DreamZero client protocol with an OpenPI websocket path suffix."""
+
+    def __init__(
+        self,
+        host: str = "127.0.0.1",
+        port: int = 8000,
+        path: str = "/v1/realtime/robot/openpi",
+    ) -> None:
+        self._uri = f"ws://{host}:{port}{path}"
+        self._packer = msgpack_numpy.Packer()
+        self._ws, self._server_metadata = self._wait_for_server()
+
+
+def _vllm_executable() -> str:
+    fallback = Path(sys.executable).with_name("vllm")
+    if fallback.exists():
+        return str(fallback)
+    exe = shutil.which("vllm")
+    if exe:
+        return exe
+    raise FileNotFoundError("Unable to locate `vllm` executable in current environment.")
+
+
+def _cfg_parallel_size() -> int:
+    return int(os.environ.get("OPENPI_E2E_CFG_PARALLEL_SIZE", "2"))
+
+
+def _pick_test_gpus() -> list[str]:
+    cfg_parallel_size = _cfg_parallel_size()
+    override = os.environ.get("OPENPI_E2E_GPUS") or os.environ.get("OPENPI_E2E_GPU")
+    if override is not None:
+        gpus = [part.strip() for part in override.split(",") if part.strip()]
+        if not gpus:
+            raise ValueError("OPENPI_E2E_GPUS is set but empty.")
+        if len(gpus) < cfg_parallel_size:
+            raise RuntimeError(f"Need {cfg_parallel_size} GPUs, but OPENPI_E2E_GPUS only provided {gpus}.")
+        return gpus
+
+    query = subprocess.check_output(
+        [
+            "nvidia-smi",
+            "--query-gpu=index,memory.used",
+            "--format=csv,noheader,nounits",
+        ],
+        text=True,
+    )
+    gpu_rows = []
+    for line in query.strip().splitlines():
+        gpu_index, used_mb = [part.strip() for part in line.split(",", maxsplit=1)]
+        gpu_rows.append((int(used_mb), gpu_index))
+    gpu_rows.sort()
+    gpus = [gpu_index for _, gpu_index in gpu_rows[: max(cfg_parallel_size, 1)]]
+    if len(gpus) < cfg_parallel_size:
+        raise RuntimeError(
+            f"Need {cfg_parallel_size} GPUs for cfg_parallel_size={cfg_parallel_size}, "
+            f"but found only {len(gpus)} candidates."
+        )
+    return gpus
+
+
+def _torchrun_argv(script: str, port: int) -> list[str]:
+    return [
+        sys.executable,
+        "-m",
+        "torch.distributed.run",
+        "--standalone",
+        "--nproc_per_node",
+        str(_cfg_parallel_size()),
+        script,
+        "--port",
+        str(port),
+        "--model_path",
+        str(CHECKPOINT_DIR),
+    ]
+
+
+def _run_upstream_service(port: int, log_path: Path) -> subprocess.Popen[str]:
+    env = os.environ.copy()
+    env.setdefault("PYTHONPATH", "")
+    env["PYTHONPATH"] = f"{Path.cwd()}:{DREAMZERO_REPO}:{env['PYTHONPATH']}".rstrip(":")
+    env["CUDA_VISIBLE_DEVICES"] = ",".join(_pick_test_gpus())
+    env.setdefault("NO_ALBUMENTATIONS_UPDATE", "1")
+    env["ATTENTION_BACKEND"] = "torch"
+    env.setdefault("ENABLE_TENSORRT", "false")
+    env["ENABLE_DIT_CACHE"] = "false"
+    env["NUM_DIT_STEPS"] = "16"
+    env["DYNAMIC_CACHE_SCHEDULE"] = "false"
+    argv = _torchrun_argv(
+        str(Path("tests/dreamzero/upstream/upstream_socket_server_no_compile.py")),
+        port,
+    )
+    log_file = log_path.open("w")
+    proc = subprocess.Popen(
+        argv,
+        stdout=log_file,
+        stderr=subprocess.STDOUT,
+        text=True,
+        cwd=str(Path.cwd()),
+        env=env,
+    )
+    proc._codex_log_file = log_file  # type: ignore[attr-defined]
+    return proc
+
+
+def _run_vllm_service(port: int, log_path: Path) -> subprocess.Popen[str]:
+    env = os.environ.copy()
+    gpus = _pick_test_gpus()
+    cfg_parallel_size = _cfg_parallel_size()
+    if cfg_parallel_size > len(gpus):
+        raise RuntimeError(
+            f"cfg_parallel_size={cfg_parallel_size} requires at least {cfg_parallel_size} GPUs, but only got {gpus}."
+        )
+    env["CUDA_VISIBLE_DEVICES"] = ",".join(gpus[:cfg_parallel_size])
+    env.setdefault("ATTENTION_BACKEND", "torch")
+    env.setdefault("DIFFUSION_ATTENTION_BACKEND", "TORCH_SDPA")
+    env.setdefault("MASTER_PORT", str(get_open_port()))
+    argv = [
+        _vllm_executable(),
+        "serve",
+        VLLM_MODEL,
+        "--omni",
+        "--host",
+        "127.0.0.1",
+        "--port",
+        str(port),
+        "--served-model-name",
+        "dreamzero-droid",
+        "--enforce-eager",
+    ]
+    if cfg_parallel_size > 1:
+        argv.extend(["--cfg-parallel-size", str(cfg_parallel_size)])
+    log_file = log_path.open("w")
+    proc = subprocess.Popen(
+        argv,
+        stdout=log_file,
+        stderr=subprocess.STDOUT,
+        text=True,
+        env=env,
+        cwd=str(Path.cwd()),
+    )
+    proc._codex_log_file = log_file  # type: ignore[attr-defined]
+    return proc
+
+
+def _stop_process(proc: subprocess.Popen[str]) -> None:
+    log_file = getattr(proc, "_codex_log_file", None)
+    if proc.poll() is None:
+        proc.terminate()
+        try:
+            proc.wait(timeout=30)
+        except subprocess.TimeoutExpired:  # pragma: no cover - cleanup path
+            proc.kill()
+            proc.wait(timeout=10)
+    if log_file is not None:
+        log_file.close()
+
+
+def _build_obs_sequence() -> tuple[dict, dict]:
+    camera_frames = dreamzero_client.load_camera_frames()
+    chunks = dreamzero_client.build_frame_schedule(
+        min(v.shape[0] for v in camera_frames.values()),
+        1,
+    )
+    obs0 = dreamzero_client._make_obs_from_video(camera_frames, [0], PROMPT, SESSION_ID)
+    obs1 = dreamzero_client._make_obs_from_video(camera_frames, chunks[0], PROMPT, SESSION_ID)
+    return obs0, obs1
+
+
+def _wait_for_client_ready(client_factory, timeout_s: float, proc=None, log_path: Path | None = None):
+    deadline = time.time() + timeout_s
+    last_err: Exception | None = None
+    while time.time() < deadline:
+        if proc is not None and proc.poll() is not None:
+            details = ""
+            if log_path is not None and log_path.exists():
+                details = log_path.read_text(errors="replace")[-8000:]
+            raise RuntimeError(f"Service exited before becoming ready with code {proc.returncode}.\n{details}")
+        try:
+            return client_factory()
+        except Exception as exc:  # pragma: no cover - retry path
+            last_err = exc
+            time.sleep(1)
+    raise TimeoutError(f"Timed out waiting for websocket service: {last_err}")
+
+
+def _collect_outputs_with_client(client) -> tuple[dict, list[np.ndarray]]:
+    metadata = client.get_server_metadata()
+    obs0, obs1 = _build_obs_sequence()
+    outputs = [
+        client.infer(dict(obs0)),
+        client.infer(dict(obs1)),
+    ]
+    assert _normalize_reset_response(client.reset({})) == "reset successful"
+    outputs.append(client.infer(dict(obs0)))
+    client._ws.close()
+    return metadata, outputs
+
+
+def _normalize_reset_response(response) -> str:
+    if isinstance(response, str):
+        return response
+    decoded = msgpack_numpy.unpackb(response)
+    if isinstance(decoded, dict):
+        return str(decoded.get("status"))
+    return str(decoded)
+
+
+def _normalize_metadata(metadata: dict) -> dict:
+    normalized = dict(metadata)
+    if isinstance(normalized.get("image_resolution"), tuple):
+        normalized["image_resolution"] = list(normalized["image_resolution"])
+    return normalized
+
+
+def _assert_logs_clean(log_path: Path) -> None:
+    text = log_path.read_text(errors="replace")
+    if "SignalException: Process" in text and "got signal: 15" in text:
+        text = text.split("Traceback (most recent call last):", 1)[0]
+    assert "Traceback" not in text, text
+    assert "RuntimeError:" not in text, text
+
+
+def _assert_upstream_log_matches_vllm_baseline(log_path: Path) -> None:
+    text = log_path.read_text(errors="replace")
+    assert "DIT Compute Steps 8 steps" not in text, text
+    assert "DIT Compute Steps 16 steps" in text, text
+
+
+def test_openpi_service_matches_upstream_server_noncompile(tmp_path: Path) -> None:
+    expected_metadata = {
+        "image_resolution": [180, 320],
+        "n_external_cameras": 2,
+        "needs_wrist_camera": True,
+        "needs_stereo_camera": False,
+        "needs_session_id": True,
+        "action_space": "joint_position",
+    }
+
+    upstream_port = get_open_port()
+    upstream_log = tmp_path / "dreamzero_upstream.log"
+    upstream_proc = _run_upstream_service(upstream_port, upstream_log)
+    try:
+        upstream_client = _wait_for_client_ready(
+            lambda: WebsocketClientPolicy(host="127.0.0.1", port=upstream_port),
+            timeout_s=SERVICE_READY_TIMEOUT_S,
+            proc=upstream_proc,
+            log_path=upstream_log,
+        )
+        upstream_metadata, upstream_outputs = _collect_outputs_with_client(upstream_client)
+    finally:
+        _stop_process(upstream_proc)
+    _assert_logs_clean(upstream_log)
+    _assert_upstream_log_matches_vllm_baseline(upstream_log)
+
+    vllm_port = get_open_port()
+    vllm_log = tmp_path / "vllm_openpi.log"
+    vllm_proc = _run_vllm_service(vllm_port, vllm_log)
+    try:
+        vllm_client = _wait_for_client_ready(
+            lambda: OpenPIWebsocketClientPolicy(host="127.0.0.1", port=vllm_port),
+            timeout_s=SERVICE_READY_TIMEOUT_S,
+            proc=vllm_proc,
+            log_path=vllm_log,
+        )
+        vllm_metadata, vllm_outputs = _collect_outputs_with_client(vllm_client)
+    finally:
+        _stop_process(vllm_proc)
+    _assert_logs_clean(vllm_log)
+
+    assert _normalize_metadata(upstream_metadata) == expected_metadata
+    assert _normalize_metadata(vllm_metadata) == expected_metadata
+
+    for idx, (actual, expected) in enumerate(zip(vllm_outputs, upstream_outputs, strict=True)):
+        np.testing.assert_allclose(
+            actual,
+            expected,
+            rtol=1e-2,
+            atol=1e-3,
+            err_msg=f"OpenPI step {idx} output mismatch",
+        )

tests/dreamzero/upstream/upstream_socket_server_no_compile.py

@@ -0,0 +1,184 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Launch the upstream DreamZero websocket server with `torch.compile` disabled.
+
+This wrapper is meant for formal parity tests against `vllm serve --omni`.
+It monkeypatches `torch.compile` before importing DreamZero modules so that
+all import-time decorators and `post_initialize()` compile calls become eager.
+
+For the current DreamZero port baseline we also disable all upstream DiT cache
+and step-skipping behavior so the reference server matches vLLM's current
+eager/no-skip implementation:
+- `ENABLE_DIT_CACHE=false`
+- `NUM_DIT_STEPS=16`
+- `DYNAMIC_CACHE_SCHEDULE=false`
+
+In our CI/dev environment we also do not install Transformer Engine or
+FlashAttention, so upstream's default `socket_test_optimized_AR.main()` path
+(`ATTENTION_BACKEND=TE` and then fallback to `FA2`) cannot execute. To keep the
+formal server-vs-server test runnable without pulling in those heavyweight
+optional deps, this wrapper reproduces upstream `main()` but pins attention to
+PyTorch SDPA (`ATTENTION_BACKEND=torch`) for this subprocess only.
+
+Usage:
+    PYTHONPATH="${DREAMZERO_REPO}" \\
+      .venv/bin/python -m torch.distributed.run --standalone --nproc_per_node=2 \\
+      tests/dreamzero/upstream/upstream_socket_server_no_compile.py --port 18081 \\
+      --model_path "${DREAMZERO_REPO}/checkpoints/dreamzero"
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+
+
+def _identity_compile(*args, **kwargs):
+    if args and callable(args[0]) and len(args) == 1 and not kwargs:
+        return args[0]
+
+    def deco(fn):
+        return fn
+
+    return deco
+
+
+torch.compile = _identity_compile
+
+os.environ.setdefault("NO_ALBUMENTATIONS_UPDATE", "1")
+
+DREAMZERO_REPO_ENV = os.environ.get("DREAMZERO_REPO")
+if not DREAMZERO_REPO_ENV:
+    raise RuntimeError("Set DREAMZERO_REPO to an upstream DreamZero checkout before launching this helper.")
+DREAMZERO_REPO = Path(DREAMZERO_REPO_ENV).expanduser()
+if str(DREAMZERO_REPO) not in sys.path:
+    sys.path.insert(0, str(DREAMZERO_REPO))
+
+import socket_test_optimized_AR as upstream  # noqa: E402
+import tyro  # noqa: E402
+from groot.vla.model.dreamzero.modules import attention as upstream_attention  # noqa: E402
+from groot.vla.model.dreamzero.modules import wan2_1_submodule as upstream_submodule  # noqa: E402
+from groot.vla.model.dreamzero.modules import wan_video_dit as upstream_wan_video_dit  # noqa: E402
+
+
+def _torch_varlen_flash_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    q_lens: torch.Tensor | None = None,
+    k_lens: torch.Tensor | None = None,
+    dropout_p: float = 0.0,
+    causal: bool = False,
+    dtype: torch.dtype = torch.bfloat16,
+    **_: object,
+) -> torch.Tensor:
+    if q_lens is not None or k_lens is not None:
+        upstream_attention.warnings.warn(
+            "Padding mask is disabled in the test-only SDPA fallback.",
+        )
+    out_dtype = q.dtype
+    q = q.transpose(1, 2).to(dtype)
+    k = k.transpose(1, 2).to(dtype)
+    v = v.transpose(1, 2).to(dtype)
+    out = F.scaled_dot_product_attention(
+        q,
+        k,
+        v,
+        attn_mask=None,
+        is_causal=causal,
+        dropout_p=dropout_p,
+    )
+    return out.transpose(1, 2).contiguous().to(out_dtype)
+
+
+upstream_attention.flash_attention = _torch_varlen_flash_attention
+upstream_submodule.flash_attention = _torch_varlen_flash_attention
+upstream_wan_video_dit.FLASH_ATTN_3_AVAILABLE = False
+upstream_wan_video_dit.FLASH_ATTN_2_AVAILABLE = False
+upstream_wan_video_dit.SAGE_ATTN_AVAILABLE = False
+
+
+def main(args: upstream.Args) -> None:
+    os.environ["ENABLE_DIT_CACHE"] = "false"
+    os.environ["ATTENTION_BACKEND"] = "torch"
+    os.environ["NUM_DIT_STEPS"] = "16"
+    os.environ["DYNAMIC_CACHE_SCHEDULE"] = "false"
+    torch._dynamo.config.recompile_limit = 800
+
+    embodiment_tag = "oxe_droid"
+    model_path = args.model_path
+
+    device_mesh = upstream.init_mesh()
+    rank = upstream.dist.get_rank()
+
+    timeout_delta = upstream.datetime.timedelta(seconds=args.timeout_seconds)
+    signal_group = upstream.dist.new_group(backend="gloo", timeout=timeout_delta)
+    upstream.logger.info("Rank %s initialized signal_group (gloo)", rank)
+
+    policy = upstream.GrootSimPolicy(
+        embodiment_tag=upstream.EmbodimentTag(embodiment_tag),
+        model_path=model_path,
+        device="cuda" if torch.cuda.is_available() else "cpu",
+        device_mesh=device_mesh,
+    )
+
+    hostname = upstream.socket.gethostname()
+    local_ip = upstream.socket.gethostbyname(hostname)
+
+    if rank == 0:
+        upstream.logging.info("Creating server (host: %s, ip: %s)", hostname, local_ip)
+        parent_dir = os.path.dirname(model_path)
+        date_suffix = upstream.datetime.datetime.now().strftime("%Y%m%d")
+        checkpoint_name = os.path.basename(model_path)
+        output_dir = os.path.join(
+            parent_dir,
+            f"real_world_eval_gen_{date_suffix}_{args.index}",
+            checkpoint_name,
+        )
+        os.makedirs(output_dir, exist_ok=True)
+        upstream.logging.info("Videos will be saved to: %s", output_dir)
+    else:
+        output_dir = None
+        upstream.logging.info("Rank %s starting as worker for distributed inference...", rank)
+
+    wrapper_policy = upstream.ARDroidRoboarenaPolicy(
+        groot_policy=policy,
+        signal_group=signal_group,
+        output_dir=output_dir,
+    )
+
+    server_config = upstream.PolicyServerConfig(
+        image_resolution=(180, 320),
+        needs_wrist_camera=True,
+        n_external_cameras=2,
+        needs_stereo_camera=False,
+        needs_session_id=True,
+        action_space="joint_position",
+    )
+
+    if rank == 0:
+        upstream.logging.info("Using roboarena policy server interface")
+        upstream.logging.info("Server config: %s", server_config)
+        roboarena_server = upstream.RoboarenaServer(
+            policy=wrapper_policy,
+            server_config=server_config,
+            host="0.0.0.0",
+            port=args.port,
+        )
+        roboarena_server.serve_forever()
+    else:
+        worker = upstream.DistributedWorker(
+            policy=policy,
+            signal_group=signal_group,
+        )
+        upstream.asyncio.run(worker.run())
+
+
+if __name__ == "__main__":
+    upstream.logging.basicConfig(level=upstream.logging.INFO, force=True)
+    main(tyro.cli(upstream.Args))

tests/e2e/online_serving/test_dreamzero_expansion.py

@@ -0,0 +1,90 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""E2E online serving test for DreamZero OpenPI websocket serving."""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+from tests.dreamzero import openpi_client_helper as openpi_client
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServerParams, get_open_port
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
+
+MODEL = "GEAR-Dreams/DreamZero-DROID"
+
+
+test_params = [
+    OmniServerParams(
+        model=MODEL,
+        port=8091,
+        server_args=[
+            "--deploy-config",
+            "vllm_omni/deploy/dreamzero_tp1_cfg2.yaml",
+            "--enforce-eager",
+            "--disable-log-stats",
+        ],
+        env_dict={
+            "ATTENTION_BACKEND": "torch",
+            "DIFFUSION_ATTENTION_BACKEND": "TORCH_SDPA",
+            "VLLM_DISABLE_COMPILE_CACHE": "1",
+            "MASTER_PORT": str(get_open_port()),
+        },
+    )
+]
+
+
+def _write_synthetic_video(path: Path, cv2_module, *, channel: int) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    height, width, num_frames = 180, 320, 24
+    writer = cv2_module.VideoWriter(str(path), cv2_module.VideoWriter_fourcc(*"mp4v"), 15.0, (width, height))
+    if not writer.isOpened():
+        raise RuntimeError(f"Failed to open video writer for {path}")
+    try:
+        for frame_idx in range(num_frames):
+            frame = np.zeros((height, width, 3), dtype=np.uint8)
+            frame[..., channel] = (frame_idx * 7) % 255
+            frame[..., (channel + 1) % 3] = 64
+            writer.write(cv2_module.cvtColor(frame, cv2_module.COLOR_RGB2BGR))
+    finally:
+        writer.release()
+
+
+def _write_synthetic_dreamzero_videos(client_mod, video_dir: Path) -> None:
+    for channel, file_name in enumerate(client_mod.CAMERA_FILES.values()):
+        _write_synthetic_video(video_dir / file_name, client_mod.cv2, channel=channel)
+
+
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@pytest.mark.distributed_cuda
+@hardware_test(res={"cuda": "H100"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_dreamzero_openpi_online(omni_server, tmp_path: Path) -> None:
+    try:
+        openpi_client.require_dependencies()
+    except ModuleNotFoundError as exc:
+        pytest.skip(str(exc))
+
+    video_dir = tmp_path / "dreamzero_videos"
+    _write_synthetic_dreamzero_videos(openpi_client, video_dir)
+    result = openpi_client.run_policy_session(
+        host=omni_server.host,
+        port=omni_server.port,
+        video_dir=video_dir,
+        session_id="dreamzero-online-e2e",
+    )
+
+    openpi_client.validate_session_result(result)
+
+    metadata = result["metadata"]
+    assert metadata["needs_session_id"] is True
+    assert metadata["needs_stereo_camera"] is False
+    assert tuple(metadata["image_resolution"]) == (180, 320)

tests/entrypoints/test_omni_entrypoints.py

@@ -413,6 +413,19 @@ def test_openai_serving_models_can_consume_async_omni_compat_attrs():
 
 
 def test_get_diffusion_od_config_returns_diffusion_stage_config():
+    diffusion_od_config = object()
+    omni = object.__new__(AsyncOmni)
+    omni.engine = SimpleNamespace(
+        stage_clients=[
+            SimpleNamespace(stage_type="llm"),
+            SimpleNamespace(stage_type="diffusion", od_config=diffusion_od_config),
+        ]
+    )
+
+    assert omni.get_diffusion_od_config() is diffusion_od_config
+
+
+def test_get_diffusion_od_config_falls_back_to_inner_engine():
     diffusion_od_config = object()
     omni = object.__new__(AsyncOmni)
     omni.engine = SimpleNamespace(

tests/entrypoints/test_resolve_dreamzero_config.py

@@ -0,0 +1,56 @@
+import pytest
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc
+from vllm_omni.entrypoints.utils import load_stage_configs_from_model, resolve_model_config_path
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_dreamzero_vla_resolves_to_dreamzero_config(monkeypatch):
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.utils.get_config",
+        lambda _model, trust_remote_code=True: type("Cfg", (), {"model_type": "vla"})(),
+    )
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.utils._looks_like_dreamzero",
+        lambda _model: True,
+    )
+    result = resolve_model_config_path("GEAR-Dreams/DreamZero-DROID")
+
+    assert result is not None
+    assert result.endswith("vllm_omni/deploy/dreamzero.yaml")
+
+
+def test_dreamzero_config_sets_model_class_and_policy_config(monkeypatch):
+    monkeypatch.setattr(
+        "vllm_omni.config.stage_config.StageConfigFactory._auto_detect_model_type",
+        classmethod(lambda _cls, _model, trust_remote_code=True: ("vla", None)),
+    )
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.utils.hf_utils._looks_like_dreamzero",
+        lambda _model: True,
+    )
+
+    stage_configs = load_stage_configs_from_model("GEAR-Dreams/DreamZero-DROID")
+    engine_args = stage_configs[0].engine_args
+
+    assert engine_args.model_class_name == "DreamZeroPipeline"
+    assert engine_args.model_config.policy_server_config.action_space == "joint_position"
+
+
+def test_dreamzero_enrich_config_preserves_explicit_model_class_name(monkeypatch):
+    monkeypatch.setattr(
+        "vllm.transformers_utils.config.get_hf_file_to_dict",
+        lambda path, _model: None if path == "model_index.json" else {"model_type": "vla", "architectures": ["VLA"]},
+    )
+
+    od_config = OmniDiffusionConfig(
+        model="GEAR-Dreams/DreamZero-DROID",
+        model_class_name="DreamZeroPipeline",
+    )
+    proc = StageDiffusionProc(od_config.model, od_config)
+
+    proc._enrich_config()
+
+    assert od_config.model_class_name == "DreamZeroPipeline"