This paper proposes a new method for differentiating through optimal
trajectories arising from non-convex, constrained discrete-time optimal control
(COC) problems using the implicit function theorem (IFT). Previous works solve
a differential Karush-Kuhn-Tucker (KKT) system for the trajectory derivative,
and achieve this efficiently by solving an auxiliary Linear Quadratic Regulator
(LQR) problem. In contrast, we directly evaluate the matrix equations which
arise from applying variable elimination on the Lagrange multiplier terms in
the (differential) KKT system. By appropriately accounting for the structure of
the terms within the resulting equations, we show that the trajectory
derivatives scale linearly with the number of timesteps. Furthermore, our
approach allows for easy parallelization, significantly improved scalability
with model size, direct computation of vector-Jacobian products and improved
numerical stability compared to prior works. As an additional contribution, we
unify prior works, addressing claims that computing trajectory derivatives
using IFT scales quadratically with the number of timesteps. We evaluate our
method on a both synthetic benchmark and four challenging, learning from
demonstration benchmarks including a 6-DoF maneuvering quadrotor and 6-DoF
rocket powered landing.Comment: Accepted to NeurIPS 2023 (poster