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.
Enables Batched Trajectory Optimisation (iLQR/ Adam)
Finger
Hand reorientation\({}^{*}\)
Franka pushing\({}^{*}\)
* Trajectories computed in Batch - visualised individually
Unitree locomotion\({}^{*}\)
Amortised Trajectory Optimisation
Training
Evaluation (Epochs 5, 10, 15, 20, 25)
AmortTO
By Aditya Kamireddypalli
AmortTO
- 3
