Amortising Trajectory Optimisation for Residual MPC via Implicit Contact Differentiation

Motivation: Efficient derivatives through Contact

Core Question: How can we make querying these gradients memory and batch efficient ?

Differentiable Simulation

Useful for "greybox" inference to estimate parameters of trajectory, dynamics, policy.
Existing sims utilise Autodiff, example MJX.
Often involves accumulating gradients through solver states.

Planning and Learning need gradients

Trajectory Optimisation
Policy Learning

Requires gradients of task cost w.r.t controls, states or policy parameters.

Motivation: Efficient derivatives through Contact

\( \mathbf{x}_t, \mathbf{u}_t\)

\( \mathbf{x}_{t+1} \)

Rigid-body Dynamics

Contact geometry / Jacobians

Iterative Contact Solve

Integration

One Simulator Rollout

Planning,

Learning,

Inference

in Robotics

Accumulated

task/learning

Cost

Differentiate through rollout

Loss gradients backpropagate through rollout

Simulator step \(f(\mathbf{x}_t, \mathbf{u}_t)\)

Motivation: Efficient derivatives through Contact

One Simulator Rollout

Planning,

Learning,

Inference

in Robotics

\( \theta \)

Accumulated

task/learning

Cost, \( \mathcal{L}\)

Differentiate through rollout

Loss gradients backpropagate through rollout \( \frac{\partial{\mathcal{L}}}{{\partial{\theta}}}\)

Simulator step \(f_{\theta}(\mathbf{x}_t, \mathbf{u}_t)\)

Contact geometry / Jacobians

\( \mathbf{x}_t, \mathbf{u}_t\)

\( \mathbf{x}_{t+1} \)

Rigid-body Dynamics

Iterative Contact Solve

Integration

Bottleneck: backprop through

iterative contact solve using Automatic Differentiation

Insight: Implicit Contact Differentiation

Unrolled Autodiff (AD)

Implicit Contact Solve

Iter 1

Iter 2

Iter 3

Iter K

\( \theta \)

\( a^{*} \)

Solve forward

...

Iter 1

Iter 2

Iter 3

Iter K

\( \theta \)

\( a^{*} \)

Solve forward

...

AD unrolled backprop

Through every solver iteration \( k \in K \).
Stores the full solver trace.
Memory grows with Solver effort and Batch.

Utilise Implicit Function Theorem

Same forward solver.
Backward pass uses \( F(a^{*}, \theta) = 0 \).
No per-iteration trace in backward pass.

Contributions

Implicit Contact Differentiation Implementation

Amortized Trajectory Optimization for Residual MPC

An implementation of Implicit Function Theorem (IFT) based backward pass through contact in MJX.
Solve batched Trajectory Optimisation (iLQR) for common contact tasks using proposed implemented gradients.

Distil full-horizon optimised actions into
a policy paired with short-horizon residual iLQR, and evaluate closed-loop behaviour on Finger, Franka, and Unitree.