feat: mujoco 后端执行器切换到 mjbatch-uni(roadmap #1552) - #1562
Merged
Merged
Conversation
* feat: switch MuJoCo backend dependency from mujoco-uni-runtime to mjbatch (#1553) Replace the mujoco-uni-runtime dependency (import mujoco_uni) with the mjbatch native batch engine across the repository, clean up the dead config knobs it owned, and repoint every test guard at the new module so the repo is ready for the #1554 validation gate. Roadmap: #1552. Closes #1553. * pyproject.toml / pyproject.rocm.toml: mujoco extra now pins mjbatch @ git+https://github.com/unilabsim/mjbatch.git@473bba9 (integration fork commit carrying the cpu_ids / hfield-yaw / native-callback patches); drop pybind11/wheel, the sdist build-fallback comments, and [tool.uv] no-build-isolation-package. mujoco~=3.11.0 stays. The final distribution identity (PyPI name vs git pin / prebuilt wheels) is the roadmap's open maintainer item. * uv.lock / uv.rocm.lock regenerated and committed (CI sets UV_FROZEN=1); uv sync --extra mujoco --extra motrix --locked verified. * Makefile: delete the mujoco MJ= version-switch target, the check-cxx-toolchain preflight, and setup-mujoco; plain setup/sync remain. * src/unilab/cli.py: the sim=mujoco runtime gate now probes mjbatch (plain mujoco still arrives via the superdex extra; fail-closed SystemExit kept). * EnvCfg/Hydra cleanup: delete post_step_forward_sensor, adaptive_chunk_size, chunk_size and the bench_nsteps backend kwarg from EnvCfg, env_backend_kwargs, the five algo owner configs, and the task-level overrides (go2_footstand, go2_joystick_flat/superdex, fr3_joint_target/superdex). cpu_ids stays. * pyright excludes: drop the stale mujoco-uni-runtime/MjSpec rationales; src/unilab has no such usage. terrains/ no longer needs an exclude (verified 0 errors); envs/ exclusion kept with an accurate reason. * Benchmarks: benchmark_mujoco_backend_step_detail and benchmark_mujoco_pool_thread_scaling ported to the mjbatch live-view API; benchmark_forward_reset_methods and benchmark_mujoco_single_dispatch_callback deleted (chunk-size studies / answered by the native callback); benchmark_np_env_shard_throughput drops the removed adaptive_chunk_size knob. * Tests: swap every mujoco_uni import guard to mjbatch/backend-availability equivalents so suites skip cleanly until the new unisim-core adapter (unisim#58) lands; delete test_mujoco_batch_env_randomization.py, test_mujoco_batch_env_jacobian.py (direct mujoco_uni API tests, replaced by adapter-level tests) and test_mujoco_chunk_size_wiring.py (chunk wiring no longer exists); delete the config tests asserting the removed EnvCfg fields. * Docs (en + zh_CN): installation, MuJoCo backend page (version switching now = mjbatch rebuild), DR contract (expand/set_const instead of set_field_indexed), domain randomization, backend swap, superdex, CLI reference, contributing, api reference, and a changelog entry recording the dependency swap, the accepted-drift statement (characterized by #1554), and the variants/Windows/chunk_size deprecations. Numeric equivalence before/after the swap is not guaranteed; the accepted drift is characterized by the #1554 drift baseline. * chore: repin mjbatch to e6c19ba (sample_hfield semantics fix, mjbatch#21)
) * feat(tools): capture mujoco_uni drift-characterization baseline (#1554) Roadmap #1552 replaces the mujoco backend's native executor (mujoco-uni-runtime) with mjbatch. Record the BEFORE half of the drift baseline while the repo still runs on mujoco_uni: fixed model/seed/action-sequence trajectories as a regression reference (not a bit-exact gate). scripts/tools/capture_mujoco_drift_baseline.py builds each task through the trainer's Hydra owner path (conf/ppo -> BackendAdapter -> registry.make), pins executor determinism (seed, cpu_ids, chunk_size, adaptive tuning off), drives a recorded seeded action sequence (iid uniform plus saturated held envs so termination, autoreset, and backend set_state are exercised), and records per-step obs groups, rewards, done flags, and qpos/qvel via the public SimBackend.get_state interface. Covered tasks (mujoco backend, 8 envs, 300 steps, seed 42, action seed 1234): - Go2WJoystickFlat: per-substep state-feedback control path (Go2WMixedAction via set_pre_step_control), the highest-risk migration surface. - Go2JoystickFlat: plain position-action path on the same robot. Two identical runs produce bitwise-identical arrays (verified via per-array SHA-256 recorded in metadata). AFTER capture: re-run with --output scripts/tools/drift_baseline/after post-switch. * test: rewire mujoco backend tests to the mjbatch executor (#1554) Deferred numeric/sentinel rewrites from #1555: - test_backend_pre_step_control.py: replace the fake mujoco_uni BatchEnvPool with a fake mjbatch Batch implementing the step-with-callback protocol; assert single dispatch, k=0 no-refresh, per-k ctrl writes, xfrc absolute write + pending clear, callback_sensordata=False. Drop the deleted post_step_forward_sensor flag test. - test_sim_backend.py: delete the two per-env model-variant tests (mjbatch has no model variants); rewrite the six DR tests to read pool.expand(...) views and the two xfrc interval-push tests to spy on the new step path (_xfrc_view contents + _pending_xfrc_applied clear). - test_mujoco_cpu_affinity_wiring.py: drop the deleted adaptive_chunk_size kwarg; reduce pool pinning tests to ctor/wiring assertions (mjbatch pins workers at Batch construction and exposes no per-worker introspection); move the hot-path-no-XML-parse test here from the deleted chunk-wiring file. - test_mujoco_site_jacobian.py: build the serial reference from the new state accessors (_time_view/_qpos_view/_qvel_view) instead of FULLPHYSICS rows and per-env model variants. - test_sim_backend_smoke.py: drop _idx_qpos/_idx_qvel reads (removed internals) and the old pool.forward(state) sensor refresh. - test_sim_backend_set_state_timing.py: set_state_state_scatter_ms stays populated at 0.0 — the mjbatch executor keeps canonical state in bound views, so there is no host scatter. - new test_mujoco_backend_sentinel.py: unguarded fast-lane sentinels that fail (never skip) when the mjbatch+adapter guard would silently disable the gated suites, plus a real one-env materialize/step under that guard. * test: AFTER drift baseline + characterization report for the mjbatch swap (#1554) - capture_mujoco_drift_baseline.py: drop the deleted chunk_size/ adaptive_chunk_size overrides and metadata entries (knobs no longer exist in EnvCfg); record mjbatch instead of mujoco-uni-runtime in METADATA_PACKAGES; keep cpu_ids and every other BEFORE knob identical. - New compare_mujoco_drift_baseline.py: per-array max/mean abs diff, reference magnitude, first-divergence step + magnitude at that step, rendered as a markdown characterization report (not a pass/fail gate). - AFTER artifacts: scripts/tools/drift_baseline/after/*, captured with the updated script on the mjbatch executor (mujoco 3.11.0, mjbatch fork, unisim adapter) after fixing cold-path sensor-view staleness in the adapter (bind_sensor_data now resolves views at read time). - drift_report.md: reset state and the recorded action sequence are bit-identical; both tasks diverge from float roundoff at step 0/1 (<=6e-5 qvel, <=5e-7 qpos) and chaotically amplify over 300 steps on the toppling envs (qvel up to ~1.1e1, qpos ~1.1, reward <=8.1e-2 mean 4.2e-3). * chore: ruff format drift comparison script * test: rewrite per-env playback-model export test to the single-model contract The mjbatch executor has no per-env model variants: playback video export resolves ONE visual model file for every env. Rewrite test_render_play_mode_uses_visualized_per_env_playback_models_for_video_export to assert the new contract (single str model file, visual scene geometry); the per-env variant assertions tested the deleted machinery. * test: refresh AFTER drift artifacts on the final adapter stack Re-captured after the adapter set_state reorder (reset-then-upload) and the materialize_visual_playback_model restore; trajectories and drift magnitudes are unchanged (the fix only affects the dropped-identical- upload corner case), confirming the characterization is stable. * docs: record rsl_rl PPO training smoke results for the mjbatch gate (#1554)
- _FakeMjBatch 回调协议简化为 fn(k, state, ctrl):删除 callback_sensordata/steps_done/stop_on_warning 形参与相关断言, 保留单子步 dispatch 与 sensordata 落后一个子步的 copy-out 语义。 - mjbatch 钉版 e6c19ba -> cf4a83d(消融后 fork),pyproject.toml 与 pyproject.rocm.toml 同步,双 lockfile 重新生成。 - 文档扫尾:changelog 记录消融后的最终 API 形态;安装文档(en/zh_CN) 的 mjbatch 示例钉版更新至 cf4a83d;training_smoke_1554.md 注明回调 已精简为 fn(k, state, ctrl)。 - 复验:make check 绿;pytest -m 'not slow' 1587 passed / 26 skipped (与基线一致);benchmark smoke 32/33 + 33/34(mlx 平台跳过); 漂移 AFTER 重采集与已提交工件 bit-identical,drift_report.md 数值不变。 Closes #1557
… (#1561) unisim-core 1.2.1 is the first released line whose mujoco adapter executes on the unilabsim mjbatch fork, and the fork is published on PyPI as mjbatch-uni 0.1.0 (cp310-cp313 wheels for linux x86_64/aarch64 and macOS arm64). Replace the interim git direct reference in the mujoco extra with mjbatch-uni~=0.1.0 and raise the unisim-core floor to >=1.2.1 in both manifests; regenerate uv.lock and uv.rocm.lock. This retires the local-checkout development setup used during the migration.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Proposes the complete roadmap #1552 migration to
main: the mujoco backend executor moves frommujoco-uni-runtime(BatchEnvPool) to the unilabsim mjbatch fork, with the whole dependency chain served from PyPI. All commits were developed and reviewed on the integration branchdev/issue-1552-mujoco-mjbatch; merge is left to the maintainer.Included PRs (4 commits):
feat: mujoco 后端依赖切换 mjbatch + 配置清理 (feat: mujoco 后端依赖切换 mjbatch + 配置清理 (roadmap #1552 child #6) #1553)test: mujoco 后端测试重指收口 + 漂移基线 + 全量验证 gate (Dev/mujoco mlx #7)test: 跟随 mjbatch API 消融精简 fake pool 并复验漂移基线 (test: 跟随 mjbatch API 消融精简调整 fake/sentinel(roadmap #1552 后续) #1557)chore(deps): follow unisim-core 1.2.1 and mjbatch-uni from PyPI (roadmap: unisim mujoco 后端批执行器从 mujoco-uni-runtime 替换为 mjbatch #1552)Dependency chain after merge (all PyPI, no git pins, no local checkouts):
unisim-core>=1.2.1— mujoco adapter executes on the mjbatch fork (unisim.backend.mujoco)mjbatch-uni~=0.1.0— the fork engine, published with cp310-cp313 wheels for linux x86_64/aarch64 and macOS arm64; upstream history aligned with Accept any history array numpy considers C-contiguous. kevinzakka/mjbatch#3Key behavior notes:
jac_site/sample_hfieldrun as native live-state query ops (A/B benchmark vs user-land assembly: bit-exact, performance parity)post_step_forward_sensorremoved end to endValidation
669d69a5(pure PyPI environment, zero editable/local packages):make check— ruff / mypy / pyright all passmake test— 1587 passed, 26 skippedmake test-all— exit 0 (script-mode 33/34, 1 platform-optional mlx skip)Refs #1552.