[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>

vllm_omni/diffusion/models/dreamzero/transform/__init__.py

@@ -0,0 +1,31 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import importlib
+
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.models.dreamzero.transform.base import TRANSFORMS
+
+logger = init_logger(__name__)
+
+DEFAULT_EMBODIMENT = "roboarena"
+_BUILTIN_TRANSFORM_MODULES = (
+    "vllm_omni.diffusion.models.dreamzero.transform.droid",
+    "vllm_omni.diffusion.models.dreamzero.transform.roboarena",
+)
+
+
+def ensure_transforms_loaded() -> None:
+    """Import DreamZero transform modules and verify registration."""
+    for module_name in _BUILTIN_TRANSFORM_MODULES:
+        try:
+            importlib.import_module(module_name)
+        except Exception as exc:
+            logger.exception("Failed to import DreamZero transform module %s", module_name)
+            raise RuntimeError(f"Failed to import DreamZero transform module '{module_name}'.") from exc
+
+    if DEFAULT_EMBODIMENT not in TRANSFORMS:
+        raise RuntimeError(f"Built-in DreamZero transform '{DEFAULT_EMBODIMENT}' is not registered after import.")

vllm_omni/diffusion/models/dreamzero/transform/base.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Base transform interface for DreamZero robot policy serving.
+
+Transforms handle dataset-specific concerns ONLY:
+  - Observation key mapping
+  - Multi-view stitching (embodiment-specific layout)
+  - Language template wrapping (embodiment-specific)
+  - Raw state extraction (dataset-specific keys)
+  - Output action slicing (to actual action_dim)
+
+Model-specific concerns belong in the pipeline:
+  - Tokenization (pipeline owns tokenizer)
+  - State padding (pipeline knows MAX_STATE_DIM)
+  - Negative prompt (pipeline owns the string)
+  - Noise generation, encoding, decoding
+
+Flow:
+  raw obs (dataset format)
+    → DreamZeroPipeline selects transform by embodiment
+      → unified dict (stitched video, templated prompt str, raw state)
+        → tokenize, pad, encode, denoise
+          → transform_action_output()
+            → ndarray (N, action_dim)
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+
+
+class RobotPolicyTransform:
+    """Base class for dataset-specific observation transforms.
+
+    Subclasses MUST define:
+        IMAGE_KEY_MAP: dict — dataset obs keys → unified keys
+        EMBODIMENT_NAME: str — embodiment identity (pipeline maps to numeric ID)
+        ACTION_DIM: int — actual action dimensions (for output slicing)
+
+    Subclasses MUST override:
+        _stitch_views()      — multi-view → single stitched image
+        _language_template() — prompt → embodiment-aware template
+        _extract_raw_state() — obs → raw state ndarray
+    """
+
+    IMAGE_KEY_MAP: dict[str, str]
+    EMBODIMENT_NAME: str
+    ACTION_DIM: int
+
+    def transform_input(self, obs: dict) -> dict:
+        """Dataset-specific transform: key map → stitch → template → state."""
+        # 1. Map image keys → unified keys
+        images: dict[str, np.ndarray] = {}
+        for src_key, dst_key in self.IMAGE_KEY_MAP.items():
+            if src_key in obs:
+                images[dst_key] = np.asarray(obs[src_key])
+
+        # 2. Multi-view stitching
+        stitched = self._stitch_views(images)
+
+        # 3. Language template (string only, pipeline tokenizes)
+        prompt = obs.get("prompt", "")
+        templated_prompt = self._language_template(prompt)
+
+        # 4. Raw state extraction (pipeline pads)
+        raw_state = self._extract_raw_state(obs)
+
+        # 5. Build unified output
+        unified: dict[str, Any] = {
+            "images": stitched,  # ndarray (T, H_out, W_out, 3)
+            "prompt": templated_prompt,  # str (templated, not tokenized)
+            "state": raw_state,  # ndarray (state_dim,) — pipeline pads
+            "embodiment_name": self.EMBODIMENT_NAME,
+        }
+        if "session_id" in obs:
+            unified["session_id"] = obs["session_id"]
+        return unified
+
+    def transform_action_output(self, actions: Any) -> np.ndarray:
+        """Adapt model action output to this transform's action dimensions."""
+        actions = np.asarray(actions, dtype=np.float32)
+        # Handle any remaining batch dims: squeeze to 2D (horizon, dim)
+        while actions.ndim > 2:
+            actions = actions[0]
+        # Slice padded dim to actual ACTION_DIM
+        if actions.ndim == 2 and actions.shape[-1] > self.ACTION_DIM:
+            actions = actions[:, : self.ACTION_DIM]
+        return actions
+
+    # ------------------------------------------------------------------
+    # Subclass MUST override
+    # ------------------------------------------------------------------
+
+    def _stitch_views(self, images: dict[str, np.ndarray]) -> np.ndarray:
+        """Stitch camera views into single image.
+        Input: unified key → ndarray (H,W,3) or (T,H,W,3).
+        Output: ndarray (T, H_out, W_out, 3).
+        """
+        raise NotImplementedError
+
+    def _language_template(self, prompt: str) -> str:
+        """Wrap prompt in embodiment-specific template string."""
+        raise NotImplementedError
+
+    def _extract_raw_state(self, obs: dict) -> np.ndarray:
+        """Extract raw state vector from obs.
+        Returns: ndarray (state_dim,) float64. Pipeline handles padding.
+        """
+        raise NotImplementedError
+
+
+# Transform registry — keyed by embodiment/dataset name
+TRANSFORMS: dict[str, RobotPolicyTransform] = {}
+
+
+def register_transform(name: str, transform: RobotPolicyTransform) -> None:
+    TRANSFORMS[name] = transform
+
+
+def get_transform(name: str) -> RobotPolicyTransform:
+    if name not in TRANSFORMS:
+        raise KeyError(f"Unknown transform '{name}'. Available: {list(TRANSFORMS.keys())}")
+    return TRANSFORMS[name]

vllm_omni/diffusion/models/dreamzero/transform/droid.py

@@ -0,0 +1,138 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""DROID dataset transform.
+
+DROID uses 1-indexed exterior cameras, 3 views total (OXE_DROID embodiment).
+Stitching layout (same as RoboArena — both are OXE_DROID):
+    ┌─────────────────────────┐
+    │   wrist (2x width)      │  ← pixel-repeat along width
+    ├────────────┬────────────┤
+    │  left ext  │ right ext  │
+    └────────────┴────────────┘
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+import torchvision.transforms.v2 as T
+
+from vllm_omni.diffusion.models.dreamzero.transform.base import (
+    RobotPolicyTransform,
+    register_transform,
+)
+
+
+class DroidTransform(RobotPolicyTransform):
+    """Transform for DROID dataset (OXE_DROID embodiment).
+
+    DROID observation keys (1-indexed exterior cameras):
+        observation/exterior_image_1_left  → left exterior
+        observation/exterior_image_2_left  → right exterior
+        observation/wrist_image_left       → wrist
+    """
+
+    IMAGE_KEY_MAP = {
+        "observation/exterior_image_1_left": "images/exterior_0",
+        "observation/exterior_image_2_left": "images/exterior_1",
+        "observation/wrist_image_left": "images/wrist",
+    }
+    EMBODIMENT_NAME = "oxe_droid"
+    ACTION_DIM = 8  # 7 joint + 1 gripper
+    _VIDEO_CROP_SCALE = 0.95
+    _VIDEO_RESIZE_HW = (176, 320)
+
+    @classmethod
+    def _preprocess_view(cls, arr: np.ndarray) -> np.ndarray:
+        """Apply per-view crop and resize before stitching."""
+        frames = torch.from_numpy(arr).to(torch.float32).permute(0, 3, 1, 2) / 255.0
+        crop_h = int(arr.shape[1] * cls._VIDEO_CROP_SCALE)
+        crop_w = int(arr.shape[2] * cls._VIDEO_CROP_SCALE)
+        frames = T.CenterCrop((crop_h, crop_w))(frames)
+        frames = T.Resize(
+            cls._VIDEO_RESIZE_HW,
+            interpolation=T.InterpolationMode.BILINEAR,
+            antialias=True,
+        )(frames)
+        return (frames.permute(0, 2, 3, 1) * 255.0).to(torch.uint8).cpu().numpy()
+
+    def _stitch_views(self, images: dict[str, np.ndarray]) -> np.ndarray:
+        """OXE_DROID 2x2 stitching: wrist top (2x wide), exteriors bottom."""
+        left_ext = images.get("images/exterior_0")
+        right_ext = images.get("images/exterior_1")
+        wrist = images.get("images/wrist")
+
+        # Ensure 4D: (T, H, W, C)
+        def ensure_4d(arr: np.ndarray | None) -> np.ndarray | None:
+            if arr is None:
+                return None
+            return arr if arr.ndim == 4 else arr[np.newaxis]
+
+        left_ext = ensure_4d(left_ext)
+        right_ext = ensure_4d(right_ext)
+        wrist = ensure_4d(wrist)
+
+        # Determine shape from first available view.
+        ref = next((v for v in [wrist, left_ext, right_ext] if v is not None), None)
+        if ref is None:
+            return np.zeros((1, 352, 640, 3), dtype=np.uint8)
+
+        # Apply per-view crop + resize before stitching.
+        def maybe_preprocess(arr: np.ndarray | None) -> np.ndarray | None:
+            if arr is None:
+                return None
+            return self._preprocess_view(arr)
+
+        left_ext = maybe_preprocess(left_ext)
+        right_ext = maybe_preprocess(right_ext)
+        wrist = maybe_preprocess(wrist)
+        ref = next((v for v in [wrist, left_ext, right_ext] if v is not None), None)
+        if ref is None:
+            raise RuntimeError("Expected at least one DROID camera view after preprocessing.")
+        t, h, w, c = ref.shape
+
+        out = np.zeros((t, 2 * h, 2 * w, c), dtype=ref.dtype)  # (T, 2H, 2W, C)
+
+        # Top row: wrist repeated 2x along width.
+        if wrist is not None:
+            wrist_wide = np.repeat(wrist, 2, axis=2)  # (T, H, 2W, C)
+            out[:, :h, :] = wrist_wide
+
+        # Bottom row: left exterior | right exterior.
+        if left_ext is not None:
+            out[:, h:, :w] = left_ext
+        if right_ext is not None:
+            out[:, h:, w:] = right_ext
+
+        return out
+
+    def _language_template(self, prompt: str) -> str:
+        """Expand the language prompt for the OXE_DROID multi-view format."""
+        prompt = (prompt or "Perform the default behavior.").strip()
+        prompt_lower = prompt.lower()
+        return (
+            "A multi-view video shows that a robot "
+            + prompt_lower
+            + " The video is split into three views: The top view shows the "
+            + "camera view from the robot's wrist, the bottom-left view shows "
+            + "the camera view from the left exterior camera, and the "
+            + "bottom-right view shows the camera view from the right exterior "
+            + "camera. During training, one of the two bottom exterior views "
+            + "may be a black screen (dropped view). The robot "
+            + prompt_lower
+        )
+
+    def _extract_raw_state(self, obs: dict) -> np.ndarray:
+        """OXE_DROID state: 7 joint + 1 gripper = 8 dims."""
+        parts = []
+        if "observation/joint_position" in obs:
+            parts.append(np.asarray(obs["observation/joint_position"], dtype=np.float64).flatten())
+        if "observation/gripper_position" in obs:
+            parts.append(np.asarray(obs["observation/gripper_position"], dtype=np.float64).flatten())
+        if parts:
+            return np.concatenate(parts)
+        return np.zeros(8, dtype=np.float64)
+
+
+register_transform("droid", DroidTransform())

vllm_omni/diffusion/models/dreamzero/transform/roboarena.py

@@ -0,0 +1,41 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""RoboArena dataset transform.
+
+RoboArena uses 0-indexed exterior cameras, 3 views total (OXE_DROID embodiment).
+Same stitching layout as DROID — both map to OXE_DROID in DreamZero.
+
+"""
+
+from __future__ import annotations
+
+from vllm_omni.diffusion.models.dreamzero.transform.base import (
+    register_transform,
+)
+from vllm_omni.diffusion.models.dreamzero.transform.droid import (
+    DroidTransform,
+)
+
+
+class RoboArenaTransform(DroidTransform):
+    """Transform for RoboArena dataset.
+
+    Same embodiment as DROID (OXE_DROID), same stitching and template.
+    Only difference: 0-indexed exterior camera keys.
+
+    RoboArena observation keys (0-indexed):
+        observation/exterior_image_0_left  → left exterior
+        observation/exterior_image_1_left  → right exterior
+        observation/wrist_image_left       → wrist
+
+    """
+
+    IMAGE_KEY_MAP = {
+        "observation/exterior_image_0_left": "images/exterior_0",
+        "observation/exterior_image_1_left": "images/exterior_1",
+        "observation/wrist_image_left": "images/wrist",
+    }
+
+
+register_transform("roboarena", RoboArenaTransform())