From 9d4ca505a9ac496f1ae5adc2573d1a8a4548e509 Mon Sep 17 00:00:00 2001
From: yuzhud <yuzhudong@gmail.com>
Date: Wed, 3 Jun 2026 19:34:42 -0700
Subject: [PATCH 01/12] support DROID policy server

Signed-off-by: yuzhud <yuzhudong@gmail.com>
---
 requirements/common.txt                       |   1 +
 .../models/cosmos3/pipeline_cosmos3.py        | 472 ++++++++++++++++--
 vllm_omni/entrypoints/openai/api_server.py    |   1 +
 vllm_omni/entrypoints/openpi/connection.py    |  53 +-
 vllm_omni/entrypoints/openpi/serving.py       |  11 +-
 5 files changed, 494 insertions(+), 44 deletions(-)

diff --git a/requirements/common.txt b/requirements/common.txt
index fb435218a45..6ca11ac7068 100644
--- a/requirements/common.txt
+++ b/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
diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py
index df9eb835dda..885b47425ae 100644
--- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py
+++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py
@@ -26,6 +26,7 @@
 import numpy as np
 import PIL.Image
 import torch
+import torch.nn.functional as F
 from diffusers import UniPCMultistepScheduler
 from diffusers.utils.torch_utils import randn_tensor
 from diffusers.video_processor import VideoProcessor
@@ -101,6 +102,24 @@
 COSMOS3_T2I_DEFAULT_FLOW_SHIFT = 3.0
 COSMOS3_T2I_DEFAULT_GUIDANCE_INTERVAL: tuple[float, float] = (400.0, 1000.0)
 
+ROBOLAB_DEFAULT_CONDITIONING_FPS = 15.0
+ROBOLAB_DEFAULT_ACTION_CHUNK_SIZE = 32
+ROBOLAB_DEFAULT_IMAGE_HEIGHT = 540
+ROBOLAB_DEFAULT_IMAGE_WIDTH = 640
+ROBOLAB_DEFAULT_RAW_ACTION_DIM = 8
+ROBOLAB_DEFAULT_DOMAIN_NAME = "droid_lerobot"
+ROBOLAB_DEFAULT_RESOLUTION = "480"
+ROBOLAB_DEFAULT_GUIDANCE_SCALE = 3.0
+ROBOLAB_DEFAULT_NUM_INFERENCE_STEPS = 4
+ROBOLAB_DEFAULT_FLOW_SHIFT = 5.0
+ROBOLAB_DEFAULT_SEED = 0
+ROBOLAB_DEFAULT_ACTION_SPACE = "joint_pos"
+ROBOLAB_CONCAT_VIEW_DESCRIPTION = (
+    "The top row is from the wrist-mounted camera. "
+    "The bottom row contains two horizontally concatenated third-person perspective views of the scene from opposite "
+    "sides, with the robot visible."
+)
+
 # Truncation cap on the prompt token count (shared by T2I and T2V).  Prompts
 # are tokenized to their natural length (no padding); this only bounds the
 # UND pathway / GEN cross-attention cost for pathologically long prompts.
@@ -144,6 +163,121 @@ def _normalize_condition_video_keep(value: Any) -> str:
     return keep
 
 
+def _ensure_rgb_uint8_image(value: Any, key: str) -> np.ndarray:
+    image = np.asarray(value)
+    if image.ndim != 3 or image.shape[-1] != 3:
+        raise ValueError(f"{key!r} must have shape [H, W, 3], got {image.shape}.")
+    if image.dtype != np.uint8:
+        image = np.clip(image, 0, 255).astype(np.uint8)
+    return np.ascontiguousarray(image)
+
+
+def _ensure_2d_float_array(value: Any, key: str, width: int | None = None) -> np.ndarray:
+    array = np.asarray(value, dtype=np.float32)
+    if array.ndim == 1:
+        array = array[None, :]
+    if array.ndim != 2:
+        raise ValueError(f"{key!r} must have shape [T, D] or [D], got {array.shape}.")
+    if width is not None and array.shape[-1] != width:
+        raise ValueError(f"{key!r} must have width {width}, got {array.shape[-1]}.")
+    return np.ascontiguousarray(array)
+
+
+def _ensure_gripper_array(value: Any) -> np.ndarray:
+    array = np.asarray(value, dtype=np.float32)
+    if array.ndim == 0:
+        array = array.reshape(1, 1)
+    elif array.ndim == 1:
+        array = array[:, None]
+    if array.ndim != 2 or array.shape[-1] != 1:
+        raise ValueError(f"'observation/gripper_position' must have shape [T, 1], [T], or scalar, got {array.shape}.")
+    return np.ascontiguousarray(array)
+
+
+def _resize_rgb_uint8(image: np.ndarray, size: tuple[int, int]) -> np.ndarray:
+    tensor = torch.from_numpy(image).permute(2, 0, 1).unsqueeze(0).float()
+    resized = F.interpolate(tensor, size=size, mode="bilinear", align_corners=False)
+    return resized.squeeze(0).permute(1, 2, 0).numpy().astype(np.uint8)
+
+
+def _compose_robolab_views(obs: dict[str, Any]) -> np.ndarray | None:
+    required_keys = (
+        "observation/wrist_image_left",
+        "observation/exterior_image_1_left",
+        "observation/exterior_image_2_left",
+    )
+    if not all(key in obs for key in required_keys):
+        return None
+
+    wrist = _ensure_rgb_uint8_image(obs["observation/wrist_image_left"], "observation/wrist_image_left")
+    left_raw = _ensure_rgb_uint8_image(obs["observation/exterior_image_1_left"], "observation/exterior_image_1_left")
+    right_raw = _ensure_rgb_uint8_image(obs["observation/exterior_image_2_left"], "observation/exterior_image_2_left")
+    half_h, half_w = wrist.shape[0] // 2, wrist.shape[1] // 2
+    left = _resize_rgb_uint8(left_raw, (half_h, half_w))
+    right = _resize_rgb_uint8(right_raw, (half_h, half_w))
+    return np.concatenate([wrist, np.concatenate([left, right], axis=1)], axis=0)
+
+
+def _extract_robolab_image(obs: dict[str, Any]) -> np.ndarray:
+    if "observation/image" in obs:
+        return _ensure_rgb_uint8_image(obs["observation/image"], "observation/image")
+    image = _compose_robolab_views(obs)
+    if image is not None:
+        return image
+    raise ValueError("Observation must contain 'observation/image' or RoboLab wrist/exterior image keys.")
+
+
+def _lazy_action_transform_pipeline(max_action_dim: int):
+    try:
+        from cosmos_framework.data.vfm.action.transforms import ActionTransformPipeline
+    except ModuleNotFoundError as exc:
+        raise ModuleNotFoundError(
+            "Cosmos3 RoboLab policy serving requires cosmos_framework on PYTHONPATH so the "
+            "golden ActionTransformPipeline can be reused."
+        ) from exc
+    return ActionTransformPipeline(max_action_dim=max_action_dim, cfg_dropout_rate=0.0)
+
+
+def _convert_midtrain_rotation(value: Any, src: str, dst: str) -> np.ndarray:
+    try:
+        from cosmos_framework.data.vfm.action.pose_utils import convert_rotation
+    except ModuleNotFoundError as exc:
+        raise ModuleNotFoundError(
+            "Cosmos3 RoboLab midtrain action serving requires cosmos_framework pose_utils on PYTHONPATH."
+        ) from exc
+    return convert_rotation(value, src, dst)
+
+
+def _pose_abs_to_rel(*args, **kwargs) -> np.ndarray:
+    try:
+        from cosmos_framework.data.vfm.action.pose_utils import pose_abs_to_rel
+    except ModuleNotFoundError as exc:
+        raise ModuleNotFoundError(
+            "Cosmos3 RoboLab midtrain action serving requires cosmos_framework pose_utils on PYTHONPATH."
+        ) from exc
+    return pose_abs_to_rel(*args, **kwargs)
+
+
+def _pose_rel_to_abs(*args, **kwargs) -> np.ndarray:
+    try:
+        from cosmos_framework.data.vfm.action.pose_utils import pose_rel_to_abs
+    except ModuleNotFoundError as exc:
+        raise ModuleNotFoundError(
+            "Cosmos3 RoboLab midtrain action serving requires cosmos_framework pose_utils on PYTHONPATH."
+        ) from exc
+    return pose_rel_to_abs(*args, **kwargs)
+
+
+def _build_abs_pose_from_components(*args, **kwargs) -> np.ndarray:
+    try:
+        from cosmos_framework.data.vfm.action.pose_utils import build_abs_pose_from_components
+    except ModuleNotFoundError as exc:
+        raise ModuleNotFoundError(
+            "Cosmos3 RoboLab midtrain action serving requires cosmos_framework pose_utils on PYTHONPATH."
+        ) from exc
+    return build_abs_pose_from_components(*args, **kwargs)
+
+
 # ---------------------------------------------------------------------------
 # Post-process function (registered in registry.py)
 # ---------------------------------------------------------------------------
