MJX

用 JAX 实现的 MuJoCo，在 GPU/TPU 上批量并行仿真，可求导。

MJX（MuJoCo XLA）是 Google DeepMind 随 MuJoCo 3.0 在 2023 年 10 月推出的 JAX 接口：用 JAX 重新实现 MuJoCo 物理，经 XLA 编译器编译后能跑在 NVIDIA 和 AMD 显卡、Apple 芯片和谷歌 TPU 上。配合 JAX 的 vmap 可以把几千上万个场景打包一起算，适合大规模强化学习；但只跑一个场景时可能比原版 MuJoCo 慢 10 倍左右。纯 JAX 版支持自动求导，可做可微仿真。现在 MJX 有两个后端：MJX-JAX 和调用 MuJoCo Warp 的 MJX-Warp，后者在 NVIDIA 显卡和接触多的场景上更快，但不能求导。Brax 训练库和 MuJoCo Playground 都建在它之上。

例子典型三步：mjx.put_model 把模型放上加速器，mjx.make_data 建状态，再用 jax.vmap(mjx.step) 让几千个机器人同时前进一步，接 Brax 的 PPO 训练行走策略。

也叫

MuJoCo MJX、mujoco-mjx、MJX-JAX

相关

MuJoCo、JAX、MuJoCo Warp、Brax、MuJoCo Playground、可微仿真

来源

MuJoCo XLA (MJX) documentation
mujoco-mjx (PyPI)

信息截至

2026-09