⚡ electrical engineering

Amortising Trajectory Optimisation for Residual MPC via Implicit Contact Differentiation

This paper introduces an efficient Implicit Function Theorem-based differentiable simulation method for contact-rich trajectory optimization that drastically reduces memory usage compared to unrolled automatic differentiation, and combines it with optimiser distillation to significantly improve the success rate of residual MPC in complex robotic tasks.

Original authors: Daniel Layeghi, Thomas Corbères, Calum Arnott, Aditya Kamireddypalli, Hashim Al-Obaidi, Steve Tonneau, Michael Mistry

Published 2026-07-29

📖 7 min read🧠 Deep dive

CC BY 4.0

Original authors: Daniel Layeghi, Thomas Corbères, Calum Arnott, Aditya Kamireddypalli, Hashim Al-Obaidi, Steve Tonneau, Michael Mistry

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot to walk, juggle, or play soccer. To do this, the robot needs a "brain" that can simulate the future in its head, trying out millions of different moves to see which one works best. This is called trajectory optimisation. The tricky part is contact. When a robot's foot hits the ground, a ball bounces off a wall, or a hand grabs a cup, the physics get messy and unpredictable. It's like trying to predict exactly how a pile of Jenga blocks will fall when you pull one out; tiny changes in force lead to huge changes in outcome.

To make these predictions, scientists use differentiable simulation. Think of this as a super-powered video game engine that doesn't just show you the next frame, but also tells you exactly how the game would change if you nudged the controls a tiny bit. This "nudge sensitivity" allows the robot to learn from its mistakes instantly. However, there's a catch: calculating these nudges for contact-heavy tasks is incredibly expensive. It's like trying to film a movie in slow motion, but every time the camera zooms in on a collision, the film reel gets longer and longer, eventually filling up your hard drive before you can finish the scene. This paper tackles that memory problem and shows how to make robots learn complex contact skills much faster and more reliably.

The Problem: The "Memory Monster" in Robot Brains

Imagine you are trying to solve a maze. The standard way to teach a robot to solve it is to let it walk through the maze, hit a wall, and then rewind the tape to see exactly where it went wrong. In the world of robot physics, this "rewinding" is called unrolled automatic differentiation.

The problem arises when the robot hits a wall (or a floor, or another object). To figure out the physics of that bounce, the computer has to run a complex calculation many times, like a detective checking and re-checking clues until the answer is perfect. If the computer needs to check the clues 10 times to get it right, the "rewind tape" has to store the memory of all 10 checks. If you want the answer to be even more perfect, you might need 100 checks. Suddenly, the memory tape becomes 100 times longer.

This creates a terrible trade-off. If you want the robot to be precise (check the clues 100 times), you can only run a few robots at once because your computer runs out of memory. If you want to run thousands of robots at once to learn faster, you have to cut the checks short (maybe only 5 times), which means the robot learns from a sloppy, inaccurate answer. It's like trying to learn a dance by only watching the first five seconds of the video; you might get the steps, but you'll miss the crucial spin at the end.

The Solution: The "Magic Snapshot"

The authors of this paper, working with the MuJoCo physics simulator (a popular tool for robot research), found a clever way to bypass the memory monster. Instead of rewinding the entire tape of the detective's 100 checks, they used a mathematical trick called the Implicit Function Theorem (IFT).

Think of it this way: Imagine a detective who has solved a mystery. Instead of showing you the 100 pages of notes they took to get there, they just hand you the final solved case file and a "magic snapshot" of the solution. This snapshot tells you exactly how the solution would change if you tweaked one tiny detail, without needing to see the messy notes.

In technical terms, the paper introduces a method that differentiates the stationarity residual (a fancy way of saying "the point where the math says we are done") rather than the steps taken to get there.

The Old Way (Unrolled AD): Stores every single step of the solver. If you go from 1 step to 10 steps, your memory usage jumps by 10.6 times.
The New Way (IFT): Stores a nearly constant amount of memory. Even if you increase the solver effort from 1 step to 10 steps, the memory usage only changes by less than 4%.

This is a game-changer. It means the computer can demand a super-precise answer (checking the clues 100 times) without running out of memory. In fact, when the paper tested this with 256 active contacts (like a robot with many fingers touching a table), the new method used 20 times less memory than the old way. With 16 contacts and a complex robot model, it used 6 times less memory.

