SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ming Xu
School of Computing Australian National University mingda.xu@anu.edu.au
Timothy Molloy
School of Engineering Australian National University timothy.molloy@anu.edu.au
School of Computing Australian National University stephen.gould@anu.edu.au
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.
This paper addresses end-to-end learning problems that arise in the context of constrained optimal control, including trajectory optimization, inverse optimal control, and system identification. We propose a novel, computationally efficient approach to computing analytical derivatives of state and control trajectories that solve constrained optimal control (COC) problems with respect to underlying parameters in cost functions, system dynamics, and constraints (e.g., state and control limits).
The efficient computation of these trajectory derivatives is important in the (open-loop) solution of optimal control problems for both trajectory optimization and model predictive control (MPC). Such derivatives are also crucial in inverse optimal control (also called inverse reinforcement learning or learning from demonstration), where given expert demonstration trajectories, the objective is to compute parameters of the cost function that best explain these trajectories. Extensions of inverse optimal control that involve inferring parameters of system dynamics in addition to cost functions also subsume system identification, and provide further motivation for the efficient computation of trajectory derivatives. Our proposed method of computing trajectory derivatives enables direct
(a) System Figure
(b) Imitation Learning: Rocket Landing
Figure 1: Left a): Our approach addresses learning problems in optimal control such as imitation learning. IDOC provides a method for computing the trajectory derivative or alternatively, vector-Jacobian products with respect to an outer-level loss, important for solving this problem using an end-to-end learning approach. Right b): IDOC yields superior trajectory derivatives for the imitation learning task when inequality constraints (rocket tilt and thrust limits) are present.
minimization of a loss function defined over optimal trajectories (e.g., imitation loss) using first-order descent techniques and is a direct alternative to existing works [2, 20, 22], including those based on bespoke derivations of Pontryagin's maximum principle (PMP) in discrete time [20, 22].
Analytical trajectory derivatives for COC problems are derived by differentiating through the optimality conditions of the underlying optimization problem and applying Dini's implicit function theorem. Derivatives are recovered as the solution to a system of linear equations commonly referred to as the (differential) KKT system. This framework underlies all existing works [2, 20, 22] as well as our proposed approach. Similarly to Amos et al. [2], we construct the differential KKT system, however we use the identities from Gould et al. [14] that apply variable elimination on the Lagrange multipliers relating to the dynamics and constraints. We show how to exploit the block-sparse structure of the resultant matrix equations to achieve linear complexity in computing the trajectory derivative with trajectory length. Furthermore, we can parallelize the computation, yielding superior scalability and numerical stability on multithreaded systems compared to methods derived from the PMP [20, 22].
Our specific contributions are as follows. First, we derive the analytical trajectory derivatives for a broad class of constrained (discrete-time) optimal control problems with additive cost functions and dynamics described by first-order difference equations. Furthermore, we show that the computation of these derivatives is linear in trajectory length by exploiting sparsity in the resulting matrix expressions. Second, we describe how to parallelize the computation of the trajectory derivatives, yielding lower computation time and superior numerical stability for long trajectories. Third, we show how to directly compute vector-Jacobian products (VJPs) with respect to some outer-level loss over optimal trajectories, yielding further improvements to computation time in the context of bi-level optimization. This setting commonly arises in the direct solution of inverse optimal control problems [17, 33, 34]. Finally, we provide discussion unifying existing methods for computing trajectory derivatives [2, 20, 22]. We dub our method IDOC (Implicit Differentiation for Optimal Control) and validate it across numerous experiments, showing a consistent speedup over existing approaches derived from the PMP. Furthermore, for constrained problems, IDOC provides significantly improved trajectory derivatives, resulting in superior performance for a general learning from demonstration (LfD) task. $^{1}$ .
Differentiating Through Optimal Control Problems. An optimal control (OC) problem consists of (system) dynamics, a cost function, an optimal control policy, and a set of state and/or control constraints. Learning problems in optimal control involve learning some or all of these aspects, with the problem of learning dynamics referred to as system identification [29], the problem of learning the cost function referred to as inverse optimal control (or inverse reinforcement learning or learning from demonstration) [34, 17, 21, 33, 37, 51], and the problem of learning the control policy referred to as reinforcement learning [4]. Recent work [20] has shown that solving these
$^{1}$ Code available at https://github.com/mingu6/Implicit-Diff-Optimal-Control
learning problems can be approached in a unified end-to-end fashion by minimizing task-specific loss functions with respect to (unknown) parameters of the associated OC problem. This unified formulation places great importance on the efficient computation of derivatives of optimal trajectories with respect to said parameters. Methods reliant on computing trajectory derivatives have, however, been mostly avoided (and argued against) in the robotics and control literature due to concerns about computational tractability. For example, bi-level methods of inverse optimal control [34, 33] have mostly been argued against in favor of methods that avoid derivative calculations by instead seeking to satisfy optimality conditions derived from KKT [24, 12] or PMP conditions [17, 32, 23, 21].
Analytical Trajectory Derivatives. Our method is a direct alternative to methods such as DiffMPC [2], PDP [20] and its extension Safe-PDP [22], which differentiate through optimality conditions to derive trajectory derivatives. Common to these methods is identifying that trajectory derivatives can be computed by solving an auxiliary affine-quadratic OC problem, and furthermore, that this can be done efficiently using a matrix Riccati equation. While our approach still differentiates the optimality conditions, we avoid solving matrix Riccati equations and show that this enables easy computation of VJPs, parallelization across timesteps and improved numerical stability. Like Safe-PDP, IDOC computes derivatives through COC problems with smooth inequality constraints, however we show in our experiments that our derivatives are significantly more stable during training. Section 4.3 provides a detailed description of the differences between IDOC and existing methods.
Differentiable Optimization. Our work falls under the area of differentiable optimization, which aims to embed optimization problems within end-to-end learning frameworks. Gould et al. [14] proposed the deep declarative networks framework, and provide identities for differentiating through continuous, inequality-constrained optimization problems. Combinatorial problems [31, 44], as well as specialized algorithms for convex problems [1] have also been addressed. Differentiable optimization has been applied to end-to-end learning of state estimation [49, 42] and motion planning problems [6, 2, 26] in robotics, see Pineda et al. [35] for a recent survey. In addition, numerous machine learning and computer vision tasks such as pose estimation [9, 40], meta-learning [28], sorting [11, 7] and aligning time series [47] have been investigated under this setting.
3 Constrained Optimal Control Formulation
In this section, we provide an overview of the COC problems we consider in this paper. These problems involve minimizing a cost function with an additive structure, subject to constraints imposed by system dynamics and additional arbitrary constraints such as state and control limits. As a result, they can be formulated as a non-linear program (NLP). We will discuss how to differentiate through these COC problems in Section 4.
First, we introduce notation around differentiating vector-valued functions with respect to vector arguments, consistent with Gould et al. [14]. Let  $f: \mathbb{R}^n \to \mathbb{R}^m$  be a vector-valued function with vector arguments and let  $\mathrm{D}f \in \mathbb{R}^{n \times m}$  be the (matrix-valued) derivative where elements are given by
$$
\left(\mathrm {D} f (x)\right) _ {i j} = \frac {\partial f _ {i}}{\partial x _ {j}} (x). \tag {1}
$$
For a scalar-valued function with (multiple) vector-valued inputs  $f: \mathbb{R}^n \times \mathbb{R}^m \to \mathbb{R}$  evaluated as  $f(x, y)$ , let the second order derivatives  $\mathrm{D}_{XY}^2 = \mathrm{D}_X(\mathrm{D}_Yf)^\top$ . See Gould et al. [14] for more details.
3.2 Optimization Formulation for Constrained Optimal Control
We can formulate the (discrete-time) COC problem as finding a cost-minimizing trajectory subject to constraints imposed by (possibly non-linear) system dynamics. In addition, trajectories may be subject to further constraints such as state and control limits.
To begin, let the dynamics governing the COC system at time  $t$  be given by  $x_{t + 1} = f_t(x_t,u_t;\theta)$  where  $x_{t}\in \mathbb{R}^{n}$ ,  $u_{t}\in \mathbb{R}^{m}$  denotes state and control variables, respectively<sup>2</sup>. For time horizon  $T > 0$ ,
The subscript  $f_{t}$  allows for time-varying dynamics in principle, though we do not evaluate IDOC on any optimal control problems that have this property in our experiments.
let  $x \triangleq (x_0, x_1, \ldots, x_T)$  and  $u \triangleq (u_0, u_1, \ldots, u_{T-1})$  denote the state and control trajectories, respectively. Furthermore, let  $\theta \in \mathbb{R}^d$  be the parameter vector that parameterizes our COC problem. Our COC problem is then equivalent to the following constrained optimization problem
$$
\text {m i n i m i z e} J (x, u; \theta) \triangleq \sum_ {t = 0} ^ {T - 1} c _ {t} \left(x _ {t}, u _ {t}; \theta\right) + c _ {T} \left(x _ {T}; \theta\right)
$$
$$
\begin{array}{l} x _ {0} = x _ {\mathrm {i n i t}} \\ x _ {t + 1} - f _ {t} \left(x _ {t}, u _ {t}; \theta\right) = 0 \quad \forall t \in \{0, \dots , T - 1 \} \quad \text {(d y n a m i c s)} \tag {2} \\ \begin{array}{l l} g _ {t} (x _ {t}, u _ {t}; \theta) \leq 0 & \quad \forall t \in \{0, \ldots , T - 1 \}, \\ h _ {t} (x _ {t}, u _ {t}; \theta) = 0 & \quad \forall t \in \{0, \ldots , T - 1 \}, \end{array} \qquad \begin{array}{l l} (\text {p a t h i n e q . c o n s t r a i n t s}) \\ (\text {p a t h e q . c o n s t r a i n t s}) \end{array} \\ g _ {T} (x _ {T}; \theta) \leq 0 \\ h _ {T} \left(x _ {T}; \theta\right) = 0, \quad (\text {t e r m i n a l} \\ \end{array}
$$
where  $f_{t}:\mathbb{R}^{n}\times \mathbb{R}^{m}\times \mathbb{R}^{d}\to \mathbb{R}^{n}$  describe the dynamics, and  $c_{t}:\mathbb{R}^{n}\times \mathbb{R}^{m}\times \mathbb{R}^{d}\to \mathbb{R}$ ,  $c_{T}:\mathbb{R}^{n}\times \mathbb{R}^{d}\to \mathbb{R}$  are the instantaneous and terminal costs, respectively. Furthermore, the COC system may be subject to (vector-valued) inequality constraints  $g_{t}:\mathbb{R}^{n}\times \mathbb{R}^{m}\times \mathbb{R}^{d}\to \mathbb{R}^{q_{t}}$  and  $g_{T}:\mathbb{R}^{n}\times \mathbb{R}^{d}\to \mathbb{R}^{q_{T}}$  such as control limits and state constraints, for example. Additional equality constraints  $h_t:\mathbb{R}^n\times \mathbb{R}^m\times \mathbb{R}^d\rightarrow \mathbb{R}^{s_t}$  and  $h_T:\mathbb{R}^n\times \mathbb{R}^d\to \mathbb{R}^{s_T}$  can also be included.
The decision variables of Equation 2 are the state and control trajectory  $\xi \triangleq (x,u)$ , whereas parameters  $\theta$  are assumed to be fixed. Equation 2 can be interpreted as an inequality-constrained optimization problem with an additive cost structure and vector-valued constraints for the initial state  $x_0$ , dynamics, etc. over all timesteps. Let  $\xi (\theta) \triangleq (x^{\star},u^{\star})$  be an optimal solution to Equation 2. We can treat the optimal trajectory  $\xi (\theta)$  as an implicit function of parameters  $\theta$ , since Equation 2 can be solved to yield an optimal  $\xi (\theta)$  for any (valid)  $\theta$ .
We can solve Equation 2 for the optimal trajectory  $\xi (\theta)$  using a number of techniques. We can use general purpose solvers [13, 45], as well as specialized solvers designed to exploit the structure of COC problems, e.g., ones based on differential dynamic programming [30, 19]. Regardless, for the purposes of computing analytical trajectory derivatives using our method (as well as methods that differentiate through optimality conditions [20, 22, 2]), we only need to ensure that our solver returns a vector  $\xi (\theta)$  which is a (local) minimizer to Equation 2. We now describe how to differentiate through these optimal control problems, important in the end-to-end learning context.
4 Trajectory Derivatives using Implicit Differentiation
In this section, we present our identities for computing trajectory derivatives  $\mathrm{D}\xi (\theta)$  based on those derived in Gould et al. [14] for general optimization problems by leveraging first-order optimality conditions and the implicit function theorem. We exploit the block structure of the matrices in these identities that arises in optimal control problems to enable efficient computation and furthermore, show that computing trajectory derivatives is linear in the trajectory length  $T$ .
Before we present the main analytical result, recall the motivation for computing  $\mathrm{D}\xi (\theta)$ . Suppose in the LfD context we have a demonstration trajectory  $\xi^{\mathrm{demo}}$  (we can extend this to multiple trajectories, but choose not to for notational simplicity). Furthermore, define a loss  $\mathcal{L}(\xi (\theta),\xi^{\mathrm{demo}})$  which measures the deviation from predicted trajectory  $\xi (\theta)$  to demonstration trajectory  $\xi^{\mathrm{demo}}$ . Ultimately, if we wish to minimize the loss  $\mathcal{L}$  with respect to parameters  $\theta$  using a first-order decent method, we need to compute  $\mathrm{D}_{\theta}\mathcal{L}(\xi (\theta),\xi^{\mathrm{demo}})$  by applying the chain rule. Specifically, we compute  $\mathrm{D}_{\theta}\mathcal{L}(\xi) = \mathrm{D}_{\xi}\mathcal{L}(\xi)\mathrm{D}\xi (\theta)$ , which requires trajectory derivative  $\mathrm{D}\xi (\theta)$ .
First, we first reorder  $\xi$  such that  $\xi = (x_0, u_0, x_1, u_1, \ldots, x_T) \in \mathbb{R}^{n_\xi \times 1}$ , where  $n_\xi = (n + m)T + n$ . This grouping of decision variable blocks w.r.t. timesteps is essential for showing linear time complexity of the computation of the derivative. Let  $\xi_t \in \mathbb{R}^{n + m}$  represent the subset of variables in  $\xi$  associated to time  $t$ , with final state  $\xi_T = x_T \in \mathbb{R}^n$ . Next, we stack all constraints defined in Equation 2 into a single vector-valued constraint  $r$  comprised of  $T + 2$  blocks. Specifically, let  $r(\xi; \theta) \triangleq (r_{-1}, r_0, r_1, \ldots, r_T) \in \mathbb{R}^{n_r}$ , where  $n_r = (T + 1)n + s + q$ . The first block is given by
$r_{-1} = x_0$ , and subsequent blocks are given by
$$
r _ {t} = \left\{ \begin{array}{l l} \left(\tilde {g} _ {t} \left(x _ {t}, u _ {t}; \theta\right), h _ {t} \left(x _ {t}, u _ {t}; \theta\right), x _ {t + 1} - f \left(x _ {t}, u _ {t}; \theta\right)\right) & \text {f o r} 0 \leq t \leq T - 1 \\ \left(\tilde {g} _ {T} \left(x _ {T}; \theta\right), h _ {T} \left(x _ {T}; \theta\right)\right) & \text {f o r} t = T. \end{array} \right. \tag {3}
$$
Here,  $\tilde{g}_t$  are the subset of active inequality constraints for  $\xi$  (detected numerically using a threshold  $\epsilon$ ) at timestep  $t$ . Note,  $s = \sum_{t=0}^{T} s_t$ ,  $q = \sum_{t=0}^{T} |\tilde{g}_t|$  are the total number of additional equality and active inequality constraints (on top of the  $Tn$  dynamics and  $n$  initial state constraints). Each block represents a group of constraints associated with a particular timestep (except the first block  $r_{-1}$ ). As we will see shortly, grouping decision variables and constraints in this way will admit a favorable block-sparse matrix structure for cost/constraint Jacobians and Hessians required to compute  $\mathrm{D}\xi(\theta)$ .
4.2 Analytical Results for Trajectory Derivatives
Proposition 1 (IDOC). Consider the optimization problem defined in Equation 2. Suppose  $\xi(\theta)$  exists which minimizes Equation 2. Furthermore, assume  $f_t$ ,  $c_t$ ,  $g_t$  and  $h_t$  are twice differentiable for all  $t$  in a neighborhood of  $(\theta, \xi)$ . If  $\text{rank}(A) = n_r$  and furthermore,  $H$  is non-singular, then
$$
D \xi (\theta) = H ^ {- 1} A ^ {\top} \left(A H ^ {- 1} A ^ {\top}\right) ^ {- 1} \left(A H ^ {- 1} B - C\right) - H ^ {- 1} B, \tag {4}
$$
$$
A = D _ {\xi} r (\xi ; \theta) \in \mathbb {R} ^ {n _ {r} \times n _ {\xi}}
$$
$$
B = D _ {\theta \xi} ^ {2} J (\xi ; \theta) - \sum_ {t = - 1} ^ {T} \sum_ {i = 1} ^ {| r _ {t} |} \lambda_ {t, i} D _ {\theta \xi} ^ {2} r _ {t} (\xi ; \theta) _ {i} \in \mathbb {R} ^ {n _ {\xi} \times d}
$$
$$
C = D _ {\theta} r (\xi ; \theta) \in \mathbb {R} ^ {n _ {r} \times d}
$$
$$
H = D _ {\xi \xi} ^ {2} J (\xi ; \theta) - \sum_ {t = - 1} ^ {T} \sum_ {i = 1} ^ {| r _ {t} |} \lambda_ {t, i} D _ {\xi \xi} ^ {2} r _ {t} (\xi ; \theta) _ {i} \in \mathbb {R} ^ {n _ {\xi} \times n _ {\xi}},
$$
and Lagrange multipliers  $\lambda \in \mathbb{R}^{n_r}$  satisfies  $\lambda^\top A = D_\xi J(\xi ;\theta)$ .
Proof. This is a direct application of Proposition 4.5 in Gould et al. [14].
Remark 1. Matrix  $H$  is block diagonal with  $T + 1$  blocks.
To see this,  $\hat{H} = \mathrm{D}_{\xi \xi}^{2}J(\xi ;\theta)$  has a block diagonal structure with blocks given by  $\mathrm{D}_{\xi_t\xi_t}^2 c_t(\xi_t;\theta)$ , which follows from the assumption of additive costs in the COC problem described in Equation 2. In addition, constraints  $g_{t}$ ,  $h_t$  depend only on  $\xi_{t}$  and furthermore, the dynamics constraint is first order in  $\xi_{t + 1}$ . Therefore,  $\mathrm{D}_{\xi \xi}^{2}r_{t}(\xi ;\dot{\theta})_{i}$  is only non-zero in the block relating to  $\xi_{t}$  for all  $t$  and  $i$ . We plot the block-sparsity structure of  $H$  in Figure 2a.
Remark 2. Matrix  $A$  is a block-banded matrix that is two blocks wide.
This is shown by noting that the only non-zero blocks in  $A$  correspond to  $\mathrm{D}_{\xi_t}r_t(\xi ;\theta)$  and  $\mathrm{D}_{\xi_{t + 1}}r_t(\xi ;\theta)$ . The former relates to  $g_{t}$ ,  $h_t$  and the  $-f_{t}(\xi_{t};\theta)$  component of the dynamics constraint. The latter only relates to the  $x_{t + 1}$  component of the dynamics constraint. Finally, the initial condition block  $r_{-1}$  only depends on  $\xi_0$ , hence the only non-zero block is given by  $\mathrm{D}_{\xi_0}r_{-1}(\xi_0;\theta)$ . We plot the block-sparsity structure of  $A$  in Figure 2b.
Proposition 2. Evaluating Equation 4 has  $O(T)$  time complexity.
Proof. From Remarks 1 and 2,  $H$  and  $A$  have a block diagonal and block-banded structure with  $T + 1$  and  $2T + 2$  blocks, respectively. Therefore, we can evaluate  $H^{-1}A^{\top}$  in  $O(T)$  time and furthermore,  $H^{-1}A^{\top}$  has the same structure as  $A$ . It follows that we can also evaluate  $AH^{-1}B - C$  in  $O(T)$  time after partitioning  $B$  and  $C$  into blocks based on groupings of  $\xi$  and  $r$ .
Next, observe that  $AH^{-1}A^{\top}$  yields a block tridiagonal matrix with  $T + 2$  blocks on the main diagonal (see Figure 2c for a visualization of the block-sparse structure). As a result, we can evaluate  $(AH^{-1}A^{\top})^{-1}(AH^{-1}B)$  by solving the linear system  $(AH^{-1}A^{\top})Y = AH^{-1}B$  for  $Y$  in  $O(T)$
(a) Matrices  $H$  and  $H^{-1}$
(b) Matrix  $A$
Figure 2: Block-sparse structure for matrices  $H$ ,  $A$ , and  $AH^{-1}A^{\top}$  assuming  $T = 3$ . For  $A$ , we combined inequality and equality constraint groups  $g_{t}$  and  $h_t$  into a single group for simplicity. Shaded regions indicate (possible) non-zero blocks.
(c) Matrix  $AH^{-1}A^{\top}$
time using any linear time (w.r.t. number of diagonal blocks) block tridiagonal linear solver. A simple example is the block tridiagonal matrix algorithm, easily derived from block Gaussian elimination.
Finally,  $H^{-1}B$  can easily be evaluated in  $O(T)$  time by solving  $T + 1$  linear systems comprised of the blocks of  $H$  and  $B$ . We conclude that evaluating Equation 4 requires  $O(T)$  operations.
4.3 Comparison to DiffMPC and PDP
All existing methods for computing analytical trajectory derivatives involve solving the linear system
$$
\underbrace {\left[ \begin{array}{l l} H & A ^ {\top} \\ A & 0 \end{array} \right]} _ {K} \left[ \begin{array}{l} D \xi \\ - D \lambda \end{array} \right] = \left[ \begin{array}{l} - B \\ - C \end{array} \right], \tag {5}
$$
where blocks  $H, A, B, C$  are defined in Section 4.2 and  $\mathrm{D}\xi$  is the desired trajectory derivative. We call Equation 5 the (differential) KKT system and the matrix  $K$  the (differential) KKT matrix. Note,  $\mathrm{D}\xi$  is unique if  $K$  is non-singular<sup>3</sup>. To better contextualize IDOC, we now briefly compare and contrast how previous methods such as DiffMPC [2] and PDP [20, 22] solve this KKT system.
DiffMPC. DiffMPC [2] proposes a method for differentiating through Linear Quadratic Regulator (LQR) problems, which are OC problems where the cost function is (convex) quadratic and the dynamics are affine. The authors show that solving the LQR problem using matrix Riccati equations can be viewed as an efficient method for solving a KKT system which encodes the optimality conditions. In addition, they show that computing trajectory derivatives involves solving a similar KKT system, motivating efficient computation by solving an auxiliary LQR problem in the backward pass. The matrix equation interpretation of the backward pass allows the efficient computation of VJPs for some downstream loss  $\mathcal{L}(\xi)$  in a bi-level optimization context.