SlowGuess/NeurIPS_2023_OCR huggingface dataset:
DiffMPC extends to handling non-convex OC problems with box constraints on control inputs via an iterative LQR-based approach. Specifically, such problems are handled by first computing VJPs w.r.t. the parameters of an LQR approximation to the non-convex problem around the optimal trajectory  $\xi, \lambda$ . Next, the parameters of the approximation are differentiated w.r.t. the underlying parameters  $\theta$  and combined using the chain rule. However, differentiating the quadratic cost approximation and dynamics requires evaluating costly higher-order derivatives, e.g.,  $\frac{d}{d\theta} \frac{d^2c(\xi)}{d\xi^2}$ , which are 3D tensors. These tensors are dense in general, although problem specific sparsity structures may exist.
PDP/Safe-PDP. PDP [20] and its extension, Safe-PDP [22], take a similar approach to DiffMPC in deriving trajectory derivatives. PDP derives trajectory derivatives by starting with PMP, which is well-known in the control community and applies to non-convex OC problems. Furthermore, Jin et al. [20] shows that the PMP and KKT conditions are equivalent in the discrete-time COC setting. Trajectory derivatives are obtained by differentiating the PMP conditions, yielding a new set of PMP conditions for an auxiliary LQR problem, which can be solved using matrix Riccati equations. This
See Section 10.1 in [8] for further discussion on conditions for non-singular  $K$
approach is fundamentally identical to DiffMPC, with only superficial differences for LQR problems (solving a differential KKT system versus differentiating PMP conditions) $^4$ .
For non-convex problems, PDP does not evaluate higher-order derivatives, unlike DiffMPC. However, it is not clear how to compute VJPs under approaches derived using the PDP framework. Safe-PDP extended PDP to handle arbitrary, smooth inequality constraints using the constrained PMP conditions, which extends the capability of DiffMPC to handle box constraints. Trajectory derivatives are computed by solving an auxiliary equality constrained LQR problem. Experiments in both Jin et al. [22] and Section 6 show that Safe-PDP yields unreliable derivatives for COC problems. While Jin et al. [22] conjected that instability was due to the set of active constraints changing between iterations, we find that IDOC still learns reliably in the presence of constraint switching. We instead observed that the poor gradient quality arises naturally from the equality constrained LQR solver used for the backward pass [25] not matching the solution obtained by directly solving the KKT system.
IDOC. In contrast, we solve Equation 5 for  $\mathrm{D}\xi$  by first applying variable elimination on  $\mathrm{D}\lambda$ , yielding Proposition 1. We will discuss in Section 5 how the equations resulting from this approach admit parallelization and favorable numerical stability compared to DiffMPC and PDP's auxiliary LQR approach. Furthermore, VJPs can be easily computed (unlike PDP) without evaluating higher-order derivatives (unlike DiffMPC). IDOC handles arbitrary smooth constraints like Safe-PDP, and we show in Section 6 that our derivatives for COC problems are reliable during training. However, we require the additional assumption that  $H$  is non-singular (which holds if and only if all blocks in  $H$  are non-singular), which is not required for  $K$  to be non-singular. An oftentimes effective solution when  $H$  is singular, initially proposed by Russell et al. [39], is to set  $H = H + \frac{\delta}{2} I$  for small  $\delta$ , which is analogous to adding a proximal term to the cost function in Equation 2. We will now describe how to evaluate Equation 4 in linear time by exploiting the block-sparse structure of the matrix equation.
5 Algorithmic Implications of IDOC
5.1 Parallelization and Numerical Stability
Parallelization. To evaluate Equation 4, we can leverage the block diagonal structure of  $H$  and block-banded structure of  $A$  to compute  $H^{-1}A^{\top}$  and  $H^{-1}B$  in parallel across all timesteps. Specifically, we evaluate the matrix product block-wise, e.g.,  $H_{t}^{-1}B_{t}$  for all  $t$  in parallel (similarly with  $A$ ). Following this, we can then compute  $A(H^{-1}A^{\top})$  in parallel also using the same argument.
In addition, we can solve the block tridiagonal linear systems involving  $AH^{-1}A^{\top}$  by using a specialized parallel solver [5, 36, 27, 41, 15, 18]. Unfortunately, none of these methods have open source code available, and so in our experiments, we implement the simple block tridiagonal matrix algorithm. Implementing a robust parallel solver is an important direction for future work, and will further improve the scalability and numerical stability of IDOC.
Numerical Stability. In addition to parallelization, another benefit of avoiding Riccati-style recursions for computing derivatives is improved numerical stability. For long trajectories and/or poorly conditioned COC problems (e.g., stiff dynamics), IDOC reduces the rounding errors that accumulate in recursive approaches. We show superior numerical stability in our experiments compared to PDP in Section 6, despite using the simple recursive block tridiagonal matrix algorithm for solving  $AH^{-1}A^{\top}$ . Replacing this recursion with a more sophisticated block tridiagonal solver should further improve the stability of the backwards pass and is left as future work.
5.2 Vector Jacobian Products
Another benefit of explicitly writing out the matrix equations for  $\mathrm{D}\xi (\theta)$  given in Equation 4 is that we can now directly compute VJPs given some outer loss  $\mathcal{L}(\xi)\in \mathbb{R}$  over the optimal trajectory. Let  $v\triangleq \mathrm{D}_{\xi}\mathcal{L}(\xi)^{\top}\in \mathbb{R}^{n_{\xi}\times 1}$  be the gradient of the loss w.r.t. trajectory  $\xi (\theta)$ . The desired gradient  $\mathrm{D}_{\theta}\mathcal{L}(\xi (\theta))$  is then given by  $\mathrm{D}_{\theta}\mathcal{L}(\xi) = v^{\top}\mathrm{D}\xi (\theta)$  using the chain rule. The resultant expression is
$$
\mathrm {D} _ {\theta} \mathcal {L} (\xi (\theta)) = v ^ {\top} \left(H ^ {- 1} A ^ {\top} \left(A H ^ {- 1} A ^ {\top}\right) ^ {- 1} \left(A H ^ {- 1} B - C\right) - H ^ {- 1} B\right). \tag {6}
$$
$^4$ Jin et al., [20] claim that the backward pass for DiffMPC is  $O(T^2)$  due to direct inversion of the KKT matrix. However, DiffMPC actually solves an auxiliary LQR problem for  $O(T)$  complexity.
(b) 6-DoF Quadrotor
(c) Robotarm
(d) 6-DoF Rocket
Figure 3: Learning curves for the imitation learning task over five trials. Bold lines represent mean loss, lighter lines represent individual trials. IDOC yields more stable gradients and lower final imitation loss across all environments and trials compared to both Safe-PDP (S-PDP) and Safe-PDP with log-barrier functions (S-PDP (b)).
The simple observation here is that we do not need to construct  $D_{\theta}\xi^{\star}(\theta)$  explicitly, and can instead evaluate the VJP directly. We propose evaluating Equation 6 from left to right and block-wise, which will reduce computation time compared to explicitly constructing  $\mathrm{D}\xi (\theta)$  and then multiplying with  $v$
To see this, we follow the example in Gould et al. [14] and assume the blocks in  $H$  have been factored. Then for a single block (ignoring constraints for simplicity), evaluating  $v^{\top}(H_t^{-1}B_t)$  is  $O((n + m)^2p)$  while evaluating  $(v^{\top}H_t^{-1})B_t$  is  $O((n + m)^2 + (n + m)p)$ . Evaluating the VJP directly significantly reduces computation time compared to constructing the full trajectory derivative for problems with higher numbers of state/control variables and tunable parameters.
While our contributions are largely analytical, we have implemented the identities in Equation 4 to verify our claims around numerical stability and computational efficiency on a number of simulated COC environments. We will show that IDOC is able to compute trajectory derivatives significantly faster compared to its direct alternative PDP [20] and Safe-PDP [22] with superior numerical stability.
We evaluate IDOC against PDP [20] and Safe-PDP [22] in an LfD setting across four simulation environments proposed in Jin et al. [22], as well as a synthetic experiment. The simulation environments showcase the gradient quality for a realistic learning task, whereas the synthetic experiment is designed to measure numerical stability and computation times. For Safe-PDP, we evaluate against the method proposed for inequality constrained problems (Safe-PDP), as well as using an approximate log-barrier problem in place of the full problem (Safe-PDP (b)). The IPOPT solver [45] is used in the forward pass to solve the COC problem, and Lagrange multipliers  $\lambda$  are extracted from the solver output. More generally however, note that the method proposed in Gould et al. [14] can be used to recover  $\lambda$  if another solver is used where  $\lambda$  is not provided.
Imitation Learning/LfD. The LfD setting involves recovering the model parameters  $\theta$  that minimizes the mean-squared imitation error to a set of  $N$  demonstration trajectories, given by
(a) Horizon length scalability.
(b) Model size scalability.
Figure 4: Synthetic experiments measuring a) scalability with horizon length  $T$ , with fixed parameter size  $|\theta| = 10\mathrm{k}$ , b) scalability with number of parameters  $d$  with fixed horizon length  $T = 1\mathrm{k}$ , c) numerical stability with varying block condition number  $\kappa(H_{t})$ . Error bars measure standard error across 5 and 25 samples for computation time (a, b) and numerical stability (c), respectively.
(c) Numerical stability.
$\mathcal{L}(\xi (\theta),\xi^{\mathrm{demo}})\triangleq \frac{1}{N}\sum_{i}\| \xi (\theta)_{i} - \xi_{i}^{\mathrm{demo}}\|^{2}$ . We include all parameters in the cost, dynamics and constraint functions in  $\theta$  and furthermore, evaluate performance both with (S-PDP) and without (PDP) inequality constraints. We perform experiments in four standard simulated environments: cartpole, 6-DoF quadrotor<sup>5</sup>, 2-link robot arm and 6-DoF rocket landing. For a detailed description of the imitation learning problem as well as each COC task, see the appendix. In addition, we provide timings for all methods within the simulation environments in the appendix.
Synthetic Benchmark. The simulation experiments described above are not sufficiently large-scale to adequately measure the benefits of parallelization and numerical stability afforded by IDOC over PDP. To demonstrate these benefits, we constructed a large-scale, synthetic experiment where the blocks required to construct  $H, A, B$  and  $C$  (as discussed in Section 4) are generated randomly. Computation time against the horizon length  $T$ , as well as the number of parameters  $d$  is reported. Numerical stability is measured by comparing mean-absolute error between trajectory derivatives evaluated under 32-bit and 64-bit precision for varying condition numbers over the blocks of  $H$ . State and control dimensions are fixed at  $n = 50$  and  $m = 10$ . Experiments are run on an AMD Ryzen 9 7900X 12-core processor, 128Gb RAM and Ubuntu 20.04. See the appendix for more details.
6.2 Imitation Learning/LfD Results
In this section, we evaluate the effectiveness of gradients produced from IDOC against PDP and Safe-PDP [20, 22] for the imitation learning/LfD task. In the equality constrained setting, IDOC and PDP yield identical (up to machine precision) trajectory derivatives and imitation loss throughout learning; see the appendix for more detailed results and analysis. However, as shown in Figure 3, when inequality constraints are introduced, the derivatives produced by Safe-PDP and IDOC differ significantly. Safe-PDP fails to reduce the imitation loss due to unreliable gradients, whereas IDOC successively decreases the imitation loss. While Safe-PDP (b), which differentiates through a log-barrier problem, also provides stable learning, it ultimately yields higher imitation loss compared to IDOC. This is because
the barrier problem is an approximation to the true COC problem. Further discussion around using log-barrier methods and COC problems is provided in Section 7.1.
Figure 5: Control Limit Learning.
6.3 Synthetic Benchmark Results
Figure 4 illustrates the results of the synthetic benchmark. We observed that IDOC and PDP have similar computation time and scalability with problem size when computing full trajectory derivatives. However, computing VJPs with IDOC significantly reduces computation time and improves scalability
<sup>5</sup>Interestingly,  $H$  is singular in this case, relating to block  $H_{T}$ , so we use the trick proposed in Section 4.3 with  $\delta = 10^{-6}$  to compute Equation 4. More details around this are provided in the appendix.
by an order of magnitude with model size (i.e.,  $d = |\theta|$ ). Numerical stability experiments show that IDOC (full and VJP) yields lower round-off errors and improved stability compared to PDP. Using more sophisticated block tridiagonal solvers will further improve stability.
7.1 Differentiating through Log-Barrier Methods
In the imitation learning setting, we assume that demonstration data  $\xi^{\mathrm{demo}}$  is an optimal solution to a COC problem with (unknown) parameters  $\theta^{\star}$ . Furthermore, we assume that a subset of timesteps  $\mathcal{A} \subseteq \{1, \dots, T\}$  yield active constraints. Under the log-barrier formulation, the constraint boundaries must be relaxed beyond their true values during learning to minimize the imitation loss. Concretely, under  $\theta^{\star}$ ,  $\log g_t(x_t^{\mathrm{demo}}, u_t^{\mathrm{demo}}; \theta^{\star})$  are undefined for  $t \in \mathcal{A}$ . Therefore to recover  $\xi(\theta) = \xi^{\mathrm{demo}}$ , we must have  $g_t(x_t(\theta), u_t(\theta)) < 0$  for  $t \in \mathcal{A}$ , i.e., we cannot recover the true constraint function through minimizing the imitation loss. We verify this using the robot arm environment, and present the results in Figure 5, plotting the error between the estimated and true constraint value. We observe that IDOC recovers the true constraint value more closely compared to Safe-PDP (b).
Generality of IFT. The generality of the differentiable optimization framework, and the matrix equation formulation for trajectory derivatives given in Equation 4 yield additional conceptual benefits which may help us tackle even broader classes of COC problems. For example, we can relax the additive cost assumption and add a final cost defined over the full trajectory such as
$$
h _ {T} (\xi ; \theta) = \xi^ {\top} Q \xi , \tag {7}
$$
where  $Q = CC^{\top}$  and  $C\in \mathbb{R}^{n_{\xi}\times k}$  for  $k\ll n_{\xi}$  is full rank. For IDOC, we can simply use the matrix inversion lemma to invert  $H$  in  $O(T)$ . However, this is more complicated for methods derived from the PMP which rely on the specific additive cost structure of the underlying COC problem.
Dynamic Games. A generalization of our work is to apply the differentiable optimization framework to handle learning problems that arise in dynamic games [3] and have recently begun to employ PDP-based approaches [10]. Being able to differentiate through solution and equilibrium concepts that arise in dynamic games enables the solution of problems ranging from minimax robust control [3] to inverse dynamic games [10, 33], and is a promising direction for future work.
Fast forward passes. While we have proposed a more efficient way of computing trajectory derivatives (i.e., the backward pass), we observe that the forward pass to solve the COC problem bottlenecks the learning process. Significant effort must be placed on developing fast solvers for COC problems for hardware accelerators to allow learning to scale to larger problem sizes.
Combining MPC and Deep Learning. A recent application for computing trajectory derivatives is combining ideas from deep learning (including reinforcement learning) with MPC [38, 50, 46, 48]. Common to these approaches is using a learned model to select the parameters for the MPC controller. By allowing for robust differentiation through a broader class COC problems compared to previous approaches, we hope that IDOC will allow for future development in this avenue of research.
In this paper, we present IDOC, a novel approach to differentiating through COC problems. Trajectory derivatives are evaluated by differentiating KKT conditions and using the implicit function theorem. Contrary to prior works, we do not solve an auxiliary LQR problem to efficiently solve the (differential) KKT system for the trajectory derivative. Instead, we apply variable elimination on the KKT system and solve the resultant matrix equations directly. We show that linear time evaluation is possible by appropriately considering equation structure. In fact, we show that IDOC is faster and more numerically stable in practice compared to methods derived from the PMP. We hope that our discussion connecting the fields of inverse optimal control and differentiable optimization will lead to future work which enables differentiability of a broader class of COC problems.
[39] C. Russell, M. Toso, and N. Campbell. Fixing implicit derivatives: Trust-region based learning of continuous energy functions. In Advances in Neural Information Processing Systems, volume 32, 2019.
[40] P.-E. Sarlin, D. DeTone, T. Malisiewicz, and A. Rabinovich. Superglue: Learning feature matching with graph neural networks. In Proc. of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), June 2020.
[41] S. K. Seal, K. S. Perumalla, and S. P. Hirshman. Revisiting parallel cyclic reduction and parallel prefix-based algorithms for block tridiagonal system of equations. Journal of Parallel and Distributed Computing, 73, 2013.
[42] P. Sodhi, E. Dexheimer, M. Mukadam, S. Anderson, and M. Kaess. LEO: Learning energy-based models in factor graph optimization. In Proc. of the Conference on Robot Learning (CoRL), 2021.
[43] R. Tedrake. Underactuated Robotics. 2023. URL https://underactuated.csail.mit.edu.
[44] M. Vlastelica, A. Paulus, V. Musil, G. Martius, and M. Rolínek. Differentiation of blackbox combinatorial solvers. In Proc. of the International Conference on Learning Representations (ICLR), 2020. *Equal Contribution.
[45] A. Wächter and L. Biegler. On the implementation of an interior-point filter line-search algorithm for large-scale nonlinear programming. Mathematical Programming, 106:25-57, 2006.
[46] W. Xiao, T.-H. Wang, R. Hasani, M. Chahine, A. Amini, X. Li, and D. Rus. Barriernet: Differentiable control barrier functions for learning of safe robot control. IEEE Transactions on Robotics, 39:2289-2307, 2023.
[47] M. Xu, S. Garg, M. Milford, and S. Gould. Deep declarative dynamic time warping for end-to-end learning of alignment paths. In Proc. of the International Conference on Learning Representations (ICLR), 2023.
[48] F. Yang, C. Wang, C. Cadena, and M. Hutter. iPlanner: Imperative Path Planning. In Proceedings of Robotics: Science and Systems, Daegu, Republic of Korea, July 2023. doi: 10.15607/RSS.2023.XIX.064.
[49] B. Yi, M. Lee, A. Kloss, R. Martin-Martin, and J. Bohg. Differentiable factor graph optimization for learning smoothers. In Proc. of the IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), 2021.
[50] M. Zanon and S. Gros. Safe Reinforcement Learning Using Robust MPC. IEEE Transactions on Automatic Control, 66:3638-3652, 2021.
[51] B. D. Ziebart, A. L. Maas, J. A. Bagnell, and A. K. Dey. Maximum entropy inverse reinforcement learning. In AAAI Conference on Artificial Intelligence, pages 1433–1438, 2008.
A Experimental Setup for Imitation Learning/LfD Task
We use the code and simulation environments provided by [22] and follow a similar experimental setup, described in more detail in this section. The LfD experimental setting without inequality constraints is referred to in [20] as the inverse reinforcement learning (IRL) task, whereas including inequality constraints is referred to in [22] as the constrained inverse optimal control (CIOC) task.
**Demonstration Trajectories.** Up to five demonstration trajectories are used in the LfD setting with no inequality constraints. Each demonstration trajectory is generated by solving a COC problem using the same underlying parameters  $\theta$ , however the initial conditions differ across trajectories and are sampled randomly. For the setting with inequality constraints present, we use only one demonstration trajectory.
Initialization. Consistent with [20, 22], we run five imitation learning trials for all LfD experiments, where for each trial,  $\theta$  is initialized by adding uniform noise to the true value. For all methods, we use the gradient descent with the same learning rate for a given environment.
Environment Specifications. Table 1 provides summary specifications of the simulation environments used in our experiments, which are consistent with prior works [20, 22]. Recall that  $n$ ,  $m$  denote the number of states and controls, respectively. The number of parameters is denoted  $d = |\theta|$  and  $T$  is the horizon length.
Table 1: Description of environments
Environment	n	m	d	TIRL	TCIOC
Cartpole	4	1	9	30	35
Quadrotor	13	4	11	50	25
Robotarm	4	2	10	35	25
Rocket	13	3	12	40	40
Additional Environment Parameters. Additional specifications are presented in Table 2. These include log-barrier parameter  $\gamma$  used for Safe-PDP (b), time discretization for Forward Euler to discretize continuous time dynamics, the learning rate for minimizing the imitation loss and finally the number of demonstration trajectories per environment. (E) refers to experiments without constraints, whereas (I) means inequality constraints are present.
Table 2: Additional hyperparameters for LfD experiments
Environment	γ	Δ	lr	ndemos
Cartpole (E)	-	0.1	10-4	5
Quadrotor (E)	-	0.1	10-4	2
Robotarm (E)	-	0.1	10-4	4
Rocket (E)	-	0.1	3 × 10-4	1
Cartpole (I)	0.01	0.1	8 × 10-5	1
Quadrotor (I)	0.01	0.15	2 × 10-4	1
Robotarm (I)	0.01	0.2	2 × 10-3	1
Rocket (I)	1	0.1	10-5	1
A.1 Cartpole
The cartpole environment relates to the swing-up and balance task, where a simple pendulum (i.e., massless pole and point mass on the end of the pole) is swinging on a cart and the objective is to balance the pendulum above the cart. Only a horizontal force can be applied to drive the cart.
Dynamics. The state variable  $x = [y, q, \dot{y}, \dot{q}]^{\top} \in \mathbb{R}^{4}$  is comprised of the horizontal position of the cart  $y$ , counter-clockwise angle of the pendulum  $\theta$  (from the hanging position) and their respective velocities. Control  $u \in \mathbb{R}$  relates to the horizontal force applied to the cart. Dynamics parameters  $\theta_{\mathrm{dyn}} = [m_c, m_lp, \ell]^{\top} \in \mathbb{R}^{3}$  relate to the mass of the cart  $m_c$ , mass of the point mass at the end of the pole  $m_p$  and length of the pole  $\ell$ . Gravity is set at  $g = 10$ . The continuous time dynamics are given by
$$
\ddot {y} = \frac {1}{m _ {c} + m _ {c} \sin^ {2} q} [ u + m _ {p} \sin q (\ell \dot {q} ^ {2} + g \cos q) ] \tag {8}
$$
$$
\ddot {q} = \frac {1}{\ell (m _ {c} + m _ {p} \sin^ {2} q)} \left[ - u \cos q - m _ {p} \ell \dot {q} ^ {2} \cos q \sin q - (m _ {c} + m _ {p}) g \sin q \right], \tag {9}
$$
with a detailed derivation provided in [43, Ch. 3].
Cost Functions. The cost function is a quadratic cost over states and controls. Specifically,
$$
c _ {t} \left(x _ {t}, u _ {t}; \theta\right) = \left(x _ {t} - x _ {\text {g o a l}}\right) ^ {\top} W \left(x _ {t} - x _ {\text {g o a l}}\right) + w _ {u} u ^ {2}, \tag {10}
$$
$$
c _ {T} \left(x _ {T}, u _ {T}; \theta\right) = \left(x _ {T} - x _ {\text {g o a l}}\right) ^ {\top} W \left(x _ {T} - x _ {\text {g o a l}}\right), \tag {11}
$$
with  $x_{\mathrm{goal}} = [0,0, - \pi ,0]^{\top}$ . We specify that  $W\triangleq \mathrm{diag}(w)$ , where  $w = [w_y,w_q,w_\dot{y},w_\dot{q} ]\geq 0$  and furthermore,  $w_{u}\geq 0$ . We have that  $\theta_{\mathrm{cost}} = [w^{\top},w_u^{\top}]^{\top}\in \mathbb{R}^{5}$ .