fix: split mujoco and motrix backend extras

.github/workflows/ci.yml

@@ -107,7 +107,7 @@ jobs:
         run: sudo apt-get update && sudo apt-get install -y libwayland-client0
       - name: Install dependencies
         run: |
-          uv sync --extra motrix \
+          uv sync --extra mujoco --extra motrix \
             --no-install-package torch \
             --no-install-package nvidia-cublas-cu12 \
             --no-install-package nvidia-cuda-cupti-cu12 \

Makefile

@@ -1,10 +1,15 @@
 .PHONY: sync
 sync:
-	uv sync
+	uv sync --extra mujoco --extra motrix
 
 .PHONY: setup
 setup:
-	uv sync
+	uv sync --extra mujoco --extra motrix
+	uv run --no-sync unilab-complete install
+
+.PHONY: setup-mujoco
+setup-mujoco:
+	uv sync --extra mujoco
 	uv run --no-sync unilab-complete install
 
 .PHONY: setup-motrix
@@ -20,12 +25,12 @@ install-completion:
 sync-rocm:
 	@cp pyproject.rocm.toml pyproject.toml
 	@if [ -f uv.rocm.lock ]; then cp uv.rocm.lock uv.lock; fi
-	uv sync --extra motrix
+	uv sync --extra mujoco --extra motrix
 	cp uv.lock uv.rocm.lock
 
 .PHONY: sync-xpu
 sync-xpu:
-	uv sync --extra motrix --no-install-package torch
+	uv sync --extra mujoco --extra motrix --no-install-package torch
 	uv pip install torch==2.7.0 --torch-backend xpu
 
 .PHONY: format
@@ -71,4 +76,4 @@ clean:
 	find src/unilab/assets/.cache -type f ! -name '.gitkeep' -delete 2>/dev/null || true
 	find src/unilab/assets/caches -type f ! -name '.gitkeep' -delete 2>/dev/null || true
 	find src/unilab/assets/checkpoints -type f ! -name '.gitkeep' -delete 2>/dev/null || true
-	find src/unilab/assets/scenes -type f ! -name '.gitkeep' -delete 2>/dev/null || true
+	find src/unilab/assets/scenes -type f ! -name '.gitkeep' -delete 2>/dev/null || true

README.md

@@ -91,16 +91,19 @@ cd UniLab
 # Pick the setup command for your platform.
 
 # Linux CUDA or macOS
-make setup-motrix
-# Without shell completion setup: uv sync --extra motrix
-# If `make` is not installed: uv sync --extra motrix && uv run --no-sync unilab-complete install
+make setup
 
 # Linux AMD / ROCm
 # make sync-rocm
 
 # Linux Intel Arc / iGPU
 # make sync-xpu
 
+# Without shell completion setup:
+# uv sync --extra mujoco --extra motrix
+# If `make` is not installed:
+# uv sync --extra mujoco --extra motrix && uv run --no-sync unilab-complete install
+
 # 3. Pre-trained checkpoint playback (downloads from Hugging Face on first run)
 uv run demo dance
 ```

README_zh.md

@@ -91,16 +91,19 @@ cd UniLab
 # 请按你的平台选择对应的安装命令。
 
 # Linux CUDA 或 macOS
-make setup-motrix
-# 不使用 shell completion 设置时：uv sync --extra motrix
-# 如果没有安装 `make`：uv sync --extra motrix && uv run --no-sync unilab-complete install
+make setup
 
 # Linux AMD / ROCm
 # make sync-rocm
 
 # Linux Intel Arc / iGPU
 # make sync-xpu
 
+# 不使用 shell completion 设置时：
+# uv sync --extra mujoco --extra motrix
+# 如果没有安装 `make`：
+# uv sync --extra mujoco --extra motrix && uv run --no-sync unilab-complete install
+
 # 3. 预训练 checkpoint 回放（首次运行会从 Hugging Face 下载）
 uv run demo dance
 ```

pyproject.rocm.toml

