robfiras/loco-mujoco: Whole-Body Locomotion Learning

What LocoMuJoCo Models For Whole-Body Control

LocoMuJoCo targets a specific research problem: learning coordinated movement for an entire body rather than optimizing one isolated joint or action. The README describes an imitation learning benchmark for whole-body control, with quadrupeds, humanoids, and musculoskeletal human models. Its intended users are researchers and engineers who need a repeatable environment for studying locomotion from motion capture, comparing learning algorithms, or experimenting with custom reward functions.

The benchmark covers more than humanoid walking. The README lists 12 humanoid and 4 quadruped environments, including four biomechanical human models. It also includes motion data for humanoid tasks, with more than 22,000 samples per humanoid, based on AMASS, LAFAN1, and native LocoMuJoCo datasets. Those datasets are retargeted for each humanoid, which matters because motion learned for one body does not automatically transfer to another skeleton.

LocoMuJoCo is not limited to imitation learning. The project supports custom reward function classes, so the same environment and simulation stack can also be used for pure reinforcement learning. That flexibility is useful when a research question begins with a reference motion but later needs a task-specific objective. The relevant tradeoff is scope. This is a benchmark and environment collection with datasets, factories, metrics, algorithms, and examples, not a small package intended solely to load a model file.

The repository also supports a Gymnasium interface. This makes the environments easier to connect to code that expects Gymnasium-style environments, although the exact observation and action configuration still depends on the selected factory and environment.

The MuJoCo, MJX, and MjWarp Execution Path

LocoMuJoCo offers two simulation paths described in the README: MuJoCo for a single environment and MJX for parallel environments. Version 1.1 adds support for the MjWarp backend on compatible GPUs. The standard path is appropriate for a single interactive trajectory, while MJX and MjWarp target workloads where many environments run in parallel.

The project combines environment construction and training into a JIT-compiled function according to the README. Its reference algorithms are written in single files and include PPO, GAIL, AMP, and DeepMimic. That arrangement favors short experiments and direct comparison between algorithms, but it does not mean every component is interchangeable. The environment defines the physical model, observation configuration, reward logic, terminal-state handling, and randomization, while the algorithm determines how actions are selected and how parameters are updated.

The modular design is an important part of the architecture. The README names observation types, reward functions, terminal state handlers, and domain randomization as components that can be defined, swapped, and reused. Domain and terrain randomization are built in, which supports experiments over changing dynamics or terrain conditions.

Trajectory comparison is also included. The README documents dynamic time warping and discrete Fréchet distance in JAX. These metrics are intended for comparing motion trajectories when their time alignment or full shape matters. The repository does not present them as a universal replacement for task success metrics, so a study should decide whether trajectory similarity, reward, stability, or another measurement matches its objective.

The MjWarp path adds a further hardware choice. According to the README, callers can enable it from the low-level `Mjx` API with `use_mjwarp=True`, and high-level factories can forward the same flag. The documentation says this can further accelerate parallel simulation on compatible GPUs. The repository does not establish compatibility for every GPU or operating system, so hardware support must be checked before selecting that backend.

Installing LocoMuJoCo and Running a First Trajectory

The README gives an editable installation for a clone of the repository. This is the first path to use when the goal is to inspect the examples or work with the checked-out code.

cd loco-mujoco
pip install -e .

The package metadata requires Python 3.10 or later and lists dependencies including JAX, MuJoCo, MJX, MuJoCo Warp, Flax, Hydra, Gymnasium, OpenCV, and the `loco-mujoco-models` Git dependency. A default installation uses the CPU version of JAX according to the README. For JAX on CUDA 12, the documented installation is:

pip install jax["cuda12"]

The next step depends on the environment. LocoMuJoCo can load its default or LAFAN1 datasets through the Hugging Face dataset repository and cache them automatically. The README gives this visualization example for `UnitreeH1`:

from loco_mujoco.task_factories import ImitationFactory, LAFAN1DatasetConf, DefaultDatasetConf, AMASSDatasetConf


env = ImitationFactory.make("UnitreeH1",
                            default_dataset_conf=DefaultDatasetConf(["squat"]),
                            lafan1_dataset_conf=LAFAN1DatasetConf(["dance2_subject4", "walk1_subject1"]),
                            )

env.play_trajectory(n_episodes=3, n_steps_per_episode=500, render=True)

This creates an imitation environment from the named human model, selects a default squat dataset and two LAFAN1 trajectories, then plays three episodes with rendering enabled. Rendering can require a graphical environment or suitable display configuration, which the README does not document separately. A command-line example should therefore be used when the goal is only to inspect data loading or verify installation.

For repeated dataset loading, the README documents an optional cache command. The path is supplied by the user, and the project stores calculated forward-kinematics results for reuse.

loco-mujoco-set-all-caches --path "$HOME/.loco-mujoco-caches"

The repository does not specify the exact output or cache format in its documentation, so this command should be treated as an optimization rather than a required installation step.

