Cosmos3 action policy improvements (#6460)

Signed-off-by: Maciej Bala <mbala@nvidia.com>
Signed-off-by: MaciejBalaNV <mbala@nvidia.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

recipes/cosmos3/Cosmos3-Nano.md

@@ -50,9 +50,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.
+- **DROID policy server** — serve `nvidia/Cosmos3-Nano-Policy-DROID` and connect
+  RoboLab or another OpenPI-compatible client to
+  `/v1/realtime/robot/openpi`. This endpoint 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`
@@ -314,40 +315,71 @@ VIDEO_ID=$(curl -sS -X POST http://localhost:8000/v1/videos \
 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
+# DROID websocket policy server. Use the tested cosmos-framework revision
+# directly from source; installing the package can introduce dependency conflicts.
+export COSMOS_FRAMEWORK_ROOT=/path/to/cosmos-framework
+git -C "$COSMOS_FRAMEWORK_ROOT" checkout c14617c2bc93dacbf69674fb964eec93182933d9
+export PYTHONPATH="$COSMOS_FRAMEWORK_ROOT"
 
+# Validate every cosmos-framework symbol used by the action-policy pipeline
+# before allocating the model.
+python - <<'PY'
+from vllm_omni.diffusion.models.cosmos3.utils import (
+    get_robolab_domain_id,
+    preflight_cosmos3_action_framework_imports,
+)
+
+preflight_cosmos3_action_framework_imports()
+print("cosmos-framework action imports OK; DROID domain:", get_robolab_domain_id("droid_lerobot"))
+PY
+
+# The bundled deploy config (vllm_omni/deploy/cosmos3_policy_droid.yaml)
+# selects the registered cosmos3_policy pipeline and carries the DROID
+# checkpoint's serving defaults: JSON prompt formatting (a property of this
+# checkpoint's training recipe, not a generic server default), the
+# model-specific OpenPI binary handshake metadata (policy_server_config), and
+# guardrails off (policy serving emits robot actions; the guardrail stack is
+# not part of it, so --no-guardrails is implied).
+# Policy checkpoints cannot be auto-detected — they share their HF metadata
+# with the T2I/video Cosmos3 checkpoints — so --deploy-config is required.
+# Per-checkpoint model_config tweaks can be layered on top with
+# --stage-overrides; dict overrides deep-merge with the deploy yaml.
+export VLLM_OMNI_ROOT=/absolute/path/to/vllm-omni
 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)"
+  --deploy-config "$VLLM_OMNI_ROOT/vllm_omni/deploy/cosmos3_policy_droid.yaml" \
+  --robot-openpi-idle-timeout 0
 
-# 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.
+# From a RoboLab checkout:
+python policies/cosmos3/run.py \
+  --remote-uri ws://localhost:8000/v1/realtime/robot/openpi \
+  --task BananaInBowlTask
+
+# For a server started with "--api-key $POLICY_API_KEY", pass the same value as
+# a Bearer token. --remote-token takes precedence over COSMOS3_API_TOKEN.
+COSMOS3_API_TOKEN="$POLICY_API_KEY" python policies/cosmos3/run.py \
+  --remote-uri wss://policy.example/v1/realtime/robot/openpi \
+  --task BananaInBowlTask
 ```
 
+The endpoint sends `policy_server_config` as its initial binary MsgPack
+message and returns the action array directly. RoboLab accepts that direct
+array as well as proxy responses shaped as `{"action": ...}` or
+`{"actions": ...}`, and turns structured server errors into client exceptions.
+
+RoboLab supplies `session_id=robolab-episode-<episode>-env-<env_id>`, so
+parallel environments have independent policy state even though they share
+one WebSocket. The server tracks interleaved session IDs independently.
+The default receive-idle timeout is 30 seconds; set
+`--robot-openpi-idle-timeout 0` when simulator steps between replans can take
+longer, or set another non-negative timeout in seconds.
+
+If vLLM-Omni is started with `--api-key` or `VLLM_API_KEY`, the standard
+OpenPI route uses the normal API authentication middleware. RoboLab's
+`--remote-token` and `COSMOS3_API_TOKEN` send the required
+`Authorization: Bearer ...` header.
+
 #### Notes
 
 - **Measured latency (1x B300, bf16, guardrails off):**

tests/config/test_config_factory.py

@@ -197,6 +197,25 @@ class TestStageConfig:
         omega_config = config.to_omegaconf()
         assert omega_config.engine_args.max_num_seqs == 32
 
+    def test_to_omegaconf_dict_override_deep_merges_yaml_dict(self):
+        """A partial dict CLI override (e.g. --no-guardrails riding on
+        model_config) must layer onto the YAML dict, not clobber siblings."""
+        config = StageConfig(
+            stage_id=0,
+            model_stage="diffusion",
+            stage_type=StageType.DIFFUSION,
+            yaml_engine_args={
+                "model_config": {
+                    "guardrails": True,
+                    "policy_server_config": {"action_space": "joint_position"},
+                }
+            },
+            runtime_overrides={"model_config": {"guardrails": False}},
+        )
+        omega_config = config.to_omegaconf()
+        assert omega_config.engine_args.model_config.guardrails is False
+        assert omega_config.engine_args.model_config.policy_server_config.action_space == "joint_position"
+
     def test_to_omegaconf_diffusion_parallel_overrides_replace_nested_values(self):
         config = StageConfig(
             stage_id=1,
@@ -623,6 +642,58 @@ class TestPipelineDiscovery:
         assert p.hf_architectures == ("SomeCollidingArch",)
 
 
+class TestCosmos3PolicyPipeline:
+    """Cosmos3 policy serving resolves via deploy yaml selection (π0 precedent)."""
+
+    def test_registered_without_capturing_base_cosmos3_checkpoints(self):
+        assert "cosmos3_policy" in OMNI_PIPELINES
+        # T2I/video Cosmos3 checkpoints report model_type=cosmos3_omni and the
+        # same model_index.json _class_name as policy checkpoints; they must
+        # keep resolving through the single-stage diffusion fallback, so the
+        # policy pipeline must not be reachable by auto-detection.
+        assert "cosmos3_omni" not in OMNI_PIPELINES
+        pipeline = OMNI_PIPELINES["cosmos3_policy"]
+        assert pipeline.hf_architectures == ()
+        assert pipeline.diffusers_class_name is None
+
+    def test_droid_deploy_yaml_carries_policy_server_config(self):
+        deploy = load_deploy_config(get_deploy_config_path("cosmos3_policy_droid.yaml"))
+        assert deploy.pipeline == "cosmos3_policy"
+
+        stages = merge_pipeline_deploy(OMNI_PIPELINES["cosmos3_policy"], deploy)
+        assert len(stages) == 1
+        stage = stages[0].to_omegaconf()
+
+        assert stage.stage_type == "diffusion"
+        assert stage.final_output_type == "action"
+        assert stage.engine_args.model_class_name == "Cosmos3OmniDiffusersPipeline"
+        # OpenPI websocket handshake metadata must reach the serving layer via
+        # engine_args.model_config (ServingRealtimeRobotOpenPI reads it there).
+        psc = stage.engine_args.model_config.policy_server_config
+        assert list(psc.image_resolution) == [540, 640]
+        assert psc.n_external_cameras == 2
+        assert psc.needs_wrist_camera is True
+        assert psc.needs_stereo_camera is False
+        assert psc.needs_session_id is True
+        assert psc.action_space == "joint_position"
+        # The DROID checkpoint's training recipe expects JSON-formatted prompts.
+        assert stage.default_sampling_params.extra_args.format_prompt_as_json is True
+
+    def test_no_guardrails_cli_flag_keeps_policy_server_config(self):
+        """--no-guardrails becomes a partial model_config CLI override; it must
+        deep-merge with the deploy yaml's model_config, not clobber the
+        policy_server_config the OpenPI handshake depends on."""
+        deploy = load_deploy_config(get_deploy_config_path("cosmos3_policy_droid.yaml"))
+        stages, _ = StageConfigFactory._create_legacy_from_registry(
+            OMNI_PIPELINES["cosmos3_policy"],
+            {"model_config": {"guardrails": False}},
+            user_deploy_config=deploy,
+        )
+        stage = stages[0].to_omegaconf()
+        assert stage.engine_args.model_config.guardrails is False
+        assert stage.engine_args.model_config.policy_server_config.action_space == "joint_position"
+
+
 class TestStagePipelineConfig:
     def test_frozen(self):
         s = StagePipelineConfig(stage_id=0, model_stage="a")

tests/diffusion/models/cosmos3/test_cosmos3_framework_compat.py

@@ -0,0 +1,144 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from importlib.util import find_spec
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+_FRAMEWORK_AVAILABLE = find_spec("cosmos_framework") is not None
+requires_cosmos_framework = pytest.mark.skipif(
+    not _FRAMEWORK_AVAILABLE,
+    reason="cosmos_framework source is not present on PYTHONPATH",
+)
+
+
+def test_framework_preflight_checks_every_policy_symbol(monkeypatch) -> None:
+    from vllm_omni.diffusion.models.cosmos3 import utils as cosmos3_utils
+
+    imported = []
+
+    def record_import(module_name: str, symbol_name: str, error_message: str):
+        imported.append((module_name, symbol_name, error_message))
+        return object()
+
+    monkeypatch.setattr(cosmos3_utils, "lazy_import", record_import)
+
+    cosmos3_utils.preflight_cosmos3_action_framework_imports()
+
+    expected_symbols = {
+        "ActionTransformPipeline",
+        "FlowUniPCMultistepScheduler",
+        "build_abs_pose_from_components",
+        "convert_rotation",
+        "get_domain_id",
+        "pose_abs_to_rel",
+        "pose_rel_to_abs",
+    }
+    assert {symbol_name for _, symbol_name, _ in imported} == expected_symbols
+    assert [(module_name, symbol_name) for module_name, symbol_name, _ in imported] == list(
+        cosmos3_utils._ROBOLAB_FRAMEWORK_IMPORTS.values()
+    )
+
+
+def test_lazy_import_reports_framework_symbol_skew(monkeypatch) -> None:
+    from vllm_omni.diffusion.models.cosmos3 import utils as cosmos3_utils
+
+    monkeypatch.setattr(cosmos3_utils, "import_module", lambda _module_name: object())
+
+    with pytest.raises(ImportError, match="MissingSymbol") as exc_info:
+        cosmos3_utils.lazy_import(
+            "cosmos_framework.fake_module",
+            "MissingSymbol",
+            "Cosmos Framework compatibility error.",
+        )
+
+    assert isinstance(exc_info.value.__cause__, AttributeError)
+
+
+@requires_cosmos_framework
+def test_action_domain_table_matches_current_framework_and_legacy_aliases() -> None:
+    from cosmos_framework.data.generator.action.utils.domain_utils import (
+        EMBODIMENT_TO_DOMAIN_ID as FRAMEWORK_DOMAIN_IDS,
+    )
+
+    from vllm_omni.diffusion.models.cosmos3.action import EMBODIMENT_TO_DOMAIN_ID as VLLM_DOMAIN_IDS
+    from vllm_omni.diffusion.models.cosmos3.utils import _LEGACY_EMBODIMENT_ALIASES
+
+    canonical_vllm_domain_ids = {
+        name: domain_id for name, domain_id in VLLM_DOMAIN_IDS.items() if name not in _LEGACY_EMBODIMENT_ALIASES
+    }
+
+    assert canonical_vllm_domain_ids == FRAMEWORK_DOMAIN_IDS
+    assert {alias: VLLM_DOMAIN_IDS[alias] for alias in _LEGACY_EMBODIMENT_ALIASES} == {
+        alias: FRAMEWORK_DOMAIN_IDS[canonical_name] for alias, canonical_name in _LEGACY_EMBODIMENT_ALIASES.items()
+    }
+
+
+@requires_cosmos_framework
+def test_current_cosmos_framework_policy_modules_import_and_instantiate() -> None:
+    from cosmos_framework.data.generator.action.utils.domain_utils import get_domain_id
+    from cosmos_framework.data.generator.action.utils.pose_utils import (
+        build_abs_pose_from_components,
+        convert_rotation,
+        pose_abs_to_rel,
+        pose_rel_to_abs,
+    )
+    from cosmos_framework.data.generator.action.utils.transforms import ActionTransformPipeline
+    from cosmos_framework.model.generator.diffusion.samplers.fm_solvers_unipc import (
+        FlowUniPCMultistepScheduler,
+    )
+
+    transform = ActionTransformPipeline(
+        max_action_dim=64,
+        cfg_dropout_rate=0.0,
+        format_prompt_as_json=True,
+    )
+    scheduler = FlowUniPCMultistepScheduler(
+        num_train_timesteps=1000,
+        shift=1.0,
+        use_dynamic_shifting=False,
+    )
+    scheduler.set_timesteps(4, device=torch.device("cpu"), shift=5.0)
+
+    assert transform.prompt_json_formatter is not None
+    assert len(scheduler.timesteps) == 4
+    assert get_domain_id("droid_lerobot") == 8
+    assert convert_rotation([0.0, 0.0, 0.0, 1.0], "quat_xyzw", "rot6d").shape == (6,)
+    assert callable(pose_abs_to_rel)
+    assert callable(pose_rel_to_abs)
+    assert callable(build_abs_pose_from_components)
+
+
+@requires_cosmos_framework
+def test_real_action_transform_pipeline_matches_droid_policy_contract() -> None:
+    from cosmos_framework.data.generator.action.utils.transforms import ActionTransformPipeline
+
+    transform = ActionTransformPipeline(
+        max_action_dim=64,
+        cfg_dropout_rate=0.0,
+        format_prompt_as_json=True,
+    )
+    sample = {
+        "ai_caption": "Pick up the cube.",
+        "video": torch.zeros((3, 3, 192, 320), dtype=torch.uint8),
+        "action": torch.zeros((3, 8), dtype=torch.float32),
+        "conditioning_fps": torch.tensor(15, dtype=torch.long),
+        "mode": "wam",
+        "domain_id": torch.tensor(8, dtype=torch.long),
+        "viewpoint": "concat_view",
+        "additional_view_description": "Wrist view above two exterior views.",
+    }
+
+    transformed = transform(sample, "256")
+
+    assert isinstance(transformed["ai_caption"], dict)
+    assert transformed["action"].shape == (3, 64)
+    assert int(transformed["raw_action_dim"].item()) == 8
+    assert transformed["sequence_plan"].has_action is True
+    assert transformed["mode"] == "wam"
+    assert int(transformed["domain_id"].item()) == 8

tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py

@@ -3,6 +3,7 @@
 
 from __future__ import annotations
 
+import json
 import sys
 import types
 from dataclasses import dataclass
@@ -21,6 +22,20 @@ from vllm_omni.experimental.world_models.session_state import SessionStateManage
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
 
+@pytest.mark.parametrize(
+    ("alias", "canonical_name"),
+    [
+        ("galbot", "embodiment_b"),
+        ("agibot_gear_gripper", "embodiment_c_gripper"),
+        ("agibot_gear_gripper_ext", "embodiment_c_gripper_ext"),
+    ],
+)
+def test_action_domain_table_preserves_legacy_aliases(alias: str, canonical_name: str) -> None:
+    from vllm_omni.diffusion.models.cosmos3.action import resolve_domain_id
+
+    assert resolve_domain_id(domain_name=alias) == resolve_domain_id(domain_name=canonical_name)
+
+
 def test_pipeline_declares_layerwise_offload_components() -> None:
     from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline
 
@@ -479,6 +494,69 @@ def test_forward_threads_request_id_to_robolab(make_cosmos3_pipeline) -> None:
     assert captured["session_id"] == "robolab-request-7"
 
 
+@pytest.mark.parametrize("format_prompt_as_json", [False, True])
+def test_robolab_input_builder_threads_prompt_format_and_uses_wam(
+    make_cosmos3_pipeline,
+    monkeypatch: pytest.MonkeyPatch,
+    format_prompt_as_json: bool,
+) -> None:
+    from vllm_omni.diffusion.models.cosmos3 import pipeline_cosmos3
+
+    pipeline = make_cosmos3_pipeline()
+    pipeline.transformer = StubCosmos3Transformer(action_gen=True, action_dim=64)
+    captured: dict[str, Any] = {}
+
+    def fake_transform(sample, resolution):
+        captured["sample_mode"] = sample["mode"]
+        captured["resolution"] = resolution
+        sample["sequence_plan"] = SimpleNamespace(
+            condition_frame_indexes_action=[0],
+            action_start_frame_offset=1,
+        )
+        sample["raw_action_dim"] = torch.tensor(8)
+        sample["image_size"] = torch.tensor([16, 16, 16, 16])
+        if format_prompt_as_json:
+            sample["ai_caption"] = {"actions": {"instruction": sample["ai_caption"]}}
+        return sample
+
+    def fake_get_transform(*, format_prompt_as_json: bool):
+        captured["format_prompt_as_json"] = format_prompt_as_json
+        return fake_transform
+
+    pipeline._get_robolab_transform = fake_get_transform
+    monkeypatch.setattr(pipeline_cosmos3, "get_robolab_domain_id", lambda name: 8)
+    obs = {
+        "prompt": "Pick up the cube.",
+        "observation/image": np.zeros((16, 16, 3), dtype=np.uint8),
+        "observation/joint_position": np.zeros(7, dtype=np.float32),
+        "observation/gripper_position": np.zeros(1, dtype=np.float32),
+    }
+    sampling_params = make_sampling_params(
+        extra_args={
+            "robot_obs": obs,
+            "action_chunk_size": 2,
+            "image_height": 16,
+            "image_width": 16,
+            "format_prompt_as_json": format_prompt_as_json,
+        }
+    )
+
+    inputs = pipeline._build_robolab_policy_inputs(sampling_params, request_id="request-1")
+
+    assert inputs is not None
+    assert captured == {
+        "sample_mode": "wam",
+        "resolution": "480",
+        "format_prompt_as_json": format_prompt_as_json,
+    }
+    assert inputs.domain_id == 8
+    assert inputs.raw_action_dim == 8
+    if format_prompt_as_json:
+        assert json.loads(inputs.prompt) == {"actions": {"instruction": "Pick up the cube."}}
+    else:
+        assert inputs.prompt == "Pick up the cube."
+
+
 @pytest.mark.parametrize(
     ("prompt", "sampling_params", "message"),
     [

tests/entrypoints/openai_api/test_openpi_connection.py

@@ -310,6 +310,29 @@ def test_handle_connection_closes_websocket_on_idle_timeout(monkeypatch):
     serving.infer.assert_not_called()
 
 
+def test_handle_connection_can_disable_idle_timeout(monkeypatch):
+    monkeypatch.setattr(openpi_connection, "_pack", lambda obj: obj)
+
+    async def fail_wait_for(*_args, **_kwargs):
+        pytest.fail("asyncio.wait_for should not be used when the idle timeout is disabled")
+
+    monkeypatch.setattr(openpi_connection.asyncio, "wait_for", fail_wait_for)
+    websocket = FakeWebSocket([{"type": "websocket.disconnect"}])
+    serving = _serving_mock()
+
+    asyncio.run(
+        openpi_connection.RobotRealtimeConnection(
+            websocket,
+            serving,
+            idle_timeout=None,
+        ).handle_connection()
+    )
+
+    assert websocket.accepted is True
+    assert websocket.closed is False
+    serving.infer.assert_not_called()
+
+
 def test_handle_connection_keeps_session_state_per_websocket(monkeypatch):
     monkeypatch.setattr(openpi_connection, "_pack", lambda obj: obj)
     requests = {
@@ -343,6 +366,73 @@ def test_handle_connection_keeps_session_state_per_websocket(monkeypatch):
     assert calls[2].kwargs == {"session_id": "session-b", "reset": True}
 
 
+def test_handle_connection_keeps_interleaved_parallel_session_state(monkeypatch):
+    monkeypatch.setattr(openpi_connection, "_pack", lambda obj: obj)
+    requests = {
+        b"a1": {"prompt": "first-a", "session_id": "session-a"},
+        b"b1": {"prompt": "first-b", "session_id": "session-b"},
+        b"a2": {"prompt": "second-a", "session_id": "session-a"},
+        b"b2": {"prompt": "second-b", "session_id": "session-b"},
+    }
+    monkeypatch.setattr(openpi_connection, "_unpack", lambda data: dict(requests[data]))
+    serving = _serving_mock()
+    websocket = FakeWebSocket(
+        [{"type": "websocket.receive", "bytes": key} for key in (b"a1", b"b1", b"a2", b"b2")]
+        + [{"type": "websocket.disconnect"}]
+    )
+
+    asyncio.run(openpi_connection.RobotRealtimeConnection(websocket, serving).handle_connection())
+
+    assert [call.kwargs for call in serving.infer.await_args_list] == [
+        {"session_id": "session-a", "reset": True},
+        {"session_id": "session-b", "reset": True},
+        {"session_id": "session-a", "reset": False},
+        {"session_id": "session-b", "reset": False},
+    ]
+
+
+def test_handle_connection_bounds_tracked_sessions_without_resetting_live_ones(monkeypatch):
+    """A client may hold one socket across many episodes, minting a fresh
+    session id per episode and never sending ``reset`` (RoboLab does exactly
+    this). Session tracking must stay bounded, and evicting finished episodes
+    must not disturb the sessions that are still running.
+    """
+    monkeypatch.setattr(openpi_connection, "_pack", lambda obj: obj)
+    monkeypatch.setattr(openpi_connection, "MAX_TRACKED_SESSIONS", 8)
+
+    # One long-lived session interleaved with far more short episodes than the cap.
+    episodes = 40
+    keys = []
+    requests = {}
+    for episode in range(episodes):
+        for session_id in (f"episode-{episode}", "live-env"):
+            key = f"{session_id}@{episode}".encode()
+            keys.append(key)
+            requests[key] = {"prompt": "pick", "session_id": session_id}
+
+    monkeypatch.setattr(openpi_connection, "_unpack", lambda data: dict(requests[data]))
+    serving = _serving_mock()
+    websocket = FakeWebSocket(
+        [{"type": "websocket.receive", "bytes": key} for key in keys] + [{"type": "websocket.disconnect"}]
+    )
+    connection = openpi_connection.RobotRealtimeConnection(websocket, serving)
+
+    asyncio.run(connection.handle_connection())
+
+    assert len(connection._seen_sessions) <= 8
+
+    resets_by_session: dict[str, list[bool]] = {}
+    for call in serving.infer.await_args_list:
+        resets_by_session.setdefault(call.kwargs["session_id"], []).append(call.kwargs["reset"])
+
+    # Each new episode resets exactly once, on its first observation.
+    for episode in range(episodes):
+        assert resets_by_session[f"episode-{episode}"] == [True]
+    # The continuously active session is never evicted, so it resets only once
+    # despite ``episodes`` other sessions churning past the cap.
+    assert resets_by_session["live-env"] == [True] + [False] * (episodes - 1)
+
+
 def test_handle_connection_reset_endpoint_resets_next_infer(monkeypatch):
     monkeypatch.setattr(openpi_connection, "_pack", lambda obj: obj)
     requests = {

tests/entrypoints/openai_api/test_openpi_serving.py

@@ -15,6 +15,7 @@ from starlette.testclient import TestClient
 
 from vllm_omni.entrypoints.openpi import connection as openpi_connection
 from vllm_omni.entrypoints.openpi import serving as openpi_serving
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -177,6 +178,22 @@ def test_create_policy_server_returns_none_without_policy_config():
     assert serving is None
 
 
+def test_legacy_policy_request_defaults_key_no_longer_controls_serving():
+    od_config = SimpleNamespace(
+        model_config={
+            "policy_server_config": {},
+            "policy_request_defaults": "malformed legacy value",
+        }
+    )
+
+    serving = openpi_serving.ServingRealtimeRobotOpenPI.create_policy_server(
+        engine_client=SimpleNamespace(get_diffusion_od_config=lambda: od_config),
+        model_name="generic-policy",
+    )
+
+    assert serving is not None
+
+
 def test_policy_server_config_allows_explicit_empty_config():
     serving = openpi_serving.ServingRealtimeRobotOpenPI(
         engine_client=_engine_with_policy_config(policy_config={}),
@@ -221,6 +238,44 @@ def test_build_request_uses_unique_engine_request_id_per_inference():
     assert request_a.request_id != request_b.request_id
 
 
+def test_build_request_clones_stage_defaults_before_protocol_fields():
+    default_params = OmniDiffusionSamplingParams(
+        extra_args={
+            "format_prompt_as_json": True,
+            "session_id": "configured-session-must-not-win",
+            "reset": False,
+            "nested": {"values": []},
+        }
+    )
+    od_config = SimpleNamespace(model_config={"policy_server_config": {}})
+    engine_client = SimpleNamespace(
+        get_diffusion_od_config=lambda: od_config,
+        default_sampling_params_list=[default_params],
+    )
+    serving = openpi_serving.ServingRealtimeRobotOpenPI(
+        engine_client=engine_client,
+        model_name="custom-policy",
+    )
+
+    request_a = serving._build_request(
+        {"prompt": "pick up the object"},
+        session_id="wire-session",
+        reset=True,
+    )
+    request_b = serving._build_request(
+        {"prompt": "pick up the object"},
+        session_id="wire-session",
+        reset=False,
+    )
+
+    assert request_a.sampling_params.extra_args["format_prompt_as_json"] is True
+    assert request_a.sampling_params.extra_args["session_id"] == "wire-session"
+    assert request_a.sampling_params.extra_args["reset"] is True
+    request_a.sampling_params.extra_args["nested"]["values"].append("request-a")
+    assert request_b.sampling_params.extra_args["nested"] == {"values": []}
+    assert default_params.extra_args["nested"] == {"values": []}
+
+
 def test_build_request_forwards_seed_to_sampling_params():
     """``seed`` in the inference message is the engine-level seed; omitted, the request auto-seeds."""
     serving = openpi_serving.ServingRealtimeRobotOpenPI(engine_client=_engine_with_policy_config())
@@ -257,6 +312,29 @@ def test_infer_keeps_session_state_but_uses_unique_engine_request_ids():
     assert sampling_params_b.extra_args["reset"] is False
 
 
+def test_infer_yields_event_loop_while_engine_is_running():
+    class DelayedEngine(RecordingEngine):
+        def generate(self, *, prompt, request_id, sampling_params_list):
+            async def _generate():
+                await asyncio.sleep(0.03)
+                yield SimpleNamespace(multimodal_output={"actions": [0.0]})
+
+            return _generate()
+
+    serving = openpi_serving.ServingRealtimeRobotOpenPI(engine_client=DelayedEngine())
+
+    async def run_test():
+        infer_task = asyncio.create_task(serving.infer({"prompt": "pick"}, session_id="session-a", reset=True))
+        ticks = 0
+        while not infer_task.done():
+            await asyncio.sleep(0.005)
+            ticks += 1
+        await infer_task
+        return ticks
+
+    assert asyncio.run(run_test()) >= 2
+
+
 def test_two_websocket_clients_without_session_id_do_not_conflict(monkeypatch):
     monkeypatch.setattr(openpi_connection, "_pack", _json_pack)
     monkeypatch.setattr(openpi_connection, "_unpack", _json_unpack)

tests/entrypoints/test_serve.py

@@ -112,6 +112,26 @@ def test_serve_parser_accepts_four_way_cfg_parallelism() -> None:
     assert args.cfg_parallel_size == 4
 
 
+def test_serve_parser_accepts_robot_openpi_idle_timeout() -> None:
+    parser = TrackingArgumentParser()
+    subparsers = parser.add_subparsers(dest="subcommand")
+    OmniServeCommand().subparser_init(subparsers)
+
+    args = parser.parse_args(["serve", "fake-model", "--omni", "--robot-openpi-idle-timeout", "0"])
+
+    assert args.robot_openpi_idle_timeout == 0
+    assert args.get_explicit_kwargs_dict()["robot_openpi_idle_timeout"] == 0
+
+
+def test_serve_parser_rejects_negative_robot_openpi_idle_timeout() -> None:
+    parser = TrackingArgumentParser()
+    subparsers = parser.add_subparsers(dest="subcommand")
+    OmniServeCommand().subparser_init(subparsers)
+
+    with pytest.raises(SystemExit):
+        parser.parse_args(["serve", "fake-model", "--omni", "--robot-openpi-idle-timeout", "-1"])
+
+
 def test_serve_parser_accepts_ulysses_a2a_permute() -> None:
     parser = TrackingArgumentParser()
     subparsers = parser.add_subparsers(dest="subcommand")

tests/model_extras/test_model_extras.py

@@ -177,6 +177,7 @@ def test_cosmos3_extra_registry_declares_request_and_response_params(pipeline_na
             "resolution",
             "image_size",
             "use_state",
+            "format_prompt_as_json",
             "observation",
             "robot_obs",
             "deterministic_seed",

vllm_omni/config/pipeline_registry.py

@@ -46,6 +46,7 @@ from vllm_omni.model_executor.models.bagel.pipeline import (
     BAGEL_SINGLE_STAGE_PIPELINE,
     BAGEL_THINK_PIPELINE,
 )
+from vllm_omni.model_executor.models.cosmos3.pipeline import COSMOS3_POLICY_PIPELINE
 from vllm_omni.model_executor.models.cosyvoice3.pipeline import COSYVOICE3_PIPELINE
 from vllm_omni.model_executor.models.covo_audio.pipeline import COVO_AUDIO_PIPELINE
 from vllm_omni.model_executor.models.dots_tts.pipeline import DOTS_TTS_PIPELINE
@@ -149,6 +150,11 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
     "lingbot_world": LINGBOT_WORLD_PIPELINE,
     "Gr00tN1d7": GR00T_N1D7_PIPELINE,
     "pi0": PI0_PIPELINE,
+    # Cosmos3 policy checkpoints share HF metadata with the T2I/video Cosmos3
+    # checkpoints (which stay on the single-stage diffusion fallback), so this
+    # entry is only reachable through a deploy yaml's ``pipeline:`` key
+    # (see deploy/cosmos3_policy_droid.yaml).
+    "cosmos3_policy": COSMOS3_POLICY_PIPELINE,
     "gepard": GEPARD_PIPELINE,
     "glm_image": GLM_IMAGE_PIPELINE,
     "hunyuan_image_3_moe": HUNYUAN_IMAGE3_PIPELINE,

vllm_omni/config/stage_config.py

@@ -1130,10 +1130,22 @@ class StageConfig:
             _apply_diffusion_parallel_runtime_overrides(engine_args, runtime_overrides)
             reconcile_diffusion_attention_overrides(engine_args, runtime_overrides)
 
-        # CLI overrides take precedence over YAML defaults
+        # CLI overrides take precedence over YAML defaults. Most dict-valued
+        # overrides are deep-merged so a partial CLI dict (e.g. --no-guardrails
+        # riding on ``model_config``) layers onto the deploy YAML instead of
+        # clobbering sibling keys such as ``policy_server_config`` — the same
+        # rationale as the platform-overlay deep-merge. Legacy atomic mappings
+        # are handled explicitly below.
         for key, value in runtime_overrides.items():
             if value is not None and key not in ("devices", "max_batch_size", "num_replicas"):
-                engine_args[key] = value
+                existing = engine_args.get(key)
+                # ``omni_kv_config`` is an atomic legacy override: callers use
+                # a partial mapping to replace the topology-provided transfer
+                # role, rather than to add fields to it.
+                if key != "omni_kv_config" and isinstance(existing, dict) and isinstance(value, dict):
+                    engine_args[key] = _get_recursively_merged_dict(existing, value)
+                else:
+                    engine_args[key] = value
 
         # Build runtime config from YAML defaults + CLI overrides
         runtime: dict[str, Any] = dict(self.yaml_runtime)

vllm_omni/deploy/cosmos3_policy_droid.yaml

@@ -0,0 +1,49 @@
+# Cosmos3 policy (DROID) deploy: single diffusion stage serving the OpenPI
+# realtime API (/v1/realtime/robot/openpi).
+#
+# Topology (COSMOS3_POLICY_PIPELINE) is declared in
+# vllm_omni/model_executor/models/cosmos3/pipeline.py. Policy checkpoints share
+# their HF metadata (model_type=cosmos3_omni, model_index.json _class_name=
+# Cosmos3OmniDiffusersPipeline) with the T2I/video Cosmos3 checkpoints, which
+# must keep using the default single-stage diffusion fallback — so this
+# pipeline cannot be auto-detected and must be selected explicitly:
+#
+#   vllm serve nvidia/Cosmos3-Nano-Policy-DROID --omni \
+#     --deploy-config /absolute/path/to/vllm-omni/vllm_omni/deploy/cosmos3_policy_droid.yaml
+#
+# Per-checkpoint model_config tweaks (e.g. policy_server_config values) can
+# still be layered on top with --stage-overrides; dict overrides deep-merge
+# with this file.
+
+pipeline: cosmos3_policy
+# async_chunk defaults to True, which the deploy validator rejects for a
+# single-stage pipeline (no next-stage processor). Must be false here.
+async_chunk: false
+distributed_executor_backend: mp
+
+stages:
+  - stage_id: 0
+    devices: "0"
+    # OpenPI policy serving is one observation -> one action reply.
+    max_num_seqs: 1
+    model_class_name: Cosmos3OmniDiffusersPipeline
+    # JSON prompt formatting is a property of this DROID checkpoint's training
+    # recipe, not a generic server default.
+    default_sampling_params:
+      extra_args:
+        format_prompt_as_json: true
+    model_config:
+      # Policy serving emits robot actions, not text/video; the Cosmos3
+      # guardrail stack (gated nvidia/Cosmos-1.0-Guardrail repo) is not part
+      # of it. Equivalent to launching with --no-guardrails.
+      guardrails: false
+      # OpenPI handshake metadata advertised to the robot client. Sent verbatim
+      # (the pipeline never reads these); they describe the camera inputs and
+      # action space the client must produce for this checkpoint.
+      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

vllm_omni/diffusion/models/cosmos3/action.py

@@ -38,13 +38,24 @@ EMBODIMENT_TO_DOMAIN_ID: dict[str, int] = {
     "bridge_orig_lerobot": 7,
     "droid_lerobot": 8,
     "robomind-franka": 8,
+    "embodiment_b": 9,
+    # Backward-compatible alias for the pre-generator Cosmos Framework name.
     "galbot": 9,
     "robomind-franka-dual": 12,
     "robomind-ur": 13,
     "agibotworld": 15,
+    "embodiment_c_gripper": 15,
+    "embodiment_c_gripper_ext": 15,
+    # Backward-compatible aliases for the pre-generator Cosmos Framework names.
     "agibot_gear_gripper": 15,
     "agibot_gear_gripper_ext": 15,
+    "xdof_yam": 16,
+    "molmoact2_yam": 16,
+    "abc_yam": 16,
     "fractal": 20,
+    "drawanything": 21,
+    "behavior1k_lerobot": 22,
+    "maniparena": 23,
 }
 
 

vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py

@@ -138,6 +138,7 @@ from .utils import (
     ensure_gripper_array,
     extract_robolab_image,
     extract_robolab_prompt_image,
+    get_robolab_domain_id,
     lazy_action_transform_pipeline,
     make_robolab_action_postprocess_inputs,
     next_robolab_seed,
@@ -1075,7 +1076,7 @@ class Cosmos3OmniDiffusersPipeline(
         self._guidance_scale = None
         self._num_timesteps = None
         self._cosmos3_branch_caches: dict[str, tuple[Any, Any]] | None = None
-        self._robolab_transform = None
+        self._robolab_transforms: dict[bool, Any] = {}
 
         # Set True by ``enable_cache_for_cosmos3`` when cache-dit is enabled on
         # this pipeline. Tells the sequential-CFG loop to keep paired
@@ -1394,11 +1395,18 @@ class Cosmos3OmniDiffusersPipeline(
             return val
         return default
 
-    def _get_robolab_transform(self):
-        if self._robolab_transform is None:
+    def _get_robolab_transform(self, *, format_prompt_as_json: bool = False):
+        transforms = getattr(self, "_robolab_transforms", None)
+        if transforms is None:
+            transforms = {}
+            self._robolab_transforms = transforms
+        if format_prompt_as_json not in transforms:
             action_dim = int(getattr(self.transformer, "action_dim", 64))
-            self._robolab_transform = lazy_action_transform_pipeline(action_dim)
-        return self._robolab_transform
+            transforms[format_prompt_as_json] = lazy_action_transform_pipeline(
+                action_dim,
+                format_prompt_as_json=format_prompt_as_json,
+            )
+        return transforms[format_prompt_as_json]
 
     def _build_robolab_policy_inputs(
         self,
@@ -1443,7 +1451,8 @@ class Cosmos3OmniDiffusersPipeline(
         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)
+        domain_id = get_robolab_domain_id(domain_name)
+        format_prompt_as_json = self._truthy(extra_param("format_prompt_as_json", False))
 
         if use_state and history_length < 1:
             raise ValueError("RoboLab history_length must be >= 1 when use_state is true.")
@@ -1506,7 +1515,9 @@ class Cosmos3OmniDiffusersPipeline(
             "action": action,
             # Cosmos Framework consumes this as an integer conditioning bucket.
             "conditioning_fps": torch.tensor(fps, dtype=torch.long),
-            "mode": ACTION_MODE_POLICY,
+            # Cosmos Framework 1.2 renamed the internal policy transform mode to
+            # ``wam``. The public vLLM action_mode remains ``policy``.
+            "mode": "wam",
             "domain_id": torch.tensor(domain_id, dtype=torch.long),
             "viewpoint": "concat_view",
             "additional_view_description": ROBOLAB_CONCAT_VIEW_DESCRIPTION,
@@ -1514,7 +1525,9 @@ class Cosmos3OmniDiffusersPipeline(
         if history_action is not None:
             sample["history_action"] = history_action
 
-        sample = self._get_robolab_transform()(sample, resolution)
+        sample = self._get_robolab_transform(format_prompt_as_json=format_prompt_as_json)(sample, resolution)
+        if isinstance(sample.get("ai_caption"), dict):
+            sample["ai_caption"] = json.dumps(sample["ai_caption"])
         sequence_plan = sample["sequence_plan"]
         video_tensor = sample["video"].float() / 127.5 - 1.0
         raw_action_dim_tensor = sample.get("raw_action_dim")

vllm_omni/diffusion/models/cosmos3/utils.py

@@ -78,6 +78,12 @@ ROBOLAB_CONCAT_VIEW_DESCRIPTION = (
     "sides, with the robot visible."
 )
 
+_LEGACY_EMBODIMENT_ALIASES = {
+    "galbot": "embodiment_b",
+    "agibot_gear_gripper": "embodiment_c_gripper",
+    "agibot_gear_gripper_ext": "embodiment_c_gripper_ext",
+}
+
 
 @dataclass(frozen=True)
 class RoboLabPolicyInputs:
@@ -252,33 +258,83 @@ def extract_robolab_prompt_image(prompt_data: Any | None) -> np.ndarray | None:
     return ensure_rgb_uint8_image(image, "multi_modal_data.image")
 
 
+_ROBOLAB_FRAMEWORK_IMPORT_ERROR = (
+    "Cosmos3 RoboLab policy serving requires a compatible cosmos_framework source on PYTHONPATH."
+)
+_ROBOLAB_FRAMEWORK_IMPORTS: dict[str, tuple[str, str]] = {
+    "ActionTransformPipeline": (
+        "cosmos_framework.data.generator.action.utils.transforms",
+        "ActionTransformPipeline",
+    ),
+    "convert_rotation": (
+        "cosmos_framework.data.generator.action.utils.pose_utils",
+        "convert_rotation",
+    ),
+    "pose_abs_to_rel": (
+        "cosmos_framework.data.generator.action.utils.pose_utils",
+        "pose_abs_to_rel",
+    ),
+    "pose_rel_to_abs": (
+        "cosmos_framework.data.generator.action.utils.pose_utils",
+        "pose_rel_to_abs",
+    ),
+    "build_abs_pose_from_components": (
+        "cosmos_framework.data.generator.action.utils.pose_utils",
+        "build_abs_pose_from_components",
+    ),
+    "FlowUniPCMultistepScheduler": (
+        "cosmos_framework.model.generator.diffusion.samplers.fm_solvers_unipc",
+        "FlowUniPCMultistepScheduler",
+    ),
+    "get_domain_id": (
+        "cosmos_framework.data.generator.action.utils.domain_utils",
+        "get_domain_id",
+    ),
+}
+
+
 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)
+        return getattr(module, symbol_name)
+    except (ImportError, AttributeError) as exc:
+        raise ImportError(
+            f"{error_message} Could not import required symbol {symbol_name!r} from {module_name!r}."
+        ) from exc
 
 
-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.",
+def _load_robolab_framework_symbol(symbol_name: str):
+    module_name, exported_name = _ROBOLAB_FRAMEWORK_IMPORTS[symbol_name]
+    return lazy_import(
+        module_name,
+        exported_name,
+        _ROBOLAB_FRAMEWORK_IMPORT_ERROR,
     )
-    return ActionTransformPipeline(max_action_dim=max_action_dim, cfg_dropout_rate=0.0)
+
+
+def preflight_cosmos3_action_framework_imports() -> None:
+    """Validate every Cosmos Framework symbol used by action policy serving."""
+    for symbol_name in _ROBOLAB_FRAMEWORK_IMPORTS:
+        _load_robolab_framework_symbol(symbol_name)
+
+
+def lazy_action_transform_pipeline(max_action_dim: int, *, format_prompt_as_json: bool = False):
+    ActionTransformPipeline = _load_robolab_framework_symbol("ActionTransformPipeline")
+    return ActionTransformPipeline(
+        max_action_dim=max_action_dim,
+        cfg_dropout_rate=0.0,
+        format_prompt_as_json=format_prompt_as_json,
+    )
+
+
+def get_robolab_domain_id(domain_name: str) -> int:
+    get_domain_id = _load_robolab_framework_symbol("get_domain_id")
+    canonical_name = _LEGACY_EMBODIMENT_ALIASES.get(domain_name.lower().strip(), domain_name)
+    return int(get_domain_id(canonical_name))
 
 
 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."
-        ),
-    )
+    FlowUniPCMultistepScheduler = _load_robolab_framework_symbol("FlowUniPCMultistepScheduler")
 
     scheduler = FlowUniPCMultistepScheduler(
         num_train_timesteps=1000,
@@ -290,38 +346,22 @@ def build_robolab_unipc_scheduler(num_steps: int, shift: float, device: torch.de
 
 
 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.",
-    )
+    convert_rotation = _load_robolab_framework_symbol("convert_rotation")
     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.",
-    )
+    pose_abs_to_rel_func = _load_robolab_framework_symbol("pose_abs_to_rel")
     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.",
-    )
+    pose_rel_to_abs_func = _load_robolab_framework_symbol("pose_rel_to_abs")
     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.",
-    )
+    build_abs_pose_from_components_func = _load_robolab_framework_symbol("build_abs_pose_from_components")
     return build_abs_pose_from_components_func(*args, **kwargs)
 
 

vllm_omni/entrypoints/cli/serve.py

@@ -930,6 +930,15 @@ class OmniServeCommand(CLISubcommand):
             action="store_true",
             help="Disable Cosmos3 text/video safety guardrails for this server.",
         )
+        omni_config_group.add_argument(
+            "--robot-openpi-idle-timeout",
+            type=_nonneg_finite_float,
+            default=30.0,
+            help=(
+                "Seconds the /v1/realtime/robot/openpi endpoint waits for the next request "
+                "before closing an idle WebSocket (default: 30). Set to 0 to disable the timeout."
+            ),
+        )
 
         # Enable diffusion pipeline profiling
         omni_config_group.add_argument(
@@ -1011,6 +1020,8 @@ def run_headless(args: TrackingNamespace) -> None:
 
     # Filter down to a dict of things explicitly requested by the user
     args_dict = args.get_explicit_kwargs_dict()
+    # This controls only the API-process WebSocket and is not a stage-engine option.
+    args_dict.pop("robot_openpi_idle_timeout", None)
     # This CLI-only negative alias is consumed below when selecting the
     # launcher log_stats value; it is not a per-stage config override.
     args_dict.pop("disable_log_stats", None)

vllm_omni/entrypoints/openai/api_server.py

@@ -485,6 +485,8 @@ async def build_async_omni_from_stage_config(
 
     try:
         kwargs = args.get_explicit_kwargs_dict()
+        # This controls only the API-process WebSocket and is not an AsyncOmni option.
+        kwargs.pop("robot_openpi_idle_timeout", None)
         model = kwargs.pop("model", None) or args.model
         kwargs.setdefault("log_stats", not args.disable_log_stats)
         async_omni = AsyncOmni(model=model, **kwargs)
@@ -1494,11 +1496,24 @@ async def realtime_robot_openpi(websocket: WebSocket):
 
     serving = getattr(websocket.app.state, "openai_serving_realtime_robot", None)
     if serving is None:
+        logger.warning(
+            "Rejecting robot OpenPI WebSocket: policy serving is disabled. "
+            "The diffusion stage must provide model_config.policy_server_config."
+        )
         await websocket.accept()
-        await websocket.send_json({"type": "error", "error": "Robot policy not available", "code": "unsupported"})
+        await websocket.send_json(
+            {
+                "type": "error",
+                "error": "Robot policy not available",
+                "code": "unsupported",
+            }
+        )
         await websocket.close()
         return
-    connection = RobotRealtimeConnection(websocket, serving)
+    state_args = getattr(websocket.app.state, "args", None)
+    configured_timeout = getattr(state_args, "robot_openpi_idle_timeout", 30.0)
+    idle_timeout = None if configured_timeout == 0 else configured_timeout
+    connection = RobotRealtimeConnection(websocket, serving, idle_timeout=idle_timeout)
     await connection.handle_connection()
 
 

vllm_omni/entrypoints/openpi/connection.py

@@ -29,6 +29,7 @@ mappings.
 from __future__ import annotations
 
 import asyncio
+from collections import OrderedDict
 from typing import Any
 
 import msgspec
@@ -45,10 +46,12 @@ logger = init_logger(__name__)
 _DEFAULT_IDLE_TIMEOUT = 30.0
 MAX_OPENPI_PAYLOAD_BYTES = 64 * 1024 * 1024
 _MISSING = object()
+# Upper bound on the per-connection set of seen session ids.
+MAX_TRACKED_SESSIONS = 1024
 
 
 def _pack_numpy(obj: Any) -> Any:
-    if isinstance(obj, (np.ndarray, np.generic)) and obj.dtype.kind in ("V", "O", "c"):
+    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:
@@ -165,17 +168,38 @@ class RobotRealtimeConnection:
         self,
         websocket: WebSocket,
         serving: ServingRealtimeRobotOpenPI,
-        idle_timeout: float = _DEFAULT_IDLE_TIMEOUT,
+        idle_timeout: float | None = _DEFAULT_IDLE_TIMEOUT,
     ) -> None:
         self.websocket = websocket
         self.serving = serving
         self._idle_timeout = idle_timeout
         self._current_session_id: str | None = None
-        self._call_count = 0
+        # Session ids seen on this connection, most-recently-used last.
+        self._seen_sessions: OrderedDict[str, None] = OrderedDict()
 
     def reset(self) -> None:
         self._current_session_id = None
-        self._call_count = 0
+        self._seen_sessions.clear()
+
+    def _mark_session_seen(self, session_id: str) -> bool:
+        """Record ``session_id``; return True when it is the first sighting.
+
+        A session's first observation is the one that carries ``reset`` to the
+        policy, so this needs per-session memory.
+        """
+        if session_id in self._seen_sessions:
+            self._seen_sessions.move_to_end(session_id)
+            return False
+
+        self._seen_sessions[session_id] = None
+        if len(self._seen_sessions) > MAX_TRACKED_SESSIONS:
+            evicted, _ = self._seen_sessions.popitem(last=False)
+            logger.debug(
+                "Robot OpenPI connection is tracking more than %d sessions; dropped least-recently-used %s",
+                MAX_TRACKED_SESSIONS,
+                evicted,
+            )
+        return True
 
     async def _send_error(self, message: str) -> None:
         await self.websocket.send_bytes(_pack({"type": "error", "message": message}))
@@ -199,18 +223,22 @@ class RobotRealtimeConnection:
             await self.websocket.send_bytes(_pack(metadata))
 
             while True:
-                try:
-                    msg = await asyncio.wait_for(
-                        self.websocket.receive(),
-                        timeout=self._idle_timeout,
-                    )
-                except asyncio.TimeoutError:
-                    logger.info("Robot OpenPI connection idle timeout after %.1f seconds", self._idle_timeout)
+                idle_timeout = self._idle_timeout
+                if idle_timeout is None:
+                    msg = await self.websocket.receive()
+                else:
                     try:
-                        await self.websocket.close()
-                    except Exception:
-                        logger.debug("Failed to close idle robot OpenPI websocket", exc_info=True)
-                    return
+                        msg = await asyncio.wait_for(
+                            self.websocket.receive(),
+                            timeout=idle_timeout,
+                        )
+                    except asyncio.TimeoutError:
+                        logger.info("Robot OpenPI connection idle timeout after %.1f seconds", idle_timeout)
+                        try:
+                            await self.websocket.close()
+                        except Exception:
+                            logger.debug("Failed to close idle robot OpenPI websocket", exc_info=True)
+                        return
 
                 if msg.get("type") == "websocket.disconnect":
                     break
@@ -245,13 +273,12 @@ class RobotRealtimeConnection:
                                     session_id,
                                 )
                             self._current_session_id = session_id
-                            self._call_count = 0
 
-                        self._call_count += 1
+                        reset = self._mark_session_seen(session_id)
                         actions = await self.serving.infer(
                             obs,
                             session_id=session_id,
-                            reset=self._call_count <= 1,
+                            reset=reset,
                         )
                         await self.websocket.send_bytes(_pack(actions))
                 except Exception:

vllm_omni/entrypoints/openpi/serving.py

@@ -148,16 +148,32 @@ class ServingRealtimeRobotOpenPI:
         `AsyncOmni.generate()` and routed to the diffusion stage.
         """
         from vllm_omni.diffusion.request import OmniDiffusionRequest
+        from vllm_omni.entrypoints.openai.stage_params import (
+            clone_sampling_params,
+            get_default_sampling_params_list,
+        )
         from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
-        # An optional integer ``seed`` in the inference message becomes the engine request's
-        # ``sampling_params.seed``; omitted, ``OmniDiffusionRequest`` assigns a random one.
+        # The engine applies stage default_sampling_params only to requests
+        # that carry no explicit params; this endpoint always passes explicit
+        # params, so start from a clone of the diffusion stage's defaults
+        # (e.g. a policy deploy yaml's ``extra_args``) and layer the OpenPI
+        # protocol fields on top.
         seed = obs.pop("seed", None)
-        extra_args = {
-            "reset": reset,
-            "session_id": session_id,
-            "robot_obs": obs,
-        }
+        sampling_params = OmniDiffusionSamplingParams()
+        for default_params in get_default_sampling_params_list(self.engine_client):
+            if isinstance(default_params, OmniDiffusionSamplingParams):
+                sampling_params = clone_sampling_params(default_params)
+                break
+
+        extra_args = sampling_params.extra_args or {}
+        extra_args.update(
+            {
+                "reset": reset,
+                "session_id": session_id,
+                "robot_obs": obs,
+            }
+        )
 
         prompt = obs.get("prompt", "")
         sampling_params = OmniDiffusionSamplingParams(

vllm_omni/model_executor/models/cosmos3/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm_omni/model_executor/models/cosmos3/pipeline.py

@@ -0,0 +1,41 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Cosmos3 policy single-stage topology for online OpenPI serving.
+
+Cosmos3 policy checkpoints (e.g. ``nvidia/Cosmos3-Nano-Policy-DROID``) are one
+diffusion stage: robot observation -> action chunk. They are
+registered in ``OMNI_PIPELINES`` because online OpenPI serving
+(``/v1/realtime/robot/openpi``) requires a registered pipeline so the deploy
+yaml's ``model_config.policy_server_config`` reaches the websocket handshake.
+
+This pipeline declares neither ``hf_architectures`` nor
+``diffusers_class_name``: policy checkpoints share their HF metadata
+(``model_type=cosmos3_omni``, ``model_index.json`` ``_class_name=
+Cosmos3OmniDiffusersPipeline``) with the T2I/video Cosmos3 checkpoints, which
+must keep resolving through the default single-stage diffusion fallback.
+Select this pipeline explicitly via a deploy yaml ``pipeline:`` key, e.g.
+``vllm serve nvidia/Cosmos3-Nano-Policy-DROID --omni --deploy-config
+/absolute/path/to/vllm-omni/vllm_omni/deploy/cosmos3_policy_droid.yaml``.
+"""
+
+from vllm_omni.config.stage_config import (
+    PipelineConfig,
+    StageExecutionType,
+    StagePipelineConfig,
+)
+
+COSMOS3_POLICY_PIPELINE = PipelineConfig(
+    model_type="cosmos3_policy",
+    model_arch="Cosmos3OmniDiffusersPipeline",
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="diffusion",
+            execution_type=StageExecutionType.DIFFUSION,
+            input_sources=(),
+            final_output=True,
+            final_output_type="action",
+            model_arch="Cosmos3OmniDiffusersPipeline",
+        ),
+    ),
+)

vllm_omni/model_extras/cosmos3.py

@@ -35,6 +35,7 @@ COSMOS3_EXTRA_BODY_PARAMS = frozenset(
         "resolution",
         "image_size",
         "use_state",
+        "format_prompt_as_json",
         "observation",
         "robot_obs",
         "deterministic_seed",