diff --git a/Makefile b/Makefile index 9b052837a..fbdb938a6 100644 --- a/Makefile +++ b/Makefile @@ -2,52 +2,11 @@ sync: uv sync --extra mujoco --extra motrix -# Switch the MuJoCo solver version (support window: >=3.5,<3.12), e.g. -# make mujoco MJ=3.10.0 -# This is the explicit sdist fallback path: the prebuilt mujoco-uni-runtime -# wheels only bind the default mujoco==3.11.0, so any other version requires -# recompiling the native extension from source against the requested mujoco. -# The mujoco extra declares `mujoco~=3.11.0`, which a lock upgrade can never -# leave, so the override operates on the environment directly (uv pip, no -# re-lock): install the requested mujoco plus the runtime's build -# requirements, then force an in-env sdist rebuild of the runtime. The -# override is environment-local — uv.lock is untouched, and -# `uv sync --extra mujoco --reinstall-package mujoco-uni-runtime` reverts to -# the locked default (mujoco 3.11.0 + prebuilt wheel), which is the -# switch-back path. The reinstall flag is required there: plain `uv sync` -# restores mujoco but keeps the locally rebuilt extension, which then fails -# to load against the reverted mujoco. The default `make setup` path installs -# the prebuilt wheel and needs no compiler. -.PHONY: check-cxx-toolchain -check-cxx-toolchain: - @command -v c++ >/dev/null 2>&1 || { \ - echo "error: building mujoco-uni-runtime from source requires a C++ toolchain, but 'c++' was not found."; \ - echo " Debian/Ubuntu: sudo apt-get install build-essential"; \ - echo " macOS: xcode-select --install"; \ - echo " Fedora/RHEL: sudo dnf install gcc-c++ make"; \ - exit 1; \ - } - -.PHONY: mujoco -mujoco: - @test -n "$(MJ)" || (echo "usage: make mujoco MJ=3.10.0" && exit 1) - @$(MAKE) --no-print-directory check-cxx-toolchain - uv pip install "mujoco==$(MJ)" pybind11 wheel setuptools - uv cache clean mujoco-uni-runtime - ver=$$(uv pip show mujoco-uni-runtime | sed -n 's/^Version: //p') && \ - uv pip install --force-reinstall --no-deps --no-build-isolation \ - --no-binary mujoco-uni-runtime "mujoco-uni-runtime==$$ver" - .PHONY: setup setup: 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 - # Installs the Python extra and builds DrakeUni's native extension. By default # the host-compatible official tarball is downloaded; use DRAKE_HOME= # to build against an existing installation. diff --git a/docs/sphinx/source/api_reference/backend/index.md b/docs/sphinx/source/api_reference/backend/index.md index da562fb28..bba80cabc 100644 --- a/docs/sphinx/source/api_reference/backend/index.md +++ b/docs/sphinx/source/api_reference/backend/index.md @@ -5,7 +5,7 @@ only assembles task-owned scenes through `unilab.base.backend_factory`. | Backend | Strengths | Notes | |---|---|---| -| **MuJoCo** (`mujoco` + `mujoco-uni-runtime`) | Mature, broad asset support, deterministic | Default for research | +| **MuJoCo** (`mujoco` + `mjbatch`) | Mature, broad asset support, deterministic | Default for research | | **Motrix** (`motrixsim-core`) | High-throughput, multithread step, snapshot/playback | Cross-platform; required for video export on macOS | Pick a backend per task via the top-level `--sim ` CLI flag — see diff --git a/docs/sphinx/source/changelog.md b/docs/sphinx/source/changelog.md index 7fbc57026..15a24e307 100644 --- a/docs/sphinx/source/changelog.md +++ b/docs/sphinx/source/changelog.md @@ -13,6 +13,54 @@ UniLab 遵循[语义化版本](https://semver.org/)。本共享页面以中英 ## Unreleased / 未发布 +- Replace the `mujoco-uni-runtime` dependency (`mujoco_uni` import) with the + `mjbatch` native batch engine across the repository (roadmap + [#1552](https://github.com/unilabsim/UniLab/issues/1552), + [#1553](https://github.com/unilabsim/UniLab/issues/1553)). The `mujoco` + extra now installs `mujoco~=3.11.0` plus `mjbatch` pinned to the + [integration fork](https://github.com/unilabsim/mjbatch); the fork's final + distribution identity (PyPI package vs git pin, and prebuilt wheels) is the + roadmap's open maintainer item. The `sim=mujoco` CLI runtime check now gates + on the `mjbatch` module. A post-swap ablation slimmed the pinned fork's + API ([#1557](https://github.com/unilabsim/UniLab/issues/1557)): the + per-substep callback is `fn(k, state, ctrl)` (no `callback_sensordata` + argument), `steps_done` / `stop_on_warning` are gone, and the hfield + scanner is height-only with its own validation. Numerical equivalence + before and after the swap is **not** guaranteed; the accepted drift is + characterized by the #1554 drift baseline. + 全仓库将 `mujoco-uni-runtime` 依赖（`mujoco_uni` 导入）替换为 `mjbatch` + 原生 batch 引擎（roadmap #1552、#1553）。`mujoco` extra 现安装 + `mujoco~=3.11.0` 加钉住的 [集成 fork](https://github.com/unilabsim/mjbatch) + `mjbatch`；fork 的最终分发身份（PyPI package 还是 git 钉版、是否提供预编译 + wheel）是 roadmap 上的待定维护事项。`sim=mujoco` 的 CLI 运行时检查改为检查 + `mjbatch` 模块。替换后的消融精简了钉住 fork 的 API（#1557）：per-substep + 回调为 `fn(k, state, ctrl)`（不再有 `callback_sensordata` 参数）， + `steps_done` / `stop_on_warning` 已移除，hfield 扫描器只输出高度并自带 + 校验。替换前后数值不保证一致；接受的漂移由 #1554 漂移基线表征。 + +- Deprecate and remove the MuJoCo chunk/forward knobs: `EnvCfg` fields + `post_step_forward_sensor`, `adaptive_chunk_size`, and `chunk_size`, the + `bench_nsteps` backend kwarg, the matching Hydra owner keys, and the + `make mujoco MJ=` / `check-cxx-toolchain` / `setup-mujoco` Makefile + targets are gone. `mjbatch` schedules per-simulation work without a chunk + knob, and step ends one substep behind the state by default, matching the + previous `post_step_forward_sensor=False` semantics; the per-env model + variant machinery (`ModelVariantSpec` materialization on the MuJoCo backend) + is no longer supported there — init-lifecycle geometry overrides move to the + remaining variant-capable backends. Windows support is unchanged: the + MuJoCo physics backend stays Linux/macOS-only because `mjbatch` ships no + Windows wheels. + 弃用并移除 MuJoCo chunk/forward 旋钮：`EnvCfg` 字段 + `post_step_forward_sensor`、`adaptive_chunk_size`、`chunk_size`、`bench_nsteps` + 后端 kwarg、对应的 Hydra owner 键，以及 Makefile 目标 + `make mujoco MJ=` / `check-cxx-toolchain` / `setup-mujoco` 均已删除。 + `mjbatch` 在没有 chunk 旋钮的情况下调度 per-simulation 工作，且 sensordata + 默认落后一个子步，与之前的 `post_step_forward_sensor=False` 语义一致； + MuJoCo 后端不再支持 per-env 模型 variants（`ModelVariantSpec` + materialization）——init-lifecycle 几何覆盖改由仍支持 variants 的后端提供。 + Windows 支持不变：由于 `mjbatch` 不提供 Windows wheel，MuJoCo 物理后端 + 仍然只支持 Linux/macOS。 + ## 1.2.0 (2026-09-10) - Update the required `unisim-core` release to `>=1.2.0` and the pinned diff --git a/docs/sphinx/source/en/1-getting_started/2-installation.md b/docs/sphinx/source/en/1-getting_started/2-installation.md index 3fedd9ea8..73f6aa118 100644 --- a/docs/sphinx/source/en/1-getting_started/2-installation.md +++ b/docs/sphinx/source/en/1-getting_started/2-installation.md @@ -10,14 +10,15 @@ live in the getting-started and algorithm pages. - Git and `curl`, used to clone the repository and fetch runtime assets. - `cmake`, required when building the Drake native batch extension. The Drake setup script uses CMake and a C++ toolchain. -- For the `mujoco` extra: the default install path uses the prebuilt - `mujoco-uni-runtime` wheel (bound to `mujoco==3.11.0`), so - `make setup` / `uv sync --extra mujoco` needs no compiler. A C++17 toolchain - and Python development headers are only required on the explicit - source-rebuild path (switching the MuJoCo version; see "Switching The Local - MuJoCo Version"); without them the build fails with errors such as - `fatal error: Python.h: No such file or directory` (see "Install Error - Signatures" for the full lookup table). +- For the `mujoco` extra: the MuJoCo physics backend executes on the + `mjbatch` native batch engine, currently consumed from the pinned + integration fork (`unilabsim/mjbatch`). uv builds it from the pinned git + source with an isolated build (scikit-build-core + nanobind), which requires + a C++17 toolchain and Python development headers; the build binds + `mujoco==3.11.0` and the engine refuses to import against any other mujoco + version. A no-compiler install path depends on the fork's final distribution + channel (prebuilt wheels vs the git pin), which is the roadmap's open item + (see "Switching The Local MuJoCo Version" and "Install Error Signatures"). - macOS: `xcode-select --install` - Ubuntu / Debian: `sudo apt-get install build-essential python3-dev` - Fedora / RHEL: `sudo dnf install gcc-c++ make python3-devel` @@ -61,16 +62,12 @@ make setup # Fastest path for the first Motrix demo. # make setup-motrix - -# MuJoCo only. -# make setup-mujoco ``` `make setup` runs `uv sync --extra mujoco --extra motrix` and installs shell completion. `make setup-motrix` runs `uv sync --extra motrix` and installs the -same completion entry. `make setup-mujoco` runs `uv sync --extra mujoco` and -installs completion. Run only one of these paths. If `make` is unavailable, run -the matching commands directly: +same completion entry. Run only one of these paths. If `make` is unavailable, +run the matching commands directly: ```bash # Full default setup: @@ -79,16 +76,13 @@ uv run --no-sync unilab-complete install # Motrix only: # uv sync --extra motrix && uv run --no-sync unilab-complete install - -# MuJoCo only: -# uv sync --extra mujoco && uv run --no-sync unilab-complete install ``` ## Conda And Pip The recommended path is still the in-repo `make setup` / `make setup-motrix` (or -`uv`) workflow. Use `make setup-mujoco` when Motrix is not needed. Conda can -serve as an outer environment for Python, CUDA, or system-library isolation, +`uv`) workflow. Use `uv sync --extra mujoco` when Motrix is not needed. Conda +can serve as an outer environment for Python, CUDA, or system-library isolation, but once the environment is active keep using the repository's `make` / `uv` commands inside it: @@ -101,8 +95,8 @@ cd UniLab make setup-motrix ``` -Use `make setup-mujoco` if you do not need Motrix. ROCm and XPU still go through -the platform-specific `make` targets below. +Use `uv sync --extra mujoco` if you do not need Motrix. ROCm and XPU still go +through the platform-specific `make` targets below. From a source checkout, pip is a fallback path. Install the package first, then add optional runtimes explicitly: @@ -117,19 +111,17 @@ pip install -e . # Motrix, when needed: pip install motrixsim-core==0.8.2 -# MuJoCo, when needed (the default install resolves the prebuilt wheel bound -# to mujoco==3.11.0): -pip install "mujoco~=3.11.0" "mujoco-uni-runtime==0.5.0" +# MuJoCo, when needed (resolves the pinned mjbatch integration fork, built +# against mujoco==3.11.0): +pip install "mujoco~=3.11.0" "mjbatch @ git+https://github.com/unilabsim/mjbatch.git@cf4a83d" ``` The editable install points at the checkout; the regular install copies the package and its task configs (`unilab/conf/`) into the environment. In both cases, `train`, `eval`, and `demo` work from any directory, while logs and -checkpoints are written under the current working directory. The prebuilt -`mujoco-uni-runtime` wheel installs directly through pip with no build step; -`pybind11` / `wheel` and `--no-build-isolation` are only needed when forcing an -sdist rebuild against a non-default mujoco version (see "Switching The Local -MuJoCo Version"). For MJWarp, Genesis, platform-specific torch indexes, and +checkpoints are written under the current working directory. The `mjbatch` +engine builds against the pinned `mujoco==3.11.0`; for MJWarp, Genesis, +platform-specific torch indexes, and ROCm/XPU profiles, prefer the uv paths above. Robot meshes and textures are intentionally excluded from the wheel and downloaded on the cold path from the `unilabsim/unilab-robots` dataset. Ensure the installed package @@ -178,7 +170,7 @@ uv sync --extra mujoco --extra mjwarp --extra newton | Backend | Install path | Important prerequisites | | --- | --- | --- | -| MuJoCo | `make setup-mujoco` or `uv sync --extra mujoco` | Prebuilt wheel (bound to `mujoco==3.11.0`), no compiler needed; a C++17 toolchain and Python development headers are only required when switching versions (`make mujoco MJ=`) | +| MuJoCo | `make setup` or `uv sync --extra mujoco` | Builds the pinned `mjbatch` fork (bound to `mujoco==3.11.0`) from source; a C++17 toolchain and Python development headers are required until prebuilt wheels exist (roadmap open item) | | Motrix | `make setup-motrix` or `uv sync --extra motrix` | Motrix runtime is installed from the pinned Python package | | MJWarp | `uv sync --extra mujoco --extra mjwarp` | NVIDIA CUDA and an explicit CUDA process device | | Genesis | `uv sync --extra genesis` | The validated path uses Linux x86_64, an NVIDIA GPU, and the pinned torch/Genesis versions | @@ -205,44 +197,23 @@ runtime variables, renderer requirements, and verification commands: ## Switching The Local MuJoCo Version -The default install path of the `mujoco` extra uses the prebuilt -`mujoco-uni-runtime==0.5.0` wheel. Each runtime release carries exactly one -prebuilt MuJoCo binding: the 0.5.0 wheels are compiled against -`mujoco==3.11.0`, and the native extension records its build-time mujoco -version and refuses to load on a mismatch (see the watchdog row in "Install -Error Signatures"). A bump of the default MuJoCo version therefore always -ships with a new runtime release; the coordination details live in the -mujoco-uni-runtime repository's `docs/release-coordination.md`. - -Switching the MuJoCo version inside the support window `>=3.5,<3.12` always -takes the source-rebuild path: +The `mujoco` extra declares `mujoco~=3.11.0`, and the `mjbatch` batch engine +is built against `mujoco==3.11.0`: it records its build-time mujoco version +and refuses to import against a different one (see the watchdog row in +"Install Error Signatures"). Switching the local MuJoCo version therefore +requires an `mjbatch` build against that version — it is not a UniLab config +change: -```bash -make mujoco MJ=3.10.0 -``` +1. bump the `mujoco` bound and the `mjbatch` source pin in `pyproject.toml` + (and mirror `pyproject.rocm.toml`), +2. re-lock (`uv lock`, plus the ROCm lockfile via `make sync-rocm`) and + re-sync (`uv sync --extra mujoco`). -The `mujoco` extra declares `mujoco~=3.11.0`, which a re-lock can never leave, -so the target operates on the environment directly (`uv pip`, without touching -`uv.lock`). It runs, in order: - -1. the `check-cxx-toolchain` preflight: fails fast when no C++ compiler is - found and prints per-platform install commands; -2. `uv pip install "mujoco==3.10.0" pybind11 wheel setuptools`: installs the - requested mujoco plus the runtime's build requirements into the current - environment; -3. `uv cache clean mujoco-uni-runtime`: drops the build cache (uv's cache - cannot see that the extension depends on the mujoco version); -4. `uv pip install --force-reinstall --no-deps --no-build-isolation - --no-binary mujoco-uni-runtime "mujoco-uni-runtime=="`: - recompiles the native extension from the sdist against the new mujoco. - -The override is **environment-local**: `uv.lock` stays unchanged. The -switch-back path is `uv sync --extra mujoco --reinstall-package -mujoco-uni-runtime`, which restores the locked default (mujoco 3.11.0 + -prebuilt wheel); the `--reinstall-package` flag is required because a plain -`uv sync` restores mujoco but keeps the locally rebuilt extension, which then -fails to load. The source rebuild requires a C++17 toolchain and Python -development headers (see "Requirements"). +The fork's build pins `mujoco==3.11.0` at build time, so the isolated build +always compiles against the matching mujoco. Until the fork's distribution +identity is decided (PyPI package vs git pin — the roadmap's open item), +coordinate version bumps with the +[fork](https://github.com/unilabsim/mjbatch). ## Install Error Signatures @@ -250,12 +221,9 @@ Reverse-lookup from error text to cause and fix. | Error signature | Where it comes from | Fix | | --- | --- | --- | -| `error: building mujoco-uni-runtime from source requires a C++ toolchain, but 'c++' was not found.` | The `check-cxx-toolchain` preflight of `make mujoco MJ=` | Install a C++ toolchain and retry: Debian/Ubuntu `sudo apt-get install build-essential`; macOS `xcode-select --install`; Fedora/RHEL `sudo dnf install gcc-c++ make` | -| `error: [Errno 2] No such file or directory: 'c++'` (or `c++: No such file or directory`) | Building `mujoco-uni-runtime` from the sdist without a compiler; only occurs on the source-rebuild path — the default wheel path never compiles | Same toolchain install as above; or do not switch versions and use the default wheel path `uv sync --extra mujoco` | -| `fatal error: Python.h: No such file or directory` | Missing Python development headers during a source rebuild | A uv-managed Python (`uv python install`) bundles the headers; system Pythons need `python3-dev` (Debian/Ubuntu) or `python3-devel` (Fedora/RHEL) | -| `MuJoCoUni native batch extension was built against mujoco '3.11.0', but loaded mujoco is '...'` | Version watchdog: the extension's recorded build-time mujoco version does not match the loaded mujoco | Install the mujoco version the extension binds (the prebuilt wheel binds `3.11.0`: `uv sync --extra mujoco --reinstall-package mujoco-uni-runtime`); or rebuild from source against the active mujoco: `make mujoco MJ=` | -| `mujoco_uni 0.5.0 supports official mujoco>=3.5,<3.12; found mujoco '...'` | The installed mujoco is outside the runtime's support window | Install a mujoco version inside `>=3.5,<3.12` (`make mujoco MJ=`) | -| `MuJoCoUni native batch extension has not been built` | The native extension failed to import (`mujoco_uni.batch_available()` returns `False`); a common cause is a plain `uv sync` after a version switch, which restores mujoco but keeps the locally rebuilt extension linked to the old `libmujoco.so` | Run `uv run python -c "import mujoco_uni; print(mujoco_uni.batch_import_error())"` for the underlying cause; after a version switch, restore the prebuilt wheel with `uv sync --extra mujoco --reinstall-package mujoco-uni-runtime` | +| `fatal error: Python.h: No such file or directory` | Missing Python development headers while uv builds the `mjbatch` fork from source | A uv-managed Python (`uv python install`) bundles the headers; system Pythons need `python3-dev` (Debian/Ubuntu) or `python3-devel` (Fedora/RHEL) | +| `error: [Errno 2] No such file or directory: 'c++'` (or `c++: No such file or directory`) | Building `mjbatch` from the pinned git source without a compiler | Install a C++ toolchain (see "Requirements") and retry `uv sync --extra mujoco` | +| `mjbatch was built against MuJoCo but is installed` | Version watchdog: the engine's build-time mujoco pin does not match the installed mujoco | Restore the locked pair with `uv sync --extra mujoco`; a different mujoco version requires an `mjbatch` rebuild (see "Switching The Local MuJoCo Version") | ## Platform Profiles @@ -269,10 +237,12 @@ uses the `mjpython` application bundled by the official MuJoCo wheel. Torch's `cuda` alias resolves to MPS when CUDA is absent. On Windows, use the direct `uv sync` commands from above unless GNU `make` and -Bash are available. The default install uses the prebuilt wheel; MSVC Build -Tools and Python development headers are only needed when rebuilding the MuJoCo -native extension from source (version switch). If you want to use the Makefile, -install GNU Make and Bash separately (for example through Chocolatey or WSL). +Bash are available. The `mjbatch` engine does not ship Windows wheels (the +official `mujoco.dll` provides no import library), so the MuJoCo physics +backend stays Linux/macOS-only there; plain `mujoco` (MJCF conversion, +playback rendering via other backends) still installs. If you want to use the +Makefile, install GNU Make and Bash separately (for example through Chocolatey +or WSL). ROCm and Intel XPU have explicit Makefile targets: diff --git a/docs/sphinx/source/en/2-user_guide/1-training/1-cli_reference.md b/docs/sphinx/source/en/2-user_guide/1-training/1-cli_reference.md index 030ea438b..11a3291b2 100644 --- a/docs/sphinx/source/en/2-user_guide/1-training/1-cli_reference.md +++ b/docs/sphinx/source/en/2-user_guide/1-training/1-cli_reference.md @@ -45,16 +45,14 @@ completion: # Full default environment (MuJoCo + Motrix): make setup -# MuJoCo only: -# make setup-mujoco - # Motrix only (the shortest path for Motrix demos): # make setup-motrix ``` `make setup` runs `uv sync --extra mujoco --extra motrix` followed by -`uv run --no-sync unilab-complete install`; `make setup-mujoco` and -`make setup-motrix` select only their named extra and install the same completion. +`uv run --no-sync unilab-complete install`; `make setup-motrix` selects only +the Motrix extra and installs the same completion; use bare +`uv sync --extra mujoco` for a MuJoCo-only environment. Choose one setup path for an environment. The install command picks Bash or Zsh from `$SHELL` / platform and only writes user-level rc files. The current shell is not auto-activated; reopen the terminal or source the rc file to apply. diff --git a/docs/sphinx/source/en/2-user_guide/3-backends/1-mujoco.md b/docs/sphinx/source/en/2-user_guide/3-backends/1-mujoco.md index 27322bcc7..3ac63f7e8 100644 --- a/docs/sphinx/source/en/2-user_guide/3-backends/1-mujoco.md +++ b/docs/sphinx/source/en/2-user_guide/3-backends/1-mujoco.md @@ -2,9 +2,10 @@ MuJoCo is the default backend path in the committed owner configs. The Python dependencies are the official `mujoco` package (`~=3.11.0`, with the exact -default version pinned by the committed `uv.lock`) plus -`mujoco-uni-runtime` in `pyproject.toml`, and the adapter lives -under `unisim.backend.mujoco`. +default version pinned by the committed `uv.lock`) plus the +`mjbatch` native batch engine (currently pinned to the +[integration fork](https://github.com/unilabsim/mjbatch) in +`pyproject.toml`), and the adapter lives under `unisim.backend.mujoco`. ## When To Use It @@ -31,42 +32,19 @@ rather than opening the Motrix native interactive renderer. pyproject constrains `mujoco~=3.11.0`; the committed `uv.lock` pins the exact default version, and uv's prefer-locked semantics keep ordinary relocks from -drifting. The default install path uses the prebuilt `mujoco-uni-runtime` -wheel, which binds `mujoco==3.11.0` — no compiler is needed. The support -window is `>=3.5,<3.12`; switching to any version other than the wheel's -binding always takes the source-rebuild path. The -`mujoco-uni-runtime` native extension records its build-time `mujoco` -version and refuses to load against any other version, -so switching versions means installing the requested `mujoco` and rebuilding -the extension from source (a C++17 toolchain and Python development headers -are required; the target's `check-cxx-toolchain` preflight fails fast with -per-platform install commands when no compiler is found): - -```bash -make mujoco MJ=3.10.0 -``` - -Because the `~=` bound can never be re-locked to another line, the target -operates on the environment directly (`uv pip`, without touching `uv.lock`): -it installs `mujoco==3.10.0` plus the runtime's build requirements -(`pybind11`, `wheel`, `setuptools`), clears uv's build -cache for `mujoco-uni-runtime` (the cache cannot see that the extension -depends on the mujoco version), and forces an in-env sdist rebuild of the -runtime. -Without the Makefile shortcut, the equivalent is: - -```bash -uv pip install "mujoco==3.10.0" pybind11 wheel setuptools -uv cache clean mujoco-uni-runtime -uv pip install --force-reinstall --no-deps --no-build-isolation \ - --no-binary mujoco-uni-runtime "mujoco-uni-runtime==0.5.0" -``` - -Skipping the cache clean or the forced reinstall lets uv reuse a cached -extension built against the previous mujoco version, which then fails to -import with a version-watchdog error (fail-closed, never a silent behavior -change). The override is environment-local: to return to the default -prebuilt-wheel path, run `uv sync --extra mujoco --reinstall-package -mujoco-uni-runtime`. The `--reinstall-package` flag is required — a plain -`uv sync` restores mujoco but keeps the locally rebuilt extension, which -then fails to load against the reverted mujoco. +drifting. The `mjbatch` engine is built against `mujoco==3.11.0` and records +its build-time `mujoco` version, refusing to import against any other version +(fail-closed, never a silent behavior change). Switching versions therefore +requires an `mjbatch` build against the requested version: + +1. bump the `mujoco` bound and the `mjbatch` source pin in + `pyproject.toml` (and mirror `pyproject.rocm.toml`), +2. re-lock (`uv lock`, plus the ROCm lockfile via `make sync-rocm`) and + re-sync (`uv sync --extra mujoco`). + +The fork's build pins `mujoco==3.11.0` at build time, so the isolated build +always compiles against the matching mujoco. The engine is currently consumed +from the pinned integration fork; its final distribution identity (PyPI +package vs git pin, and prebuilt wheels) is the roadmap's open item — until +then, coordinate version bumps with the +[fork](https://github.com/unilabsim/mjbatch). diff --git a/docs/sphinx/source/en/2-user_guide/3-backends/8-superdex.md b/docs/sphinx/source/en/2-user_guide/3-backends/8-superdex.md index 73f1d13c8..63a3cc0e6 100644 --- a/docs/sphinx/source/en/2-user_guide/3-backends/8-superdex.md +++ b/docs/sphinx/source/en/2-user_guide/3-backends/8-superdex.md @@ -30,7 +30,7 @@ pip install "unilab[superdex]" The extra delegates version pins to UniSim through `unisim-core[superdex]`. UniSim's superdex extra already carries the plain `mujoco` package (used for -MJCF conversion and the offline playback renderer); `mujoco-uni-runtime` is +MJCF conversion and the offline playback renderer); `mjbatch` is not required — only the MuJoCo physics backend needs it. The current wheels are a temporary unilabsim build (`superdex-physics-uni` / `superdex-robotics-uni`); once the upstream project_superdex release publishes diff --git a/docs/sphinx/source/en/2-user_guide/5-domain_randomization/0-index.md b/docs/sphinx/source/en/2-user_guide/5-domain_randomization/0-index.md index b559859ae..823bca27b 100644 --- a/docs/sphinx/source/en/2-user_guide/5-domain_randomization/0-index.md +++ b/docs/sphinx/source/en/2-user_guide/5-domain_randomization/0-index.md @@ -144,7 +144,8 @@ Notes: - `gravity_range` must be convertible into a `(2, 3)` array; otherwise reset will raise an error when constructing the payload. - This term does not call `mj_setConst`; MuJoCo step / forward reads `mjModel.opt.gravity` directly. - Do not enable this term under the Motrix backend; the current Motrix capability does not include `gravity`. -- If your current environment still has a `mujoco-uni-runtime` package installed that does not include the `gravity` field, MuJoCo reset will raise unsupported field; you need to use a `mujoco-uni-runtime` build/release that includes the field. +- The MuJoCo backend writes gravity through the `mjbatch` per-simulation model + expansion (`expand("gravity")`), which the pinned `mjbatch` build ships. - During training it is recommended to start from a small tilt range; otherwise sampling a too-large horizontal gravity early on may degrade the task into being unlearnable. ## Interval push Usage diff --git a/docs/sphinx/source/en/3-deployment/2-sim_to_sim/1-backend_swap.md b/docs/sphinx/source/en/3-deployment/2-sim_to_sim/1-backend_swap.md index 268d444a5..2d7f6e3b0 100644 --- a/docs/sphinx/source/en/3-deployment/2-sim_to_sim/1-backend_swap.md +++ b/docs/sphinx/source/en/3-deployment/2-sim_to_sim/1-backend_swap.md @@ -1,7 +1,7 @@ # Backend Swap UniLab supports two CPU physics backends: **MuJoCo** (via the official `mujoco` -package plus `mujoco-uni-runtime`) and +package plus the `mjbatch` batch engine) and **Motrix** (via `motrixsim-core`). Both implement the same `SimBackend` contract and the same env contract. Backend-specific behavior is exposed through explicit methods and capability records. diff --git a/docs/sphinx/source/en/4-developer_guide/2-contracts/4-dr_contract.md b/docs/sphinx/source/en/4-developer_guide/2-contracts/4-dr_contract.md index 349c8bf12..c7f15979b 100644 --- a/docs/sphinx/source/en/4-developer_guide/2-contracts/4-dr_contract.md +++ b/docs/sphinx/source/en/4-developer_guide/2-contracts/4-dr_contract.md @@ -77,12 +77,13 @@ optional `body_ids`) from `unisim.dr.interval`, re-exported through providers should populate `ops`; the legacy fields will be removed in the next unisim-core major release. -## MuJoCo BatchEnvPool Snapshot +## MuJoCo mjbatch Snapshot -Current MuJoCo reset randomization uses `BatchEnvPool.reset(..., -randomization=...)` with a fixed field whitelist. Indexed reads and writes are -available through `get_field_indexed(...)` and `set_field_indexed(...)`. This -interface lives in the `mujoco-uni-runtime` package (`mujoco_uni.batch_env`), not in this +Current MuJoCo reset randomization writes the nine supported fields through +`mjbatch` per-simulation model views: the backend expands a field with +`Batch.expand(name)`, writes the targeted env rows, then refreshes derived +constants with `Batch.set_const(ids)` before the fused reset runs `mj_forward`. +This interface lives in the `mjbatch` package, not in this repository; the reset-term constants that map onto it are in `unisim.dr.types`. @@ -109,14 +110,13 @@ not. Two caveats: -- `geom_size` is not in `SUPPORTED_FIELDS`. Geometry size is expressed through - init-lifecycle model materialization (see `GeomSizeOverride` / +- `geom_size` is not in the supported reset fields. Geometry size is expressed + through init-lifecycle model materialization (see `GeomSizeOverride` / `ModelVariantSpec` in `unisim.dr.types`), not reset randomization. -- `gravity` reset randomization requires a `mujoco-uni-runtime` build that ships - it. This repository depends on the official `mujoco` package (`>=3.5`, with - the default version pinned by `uv.lock`) - plus `mujoco-uni-runtime`, whose `SUPPORTED_FIELDS` includes `gravity`; older - batch-env packages such as `mujoco-uni==3.6.0.post6` do not. +- `gravity` reset randomization requires an `mjbatch` build that ships it + (`expand("gravity")` covers the `mjOption` vector). This repository depends + on the official `mujoco` package (`~=3.11.0`, with the default version pinned + by `uv.lock`) plus `mjbatch`, whose expandable fields include `gravity`. ## Motor Control Extension diff --git a/docs/sphinx/source/en/4-developer_guide/4-contributing.md b/docs/sphinx/source/en/4-developer_guide/4-contributing.md index 7b153e6bb..064650beb 100644 --- a/docs/sphinx/source/en/4-developer_guide/4-contributing.md +++ b/docs/sphinx/source/en/4-developer_guide/4-contributing.md @@ -8,7 +8,7 @@ architecture details live in {doc}`1-architecture/1-overview`. Install dependencies for your platform. The setup targets also install the optional simulator extras used by the repository's checks: -- macOS (MPS, PyPI torch wheel): `make setup-motrix` (or `make setup-mujoco`) +- macOS (MPS, PyPI torch wheel): `make setup-motrix` (or `uv sync --extra mujoco`) - Linux with NVIDIA (PyTorch cu128 wheel): `make setup` - Linux AMD / ROCm: `make sync-rocm`, then run commands with `uv run ...`. To return to the default CUDA / macOS profile, `git restore -- pyproject.toml @@ -21,7 +21,6 @@ optional simulator extras used by the repository's checks: ```bash # Choose one core setup path: make setup -# make setup-mujoco # make setup-motrix make sync-rocm make sync-xpu diff --git a/docs/sphinx/source/zh_CN/1-getting_started/2-installation.md b/docs/sphinx/source/zh_CN/1-getting_started/2-installation.md index ba055330d..67bcee9f0 100644 --- a/docs/sphinx/source/zh_CN/1-getting_started/2-installation.md +++ b/docs/sphinx/source/zh_CN/1-getting_started/2-installation.md @@ -9,11 +9,12 @@ - Git 和 `curl`，用于克隆仓库及下载 runtime asset。 - `cmake`，构建 Drake 原生 batch extension 时需要。Drake setup 脚本使用 CMake 和 C++ 工具链。 -- 使用 `mujoco` extra 时：默认安装路径使用 `mujoco-uni-runtime` 的预编译 wheel - （绑定 `mujoco==3.11.0`），`make setup` / `uv sync --extra mujoco` 不需要编译器。 - 只有显式的源码重建路径（切换 MuJoCo 版本，见「切换本地 MuJoCo 版本」）才需要 - C++17 工具链和 Python 开发头文件；缺少时会以 - `fatal error: Python.h: No such file or directory` 等错误失败（完整对照见 +- 使用 `mujoco` extra 时：MuJoCo 物理后端运行在 `mjbatch` 原生 batch 引擎上， + 当前从钉住的集成 fork（`unilabsim/mjbatch`）源码构建。uv 以隔离构建 + （scikit-build-core + nanobind）编译该 git 源码，需要 C++17 工具链和 + Python 开发头文件；构建绑定 `mujoco==3.11.0`，引擎在 mujoco 版本不一致时 + 拒绝导入。免编译器的安装路径取决于 fork 的最终分发渠道（预编译 wheel 还是 + git 钉版），这是 roadmap 上的待定事项（见「切换本地 MuJoCo 版本」和 「安装错误特征对照表」）。 - macOS：`xcode-select --install` - Ubuntu / Debian：`sudo apt-get install build-essential python3-dev` @@ -58,15 +59,11 @@ make setup # 运行第一次 Motrix demo 的最快路径。 # make setup-motrix - -# 仅安装 MuJoCo。 -# make setup-mujoco ``` `make setup` 会运行 `uv sync --extra mujoco --extra motrix` 并安装 shell 自动补全。 `make setup-motrix` 会运行 `uv sync --extra motrix` 并安装相同的补全条目。 -`make setup-mujoco` 会运行 `uv sync --extra mujoco` 并安装补全。三条路径只选择一条。 -如果 `make` 不可用，可运行对应的底层命令： +两条路径只选择一条。如果 `make` 不可用，可运行对应的底层命令： ```bash # 完整默认环境： @@ -75,9 +72,6 @@ uv run --no-sync unilab-complete install # 仅 Motrix： # uv sync --extra motrix && uv run --no-sync unilab-complete install - -# 仅 MuJoCo： -# uv sync --extra mujoco && uv run --no-sync unilab-complete install ``` ## conda 与 pip @@ -95,7 +89,7 @@ cd UniLab make setup-motrix ``` -如果不需要 Motrix，可使用 `make setup-mujoco`；ROCm / XPU 仍走下方专用的 `make` 路径。 +如果不需要 Motrix，可使用 `uv sync --extra mujoco`；ROCm / XPU 仍走下方专用的 `make` 路径。 从源码 checkout 使用 pip 时，这是备用路径。请先安装 package，再显式添加所需 runtime： @@ -109,17 +103,15 @@ pip install -e . # 需要 Motrix 时： pip install motrixsim-core==0.8.2 -# 需要 MuJoCo 时（默认安装绑定 mujoco==3.11.0 的预编译 wheel）： -pip install "mujoco~=3.11.0" "mujoco-uni-runtime==0.5.0" +# 需要 MuJoCo 时（解析钉住的 mjbatch 集成 fork，针对 mujoco==3.11.0 构建）： +pip install "mujoco~=3.11.0" "mjbatch @ git+https://github.com/unilabsim/mjbatch.git@cf4a83d" ``` editable install 会指向源码 checkout；常规安装会把 package 和任务配置 （`unilab/conf/`）复制进环境。两种方式都支持在任意目录运行 `train` / `eval` / `demo`， -日志与 checkpoint 写入当前工作目录。MuJoCo runtime 的预编译 wheel 直接通过 pip -安装即可，无需额外构建步骤；只有强制从 sdist 重建（绑定非默认 mujoco 版本）时 -才需要 `pybind11` / `wheel` 与 `--no-build-isolation`，见「切换本地 MuJoCo 版本」。 -MJWarp、Genesis、平台相关 torch index、ROCm / XPU profile 请优先使用上面的 -uv 路径。机器人 mesh 和纹理不会打进 wheel，而是在 cold path 从 +日志与 checkpoint 写入当前工作目录。`mjbatch` 引擎针对钉住的 +`mujoco==3.11.0` 构建；MJWarp、Genesis、平台相关 torch index、ROCm / XPU +profile 请优先使用上面的 uv 路径。机器人 mesh 和纹理不会打进 wheel，而是在 cold path 从 `unilabsim/unilab-robots` 数据集下载。请确保安装位置可写，或从源码 checkout 使用 `uv run unilab-pull-assets` 预拉取。isaacgym / isaacsim 后端仍假设 源码 checkout；外部后端请使用下方专用安装页。 @@ -160,7 +152,7 @@ uv sync --extra mujoco --extra mjwarp --extra newton | 后端 | 安装路径 | 重要前置条件 | | --- | --- | --- | -| MuJoCo | `make setup-mujoco` 或 `uv sync --extra mujoco` | 默认使用预编译 wheel（绑定 `mujoco==3.11.0`），无需编译器；仅切换版本（`make mujoco MJ=`）时需要 C++17 工具链和 Python 开发头文件 | +| MuJoCo | `make setup` 或 `uv sync --extra mujoco` | 从源码构建钉住的 `mjbatch` fork（绑定 `mujoco==3.11.0`）；在预编译 wheel 可用之前（roadmap 待定事项）需要 C++17 工具链和 Python 开发头文件 | | Motrix | `make setup-motrix` 或 `uv sync --extra motrix` | 从固定版本 Python package 安装 Motrix runtime | | MJWarp | `uv sync --extra mujoco --extra mjwarp` | NVIDIA CUDA 和显式 CUDA process device | | Genesis | `uv sync --extra genesis` | 已验证路径使用 Linux x86_64、NVIDIA GPU 及固定版本 torch/Genesis | @@ -186,36 +178,20 @@ Drake、IsaacGym 和 IsaacSim 的 setup 脚本会将外部 runtime 安装到仓 ## 切换本地 MuJoCo 版本 -`mujoco` extra 的默认安装路径使用 `mujoco-uni-runtime==0.5.0` 的预编译 wheel。 -每个 runtime 发布版本只携带一个预编译 MuJoCo 绑定：0.5.0 的 wheel 针对 -`mujoco==3.11.0` 编译，原生扩展会记录编译时的 mujoco 版本，加载时检测到不一致 -会拒绝工作（见「安装错误特征对照表」中的 watchdog 行）。因此 MuJoCo 默认版本的 -升级总是伴随一次新的 runtime 发布；发布协调细节见 mujoco-uni-runtime 仓库的 -`docs/release-coordination.md`。 - -在支持窗口 `>=3.5,<3.12` 内切换 MuJoCo 版本总是走源码重建路径： - -```bash -make mujoco MJ=3.10.0 -``` - -`mujoco` extra 声明的 `mujoco~=3.11.0` 边界意味着 relock 无法选出窗口内的旧版本， -所以该目标直接操作当前环境（`uv pip`，不改动 `uv.lock`），依次执行： +`mujoco` extra 声明 `mujoco~=3.11.0`，`mjbatch` batch 引擎针对 +`mujoco==3.11.0` 构建：引擎会记录编译时的 mujoco 版本，加载时检测到不一致 +会拒绝导入（见「安装错误特征对照表」中的 watchdog 行）。因此切换本地 +MuJoCo 版本需要对应版本的 `mjbatch` 构建，而不是修改 UniLab 配置： -1. `check-cxx-toolchain` 预检：缺少 C++ 编译器时立即失败，并打印各平台的安装命令； -2. `uv pip install "mujoco==3.10.0" pybind11 wheel setuptools`：把请求的 mujoco - 和 runtime 的构建依赖装进当前环境； -3. `uv cache clean mujoco-uni-runtime`：清除构建缓存（uv 的缓存无法感知该扩展 - 依赖 mujoco 版本）； -4. `uv pip install --force-reinstall --no-deps --no-build-isolation - --no-binary mujoco-uni-runtime "mujoco-uni-runtime==<当前版本>"`：从 sdist - 针对新 mujoco 就地重新编译原生扩展。 +1. 在 `pyproject.toml`（并同步镜像 `pyproject.rocm.toml`）中提升 `mujoco` + 边界和 `mjbatch` 源码钉版； +2. 重新锁定（`uv lock`，ROCm lockfile 通过 `make sync-rocm`）并重新同步 + （`uv sync --extra mujoco`）。 -这个覆盖是**环境本地**的：`uv.lock` 不变。切回默认版本的路径是 -`uv sync --extra mujoco --reinstall-package mujoco-uni-runtime`，恢复到 lock -钉住的默认状态（mujoco 3.11.0 + 预编译 wheel）；`--reinstall-package` 不可省略， -因为裸 `uv sync` 只恢复 mujoco 而保留本地重编的扩展，扩展会因此无法加载。 -源码重建需要 C++17 工具链和 Python 开发头文件（见「环境要求」）。 +fork 的构建在编译期钉住 `mujoco==3.11.0`，因此隔离构建总是针对匹配的 mujoco +编译。在 fork 的最终分发身份确定之前（PyPI package 还是 git 钉版——roadmap +待定事项），版本提升需要与 [fork](https://github.com/unilabsim/mjbatch) +维护者协调。 ## 安装错误特征对照表 @@ -223,12 +199,9 @@ make mujoco MJ=3.10.0 | 报错特征 | 触发场景 | 修复 | | --- | --- | --- | -| `error: building mujoco-uni-runtime from source requires a C++ toolchain, but 'c++' was not found.` | `make mujoco MJ=` 的 `check-cxx-toolchain` 预检失败 | 安装 C++ 工具链后重试：Debian/Ubuntu `sudo apt-get install build-essential`；macOS `xcode-select --install`；Fedora/RHEL `sudo dnf install gcc-c++ make` | -| `error: [Errno 2] No such file or directory: 'c++'`（或 `c++: No such file or directory`） | 从 sdist 编译 `mujoco-uni-runtime` 时缺少编译器；只会出现在源码重建路径，默认 wheel 路径不会编译 | 同上安装工具链；或者不切换版本，直接走默认 wheel 路径 `uv sync --extra mujoco` | -| `fatal error: Python.h: No such file or directory` | 源码重建时缺少 Python 开发头文件 | uv 托管的 Python（`uv python install`）自带头文件；系统 Python 需安装 `python3-dev`（Debian/Ubuntu）或 `python3-devel`（Fedora/RHEL） | -| `MuJoCoUni native batch extension was built against mujoco '3.11.0', but loaded mujoco is '...'` | 版本 watchdog：扩展记录的编译期 mujoco 版本与当前加载的 mujoco 不一致 | 安装扩展绑定的 mujoco 版本（预编译 wheel 绑定 `3.11.0`：`uv sync --extra mujoco --reinstall-package mujoco-uni-runtime`）；或针对当前 mujoco 从源码重建：`make mujoco MJ=` | -| `mujoco_uni 0.5.0 supports official mujoco>=3.5,<3.12; found mujoco '...'` | 环境中的 mujoco 超出 runtime 支持窗口 | 换装窗口 `>=3.5,<3.12` 内的 mujoco 版本（`make mujoco MJ=`） | -| `MuJoCoUni native batch extension has not been built` | 原生扩展导入失败（`mujoco_uni.batch_available()` 返回 `False`）；常见诱因是版本切换后裸跑 `uv sync`：mujoco 被恢复而本地重编的扩展仍链接旧版 `libmujoco.so` | 运行 `uv run python -c "import mujoco_uni; print(mujoco_uni.batch_import_error())"` 查看底层原因；版本切换后用 `uv sync --extra mujoco --reinstall-package mujoco-uni-runtime` 恢复预编译 wheel | +| `fatal error: Python.h: No such file or directory` | uv 从源码构建 `mjbatch` fork 时缺少 Python 开发头文件 | uv 托管的 Python（`uv python install`）自带头文件；系统 Python 需安装 `python3-dev`（Debian/Ubuntu）或 `python3-devel`（Fedora/RHEL） | +| `error: [Errno 2] No such file or directory: 'c++'`（或 `c++: No such file or directory`） | 缺少编译器，无法从钉住的 git 源码构建 `mjbatch` | 安装 C++ 工具链（见「环境要求」）后重试 `uv sync --extra mujoco` | +| `mjbatch was built against MuJoCo but is installed` | 版本 watchdog：引擎的编译期 mujoco 钉版与已安装的 mujoco 不一致 | 用 `uv sync --extra mujoco` 恢复 lock 钉住的组合；更换 mujoco 版本需要重建 `mjbatch`（见「切换本地 MuJoCo 版本」） | ## 平台配置档 @@ -241,9 +214,10 @@ wheel 来源是在 `pyproject.toml` 中配置的 PyTorch `cu128` 索引。 没有 CUDA 时，`cuda` alias 会解析到 MPS。 在 Windows 上，如果没有 GNU `make` 和 Bash，请使用上面的直接 `uv sync` 命令。 -默认安装使用预编译 wheel；只有源码重建 MuJoCo 原生扩展（切换版本）时才需要 -MSVC Build Tools 和 Python 开发头文件。如果要使用 -Makefile，请另行安装 GNU Make 和 Bash（例如通过 Chocolatey 或 WSL）。 +`mjbatch` 引擎不提供 Windows wheel（官方的 `mujoco.dll` 不带 import +library），因此 MuJoCo 物理后端在 Windows 上仍不可用；纯 `mujoco` +（MJCF 转换、其他后端的回放渲染）仍可安装。如果要使用 Makefile，请另行安装 +GNU Make 和 Bash（例如通过 Chocolatey 或 WSL）。 ROCm 和 Intel XPU 有各自显式的 Makefile 目标： diff --git a/docs/sphinx/source/zh_CN/2-user_guide/1-training/1-cli_reference.md b/docs/sphinx/source/zh_CN/2-user_guide/1-training/1-cli_reference.md index e66b458e0..a56c477f9 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/1-training/1-cli_reference.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/1-training/1-cli_reference.md @@ -43,16 +43,14 @@ choices，不改变命令行为。在新 checkout 上，可用一条 setup 命 # 完整默认环境（MuJoCo + Motrix）： make setup -# 仅 MuJoCo： -# make setup-mujoco - # 仅 Motrix（运行 Motrix demo 的最快路径）： # make setup-motrix ``` `make setup` 会执行 `uv sync --extra mujoco --extra motrix`，随后运行 -`uv run --no-sync unilab-complete install`；`make setup-mujoco` 和 -`make setup-motrix` 只选择对应的 extra，并安装相同的补全。一个环境只选择一条安装路径。 +`uv run --no-sync unilab-complete install`；`make setup-motrix` 只选择 +Motrix extra，并安装相同的补全；仅需要 MuJoCo 时使用裸 +`uv sync --extra mujoco`。一个环境只选择一条安装路径。 安装命令会按 `$SHELL` / 平台选择 Bash 或 Zsh，只写入用户级 rc 文件。当前终端不会被自动 激活，重新打开终端或 source 对应 rc 文件后生效。 diff --git a/docs/sphinx/source/zh_CN/2-user_guide/3-backends/1-mujoco.md b/docs/sphinx/source/zh_CN/2-user_guide/3-backends/1-mujoco.md index fca4beaa2..16d31f781 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/3-backends/1-mujoco.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/3-backends/1-mujoco.md @@ -1,7 +1,9 @@ # MuJoCo 后端 MuJoCo 是已提交 owner 配置中的默认后端路径。其 Python 依赖为官方 -`mujoco` 包（`~=3.11.0`，默认版本由已提交的 `uv.lock` 精确钉住）加 `mujoco-uni-runtime`（见 `pyproject.toml`），适配层位于 +`mujoco` 包（`~=3.11.0`，默认版本由已提交的 `uv.lock` 精确钉住）加 +`mjbatch` 原生 batch 引擎（当前在 `pyproject.toml` 中钉住 +[集成 fork](https://github.com/unilabsim/mjbatch)），适配层位于 `unisim.backend.mujoco` 下。 ## 何时使用 @@ -26,34 +28,17 @@ MuJoCo 在 `unisim.backend.mujoco.backend` 中声明对物理状态回放的支 ## 切换 MuJoCo 版本 pyproject 约束 `mujoco~=3.11.0`；默认版本由已提交的 `uv.lock` 精确钉住， -uv 的 prefer-locked 语义保证普通 relock 不会漂移。默认安装路径使用 -`mujoco-uni-runtime` 的预编译 wheel（绑定 `mujoco==3.11.0`），不需要编译器。 -支持窗口为 `>=3.5,<3.12`；切换到 wheel 绑定以外的任何版本都必须走源码重建 -路径。`mujoco-uni-runtime` 的原生扩展会记录编译时的 `mujoco` 版本，且拒绝在 -其它版本下加载，因此切换版本 = 安装指定 `mujoco` + 从源码重编扩展（需要 C++17 -工具链和 Python 开发头文件；缺少编译器时，该目标的 `check-cxx-toolchain` -预检会立即失败并打印各平台的安装命令）： - -```bash -make mujoco MJ=3.10.0 -``` - -`~=` 边界无法 relock 到其它版本线，因此该目标直接操作当前环境 -（`uv pip`，不改动 `uv.lock`）：装入 `mujoco==3.10.0` 和 runtime 的构建依赖 -（`pybind11`、`wheel`、`setuptools`）、清除 uv 对 -`mujoco-uni-runtime` 的构建缓存（缓存无法感知扩展对 mujoco 版本的依赖）、 -并强制从 sdist 在本环境内重新编译 runtime。不用 Makefile 时的等价命令： - -```bash -uv pip install "mujoco==3.10.0" pybind11 wheel setuptools -uv cache clean mujoco-uni-runtime -uv pip install --force-reinstall --no-deps --no-build-isolation \ - --no-binary mujoco-uni-runtime "mujoco-uni-runtime==0.5.0" -``` - -如果省略清缓存或强制重装，uv 可能复用按旧版本 mujoco 编译的扩展， -import 时会以版本 watchdog 错误失败（fail-closed，不会静默出错行为）。 -该覆盖是环境本地的：切回默认预编译 wheel 路径用 -`uv sync --extra mujoco --reinstall-package mujoco-uni-runtime`。 -`--reinstall-package` 不可省略——裸 `uv sync` 只恢复 mujoco 而保留本地重编的 -扩展，扩展会因此无法加载。 +uv 的 prefer-locked 语义保证普通 relock 不会漂移。`mjbatch` 引擎针对 +`mujoco==3.11.0` 构建，并记录编译时的 `mujoco` 版本，拒绝在其它版本下导入 +（fail-closed，不会静默出错行为）。因此切换版本需要对应版本的 `mjbatch` +构建： + +1. 在 `pyproject.toml`（并同步镜像 `pyproject.rocm.toml`）中提升 `mujoco` + 边界和 `mjbatch` 源码钉版； +2. 重新锁定（`uv lock`，ROCm lockfile 通过 `make sync-rocm`）并重新同步 + （`uv sync --extra mujoco`）。 + +fork 的构建在编译期钉住 `mujoco==3.11.0`，因此隔离构建总是针对匹配的 mujoco +编译。引擎当前从钉住的集成 fork 消费；其最终分发身份（PyPI package 还是 +git 钉版、是否提供预编译 wheel）是 roadmap 上的待定事项——在此之前，版本提升 +需要与 [fork](https://github.com/unilabsim/mjbatch) 维护者协调。 diff --git a/docs/sphinx/source/zh_CN/2-user_guide/3-backends/8-superdex.md b/docs/sphinx/source/zh_CN/2-user_guide/3-backends/8-superdex.md index 1f802bab6..c4ce4de21 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/3-backends/8-superdex.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/3-backends/8-superdex.md @@ -28,7 +28,7 @@ pip install "unilab[superdex]" 该 extra 通过 `unisim-core[superdex]` 委托 UniSim 钉定版本；UniSim 的 superdex extra 已包含普通 `mujoco` 包（MJCF 转换与离线回放渲染使用），不需要 -`mujoco-uni-runtime`（仅 MuJoCo 物理后端需要）。当前 wheel 为临时 +`mjbatch`（仅 MuJoCo 物理后端需要）。当前 wheel 为临时 unilabsim 构建（`superdex-physics-uni`/`superdex-robotics-uni`）；上游 project_superdex 发布正式 `superdex-physics`/`superdex-robotics` wheel 后， UniSim 会切换包名，UniLab 侧无需改动。 diff --git a/docs/sphinx/source/zh_CN/2-user_guide/5-domain_randomization/0-index.md b/docs/sphinx/source/zh_CN/2-user_guide/5-domain_randomization/0-index.md index 55c07ff78..3cc68e0f6 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/5-domain_randomization/0-index.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/5-domain_randomization/0-index.md @@ -140,7 +140,8 @@ env: - `gravity_range` 必须可转换为 `(2, 3)` 数组；否则 reset 在构造 payload 时会抛出错误。 - 该项不调用 `mj_setConst`；MuJoCo step / forward 直接读取 `mjModel.opt.gravity`。 - 不要在 Motrix 后端下启用该项；当前 Motrix 能力不包含 `gravity`。 -- 如果你当前的环境仍安装了不包含 `gravity` 字段的 `mujoco-uni-runtime` 包，MuJoCo reset 会抛出 unsupported field；你需要使用包含该字段的 `mujoco-uni-runtime` 构建/发布版本。 +- MuJoCo 后端通过 `mjbatch` 的 per-simulation 模型展开（`expand("gravity")`） + 写入 gravity，钉住的 `mjbatch` 构建已包含该字段。 - 在训练期间，建议从较小的倾斜范围开始；否则在早期采样到过大的水平重力，可能会使任务退化为不可学习。 ## Interval push 用法 diff --git a/docs/sphinx/source/zh_CN/3-deployment/2-sim_to_sim/1-backend_swap.md b/docs/sphinx/source/zh_CN/3-deployment/2-sim_to_sim/1-backend_swap.md index 9a879eda7..b92d4ee31 100644 --- a/docs/sphinx/source/zh_CN/3-deployment/2-sim_to_sim/1-backend_swap.md +++ b/docs/sphinx/source/zh_CN/3-deployment/2-sim_to_sim/1-backend_swap.md @@ -1,7 +1,7 @@ # 切换后端 UniLab 支持两个 CPU 物理后端：**MuJoCo**（通过官方 `mujoco` 包加 -`mujoco-uni-runtime`）和 +`mjbatch` batch 引擎）和 **Motrix**（通过 `motrixsim-core`）。两者实现了相同的 `SimBackend` contract 和相同的 env contract。后端特有的行为通过显式方法和能力记录 （capability records）暴露出来。 diff --git a/docs/sphinx/source/zh_CN/4-developer_guide/2-contracts/4-dr_contract.md b/docs/sphinx/source/zh_CN/4-developer_guide/2-contracts/4-dr_contract.md index 74d080650..aa7f118f5 100644 --- a/docs/sphinx/source/zh_CN/4-developer_guide/2-contracts/4-dr_contract.md +++ b/docs/sphinx/source/zh_CN/4-developer_guide/2-contracts/4-dr_contract.md @@ -70,12 +70,12 @@ Interval plan 基于 term 描述符：`IntervalRandomizationPlan.ops` 携带一 op，布尔位仍作为能力回退。新 provider 应填充 `ops`；旧字段将在下一个 unisim-core major release 中移除。 -## MuJoCo BatchEnvPool 快照 +## MuJoCo mjbatch 快照 -当前 MuJoCo 的 reset 随机化使用 `BatchEnvPool.reset(..., randomization=...)`， -并带有固定的字段白名单。带索引的读写可通过 `get_field_indexed(...)` 与 -`set_field_indexed(...)` 实现。该接口位于 `mujoco-uni-runtime` 包 -（`mujoco_uni.batch_env`），不在本仓库中；映射到它的 reset-term 常量定义在 +当前 MuJoCo 的 reset 随机化通过 `mjbatch` 的 per-simulation 模型视图写入九个 +受支持字段：backend 先用 `Batch.expand(name)` 展开字段、写入目标 env 行，然后在 +融合 reset 运行 `mj_forward` 之前用 `Batch.set_const(ids)` 刷新派生常量。该接口 +位于 `mjbatch` 包，不在本仓库中；映射到它的 reset-term 常量定义在 `unisim.dr.types`。 支持的 reset 字段及其每 env 整块形状如下。首维始终是 `len(env_ids)`；尾部 @@ -99,12 +99,13 @@ refresh 行为由 backend 固定：`body_mass`、`body_ipos`、`body_iquat`、 两点注意： -- `geom_size` 不在 `SUPPORTED_FIELDS` 里。几何尺寸通过 init-lifecycle 的模型 +- `geom_size` 不在受支持的 reset 字段里。几何尺寸通过 init-lifecycle 的模型 materialization 表达（见 `unisim.dr.types` 中的 `GeomSizeOverride` / `ModelVariantSpec`），不走 reset 随机化。 -- `gravity` 的 reset 随机化需要包含它的 `mujoco-uni-runtime` 构建。本仓库依赖 - 官方 `mujoco` 包（`>=3.5`，默认版本由 `uv.lock` 钉住）加 `mujoco-uni-runtime`，其 `SUPPORTED_FIELDS` - 包含 `gravity`；更旧的 batch-env 包（例如 `mujoco-uni==3.6.0.post6`）则没有。 +- `gravity` 的 reset 随机化需要包含它的 `mjbatch` 构建 + （`expand("gravity")` 覆盖 `mjOption` 向量）。本仓库依赖官方 `mujoco` + 包（`~=3.11.0`，默认版本由 `uv.lock` 钉住）加 `mjbatch`，其可展开字段包含 + `gravity`。 ## 电机控制扩展 diff --git a/docs/sphinx/source/zh_CN/4-developer_guide/4-contributing.md b/docs/sphinx/source/zh_CN/4-developer_guide/4-contributing.md index 1eba676ae..22e8f9f83 100644 --- a/docs/sphinx/source/zh_CN/4-developer_guide/4-contributing.md +++ b/docs/sphinx/source/zh_CN/4-developer_guide/4-contributing.md @@ -7,7 +7,7 @@ 按平台安装依赖。setup target 也会安装仓库检查所需的可选仿真器 extra： -- macOS（MPS，PyPI torch wheel）：`make setup-motrix`（或 `make setup-mujoco`） +- macOS（MPS，PyPI torch wheel）：`make setup-motrix`（或 `uv sync --extra mujoco`） - Linux NVIDIA（PyTorch cu128 wheel）：`make setup` - Linux AMD / ROCm：`make sync-rocm`，随后用 `uv run ...` 运行命令。要切回默认 CUDA / macOS profile，执行 `git restore -- pyproject.toml uv.lock` 后重新 @@ -19,7 +19,6 @@ ```bash # 选择一条核心安装路径： make setup -# make setup-mujoco # make setup-motrix make sync-rocm make sync-xpu diff --git a/pyproject.rocm.toml b/pyproject.rocm.toml index bcba667b9..90f61112a 100644 --- a/pyproject.rocm.toml +++ b/pyproject.rocm.toml @@ -70,17 +70,16 @@ drake = [ mujoco = [ # MuJoCo 3.11 alignment (unilabsim/UniLab#1515): the bound is a # compatible-release specifier; uv.rocm.lock pins the exact MuJoCo - # version, and the `mujoco-uni-runtime` pin is exact because each runtime - # release carries exactly one prebuilt MuJoCo binding (its 0.5.0 wheels - # bind mujoco==3.11.0; the runtime supports >=3.5,<3.12 via sdist - # rebuilds). Switch via `make mujoco MJ=` (source rebuild path). + # version. The MuJoCo physics backend (unisim.backend.mujoco.backend) + # executes on the mjbatch native batch engine, which is built and + # version-checked against mujoco==3.11.0 — switching MuJoCo versions + # requires an mjbatch rebuild, not a UniLab config change. The batch + # engine is the unilabsim mjbatch fork, published on PyPI as + # mjbatch-uni (#1552); its build backend pins mujoco==3.11.0 at build + # time, so isolated builds are correct and no compiler preflight is + # needed. "mujoco~=3.11.0", - "mujoco-uni-runtime==0.5.0", - # pybind11/wheel are build requirements of the mujoco-uni-runtime sdist - # fallback build (see no-build-isolation-package below); the default - # install path uses the prebuilt wheel and never invokes them. - "pybind11>=2.12", - "wheel", + "mjbatch-uni~=0.1.0", ] motrix = ["motrixsim-core==0.8.2"] viser = ["viser>=1.0.26", "trimesh>=3.21.7"] @@ -109,15 +108,6 @@ exclude-dependencies = [ "nvidia-cuda-runtime-cu12", "nvidia-cudnn-cu12", ] -# Only relevant if mujoco-uni-runtime is ever built from its sdist inside -# this project: the native extension must be compiled against the mujoco -# version installed in this environment rather than in an isolated build env, -# because it records the build-time mujoco version and refuses to load on -# mismatch. The sdist fallback path (`make mujoco MJ=`) builds via -# `uv pip install --no-build-isolation` explicitly; this setting covers any -# other in-project build. The default install path resolves the prebuilt -# wheel (bound to mujoco==3.11.0) and never builds anything. -no-build-isolation-package = ["mujoco-uni-runtime"] required-environments = [ "sys_platform == 'darwin' and platform_machine == 'arm64'", "sys_platform == 'linux' and platform_machine == 'x86_64'", @@ -170,11 +160,9 @@ venvPath = "." venv = ".venv" include = ["src/unilab"] exclude = [ - "src/unilab/base/backend/", - "src/unilab/envs/", - "src/unilab/terrains/", - "src/unilab/utils/monitoring.py", - "src/unilab/visualization/", + "src/unilab/envs/", # lazy optional `mujoco` import; not installed in the pyright env + "src/unilab/utils/monitoring.py", # optional pynvml/psutil deps + "src/unilab/visualization/", # direct mujoco C bindings + optional viser deps ] reportMissingImports = "warning" reportMissingModuleSource = "none" diff --git a/pyproject.toml b/pyproject.toml index 6e0ad4f97..16a7ce0d3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,8 +41,9 @@ requires-python = ">=3.10,<3.14" dependencies = [ "numpy", # Physics implementations are provided by the independently released - # unisim-core package from the production PyPI index. - "unisim-core>=1.2.0", + # unisim-core package from the production PyPI index. >=1.2.1 is the + # first line whose mujoco adapter executes on the mjbatch-uni engine. + "unisim-core>=1.2.1", # RL algorithms and async runtimes (PPO/APPO/SAC/TD3 runners, # collectors, IPC, logging) live in the independently released uni-rl # package (distribution name ``unilab-rl``), consumed via the injected @@ -114,20 +115,18 @@ mujoco = [ # MuJoCo 3.11 alignment (unilabsim/UniLab#1515): the mujoco, mjwarp, and # newton extras share one MuJoCo 3.11 / MuJoCo-Warp 3.11 / warp-lang 1.16.0 # line and are jointly installable. The `mujoco` bound is a compatible- - # release specifier; `uv.lock` pins the exact MuJoCo version, and the - # `mujoco-uni-runtime` pin is exact because each runtime release carries - # exactly one prebuilt MuJoCo binding (its 0.5.0 wheels bind - # mujoco==3.11.0; the runtime supports >=3.5,<3.12 via sdist rebuilds). - # Switch to a non-default MuJoCo version via `make mujoco MJ=`, - # which forces a source rebuild of the runtime. + # release specifier; `uv.lock` pins the exact MuJoCo version. The MuJoCo + # physics backend (unisim.backend.mujoco.backend) executes on the mjbatch + # native batch engine, which is built and version-checked against + # mujoco==3.11.0 — switching MuJoCo versions requires an mjbatch rebuild, + # not a UniLab config change. "mujoco~=3.11.0", - "mujoco-uni-runtime==0.5.0", - # pybind11/wheel are build requirements of the mujoco-uni-runtime sdist - # fallback build (see no-build-isolation-package below). The default - # install path uses the prebuilt wheel and never invokes them, but they - # must already be in the env when `make mujoco MJ=` recompiles. - "pybind11>=2.12", - "wheel", + # The batch engine is the unilabsim mjbatch fork, published on PyPI as + # mjbatch-uni (#1552). Its build backend pins mujoco==3.11.0 at build + # time, so isolated builds are correct and no compiler preflight is + # needed; prebuilt wheels cover cp310-cp313 on linux x86_64/aarch64 and + # macOS arm64. + "mjbatch-uni~=0.1.0", ] mjwarp = [ # Keep the Warp backend on the same MuJoCo minor line as the host backend. @@ -169,10 +168,10 @@ viser = ["viser>=1.0.26", "trimesh>=3.21.7"] # the temporary superdex-uni build names are owned by unisim-core's superdex # extra (they revert to superdex-physics/-robotics once the upstream # project_superdex release lands). That extra already carries plain `mujoco` -# for MJCF conversion and the offline playback renderer; mujoco-uni-runtime is +# for MJCF conversion and the offline playback renderer; mjbatch is # deliberately NOT bound here, since only the MuJoCo physics backend # (unisim.backend.mujoco.backend) needs it and UniLab guards that backend on -# the `mujoco_uni` module. The wheels are CPython 3.12/3.13 Linux x86_64 only, +# the `mjbatch` module. The wheels are CPython 3.12/3.13 Linux x86_64 only, # so the marker keeps `uv lock` resolvable for the other # required-environments; elsewhere the extra is empty and the CLI reports a # targeted runtime diagnostic. @@ -202,16 +201,6 @@ torch = [ [tool.uv] exclude-dependencies = ["torchvision"] -# Only relevant if mujoco-uni-runtime is ever built from its sdist inside -# this project: the native extension must be compiled against the mujoco -# version installed in this environment rather than in an isolated build env, -# because it records the build-time mujoco version and refuses to load on -# mismatch. The sdist fallback path (`make mujoco MJ=`) builds via -# `uv pip install --no-build-isolation` explicitly; this setting covers any -# other in-project build. The default install path resolves the prebuilt -# wheel (bound to mujoco==3.11.0) and never builds anything, so no compiler -# is needed for `make setup`. -no-build-isolation-package = ["mujoco-uni-runtime"] required-environments = [ "sys_platform == 'darwin' and platform_machine == 'arm64'", "sys_platform == 'linux' and platform_machine == 'x86_64'", @@ -265,8 +254,7 @@ venvPath = "." venv = ".venv" include = ["src/unilab"] exclude = [ - "src/unilab/envs/", # mujoco-uni-runtime internal API, stubs mismatch - "src/unilab/terrains/", # mujoco-uni-runtime MjSpec API, stubs mismatch + "src/unilab/envs/", # lazy optional `mujoco` import; not installed in the pyright env "src/unilab/utils/monitoring.py", # optional pynvml/psutil deps "src/unilab/visualization/", # direct mujoco C bindings + optional viser deps ] diff --git a/scripts/benchmark/env/benchmark_mujoco_pool_thread_scaling.py b/scripts/benchmark/env/benchmark_mujoco_pool_thread_scaling.py index 65ff96343..f1983a512 100644 --- a/scripts/benchmark/env/benchmark_mujoco_pool_thread_scaling.py +++ b/scripts/benchmark/env/benchmark_mujoco_pool_thread_scaling.py @@ -1,12 +1,12 @@ -"""MuJoCo BatchEnvPool thread-count scaling probe (issue #1328). +"""MuJoCo mjbatch thread-count scaling probe (issue #1328). -Steps a raw ``BatchEnvPool`` (no env semantics, no learner) on the G1 flat -scene with several ``nthread`` / ``cpu_ids`` configurations and reports, per -configuration, wall time per ``pool.step`` and the average number of cores the -process kept busy (process CPU time / wall time via ``os.times``). +Steps a raw ``mjbatch.Batch`` (no env semantics, no learner) on the G1 flat +scene with several ``num_threads`` / ``cpu_ids`` configurations and reports, +per configuration, wall time per ``batch.step`` and the average number of +cores the process kept busy (process CPU time / wall time via ``os.times``). Used to separate two effects of the default -``nthread = min(num_envs, 2 * cpu_count)`` pool sizing: +``num_threads = min(num_envs, 2 * cpu_count)`` pool sizing: - thread count vs. pinning (``cpu_ids``): on the reference 16C/32T host the 32-thread unpinned and pinned rows match, so the 2x-oversubscription loss @@ -19,7 +19,7 @@ # subset + tuning: uv run scripts/benchmark/env/benchmark_mujoco_pool_thread_scaling.py \ - --num-envs 4096 --nstep 3 --chunk-size 6 \ + --num-envs 4096 --nstep 3 \ --configs 64:unpinned,32:unpinned,32:pinned,16:pinned """ @@ -43,8 +43,8 @@ def _cpu_time() -> float: return t.user + t.system -def build_state(model, nenvs: int) -> np.ndarray: - """Tile the ``stand`` keyframe (or a plain forward) into a full-batch state.""" +def build_state(model, nenvs: int) -> tuple[np.ndarray, np.ndarray]: + """Tile the ``stand`` keyframe (or a plain forward) qpos/qvel for the batch.""" import mujoco data = mujoco.MjData(model) @@ -52,43 +52,42 @@ def build_state(model, nenvs: int) -> np.ndarray: if key_id >= 0: mujoco.mj_resetDataKeyframe(model, data, key_id) mujoco.mj_forward(model, data) - spec = int(mujoco.mjtState.mjSTATE_FULLPHYSICS) - row = np.empty(mujoco.mj_stateSize(model, spec), dtype=np.float64) - mujoco.mj_getState(model, data, row, spec) - return np.tile(row, (nenvs, 1)).copy() + qpos = np.tile(np.asarray(data.qpos, dtype=np.float64), (nenvs, 1)).copy() + qvel = np.tile(np.asarray(data.qvel, dtype=np.float64), (nenvs, 1)).copy() + return qpos, qvel def bench_config( model, - state0: np.ndarray, + qpos0: np.ndarray, + qvel0: np.ndarray, *, nthread: int, pinned: bool, nstep: int, - chunk_size: int | None, warmup: int, iters: int, ) -> tuple[float, float]: - """Return (wall ms/step, busy cores) for one pool configuration.""" - from mujoco_uni.batch_env import BatchEnvPool + """Return (wall ms/step, busy cores) for one batch configuration.""" + import mjbatch cpu_ids = list(range(nthread)) if pinned else None - pool = BatchEnvPool(model, nbatch=state0.shape[0], nthread=nthread, cpu_ids=cpu_ids) - nenvs = state0.shape[0] - ctrl = np.zeros((nenvs, nstep, model.nu), dtype=np.float64) - st = state0.copy() - try: - for _ in range(warmup): - st = pool.step(st, nstep=nstep, control=ctrl, chunk_size=chunk_size) - t0 = time.perf_counter() - c0 = _cpu_time() - for _ in range(iters): - st = pool.step(st, nstep=nstep, control=ctrl, chunk_size=chunk_size) - wall_ms = (time.perf_counter() - t0) / iters * 1000.0 - cores = (_cpu_time() - c0) / iters * 1000.0 / wall_ms - return wall_ms, cores - finally: - pool.close() + batch = mjbatch.Batch(model, qpos0.shape[0], num_threads=nthread, cpu_ids=cpu_ids) + ctrl_view = batch.bind("ctrl") + ctrl_view[:] = 0.0 + qpos_view = batch.bind("qpos") + qvel_view = batch.bind("qvel") + qpos_view[:] = qpos0 + qvel_view[:] = qvel0 + for _ in range(warmup): + batch.step(nstep=nstep) + t0 = time.perf_counter() + c0 = _cpu_time() + for _ in range(iters): + batch.step(nstep=nstep) + wall_ms = (time.perf_counter() - t0) / iters * 1000.0 + cores = (_cpu_time() - c0) / iters * 1000.0 / wall_ms + return wall_ms, cores def _parse_configs(spec: str) -> list[tuple[int, bool]]: @@ -105,8 +104,7 @@ def main(argv: Sequence[str] | None = None) -> int: parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) parser.add_argument("--model", default=DEFAULT_MODEL, help="MuJoCo XML scene path") parser.add_argument("--num-envs", type=int, default=4096) - parser.add_argument("--nstep", type=int, default=3, help="sim substeps per pool.step") - parser.add_argument("--chunk-size", type=int, default=6) + parser.add_argument("--nstep", type=int, default=3, help="sim substeps per batch.step") parser.add_argument("--warmup", type=int, default=5) parser.add_argument("--iters", type=int, default=30) parser.add_argument( @@ -119,21 +117,21 @@ def main(argv: Sequence[str] | None = None) -> int: import mujoco model = mujoco.MjModel.from_xml_path(args.model) - state0 = build_state(model, args.num_envs) + qpos0, qvel0 = build_state(model, args.num_envs) print( f"model={os.path.basename(args.model)} nu={model.nu} nv={model.nv} " - f"nstate={state0.shape[1]} num_envs={args.num_envs} host_cpus={os.cpu_count()} " - f"nstep={args.nstep} chunk_size={args.chunk_size}" + f"nq={model.nq} num_envs={args.num_envs} host_cpus={os.cpu_count()} " + f"nstep={args.nstep}" ) print(f"{'config':>18s} | {'ms/step':>8s} | {'cores':>6s}") for nthread, pinned in _parse_configs(args.configs): wall_ms, cores = bench_config( model, - state0, + qpos0, + qvel0, nthread=nthread, pinned=pinned, nstep=args.nstep, - chunk_size=args.chunk_size, warmup=args.warmup, iters=args.iters, ) diff --git a/scripts/benchmark/env/benchmark_np_env_shard_throughput.py b/scripts/benchmark/env/benchmark_np_env_shard_throughput.py index 4cb08bc69..314d76221 100644 --- a/scripts/benchmark/env/benchmark_np_env_shard_throughput.py +++ b/scripts/benchmark/env/benchmark_np_env_shard_throughput.py @@ -195,10 +195,6 @@ def _shard_worker( cfg = task_config.build_cfg(BACKEND) task_config.finalize_cfg(cfg, BACKEND) cfg.validate() - # Skip the adaptive chunk_size sweep: it re-probes on every env - # materialization and dominates shard setup time. Use the native - # default chunk_size instead (issue #960). - cfg.adaptive_chunk_size = False env = task_config.env_cls_factory()(cfg, num_envs=shard_rows, backend_type=BACKEND) try: env.init_state() diff --git a/scripts/benchmark/physics/benchmark_forward_reset_methods.py b/scripts/benchmark/physics/benchmark_forward_reset_methods.py deleted file mode 100644 index 76644ab30..000000000 --- a/scripts/benchmark/physics/benchmark_forward_reset_methods.py +++ /dev/null @@ -1,324 +0,0 @@ -#!/usr/bin/env python3 -"""Benchmark C++ batch forward vs Python forward methods.""" - -import argparse -import sys -from dataclasses import asdict, dataclass -from multiprocessing import cpu_count -from pathlib import Path - -import mujoco -import numpy as np -from mujoco_uni.batch_env import BatchEnvPool - -ROOT_DIR = Path(__file__).resolve().parents[3] -if str(ROOT_DIR) not in sys.path: - sys.path.append(str(ROOT_DIR)) - -import matplotlib - -from scripts.benchmark.core.device_info import get_device_info_dict, get_device_info_line -from scripts.benchmark.core.task_names import ( - canonical_locomotion_task_ids, - locomotion_task_model_file, - normalize_locomotion_task_id, -) - -matplotlib.use("Agg") -import matplotlib.pyplot as plt - -CONSISTENCY_TASKS = canonical_locomotion_task_ids() -DEFAULT_RESET_TASK = "go2_joystick_flat" - - -@dataclass -class ConsistencyRecord: - task: str - batch_size: int - max_abs_diff: float - mean_abs_diff: float - allclose: bool - - -@dataclass -class SpeedRecord: - task: str - method: str - env_num: int - elapsed_sec: float - us_per_env: float - - -def load_model_for_task(task: str) -> mujoco.MjModel: - task_key = normalize_locomotion_task_id(task) - return mujoco.MjModel.from_xml_path(locomotion_task_model_file(task_key)) - - -def _set_state_for_forward(model, data, state): - mujoco.mj_setState(model, data, state, mujoco.mjtState.mjSTATE_FULLPHYSICS) - data.ctrl[:] = 0.0 - data.qfrc_applied[:] = 0.0 - data.xfrc_applied[:] = 0.0 - data.qacc_warmstart[:] = 0.0 - - -def python_forward_for_loop(model, states): - out = np.empty((states.shape[0], model.nsensordata), dtype=np.float64) - data = mujoco.MjData(model) - for i in range(states.shape[0]): - _set_state_for_forward(model, data, states[i]) - mujoco.mj_forward(model, data) - out[i] = data.sensordata - return out - - -def python_forward_chunk_loop(model, states, chunk_size): - out = np.empty((states.shape[0], model.nsensordata), dtype=np.float64) - data = mujoco.MjData(model) - for start in range(0, states.shape[0], chunk_size): - end = min(start + chunk_size, states.shape[0]) - for i in range(start, end): - _set_state_for_forward(model, data, states[i]) - mujoco.mj_forward(model, data) - out[i] = data.sensordata - return out - - -def batch_env_forward(pool: BatchEnvPool, states: np.ndarray) -> np.ndarray: - sensordata = pool.forward(states) - return np.asarray(sensordata, dtype=np.float64) - - -def collect_random_reset_states(task, batch_size, random_rounds, seed): - model = load_model_for_task(task) - rng = np.random.default_rng(seed) - data = mujoco.MjData(model) - nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) - - if model.nkey > 0: - mujoco.mj_resetDataKeyframe(model, data, 0) - else: - mujoco.mj_resetData(model, data) - - ctrl_low = np.full((model.nu,), -1.0) - ctrl_high = np.full((model.nu,), 1.0) - if model.nu > 0 and hasattr(model, "actuator_ctrllimited"): - limited = np.asarray(model.actuator_ctrllimited, dtype=bool) - if np.any(limited): - ranges = np.asarray(model.actuator_ctrlrange) - ctrl_low[limited] = ranges[limited, 0] - ctrl_high[limited] = ranges[limited, 1] - - states = np.empty((batch_size, nstate)) - round_steps = max(2, int(random_rounds)) - for i in range(batch_size): - if i % round_steps == 0: - if model.nkey > 0: - mujoco.mj_resetDataKeyframe(model, data, 0) - else: - mujoco.mj_resetData(model, data) - - nsteps = int(rng.integers(1, round_steps + 1)) - for _ in range(nsteps): - if model.nu > 0: - data.ctrl[:] = rng.uniform(ctrl_low, ctrl_high) - mujoco.mj_step(model, data) - - state_buf = np.empty((nstate,)) - mujoco.mj_getState(model, data, state_buf, mujoco.mjtState.mjSTATE_FULLPHYSICS) - states[i] = state_buf - - return model, states - - -def run_consistency(tasks, batch_size, random_rounds, seed, atol, rtol, chunk_size): - _ = chunk_size # BatchEnvPool.forward does not expose chunk-size tuning. - records = [] - for task in tasks: - task_key = normalize_locomotion_task_id(task) - model, states = collect_random_reset_states( - task_key, batch_size, random_rounds, seed + abs(hash(task_key)) % 10000 - ) - nthread = min(batch_size, cpu_count()) - py_sensor = python_forward_for_loop(model, states) - with BatchEnvPool(model, nbatch=states.shape[0], nthread=nthread) as pool: - batch_sensor = batch_env_forward(pool, states) - diff = np.abs(py_sensor - batch_sensor) - max_abs = float(np.max(diff)) - mean_abs = float(np.mean(diff)) - ok = bool(np.allclose(py_sensor, batch_sensor, atol=atol, rtol=rtol)) - records.append(ConsistencyRecord(task_key, batch_size, max_abs, mean_abs, ok)) - print( - f"[Consistency] {task_key}: allclose={ok}, max_abs={max_abs:.3e}, mean_abs={mean_abs:.3e}" - ) - return records - - -def _bench_method(func, repeats): - import time - - samples = [time.perf_counter() or func() or time.perf_counter() for _ in range(repeats)] - samples = [] - for _ in range(repeats): - t0 = time.perf_counter() - func() - samples.append(time.perf_counter() - t0) - return float(np.median(samples)) - - -def run_reset_speed(task, env_nums, random_rounds, chunk_size, repeats, seed): - _ = chunk_size # BatchEnvPool.forward does not expose chunk-size tuning. - task_key = normalize_locomotion_task_id(task) - records = [] - for env_num in env_nums: - model, states = collect_random_reset_states( - task_key, env_num, random_rounds, seed + env_num - ) - nthread = min(env_num, cpu_count()) - - t_for = _bench_method(lambda: python_forward_for_loop(model, states), repeats) - t_chunk = _bench_method( - lambda: python_forward_chunk_loop(model, states, chunk_size), repeats - ) - with BatchEnvPool(model, nbatch=states.shape[0], nthread=nthread) as pool: - t_batch = _bench_method(lambda: batch_env_forward(pool, states), repeats) - - for method, elapsed in [ - ("for_loop", t_for), - ("chunk_for_loop", t_chunk), - ("batch_env_forward", t_batch), - ]: - records.append( - SpeedRecord(task_key, method, env_num, elapsed, elapsed * 1e6 / max(env_num, 1)) - ) - - print( - f"[Speed] {task_key} env={env_num}: for={t_for * 1e3:.3f}ms, chunk={t_chunk * 1e3:.3f}ms, batch_env={t_batch * 1e3:.3f}ms" - ) - return records - - -def plot_speed(records, out_png, reset_task): - out_png.parent.mkdir(parents=True, exist_ok=True) - env_nums = sorted({r.env_num for r in records}) - methods = ["for_loop", "chunk_for_loop", "batch_env_forward"] - colors = {"for_loop": "#3B82F6", "chunk_for_loop": "#10B981", "batch_env_forward": "#F59E0B"} - - fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5), sharex=True) - x = np.arange(len(env_nums)) - width = 0.24 - - for i, method in enumerate(methods): - total_ms = [ - next(r.elapsed_sec for r in records if r.method == method and r.env_num == n) * 1e3 - for n in env_nums - ] - throughput = [ - n - / max( - next(r.elapsed_sec for r in records if r.method == method and r.env_num == n), 1e-12 - ) - for n in env_nums - ] - offset = (i - 1) * width - ax1.bar(x + offset, total_ms, width, label=method, color=colors[method], alpha=0.95) - ax2.bar(x + offset, throughput, width, label=method, color=colors[method], alpha=0.95) - - for ax in (ax1, ax2): - ax.set_xticks(x) - ax.set_xticklabels([str(v) for v in env_nums]) - ax.set_xlabel("num_env") - ax.grid(axis="y", alpha=0.3) - - ax1.set_yscale("log") - ax1.set_ylabel("total time (ms)") - ax1.set_title("Total Time") - ax2.set_yscale("log") - ax2.set_ylabel("throughput (env/s)") - ax2.set_title("Throughput") - fig.suptitle( - f"Reset-forward method speed on {normalize_locomotion_task_id(reset_task)}\n{get_device_info_line()}" - ) - ax2.legend() - fig.tight_layout() - fig.savefig(out_png, dpi=180) - plt.close(fig) - - -def main(): - parser = argparse.ArgumentParser(description="Benchmark C++ batch forward vs Python forward") - parser.add_argument("--consistency-tasks", type=str, default=",".join(CONSISTENCY_TASKS)) - parser.add_argument("--reset-task", type=str, default=DEFAULT_RESET_TASK) - parser.add_argument("--consistency-batch", type=int, default=1024) - parser.add_argument("--env-nums", type=str, default="256,512,1024,2048,4096,8192,16384") - parser.add_argument("--random-rounds", type=int, default=8) - parser.add_argument( - "--chunk-size", - type=int, - default=64, - help="Legacy no-op retained for CLI compatibility; BatchEnvPool.forward does not use chunk_size.", - ) - parser.add_argument("--repeats", type=int, default=8) - parser.add_argument("--seed", type=int, default=42) - parser.add_argument("--atol", type=float, default=1e-9) - parser.add_argument("--rtol", type=float, default=1e-8) - parser.add_argument( - "--out-json", type=str, default="scripts/benchmark/outputs/reset_forward_batch/results.json" - ) - parser.add_argument( - "--out-png", - type=str, - default="scripts/benchmark/outputs/reset_forward_batch/speed_plot.png", - ) - args = parser.parse_args() - - consistency_tasks = [ - normalize_locomotion_task_id(t) for t in args.consistency_tasks.split(",") if t.strip() - ] - reset_task = normalize_locomotion_task_id(args.reset_task) - env_nums = [int(x.strip()) for x in args.env_nums.split(",") if x.strip()] - - consistency = run_consistency( - consistency_tasks, - args.consistency_batch, - args.random_rounds, - args.seed, - args.atol, - args.rtol, - args.chunk_size, - ) - speed = run_reset_speed( - reset_task, env_nums, args.random_rounds, args.chunk_size, args.repeats, args.seed - ) - - out_json = Path(args.out_json) - out_json.parent.mkdir(parents=True, exist_ok=True) - - import json - from datetime import datetime, timezone - - with out_json.open("w") as f: - json.dump( - { - "meta": { - "timestamp": datetime.now(timezone.utc).isoformat(), - "device_info": get_device_info_dict(), - "consistency_tasks": consistency_tasks, - "reset_task": reset_task, - "env_nums": env_nums, - "chunk_size": args.chunk_size, - }, - "consistency": [asdict(r) for r in consistency], - "speed": [asdict(r) for r in speed], - }, - f, - indent=2, - ) - - plot_speed(speed, Path(args.out_png), reset_task=reset_task) - print(f"Consistency: {all(r.allclose for r in consistency)}") - print(f"Saved: {out_json}, {args.out_png}") - - -if __name__ == "__main__": - main() diff --git a/scripts/benchmark/physics/benchmark_mujoco_backend_step_detail.py b/scripts/benchmark/physics/benchmark_mujoco_backend_step_detail.py index d82e61d38..be5e94e1f 100644 --- a/scripts/benchmark/physics/benchmark_mujoco_backend_step_detail.py +++ b/scripts/benchmark/physics/benchmark_mujoco_backend_step_detail.py @@ -5,10 +5,10 @@ This benchmark mirrors the hot path inside `MujocoBackend.step()` and splits it into: - 1. control broadcast (`set_ctrl`) - 2. `BatchEnvPool.step` (`pool_step`) + 1. control upload (`set_ctrl`) + 2. `mjbatch.Batch.step` (`pool_step`) 3. physics-state cast/copy (`state_copy`) - 4. `BatchEnvPool.forward` (`forward`) + 4. `mjbatch.Batch.forward` (`forward`) 5. sensor-data cast/copy (`sensor_copy`) It sweeps current locomotion owner tasks across MuJoCo only, with environment @@ -40,10 +40,10 @@ from typing import Sequence import matplotlib +import mjbatch import mujoco import numpy as np from matplotlib.patches import Rectangle -from mujoco_uni.batch_env import BatchEnvPool from unisim.backend.mujoco.xml import create_discardvisual_xml from unilab.dtype_config import get_global_dtype @@ -140,16 +140,17 @@ def _parse_csv_tasks(text: str) -> list[str]: return values -def _keyframe0_state_and_ctrl(model: mujoco.MjModel) -> tuple[np.ndarray, np.ndarray]: +def _keyframe0_state_and_ctrl( + model: mujoco.MjModel, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: data = mujoco.MjData(model) if model.nkey > 0: mujoco.mj_resetDataKeyframe(model, data, 0) else: - mujoco.mj_resetData(model, data) + mujoco.mj_resetData(model) - nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) - state0 = np.empty((nstate,), dtype=np.float64) - mujoco.mj_getState(model, data, state0, mujoco.mjtState.mjSTATE_FULLPHYSICS) + qpos0 = np.array(data.qpos, dtype=np.float64) + qvel0 = np.array(data.qvel, dtype=np.float64) if model.nu == 0: ctrl0 = np.empty((0,), dtype=np.float64) @@ -157,7 +158,7 @@ def _keyframe0_state_and_ctrl(model: mujoco.MjModel) -> tuple[np.ndarray, np.nda ctrl0 = np.asarray(model.key_ctrl[0], dtype=np.float64).copy() else: ctrl0 = np.zeros((model.nu,), dtype=np.float64) - return state0, ctrl0 + return qpos0, qvel0, ctrl0 def _load_discardvisual_model(model_file: str) -> mujoco.MjModel: @@ -211,7 +212,7 @@ def _benchmark_one( model = _load_discardvisual_model(locomotion_task_model_file(task)) np_dtype = get_global_dtype() - state0, _ = _keyframe0_state_and_ctrl(model) + qpos0, qvel0, _ = _keyframe0_state_and_ctrl(model) ctrl_low, ctrl_high = _control_limits(model) nthread = min(env_num, cpu_count() * 2) rng = np.random.default_rng(seed) @@ -225,44 +226,50 @@ def _benchmark_one( forward_samples: list[float] = [] sensor_copy_samples: list[float] = [] - with BatchEnvPool(model, nbatch=env_num, nthread=nthread) as pool: - physics_state = np.broadcast_to(state0.astype(np_dtype), (env_num, state0.shape[0])).copy() - sensor_data = np.zeros((env_num, model.nsensordata), dtype=np_dtype) - sensor_init = pool.forward(physics_state) - sensor_data[:] = sensor_init.astype(np_dtype) - - for iteration_idx, ctrl in enumerate(controls): - t0 = time.perf_counter() - control_traj = np.broadcast_to(ctrl[:, None, :], (env_num, nstep, ctrl.shape[-1])) - set_ctrl_ms = (time.perf_counter() - t0) * 1000.0 - - t0 = time.perf_counter() - state_np = pool.step( - physics_state, - nstep=nstep, - control=control_traj, - control_spec=int(mujoco.mjtState.mjSTATE_CTRL), - ) - pool_step_ms = (time.perf_counter() - t0) * 1000.0 - - t0 = time.perf_counter() - physics_state[:] = state_np.astype(np_dtype) - state_copy_ms = (time.perf_counter() - t0) * 1000.0 - - t0 = time.perf_counter() - sensor_np = pool.forward(physics_state) - forward_ms = (time.perf_counter() - t0) * 1000.0 - - t0 = time.perf_counter() - sensor_data[:] = sensor_np.astype(np_dtype) - sensor_copy_ms = (time.perf_counter() - t0) * 1000.0 - - if iteration_idx >= warmup: - set_ctrl_samples.append(set_ctrl_ms) - pool_step_samples.append(pool_step_ms) - state_copy_samples.append(state_copy_ms) - forward_samples.append(forward_ms) - sensor_copy_samples.append(sensor_copy_ms) + batch = mjbatch.Batch(model, env_num, num_threads=nthread) + time_view = batch.bind("time", np_dtype) + qpos_view = batch.bind("qpos", np_dtype) + qvel_view = batch.bind("qvel", np_dtype) + ctrl_view = batch.bind("ctrl", np_dtype) + sensor_view = batch.bind("sensordata", np_dtype) + + qpos_view[:] = qpos0 + qvel_view[:] = qvel0 + state_dim = 1 + model.nq + model.nv + physics_state = np.empty((env_num, state_dim), dtype=np_dtype) + sensor_data = np.empty((env_num, model.nsensordata), dtype=np_dtype) + batch.forward() + sensor_data[:] = sensor_view + + for iteration_idx, ctrl in enumerate(controls): + t0 = time.perf_counter() + ctrl_view[:] = ctrl + set_ctrl_ms = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() + batch.step(nstep=nstep) + pool_step_ms = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() + physics_state[:, 0] = time_view + physics_state[:, 1 : 1 + model.nq] = qpos_view + physics_state[:, 1 + model.nq :] = qvel_view + state_copy_ms = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() + batch.forward() + forward_ms = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() + sensor_data[:] = sensor_view + sensor_copy_ms = (time.perf_counter() - t0) * 1000.0 + + if iteration_idx >= warmup: + set_ctrl_samples.append(set_ctrl_ms) + pool_step_samples.append(pool_step_ms) + state_copy_samples.append(state_copy_ms) + forward_samples.append(forward_ms) + sensor_copy_samples.append(sensor_copy_ms) set_ctrl_ms = _median_ms(set_ctrl_samples) pool_step_ms = _median_ms(pool_step_samples) @@ -281,7 +288,7 @@ def _benchmark_one( warmup=warmup, iters=iters, control_dim=model.nu, - state_dim=state0.shape[0], + state_dim=state_dim, sensor_dim=model.nsensordata, set_ctrl_ms=set_ctrl_ms, pool_step_ms=pool_step_ms, diff --git a/scripts/benchmark/physics/benchmark_mujoco_single_dispatch_callback.py b/scripts/benchmark/physics/benchmark_mujoco_single_dispatch_callback.py deleted file mode 100644 index 20e15c9a3..000000000 --- a/scripts/benchmark/physics/benchmark_mujoco_single_dispatch_callback.py +++ /dev/null @@ -1,212 +0,0 @@ -#!/usr/bin/env python3 -"""One-off microbenchmark for issue #1262. - -Question: can one ``BatchEnvPool.step(nstep=N)`` dispatch reproduce the MBA -per-substep control semantics that today require N ``pool.step(nstep=1)`` -dispatches (``MuJoCoBackend._step_with_pre_step_control``)? - -Measured on the g1_walk_flat MuJoCo contract (scene_flat.xml + injected body -tracking sensors, sim_dt=1/150, sim_substeps=3): - -- ``multi``: current MBA path — N dispatches of nstep=1, with a Python-side - per-substep control recompute (JointPositionAction-style: - ``target = processed - encoder_bias``, constant within one - control step) between dispatches. -- ``traj``: single dispatch, nstep=N, control baked as an (nbatch, N, nu) - trajectory with identical rows. -- ``const``: single dispatch, nstep=N, control passed as one (nbatch, nu) - array (native ``control_is_constant`` fast path). - -All paths start from the same contact-rich state (standing keyframe settled -with foot contact) and use the same pool, model, sensors, and chunk size. -The numerical check compares final full-physics state and sensordata between -``multi`` and the single-dispatch paths. -""" - -from __future__ import annotations - -import argparse -import os -import time - -import mujoco -import numpy as np -from mujoco_uni import BatchEnvPool -from unisim.backend.mujoco.xml import ( - create_discardvisual_xml, - inject_mujoco_tracking_sensors, -) - -SCENE_XML = "src/unilab/assets/robots/g1/scene_flat.xml" -BASE_BODY = "pelvis" -KEYFRAME = "stand" -SIM_DT = 1.0 / 150.0 -SUBSTEPS = 3 # ctrl_dt 0.02 / sim_dt 0.006667, matches g1_walk_flat -CTRL_SPEC = int(mujoco.mjtState.mjSTATE_CTRL) -FULLPHYSICS = mujoco.mjtState.mjSTATE_FULLPHYSICS - - -def build_model() -> mujoco.MjModel: - path = create_discardvisual_xml(SCENE_XML) - path, _, _ = inject_mujoco_tracking_sensors(path, baselink_name=BASE_BODY) - model = mujoco.MjModel.from_xml_path(path) - model.opt.timestep = SIM_DT - return model - - -def keyframe_state(model: mujoco.MjModel) -> np.ndarray: - data = mujoco.MjData(model) - kid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_KEY, KEYFRAME) - if kid < 0: - raise ValueError(f"keyframe '{KEYFRAME}' not found in {SCENE_XML}") - mujoco.mj_resetDataKeyframe(model, data, kid) - mujoco.mj_forward(model, data) - state = np.zeros(mujoco.mj_stateSize(model, FULLPHYSICS), dtype=np.float64) - mujoco.mj_getState(model, data, state, FULLPHYSICS) - return state - - -def make_ctrl(model: mujoco.MjModel, state0: np.ndarray, nbatch: int) -> np.ndarray: - """Stand-pose position targets plus per-env jitter (deterministic).""" - nq = model.nq - qpos = state0[1 : 1 + nq] # FULLPHYSICS: time, qpos, qvel, act, ... - stand = qpos[-model.nu :] # free root (7) precedes actuated joints - rng = np.random.default_rng(0) - ctrl = stand[None, :] + rng.uniform(-0.05, 0.05, size=(nbatch, model.nu)) - return np.ascontiguousarray(ctrl, dtype=np.float64) - - -def settle(pool: BatchEnvPool, state0: np.ndarray, ctrl: np.ndarray, chunk_size: int) -> np.ndarray: - """Roll out a few control steps so feet are in steady contact.""" - state = state0 - for _ in range(40): - state = pool.step( - state, - nstep=SUBSTEPS, - control=ctrl, - control_spec=CTRL_SPEC, - chunk_size=chunk_size, - ) - return state - - -def run_multi(pool, state, ctrl, bias, chunk_size): - """Current MBA path: SUBSTEPS dispatches of nstep=1 with callback between.""" - pool_ms = 0.0 - callback_ms = 0.0 - for _ in range(SUBSTEPS): - t0 = time.perf_counter() - native_ctrl = np.subtract(ctrl, bias) # JointPositionAction-style recompute - callback_ms += (time.perf_counter() - t0) * 1e3 - t0 = time.perf_counter() - state, _sensor = pool.step( - state, - nstep=1, - control=native_ctrl[:, None, :], - control_spec=CTRL_SPEC, - chunk_size=chunk_size, - return_sensor=True, - ) - pool_ms += (time.perf_counter() - t0) * 1e3 - return state, pool_ms, callback_ms - - -def run_traj(pool, state, ctrl, bias, chunk_size): - """Single dispatch, baked (nbatch, N, nu) trajectory with identical rows.""" - t0 = time.perf_counter() - native_ctrl = np.subtract(ctrl, bias) - callback_ms = (time.perf_counter() - t0) * 1e3 - t0 = time.perf_counter() - traj = np.broadcast_to(native_ctrl[:, None, :], (ctrl.shape[0], SUBSTEPS, ctrl.shape[1])) - state, _sensor = pool.step( - state, - nstep=SUBSTEPS, - control=traj, - control_spec=CTRL_SPEC, - chunk_size=chunk_size, - return_sensor=True, - ) - return state, (time.perf_counter() - t0) * 1e3, callback_ms - - -def run_const(pool, state, ctrl, bias, chunk_size): - """Single dispatch, baked (nbatch, nu) constant control (native fast path).""" - t0 = time.perf_counter() - native_ctrl = np.subtract(ctrl, bias) - callback_ms = (time.perf_counter() - t0) * 1e3 - t0 = time.perf_counter() - state, _sensor = pool.step( - state, - nstep=SUBSTEPS, - control=native_ctrl, - control_spec=CTRL_SPEC, - chunk_size=chunk_size, - return_sensor=True, - ) - return state, (time.perf_counter() - t0) * 1e3, callback_ms - - -def main() -> None: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--num-envs", type=int, default=8192) - parser.add_argument("--chunk-sizes", type=int, nargs="+", default=[13, 41]) - parser.add_argument("--repeats", type=int, default=15) - parser.add_argument("--warmup", type=int, default=3) - args = parser.parse_args() - - model = build_model() - nthread = min(args.num_envs, (os.cpu_count() or 1) * 2) - print( - f"model: nu={model.nu} nq={model.nq} nv={model.nv} nsensor={model.nsensor} " - f"nsensordata={model.nsensordata}; nbatch={args.num_envs} nthread={nthread}" - ) - pool = BatchEnvPool(model, nbatch=args.num_envs, nthread=nthread) - try: - state0 = np.broadcast_to(keyframe_state(model)[None, :], (args.num_envs, pool.nstate)) - state0 = np.ascontiguousarray(state0, dtype=np.float64) - ctrl = make_ctrl(model, state0[0], args.num_envs) - bias = np.zeros_like(ctrl) # encoder_bias placeholder (zeros on g1_walk_flat) - contact_state = settle(pool, state0, ctrl, chunk_size=args.chunk_sizes[0]) - - paths = {"multi": run_multi, "traj": run_traj, "const": run_const} - - # Numerical comparison from the same contact-rich state. - finals = {} - for name, fn in paths.items(): - out = fn(pool, contact_state, ctrl, bias, args.chunk_sizes[0]) - finals[name] = out[0] - for name in ("traj", "const"): - diff = np.abs(finals[name] - finals["multi"]) - bitwise = float(np.mean(finals[name] == finals["multi"])) - print( - f"numerics {name} vs multi: max_abs_diff={diff.max():.3e} " - f"bitwise_equal={bitwise:.4f}" - ) - print( - f"numerics traj vs const: bitwise={bool(np.array_equal(finals['traj'], finals['const']))}" - ) - - # Timing. - for chunk_size in args.chunk_sizes: - print( - f"--- chunk_size={chunk_size} (ms per control step, median of {args.repeats}) ---" - ) - for name, fn in paths.items(): - for _ in range(args.warmup): - fn(pool, contact_state, ctrl, bias, chunk_size) - pool_ts, cb_ts = [], [] - for _ in range(args.repeats): - _s, pool_ms, cb_ms = fn(pool, contact_state, ctrl, bias, chunk_size) - pool_ts.append(pool_ms) - cb_ts.append(cb_ms) - print( - f" {name:>5}: pool={np.median(pool_ts):7.2f} " - f"callback={np.median(cb_ts):5.2f} " - f"total={np.median(pool_ts) + np.median(cb_ts):7.2f}" - ) - finally: - pool.close() - - -if __name__ == "__main__": - main() diff --git a/scripts/tools/capture_mujoco_drift_baseline.py b/scripts/tools/capture_mujoco_drift_baseline.py new file mode 100644 index 000000000..c01947619 --- /dev/null +++ b/scripts/tools/capture_mujoco_drift_baseline.py @@ -0,0 +1,303 @@ +#!/usr/bin/env python3 +"""Capture a deterministic drift-characterization baseline for the mujoco backend. + +Roadmap issue #1552 replaces the mujoco backend's native executor +(mujoco-uni-runtime, import name ``mujoco_uni``) with mjbatch. Numerical drift +between the two executors is expected and accepted; the quality gate is a +fixed model/seed/action-sequence trajectory diff against this baseline +(issue #1554), not a bit-exact gate. + +This script builds each task's training environment through the exact Hydra +owner path used by the trainer (``conf/ppo`` task config -> ``BackendAdapter`` +-> ``registry.make``), seeds every consumed RNG, pins the executor thread +count via ``cpu_ids``, drives the env with a fixed pseudo-random +action sequence (seeded ``numpy`` Generator, stored in the artifact; no neural +network), and records per-step observations (all obs groups), rewards, done +flags, and backend qpos/qvel read through the public ``SimBackend.get_state`` +interface. + +Per task it writes ``/.npz`` plus ``/.metadata.json`` +(commit hash, package versions, task/seed/step configuration, per-array +SHA-256 digests). The default ``--output`` is the committed BEFORE baseline +location; after the dependency switch, re-run with the identical configuration +and a different ``--output`` (e.g. ``.../drift_baseline/after``) to produce the +AFTER artifact, then diff the two with any npz-aware comparison. + +Note: the .npz zip container embeds file timestamps, so two runs are compared +on array contents (the metadata records per-array SHA-256 digests for exactly +this), not on container bytes. + +Usage: + # BEFORE capture (writes the committed baseline): + uv run scripts/tools/capture_mujoco_drift_baseline.py + + # AFTER capture (identical configuration, different output dir): + uv run scripts/tools/capture_mujoco_drift_baseline.py \ + --output scripts/tools/drift_baseline/after + + # Single task / smoke run: + uv run scripts/tools/capture_mujoco_drift_baseline.py \ + --tasks go2w_joystick_flat/mujoco --steps 50 --output /tmp/drift_smoke + + # Determinism check: run twice into different dirs, compare array contents. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import subprocess +import sys +from datetime import datetime, timezone +from importlib.metadata import PackageNotFoundError +from importlib.metadata import version as package_version +from pathlib import Path +from typing import Any + +import numpy as np + +ROOT_DIR = Path(__file__).resolve().parents[2] +SRC_DIR = ROOT_DIR / "src" +for path in (str(SRC_DIR), str(ROOT_DIR)): + if path not in sys.path: + sys.path.insert(0, path) + +from hydra import compose, initialize_config_dir +from hydra.core.global_hydra import GlobalHydra + +from unilab.base import registry +from unilab.base.config_adapter import BackendAdapter +from unilab.training import ensure_registries + +CONF_DIR = ROOT_DIR / "src" / "unilab" / "conf" +DEFAULT_OUTPUT_DIR = ROOT_DIR / "scripts" / "tools" / "drift_baseline" / "before" + +# Default tasks: Go2WJoystickFlat exercises the per-substep state-feedback +# control path (Go2WMixedAction via SimBackend.set_pre_step_control), the +# highest-risk surface for the executor swap; Go2JoystickFlat covers the plain +# position-action path on the same robot family. +DEFAULT_TASKS = ("go2w_joystick_flat/mujoco", "go2_joystick_flat/mujoco") + +# Executor determinism contract: pin the pool worker count via cpu_ids so +# every run steps an identical partition. Recorded in metadata for the AFTER +# comparison. (The mujoco_uni-era chunk_size/adaptive_chunk_size knobs no +# longer exist in EnvCfg; mjbatch schedules per-sim work without a chunk +# knob, so there is nothing else to pin.) +DEFAULT_CPU_IDS = (0, 1, 2, 3) + +METADATA_PACKAGES = ("mujoco", "mjbatch", "unisim-core", "numpy") + + +def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--tasks", + nargs="+", + default=list(DEFAULT_TASKS), + metavar="TASK/BACKEND", + help="Hydra task config paths under conf/ppo/task (default: %(default)s).", + ) + parser.add_argument( + "--output", + type=Path, + default=DEFAULT_OUTPUT_DIR, + help="Artifact directory for .npz + .metadata.json (default: %(default)s).", + ) + parser.add_argument("--num-envs", type=int, default=8) + parser.add_argument("--steps", type=int, default=300) + parser.add_argument( + "--seed", + type=int, + default=42, + help="Env seed: ManagerBasedRlEnvCfg.seed and env.reset(seed=...) (default: 42).", + ) + parser.add_argument( + "--action-seed", + type=int, + default=1234, + help="Seed of the recorded numpy Generator action sequence (default: 1234).", + ) + parser.add_argument( + "--cpu-ids", + type=int, + nargs="+", + default=list(DEFAULT_CPU_IDS), + help="Explicit CPU ids; also fixes the executor pool worker count (default: %(default)s).", + ) + return parser.parse_args(argv) + + +def _git_commit() -> str: + result = subprocess.run( + ["git", "rev-parse", "HEAD"], + cwd=ROOT_DIR, + check=True, + capture_output=True, + text=True, + ) + return result.stdout.strip() + + +def _package_versions() -> dict[str, str]: + versions: dict[str, str] = {} + for name in METADATA_PACKAGES: + try: + versions[name] = package_version(name) + except PackageNotFoundError: + versions[name] = "" + return versions + + +def _compose_task_cfg(task_path: str): + GlobalHydra.instance().clear() + with initialize_config_dir(config_dir=str(CONF_DIR / "ppo"), version_base="1.3"): + return compose("config", overrides=[f"task={task_path}"]) + + +def _build_env(task_path: str, args: argparse.Namespace): + hydra_cfg = _compose_task_cfg(task_path) + task_name = str(hydra_cfg.training.task_name) + env_cfg_override = BackendAdapter(hydra_cfg, root_dir=ROOT_DIR).build_task_env_cfg_override() + env_cfg_override.update( + { + "seed": args.seed, + "cpu_ids": list(args.cpu_ids), + } + ) + env = registry.make( + task_name, + sim_backend=str(hydra_cfg.training.sim_backend), + env_cfg_override=env_cfg_override, + num_envs=args.num_envs, + ) + return task_name, env + + +def _sha256(array: np.ndarray) -> str: + return hashlib.sha256(np.ascontiguousarray(array).tobytes()).hexdigest() + + +def make_action_sequence(steps: int, num_envs: int, action_dim: int, seed: int) -> np.ndarray: + """Fixed recorded action sequence: seeded Generator, no policy network. + + Base: iid uniform(-1, 1) for every env. Every 4th env (indices 3, 7, ...) + instead holds a seeded saturated +/-1 action vector for the first half of + the run and its negation for the second half; the saturated pose plus the + mid-run reversal topples those envs, so the baseline also exercises + termination, autoreset, and the backend set_state path mid-trajectory. + """ + rng = np.random.default_rng(seed) + actions = rng.uniform( + low=-1.0, + high=1.0, + size=(steps, num_envs, action_dim), + ).astype(np.float32) + held_envs = [i for i in range(num_envs) if i % 4 == 3] + held = rng.choice([-1.0, 1.0], size=(len(held_envs), action_dim)).astype(np.float32) + half = steps // 2 + for row, env_index in enumerate(held_envs): + actions[:half, env_index] = held[row] + actions[half:, env_index] = -held[row] + return actions + + +def capture_task(task_path: str, args: argparse.Namespace) -> dict[str, Any]: + task_name, env = _build_env(task_path, args) + try: + action_dim = int(env.action_space.shape[0]) + actions = make_action_sequence(args.steps, args.num_envs, action_dim, args.action_seed) + + # Belt and braces for any global-numpy consumer on the reset/step path; + # the authoritative env RNG is seeded explicitly below and via cfg.seed. + np.random.seed(args.seed) + reset_obs, _ = env.reset(seed=args.seed) + + backend_state = env._backend.get_state + obs_groups = sorted(env.obs_groups_spec) + init_backend = backend_state(("qpos", "qvel")) + + traces: dict[str, list[np.ndarray]] = { + "reward": [], + "terminated": [], + "truncated": [], + "qpos": [], + "qvel": [], + } + traces.update({f"obs/{name}": [] for name in obs_groups}) + + for step in range(args.steps): + state = env.step(actions[step]) + backend = backend_state(("qpos", "qvel")) + for name in obs_groups: + traces[f"obs/{name}"].append(np.asarray(state.obs[name]).copy()) + traces["reward"].append(np.asarray(state.reward).copy()) + traces["terminated"].append(np.asarray(state.terminated).copy()) + traces["truncated"].append(np.asarray(state.truncated).copy()) + traces["qpos"].append(np.asarray(backend["qpos"]).copy()) + traces["qvel"].append(np.asarray(backend["qvel"]).copy()) + if (step + 1) % 100 == 0: + print(f"[capture] {task_name}: {step + 1}/{args.steps} steps", flush=True) + + arrays: dict[str, np.ndarray] = {name: np.stack(values) for name, values in traces.items()} + arrays["actions"] = actions + for name in obs_groups: + arrays[f"obs_init/{name}"] = np.asarray(reset_obs[name]).copy() + arrays["qpos_init"] = np.asarray(init_backend["qpos"]).copy() + arrays["qvel_init"] = np.asarray(init_backend["qvel"]).copy() + + metadata: dict[str, Any] = { + "task_name": task_name, + "hydra_task_path": task_path, + "sim_backend": env._backend.backend_type, + "num_envs": args.num_envs, + "num_steps": args.steps, + "seed": args.seed, + "action_seed": args.action_seed, + "action_generator": ( + "make_action_sequence: default_rng(action_seed) iid uniform(-1, 1); " + "envs i%4==3 hold a seeded +/-1 vector for the first half, negated after" + ), + "cpu_ids": list(args.cpu_ids), + "sim_dt": float(env.cfg.sim_dt), + "ctrl_dt": float(env.cfg.ctrl_dt), + "obs_groups_spec": {k: int(v) for k, v in env.obs_groups_spec.items()}, + "action_dim": action_dim, + "git_commit": _git_commit(), + "package_versions": _package_versions(), + "timestamp_utc": datetime.now(timezone.utc).isoformat(), + "array_sha256": {name: _sha256(value) for name, value in arrays.items()}, + } + return {"task_name": task_name, "arrays": arrays, "metadata": metadata} + finally: + env.close() + + +def main(argv: list[str] | None = None) -> int: + args = _parse_args(argv) + if args.num_envs < 1: + raise SystemExit(f"--num-envs must be >= 1, got {args.num_envs}") + if args.steps < 1: + raise SystemExit(f"--steps must be >= 1, got {args.steps}") + + ensure_registries() + args.output.mkdir(parents=True, exist_ok=True) + + for task_path in args.tasks: + print(f"[capture] task={task_path} num_envs={args.num_envs} steps={args.steps}", flush=True) + result = capture_task(task_path, args) + task_name = result["task_name"] + npz_path = args.output / f"{task_name}.npz" + metadata_path = args.output / f"{task_name}.metadata.json" + np.savez_compressed(npz_path, **result["arrays"]) + metadata_path.write_text(json.dumps(result["metadata"], indent=2) + "\n") + print( + f"[capture] wrote {npz_path} ({npz_path.stat().st_size / 1e6:.2f} MB) " + f"and {metadata_path.name}", + flush=True, + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/tools/compare_mujoco_drift_baseline.py b/scripts/tools/compare_mujoco_drift_baseline.py new file mode 100644 index 000000000..c600daeef --- /dev/null +++ b/scripts/tools/compare_mujoco_drift_baseline.py @@ -0,0 +1,201 @@ +#!/usr/bin/env python3 +"""Characterize executor-swap drift between two mujoco drift-baseline captures. + +Roadmap issue #1552 replaces the mujoco backend's native executor +(mujoco-uni-runtime) with mjbatch. Numerical drift between the two is expected +and accepted; this tool quantifies it as a regression reference (issue #1554), +not as a pass/fail gate — it always exits 0 when the comparison completes. + +For every task present in both ``--before`` and ``--after`` it loads the .npz +artifacts written by ``capture_mujoco_drift_baseline.py`` and reports, per +array and per observation group: + +- max / mean absolute difference, +- the magnitude of the recorded values (so relative scale is visible), +- the first diverging step (leading-dim index of the first step whose slice + contains any difference; -1 when identical). + +Per-step arrays have a leading ``steps`` dimension; ``*_init`` snapshots and +``actions`` are compared as single snapshots (the action generator is seeded +and must be bit-identical — a mismatch there invalidates the comparison). + +Usage: + uv run scripts/tools/compare_mujoco_drift_baseline.py \ + --before scripts/tools/drift_baseline/before \ + --after scripts/tools/drift_baseline/after \ + --report scripts/tools/drift_baseline/drift_report.md +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +import numpy as np + +ROOT_DIR = Path(__file__).resolve().parents[2] +DEFAULT_BEFORE_DIR = ROOT_DIR / "scripts" / "tools" / "drift_baseline" / "before" +DEFAULT_AFTER_DIR = ROOT_DIR / "scripts" / "tools" / "drift_baseline" / "after" +DEFAULT_REPORT = ROOT_DIR / "scripts" / "tools" / "drift_baseline" / "drift_report.md" + + +def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--before", + type=Path, + default=DEFAULT_BEFORE_DIR, + help="BEFORE artifact directory (default: %(default)s).", + ) + parser.add_argument( + "--after", + type=Path, + default=DEFAULT_AFTER_DIR, + help="AFTER artifact directory (default: %(default)s).", + ) + parser.add_argument( + "--report", + type=Path, + default=DEFAULT_REPORT, + help="Markdown report output path (default: %(default)s).", + ) + return parser.parse_args(argv) + + +def _load(dir_path: Path) -> dict[str, dict[str, np.ndarray]]: + tasks: dict[str, dict[str, np.ndarray]] = {} + for npz_path in sorted(dir_path.glob("*.npz")): + with np.load(npz_path) as data: + tasks[npz_path.stem] = {name: data[name].copy() for name in data.files} + return tasks + + +def _compare_array(before: np.ndarray, after: np.ndarray) -> dict[str, object]: + if before.shape != after.shape: + return {"shape_before": before.shape, "shape_after": after.shape, "mismatch": True} + diff = np.abs(after.astype(np.float64) - before.astype(np.float64)) + result: dict[str, object] = { + "shape": before.shape, + "max_abs_diff": float(diff.max()) if diff.size else 0.0, + "mean_abs_diff": float(diff.mean()) if diff.size else 0.0, + "ref_max_abs": float(np.abs(before).max()) if before.size else 0.0, + "first_divergence": -1, + } + if diff.ndim >= 1 and diff.shape[0] > 1 and diff.size: + per_step = diff.reshape(diff.shape[0], -1).max(axis=1) + nonzero = np.flatnonzero(per_step > 0.0) + if nonzero.size: + result["first_divergence"] = int(nonzero[0]) + result["first_divergence_max_abs_diff"] = float(per_step[nonzero[0]]) + elif diff.size and diff.max() > 0.0: + result["first_divergence"] = 0 + result["first_divergence_max_abs_diff"] = float(diff.max()) + return result + + +def _fmt(value: object) -> str: + if isinstance(value, float): + return f"{value:.6e}" + return str(value) + + +def compare(before_dir: Path, after_dir: Path) -> dict[str, dict[str, dict[str, object]]]: + before_tasks = _load(before_dir) + after_tasks = _load(after_dir) + common = sorted(set(before_tasks) & set(after_tasks)) + if not common: + raise SystemExit(f"no task artifacts common to {before_dir} and {after_dir}") + + report: dict[str, dict[str, dict[str, object]]] = {} + for task in common: + before_arrays = before_tasks[task] + after_arrays = after_tasks[task] + arrays: dict[str, dict[str, object]] = {} + for name in sorted(set(before_arrays) | set(after_arrays)): + if name not in before_arrays or name not in after_arrays: + arrays[name] = {"mismatch": True, "note": "present in only one artifact"} + continue + arrays[name] = _compare_array(before_arrays[name], after_arrays[name]) + report[task] = arrays + return report + + +def render_markdown( + report: dict[str, dict[str, dict[str, object]]], + before_dir: Path, + after_dir: Path, +) -> str: + lines = [ + "# MuJoCo executor-swap drift characterization (#1554)", + "", + "BEFORE: mujoco-uni-runtime 0.5.0 executor, captured from" + " `scripts/tools/drift_baseline/before/`.", + "AFTER: mjbatch executor (unilabsim fork), captured from" + " `scripts/tools/drift_baseline/after/`.", + "", + "Both captures use the identical configuration: 8 envs, 300 steps," + " env seed 42, action seed 1234, `cpu_ids=[0, 1, 2, 3]`, fixed" + " pseudo-random action sequence (no policy network). Drift between the" + " two executors is expected and accepted; this report is a regression" + " reference, not a pass/fail gate.", + "", + f"- BEFORE dir: `{before_dir}`", + f"- AFTER dir: `{after_dir}`", + "", + ] + for task, arrays in report.items(): + lines.append(f"## {task}") + lines.append("") + lines.append( + "| array | shape | max abs diff | mean abs diff | ref max abs " + "| first divergence step | max abs diff at first divergence |" + ) + lines.append("|---|---|---|---|---|---|") + step_arrays = { + name: res + for name, res in arrays.items() + if not res.get("mismatch") and int(res.get("first_divergence", -1)) >= 0 + } + first_overall = ( + min(int(res["first_divergence"]) for res in step_arrays.values()) if step_arrays else -1 + ) + for name, res in arrays.items(): + if res.get("mismatch"): + lines.append(f"| `{name}` | — | — | — | — | present in only one artifact |") + continue + lines.append( + f"| `{name}` | {res['shape']} | {_fmt(res['max_abs_diff'])} " + f"| {_fmt(res['mean_abs_diff'])} | {_fmt(res['ref_max_abs'])} " + f"| {res['first_divergence']} " + f"| {_fmt(res.get('first_divergence_max_abs_diff', 0.0))} |" + ) + lines.append("") + lines.append(f"First divergence step (any array): **{first_overall}**") + lines.append("") + return "\n".join(lines) + + +def main(argv: list[str] | None = None) -> int: + args = _parse_args(argv) + report = compare(args.before, args.after) + args.report.parent.mkdir(parents=True, exist_ok=True) + args.report.write_text(render_markdown(report, args.before, args.after) + "\n") + + for task, arrays in report.items(): + print(f"[compare] {task}") + for name, res in arrays.items(): + if res.get("mismatch"): + print(f" {name}: present in only one artifact") + continue + print( + f" {name}: max={_fmt(res['max_abs_diff'])} " + f"mean={_fmt(res['mean_abs_diff'])} first_divergence={res['first_divergence']}" + ) + print(f"[compare] wrote {args.report}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/tools/drift_baseline/after/Go2JoystickFlat.metadata.json b/scripts/tools/drift_baseline/after/Go2JoystickFlat.metadata.json new file mode 100644 index 000000000..098790093 --- /dev/null +++ b/scripts/tools/drift_baseline/after/Go2JoystickFlat.metadata.json @@ -0,0 +1,45 @@ +{ + "task_name": "Go2JoystickFlat", + "hydra_task_path": "go2_joystick_flat/mujoco", + "sim_backend": "mujoco", + "num_envs": 8, + "num_steps": 300, + "seed": 42, + "action_seed": 1234, + "action_generator": "make_action_sequence: default_rng(action_seed) iid uniform(-1, 1); envs i%4==3 hold a seeded +/-1 vector for the first half, negated after", + "cpu_ids": [ + 0, + 1, + 2, + 3 + ], + "sim_dt": 0.01, + "ctrl_dt": 0.02, + "obs_groups_spec": { + "obs": 49, + "critic": 52 + }, + "action_dim": 12, + "git_commit": "fc64b6ae8b06eae89badc36cbc309d61e8f363af", + "package_versions": { + "mujoco": "3.11.0", + "mjbatch": "0.1.0", + "unisim-core": "1.2.0", + "numpy": "2.4.4" + }, + "timestamp_utc": "2026-09-11T17:43:31.433930+00:00", + "array_sha256": { + "reward": "b73f1549689a9a2fb263d70e69161b5e74b4b3ed96a91cf00b98dd643c22e40b", + "terminated": "9c37d0493149d133b7589ca5dfb535da71b39c53279dcfe4a860c3a6247239f4", + "truncated": "a0ee989ed2a0a2e3626520afa4032e06144865c8c8f6357293c9f4cd2069eaf2", + "qpos": "1a9ddfca94effa892949a1380a9132cb40a45e93048e4e488d563caa7a2bbaee", + "qvel": "541a22c268487af0e4c8648c3e518d7eb1e2113a74d9f1dfc70c6867e22f0dc7", + "obs/critic": "a2f1b86561d68722a9451edc98ef2d5b17f47c3e3aa917906b4f9695d1f798a2", + "obs/obs": "95d920fd9d5fffa72c46eaf181d290f51899e826ec654a6150031cb5c47f7ef4", + "actions": "08a768116d366f73238a70500f388fcb29d6743543b9693e271472aa280dda55", + "obs_init/critic": "72a29ca56490d1956105fbf53db63ff904c4fc8af693e7df7cb86e65bafdd185", + "obs_init/obs": "e69c020cfa82f501a98ccc336ae93b6edc1c608c57f825484c44c00ba88de6e4", + "qpos_init": "1f0c0fae9eaa1ab7809218486c84eabfc16ec2bc9ffc672008d634d4bba27f18", + "qvel_init": "c85bb78cef946b274d9682041ef5e8b0f29d18f50a1e0ba282cb1f8a156803ec" + } +} diff --git a/scripts/tools/drift_baseline/after/Go2JoystickFlat.npz b/scripts/tools/drift_baseline/after/Go2JoystickFlat.npz new file mode 100644 index 000000000..872754ead Binary files /dev/null and b/scripts/tools/drift_baseline/after/Go2JoystickFlat.npz differ diff --git a/scripts/tools/drift_baseline/after/Go2WJoystickFlat.metadata.json b/scripts/tools/drift_baseline/after/Go2WJoystickFlat.metadata.json new file mode 100644 index 000000000..2d1308713 --- /dev/null +++ b/scripts/tools/drift_baseline/after/Go2WJoystickFlat.metadata.json @@ -0,0 +1,45 @@ +{ + "task_name": "Go2WJoystickFlat", + "hydra_task_path": "go2w_joystick_flat/mujoco", + "sim_backend": "mujoco", + "num_envs": 8, + "num_steps": 300, + "seed": 42, + "action_seed": 1234, + "action_generator": "make_action_sequence: default_rng(action_seed) iid uniform(-1, 1); envs i%4==3 hold a seeded +/-1 vector for the first half, negated after", + "cpu_ids": [ + 0, + 1, + 2, + 3 + ], + "sim_dt": 0.005, + "ctrl_dt": 0.02, + "obs_groups_spec": { + "obs": 53, + "critic": 72 + }, + "action_dim": 16, + "git_commit": "fc64b6ae8b06eae89badc36cbc309d61e8f363af", + "package_versions": { + "mujoco": "3.11.0", + "mjbatch": "0.1.0", + "unisim-core": "1.2.0", + "numpy": "2.4.4" + }, + "timestamp_utc": "2026-09-11T17:43:31.120882+00:00", + "array_sha256": { + "reward": "b5ea168ca93c3280f4f461d32f8e21d5a0445f5962276236ff04710b8cc24c22", + "terminated": "1b4fa292d26b5857652ffa862089156a03091bb14d82fc107d3ec9bcebc4c866", + "truncated": "a0ee989ed2a0a2e3626520afa4032e06144865c8c8f6357293c9f4cd2069eaf2", + "qpos": "185ae7ba7075296515e3ea0da4dc7ec8806e4d5eea69ae49b95fafeacbf368f4", + "qvel": "a66798bf2ab1f9828f01540e4b2b8c0643246fe519907e0df3aacd2f8ec940ce", + "obs/critic": "b2d6f68504ab0ef62dd44671b4c3e4e0cfad04a480073760d58029fffbfed59e", + "obs/obs": "fa658e25ded3bc88704edb7b2827f7d37aa2fdd2b0a47af1798c940e5d800ef4", + "actions": "4c2df924382b83d34cb56a188cb966c23d27d30fd4c8a6973e5d92bc6e7e7a8f", + "obs_init/critic": "149379b215c40beb65da5e4301e064a3e92f39a112a2783fec577dfc8c34518d", + "obs_init/obs": "d4b73148fbb776fd7912c20dbd84fd3c0e8cdf9757868b8e30b5b92e406b8e72", + "qpos_init": "924ca2ab30273642e3a284c0b57718bb61df707e3df00b7d32792826b074db81", + "qvel_init": "c50170fa7c86bb5bc265f44bfe0d21e86d5f112feeab31725641b981d1e42c54" + } +} diff --git a/scripts/tools/drift_baseline/after/Go2WJoystickFlat.npz b/scripts/tools/drift_baseline/after/Go2WJoystickFlat.npz new file mode 100644 index 000000000..bfeb8e49d Binary files /dev/null and b/scripts/tools/drift_baseline/after/Go2WJoystickFlat.npz differ diff --git a/scripts/tools/drift_baseline/before/Go2JoystickFlat.metadata.json b/scripts/tools/drift_baseline/before/Go2JoystickFlat.metadata.json new file mode 100644 index 000000000..2e1037467 --- /dev/null +++ b/scripts/tools/drift_baseline/before/Go2JoystickFlat.metadata.json @@ -0,0 +1,47 @@ +{ + "task_name": "Go2JoystickFlat", + "hydra_task_path": "go2_joystick_flat/mujoco", + "sim_backend": "mujoco", + "num_envs": 8, + "num_steps": 300, + "seed": 42, + "action_seed": 1234, + "action_generator": "make_action_sequence: default_rng(action_seed) iid uniform(-1, 1); envs i%4==3 hold a seeded +/-1 vector for the first half, negated after", + "cpu_ids": [ + 0, + 1, + 2, + 3 + ], + "chunk_size": 2, + "adaptive_chunk_size": false, + "sim_dt": 0.01, + "ctrl_dt": 0.02, + "obs_groups_spec": { + "obs": 49, + "critic": 52 + }, + "action_dim": 12, + "git_commit": "db1a6e5b3dbe4dd096b16a517837f2d46dba8164", + "package_versions": { + "mujoco": "3.11.0", + "mujoco-uni-runtime": "0.5.0", + "unisim-core": "1.2.0", + "numpy": "2.4.4" + }, + "timestamp_utc": "2026-09-11T15:32:47.796881+00:00", + "array_sha256": { + "reward": "498775ce022146f2d5625e4ea7be941ec7b983e5cd3876d62d849e42827ae592", + "terminated": "9c37d0493149d133b7589ca5dfb535da71b39c53279dcfe4a860c3a6247239f4", + "truncated": "a0ee989ed2a0a2e3626520afa4032e06144865c8c8f6357293c9f4cd2069eaf2", + "qpos": "3103e25f39fc1ae578989f789e48e42fe446ff8864dd5719c376536de6d08b0e", + "qvel": "44d09f283ecc0c6cdd6e5ac2a8e50256572a2b5f04fd5ba2f5923b32f2f36d76", + "obs/critic": "ae084ceb54d60cda05e4b81423df6cb9108bfbbd78409f544d6460d6add596b9", + "obs/obs": "d252ea7c5cc081adeb7133b35dbcd7ab420bf6b5f58e1751867ea10f21f19bc7", + "actions": "08a768116d366f73238a70500f388fcb29d6743543b9693e271472aa280dda55", + "obs_init/critic": "72a29ca56490d1956105fbf53db63ff904c4fc8af693e7df7cb86e65bafdd185", + "obs_init/obs": "e69c020cfa82f501a98ccc336ae93b6edc1c608c57f825484c44c00ba88de6e4", + "qpos_init": "1f0c0fae9eaa1ab7809218486c84eabfc16ec2bc9ffc672008d634d4bba27f18", + "qvel_init": "c85bb78cef946b274d9682041ef5e8b0f29d18f50a1e0ba282cb1f8a156803ec" + } +} diff --git a/scripts/tools/drift_baseline/before/Go2JoystickFlat.npz b/scripts/tools/drift_baseline/before/Go2JoystickFlat.npz new file mode 100644 index 000000000..a8aa274bb Binary files /dev/null and b/scripts/tools/drift_baseline/before/Go2JoystickFlat.npz differ diff --git a/scripts/tools/drift_baseline/before/Go2WJoystickFlat.metadata.json b/scripts/tools/drift_baseline/before/Go2WJoystickFlat.metadata.json new file mode 100644 index 000000000..8cacf0a5c --- /dev/null +++ b/scripts/tools/drift_baseline/before/Go2WJoystickFlat.metadata.json @@ -0,0 +1,47 @@ +{ + "task_name": "Go2WJoystickFlat", + "hydra_task_path": "go2w_joystick_flat/mujoco", + "sim_backend": "mujoco", + "num_envs": 8, + "num_steps": 300, + "seed": 42, + "action_seed": 1234, + "action_generator": "make_action_sequence: default_rng(action_seed) iid uniform(-1, 1); envs i%4==3 hold a seeded +/-1 vector for the first half, negated after", + "cpu_ids": [ + 0, + 1, + 2, + 3 + ], + "chunk_size": 2, + "adaptive_chunk_size": false, + "sim_dt": 0.005, + "ctrl_dt": 0.02, + "obs_groups_spec": { + "obs": 53, + "critic": 72 + }, + "action_dim": 16, + "git_commit": "db1a6e5b3dbe4dd096b16a517837f2d46dba8164", + "package_versions": { + "mujoco": "3.11.0", + "mujoco-uni-runtime": "0.5.0", + "unisim-core": "1.2.0", + "numpy": "2.4.4" + }, + "timestamp_utc": "2026-09-11T15:32:47.470709+00:00", + "array_sha256": { + "reward": "469cd45dcd4efd4af7d549b7b339c062a444de8d55ee311316f867831851da09", + "terminated": "1b4fa292d26b5857652ffa862089156a03091bb14d82fc107d3ec9bcebc4c866", + "truncated": "a0ee989ed2a0a2e3626520afa4032e06144865c8c8f6357293c9f4cd2069eaf2", + "qpos": "fb4e78b67a0e26df9a483f0c15ba72883c3226ddc6e4988c9277cfcbf99b98e9", + "qvel": "d68d669c9efb7a5a269cc672d70a29229e16276f3573ad16ffbee9240701cfb0", + "obs/critic": "f7445001580e6fe5499c80b01bd07751309b632e621c1fd191bf0fb857cf355d", + "obs/obs": "ec7b82b63128e7ab85d7da08eaabbcab72b41b20f381c4574b58bf7450bc97c1", + "actions": "4c2df924382b83d34cb56a188cb966c23d27d30fd4c8a6973e5d92bc6e7e7a8f", + "obs_init/critic": "149379b215c40beb65da5e4301e064a3e92f39a112a2783fec577dfc8c34518d", + "obs_init/obs": "d4b73148fbb776fd7912c20dbd84fd3c0e8cdf9757868b8e30b5b92e406b8e72", + "qpos_init": "924ca2ab30273642e3a284c0b57718bb61df707e3df00b7d32792826b074db81", + "qvel_init": "c50170fa7c86bb5bc265f44bfe0d21e86d5f112feeab31725641b981d1e42c54" + } +} diff --git a/scripts/tools/drift_baseline/before/Go2WJoystickFlat.npz b/scripts/tools/drift_baseline/before/Go2WJoystickFlat.npz new file mode 100644 index 000000000..fdebe46fb Binary files /dev/null and b/scripts/tools/drift_baseline/before/Go2WJoystickFlat.npz differ diff --git a/scripts/tools/drift_baseline/drift_report.md b/scripts/tools/drift_baseline/drift_report.md new file mode 100644 index 000000000..575c5f326 --- /dev/null +++ b/scripts/tools/drift_baseline/drift_report.md @@ -0,0 +1,48 @@ +# MuJoCo executor-swap drift characterization (#1554) + +BEFORE: mujoco-uni-runtime 0.5.0 executor, captured from `scripts/tools/drift_baseline/before/`. +AFTER: mjbatch executor (unilabsim fork), captured from `scripts/tools/drift_baseline/after/`. + +Both captures use the identical configuration: 8 envs, 300 steps, env seed 42, action seed 1234, `cpu_ids=[0, 1, 2, 3]`, fixed pseudo-random action sequence (no policy network). Drift between the two executors is expected and accepted; this report is a regression reference, not a pass/fail gate. + +- BEFORE dir: `/home/user/ws/unilabsim2/UniLab/scripts/tools/drift_baseline/before` +- AFTER dir: `/home/user/ws/unilabsim2/UniLab/scripts/tools/drift_baseline/after` + +## Go2JoystickFlat + +| array | shape | max abs diff | mean abs diff | ref max abs | first divergence step | max abs diff at first divergence | +|---|---|---|---|---|---| +| `actions` | (300, 8, 12) | 0.000000e+00 | 0.000000e+00 | 1.000000e+00 | -1 | 0.000000e+00 | +| `obs/critic` | (300, 8, 52) | 2.768040e-03 | 4.664116e-06 | 1.162892e+01 | 1 | 1.633167e-05 | +| `obs/obs` | (300, 8, 49) | 2.768040e-03 | 4.869826e-06 | 1.162892e+01 | 1 | 1.633167e-05 | +| `obs_init/critic` | (8, 52) | 0.000000e+00 | 0.000000e+00 | 1.000000e+00 | -1 | 0.000000e+00 | +| `obs_init/obs` | (8, 49) | 0.000000e+00 | 0.000000e+00 | 1.000000e+00 | -1 | 0.000000e+00 | +| `qpos` | (300, 8, 19) | 1.286268e-04 | 1.379656e-06 | 2.227637e+00 | 1 | 4.768372e-07 | +| `qpos_init` | (8, 19) | 0.000000e+00 | 0.000000e+00 | 1.500000e+00 | -1 | 0.000000e+00 | +| `qvel` | (300, 8, 18) | 2.768040e-03 | 1.204221e-05 | 1.162892e+01 | 1 | 1.633167e-05 | +| `qvel_init` | (8, 18) | 0.000000e+00 | 0.000000e+00 | 5.306464e-01 | -1 | 0.000000e+00 | +| `reward` | (300, 8) | 4.128367e-05 | 1.929190e-07 | 1.274497e-01 | 1 | 1.303852e-08 | +| `terminated` | (300, 8) | 0.000000e+00 | 0.000000e+00 | 1.000000e+00 | -1 | 0.000000e+00 | +| `truncated` | (300, 8) | 0.000000e+00 | 0.000000e+00 | 0.000000e+00 | -1 | 0.000000e+00 | + +First divergence step (any array): **1** + +## Go2WJoystickFlat + +| array | shape | max abs diff | mean abs diff | ref max abs | first divergence step | max abs diff at first divergence | +|---|---|---|---|---|---| +| `actions` | (300, 8, 16) | 0.000000e+00 | 0.000000e+00 | 1.000000e+00 | -1 | 0.000000e+00 | +| `obs/critic` | (300, 8, 72) | 3.289003e+01 | 1.981591e-01 | 4.543000e+01 | 0 | 5.564094e-05 | +| `obs/obs` | (300, 8, 53) | 1.065044e+01 | 8.987349e-02 | 2.687317e+01 | 0 | 5.564094e-05 | +| `obs_init/critic` | (8, 72) | 0.000000e+00 | 0.000000e+00 | 1.000000e+00 | -1 | 0.000000e+00 | +| `obs_init/obs` | (8, 53) | 0.000000e+00 | 0.000000e+00 | 1.000000e+00 | -1 | 0.000000e+00 | +| `qpos` | (300, 8, 23) | 1.096748e+00 | 2.048076e-02 | 1.769883e+01 | 0 | 2.784654e-07 | +| `qpos_init` | (8, 23) | 0.000000e+00 | 0.000000e+00 | 1.500000e+00 | -1 | 0.000000e+00 | +| `qvel` | (300, 8, 22) | 1.065044e+01 | 2.105295e-01 | 2.687317e+01 | 0 | 5.564094e-05 | +| `qvel_init` | (8, 22) | 0.000000e+00 | 0.000000e+00 | 5.306464e-01 | -1 | 0.000000e+00 | +| `reward` | (300, 8) | 8.046474e-02 | 4.191729e-03 | 7.397851e-01 | 0 | 1.490116e-08 | +| `terminated` | (300, 8) | 0.000000e+00 | 0.000000e+00 | 1.000000e+00 | -1 | 0.000000e+00 | +| `truncated` | (300, 8) | 0.000000e+00 | 0.000000e+00 | 0.000000e+00 | -1 | 0.000000e+00 | + +First divergence step (any array): **0** + diff --git a/scripts/tools/drift_baseline/training_smoke_1554.md b/scripts/tools/drift_baseline/training_smoke_1554.md new file mode 100644 index 000000000..f1cf2b045 --- /dev/null +++ b/scripts/tools/drift_baseline/training_smoke_1554.md @@ -0,0 +1,38 @@ +# Training smoke — mjbatch executor validation (#1554) + +Short rsl_rl PPO runs on the mujoco backend (RTX 4090, local integration +branches: mjbatch fork + unisim adapter), exercising both action paths: + +- `Go2JoystickFlat` — plain position-action path. +- `Go2WJoystickFlat` — wheeled-leg task; its mixed action runs through + `SimBackend.set_pre_step_control`, i.e. mjbatch's native per-substep + callback (slimmed to `fn(k, state, ctrl)` by the post-swap ablation; this + smoke ran against the pre-ablation build with `callback_sensordata=False`). + +Both runs: `algo.num_envs=32 algo.max_iterations=100 algo.seed=42`, zero +NaN/Inf, playback video rendered after training. + +## Commands + +```bash +uv run train --algo ppo --task go2_joystick_flat --sim mujoco -- \ + algo.num_envs=32 algo.max_iterations=100 algo.seed=42 +uv run train --algo ppo --task go2w_joystick_flat --sim mujoco -- \ + algo.num_envs=32 algo.max_iterations=100 algo.seed=42 +``` + +(Executed with `UV_NO_SYNC=1` and the local mjbatch/unisim checkouts +installed editable; `UNILAB_LOCAL_UNISIM` set for the dependency-source +sentinel.) + +## Go2JoystickFlat (logs/rsl_rl_ppo/Go2JoystickFlat/2026-09-12_01-44-50_mujoco) + +- `Train/mean_reward`: 0.4275 (iter 0) → 17.96 (iter 99), monotone-ish, no NaN. +- 100 iterations, ~0.07 s/iteration, playback video rendered. + +## Go2WJoystickFlat (logs/rsl_rl_ppo/Go2WJoystickFlat/2026-09-12_01-45-37_mujoco) + +- `Train/mean_reward`: 1.3125 (iter 0) → 42.99 (iter 99), monotone-ish, no NaN. +- run_summary: 99 completed iterations, 76,800 env steps, + ~8,662 env-steps/s, final_mean_reward 42.99, mean_episode_length ~780 + (ctrl steps; no early termination mass), playback video rendered. diff --git a/src/unilab/base/backend_factory.py b/src/unilab/base/backend_factory.py index ba8e6fbf6..55dcee972 100644 --- a/src/unilab/base/backend_factory.py +++ b/src/unilab/base/backend_factory.py @@ -44,16 +44,12 @@ def _legacy_genesis_device_option_error(exc: TypeError) -> bool: def env_backend_kwargs(cfg: "EnvCfg") -> dict[str, Any]: """Translate ``EnvCfg`` backend knobs into UniSim adapter options.""" result: dict[str, Any] = { - "post_step_forward_sensor": cfg.post_step_forward_sensor, "superdex_num_workers": cfg.superdex_num_workers, "superdex_assets_root": cfg.superdex_assets_root, "superdex_effort_limits": cfg.superdex_effort_limits, "superdex_allow_contact_approximation": cfg.superdex_allow_contact_approximation, "motrix_max_iterations": cfg.motrix_max_iterations, - "chunk_size": cfg.chunk_size, - "adaptive_chunk_size": cfg.adaptive_chunk_size, "cpu_ids": cfg.cpu_ids, - "bench_nsteps": cfg.sim_substeps, "mjwarp_nconmax": cfg.mjwarp_nconmax, "mjwarp_njmax": cfg.mjwarp_njmax, "newton_device": cfg.newton_device, diff --git a/src/unilab/base/base.py b/src/unilab/base/base.py index 25a864c5a..91c4205d2 100644 --- a/src/unilab/base/base.py +++ b/src/unilab/base/base.py @@ -50,9 +50,6 @@ class EnvCfg: superdex_effort_limits: Optional[list[float]] = None superdex_allow_contact_approximation: bool = False motrix_max_iterations: Optional[int] = None - post_step_forward_sensor: bool = False - adaptive_chunk_size: bool = True - chunk_size: Optional[int] = None # Explicit CPU block owned by this env's process (Linux affinity only). # ``cpu_ids[i]`` pins MuJoCo BatchEnvPool worker thread ``i`` to one CPU; # env construction also confines the owning process to the same block and diff --git a/src/unilab/cli.py b/src/unilab/cli.py index acb288204..6d705d4e2 100644 --- a/src/unilab/cli.py +++ b/src/unilab/cli.py @@ -110,9 +110,9 @@ def _check_load_run(load_run: str) -> None: def _check_runtime_requirements(algo: str, sim: str) -> None: # The MuJoCo physics backend (unisim.backend.mujoco.backend) needs the - # mujoco-uni-runtime native binding; plain `mujoco` can also arrive via - # other extras (e.g. superdex), so gate on `mujoco_uni` here. - if sim == "mujoco" and (find_spec("mujoco") is None or find_spec("mujoco_uni") is None): + # mjbatch native batch engine; plain `mujoco` can also arrive via + # other extras (e.g. superdex), so gate on `mjbatch` here. + if sim == "mujoco" and (find_spec("mujoco") is None or find_spec("mjbatch") is None): raise SystemExit( "sim=mujoco requires the MuJoCo extra. Install it with " "`pip install unilab[mujoco]` (or `uv sync --extra mujoco` in a source checkout)." diff --git a/src/unilab/conf/appo/config.yaml b/src/unilab/conf/appo/config.yaml index 799fee5dc..6334e39b7 100644 --- a/src/unilab/conf/appo/config.yaml +++ b/src/unilab/conf/appo/config.yaml @@ -114,13 +114,6 @@ interactive: keyboard_step_lin: 0.1 keyboard_step_ang: 0.2 -env: - post_step_forward_sensor: false - # adaptive_chunk_size: auto-tune the MuJoCo BatchEnvPool chunk_size at materialize - # (cache-backed). chunk_size (int) manually overrides and wins; null => use default. - adaptive_chunk_size: true - chunk_size: null - hydra: run: dir: . diff --git a/src/unilab/conf/flashsac/config.yaml b/src/unilab/conf/flashsac/config.yaml index 27b3064cb..14e9a6173 100644 --- a/src/unilab/conf/flashsac/config.yaml +++ b/src/unilab/conf/flashsac/config.yaml @@ -134,13 +134,6 @@ interactive: keyboard_step_lin: 0.1 keyboard_step_ang: 0.2 -env: - post_step_forward_sensor: false - # adaptive_chunk_size: auto-tune the MuJoCo BatchEnvPool chunk_size at materialize - # (cache-backed). chunk_size (int) manually overrides and wins; null => use default. - adaptive_chunk_size: true - chunk_size: null - hydra: run: dir: . diff --git a/src/unilab/conf/ppo/config.yaml b/src/unilab/conf/ppo/config.yaml index 16a370b07..c7f5febe2 100644 --- a/src/unilab/conf/ppo/config.yaml +++ b/src/unilab/conf/ppo/config.yaml @@ -132,13 +132,6 @@ viser: display_mode: all max_envs: 16 -env: - post_step_forward_sensor: false - # adaptive_chunk_size: auto-tune the MuJoCo BatchEnvPool chunk_size at materialize - # (cache-backed). chunk_size (int) manually overrides and wins; null => use default. - adaptive_chunk_size: true - chunk_size: null - hydra: run: dir: . diff --git a/src/unilab/conf/ppo/task/fr3_joint_target/superdex.yaml b/src/unilab/conf/ppo/task/fr3_joint_target/superdex.yaml index 33960298c..c13ec2714 100644 --- a/src/unilab/conf/ppo/task/fr3_joint_target/superdex.yaml +++ b/src/unilab/conf/ppo/task/fr3_joint_target/superdex.yaml @@ -26,7 +26,6 @@ algo: num_mini_batches: 1 env: - adaptive_chunk_size: false superdex_num_workers: 0 superdex_assets_root: null superdex_effort_limits: [20.0, 20.0, 20.0, 20.0, 5.0, 5.0, 5.0] diff --git a/src/unilab/conf/ppo/task/go2_footstand/base.yaml b/src/unilab/conf/ppo/task/go2_footstand/base.yaml index 5478fc5b5..eafdd6826 100644 --- a/src/unilab/conf/ppo/task/go2_footstand/base.yaml +++ b/src/unilab/conf/ppo/task/go2_footstand/base.yaml @@ -52,7 +52,6 @@ env: sim_dt: 0.004 ctrl_dt: 0.02 max_episode_seconds: 10.0 - adaptive_chunk_size: false observations: policy: enable_corruption: true diff --git a/src/unilab/conf/ppo/task/go2_joystick_flat/superdex.yaml b/src/unilab/conf/ppo/task/go2_joystick_flat/superdex.yaml index 3fcfb625e..4abc4f763 100644 --- a/src/unilab/conf/ppo/task/go2_joystick_flat/superdex.yaml +++ b/src/unilab/conf/ppo/task/go2_joystick_flat/superdex.yaml @@ -26,7 +26,6 @@ env: lin_vel_x: [0.5, 0.5] lin_vel_y: [0.0, 0.0] ang_vel_z: [0.0, 0.0] - adaptive_chunk_size: false superdex_num_workers: 0 superdex_allow_contact_approximation: true events: diff --git a/src/unilab/conf/sac/config.yaml b/src/unilab/conf/sac/config.yaml index 8f71876f4..85f3dbe06 100644 --- a/src/unilab/conf/sac/config.yaml +++ b/src/unilab/conf/sac/config.yaml @@ -128,13 +128,6 @@ interactive: keyboard_step_lin: 0.1 keyboard_step_ang: 0.2 -env: - post_step_forward_sensor: false - # adaptive_chunk_size: auto-tune the MuJoCo BatchEnvPool chunk_size at materialize - # (cache-backed). chunk_size (int) manually overrides and wins; null => use default. - adaptive_chunk_size: true - chunk_size: null - hydra: run: dir: . diff --git a/src/unilab/conf/sac/task/go2_footstand/base.yaml b/src/unilab/conf/sac/task/go2_footstand/base.yaml index 08095758c..250f5d950 100644 --- a/src/unilab/conf/sac/task/go2_footstand/base.yaml +++ b/src/unilab/conf/sac/task/go2_footstand/base.yaml @@ -52,7 +52,6 @@ env: sim_dt: 0.004 ctrl_dt: 0.02 max_episode_seconds: 10.0 - adaptive_chunk_size: false observations: policy: enable_corruption: true diff --git a/src/unilab/conf/td3/config.yaml b/src/unilab/conf/td3/config.yaml index 909355b08..317a99b55 100644 --- a/src/unilab/conf/td3/config.yaml +++ b/src/unilab/conf/td3/config.yaml @@ -119,13 +119,6 @@ interactive: keyboard_step_lin: 0.1 keyboard_step_ang: 0.2 -env: - post_step_forward_sensor: false - # adaptive_chunk_size: auto-tune the MuJoCo BatchEnvPool chunk_size at materialize - # (cache-backed). chunk_size (int) manually overrides and wins; null => use default. - adaptive_chunk_size: true - chunk_size: null - hydra: run: dir: . diff --git a/tests/algos/test_appo_runner.py b/tests/algos/test_appo_runner.py index e650d188a..0ab1ad143 100644 --- a/tests/algos/test_appo_runner.py +++ b/tests/algos/test_appo_runner.py @@ -14,6 +14,10 @@ from hydra.core.global_hydra import GlobalHydra pytest.importorskip("mujoco") +pytest.importorskip( + "unisim.backend.mujoco.backend", + reason="unisim-core MuJoCo adapter (mjbatch build) not available", +) from uni_rl.algos.appo.runner import APPORunner diff --git a/tests/algos/test_rsl_rl_runner.py b/tests/algos/test_rsl_rl_runner.py index a80e6188a..a86020b35 100644 --- a/tests/algos/test_rsl_rl_runner.py +++ b/tests/algos/test_rsl_rl_runner.py @@ -15,6 +15,10 @@ from hydra.core.global_hydra import GlobalHydra pytest.importorskip("mujoco") +pytest.importorskip( + "unisim.backend.mujoco.backend", + reason="unisim-core MuJoCo adapter (mjbatch build) not available", +) rsl_rl = pytest.importorskip("rsl_rl") import numpy as np diff --git a/tests/base/backend/test_mujoco_backend_sentinel.py b/tests/base/backend/test_mujoco_backend_sentinel.py new file mode 100644 index 000000000..4aef0b86e --- /dev/null +++ b/tests/base/backend/test_mujoco_backend_sentinel.py @@ -0,0 +1,48 @@ +"""Guard-integrity sentinels for the mujoco backend test safety net (#1554). + +Every heavy mujoco suite gates itself on ``mjbatch`` + the unisim MuJoCo +adapter being importable, so a broken or missing install would silently skip +the whole safety net and leave the executor swap unverified. These sentinels +run unguarded in the fast lane and fail (never skip) when the guard cannot +admit a real backend, and prove the guard path end to end with a one-env +materialize/step — the same operations the gated suites exist to protect. +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np + +_MODEL_FILE = str( + Path(__file__).resolve().parents[2] / "fixtures" / "mjlab_cartpole" / "cartpole.xml" +) + + +def test_mjbatch_and_unisim_mujoco_adapter_importable() -> None: + """The exact guard expression that admits the gated mujoco suites.""" + import mjbatch # noqa: F401 + import unisim.backend.mujoco.backend # noqa: F401 + + +def test_mujoco_backend_materializes_and_steps_under_suite_guard() -> None: + """A real one-env backend materialize/step must succeed under that guard. + + Deliberately NOT importorskip-gated: if this errors, the gated heavy + suites would have silently skipped, which is the failure mode #1554 pins. + """ + from unisim.backend.mujoco.backend import MuJoCoBackend + + from unilab.base.scene import SceneCfg + + backend = MuJoCoBackend( + SceneCfg(model_file=_MODEL_FILE), + num_envs=1, + sim_dt=0.01, + base_name="cart", + ) + backend.materialize() + nu = backend.model.nu + backend.step(np.zeros((1, nu), dtype=np.float64), nsteps=2) + assert np.all(np.isfinite(backend.get_dof_pos())) + assert np.all(np.isfinite(backend.get_dof_vel())) diff --git a/tests/base/backend/test_mujoco_chunk_size_wiring.py b/tests/base/backend/test_mujoco_chunk_size_wiring.py deleted file mode 100644 index c88540d63..000000000 --- a/tests/base/backend/test_mujoco_chunk_size_wiring.py +++ /dev/null @@ -1,142 +0,0 @@ -from pathlib import Path - -import numpy as np -import pytest - -from unilab.base.backend_factory import env_backend_kwargs -from unilab.base.base import EnvCfg - -pytest.importorskip("mujoco", reason="mujoco not installed") - -try: - from mujoco_uni.batch_env import BatchEnvPool # noqa: F401 -except Exception: - pytest.skip( - "mujoco_uni.batch_env not available (platform/libstdc++ issue)", allow_module_level=True - ) - -import mujoco -from unisim.backend.mujoco.backend import MuJoCoBackend - -from unilab.base.scene import SceneCfg - -_MODEL_FILE = str(Path(__file__).resolve().parents[2] / "fixtures/free_chain.xml") -_NUM_ENVS = 4 - - -def test_envcfg_chunk_size_defaults(): - cfg = EnvCfg() - assert cfg.adaptive_chunk_size is True - assert cfg.chunk_size is None - - -def test_envcfg_chunk_size_overridable(): - cfg = EnvCfg(adaptive_chunk_size=False, chunk_size=128) - assert cfg.adaptive_chunk_size is False - assert cfg.chunk_size == 128 - - -def test_env_backend_kwargs_maps_fields(): - cfg = EnvCfg(ctrl_dt=0.02, sim_dt=0.005, chunk_size=64, adaptive_chunk_size=False) - kw = env_backend_kwargs(cfg) - assert kw["chunk_size"] == 64 - assert kw["adaptive_chunk_size"] is False - assert kw["bench_nsteps"] == cfg.sim_substeps == 4 # round(0.02/0.005) - assert kw["post_step_forward_sensor"] == cfg.post_step_forward_sensor - assert "motrix_max_iterations" in kw - - -def test_env_backend_kwargs_maps_drake_fields(): - kw = env_backend_kwargs(EnvCfg(drake_backend_mode="batch", drake_nthread=8)) - assert kw["drake_backend_mode"] == "batch" - assert kw["drake_nthread"] == 8 - default_kw = env_backend_kwargs(EnvCfg()) - assert default_kw["drake_backend_mode"] == "batch" - assert default_kw["drake_nthread"] == 0 - - -def _build_small_backend(**backend_kwargs): - """Build a minimal real MuJoCoBackend. - - Mirrors the SceneCfg + construction used in - ``tests/base/backend/test_mujoco_site_jacobian.py``, forwarding - ``**backend_kwargs`` into ``MuJoCoBackend(...)`` and calling ``.materialize()``. - """ - backend = MuJoCoBackend( - SceneCfg(model_file=_MODEL_FILE), - num_envs=_NUM_ENVS, - sim_dt=0.01, - base_name="base", - **backend_kwargs, - ) - backend.materialize() - return backend - - -def test_step_passes_resolved_chunk_size(monkeypatch): - backend = _build_small_backend(chunk_size=7, adaptive_chunk_size=False) - assert backend._chunk_size == 7 - - seen = {} - real_step = backend._pool.step - - def _spy(state, **kw): - seen["chunk_size"] = kw.get("chunk_size") - return real_step(state, **kw) - - monkeypatch.setattr(backend._pool, "step", _spy) - nu = backend._model.nu - backend.step(np.zeros((backend.num_envs, nu), dtype=np.float64), nsteps=1) - assert seen["chunk_size"] == 7 - - -def test_hot_path_does_no_xml_parse(monkeypatch): - """Acceptance ③: step/reset must not parse asset/XML (any entrypoint).""" - backend = _build_small_backend(adaptive_chunk_size=False) # cold path done - - # Install all parse spies AFTER the cold-path materialize so only hot-path - # (step) parses are counted. Spy multiple XML entrypoints, not just MjSpec. - spec_calls = {"n": 0} - model_calls = {"n": 0} - orig_from_file = mujoco.MjSpec.from_file - orig_from_xml_path = mujoco.MjModel.from_xml_path - - def _counting_from_file(*a, **k): - spec_calls["n"] += 1 - return orig_from_file(*a, **k) - - def _counting_from_xml_path(*a, **k): - model_calls["n"] += 1 - return orig_from_xml_path(*a, **k) - - monkeypatch.setattr(mujoco.MjSpec, "from_file", staticmethod(_counting_from_file)) - monkeypatch.setattr(mujoco.MjModel, "from_xml_path", staticmethod(_counting_from_xml_path)) - nu = backend._model.nu - backend.step(np.zeros((backend.num_envs, nu), dtype=np.float64), nsteps=1) - assert spec_calls["n"] == 0 - assert model_calls["n"] == 0 - - -def test_benchmark_runs_and_logs_table_on_adaptive(caplog, monkeypatch, tmp_path): - """Acceptance ④: adaptive path benchmarks and emits a per-candidate table. - - Point the chunk_size cache at an empty ``tmp_path`` file so the resolve is a - guaranteed MISS (no warm-cache short-circuit) and never pollutes the real - ``~/.cache/unisim/chunk_size.json``. A clean miss forces the benchmark path, - so we can assert the benchmark-table INFO record specifically. - """ - import logging - - from unisim.backend.mujoco import backend as backend_mod - - # Force nthread < num_envs so there is genuinely something to tune; otherwise the - # resolve short-circuits (num_envs <= nthread => one chunk => no benchmark). With - # _effective_cpu_count()==1, nthread = min(_NUM_ENVS, 1) = 1 < _NUM_ENVS, - # deterministically. - monkeypatch.setattr(backend_mod, "_effective_cpu_count", lambda: 1) - monkeypatch.setenv("UNISIM_CHUNK_SIZE_CACHE", str(tmp_path / "chunk_size.json")) - with caplog.at_level(logging.INFO, logger="unisim.backend.mujoco.chunk_tuner"): - backend = _build_small_backend(adaptive_chunk_size=True, chunk_size=None) - assert backend._chunk_size is None or isinstance(backend._chunk_size, int) - # Forced miss -> _log_benchmark_table emits the "chunk_size benchmark" record. - assert any("chunk_size benchmark" in r.message for r in caplog.records) diff --git a/tests/base/backend/test_mujoco_cpu_affinity_wiring.py b/tests/base/backend/test_mujoco_cpu_affinity_wiring.py index f6cd8c839..2c23afeb8 100644 --- a/tests/base/backend/test_mujoco_cpu_affinity_wiring.py +++ b/tests/base/backend/test_mujoco_cpu_affinity_wiring.py @@ -1,14 +1,21 @@ -"""CPU affinity wiring tests for the MuJoCo BatchEnvPool adapter (issue #959). +"""CPU affinity wiring tests for the MuJoCo backend on the mjbatch engine (#959, #1554). Covers the UniLab side of the contract: ``EnvCfg.cpu_ids`` validation, ``env_backend_kwargs``/``create_backend`` routing, cold-path validation in -``MuJoCoBackend``, and the actual worker pinning exposed by mujoco-uni. +``MuJoCoBackend``, and the pool wiring mjbatch exposes. + +The mjbatch fork pins workers inside ``Batch`` construction and exposes no +per-worker introspection (no ``pool.cpu_ids``/``worker_cpu_ids()``), so the +pool-level tests assert constructor/wiring behavior — materializing a pool +with pins succeeds and the worker count follows the pin list — rather than +observed per-thread affinity. """ import inspect import os from pathlib import Path +import numpy as np import pytest from unilab.base.backend_factory import create_backend, env_backend_kwargs @@ -17,20 +24,20 @@ pytest.importorskip("mujoco", reason="mujoco not installed") try: - from mujoco_uni.batch_env import BatchEnvPool + import mjbatch + from unisim.backend.mujoco.backend import MuJoCoBackend except Exception: pytest.skip( - "mujoco_uni.batch_env not available (platform/libstdc++ issue)", allow_module_level=True + "mjbatch/unisim MuJoCo backend not available (platform/build issue)", + allow_module_level=True, ) -if "cpu_ids" not in inspect.signature(BatchEnvPool.__init__).parameters: +if "cpu_ids" not in inspect.signature(mjbatch.Batch.__init__).parameters: pytest.skip( - "installed mujoco-uni-runtime has no cpu_ids support (pre-0.3.1)", + "installed mjbatch has no cpu_ids support", allow_module_level=True, ) -from unisim.backend.mujoco.backend import MuJoCoBackend - from unilab.base.scene import SceneCfg _MODEL_FILE = str( @@ -54,7 +61,6 @@ def _build_small_backend(**backend_kwargs): num_envs=_NUM_ENVS, sim_dt=0.01, base_name=_BASE_NAME, - adaptive_chunk_size=False, **backend_kwargs, ) backend.materialize() @@ -103,7 +109,6 @@ def test_create_backend_routes_cpu_ids(): _NUM_ENVS, 0.01, base_name=_BASE_NAME, - adaptive_chunk_size=False, cpu_ids=cpu_ids, ) assert isinstance(backend, MuJoCoBackend) @@ -124,14 +129,19 @@ def test_backend_rejects_invalid_cpu_ids_on_cold_path(cpu_ids): def test_workers_pinned_to_configured_cpus(): + """Materializing a pool with pins succeeds and the worker count follows. + + mjbatch applies the pinning inside ``Batch`` construction (worker i to + ``cpu_ids[i]``) and exposes no per-worker query, so the observable contract + here is wiring-level: the pool builds with exactly ``len(cpu_ids)`` + threads and stays usable. + """ cpu_ids = _AVAILABLE_CPUS[:2] backend = _build_small_backend(cpu_ids=cpu_ids) - try: - # Configured mapping is queryable and workers were observed on those CPUs. - assert tuple(backend._pool.cpu_ids) == tuple(cpu_ids) - assert backend._pool.worker_cpu_ids() == tuple(cpu_ids) - finally: - backend._pool.close() + assert backend._pool.num_threads == len(cpu_ids) + nu = backend._model.nu + backend.step(np.zeros((_NUM_ENVS, nu), dtype=np.float64), nsteps=1) + assert np.all(np.isfinite(backend.get_dof_pos())) def test_unavailable_cpu_id_fails_at_pool_creation(): @@ -149,12 +159,11 @@ def test_unavailable_cpu_id_fails_at_pool_creation(): def test_default_path_keeps_os_scheduling(): backend = _build_small_backend() - try: - assert backend._cpu_ids is None - assert backend._pool.cpu_ids is None - assert backend._pool.worker_cpu_ids() == () - finally: - backend._pool.close() + assert backend._cpu_ids is None + assert backend._pool.num_threads == backend._n_threads + nu = backend._model.nu + backend.step(np.zeros((_NUM_ENVS, nu), dtype=np.float64), nsteps=1) + assert np.all(np.isfinite(backend.get_dof_pos())) def test_default_nthread_sized_to_effective_cpus(): @@ -180,3 +189,35 @@ def test_default_nthread_capped_by_num_envs(): # num_envs caps the pool size, but the effective-CPU cap wins first on # single-core hosts (e.g. ubuntu-slim CI runners). assert backend._n_threads == min(len(os.sched_getaffinity(0)), 2) + + +def test_hot_path_does_no_xml_parse(monkeypatch): + """Step/reset must not parse asset/XML — an architecture contract. + + Moved here from the deleted ``test_mujoco_chunk_size_wiring.py`` (#1554). + Install all parse spies AFTER the cold-path materialize so only hot-path + (step) parses are counted. Spy multiple XML entrypoints, not just MjSpec. + """ + import mujoco + + backend = _build_small_backend() # cold path done + + spec_calls = {"n": 0} + model_calls = {"n": 0} + orig_from_file = mujoco.MjSpec.from_file + orig_from_xml_path = mujoco.MjModel.from_xml_path + + def _counting_from_file(*a, **k): + spec_calls["n"] += 1 + return orig_from_file(*a, **k) + + def _counting_from_xml_path(*a, **k): + model_calls["n"] += 1 + return orig_from_xml_path(*a, **k) + + monkeypatch.setattr(mujoco.MjSpec, "from_file", staticmethod(_counting_from_file)) + monkeypatch.setattr(mujoco.MjModel, "from_xml_path", staticmethod(_counting_from_xml_path)) + nu = backend._model.nu + backend.step(np.zeros((backend.num_envs, nu), dtype=np.float64), nsteps=1) + assert spec_calls["n"] == 0 + assert model_calls["n"] == 0 diff --git a/tests/base/backend/test_mujoco_scene_context_visual.py b/tests/base/backend/test_mujoco_scene_context_visual.py index 0c7182d31..16b22358d 100644 --- a/tests/base/backend/test_mujoco_scene_context_visual.py +++ b/tests/base/backend/test_mujoco_scene_context_visual.py @@ -11,7 +11,11 @@ import pytest -pytest.importorskip("mujoco_uni") +pytest.importorskip("mjbatch") +pytest.importorskip( + "unisim.backend.mujoco.backend", + reason="unisim-core MuJoCo adapter (mjbatch build) not available", +) from unisim.backend.mujoco.backend import _build_mujoco_scene_context diff --git a/tests/base/backend/test_mujoco_site_jacobian.py b/tests/base/backend/test_mujoco_site_jacobian.py index 2e20412fb..08dddbd50 100644 --- a/tests/base/backend/test_mujoco_site_jacobian.py +++ b/tests/base/backend/test_mujoco_site_jacobian.py @@ -10,14 +10,14 @@ pytest.importorskip("mujoco", reason="mujoco not installed") try: - from mujoco_uni.batch_env import BatchEnvPool # noqa: F401 + import mjbatch # noqa: F401 + from unisim.backend.mujoco.backend import MuJoCoBackend except Exception: pytest.skip( - "mujoco_uni.batch_env not available (platform/libstdc++ issue)", allow_module_level=True + "mjbatch/unisim MuJoCo backend not available (platform/build issue)", + allow_module_level=True, ) -from unisim.backend.mujoco.backend import MuJoCoBackend - from unilab.base.scene import SceneCfg MODEL_FILE = str(Path(__file__).resolve().parents[2] / "fixtures/free_chain.xml") @@ -108,15 +108,17 @@ def test_get_site_jacobian_matches_serial(backend): jacp_par, jacr_par = backend.get_site_jacobian_w(site_id, dof_indices) - # Serial reference. + # Serial reference built from the backend's state accessors (the mjbatch + # executor keeps canonical state in bound per-field views, not FULLPHYSICS + # rows, and has no per-env model variants). + model = backend._model jacp_ser = np.zeros((NUM_ENVS, 3, 6), dtype=np.float64) jacr_ser = np.zeros((NUM_ENVS, 3, 6), dtype=np.float64) for env_idx in range(NUM_ENVS): - variant_idx = int(backend._model_assignments[env_idx]) - model = backend._model_variants[variant_idx] data = mujoco.MjData(model) - state = np.asarray(backend._physics_state[env_idx], dtype=np.float64) - mujoco.mj_setState(model, data, state, int(mujoco.mjtState.mjSTATE_FULLPHYSICS)) + data.time = float(backend._time_view[env_idx]) + data.qpos[:] = backend._qpos_view[env_idx] + data.qvel[:] = backend._qvel_view[env_idx] mujoco.mj_forward(model, data) jacp_full = np.zeros((3, model.nv), dtype=np.float64) jacr_full = np.zeros((3, model.nv), dtype=np.float64) diff --git a/tests/base/test_backend_conformance.py b/tests/base/test_backend_conformance.py index 16365a7f1..f7ec1b44d 100644 --- a/tests/base/test_backend_conformance.py +++ b/tests/base/test_backend_conformance.py @@ -116,6 +116,10 @@ def _newton_runtime_available() -> bool: def _require_backend(backend_type: str) -> None: if backend_type == "mujoco": pytest.importorskip("mujoco", reason="mujoco not installed") + pytest.importorskip( + "unisim.backend.mujoco.backend", + reason="unisim-core MuJoCo adapter (mjbatch build) not available", + ) elif backend_type == "motrix": pytest.importorskip("motrixsim", reason="motrixsim not installed") elif backend_type == "mjwarp": @@ -397,6 +401,11 @@ def test_root_qvel_body_angular_contract_reads_back_world_velocity(backend_type: def test_mujoco_root_layout_resolves_a_nonfirst_free_joint() -> None: import mujoco + + pytest.importorskip( + "unisim.backend.mujoco.backend", + reason="unisim-core MuJoCo adapter (mjbatch build) not available", + ) from unisim.backend.mujoco.backend import MuJoCoBackend model = mujoco.MjModel.from_xml_string( diff --git a/tests/base/test_backend_imports.py b/tests/base/test_backend_imports.py index a21e0eb91..24258ba2d 100644 --- a/tests/base/test_backend_imports.py +++ b/tests/base/test_backend_imports.py @@ -86,11 +86,12 @@ def test_mujoco_backend_import_path_does_not_eagerly_import_motrix() -> None: print("mujoco_runtime", "mujoco" in sys.modules) print("mujoco_backend", "unisim.backend.mujoco.backend" in sys.modules) - if importlib.util.find_spec("mujoco_uni") is not None: + try: import unisim.backend.mujoco.backend - print("mujoco_backend imported") - else: + except ImportError: print("mujoco_backend skipped") + else: + print("mujoco_backend imported") print("motrix_backend", "unisim.backend.motrix.backend" in sys.modules) print("motrixsim", "motrixsim" in sys.modules) diff --git a/tests/base/test_backend_pre_step_control.py b/tests/base/test_backend_pre_step_control.py index 62be70300..b1114b5ab 100644 --- a/tests/base/test_backend_pre_step_control.py +++ b/tests/base/test_backend_pre_step_control.py @@ -41,78 +41,101 @@ def test_pre_step_control_rejects_shape_mismatch() -> None: SimBackend._apply_pre_step_control(backend, ctrl) # type: ignore[arg-type] -class _FakeMuJoCoPool: - def __init__(self) -> None: +class _FakeMjBatch: + """Minimal mjbatch ``Batch`` stand-in for the step-with-callback protocol. + + Emulates the documented mjbatch semantics the adapter relies on: the + callback runs on the calling thread before every substep as + ``callback(k, state, ctrl)``, one dispatch per ``step()`` call, and an + end-of-call copy-out that leaves the bound qpos/qvel views at the final + state and sensordata one substep behind (matching ``mj_step``). Each + substep adds 1.0 to every state row. + """ + + def __init__(self, nq: int = 1, nv: int = 1, nu: int = 2, nbody: int = 1) -> None: + self.num_sims = 1 + self._nq, self._nv = nq, nv + self.nact = 0 + self.qpos_slice = slice(0, nq) + self.qvel_slice = slice(nq, nq + nv) + self.act_slice = slice(nq + nv, nq + nv) + self._state = np.zeros((1, nq + nv)) + self._bufs: dict[str, np.ndarray] = { + "qpos": np.zeros((1, nq)), + "qvel": np.zeros((1, nv)), + "act": np.zeros((1, 0)), + "ctrl": np.zeros((1, nu)), + "xfrc_applied": np.zeros((1, 6 * nbody)), + "sensordata": np.zeros((1, 1)), + } self.step_calls: list[dict] = [] - self.forward_calls: list[np.ndarray] = [] self.callback_controls: list[np.ndarray] = [] + def bind(self, name: str, dtype=None) -> np.ndarray: + return self._bufs[name] + def step( self, - state, + ids=None, + nstep: int = 1, + history=None, *, - nstep, - control=None, - control_spec, - return_sensor=False, - post_step_forward_sensor=False, - chunk_size=None, - control_callback=None, - callback_sensordata=True, - ): + callback=None, + ) -> None: self.step_calls.append( { "nstep": nstep, - "control": None if control is None else np.array(control, copy=True), - "control_spec": control_spec, - "return_sensor": return_sensor, - "post_step_forward_sensor": post_step_forward_sensor, - "chunk_size": chunk_size, - "control_callback": control_callback, - "callback_sensordata": callback_sensordata, + "callback": callback, } ) - state_out = np.ascontiguousarray(np.asarray(state), dtype=np.float64) - if control_callback is not None: - # Emulate the upstream per-substep control_callback protocol: - # callback(0) sees the initial state and sensordata=None; each - # substep adds 1.0; callback(t>0) sees fresh state and sensordata - # only when callback_sensordata is true. - for t in range(nstep): - sensor_arg = None - if t > 0 and callback_sensordata: - sensor_arg = state_out[:, :1] - cb_control = control_callback(t, state_out, sensor_arg) - self.callback_controls.append(np.array(cb_control, copy=True)) - state_out = state_out + 1.0 - else: - state_out = state_out + 1.0 - if return_sensor: - return state_out, state_out[:, :1] - return state_out - - def forward(self, state): - state_np = np.asarray(state) - self.forward_calls.append(state_np.copy()) - return state_np[:, :1] - - -def _fake_mujoco_backend(pre_step_control_fn=None, post_step_forward_sensor=False): + state = self._state + ctrl_buf = self._bufs["ctrl"] + for k in range(nstep): + if callback is not None: + callback(k, state, ctrl_buf) + self.callback_controls.append(ctrl_buf.copy()) + state += 1.0 + # End-of-call copy-out of the bound input/derived fields. + self._bufs["qpos"][:] = state[:, self.qpos_slice] + self._bufs["qvel"][:] = state[:, self.qvel_slice] + self._bufs["sensordata"][:] = state[:, :1] - 1.0 # one substep behind + + +def _fake_mujoco_backend(pre_step_control_fn=None): + pytest.importorskip( + "mjbatch", + reason="mjbatch native batch engine not available", + ) + pytest.importorskip( + "unisim.backend.mujoco.backend", + reason="unisim-core MuJoCo adapter (mjbatch build) not available", + ) try: from unisim.backend.mujoco.backend import MuJoCoBackend except Exception as exc: pytest.skip(f"MuJoCo backend import unavailable: {exc}") + pool = _FakeMjBatch() backend = object.__new__(MuJoCoBackend) backend._pre_step_control_fn = pre_step_control_fn backend._num_envs = 1 backend._np_dtype = np.float32 - backend._physics_state = np.zeros((1, 1), dtype=np.float32) - backend._sensor_data = np.zeros((1, 1), dtype=np.float32) - backend._pending_xfrc_applied = np.zeros((1, 0), dtype=np.float64) - backend._post_step_forward_sensor = post_step_forward_sensor - backend._chunk_size = None - backend._pool = _FakeMuJoCoPool() + backend.nq = pool._nq + backend.nv = pool._nv + backend.nact = pool.nact + backend._root_qpos_dim = 7 + backend._root_qvel_dim = 6 + backend._state_layout = SimpleNamespace( + qpos=pool.qpos_slice, qvel=pool.qvel_slice, act=pool.act_slice + ) + backend._qpos_view = pool.bind("qpos", np.float64) + backend._qvel_view = pool.bind("qvel", np.float64) + backend._act_view = pool.bind("act", np.float64) + backend._ctrl_view = pool.bind("ctrl", np.float64) + backend._xfrc_view = pool.bind("xfrc_applied") + backend._sensor_data = pool.bind("sensordata", np.float64) + backend._pending_xfrc_applied = np.zeros((1, 6), dtype=np.float64) + backend._pool = pool return backend @@ -120,76 +143,85 @@ def test_mujoco_step_without_pre_step_control_keeps_batched_nsteps() -> None: backend = _fake_mujoco_backend() ctrl = np.array([[0.5, -0.5]], dtype=np.float32) - backend.step(ctrl, nsteps=3) - - assert len(backend._pool.step_calls) == 1 - assert backend._pool.step_calls[0]["nstep"] == 3 - assert backend._pool.step_calls[0]["return_sensor"] is True - assert backend._pool.step_calls[0]["post_step_forward_sensor"] is False - assert backend._pool.step_calls[0]["chunk_size"] is None - assert backend._pool.forward_calls == [] - expected_control = np.broadcast_to(ctrl[:, None, :], (1, 3, ctrl.shape[-1])) - np.testing.assert_allclose(backend._pool.step_calls[0]["control"], expected_control) - np.testing.assert_allclose(backend._physics_state, [[1.0]]) - np.testing.assert_allclose(backend._sensor_data, [[1.0]]) + result = backend.step(ctrl, nsteps=3) - -def test_mujoco_step_honors_post_step_forward_sensor_flag() -> None: - backend = _fake_mujoco_backend(post_step_forward_sensor=True) - ctrl = np.array([[0.5, -0.5]], dtype=np.float32) - - backend.step(ctrl, nsteps=3) - - assert backend._pool.step_calls[0]["return_sensor"] is True - assert backend._pool.step_calls[0]["post_step_forward_sensor"] is True + pool = backend._pool + assert len(pool.step_calls) == 1 + call = pool.step_calls[0] + assert call["nstep"] == 3 + assert call["callback"] is None + np.testing.assert_allclose(backend._ctrl_view, ctrl) + # Nothing staged: the persistent xfrc channel is written with zeros. + np.testing.assert_allclose(backend._xfrc_view.reshape(1, -1), [[0.0] * 6]) + np.testing.assert_allclose(backend._pending_xfrc_applied, [[0.0] * 6]) + np.testing.assert_allclose(backend._qpos_view, [[3.0]]) + assert set(result["timing"]) == {"set_ctrl_ms", "physics_ms", "refresh_cache_ms"} def test_mujoco_step_with_pre_step_control_uses_single_dispatch_callback() -> None: - seen_states: list[np.ndarray] = [] + seen_qpos: list[np.ndarray] = [] seen_sensors: list[np.ndarray] = [] - backend = _fake_mujoco_backend(post_step_forward_sensor=True) + backend = _fake_mujoco_backend() def hook(current_backend, owner_ctrl: np.ndarray) -> np.ndarray: - seen_states.append(current_backend._physics_state.copy()) + seen_qpos.append(current_backend._qpos_view.copy()) seen_sensors.append(current_backend._sensor_data.copy()) - return owner_ctrl + len(seen_states) + return owner_ctrl + len(seen_qpos) backend.set_pre_step_control(hook) ctrl = np.array([[0.5, -0.5]], dtype=np.float32) - backend.step(ctrl, nsteps=3) + result = backend.step(ctrl, nsteps=3) pool = backend._pool assert len(pool.step_calls) == 1 call = pool.step_calls[0] assert call["nstep"] == 3 - assert call["control"] is None - assert call["control_callback"] is not None - assert call["callback_sensordata"] is False - assert call["return_sensor"] is True - assert call["post_step_forward_sensor"] is True - assert call["chunk_size"] is None - assert pool.forward_calls == [] - # The hook sees the physics state refreshed before every substep. - np.testing.assert_allclose(seen_states, [[[0.0]], [[1.0]], [[2.0]]]) - # _sensor_data is no longer refreshed per substep (action terms only read - # physics-state-backed getters); it is refreshed once from the final - # return below. + assert call["callback"] is not None + # k=0 refreshes nothing (the bound views already hold the pre-call state); + # k>0 receives the state after substep k-1. + np.testing.assert_allclose(seen_qpos, [[[0.0]], [[1.0]], [[2.0]]]) + # The slim callback protocol carries no sensordata argument, so the hook + # never sees a sensor refresh; the bound sensordata view only catches up + # via the end-of-call copy-out. np.testing.assert_allclose(seen_sensors, [[[0.0]], [[0.0]], [[0.0]]]) assert len(pool.callback_controls) == 3 np.testing.assert_allclose(pool.callback_controls[0], ctrl + 1) np.testing.assert_allclose(pool.callback_controls[1], ctrl + 2) np.testing.assert_allclose(pool.callback_controls[2], ctrl + 3) assert pool.callback_controls[0].dtype == np.float64 - np.testing.assert_allclose(backend._physics_state, [[3.0]]) - np.testing.assert_allclose(backend._sensor_data, [[3.0]]) + np.testing.assert_allclose(backend._qpos_view, [[3.0]]) + np.testing.assert_allclose(backend._sensor_data, [[2.0]]) # one substep behind + assert set(result["timing"]) == {"set_ctrl_ms", "physics_ms", "refresh_cache_ms"} + + +def test_mujoco_step_writes_xfrc_absolutely_and_clears_pending() -> None: + backend = _fake_mujoco_backend() + backend._pending_xfrc_applied = np.full((1, 6), 7.0, dtype=np.float64) + ctrl = np.array([[0.5, -0.5]], dtype=np.float32) + + backend.step(ctrl, nsteps=2) + + pool = backend._pool + assert len(pool.step_calls) == 1 + # The staged wrench reached the persistent channel and survives the call; + # the pending staging buffer is cleared. + np.testing.assert_allclose(backend._xfrc_view.reshape(1, -1), [[7.0] * 6]) + np.testing.assert_allclose(backend._pending_xfrc_applied, [[0.0] * 6]) + + # An idle second dispatch must overwrite the channel with zeros, not leave + # the previous wrench to decay on its own. + backend.step(ctrl, nsteps=1) + + assert len(pool.step_calls) == 2 + np.testing.assert_allclose(backend._xfrc_view.reshape(1, -1), [[0.0] * 6]) + np.testing.assert_allclose(backend._pending_xfrc_applied, [[0.0] * 6]) -def test_mujoco_step_with_pre_step_control_appends_pending_xfrc_each_substep() -> None: - mujoco = pytest.importorskip("mujoco") +def test_mujoco_step_with_pre_step_control_stages_xfrc_before_dispatch() -> None: backend = _fake_mujoco_backend() - backend._pending_xfrc_applied = np.full((1, 2), 7.0, dtype=np.float64) + backend._pending_xfrc_applied = np.full((1, 6), 7.0, dtype=np.float64) backend.set_pre_step_control(lambda current_backend, owner_ctrl: owner_ctrl + 1.0) ctrl = np.array([[0.5, -0.5]], dtype=np.float32) @@ -198,11 +230,12 @@ def test_mujoco_step_with_pre_step_control_appends_pending_xfrc_each_substep() - pool = backend._pool assert len(pool.step_calls) == 1 - assert pool.step_calls[0]["control_spec"] & int(mujoco.mjtState.mjSTATE_XFRC_APPLIED) + # The callback protocol writes ctrl only; the wrench rides its own channel. assert len(pool.callback_controls) == 2 for cb_control in pool.callback_controls: - np.testing.assert_allclose(cb_control, [[1.5, 0.5, 7.0, 7.0]]) - np.testing.assert_allclose(backend._pending_xfrc_applied, [[0.0, 0.0]]) + np.testing.assert_allclose(cb_control, [[1.5, 0.5]]) + np.testing.assert_allclose(backend._xfrc_view.reshape(1, -1), [[7.0] * 6]) + np.testing.assert_allclose(backend._pending_xfrc_applied, [[0.0] * 6]) class _FakeMotrixModel: diff --git a/tests/base/test_entity_facade.py b/tests/base/test_entity_facade.py index 803dfdf97..e4d7e1ef8 100644 --- a/tests/base/test_entity_facade.py +++ b/tests/base/test_entity_facade.py @@ -743,6 +743,10 @@ def test_entity_facade_has_no_backend_model_or_asset_access() -> None: def test_real_mujoco_entity_selector_and_numpy_state_smoke() -> None: + pytest.importorskip( + "unisim.backend.mujoco.backend", + reason="unisim-core MuJoCo adapter (mjbatch build) not available", + ) from unisim.backend.mujoco.backend import MuJoCoBackend joint_names = ( diff --git a/tests/base/test_mujoco_batch_env_jacobian.py b/tests/base/test_mujoco_batch_env_jacobian.py deleted file mode 100644 index d2c6df55b..000000000 --- a/tests/base/test_mujoco_batch_env_jacobian.py +++ /dev/null @@ -1,146 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass -from typing import Any, Iterator - -import numpy as np -import pytest - -pytest.importorskip("mujoco", reason="mujoco not installed") - -try: - import mujoco - from mujoco_uni.batch_env import BatchEnvPool -except Exception: - pytest.skip( - "mujoco_uni.batch_env not available (platform/libstdc++ issue)", allow_module_level=True - ) - -if not hasattr(BatchEnvPool, "compute_site_jacobians"): - pytest.skip( - "BatchEnvPool.compute_site_jacobians requires a mujoco-uni-runtime build that ships it", - allow_module_level=True, - ) - -from unilab.assets import ASSETS_ROOT_PATH - -mj: Any = mujoco - -GO2_SITE_NAME = "imu" - - -@dataclass -class _PoolCtx: - model: Any - pool: BatchEnvPool - initial_state: np.ndarray - site_id: int - - -def _xml(robot: str, scene: str = "scene_flat.xml") -> str: - return str(ASSETS_ROOT_PATH / "robots" / robot / scene) - - -def _make_initial_state(model: Any, nbatch: int, rng: np.random.Generator) -> np.ndarray: - nstate = mj.mj_stateSize(model, mj.mjtState.mjSTATE_FULLPHYSICS) - state0 = np.zeros((nbatch, nstate), dtype=np.float64) - state0[:, 1 : 1 + model.nq] = model.qpos0 - # Perturb dof positions / velocities so each env produces a distinct Jacobian. - state0[:, 1 + 7 : 1 + model.nq] += 0.05 * rng.standard_normal((nbatch, model.nq - 7)) - state0[:, 1 + model.nq + 6 : 1 + model.nq + model.nv] += 0.02 * rng.standard_normal( - (nbatch, model.nv - 6) - ) - return state0 - - -def _reference_site_jacobian( - model: Any, state_row: np.ndarray, site_id: int, jacp: bool, jacr: bool -) -> tuple[np.ndarray | None, np.ndarray | None]: - data = mj.MjData(model) - mj.mj_setState(model, data, state_row, int(mj.mjtState.mjSTATE_FULLPHYSICS)) - mj.mj_kinematics(model, data) - mj.mj_comPos(model, data) - jp = np.zeros((3, model.nv), dtype=np.float64) if jacp else None - jr = np.zeros((3, model.nv), dtype=np.float64) if jacr else None - mj.mj_jacSite(model, data, jp, jr, site_id) - return jp, jr - - -@pytest.fixture -def pool_ctx() -> Iterator[_PoolCtx]: - rng = np.random.default_rng(0) - model = mj.MjModel.from_xml_path(_xml("go2")) - site_id = mj.mj_name2id(model, int(mj.mjtObj.mjOBJ_SITE), GO2_SITE_NAME) - assert site_id >= 0, f"go2 model is expected to expose site '{GO2_SITE_NAME}'" - nbatch = 3 - pool = BatchEnvPool(model, nbatch=nbatch, nthread=2) - try: - yield _PoolCtx( - model=model, - pool=pool, - initial_state=_make_initial_state(model, nbatch, rng), - site_id=site_id, - ) - finally: - pool.close() - - -def test_compute_site_jacobians_matches_reference_jacp_only(pool_ctx: _PoolCtx) -> None: - jp, jr = pool_ctx.pool.compute_site_jacobians( - pool_ctx.initial_state, [pool_ctx.site_id], jacp=True, jacr=False - ) - assert jr is None - assert jp.shape == (pool_ctx.initial_state.shape[0], 1, 3, pool_ctx.model.nv) - for i in range(pool_ctx.initial_state.shape[0]): - ref_jp, _ = _reference_site_jacobian( - pool_ctx.model, pool_ctx.initial_state[i], pool_ctx.site_id, True, False - ) - np.testing.assert_allclose(jp[i, 0], ref_jp, atol=1e-12, rtol=0) - - -def test_compute_site_jacobians_matches_reference_jacp_and_jacr(pool_ctx: _PoolCtx) -> None: - jp, jr = pool_ctx.pool.compute_site_jacobians( - pool_ctx.initial_state, [pool_ctx.site_id], jacp=True, jacr=True - ) - assert jp.shape == (pool_ctx.initial_state.shape[0], 1, 3, pool_ctx.model.nv) - assert jr.shape == (pool_ctx.initial_state.shape[0], 1, 3, pool_ctx.model.nv) - for i in range(pool_ctx.initial_state.shape[0]): - ref_jp, ref_jr = _reference_site_jacobian( - pool_ctx.model, pool_ctx.initial_state[i], pool_ctx.site_id, True, True - ) - np.testing.assert_allclose(jp[i, 0], ref_jp, atol=1e-12, rtol=0) - np.testing.assert_allclose(jr[i, 0], ref_jr, atol=1e-12, rtol=0) - - -def test_compute_site_jacobians_scalar_site_squeezes_k_dim(pool_ctx: _PoolCtx) -> None: - jp, jr = pool_ctx.pool.compute_site_jacobians( - pool_ctx.initial_state, pool_ctx.site_id, jacp=True, jacr=True - ) - assert jp.shape == (pool_ctx.initial_state.shape[0], 3, pool_ctx.model.nv) - assert jr.shape == (pool_ctx.initial_state.shape[0], 3, pool_ctx.model.nv) - for i in range(pool_ctx.initial_state.shape[0]): - ref_jp, ref_jr = _reference_site_jacobian( - pool_ctx.model, pool_ctx.initial_state[i], pool_ctx.site_id, True, True - ) - np.testing.assert_allclose(jp[i], ref_jp, atol=1e-12, rtol=0) - np.testing.assert_allclose(jr[i], ref_jr, atol=1e-12, rtol=0) - - -def test_compute_site_jacobians_requires_at_least_one_flag(pool_ctx: _PoolCtx) -> None: - with pytest.raises(ValueError): - pool_ctx.pool.compute_site_jacobians( - pool_ctx.initial_state, [pool_ctx.site_id], jacp=False, jacr=False - ) - - -def test_compute_site_jacobians_rejects_invalid_site_id(pool_ctx: _PoolCtx) -> None: - pytest.xfail( - "current mujoco-uni-runtime BatchEnvPool aborts on invalid site ids; " - "UniLab validates ids before native calls at the backend boundary" - ) - - -def test_compute_site_jacobians_rejects_wrong_state_shape(pool_ctx: _PoolCtx) -> None: - bad_state = pool_ctx.initial_state[:-1] - with pytest.raises(ValueError): - pool_ctx.pool.compute_site_jacobians(bad_state, [pool_ctx.site_id], jacp=True) diff --git a/tests/base/test_mujoco_batch_env_randomization.py b/tests/base/test_mujoco_batch_env_randomization.py deleted file mode 100644 index c9228689a..000000000 --- a/tests/base/test_mujoco_batch_env_randomization.py +++ /dev/null @@ -1,139 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass -from typing import Any, Iterator - -import numpy as np -import pytest - -pytest.importorskip("mujoco", reason="mujoco not installed") - -try: - import mujoco - from mujoco_uni.batch_env import SUPPORTED_FIELDS, BatchEnvPool -except Exception: - pytest.skip( - "mujoco_uni.batch_env not available (platform/libstdc++ issue)", allow_module_level=True - ) - -from unilab.assets import ASSETS_ROOT_PATH - -mj: Any = mujoco - -EXPECTED_SUPPORTED_FIELDS = { - "body_mass", - "body_ipos", - "gravity", - "body_iquat", - "body_inertia", - "dof_armature", - "geom_friction", - "kp", - "kd", -} - - -@dataclass -class _PoolCtx: - model: Any - pool: BatchEnvPool - initial_state: np.ndarray - - -def _xml(robot: str, scene: str = "scene_flat.xml") -> str: - return str(ASSETS_ROOT_PATH / "robots" / robot / scene) - - -def _make_initial_state(model: Any) -> np.ndarray: - nstate = mj.mj_stateSize(model, mj.mjtState.mjSTATE_FULLPHYSICS) - initial_state = np.zeros((1, nstate), dtype=np.float64) - initial_state[:, 1 : 1 + model.nq] = model.qpos0 - return initial_state - - -@pytest.fixture -def pool_ctx() -> Iterator[_PoolCtx]: - model = mj.MjModel.from_xml_path(_xml("go2")) - pool = BatchEnvPool(model, nbatch=2, nthread=1) - try: - yield _PoolCtx(model=model, pool=pool, initial_state=_make_initial_state(model)) - finally: - pool.close() - - -def _reset_and_assert_field_applied( - pool_ctx: _PoolCtx, field_name: str, updated: np.ndarray -) -> None: - original_0 = pool_ctx.pool.get_field(0, field_name).copy() - original_1 = pool_ctx.pool.get_field(1, field_name).copy() - - assert updated.shape == original_1.shape - assert not np.allclose(updated, original_1) - - pool_ctx.pool.reset( - env_ids=[1], - initial_state=pool_ctx.initial_state.copy(), - randomization={field_name: updated[None, :]}, - ) - - np.testing.assert_array_equal(pool_ctx.pool.get_field(0, field_name), original_0) - np.testing.assert_allclose(pool_ctx.pool.get_field(1, field_name), updated) - - -def test_batch_env_supported_fields_match_documented_reset_randomization_fields() -> None: - assert set(SUPPORTED_FIELDS) == EXPECTED_SUPPORTED_FIELDS - - -def test_batch_env_reset_applies_body_mass_randomization(pool_ctx: _PoolCtx) -> None: - updated = pool_ctx.pool.get_field(1, "body_mass").copy() - updated[1] += 0.25 - _reset_and_assert_field_applied(pool_ctx, "body_mass", updated) - - -def test_batch_env_reset_applies_body_ipos_randomization(pool_ctx: _PoolCtx) -> None: - updated = pool_ctx.pool.get_field(1, "body_ipos").reshape(pool_ctx.model.nbody, 3).copy() - updated[1] += np.array([0.01, -0.02, 0.03], dtype=np.float64) - _reset_and_assert_field_applied(pool_ctx, "body_ipos", updated.reshape(-1)) - - -def test_batch_env_reset_applies_gravity_randomization(pool_ctx: _PoolCtx) -> None: - updated = pool_ctx.pool.get_field(1, "gravity").copy() - updated += np.array([0.2, -0.1, 0.4], dtype=np.float64) - _reset_and_assert_field_applied(pool_ctx, "gravity", updated) - - -def test_batch_env_reset_applies_body_iquat_randomization(pool_ctx: _PoolCtx) -> None: - updated = pool_ctx.pool.get_field(1, "body_iquat").reshape(pool_ctx.model.nbody, 4).copy() - quat = np.array([0.92387953, 0.0, 0.38268343, 0.0], dtype=np.float64) - updated[1] = quat / np.linalg.norm(quat) - _reset_and_assert_field_applied(pool_ctx, "body_iquat", updated.reshape(-1)) - - -def test_batch_env_reset_applies_body_inertia_randomization(pool_ctx: _PoolCtx) -> None: - updated = pool_ctx.pool.get_field(1, "body_inertia").reshape(pool_ctx.model.nbody, 3).copy() - updated[1] *= 1.25 - _reset_and_assert_field_applied(pool_ctx, "body_inertia", updated.reshape(-1)) - - -def test_batch_env_reset_applies_dof_armature_randomization(pool_ctx: _PoolCtx) -> None: - updated = pool_ctx.pool.get_field(1, "dof_armature").copy() - updated += np.linspace(0.01, 0.03, updated.size, dtype=np.float64) - _reset_and_assert_field_applied(pool_ctx, "dof_armature", updated) - - -def test_batch_env_reset_applies_geom_friction_randomization(pool_ctx: _PoolCtx) -> None: - updated = pool_ctx.pool.get_field(1, "geom_friction").reshape(pool_ctx.model.ngeom, 3).copy() - updated[0] += np.array([0.1, 0.002, 0.0002], dtype=np.float64) - _reset_and_assert_field_applied(pool_ctx, "geom_friction", updated.reshape(-1)) - - -def test_batch_env_reset_applies_kp_randomization(pool_ctx: _PoolCtx) -> None: - updated = pool_ctx.pool.get_field(1, "kp").copy() - updated += 1.25 - _reset_and_assert_field_applied(pool_ctx, "kp", updated) - - -def test_batch_env_reset_applies_kd_randomization(pool_ctx: _PoolCtx) -> None: - updated = pool_ctx.pool.get_field(1, "kd").copy() - updated += 0.25 - _reset_and_assert_field_applied(pool_ctx, "kd", updated) diff --git a/tests/base/test_reward_override.py b/tests/base/test_reward_override.py index 6b635ee1e..e378ddbc5 100644 --- a/tests/base/test_reward_override.py +++ b/tests/base/test_reward_override.py @@ -3,6 +3,7 @@ from pathlib import Path from typing import Any, cast +import pytest from hydra import compose, initialize_config_dir from hydra.core.global_hydra import GlobalHydra @@ -15,6 +16,10 @@ def test_reward_override_g1(): """Test G1 manager reward override through the registry.""" + pytest.importorskip( + "unisim.backend.mujoco.backend", + reason="unisim-core MuJoCo adapter (mjbatch build) not available", + ) ensure_registries() GlobalHydra.instance().clear() diff --git a/tests/base/test_sim_backend.py b/tests/base/test_sim_backend.py index d2ac7d527..dc9c2472b 100644 --- a/tests/base/test_sim_backend.py +++ b/tests/base/test_sim_backend.py @@ -15,10 +15,7 @@ from unilab.assets import ASSETS_ROOT_PATH from unilab.base.scene import SceneCfg from unilab.dr import ( - GeomSizeOverride, - InitRandomizationPlan, IntervalRandomizationPlan, - ModelVariantSpec, ResetRandomizationPayload, ) @@ -161,72 +158,6 @@ def bkd(self, request): def test_num_envs(self, bkd): assert bkd.num_envs == NUM_ENVS - def test_apply_init_randomization_sets_variants_before_materialization(self): - from unisim.backend.mujoco.backend import MuJoCoBackend - - bkd = MuJoCoBackend( - SceneCfg(model_file=_G1["model_file"]), 4, SIM_DT, base_name=_G1["base_name"] - ) - assert bkd._pool is None - mujoco = _mujoco_module() - geom_id = mujoco.mj_name2id(bkd.model, mujoco.mjtObj.mjOBJ_GEOM, "floor") - base_size = np.asarray(bkd.model.geom_size[geom_id], dtype=np.float64).copy() - - bkd.apply_init_randomization( - InitRandomizationPlan( - model_assignments=np.array([0, 1, 0, 1], dtype=np.int32), - model_variants=( - ModelVariantSpec( - geom_size_overrides=(GeomSizeOverride("floor", tuple(base_size * 0.5)),) - ), - ModelVariantSpec( - geom_size_overrides=(GeomSizeOverride("floor", tuple(base_size * 0.75)),) - ), - ), - ) - ) - - assert bkd._pool is None - np.testing.assert_array_equal( - bkd._model_assignments, - np.array([0, 1, 0, 1], dtype=np.int32), - ) - np.testing.assert_allclose(bkd._model_variants[0].geom_size[geom_id], base_size * 0.5) - np.testing.assert_allclose(bkd._model_variants[1].geom_size[geom_id], base_size * 0.75) - - bkd.materialize() - assert bkd._pool is not None - - def test_get_playback_model_returns_env_specific_variant(self): - from unisim.backend.mujoco.backend import MuJoCoBackend - - bkd = MuJoCoBackend( - SceneCfg(model_file=_G1["model_file"]), 4, SIM_DT, base_name=_G1["base_name"] - ) - mujoco = _mujoco_module() - geom_id = mujoco.mj_name2id(bkd.model, mujoco.mjtObj.mjOBJ_GEOM, "floor") - base_size = np.asarray(bkd.model.geom_size[geom_id], dtype=np.float64).copy() - - bkd.apply_init_randomization( - InitRandomizationPlan( - model_assignments=np.array([0, 1, 0, 1], dtype=np.int32), - model_variants=( - ModelVariantSpec( - geom_size_overrides=(GeomSizeOverride("floor", tuple(base_size * 0.5)),) - ), - ModelVariantSpec( - geom_size_overrides=(GeomSizeOverride("floor", tuple(base_size * 0.75)),) - ), - ), - ) - ) - - model0 = bkd.get_playback_model(0) - model1 = bkd.get_playback_model(1) - - np.testing.assert_allclose(model0.geom_size[geom_id], base_size * 0.5) - np.testing.assert_allclose(model1.geom_size[geom_id], base_size * 0.75) - # simulation control def test_set_state_only_affects_target_envs(self, bkd): @@ -238,7 +169,7 @@ def test_set_state_only_affects_target_envs(self, bkd): def test_set_state_randomization_only_affects_target_envs(self, bkd): pool = bkd._pool - original = [pool.get_field(i, "body_mass").copy() for i in range(NUM_ENVS)] + original = pool.expand("body_mass").copy() qpos = _identity_qpos_mujoco(bkd.model.nq) qvel = np.zeros((1, bkd.model.nv)) base_body_id = bkd._base_body_id @@ -247,8 +178,8 @@ def test_set_state_randomization_only_affects_target_envs(self, bkd): bkd.set_state(np.array([1]), qpos, qvel, randomization=randomization) - np.testing.assert_array_equal(pool.get_field(0, "body_mass"), original[0]) - updated = pool.get_field(1, "body_mass") + np.testing.assert_array_equal(pool.expand("body_mass")[0], original[0]) + updated = pool.expand("body_mass")[1] np.testing.assert_allclose(updated[:base_body_id], original[1][:base_body_id]) np.testing.assert_allclose(updated[base_body_id], original[1][base_body_id] + delta[0]) np.testing.assert_allclose(updated[base_body_id + 1 :], original[1][base_body_id + 1 :]) @@ -267,61 +198,30 @@ def _fake_push_robots(force_range): np.testing.assert_allclose(called["force_range"], limit) def test_step_uses_xfrc_applied_for_interval_push(self, bkd, monkeypatch: pytest.MonkeyPatch): - mujoco = _mujoco_module() sampled = np.array([[0.5, -0.25, 0.1], [-0.1, 0.8, -0.4]], dtype=np.float64) limit = np.array([0.7, 0.3, 0.2], dtype=np.float64) - calls: list[dict[str, Any]] = [] monkeypatch.setattr(np.random, "uniform", lambda low, high, size: sampled.copy()) - def _fake_step( - initial_state, - *, - nstep, - control_spec=0, - control=None, - return_sensor=False, - **kwargs, - ): - calls.append( - { - "nstep": nstep, - "control_spec": control_spec, - "control": None if control is None else np.array(control, copy=True), - } - ) - state = np.array(initial_state, copy=True) - if return_sensor: - return state, np.array(bkd._sensor_data, copy=True) - return state - - monkeypatch.setattr(bkd._pool, "step", _fake_step) - bkd.apply_interval_randomization(IntervalRandomizationPlan(push_perturbation_limit=limit)) ctrl = np.zeros((NUM_ENVS, bkd.model.nu), dtype=np.float64) bkd.step(ctrl, nsteps=2) - bkd.step(ctrl, nsteps=1) expected_xfrc = np.zeros((NUM_ENVS, 6 * bkd.model.nbody), dtype=np.float64) start = 6 * bkd._base_body_id expected_xfrc[:, start : start + 3] = sampled * limit[None, :] - expected_xfrc_traj = np.broadcast_to( - expected_xfrc[:, None, :], - (NUM_ENVS, 2, 6 * bkd.model.nbody), - ) - expected_ctrl_traj = np.broadcast_to(ctrl[:, None, :], (NUM_ENVS, 2, bkd.model.nu)) - assert len(calls) == 2 - assert calls[0]["nstep"] == 2 - assert calls[0]["control_spec"] & int(mujoco.mjtState.mjSTATE_CTRL) - assert calls[0]["control_spec"] & int(mujoco.mjtState.mjSTATE_XFRC_APPLIED) - np.testing.assert_allclose(calls[0]["control"][:, :, : bkd.model.nu], expected_ctrl_traj) - np.testing.assert_allclose(calls[0]["control"][:, :, bkd.model.nu :], expected_xfrc_traj) + # The staged wrench reached the persistent channel and survives the + # call; the pending staging buffer is cleared. + np.testing.assert_allclose(np.asarray(bkd._xfrc_view).reshape(NUM_ENVS, -1), expected_xfrc) + np.testing.assert_allclose(bkd._pending_xfrc_applied, 0.0) - assert calls[1]["nstep"] == 1 - assert calls[1]["control_spec"] == int(mujoco.mjtState.mjSTATE_CTRL) - assert calls[1]["control"].shape == (NUM_ENVS, 1, bkd.model.nu) + # An idle second dispatch rewrites the channel absolutely with zeros, + # instead of leaving the previous wrench to decay on its own. + bkd.step(ctrl, nsteps=1) + np.testing.assert_allclose(np.asarray(bkd._xfrc_view).reshape(NUM_ENVS, -1), 0.0) + np.testing.assert_allclose(bkd._pending_xfrc_applied, 0.0) def test_interval_push_uses_configured_body(self, monkeypatch: pytest.MonkeyPatch): from unisim.backend.mujoco.backend import MuJoCoBackend @@ -337,44 +237,22 @@ def test_interval_push_uses_configured_body(self, monkeypatch: pytest.MonkeyPatc bkd.materialize() sampled = np.array([[0.5, -0.25, 0.1], [-0.1, 0.8, -0.4]], dtype=np.float64) limit = np.array([0.7, 0.3, 0.2], dtype=np.float64) - calls: list[dict[str, Any]] = [] monkeypatch.setattr(np.random, "uniform", lambda low, high, size: sampled.copy()) - def _fake_step( - initial_state, - *, - nstep, - control_spec=0, - control=None, - return_sensor=False, - **kwargs, - ): - calls.append( - { - "control_spec": control_spec, - "control": None if control is None else np.array(control, copy=True), - } - ) - state = np.array(initial_state, copy=True) - if return_sensor: - return state, np.array(bkd._sensor_data, copy=True) - return state - - monkeypatch.setattr(bkd._pool, "step", _fake_step) - bkd.apply_interval_randomization(IntervalRandomizationPlan(push_perturbation_limit=limit)) bkd.step(np.zeros((NUM_ENVS, bkd.model.nu), dtype=np.float64)) base_start = 6 * bkd._base_body_id push_start = 6 * mujoco.mj_name2id(bkd.model, mujoco.mjtObj.mjOBJ_BODY, "torso_link") - xfrc = calls[0]["control"][:, 0, bkd.model.nu :] + xfrc = np.asarray(bkd._xfrc_view).reshape(NUM_ENVS, -1) np.testing.assert_allclose(xfrc[:, push_start : push_start + 3], sampled * limit[None, :]) np.testing.assert_allclose(xfrc[:, base_start : base_start + 3], 0.0) + np.testing.assert_allclose(bkd._pending_xfrc_applied, 0.0) def test_set_state_body_iquat_randomization_only_affects_target_envs(self, bkd): pool = bkd._pool - original = [pool.get_field(i, "body_iquat").copy() for i in range(NUM_ENVS)] + original = pool.expand("body_iquat").copy() qpos = _identity_qpos_mujoco(bkd.model.nq) qvel = np.zeros((1, bkd.model.nv)) updated = original[1].reshape(bkd.model.nbody, 4).copy() @@ -387,14 +265,12 @@ def test_set_state_body_iquat_randomization_only_affects_target_envs(self, bkd): randomization=ResetRandomizationPayload(body_iquat=updated[None, :, :]), ) - np.testing.assert_array_equal(pool.get_field(0, "body_iquat"), original[0]) - np.testing.assert_allclose( - pool.get_field(1, "body_iquat").reshape(bkd.model.nbody, 4), updated - ) + np.testing.assert_array_equal(pool.expand("body_iquat")[0], original[0]) + np.testing.assert_allclose(pool.expand("body_iquat")[1], updated) def test_set_state_gravity_randomization_only_affects_target_envs(self, bkd): pool = bkd._pool - original = [pool.get_field(i, "gravity").copy() for i in range(NUM_ENVS)] + original = pool.expand("gravity").copy() qpos = _identity_qpos_mujoco(bkd.model.nq) qvel = np.zeros((1, bkd.model.nv)) updated = original[1].copy() @@ -407,12 +283,12 @@ def test_set_state_gravity_randomization_only_affects_target_envs(self, bkd): randomization=ResetRandomizationPayload(gravity=updated[None, :]), ) - np.testing.assert_array_equal(pool.get_field(0, "gravity"), original[0]) - np.testing.assert_allclose(pool.get_field(1, "gravity"), updated) + np.testing.assert_array_equal(pool.expand("gravity")[0], original[0]) + np.testing.assert_allclose(pool.expand("gravity")[1], updated) def test_set_state_body_inertia_randomization_only_affects_target_envs(self, bkd): pool = bkd._pool - original = [pool.get_field(i, "body_inertia").copy() for i in range(NUM_ENVS)] + original = pool.expand("body_inertia").copy() qpos = _identity_qpos_mujoco(bkd.model.nq) qvel = np.zeros((1, bkd.model.nv)) updated = original[1].reshape(bkd.model.nbody, 3).copy() @@ -425,14 +301,12 @@ def test_set_state_body_inertia_randomization_only_affects_target_envs(self, bkd randomization=ResetRandomizationPayload(body_inertia=updated[None, :, :]), ) - np.testing.assert_array_equal(pool.get_field(0, "body_inertia"), original[0]) - np.testing.assert_allclose( - pool.get_field(1, "body_inertia").reshape(bkd.model.nbody, 3), updated - ) + np.testing.assert_array_equal(pool.expand("body_inertia")[0], original[0]) + np.testing.assert_allclose(pool.expand("body_inertia")[1], updated) def test_set_state_dof_armature_randomization_only_affects_target_envs(self, bkd): pool = bkd._pool - original = [pool.get_field(i, "dof_armature").copy() for i in range(NUM_ENVS)] + original = pool.expand("dof_armature").copy() qpos = _identity_qpos_mujoco(bkd.model.nq) qvel = np.zeros((1, bkd.model.nv)) updated = original[1].copy() @@ -445,17 +319,17 @@ def test_set_state_dof_armature_randomization_only_affects_target_envs(self, bkd randomization=ResetRandomizationPayload(dof_armature=updated[None, :]), ) - np.testing.assert_array_equal(pool.get_field(0, "dof_armature"), original[0]) - np.testing.assert_allclose(pool.get_field(1, "dof_armature"), updated) + np.testing.assert_array_equal(pool.expand("dof_armature")[0], original[0]) + np.testing.assert_allclose(pool.expand("dof_armature")[1], updated) def test_set_state_kp_kd_randomization_only_affects_target_envs(self, bkd): pool = bkd._pool - original_kp = [pool.get_field(i, "kp").copy() for i in range(NUM_ENVS)] - original_kd = [pool.get_field(i, "kd").copy() for i in range(NUM_ENVS)] + original_gain = pool.expand("actuator_gainprm").copy() + original_bias = pool.expand("actuator_biasprm").copy() qpos = _identity_qpos_mujoco(bkd.model.nq) qvel = np.zeros((1, bkd.model.nv)) - new_kp = original_kp[1] + 1.25 - new_kd = np.maximum(original_kd[1] + 0.25, 0.25) + new_kp = original_gain[1, :, 0] + 1.25 + new_kd = np.maximum(-original_bias[1, :, 2] + 0.25, 0.25) bkd.set_state( np.array([1]), @@ -464,10 +338,13 @@ def test_set_state_kp_kd_randomization_only_affects_target_envs(self, bkd): randomization=ResetRandomizationPayload(kp=new_kp[None, :], kd=new_kd[None, :]), ) - np.testing.assert_array_equal(pool.get_field(0, "kp"), original_kp[0]) - np.testing.assert_array_equal(pool.get_field(0, "kd"), original_kd[0]) - np.testing.assert_allclose(pool.get_field(1, "kp"), new_kp) - np.testing.assert_allclose(pool.get_field(1, "kd"), new_kd) + gain = pool.expand("actuator_gainprm") + bias = pool.expand("actuator_biasprm") + np.testing.assert_array_equal(gain[0], original_gain[0]) + np.testing.assert_array_equal(bias[0], original_bias[0]) + np.testing.assert_allclose(gain[1, :, 0], new_kp) + np.testing.assert_allclose(-bias[1, :, 2], new_kd) + np.testing.assert_allclose(bias[1, :, 1], -new_kp) # base kinematics diff --git a/tests/base/test_sim_backend_set_state_timing.py b/tests/base/test_sim_backend_set_state_timing.py index fdd4f70bd..d0a313855 100644 --- a/tests/base/test_sim_backend_set_state_timing.py +++ b/tests/base/test_sim_backend_set_state_timing.py @@ -31,6 +31,10 @@ from unilab.base.scene import SceneCfg pytest.importorskip("mujoco", reason="mujoco not installed") +pytest.importorskip( + "unisim.backend.mujoco.backend", + reason="unisim-core MuJoCo adapter (mjbatch build) not available", +) NUM_ENVS = 2 @@ -95,7 +99,9 @@ def test_mujoco_set_state_returns_schema_conformant_timing() -> None: _assert_gap_bounded(timing) # MuJoCo-specific sub-keys must be populated on the mujoco path. assert timing["set_state_pool_reset_ms"] > 0.0 - assert timing["set_state_state_scatter_ms"] > 0.0 + # mjbatch executor: canonical state lives in the batch's bound views, so + # there is no host scatter — the key stays populated at 0.0 (#1554). + assert timing["set_state_state_scatter_ms"] == 0.0 # Motrix-only sub-keys report 0.0 on the mujoco backend. assert timing["set_state_mask_ms"] == 0.0 assert timing["set_state_data_slice_ms"] == 0.0 diff --git a/tests/base/test_sim_backend_smoke.py b/tests/base/test_sim_backend_smoke.py index 51d79a94a..a6227ab32 100644 --- a/tests/base/test_sim_backend_smoke.py +++ b/tests/base/test_sim_backend_smoke.py @@ -17,6 +17,10 @@ from unilab.base.scene import SceneCfg pytest.importorskip("mujoco", reason="mujoco not installed") +pytest.importorskip( + "unisim.backend.mujoco.backend", + reason="unisim-core MuJoCo adapter (mjbatch build) not available", +) def _xml(robot: str, scene: str = "scene_flat.xml") -> str: @@ -151,7 +155,12 @@ def test_mujoco_interval_root_velocity_kick_is_row_selective_and_refreshes_senso bkd.set_state(np.arange(NUM_ENVS, dtype=np.int32), qpos, qvel) bkd.step(np.zeros((NUM_ENVS, bkd.model.nu)), nsteps=1) assert bkd._pool is not None - bkd._sensor_data[:] = bkd._pool.forward(bkd.get_physics_state()) + # Refresh the derived fields (sensors one substep behind after step()). + bkd._pool.forward() + + # get_physics_state rows are [time, qpos, qvel]. + idx_qpos = 1 + idx_qvel = 1 + bkd.nq body_ids = bkd.get_body_ids(("base",)) state_before = bkd.get_physics_state().copy() @@ -171,18 +180,18 @@ def test_mujoco_interval_root_velocity_kick_is_row_selective_and_refreshes_senso np.testing.assert_array_equal(state_after[0], state_before[0]) np.testing.assert_allclose(state_after[1, 0], state_before[1, 0], atol=1e-12) np.testing.assert_allclose( - state_after[1, bkd._idx_qpos : bkd._idx_qvel], - state_before[1, bkd._idx_qpos : bkd._idx_qvel], + state_after[1, idx_qpos:idx_qvel], + state_before[1, idx_qpos:idx_qvel], atol=1e-12, ) np.testing.assert_allclose( - state_after[1, bkd._idx_qvel + 3 :], - state_before[1, bkd._idx_qvel + 3 :], + state_after[1, idx_qvel + 3 :], + state_before[1, idx_qvel + 3 :], atol=1e-12, ) np.testing.assert_allclose( - state_after[1, bkd._idx_qvel : bkd._idx_qvel + 3], - state_before[1, bkd._idx_qvel : bkd._idx_qvel + 3] + delta[1, 0], + state_after[1, idx_qvel : idx_qvel + 3], + state_before[1, idx_qvel : idx_qvel + 3] + delta[1, 0], atol=1e-7, ) np.testing.assert_allclose( @@ -284,15 +293,16 @@ def test_mujoco_interval_root_angular_velocity_kick_converts_to_body_frame(): ) state_after = bkd.get_physics_state() + idx_qvel = 1 + bkd.nq # The free-root qvel angular channels are body-frame: yaw(+90deg) maps the # world-frame +x kick onto body-frame -y. np.testing.assert_allclose( - state_after[0, bkd._idx_qvel + 3 : bkd._idx_qvel + 6], + state_after[0, idx_qvel + 3 : idx_qvel + 6], 0.0, atol=1e-12, ) np.testing.assert_allclose( - state_after[1, bkd._idx_qvel + 3 : bkd._idx_qvel + 6], + state_after[1, idx_qvel + 3 : idx_qvel + 6], [0.0, -0.5, 0.0], atol=1e-7, ) @@ -342,7 +352,8 @@ def test_mujoco_interval_body_force_and_torque_have_observable_effect(tmp_path): bkd.set_state(ids, qpos, qvel0) bkd.step(ctrl, nsteps=1) - baseline = bkd.get_physics_state()[:, bkd._idx_qvel : bkd._idx_qvel + 6].copy() + idx_qvel = 1 + bkd.nq + baseline = bkd.get_physics_state()[:, idx_qvel : idx_qvel + 6].copy() force = np.zeros((NUM_ENVS, 1, 3), dtype=np.float64) force[:, 0, 2] = 10.0 # N, world +z @@ -358,7 +369,7 @@ def test_mujoco_interval_body_force_and_torque_have_observable_effect(tmp_path): ) ) bkd.step(ctrl, nsteps=1) - kicked = bkd.get_physics_state()[:, bkd._idx_qvel : bkd._idx_qvel + 6] + kicked = bkd.get_physics_state()[:, idx_qvel : idx_qvel + 6] # Free fall is common to both runs, so the staged wrench adds exactly # F/m * dt to linear z and tau/I * dt to angular z. diff --git a/tests/config/test_config_system.py b/tests/config/test_config_system.py index 6102c15ed..4d7a47cbf 100644 --- a/tests/config/test_config_system.py +++ b/tests/config/test_config_system.py @@ -148,16 +148,6 @@ def test_task_files_keep_full_identity_without_hidden_backend_marker(): assert "sim_backend" in training_raw, f"task missing sim_backend: {path}" -def test_motrix_task_files_do_not_declare_post_step_forward_sensor(): - for path in sorted(CONF_DIR.glob("*/task/**/*motrix*.yaml")): - cfg = OmegaConf.load(path) - - assert OmegaConf.select(cfg, "env.post_step_forward_sensor") is None, ( - "post_step_forward_sensor is routed only to MuJoCo backends: " - f"{path.relative_to(CONF_DIR)}" - ) - - @pytest.mark.parametrize( "algo_dir,config_name,task,backend,task_file,overrides", _supported_task_cases(), @@ -288,30 +278,6 @@ def test_ppo_g1_backend_specific_hyperparams_remain_separate(): assert motrix_cfg.env.terminations.tilt.params.max_tilt_deg == pytest.approx(35.0) -@pytest.mark.parametrize( - ("algo_dir", "overrides"), - [ - ("ppo", ["task=g1_walk_flat/mujoco"]), - ("appo", ["task=g1_walk_flat/mujoco"]), - ("sac", ["task=g1_walk_flat/mujoco"]), - ("flashsac", ["task=g1_walk_flat/mujoco"]), - ], -) -def test_post_step_forward_sensor_defaults_false(algo_dir: str, overrides: list[str]): - cfg = _compose(algo_dir, overrides=overrides) - - assert cfg.env.post_step_forward_sensor is False - - -def test_mujoco_post_step_forward_sensor_can_be_overridden(): - override_cfg = _compose( - "ppo", - overrides=["task=g1_walk_flat/mujoco", "env.post_step_forward_sensor=true"], - ) - - assert override_cfg.env.post_step_forward_sensor is True - - def test_appo_adaptive_lr_factors_are_overridden_only_by_dex_hand_owners(): g1_cfg = _compose("appo", overrides=["task=g1_walk_flat/mujoco"]) allegro_cfg = _compose("appo", overrides=["task=allegro_inhand/mujoco"]) diff --git a/tests/config/test_locomotion_params.py b/tests/config/test_locomotion_params.py index 26d1fc045..7a0bc5e94 100644 --- a/tests/config/test_locomotion_params.py +++ b/tests/config/test_locomotion_params.py @@ -370,7 +370,6 @@ def test_ppo_go2_footstand_uses_hydra_owned_manager_task(): assert cfg.env.sim_dt == pytest.approx(0.004) assert cfg.env.ctrl_dt == pytest.approx(0.02) assert cfg.env.max_episode_seconds == pytest.approx(10.0) - assert cfg.env.adaptive_chunk_size is False assert cfg.env.observations.policy.terms.frame.history_length == 15 assert cfg.env.observations.critic.terms.frame.history_length == 15 assert cfg.env.actions.joint_pos.action_scale == pytest.approx(0.3) diff --git a/tests/envs/locomotion/a2/test_a2_joystick_contract.py b/tests/envs/locomotion/a2/test_a2_joystick_contract.py index c23d33e6f..dfcc3664c 100644 --- a/tests/envs/locomotion/a2/test_a2_joystick_contract.py +++ b/tests/envs/locomotion/a2/test_a2_joystick_contract.py @@ -290,10 +290,11 @@ def test_a2_registry_has_no_legacy_config_or_runtime_fallback() -> None: def test_a2_registry_executes_real_manager_runtime() -> None: pytest.importorskip("mujoco") - try: - from mujoco_uni.batch_env import BatchEnvPool # noqa: F401 - except Exception: - pytest.skip("mujoco_uni.batch_env not available") + pytest.importorskip("mjbatch", reason="mjbatch not installed") + pytest.importorskip( + "unisim.backend.mujoco.backend", + reason="unisim-core MuJoCo adapter (mjbatch build) not available", + ) hydra_cfg, _, env_override = _materialize() env = registry.make( diff --git a/tests/envs/locomotion/go2/test_manager_based_cfg.py b/tests/envs/locomotion/go2/test_manager_based_cfg.py index 10f975203..701add610 100644 --- a/tests/envs/locomotion/go2/test_manager_based_cfg.py +++ b/tests/envs/locomotion/go2/test_manager_based_cfg.py @@ -292,6 +292,11 @@ def test_go2_flat_registry_has_no_legacy_config_fallback() -> None: (("mujoco", "task=go2_joystick_flat/mujoco"), ("motrix", "task=go2_joystick_flat/motrix")), ) def test_go2_flat_registry_executes_real_manager_runtime(backend: str, owner: str) -> None: + if backend == "mujoco": + pytest.importorskip( + "unisim.backend.mujoco.backend", + reason="unisim-core MuJoCo adapter (mjbatch build) not available", + ) registry.ensure_registries() hydra_cfg, _, env_override = _materialize("ppo", (owner,)) env = registry.make( diff --git a/tests/envs/locomotion/go2w/test_go2w_manager_based_flat_cfg.py b/tests/envs/locomotion/go2w/test_go2w_manager_based_flat_cfg.py index 1adca205d..8ce0176f1 100644 --- a/tests/envs/locomotion/go2w/test_go2w_manager_based_flat_cfg.py +++ b/tests/envs/locomotion/go2w/test_go2w_manager_based_flat_cfg.py @@ -266,6 +266,10 @@ def test_go2w_flat_registry_executes_real_manager_runtime(backend: str) -> None: def test_go2w_flat_dead_zone_and_motor_gain_overrides_are_manager_owned() -> None: + pytest.importorskip( + "unisim.backend.mujoco.backend", + reason="unisim-core MuJoCo adapter (mjbatch build) not available", + ) _, env_cfg, _ = _materialize("mujoco") command = env_cfg.commands["twist"] command.ranges.lin_vel_x = (0.1, 0.1) @@ -288,6 +292,10 @@ def test_go2w_flat_dead_zone_and_motor_gain_overrides_are_manager_owned() -> Non def test_go2w_flat_incomplete_motor_selection_fails_closed() -> None: + pytest.importorskip( + "unisim.backend.mujoco.backend", + reason="unisim-core MuJoCo adapter (mjbatch build) not available", + ) _, env_cfg, _ = _materialize("mujoco") action = env_cfg.actions["motor"] assert isinstance(action, Go2WMixedActionCfg) diff --git a/tests/envs/locomotion/test_go2_footstand.py b/tests/envs/locomotion/test_go2_footstand.py index 22e119f82..5b820f272 100644 --- a/tests/envs/locomotion/test_go2_footstand.py +++ b/tests/envs/locomotion/test_go2_footstand.py @@ -205,7 +205,6 @@ def test_footstand_owner_materializes_complete_plain_manager_cfg( assert env_cfg.sim_dt == pytest.approx(0.004) assert env_cfg.ctrl_dt == pytest.approx(0.02) assert env_cfg.max_episode_seconds == pytest.approx(10.0) - assert env_cfg.adaptive_chunk_size is False assert env_cfg.policy_observation_group == "policy" assert env_cfg.critic_observation_group == "critic" @@ -302,6 +301,11 @@ def test_footstand_registry_has_no_legacy_config_or_factory() -> None: def test_footstand_real_runtime_preserves_history_action_and_partial_reset( backend: str, ) -> None: + if backend == "mujoco": + pytest.importorskip( + "unisim.backend.mujoco.backend", + reason="unisim-core MuJoCo adapter (mjbatch build) not available", + ) registry.ensure_registries() env = _make_env(backend) try: @@ -363,6 +367,10 @@ def test_footstand_real_runtime_preserves_history_action_and_partial_reset( def test_footstand_termination_uses_grace_boundary() -> None: + pytest.importorskip( + "unisim.backend.mujoco.backend", + reason="unisim-core MuJoCo adapter (mjbatch build) not available", + ) registry.ensure_registries() env = _make_env("mujoco") try: @@ -385,6 +393,10 @@ def test_footstand_termination_uses_grace_boundary() -> None: def test_footstand_reward_clips_aggregate_before_dt_scaling() -> None: + pytest.importorskip( + "unisim.backend.mujoco.backend", + reason="unisim-core MuJoCo adapter (mjbatch build) not available", + ) registry.ensure_registries() env = _make_env("mujoco") try: diff --git a/tests/envs/mdp/test_joint_position_action.py b/tests/envs/mdp/test_joint_position_action.py index 98d08e79a..fb57af871 100644 --- a/tests/envs/mdp/test_joint_position_action.py +++ b/tests/envs/mdp/test_joint_position_action.py @@ -217,6 +217,11 @@ def test_relative_joint_position_action_rejects_nonzero_offsets() -> None: def test_go2_joint_targets_are_mapped_to_backend_control_order(backend_type: str) -> None: if backend_type == "motrix": pytest.importorskip("motrixsim") + else: + pytest.importorskip( + "unisim.backend.mujoco.backend", + reason="unisim-core MuJoCo adapter (mjbatch build) not available", + ) joint_names = ( "FL_hip_joint", "FL_thigh_joint", diff --git a/tests/envs/test_env_configs.py b/tests/envs/test_env_configs.py index 4f24ecd40..01e6d9674 100644 --- a/tests/envs/test_env_configs.py +++ b/tests/envs/test_env_configs.py @@ -24,10 +24,11 @@ def _require_mujoco_runtime() -> None: pytest.importorskip("mujoco", reason="mujoco not installed") - try: - from mujoco_uni.batch_env import BatchEnvPool as _ # noqa: F401 - except Exception: - pytest.skip("mujoco_uni.batch_env not available (platform/libstdc++ issue)") + pytest.importorskip("mjbatch", reason="mjbatch not installed") + pytest.importorskip( + "unisim.backend.mujoco.backend", + reason="unisim-core MuJoCo adapter (mjbatch build) not available", + ) def _allegro_manager_override( diff --git a/tests/envs/test_manager_based_rl_env.py b/tests/envs/test_manager_based_rl_env.py index 0cb4af7c4..60f050bc4 100644 --- a/tests/envs/test_manager_based_rl_env.py +++ b/tests/envs/test_manager_based_rl_env.py @@ -938,6 +938,10 @@ def test_named_keyframe_snapshot_is_shared_by_entity_and_reset_cold_path() -> No def test_real_mujoco_backend_is_materialized_before_first_reset() -> None: + pytest.importorskip( + "unisim.backend.mujoco.backend", + reason="unisim-core MuJoCo adapter (mjbatch build) not available", + ) scene = SceneCfg( model_file=str(ASSETS_ROOT_PATH / "robots" / "go2" / "scene_flat.xml"), entities={"robot": EntityCfg(root_body_name="base")}, @@ -990,6 +994,10 @@ def test_real_mujoco_default_state_matches_qpos0_or_named_home( expected_root_z: float, expected_joint_pos: np.ndarray, ) -> None: + pytest.importorskip( + "unisim.backend.mujoco.backend", + reason="unisim-core MuJoCo adapter (mjbatch build) not available", + ) joint_names = ( "FL_hip_joint", "FL_thigh_joint", diff --git a/tests/envs/test_motion_command_partial_reset.py b/tests/envs/test_motion_command_partial_reset.py index 76584cbe2..6956b75de 100644 --- a/tests/envs/test_motion_command_partial_reset.py +++ b/tests/envs/test_motion_command_partial_reset.py @@ -115,10 +115,11 @@ def test_motion_command_partial_reset_row_parity( config_root: str, task: str, identity: str ) -> None: pytest.importorskip("mujoco") - try: - from mujoco_uni.batch_env import BatchEnvPool as _ # noqa: F401 - except Exception: - pytest.skip("mujoco_uni.batch_env not available") + pytest.importorskip("mjbatch", reason="mjbatch not installed") + pytest.importorskip( + "unisim.backend.mujoco.backend", + reason="unisim-core MuJoCo adapter (mjbatch build) not available", + ) num_envs = 4 env = _make_env(config_root, task, "mujoco", identity, num_envs) diff --git a/tests/envs/test_motion_profiles.py b/tests/envs/test_motion_profiles.py index 3ea2ac7bc..35a9892dd 100644 --- a/tests/envs/test_motion_profiles.py +++ b/tests/envs/test_motion_profiles.py @@ -508,10 +508,11 @@ def test_representative_motion_profiles_reset_and_step( ) -> None: if backend == "mujoco": pytest.importorskip("mujoco") - try: - from mujoco_uni.batch_env import BatchEnvPool as _ # noqa: F401 - except Exception: - pytest.skip("mujoco_uni.batch_env not available") + pytest.importorskip("mjbatch", reason="mjbatch not installed") + pytest.importorskip( + "unisim.backend.mujoco.backend", + reason="unisim-core MuJoCo adapter (mjbatch build) not available", + ) else: pytest.importorskip("motrixsim") diff --git a/tests/envs/test_observation_partial_reset.py b/tests/envs/test_observation_partial_reset.py index fddab62ee..fcee8c458 100644 --- a/tests/envs/test_observation_partial_reset.py +++ b/tests/envs/test_observation_partial_reset.py @@ -50,10 +50,11 @@ def _make_env(config_root: str, task: str, backend: str, identity: str, num_envs def test_observation_partial_reset_row_contract() -> None: pytest.importorskip("mujoco") - try: - from mujoco_uni.batch_env import BatchEnvPool as _ # noqa: F401 - except Exception: - pytest.skip("mujoco_uni.batch_env not available") + pytest.importorskip("mjbatch", reason="mjbatch not installed") + pytest.importorskip( + "unisim.backend.mujoco.backend", + reason="unisim-core MuJoCo adapter (mjbatch build) not available", + ) num_envs = 4 env = _make_env("sac", "g1_motion_tracking", "mujoco", "G1MotionTrackingSAC", num_envs) diff --git a/tests/managers/test_isaac_lab_migration_fixture.py b/tests/managers/test_isaac_lab_migration_fixture.py index 6043d5537..79c52b240 100644 --- a/tests/managers/test_isaac_lab_migration_fixture.py +++ b/tests/managers/test_isaac_lab_migration_fixture.py @@ -81,6 +81,10 @@ def test_fixture_hydra_owner_materializes_source_order_as_plain_manager_cfg() -> def test_fixture_real_mujoco_reset_step_and_partial_reset() -> None: + pytest.importorskip( + "unisim.backend.mujoco.backend", + reason="unisim-core MuJoCo adapter (mjbatch build) not available", + ) _, _, override = _materialize_fixture() env = registry.make( FIXTURE_ENV_NAME, @@ -124,6 +128,10 @@ def test_fixture_real_mujoco_reset_step_and_partial_reset() -> None: def test_fixture_missing_actuator_fails_during_cold_path_binding() -> None: + pytest.importorskip( + "unisim.backend.mujoco.backend", + reason="unisim-core MuJoCo adapter (mjbatch build) not available", + ) _, _, override = _materialize_fixture() override["actions"]["joint_effort"]["actuator_names"] = ["missing_actuator"] diff --git a/tests/managers/test_mjlab_migration_fixture.py b/tests/managers/test_mjlab_migration_fixture.py index baefc8d59..085cb2129 100644 --- a/tests/managers/test_mjlab_migration_fixture.py +++ b/tests/managers/test_mjlab_migration_fixture.py @@ -70,6 +70,10 @@ def test_mjlab_fixture_hydra_materializes_source_structure() -> None: def test_mjlab_fixture_real_mujoco_reset_step_and_reward() -> None: + pytest.importorskip( + "unisim.backend.mujoco.backend", + reason="unisim-core MuJoCo adapter (mjbatch build) not available", + ) _, override = _materialize() env = registry.make( FIXTURE_ENV_NAME, @@ -112,6 +116,10 @@ def test_mjlab_fixture_real_mujoco_reset_step_and_reward() -> None: def test_mjlab_fixture_missing_actuator_fails_on_cold_path() -> None: + pytest.importorskip( + "unisim.backend.mujoco.backend", + reason="unisim-core MuJoCo adapter (mjbatch build) not available", + ) _, override = _materialize() override["actions"]["effort"]["actuator_names"] = ["missing_actuator"] with pytest.raises(ValueError, match="regular expressions matched.*missing_actuator"): diff --git a/tests/managers/test_stage_curriculum_demo.py b/tests/managers/test_stage_curriculum_demo.py index f65d40801..517c6f393 100644 --- a/tests/managers/test_stage_curriculum_demo.py +++ b/tests/managers/test_stage_curriculum_demo.py @@ -62,6 +62,10 @@ def test_stage_curriculum_demo_owner_materializes_declared_ladders() -> None: def test_stage_curriculum_demo_runtime_ramps_with_step_counter() -> None: + pytest.importorskip( + "unisim.backend.mujoco.backend", + reason="unisim-core MuJoCo adapter (mjbatch build) not available", + ) _, override = _materialize() env = registry.make( FIXTURE_ENV_NAME, @@ -102,6 +106,10 @@ def test_stage_curriculum_demo_runtime_ramps_with_step_counter() -> None: def test_stage_curriculum_demo_owner_invalid_stages_fail_closed() -> None: + pytest.importorskip( + "unisim.backend.mujoco.backend", + reason="unisim-core MuJoCo adapter (mjbatch build) not available", + ) _, override = _materialize() override["curriculum"]["smooth_reward_weight"]["params"]["stages"] = [ {"step": 4, "weight": 0.5}, diff --git a/tests/scripts/test_train_script_configs.py b/tests/scripts/test_train_script_configs.py index 229e1c9bd..85a8a7618 100644 --- a/tests/scripts/test_train_script_configs.py +++ b/tests/scripts/test_train_script_configs.py @@ -15,6 +15,10 @@ import torch pytest.importorskip("mujoco") +pytest.importorskip( + "unisim.backend.mujoco.backend", + reason="unisim-core MuJoCo adapter (mjbatch build) not available", +) ROOT_DIR = Path(__file__).resolve().parents[2] diff --git a/tests/scripts/test_train_scripts.py b/tests/scripts/test_train_scripts.py index bdce53684..b5c34430c 100644 --- a/tests/scripts/test_train_scripts.py +++ b/tests/scripts/test_train_scripts.py @@ -558,18 +558,6 @@ def test_build_ppo_env_cfg_override_carries_motrix_max_iterations_override( assert env_cfg_override["motrix_max_iterations"] == 9 -def test_build_ppo_env_cfg_override_carries_post_step_forward_sensor_override( - monkeypatch: pytest.MonkeyPatch, -): - mod = _train_rsl_rl(monkeypatch) - for value in (True, False): - cfg = _ppo_cfg(["task=g1_walk_flat/mujoco", f"env.post_step_forward_sensor={value}"]) - - env_cfg_override = mod.build_ppo_env_cfg_override(cfg) - - assert env_cfg_override["post_step_forward_sensor"] is value - - def test_offpolicy_g1_walk_flat_motrix_env_cfg_override_disables_pd_gains(): cfg = _offpolicy_cfg(["task=g1_walk_flat/motrix"]) diff --git a/tests/test_cli.py b/tests/test_cli.py index ce5bac489..2a3b936b9 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -224,7 +224,7 @@ def test_eval_mujoco_interactive_routes_to_dedicated_viewer( monkeypatch.setattr( cli, "find_spec", - lambda name: ModuleSpec(name, loader=None) if name in {"mujoco", "mujoco_uni"} else None, + lambda name: ModuleSpec(name, loader=None) if name in {"mujoco", "mjbatch"} else None, ) command = cli.build_command( @@ -270,7 +270,7 @@ def test_eval_mujoco_interactive_honors_profile_and_render_override( monkeypatch.setattr( cli, "find_spec", - lambda name: ModuleSpec(name, loader=None) if name in {"mujoco", "mujoco_uni"} else None, + lambda name: ModuleSpec(name, loader=None) if name in {"mujoco", "mjbatch"} else None, ) command = cli.build_command( @@ -398,7 +398,7 @@ def test_eval_mujoco_interactive_preserves_explicit_action_mode( monkeypatch.setattr( cli, "find_spec", - lambda name: ModuleSpec(name, loader=None) if name in {"mujoco", "mujoco_uni"} else None, + lambda name: ModuleSpec(name, loader=None) if name in {"mujoco", "mjbatch"} else None, ) command = cli.build_command( @@ -554,7 +554,7 @@ def test_eval_mujoco_interactive_falls_back_to_sibling_owner( monkeypatch.setattr( cli, "find_spec", - lambda name: ModuleSpec(name, loader=None) if name in {"mujoco", "mujoco_uni"} else None, + lambda name: ModuleSpec(name, loader=None) if name in {"mujoco", "mjbatch"} else None, ) command = cli.build_command( diff --git a/tests/test_cli_runtime_requirements.py b/tests/test_cli_runtime_requirements.py index 572361b8c..de8fbb7af 100644 --- a/tests/test_cli_runtime_requirements.py +++ b/tests/test_cli_runtime_requirements.py @@ -12,12 +12,12 @@ def test_check_runtime_requirements_requires_mujoco_extra(monkeypatch: pytest.Mo cli._check_runtime_requirements("ppo", "mujoco") -def test_check_runtime_requirements_mujoco_needs_the_uni_runtime( +def test_check_runtime_requirements_mujoco_needs_mjbatch( monkeypatch: pytest.MonkeyPatch, ) -> None: # Plain `mujoco` can arrive via other extras (e.g. superdex); the MuJoCo - # physics backend is only usable with the mujoco-uni-runtime binding. - monkeypatch.setattr(cli, "find_spec", lambda name: None if name == "mujoco_uni" else object()) + # physics backend is only usable with the mjbatch batch engine. + monkeypatch.setattr(cli, "find_spec", lambda name: None if name == "mjbatch" else object()) with pytest.raises(SystemExit, match="sim=mujoco requires the MuJoCo extra"): cli._check_runtime_requirements("ppo", "mujoco") diff --git a/tests/training/test_training_helpers.py b/tests/training/test_training_helpers.py index c0afecf8d..47bf735a3 100644 --- a/tests/training/test_training_helpers.py +++ b/tests/training/test_training_helpers.py @@ -9,7 +9,11 @@ from hydra.core.global_hydra import GlobalHydra from omegaconf import OmegaConf -pytest.importorskip("mujoco_uni") +pytest.importorskip("mjbatch") +pytest.importorskip( + "unisim.backend.mujoco.backend", + reason="unisim-core MuJoCo adapter (mjbatch build) not available", +) from unisim.backend.base import RenderClosedError from unisim.backend.motrix.backend import MotrixBackend @@ -764,9 +768,14 @@ def _render_states_get_frames(state_list, model_file, **kwargs): assert captured["fps"] == 20 -def test_render_play_mode_uses_visualized_per_env_playback_models_for_video_export( +def test_render_play_mode_uses_visual_playback_model_for_video_export( monkeypatch: pytest.MonkeyPatch, tmp_path: Path ): + """Video export resolves ONE visual model for every env (#1554). + + The mjbatch executor has no per-env model variants, so playback renders + against a single visual model file instead of a per-env list. + """ import mujoco captured: dict[str, object] = {} @@ -794,22 +803,13 @@ def __init__(self): {"ctrl_dt": 0.05, "scene": SceneCfg(model_file=str(visual_model_path))}, )() self.snapshot_calls = 0 - self._models = [ - mujoco.MjModel.from_xml_string( - "" - ), - mujoco.MjModel.from_xml_string( - "" - ), - ] def get_physics_state_snapshot(self) -> np.ndarray: self.snapshot_calls += 1 return np.full((2, 2), self.snapshot_calls, dtype=np.float32) def get_playback_model(self, env_index: int | None = None): - idx = 0 if env_index is None else int(env_index) - return self._models[idx] + raise AssertionError("one visual model serves every env; no per-env lookup") def run_playback(self, **kwargs): kwargs = _resolve_low_level_playback_flags(kwargs) @@ -821,21 +821,12 @@ def _render_states_get_frames(state_list, model_file, **kwargs): del kwargs captured["states"] = state_list captured["model_file"] = model_file - assert isinstance(model_file, list) - model0 = mujoco.MjModel.from_binary_path(model_file[0]) - model1 = mujoco.MjModel.from_binary_path(model_file[1]) - object0 = mujoco.mj_name2id(model0, mujoco.mjtObj.mjOBJ_GEOM, "object") - object1 = mujoco.mj_name2id(model1, mujoco.mjtObj.mjOBJ_GEOM, "object") - hand0 = mujoco.mj_name2id(model0, mujoco.mjtObj.mjOBJ_GEOM, "hand_geom") - hand1 = mujoco.mj_name2id(model1, mujoco.mjtObj.mjOBJ_GEOM, "hand_geom") - ground0 = mujoco.mj_name2id(model0, mujoco.mjtObj.mjOBJ_GEOM, "ground") - ground1 = mujoco.mj_name2id(model1, mujoco.mjtObj.mjOBJ_GEOM, "ground") - captured["object0_size"] = model0.geom_size[object0].copy() - captured["object1_size"] = model1.geom_size[object1].copy() - captured["hand0_size"] = model0.geom_size[hand0].copy() - captured["hand1_size"] = model1.geom_size[hand1].copy() - captured["ground0_size"] = model0.geom_size[ground0].copy() - captured["ground1_size"] = model1.geom_size[ground1].copy() + assert isinstance(model_file, str) + rendered = mujoco.MjModel.from_xml_path(model_file) + hand = mujoco.mj_name2id(rendered, mujoco.mjtObj.mjOBJ_GEOM, "hand_geom") + ground = mujoco.mj_name2id(rendered, mujoco.mjtObj.mjOBJ_GEOM, "ground") + captured["hand_size"] = rendered.geom_size[hand].copy() + captured["ground_size"] = rendered.geom_size[ground].copy() return [np.zeros((2, 2, 3), dtype=np.uint8)] monkeypatch.setattr( @@ -862,13 +853,9 @@ def _render_states_get_frames(state_list, model_file, **kwargs): assert result == str(output_path) assert env.snapshot_calls == 2 - assert isinstance(captured["model_file"], list) - model_files = captured["model_file"] - assert len(model_files) == 2 - np.testing.assert_allclose(captured["object0_size"], [0.1, 0.1, 0.1]) - np.testing.assert_allclose(captured["object1_size"], [0.2, 0.2, 0.2]) - np.testing.assert_allclose(captured["hand0_size"], captured["hand1_size"]) - np.testing.assert_allclose(captured["ground0_size"], captured["ground1_size"]) + assert captured["model_file"] == str(visual_model_path) + np.testing.assert_allclose(captured["hand_size"], [0.05, 0.05, 0.05]) + np.testing.assert_allclose(captured["ground_size"], [2.0, 2.0, 0.1]) def test_render_play_mode_requires_env_snapshot_contract_for_video_export(tmp_path: Path): diff --git a/uv.lock b/uv.lock index 09d845fde..88e9b37b8 100644 --- a/uv.lock +++ b/uv.lock @@ -2134,6 +2134,31 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/37/8c/52f0299f1675cdfa1ab39a6028a2e5adf9032ae1118c9895c84b08af162b/mediapy-1.2.6-py3-none-any.whl", hash = "sha256:0a0ea00eb0da83c3c54d588b49c49a41ba456174aa33e530ffe13e17269c9072", size = 27494, upload-time = "2026-02-03T10:29:30.245Z" }, ] +[[package]] +name = "mjbatch-uni" +version = "0.1.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mujoco" }, + { 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'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/bd/5f/284d7e603033b977b825cbb73282f9233d071a4b43da1ef98335fae610db/mjbatch_uni-0.1.0.tar.gz", hash = "sha256:e2856c55cb6179d256360aa56c76c53207cfaf6dee5ba6130bd83fd9b0778a6d", size = 28999, upload-time = "2026-09-13T06:49:56.381Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7f/b5/d737ce697d40763b32d4c3ddf7a2188a2f1abd6bc2a8fb2396006048d6f9/mjbatch_uni-0.1.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:40f787b39f3a4c4214bc41d6928fa6d3ea7737d012dd72e64c39c533fdaaf197", size = 140421, upload-time = "2026-09-13T06:49:38.906Z" }, + { url = "https://files.pythonhosted.org/packages/cc/18/1ec66faa8c51e72c84696c7432cd17ab45c32834e9fea580488d10c202a2/mjbatch_uni-0.1.0-cp310-cp310-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:07310303cc9a4f67ba64965325a4456783556586aa69a0f18756ae686cdf8430", size = 164953, upload-time = "2026-09-13T06:49:40.537Z" }, + { url = "https://files.pythonhosted.org/packages/12/b8/f04152b71065d6ba14363f8e66bd2f0ade429af058b19dc75ae691dd8c58/mjbatch_uni-0.1.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bf8ab61861d59d4e537ddca47c1e5df1fc5e80c5a81d753f7d7608a528a99dfb", size = 175346, upload-time = "2026-09-13T06:49:42.052Z" }, + { url = "https://files.pythonhosted.org/packages/4d/a5/734641dfb1d442b3c53c9024b050be59652a00d6ac2d830502350fcf6233/mjbatch_uni-0.1.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:b4b02e5c3829225e697207719b22061004c255f332f724cdc220fc2c633f233e", size = 140003, upload-time = "2026-09-13T06:49:43.424Z" }, + { url = "https://files.pythonhosted.org/packages/fb/59/11885b99ef0797e1a286a5edd5017fec074a6a145c6dd8d59599cc6bda95/mjbatch_uni-0.1.0-cp311-cp311-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:cb2fd277ea5f9c045d263c79425ea92fbce0bc7b5d79e648dfd4ea18d1501282", size = 164584, upload-time = "2026-09-13T06:49:44.675Z" }, + { url = "https://files.pythonhosted.org/packages/af/12/9ca93e335f9a6f7d0e62e8eff24b8771d73ed4c38273c8950856133d3ddf/mjbatch_uni-0.1.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0476b1dd333d95118123c9b4c05d6e42dc652cb8c69fc5b0960aed1dbf0ec09e", size = 175091, upload-time = "2026-09-13T06:49:46.473Z" }, + { url = "https://files.pythonhosted.org/packages/76/d1/5f46d25aa730bc2dbf1e90a5cfc9c0744d19a4be766ff79988db9bf0e518/mjbatch_uni-0.1.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:23bb9f101fcb02c04482eea4285fedccea19450550902151c76a7739491cf54c", size = 139197, upload-time = "2026-09-13T06:49:47.923Z" }, + { url = "https://files.pythonhosted.org/packages/32/ee/8308873ed1a1697a44454ebe9f03a3b2c962aa0d1786ed5c79939a5a7d57/mjbatch_uni-0.1.0-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e3c65e2a0936bb210f4b6c5229512490e07ac8e7be42cc14e464b5d4cef93de2", size = 163419, upload-time = "2026-09-13T06:49:49.311Z" }, + { url = "https://files.pythonhosted.org/packages/9b/7f/9b5d8f6c51abb8a3f02d9a1f4bd2d49e7b74a238921840cf9ab747b8372f/mjbatch_uni-0.1.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:25fbf20cbb9e72bc98e37c1a65ae621d9cea5904d3e3a8e8aad5bbcd59faa1ed", size = 174676, upload-time = "2026-09-13T06:49:50.994Z" }, + { url = "https://files.pythonhosted.org/packages/4b/95/13657cf8311dbb21d78ca219c07675899a45fa27f5eedfb6e56711fb5452/mjbatch_uni-0.1.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:87296ee8012d4435e1652d80e4679d5196ecd196c0ed05bb23601490bb2e1428", size = 139260, upload-time = "2026-09-13T06:49:52.406Z" }, + { url = "https://files.pythonhosted.org/packages/ee/1a/4d456c39d58b581df6b157f3f2df59cd9a5a0a876cda7f0765d152c34bc2/mjbatch_uni-0.1.0-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d6cb335cd3c745228c33d372b6e95115d9fe68249ebb8dc09fbc78de000e1a27", size = 163355, upload-time = "2026-09-13T06:49:53.777Z" }, + { url = "https://files.pythonhosted.org/packages/67/da/d5a2d73db2a21d884615a62c169733c0ce7ec24a3633515bd8a5ebb95ab9/mjbatch_uni-0.1.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:838e207355095053f70d571803de2b7958a95cafed9e7df19c6aa2cbf3ef135a", size = 174664, upload-time = "2026-09-13T06:49:55.173Z" }, +] + [[package]] name = "ml-dtypes" version = "0.5.4" @@ -2299,31 +2324,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/bb/25/2a518a33b359b0798a9e8b55f30a4c9ab0cadd0cab8d85872ff9da6ca309/mujoco-3.11.0-cp313-cp313-win_amd64.whl", hash = "sha256:d48b4a18d409ce4c746f32b3b402a8419b7450d9a89de68b81d73a3ef333d0e4", size = 15792231, upload-time = "2026-07-28T01:23:05.798Z" }, ] -[[package]] -name = "mujoco-uni-runtime" -version = "0.5.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "mujoco" }, - { 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'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/e9/ce/476cf50ec4f76b86b855b861e01f106b330c7d62a61574655cfe90336a07/mujoco_uni_runtime-0.5.0.tar.gz", hash = "sha256:a0c774da3607b6111d7c91155867377f9eb6d91ab6e0bbb97f15247df1ead9d7", size = 55656, upload-time = "2026-09-05T19:29:40.265Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/68/ee/61348995acb39aef26b5a1d9af57d61d6a2d042002063ab1b7bf28b0fdc8/mujoco_uni_runtime-0.5.0-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:755146199fcf8958e5dd3b74c32d93419502a6dd89751c445c5fe70cf68a0847", size = 401266, upload-time = "2026-09-05T19:29:23.366Z" }, - { url = "https://files.pythonhosted.org/packages/87/23/7ff0ccefecbbcfbaffce7e5959ec35023c354813546424eca4ccdc12cabd/mujoco_uni_runtime-0.5.0-cp310-cp310-manylinux_2_35_aarch64.whl", hash = "sha256:3a4d8c9471298b0c30a5526d49fe623aa8a8682001e208ea4bd3421318a9de3d", size = 2332548, upload-time = "2026-09-05T19:29:24.831Z" }, - { url = "https://files.pythonhosted.org/packages/d8/ed/c438dd170d7a6bbbbbc31cf6188ec94bd6e8cdc384e1e1619c6c24e27d7d/mujoco_uni_runtime-0.5.0-cp310-cp310-manylinux_2_35_x86_64.whl", hash = "sha256:d00015521610a39a473a15ef350f015914f1b82372a4eae1f580934be44bbb94", size = 2392108, upload-time = "2026-09-05T19:29:26.258Z" }, - { url = "https://files.pythonhosted.org/packages/45/a9/3cc6a7e31140bb48749a64165698358af4422a8b948293397be4621b81a9/mujoco_uni_runtime-0.5.0-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:f86a5bc75e1758b05816777b41bb3232015ce1a26d1f63bdb2da4f3b7ae87030", size = 403670, upload-time = "2026-09-05T19:29:27.892Z" }, - { url = "https://files.pythonhosted.org/packages/de/ed/a66e735512317bd135b2b9c806a882ff9f01f8503ce53b59874a6f4c685e/mujoco_uni_runtime-0.5.0-cp311-cp311-manylinux_2_35_aarch64.whl", hash = "sha256:746caa1ed23237ba4a0e2d7d7afb6afa0e3c6223b8000f22bc0c6da64a538fef", size = 2346314, upload-time = "2026-09-05T19:29:29.198Z" }, - { url = "https://files.pythonhosted.org/packages/46/ad/f8edda197f41b2dee678faca8964cf307e6bfbdc7b2bea6b907f5bf5fdf8/mujoco_uni_runtime-0.5.0-cp311-cp311-manylinux_2_35_x86_64.whl", hash = "sha256:3ccc49063dc34f3d657ffa7afe12c62d6a88615edcb2cc21b60ec59272dea3e3", size = 2406269, upload-time = "2026-09-05T19:29:30.476Z" }, - { url = "https://files.pythonhosted.org/packages/dc/fa/fe3e5173e804ac4eb1acbbd0fc9df81fe079fab7ae3d9813cd92cdf79b23/mujoco_uni_runtime-0.5.0-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:9597385aaafc77ae7d6f89aedfb677ad9f67a0874afd5962330c08cebbad56b6", size = 396860, upload-time = "2026-09-05T19:29:32.052Z" }, - { url = "https://files.pythonhosted.org/packages/7d/9d/2b14e54f2a0c70fdd95091f2648f9c50203bc87ac238370f5f51f9dac170/mujoco_uni_runtime-0.5.0-cp312-cp312-manylinux_2_35_aarch64.whl", hash = "sha256:61539c6287ca62600836062aefeb458e550f4ea6782d8f24e501027824aac546", size = 2366290, upload-time = "2026-09-05T19:29:33.248Z" }, - { url = "https://files.pythonhosted.org/packages/0a/25/7633b376c3cadffd35420f7f13a1a35ddfed86cae8842cf1206c3555b516/mujoco_uni_runtime-0.5.0-cp312-cp312-manylinux_2_35_x86_64.whl", hash = "sha256:a0117a107319661dccc6657e3cc0d2545f58eb0f9b0c725acba1d4047f507847", size = 2427155, upload-time = "2026-09-05T19:29:34.677Z" }, - { url = "https://files.pythonhosted.org/packages/de/33/c4c2d2fb73b43c8a4e331fe065d4f49f1ae70b17d4327255ccb733f917d9/mujoco_uni_runtime-0.5.0-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:1831182949ace71f48e9ee7c9ef7406b4b1ef2bba9884c87e4932ba92862041b", size = 396959, upload-time = "2026-09-05T19:29:36.208Z" }, - { url = "https://files.pythonhosted.org/packages/e7/1c/ff0cb5a53311f94200a1ecd9645d9e48442b994bb866a04e5c52f3f28537/mujoco_uni_runtime-0.5.0-cp313-cp313-manylinux_2_35_aarch64.whl", hash = "sha256:9d22628e59ae205610e1eeba2429c21801cf8ea25d99de6d66733621c44086c2", size = 2365046, upload-time = "2026-09-05T19:29:37.558Z" }, - { url = "https://files.pythonhosted.org/packages/56/83/df50eae1b379198bf0c2c58d36f768afe30db7dad823b9210d21c6399eb0/mujoco_uni_runtime-0.5.0-cp313-cp313-manylinux_2_35_x86_64.whl", hash = "sha256:84f1c81819752622cb794458a731d8b1b98ec2a589380fedf55999b3fe89faed", size = 2427855, upload-time = "2026-09-05T19:29:38.944Z" }, -] - [[package]] name = "mujoco-warp" version = "3.11.0" @@ -3511,15 +3511,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/98/9e/a112df5cfd5a68cb1d9fc31cfe38c28d5aec9f10865ce37ecef2e4450873/pyarrow-25.0.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:6943e2fe7954d29d84de45d29d34c8dc36ce96570e67d89aa9976e650a4a9138", size = 53144784, upload-time = "2026-08-10T12:39:20.503Z" }, ] -[[package]] -name = "pybind11" -version = "3.0.4" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/cc/f0/35145a3c3baffeef55d4b8324caa33abaa8fa56ab345ecd4b2211d09163e/pybind11-3.0.4.tar.gz", hash = "sha256:3286b59c8a774b9ee650169302dd5a4eedc30a8617905a0560dd8ee44775130c", size = 589533, upload-time = "2026-04-19T03:08:15.925Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/b3/06/c3a23c9a0263b136c519f033a58d4641e73065fefc7754e9667ec206d992/pybind11-3.0.4-py3-none-any.whl", hash = "sha256:961720ee652da51d531b7b2451a6bd2bc042b0106e6d9baa48ecb7d58034ce63", size = 314166, upload-time = "2026-04-19T03:08:14.091Z" }, -] - [[package]] name = "pycollada" version = "0.9.3" @@ -5177,10 +5168,8 @@ motrix = [ { name = "motrixsim-core" }, ] mujoco = [ + { name = "mjbatch-uni" }, { name = "mujoco" }, - { name = "mujoco-uni-runtime" }, - { name = "pybind11" }, - { name = "wheel" }, ] newton = [ { name = "imgui-bundle" }, @@ -5221,11 +5210,11 @@ requires-dist = [ { name = "imgui-bundle", marker = "extra == 'newton'", specifier = ">=1.92.0" }, { name = "lark", specifier = ">=1.3.1" }, { name = "mediapy" }, + { name = "mjbatch-uni", marker = "extra == 'mujoco'", specifier = "~=0.1.0" }, { name = "motrixsim-core", marker = "extra == 'motrix'", specifier = "==0.8.2" }, { name = "mujoco", marker = "extra == 'drake'", specifier = ">=3.5" }, { name = "mujoco", marker = "extra == 'mujoco'", specifier = "~=3.11.0" }, { name = "mujoco", marker = "extra == 'newton'", specifier = "==3.11.0" }, - { name = "mujoco-uni-runtime", marker = "extra == 'mujoco'", specifier = "==0.5.0" }, { name = "mujoco-warp", marker = "extra == 'mjwarp'", specifier = "~=3.11.0" }, { name = "mujoco-warp", marker = "extra == 'newton'", specifier = "==3.11.0" }, { name = "newton", marker = "extra == 'newton'", specifier = "==1.5.1" }, @@ -5236,7 +5225,6 @@ requires-dist = [ { name = "onnxruntime", marker = "python_full_version >= '3.11'", specifier = ">=1.20" }, { name = "packaging" }, { name = "prettytable", specifier = ">=3.10" }, - { name = "pybind11", marker = "extra == 'mujoco'", specifier = ">=2.12" }, { name = "pyglet", marker = "extra == 'newton'", specifier = ">=2.1.6,<3" }, { name = "rich" }, { name = "rsl-rl-lib", specifier = ">=5.0.0" }, @@ -5252,13 +5240,12 @@ requires-dist = [ { name = "trimesh", marker = "extra == 'viser'", specifier = ">=3.21.7" }, { name = "typing-extensions" }, { name = "unilab-rl", specifier = "==1.2.0" }, - { name = "unisim-core", specifier = ">=1.2.0" }, + { name = "unisim-core", specifier = ">=1.2.1" }, { name = "unisim-core", extras = ["superdex"], marker = "python_full_version >= '3.12' and platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'superdex'" }, { name = "viser", marker = "extra == 'viser'", specifier = ">=1.0.26" }, { name = "wandb" }, { name = "warp-lang", marker = "extra == 'mjwarp'", specifier = "==1.16.0" }, { name = "warp-lang", marker = "extra == 'newton'", specifier = "==1.16.0" }, - { name = "wheel", marker = "extra == 'mujoco'" }, ] provides-extras = ["drake", "mujoco", "mjwarp", "newton", "motrix", "genesis", "viser", "superdex"] @@ -5297,13 +5284,13 @@ wheels = [ [[package]] name = "unisim-core" -version = "1.2.0" +version = "1.2.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { 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'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/bb/b7/2cf8626884236d9e14256fa4330b9e04a8a41598867ced1ac4030f827989/unisim_core-1.2.0.tar.gz", hash = "sha256:fccd67edbde98eeb16d2e6f9f5f8ce608dd2790266a673788897cbd9a42fa4ce", size = 253069, upload-time = "2026-09-10T05:19:05.388Z" } +sdist = { url = "https://files.pythonhosted.org/packages/8c/46/cfc76e05dc6770c7296849ebb2d91aba63ae5c137a964597b5d8c5396015/unisim_core-1.2.1.tar.gz", hash = "sha256:0fe214134e66cf0a06913689ed9061e0301217f51a245571b883541bd5860206", size = 250011, upload-time = "2026-09-13T07:02:51.499Z" } [package.optional-dependencies] superdex = [ @@ -5545,18 +5532,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/93/8c/2e650f2afeb7ee576912636c23ddb621c91ac6a98e66dc8d29c3c69446e1/werkzeug-3.1.8-py3-none-any.whl", hash = "sha256:63a77fb8892bf28ebc3178683445222aa500e48ebad5ec77b0ad80f8726b1f50", size = 226459, upload-time = "2026-04-02T18:49:12.72Z" }, ] -[[package]] -name = "wheel" -version = "0.47.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "packaging" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/39/62/75f18a0f03b4219c456652c7780e4d749b929eb605c098ce3a5b6b6bc081/wheel-0.47.0.tar.gz", hash = "sha256:cc72bd1009ba0cf63922e28f94d9d83b920aa2bb28f798a31d0691b02fa3c9b3", size = 63854, upload-time = "2026-04-22T15:51:27.727Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/87/1b/9e33c09813d65e248f7f773119148a612516a4bea93e9c6f545f78455b7c/wheel-0.47.0-py3-none-any.whl", hash = "sha256:212281cab4dff978f6cedd499cd893e1f620791ca6ff7107cf270781e587eced", size = 32218, upload-time = "2026-04-22T15:51:26.296Z" }, -] - [[package]] name = "wrapt" version = "2.4.0" diff --git a/uv.rocm.lock b/uv.rocm.lock index e88a53bc7..06b35c9cc 100644 --- a/uv.rocm.lock +++ b/uv.rocm.lock @@ -1684,6 +1684,31 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/37/8c/52f0299f1675cdfa1ab39a6028a2e5adf9032ae1118c9895c84b08af162b/mediapy-1.2.6-py3-none-any.whl", hash = "sha256:0a0ea00eb0da83c3c54d588b49c49a41ba456174aa33e530ffe13e17269c9072", size = 27494, upload-time = "2026-02-03T10:29:30.245Z" }, ] +[[package]] +name = "mjbatch-uni" +version = "0.1.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mujoco" }, + { 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'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/bd/5f/284d7e603033b977b825cbb73282f9233d071a4b43da1ef98335fae610db/mjbatch_uni-0.1.0.tar.gz", hash = "sha256:e2856c55cb6179d256360aa56c76c53207cfaf6dee5ba6130bd83fd9b0778a6d", size = 28999, upload-time = "2026-09-13T06:49:56.381Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7f/b5/d737ce697d40763b32d4c3ddf7a2188a2f1abd6bc2a8fb2396006048d6f9/mjbatch_uni-0.1.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:40f787b39f3a4c4214bc41d6928fa6d3ea7737d012dd72e64c39c533fdaaf197", size = 140421, upload-time = "2026-09-13T06:49:38.906Z" }, + { url = "https://files.pythonhosted.org/packages/cc/18/1ec66faa8c51e72c84696c7432cd17ab45c32834e9fea580488d10c202a2/mjbatch_uni-0.1.0-cp310-cp310-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:07310303cc9a4f67ba64965325a4456783556586aa69a0f18756ae686cdf8430", size = 164953, upload-time = "2026-09-13T06:49:40.537Z" }, + { url = "https://files.pythonhosted.org/packages/12/b8/f04152b71065d6ba14363f8e66bd2f0ade429af058b19dc75ae691dd8c58/mjbatch_uni-0.1.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bf8ab61861d59d4e537ddca47c1e5df1fc5e80c5a81d753f7d7608a528a99dfb", size = 175346, upload-time = "2026-09-13T06:49:42.052Z" }, + { url = "https://files.pythonhosted.org/packages/4d/a5/734641dfb1d442b3c53c9024b050be59652a00d6ac2d830502350fcf6233/mjbatch_uni-0.1.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:b4b02e5c3829225e697207719b22061004c255f332f724cdc220fc2c633f233e", size = 140003, upload-time = "2026-09-13T06:49:43.424Z" }, + { url = "https://files.pythonhosted.org/packages/fb/59/11885b99ef0797e1a286a5edd5017fec074a6a145c6dd8d59599cc6bda95/mjbatch_uni-0.1.0-cp311-cp311-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:cb2fd277ea5f9c045d263c79425ea92fbce0bc7b5d79e648dfd4ea18d1501282", size = 164584, upload-time = "2026-09-13T06:49:44.675Z" }, + { url = "https://files.pythonhosted.org/packages/af/12/9ca93e335f9a6f7d0e62e8eff24b8771d73ed4c38273c8950856133d3ddf/mjbatch_uni-0.1.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0476b1dd333d95118123c9b4c05d6e42dc652cb8c69fc5b0960aed1dbf0ec09e", size = 175091, upload-time = "2026-09-13T06:49:46.473Z" }, + { url = "https://files.pythonhosted.org/packages/76/d1/5f46d25aa730bc2dbf1e90a5cfc9c0744d19a4be766ff79988db9bf0e518/mjbatch_uni-0.1.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:23bb9f101fcb02c04482eea4285fedccea19450550902151c76a7739491cf54c", size = 139197, upload-time = "2026-09-13T06:49:47.923Z" }, + { url = "https://files.pythonhosted.org/packages/32/ee/8308873ed1a1697a44454ebe9f03a3b2c962aa0d1786ed5c79939a5a7d57/mjbatch_uni-0.1.0-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e3c65e2a0936bb210f4b6c5229512490e07ac8e7be42cc14e464b5d4cef93de2", size = 163419, upload-time = "2026-09-13T06:49:49.311Z" }, + { url = "https://files.pythonhosted.org/packages/9b/7f/9b5d8f6c51abb8a3f02d9a1f4bd2d49e7b74a238921840cf9ab747b8372f/mjbatch_uni-0.1.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:25fbf20cbb9e72bc98e37c1a65ae621d9cea5904d3e3a8e8aad5bbcd59faa1ed", size = 174676, upload-time = "2026-09-13T06:49:50.994Z" }, + { url = "https://files.pythonhosted.org/packages/4b/95/13657cf8311dbb21d78ca219c07675899a45fa27f5eedfb6e56711fb5452/mjbatch_uni-0.1.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:87296ee8012d4435e1652d80e4679d5196ecd196c0ed05bb23601490bb2e1428", size = 139260, upload-time = "2026-09-13T06:49:52.406Z" }, + { url = "https://files.pythonhosted.org/packages/ee/1a/4d456c39d58b581df6b157f3f2df59cd9a5a0a876cda7f0765d152c34bc2/mjbatch_uni-0.1.0-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d6cb335cd3c745228c33d372b6e95115d9fe68249ebb8dc09fbc78de000e1a27", size = 163355, upload-time = "2026-09-13T06:49:53.777Z" }, + { url = "https://files.pythonhosted.org/packages/67/da/d5a2d73db2a21d884615a62c169733c0ce7ec24a3633515bd8a5ebb95ab9/mjbatch_uni-0.1.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:838e207355095053f70d571803de2b7958a95cafed9e7df19c6aa2cbf3ef135a", size = 174664, upload-time = "2026-09-13T06:49:55.173Z" }, +] + [[package]] name = "ml-dtypes" version = "0.5.4" @@ -1830,31 +1855,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/bb/25/2a518a33b359b0798a9e8b55f30a4c9ab0cadd0cab8d85872ff9da6ca309/mujoco-3.11.0-cp313-cp313-win_amd64.whl", hash = "sha256:d48b4a18d409ce4c746f32b3b402a8419b7450d9a89de68b81d73a3ef333d0e4", size = 15792231, upload-time = "2026-07-28T01:23:05.798Z" }, ] -[[package]] -name = "mujoco-uni-runtime" -version = "0.5.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "mujoco" }, - { 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'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/e9/ce/476cf50ec4f76b86b855b861e01f106b330c7d62a61574655cfe90336a07/mujoco_uni_runtime-0.5.0.tar.gz", hash = "sha256:a0c774da3607b6111d7c91155867377f9eb6d91ab6e0bbb97f15247df1ead9d7", size = 55656, upload-time = "2026-09-05T19:29:40.265Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/68/ee/61348995acb39aef26b5a1d9af57d61d6a2d042002063ab1b7bf28b0fdc8/mujoco_uni_runtime-0.5.0-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:755146199fcf8958e5dd3b74c32d93419502a6dd89751c445c5fe70cf68a0847", size = 401266, upload-time = "2026-09-05T19:29:23.366Z" }, - { url = "https://files.pythonhosted.org/packages/87/23/7ff0ccefecbbcfbaffce7e5959ec35023c354813546424eca4ccdc12cabd/mujoco_uni_runtime-0.5.0-cp310-cp310-manylinux_2_35_aarch64.whl", hash = "sha256:3a4d8c9471298b0c30a5526d49fe623aa8a8682001e208ea4bd3421318a9de3d", size = 2332548, upload-time = "2026-09-05T19:29:24.831Z" }, - { url = "https://files.pythonhosted.org/packages/d8/ed/c438dd170d7a6bbbbbc31cf6188ec94bd6e8cdc384e1e1619c6c24e27d7d/mujoco_uni_runtime-0.5.0-cp310-cp310-manylinux_2_35_x86_64.whl", hash = "sha256:d00015521610a39a473a15ef350f015914f1b82372a4eae1f580934be44bbb94", size = 2392108, upload-time = "2026-09-05T19:29:26.258Z" }, - { url = "https://files.pythonhosted.org/packages/45/a9/3cc6a7e31140bb48749a64165698358af4422a8b948293397be4621b81a9/mujoco_uni_runtime-0.5.0-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:f86a5bc75e1758b05816777b41bb3232015ce1a26d1f63bdb2da4f3b7ae87030", size = 403670, upload-time = "2026-09-05T19:29:27.892Z" }, - { url = "https://files.pythonhosted.org/packages/de/ed/a66e735512317bd135b2b9c806a882ff9f01f8503ce53b59874a6f4c685e/mujoco_uni_runtime-0.5.0-cp311-cp311-manylinux_2_35_aarch64.whl", hash = "sha256:746caa1ed23237ba4a0e2d7d7afb6afa0e3c6223b8000f22bc0c6da64a538fef", size = 2346314, upload-time = "2026-09-05T19:29:29.198Z" }, - { url = "https://files.pythonhosted.org/packages/46/ad/f8edda197f41b2dee678faca8964cf307e6bfbdc7b2bea6b907f5bf5fdf8/mujoco_uni_runtime-0.5.0-cp311-cp311-manylinux_2_35_x86_64.whl", hash = "sha256:3ccc49063dc34f3d657ffa7afe12c62d6a88615edcb2cc21b60ec59272dea3e3", size = 2406269, upload-time = "2026-09-05T19:29:30.476Z" }, - { url = "https://files.pythonhosted.org/packages/dc/fa/fe3e5173e804ac4eb1acbbd0fc9df81fe079fab7ae3d9813cd92cdf79b23/mujoco_uni_runtime-0.5.0-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:9597385aaafc77ae7d6f89aedfb677ad9f67a0874afd5962330c08cebbad56b6", size = 396860, upload-time = "2026-09-05T19:29:32.052Z" }, - { url = "https://files.pythonhosted.org/packages/7d/9d/2b14e54f2a0c70fdd95091f2648f9c50203bc87ac238370f5f51f9dac170/mujoco_uni_runtime-0.5.0-cp312-cp312-manylinux_2_35_aarch64.whl", hash = "sha256:61539c6287ca62600836062aefeb458e550f4ea6782d8f24e501027824aac546", size = 2366290, upload-time = "2026-09-05T19:29:33.248Z" }, - { url = "https://files.pythonhosted.org/packages/0a/25/7633b376c3cadffd35420f7f13a1a35ddfed86cae8842cf1206c3555b516/mujoco_uni_runtime-0.5.0-cp312-cp312-manylinux_2_35_x86_64.whl", hash = "sha256:a0117a107319661dccc6657e3cc0d2545f58eb0f9b0c725acba1d4047f507847", size = 2427155, upload-time = "2026-09-05T19:29:34.677Z" }, - { url = "https://files.pythonhosted.org/packages/de/33/c4c2d2fb73b43c8a4e331fe065d4f49f1ae70b17d4327255ccb733f917d9/mujoco_uni_runtime-0.5.0-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:1831182949ace71f48e9ee7c9ef7406b4b1ef2bba9884c87e4932ba92862041b", size = 396959, upload-time = "2026-09-05T19:29:36.208Z" }, - { url = "https://files.pythonhosted.org/packages/e7/1c/ff0cb5a53311f94200a1ecd9645d9e48442b994bb866a04e5c52f3f28537/mujoco_uni_runtime-0.5.0-cp313-cp313-manylinux_2_35_aarch64.whl", hash = "sha256:9d22628e59ae205610e1eeba2429c21801cf8ea25d99de6d66733621c44086c2", size = 2365046, upload-time = "2026-09-05T19:29:37.558Z" }, - { url = "https://files.pythonhosted.org/packages/56/83/df50eae1b379198bf0c2c58d36f768afe30db7dad823b9210d21c6399eb0/mujoco_uni_runtime-0.5.0-cp313-cp313-manylinux_2_35_x86_64.whl", hash = "sha256:84f1c81819752622cb794458a731d8b1b98ec2a589380fedf55999b3fe89faed", size = 2427855, upload-time = "2026-09-05T19:29:38.944Z" }, -] - [[package]] name = "mypy" version = "1.20.0" @@ -2668,15 +2668,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/8e/37/efad0257dc6e593a18957422533ff0f87ede7c9c6ea010a2177d738fb82f/pure_eval-0.2.3-py3-none-any.whl", hash = "sha256:1db8e35b67b3d218d818ae653e27f06c3aa420901fa7b081ca98cbedc874e0d0", size = 11842, upload-time = "2024-07-21T12:58:20.04Z" }, ] -[[package]] -name = "pybind11" -version = "3.0.4" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/cc/f0/35145a3c3baffeef55d4b8324caa33abaa8fa56ab345ecd4b2211d09163e/pybind11-3.0.4.tar.gz", hash = "sha256:3286b59c8a774b9ee650169302dd5a4eedc30a8617905a0560dd8ee44775130c", size = 589533, upload-time = "2026-04-19T03:08:15.925Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/b3/06/c3a23c9a0263b136c519f033a58d4641e73065fefc7754e9667ec206d992/pybind11-3.0.4-py3-none-any.whl", hash = "sha256:961720ee652da51d531b7b2451a6bd2bc042b0106e6d9baa48ecb7d58034ce63", size = 314166, upload-time = "2026-04-19T03:08:14.091Z" }, -] - [[package]] name = "pycollada" version = "0.9.3" @@ -3751,10 +3742,8 @@ motrix = [ { name = "motrixsim-core" }, ] mujoco = [ + { name = "mjbatch-uni" }, { name = "mujoco" }, - { name = "mujoco-uni-runtime" }, - { name = "pybind11" }, - { name = "wheel" }, ] viser = [ { name = "trimesh" }, @@ -3781,15 +3770,14 @@ requires-dist = [ { name = "imageio-ffmpeg", specifier = ">=0.6.0" }, { name = "lark", specifier = ">=1.3.1" }, { name = "mediapy" }, + { name = "mjbatch-uni", marker = "extra == 'mujoco'", specifier = "~=0.1.0" }, { name = "motrixsim-core", marker = "extra == 'motrix'", specifier = "==0.8.2" }, { name = "mujoco", marker = "extra == 'mujoco'", specifier = "~=3.11.0" }, - { name = "mujoco-uni-runtime", marker = "extra == 'mujoco'", specifier = "==0.5.0" }, { name = "ninja", marker = "sys_platform == 'linux'" }, { name = "numpy" }, { name = "onnxruntime", marker = "python_full_version < '3.11'", specifier = "<1.20" }, { name = "onnxruntime", marker = "python_full_version >= '3.11'", specifier = ">=1.20" }, { name = "packaging" }, - { name = "pybind11", marker = "extra == 'mujoco'", specifier = ">=2.12" }, { name = "rich" }, { name = "rsl-rl-lib", specifier = ">=5.0.0" }, { name = "setuptools", specifier = "<70" }, @@ -3805,7 +3793,6 @@ requires-dist = [ { name = "unisim-core", specifier = ">=1.2.0" }, { name = "viser", marker = "extra == 'viser'", specifier = ">=1.0.26" }, { name = "wandb" }, - { name = "wheel", marker = "extra == 'mujoco'" }, ] provides-extras = ["drake", "mujoco", "motrix", "viser"] @@ -4023,18 +4010,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/93/8c/2e650f2afeb7ee576912636c23ddb621c91ac6a98e66dc8d29c3c69446e1/werkzeug-3.1.8-py3-none-any.whl", hash = "sha256:63a77fb8892bf28ebc3178683445222aa500e48ebad5ec77b0ad80f8726b1f50", size = 226459, upload-time = "2026-04-02T18:49:12.72Z" }, ] -[[package]] -name = "wheel" -version = "0.47.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "packaging" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/39/62/75f18a0f03b4219c456652c7780e4d749b929eb605c098ce3a5b6b6bc081/wheel-0.47.0.tar.gz", hash = "sha256:cc72bd1009ba0cf63922e28f94d9d83b920aa2bb28f798a31d0691b02fa3c9b3", size = 63854, upload-time = "2026-04-22T15:51:27.727Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/87/1b/9e33c09813d65e248f7f773119148a612516a4bea93e9c6f545f78455b7c/wheel-0.47.0-py3-none-any.whl", hash = "sha256:212281cab4dff978f6cedd499cd893e1f620791ca6ff7107cf270781e587eced", size = 32218, upload-time = "2026-04-22T15:51:26.296Z" }, -] - [[package]] name = "xxhash" version = "3.6.0"