diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 016ff47ae..fa04c0658 100644 --- a/.github/workflows/ci.yml +++ b/.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 \ diff --git a/Makefile b/Makefile index b830166e4..bde2f927b 100644 --- a/Makefile +++ b/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 \ No newline at end of file + find src/unilab/assets/scenes -type f ! -name '.gitkeep' -delete 2>/dev/null || true diff --git a/README.md b/README.md index 745aef8bc..8490e3a9b 100644 --- a/README.md +++ b/README.md @@ -91,9 +91,7 @@ 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 @@ -101,6 +99,11 @@ make setup-motrix # 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 ``` diff --git a/README_zh.md b/README_zh.md index 5aa3f04d7..776f3c11d 100644 --- a/README_zh.md +++ b/README_zh.md @@ -91,9 +91,7 @@ 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 @@ -101,6 +99,11 @@ make setup-motrix # 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 ``` diff --git a/pyproject.rocm.toml b/pyproject.rocm.toml index 741677edb..5b32bb3b1 100644 --- a/pyproject.rocm.toml +++ b/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"] diff --git a/pyproject.toml b/pyproject.toml index 29675a65f..74f7153c0 100644 --- a/pyproject.toml +++ b/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"] diff --git a/src/unilab/base/backend/__init__.py b/src/unilab/base/backend/__init__.py index 4563451e9..ba49b4eca 100644 --- a/src/unilab/base/backend/__init__.py +++ b/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}") diff --git a/src/unilab/cli.py b/src/unilab/cli.py index a76e9569a..e3b03897e 100644 --- a/src/unilab/cli.py +++ b/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`." diff --git a/tests/base/test_backend_imports.py b/tests/base/test_backend_imports.py index dea10e982..8461c328e 100644 --- a/tests/base/test_backend_imports.py +++ b/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"] diff --git a/tests/test_cli_runtime_requirements.py b/tests/test_cli_runtime_requirements.py new file mode 100644 index 000000000..77d174ff4 --- /dev/null +++ b/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") diff --git a/uv.lock b/uv.lock index e3f2e85a3..5bb321cf1 100644 --- a/uv.lock +++ b/uv.lock @@ -1120,7 +1120,7 @@ wheels = [ [[package]] name = "huggingface-hub" -version = "1.18.0" +version = "1.20.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "click" }, @@ -1134,9 +1134,9 @@ dependencies = [ { name = "typer" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/fb/d8/748ea0a47f0fa15227fe682f7a80826b4b7c096e4818044b8f56d6cb66d6/huggingface_hub-1.18.0.tar.gz", hash = "sha256:f0c5ecd1ef8c6a60f86f61ee278f2c1570ba9e279c9f54de9094210723b3613b", size = 812699, upload-time = "2026-06-05T09:26:33.401Z" } +sdist = { url = "https://files.pythonhosted.org/packages/e6/7e/fad82ad491b226e832d2da90a1a59f36acd4526cda8c726f639834754aa4/huggingface_hub-1.20.1.tar.gz", hash = "sha256:9f6d63bfbeab2d2a8357200a9bc4f18cd2c8bfac9579f792f5922e77bf6471d0", size = 859910, upload-time = "2026-06-18T22:06:53.348Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/0b/03/40a05316cb6616e5b7efd7773656441ab04b4b022c2199e79bb4622a92a3/huggingface_hub-1.18.0-py3-none-any.whl", hash = "sha256:729be4a976fb706dcc02d176bcda8a3f32bdf21a294e8f4b3dda6fbcbc9c1ab1", size = 684411, upload-time = "2026-06-05T09:26:31.48Z" }, + { url = "https://files.pythonhosted.org/packages/8e/b5/ff8516e74b459da3dce9567540c39f2d305ee7a2655109f6802873ff1588/huggingface_hub-1.20.1-py3-none-any.whl", hash = "sha256:274448a45c1ba6f112fe2fb168ead05574c654faa156904157a84085cfae14bd", size = 719837, upload-time = "2026-06-18T22:06:51.486Z" }, ] [[package]] @@ -4560,7 +4560,6 @@ dependencies = [ { name = "lark" }, { name = "mediapy" }, { name = "mlx", marker = "sys_platform == 'darwin'" }, - { name = "mujoco-uni" }, { name = "ninja", marker = "sys_platform == 'linux'" }, { name = "notebook" }, { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, @@ -4584,6 +4583,9 @@ dependencies = [ motrix = [ { name = "motrixsim-core" }, ] +mujoco = [ + { name = "mujoco-uni" }, +] viser = [ { name = "trimesh" }, { name = "viser" }, @@ -4609,7 +4611,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 = "ninja", marker = "sys_platform == 'linux'" }, { name = "notebook", specifier = ">=7.5.5" }, { name = "numpy" }, @@ -4629,7 +4631,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 = [ diff --git a/uv.rocm.lock b/uv.rocm.lock index 28111b332..22885937c 100644 --- a/uv.rocm.lock +++ b/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 = [