[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/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())