The Result: Teaching Robots to "Distill" Wisdom

With this new, memory-efficient tool, the authors didn't just stop at making the math faster; they used it to teach robots better. They created a system they call Optimiser Distillation.

Imagine a master chef (the "teacher") who spends hours perfecting a complex recipe. This chef is slow but incredibly accurate. Then, you have a sous-chef (the "student" or policy) who is fast but needs guidance.

The Teacher: The computer runs a full, long-horizon optimisation (like the master chef planning the whole meal) to find the perfect sequence of moves. This is done in batches, thanks to the new memory-saving trick.
The Student: The robot learns from these perfect sequences, creating a "policy" (a set of instincts) that knows the general plan.
The Hybrid: When the robot is actually doing the task, it doesn't just blindly follow the policy. It uses the policy for the big picture (the long-term plan) but adds a quick, local "residual" correction (a short-horizon optimiser) to handle sudden bumps or slips.

The paper tested this on three different robots:

Finger: A small arm spinning a top.
Franka: A large arm pushing a box.
Unitree: A four-legged dog-like robot running.

The results were impressive. When the planning horizon (how far ahead the robot looks) was short (just 6 steps), the standard method (iLQR) often failed. But with the new "distilled" policy guiding the robot, the success rate jumped dramatically:

Across all three tasks (Finger, Franka, and Unitree): Success increased by 28 to 98 percentage points compared to standard iLQR.

For the Franka robot pushing a box, the standard short-sighted robot barely succeeded at all, while the new hybrid robot succeeded with much less "lookahead," proving that the policy provided the long-term strategy while the local optimiser handled the tricky contact moments.

Why This Matters

This paper doesn't just suggest a theoretical idea; it provides a working, open-source tool that changes the rules of the game. By proving that you can get high-precision contact derivatives without the massive memory cost, the authors have removed a major bottleneck in robot learning. They showed that you don't have to choose between "fast but sloppy" and "slow but precise." You can have both.

The authors are confident in these findings, having validated them against standard numerical methods (finite differences) and showing that their new method matches the accuracy of the old methods while using a fraction of the resources. They have even released their code to the public, inviting others to build faster, smarter, and more dexterous robots that can handle the messy, real-world physics of contact without getting stuck in a memory loop.

Technical Summary: Amortising Trajectory Optimisation for Residual MPC via Implicit Contact Differentiation

Problem Statement

Differentiable simulation is essential for contact-rich trajectory optimisation and learning, as it exposes local sensitivities of task outcomes to controls. However, scaling this pipeline is hindered by the computational cost of differentiating iterative contact solvers. Existing approaches face a fundamental trade-off between convergence and parallelism:

Finite Differences (FD): Expensive and sensitive to step size.
Unrolled Automatic Differentiation (AD): Differentiates the entire solver trace. While accurate, it stores a growing computation trace for every solver iteration. Tightening the contact solution (increasing iterations) consumes significantly more memory, limiting the number of parallel trajectories that can be processed.
KKT Sensitivity Derivations: Require intricate, solver-specific hand-assembled derivations.

The core challenge is that standard reverse-mode AD differentiates the executed solver trace rather than the converged solution. This creates a memory bottleneck where increasing solver effort (to improve accuracy) directly reduces batch capacity, preventing efficient full-horizon optimisation.

Methodology

The authors propose a framework that decouples the forward solver from its backward derivative using the Implicit Function Theorem (IFT).

1. Implicit Differentiation via IFT

Instead of unrolling the solver iterations, the method differentiates the stationarity residual of the regularised smooth contact problem at the tolerance-converged solution.

Formulation: The contact problem in MuJoCo (MJX) is formulated as a regularised convex optimisation. The solution $a^\star$ is defined as the root of the stationarity residual $F(a, \theta) = 0$ , where $\theta$ represents system parameters (state, controls, friction, etc.).
Derivation: By applying the IFT to the residual $F$ , the derivative of the solution with respect to parameters is derived as:
$\frac{da^\star}{d\theta} = -F_a^{-1} F_\theta$
where $F_a$ is the Jacobian of the residual with respect to acceleration.
Implementation: The method uses jax.lax.custom_root to implement this backward pass. It requires only one transposed linear solve and one vector-Jacobian product. Crucially, it does not store the intermediate states of the $K$ solver iterations used in the forward pass.
Scope: The derivative is local to fixed smooth contact regions and may change across contact-set or friction-mode transitions, but within a fixed regime, it provides exact gradients of the converged solution.

