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

TL;DR

The paper presents Web of Affine Spaces (WASP) derivatives as a drop-in replacement for finite differences in MuJoCo-based MPC (MJPC), addressing the time-consuming derivative computations that bottleneck derivative-based planners like iLQG. WASP leverages coherence across sequential derivative estimates to produce accurate yet efficient Jacobians, achieving up to 2x faster derivative evaluations and often better task performance, while maintaining a black-box MuJoCo interface. Across a diverse set of locomotion tasks, WASP-enabled MJPC outperforms stochastic planners in speed and reliability, with notable improvements in highly dynamic, contact-rich scenarios. The work includes an open-source implementation, enabling practitioners to adopt WASP derivatives in real-time MPC and motivating further exploration of coherence-based derivatives in robotics control.

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 (3)

• Figure 1: Visuals of nine 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.
• Figure 2: A screenshot of GUI components for users to activate WASP and control WASP's accuracy interactively
• Figure 3: Parameter sensitivity analysis of WASP on the quadruped trotting task. (a) Cost convergence across parameter settings, showing that task performance remains stable except when state accuracy is severely reduced. (b) Average computation time per derivative evaluation, consistently lower for WASP than FD. (c) Number of forward simulations required, illustrating how WASP adapts resource usage based on accuracy requirements. Overall, WASP achieves stable performance and computational savings across a wide range of settings, with state accuracy proving more critical than control accuracy.