@@ -654,6 +788,8 @@ def __init__(
 
         self._guidance_scale = None
         self._num_timesteps = None
+        self._robolab_transform = None
+        self._robolab_rng = np.random.default_rng(ROBOLAB_DEFAULT_SEED)
 
         # Set True by ``enable_cache_for_cosmos3`` when cache-dit is enabled on
         # this pipeline. Tells the sequential-CFG loop to keep paired
@@ -875,6 +1011,212 @@ def _get_sp_param(sp: OmniDiffusionSamplingParams, key: str, default: Any = None
             return val
         return default
 
+    @staticmethod
+    def _normalize_robolab_action_space(value: Any) -> str:
+        action_space = str(value or ROBOLAB_DEFAULT_ACTION_SPACE).strip().lower()
+        aliases = {
+            "jointpos": "joint_pos",
+            "joint_pos": "joint_pos",
+            "abs_ik": "midtrain",
+            "midtrain": "midtrain",
+        }
+        if action_space not in aliases:
+            raise ValueError(
+                f"Unsupported RoboLab action_space={value!r}; expected joint_pos/jointpos or midtrain/abs_ik."
+            )
+        return aliases[action_space]
+
+    def _get_robolab_transform(self):
+        if self._robolab_transform is None:
+            action_dim = int(getattr(self.transformer, "action_dim", 64))
+            self._robolab_transform = _lazy_action_transform_pipeline(action_dim)
+        return self._robolab_transform
+
+    def _next_robolab_seed(self, extra: dict[str, Any]) -> int:
+        base_seed = int(extra.get("robolab_seed") or ROBOLAB_DEFAULT_SEED)
+        if self._truthy(extra.get("deterministic_seed")):
+            return base_seed
+        explicit_seed = extra.get("seed")
+        if explicit_seed is not None:
+            return int(explicit_seed)
+        return int(self._robolab_rng.integers(0, 2**31))
+
+    def _build_robolab_policy_inputs(
+        self,
+        sp: OmniDiffusionSamplingParams,
+    ) -> dict[str, Any] | None:
+        extra = sp.extra_args if isinstance(sp.extra_args, dict) else {}
+        obs = extra.get("robot_obs")
+        if obs is None:
+            obs = extra.get("observation")
+        if obs is None:
+            return None
+        if not isinstance(obs, dict):
+            raise TypeError(f"Cosmos3 RoboLab observation must be a dict, got {type(obs)!r}.")
+
+        prompt = obs.get("prompt")
+        if not isinstance(prompt, str):
+            raise ValueError("RoboLab observation must contain string key 'prompt'.")
+
+        def extra_param(key: str, default: Any) -> Any:
+            value = extra.get(key)
+            return default if value is None else value
+
+        def extra_param_alias(primary_key: str, alias_key: str, default: Any) -> Any:
+            value = extra.get(primary_key)
+            if value is not None:
+                return value
+            value = extra.get(alias_key)
+            return default if value is None else value
+
+        action_space = self._normalize_robolab_action_space(extra_param("action_space", ROBOLAB_DEFAULT_ACTION_SPACE))
+        action_chunk_size = int(extra_param("action_chunk_size", ROBOLAB_DEFAULT_ACTION_CHUNK_SIZE))
+        raw_action_dim_default = ROBOLAB_DEFAULT_RAW_ACTION_DIM if action_space == "joint_pos" else 10
+        raw_action_dim = int(extra_param("raw_action_dim", raw_action_dim_default))
+        image_h = int(extra_param("image_height", ROBOLAB_DEFAULT_IMAGE_HEIGHT))
+        image_w = int(extra_param("image_width", ROBOLAB_DEFAULT_IMAGE_WIDTH))
+        history_length = int(extra_param("history_length", 1))
+        use_state = self._truthy(extra_param("use_state", True))
+        resolution = str(extra_param("resolution", ROBOLAB_DEFAULT_RESOLUTION))
+        fps = float(extra_param("conditioning_fps", ROBOLAB_DEFAULT_CONDITIONING_FPS))
+        domain_name = str(extra_param("domain_name", ROBOLAB_DEFAULT_DOMAIN_NAME))
+        domain_id = resolve_domain_id(domain_name=domain_name, require_explicit=True)
+
+        if history_length < (1 if use_state else 0):
+            raise ValueError("RoboLab history_length must be >= 1 when use_state is true.")
+        if action_chunk_size <= 0:
+            raise ValueError(f"RoboLab action_chunk_size must be positive, got {action_chunk_size}.")
+        if raw_action_dim <= 0:
+            raise ValueError(f"RoboLab raw_action_dim must be positive, got {raw_action_dim}.")
+
+        image = _extract_robolab_image(obs)
+        if image.shape[:2] != (image_h, image_w):
+            image = _resize_rgb_uint8(image, (image_h, image_w))
+
+        t_frames = action_chunk_size + 1
+        video = torch.zeros((3, t_frames, image_h, image_w), dtype=torch.uint8)
+        video[:, 0] = torch.from_numpy(image.copy()).permute(2, 0, 1)
+
+        use_state_rows = 1 if use_state else 0
+        action = torch.zeros((action_chunk_size + use_state_rows, raw_action_dim), dtype=torch.float32)
+        history_action = None
+        num_history_rows = history_length - use_state_rows
+        gripper_position = 1.0 - _ensure_gripper_array(obs["observation/gripper_position"])
+
+        if action_space == "joint_pos":
+            joint_position = _ensure_2d_float_array(obs["observation/joint_position"], "observation/joint_position", 7)
+            if use_state:
+                action[0] = torch.from_numpy(np.concatenate((joint_position[-1], gripper_position[-1])))
+            if num_history_rows > 0:
+                if len(joint_position) < num_history_rows + 1:
+                    raise ValueError("Not enough joint_position rows for requested history_length.")
+                history_np = np.concatenate(
+                    (joint_position[-num_history_rows - 1 : -1], gripper_position[-num_history_rows - 1 : -1]),
+                    axis=-1,
+                )
+                history_action = torch.from_numpy(history_np).float()
+        else:
+            eef_pos = _ensure_2d_float_array(obs["observation/eef_pos"], "observation/eef_pos", 3)
+            eef_quat = _ensure_2d_float_array(obs["observation/eef_quat"], "observation/eef_quat", 4)
+            if use_state:
+                rot6d = _convert_midtrain_rotation(eef_quat[-1], "quat_xyzw", "rot6d")
+                action[0] = torch.from_numpy(np.concatenate((eef_pos[-1], rot6d, gripper_position[-1])))
+            if num_history_rows > 0:
+                if len(eef_pos) < num_history_rows + 1 or len(eef_quat) < num_history_rows + 1:
+                    raise ValueError("Not enough eef_pos/eef_quat rows for requested history_length.")
+                poses_abs = _build_abs_pose_from_components(eef_pos, eef_quat, "quat_xyzw")
+                poses_rel = _pose_abs_to_rel(poses_abs, rotation_format="rot6d", pose_convention="backward_framewise")
+                history_np = np.concatenate(
+                    [poses_rel[-num_history_rows:], gripper_position[-num_history_rows:]],
+                    axis=-1,
+                )
+                history_action = torch.from_numpy(history_np).float()
+
+        sample: dict[str, Any] = {
+            "ai_caption": prompt,
+            "video": video,
+            "action": action,
+            "conditioning_fps": torch.tensor(fps, dtype=torch.long),
+            "mode": ACTION_MODE_POLICY,
+            "domain_id": torch.tensor(domain_id, dtype=torch.long),
+            "viewpoint": "concat_view",
+            "additional_view_description": ROBOLAB_CONCAT_VIEW_DESCRIPTION,
+        }
+        if history_action is not None:
+            sample["history_action"] = history_action
+
+        sample = self._get_robolab_transform()(sample, resolution)
+        sequence_plan = sample["sequence_plan"]
+        video_tensor = sample["video"].float() / 127.5 - 1.0
+        raw_action_dim_tensor = sample.get("raw_action_dim")
+        if isinstance(raw_action_dim_tensor, torch.Tensor):
+            transformed_raw_action_dim = int(raw_action_dim_tensor.item())
+        else:
+            transformed_raw_action_dim = raw_action_dim
+
+        return {
+            "prompt": sample["ai_caption"],
+            "video_tensor": video_tensor.unsqueeze(0),
+            "action_tensor": sample["action"].float(),
+            "action_condition_indexes": list(getattr(sequence_plan, "condition_frame_indexes_action", []) or []),
+            "action_start_frame_offset": int(getattr(sequence_plan, "action_start_frame_offset", 1)),
+            "raw_action_dim": transformed_raw_action_dim,
+            "domain_id": domain_id,
+            "fps": fps,
+            "height": int(video_tensor.shape[-2]),
+            "width": int(video_tensor.shape[-1]),
+            "num_frames": int(video_tensor.shape[1]),
+            "num_inference_steps": int(
+                extra_param_alias("num_inference_steps", "num_steps", ROBOLAB_DEFAULT_NUM_INFERENCE_STEPS)
+            ),
+            "guidance_scale": float(extra_param_alias("guidance_scale", "guidance", ROBOLAB_DEFAULT_GUIDANCE_SCALE)),
+            "flow_shift": float(extra_param_alias("flow_shift", "shift", ROBOLAB_DEFAULT_FLOW_SHIFT)),
+            "seed": self._next_robolab_seed(extra),
+            "history_length": history_length,
+            "action_space": action_space,
+            "observation": obs,
+        }
+
+    @staticmethod
+    def _build_action_condition_mask_from_indexes(
+        indexes: list[int],
+        action_length: int,
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        mask = torch.zeros(1, action_length, 1, device=device, dtype=dtype)
+        for idx in indexes:
+            if idx < 0 or idx >= action_length:
+                raise ValueError(f"Action condition index {idx} is out of range for action length {action_length}.")
+            mask[:, idx, :] = 1.0
+        return mask
+
+    def _postprocess_robolab_action(self, action: torch.Tensor, inputs: dict[str, Any]) -> np.ndarray:
+        action_np = action[0].float().cpu().numpy()
+        history_length = int(inputs["history_length"])
+        action_np = action_np[history_length:]
+        action_np[:, -1] = 1.0 - action_np[:, -1]
+
+        if inputs["action_space"] == "midtrain":
+            obs = inputs["observation"]
+            eef_pos = _ensure_2d_float_array(obs["observation/eef_pos"], "observation/eef_pos", 3)
+            eef_quat = _ensure_2d_float_array(obs["observation/eef_quat"], "observation/eef_quat", 4)
+            initial_pose = np.eye(4, dtype=np.float32)
+            initial_pose[:3, :3] = _convert_midtrain_rotation(eef_quat[-1], "quat_xyzw", "matrix")
+            initial_pose[:3, 3] = eef_pos[-1]
+            abs_pose = _pose_rel_to_abs(
+                action_np[:, :9],
+                rotation_format="rot6d",
+                pose_convention="backward_framewise",
+                initial_pose=initial_pose,
+            )
+            position = abs_pose[1:, :3, 3]
+            quat_xyzw = _convert_midtrain_rotation(abs_pose[1:, :3, :3], "matrix", "quat_xyzw")
+            action_np = np.concatenate([position, quat_xyzw, action_np[:, 9:]], axis=-1)
+
+        return np.asarray(action_np, dtype=np.float32)
+
     @staticmethod
     def _truthy(value) -> bool:
         if isinstance(value, str):
@@ -1484,9 +1826,25 @@ def _prepare_action_latents(
         raw_action_dim: int | None,
         generator: torch.Generator,
         sp,
+        clean_action: torch.Tensor | None = None,
+        condition_indexes: list[int] | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
         action_dim = int(getattr(self.transformer, "action_dim", 64))
-        if mode == ACTION_MODE_FORWARD_DYNAMICS:
+        if clean_action is not None:
+            action = clean_action.detach().to(dtype=torch.float32)
+            if action.ndim == 3 and action.shape[0] == 1:
+                action = action.squeeze(0)
+            if action.ndim != 2:
+                raise ValueError(f"Cosmos3 clean action must have shape [T, D], got {tuple(action.shape)}.")
+            if action.shape[0] < action_chunk_size:
+                pad = action[-1:].repeat(action_chunk_size - action.shape[0], 1)
+                action = torch.cat([action, pad], dim=0)
+            elif action.shape[0] > action_chunk_size:
+                action = action[:action_chunk_size]
+            if raw_action_dim is None:
+                raw_action_dim = int(action.shape[-1])
+            clean_action = pad_action_to_dim(action, action_dim)
+        elif mode == ACTION_MODE_FORWARD_DYNAMICS:
             action = load_action_tensor(self._get_sp_param(sp, "action", None))
             if action.shape[0] < action_chunk_size:
                 pad = action[-1:].repeat(action_chunk_size - action.shape[0], 1)
@@ -1508,12 +1866,20 @@ def _prepare_action_latents(
             raise ValueError(f"Cosmos3 raw_action_dim must be in [1, {action_dim}], got {raw_action_dim}.")
 
         clean_action = clean_action.to(device=self.device, dtype=self.dtype).unsqueeze(0)
-        condition_mask = build_action_condition_mask(
-            mode,
-            action_chunk_size,
-            device=self.device,
-            dtype=self.dtype,
-        )
+        if condition_indexes is None:
+            condition_mask = build_action_condition_mask(
+                mode,
+                action_chunk_size,
+                device=self.device,
+                dtype=self.dtype,
+            )
+        else:
+            condition_mask = self._build_action_condition_mask_from_indexes(
+                condition_indexes,
+                action_chunk_size,
+                device=self.device,
+                dtype=self.dtype,
+            )
         noise = randn_tensor(
             (1, action_chunk_size, action_dim),
             generator=generator,
@@ -1811,8 +2177,15 @@ def forward(
         if len(req.prompts) > 1:
             raise ValueError("Cosmos3OmniDiffusersPipeline currently supports a single prompt per request.")
 
+        sp = req.sampling_params
         prompt_data = req.prompts[0]
-        if isinstance(prompt_data, str):
+        robolab_inputs = self._build_robolab_policy_inputs(sp)
+        if robolab_inputs is not None:
+            prompt = robolab_inputs["prompt"]
+            negative_prompt = None
+            image_tensor = None
+            action_video_tensor = robolab_inputs["video_tensor"]
+        elif isinstance(prompt_data, str):
             prompt = prompt_data
             negative_prompt = None
             image_tensor = None
@@ -1824,10 +2197,9 @@ def forward(
             image_tensor = additional_info.get("preprocessed_image")
             video_tensor = additional_info.get("preprocessed_video")
 
-        sp = req.sampling_params
-        is_t2i = self._is_t2i_request(req)
-        sound_enabled = self._is_sound_request(prompt_data, sp)
-        action_mode = self._get_action_mode(prompt_data, sp)
+        is_t2i = False if robolab_inputs is not None else self._is_t2i_request(req)
+        sound_enabled = False if robolab_inputs is not None else self._is_sound_request(prompt_data, sp)
+        action_mode = ACTION_MODE_POLICY if robolab_inputs is not None else self._get_action_mode(prompt_data, sp)
         action_enabled = action_mode is not None
         action_video_tensor = video_tensor if action_enabled else None
         if action_enabled and is_t2i:
@@ -1885,7 +2257,16 @@ def forward(
             default_guidance_interval = None
             batch_size = 1  # Existing video pipeline assumes B=1.
 
-        if action_enabled:
+        if robolab_inputs is not None:
+            height = robolab_inputs["height"]
+            width = robolab_inputs["width"]
+            num_frames = robolab_inputs["num_frames"]
+            action_chunk_size = int(robolab_inputs["action_tensor"].shape[0])
+            num_inference_steps = robolab_inputs["num_inference_steps"]
+            guidance_scale = float(robolab_inputs["guidance_scale"])
+            default_flow_shift = robolab_inputs["flow_shift"]
+            default_guidance_interval = None
+        elif action_enabled:
             action_chunk_param = self._get_sp_param(sp, "action_chunk_size", None)
             if action_chunk_param is not None:
                 action_chunk_size = int(action_chunk_param)
@@ -1909,18 +2290,28 @@ def forward(
 
         domain_id = None
         if action_enabled:
-            domain_id = resolve_domain_id(
-                domain_id=self._get_sp_param(sp, "domain_id", None),
-                domain_name=self._get_sp_param(sp, "domain_name", None),
-                require_explicit=True,
-            )
+            if robolab_inputs is not None:
+                domain_id = int(robolab_inputs["domain_id"])
+            else:
+                domain_id = resolve_domain_id(
+                    domain_id=self._get_sp_param(sp, "domain_id", None),
+                    domain_name=self._get_sp_param(sp, "domain_name", None),
+                    require_explicit=True,
+                )
 
         # Runtime controls: prefer ``extra_args`` (OpenAI endpoints write
         # there) over direct attrs.
-        flow_shift_target = float(self._get_sp_param(sp, "flow_shift", default_flow_shift))
+        if robolab_inputs is not None:
+            flow_shift_target = float(default_flow_shift)
+        else:
+            flow_shift_target = float(self._get_sp_param(sp, "flow_shift", default_flow_shift))
         guidance_interval = self._get_sp_param(sp, "guidance_interval", default_guidance_interval)
 
-        frame_rate = self._get_sp_param(sp, "resolved_frame_rate") or self._get_sp_param(sp, "frame_rate") or 24.0
+        frame_rate = (
+            self._get_sp_param(sp, "resolved_frame_rate")
+            or self._get_sp_param(sp, "frame_rate")
+            or (robolab_inputs["fps"] if robolab_inputs is not None else 24.0)
+        )
         max_sequence_length = (
             self._get_sp_param(sp, "max_sequence_length", COSMOS3_DEFAULT_MAX_SEQUENCE_LENGTH)
             or COSMOS3_DEFAULT_MAX_SEQUENCE_LENGTH
@@ -1954,7 +2345,10 @@ def forward(
 
         generator = sp.generator
         if generator is None:
-            seed = sp.seed if sp.seed is not None else 42
+            if robolab_inputs is not None:
+                seed = int(robolab_inputs["seed"])
+            else:
+                seed = sp.seed if sp.seed is not None else 42
             generator = torch.Generator(device=self.device).manual_seed(seed)
 
         # --- Format prompts & tokenize (B=1; reused across loop iterations
@@ -2008,17 +2402,29 @@ def forward(
                     )
                 image_tensor = action_video_tensor[:, :, 0]
 
-            raw_action_dim_param = self._get_sp_param(sp, "raw_action_dim", None)
-            raw_action_dim = int(raw_action_dim_param) if raw_action_dim_param is not None else None
+            if robolab_inputs is not None:
+                raw_action_dim = int(robolab_inputs["raw_action_dim"])
+                clean_action = robolab_inputs["action_tensor"]
+                action_condition_indexes = robolab_inputs["action_condition_indexes"]
+            else:
+                raw_action_dim_param = self._get_sp_param(sp, "raw_action_dim", None)
+                raw_action_dim = int(raw_action_dim_param) if raw_action_dim_param is not None else None
+                clean_action = None
+                action_condition_indexes = None
             action_prepared = self._prepare_action_latents(
                 mode=action_mode,
                 action_chunk_size=action_chunk_size,
                 raw_action_dim=raw_action_dim,
                 generator=generator,
                 sp=sp,
+                clean_action=clean_action,
+                condition_indexes=action_condition_indexes,
             )
             action_latents, action_velocity_mask, action_condition_latents, raw_action_dim = action_prepared
-            action_offset = action_start_frame_offset(action_mode, action_chunk_size, num_frames)
+            if robolab_inputs is not None:
+                action_offset = int(robolab_inputs["action_start_frame_offset"])
+            else:
+                action_offset = action_start_frame_offset(action_mode, action_chunk_size, num_frames)
 
         if action_enabled and action_video_tensor is not None:
             latents, velocity_mask, condition_latents = self._prepare_latents_action_video(
@@ -2153,14 +2559,18 @@ def _run_diffusion(start_latents):
             if action_latents is None or raw_action_dim is None or domain_id is None:
                 raise ValueError("Cosmos3 action generation finished without action latents.")
             action = action_latents[:, :, :raw_action_dim].detach().cpu()
+            custom_action_output: dict[str, Any] = {
+                "action": action,
+                "raw_action_dim": raw_action_dim,
+                "action_mode": action_mode,
+                "domain_id": domain_id,
+            }
+            output = {"video": video}
+            if robolab_inputs is not None:
+                output["actions"] = self._postprocess_robolab_action(action, robolab_inputs)
             return DiffusionOutput(
-                output={"video": video},
-                custom_output={
-                    "action": action,
-                    "raw_action_dim": raw_action_dim,
-                    "action_mode": action_mode,
-                    "domain_id": domain_id,
-                },
+                output=output,
+                custom_output=custom_action_output,
             )
 
         return DiffusionOutput(output={"image": video} if is_t2i else {"video": video})
diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py
index c11db925f71..6ff9dcc4818 100755
--- a/vllm_omni/entrypoints/openai/api_server.py
+++ b/vllm_omni/entrypoints/openai/api_server.py
@@ -1531,6 +1531,7 @@ async def realtime_websocket(websocket: WebSocket):
     await connection.handle_connection()
 
 
+@router.websocket("/")
 @router.websocket("/v1/realtime/robot/openpi")
 async def realtime_robot_openpi(websocket: WebSocket):
     """WebSocket endpoint for robot policy inference via OpenPI messages."""
diff --git a/vllm_omni/entrypoints/openpi/connection.py b/vllm_omni/entrypoints/openpi/connection.py
index 670c37a9233..a47781037f7 100644
--- a/vllm_omni/entrypoints/openpi/connection.py
+++ b/vllm_omni/entrypoints/openpi/connection.py
@@ -14,6 +14,8 @@
 import asyncio
 from typing import Any
 
+import msgpack
+import numpy as np
 from fastapi import WebSocket
 from starlette.websockets import WebSocketDisconnect
 from vllm.logger import init_logger
@@ -27,24 +29,51 @@
 MAX_OPENPI_PAYLOAD_BYTES = 64 * 1024 * 1024
 
 
-def _get_msgpack_numpy() -> Any:
-    try:
-        from openpi_client import msgpack_numpy
-    except ImportError as exc:
-        raise ImportError(
-            "The `/v1/realtime/robot/openpi` endpoint requires the optional "
-            "`openpi-client` dependency. Install it with `pip install openpi-client`."
-        ) from exc
-
-    return msgpack_numpy
+def _pack_array(obj: Any) -> Any:
+    if isinstance(obj, (np.ndarray, np.generic)) and obj.dtype.kind in ("V", "O", "c"):
+        raise ValueError(f"Unsupported dtype: {obj.dtype}")
+    if isinstance(obj, np.ndarray):
+        return {
+            b"__ndarray__": True,
+            b"data": obj.tobytes(),
+            b"dtype": obj.dtype.str,
+            b"shape": obj.shape,
+        }
+    if isinstance(obj, np.generic):
+        return {
+            b"__npgeneric__": True,
+            b"data": obj.item(),
+            b"dtype": obj.dtype.str,
+        }
+    return obj
+
+
+def _unpack_array(obj: dict[Any, Any]) -> Any:
+    if b"__ndarray__" in obj:
+        return np.ndarray(
+            buffer=obj[b"data"],
+            dtype=np.dtype(obj[b"dtype"]),
+            shape=tuple(obj[b"shape"]),
+        )
+    if "__ndarray__" in obj:
+        return np.ndarray(
+            buffer=obj["data"],
+            dtype=np.dtype(obj["dtype"]),
+            shape=tuple(obj["shape"]),
+        )
+    if b"__npgeneric__" in obj:
+        return np.dtype(obj[b"dtype"]).type(obj[b"data"])
+    if "__npgeneric__" in obj:
+        return np.dtype(obj["dtype"]).type(obj["data"])
+    return obj
 
 
 def _pack(obj: Any) -> bytes:
-    return _get_msgpack_numpy().packb(obj)
+    return msgpack.packb(obj, default=_pack_array)
 
 
 def _unpack(data: bytes) -> Any:
-    return _get_msgpack_numpy().unpackb(data)
+    return msgpack.unpackb(data, object_hook=_unpack_array)
 
 
 class RobotRealtimeConnection:
diff --git a/vllm_omni/entrypoints/openpi/serving.py b/vllm_omni/entrypoints/openpi/serving.py
index ec5c88408ea..c2cfe3a4429 100644
--- a/vllm_omni/entrypoints/openpi/serving.py
+++ b/vllm_omni/entrypoints/openpi/serving.py
@@ -21,6 +21,10 @@
 logger = init_logger(__name__)
 
 ActionOutput = np.ndarray | dict[str, np.ndarray]
+COSMOS3_DROID_POLICY_MODELS = {
+    "Cosmos3-Nano-Policy-DROID",
+    "nvidia/Cosmos3-Nano-Policy-DROID",
+}
 
 
 def _to_builtin_container(value: Any) -> Any:
@@ -74,7 +78,12 @@ def __init__(
     ) -> None:
         self.engine_client = engine_client
         self.model_name = model_name
-        self.policy_server_config = self._get_policy_server_config(engine_client)
+        try:
+            self.policy_server_config = self._get_policy_server_config(engine_client)
+        except ValueError:
+            if model_name not in COSMOS3_DROID_POLICY_MODELS:
+                raise
+            self.policy_server_config = PolicyServerConfig({})
         self._request_counter = count()
 
     @classmethod

From 663f1a5ecdb265e271c7b08d8fdd8e24f862c890 Mon Sep 17 00:00:00 2001
From: yuzhud <yuzhudong@gmail.com>
Date: Sun, 7 Jun 2026 12:25:09 -0700
Subject: [PATCH 02/12] parity bridge

Signed-off-by: yuzhud <yuzhudong@gmail.com>
---
 vllm_omni/diffusion/diffusion_engine.py       |  14 ++-
 .../models/cosmos3/pipeline_cosmos3.py        | 105 ++++++++++++++----
 .../models/cosmos3/transformer_cosmos3.py     |  46 +++++---
 3 files changed, 124 insertions(+), 41 deletions(-)

diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py
index ce1185c6689..81eb14a3c8a 100644
--- a/vllm_omni/diffusion/diffusion_engine.py
+++ b/vllm_omni/diffusion/diffusion_engine.py
@@ -261,8 +261,14 @@ async def step(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]:
         if self.od_config.enable_cpu_offload:
             output_data = _move_tensor_tree_to_cpu(output_data)
 
+        custom_output = output.custom_output or {}
+        action_payload = custom_output.get("actions")
+        action_only_output = action_payload is not None and isinstance(output_data, dict) and not output_data
+
         postprocess_start_time = time.perf_counter()
-        if self.post_process_func is not None:
+        if action_only_output:
+            outputs = {}
+        elif self.post_process_func is not None:
             # Some video pipelines need request-level controls during
             # postprocess (for example worker-side frame interpolation).
             if self._post_process_accepts_sampling_params:
@@ -272,17 +278,17 @@ async def step(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]:
         else:
             outputs = output_data
         audio_payload = None
-        custom_output = output.custom_output or {}
         model_audio_sample_rate = None
         model_fps = None
-        action_payload = None
         if isinstance(outputs, dict):
             audio_payload = outputs.get("audio")
-            action_payload = outputs.get("actions")
+            action_payload = outputs.get("actions", action_payload)
             custom_output.update(outputs.get("custom_output") or {})
             model_audio_sample_rate = outputs.get("audio_sample_rate")
             model_fps = outputs.get("fps")
             outputs = outputs.get("video", outputs)
+        if action_payload is None:
+            action_payload = custom_output.get("actions")
         postprocess_time = time.perf_counter() - postprocess_start_time
         logger.debug("Post-processing completed in %.4f seconds", postprocess_time)
 
diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py
index 885b47425ae..a948c545aa9 100644
--- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py
+++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py
@@ -238,6 +238,24 @@ def _lazy_action_transform_pipeline(max_action_dim: int):
     return ActionTransformPipeline(max_action_dim=max_action_dim, cfg_dropout_rate=0.0)
 
 
+def _build_robolab_unipc_scheduler(num_steps: int, shift: float, device: torch.device):
+    try:
+        from cosmos_framework.model.vfm.diffusion.samplers.fm_solvers_unipc import FlowUniPCMultistepScheduler
+    except ModuleNotFoundError as exc:
+        raise ModuleNotFoundError(
+            "Cosmos3 RoboLab policy serving requires cosmos_framework on PYTHONPATH so the "
+            "golden FlowUniPCMultistepScheduler can be reused."
+        ) from exc
+
+    scheduler = FlowUniPCMultistepScheduler(
+        num_train_timesteps=1000,
+        shift=1.0,
+        use_dynamic_shifting=False,
+    )
+    scheduler.set_timesteps(num_steps, device=device, shift=float(shift))
+    return scheduler
+
+
 def _convert_midtrain_rotation(value: Any, src: str, dst: str) -> np.ndarray:
     try:
         from cosmos_framework.data.vfm.action.pose_utils import convert_rotation
@@ -1165,6 +1183,7 @@ def extra_param_alias(primary_key: str, alias_key: str, default: Any) -> Any:
             "fps": fps,
             "height": int(video_tensor.shape[-2]),
             "width": int(video_tensor.shape[-1]),
+            "image_size": sample.get("image_size"),
             "num_frames": int(video_tensor.shape[1]),
             "num_inference_steps": int(
                 extra_param_alias("num_inference_steps", "num_steps", ROBOLAB_DEFAULT_NUM_INFERENCE_STEPS)
@@ -1655,7 +1674,28 @@ def _encode_conditioning_video(
 
         return latent.to(self.dtype)
 
-    def _encode_video_tensor(self, video_tensor: torch.Tensor) -> torch.Tensor:
+    def _latent_hw_from_image_size(self, image_size: Any | None) -> tuple[int, int] | None:
+        if image_size is None:
+            return None
+        if isinstance(image_size, torch.Tensor):
+            frame_size = image_size.detach().cpu().flatten()
+        else:
+            frame_size = torch.as_tensor(image_size).flatten()
+        if frame_size.numel() < 4:
+            return None
+        orig_h = int(frame_size[2].item())
+        orig_w = int(frame_size[3].item())
+        spatial_factor = int(self.vae_scale_factor_spatial)
+        return max(orig_h // spatial_factor, 1), max(orig_w // spatial_factor, 1)
+
+    def _crop_latent_to_image_size(self, latent: torch.Tensor, image_size: Any | None) -> torch.Tensor:
+        latent_hw = self._latent_hw_from_image_size(image_size)
+        if latent_hw is None:
+            return latent
+        h_latent, w_latent = latent_hw
+        return latent[:, :, :, :h_latent, :w_latent].contiguous()
+
+    def _encode_video_tensor(self, video_tensor: torch.Tensor, image_size: Any | None = None) -> torch.Tensor:
         """VAE-encode a preprocessed pixel video [1, 3, T, H, W]."""
         if video_tensor.ndim == 4:
             video_tensor = video_tensor.unsqueeze(0)
@@ -1677,6 +1717,7 @@ def _encode_video_tensor(self, video_tensor: torch.Tensor) -> torch.Tensor:
             scaling_factor = getattr(self.vae.config, "scaling_factor", 1.0)
             latent = latent * scaling_factor
 
+        latent = self._crop_latent_to_image_size(latent, image_size)
         return latent.to(self.dtype)
 
     def _prepare_latents_i2v(
@@ -1787,13 +1828,18 @@ def _prepare_latents_action_video(
         width: int,
         num_frames: int,
         generator: torch.Generator,
+        image_size: Any | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """Prepare video latents for action modes with mode-specific conditioning."""
         del height, width
         C = self.transformer.latent_channel_size
         T_lat = (num_frames - 1) // self.vae_scale_factor_temporal + 1
-        H_lat = video_tensor.shape[-2] // self.vae_scale_factor_spatial
-        W_lat = video_tensor.shape[-1] // self.vae_scale_factor_spatial
+        latent_hw = self._latent_hw_from_image_size(image_size)
+        if latent_hw is None:
+            H_lat = video_tensor.shape[-2] // self.vae_scale_factor_spatial
+            W_lat = video_tensor.shape[-1] // self.vae_scale_factor_spatial
+        else:
+            H_lat, W_lat = latent_hw
 
         noise = randn_tensor(
             (1, C, T_lat, H_lat, W_lat),
@@ -1801,7 +1847,7 @@ def _prepare_latents_action_video(
             device=self.device,
             dtype=self.dtype,
         )
-        cond_latent = self._encode_video_tensor(video_tensor)
+        cond_latent = self._encode_video_tensor(video_tensor, image_size=image_size)
         if cond_latent.shape[2:] != noise.shape[2:]:
             raise ValueError(
                 "Cosmos3 action video latent shape mismatch: "
@@ -1914,6 +1960,7 @@ def diffuse(
         condition_latents: torch.Tensor | None = None,
         guidance_interval: tuple[float, float] | None = None,
         raw_action_dim: int | None = None,
+        scheduler: Any | None = None,
     ) -> torch.Tensor | tuple[torch.Tensor, ...]:
         """Denoising loop with 3-mode CFG support (parallel, sequential, none).
 
@@ -1944,6 +1991,7 @@ def diffuse(
         """
         do_cfg = guidance_scale > 1.0
         cfg_parallel = self._cfg_parallel_active() and do_cfg
+        step_scheduler = scheduler if scheduler is not None else self.scheduler
         self.transformer.reset_cache()
 
         def _cfg_active_at(t: torch.Tensor) -> bool:
@@ -2021,11 +2069,11 @@ def _step(
                 if raw_action_dim is not None and 0 < raw_action_dim < action_pred.shape[-1]:
                     action_pred[..., raw_action_dim:] = 0
             if action_latents is None and sound_latents is None:
-                latents = self.scheduler.step(video_pred, t, latents, return_dict=False)[0]
+                latents = step_scheduler.step(video_pred, t, latents, return_dict=False)[0]
             else:
                 packed_noise, shapes, numels = _pack_joint(video_pred, action_pred, sound_pred)
                 packed_latents, _, _ = _pack_joint(latents, action_latents, sound_latents)
-                packed_next = self.scheduler.step(packed_noise, t, packed_latents, return_dict=False)[0]
+                packed_next = step_scheduler.step(packed_noise, t, packed_latents, return_dict=False)[0]
                 unpacked = _unpack_joint(packed_next, shapes, numels)
                 latents = unpacked[0]
                 idx = 1
@@ -2339,14 +2387,17 @@ def forward(
         self._num_timesteps = num_inference_steps
 
         # Always resolve to a concrete target shift for this request, then
-        # update the scheduler.  This is what guarantees mode-to-mode
-        # transitions restore the right schedule (no T2I to T2V leak).
-        self._set_flow_shift(flow_shift_target)
+        # update the shared Diffusers scheduler for non-RoboLab paths. RoboLab
+        # uses a request-local Cosmos UniPC scheduler to match the golden
+        # action server.
+        if robolab_inputs is None:
+            self._set_flow_shift(flow_shift_target)
 
         generator = sp.generator
+        robolab_seed = int(robolab_inputs["seed"]) if robolab_inputs is not None else None
         if generator is None:
-            if robolab_inputs is not None:
-                seed = int(robolab_inputs["seed"])
+            if robolab_seed is not None:
+                seed = robolab_seed
             else:
                 seed = sp.seed if sp.seed is not None else 42
             generator = torch.Generator(device=self.device).manual_seed(seed)
@@ -2434,6 +2485,7 @@ def forward(
                 width,
                 num_frames,
                 generator,
+                image_size=robolab_inputs.get("image_size") if robolab_inputs is not None else None,
             )
             image_latent = condition_latents[:, :, 0:1]
         elif is_v2v:
@@ -2493,10 +2545,14 @@ def forward(
             )
 
         def _run_diffusion(start_latents):
-            self.scheduler.set_timesteps(num_inference_steps, device=self.device)
+            if robolab_inputs is not None:
+                scheduler = _build_robolab_unipc_scheduler(num_inference_steps, flow_shift_target, self.device)
+            else:
+                self.scheduler.set_timesteps(num_inference_steps, device=self.device)
+                scheduler = self.scheduler
             return self.diffuse(
                 latents=start_latents,
-                timesteps=self.scheduler.timesteps,
+                timesteps=scheduler.timesteps,
                 cond_ids=cond_ids,
                 cond_mask=cond_mask,
                 uncond_ids=uncond_ids,
@@ -2512,6 +2568,7 @@ def _run_diffusion(start_latents):
                 condition_latents=condition_latents,
                 guidance_interval=guidance_interval,
                 raw_action_dim=raw_action_dim,
+                scheduler=scheduler,
             )
 
         if is_t2i and batch_size > 1:
@@ -2539,12 +2596,17 @@ def _run_diffusion(start_latents):
                 latents = diffusion_output
 
         # --- Decode ---
-        if _is_rank_zero():
-            logger.info("Decoding video...")
-        decode_start = time.time()
-        video = self._decode_latents(latents)
-        if _is_rank_zero():
-            logger.info("Video decoded in %.2fs", time.time() - decode_start)
+        video = None
+        should_decode_video = robolab_inputs is None
+        if should_decode_video:
+            if _is_rank_zero():
+                logger.info("Decoding video...")
+            decode_start = time.time()
+            video = self._decode_latents(latents)
+            if _is_rank_zero():
+                logger.info("Video decoded in %.2fs", time.time() - decode_start)
+                logger.info("Total pipeline time: %.2fs", time.time() - pipeline_start)
+        elif _is_rank_zero():
             logger.info("Total pipeline time: %.2fs", time.time() - pipeline_start)
 
         if sound_enabled:
@@ -2565,11 +2627,10 @@ def _run_diffusion(start_latents):
                 "action_mode": action_mode,
                 "domain_id": domain_id,
             }
-            output = {"video": video}
             if robolab_inputs is not None:
-                output["actions"] = self._postprocess_robolab_action(action, robolab_inputs)
+                custom_action_output["actions"] = self._postprocess_robolab_action(action, robolab_inputs)
             return DiffusionOutput(
-                output=output,
+                output={} if robolab_inputs is not None else {"video": video},
                 custom_output=custom_action_output,
             )
 
diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py
index 515e2b258cc..58c5010a01e 100644
--- a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py
+++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py
@@ -39,6 +39,22 @@
 logger = init_logger(__name__)
 
 
+class Cosmos3RMSNorm(nn.Module):
+    """RMSNorm matching Cosmos Framework's Qwen3VLTextRMSNorm."""
+
+    def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+
 def _get_ulysses_state() -> tuple[int, int, dist.ProcessGroup | None]:
     """Return (ulysses_size, ulysses_rank, ulysses_pg) from vllm-omni parallel state.
 
@@ -511,8 +527,8 @@ def __init__(
             prefix=f"{prefix}.to_out",
         )
 
-        self.norm_q = RMSNorm(self.head_dim, eps=rms_norm_eps)
-        self.norm_k = RMSNorm(self.head_dim, eps=rms_norm_eps)
+        self.norm_q = Cosmos3RMSNorm(self.head_dim, eps=rms_norm_eps)
+        self.norm_k = Cosmos3RMSNorm(self.head_dim, eps=rms_norm_eps)
 
         # skip_sequence_parallel=True because the UND pathway is
         # computed once and replicated across SP ranks.
@@ -539,8 +555,8 @@ def forward(
         v = self.to_v(hidden_states).view(B, S, self.num_kv_heads_local, self.head_dim)
 
         # Per-head QK norm
-        q = F.rms_norm(q, (q.shape[-1],), self.norm_q.weight, self.norm_q.variance_epsilon)
-        k = F.rms_norm(k, (k.shape[-1],), self.norm_k.weight, self.norm_k.variance_epsilon)
+        q = self.norm_q(q)
+        k = self.norm_k(k)
 
         # Qwen3-style RoPE
         q, k = _apply_rotary_pos_emb(q, k, freqs_cos, freqs_sin)
@@ -619,8 +635,8 @@ def __init__(
             prefix=f"{prefix}.to_out",
         )
 
-        self.norm_q = RMSNorm(self.head_dim, eps=rms_norm_eps)
-        self.norm_k = RMSNorm(self.head_dim, eps=rms_norm_eps)
+        self.norm_q = Cosmos3RMSNorm(self.head_dim, eps=rms_norm_eps)
+        self.norm_k = Cosmos3RMSNorm(self.head_dim, eps=rms_norm_eps)
 
         self.attn = FrameworkAttention(
             num_heads=self.num_heads,
@@ -704,8 +720,8 @@ def forward(
         v = self.to_v(hidden_states).view(B, S_gen, self.num_kv_heads_local, self.head_dim)
 
         # Per-head QK norm
-        q = F.rms_norm(q, (q.shape[-1],), self.norm_q.weight, self.norm_q.variance_epsilon)
-        k = F.rms_norm(k, (k.shape[-1],), self.norm_k.weight, self.norm_k.variance_epsilon)
+        q = self.norm_q(q)
+        k = self.norm_k(k)
 
         # Qwen3-style RoPE
         q, k = _apply_rotary_pos_emb(q, k, freqs_cos, freqs_sin)
@@ -746,8 +762,8 @@ def __init__(
             quant_config=quant_config,
             prefix=f"{prefix}.self_attn",
         )
-        self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
-        self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
+        self.input_layernorm = Cosmos3RMSNorm(hidden_size, eps=rms_norm_eps)
+        self.post_attention_layernorm = Cosmos3RMSNorm(hidden_size, eps=rms_norm_eps)
         self.mlp = Cosmos3GatedMLP(
             hidden_size=hidden_size,
             intermediate_size=intermediate_size,
@@ -802,8 +818,8 @@ def __init__(
             quant_config=quant_config,
             prefix=f"{prefix}.cross_attention",
         )
-        self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
-        self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
+        self.input_layernorm = Cosmos3RMSNorm(hidden_size, eps=rms_norm_eps)
+        self.post_attention_layernorm = Cosmos3RMSNorm(hidden_size, eps=rms_norm_eps)
         self.mlp = Cosmos3GatedMLP(
             hidden_size=hidden_size,
             intermediate_size=intermediate_size,
@@ -896,7 +912,7 @@ def __init__(
             ]
         )
         # TODO: Not used right now, will be used in the future for prompt upsampler.
-        self.norm = RMSNorm(hidden_size, eps=rms_norm_eps)
+        self.norm = Cosmos3RMSNorm(hidden_size, eps=rms_norm_eps)
 
     def forward(
         self,
@@ -1111,7 +1127,7 @@ def __init__(
             ]
         )
 
-        self.norm_moe_gen = RMSNorm(self.hidden_size, eps=self.rms_norm_eps)
+        self.norm_moe_gen = Cosmos3RMSNorm(self.hidden_size, eps=self.rms_norm_eps)
         self.gen_sp_prepare = Cosmos3GenSPPrepare()
         self.gen_sp_gather = nn.Identity()
 
@@ -1446,7 +1462,7 @@ def forward(
                         "Cosmos3 action_noisy_mask must have shape [B, T_action, 1], "
                         f"got {tuple(action_noisy_mask.shape)}."
                     )
-                hidden_action = hidden_action + time_embed.unsqueeze(1) * action_noisy_mask.to(hidden_action.dtype)
+                hidden_action = hidden_action + time_embed.unsqueeze(1) * action_noisy_mask
 
         if hidden_sound is not None:
             hidden_sound = hidden_sound + time_embed.unsqueeze(1)

From 9513e96a3c004451882b8c99c4e893a746157b28 Mon Sep 17 00:00:00 2001
From: yuzhud <yuzhudong@gmail.com>
Date: Mon, 8 Jun 2026 10:15:01 -0700
Subject: [PATCH 03/12] Address OpenPI review comments

Signed-off-by: yuzhud <yuzhudong@gmail.com>
---
 .../openai_api/test_openpi_connection.py      | 73 ++++++++-----------
 .../openai_api/test_openpi_serving.py         | 13 +++-
 .../models/cosmos3/transformer_cosmos3.py     | 36 +++------
 vllm_omni/entrypoints/openai/api_server.py    |  1 -
 vllm_omni/entrypoints/openpi/connection.py    | 62 ++++++++--------
 vllm_omni/entrypoints/openpi/serving.py       | 11 +--
 6 files changed, 87 insertions(+), 109 deletions(-)

diff --git a/tests/entrypoints/openai_api/test_openpi_connection.py b/tests/entrypoints/openai_api/test_openpi_connection.py
index ba5a45c3454..30c99125e68 100644
--- a/tests/entrypoints/openai_api/test_openpi_connection.py
+++ b/tests/entrypoints/openai_api/test_openpi_connection.py
@@ -1,9 +1,7 @@
 import asyncio
-import builtins
-import sys
-import types
 from unittest.mock import AsyncMock, MagicMock
 
+import numpy as np
 import pytest
 
 from vllm_omni.entrypoints.openpi import connection as openpi_connection
@@ -52,48 +50,39 @@ def _serving_mock():
     return serving
 
 
-def test_pack_reports_clear_error_when_openpi_client_is_missing(monkeypatch):
-    real_import = builtins.__import__
-
-    def import_without_openpi_client(name, globals=None, locals=None, fromlist=(), level=0):
-        if name == "openpi_client":
-            raise ModuleNotFoundError("No module named 'openpi_client'", name="openpi_client")
-        return real_import(name, globals, locals, fromlist, level)
-
-    monkeypatch.setattr(builtins, "__import__", import_without_openpi_client)
-
-    with pytest.raises(ImportError) as exc_info:
-        openpi_connection._pack({"prompt": "pick up the object"})
-
-    message = str(exc_info.value)
-    assert "/v1/realtime/robot/openpi" in message
-    assert "pip install openpi-client" in message
-
-
-def test_pack_and_unpack_delegate_to_openpi_msgpack_numpy(monkeypatch):
-    calls = []
-
-    class FakeMsgpackNumpy:
-        @staticmethod
-        def packb(obj):
-            calls.append(("packb", obj))
-            return b"packed"
+def test_pack_and_unpack_round_trip_numpy_values():
+    payload = {
+        "image": np.arange(6, dtype=np.uint8).reshape(2, 3),
+        "action": np.asarray([[1.0, 2.0]], dtype=np.float32),
+        "scalar": np.float32(3.5),
+        "nested": [{"done": np.bool_(True)}],
+    }
 
-        @staticmethod
-        def unpackb(data):
-            calls.append(("unpackb", data))
-            return {"unpacked": data}
+    decoded = openpi_connection._unpack(openpi_connection._pack(payload))
+
+    np.testing.assert_array_equal(decoded["image"], payload["image"])
+    np.testing.assert_allclose(decoded["action"], payload["action"])
+    assert decoded["image"].dtype == np.uint8
+    assert decoded["action"].dtype == np.float32
+    assert decoded["scalar"] == np.float32(3.5)
+    assert decoded["nested"][0]["done"] == np.bool_(True)
+
+
+def test_unpack_accepts_msgpack_numpy_marker_dicts():
+    action = np.asarray([[1.0, 2.0]], dtype=np.float32)
+    payload = {
+        b"actions": {
+            b"nd": True,
+            b"type": action.dtype.str,
+            b"kind": action.dtype.kind,
+            b"shape": action.shape,
+            b"data": action.tobytes(),
+        }
+    }
 
-    fake_openpi_client = types.ModuleType("openpi_client")
-    fake_openpi_client.msgpack_numpy = FakeMsgpackNumpy
-    monkeypatch.setitem(sys.modules, "openpi_client", fake_openpi_client)
+    decoded = openpi_connection._unpack_numpy(payload)
 
-    assert openpi_connection._pack({"x": 1}) == b"packed"
-    assert openpi_connection._unpack(b"payload") == {"unpacked": b"payload"}
-    assert calls == [
-        ("packb", {"x": 1}),
-        ("unpackb", b"payload"),
-    ]
+    np.testing.assert_allclose(decoded[b"actions"], action)
 
 
 def test_handle_connection_returns_structured_error_for_invalid_payload(monkeypatch):
diff --git a/tests/entrypoints/openai_api/test_openpi_serving.py b/tests/entrypoints/openai_api/test_openpi_serving.py
index 9eb1d5bfe0b..cdd833db083 100644
--- a/tests/entrypoints/openai_api/test_openpi_serving.py
+++ b/tests/entrypoints/openai_api/test_openpi_serving.py
@@ -40,7 +40,9 @@ def _json_unpack(data):
 
 
 def _engine_with_policy_config(policy_config=None):
-    od_config = SimpleNamespace(model_config={"policy_server_config": policy_config or TEST_POLICY_SERVER_CONFIG})
+    if policy_config is None:
+        policy_config = TEST_POLICY_SERVER_CONFIG
+    od_config = SimpleNamespace(model_config={"policy_server_config": policy_config})
     return SimpleNamespace(get_diffusion_od_config=lambda: od_config)
 
 
@@ -172,6 +174,15 @@ def test_create_policy_server_returns_none_without_policy_config():
     assert serving is None
 
 
+def test_policy_server_config_allows_explicit_empty_config():
+    serving = openpi_serving.ServingRealtimeRobotOpenPI(
+        engine_client=_engine_with_policy_config(policy_config={}),
+        model_name="nvidia/Cosmos3-Nano-Policy-DROID",
+    )
+
+    assert serving.policy_server_config.to_dict() == {}
+
+
 def test_policy_server_config_reads_engine_model_config():
     policy_config = {"custom_model_key": "custom-value"}
     engine_client = SimpleNamespace(model_config=SimpleNamespace(policy_server_config=policy_config))
diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py
index 58c5010a01e..a333dc6dd70 100644
--- a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py
+++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py
@@ -39,22 +39,6 @@
 logger = init_logger(__name__)
 
 
-class Cosmos3RMSNorm(nn.Module):
-    """RMSNorm matching Cosmos Framework's Qwen3VLTextRMSNorm."""
-
-    def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
-        super().__init__()
-        self.weight = nn.Parameter(torch.ones(hidden_size))
-        self.variance_epsilon = eps
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        input_dtype = hidden_states.dtype
-        hidden_states = hidden_states.to(torch.float32)
-        variance = hidden_states.pow(2).mean(-1, keepdim=True)
-        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
-        return self.weight * hidden_states.to(input_dtype)
-
-
 def _get_ulysses_state() -> tuple[int, int, dist.ProcessGroup | None]:
     """Return (ulysses_size, ulysses_rank, ulysses_pg) from vllm-omni parallel state.
 
@@ -527,8 +511,8 @@ def __init__(
             prefix=f"{prefix}.to_out",
         )
 
-        self.norm_q = Cosmos3RMSNorm(self.head_dim, eps=rms_norm_eps)
-        self.norm_k = Cosmos3RMSNorm(self.head_dim, eps=rms_norm_eps)
+        self.norm_q = RMSNorm(self.head_dim, eps=rms_norm_eps)
+        self.norm_k = RMSNorm(self.head_dim, eps=rms_norm_eps)
 
         # skip_sequence_parallel=True because the UND pathway is
         # computed once and replicated across SP ranks.
@@ -635,8 +619,8 @@ def __init__(
             prefix=f"{prefix}.to_out",
         )
 
-        self.norm_q = Cosmos3RMSNorm(self.head_dim, eps=rms_norm_eps)
-        self.norm_k = Cosmos3RMSNorm(self.head_dim, eps=rms_norm_eps)
+        self.norm_q = RMSNorm(self.head_dim, eps=rms_norm_eps)
+        self.norm_k = RMSNorm(self.head_dim, eps=rms_norm_eps)
 
         self.attn = FrameworkAttention(
             num_heads=self.num_heads,
@@ -762,8 +746,8 @@ def __init__(
             quant_config=quant_config,
             prefix=f"{prefix}.self_attn",
         )
-        self.input_layernorm = Cosmos3RMSNorm(hidden_size, eps=rms_norm_eps)
-        self.post_attention_layernorm = Cosmos3RMSNorm(hidden_size, eps=rms_norm_eps)
+        self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
         self.mlp = Cosmos3GatedMLP(
             hidden_size=hidden_size,
             intermediate_size=intermediate_size,
@@ -818,8 +802,8 @@ def __init__(
             quant_config=quant_config,
             prefix=f"{prefix}.cross_attention",
         )
-        self.input_layernorm = Cosmos3RMSNorm(hidden_size, eps=rms_norm_eps)
-        self.post_attention_layernorm = Cosmos3RMSNorm(hidden_size, eps=rms_norm_eps)
+        self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
         self.mlp = Cosmos3GatedMLP(
             hidden_size=hidden_size,
             intermediate_size=intermediate_size,
@@ -912,7 +896,7 @@ def __init__(
             ]
         )
         # TODO: Not used right now, will be used in the future for prompt upsampler.
-        self.norm = Cosmos3RMSNorm(hidden_size, eps=rms_norm_eps)
+        self.norm = RMSNorm(hidden_size, eps=rms_norm_eps)
 
     def forward(
         self,
@@ -1127,7 +1111,7 @@ def __init__(
             ]
         )
 
-        self.norm_moe_gen = Cosmos3RMSNorm(self.hidden_size, eps=self.rms_norm_eps)
+        self.norm_moe_gen = RMSNorm(self.hidden_size, eps=self.rms_norm_eps)
         self.gen_sp_prepare = Cosmos3GenSPPrepare()
         self.gen_sp_gather = nn.Identity()
 
diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py
index 6ff9dcc4818..c11db925f71 100755
--- a/vllm_omni/entrypoints/openai/api_server.py
+++ b/vllm_omni/entrypoints/openai/api_server.py
@@ -1531,7 +1531,6 @@ async def realtime_websocket(websocket: WebSocket):
     await connection.handle_connection()
 
 
-@router.websocket("/")
 @router.websocket("/v1/realtime/robot/openpi")
 async def realtime_robot_openpi(websocket: WebSocket):
     """WebSocket endpoint for robot policy inference via OpenPI messages."""
diff --git a/vllm_omni/entrypoints/openpi/connection.py b/vllm_omni/entrypoints/openpi/connection.py
index a47781037f7..defc82deaa0 100644
--- a/vllm_omni/entrypoints/openpi/connection.py
+++ b/vllm_omni/entrypoints/openpi/connection.py
@@ -14,7 +14,7 @@
 import asyncio
 from typing import Any
 
-import msgpack
+import msgspec
 import numpy as np
 from fastapi import WebSocket
 from starlette.websockets import WebSocketDisconnect
@@ -29,51 +29,55 @@
 MAX_OPENPI_PAYLOAD_BYTES = 64 * 1024 * 1024
 
 
-def _pack_array(obj: Any) -> Any:
+def _pack_numpy(obj: Any) -> Any:
     if isinstance(obj, (np.ndarray, np.generic)) and obj.dtype.kind in ("V", "O", "c"):
         raise ValueError(f"Unsupported dtype: {obj.dtype}")
     if isinstance(obj, np.ndarray):
+        if not obj.flags.c_contiguous:
+            obj = np.ascontiguousarray(obj)
         return {
-            b"__ndarray__": True,
+            b"nd": True,
             b"data": obj.tobytes(),
-            b"dtype": obj.dtype.str,
+            b"type": obj.dtype.str,
+            b"kind": obj.dtype.kind,
             b"shape": obj.shape,
         }
     if isinstance(obj, np.generic):
         return {
-            b"__npgeneric__": True,
-            b"data": obj.item(),
-            b"dtype": obj.dtype.str,
+            b"nd": False,
+            b"data": obj.tobytes(),
+            b"type": obj.dtype.str,
+            b"kind": obj.dtype.kind,
         }
-    return obj
-
-
-def _unpack_array(obj: dict[Any, Any]) -> Any:
-    if b"__ndarray__" in obj:
-        return np.ndarray(
-            buffer=obj[b"data"],
-            dtype=np.dtype(obj[b"dtype"]),
-            shape=tuple(obj[b"shape"]),
-        )
-    if "__ndarray__" in obj:
-        return np.ndarray(
-            buffer=obj["data"],
-            dtype=np.dtype(obj["dtype"]),
-            shape=tuple(obj["shape"]),
-        )
-    if b"__npgeneric__" in obj:
-        return np.dtype(obj[b"dtype"]).type(obj[b"data"])
-    if "__npgeneric__" in obj:
-        return np.dtype(obj["dtype"]).type(obj["data"])
+    raise TypeError(f"Unsupported type: {type(obj)!r}")
+
+
+def _mapping_get(obj: dict[Any, Any], key: str, default: Any = None) -> Any:
+    return obj.get(key, obj.get(key.encode(), default))
+
+
+def _unpack_numpy(obj: Any) -> Any:
+    if isinstance(obj, dict):
+        nd = _mapping_get(obj, "nd")
+        dtype = _mapping_get(obj, "type")
+        data = _mapping_get(obj, "data")
+        if nd is not None and dtype is not None and data is not None:
+            array = np.frombuffer(data, dtype=np.dtype(dtype))
+            if nd:
+                return array.reshape(tuple(_mapping_get(obj, "shape", ())))
+            return array[0]
+        return {key: _unpack_numpy(value) for key, value in obj.items()}
+    if isinstance(obj, list):
+        return [_unpack_numpy(value) for value in obj]
     return obj
 
 
 def _pack(obj: Any) -> bytes:
-    return msgpack.packb(obj, default=_pack_array)
+    return msgspec.msgpack.encode(obj, enc_hook=_pack_numpy)
 
 
 def _unpack(data: bytes) -> Any:
-    return msgpack.unpackb(data, object_hook=_unpack_array)
+    return _unpack_numpy(msgspec.msgpack.decode(data))
 
 
 class RobotRealtimeConnection:
diff --git a/vllm_omni/entrypoints/openpi/serving.py b/vllm_omni/entrypoints/openpi/serving.py
index c2cfe3a4429..ec5c88408ea 100644
--- a/vllm_omni/entrypoints/openpi/serving.py
+++ b/vllm_omni/entrypoints/openpi/serving.py
@@ -21,10 +21,6 @@
 logger = init_logger(__name__)
 
 ActionOutput = np.ndarray | dict[str, np.ndarray]
-COSMOS3_DROID_POLICY_MODELS = {
-    "Cosmos3-Nano-Policy-DROID",
-    "nvidia/Cosmos3-Nano-Policy-DROID",
-}
 
 
 def _to_builtin_container(value: Any) -> Any:
@@ -78,12 +74,7 @@ def __init__(
     ) -> None:
         self.engine_client = engine_client
         self.model_name = model_name
-        try:
-            self.policy_server_config = self._get_policy_server_config(engine_client)
-        except ValueError:
-            if model_name not in COSMOS3_DROID_POLICY_MODELS:
-                raise
-            self.policy_server_config = PolicyServerConfig({})
+        self.policy_server_config = self._get_policy_server_config(engine_client)
         self._request_counter = count()
 
     @classmethod

From 8c61bba2d2b23244dadc8c190c401ae0cb614d73 Mon Sep 17 00:00:00 2001
From: yuzhud <yuzhudong@gmail.com>
Date: Tue, 9 Jun 2026 13:12:17 -0700
Subject: [PATCH 04/12] cosmos rms

Signed-off-by: yuzhud <yuzhudong@gmail.com>
---
 .../models/cosmos3/pipeline_cosmos3.py        | 26 ++++++++++++--
 .../models/cosmos3/transformer_cosmos3.py     | 36 +++++++++++++------
 vllm_omni/entrypoints/openai/serving_video.py | 31 ++++++++++------
 3 files changed, 69 insertions(+), 24 deletions(-)

diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py
index a948c545aa9..35b6502083d 100644
--- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py
+++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py
@@ -1062,6 +1062,7 @@ def _next_robolab_seed(self, extra: dict[str, Any]) -> int:
     def _build_robolab_policy_inputs(
         self,
         sp: OmniDiffusionSamplingParams,
+        prompt_data: Any | None = None,
     ) -> dict[str, Any] | None:
         extra = sp.extra_args if isinstance(sp.extra_args, dict) else {}
         obs = extra.get("robot_obs")
@@ -1107,7 +1108,12 @@ def extra_param_alias(primary_key: str, alias_key: str, default: Any) -> Any:
         if raw_action_dim <= 0:
             raise ValueError(f"RoboLab raw_action_dim must be positive, got {raw_action_dim}.")
 
-        image = _extract_robolab_image(obs)
+        try:
+            image = _extract_robolab_image(obs)
+        except ValueError as exc:
+            image = self._extract_robolab_prompt_image(prompt_data)
+            if image is None:
+                raise exc
         if image.shape[:2] != (image_h, image_w):
             image = _resize_rgb_uint8(image, (image_h, image_w))
 
@@ -1196,6 +1202,18 @@ def extra_param_alias(primary_key: str, alias_key: str, default: Any) -> Any:
             "observation": obs,
         }
 
+    @staticmethod
+    def _extract_robolab_prompt_image(prompt_data: Any | None) -> np.ndarray | None:
+        if not isinstance(prompt_data, dict):
+            return None
+        multi_modal_data = prompt_data.get("multi_modal_data", {}) or {}
+        image = multi_modal_data.get("image")
+        if image is None:
+            return None
+        if isinstance(image, PIL.Image.Image):
+            return np.asarray(image.convert("RGB"))
+        return _ensure_rgb_uint8_image(image, "multi_modal_data.image")
+
     @staticmethod
     def _build_action_condition_mask_from_indexes(
         indexes: list[int],
@@ -2227,11 +2245,12 @@ def forward(
 
         sp = req.sampling_params
         prompt_data = req.prompts[0]
-        robolab_inputs = self._build_robolab_policy_inputs(sp)
+        robolab_inputs = self._build_robolab_policy_inputs(sp, prompt_data)
         if robolab_inputs is not None:
             prompt = robolab_inputs["prompt"]
             negative_prompt = None
             image_tensor = None
+            video_tensor = None
             action_video_tensor = robolab_inputs["video_tensor"]
         elif isinstance(prompt_data, str):
             prompt = prompt_data
@@ -2249,7 +2268,8 @@ def forward(
         sound_enabled = False if robolab_inputs is not None else self._is_sound_request(prompt_data, sp)
         action_mode = ACTION_MODE_POLICY if robolab_inputs is not None else self._get_action_mode(prompt_data, sp)
         action_enabled = action_mode is not None
-        action_video_tensor = video_tensor if action_enabled else None
+        if robolab_inputs is None:
+            action_video_tensor = video_tensor if action_enabled else None
         if action_enabled and is_t2i:
             raise ValueError("Cosmos3 action generation is supported only for video outputs.")
         if action_enabled and sound_enabled:
diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py
index a333dc6dd70..58c5010a01e 100644
--- a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py
+++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py
@@ -39,6 +39,22 @@
 logger = init_logger(__name__)
 
 
+class Cosmos3RMSNorm(nn.Module):
+    """RMSNorm matching Cosmos Framework's Qwen3VLTextRMSNorm."""
+
+    def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+
 def _get_ulysses_state() -> tuple[int, int, dist.ProcessGroup | None]:
     """Return (ulysses_size, ulysses_rank, ulysses_pg) from vllm-omni parallel state.
 
@@ -511,8 +527,8 @@ def __init__(
             prefix=f"{prefix}.to_out",
         )
 
-        self.norm_q = RMSNorm(self.head_dim, eps=rms_norm_eps)
-        self.norm_k = RMSNorm(self.head_dim, eps=rms_norm_eps)
+        self.norm_q = Cosmos3RMSNorm(self.head_dim, eps=rms_norm_eps)
+        self.norm_k = Cosmos3RMSNorm(self.head_dim, eps=rms_norm_eps)
 
         # skip_sequence_parallel=True because the UND pathway is
         # computed once and replicated across SP ranks.
@@ -619,8 +635,8 @@ def __init__(
             prefix=f"{prefix}.to_out",
         )
 
-        self.norm_q = RMSNorm(self.head_dim, eps=rms_norm_eps)
-        self.norm_k = RMSNorm(self.head_dim, eps=rms_norm_eps)
+        self.norm_q = Cosmos3RMSNorm(self.head_dim, eps=rms_norm_eps)
+        self.norm_k = Cosmos3RMSNorm(self.head_dim, eps=rms_norm_eps)
 
         self.attn = FrameworkAttention(
             num_heads=self.num_heads,
@@ -746,8 +762,8 @@ def __init__(
             quant_config=quant_config,
             prefix=f"{prefix}.self_attn",
         )
-        self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
-        self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
+        self.input_layernorm = Cosmos3RMSNorm(hidden_size, eps=rms_norm_eps)
+        self.post_attention_layernorm = Cosmos3RMSNorm(hidden_size, eps=rms_norm_eps)
         self.mlp = Cosmos3GatedMLP(
             hidden_size=hidden_size,
             intermediate_size=intermediate_size,
@@ -802,8 +818,8 @@ def __init__(
             quant_config=quant_config,
             prefix=f"{prefix}.cross_attention",
         )
-        self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
-        self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
+        self.input_layernorm = Cosmos3RMSNorm(hidden_size, eps=rms_norm_eps)
+        self.post_attention_layernorm = Cosmos3RMSNorm(hidden_size, eps=rms_norm_eps)
         self.mlp = Cosmos3GatedMLP(
             hidden_size=hidden_size,
             intermediate_size=intermediate_size,
@@ -896,7 +912,7 @@ def __init__(
             ]
         )
         # TODO: Not used right now, will be used in the future for prompt upsampler.
-        self.norm = RMSNorm(hidden_size, eps=rms_norm_eps)
+        self.norm = Cosmos3RMSNorm(hidden_size, eps=rms_norm_eps)
 
     def forward(
         self,
@@ -1111,7 +1127,7 @@ def __init__(
             ]
         )
 
-        self.norm_moe_gen = RMSNorm(self.hidden_size, eps=self.rms_norm_eps)
+        self.norm_moe_gen = Cosmos3RMSNorm(self.hidden_size, eps=self.rms_norm_eps)
         self.gen_sp_prepare = Cosmos3GenSPPrepare()
         self.gen_sp_gather = nn.Identity()
 
diff --git a/vllm_omni/entrypoints/openai/serving_video.py b/vllm_omni/entrypoints/openai/serving_video.py
index ac6038f9285..712f6647ea0 100644
--- a/vllm_omni/entrypoints/openai/serving_video.py
+++ b/vllm_omni/entrypoints/openai/serving_video.py
@@ -195,18 +195,21 @@ async def _run_and_extract(
             # Merge extra_params into extra_args
             gen_params.extra_args.update(request.extra_params)
 
-            # Redact the inline ``action`` array (hundreds of floats) when
-            # logging so it doesn't flood the logs; everything else is logged
-            # verbatim.
+            # Redact inline arrays when logging so RoboLab policy requests do
+            # not flood the server log with image/state payloads.
             loggable = request.extra_params
-            action_val = loggable.get("action")
-            if action_val is not None:
-                summary = (
-                    f"<{type(action_val).__name__} len={len(action_val)}>"
-                    if hasattr(action_val, "__len__")
-                    else f"<{type(action_val).__name__}>"
+            redacted = {}
+            for key in ("action", "robot_obs", "observation"):
+                value = loggable.get(key)
+                if value is None:
+                    continue
+                redacted[key] = (
+                    f"<{type(value).__name__} len={len(value)}>"
+                    if hasattr(value, "__len__")
+                    else f"<{type(value).__name__}>"
                 )
-                loggable = {**loggable, "action": summary}
+            if redacted:
+                loggable = {**loggable, **redacted}
             logger.info("Applied extra_params: %s", loggable)
 
         self._apply_lora(request.lora, gen_params)
@@ -314,6 +317,11 @@ async def generate_video_bytes(
             if "video_codec_options" in request.extra_params:
                 video_codec_options = request.extra_params["video_codec_options"]
 
+        action = artifacts.actions[0]
+        if action is not None and isinstance(artifacts.videos[0], dict):
+            logger.info("Action-only video request %s completed; skipping MP4 encoding.", reference_id)
+            return b"", artifacts.stage_durations, artifacts.peak_memory_mb, action
+
         _t_encode_start = time.perf_counter()
         video_bytes = _encode_video_bytes(
             artifacts.videos[0],
@@ -510,7 +518,8 @@ def _extract_action_outputs(cls, result: Any, expected_count: int) -> list[Video
         if not custom_output or "action" not in custom_output:
             return [None] * expected_count
 
-        action_items = cls._split_action_payload(custom_output["action"], expected_count)
+        action_payload = custom_output.get("actions", custom_output["action"])
+        action_items = cls._split_action_payload(action_payload, expected_count)
         return [
             cls._make_video_action(action_item, custom_output) if action_item is not None else None
             for action_item in action_items

From f350c45e6efe5b0bee78346a5151c0e56e5287d4 Mon Sep 17 00:00:00 2001
From: yuzhud <yuzhudong@gmail.com>
Date: Tue, 9 Jun 2026 21:02:11 -0700
Subject: [PATCH 05/12] Address Cosmos3 OpenPI review comments

Signed-off-by: yuzhud <yuzhudong@gmail.com>
---
 recipes/cosmos3/Cosmos3-Nano.md               |  44 +++++++
 .../models/cosmos3/test_cosmos3_pipeline.py   |   1 +
 tests/diffusion/test_diffusion_scheduler.py   |  65 ++++++++++
 .../openai_api/test_openpi_connection.py      |  17 +++
 vllm_omni/diffusion/diffusion_engine.py       |   6 +-
 .../models/cosmos3/pipeline_cosmos3.py        | 119 +++++++++++-------
 vllm_omni/entrypoints/openpi/connection.py    |  34 ++++-
 7 files changed, 230 insertions(+), 56 deletions(-)

diff --git a/recipes/cosmos3/Cosmos3-Nano.md b/recipes/cosmos3/Cosmos3-Nano.md
index 2ee9c2ed7d3..0d51a1a3a5a 100644
--- a/recipes/cosmos3/Cosmos3-Nano.md
+++ b/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
diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py
index c51f3ff0805..3eebda1b639 100644
--- a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py
+++ b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py
@@ -762,6 +762,7 @@ 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
 
     @pytest.mark.parametrize(
         ("prompt", "sampling_params", "message"),
diff --git a/tests/diffusion/test_diffusion_scheduler.py b/tests/diffusion/test_diffusion_scheduler.py
index 76234b7fa42..1b5e813ee32 100644
--- a/tests/diffusion/test_diffusion_scheduler.py
+++ b/tests/diffusion/test_diffusion_scheduler.py
@@ -567,6 +567,71 @@ async def test_step_multi_request_reuses_multimodal_slice_logic(self, mocker: Mo
             torch.tensor([3.0, 4.0]),
         )
 
+    @pytest.mark.asyncio
+    async def test_step_empty_dict_output_still_runs_postprocess(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 = mocker.Mock(return_value={"video": ["processed"]})
+        engine._post_process_accepts_sampling_params = False
+        engine._check_and_start_background_loop = mocker.AsyncMock()
+        engine.async_add_req_and_wait_for_response = mocker.AsyncMock(
+            return_value=DiffusionOutput(
+                output={},
+                custom_output={"actions": torch.tensor([[1.0, 2.0]])},
+            )
+        )
+
+        request = OmniDiffusionRequest(
+            prompts=["prompt"],
+            sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+            request_id="req-action",
+        )
+
+        mocker.patch("vllm_omni.diffusion.diffusion_engine.supports_audio_output", return_value=False)
+        outputs = await engine.step(request)
+
+        engine.post_process_func.assert_called_once_with({})
+        assert outputs[0].images == ["processed"]
+        torch.testing.assert_close(outputs[0].multimodal_output["actions"], torch.tensor([[1.0, 2.0]]))
+
+    @pytest.mark.asyncio
+    async def test_step_action_only_flag_skips_postprocess(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 = mocker.Mock(side_effect=AssertionError("postprocess should be skipped"))
+        engine._post_process_accepts_sampling_params = False
+        engine._check_and_start_background_loop = mocker.AsyncMock()
+        engine.async_add_req_and_wait_for_response = mocker.AsyncMock(
+            return_value=DiffusionOutput(
+                output={},
+                custom_output={
+                    "actions": torch.tensor([[1.0, 2.0]]),
+                    "action_only_output": True,
+                },
+            )
+        )
+
+        request = OmniDiffusionRequest(
+            prompts=["prompt"],
+            sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+            request_id="req-action",
+        )
+
+        mocker.patch("vllm_omni.diffusion.diffusion_engine.supports_audio_output", return_value=False)
+        outputs = await engine.step(request)
+
+        engine.post_process_func.assert_not_called()
+        assert outputs[0].images == []
+        torch.testing.assert_close(outputs[0].multimodal_output["actions"], torch.tensor([[1.0, 2.0]]))
+
 
 class TestStepScheduler:
     def setup_method(self) -> None:
diff --git a/tests/entrypoints/openai_api/test_openpi_connection.py b/tests/entrypoints/openai_api/test_openpi_connection.py
index 30c99125e68..83d17a28cd3 100644
--- a/tests/entrypoints/openai_api/test_openpi_connection.py
+++ b/tests/entrypoints/openai_api/test_openpi_connection.py
@@ -83,6 +83,23 @@ def test_unpack_accepts_msgpack_numpy_marker_dicts():
     decoded = openpi_connection._unpack_numpy(payload)
 
     np.testing.assert_allclose(decoded[b"actions"], action)
+    assert decoded[b"actions"].flags.writeable is True
+    decoded[b"actions"][:, -1] = 0.0
+    np.testing.assert_allclose(decoded[b"actions"], np.asarray([[1.0, 0.0]], dtype=np.float32))
+
+
+def test_unpack_leaves_user_dict_without_numpy_kind_marker_unchanged():
+    payload = {
+        "metadata": {
+            "nd": True,
+            "type": "<f4",
+            "data": b"user payload",
+        }
+    }
+
+    decoded = openpi_connection._unpack_numpy(payload)
+
+    assert decoded == payload
 
 
 def test_handle_connection_returns_structured_error_for_invalid_payload(monkeypatch):
diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py
index 81eb14a3c8a..344847ed334 100644
--- a/vllm_omni/diffusion/diffusion_engine.py
+++ b/vllm_omni/diffusion/diffusion_engine.py
@@ -263,11 +263,11 @@ async def step(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]:
 
         custom_output = output.custom_output or {}
         action_payload = custom_output.get("actions")
-        action_only_output = action_payload is not None and isinstance(output_data, dict) and not output_data
+        action_only_output = bool(custom_output.get("action_only_output"))
 
         postprocess_start_time = time.perf_counter()
         if action_only_output:
-            outputs = {}
+            outputs = []
         elif self.post_process_func is not None:
             # Some video pipelines need request-level controls during
             # postprocess (for example worker-side frame interpolation).
@@ -287,8 +287,6 @@ async def step(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]:
             model_audio_sample_rate = outputs.get("audio_sample_rate")
             model_fps = outputs.get("fps")
             outputs = outputs.get("video", outputs)
-        if action_payload is None:
-            action_payload = custom_output.get("actions")
         postprocess_time = time.perf_counter() - postprocess_start_time
         logger.debug("Post-processing completed in %.4f seconds", postprocess_time)
 
diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py
index 35b6502083d..4195abd24b2 100644
--- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py
+++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py
@@ -21,6 +21,7 @@
 import os
 import time
 from collections.abc import Iterable
+from dataclasses import dataclass
 from typing import Any, ClassVar
 
 import numpy as np
@@ -114,6 +115,7 @@
 ROBOLAB_DEFAULT_FLOW_SHIFT = 5.0
 ROBOLAB_DEFAULT_SEED = 0
 ROBOLAB_DEFAULT_ACTION_SPACE = "joint_pos"
+ROBOLAB_MIDTRAIN_POSE_ACTION_DIM = 9  # xyz position + rot6d orientation
 ROBOLAB_CONCAT_VIEW_DESCRIPTION = (
     "The top row is from the wrist-mounted camera. "
     "The bottom row contains two horizontally concatenated third-person perspective views of the scene from opposite "
@@ -126,6 +128,29 @@
 COSMOS3_DEFAULT_MAX_SEQUENCE_LENGTH = 4096
 
 
+@dataclass(frozen=True)
+class RoboLabPolicyInputs:
+    prompt: str
+    video_tensor: torch.Tensor
+    action_tensor: torch.Tensor
+    action_condition_indexes: list[int]
+    action_start_frame_offset: int
+    raw_action_dim: int
+    domain_id: int
+    fps: float
+    height: int
+    width: int
+    image_size: Any
+    num_frames: int
+    num_inference_steps: int
+    guidance_scale: float
+    flow_shift: float
+    seed: int
+    history_length: int
+    action_space: str
+    observation: dict[str, Any]
+
+
 def _normalize_condition_frame_indexes_vision(value: Any) -> tuple[int, ...]:
     """Normalize Cosmos3 vision-conditioning latent frame indexes."""
     if value is None:
@@ -1063,7 +1088,7 @@ def _build_robolab_policy_inputs(
         self,
         sp: OmniDiffusionSamplingParams,
         prompt_data: Any | None = None,
-    ) -> dict[str, Any] | None:
+    ) -> RoboLabPolicyInputs | None:
         extra = sp.extra_args if isinstance(sp.extra_args, dict) else {}
         obs = extra.get("robot_obs")
         if obs is None:
@@ -1101,7 +1126,7 @@ def extra_param_alias(primary_key: str, alias_key: str, default: Any) -> Any:
         domain_name = str(extra_param("domain_name", ROBOLAB_DEFAULT_DOMAIN_NAME))
         domain_id = resolve_domain_id(domain_name=domain_name, require_explicit=True)
 
-        if history_length < (1 if use_state else 0):
+        if use_state and history_length < 1:
             raise ValueError("RoboLab history_length must be >= 1 when use_state is true.")
         if action_chunk_size <= 0:
             raise ValueError(f"RoboLab action_chunk_size must be positive, got {action_chunk_size}.")
@@ -1160,6 +1185,7 @@ def extra_param_alias(primary_key: str, alias_key: str, default: Any) -> Any:
             "ai_caption": prompt,
             "video": video,
             "action": action,
+            # Cosmos Framework consumes this as an integer conditioning bucket.
             "conditioning_fps": torch.tensor(fps, dtype=torch.long),
             "mode": ACTION_MODE_POLICY,
             "domain_id": torch.tensor(domain_id, dtype=torch.long),
@@ -1178,29 +1204,29 @@ def extra_param_alias(primary_key: str, alias_key: str, default: Any) -> Any:
         else:
             transformed_raw_action_dim = raw_action_dim
 
-        return {
-            "prompt": sample["ai_caption"],
-            "video_tensor": video_tensor.unsqueeze(0),
-            "action_tensor": sample["action"].float(),
-            "action_condition_indexes": list(getattr(sequence_plan, "condition_frame_indexes_action", []) or []),
-            "action_start_frame_offset": int(getattr(sequence_plan, "action_start_frame_offset", 1)),
-            "raw_action_dim": transformed_raw_action_dim,
-            "domain_id": domain_id,
-            "fps": fps,
-            "height": int(video_tensor.shape[-2]),
-            "width": int(video_tensor.shape[-1]),
-            "image_size": sample.get("image_size"),
-            "num_frames": int(video_tensor.shape[1]),
-            "num_inference_steps": int(
+        return RoboLabPolicyInputs(
+            prompt=sample["ai_caption"],
+            video_tensor=video_tensor.unsqueeze(0),
+            action_tensor=sample["action"].float(),
+            action_condition_indexes=list(getattr(sequence_plan, "condition_frame_indexes_action", []) or []),
+            action_start_frame_offset=int(getattr(sequence_plan, "action_start_frame_offset", 1)),
+            raw_action_dim=transformed_raw_action_dim,
+            domain_id=domain_id,
+            fps=fps,
+            height=int(video_tensor.shape[-2]),
+            width=int(video_tensor.shape[-1]),
+            image_size=sample.get("image_size"),
+            num_frames=int(video_tensor.shape[1]),
+            num_inference_steps=int(
                 extra_param_alias("num_inference_steps", "num_steps", ROBOLAB_DEFAULT_NUM_INFERENCE_STEPS)
             ),
-            "guidance_scale": float(extra_param_alias("guidance_scale", "guidance", ROBOLAB_DEFAULT_GUIDANCE_SCALE)),
-            "flow_shift": float(extra_param_alias("flow_shift", "shift", ROBOLAB_DEFAULT_FLOW_SHIFT)),
-            "seed": self._next_robolab_seed(extra),
-            "history_length": history_length,
-            "action_space": action_space,
-            "observation": obs,
-        }
+            guidance_scale=float(extra_param_alias("guidance_scale", "guidance", ROBOLAB_DEFAULT_GUIDANCE_SCALE)),
+            flow_shift=float(extra_param_alias("flow_shift", "shift", ROBOLAB_DEFAULT_FLOW_SHIFT)),
+            seed=self._next_robolab_seed(extra),
+            history_length=history_length,
+            action_space=action_space,
+            observation=obs,
+        )
 
     @staticmethod
     def _extract_robolab_prompt_image(prompt_data: Any | None) -> np.ndarray | None:
@@ -1229,28 +1255,28 @@ def _build_action_condition_mask_from_indexes(
             mask[:, idx, :] = 1.0
         return mask
 
-    def _postprocess_robolab_action(self, action: torch.Tensor, inputs: dict[str, Any]) -> np.ndarray:
+    def _postprocess_robolab_action(self, action: torch.Tensor, inputs: RoboLabPolicyInputs) -> np.ndarray:
         action_np = action[0].float().cpu().numpy()
-        history_length = int(inputs["history_length"])
+        history_length = int(inputs.history_length)
         action_np = action_np[history_length:]
         action_np[:, -1] = 1.0 - action_np[:, -1]
 
-        if inputs["action_space"] == "midtrain":
-            obs = inputs["observation"]
+        if inputs.action_space == "midtrain":
+            obs = inputs.observation
             eef_pos = _ensure_2d_float_array(obs["observation/eef_pos"], "observation/eef_pos", 3)
             eef_quat = _ensure_2d_float_array(obs["observation/eef_quat"], "observation/eef_quat", 4)
             initial_pose = np.eye(4, dtype=np.float32)
             initial_pose[:3, :3] = _convert_midtrain_rotation(eef_quat[-1], "quat_xyzw", "matrix")
             initial_pose[:3, 3] = eef_pos[-1]
             abs_pose = _pose_rel_to_abs(
-                action_np[:, :9],
+                action_np[:, :ROBOLAB_MIDTRAIN_POSE_ACTION_DIM],
                 rotation_format="rot6d",
                 pose_convention="backward_framewise",
                 initial_pose=initial_pose,
             )
             position = abs_pose[1:, :3, 3]
             quat_xyzw = _convert_midtrain_rotation(abs_pose[1:, :3, :3], "matrix", "quat_xyzw")
-            action_np = np.concatenate([position, quat_xyzw, action_np[:, 9:]], axis=-1)
+            action_np = np.concatenate([position, quat_xyzw, action_np[:, ROBOLAB_MIDTRAIN_POSE_ACTION_DIM:]], axis=-1)
 
         return np.asarray(action_np, dtype=np.float32)
 
@@ -2247,11 +2273,11 @@ def forward(
         prompt_data = req.prompts[0]
         robolab_inputs = self._build_robolab_policy_inputs(sp, prompt_data)
         if robolab_inputs is not None:
-            prompt = robolab_inputs["prompt"]
+            prompt = robolab_inputs.prompt
             negative_prompt = None
             image_tensor = None
             video_tensor = None
-            action_video_tensor = robolab_inputs["video_tensor"]
+            action_video_tensor = robolab_inputs.video_tensor
         elif isinstance(prompt_data, str):
             prompt = prompt_data
             negative_prompt = None
@@ -2326,13 +2352,13 @@ def forward(
             batch_size = 1  # Existing video pipeline assumes B=1.
 
         if robolab_inputs is not None:
-            height = robolab_inputs["height"]
-            width = robolab_inputs["width"]
-            num_frames = robolab_inputs["num_frames"]
-            action_chunk_size = int(robolab_inputs["action_tensor"].shape[0])
-            num_inference_steps = robolab_inputs["num_inference_steps"]
-            guidance_scale = float(robolab_inputs["guidance_scale"])
-            default_flow_shift = robolab_inputs["flow_shift"]
+            height = robolab_inputs.height
+            width = robolab_inputs.width
+            num_frames = robolab_inputs.num_frames
+            action_chunk_size = int(robolab_inputs.action_tensor.shape[0])
+            num_inference_steps = robolab_inputs.num_inference_steps
+            guidance_scale = float(robolab_inputs.guidance_scale)
+            default_flow_shift = robolab_inputs.flow_shift
             default_guidance_interval = None
         elif action_enabled:
             action_chunk_param = self._get_sp_param(sp, "action_chunk_size", None)
@@ -2359,7 +2385,7 @@ def forward(
         domain_id = None
         if action_enabled:
             if robolab_inputs is not None:
-                domain_id = int(robolab_inputs["domain_id"])
+                domain_id = int(robolab_inputs.domain_id)
             else:
                 domain_id = resolve_domain_id(
                     domain_id=self._get_sp_param(sp, "domain_id", None),
@@ -2378,7 +2404,7 @@ def forward(
         frame_rate = (
             self._get_sp_param(sp, "resolved_frame_rate")
             or self._get_sp_param(sp, "frame_rate")
-            or (robolab_inputs["fps"] if robolab_inputs is not None else 24.0)
+            or (robolab_inputs.fps if robolab_inputs is not None else 24.0)
         )
         max_sequence_length = (
             self._get_sp_param(sp, "max_sequence_length", COSMOS3_DEFAULT_MAX_SEQUENCE_LENGTH)
@@ -2414,7 +2440,7 @@ def forward(
             self._set_flow_shift(flow_shift_target)
 
         generator = sp.generator
-        robolab_seed = int(robolab_inputs["seed"]) if robolab_inputs is not None else None
+        robolab_seed = int(robolab_inputs.seed) if robolab_inputs is not None else None
         if generator is None:
             if robolab_seed is not None:
                 seed = robolab_seed
@@ -2474,9 +2500,9 @@ def forward(
                 image_tensor = action_video_tensor[:, :, 0]
 
             if robolab_inputs is not None:
-                raw_action_dim = int(robolab_inputs["raw_action_dim"])
-                clean_action = robolab_inputs["action_tensor"]
-                action_condition_indexes = robolab_inputs["action_condition_indexes"]
+                raw_action_dim = int(robolab_inputs.raw_action_dim)
+                clean_action = robolab_inputs.action_tensor
+                action_condition_indexes = robolab_inputs.action_condition_indexes
             else:
                 raw_action_dim_param = self._get_sp_param(sp, "raw_action_dim", None)
                 raw_action_dim = int(raw_action_dim_param) if raw_action_dim_param is not None else None
@@ -2493,7 +2519,7 @@ def forward(
             )
             action_latents, action_velocity_mask, action_condition_latents, raw_action_dim = action_prepared
             if robolab_inputs is not None:
-                action_offset = int(robolab_inputs["action_start_frame_offset"])
+                action_offset = int(robolab_inputs.action_start_frame_offset)
             else:
                 action_offset = action_start_frame_offset(action_mode, action_chunk_size, num_frames)
 
@@ -2505,7 +2531,7 @@ def forward(
                 width,
                 num_frames,
                 generator,
-                image_size=robolab_inputs.get("image_size") if robolab_inputs is not None else None,
+                image_size=robolab_inputs.image_size if robolab_inputs is not None else None,
             )
             image_latent = condition_latents[:, :, 0:1]
         elif is_v2v:
@@ -2648,6 +2674,7 @@ def _run_diffusion(start_latents):
                 "domain_id": domain_id,
             }
             if robolab_inputs is not None:
+                custom_action_output["action_only_output"] = True
                 custom_action_output["actions"] = self._postprocess_robolab_action(action, robolab_inputs)
             return DiffusionOutput(
                 output={} if robolab_inputs is not None else {"video": video},
diff --git a/vllm_omni/entrypoints/openpi/connection.py b/vllm_omni/entrypoints/openpi/connection.py
index defc82deaa0..69f722ac824 100644
--- a/vllm_omni/entrypoints/openpi/connection.py
+++ b/vllm_omni/entrypoints/openpi/connection.py
@@ -7,6 +7,10 @@
     Connect  -> server sends msgpack(PolicyServerConfig fields)
     Infer    -> client sends msgpack(obs), server sends msgpack(ndarray)
     Reset    -> client sends msgpack({endpoint:reset}), server sends msgpack(status)
+
+NumPy values use the msgpack-numpy marker mapping:
+    ndarray -> {nd: true, type, kind, shape, data}
+    scalar  -> {nd: false, type, kind, data}
 """
 
 from __future__ import annotations
@@ -27,6 +31,7 @@
 logger = init_logger(__name__)
 _DEFAULT_IDLE_TIMEOUT = 30.0
 MAX_OPENPI_PAYLOAD_BYTES = 64 * 1024 * 1024
+_MISSING = object()
 
 
 def _pack_numpy(obj: Any) -> Any:
@@ -56,15 +61,32 @@ def _mapping_get(obj: dict[Any, Any], key: str, default: Any = None) -> Any:
     return obj.get(key, obj.get(key.encode(), default))
 
 
+def _decode_marker_text(value: Any) -> str:
+    if isinstance(value, bytes):
+        return value.decode()
+    return str(value)
+
+
 def _unpack_numpy(obj: Any) -> Any:
     if isinstance(obj, dict):
-        nd = _mapping_get(obj, "nd")
-        dtype = _mapping_get(obj, "type")
-        data = _mapping_get(obj, "data")
-        if nd is not None and dtype is not None and data is not None:
-            array = np.frombuffer(data, dtype=np.dtype(dtype))
+        nd = _mapping_get(obj, "nd", _MISSING)
+        dtype = _mapping_get(obj, "type", _MISSING)
+        kind = _mapping_get(obj, "kind", _MISSING)
+        data = _mapping_get(obj, "data", _MISSING)
+        if nd is not _MISSING and dtype is not _MISSING and kind is not _MISSING and data is not _MISSING:
+            dtype_obj = np.dtype(_decode_marker_text(dtype))
+            kind_text = _decode_marker_text(kind)
+            if dtype_obj.kind != kind_text:
+                raise ValueError(f"NumPy dtype marker kind mismatch: {dtype_obj.kind!r} != {kind_text!r}")
+            if dtype_obj.kind in ("V", "O", "c"):
+                raise ValueError(f"Unsupported dtype: {dtype_obj}")
+
+            array = np.frombuffer(data, dtype=dtype_obj).copy()
             if nd:
-                return array.reshape(tuple(_mapping_get(obj, "shape", ())))
+                shape = _mapping_get(obj, "shape", _MISSING)
+                if shape is _MISSING:
+                    raise ValueError("NumPy ndarray marker is missing shape")
+                return array.reshape(tuple(shape))
             return array[0]
         return {key: _unpack_numpy(value) for key, value in obj.items()}
     if isinstance(obj, list):

From 566470e612087aa36e831b5ff78b29286e7d056e Mon Sep 17 00:00:00 2001
From: yuzhud <yuzhudong@gmail.com>
Date: Wed, 10 Jun 2026 15:17:31 -0700
Subject: [PATCH 06/12] PR comments

Signed-off-by: yuzhud <yuzhudong@gmail.com>
---
 .../models/cosmos3/test_cosmos3_pipeline.py   |  86 ++++
 vllm_omni/diffusion/diffusion_engine.py       |   6 +-
 .../models/cosmos3/pipeline_cosmos3.py        | 383 ++++++++++++------
 .../models/cosmos3/transformer_cosmos3.py     |  49 +--
 vllm_omni/entrypoints/openai/serving_video.py |   4 +-
 5 files changed, 371 insertions(+), 157 deletions(-)

diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py
index 3eebda1b639..cf513d75f6b 100644
--- a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py
+++ b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py
@@ -764,6 +764,92 @@ def test_forward_i2v_sound_and_action_routes(self, make_cosmos3_pipeline) -> Non
         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 output.custom_output["actions"].shape == (1, 2)
+
     @pytest.mark.parametrize(
         ("prompt", "sampling_params", "message"),
         [
diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py
index 344847ed334..196be527479 100644
--- a/vllm_omni/diffusion/diffusion_engine.py
+++ b/vllm_omni/diffusion/diffusion_engine.py
@@ -262,7 +262,7 @@ async def step(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]:
             output_data = _move_tensor_tree_to_cpu(output_data)
 
         custom_output = output.custom_output or {}
-        action_payload = custom_output.get("actions")
+        action_payload = None
         action_only_output = bool(custom_output.get("action_only_output"))
 
         postprocess_start_time = time.perf_counter()
@@ -282,11 +282,13 @@ async def step(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]:
         model_fps = None
         if isinstance(outputs, dict):
             audio_payload = outputs.get("audio")
-            action_payload = outputs.get("actions", action_payload)
+            action_payload = outputs.get("actions")
             custom_output.update(outputs.get("custom_output") or {})
             model_audio_sample_rate = outputs.get("audio_sample_rate")
             model_fps = outputs.get("fps")
             outputs = outputs.get("video", outputs)
+        if action_payload is None:
+            action_payload = custom_output.get("actions")
         postprocess_time = time.perf_counter() - postprocess_start_time
         logger.debug("Post-processing completed in %.4f seconds", postprocess_time)
 
diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py
index 4195abd24b2..fa0d8a94917 100644
--- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py
+++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py
@@ -20,8 +20,10 @@
 import math
 import os
 import time
+import zlib
 from collections.abc import Iterable
 from dataclasses import dataclass
+from importlib import import_module
 from typing import Any, ClassVar
 
 import numpy as np
@@ -115,6 +117,7 @@
 ROBOLAB_DEFAULT_FLOW_SHIFT = 5.0
 ROBOLAB_DEFAULT_SEED = 0
 ROBOLAB_DEFAULT_ACTION_SPACE = "joint_pos"
+ROBOLAB_MIDTRAIN_RAW_ACTION_DIM = 10
 ROBOLAB_MIDTRAIN_POSE_ACTION_DIM = 9  # xyz position + rot6d orientation
 ROBOLAB_CONCAT_VIEW_DESCRIPTION = (
     "The top row is from the wrist-mounted camera. "
@@ -222,7 +225,7 @@ def _ensure_gripper_array(value: Any) -> np.ndarray:
 def _resize_rgb_uint8(image: np.ndarray, size: tuple[int, int]) -> np.ndarray:
     tensor = torch.from_numpy(image).permute(2, 0, 1).unsqueeze(0).float()
     resized = F.interpolate(tensor, size=size, mode="bilinear", align_corners=False)
-    return resized.squeeze(0).permute(1, 2, 0).numpy().astype(np.uint8)
+    return np.clip(np.round(resized.squeeze(0).permute(1, 2, 0).numpy()), 0, 255).astype(np.uint8)
 
 
 def _compose_robolab_views(obs: dict[str, Any]) -> np.ndarray | None:
@@ -252,25 +255,33 @@ def _extract_robolab_image(obs: dict[str, Any]) -> np.ndarray:
     raise ValueError("Observation must contain 'observation/image' or RoboLab wrist/exterior image keys.")
 
 
-def _lazy_action_transform_pipeline(max_action_dim: int):
+def _lazy_import(module_name: str, symbol_name: str, error_message: str):
     try:
-        from cosmos_framework.data.vfm.action.transforms import ActionTransformPipeline
+        module = import_module(module_name)
     except ModuleNotFoundError as exc:
-        raise ModuleNotFoundError(
-            "Cosmos3 RoboLab policy serving requires cosmos_framework on PYTHONPATH so the "
-            "golden ActionTransformPipeline can be reused."
-        ) from exc
+        raise ModuleNotFoundError(error_message) from exc
+    return getattr(module, symbol_name)
+
+
+def _lazy_action_transform_pipeline(max_action_dim: int):
+    ActionTransformPipeline = _lazy_import(
+        "cosmos_framework.data.vfm.action.transforms",
+        "ActionTransformPipeline",
+        "Cosmos3 RoboLab policy serving requires cosmos_framework on PYTHONPATH so the "
+        "golden ActionTransformPipeline can be reused.",
+    )
     return ActionTransformPipeline(max_action_dim=max_action_dim, cfg_dropout_rate=0.0)
 
 
 def _build_robolab_unipc_scheduler(num_steps: int, shift: float, device: torch.device):
-    try:
-        from cosmos_framework.model.vfm.diffusion.samplers.fm_solvers_unipc import FlowUniPCMultistepScheduler
-    except ModuleNotFoundError as exc:
-        raise ModuleNotFoundError(
+    FlowUniPCMultistepScheduler = _lazy_import(
+        "cosmos_framework.model.vfm.diffusion.samplers.fm_solvers_unipc",
+        "FlowUniPCMultistepScheduler",
+        (
             "Cosmos3 RoboLab policy serving requires cosmos_framework on PYTHONPATH so the "
             "golden FlowUniPCMultistepScheduler can be reused."
-        ) from exc
+        ),
+    )
 
     scheduler = FlowUniPCMultistepScheduler(
         num_train_timesteps=1000,
@@ -282,42 +293,38 @@ def _build_robolab_unipc_scheduler(num_steps: int, shift: float, device: torch.d
 
 
 def _convert_midtrain_rotation(value: Any, src: str, dst: str) -> np.ndarray:
-    try:
-        from cosmos_framework.data.vfm.action.pose_utils import convert_rotation
-    except ModuleNotFoundError as exc:
-        raise ModuleNotFoundError(
-            "Cosmos3 RoboLab midtrain action serving requires cosmos_framework pose_utils on PYTHONPATH."
-        ) from exc
+    convert_rotation = _lazy_import(
+        "cosmos_framework.data.vfm.action.pose_utils",
+        "convert_rotation",
+        "Cosmos3 RoboLab midtrain action serving requires cosmos_framework pose_utils on PYTHONPATH.",
+    )
     return convert_rotation(value, src, dst)
 
 
 def _pose_abs_to_rel(*args, **kwargs) -> np.ndarray:
-    try:
-        from cosmos_framework.data.vfm.action.pose_utils import pose_abs_to_rel
-    except ModuleNotFoundError as exc:
-        raise ModuleNotFoundError(
-            "Cosmos3 RoboLab midtrain action serving requires cosmos_framework pose_utils on PYTHONPATH."
-        ) from exc
+    pose_abs_to_rel = _lazy_import(
+        "cosmos_framework.data.vfm.action.pose_utils",
+        "pose_abs_to_rel",
+        "Cosmos3 RoboLab midtrain action serving requires cosmos_framework pose_utils on PYTHONPATH.",
+    )
     return pose_abs_to_rel(*args, **kwargs)
 
 
 def _pose_rel_to_abs(*args, **kwargs) -> np.ndarray:
-    try:
-        from cosmos_framework.data.vfm.action.pose_utils import pose_rel_to_abs
-    except ModuleNotFoundError as exc:
-        raise ModuleNotFoundError(
-            "Cosmos3 RoboLab midtrain action serving requires cosmos_framework pose_utils on PYTHONPATH."
-        ) from exc
+    pose_rel_to_abs = _lazy_import(
+        "cosmos_framework.data.vfm.action.pose_utils",
+        "pose_rel_to_abs",
+        "Cosmos3 RoboLab midtrain action serving requires cosmos_framework pose_utils on PYTHONPATH.",
+    )
     return pose_rel_to_abs(*args, **kwargs)
 
 
 def _build_abs_pose_from_components(*args, **kwargs) -> np.ndarray:
-    try:
-        from cosmos_framework.data.vfm.action.pose_utils import build_abs_pose_from_components
-    except ModuleNotFoundError as exc:
-        raise ModuleNotFoundError(
-            "Cosmos3 RoboLab midtrain action serving requires cosmos_framework pose_utils on PYTHONPATH."
-        ) from exc
+    build_abs_pose_from_components = _lazy_import(
+        "cosmos_framework.data.vfm.action.pose_utils",
+        "build_abs_pose_from_components",
+        "Cosmos3 RoboLab midtrain action serving requires cosmos_framework pose_utils on PYTHONPATH.",
+    )
     return build_abs_pose_from_components(*args, **kwargs)
 
 
@@ -832,7 +839,6 @@ def __init__(
         self._guidance_scale = None
         self._num_timesteps = None
         self._robolab_transform = None
-        self._robolab_rng = np.random.default_rng(ROBOLAB_DEFAULT_SEED)
 
         # Set True by ``enable_cache_for_cosmos3`` when cache-dit is enabled on
         # this pipeline. Tells the sequential-CFG loop to keep paired
@@ -1075,19 +1081,29 @@ def _get_robolab_transform(self):
             self._robolab_transform = _lazy_action_transform_pipeline(action_dim)
         return self._robolab_transform
 
-    def _next_robolab_seed(self, extra: dict[str, Any]) -> int:
+    def _next_robolab_seed(self, extra: dict[str, Any], obs: dict[str, Any], request_id: str | None) -> int:
         base_seed = int(extra.get("robolab_seed") or ROBOLAB_DEFAULT_SEED)
         if self._truthy(extra.get("deterministic_seed")):
             return base_seed
         explicit_seed = extra.get("seed")
         if explicit_seed is not None:
             return int(explicit_seed)
-        return int(self._robolab_rng.integers(0, 2**31))
+        seed_key = "|".join(
+            str(part)
+            for part in (
+                base_seed,
+                extra.get("session_id", ""),
+                request_id or "",
+                obs.get("prompt", ""),
+            )
+        )
+        return zlib.crc32(seed_key.encode("utf-8")) & 0x7FFFFFFF
 
     def _build_robolab_policy_inputs(
         self,
         sp: OmniDiffusionSamplingParams,
         prompt_data: Any | None = None,
+        request_id: str | None = None,
     ) -> RoboLabPolicyInputs | None:
         extra = sp.extra_args if isinstance(sp.extra_args, dict) else {}
         obs = extra.get("robot_obs")
@@ -1115,7 +1131,9 @@ def extra_param_alias(primary_key: str, alias_key: str, default: Any) -> Any:
 
         action_space = self._normalize_robolab_action_space(extra_param("action_space", ROBOLAB_DEFAULT_ACTION_SPACE))
         action_chunk_size = int(extra_param("action_chunk_size", ROBOLAB_DEFAULT_ACTION_CHUNK_SIZE))
-        raw_action_dim_default = ROBOLAB_DEFAULT_RAW_ACTION_DIM if action_space == "joint_pos" else 10
+        raw_action_dim_default = (
+            ROBOLAB_DEFAULT_RAW_ACTION_DIM if action_space == "joint_pos" else ROBOLAB_MIDTRAIN_RAW_ACTION_DIM
+        )
         raw_action_dim = int(extra_param("raw_action_dim", raw_action_dim_default))
         image_h = int(extra_param("image_height", ROBOLAB_DEFAULT_IMAGE_HEIGHT))
         image_w = int(extra_param("image_width", ROBOLAB_DEFAULT_IMAGE_WIDTH))
@@ -1222,7 +1240,7 @@ def extra_param_alias(primary_key: str, alias_key: str, default: Any) -> Any:
             ),
             guidance_scale=float(extra_param_alias("guidance_scale", "guidance", ROBOLAB_DEFAULT_GUIDANCE_SCALE)),
             flow_shift=float(extra_param_alias("flow_shift", "shift", ROBOLAB_DEFAULT_FLOW_SHIFT)),
-            seed=self._next_robolab_seed(extra),
+            seed=self._next_robolab_seed(extra, obs, request_id),
             history_length=history_length,
             action_space=action_space,
             observation=obs,
@@ -1255,8 +1273,39 @@ def _build_action_condition_mask_from_indexes(
             mask[:, idx, :] = 1.0
         return mask
 
+    @staticmethod
+    def _log_robolab_action_summary(label: str, value: Any) -> None:
+        if not _is_rank_zero():
+            return
+        if isinstance(value, torch.Tensor):
+            array = value.detach().float().cpu().numpy()
+        else:
+            array = np.asarray(value, dtype=np.float32)
+        finite = np.isfinite(array)
+        if finite.any():
+            finite_min = float(array[finite].min())
+            finite_max = float(array[finite].max())
+        else:
+            finite_min = None
+            finite_max = None
+        if array.ndim == 0:
+            head = array.reshape(1).tolist()
+        else:
+            head = array.reshape(-1, array.shape[-1])[:3].tolist()
+        logger.info(
+            "RoboLab action summary %s: shape=%s nan=%d finite=%d finite_min=%s finite_max=%s head=%s",
+            label,
+            tuple(array.shape),
+            int(np.isnan(array).sum()),
+            int(finite.sum()),
+            finite_min,
+            finite_max,
+            head,
+        )
+
     def _postprocess_robolab_action(self, action: torch.Tensor, inputs: RoboLabPolicyInputs) -> np.ndarray:
         action_np = action[0].float().cpu().numpy()
+        self._log_robolab_action_summary("raw_model_action", action_np)
         history_length = int(inputs.history_length)
         action_np = action_np[history_length:]
         action_np[:, -1] = 1.0 - action_np[:, -1]
@@ -1278,8 +1327,141 @@ def _postprocess_robolab_action(self, action: torch.Tensor, inputs: RoboLabPolic
             quat_xyzw = _convert_midtrain_rotation(abs_pose[1:, :3, :3], "matrix", "quat_xyzw")
             action_np = np.concatenate([position, quat_xyzw, action_np[:, ROBOLAB_MIDTRAIN_POSE_ACTION_DIM:]], axis=-1)
 
+        self._log_robolab_action_summary("postprocessed_robolab_action", action_np)
         return np.asarray(action_np, dtype=np.float32)
 
+    def _forward_robolab_policy(
+        self,
+        sp: OmniDiffusionSamplingParams,
+        inputs: RoboLabPolicyInputs,
+        pipeline_start: float,
+    ) -> DiffusionOutput:
+        if not getattr(self.transformer, "action_gen", False):
+            raise ValueError(
+                "Cosmos3 RoboLab policy serving was requested, but the transformer "
+                "was initialized without action modules. Check that the checkpoint "
+                "config enables action_gen and includes action weights."
+            )
+
+        action_mode = ACTION_MODE_POLICY
+        height = inputs.height
+        width = inputs.width
+        num_frames = inputs.num_frames
+        action_chunk_size = int(inputs.action_tensor.shape[0])
+        num_inference_steps = inputs.num_inference_steps
+        guidance_scale = float(inputs.guidance_scale)
+        flow_shift_target = float(inputs.flow_shift)
+        domain_id = int(inputs.domain_id)
+        frame_rate = self._get_sp_param(sp, "resolved_frame_rate") or self._get_sp_param(sp, "frame_rate") or inputs.fps
+        max_sequence_length = (
+            self._get_sp_param(sp, "max_sequence_length", COSMOS3_DEFAULT_MAX_SEQUENCE_LENGTH)
+            or COSMOS3_DEFAULT_MAX_SEQUENCE_LENGTH
+        )
+        use_system_prompt = bool(self._get_sp_param(sp, "use_system_prompt", False))
+
+        self._guidance_scale = guidance_scale
+        self._num_timesteps = num_inference_steps
+
+        generator = sp.generator
+        if generator is None:
+            generator = torch.Generator(device=self.device).manual_seed(int(inputs.seed))
+
+        cond_ids, cond_mask, uncond_ids, uncond_mask = self._format_and_tokenize_prompts(
+            inputs.prompt,
+            "",
+            num_frames,
+            frame_rate,
+            height,
+            width,
+            max_sequence_length,
+            sp,
+            use_system_prompt,
+            is_t2i=False,
+        )
+
+        action_video_tensor = inputs.video_tensor
+        if action_video_tensor.ndim == 4:
+            action_video_tensor = action_video_tensor.unsqueeze(0)
+        if action_video_tensor.ndim != 5:
+            raise ValueError(
+                "Cosmos3 RoboLab action video tensor must have shape [1, 3, T, H, W] "
+                f"or [3, T, H, W], got {tuple(action_video_tensor.shape)}."
+            )
+        if action_video_tensor.shape[2] < num_frames:
+            pad = action_video_tensor[:, :, -1:].repeat(1, 1, num_frames - action_video_tensor.shape[2], 1, 1)
+            action_video_tensor = torch.cat([action_video_tensor, pad], dim=2)
+        elif action_video_tensor.shape[2] > num_frames:
+            action_video_tensor = action_video_tensor[:, :, :num_frames]
+
+        action_latents, action_velocity_mask, action_condition_latents, raw_action_dim = self._prepare_action_latents(
+            mode=action_mode,
+            action_chunk_size=action_chunk_size,
+            raw_action_dim=int(inputs.raw_action_dim),
+            generator=generator,
+            sp=sp,
+            clean_action=inputs.action_tensor,
+            condition_indexes=inputs.action_condition_indexes,
+        )
+        action_offset = int(inputs.action_start_frame_offset)
+
+        latents, velocity_mask, condition_latents = self._prepare_latents_action_video(
+            action_video_tensor,
+            action_mode,
+            height,
+            width,
+            num_frames,
+            generator,
+            image_size=inputs.image_size,
+        )
+        image_latent = condition_latents[:, :, 0:1]
+
+        video_shape = (latents.shape[2], latents.shape[3], latents.shape[4])
+        shared_kwargs = dict(
+            video_shape=video_shape,
+            fps=frame_rate,
+            noisy_frame_mask=velocity_mask,
+            action_domain_ids=torch.tensor([domain_id], dtype=torch.long, device=self.device),
+            action_noisy_mask=action_velocity_mask,
+            action_start_frame_offset=action_offset,
+            action_fps=float(self._get_sp_param(sp, "action_fps", frame_rate) or frame_rate),
+        )
+
+        scheduler = _build_robolab_unipc_scheduler(num_inference_steps, flow_shift_target, self.device)
+        _, action_latents = self.diffuse(
+            latents=latents,
+            timesteps=scheduler.timesteps,
+            cond_ids=cond_ids,
+            cond_mask=cond_mask,
+            uncond_ids=uncond_ids,
+            uncond_mask=uncond_mask,
+            guidance_scale=guidance_scale,
+            shared_kwargs=shared_kwargs,
+            action_latents=action_latents,
+            action_velocity_mask=action_velocity_mask,
+            action_condition_latents=action_condition_latents,
+            sound_latents=None,
+            velocity_mask=velocity_mask,
+            image_latent=image_latent,
+            condition_latents=condition_latents,
+            guidance_interval=None,
+            raw_action_dim=raw_action_dim,
+            scheduler=scheduler,
+        )
+
+        if _is_rank_zero():
+            logger.info("Total pipeline time: %.2fs", time.time() - pipeline_start)
+
+        action = action_latents[:, :, :raw_action_dim].detach().cpu()
+        custom_action_output: dict[str, Any] = {
+            "action": action,
+            "raw_action_dim": raw_action_dim,
+            "action_mode": action_mode,
+            "domain_id": domain_id,
+            "action_only_output": True,
+            "actions": self._postprocess_robolab_action(action, inputs),
+        }
+        return DiffusionOutput(output={}, custom_output=custom_action_output)
+
     @staticmethod
     def _truthy(value) -> bool:
         if isinstance(value, str):
@@ -2271,14 +2453,11 @@ def forward(
 
         sp = req.sampling_params
         prompt_data = req.prompts[0]
-        robolab_inputs = self._build_robolab_policy_inputs(sp, prompt_data)
+        robolab_inputs = self._build_robolab_policy_inputs(sp, prompt_data, getattr(req, "request_id", None))
         if robolab_inputs is not None:
-            prompt = robolab_inputs.prompt
-            negative_prompt = None
-            image_tensor = None
-            video_tensor = None
-            action_video_tensor = robolab_inputs.video_tensor
-        elif isinstance(prompt_data, str):
+            return self._forward_robolab_policy(sp, robolab_inputs, pipeline_start)
+
+        if isinstance(prompt_data, str):
             prompt = prompt_data
             negative_prompt = None
             image_tensor = None
@@ -2290,12 +2469,11 @@ def forward(
             image_tensor = additional_info.get("preprocessed_image")
             video_tensor = additional_info.get("preprocessed_video")
 
-        is_t2i = False if robolab_inputs is not None else self._is_t2i_request(req)
-        sound_enabled = False if robolab_inputs is not None else self._is_sound_request(prompt_data, sp)
-        action_mode = ACTION_MODE_POLICY if robolab_inputs is not None else self._get_action_mode(prompt_data, sp)
+        is_t2i = self._is_t2i_request(req)
+        sound_enabled = self._is_sound_request(prompt_data, sp)
+        action_mode = self._get_action_mode(prompt_data, sp)
         action_enabled = action_mode is not None
-        if robolab_inputs is None:
-            action_video_tensor = video_tensor if action_enabled else None
+        action_video_tensor = video_tensor if action_enabled else None
         if action_enabled and is_t2i:
             raise ValueError("Cosmos3 action generation is supported only for video outputs.")
         if action_enabled and sound_enabled:
@@ -2351,16 +2529,7 @@ def forward(
             default_guidance_interval = None
             batch_size = 1  # Existing video pipeline assumes B=1.
 
-        if robolab_inputs is not None:
-            height = robolab_inputs.height
-            width = robolab_inputs.width
-            num_frames = robolab_inputs.num_frames
-            action_chunk_size = int(robolab_inputs.action_tensor.shape[0])
-            num_inference_steps = robolab_inputs.num_inference_steps
-            guidance_scale = float(robolab_inputs.guidance_scale)
-            default_flow_shift = robolab_inputs.flow_shift
-            default_guidance_interval = None
-        elif action_enabled:
+        if action_enabled:
             action_chunk_param = self._get_sp_param(sp, "action_chunk_size", None)
             if action_chunk_param is not None:
                 action_chunk_size = int(action_chunk_param)
@@ -2384,28 +2553,18 @@ def forward(
 
         domain_id = None
         if action_enabled:
-            if robolab_inputs is not None:
-                domain_id = int(robolab_inputs.domain_id)
-            else:
-                domain_id = resolve_domain_id(
-                    domain_id=self._get_sp_param(sp, "domain_id", None),
-                    domain_name=self._get_sp_param(sp, "domain_name", None),
-                    require_explicit=True,
-                )
+            domain_id = resolve_domain_id(
+                domain_id=self._get_sp_param(sp, "domain_id", None),
+                domain_name=self._get_sp_param(sp, "domain_name", None),
+                require_explicit=True,
+            )
 
         # Runtime controls: prefer ``extra_args`` (OpenAI endpoints write
         # there) over direct attrs.
-        if robolab_inputs is not None:
-            flow_shift_target = float(default_flow_shift)
-        else:
-            flow_shift_target = float(self._get_sp_param(sp, "flow_shift", default_flow_shift))
+        flow_shift_target = float(self._get_sp_param(sp, "flow_shift", default_flow_shift))
         guidance_interval = self._get_sp_param(sp, "guidance_interval", default_guidance_interval)
 
-        frame_rate = (
-            self._get_sp_param(sp, "resolved_frame_rate")
-            or self._get_sp_param(sp, "frame_rate")
-            or (robolab_inputs.fps if robolab_inputs is not None else 24.0)
-        )
+        frame_rate = self._get_sp_param(sp, "resolved_frame_rate") or self._get_sp_param(sp, "frame_rate") or 24.0
         max_sequence_length = (
             self._get_sp_param(sp, "max_sequence_length", COSMOS3_DEFAULT_MAX_SEQUENCE_LENGTH)
             or COSMOS3_DEFAULT_MAX_SEQUENCE_LENGTH
@@ -2433,19 +2592,12 @@ def forward(
         self._num_timesteps = num_inference_steps
 
         # Always resolve to a concrete target shift for this request, then
-        # update the shared Diffusers scheduler for non-RoboLab paths. RoboLab
-        # uses a request-local Cosmos UniPC scheduler to match the golden
-        # action server.
-        if robolab_inputs is None:
-            self._set_flow_shift(flow_shift_target)
+        # update the shared Diffusers scheduler.
+        self._set_flow_shift(flow_shift_target)
 
         generator = sp.generator
-        robolab_seed = int(robolab_inputs.seed) if robolab_inputs is not None else None
         if generator is None:
-            if robolab_seed is not None:
-                seed = robolab_seed
-            else:
-                seed = sp.seed if sp.seed is not None else 42
+            seed = sp.seed if sp.seed is not None else 42
             generator = torch.Generator(device=self.device).manual_seed(seed)
 
         # --- Format prompts & tokenize (B=1; reused across loop iterations
@@ -2499,15 +2651,10 @@ def forward(
                     )
                 image_tensor = action_video_tensor[:, :, 0]
 
-            if robolab_inputs is not None:
-                raw_action_dim = int(robolab_inputs.raw_action_dim)
-                clean_action = robolab_inputs.action_tensor
-                action_condition_indexes = robolab_inputs.action_condition_indexes
-            else:
-                raw_action_dim_param = self._get_sp_param(sp, "raw_action_dim", None)
-                raw_action_dim = int(raw_action_dim_param) if raw_action_dim_param is not None else None
-                clean_action = None
-                action_condition_indexes = None
+            raw_action_dim_param = self._get_sp_param(sp, "raw_action_dim", None)
+            raw_action_dim = int(raw_action_dim_param) if raw_action_dim_param is not None else None
+            clean_action = None
+            action_condition_indexes = None
             action_prepared = self._prepare_action_latents(
                 mode=action_mode,
                 action_chunk_size=action_chunk_size,
@@ -2518,10 +2665,7 @@ def forward(
                 condition_indexes=action_condition_indexes,
             )
             action_latents, action_velocity_mask, action_condition_latents, raw_action_dim = action_prepared
-            if robolab_inputs is not None:
-                action_offset = int(robolab_inputs.action_start_frame_offset)
-            else:
-                action_offset = action_start_frame_offset(action_mode, action_chunk_size, num_frames)
+            action_offset = action_start_frame_offset(action_mode, action_chunk_size, num_frames)
 
         if action_enabled and action_video_tensor is not None:
             latents, velocity_mask, condition_latents = self._prepare_latents_action_video(
@@ -2531,7 +2675,6 @@ def forward(
                 width,
                 num_frames,
                 generator,
-                image_size=robolab_inputs.image_size if robolab_inputs is not None else None,
             )
             image_latent = condition_latents[:, :, 0:1]
         elif is_v2v:
@@ -2591,11 +2734,8 @@ def forward(
             )
 
         def _run_diffusion(start_latents):
-            if robolab_inputs is not None:
-                scheduler = _build_robolab_unipc_scheduler(num_inference_steps, flow_shift_target, self.device)
-            else:
-                self.scheduler.set_timesteps(num_inference_steps, device=self.device)
-                scheduler = self.scheduler
+            self.scheduler.set_timesteps(num_inference_steps, device=self.device)
+            scheduler = self.scheduler
             return self.diffuse(
                 latents=start_latents,
                 timesteps=scheduler.timesteps,
@@ -2642,17 +2782,12 @@ def _run_diffusion(start_latents):
                 latents = diffusion_output
 
         # --- Decode ---
-        video = None
-        should_decode_video = robolab_inputs is None
-        if should_decode_video:
-            if _is_rank_zero():
-                logger.info("Decoding video...")
-            decode_start = time.time()
-            video = self._decode_latents(latents)
-            if _is_rank_zero():
-                logger.info("Video decoded in %.2fs", time.time() - decode_start)
-                logger.info("Total pipeline time: %.2fs", time.time() - pipeline_start)
-        elif _is_rank_zero():
+        if _is_rank_zero():
+            logger.info("Decoding video...")
+        decode_start = time.time()
+        video = self._decode_latents(latents)
+        if _is_rank_zero():
+            logger.info("Video decoded in %.2fs", time.time() - decode_start)
             logger.info("Total pipeline time: %.2fs", time.time() - pipeline_start)
 
         if sound_enabled:
@@ -2673,12 +2808,6 @@ def _run_diffusion(start_latents):
                 "action_mode": action_mode,
                 "domain_id": domain_id,
             }
-            if robolab_inputs is not None:
-                custom_action_output["action_only_output"] = True
-                custom_action_output["actions"] = self._postprocess_robolab_action(action, robolab_inputs)
-            return DiffusionOutput(
-                output={} if robolab_inputs is not None else {"video": video},
-                custom_output=custom_action_output,
-            )
+            return DiffusionOutput(output={"video": video}, custom_output=custom_action_output)
 
         return DiffusionOutput(output={"image": video} if is_t2i else {"video": video})
diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py
index 58c5010a01e..163158a39dc 100644
--- a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py
+++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py
@@ -33,26 +33,20 @@
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput, SequenceParallelOutput
 from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
-from vllm_omni.diffusion.layers.norm import RMSNorm
+from vllm_omni.diffusion.layers.norm import RMSNorm as _VllmRMSNorm
 from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
 
-class Cosmos3RMSNorm(nn.Module):
-    """RMSNorm matching Cosmos Framework's Qwen3VLTextRMSNorm."""
+class RMSNorm(_VllmRMSNorm):
+    """Cosmos3-local RMSNorm that uses the FP32 native implementation."""
 
-    def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
-        super().__init__()
-        self.weight = nn.Parameter(torch.ones(hidden_size))
-        self.variance_epsilon = eps
+    def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
+        return self.forward_native(x)
 
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        input_dtype = hidden_states.dtype
-        hidden_states = hidden_states.to(torch.float32)
-        variance = hidden_states.pow(2).mean(-1, keepdim=True)
-        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
-        return self.weight * hidden_states.to(input_dtype)
+    def forward_hip(self, x: torch.Tensor) -> torch.Tensor:
+        return self.forward_native(x)
 
 
 def _get_ulysses_state() -> tuple[int, int, dist.ProcessGroup | None]:
@@ -527,8 +521,8 @@ def __init__(
             prefix=f"{prefix}.to_out",
         )
 
-        self.norm_q = Cosmos3RMSNorm(self.head_dim, eps=rms_norm_eps)
-        self.norm_k = Cosmos3RMSNorm(self.head_dim, eps=rms_norm_eps)
+        self.norm_q = RMSNorm(self.head_dim, eps=rms_norm_eps)
+        self.norm_k = RMSNorm(self.head_dim, eps=rms_norm_eps)
 
         # skip_sequence_parallel=True because the UND pathway is
         # computed once and replicated across SP ranks.
@@ -555,8 +549,8 @@ def forward(
         v = self.to_v(hidden_states).view(B, S, self.num_kv_heads_local, self.head_dim)
 
         # Per-head QK norm
-        q = self.norm_q(q)
-        k = self.norm_k(k)
+        q = F.rms_norm(q, (self.head_dim,), self.norm_q.weight, eps=self.norm_q.variance_epsilon)
+        k = F.rms_norm(k, (self.head_dim,), self.norm_k.weight, eps=self.norm_k.variance_epsilon)
 
         # Qwen3-style RoPE
         q, k = _apply_rotary_pos_emb(q, k, freqs_cos, freqs_sin)
@@ -635,8 +629,8 @@ def __init__(
             prefix=f"{prefix}.to_out",
         )
 
-        self.norm_q = Cosmos3RMSNorm(self.head_dim, eps=rms_norm_eps)
-        self.norm_k = Cosmos3RMSNorm(self.head_dim, eps=rms_norm_eps)
+        self.norm_q = RMSNorm(self.head_dim, eps=rms_norm_eps)
+        self.norm_k = RMSNorm(self.head_dim, eps=rms_norm_eps)
 
         self.attn = FrameworkAttention(
             num_heads=self.num_heads,
@@ -720,8 +714,8 @@ def forward(
         v = self.to_v(hidden_states).view(B, S_gen, self.num_kv_heads_local, self.head_dim)
 
         # Per-head QK norm
-        q = self.norm_q(q)
-        k = self.norm_k(k)
+        q = F.rms_norm(q, (self.head_dim,), self.norm_q.weight, eps=self.norm_q.variance_epsilon)
+        k = F.rms_norm(k, (self.head_dim,), self.norm_k.weight, eps=self.norm_k.variance_epsilon)
 
         # Qwen3-style RoPE
         q, k = _apply_rotary_pos_emb(q, k, freqs_cos, freqs_sin)
@@ -762,8 +756,8 @@ def __init__(
             quant_config=quant_config,
             prefix=f"{prefix}.self_attn",
         )
-        self.input_layernorm = Cosmos3RMSNorm(hidden_size, eps=rms_norm_eps)
-        self.post_attention_layernorm = Cosmos3RMSNorm(hidden_size, eps=rms_norm_eps)
+        self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
         self.mlp = Cosmos3GatedMLP(
             hidden_size=hidden_size,
             intermediate_size=intermediate_size,
@@ -818,8 +812,8 @@ def __init__(
             quant_config=quant_config,
             prefix=f"{prefix}.cross_attention",
         )
-        self.input_layernorm = Cosmos3RMSNorm(hidden_size, eps=rms_norm_eps)
-        self.post_attention_layernorm = Cosmos3RMSNorm(hidden_size, eps=rms_norm_eps)
+        self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
         self.mlp = Cosmos3GatedMLP(
             hidden_size=hidden_size,
             intermediate_size=intermediate_size,
@@ -912,7 +906,7 @@ def __init__(
             ]
         )
         # TODO: Not used right now, will be used in the future for prompt upsampler.
-        self.norm = Cosmos3RMSNorm(hidden_size, eps=rms_norm_eps)
+        self.norm = RMSNorm(hidden_size, eps=rms_norm_eps)
 
     def forward(
         self,
@@ -1127,7 +1121,7 @@ def __init__(
             ]
         )
 
-        self.norm_moe_gen = Cosmos3RMSNorm(self.hidden_size, eps=self.rms_norm_eps)
+        self.norm_moe_gen = RMSNorm(self.hidden_size, eps=self.rms_norm_eps)
         self.gen_sp_prepare = Cosmos3GenSPPrepare()
         self.gen_sp_gather = nn.Identity()
 
@@ -1462,6 +1456,7 @@ def forward(
                         "Cosmos3 action_noisy_mask must have shape [B, T_action, 1], "
                         f"got {tuple(action_noisy_mask.shape)}."
                     )
+                action_noisy_mask = action_noisy_mask.to(dtype=hidden_action.dtype, device=hidden_action.device)
                 hidden_action = hidden_action + time_embed.unsqueeze(1) * action_noisy_mask
 
         if hidden_sound is not None:
diff --git a/vllm_omni/entrypoints/openai/serving_video.py b/vllm_omni/entrypoints/openai/serving_video.py
index 712f6647ea0..4bdbfce0217 100644
--- a/vllm_omni/entrypoints/openai/serving_video.py
+++ b/vllm_omni/entrypoints/openai/serving_video.py
@@ -223,7 +223,9 @@ async def _run_and_extract(
         )
 
         result = await self._run_generation(prompt, gen_params, reference_id)
-        videos = self._extract_video_outputs(result)
+        custom_output = self._extract_custom_output(result)
+        action_only = isinstance(custom_output, dict) and bool(custom_output.get("action_only_output"))
+        videos = [{"action_only_output": True}] if action_only else self._extract_video_outputs(result)
         audios = self._extract_audio_outputs(result, expected_count=len(videos))
         actions = self._extract_action_outputs(result, expected_count=len(videos))
         audio_sample_rate = self._resolve_audio_sample_rate(result)

From b9dd34381f1b741f8a1e267a9d7c3737eb7eeb55 Mon Sep 17 00:00:00 2001
From: yuzhud <yuzhudong@gmail.com>
Date: Thu, 11 Jun 2026 10:30:09 -0700
Subject: [PATCH 07/12] Refactor Cosmos3 action postprocessing

Signed-off-by: yuzhud <yuzhudong@gmail.com>
---
 .../models/cosmos3/test_cosmos3_pipeline.py   |  46 +-
 .../diffusion/test_diffusion_plugin_hooks.py  |   6 +
 tests/diffusion/test_diffusion_scheduler.py   |  11 +-
 vllm_omni/diffusion/diffusion_engine.py       |  31 +-
 .../diffusion/models/cosmos3/__init__.py      |   2 +
 .../models/cosmos3/pipeline_cosmos3.py        | 425 +++---------------
 vllm_omni/diffusion/models/cosmos3/utils.py   | 348 ++++++++++++++
 vllm_omni/diffusion/registry.py               |  20 +
 8 files changed, 528 insertions(+), 361 deletions(-)
 create mode 100644 vllm_omni/diffusion/models/cosmos3/utils.py

diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py
index cf513d75f6b..5ca99ef9caf 100644
--- a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py
+++ b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py
@@ -230,6 +230,7 @@ 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_MODELS,
         _DIFFUSION_POST_PROCESS_FUNCS,
         _DIFFUSION_PRE_PROCESS_FUNCS,
@@ -245,6 +246,9 @@ 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 "Cosmos3OmniDiffusersPipeline" in CUSTOM_DIT_ENABLERS
     assert "Cosmos3OmniDiffusersPipeline" in cosmos3.__all__
 
@@ -413,6 +417,43 @@ 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,
+    )
+
+    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]]])
+    processed = func(action, custom_output={"robolab_policy_inputs": inputs})
+
+    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]]))
+
+
 def test_prompt_formatting_and_checkpoint_key_remap(make_cosmos3_pipeline) -> None:
     from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline
 
@@ -817,7 +858,7 @@ def fake_prepare_action_video(*args, **kwargs):
 
         monkeypatch.setattr(
             pipeline_cosmos3,
-            "_build_robolab_unipc_scheduler",
+            "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
@@ -848,7 +889,8 @@ def fake_prepare_action_video(*args, **kwargs):
         assert output.output == {}
         assert output.custom_output["action_only_output"] is True
         assert output.custom_output["action"].shape == (1, 2, 2)
-        assert output.custom_output["actions"].shape == (1, 2)
+        assert output.custom_output["robolab_policy_inputs"] is inputs
+        assert "actions" not in output.custom_output
 
     @pytest.mark.parametrize(
         ("prompt", "sampling_params", "message"),
diff --git a/tests/diffusion/test_diffusion_plugin_hooks.py b/tests/diffusion/test_diffusion_plugin_hooks.py
index 72426693917..2a9a0a82109 100644
--- a/tests/diffusion/test_diffusion_plugin_hooks.py
+++ b/tests/diffusion/test_diffusion_plugin_hooks.py
@@ -15,6 +15,7 @@
 import pytest
 
 from vllm_omni.diffusion.registry import (
+    _DIFFUSION_ACTION_POST_PROCESS_FUNCS,
     _DIFFUSION_MODELS,
     _DIFFUSION_POST_PROCESS_FUNCS,
     _DIFFUSION_PRE_PROCESS_FUNCS,
@@ -59,6 +60,7 @@ 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()
         yield
         _DIFFUSION_MODELS.clear()
         _DIFFUSION_MODELS.update(original_models)
@@ -66,6 +68,8 @@ def cleanup_registry(self):
         _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)
 
     def test_register_new_model(self):
         """Test registering a new diffusion model with pre/post process functions."""
@@ -75,6 +79,7 @@ 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",
         )
         assert "TestPipeline" in _DIFFUSION_MODELS
         assert _DIFFUSION_MODELS["TestPipeline"] == (
@@ -84,6 +89,7 @@ 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"
 
 
 class TestWorkerUsesHook:
diff --git a/tests/diffusion/test_diffusion_scheduler.py b/tests/diffusion/test_diffusion_scheduler.py
index 1b5e813ee32..f5a4f7f7717 100644
--- a/tests/diffusion/test_diffusion_scheduler.py
+++ b/tests/diffusion/test_diffusion_scheduler.py
@@ -607,13 +607,17 @@ async def test_step_action_only_flag_skips_postprocess(self, mocker: MockerFixtu
         )
         engine.pre_process_func = None
         engine.post_process_func = mocker.Mock(side_effect=AssertionError("postprocess should be skipped"))
+        engine.action_post_process_func = mocker.Mock(return_value=torch.tensor([[3.0, 4.0]]))
         engine._post_process_accepts_sampling_params = False
+        engine._action_post_process_accepts_custom_output = True
+        engine._action_post_process_accepts_sampling_params = False
         engine._check_and_start_background_loop = mocker.AsyncMock()
+        raw_action = torch.tensor([[1.0, 2.0]])
         engine.async_add_req_and_wait_for_response = mocker.AsyncMock(
             return_value=DiffusionOutput(
                 output={},
                 custom_output={
-                    "actions": torch.tensor([[1.0, 2.0]]),
+                    "action": raw_action,
                     "action_only_output": True,
                 },
             )
@@ -629,8 +633,11 @@ async def test_step_action_only_flag_skips_postprocess(self, mocker: MockerFixtu
         outputs = await engine.step(request)
 
         engine.post_process_func.assert_not_called()
+        engine.action_post_process_func.assert_called_once()
+        assert engine.action_post_process_func.call_args.args[0] is raw_action
+        assert "custom_output" in engine.action_post_process_func.call_args.kwargs
         assert outputs[0].images == []
-        torch.testing.assert_close(outputs[0].multimodal_output["actions"], torch.tensor([[1.0, 2.0]]))
+        torch.testing.assert_close(outputs[0].multimodal_output["actions"], torch.tensor([[3.0, 4.0]]))
 
 
 class TestStepScheduler:
diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py
index 196be527479..39d51186caa 100644
--- a/vllm_omni/diffusion/diffusion_engine.py
+++ b/vllm_omni/diffusion/diffusion_engine.py
@@ -27,6 +27,7 @@
 from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
 from vllm_omni.diffusion.registry import (
     DiffusionModelRegistry,
+    get_diffusion_action_post_process_func,
     get_diffusion_post_process_func,
     get_diffusion_pre_process_func,
 )
@@ -84,6 +85,15 @@ def supports_audio_output(model_class_name: str) -> bool:
     return bool(getattr(model_cls, "support_audio_output", False))
 
 
+def _func_accepts_parameter(func: object | None, parameter_name: str) -> bool:
+    if func is None:
+        return False
+    parameters = inspect.signature(func).parameters
+    return parameter_name in parameters or any(
+        parameter.kind == inspect.Parameter.VAR_KEYWORD for parameter in parameters.values()
+    )
+
+
 def _move_tensor_tree_to_cpu(value: object) -> object:
     if isinstance(value, torch.Tensor):
         return value.cpu() if value.device.type != "cpu" else value
@@ -149,12 +159,16 @@ def __init__(
         self.od_config = od_config
 
         self.post_process_func = get_diffusion_post_process_func(od_config)
+        self.action_post_process_func = get_diffusion_action_post_process_func(od_config)
         self.pre_process_func = get_diffusion_pre_process_func(od_config)
         # Cache whether the model-specific postprocess accepts request-level
         # sampling params so step() can support both legacy and extended hooks.
-        self._post_process_accepts_sampling_params = bool(
-            self.post_process_func is not None
-            and "sampling_params" in inspect.signature(self.post_process_func).parameters
+        self._post_process_accepts_sampling_params = _func_accepts_parameter(self.post_process_func, "sampling_params")
+        self._action_post_process_accepts_sampling_params = _func_accepts_parameter(
+            self.action_post_process_func, "sampling_params"
+        )
+        self._action_post_process_accepts_custom_output = _func_accepts_parameter(
+            self.action_post_process_func, "custom_output"
         )
 
         executor_class = DiffusionExecutor.get_class(od_config)
@@ -289,6 +303,17 @@ async def step(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]:
             outputs = outputs.get("video", outputs)
         if action_payload is None:
             action_payload = custom_output.get("actions")
+        action_post_process_func = getattr(self, "action_post_process_func", None)
+        if action_post_process_func is not None:
+            raw_action_payload = custom_output.get("action", action_payload)
+            if raw_action_payload is not None:
+                action_kwargs: dict[str, Any] = {}
+                if getattr(self, "_action_post_process_accepts_custom_output", False):
+                    action_kwargs["custom_output"] = custom_output
+                if getattr(self, "_action_post_process_accepts_sampling_params", False):
+                    action_kwargs["sampling_params"] = request.sampling_params
+                action_payload = action_post_process_func(raw_action_payload, **action_kwargs)
+                custom_output["actions"] = action_payload
         postprocess_time = time.perf_counter() - postprocess_start_time
         logger.debug("Post-processing completed in %.4f seconds", postprocess_time)
 
diff --git a/vllm_omni/diffusion/models/cosmos3/__init__.py b/vllm_omni/diffusion/models/cosmos3/__init__.py
index 6df062b5c0d..43be0c37aa7 100644
--- a/vllm_omni/diffusion/models/cosmos3/__init__.py
+++ b/vllm_omni/diffusion/models/cosmos3/__init__.py
@@ -3,6 +3,7 @@
 
 from .pipeline_cosmos3 import (
     Cosmos3OmniDiffusersPipeline,
+    get_cosmos3_action_post_process_func,
     get_cosmos3_post_process_func,
     get_cosmos3_pre_process_func,
 )
@@ -10,6 +11,7 @@
 
 __all__ = [
     "Cosmos3OmniDiffusersPipeline",
+    "get_cosmos3_action_post_process_func",
     "get_cosmos3_post_process_func",
     "get_cosmos3_pre_process_func",
     "Cosmos3VFMTransformer",
diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py
index fa0d8a94917..92897d35192 100644
--- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py
+++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py
@@ -20,16 +20,12 @@
 import math
 import os
 import time
-import zlib
 from collections.abc import Iterable
-from dataclasses import dataclass
-from importlib import import_module
 from typing import Any, ClassVar
 
 import numpy as np
 import PIL.Image
 import torch
-import torch.nn.functional as F
 from diffusers import UniPCMultistepScheduler
 from diffusers.utils.torch_utils import randn_tensor
 from diffusers.video_processor import VideoProcessor
@@ -69,16 +65,43 @@
     resolve_domain_id,
 )
 from .transformer_cosmos3 import Cosmos3VFMTransformer, resolve_sound_gen
+from .utils import (
+    COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION,
+    COSMOS3_VAE_TEMPORAL_COMPRESSION,
+    ROBOLAB_CONCAT_VIEW_DESCRIPTION,
+    ROBOLAB_DEFAULT_ACTION_CHUNK_SIZE,
+    ROBOLAB_DEFAULT_ACTION_SPACE,
+    ROBOLAB_DEFAULT_CONDITIONING_FPS,
+    ROBOLAB_DEFAULT_DOMAIN_NAME,
+    ROBOLAB_DEFAULT_FLOW_SHIFT,
+    ROBOLAB_DEFAULT_GUIDANCE_SCALE,
+    ROBOLAB_DEFAULT_IMAGE_HEIGHT,
+    ROBOLAB_DEFAULT_IMAGE_WIDTH,
+    ROBOLAB_DEFAULT_NUM_INFERENCE_STEPS,
+    ROBOLAB_DEFAULT_RAW_ACTION_DIM,
+    ROBOLAB_DEFAULT_RESOLUTION,
+    ROBOLAB_MIDTRAIN_RAW_ACTION_DIM,
+    RoboLabPolicyInputs,
+    build_abs_pose_from_components,
+    build_robolab_unipc_scheduler,
+    condition_pixel_frame_count,
+    convert_midtrain_rotation,
+    ensure_2d_float_array,
+    ensure_gripper_array,
+    extract_robolab_image,
+    extract_robolab_prompt_image,
+    lazy_action_transform_pipeline,
+    next_robolab_seed,
+    normalize_condition_frame_indexes_vision,
+    normalize_condition_video_keep,
+    normalize_robolab_action_space,
+    pose_abs_to_rel,
+    postprocess_robolab_action,
+    resize_rgb_uint8,
+)
 
 logger = init_logger(__name__)
 
-COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION = (0, 1)
-COSMOS3_DEFAULT_CONDITION_VIDEO_KEEP = "first"
-# Mirrors the WAN VAE's temporal compression. Authoritative value is
-# ``self.vae.config.scale_factor_temporal`` at runtime; this constant exists so
-# off-line / API code that runs before the pipeline is constructed can compute
-# pixel-frame budgets without instantiating the VAE.
-COSMOS3_VAE_TEMPORAL_COMPRESSION = 4
 COSMOS3_DEFAULT_CONDITION_PIXEL_FRAMES = (
     max(COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION) * COSMOS3_VAE_TEMPORAL_COMPRESSION + 1
 )
@@ -105,229 +128,12 @@
 COSMOS3_T2I_DEFAULT_FLOW_SHIFT = 3.0
 COSMOS3_T2I_DEFAULT_GUIDANCE_INTERVAL: tuple[float, float] = (400.0, 1000.0)
 
-ROBOLAB_DEFAULT_CONDITIONING_FPS = 15.0
-ROBOLAB_DEFAULT_ACTION_CHUNK_SIZE = 32
-ROBOLAB_DEFAULT_IMAGE_HEIGHT = 540
-ROBOLAB_DEFAULT_IMAGE_WIDTH = 640
-ROBOLAB_DEFAULT_RAW_ACTION_DIM = 8
-ROBOLAB_DEFAULT_DOMAIN_NAME = "droid_lerobot"
-ROBOLAB_DEFAULT_RESOLUTION = "480"
-ROBOLAB_DEFAULT_GUIDANCE_SCALE = 3.0
-ROBOLAB_DEFAULT_NUM_INFERENCE_STEPS = 4
-ROBOLAB_DEFAULT_FLOW_SHIFT = 5.0
-ROBOLAB_DEFAULT_SEED = 0
-ROBOLAB_DEFAULT_ACTION_SPACE = "joint_pos"
-ROBOLAB_MIDTRAIN_RAW_ACTION_DIM = 10
-ROBOLAB_MIDTRAIN_POSE_ACTION_DIM = 9  # xyz position + rot6d orientation
-ROBOLAB_CONCAT_VIEW_DESCRIPTION = (
-    "The top row is from the wrist-mounted camera. "
-    "The bottom row contains two horizontally concatenated third-person perspective views of the scene from opposite "
-    "sides, with the robot visible."
-)
-
 # Truncation cap on the prompt token count (shared by T2I and T2V).  Prompts
 # are tokenized to their natural length (no padding); this only bounds the
 # UND pathway / GEN cross-attention cost for pathologically long prompts.
 COSMOS3_DEFAULT_MAX_SEQUENCE_LENGTH = 4096
 
 
-@dataclass(frozen=True)
-class RoboLabPolicyInputs:
-    prompt: str
-    video_tensor: torch.Tensor
-    action_tensor: torch.Tensor
-    action_condition_indexes: list[int]
-    action_start_frame_offset: int
-    raw_action_dim: int
-    domain_id: int
-    fps: float
-    height: int
-    width: int
-    image_size: Any
-    num_frames: int
-    num_inference_steps: int
-    guidance_scale: float
-    flow_shift: float
-    seed: int
-    history_length: int
-    action_space: str
-    observation: dict[str, Any]
-
-
-def _normalize_condition_frame_indexes_vision(value: Any) -> tuple[int, ...]:
-    """Normalize Cosmos3 vision-conditioning latent frame indexes."""
-    if value is None:
-        return COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION
-    if isinstance(value, str):
-        value = [item.strip() for item in value.split(",") if item.strip()]
-    elif isinstance(value, int):
-        value = [value]
-
-    if not isinstance(value, Iterable):
-        raise TypeError(
-            "Cosmos3 condition_frame_indexes_vision must be an int, comma-separated string, "
-            f"or iterable of ints; got {type(value)!r}."
-        )
-
-    indexes = tuple(sorted({int(index) for index in value}))
-    if not indexes:
-        raise ValueError("Cosmos3 condition_frame_indexes_vision must contain at least one index.")
-    if any(index < 0 for index in indexes):
-        raise ValueError(f"Cosmos3 condition_frame_indexes_vision must be non-negative, got {indexes}.")
-    return indexes
-
-
-def _condition_pixel_frame_count(
-    condition_frame_indexes_vision: Iterable[int],
-    temporal_compression: int = COSMOS3_VAE_TEMPORAL_COMPRESSION,
-) -> int:
-    return max(condition_frame_indexes_vision) * int(temporal_compression) + 1
-
-
-def _normalize_condition_video_keep(value: Any) -> str:
-    keep = str(value or COSMOS3_DEFAULT_CONDITION_VIDEO_KEEP).strip().lower()
-    if keep not in {"first", "last"}:
-        raise ValueError("Cosmos3 condition_video_keep must be either 'first' or 'last'.")
-    return keep
-
-
-def _ensure_rgb_uint8_image(value: Any, key: str) -> np.ndarray:
-    image = np.asarray(value)
-    if image.ndim != 3 or image.shape[-1] != 3:
-        raise ValueError(f"{key!r} must have shape [H, W, 3], got {image.shape}.")
-    if image.dtype != np.uint8:
-        image = np.clip(image, 0, 255).astype(np.uint8)
-    return np.ascontiguousarray(image)
-
-
-def _ensure_2d_float_array(value: Any, key: str, width: int | None = None) -> np.ndarray:
-    array = np.asarray(value, dtype=np.float32)
-    if array.ndim == 1:
-        array = array[None, :]
-    if array.ndim != 2:
-        raise ValueError(f"{key!r} must have shape [T, D] or [D], got {array.shape}.")
-    if width is not None and array.shape[-1] != width:
-        raise ValueError(f"{key!r} must have width {width}, got {array.shape[-1]}.")
-    return np.ascontiguousarray(array)
-
-
-def _ensure_gripper_array(value: Any) -> np.ndarray:
-    array = np.asarray(value, dtype=np.float32)
-    if array.ndim == 0:
-        array = array.reshape(1, 1)
-    elif array.ndim == 1:
-        array = array[:, None]
-    if array.ndim != 2 or array.shape[-1] != 1:
-        raise ValueError(f"'observation/gripper_position' must have shape [T, 1], [T], or scalar, got {array.shape}.")
-    return np.ascontiguousarray(array)
-
-
-def _resize_rgb_uint8(image: np.ndarray, size: tuple[int, int]) -> np.ndarray:
-    tensor = torch.from_numpy(image).permute(2, 0, 1).unsqueeze(0).float()
-    resized = F.interpolate(tensor, size=size, mode="bilinear", align_corners=False)
-    return np.clip(np.round(resized.squeeze(0).permute(1, 2, 0).numpy()), 0, 255).astype(np.uint8)
-
-
-def _compose_robolab_views(obs: dict[str, Any]) -> np.ndarray | None:
-    required_keys = (
-        "observation/wrist_image_left",
-        "observation/exterior_image_1_left",
-        "observation/exterior_image_2_left",
-    )
-    if not all(key in obs for key in required_keys):
-        return None
-
-    wrist = _ensure_rgb_uint8_image(obs["observation/wrist_image_left"], "observation/wrist_image_left")
-    left_raw = _ensure_rgb_uint8_image(obs["observation/exterior_image_1_left"], "observation/exterior_image_1_left")
-    right_raw = _ensure_rgb_uint8_image(obs["observation/exterior_image_2_left"], "observation/exterior_image_2_left")
-    half_h, half_w = wrist.shape[0] // 2, wrist.shape[1] // 2
-    left = _resize_rgb_uint8(left_raw, (half_h, half_w))
-    right = _resize_rgb_uint8(right_raw, (half_h, half_w))
-    return np.concatenate([wrist, np.concatenate([left, right], axis=1)], axis=0)
-
-
-def _extract_robolab_image(obs: dict[str, Any]) -> np.ndarray:
-    if "observation/image" in obs:
-        return _ensure_rgb_uint8_image(obs["observation/image"], "observation/image")
-    image = _compose_robolab_views(obs)
-    if image is not None:
-        return image
-    raise ValueError("Observation must contain 'observation/image' or RoboLab wrist/exterior image keys.")
-
-
-def _lazy_import(module_name: str, symbol_name: str, error_message: str):
-    try:
-        module = import_module(module_name)
-    except ModuleNotFoundError as exc:
-        raise ModuleNotFoundError(error_message) from exc
-    return getattr(module, symbol_name)
-
-
-def _lazy_action_transform_pipeline(max_action_dim: int):
-    ActionTransformPipeline = _lazy_import(
-        "cosmos_framework.data.vfm.action.transforms",
-        "ActionTransformPipeline",
-        "Cosmos3 RoboLab policy serving requires cosmos_framework on PYTHONPATH so the "
-        "golden ActionTransformPipeline can be reused.",
-    )
-    return ActionTransformPipeline(max_action_dim=max_action_dim, cfg_dropout_rate=0.0)
-
-
-def _build_robolab_unipc_scheduler(num_steps: int, shift: float, device: torch.device):
-    FlowUniPCMultistepScheduler = _lazy_import(
-        "cosmos_framework.model.vfm.diffusion.samplers.fm_solvers_unipc",
-        "FlowUniPCMultistepScheduler",
-        (
-            "Cosmos3 RoboLab policy serving requires cosmos_framework on PYTHONPATH so the "
-            "golden FlowUniPCMultistepScheduler can be reused."
-        ),
-    )
-
-    scheduler = FlowUniPCMultistepScheduler(
-        num_train_timesteps=1000,
-        shift=1.0,
-        use_dynamic_shifting=False,
-    )
-    scheduler.set_timesteps(num_steps, device=device, shift=float(shift))
-    return scheduler
-
-
-def _convert_midtrain_rotation(value: Any, src: str, dst: str) -> np.ndarray:
-    convert_rotation = _lazy_import(
-        "cosmos_framework.data.vfm.action.pose_utils",
-        "convert_rotation",
-        "Cosmos3 RoboLab midtrain action serving requires cosmos_framework pose_utils on PYTHONPATH.",
-    )
-    return convert_rotation(value, src, dst)
-
-
-def _pose_abs_to_rel(*args, **kwargs) -> np.ndarray:
-    pose_abs_to_rel = _lazy_import(
-        "cosmos_framework.data.vfm.action.pose_utils",
-        "pose_abs_to_rel",
-        "Cosmos3 RoboLab midtrain action serving requires cosmos_framework pose_utils on PYTHONPATH.",
-    )
-    return pose_abs_to_rel(*args, **kwargs)
-
-
-def _pose_rel_to_abs(*args, **kwargs) -> np.ndarray:
-    pose_rel_to_abs = _lazy_import(
-        "cosmos_framework.data.vfm.action.pose_utils",
-        "pose_rel_to_abs",
-        "Cosmos3 RoboLab midtrain action serving requires cosmos_framework pose_utils on PYTHONPATH.",
-    )
-    return pose_rel_to_abs(*args, **kwargs)
-
-
-def _build_abs_pose_from_components(*args, **kwargs) -> np.ndarray:
-    build_abs_pose_from_components = _lazy_import(
-        "cosmos_framework.data.vfm.action.pose_utils",
-        "build_abs_pose_from_components",
-        "Cosmos3 RoboLab midtrain action serving requires cosmos_framework pose_utils on PYTHONPATH.",
-    )
-    return build_abs_pose_from_components(*args, **kwargs)
-
-
 # ---------------------------------------------------------------------------
 # Post-process function (registered in registry.py)
 # ---------------------------------------------------------------------------
@@ -549,16 +355,16 @@ def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest:
             else:
                 assert raw_video_frames is not None
                 extra = _extra_args(request)
-                condition_frame_indexes_vision = _normalize_condition_frame_indexes_vision(
+                condition_frame_indexes_vision = normalize_condition_frame_indexes_vision(
                     extra.get(
                         "condition_frame_indexes_vision",
                         prompt.get("condition_frame_indexes_vision"),
                     )
                 )
-                keep = _normalize_condition_video_keep(
+                keep = normalize_condition_video_keep(
                     extra.get("condition_video_keep", prompt.get("condition_video_keep"))
                 )
-                max_frames = _condition_pixel_frame_count(condition_frame_indexes_vision)
+                max_frames = condition_pixel_frame_count(condition_frame_indexes_vision)
                 prompt["additional_information"]["preprocessed_video"] = _preprocess_condition_video(
                     raw_video_frames,
                     int(target_h),
@@ -675,6 +481,19 @@ def post_process_func(
     return post_process_func
 
 
+def get_cosmos3_action_post_process_func(od_config: OmniDiffusionConfig):
+    del od_config
+
+    def action_post_process_func(action: Any, custom_output: dict[str, Any] | None = None, sampling_params=None):
+        del sampling_params
+        inputs = custom_output.get("robolab_policy_inputs") if isinstance(custom_output, dict) else None
+        if isinstance(inputs, RoboLabPolicyInputs):
+            return postprocess_robolab_action(action, inputs)
+        return action
+
+    return action_post_process_func
+
+
 # ---------------------------------------------------------------------------
 # Pipeline
 # ---------------------------------------------------------------------------
@@ -733,11 +552,11 @@ def reference_video_decode_spec(
                     return ReferenceVideoDecodeSpec(max_frames=max_frames, keep="first")
             return ReferenceVideoDecodeSpec(max_frames=None, keep="first")
 
-        condition_indexes = _normalize_condition_frame_indexes_vision(extra_args.get("condition_frame_indexes_vision"))
-        max_frames = _condition_pixel_frame_count(condition_indexes)
+        condition_indexes = normalize_condition_frame_indexes_vision(extra_args.get("condition_frame_indexes_vision"))
+        max_frames = condition_pixel_frame_count(condition_indexes)
         if num_frames is not None:
             max_frames = min(max_frames, int(num_frames))
-        keep = _normalize_condition_video_keep(extra_args.get("condition_video_keep"))
+        keep = normalize_condition_video_keep(extra_args.get("condition_video_keep"))
         return ReferenceVideoDecodeSpec(max_frames=max_frames, keep=keep)
 
     def __init__(
@@ -1060,45 +879,12 @@ def _get_sp_param(sp: OmniDiffusionSamplingParams, key: str, default: Any = None
             return val
         return default
 
-    @staticmethod
-    def _normalize_robolab_action_space(value: Any) -> str:
-        action_space = str(value or ROBOLAB_DEFAULT_ACTION_SPACE).strip().lower()
-        aliases = {
-            "jointpos": "joint_pos",
-            "joint_pos": "joint_pos",
-            "abs_ik": "midtrain",
-            "midtrain": "midtrain",
-        }
-        if action_space not in aliases:
-            raise ValueError(
-                f"Unsupported RoboLab action_space={value!r}; expected joint_pos/jointpos or midtrain/abs_ik."
-            )
-        return aliases[action_space]
-
     def _get_robolab_transform(self):
         if self._robolab_transform is None:
             action_dim = int(getattr(self.transformer, "action_dim", 64))
-            self._robolab_transform = _lazy_action_transform_pipeline(action_dim)
+            self._robolab_transform = lazy_action_transform_pipeline(action_dim)
         return self._robolab_transform
 
-    def _next_robolab_seed(self, extra: dict[str, Any], obs: dict[str, Any], request_id: str | None) -> int:
-        base_seed = int(extra.get("robolab_seed") or ROBOLAB_DEFAULT_SEED)
-        if self._truthy(extra.get("deterministic_seed")):
-            return base_seed
-        explicit_seed = extra.get("seed")
-        if explicit_seed is not None:
-            return int(explicit_seed)
-        seed_key = "|".join(
-            str(part)
-            for part in (
-                base_seed,
-                extra.get("session_id", ""),
-                request_id or "",
-                obs.get("prompt", ""),
-            )
-        )
-        return zlib.crc32(seed_key.encode("utf-8")) & 0x7FFFFFFF
-
     def _build_robolab_policy_inputs(
         self,
         sp: OmniDiffusionSamplingParams,
@@ -1129,7 +915,7 @@ def extra_param_alias(primary_key: str, alias_key: str, default: Any) -> Any:
             value = extra.get(alias_key)
             return default if value is None else value
 
-        action_space = self._normalize_robolab_action_space(extra_param("action_space", ROBOLAB_DEFAULT_ACTION_SPACE))
+        action_space = normalize_robolab_action_space(extra_param("action_space", ROBOLAB_DEFAULT_ACTION_SPACE))
         action_chunk_size = int(extra_param("action_chunk_size", ROBOLAB_DEFAULT_ACTION_CHUNK_SIZE))
         raw_action_dim_default = (
             ROBOLAB_DEFAULT_RAW_ACTION_DIM if action_space == "joint_pos" else ROBOLAB_MIDTRAIN_RAW_ACTION_DIM
@@ -1152,13 +938,13 @@ def extra_param_alias(primary_key: str, alias_key: str, default: Any) -> Any:
             raise ValueError(f"RoboLab raw_action_dim must be positive, got {raw_action_dim}.")
 
         try:
-            image = _extract_robolab_image(obs)
+            image = extract_robolab_image(obs)
         except ValueError as exc:
-            image = self._extract_robolab_prompt_image(prompt_data)
+            image = extract_robolab_prompt_image(prompt_data)
             if image is None:
                 raise exc
         if image.shape[:2] != (image_h, image_w):
-            image = _resize_rgb_uint8(image, (image_h, image_w))
+            image = resize_rgb_uint8(image, (image_h, image_w))
 
         t_frames = action_chunk_size + 1
         video = torch.zeros((3, t_frames, image_h, image_w), dtype=torch.uint8)
@@ -1168,10 +954,10 @@ def extra_param_alias(primary_key: str, alias_key: str, default: Any) -> Any:
         action = torch.zeros((action_chunk_size + use_state_rows, raw_action_dim), dtype=torch.float32)
         history_action = None
         num_history_rows = history_length - use_state_rows
-        gripper_position = 1.0 - _ensure_gripper_array(obs["observation/gripper_position"])
+        gripper_position = 1.0 - ensure_gripper_array(obs["observation/gripper_position"])
 
         if action_space == "joint_pos":
-            joint_position = _ensure_2d_float_array(obs["observation/joint_position"], "observation/joint_position", 7)
+            joint_position = ensure_2d_float_array(obs["observation/joint_position"], "observation/joint_position", 7)
             if use_state:
                 action[0] = torch.from_numpy(np.concatenate((joint_position[-1], gripper_position[-1])))
             if num_history_rows > 0:
@@ -1183,16 +969,16 @@ def extra_param_alias(primary_key: str, alias_key: str, default: Any) -> Any:
                 )
                 history_action = torch.from_numpy(history_np).float()
         else:
-            eef_pos = _ensure_2d_float_array(obs["observation/eef_pos"], "observation/eef_pos", 3)
-            eef_quat = _ensure_2d_float_array(obs["observation/eef_quat"], "observation/eef_quat", 4)
+            eef_pos = ensure_2d_float_array(obs["observation/eef_pos"], "observation/eef_pos", 3)
+            eef_quat = ensure_2d_float_array(obs["observation/eef_quat"], "observation/eef_quat", 4)
             if use_state:
-                rot6d = _convert_midtrain_rotation(eef_quat[-1], "quat_xyzw", "rot6d")
+                rot6d = convert_midtrain_rotation(eef_quat[-1], "quat_xyzw", "rot6d")
                 action[0] = torch.from_numpy(np.concatenate((eef_pos[-1], rot6d, gripper_position[-1])))
             if num_history_rows > 0:
                 if len(eef_pos) < num_history_rows + 1 or len(eef_quat) < num_history_rows + 1:
                     raise ValueError("Not enough eef_pos/eef_quat rows for requested history_length.")
-                poses_abs = _build_abs_pose_from_components(eef_pos, eef_quat, "quat_xyzw")
-                poses_rel = _pose_abs_to_rel(poses_abs, rotation_format="rot6d", pose_convention="backward_framewise")
+                poses_abs = build_abs_pose_from_components(eef_pos, eef_quat, "quat_xyzw")
+                poses_rel = pose_abs_to_rel(poses_abs, rotation_format="rot6d", pose_convention="backward_framewise")
                 history_np = np.concatenate(
                     [poses_rel[-num_history_rows:], gripper_position[-num_history_rows:]],
                     axis=-1,
@@ -1240,24 +1026,12 @@ def extra_param_alias(primary_key: str, alias_key: str, default: Any) -> Any:
             ),
             guidance_scale=float(extra_param_alias("guidance_scale", "guidance", ROBOLAB_DEFAULT_GUIDANCE_SCALE)),
             flow_shift=float(extra_param_alias("flow_shift", "shift", ROBOLAB_DEFAULT_FLOW_SHIFT)),
-            seed=self._next_robolab_seed(extra, obs, request_id),
+            seed=next_robolab_seed(extra, obs, request_id),
             history_length=history_length,
             action_space=action_space,
             observation=obs,
         )
 
-    @staticmethod
-    def _extract_robolab_prompt_image(prompt_data: Any | None) -> np.ndarray | None:
-        if not isinstance(prompt_data, dict):
-            return None
-        multi_modal_data = prompt_data.get("multi_modal_data", {}) or {}
-        image = multi_modal_data.get("image")
-        if image is None:
-            return None
-        if isinstance(image, PIL.Image.Image):
-            return np.asarray(image.convert("RGB"))
-        return _ensure_rgb_uint8_image(image, "multi_modal_data.image")
-
     @staticmethod
     def _build_action_condition_mask_from_indexes(
         indexes: list[int],
@@ -1273,63 +1047,6 @@ def _build_action_condition_mask_from_indexes(
             mask[:, idx, :] = 1.0
         return mask
 
-    @staticmethod
-    def _log_robolab_action_summary(label: str, value: Any) -> None:
-        if not _is_rank_zero():
-            return
-        if isinstance(value, torch.Tensor):
-            array = value.detach().float().cpu().numpy()
-        else:
-            array = np.asarray(value, dtype=np.float32)
-        finite = np.isfinite(array)
-        if finite.any():
-            finite_min = float(array[finite].min())
-            finite_max = float(array[finite].max())
-        else:
-            finite_min = None
-            finite_max = None
-        if array.ndim == 0:
-            head = array.reshape(1).tolist()
-        else:
-            head = array.reshape(-1, array.shape[-1])[:3].tolist()
-        logger.info(
-            "RoboLab action summary %s: shape=%s nan=%d finite=%d finite_min=%s finite_max=%s head=%s",
-            label,
-            tuple(array.shape),
-            int(np.isnan(array).sum()),
-            int(finite.sum()),
-            finite_min,
-            finite_max,
-            head,
-        )
-
-    def _postprocess_robolab_action(self, action: torch.Tensor, inputs: RoboLabPolicyInputs) -> np.ndarray:
-        action_np = action[0].float().cpu().numpy()
-        self._log_robolab_action_summary("raw_model_action", action_np)
-        history_length = int(inputs.history_length)
-        action_np = action_np[history_length:]
-        action_np[:, -1] = 1.0 - action_np[:, -1]
-
-        if inputs.action_space == "midtrain":
-            obs = inputs.observation
-            eef_pos = _ensure_2d_float_array(obs["observation/eef_pos"], "observation/eef_pos", 3)
-            eef_quat = _ensure_2d_float_array(obs["observation/eef_quat"], "observation/eef_quat", 4)
-            initial_pose = np.eye(4, dtype=np.float32)
-            initial_pose[:3, :3] = _convert_midtrain_rotation(eef_quat[-1], "quat_xyzw", "matrix")
-            initial_pose[:3, 3] = eef_pos[-1]
-            abs_pose = _pose_rel_to_abs(
-                action_np[:, :ROBOLAB_MIDTRAIN_POSE_ACTION_DIM],
-                rotation_format="rot6d",
-                pose_convention="backward_framewise",
-                initial_pose=initial_pose,
-            )
-            position = abs_pose[1:, :3, 3]
-            quat_xyzw = _convert_midtrain_rotation(abs_pose[1:, :3, :3], "matrix", "quat_xyzw")
-            action_np = np.concatenate([position, quat_xyzw, action_np[:, ROBOLAB_MIDTRAIN_POSE_ACTION_DIM:]], axis=-1)
-
-        self._log_robolab_action_summary("postprocessed_robolab_action", action_np)
-        return np.asarray(action_np, dtype=np.float32)
-
     def _forward_robolab_policy(
         self,
         sp: OmniDiffusionSamplingParams,
@@ -1426,7 +1143,7 @@ def _forward_robolab_policy(
             action_fps=float(self._get_sp_param(sp, "action_fps", frame_rate) or frame_rate),
         )
 
-        scheduler = _build_robolab_unipc_scheduler(num_inference_steps, flow_shift_target, self.device)
+        scheduler = build_robolab_unipc_scheduler(num_inference_steps, flow_shift_target, self.device)
         _, action_latents = self.diffuse(
             latents=latents,
             timesteps=scheduler.timesteps,
@@ -1458,7 +1175,7 @@ def _forward_robolab_policy(
             "action_mode": action_mode,
             "domain_id": domain_id,
             "action_only_output": True,
-            "actions": self._postprocess_robolab_action(action, inputs),
+            "robolab_policy_inputs": inputs,
         }
         return DiffusionOutput(output={}, custom_output=custom_action_output)
 
@@ -2004,7 +1721,7 @@ def _prepare_latents_v2v(
         T_lat = (num_frames - 1) // self.vae_scale_factor_temporal + 1
         H_lat = video_tensor.shape[-2] // self.vae_scale_factor_spatial
         W_lat = video_tensor.shape[-1] // self.vae_scale_factor_spatial
-        indexes = _normalize_condition_frame_indexes_vision(condition_frame_indexes_vision)
+        indexes = normalize_condition_frame_indexes_vision(condition_frame_indexes_vision)
         out_of_range = [index for index in indexes if index >= T_lat]
         if out_of_range:
             raise ValueError(
@@ -2018,7 +1735,7 @@ def _prepare_latents_v2v(
             device=self.device,
             dtype=self.dtype,
         )
-        condition_pixel_frames = _condition_pixel_frame_count(indexes, self.vae_scale_factor_temporal)
+        condition_pixel_frames = condition_pixel_frame_count(indexes, self.vae_scale_factor_temporal)
         condition_video = video_tensor[:, :, :condition_pixel_frames]
         if condition_video.shape[2] < condition_pixel_frames:
             pad = condition_video[:, :, -1:].repeat(1, 1, condition_pixel_frames - condition_video.shape[2], 1, 1)
@@ -2678,7 +2395,7 @@ def forward(
             )
             image_latent = condition_latents[:, :, 0:1]
         elif is_v2v:
-            condition_frame_indexes_vision = _normalize_condition_frame_indexes_vision(
+            condition_frame_indexes_vision = normalize_condition_frame_indexes_vision(
                 self._get_sp_param(
                     sp,
                     "condition_frame_indexes_vision",
diff --git a/vllm_omni/diffusion/models/cosmos3/utils.py b/vllm_omni/diffusion/models/cosmos3/utils.py
new file mode 100644
index 00000000000..6f91674e644
--- /dev/null
+++ b/vllm_omni/diffusion/models/cosmos3/utils.py
@@ -0,0 +1,348 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import zlib
+from collections.abc import Iterable
+from dataclasses import dataclass
+from importlib import import_module
+from typing import Any
+
+import numpy as np
+import PIL.Image
+import torch
+import torch.nn.functional as F
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.models.progress_bar import _is_rank_zero
+
+logger = init_logger(__name__)
+
+COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION = (0, 1)
+COSMOS3_DEFAULT_CONDITION_VIDEO_KEEP = "first"
+# Mirrors the WAN VAE's temporal compression. Authoritative value is
+# ``self.vae.config.scale_factor_temporal`` at runtime; this constant exists so
+# off-line / API code that runs before the pipeline is constructed can compute
+# pixel-frame budgets without instantiating the VAE.
+COSMOS3_VAE_TEMPORAL_COMPRESSION = 4
+
+ROBOLAB_DEFAULT_CONDITIONING_FPS = 15.0
+ROBOLAB_DEFAULT_ACTION_CHUNK_SIZE = 32
+ROBOLAB_DEFAULT_IMAGE_HEIGHT = 540
+ROBOLAB_DEFAULT_IMAGE_WIDTH = 640
+ROBOLAB_DEFAULT_RAW_ACTION_DIM = 8
+ROBOLAB_DEFAULT_DOMAIN_NAME = "droid_lerobot"
+ROBOLAB_DEFAULT_RESOLUTION = "480"
+ROBOLAB_DEFAULT_GUIDANCE_SCALE = 3.0
+ROBOLAB_DEFAULT_NUM_INFERENCE_STEPS = 4
+ROBOLAB_DEFAULT_FLOW_SHIFT = 5.0
+ROBOLAB_DEFAULT_SEED = 0
+ROBOLAB_DEFAULT_ACTION_SPACE = "joint_pos"
+ROBOLAB_MIDTRAIN_RAW_ACTION_DIM = 10
+ROBOLAB_MIDTRAIN_POSE_ACTION_DIM = 9  # xyz position + rot6d orientation
+ROBOLAB_CONCAT_VIEW_DESCRIPTION = (
+    "The top row is from the wrist-mounted camera. "
+    "The bottom row contains two horizontally concatenated third-person perspective views of the scene from opposite "
+    "sides, with the robot visible."
+)
+
+
+@dataclass(frozen=True)
+class RoboLabPolicyInputs:
+    prompt: str
+    video_tensor: torch.Tensor
+    action_tensor: torch.Tensor
+    action_condition_indexes: list[int]
+    action_start_frame_offset: int
+    raw_action_dim: int
+    domain_id: int
+    fps: float
+    height: int
+    width: int
+    image_size: Any
+    num_frames: int
+    num_inference_steps: int
+    guidance_scale: float
+    flow_shift: float
+    seed: int
+    history_length: int
+    action_space: str
+    observation: dict[str, Any]
+
+
+def normalize_condition_frame_indexes_vision(value: Any) -> tuple[int, ...]:
+    """Normalize Cosmos3 vision-conditioning latent frame indexes."""
+    if value is None:
+        return COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION
+    if isinstance(value, str):
+        value = [item.strip() for item in value.split(",") if item.strip()]
+    elif isinstance(value, int):
+        value = [value]
+
+    if not isinstance(value, Iterable):
+        raise TypeError(
+            "Cosmos3 condition_frame_indexes_vision must be an int, comma-separated string, "
+            f"or iterable of ints; got {type(value)!r}."
+        )
+
+    indexes = tuple(sorted({int(index) for index in value}))
+    if not indexes:
+        raise ValueError("Cosmos3 condition_frame_indexes_vision must contain at least one index.")
+    if any(index < 0 for index in indexes):
+        raise ValueError(f"Cosmos3 condition_frame_indexes_vision must be non-negative, got {indexes}.")
+    return indexes
+
+
+def condition_pixel_frame_count(
+    condition_frame_indexes_vision: Iterable[int],
+    temporal_compression: int = COSMOS3_VAE_TEMPORAL_COMPRESSION,
+) -> int:
+    return max(condition_frame_indexes_vision) * int(temporal_compression) + 1
+
+
+def normalize_condition_video_keep(value: Any) -> str:
+    keep = str(value or COSMOS3_DEFAULT_CONDITION_VIDEO_KEEP).strip().lower()
+    if keep not in {"first", "last"}:
+        raise ValueError("Cosmos3 condition_video_keep must be either 'first' or 'last'.")
+    return keep
+
+
+def normalize_robolab_action_space(value: Any) -> str:
+    action_space = str(value or ROBOLAB_DEFAULT_ACTION_SPACE).strip().lower()
+    aliases = {
+        "jointpos": "joint_pos",
+        "joint_pos": "joint_pos",
+        "abs_ik": "midtrain",
+        "midtrain": "midtrain",
+    }
+    if action_space not in aliases:
+        raise ValueError(f"Unsupported RoboLab action_space={value!r}; expected joint_pos/jointpos or midtrain/abs_ik.")
+    return aliases[action_space]
+
+
+def ensure_rgb_uint8_image(value: Any, key: str) -> np.ndarray:
+    image = np.asarray(value)
+    if image.ndim != 3 or image.shape[-1] != 3:
+        raise ValueError(f"{key!r} must have shape [H, W, 3], got {image.shape}.")
+    if image.dtype != np.uint8:
+        image = np.clip(image, 0, 255).astype(np.uint8)
+    return np.ascontiguousarray(image)
+
+
+def ensure_2d_float_array(value: Any, key: str, width: int | None = None) -> np.ndarray:
+    array = np.asarray(value, dtype=np.float32)
+    if array.ndim == 1:
+        array = array[None, :]
+    if array.ndim != 2:
+        raise ValueError(f"{key!r} must have shape [T, D] or [D], got {array.shape}.")
+    if width is not None and array.shape[-1] != width:
+        raise ValueError(f"{key!r} must have width {width}, got {array.shape[-1]}.")
+    return np.ascontiguousarray(array)
+
+
+def ensure_gripper_array(value: Any) -> np.ndarray:
+    array = np.asarray(value, dtype=np.float32)
+    if array.ndim == 0:
+        array = array.reshape(1, 1)
+    elif array.ndim == 1:
+        array = array[:, None]
+    if array.ndim != 2 or array.shape[-1] != 1:
+        raise ValueError(f"'observation/gripper_position' must have shape [T, 1], [T], or scalar, got {array.shape}.")
+    return np.ascontiguousarray(array)
+
+
+def resize_rgb_uint8(image: np.ndarray, size: tuple[int, int]) -> np.ndarray:
+    tensor = torch.from_numpy(image).permute(2, 0, 1).unsqueeze(0).float()
+    resized = F.interpolate(tensor, size=size, mode="bilinear", align_corners=False)
+    return np.clip(np.round(resized.squeeze(0).permute(1, 2, 0).numpy()), 0, 255).astype(np.uint8)
+
+
+def compose_robolab_views(obs: dict[str, Any]) -> np.ndarray | None:
+    required_keys = (
+        "observation/wrist_image_left",
+        "observation/exterior_image_1_left",
+        "observation/exterior_image_2_left",
+    )
+    if not all(key in obs for key in required_keys):
+        return None
+
+    wrist = ensure_rgb_uint8_image(obs["observation/wrist_image_left"], "observation/wrist_image_left")
+    left_raw = ensure_rgb_uint8_image(obs["observation/exterior_image_1_left"], "observation/exterior_image_1_left")
+    right_raw = ensure_rgb_uint8_image(obs["observation/exterior_image_2_left"], "observation/exterior_image_2_left")
+    half_h, half_w = wrist.shape[0] // 2, wrist.shape[1] // 2
+    left = resize_rgb_uint8(left_raw, (half_h, half_w))
+    right = resize_rgb_uint8(right_raw, (half_h, half_w))
+    return np.concatenate([wrist, np.concatenate([left, right], axis=1)], axis=0)
+
+
+def extract_robolab_image(obs: dict[str, Any]) -> np.ndarray:
+    if "observation/image" in obs:
+        return ensure_rgb_uint8_image(obs["observation/image"], "observation/image")
+    image = compose_robolab_views(obs)
+    if image is not None:
+        return image
+    raise ValueError("Observation must contain 'observation/image' or RoboLab wrist/exterior image keys.")
+
+
+def extract_robolab_prompt_image(prompt_data: Any | None) -> np.ndarray | None:
+    if not isinstance(prompt_data, dict):
+        return None
+    multi_modal_data = prompt_data.get("multi_modal_data", {}) or {}
+    image = multi_modal_data.get("image")
+    if image is None:
+        return None
+    if isinstance(image, PIL.Image.Image):
+        return np.asarray(image.convert("RGB"))
+    return ensure_rgb_uint8_image(image, "multi_modal_data.image")
+
+
+def lazy_import(module_name: str, symbol_name: str, error_message: str):
+    try:
+        module = import_module(module_name)
+    except ModuleNotFoundError as exc:
+        raise ModuleNotFoundError(error_message) from exc
+    return getattr(module, symbol_name)
+
+
+def lazy_action_transform_pipeline(max_action_dim: int):
+    ActionTransformPipeline = lazy_import(
+        "cosmos_framework.data.vfm.action.transforms",
+        "ActionTransformPipeline",
+        "Cosmos3 RoboLab policy serving requires cosmos_framework on PYTHONPATH so the "
+        "golden ActionTransformPipeline can be reused.",
+    )
+    return ActionTransformPipeline(max_action_dim=max_action_dim, cfg_dropout_rate=0.0)
+
+
+def build_robolab_unipc_scheduler(num_steps: int, shift: float, device: torch.device):
+    FlowUniPCMultistepScheduler = lazy_import(
+        "cosmos_framework.model.vfm.diffusion.samplers.fm_solvers_unipc",
+        "FlowUniPCMultistepScheduler",
+        (
+            "Cosmos3 RoboLab policy serving requires cosmos_framework on PYTHONPATH so the "
+            "golden FlowUniPCMultistepScheduler can be reused."
+        ),
+    )
+
+    scheduler = FlowUniPCMultistepScheduler(
+        num_train_timesteps=1000,
+        shift=1.0,
+        use_dynamic_shifting=False,
+    )
+    scheduler.set_timesteps(num_steps, device=device, shift=float(shift))
+    return scheduler
+
+
+def convert_midtrain_rotation(value: Any, src: str, dst: str) -> np.ndarray:
+    convert_rotation = lazy_import(
+        "cosmos_framework.data.vfm.action.pose_utils",
+        "convert_rotation",
+        "Cosmos3 RoboLab midtrain action serving requires cosmos_framework pose_utils on PYTHONPATH.",
+    )
+    return convert_rotation(value, src, dst)
+
+
+def pose_abs_to_rel(*args, **kwargs) -> np.ndarray:
+    pose_abs_to_rel_func = lazy_import(
+        "cosmos_framework.data.vfm.action.pose_utils",
+        "pose_abs_to_rel",
+        "Cosmos3 RoboLab midtrain action serving requires cosmos_framework pose_utils on PYTHONPATH.",
+    )
+    return pose_abs_to_rel_func(*args, **kwargs)
+
+
+def pose_rel_to_abs(*args, **kwargs) -> np.ndarray:
+    pose_rel_to_abs_func = lazy_import(
+        "cosmos_framework.data.vfm.action.pose_utils",
+        "pose_rel_to_abs",
+        "Cosmos3 RoboLab midtrain action serving requires cosmos_framework pose_utils on PYTHONPATH.",
+    )
+    return pose_rel_to_abs_func(*args, **kwargs)
+
+
+def build_abs_pose_from_components(*args, **kwargs) -> np.ndarray:
+    build_abs_pose_from_components_func = lazy_import(
+        "cosmos_framework.data.vfm.action.pose_utils",
+        "build_abs_pose_from_components",
+        "Cosmos3 RoboLab midtrain action serving requires cosmos_framework pose_utils on PYTHONPATH.",
+    )
+    return build_abs_pose_from_components_func(*args, **kwargs)
+
+
+def next_robolab_seed(extra: dict[str, Any], obs: dict[str, Any], request_id: str | None) -> int:
+    base_seed = int(extra.get("robolab_seed") or ROBOLAB_DEFAULT_SEED)
+    deterministic_seed = str(extra.get("deterministic_seed", "")).strip().lower() in {"1", "true", "yes", "on"}
+    if deterministic_seed:
+        return base_seed
+    explicit_seed = extra.get("seed")
+    if explicit_seed is not None:
+        return int(explicit_seed)
+    seed_key = "|".join(
+        str(part)
+        for part in (
+            base_seed,
+            extra.get("session_id", ""),
+            request_id or "",
+            obs.get("prompt", ""),
+        )
+    )
+    return zlib.crc32(seed_key.encode("utf-8")) & 0x7FFFFFFF
+
+
+def log_robolab_action_summary(label: str, value: Any) -> None:
+    if not _is_rank_zero():
+        return
+    if isinstance(value, torch.Tensor):
+        array = value.detach().float().cpu().numpy()
+    else:
+        array = np.asarray(value, dtype=np.float32)
+    finite = np.isfinite(array)
+    if finite.any():
+        finite_min = float(array[finite].min())
+        finite_max = float(array[finite].max())
+    else:
+        finite_min = None
+        finite_max = None
+    if array.ndim == 0:
+        head = array.reshape(1).tolist()
+    else:
+        head = array.reshape(-1, array.shape[-1])[:3].tolist()
+    logger.info(
+        "RoboLab action summary %s: shape=%s nan=%d finite=%d finite_min=%s finite_max=%s head=%s",
+        label,
+        tuple(array.shape),
+        int(np.isnan(array).sum()),
+        int(finite.sum()),
+        finite_min,
+        finite_max,
+        head,
+    )
+
+
+def postprocess_robolab_action(action: torch.Tensor, inputs: RoboLabPolicyInputs) -> np.ndarray:
+    action_np = action[0].float().cpu().numpy()
+    log_robolab_action_summary("raw_model_action", action_np)
+    history_length = int(inputs.history_length)
+    action_np = action_np[history_length:]
+    action_np[:, -1] = 1.0 - action_np[:, -1]
+
+    if inputs.action_space == "midtrain":
+        obs = inputs.observation
+        eef_pos = ensure_2d_float_array(obs["observation/eef_pos"], "observation/eef_pos", 3)
+        eef_quat = ensure_2d_float_array(obs["observation/eef_quat"], "observation/eef_quat", 4)
+        initial_pose = np.eye(4, dtype=np.float32)
+        initial_pose[:3, :3] = convert_midtrain_rotation(eef_quat[-1], "quat_xyzw", "matrix")
+        initial_pose[:3, 3] = eef_pos[-1]
+        abs_pose = pose_rel_to_abs(
+            action_np[:, :ROBOLAB_MIDTRAIN_POSE_ACTION_DIM],
+            rotation_format="rot6d",
+            pose_convention="backward_framewise",
+            initial_pose=initial_pose,
+        )
+        position = abs_pose[1:, :3, 3]
+        quat_xyzw = convert_midtrain_rotation(abs_pose[1:, :3, :3], "matrix", "quat_xyzw")
+        action_np = np.concatenate([position, quat_xyzw, action_np[:, ROBOLAB_MIDTRAIN_POSE_ACTION_DIM:]], axis=-1)
+
+    log_robolab_action_summary("postprocessed_robolab_action", action_np)
+    return np.asarray(action_np, dtype=np.float32)
diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py
index 34ac5e31b71..6d48d3ee84f 100644
--- a/vllm_omni/diffusion/registry.py
+++ b/vllm_omni/diffusion/registry.py
@@ -512,6 +512,13 @@ def _apply_sequence_parallel_if_enabled(model, od_config: OmniDiffusionConfig) -
     "HiDreamImagePipeline": "get_hidream_image_post_process_func",
 }
 
+_DIFFUSION_ACTION_POST_PROCESS_FUNCS = {
+    # arch: action_post_process_func
+    # `action_post_process_func` function must be placed in {mod_folder}/{mod_relname}.py,
+    # where mod_folder and mod_relname are defined and mapped using `_DIFFUSION_MODELS` via the `arch` key.
+    "Cosmos3OmniDiffusersPipeline": "get_cosmos3_action_post_process_func",
+}
+
 _DIFFUSION_PRE_PROCESS_FUNCS = {
     # arch: pre_process_func
     # `pre_process_func` function must be placed in {mod_folder}/{mod_relname}.py,
@@ -543,6 +550,7 @@ def register_diffusion_model(
     class_name: str,
     pre_process_func_name: str | None = None,
     post_process_func_name: str | None = None,
+    action_post_process_func_name: str | None = None,
 ) -> None:
     """Register a diffusion model pipeline from an out-of-tree plugin.
 
@@ -561,6 +569,9 @@ def register_diffusion_model(
         post_process_func_name: Optional name of the post-process function
             located in *module_name*.  Pass ``None`` to keep the existing
             entry when replacing a built-in model.
+        action_post_process_func_name: Optional name of the action post-process
+            function located in *module_name*.  Pass ``None`` to keep the
+            existing entry when replacing a built-in model.
     """
     # Register model class in DiffusionModelRegistry
     DiffusionModelRegistry.register_model(
@@ -578,6 +589,8 @@ def register_diffusion_model(
         _DIFFUSION_PRE_PROCESS_FUNCS[model_arch] = pre_process_func_name
     if post_process_func_name is not None:
         _DIFFUSION_POST_PROCESS_FUNCS[model_arch] = post_process_func_name
+    if action_post_process_func_name is not None:
+        _DIFFUSION_ACTION_POST_PROCESS_FUNCS[model_arch] = action_post_process_func_name
 
     logger.info(
         "Registered diffusion model %s -> %s.%s",
@@ -608,6 +621,13 @@ def get_diffusion_post_process_func(od_config: OmniDiffusionConfig):
     return _load_process_func(od_config, func_name)
 
 
+def get_diffusion_action_post_process_func(od_config: OmniDiffusionConfig):
+    if od_config.model_class_name not in _DIFFUSION_ACTION_POST_PROCESS_FUNCS:
+        return None
+    func_name = _DIFFUSION_ACTION_POST_PROCESS_FUNCS[od_config.model_class_name]
+    return _load_process_func(od_config, func_name)
+
+
 def get_diffusion_pre_process_func(od_config: OmniDiffusionConfig):
     if od_config.model_class_name not in _DIFFUSION_PRE_PROCESS_FUNCS:
         return None  # Return None if no pre-processing function is registered (for backward compatibility)

From fb7cd534d1b171527ddc9e2b6bc4b0176b6343ef Mon Sep 17 00:00:00 2001
From: yuzhud <yuzhudong@gmail.com>
Date: Thu, 11 Jun 2026 12:15:07 -0700
Subject: [PATCH 08/12] Preserve RoboLab action output path

Signed-off-by: yuzhud <yuzhudong@gmail.com>
---
 tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py | 4 ++--
 vllm_omni/diffusion/diffusion_engine.py                 | 2 +-
 vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py  | 2 +-
 3 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py
index 5ca99ef9caf..59641a25ca6 100644
--- a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py
+++ b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py
@@ -889,8 +889,8 @@ def fake_prepare_action_video(*args, **kwargs):
         assert output.output == {}
         assert output.custom_output["action_only_output"] is True
         assert output.custom_output["action"].shape == (1, 2, 2)
-        assert output.custom_output["robolab_policy_inputs"] is inputs
-        assert "actions" not in output.custom_output
+        assert output.custom_output["actions"].shape == (1, 2)
+        assert "robolab_policy_inputs" not in output.custom_output
 
     @pytest.mark.parametrize(
         ("prompt", "sampling_params", "message"),
diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py
index 39d51186caa..b3405e3bb5a 100644
--- a/vllm_omni/diffusion/diffusion_engine.py
+++ b/vllm_omni/diffusion/diffusion_engine.py
@@ -304,7 +304,7 @@ async def step(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]:
         if action_payload is None:
             action_payload = custom_output.get("actions")
         action_post_process_func = getattr(self, "action_post_process_func", None)
-        if action_post_process_func is not None:
+        if action_payload is None and action_post_process_func is not None:
             raw_action_payload = custom_output.get("action", action_payload)
             if raw_action_payload is not None:
                 action_kwargs: dict[str, Any] = {}
diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py
index 92897d35192..55ed15224ce 100644
--- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py
+++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py
@@ -1175,7 +1175,7 @@ def _forward_robolab_policy(
             "action_mode": action_mode,
             "domain_id": domain_id,
             "action_only_output": True,
-            "robolab_policy_inputs": inputs,
+            "actions": postprocess_robolab_action(action, inputs),
         }
         return DiffusionOutput(output={}, custom_output=custom_action_output)
 

From 1e16227db018a48525e4c5499dfc906faa4710f8 Mon Sep 17 00:00:00 2001
From: yuzhud <yuzhudong@gmail.com>
Date: Thu, 11 Jun 2026 17:54:50 -0700
Subject: [PATCH 09/12] Use native RMSNorm for Cosmos3 actions

Signed-off-by: yuzhud <yuzhudong@gmail.com>
---
 vllm_omni/platforms/cuda/platform.py | 21 +++++++++++++++++++--
 1 file changed, 19 insertions(+), 2 deletions(-)

diff --git a/vllm_omni/platforms/cuda/platform.py b/vllm_omni/platforms/cuda/platform.py
index c96da69fca5..c4836076a2c 100644
--- a/vllm_omni/platforms/cuda/platform.py
+++ b/vllm_omni/platforms/cuda/platform.py
@@ -217,6 +217,20 @@ def get_free_memory(cls, device: torch.device | None = None) -> int:
     def get_device_name(cls, device_id: int = 0) -> str:
         return torch.cuda.get_device_name(device_id)
 
+    @classmethod
+    def _uses_cosmos3_model(cls, vllm_config: VllmConfig) -> bool:
+        model_config = getattr(vllm_config, "model_config", None)
+        model = str(getattr(model_config, "model", "") or "").lower()
+        if "cosmos3" in model:
+            return True
+
+        for config_attr in ("hf_config", "hf_text_config"):
+            hf_config = getattr(model_config, config_attr, None)
+            architectures = getattr(hf_config, "architectures", []) or []
+            if any("cosmos3" in str(arch).lower() for arch in architectures):
+                return True
+        return False
+
     @classmethod
     def get_default_ir_op_priority(cls, vllm_config: VllmConfig) -> IrOpPriorityConfig:
         """Copied from vllm/platforms/cuda/platform.py v0.20.0 with force using vllm_c kernels"""
@@ -225,8 +239,11 @@ def get_default_ir_op_priority(cls, vllm_config: VllmConfig) -> IrOpPriorityConf
         # Use oink if enabled for rms_norm
         # TODO(Laurawly/luka): remove this env var,
         #  users can just use IR op priority directly
-        rms_norm = default
-        if envs.VLLM_USE_OINK_OPS:
+        if cls._uses_cosmos3_model(vllm_config):
+            rms_norm = ["native"]
+        elif envs.VLLM_USE_OINK_OPS:
             rms_norm = ["oink"] + default
+        else:
+            rms_norm = default
 
         return IrOpPriorityConfig.with_default(default, rms_norm=rms_norm, fused_add_rms_norm=rms_norm)

From 28eeab866d8e04ae0ab02fce4cfaaffba44765f0 Mon Sep 17 00:00:00 2001
From: yuzhud <yuzhudong@gmail.com>
Date: Thu, 11 Jun 2026 22:41:05 -0700
Subject: [PATCH 10/12] Use native RMSNorm for Cosmos3 diffusion

Signed-off-by: yuzhud <yuzhudong@gmail.com>
---
 .../diffusion/worker/diffusion_worker.py      | 25 ++++++++++++++++++-
 1 file changed, 24 insertions(+), 1 deletion(-)

diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py
index cd9576437fa..c7844d5f0e5 100644
--- a/vllm_omni/diffusion/worker/diffusion_worker.py
+++ b/vllm_omni/diffusion/worker/diffusion_worker.py
@@ -20,6 +20,7 @@
 import torch
 import zmq
 from vllm.config import CompilationConfig, DeviceConfig, VllmConfig, set_current_vllm_config
+from vllm.config.kernel import IrOpPriorityConfig
 from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
 from vllm.logger import init_logger
 from vllm.profiler.wrapper import CudaProfilerWrapper, WorkerProfiler
@@ -97,6 +98,23 @@ def _make_diffusion_vllm_model_config(od_config: OmniDiffusionConfig) -> _Diffus
     )
 
 
+def _uses_cosmos3_diffusion_model(od_config: OmniDiffusionConfig, model_config: _DiffusionVllmModelConfig) -> bool:
+    identifiers = (
+        getattr(od_config, "model", None),
+        getattr(od_config, "model_class_name", None),
+        getattr(model_config, "model", None),
+    )
+    if any("cosmos3" in str(identifier).lower() for identifier in identifiers if identifier):
+        return True
+
+    for config_attr in ("hf_config", "hf_text_config"):
+        hf_config = getattr(model_config, config_attr, None)
+        architectures = getattr(hf_config, "architectures", []) or []
+        if any("cosmos3" in str(architecture).lower() for architecture in architectures):
+            return True
+    return False
+
+
 @contextmanager
 def _force_cutlass_fp8_linear_kernel(quant_config: object | None) -> Iterator[None]:
     import vllm.model_executor.layers.quantization.modelopt as vllm_modelopt
@@ -234,7 +252,12 @@ def init_device(self) -> None:
         vllm_config.quant_config = self.od_config.quantization_config
         # Since vLLM v0.20.0, IR wraps GPU ops. Set IR op priority preference to enforce GPU op fusion during wrapping.
         # Also need to log, because vLLM internally logs another line in VllmConfig.__post_init__. Avoid confusion.
-        vllm_config.kernel_config.ir_op_priority = current_omni_platform.get_default_ir_op_priority(vllm_config)
+        ir_op_priority = current_omni_platform.get_default_ir_op_priority(vllm_config)
+        if _uses_cosmos3_diffusion_model(self.od_config, vllm_config.model_config):
+            ir_op_priority = IrOpPriorityConfig.with_default(
+                ["vllm_c", "native"], rms_norm=["native"], fused_add_rms_norm=["native"]
+            )
+        vllm_config.kernel_config.ir_op_priority = ir_op_priority
         logger.info(
             "Final IR op priority after setting vLLM-Omni overrides: %s", vllm_config.kernel_config.ir_op_priority
         )

From 54e0ab164f5fe246f1b0255cbaea029347d06855 Mon Sep 17 00:00:00 2001
From: yuzhud <yuzhudong@gmail.com>
Date: Thu, 11 Jun 2026 22:42:32 -0700
Subject: [PATCH 11/12] Revert "Use native RMSNorm for Cosmos3 actions"

This reverts commit d47d2570151fc77515399131d284bd7299e2cb2f.

Signed-off-by: yuzhud <yuzhudong@gmail.com>
---
 vllm_omni/platforms/cuda/platform.py | 21 ++-------------------
 1 file changed, 2 insertions(+), 19 deletions(-)

diff --git a/vllm_omni/platforms/cuda/platform.py b/vllm_omni/platforms/cuda/platform.py
index c4836076a2c..c96da69fca5 100644
--- a/vllm_omni/platforms/cuda/platform.py
+++ b/vllm_omni/platforms/cuda/platform.py
@@ -217,20 +217,6 @@ def get_free_memory(cls, device: torch.device | None = None) -> int:
     def get_device_name(cls, device_id: int = 0) -> str:
         return torch.cuda.get_device_name(device_id)
 
-    @classmethod
-    def _uses_cosmos3_model(cls, vllm_config: VllmConfig) -> bool:
-        model_config = getattr(vllm_config, "model_config", None)
-        model = str(getattr(model_config, "model", "") or "").lower()
-        if "cosmos3" in model:
-            return True
-
-        for config_attr in ("hf_config", "hf_text_config"):
-            hf_config = getattr(model_config, config_attr, None)
-            architectures = getattr(hf_config, "architectures", []) or []
-            if any("cosmos3" in str(arch).lower() for arch in architectures):
-                return True
-        return False
-
     @classmethod
     def get_default_ir_op_priority(cls, vllm_config: VllmConfig) -> IrOpPriorityConfig:
         """Copied from vllm/platforms/cuda/platform.py v0.20.0 with force using vllm_c kernels"""
@@ -239,11 +225,8 @@ def get_default_ir_op_priority(cls, vllm_config: VllmConfig) -> IrOpPriorityConf
         # Use oink if enabled for rms_norm
         # TODO(Laurawly/luka): remove this env var,
         #  users can just use IR op priority directly
-        if cls._uses_cosmos3_model(vllm_config):
-            rms_norm = ["native"]
-        elif envs.VLLM_USE_OINK_OPS:
+        rms_norm = default
+        if envs.VLLM_USE_OINK_OPS:
             rms_norm = ["oink"] + default
-        else:
-            rms_norm = default
 
         return IrOpPriorityConfig.with_default(default, rms_norm=rms_norm, fused_add_rms_norm=rms_norm)

From 0df9a175c4f61761d764cea241f468771621b26f Mon Sep 17 00:00:00 2001
From: yuzhud <yuzhudong@gmail.com>
Date: Fri, 12 Jun 2026 17:58:43 -0700
Subject: [PATCH 12/12] Move Cosmos3 review hooks out of worker

Signed-off-by: yuzhud <yuzhudong@gmail.com>
---
 .../models/cosmos3/test_cosmos3_pipeline.py   | 38 ++++++++++++++++++-
 .../diffusion/test_diffusion_plugin_hooks.py  | 28 +++++++++++++-
 .../models/cosmos3/pipeline_cosmos3.py        | 28 ++++++++++++--
 vllm_omni/diffusion/models/cosmos3/utils.py   | 35 ++++++++++++++---
 vllm_omni/diffusion/registry.py               | 20 ++++++++++
 .../diffusion/worker/diffusion_worker.py      | 34 +++++------------
 6 files changed, 146 insertions(+), 37 deletions(-)

diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py
index 59641a25ca6..f77f63d5077 100644
--- a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py
+++ b/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
 
@@ -231,6 +232,7 @@ def test_pipeline_registered_and_exported() -> None:
     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,
@@ -249,6 +251,7 @@ def test_pipeline_registered_and_exported() -> None:
     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__
 
@@ -421,6 +424,7 @@ 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())
@@ -447,11 +451,40 @@ def test_action_postprocess_handles_robolab_policy_outputs() -> None:
     )
 
     action = torch.tensor([[[0.0, 0.25], [1.0, 0.75]]])
-    processed = func(action, custom_output={"robolab_policy_inputs": inputs})
+    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:
@@ -889,7 +922,8 @@ def fake_prepare_action_video(*args, **kwargs):
         assert output.output == {}
         assert output.custom_output["action_only_output"] is True
         assert output.custom_output["action"].shape == (1, 2, 2)
-        assert output.custom_output["actions"].shape == (1, 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(
diff --git a/tests/diffusion/test_diffusion_plugin_hooks.py b/tests/diffusion/test_diffusion_plugin_hooks.py
index 2a9a0a82109..6804d550d53 100644
--- a/tests/diffusion/test_diffusion_plugin_hooks.py
+++ b/tests/diffusion/test_diffusion_plugin_hooks.py
@@ -10,12 +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,
@@ -61,6 +63,7 @@ def cleanup_registry(self):
         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)
@@ -70,6 +73,8 @@ def cleanup_registry(self):
         _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."""
@@ -80,6 +85,7 @@ def test_register_new_model(self):
             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"] == (
@@ -90,6 +96,7 @@ 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:
@@ -114,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)
diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py
index 55ed15224ce..b0b1f7559c1 100644
--- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py
+++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py
@@ -21,6 +21,7 @@
 import os
 import time
 from collections.abc import Iterable
+from dataclasses import fields
 from typing import Any, ClassVar
 
 import numpy as np
@@ -81,6 +82,7 @@
     ROBOLAB_DEFAULT_RAW_ACTION_DIM,
     ROBOLAB_DEFAULT_RESOLUTION,
     ROBOLAB_MIDTRAIN_RAW_ACTION_DIM,
+    RoboLabActionPostprocessInputs,
     RoboLabPolicyInputs,
     build_abs_pose_from_components,
     build_robolab_unipc_scheduler,
@@ -91,6 +93,7 @@
     extract_robolab_image,
     extract_robolab_prompt_image,
     lazy_action_transform_pipeline,
+    make_robolab_action_postprocess_inputs,
     next_robolab_seed,
     normalize_condition_frame_indexes_vision,
     normalize_condition_video_keep,
@@ -486,14 +489,31 @@ def get_cosmos3_action_post_process_func(od_config: OmniDiffusionConfig):
 
     def action_post_process_func(action: Any, custom_output: dict[str, Any] | None = None, sampling_params=None):
         del sampling_params
-        inputs = custom_output.get("robolab_policy_inputs") if isinstance(custom_output, dict) else None
-        if isinstance(inputs, RoboLabPolicyInputs):
-            return postprocess_robolab_action(action, inputs)
+        inputs = custom_output.get("robolab_action_postprocess") if isinstance(custom_output, dict) else None
+        if isinstance(inputs, RoboLabActionPostprocessInputs):
+            processed_action = postprocess_robolab_action(action, inputs)
+            custom_output.pop("robolab_action_postprocess", None)
+            return processed_action
         return action
 
     return action_post_process_func
 
 
+def get_cosmos3_ir_op_priority_func(od_config: OmniDiffusionConfig):
+    del od_config
+
+    def ir_op_priority_func(ir_op_priority, vllm_config=None):
+        del vllm_config
+        from vllm.config.kernel import IrOpPriorityConfig
+
+        priority_kwargs = {field.name: list(getattr(ir_op_priority, field.name)) for field in fields(ir_op_priority)}
+        priority_kwargs["rms_norm"] = ["native"]
+        priority_kwargs["fused_add_rms_norm"] = ["native"]
+        return IrOpPriorityConfig(**priority_kwargs)
+
+    return ir_op_priority_func
+
+
 # ---------------------------------------------------------------------------
 # Pipeline
 # ---------------------------------------------------------------------------
@@ -1175,7 +1195,7 @@ def _forward_robolab_policy(
             "action_mode": action_mode,
             "domain_id": domain_id,
             "action_only_output": True,
-            "actions": postprocess_robolab_action(action, inputs),
+            "robolab_action_postprocess": make_robolab_action_postprocess_inputs(inputs),
         }
         return DiffusionOutput(output={}, custom_output=custom_action_output)
 
diff --git a/vllm_omni/diffusion/models/cosmos3/utils.py b/vllm_omni/diffusion/models/cosmos3/utils.py
index 6f91674e644..bfa3fd4db18 100644
--- a/vllm_omni/diffusion/models/cosmos3/utils.py
+++ b/vllm_omni/diffusion/models/cosmos3/utils.py
@@ -71,6 +71,30 @@ class RoboLabPolicyInputs:
     observation: dict[str, Any]
 
 
+@dataclass(frozen=True)
+class RoboLabActionPostprocessInputs:
+    history_length: int
+    action_space: str
+    eef_pos: np.ndarray | None = None
+    eef_quat: np.ndarray | None = None
+
+
+def make_robolab_action_postprocess_inputs(inputs: RoboLabPolicyInputs) -> RoboLabActionPostprocessInputs:
+    if inputs.action_space != "midtrain":
+        return RoboLabActionPostprocessInputs(
+            history_length=inputs.history_length,
+            action_space=inputs.action_space,
+        )
+
+    obs = inputs.observation
+    return RoboLabActionPostprocessInputs(
+        history_length=inputs.history_length,
+        action_space=inputs.action_space,
+        eef_pos=ensure_2d_float_array(obs["observation/eef_pos"], "observation/eef_pos", 3),
+        eef_quat=ensure_2d_float_array(obs["observation/eef_quat"], "observation/eef_quat", 4),
+    )
+
+
 def normalize_condition_frame_indexes_vision(value: Any) -> tuple[int, ...]:
     """Normalize Cosmos3 vision-conditioning latent frame indexes."""
     if value is None:
@@ -320,7 +344,7 @@ def log_robolab_action_summary(label: str, value: Any) -> None:
     )
 
 
-def postprocess_robolab_action(action: torch.Tensor, inputs: RoboLabPolicyInputs) -> np.ndarray:
+def postprocess_robolab_action(action: torch.Tensor, inputs: RoboLabActionPostprocessInputs) -> np.ndarray:
     action_np = action[0].float().cpu().numpy()
     log_robolab_action_summary("raw_model_action", action_np)
     history_length = int(inputs.history_length)
@@ -328,12 +352,11 @@ def postprocess_robolab_action(action: torch.Tensor, inputs: RoboLabPolicyInputs
     action_np[:, -1] = 1.0 - action_np[:, -1]
 
     if inputs.action_space == "midtrain":
-        obs = inputs.observation
-        eef_pos = ensure_2d_float_array(obs["observation/eef_pos"], "observation/eef_pos", 3)
-        eef_quat = ensure_2d_float_array(obs["observation/eef_quat"], "observation/eef_quat", 4)
+        if inputs.eef_pos is None or inputs.eef_quat is None:
+            raise ValueError("RoboLab midtrain action postprocess requires eef_pos and eef_quat metadata.")
         initial_pose = np.eye(4, dtype=np.float32)
-        initial_pose[:3, :3] = convert_midtrain_rotation(eef_quat[-1], "quat_xyzw", "matrix")
-        initial_pose[:3, 3] = eef_pos[-1]
+        initial_pose[:3, :3] = convert_midtrain_rotation(inputs.eef_quat[-1], "quat_xyzw", "matrix")
+        initial_pose[:3, 3] = inputs.eef_pos[-1]
         abs_pose = pose_rel_to_abs(
             action_np[:, :ROBOLAB_MIDTRAIN_POSE_ACTION_DIM],
             rotation_format="rot6d",
diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py
index 6d48d3ee84f..9cc65486a7a 100644
--- a/vllm_omni/diffusion/registry.py
+++ b/vllm_omni/diffusion/registry.py
@@ -519,6 +519,13 @@ def _apply_sequence_parallel_if_enabled(model, od_config: OmniDiffusionConfig) -
     "Cosmos3OmniDiffusersPipeline": "get_cosmos3_action_post_process_func",
 }
 
+_DIFFUSION_IR_OP_PRIORITY_FUNCS = {
+    # arch: ir_op_priority_func
+    # `ir_op_priority_func` function must be placed in {mod_folder}/{mod_relname}.py,
+    # where mod_folder and mod_relname are defined and mapped using `_DIFFUSION_MODELS` via the `arch` key.
+    "Cosmos3OmniDiffusersPipeline": "get_cosmos3_ir_op_priority_func",
+}
+
 _DIFFUSION_PRE_PROCESS_FUNCS = {
     # arch: pre_process_func
     # `pre_process_func` function must be placed in {mod_folder}/{mod_relname}.py,
@@ -551,6 +558,7 @@ def register_diffusion_model(
     pre_process_func_name: str | None = None,
     post_process_func_name: str | None = None,
     action_post_process_func_name: str | None = None,
+    ir_op_priority_func_name: str | None = None,
 ) -> None:
     """Register a diffusion model pipeline from an out-of-tree plugin.
 
@@ -572,6 +580,9 @@ def register_diffusion_model(
         action_post_process_func_name: Optional name of the action post-process
             function located in *module_name*.  Pass ``None`` to keep the
             existing entry when replacing a built-in model.
+        ir_op_priority_func_name: Optional name of the IR op priority merge
+            function located in *module_name*. Pass ``None`` to keep the
+            existing entry when replacing a built-in model.
     """
     # Register model class in DiffusionModelRegistry
     DiffusionModelRegistry.register_model(
@@ -591,6 +602,8 @@ def register_diffusion_model(
         _DIFFUSION_POST_PROCESS_FUNCS[model_arch] = post_process_func_name
     if action_post_process_func_name is not None:
         _DIFFUSION_ACTION_POST_PROCESS_FUNCS[model_arch] = action_post_process_func_name
+    if ir_op_priority_func_name is not None:
+        _DIFFUSION_IR_OP_PRIORITY_FUNCS[model_arch] = ir_op_priority_func_name
 
     logger.info(
         "Registered diffusion model %s -> %s.%s",
@@ -628,6 +641,13 @@ def get_diffusion_action_post_process_func(od_config: OmniDiffusionConfig):
     return _load_process_func(od_config, func_name)
 
 
+def get_diffusion_ir_op_priority_func(od_config: OmniDiffusionConfig):
+    if od_config.model_class_name not in _DIFFUSION_IR_OP_PRIORITY_FUNCS:
+        return None
+    func_name = _DIFFUSION_IR_OP_PRIORITY_FUNCS[od_config.model_class_name]
+    return _load_process_func(od_config, func_name)
+
+
 def get_diffusion_pre_process_func(od_config: OmniDiffusionConfig):
     if od_config.model_class_name not in _DIFFUSION_PRE_PROCESS_FUNCS:
         return None  # Return None if no pre-processing function is registered (for backward compatibility)
diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py
index c7844d5f0e5..e3adfecc69e 100644
--- a/vllm_omni/diffusion/worker/diffusion_worker.py
+++ b/vllm_omni/diffusion/worker/diffusion_worker.py
@@ -20,7 +20,6 @@
 import torch
 import zmq
 from vllm.config import CompilationConfig, DeviceConfig, VllmConfig, set_current_vllm_config
-from vllm.config.kernel import IrOpPriorityConfig
 from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
 from vllm.logger import init_logger
 from vllm.profiler.wrapper import CudaProfilerWrapper, WorkerProfiler
@@ -44,6 +43,7 @@
 from vllm_omni.diffusion.forward_context import set_forward_context
 from vllm_omni.diffusion.ipc import pack_diffusion_output_shm
 from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
+from vllm_omni.diffusion.registry import get_diffusion_ir_op_priority_func
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput
 from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
@@ -98,23 +98,6 @@ def _make_diffusion_vllm_model_config(od_config: OmniDiffusionConfig) -> _Diffus
     )
 
 
-def _uses_cosmos3_diffusion_model(od_config: OmniDiffusionConfig, model_config: _DiffusionVllmModelConfig) -> bool:
-    identifiers = (
-        getattr(od_config, "model", None),
-        getattr(od_config, "model_class_name", None),
-        getattr(model_config, "model", None),
-    )
-    if any("cosmos3" in str(identifier).lower() for identifier in identifiers if identifier):
-        return True
-
-    for config_attr in ("hf_config", "hf_text_config"):
-        hf_config = getattr(model_config, config_attr, None)
-        architectures = getattr(hf_config, "architectures", []) or []
-        if any("cosmos3" in str(architecture).lower() for architecture in architectures):
-            return True
-    return False
-
-
 @contextmanager
 def _force_cutlass_fp8_linear_kernel(quant_config: object | None) -> Iterator[None]:
     import vllm.model_executor.layers.quantization.modelopt as vllm_modelopt
@@ -177,6 +160,14 @@ def _create_diffusion_worker_vllm_config(device: torch.device, od_config: OmniDi
         return vllm_config
 
 
+def _resolve_ir_op_priority(od_config: OmniDiffusionConfig, vllm_config: VllmConfig) -> Any:
+    ir_op_priority = current_omni_platform.get_default_ir_op_priority(vllm_config)
+    ir_op_priority_func = get_diffusion_ir_op_priority_func(od_config)
+    if ir_op_priority_func is not None:
+        ir_op_priority = ir_op_priority_func(ir_op_priority, vllm_config=vllm_config)
+    return ir_op_priority
+
+
 class DiffusionWorker:
     """
     A worker that manages GPU infrastructure and delegates to the model runner.
@@ -252,12 +243,7 @@ def init_device(self) -> None:
         vllm_config.quant_config = self.od_config.quantization_config
         # Since vLLM v0.20.0, IR wraps GPU ops. Set IR op priority preference to enforce GPU op fusion during wrapping.
         # Also need to log, because vLLM internally logs another line in VllmConfig.__post_init__. Avoid confusion.
-        ir_op_priority = current_omni_platform.get_default_ir_op_priority(vllm_config)
-        if _uses_cosmos3_diffusion_model(self.od_config, vllm_config.model_config):
-            ir_op_priority = IrOpPriorityConfig.with_default(
-                ["vllm_c", "native"], rms_norm=["native"], fused_add_rms_norm=["native"]
-            )
-        vllm_config.kernel_config.ir_op_priority = ir_op_priority
+        vllm_config.kernel_config.ir_op_priority = _resolve_ir_op_priority(self.od_config, vllm_config)
         logger.info(
             "Final IR op priority after setting vLLM-Omni overrides: %s", vllm_config.kernel_config.ir_op_priority
         )