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

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

Problem Statement
MuJoCo is a widely adopted physics simulator for robotics, frequently utilized within Model Predictive Control (MPC) frameworks to optimize future actions through repeated rollouts. The open-source library MuJoCo MPC (MJPC) facilitates this by providing ready-to-use algorithms. However, MJPC currently relies exclusively on Finite Differencing (FD) to compute the derivatives required for gradient-based planners (e.g., iLQG). FD estimates derivatives by perturbing every input dimension independently, causing the computational cost to scale linearly with the dimensionality of the state and action spaces. In high-degree-of-freedom (DOF) systems or complex scenes with rich contact dynamics, this results in hundreds or thousands of simulator calls per derivative evaluation, creating a bottleneck that hinders real-time performance. While Automatic Differentiation (AD) offers an alternative, exact derivatives in MPC often suffer from numerical instability due to the sharp, ill-conditioned sensitivities of nonlinear dynamics over short horizons.

Methodology
This paper introduces the integration of Web of Affine Spaces (WASP) derivatives into MJPC as a drop-in replacement for FD. WASP is a coherence-based approach that computes sequences of approximate derivatives by reusing information from prior, related evaluations.

• Core Mechanism: Instead of computing derivatives from scratch, WASP formulates derivative estimation as a constrained least-squares problem. It seeks an approximate Jacobian that lies within an affine subspace defined by a single, freshly computed Jacobian-vector product (JVP) while aligning with cached information from previous iterations.
• Integration: The authors implemented WASP directly into the MuJoCo source code (C language) alongside existing FD implementations. In the MJPC pipeline, WASP instances are created for each time step in the planning horizon. These instances maintain tangent matrices and approximate JVP matrices ($\hat{\Delta}F$) that are updated incrementally.
• Parameterization: To balance accuracy and efficiency, the implementation exposes two tunable parameters: frac (controlling the minimum number of JVPs relative to the maximum) and tol (controlling the threshold for matching the angle and norm of ground-truth JVPs). Users can adjust these via the MJPC GUI, allowing for a continuum between full FD (maximum accuracy/cost) and highly approximate derivatives (minimum cost).

Key Contributions

1. First Application of WASP to Dynamics-based MPC: While WASP was previously evaluated on kinematics-based functions, this work extends its application to full physics simulation within MPC, addressing the challenges of contact dynamics and nonlinearities.
2. Open-Source Implementation: The authors released a fully integrated, open-source version of MJPC with WASP derivatives, enabling the community to experiment with this drop-in replacement without modifying the simulator's core source code.
3. Comprehensive Evaluation: The study evaluates WASP across a diverse suite of ten locomotion tasks (including quadrupeds, bipeds, humanoids, and aerial robots) with varying complexities and contact patterns.

Results
The evaluation compares WASP-based MPC against both FD-based MPC and stochastic sampling-based planners (e.g., Predictive Sampling, Cross-Entropy Method).

• Speedup vs. FD: When used with derivative-based planners like iLQG, WASP achieved a 1.26x to 2.08x speedup in model derivative computation time compared to FD across all tasks. In terms of total planning time, WASP-based iLQG was consistently faster than FD-based iLQG.
• Performance vs. Sampling: On contact-rich tasks (e.g., quadruped climbing, biped balance), WASP-based iLQG significantly outperformed stochastic sampling-based planners. While sampling methods often failed to complete tasks (indicated by joint limit violations or falls), WASP-based iLQG achieved speedups of up to 4.0x relative to sampling planners while maintaining successful task execution.
• Task Performance: In several instances (e.g., Quadrotor, Swimmer), WASP-based control yielded slightly better task performance (lower cost) than FD. The authors speculate that mild approximation errors may act as a regularizer, smoothing sharp gradients and aiding in escaping local minima.
• Robustness: Sensitivity analysis revealed that WASP is robust to parameter variations, though state transition accuracy (fracx) is more critical than control accuracy (fracu) for maintaining stability.

Significance and Claims
The paper claims that WASP derivatives offer a compelling balance between efficiency and robustness for iterative control settings. By replacing FD with WASP, MJPC can sustain the fast update rates required for real-time robotics, particularly in high-DOF systems where FD is prohibitively expensive. The work demonstrates that coherence-based approximate derivatives are not only computationally superior but can also enhance control reliability in complex, contact-rich environments. The authors position this as a step toward expanding practical alternatives to finite differencing in robotics, lowering the barrier for adopting structured approximate derivatives in real-time MPC stacks.

Limitations
The authors note that all experiments were conducted in simulation, primarily on locomotion benchmarks, and that sim-to-real transfer remains an open question. Additionally, both FD and WASP-based planners struggled with contact-rich manipulation tasks, suggesting that the limitations lie in the short-horizon gradient-based MPC architecture itself rather than the derivative approximation method. Finally, the accuracy parameters currently require manual tuning, though the authors suggest future work could explore adaptive selection schemes.