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>

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"),
     [