Datasets, Retargeting, and Licensing Boundaries

Dataset handling is a defining part of LocoMuJoCo rather than an incidental feature. The README identifies three sources for humanoid motion capture: the default LocoMuJoCo data, LAFAN1, and AMASS. Default and LAFAN1 data are available through the LocoMuJoCo Hugging Face dataset repository and are downloaded and cached automatically. AMASS is different because its license requires a separate download and installation. The repository points to `loco_mujoco/smpl` for that process.

The default dataset avoids the extra AMASS setup, but choosing a source still changes the experimental question. Default data can provide a direct path to a documented example. LAFAN1 requires selecting named motions such as `dance2_subject4` and `walk1_subject1`. AMASS requires additional model and licensing setup, so it is not equivalent to the automatic default or LAFAN1 path.

LocoMuJoCo stores joint positions and velocities to reduce memory use. Other attributes are calculated with forward kinematics when a dataset is loaded. The cache command is therefore tied to a concrete trade-off: smaller stored data and more loading computation versus a cache that stores those calculated results for later runs. The README does not state cache invalidation behavior, cache portability, or the exact disk cost of all supported datasets.

The optional `smpl` dependencies include `smplx`, Torch, and related packages. The package metadata also carries an MIT license for LocoMuJoCo itself. That does not remove the need to observe the licensing terms of AMASS, LAFAN1, model files, or the MyoSkeleton model. The README specifically states that `loco-mujoco-myomodel-init` is required to accept the license and download the model. For any research or redistribution plan, dataset and model permissions should be checked separately from the project license.

Where LocoMuJoCo Is the Wrong Tool

LocoMuJoCo is a poor fit for a production robot deployment that needs only a direct policy inference path. Its documented scope centers on benchmark environments, motion datasets, research algorithms, rendering, customization, and simulation backends. The README does not describe hardware deployment, controller safety logic, real-time guarantees, or a supported robot integration workflow.

It can also be more than a project needs. A simple reinforcement learning experiment on one conventional MuJoCo task may not need humanoid retargeting, biomechanical models, large motion datasets, or an MJX parallel backend. Those resources add installation and modeling complexity. A researcher who only wants a conventional Gymnasium environment should first compare the dependencies and environment choices against a smaller MuJoCo setup.

GPU-backed MjWarp execution has a similar boundary. The README limits its acceleration claim to compatible GPUs and makes it optional. It does not state that MjWarp will work on every machine. JAX itself is a significant dependency, and the package metadata lists MuJoCo, MJX, MuJoCo Warp, Flax, Hydra, and other libraries. Environments without suitable GPU support may need the documented default CPU installation instead.

The benchmark also does not replace evaluation discipline. A policy can imitate a trajectory while failing a task-specific objective, and a custom reward can change the behavior that an algorithm learns. The documentation names trajectory metrics and reward customization, but it does not claim that one metric establishes locomotion quality. Researchers must define success, compare conditions, and inspect whether the selected environment represents the body and task they intend to study.

LocoMuJoCo Compared with a Plain MuJoCo Setup

A plain MuJoCo setup is the practical alternative. It can provide the simulator, model loading, and environment logic without LocoMuJoCo's benchmark layer. The difference is not simply speed. LocoMuJoCo adds whole-body imitation datasets, retargeting across robots, reusable environment components, built-in randomization, JAX reference algorithms, trajectory metrics, and factories for many humanoid and quadruped models.

LocoMuJoCo is more convenient when a study needs to compare several bodies or motion sources through a common interface. Its `ImitationFactory` examples show that datasets are selected through configuration objects, and the README describes swapping observation types, reward functions, terminal handlers, and randomization. That structure reduces the need to assemble every experiment from separate MuJoCo code.

The trade-off is a larger dependency and licensing surface. The package requires Python 3.10 or later and installs several simulation and machine-learning libraries. AMASS and MyoSkeleton introduce additional setup. Plain MuJoCo may be preferable for a narrow simulation, a small custom environment, or a deployment-oriented prototype where benchmark extras are unnecessary.

The right choice depends on the research artifact. Choose plain MuJoCo when the simulator itself is the object of study and a small interface is valuable. Choose LocoMuJoCo when the central problem involves imitation of complex locomotion, many model configurations, motion retargeting, or comparing JAX-based algorithms under a shared benchmark.

Editorial conclusion

Adopt LocoMuJoCo when you are building imitation or reinforcement learning experiments for whole-body locomotion and can work with its Python, JAX, MuJoCo, and dataset requirements. Do not adopt it as a lightweight controller library or expect AMASS support without a separate licensed installation. Check the available model files, required dataset licenses, GPU compatibility, and the specific environment factory before selecting it for a project. The last push was on 2026-08-22.

Frequently asked questions

Official sources

1. License: MIT
2. Project website
3. README
4. Releases
5. robfiras/loco-mujoco on GitHub