OverviewEdit
MuJoCo XLA (MJX) is a specialized extension of the MuJoCo physics engine, designed to run simulations on hardware supported by the XLA (Accelerated Linear Algebra) compiler via the JAX framework. Running single threaded physics simulation on the GPU is not very efficient. The advantage with MJX is that we can run environments in parallel on a hardware accelerated device.
InstallationEdit
Install using:
pip install mujoco-mjx
Colab TutorialEdit
A detailed tutorial demonstrating the use of MJX along with reinforcement learning to train humanoid and quadruped robots to locomote is available here.