Bringing MyoSuite to MJX for GPU-Accelerated Musculoskeletal RL

The Bottleneck: Muscle Simulation on CPU

MyoSuite is a collection of musculoskeletal environments and tasks simulated in MuJoCo and wrapped in the Gym API — the go-to platform for research on muscle-driven, biomechanically realistic control (e.g. the MyoHand). It’s a fantastic testbed, but the original implementation runs on CPU-based MuJoCo, which caps how many environments you can step in parallel. For large-scale reinforcement learning — where sample efficiency is bought with sheer throughput — that’s the binding constraint.

What MJX Changes

MJX is MuJoCo compiled through XLA (Accelerated Linear Algebra), so simulation runs on the GPU/TPU inside JAX. Instead of a handful of CPU environments, you can step thousands of environments in parallel on a single accelerator, keeping the whole rollout on-device alongside the policy. For contact-rich, high-DoF musculoskeletal models, that’s the difference between hours and minutes per experiment.

The Work

myosuite-mjx is my effort — through Micropilot, in collaboration with MyoSuite creator Vittorio Caggiano (FAIR, Meta) — to give MyoSuite an MJX backend:

• Port MyoSuite’s musculoskeletal environments to run under MJX/JAX.
• Preserve muscle-tendon dynamics and contact-rich simulation fidelity.
• Keep compatibility with modern MuJoCo and the standard Gym API.
• Enable GPU-accelerated, massively parallel rollouts for large-scale RL of dexterous control.

It’s an active, work-in-progress effort — building from a source MuJoCo install up through the MJX and MyoSuite layers.

Where It Fits

This is the simulation engine behind the rest of my robotics stack: the same MyoHand model that supplies tendon-control targets in emg2tendon becomes trainable at scale once it runs on MJX. Fast, parallel musculoskeletal simulation is the prerequisite for learning robust low-level control of dexterous, anthropomorphic hands.