Robust and Efficient MuJoCo-based Model Predictive Control via Web of Affine Spaces Derivatives

Chen Liang; Daniel Rakita

cs.RODec 24, 2025

Robust and Efficient MuJoCo-based Model Predictive Control via Web of Affine Spaces Derivatives

Authors: Chen Liang, Daniel Rakita

Organizations: Department of Computer Science, Yale University, New Haven, CT 06520, USA

Abstract

MuJoCo is a powerful and efficient physics simulator widely used in robotics. One common way it is applied in practice is through Model Predictive Control (MPC), which uses repeated rollouts of the simulator to optimize future actions and generate responsive control policies in real time. To make this process more accessible, the open source library MuJoCo MPC (MJPC) provides ready-to-use MPC algorithms and implementations built directly on top of the MuJoCo simulator. However, MJPC relies on finite differencing (FD) to compute derivatives through the underlying MuJoCo simulator, which is often a key bottleneck that can make it prohibitively costly for time-sensitive tasks, especially in high-DOF systems or complex scenes. In this paper, we introduce the use of Web of Affine Spaces (WASP) derivatives within MJPC as a drop-in replacement for FD. WASP is a recently developed approach for efficiently computing sequences of accurate derivative approximations. By reusing information from prior, related derivative calculations, WASP accelerates and stabilizes the computation of new derivatives, making it especially well suited for MPC's iterative, fine-grained updates over time. We evaluate WASP across a diverse suite of MJPC tasks spanning multiple robot embodiments. Our results suggest that WASP derivatives are particularly effective in MJPC: it integrates seamlessly across tasks, delivers consistently robust performance, and achieves up to a 2 $\mathsf{x}$ speedup compared to an FD backend when used with derivative-based planners, such as iLQG. In addition, WASP-based MPC outperforms MJPC's stochastic sampling-based planners on our evaluation tasks, offering both greater efficiency and reliability. To support adoption and future research, we release an open-source implementation of MJPC with WASP derivatives fully integrated.

Figures & tables

Panel of Figure 1 — Fig. 1 : Visuals of $10$ tasks evaluated to demonstrate the efficacy of WASP-based derivatives in MJPC. Relative to finite differences, WASP accelerates the total compute time by up to 2 $\mathsf{x}$ while maintaining robust performance.

Fig. 2 : A screenshot of GUI components for users to activate WASP and control WASP’s accuracy interactively

Condition	# Samples	Predictive		Robust		Cross Entropy		Sample Gradient
Condition	# Samples	Speedup	Perf.	Speedup	Perf.	Speedup	Perf.	Speedup	Perf.
Quadrotor	30	1.32	1.36	2.39	1.40	1.35	1.93	1.18	1.08
Swimmer	32	1.38	1.58	1.90	2.03	1.31	1.65	1.26	1.24
Quadruped Stand	60	1.34	173.1 ✗	2.49	442.5 ✗	0.93	32.15	1.54	105.0 ✗
Quadruped Climb	60	2.04	90.1 ✗	3.19	115.6 ✗	1.26	64.8 ✗	1.84	54.2 ✗
Quadruped Walk	60	1.78	265.3 ✗	2.78	340.4 ✗	1.10	190.8 ✗	1.61	159.5 ✗

TABLE II : WASP-Based iLQG vs. Sampling-Based Planners (A red “X” means the sampling-based planner fails to complete the task by exceeding a joint limit, falling over, or failing to reach the goal within

30

seconds)

Fig. 3 : Parameter sensitivity analysis of WASP on the quadruped trotting task. (a) Cost convergence across settings, with performance degrading only under very low state accuracy. (b) Average computation time per derivative evaluation, consistently lower than FD. (c) Number of forward simulations, showing adaptive resource usage with accuracy. Overall, WASP maintains stable performance and computational savings across a broad parameter range, with state accuracy more critical than control accuracy.

Explore similar work

May 27, 2026cs.RO

CA-AC-MPC: CUDA-Accelerated Actor-Critic Model Predictive Control

In the literature, actor-critic model predictive control (AC-MPC) integrates MPC with reinforcement learning to enable high-performance control of complex dynamical systems. However, its differentiable MPC layer requires repeatedly solving an optimization problem in both the forward and backward passes, leading to substantial training and inference latency. This paper tackles this bottleneck introducing a CUDA-accelerated variant that significantly reduces end-to-end execution time while preserving the control performance of the baseline formulation. Simulation results on an agile drone racing task show that our approach achieves state-of-the-art lap times and near-limit dynamic behaviour with markedly reduced training and inference time.

Antoonio Buo, Vittorio Cammarota, Michele Avagnale +3

PRISMA Lab and CREATE Consortium, Department of Electrical Engineering and Information Technology, University of Naples Federico II, Naples, Italy

Model Predictive Control Soft Actor-Critic Cuda Autonomous Drones

Ask PDF

Jun 23, 2026cs.RO

TurboMPC: Fast, Scalable, and Differentiable Model Predictive Control on the GPU

Robotics increasingly relies on GPUs for parallel simulation, large-scale learning, and neural-network inference. For model predictive control (MPC) to scale with this paradigm, solvers must run efficiently on this hardware while remaining fast, differentiable, and compatible with expressive MPC formulations used in robotics. We present TurboMPC, a differentiable MPC solver that runs entirely on the GPU and supports state and control inequality constraints, implicit integrators, cross-time-coupled costs, and slack variables. TurboMPC combines sequential quadratic programming (SQP), an alternating direction method of multipliers (ADMM) inner solver, implicit differentiation, and a co-designed JAX-CUDA implementation for efficiency and ease of use. In simulation, we validate TurboMPC on constrained planning, humanoid imitation learning, and reinforcement learning with neural-network cost function tasks, achieving up to

15\times

and

58\times

speedups over state-of-the-art CPU and GPU differentiable solvers, respectively. We deploy TurboMPC on a full-scale car for minimum-time racing and find that batched, GPU-accelerated tuning of MPC parameters via Bayesian optimization yields significantly faster driving than a hand-tuned baseline. TurboMPC also scales to planning horizons of over

8000

knot points while maintaining control of the vehicle. We open-source TurboMPC at: https://github.com/ToyotaResearchInstitute/turbompc

Gabriel Bravo-Palacios, Jianghan Zhang, Zachary Pestrikov +2

Dartmouth College, Hanover, NH, USA · Toyota Research Institute, Los Altos, CA, USA

Model Predictive Control Mujoco Neural Solvers Combustion Control

Ask PDF

Mar 5, 2026cs.RO

Accelerating Sampling-Based Control via Learned Linear Koopman Dynamics

This paper presents an efficient model predictive path integral (MPPI) control framework for systems with complex nonlinear dynamics. To improve the computational efficiency of classic MPPI while preserving control performance, we replace the nonlinear dynamics used for trajectory propagation with a learned linear deep Koopman operator (DKO) model, enabling faster rollout and more efficient trajectory sampling. The DKO dynamics are learned directly from interaction data, eliminating the need for analytical system models. The resulting controller, termed MPPI-DK, is evaluated in simulation on pendulum balancing and surface vehicle navigation tasks, and validated on hardware through reference-tracking experiments on a quadruped robot. Experimental results demonstrate that MPPI-DK achieves control performance close to MPPI with true dynamics while substantially reducing computational cost, enabling efficient real-time control on robotic platforms.

Wenjian Hao, Yuxuan Fang, Zehui Lu +1

independent researcher

Model Predictive Path Integral Koopman Operator Trajectory Generation

Ask PDF