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