Revisiting Implicit Differentiation for Learning Problems in Optimal Control

TL;DR

This paper introduces a novel, scalable method for differentiating through optimal control trajectories using implicit differentiation, improving efficiency, parallelization, and stability over previous approaches, and unifying prior claims about computational complexity.

Contribution

We propose a direct matrix equation evaluation approach for implicit differentiation in optimal control, enhancing scalability, parallelization, and numerical stability, and unify previous quadratic complexity claims.

Findings

Trajectory derivatives scale linearly with timesteps.

Method outperforms prior approaches in scalability and stability.

Successfully applied to complex control benchmarks including quadrotor and rocket landing.

Abstract

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.