2. Optimiser Distillation for Residual MPC

To leverage the scalability of the implicit derivative, the authors introduce an optimiser distillation pipeline for Model Predictive Control (MPC):

Teacher: A batched, full-horizon iLQR (Iterative Linear Quadratic Regulator) uses the implicit differentiable MJX dynamics to generate high-quality, contact-rich trajectories.
Student Policy: A neural network policy is trained to mimic the actions of the full-horizon iLQR teacher.
Deployment (Hybrid Controller): At runtime, the system uses a Residual MPC approach:
- The distilled policy provides the long-horizon nominal action ( $\bar{u}_t$ ).
- A short-horizon residual iLQR computes a local correction ( $r_t$ ) and feedback gains based on the differentiable contact model.
- The final control is $u_t = \text{sg}[\pi_\theta(o_t)] + r_t$ , where $\text{sg}$ denotes a stop-gradient operation, preventing the policy from being updated by the short-horizon optimiser during deployment.

Key Contributions

AD-Assisted Implicit Backward Pass: The authors introduce a backward pass for the stationarity residual of MJX's regularised contact solve. This approach leaves the forward solver unchanged, avoids hand-assembled KKT systems, and removes the solver-iteration-dependent memory storage required by unrolled AD.
Validation and Scalability: The method is validated against central finite differences and unrolled AD. Results show that IFT matches the accuracy of unrolled AD at high iteration counts while substantially outperforming it in compiled temporary memory.
- Memory Scaling: IFT memory remains nearly constant as solver iterations increase (changing by <4% from 1 to 10 iterations), whereas unrolled AD grows by 10.6×.
- Contact/Dimension Scaling: IFT uses 20× less memory at 256 active contacts and 6× less at 16 contacts with 96 Degrees of Freedom (DoF) compared to unrolled AD.
Policy-Guided Residual MPC: The paper demonstrates a distillation pipeline where full-horizon optimised actions are distilled into a policy. This policy guides a short-horizon residual iLQR, enabling successful execution on complex tasks where standard short-horizon iLQR fails.

Experimental Results

The framework was evaluated on three contact-rich tasks: Finger (rotating a spinner), Franka (pushing a box), and Unitree (quadruped acceleration).

Trajectory Optimisation: In batched full-horizon optimisation, iLQR using the implicit derivative achieved lower mean running costs and tighter percentile ranges compared to first-order methods (Adam) across all tasks.
Distillation Performance:
- The hybrid controller (Policy + Short-Horizon Residual iLQR) significantly outperformed standard short-horizon iLQR.
- At a 6-step planning horizon, the hybrid approach improved success rates by 28 to 98 percentage points over standard iLQR across the three tasks.
- Franka: Standard short-horizon iLQR achieved near-zero success, while the hybrid controller succeeded with substantially less lookahead. The policy provided the long-horizon pushing strategy, while iLQR handled local contact corrections.
- Finger & Unitree: The policy allowed the system to succeed with shorter lookaheads than required by pure iLQR. In some regimes (e.g., Finger), a policy-only controller outperformed very short residual horizons, as the residual optimiser could perturb an already correct policy.

Significance and Claims

The paper claims that residual-based implicit differentiation makes tolerance-converged MJX contact derivatives practical at accelerator scale. By removing the memory bottleneck associated with solver iterations, the method enables:

Tighter Convergence: Solvers can be run to higher precision without sacrificing batch parallelism.
Full-Horizon Supervision: The ability to run batched full-horizon iLQR as an online teacher, which was previously infeasible due to memory constraints.
Effective Amortisation: The resulting hybrid controller achieves substantially higher success rates than standard iLQR at short planning horizons, effectively amortising expensive task-specific computation across initial conditions and goals.

The authors release the implementation as open source, providing a modular pattern for differentiable contact simulation that separates the forward solver from its derivative without requiring solver-specific analytical derivations.