Amortising Trajectory Optimisation for Residual MPC via Implicit Contact Differentiation

2026-07-27

AI Summary

This work addresses the high computational cost, rapid memory growth, and reliance on intricate analytical derivations in existing differentiable simulation methods for contact-intensive trajectory optimization. It proposes an implicit differentiation approach based on the implicit function theorem, enabling automatic differentiation of the stationarity residuals of a regularized smooth contact model at its converged solution—without unrolling solver iterations or manually constructing KKT systems. Furthermore, it introduces optimizer distillation into residual MPC for the first time, training a short-horizon policy network using full-horizon iLQR solutions to enhance planning efficiency. Experiments on Finger, Franka, and Unitree platforms demonstrate 28–98 percentage point improvements in six-step task success rates over standard iLQR, with memory consumption reduced to 1/20 of that required by conventional unrolled automatic differentiation under 256 contact points.

Abstract

Differentiable simulation can accelerate contact-rich trajectory optimisation by exposing local sensitivities of task outcomes to controls. Existing approaches either use finite differences, which are expensive and step-size sensitive; differentiate iterative contact solvers by unrolling automatic differentiation (AD), which stores a growing computation trace; or require intricate, solver-specific KKT sensitivity derivations. We introduce an AD-assisted implicit derivative for regularised smooth contacts and apply it to Mujoco MJX, based on the Implicit Function Theorem (IFT). The method differentiates the stationarity residual at the tolerance-converged solution, avoiding both solver unrolling and hand-assembled KKT systems. IFT keeps compiled temporary memory nearly constant with solver effort, changing by less than 4$\%$ from one to ten iterations versus 10.6$\times$ growth for unrolled AD. IFT memory grows slower with active contacts and model dimension, using 20$\times$ less memory at 256 contacts and 6$\times$ less at 16 contacts and 96 DoF. We further introduce optimiser distillation for residual MPC, amortising batched full-horizon iLQR into a policy that guides short-horizon residual iLQR. Across Finger, Franka, and Unitree, this raises six-step success by 28-98 percentage points over standard iLQR.

Problem

• trajectory optimisation
• differentiable simulation
• contact-rich dynamics
• model predictive control
• memory efficiency

Innovation

• implicit differentiation
• differentiable simulation
• trajectory optimization
• residual MPC
• optimizer distillation