[diffusion] feat: expose cosmos3 policies through the Action API (#34243)

docs/cookbook/diffusion/Cosmos/Cosmos3.mdx

@@ -207,27 +207,48 @@ sglang serve \
   --num-gpus 1
 ```
 
-The following request predicts a 16-step action chunk from one observation. The chunk length is `num_frames - 1`, and the completed job's `action` field contains the tensor data, shape, mode, and active action dimension.
+`policy` and `inverse_dynamics` return actions, so their canonical API is the synchronous `/v1/actions/generations` endpoint. The following request predicts a 16-step action chunk from one observation image. `action_horizon=16` maps to the model's `num_frames=17` convention.
 
-```bash Command
+```python Python
-job_id=$(curl -sS -X POST http://127.0.0.1:30010/v1/videos \
+import base64
-  --form-string "prompt=Put the pot to the left of the purple item." \
+from pathlib import Path
-  --form "input_reference=@observation.png;type=image/png" \
-  --form-string "size=832x480" \
-  --form-string "num_frames=17" \
-  --form-string "fps=5" \
-  --form-string "num_inference_steps=30" \
-  --form-string "guidance_scale=1.0" \
-  --form-string "action_mode=policy" \
-  --form-string "domain_name=droid_lerobot" \
-  | python -c 'import json, sys; print(json.load(sys.stdin)["id"])')
 
-# After the job reaches "completed":
+import requests
-curl -sS "http://127.0.0.1:30010/v1/videos/${job_id}" \
+
-  | python -c 'import json, sys; print(json.dumps(json.load(sys.stdin)["action"], indent=2))'
+image_b64 = base64.b64encode(Path("observation.png").read_bytes()).decode()
+response = requests.post(
+    "http://127.0.0.1:30010/v1/actions/generations",
+    json={
+        "input": {
+            "task": "Put the pot to the left of the purple item.",
+            "observation": {
+                "image": {"b64_json": image_b64},
+            },
+        },
+        "parameters": {
+            "action_mode": "policy",
+            "action_horizon": 16,
+            "domain_name": "droid_lerobot",
+            "height": 480,
+            "width": 832,
+            "fps": 5,
+            "num_inference_steps": 30,
+            "guidance_scale": 1.0,
+            "seed": 42,
+        },
+    },
+    timeout=300,
+)
+response.raise_for_status()
+action = response.json()["data"][0]["action"]
+print(action["shape"], action["values"])
 ```
 
-The other action modes are `forward_dynamics` (condition on an observation and an `action` JSON array to generate video) and `inverse_dynamics` (condition on a full video to predict action). Select the embodiment head with `domain_name` or `domain_id`; set `raw_action_dim` explicitly when it cannot be inferred from the domain name.
+Use `GET /v1/actions/metadata` to inspect the action modes, default horizon, padded action dimension, and accepted observation modalities. Msgpack requests and the `/v1/actions/realtime` websocket use the same action envelope.
+
+`inverse_dynamics` also uses `/v1/actions/generations`; set `action_mode="inverse_dynamics"` and pass an observation video URL or server-local path as `input.observation.video`. Select the embodiment head with `domain_name` or `domain_id`; set `raw_action_dim` explicitly when it cannot be inferred from the domain name.
+
+`forward_dynamics` is intentionally different: it consumes an action array and predicts video, so it remains on `/v1/videos`. Action-producing modes submitted to `/v1/videos` return HTTP 400 with the canonical action endpoint in the error message.
 
 ## 5. Cosmos3 Parameters
 

python/sglang/multimodal_gen/benchmarks/bench_pi05_openpi.py

@@ -27,7 +27,7 @@ SCRIPT_DIR = Path(__file__).resolve().parent
 if sys.path and Path(sys.path[0]).resolve() == SCRIPT_DIR:
     sys.path.pop(0)
 
-from sglang.multimodal_gen.runtime.entrypoints.vla.protocol import (  # noqa: E402
+from sglang.multimodal_gen.runtime.entrypoints.action.protocol import (  # noqa: E402
     pack_msgpack,
     unpack_msgpack,
 )
@@ -681,10 +681,10 @@ def create_sglang_python_pipeline(
 
 
 def _make_sglang_python_req(server_args, payload: dict[str, Any]):
-    from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request
+    from sglang.multimodal_gen.runtime.entrypoints.action.protocol import (
-    from sglang.multimodal_gen.runtime.entrypoints.vla.protocol import (
         build_action_sampling_params,
     )
+    from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request
 
     sampling_params = build_action_sampling_params(payload, server_args)
     req = prepare_request(server_args, sampling_params)
@@ -729,7 +729,9 @@ def run_sglang_python(
         model_path,
         pipeline_config_path=pipeline_config_path,
     )
-    from sglang.multimodal_gen.runtime.entrypoints.vla.protocol import action_metadata
+    from sglang.multimodal_gen.runtime.entrypoints.action.protocol import (
+        action_metadata,
+    )
 
     metadata = action_metadata(server_args)
     metadata["precision"] = sglang_precision_metadata(pipeline)

python/sglang/multimodal_gen/configs/pipeline_configs/base.py

@@ -272,6 +272,16 @@ class PipelineConfig:
 
         del server_args
 
+    def supports_action_endpoint(self) -> bool:
+        """Whether this pipeline exposes the generic action generation API."""
+
+        return self.task_type.is_action_gen()
+
+    def supports_openpi_endpoint(self) -> bool:
+        """Whether this pipeline implements the OpenPI policy websocket."""
+
+        return False
+
     # Wan2.2 TI2V parameters
     boundary_ratio: float | None = None
 

python/sglang/multimodal_gen/configs/pipeline_configs/cosmos3.py

@@ -84,3 +84,9 @@ class Cosmos3Config(PipelineConfig):
                 (num_frames - 1) // vae_scale_factor_temporal
             ) * vae_scale_factor_temporal + 1
         return num_frames
+
+    def supports_action_endpoint(self) -> bool:
+        # The public Cosmos3 family shares one pipeline/config across visual-only
+        # and action-capable checkpoints. The loaded transformer validates that
+        # an action head is actually present when an action request is submitted.
+        return True

python/sglang/multimodal_gen/configs/pipeline_configs/pi05.py

@@ -90,6 +90,9 @@ class Pi05PipelineConfig(PipelineConfig):
     def supports_native_grouped_requests(self):
         return True
 
+    def supports_openpi_endpoint(self) -> bool:
+        return True
+
     def estimate_request_cost(self, batch) -> float:
         return float(
             self.action_horizon * self.action_dim * self.default_num_inference_steps

python/sglang/multimodal_gen/configs/sample/__init__.py

@@ -1,5 +1,6 @@
 # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
 
+from sglang.multimodal_gen.configs.sample.action import ActionSamplingParams
 from sglang.multimodal_gen.configs.sample.diffusers_generic import (
     DiffusersGenericSamplingParams,
 )
@@ -9,11 +10,10 @@ from sglang.multimodal_gen.configs.sample.lingbot_video_moe import (
 )
 from sglang.multimodal_gen.configs.sample.pi05 import Pi05SamplingParams
 from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
-from sglang.multimodal_gen.configs.sample.vla import VLASamplingParams
 
 __all__ = [
     "SamplingParams",
-    "VLASamplingParams",
+    "ActionSamplingParams",
     "DiffusersGenericSamplingParams",
     "Ideogram4SamplingParams",
     "Pi05SamplingParams",

python/sglang/multimodal_gen/configs/sample/action.py

@@ -17,8 +17,8 @@ if TYPE_CHECKING:
 
 
 @dataclass
-class VLASamplingParams:
+class ActionSamplingParams:
-    """Sampling parameters for VLA/action-generation policies."""
+    """Sampling parameters for policies that generate continuous actions."""
 
     data_type: DataType = DataType.ACTION
     request_id: str | None = field(default=None, metadata={"batch_sig_exclude": True})
@@ -120,7 +120,7 @@ class VLASamplingParams:
     def _validate_with_pipeline_config(self, pipeline_config):
         if not pipeline_config.task_type.is_action_gen():
             raise ValueError(
-                f"VLASamplingParams requires an ACTION pipeline, got {pipeline_config.task_type.name}"
+                f"ActionSamplingParams requires an ACTION pipeline, got {pipeline_config.task_type.name}"
             )
 
     def _adjust(self, server_args: "ServerArgs"):
@@ -140,7 +140,7 @@ class VLASamplingParams:
 
     def _set_output_file_name(self):
         if self.output_file_name is None:
-            self.output_file_name = "vla_action"
+            self.output_file_name = "action"
         self.output_file_name = _sanitize_filename(self.output_file_name)
         self._set_output_file_ext()
 
@@ -151,7 +151,7 @@ class VLASamplingParams:
 
     def _merge_with_user_params(
         self,
-        user_params: "VLASamplingParams",
+        user_params: "ActionSamplingParams",
         explicit_fields: set[str] | None = None,
     ):
         if user_params is None:
@@ -193,7 +193,7 @@ class VLASamplingParams:
             "--prompt",
             type=str,
             nargs="+",
-            help="Language instruction(s) for the VLA policy.",
+            help="Language instruction(s) for the action policy.",
         )
         add_argument(
             "--num-inference-steps",

python/sglang/multimodal_gen/configs/sample/cosmos3.py

@@ -71,7 +71,38 @@ class Cosmos3SamplingParams(SamplingParams):
     action_stats_path: str | None = None
     action_normalization: str = "quantile"
 
+    def _adjust(self, server_args) -> None:
+        action_output = False
+        if self.action_mode is not None:
+            self.action_mode = str(self.action_mode).strip().lower()
+            if self.action_mode not in (
+                "policy",
+                "forward_dynamics",
+                "inverse_dynamics",
+            ):
+                raise ValueError(
+                    f"Unsupported action_mode={self.action_mode!r}; expected "
+                    "'policy', 'forward_dynamics', or 'inverse_dynamics'."
+                )
+            action_output = self.action_mode != "forward_dynamics"
+
+        super()._adjust(server_args)
+
+        # Policy and inverse dynamics produce actions. Forward dynamics consumes
+        # actions to produce video and therefore remains a visual request.
+        if action_output:
+            self.data_type = DataType.ACTION
+            self.save_output = False
+            self.return_file_paths_only = False
+            self.return_frames = False
+            self.output_file_name = None
+            self.output_compression = 0
+
     def _set_output_file_name(self) -> None:
+        # Action outputs never need a visual filename. This also avoids hashing
+        # in-memory observation images while base visual adjustment is running.
+        if self.action_mode in ("policy", "inverse_dynamics"):
+            return
         # The pipeline config's ``task_type=TI2V`` drives ``data_type`` to
         # VIDEO, but a single-frame request is a T2I and must pick the IMAGE
         # extension. Flip before the base derives the file name.

python/sglang/multimodal_gen/configs/sample/pi05.py

@@ -3,11 +3,11 @@
 from dataclasses import dataclass, field
 from typing import Any
 
-from sglang.multimodal_gen.configs.sample.vla import VLASamplingParams
+from sglang.multimodal_gen.configs.sample.action import ActionSamplingParams
 
 
 @dataclass
-class Pi05SamplingParams(VLASamplingParams):
+class Pi05SamplingParams(ActionSamplingParams):
     """Sampling parameters for Pi0.5 flow-matching action inference."""
 
     num_inference_steps: int = 10

python/sglang/multimodal_gen/registry.py

@@ -1077,6 +1077,7 @@ def _register_configs():
         pipeline_config_cls=Cosmos3Config,
         hf_model_paths=[
             "nvidia/Cosmos3-Nano",
+            "nvidia/Cosmos3-Nano-Policy-DROID",
             "nvidia/Cosmos3-Super",
             "nvidia/Cosmos3-Super-Text2Image",
             "nvidia/Cosmos3-Super-Image2Video",

python/sglang/multimodal_gen/runtime/entrypoints/action/api.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 
 from fastapi import APIRouter, HTTPException, Request, Response, WebSocket
 
-from sglang.multimodal_gen.runtime.entrypoints.vla.protocol import (
+from sglang.multimodal_gen.runtime.entrypoints.action.protocol import (
     action_generation_response,
     action_metadata,
     action_raw_response,
@@ -12,7 +12,7 @@ from sglang.multimodal_gen.runtime.entrypoints.vla.protocol import (
     pack_msgpack,
     unpack_msgpack,
 )
-from sglang.multimodal_gen.runtime.entrypoints.vla.ws_utils import (
+from sglang.multimodal_gen.runtime.entrypoints.action.ws_utils import (
     run_action_msgpack_ws,
 )
 from sglang.multimodal_gen.runtime.server_args import ServerArgs

python/sglang/multimodal_gen/runtime/entrypoints/action/openpi.py

@@ -6,7 +6,7 @@ from typing import Any
 
 from fastapi import APIRouter, WebSocket
 
-from sglang.multimodal_gen.runtime.entrypoints.vla.ws_utils import (
+from sglang.multimodal_gen.runtime.entrypoints.action.ws_utils import (
     run_action_msgpack_ws,
 )
 from sglang.multimodal_gen.runtime.server_args import ServerArgs

python/sglang/multimodal_gen/runtime/entrypoints/action/protocol.py

@@ -13,7 +13,10 @@ from typing import Any
 import numpy as np
 from PIL import Image
 
-from sglang.multimodal_gen.configs.sample.vla import VLASamplingParams
+from sglang.multimodal_gen.configs.pipeline_configs.cosmos3 import Cosmos3Config
+from sglang.multimodal_gen.configs.sample.action import ActionSamplingParams
+from sglang.multimodal_gen.configs.sample.cosmos3 import Cosmos3SamplingParams
+from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
 from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request
 from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
@@ -107,6 +110,9 @@ def _normalize_observation(observation: dict[str, Any]) -> dict[str, Any]:
         normalized["images"] = {
             name: _normalize_image_value(value) for name, value in images.items()
         }
+    for name in ("image", "image_path", "input_reference"):
+        if name in normalized:
+            normalized[name] = _normalize_image_value(normalized[name])
     state = normalized.get("state")
     if isinstance(state, dict):
         normalized["state"] = _decode_tensor_payload(state)
@@ -141,6 +147,53 @@ def images_from_observation(
 
 def action_metadata(server_args: ServerArgs) -> dict[str, Any]:
     pipeline_config = server_args.pipeline_config
+    if isinstance(pipeline_config, Cosmos3Config):
+        defaults = Cosmos3SamplingParams()
+        return {
+            "object": "action.metadata",
+            "model": server_args.model_id or server_args.model_path,
+            "model_path": server_args.model_path,
+            "policy_family": "cosmos3",
+            "input": {
+                "modalities": ["image", "video"],
+                "supported_resolutions": [
+                    list(resolution) for resolution in defaults.supported_resolutions
+                ],
+                "state_dim": None,
+            },
+            "output": {
+                "action_type": "continuous",
+                "action_horizon": 16,
+                "action_dim": None,
+                "padded_action_dim": pipeline_config.dit_config.arch_config.action_dim,
+                "dtype": "float32",
+            },
+            "runtime": {
+                "parallelism": {
+                    "num_gpus": server_args.num_gpus,
+                    "tp_size": server_args.tp_size,
+                    "sp_degree": server_args.sp_degree,
+                    "ulysses_degree": server_args.ulysses_degree,
+                    "ring_degree": server_args.ring_degree,
+                }
+            },
+            "defaults": {
+                "action_mode": "policy",
+                "action_horizon": 16,
+                "num_inference_steps": defaults.num_inference_steps,
+                "height": 480,
+                "width": 832,
+                "fps": 5,
+            },
+            "capabilities": {
+                "action_modes": ["policy", "inverse_dynamics"],
+                "realtime_websocket": True,
+                "openpi_websocket": False,
+                "batch_inputs": False,
+                "multiple_candidates": False,
+            },
+        }
+
     policy_family = getattr(
         pipeline_config,
         "policy_family",
@@ -226,6 +279,9 @@ def _action_request_to_observation(payload: dict[str, Any]) -> dict[str, Any]:
         observation["state"] = input_payload["state"]
     if "noise" in input_payload:
         observation["noise"] = input_payload["noise"]
+    for name in ("image", "image_path", "input_reference", "video", "video_path"):
+        if name in input_payload:
+            observation[name] = input_payload[name]
     return _normalize_observation(observation)
 
 
@@ -235,14 +291,14 @@ def _resolve_action_sampling_params_cls_cached(
     backend: str | None,
     model_id: str | None,
     pipeline_class_name: str | None,
-) -> type[VLASamplingParams]:
+) -> type[SamplingParams] | type[ActionSamplingParams]:
     if pipeline_class_name:
         from sglang.multimodal_gen.registry import get_pipeline_config_classes
 
         config_classes = get_pipeline_config_classes(pipeline_class_name)
         if config_classes is not None:
             _, sampling_params_cls = config_classes
-            if issubclass(sampling_params_cls, VLASamplingParams):
+            if issubclass(sampling_params_cls, (SamplingParams, ActionSamplingParams)):
                 return sampling_params_cls
 
     from sglang.multimodal_gen.registry import get_model_info
@@ -253,16 +309,17 @@ def _resolve_action_sampling_params_cls_cached(
         model_id=model_id,
     )
     sampling_params_cls = model_info.sampling_param_cls
-    if not issubclass(sampling_params_cls, VLASamplingParams):
+    if not issubclass(sampling_params_cls, (SamplingParams, ActionSamplingParams)):
         raise ValueError(
-            f"Action endpoint requires VLASamplingParams, got {sampling_params_cls.__name__}"
+            "Action endpoint requires SamplingParams or ActionSamplingParams, got "
+            f"{sampling_params_cls.__name__}"
         )
     return sampling_params_cls
 
 
 def _resolve_action_sampling_params_cls(
     server_args: ServerArgs,
-) -> type[VLASamplingParams]:
+) -> type[SamplingParams] | type[ActionSamplingParams]:
     return _resolve_action_sampling_params_cls_cached(
         server_args.model_path,
         getattr(server_args, "backend", None),
@@ -273,15 +330,16 @@ def _resolve_action_sampling_params_cls(
 
 @lru_cache(maxsize=32)
 def _sampling_params_field_names(
-    sampling_params_cls: type[VLASamplingParams],
+    sampling_params_cls: type[SamplingParams] | type[ActionSamplingParams],
 ) -> frozenset[str]:
     return frozenset(field.name for field in dataclasses.fields(sampling_params_cls))
 
 
-def build_action_sampling_params(
+def _build_action_model_sampling_params(
     payload: dict[str, Any],
     server_args: ServerArgs,
-) -> VLASamplingParams:
+    sampling_params_cls: type[ActionSamplingParams],
+) -> ActionSamplingParams:
     pipeline_config = server_args.pipeline_config
     observation = _action_request_to_observation(payload)
     parameters = dict(payload.get("parameters") or {})
@@ -318,7 +376,6 @@ def build_action_sampling_params(
     if output_format not in ("list", "numpy"):
         raise ValueError("output_format must be 'list' or 'numpy'")
 
-    sampling_params_cls = _resolve_action_sampling_params_cls(server_args)
     sampling_kwargs = {
         "request_id": payload.get("request_id") or payload.get("id"),
         "prompt": prompt,
@@ -366,6 +423,156 @@ def build_action_sampling_params(
     return sp
 
 
+def _cosmos3_image_from_observation(observation: dict[str, Any]) -> Any:
+    image = None
+    for name in ("image", "image_path", "input_reference"):
+        if name in observation:
+            image = observation[name]
+            break
+
+    if image is None:
+        images = observation.get("images")
+        if not images:
+            return None
+        if not isinstance(images, dict) or len(images) != 1:
+            raise ValueError(
+                "Cosmos3 policy input requires exactly one observation image"
+            )
+        image = next(iter(images.values()))
+    if isinstance(image, np.ndarray):
+        if image.dtype != np.uint8:
+            raise ValueError("Cosmos3 observation image arrays must use uint8 dtype")
+        return Image.fromarray(image)
+    return image
+
+
+def _build_cosmos3_action_sampling_params(
+    payload: dict[str, Any],
+    server_args: ServerArgs,
+    sampling_params_cls: type[Cosmos3SamplingParams],
+) -> Cosmos3SamplingParams:
+    observation = _action_request_to_observation(payload)
+    parameters = dict(payload.get("parameters") or {})
+    options = {**observation, **parameters}
+    action_mode = str(options.get("action_mode", "policy")).strip().lower()
+    if action_mode == "forward_dynamics":
+        raise ValueError(
+            "Cosmos3 forward_dynamics produces video; use /v1/videos instead"
+        )
+    if action_mode not in ("policy", "inverse_dynamics"):
+        raise ValueError(
+            "Cosmos3 action endpoint supports action_mode='policy' or "
+            "'inverse_dynamics'"
+        )
+
+    action_horizon = options.get("action_horizon")
+    num_frames = options.get("num_frames")
+    if action_horizon is None and num_frames is None:
+        action_horizon = 16
+    if action_horizon is not None:
+        action_horizon = int(action_horizon)
+        if action_horizon <= 0:
+            raise ValueError("action_horizon must be a positive integer")
+        expected_num_frames = action_horizon + 1
+        if num_frames is not None and int(num_frames) != expected_num_frames:
+            raise ValueError(
+                "Cosmos3 requires num_frames == action_horizon + 1, got "
+                f"num_frames={num_frames}, action_horizon={action_horizon}"
+            )
+        num_frames = expected_num_frames
+    else:
+        num_frames = int(num_frames)
+        if num_frames <= 1:
+            raise ValueError("Cosmos3 action num_frames must be greater than 1")
+    if (num_frames - 1) % 4 != 0:
+        raise ValueError(
+            "Cosmos3 action_horizon must be divisible by 4 so num_frames "
+            "is compatible with the temporal VAE"
+        )
+
+    image_path = _cosmos3_image_from_observation(observation)
+    video_path = options.get("video_path") or observation.get("video")
+    if action_mode == "policy" and image_path is None:
+        raise ValueError("Cosmos3 policy input requires an observation image")
+    if action_mode == "inverse_dynamics" and video_path is None:
+        raise ValueError("Cosmos3 inverse_dynamics input requires an observation video")
+    if image_path is not None and video_path is not None:
+        raise ValueError("Cosmos3 action requests accept either an image or a video")
+
+    domain_id = options.get("domain_id")
+    domain_name = options.get("domain_name")
+    raw_action_dim = options.get("raw_action_dim")
+    if domain_id is None and not domain_name:
+        raise ValueError("Cosmos3 action requests require domain_name or domain_id")
+    if domain_id is not None and not domain_name and raw_action_dim is None:
+        raise ValueError("raw_action_dim is required when only domain_id is provided")
+
+    prompt = observation.get("prompt") or observation.get("task") or ""
+    sampling_kwargs = {
+        "request_id": payload.get("request_id") or payload.get("id"),
+        "prompt": prompt,
+        "image_path": image_path,
+        "video_path": video_path,
+        "action_mode": action_mode,
+        "domain_id": domain_id,
+        "domain_name": domain_name,
+        "raw_action_dim": raw_action_dim,
+        "action_fps": options.get("action_fps"),
+        "action_view_point": options.get("action_view_point", "ego_view"),
+        "action_normalization": options.get("action_normalization", "quantile"),
+        "action_stats_path": server_args.pipeline_config.action_stats_path,
+        "num_frames": num_frames,
+        "fps": int(options.get("fps", 5)),
+        "height": int(options.get("height", 480)),
+        "width": int(options.get("width", 832)),
+        "num_inference_steps": int(options.get("num_inference_steps", 35)),
+        "guidance_scale": float(options.get("guidance_scale", 1.0)),
+        "seed": int(options.get("seed", 42)),
+        "flow_shift": options.get("flow_shift"),
+        "max_sequence_length": options.get("max_sequence_length"),
+        "condition_frame_indexes": options.get("condition_frame_indexes"),
+        "condition_video_keep": options.get("condition_video_keep", "first"),
+        "use_duration_template": False,
+        "use_system_prompt": False,
+        "use_guardrails": options.get("use_guardrails"),
+        "save_output": False,
+        "return_file_paths_only": False,
+        "return_frames": False,
+    }
+    supported_fields = _sampling_params_field_names(sampling_params_cls)
+    sp = sampling_params_cls(
+        **{
+            name: value
+            for name, value in sampling_kwargs.items()
+            if name in supported_fields and value is not None
+        }
+    )
+    sp._adjust(server_args)
+    return sp
+
+
+def build_action_sampling_params(
+    payload: dict[str, Any],
+    server_args: ServerArgs,
+) -> SamplingParams | ActionSamplingParams:
+    sampling_params_cls = _resolve_action_sampling_params_cls(server_args)
+    if issubclass(sampling_params_cls, ActionSamplingParams):
+        return _build_action_model_sampling_params(
+            payload,
+            server_args,
+            sampling_params_cls,
+        )
+    if issubclass(sampling_params_cls, Cosmos3SamplingParams):
+        return _build_cosmos3_action_sampling_params(
+            payload,
+            server_args,
+            sampling_params_cls,
+        )
+    raise ValueError(
+        f"Action endpoint is not implemented for {sampling_params_cls.__name__}"
+    )
+
+
 async def infer_action(
     payload: dict[str, Any],
     server_args: ServerArgs,
@@ -395,6 +602,22 @@ def action_generation_response(
         action_dim = len(actions[0]) if horizon and isinstance(actions[0], list) else 0
         action_shape = [horizon, action_dim]
         action_values = actions
+    action = {
+        "type": "continuous",
+        "dtype": "float32",
+        "shape": action_shape,
+        "values": action_values,
+    }
+    for name in ("action_mode", "domain_id", "raw_action_dim"):
+        if output.get(name) is not None:
+            action[name] = output[name]
+
+    pipeline_config = server_args.pipeline_config
+    if isinstance(pipeline_config, Cosmos3Config):
+        default_num_inference_steps = Cosmos3SamplingParams().num_inference_steps
+    else:
+        default_num_inference_steps = pipeline_config.default_num_inference_steps
+
     response = {
         "id": output.get("request_id") or f"act_{uuid.uuid4().hex}",
         "object": "action.generation",
@@ -405,12 +628,7 @@ def action_generation_response(
                 "index": 0,
                 "input_index": 0,
                 "candidate_index": 0,
-                "action": {
+                "action": action,
-                    "type": "continuous",
-                    "dtype": "float32",
-                    "shape": action_shape,
-                    "values": action_values,
-                },
             }
         ],
         "usage": {
@@ -418,7 +636,7 @@ def action_generation_response(
             "action_dim": action_shape[1] if len(action_shape) > 1 else 0,
             "denoise_steps": output.get("parameters", {}).get(
                 "num_inference_steps",
-                server_args.pipeline_config.default_num_inference_steps,
+                default_num_inference_steps,
             ),
             "prefix_cache_hit": bool(output.get("cache", {}).get("hit", False)),
         },

python/sglang/multimodal_gen/runtime/entrypoints/action/ws_utils.py

@@ -9,7 +9,7 @@ from typing import Any
 
 from fastapi import WebSocket, WebSocketDisconnect
 
-from sglang.multimodal_gen.runtime.entrypoints.vla.protocol import (
+from sglang.multimodal_gen.runtime.entrypoints.action.protocol import (
     action_metadata,
     infer_action,
     pack_msgpack,

python/sglang/multimodal_gen/runtime/entrypoints/http_server.py

@@ -14,6 +14,8 @@ from fastapi import APIRouter, FastAPI, Request, Response
 from fastapi.middleware.cors import CORSMiddleware
 
 from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
+from sglang.multimodal_gen.runtime.entrypoints.action import api as action_api
+from sglang.multimodal_gen.runtime.entrypoints.action import openpi
 from sglang.multimodal_gen.runtime.entrypoints.openai import image_api, video_api
 from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import (
     VertexGenerateReqInput,
@@ -30,8 +32,6 @@ from sglang.multimodal_gen.runtime.entrypoints.utils import (
     prepare_request,
     save_outputs,
 )
-from sglang.multimodal_gen.runtime.entrypoints.vla import api as vla_api
-from sglang.multimodal_gen.runtime.entrypoints.vla import openpi
 from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client
 from sglang.multimodal_gen.runtime.server_args import ServerArgs, get_global_server_args
 from sglang.multimodal_gen.runtime.server_warmup import (
@@ -423,8 +423,9 @@ def create_app(server_args: ServerArgs):
     app.include_router(image_api.router)
     app.include_router(video_api.router)
     app.include_router(realtime_video_api.router)
-    if server_args.pipeline_config.task_type.is_action_gen():
+    if server_args.pipeline_config.supports_action_endpoint():
-        app.include_router(vla_api.router)
+        app.include_router(action_api.router)
+    if server_args.pipeline_config.supports_openpi_endpoint():
         app.include_router(openpi.router)
     app.include_router(mesh_api.router)
     app.include_router(weights_api.router)

python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py

@@ -23,6 +23,7 @@ from fastapi import (
 from fastapi.responses import FileResponse
 
 from sglang.multimodal_gen.configs.sample.sampling_params import (
+    DataType,
     SamplingParams,
     generate_request_id,
 )
@@ -410,7 +411,16 @@ def _build_video_sampling_params(request_id: str, request: VideoGenerationsReque
 
     sampling_params_cls = _video_sampling_params_cls(server_args)
     kwargs = sampling_params_cls.lower_video_request_kwargs(request, kwargs)
-    return build_sampling_params(request_id, **kwargs)
+    sampling_params = build_sampling_params(request_id, **kwargs)
+    if (
+        isinstance(sampling_params, SamplingParams)
+        and sampling_params.data_type == DataType.ACTION
+    ):
+        raise ValueError(
+            "Action-producing policy and inverse-dynamics requests use "
+            "/v1/actions/generations; /v1/videos is reserved for visual outputs"
+        )
+    return sampling_params
 
 
 # extract metadata which http_server needs to know

python/sglang/multimodal_gen/runtime/launch_server.py

@@ -305,9 +305,9 @@ def launch_server(server_args: ServerArgs, launch_http_server: bool = True):
         return processes
 
     if launch_http_server:
-        if server_args.pipeline_config.task_type.is_action_gen():
+        if server_args.pipeline_config.supports_action_endpoint():
             logger.info(
-                "VLA pipeline ready: model=%s; per-request details are "
+                "Action generation endpoint ready: model=%s; per-request details are "
                 "debug-only (use --log-level debug).",
                 server_args.model_id or server_args.model_path,
             )

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py

@@ -56,7 +56,7 @@ from sglang.multimodal_gen.runtime.platforms import current_platform
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.multimodal_gen.runtime.utils.profiler import SGLDiffusionProfiler
-from sglang.multimodal_gen.runtime.utils.vision import load_video
+from sglang.multimodal_gen.runtime.utils.vision import load_image, load_video
 from sglang.srt.utils.common import get_compiler_backend
 
 logger = init_logger(__name__)
@@ -129,8 +129,8 @@ class Cosmos3ImagePreprocessStage(PipelineStage):
 
         target_h, target_w = batch.height, batch.width
 
-        if isinstance(image_path, str) and image_path:
+        if image_path is not None:
-            image = PIL.Image.open(image_path).convert("RGB")
+            image = load_image(image_path)
             image = _resize_crop_pil(image, target_w, target_h)
             batch.preprocessed_image = _pil_to_normalized_tensor(image).unsqueeze(0)
             self.log_info(f"Preprocessed conditioning image to {target_w}x{target_h}")
@@ -440,9 +440,9 @@ class Cosmos3LatentPreparationStage(PipelineStage):
 
         noise = torch.randn(shape, generator=generator, device=device, dtype=dtype)
 
-        is_video_gen = batch.data_type == DataType.VIDEO
+        uses_visual_latents = batch.data_type in (DataType.VIDEO, DataType.ACTION)
-        has_image_cond = batch.preprocessed_image is not None and is_video_gen
+        has_image_cond = batch.preprocessed_image is not None and uses_visual_latents
-        has_video_cond = batch.preprocessed_video is not None and is_video_gen
+        has_video_cond = batch.preprocessed_video is not None and uses_visual_latents
 
         if has_image_cond or has_video_cond:
             vae_dtype = next(self.vae.parameters()).dtype
@@ -1387,6 +1387,56 @@ class Cosmos3DecodingStage(PipelineStage):
             OutputBatch,
         )
 
+        action_pred = None
+        if getattr(batch, "action_latents", None) is not None:
+            raw_action_dim = batch.extra.get("raw_action_dim")
+            action_pred = batch.action_latents.float().cpu()
+            if raw_action_dim is not None:
+                action_pred = action_pred[:, :, :raw_action_dim]
+            stats_path = getattr(batch.sampling_params, "action_stats_path", None)
+            if stats_path is not None:
+                method = getattr(
+                    batch.sampling_params, "action_normalization", "quantile"
+                )
+                action_pred = denormalize_action(
+                    action_pred, method, load_action_stats(stats_path)
+                )
+            self.log_info(f"Action predictions shape: {tuple(action_pred.shape)}")
+
+        action_domain_ids = batch.extra.get("action_domain_ids")
+        action_domain_id = (
+            int(action_domain_ids[0].item()) if action_domain_ids is not None else None
+        )
+        action_metadata = {
+            "action_mode": getattr(batch.sampling_params, "action_mode", None),
+            "action_domain_id": action_domain_id,
+            "action_raw_action_dim": (
+                batch.extra.get("raw_action_dim")
+                if getattr(batch, "extra", None)
+                else None
+            ),
+        }
+        if batch.data_type == DataType.ACTION:
+            if action_pred is None:
+                raise RuntimeError("Cosmos3 action request produced no action tensor")
+            payload = {
+                "request_id": batch.request_id,
+                "actions": action_pred[0].numpy(),
+                "action_mode": action_metadata["action_mode"],
+                "domain_id": action_metadata["action_domain_id"],
+                "raw_action_dim": action_metadata["action_raw_action_dim"],
+                "parameters": {
+                    "num_inference_steps": batch.num_inference_steps,
+                    "num_frames": batch.num_frames,
+                },
+            }
+            return OutputBatch(
+                output=[payload],
+                action_pred=action_pred,
+                metrics=batch.metrics if hasattr(batch, "metrics") else None,
+                **action_metadata,
+            )
+
         is_image_gen = batch.data_type == DataType.IMAGE
         self.log_info(
             "Decoding latents to image..."
@@ -1438,33 +1488,11 @@ class Cosmos3DecodingStage(PipelineStage):
                 f"Decoded audio tensor shape: {tuple(audio.shape)} @ {audio_sample_rate} Hz"
             )
 
-        action_pred = None
-        if getattr(batch, "action_latents", None) is not None:
-            raw_action_dim = batch.extra.get("raw_action_dim")
-            action_pred = batch.action_latents.float().cpu()
-            if raw_action_dim is not None:
-                action_pred = action_pred[:, :, :raw_action_dim]
-            stats_path = getattr(batch.sampling_params, "action_stats_path", None)
-            if stats_path is not None:
-                method = getattr(
-                    batch.sampling_params, "action_normalization", "quantile"
-                )
-                action_pred = denormalize_action(
-                    action_pred, method, load_action_stats(stats_path)
-                )
-            self.log_info(f"Action predictions shape: {tuple(action_pred.shape)}")
-
         return OutputBatch(
             output=output,
             audio=audio,
             audio_sample_rate=audio_sample_rate,
             action_pred=action_pred,
-            action_mode=getattr(batch.sampling_params, "action_mode", None),
-            action_domain_id=getattr(batch.sampling_params, "domain_id", None),
-            action_raw_action_dim=(
-                batch.extra.get("raw_action_dim")
-                if getattr(batch, "extra", None)
-                else None
-            ),
             metrics=batch.metrics if hasattr(batch, "metrics") else None,
+            **action_metadata,
         )

python/sglang/multimodal_gen/test/unit/test_cosmos3.py

@@ -18,6 +18,11 @@ from sglang.multimodal_gen.registry import (
     _get_config_info,
     get_non_diffusers_pipeline_name,
 )
+from sglang.multimodal_gen.runtime.entrypoints.action.protocol import (
+    action_generation_response,
+    action_metadata,
+    build_action_sampling_params,
+)
 from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import (
     ImageGenerationsRequest,
     VideoGenerationsRequest,
@@ -39,6 +44,7 @@ from sglang.multimodal_gen.runtime.models.dits.cosmos3video import (
     compute_mrope_position_ids_vision,
 )
 from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.cosmos3 import (
+    Cosmos3DecodingStage,
     Cosmos3ImagePreprocessStage,
     Cosmos3LatentPreparationStage,
     Cosmos3TimestepPreparationStage,
@@ -56,6 +62,23 @@ def _apply(mapping_fn, key):
     return mapping_fn(key)
 
 
+def _cosmos3_server_args(config=None):
+    return types.SimpleNamespace(
+        model_id=None,
+        model_path="nvidia/Cosmos3-Nano",
+        backend=None,
+        pipeline_class_name=None,
+        output_path=None,
+        comfyui_mode=False,
+        num_gpus=1,
+        tp_size=1,
+        sp_degree=1,
+        ulysses_degree=1,
+        ring_degree=1,
+        pipeline_config=config or Cosmos3Config(),
+    )
+
+
 class TestCosmos3ParamNamesMapping(unittest.TestCase):
     """Verify diffusers → sglang weight key translations."""
 
@@ -328,6 +351,167 @@ class TestCosmos3SamplingParamsDataType(unittest.TestCase):
         params._set_output_file_name()
         self.assertEqual(params.data_type, DataType.VIDEO)
 
+    def test_policy_adjusts_to_action_output(self):
+        params = Cosmos3SamplingParams(
+            prompt="test",
+            action_mode="policy",
+            num_frames=17,
+            image_path="observation.png",
+        )
+
+        params._adjust(_cosmos3_server_args())
+
+        self.assertEqual(params.data_type, DataType.ACTION)
+        self.assertFalse(params.save_output)
+        self.assertFalse(params.return_file_paths_only)
+        self.assertIsNone(params.output_file_name)
+        self.assertEqual(params.num_frames, 17)
+
+    def test_forward_dynamics_remains_video_output(self):
+        params = Cosmos3SamplingParams(
+            prompt="test",
+            action_mode="forward_dynamics",
+            num_frames=17,
+            image_path="observation.png",
+        )
+
+        params._adjust(_cosmos3_server_args())
+
+        self.assertEqual(params.data_type, DataType.VIDEO)
+
+
+class TestCosmos3ActionEndpoint(unittest.TestCase):
+    def test_policy_request_builds_action_sampling_params(self):
+        image = torch.zeros(8, 8, 3, dtype=torch.uint8).numpy()
+        payload = {
+            "request_id": "cosmos-action-1",
+            "input": {
+                "task": "pick up the block",
+                "observation": {
+                    "image": {
+                        "dtype": "uint8",
+                        "shape": [8, 8, 3],
+                        "values": image.tolist(),
+                    }
+                },
+            },
+            "parameters": {
+                "action_mode": "policy",
+                "action_horizon": 16,
+                "domain_name": "droid_lerobot",
+                "num_inference_steps": 30,
+                "height": 480,
+                "width": 832,
+                "fps": 5,
+                "seed": 7,
+            },
+        }
+
+        params = build_action_sampling_params(payload, _cosmos3_server_args())
+
+        self.assertIsInstance(params, Cosmos3SamplingParams)
+        self.assertEqual(params.data_type, DataType.ACTION)
+        self.assertEqual(params.prompt, "pick up the block")
+        self.assertEqual(params.action_mode, "policy")
+        self.assertEqual(params.domain_name, "droid_lerobot")
+        self.assertEqual(params.num_frames, 17)
+        self.assertEqual(params.num_inference_steps, 30)
+        self.assertEqual(params.seed, 7)
+        self.assertEqual(params.image_path.size, (8, 8))
+
+    def test_inverse_dynamics_maps_video_input(self):
+        payload = {
+            "input": {
+                "task": "infer the robot motion",
+                "observation": {"video": "observation.mp4"},
+            },
+            "parameters": {
+                "action_mode": "inverse_dynamics",
+                "num_frames": 61,
+                "domain_name": "av",
+            },
+        }
+
+        params = build_action_sampling_params(payload, _cosmos3_server_args())
+
+        self.assertEqual(params.data_type, DataType.ACTION)
+        self.assertEqual(params.video_path, "observation.mp4")
+        self.assertEqual(params.num_frames, 61)
+
+    def test_forward_dynamics_is_rejected_by_action_endpoint(self):
+        payload = {
+            "input": {
+                "task": "predict the next frames",
+                "observation": {"image": "observation.png"},
+            },
+            "parameters": {"action_mode": "forward_dynamics"},
+        }
+
+        with self.assertRaisesRegex(ValueError, "/v1/videos"):
+            build_action_sampling_params(payload, _cosmos3_server_args())
+
+    def test_metadata_describes_cosmos_action_contract(self):
+        metadata = action_metadata(_cosmos3_server_args())
+
+        self.assertEqual(metadata["policy_family"], "cosmos3")
+        self.assertEqual(metadata["input"]["modalities"], ["image", "video"])
+        self.assertEqual(metadata["output"]["action_horizon"], 16)
+        self.assertEqual(metadata["output"]["padded_action_dim"], 64)
+        self.assertFalse(metadata["capabilities"]["openpi_websocket"])
+
+    def test_action_response_includes_cosmos_metadata(self):
+        output = {
+            "request_id": "cosmos-action-2",
+            "actions": torch.zeros(16, 10).numpy(),
+            "action_mode": "policy",
+            "domain_id": 8,
+            "raw_action_dim": 10,
+            "parameters": {"num_inference_steps": 30},
+        }
+
+        response = action_generation_response(output, _cosmos3_server_args())
+        action = response["data"][0]["action"]
+
+        self.assertEqual(action["shape"], [16, 10])
+        self.assertEqual(action["action_mode"], "policy")
+        self.assertEqual(action["domain_id"], 8)
+        self.assertEqual(action["raw_action_dim"], 10)
+        self.assertEqual(response["usage"]["denoise_steps"], 30)
+
+    def test_action_decode_skips_vae(self):
+        class FailIfDecoded:
+            def decode(self, _latents):
+                raise AssertionError("VAE decode must not run for action output")
+
+        stage = Cosmos3DecodingStage.__new__(Cosmos3DecodingStage)
+        stage.vae = FailIfDecoded()
+        stage.sound_tokenizer = None
+        stage._guardrails = False
+        stage.log_info = lambda *_args, **_kwargs: None
+        batch = types.SimpleNamespace(
+            data_type=DataType.ACTION,
+            action_latents=torch.arange(24, dtype=torch.float32).reshape(1, 4, 6),
+            extra={
+                "raw_action_dim": 3,
+                "action_domain_ids": torch.tensor([8]),
+            },
+            sampling_params=Cosmos3SamplingParams(
+                prompt="test",
+                action_mode="policy",
+                domain_name="droid_lerobot",
+            ),
+            request_id="cosmos-action-3",
+            num_inference_steps=30,
+            num_frames=5,
+            metrics=None,
+        )
+
+        output = stage.forward(batch, types.SimpleNamespace(vae_cpu_offload=False))
+
+        self.assertEqual(output.output[0]["actions"].shape, (4, 3))
+        self.assertEqual(output.output[0]["domain_id"], 8)
+        self.assertEqual(output.action_pred.shape, (1, 4, 3))
+
 
 class TestCosmos3ModelResolution(unittest.TestCase):
     """Verify Cosmos3 checkpoints resolve to the native SGLang pipeline."""
@@ -335,6 +519,7 @@ class TestCosmos3ModelResolution(unittest.TestCase):
     def test_hf_checkpoint_uses_registered_native_pipeline_config(self):
         for model_path in (
             "nvidia/Cosmos3-Nano",
+            "nvidia/Cosmos3-Nano-Policy-DROID",
             "nvidia/Cosmos3-Super",
             "nvidia/Cosmos3-Super-Text2Image",
             "nvidia/Cosmos3-Super-Image2Video",

python/sglang/multimodal_gen/test/unit/test_pi05_action_api.py

@@ -6,13 +6,13 @@ from types import SimpleNamespace
 import numpy as np
 
 from sglang.multimodal_gen.configs.pipeline_configs.pi05 import Pi05PipelineConfig
+from sglang.multimodal_gen.configs.sample.action import ActionSamplingParams
 from sglang.multimodal_gen.configs.sample.pi05 import Pi05SamplingParams
 from sglang.multimodal_gen.configs.sample.sampling_params import (
     DataType,
     SamplingParams,
 )
-from sglang.multimodal_gen.configs.sample.vla import VLASamplingParams
+from sglang.multimodal_gen.runtime.entrypoints.action.protocol import (
-from sglang.multimodal_gen.runtime.entrypoints.vla.protocol import (
     action_generation_response,
     action_metadata,
     action_raw_response,
@@ -38,11 +38,11 @@ def _server_args(config: Pi05PipelineConfig | None = None) -> SimpleNamespace:
     )
 
 
-def test_pi05_uses_vla_sampling_params_not_visual_sampling_params():
+def test_pi05_uses_action_sampling_params_not_visual_sampling_params():
     params = Pi05SamplingParams()
     field_names = {field.name for field in dataclasses.fields(params)}
 
-    assert isinstance(params, VLASamplingParams)
+    assert isinstance(params, ActionSamplingParams)
     assert not isinstance(params, SamplingParams)
     assert "action_horizon" in field_names
     assert "action_dim" in field_names