This repo implements the code for legged robot MPC and Trajectory Optimization all in JAX.

Installation | PrePrint | Video |

Features

MPX is a JAX library that provides:

✅ True GPU Parallelism Exploits both temporal and state-space parallel scans directly on the GPU, without approximations or offline precomputations. Lower the complexity to $\mathcal{O}(\log^2{n}\log{N} + \log^2{m})$ from the classical $\mathcal{O}(N(n + m)^3)$ where n = state dim, m = control dim, N = horizon length

✅ JAX Autodiff & Vectorization Fully differentiable solver easily integrates into learning pipelines and supports batched RL-style environments.

✅ A multiple-shooting SQP formulation solves the KKT system in parallel, maintaining exactness and fast convergence.

✅ MJX MODELS Support MJX whole body dynamics (included examples with Talos, H1, Aliengo and Go2)

The solver is wrapped by the MPCControllerWrapper class, and all the settings (such as the dynamics model and cost function to be used) can be changed in the config files. Examples for various legged robots are provided in the examples folder.

Note:
If you want to solve multiple MPC in parallel, look at the examples/multi_env.py MPCWrapper is designed to use the whole body model, if you want to use the srbd model, use mpc_wrapper_srbd.py; look at examples/srbd_quad.py

Task examples

Acrobot Swing-Up	Quadruped Trot	Humanoid Jump	Quadruped Barrel Roll	Obstacle Avoidance

Available solvers

Solver	Multiple shooting	Dynamics	Integrator	Stage Equality	Inequality constraints
primal-dual	✔️	Forward	Semi-implicit	❎	❎ (Soft Cost barrier)
fddp	✔️ (dynamics defects)	Forward	Semi-implicit	❎	❎ (Soft Cost barrier)
eq_primal_dual	✔️	Forwward or Inverse	Explicit or Implicit	✔️	❎ (Soft Cost barrier)
eq_fddp	✔️ (dynamics defects)	Forwward or Inverse	Explicit or Implicit	✔️	❎ (Soft Cost barrier)
ip_primal_dual	✔️	Forwward or Inverse	Explicit or Implicit	✔️	✔️ (interior point)

Installation

Prerequisites

• Install Pixi - a fast conda package manager

Clone the repo

git clone git@github.com:iit-DLSLab/mpx.git
cd mpx && git submodule update --init --recursive

Set Up Environment with Pixi

Pixi will automatically manage your environment with all dependencies including ROS 2 Humble, JAX, and build tools:

# Initialize the Pixi environment
pixi install

# Activate the environment
pixi shell

Alternative: Traditional Conda Setup (Deprecated)

For reference, the legacy conda setup is:

conda create -n mpx_env python=3.13 -y
conda activate mpx_env
pip install -e .

RUN example

# Using Pixi
pixi run python mpx/examples/mjx_quad.py

# Or activate the environment first
pixi shell
python mpx/examples/mjx_quad.py

Use the keyboard's arrows to control the robot.

Note:
The first time running the script it can take more than a minute to JIT the solver

Citing this work

@article{amatuccisousa26ral,
 author={Amatucci, Lorenzo and Sousa-Pinto, João and Turrisi, Giulio and Orban, Dominique and Barasuol, Victor and Semini, Claudio},
 title={Primal-Dual iLQR for GPU-Accelerated Learning and Control in Legged Robots},
 year={2026},
 volume={11},
 number={1},
 pages={1010-1017},
 journal={IEEE Robotics and Automation Letters},
 doi={10.1109/LRA.2025.3632610}
}