@@ -14,7 +14,6 @@ dependencies = [
     "numpy",
     "torch==2.11.0",
     "triton-rocm==3.6.0 ; sys_platform == 'linux' and platform_machine == 'x86_64'",
-    "mujoco-uni==3.8.0",
     "gymnasium",
     "imageio",
     "etils",
@@ -43,6 +42,7 @@ unilab-viz-nan = "unilab.tools.viz_nan:main"
 unilab-export-scene = "unilab.tools.export_scene:main"
 
 [project.optional-dependencies]
+mujoco = ["mujoco-uni==3.8.0"]
 motrix = ["motrixsim-core==0.8.2"]
 viser = ["viser>=1.0.26", "trimesh>=3.21.7"]
 

pyproject.toml

@@ -13,7 +13,6 @@ requires-python = ">=3.10,<3.14"
 dependencies = [
     "numpy",
     "torch==2.7.0",
-    "mujoco-uni==3.8.0",
     "gymnasium",
     "imageio",
     "etils",
@@ -49,6 +48,7 @@ unilab-import-robot = "unilab.tools.import_robot:main"
 unilab-pull-assets = "unilab.tools.pull_assets:main"
 
 [project.optional-dependencies]
+mujoco = ["mujoco-uni==3.8.0"]
 motrix = ["motrixsim-core==0.8.2"]
 viser = ["viser>=1.0.26", "trimesh>=3.21.7"]
 

src/unilab/base/backend/__init__.py

@@ -3,11 +3,6 @@
 from unilab.base.scene import SceneCfg
 
 from .base import SimBackend
-from .motrix.scene import (
-    add_motrix_tracking_frame_sensors,
-    materialize_motrix_hfield_attached_scene,
-    materialize_motrix_scene,
-)
 
 _MUJOCO_XML_EXPORTS = frozenset(
     {
@@ -21,6 +16,13 @@
         "processed_xml",
     }
 )
+_MOTRIX_SCENE_EXPORTS = frozenset(
+    {
+        "add_motrix_tracking_frame_sensors",
+        "materialize_motrix_hfield_attached_scene",
+        "materialize_motrix_scene",
+    }
+)
 
 
 def _load_mujoco_backend() -> Any:
@@ -35,6 +37,12 @@ def _load_motrix_backend() -> tuple[Any, bool]:
     return MotrixBackend, bool(MOTRIX_AVAILABLE)
 
 
+def _load_motrix_scene_export(name: str) -> Any:
+    from .motrix import scene
+
+    return getattr(scene, name)
+
+
 def create_backend(
     backend_type: str,
     scene: SceneCfg,
@@ -89,6 +97,8 @@ def __getattr__(name: str):
         from .mujoco import xml
 
         return getattr(xml, name)
+    if name in _MOTRIX_SCENE_EXPORTS:
+        return _load_motrix_scene_export(name)
     raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
 
 

src/unilab/cli.py

@@ -102,6 +102,10 @@ def _check_load_run(load_run: str) -> None:
 def _check_runtime_requirements(algo: str, sim: str) -> None:
     if algo == "mlx_ppo" and platform.system() != "Darwin":
         raise SystemExit("mlx_ppo is only supported on macOS; use --algo ppo for torch PPO.")
+    if sim == "mujoco" and find_spec("mujoco") is None:
+        raise SystemExit(
+            "sim=mujoco requires the MuJoCo extra. Install it with `uv sync --extra mujoco`."
+        )
     if sim == "motrix" and find_spec("motrixsim") is None:
         raise SystemExit(
             "sim=motrix requires the Motrix extra. Install it with `uv sync --extra motrix`."

tests/base/test_backend_imports.py

@@ -39,3 +39,41 @@ def test_mujoco_backend_import_path_does_not_eagerly_import_motrix() -> None:
     lines = result.stdout.splitlines()
     assert lines[0] in {"mujoco_backend imported", "mujoco_backend skipped"}
     assert lines[1:] == ["motrix_backend False", "motrixsim False"]
+
+
+def test_motrix_backend_import_path_does_not_eagerly_import_mujoco() -> None:
+    code = textwrap.dedent(
+        """
+        import importlib.util
+        import sys
+
+        from unilab.base.backend import (
+            create_backend,
+            materialize_motrix_hfield_attached_scene,
+            materialize_motrix_scene,
+        )
+
+        assert create_backend is not None
+        assert materialize_motrix_scene is not None
+        assert materialize_motrix_hfield_attached_scene is not None
+
+        if importlib.util.find_spec("motrixsim") is not None:
+            import unilab.base.backend.motrix.backend
+            print("motrix_backend imported")
+        else:
+            print("motrix_backend skipped")
+
+        print("mujoco_backend", "unilab.base.backend.mujoco.backend" in sys.modules)
+        print("mujoco", "mujoco" in sys.modules)
+        """
+    )
+    result = subprocess.run(
+        [sys.executable, "-c", code],
+        check=True,
+        capture_output=True,
+        text=True,
+    )
+
+    lines = result.stdout.splitlines()
+    assert lines[0] in {"motrix_backend imported", "motrix_backend skipped"}
+    assert lines[1:] == ["mujoco_backend False", "mujoco False"]

tests/test_cli_runtime_requirements.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+import pytest
+
+from unilab import cli
+
+
+def test_check_runtime_requirements_requires_mujoco_extra(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(cli, "find_spec", lambda name: None if name == "mujoco" else object())
+
+    with pytest.raises(SystemExit, match="sim=mujoco requires the MuJoCo extra"):
+        cli._check_runtime_requirements("ppo", "mujoco")
+
+
+def test_check_runtime_requirements_requires_motrix_extra(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(cli, "find_spec", lambda name: None if name == "motrixsim" else object())
+
+    with pytest.raises(SystemExit, match="sim=motrix requires the Motrix extra"):
+        cli._check_runtime_requirements("ppo", "motrix")

uv.rocm.lock

@@ -4536,7 +4536,6 @@ dependencies = [
     { name = "lark" },
     { name = "mediapy" },
     { name = "mlx", marker = "sys_platform == 'darwin'" },
-    { name = "mujoco-uni" },
     { name = "notebook" },
     { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
     { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" },
@@ -4560,6 +4559,9 @@ dependencies = [
 motrix = [
     { name = "motrixsim-core" },
 ]
+mujoco = [
+    { name = "mujoco-uni" },
+]
 viser = [
     { name = "trimesh" },
     { name = "viser" },
@@ -4584,7 +4586,7 @@ requires-dist = [
     { name = "mediapy" },
     { name = "mlx", marker = "sys_platform == 'darwin'" },
     { name = "motrixsim-core", marker = "extra == 'motrix'", specifier = "==0.8.2" },
-    { name = "mujoco-uni", specifier = "==3.8.0" },
+    { name = "mujoco-uni", marker = "extra == 'mujoco'", specifier = "==3.8.0" },
     { name = "notebook", specifier = ">=7.5.5" },
     { name = "numpy" },
     { name = "onnxruntime", specifier = ">=1.24.3" },
@@ -4603,7 +4605,7 @@ requires-dist = [
     { name = "viser", marker = "extra == 'viser'", specifier = ">=1.0.26" },
     { name = "wandb" },
 ]
-provides-extras = ["motrix", "viser"]
+provides-extras = ["mujoco", "motrix", "viser"]
 
 [package.metadata.requires-dev]
 dev = [