Amortising Trajectory Optimisation for Residual MPC via Implicit Contact Differentiation

This paper addresses the computational bottleneck in differentiable contact-rich simulation, where standard automatic differentiation (AD) methods struggle with the memory cost of unrolling iterative contact solvers. The authors propose an implicit differentiation approach based on the Implicit Function Theorem (IFT) applied to the MuJoCo MJX simulator. By differentiating the stationarity residual at the converged solution, the method avoids storing the entire solver trace, allowing for tighter contact solves without the memory overhead associated with traditional unrolled AD.

Instead of differentiating through the iterative steps of the contact solver, the authors define a stationarity residual that vanishes at the converged solution. By applying the Implicit Function Theorem to this residual, they derive the local sensitivity of the contact response. This approach is modular, meaning it does not require changing the underlying forward solver or deriving complex, solver-specific Karush-Kuhn-Tucker (KKT) systems. The authors implement this as a custom backward pass that uses a dense QR solve to compute the required adjoints, effectively decoupling the memory cost of the backward pass from the number of iterations performed by the forward solver.

The authors demonstrate that their IFT-based method maintains high gradient accuracy compared to finite differences while offering drastically better memory scaling. Specifically, as solver effort increases, the memory footprint of unrolled AD grows significantly, whereas the IFT method remains nearly constant. This efficiency gain is leveraged to implement an optimizer distillation framework for residual Model Predictive Control (MPC). By distilling full-horizon optimized trajectories into a policy, the system can provide long-horizon nominal actions while using short-horizon residual iLQR for local corrections. This combination significantly improves success rates on complex tasks like finger manipulation and quadrupedal locomotion compared to standard iLQR.