Primal-Dual iLQR

TL;DR

The paper proposes Primal-Dual iLQR, a SQP-based solver for unconstrained discrete-time OCPs that leverages direct multiple shooting and an $\ ho$-regularized primal-dual LQR subproblem. By exploiting the LQR structure, it achieves $O(N (n+m)^3)$ linear solves per iteration and parallelizes the primal-dual Newton-KKT system to $O(\log N)$ time, with $O(1)$ parallel time per line search step. The approach guarantees global convergence and preserves fast local convergence without second-order corrections, outperforming several existing direct multiple shooting methods. An open-source JAX implementation is provided for GPU/CPU acceleration.

Abstract

We introduce a new algorithm for solving unconstrained discrete-time optimal control problems. Our method follows a direct multiple shooting approach, and consists of applying the SQP method together with an $\ell_2$ augmented Lagrangian primal-dual merit function. We use the LQR algorithm to efficiently solve the primal-dual Newton-KKT system. As our algorithm is a specialization of NPSQP, it inherits its generic properties, including global convergence, fast local convergence, and the lack of need for second order corrections or dimension expansions, improving on existing direct multiple shooting approaches such as acados, ALTRO, GNMS, FATROP, and FDDP. The solutions of the LQR-shaped subproblems posed by our algorithm can be be parallelized to run in time logarithmic in the number of stages, states, and controls. Moreover, as our method avoids sequential rollouts of the nonlinear dynamics, it can run in $O(1)$ parallel time per line search iteration. Therefore, this paper provides a practical, theoretically sound, and highly parallelizable (for example, with a GPU) method for solving nonlinear discrete-time optimal control problems. An open-source JAX implementation of this algorithm can be found on GitHub (joaospinto/primal_dual_ilqr).

Figures (1)

• Figure 1: The solution of the quad-pendulum problem.