Revisiting Implicit Differentiation for Learning Problems in Optimal Control
Authors: Ming Xu
Timothy L Molloy
Stephen Gould
Venue: NeurIPS 2023 poster (2023)
Keywords: implicit differentiation
bi-level optimization; constrained learning and control; safe learning for control
TL;DR: We derive new identities for differentiating through constrained optimal control problems (e.g., trajectory optimisation). Our method is signficantly faster and more numerically stable than existing state-of-the-art.

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.

Revisiting Implicit Differentiation for Learning Problems in Optimal Control
Ming Xu
Timothy Molloy School of Engineering Australian National University timothy molloy@anu edu au
School of Computing Australian National University mingda. xučanu. edu.eedue... au
Stephen Gould School of Computing Australian National University s stephen. gould@anu. sephen.gould@anu. edu. au
Abstract
This paper proposes new method for differentiatingg through optimal trajectories arising from non-convex, constrained discrete-time optimal control (COC) prob- lems using the implicit function theorem (IFT). Previous works solve differential Karush-
(KKT) system for the trajectory derivative, and achieve this efficiently by solving an auxiliary Linear Quadratic Regulator (LQQR) problem. In contrast, 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, 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, learn- ing from demonstration benchmarks including a 6-DoF maneuvering quadrotor and 6-DoF rocket powered landing.
approach allows for easy parallelization,
Introduction
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 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 also crucial in inverse optimal control (also called inverse reinforeement 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
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Loss Initaaion
Cuuit
(5(0))
10a
Imitation loss
10a
Constrained Optimal Control p. IPOPT, DDP)
Trajectory VJPs (2((
10
{
10t
Parameters (e.g. NN)
100
0
100
150
lteration
(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 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 optimal- ity 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
well as proposed approach. Similarly to Amos et al. 22, we construct the differential KKT system, however use the identities from Gould et al. [4] that apply variable elimination the Lagrange multipliers relating to the dynamics and constraints. We show how to exploit the block-sparse structure of the resultant matrix equations 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-ting tisserete-tme) 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 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 general learning from demonstration (LfD) task
2 Related Work
Differentiating Through Optimal Control Problems. An optimal control (OC) problem consists of (system) dynamics, cost function, an optimal control policy, and a set of state and/or control constraints. Learning problems optimal control involve learning some or all of these aspects.
with the problem of learning dynamics referred to as system identification (2), 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 4J. Recent work 20 has shown that solving these
'Code available https//gthuu.ccommmnnng////ihiii Diaeiittiiif----iiii Diff-OOpiiallConttll

---

learning problems can be approached in 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 224 or PMP conditions [7 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 optimal- ity conditions to derive trajectory derivatives. Common to these methods is identifying that trajectory derivatives can be computed by solving an auxiliary affine-quadratie OC problem, and furthermore, that this can be done efficiently using matrix Riccati equation. While our approach still differenti- ates 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 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 al. [14 proposed the deep declarative networks framework, and provide identities for differentiating through continuous, well as specialized algorithms for convex problems [ have also been addressed. Differentiable optimization has been applied to end-to-end learning of state estimation [49]4]] and motion planning problems 6 [2, [26] in robotics, see Pineda et al. [55 for recent survey. In addition, numerous
constrained optimization problems. Combinatorial problems [31 44 as
machine learning and computer vision tasks such as pose estimation 9 40 meta-learning (28).
sorting
7 and aligning time series 47 have been investigated under this setting.
3 Constrained Optimal Control Formulation
In this section, provide overview of the COC problems 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 result, they can be formulated as a non-linear program (NLP). We will discuss how to differentiate through these COC problems in Section
3.1 Preliminaries
First, we introduce notation around differentiating vector-valued functions with respect to vector arguments, consistent with Gould et al. [14]. Let f:: R" + R" be vector-valued function with vector arguments and let Df E Raxm be the (matrix-valued) derivative where elements are given by
Ōf: (Df(x)),, axx (x).
For a scalar-valued function with (multiple) vector-valued inputs f: R" x R R evaluated as f v), let the second order derivatives D Dx (Dy T. 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 be given by I++1 = fa,, 44: f). where 2, E R", 24 Rm denotes state and control variables, respectively" For time horizon T > 0.
'The subscript fi allows for vime-varying dynamics in principle, though we do not evaluate IDOC on any optimal control problems that have this property in our experiments.
3

---

let (I......... and u _ (uo, vi...... "i....,uT-11 uT-1)
denote the state and control trajectories, respectively. Furthermore, let & E IR' be the parameter vector that parameterizes our COC problem. Our COC problem is then equivalent to the following constrained optimization problem
minimize subject to
JIz 0) A E Ca 4:00 + (IT;
zo Finit
(initial state)
F(+1 fe (2,, u4; 0) =0 A1 E (0....,7
(dynamics) (path ineq. constraints) (path eq. constraints) (terminal ineq. constraints) (terminal eq. constraints)
0r (2 Uci 0) 2
Vt E {0, (0..... T 13
he(ar, u4; 0)
Vt E T
gT (rTT
hr (x:r; 8) O,
where ft. : IR" x Rm x Rd - R" describe the dynamics, and c : R" x R x Rd + RR CT R" x Rd + R the instantaneous and terminal costs, respectively. Furthermore, the COC system may be subject to pector-valued, inequality constraints
IR" x R Rd + R and @r : R" Rd + RR such as control limits and state constraints, for example. Additional equality constraints h4 Rd x R x Rd + Rº* and hT IR" x IR + Rer can also be included.
The decision variables of Equation 2 are the state and control trajectory & - (z. whereas parameters are assumed to be fixed. Equation 2 can be interpreted as an inequality-constren optimization problem with additive cost structure and vector-valued constraints for the initial state To. dynamics, etc. over all timesteps. Let c(0) & (a* u²) be an optimal solution to Equation 2 We can treat the optimal trajectory E(0) as an implicit function of parameters 0, since Equation can be solved to yield optimal {(0) for any (valid) 0.
We can solve Equation for the optimal trajectory {()) 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 3D 19]. Regardless, for the purposes of computing analytical trajectory derivatives using our method (as well as methods that differentiate through optimality conditions (20 22 ) we only need to ensure that our solver returns vector (()) which is a (local) minimizer to Equation 2] We now describe how to differentiate
through these optimal control problems, important the end-to-end learning context.
4 Trajectory Derivatives using Implicit Differentiation
In this section, we present our identities for computing trajectory derivatives DE based on those derived in Gould 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 enable efficient computation and furthermore. show that computing trajectory derivatives is linear in the trajectory length T.
Before present the main analytical result, recall the motivation for computing (E))). Suppose in the LfD context we have a demonstration trajectory edemo (we can extend this to multiple trajectories, but choose not for notational simplicity). Furthermore, define a loss L(E(O), gdemo) which measures the deviation from predicted trajectory E(0) to demonstration trajectory gdemo Ultimately, we wish to minimize the loss C with respect to parameters 0 using a first-order decent method, we need to compute DzL(g( edemo) by applying the chain rule. Specifically, we compute D₂L(E) (E)),, which requires trajectory derivative (f).
4.1 Preliminaries
First, we first reorder & such that { - (mo. ti. Z1, Mi....,, .T) ) E R where nę = (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 E + represent the subset of variables in & associated to time I, with final state ŞT = IT E _2. Next, we stack all constraints defined in Equation 2 into single vector-valued constraint comprised of T + 2 blocks. Specifically, let r(E: 0) & (r-1, To, (T) .T) Rr. where n, (T + 1)n + + q. The first block is given by

---

20, and subsequent blocks are given by
(0r(21 ui 0), (a uei 0), 3++1 f( uy; for
for 0 215
(3)
(Gr((::: 0), (rr())))))))))
Here, gt. are the subset of active inequality constraints for & (detected numerically using a threshold c) at timestep Note, - 8t, q Σi= [ōel are the total number of additional equality and active inequality constraints (on top of the Tn dynamics and initial state constraints), Each block represents group of constraints associated with a particular timestep (except the first block r- As will see shortly, grouping decision variables and constraints in this way will admit favorable block-sparse matrix structure for cost/constraint Jacobians and Hessians required to compute DE(
4.2 Analytical Results for Trajectory Derivatives
Proposition (IDOC). Consider the optimization problem defined in Equation 2 Suppose {(O) exists which minimizes Equation 2 Furthermore, assume ft. Ct 9t and hy are twice differentiable for all t in neighborhood of (0. )). rank( (A) = n, and furthermore, H is non- non-singular, then
DE(O) = 1-111'A'
(AH -C) B
(4)
where
A Der(g; XnG
B = (E:00 Σ
A Dhereer e Reed
1-11
C Dor(5:0) R=.xd
rd
H (E:00 ΣΣ
E Ar, ri E R
and Lagrange multipliers X E R" satisfies A A (5:
Proof.
This is direct application of Proposition 4.5 in Gould et al. [14]
Remark 1. Matrix H is block diagonal with T + blocks.
To see this, Ā- D J(E 0) has block diagonal structure with blocks given by Dē.,, Ce 0). which follows from the assumption additive costs the COC problem described in Equation 2 In addition, constraints gt. fu depend only & and furthermore, the dynamics constraint is first order in f++1- Therefore, re(; 0); is only non-zero in the block relating to & for all 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 D rE; 0). The former relates to St. h and the (E( f) component of the dynamics con- straint. The latter only relates the zt-1 component of the dynamics constraint. Finally, the initial condition block r-1 only depends on Eo. hence the only non-zero block is given by D (5o: f). We plot the block-sparsity structure of in Figure 2b
correspond to D re.re(E: 8) and
Proposition 2. Evaluating Equation has O(T) time complexity.
Proof. From Remarks and 2, H and A have block diagonal and block-banded banded structure with T' and 2T' blocks, respectively. Therefore, we can evaluate - in O(T') time and furthermore. H AA has the same structure A. It follows that after partitioning and C into blocks based on groupings of & and r.
also evaluate AH B in O(T) time
Next, observe that AH-1AT yields block tridiagonal matrix with T + 2 blocks on the main diagonal (see Figure for a visualization of the block-sparse structure). As result, we can evaluate (AH- A (AH -B) by solving the linear system (AH--A - AT)Y AH- for Y in O(T)
5

---

Eo
Eo
E1
E2
T-1
0
7-11
F0
F1
72
Eo
0
0
0
7-1
0
0
10
D2.Fo
0
0
Y0
81
0
F1
Dgy
0
F2
Fg -Da
122
63
0 0
13
0
(a) Matrices H and H
(b) Matrix A
(c) Matrix AH-1
Figure 2: Block-sparse structure for matrices H, A, and ATT assuming T 3. For A, we combined inequality and equality constraint groups Ir and h, into single group for simplicity. Shaded regions indicate (possible) non-zero blocks.
time using any linear time (w.rt. 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 ++ linear systems comprised of the blocks of and B. We conclude that evaluating Equation requires (T) operations.
4.3 Comparison to DifTMPC and PDP
All cxisting methods for computing analytical trajectory derivatives involve solving the linear system
H
A
DE DA
- = -
(5)
A
0
K
where blocks A, B,C C' are defined in Section 4.2 and DE is the desired trajectory derivative. We call Equation 5 the (differential) KK7 system and the matrix K the (differential) KKT matrix. Note, DE is unique if K is non-singulaf To better contextualize IDOC, now briefly compare and contrast how previous methods such as DiffMPC and PDP 20
solve this KKT system.
DifMMPC DifIMPC [2] proposes method for differentiating through Linear Quadratic Regulator (LQR) problems, which are OC problems where the cost function is (convex) quadratic and the dynamies 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 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 L(E) in bi-level optimization context.
DiffMPC extends handling non-convex OC problems with box constraints control inputs via an iterative LQR-based approach. Specificially, such problems are handled by first computing VJPs the parameters of an LQQ approximation to the non-convex problem around the optimal trajec- tory e, Next, the parameters of the approximation are differentiated w.r.t. the underlying parameters 0 and combined using the chain rule. However, differentiating the quadratic cost approximation and d ddr(E)
dynamics requires evaluating costly higher-order derivatives, c.g., do dga 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, Sa [2], take similar approach to DiffMPC in deriving trajectory derivatives. PDP derives trajectory derivatives by starting with PMP, which 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 g for further discussion conditions for non-singular K

---

approach is fundamentally identical to DiftMPC, with only superficial differences for LQR problems (solving differential KKT system versus differentiating PMP conditions
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 DifIMPC 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 O 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 25 not matching the solution obtained by directly solving the KKT system.
backward pass
IDC. In contrast, we solve Equation 5 for DE by first applying variable elimination on DA, yielding Proposition[] [ We will discuss in Section 5 Show the equations resulting from this approach admit parallelization and favorable numerical stability compared to DifiMPC 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 O 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 be singulag An oftentimes effective solution when H/ is singular, initially proposed by Russell et al. [39], is to set H H + & for small 6, which is analogous to adding proximal term to the cost function in Equation 2 We will now describe how to evaluate Equation in linear time by exploiting the slock-sparse structure of the matrix equation.
5
Algorithmic Implications of IDOC
5.1 Parallelization and Numerical Stability
Parallelization. To evaluate Equation 4 can leverage the block diagonal structure of H and block-banded structure of A to compute "AT and H- in parallel across all timesteps. Specifi- cally, we evaluate the matrix product bloek-wise, e.g., H B, for all / parallel (similarly with A). Following this, we can then compute A(H-AAA in parallel also using the same argument.
In addition, we can solve the block tridiagonal linear systems involving AH--1A by using specialized parallel solver (5 [3, 27 [5] Unfortunately, none 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 recur- sions 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 Section 6 despite using the simple recursive block tridiagonal matrix algorithm for solving AH-1 AT. 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 ( given in Equation is that we can now directly compute VJPs given some outer loss L(E) E R over the optimal trajectory. Let A DeC(E)T== x be the gradient of the loss w.r.t. trajectory ((). The desired gradient ((E(O)) then given by DEL(E)= = DE()) using the chain rule. The resultant expression is
DoL(E(H)) = v (H AA (AH-'AAA - (AH B C) H BB.
(6)
'Jin et al.. 20 claim that the backward pass for DifIMPC is O due to direct inversion of the KKT matrix. However, DifiMPC actually solves auxiliary LQR problem for ((1') complexity.
7

---

10a
5
Imitation loss
104
110"10
Imitation loss
104
10
101
109
109
10"
10-;
10-1
10-2
100
150 lteration
200
250
300
250
500
750 Iteration
1000
1250 1500
(a) Caripole
(b) 6-DoF Quadrotor
Imitation loss
10a
109:
Imitation
10-1
10a
10-2
250
500 750 1000 1250 1500 Iteration
60
120
180
240
300
Iteration
(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 not need to construct DoE (0) 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 DE (e) and then multiplying with v.
To see this, we follow the example in Gould et al. [4] and assume the blocks in H have been factored. Then for a single block (ignoring constraints for simplicity). evaluating T (H, B,) is O( + m)?p) while evaluating (v" H, B, is O(In + m)² + (n + m)p). Evaluating the VJP directly significantly reduces computation time compared constructing the full trajectory derivative for problems with higher numbers of state/control variables and tunable parameters.
Experiments
While contributions are largely analytical, have implemented the identities in Equation to verify our claims around numerical stability and computational efficiency on a number of simulated COC environments. We will show that IDOC is able compute trajectory derivatives significantly faster compared to its direct alternative PDP 20 and Safe-PDP 22 with superior numerical stability.
6.1 Experimental Setup
We evaluate IDOC against PDP 20 and Safe-PDP 22 in an LfD setting across four simulation environments proposed Jin al. [22 as well synthetic experiment. The simulation environ- ments showcase the gradient quality for 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 incquality constrained problems (Safe-PDP), as well using approximate log-barrics problem in place of the full problem (Safe-PDP (b)). The IPOPT solver [45) is used the forward pass solve the COC problem, and Lagrange multipliers À are extracted from the solver output. More generally however, note that the method proposed Gould et al. 14 can be used to recover À if another solver is used where À is not provided.
Imitation Learning/I.MD. The LfD setting involves recovering the model parameters that minimizes the mean-squared imitation error to set of N demonstration trajectories, given by

---

101
101
Computation Time (8)
PDP
10"
IDOC Full
10
IDOC vjp
10
PDP
PDP
10
10-1
10-1
IDOC
Full
10a
IDOC VJP
IDOC VJP
10a
2000
400
(00
800
100
2000
4000
B100
900
10000
6 log(x(H))
Horizon
Num. Parameters (4)
(a) Horizon length scalability,
(b) Model size scalability.
(c) Numerical stability.
Figure 4: Synthetic experiments measuring a) scalability with horizon length T, with fixed parameter size |0) = 10k, b) scalability with number of parameters with fixed horizon length T = 1k. c) numerical stability with varying block condition number K k(H)). Error bars measure standard error across 5 and 25 samples for computation time (a, b) and numerical stability (c), respectively.
C(E(0), )demm) N E; - |2 We include all parameters in the cost, dynamics and constraint functions in O and furthermore, evaluate performance both with (S-PDP) and without (PDP) inequality constraints. We perform experiments four standard simulated environments: cartpole, 6-DoF quadroto 5 2-link robot arm and 6-DoF rocket landing. For detailed description of the imitation learning problem as well as each COC task, see the appendix. 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 is reported. Numerical stability 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 fixed
50 and 10. Experiments 7900X 12-core processor, 128Gb RAM and Ubuntu 20.04. See the appendix for more details.
AMD Ryzen
6.2 Imitation Learning/LfD Results
In this section, evaluate the effectiveness of gradients pro- duced from IDOC against PDP and Safe-PDP [20 [22) for the imitation learning/LiD task. In the equality constrained setting, IDOC and PDP yield identical (up to machine precision) tra- jectory derivatives and imitation loss throughout learning: the appendix for more detailed results and analysis. However, shown in Figure 3 when inequality constraints intro- duced, the derivatives produced by Safe-PDP and IDOC differ significantly. Safe-PDP fails reduce the imitation loss due unreliable gradients, whereas IDOC successively decreases the imitation loss. While Sa (b), which differentiates through a log-barrier problem, also provides stable learning, it ultimately yields higher imitation loss compared IDOC. This is because the barrier problem is an approximation to the true COC problem. Further discussion around using
10-1
c10 - gtrue
5.00 Ja1
20-5
10-7
250
500
1250 150x Iteration
750 1000
Figure 5: Control Limit Learning.
log barrier methods and COC problems provided in Section
6.3 Synthetic Benchmark Results
Figure 4 iilustrates the results 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
'Interestingly. H is singular in this case, relating to block Hr. we use the trick proposed in Section 4.3 with 8 10 to compute Equation More details around this are provided in the appendix.

---

by an order of magnitude with model size (i.e., d = |01). Numerical stability experiments show that IDOC (full and VJP) yields lower round-off errors and improved stability compared PDP. Using more sophisticated block tridiagonal solvers will further improve stability.
Discussion and Future Work
7.1 Differentiating through Log-Barrias Methods
In the imitation learning setting, we assume that demonstration data edemo is an optimal solution to COC problem with (unknown) parameters 0*. Furthermore, we assume that a subset of timesteps A S (I,...,T") 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 8*, log 90 (z:(mm 246 uelemo 8*)) are undefined for E A. Therefore to recover {(0) galemo,
must have gi (0), (0)) < for t E A. i.c., 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 yield additional conceptual benefits which may help tackle even broader classes of COC problems. For example, can relax the additive cost assumption and add a final cost defined over the full trajectory such as
ht (E: QE,
where Q= CC and C' E R for & & ng is full rank. For IDOC, we can simply use the matrix inversion lemma to invert 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 frame- work to handle learning problems that arise in dynamic games 3 and have recently begun to employ PDP- 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 and is a promising direction for future work.
3 to inverse dynamic games [0 33
Fast forward passes. While we have proposed 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 allow learning to scale 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 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.
8 Conclusion
In this paper, we present IDOC, a novel approach to differentiating through COC problems. Tra- jectory 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 broader class of COC problems.
10

---

References
[1] A. Agrawal, B. Amos, S. Barratt, S. Boyd, S. Diamond, and Z. Kolter. Differentiable convex optimization layers. In Advances in Neural Information Processing Systems, 2019.
[2] B. Amos, I. Jimenez. J. Sacks, B. Boots, and J. Z. Kolter. Differentiable MPC for End-to-end Planning and Control. Advances in Neural Information Processing Systems, 2018.
[3] T. Başar and G. J. Olsder. Dynamic noncooperative game theory. Academic Press, London/New York, 1982. ISBN 978-0-08--996666-4
[4] D. Bertsekas. Reinforcement learning and optimal control. Athena Scientific, 2019.
[5] R. Bevilacqua, B. Codenotti, and Romani. Parallel solution of block tridiagonal linear systems. Linear Algebra and its Applications, 104:39-57, 1988.
[6] M. Bhardwaj. B. Boots, and M. Mukadam. Differentiable Gaussian process motion planning. Proc. of the IEEE International Conference on Robotics and Automation (ICRA), 2020.
[7] M. Blondel, O. Teboul, Q. Berthet, and J. Djolonga. Fast differentiable sorting and ranking. In Proc. of the International Conference on Machine Learming (ICML), 2020.
[8] S. Boyd and L. Vandenberghe. Convex Optimization. Cambridge University Press, 2004, doi: 10. 10.1017/CB0978051 1804441.
[9] D. Campbell, I. Liu, and S. Gould. Solving the Blind Perspective-n-Poinn Problem End-To-Ein With Robust Differentiable Geometric Optimization. In Proc. of the European Conference on Computer Vision (ECCV), 2020.
[10] K. Cao and L. Xie. Theoretic Inverse Reinforcement Learning: A Differential Pontrya- gin's 's Maximum Principle Approach. IEEE Transactions on Neural Networks and Learning Systems, pages 1-8, 2022. doi:
[11] A. Cherian, B. Fernando, M. Harandi, and S. Gould. Generalized rank pooling for activity recognition. In Proc. of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2017.
[12] P. Englert, N. A. Vien, and M. Toussaint. Inverse KKT: Learning cost functions of manipulation tasks from demonstrations. The International Journal of Robotics Research, 36(13-14):1474 1488, 2017. doi: 10. 00.117702783349114599988
[13] E. Gill, W. Murray, and M. A. Saunders. SNOPT: An SQP Algorithm for Large-Scarl Constrained Optimization. SIAM Review, page 99-1919 2005.
[14] S. Gould, R. Hartley and D. Campbell. Deep Declarative Networks. IEEE Transactions on Puttern Analysis and Machine Intelligence, 44:3988-400X Aug 2022. doi: 10.1 109TTAMM.. 2021 1. 1335944....
[15] I. N. Hajj and S. Skelboe. A multilevel parallel solver for block tridiagonal and banded linear systems. Parallel Computing, 15:21-45, 1990.
[16] C. R. Harris, K. J. Millman, S. J. van der Walt, R. Gommers, Virtanen, D. Cournapeau, E. Wieser, J. Taylor, S. Berg, N. J. Smith, R. Kern, M. Picus, S. Hoyer, M. H. van Kerkwijk, del Río, M. Wiebe, Peterson, T. Reddy. W. Weckesser, H. Abbasi, C. Gohlke, and T. E. Oliphant. Array programming with NumPy, Nature, 585:357-362, 2020.
M. Brett, Haldane,
K. Sheppard,
[17] K. Hatz, J. P. Schlöder, and H. G. Bock. Estimating parameters in optimal control problems. SIAM Journal on Scientific Computing, 34(3)A1107-A1122, 2012. doi: 10. 137/1 10823390.
[18] S. Hirshman, K. Perumalla, V. Lynch, and R. Sanchez. Beyclic: A parallel block tridiagonal matrix cyclic solver. Journal of Computational Physics, 229:6392-6400 2010.
[19] T. A. Howell, B. E. Jackson, and Z. Manchester. ALTRO: A Fast Solver for Constrained Trajectory Optimization. In Proc. of the IEEF/RSJ International Conference on Intelligent Robots and Systems (IROS), pages 7674-7679, 2019.

---

[20] W. Jin, Z. Wang, Z. Yang, and S. Mou. Pontryagin Differentiable Programming: An End-to-End Learning and Control Framework. Advances Neural Information Processing Systems, 33: 7979-7992, 2020.
[21] Jin, D. Kulić, S. Mou, and S. Hirche. Inverse optimal control from incomplete trajectory observations. The International Journal of Robotics Research, 40(6-7):848 865, 2021. doi: 10.1 111777007783492219966
[22] W. Jin, S. Mou, and G. J. Pappas. Safe Pontryagin Differentiable Programming. Advances in Neural Information Processing Systems, 1011666465000 2021.
[23] M. Johnson, N. Aghasadeghi, and T. Bretl. Inverse optimal control for deterministic continuous- time nonlinear systems. In Decision and Control (CDC), 2013 IEEE 52nd Annual Conference on, pages 2906-2913, Dee 2013.
[24] A. Keshavarz, Wang, and S. Boyd. Imputing a convex objective function. Intelligent Control (ISIC), 2011 IEEE International Symposium on, pages 613-619. IEEE, 2011.
[25] F. Laine and C. Tomlin. Efficient computation of feedback control for equality lqr. Proc. of the IEEE International Conference on Robotics and Automation (ICRA), 2019. doi: 11109/CCA..2199.8775..
[26] B. Landry, Z. Manchester, and M. Pavone. A differentiable augmented lagrangian method for bilevel nonlinear optimization. Robotics: Science and Systems, 2019.
[27] E. László, M. Giles, and Appleyard. Manycore algorithms for batch scalar and block tridiagonal solvers. ACM Transactions on Mathematical Software, 42(4), 2016.
[28] K. Lee, S. Maji, A. Ravichandran, and S. Soatto. Meta-learning with differentiable convex optimization. In Proc. of the IEEE/CVF Conference on Computer Vision and Puttern Recognition (CVPR), 2019.
[29] L. Ljung. System identification: theory for the user. Prentice Hall PTR, Upper Saddle River, NJ, 1999.
[30] C. Mastalli, R. Budhiraja, W. Merkt, G. Saurel, B. Hammoud, M. Naveau, J. Carpentier, L. Righetti, S. Vijayakumar, and N. Mansard. Crocoddyl: An Efficient and Versatile Framework for Multi-Contact Optimal Control. In Proc. of the IEEE International Conference on Robotics and Automation (ICRA), 2020.
[31] A. Mensch and M. Blondel. Differentiable dynamic programming for structured prediction and attention. In Proc. of the International Conference on Machine Learning (ICML), 2018.
[32] L. Molloy, J. J. Ford, and Perez. Finite-horizon inverse optimal control for discrete-time nonlinear systems. Automatica, 87:442 - 446, 2018. doi:
[33] T. L. Molloy, J. Inga Charaja, S. Hohmann, and T. Perez. Inverse Optimal Control and Inverse Noncooperative Dynamic Game Theory; Minimum- Princppe- Approach. Springer International Publishing, 2022.
[34] K. Mombaur, A. Truong, J.- J.PP. Laumond. From human to humanoid locomotion- an inverse optimal control approach. Autonomous robots, 28(3):369-383 2010.
[35] L. Pineda, T. Fan, M. Monge, Venkataraman, P. Sodhi, R. T. Chen, Ortiz, D. DeTone, A. Wang, S. Anderson, J. Dong, B. Amos, and M. Mukadam. Theseus: A Library for Dif- ferentiable Nonlinear Optimization. Advances in Neural Information Processing Systems, 2022.
[36] E. Polizzi and A. H. Sameh. A parallel hybrid banded system solver: the spike algorithm. Parallel Computing, 3272):177-194, 2006.
[37] N. D. Ratifff, J. A. Bagnell, and M. A. Zinkevich. Maximum margin planning. In Proceedings of the 23rd international conference Machine learning, pages 729-736. ACM, 2006.
[38] A. Romero, Y. Song. and D. Scaramuzza. Actor-critic model predictive control, 2023.
12

---

[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] Sodhi, E. Dexheimer, M. Mukadam, Anderson, and M. Kaess. LEO: Learning energy- models in factor graph optimization. In Proc. of the Conference on Robot Learning (CoRL), 2021.
[43] R. Tedrake. edu
Underactuated Robotics. 2023. URL http:////mddaaaattttteddd. csail. mit
[44] M. Vlastelica, A. Paulus, Musil. G. Martius, and M. Rolinck. 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 Gould. Deep declarative dynamic time warping for end- to-end learning of alignment paths. In Proc. of the International Conference Learning Representations (ICLR), 2023.
[48] 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. 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 Artificial Intelligence, pages 1433-1438, 2008.
13

---

A
Experimental Setup for Imitation Learning/LID Task
We use the code and simulation environments provided by [22 and follow 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 22 as the constrained inverse optimal control (CIOC) task.
Demonstration Trajectories. Up 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 0, 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, 0 is initialized by adding uniform noise to the true value. For all methods, use the gradient descent with the same learning rate for a given environment.
Environment Specifications. Table I provides summary specifications of the simulation environ- ments used 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 [0] and T is the horizon length.
Table 1: Description of environments
Environment
d
TiRL
Teioc
Cartpole
1
9
30
35
Quadrotor
13
4
11
50
25
Robotarm
4
2
10
35
25
Rocket
12
40
40
Additional Environment Parameters. Additional specifications are presented in Table 2 These include log-barrier parameter y used for Safe-PDP (b), time discretization for Forward Fuler to discretize continuous time dynamics, the learning rate for minimizing the imitation loss and finally the number of demonstration trajectories per environment. (E) refers experiments without constraints, whereas (I) means inequality constraints are present.
Table 2: Additional hyperparameters for LfD experiments
Environment
A
lr
7 delmms
Cartpole (E)
0.1
10-1
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
Quadrotor (I)
0.01
0.15
2 10-4
Robotarm (I)
0.01
0.2
10-3
Rocket (I)
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 horizontal force can be applied to drive the cart.
14

---

Dynamics. The state variable ly. 9. v. q] E E+ is comprised of the horizontal position of the cart y. counter-clockwise angle of the pendulum 9 (from the hanging position) and their respective velocities. Control E R relates to the horizontal force applied to the cart. Dynamics parameters Oayn = Mee mip, 8] R23 relate to the the mass of the cart m., mass the point mass at the end of the pole my and length of the pole f. Gravity is set at 10. The continuous time dynamics are given by
i = me + me sin
[u + mp sin q(q"² cos q)]
(8)
q= E(me + mp sin q)
[-ucosq mpf cos qii (m, + mplgsing)
with detailed derivation provided in
43 Ch. 3].
Cost Functions.
The cost function a quadratic cost over states and controls. Specifically,
ce (1 uz; 0) = (ate Igooll TW (xx: Fggaa) + wwu?
(10)
and
cr (2r, ur: 0) = (ET Igoal) "W(x Zggua);
(11)
with Zigoal 00 -TT, OOT. We specify that W diag(w), where [Wy, Wyy Wyj wal 2 0 and furthermore, Ww O. We have that Ooost = [w" w. R
Constraints. For the setting with constraints, we apply box constraints for the cart position and controls. Specifically, enforce lyl S 3max and |u| S Umax. We let constraint parameters Ownsstr
[Ymax: Umax] The final parameter vector is given by 0 [0 0 0
A.2 6-DoF Quadrotor Manouvering
See Section E of the appendix in 20 for the full definition of the quadrotor maneuvering problem. The objective is to drive the quadrotor using propeller thrusts to a goal configuration. The state is given by [p T, .p' .q , JT 1R where p E R represents the position of the center-of-mass of the quadrotor, E RR represents linear velocity, E R4 is a unit quaternion representation of the quadrotor attitude and R represents angular velocity. For the inequality constrained setting, non-linear (squared-norm) constraints on quadrotor position, given by l!pll S T'ax and thrust limits are applied.
A.3 2-Link Robot Arm
See Section E in the appendix in 20 for reference the full description of the robot problem. The objective of this OC problem is to move the robot arm to a goal configuration. The state given by the orientation of the base and elbow joint, well respective velocities 1Q1-92-00 41, 42)" R4. The control inputs inequality constrained setting relate to box constraints on the joint positions, as well as torque limits. The cost function is similar to cartpole, i.c., quadratic penalty to goal state and over the controls.
[ug. Ug] T IR² relate to applying torques each joint. Constraints for the
A.4 6-DoF Rocket Landing
See Section in the appendix in 20 for the full definition of the 6-DoF rocket landing problem. The objective this task is to land rocket modeled as a rigid body (softly) onto goal position, using thrusters located at the tail. A non-linear constraint over the tilt angle is applied in the inequality constrained setting. The cost function penalizes deviations from the goal configuration, as well as fuel cost.
15

---

1m-13
6004
3080
D800
X0x
(a) Cartpole
(b) 6-DoF Quadrotor
(c) Robotarm
(d) 6-DoF Rocket
Figure 6: Results for the imitation learning task with inequality constraints present. The y-axis the (percentage) relative difference the imitation loss, defined to (IDOC loss PDP loss) PDP loss (lower is better). For Cartpole, locations where the PDP gradient fails is marked by stars. Shaded regions indicate min/max values over five trials.
B Experimental Setup for Synthetic Benchmark
We discuss in more detail how blocks for H, A, B, C are generated in this section. Elements for the Jacobian and Hessian blocks for the objective function and constraints are sampled independently and identically distributed (ii.d.) from standard Gaussian. Hessian blocks are made symmetric using 4 (H + )/2, where is the initially generated random matrix. Each block's condition number is modified by applying an SVD to each block and adjusting the diagonal entries. Elements of the downstream loss gradients = DgL(E) are also sampled i...d. from a standard Gaussian.
C Additional Results for the Imitation Learning/LID Tasks
In this section, we provide additional detailed analysis around the imitation learning/L.fD setting with no inequality constraints. Figure 6 illustrates the results of these experiments. We plot the percentage difference in loss across the full training curve between IDOC and PDP for five random initializations for For the quadrotor, robotarm and rocket environments, we observe almost no difference in learning curves. This is expected since IDOC and PDP computing the same trajectory derivative. For cartpole however, we see that PDP fails catastrophically (resulting in undesirable spikes in the imitation loss) on two occasions across the five trials. We suspect this is due to numerical instability of the Riccati equations, arising from the local geometric properties of the dynamics at solved trajectories.
D Additional Timings for Simulation Experiments
In this section, we provide additional compute time experiments for the imitation learning/L.fDD tasks evaluated in the simulation environments. Given the relatively small size for the COC problems, there very limited advantage with respect to compute time for parallel evaluation of trajectory derivatives. However, proxy parallelization, IDOC can be vectorized across timesteps using the numerical linear algebra library Numpy [6]. All experiments in this section are run on single thread of an AMD Ryzen 7900X 4.7Ghz desktop CPU.
The experimental setting where inequality constraints are present challenging to vectorize compared to the setting with only equality constraints, because different timesteps may have different numbers of active inequality constraints. While matrices A and presented Equation are still block structured, the blocks may not be of uniform size in this setting. We present in Figure 7a and /bb the mean computation time (along with the upper/lower bound) per itcration across five trials with 10k iterations for each trial for all environments except for rocket (Ik). We observe around 2 speedup over PDP by computing our IFT/DDN gradients, even before direct computation of vector-Jacobian products, Computing VJPs directly yields further (approx.) 10% saving in compute time over constructing the trajectory derivative explicitly.
In our non-optimized Python implementation of IDOC, we batch together and vectorize computations involving blocks with an identical number active constraints. As expected, trajectory derivatives for inequality constrained problems are slightly slower to compute compared to the equivalent without (hard) inequality constraints (log-barrier approximation), due the necessary computational overheads required for identifying and batching blocks.
16

---

Note that Safe-PDP [22 (for hard inequality constraints) and PDP [20 (used for differentiating through log-barn bog-brrrir problems) derivatives are implemented using a custom LQR and equality con- strained LQR solver, respectively. These solvers have markedly different implementations, which
explains differences in computation time between Safe-PDP and PDP.
Compe Lakin Eime (mn)
Computation (ma)
2.5
Sate-FoP
2.5
2.0
00
Carbole
50
Cartuole
(a) Hard Constraints
(b) Log- Barrier
Figure 7: Computation times for trajectory derivatives (per trajectory, per iteration) for the CIOC task. As expected, the log-barrier method with no "hard" inequality constraints affords faster derivative computation.
17


Reviews:
[{"review_id": "Df5ZZ0jOwh", "reviewer": "Reviewer_ZYtx", "summary": "The paper introduces a method for calculating analytic trajectory gradients in constrained optimal control problems using implicit differentiation, with following contributions:\n* Shows that computation of these derivatives can be linear in trajectory time-steps, utilizing the structure in the matrices in the gradient computation.\n* Shows how to parallelize the gradient computation for reduced overall computation time and better numerical stability. \n* Shows direct computation of vector-jacobian products for find optimal trajectories with some outer loss. \n\nThe paper utilizes results from previous works [15,18] and uses insights from the those results to build incremental contributions. Though incremental, the insights can be practically useful.", "questions": "I have a few questions related to the significance of this work. If you can help me with following questions to better place the paper in the field of optimal control and RL. \n1. Is the method introduced only limited to LfD setting or are there any other setting in RL or optimal control this method can be used for? \n2. Can you place this work with respect to model-based/model-free/inverse RL methods? i.e if this method or any part of this method can be used for policy learning or planning or model learning? ", "limitations": "Limitations:\n* No clear experiments about how the method scales with more challenging and complex environments such as DM control suit, etc. ", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "Strenghts:\nThe contribution of the method can be practically useful for trajectory gradient methods allowing for faster computation.    \nOriginality - The insights and the contributions for the paper are original to best of my knowledge. Using block sparse structure of gradient computation for linear time complexity in time-steps is the main original idea.   \nQuality - The quality of contribution is moderate. Since the paper is mainly is analytical, built on insights based on previous results and the speed-up is only 2X.  \nClarity - I believe the writing of the paper can be more clear, especially section 3 onwards. Subsection 4.2 needs more work to be easy to parse.   \nSignificance - Significance is moderate. Since the speed is only 2X, and also in figure 3, the improvements over baselines are minor in 3 out of 4 environments. In Figure 4 as well, the final final imitation loss seems very close for different methods, though the gradient quality indeed seems better. ", "weaknesses": "Major Weaknesses:\n1. The paper lacks comparison on other experimental setting such as sysID and control from the baseline paper [18]. \n2. It seems the relevance for the paper is mostly in the setting where trajectory gradients wrt to model parameters or objective parameters is required such as LfD. For a holistic comparison, there should be benchmarking wrt to current important methods specialized for SysID (model learning in MBRL) or inverse RL. \nMinor Weaknesses:\n1. Improve the legends in figure 3, and enlarge figure 1(a). In general, all the figures needs improvement. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. We will gladly address editorial comments around legends and writing. In addition, we will address the concerns you have raised under both **Strengths** and **Weaknesses**.\n\n> **Quality - The quality of contribution is moderate. Since the paper is mainly is analytical, built on insights based on previous results and the speed-up is only 2X.**\n\nWhilst acknowledging that our contribution is largely analytical, we note that it holds for a very general class of problems with the state-of-the-art split over two papers (PDP and Safe-PDP). Furthermore, we believe clarifying misconceptions around the implicit function theorem (see our response to reviewer tsv7) is of value to both the learning and control communities.\n\nWe hope our new numerical experiments in the global response changes your opinion around the relative speed up!\n\n> **Significance - Significance is moderate. Since the speed is only 2X, and also in figure 3, the improvements over baselines are minor in 3 out of 4 environments.**\n\nFurthermore, our extra numerical results presented in the global response now clearly demonstrate a significant numerical performance improvement over the state-of-the-art on large-scale challenging problems (now an approximate 10x speedup, with further gains possible for larger problem sizes).\n\nSimilarly, our extra numerical results further emphasise the significant gains in gradient quality over state-of-the-art achieved by differentiating through problems with inequality constraints without resorting to approximate log-barrier problems. These stability improvements are on top of the observed speedups.  \n\n> **In Figure 4 as well, the final final imitation loss seems very close for different methods, though the gradient quality indeed seems better.**\n\nAdditional analysis exploring the benefits of avoiding the use of log-barrier approximations inherent to Safe-PDP is provided in the global response. Importantly, this analysis extends beyond imitation loss, highlighting that avoiding the log-barrier allows for the correct recovery of the true constraints and system parameters, which cannot be observed from plots of imitation loss alone.\n\n> **It seems the relevance for the paper is mostly in the setting where trajectory gradients wrt to model parameters or objective parameters is required such as LfD. For a holistic comparison, there should be benchmarking wrt to current important methods specialized for SysID (model learning in MBRL) or inverse RL.**\n\nWe believe the CIOC problem we evaluate on where the dynamics and the cost are jointly learned is more difficult than just the SysID and the IRL tasks proposed in the PDP paper. In fact, the gradient computations which arise from the SysID task amounts to just backpropagating through the dynamics equations. This renders differentiating through the KKT/PMP conditions completely unnecessary. This observation also holds for the control and planning tasks. However, we are happy to add additional results to the appendix around IRL (known dynamics) and SysID (known cost) for completeness.\n\n> **Can you place this work with respect to model-based/model-free/inverse RL methods? i.e if this method or any part of this method can be used for policy learning or planning or model learning?**\n\nThank you for suggesting this. LfD is the most obvious task to apply differentiable COC, however we have added some additional commentary in the global response. We believe our method can be used anywhere a trajectory derivative is required. For example, we provide a reference to an RL work which learns an implicit policy, defined as the solution to an optimal control problem. We hope our method will further encourage and enable the learning of these so-called implicit policies."}, {"review_id": "0ffzAsLeiz", "reviewer": "Reviewer_kJe4", "summary": "Differentiating through optimal control problems to learn various components, such as the dynamics model or cost function, is a promising method for inverse reinforcement learning (IRL) or incorporating more structure in learned control policies. The central component of these approaches is to differentiate through optimality conditions, such as the KKT conditions or Pontryagin's maximum principle (PMP). Methods which differentiate through the PMP scale linearly with the planning horizon by constructing an auxiliary control system solved in the backward pass. Prior work has argued that this is not the case for methods which differentiate through the KKT conditions due to a large matrix inversion. Instead, this paper shows that this is not true if one properly accounts for the block structure and sparsity patterns of the matrices. The authors show that their approach also scales linearly with horizon and has more opportunity for parallelism to enable faster gradient computation. They compare to Pontryagin Differentiable Programming (PDP), a method derived from PMP, on a number of standard benchmark problems in the case of inverse reinforcement learning. Both methods perform similarly most of the time, with PDP sometimes failing, potentially due to issues with numerical stability. They also show this gap widens when inequality constraints are introduced, with their method yielding significantly better gradients and imitation loss.", "questions": "- How does the gradient stability of the proposed method compare to PDP and its related methods as horizon increases? \n- How does the proposed method fair with much longer horizons than those considered in the experiments?\n- What is the solver used in the forward pass and how are the Lagrange multipliers found?", "limitations": "There is some discussion of how the bottleneck on speed is now in the forward pass, given that the backward pass has been significantly sped up. The speed of the forward pass will also depend on the choice of solver for the constrained optimization problem. And there is some discussion of how the opportunities for parallelism in the backward pass are limited due to the availability of solvers for general block tridiagonal linear systems.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- Improving the scalability of implicit differentiation for learning parameters of optimal control problems is an important problem with applications in IRL, system identification, and structured feedback policy classes for reinforcement and imitation learning.\n- The paper is well organized and clearly written. It does a good job explaining the novelty and results and provides enough information to support its claims.\n- This paper shows that computing gradients through the KKT conditions for general optimal control problems with inequality and equality constraints can also scale linearly with horizon when we properly account for the matrix structure. Additionally, they show that methods which differentiate through the PMP conditions are equivalent to their approach, only differing in the use of a recursive rule for gradient computation.\n- Parts of the gradient computation in their method is parallel across time steps, unlike prior work which is entirely sequential. This allows them to compute gradients and vector-Jacobian products much more quickly.\n- Gradients appear to be more stable, especially in the case of inequality constraints, compared to PDP and its extensions. The computation time for the backwards pass is also significantly faster for the proposed method over PDP. This may enable scalability to longer horizon problems.\n- Unlike prior methods which efficiently differentiate through the KKT conditions, such as [2], they do not rely on differentiating through an approximation of the non-convex problem. Instead, they compute gradients through the original problem, which may have benefits in terms of gradient quality.", "weaknesses": "- The central goal of the paper is to improve the scalability and numerical stability of gradient computation for long horizons. However, there is no evaluation of runtime or gradient stability across different horizon lengths. Instead, the benchmarks use a fixed horizon, which is of moderate length. It would strengthen the paper to see a breakdown of how the improvements over PDP scale with the horizon length of the problem, and if these trends carry over to even longer horizons than currently considered in the paper. If PDP truly scales worse due to numerical stability issues, then its performance should get worse with longer horizons while the proposed method does not.\n- There is no discussion of what solver is used in the forward pass for the experiments and how the Lagrange multipliers are found for gradient computation. Even if the paper reuses the methods from the PDP paper, the paper should still be stand-alone in that it contains these details in the appendix.\n- The paper says that differentiable MPC [2] is limited to affine-quadratic systems, which is not true. By using iLQR, it is able to handle nonlinear dynamics and non-convex cost functions. The paper also shows how to incorporate box constraints on controls. However, the proposed approach is more general in that it can handle arbitrary constraints. This should be fixed in the final paper.\n- This method is not the first to exploit structure in the KKT conditions to scale linearly with horizon. Despite the arguments in [18, 20], differentiable MPC [2] does not involve a large matrix inversion. Instead, it scales linearly with horizon by solving an auxiliary LQR in the backwards pass, similar in spirit to PDP in [18, 20]. While the proposed method is more general, there should be some discussion of this relationship. It would also strengthen the paper to include it as a benchmark given that it is also derived from the KKT conditions. It would especially be interesting to see how the quality of gradients and runtime of the backwards passes compare.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your appreciation of the contributions proposed in our work. We would like to address your concerns.\n\n> **It would strengthen the paper to see a breakdown of how the improvements over PDP scale with the horizon length of the problem, and if these trends carry over to even longer horizons than currently considered in the paper.**\n\nWe have provided additional insight and analysis around runtime scalability in the global response. In short, we believe we have firmly demonstrated this by evaluating both computation time and numerical stability on a synthetic example with varying $T$ up to $T=1000$.\n\n> **There is no discussion of what solver is used in the forward pass for the experiments and how the Lagrange multipliers are found for gradient computation.**\n\nWe are currently using the primal dual interior point solver IPOPT which solves for optimal trajectories and provides Lagrange multipliers. We will include these details in the paper and make our paper more stand-alone. We will also mention that we can compute the Lagrange multipliers given the solution as in PDP and Safe-PDP.\n\n> **Despite the arguments in [18, 20], differentiable MPC [2] does not involve a large matrix inversion.**\n\nThank you for mentioning this, we agree completely after revisiting the paper. We will reflect these changes in the manuscript, notably that the advantage of IDOC and Safe-PDP is around handling more general constraint functions and furthermore, that DiffMPC is very similar to PDP in that an auxiliary LQR problem is solved.\n\nInterestingly, the method proposed in DiffMPC solves the VJP directly, and as such as is a worthwhile additional comparison. However, we note that in Module 2 in the DiffMPC paper for handling non-convex problems (under Backward pass, steps 4-5), a **third-order derivative** must be computed. Specifically, this is the derivative of (w.r.t. $\\theta$) of the derivative of $\\ell$ w.r.t. $H^n_\\theta$. This will significantly affect computation time of DiffMPC for large-scale problems.\n\nWhile we have not had sufficient time to perform the comparison during the rebuttal period, we will attempt to perform a comparison during the discussion period."}, {"review_id": "BWBZQ9F5Hw", "reviewer": "Reviewer_CbeP", "summary": "Prior work shows that computing trajectory derivatives scales quadratically w.r.t to timesteps. This work proposes that trajectory derivatives scales linearly w.r.t. to timesteps, which can be parallelized, resulting in decreased computation time and increased numerical stability.", "questions": "-How does this method scale with longer and longer trajectories? Need to see X axis as trajectory length and Y axis as solving time.\n- Confused, why prior work fail to identify the block structure? It seems to directly follow from prior work (Gould et al.)", "limitations": "Yes", "rating": 4, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "- Good paper structure, reasonable flow.\n- Evaluation shows that the gradient computation is numerically stable, which leads to stable training for inverse RL and imitation learning tasks.", "weaknesses": "- Section 3.1 is not needed; readers are familiar with matrix and vector derivatives.\n- Good to move derivative w.r.t to trajectory in Section 4 as a motivation for this work. (i.e. move to Section 3)\n- Need to better structure Section 4; state in beginning of 4.2 the goal is prove that  computing the trajectory derivative relies on a block matrixes.\n- Assumes theta is a vector, does not apply for non-linear approximations such as DNNs\n- The main benefit of this work is being able to parallelize the computation of trajectory derivatives. There seems to be a lack of evaluation on this. (outside of 1(b)). Need to show total training time.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. We would like to emphasize not only the decreased computation time, but the significant performance gains on differentiating through inequality constrained problems without resorting to a log-barrier approximation. Additional compute and scalability experiments are presented in the global response. We will make your suggested changes to the document structure to enhance readability.\n\n> **Assumes theta is a vector, does not apply for non-linear approximations such as DNNs**\n\nWhile we assume $\\theta$ is a vector, this does not restrict its generality. For the DNN case you mention, $\\theta$ would represent the parameters of a neural network \"unrolled\" into a vector. Importantly, we can define expressions such as $df/d\\theta$ even for a DNN, which can then be used to compute the expressions in Eq. 4 in the original manuscript. Does this address your concern?\n\n> **Need to show total training time**\n\nThe only difference between IDOC and PDP/Safe-PDP is the computation of the backward pass. The total training time is heavily dependent on the solver used in the forward pass (e.g., IPOPT, DDP/iLQR solver, MPC PyTorch solver that runs on GPU). We believe our current comparison between methods is fair (with a very similar implementation in Python/Numpy across methods), and do not wish to muddy the comparison by specifying a full learning system. However, to make our comparisons even more comprehensive, we have provided additional compute and scalability experiments in the global response. \n\n> **Confused, why prior work fail to identify the block structure? It seems to directly follow from prior work (Gould et al.)**\n\nWe're not sure how to answer this question, but we believe we are the first to leverage the Gould et al. identities. Admittedly, there is a bit of confusion when reflecting on prior works. Differentiable MPC [2] did correctly identify the block sparse structure in the KKT system and formulated an auxiliary LQR problem to solve it, as reviewer kJe4 correctly pointed out. The Gould et al. identities are different from solving the KKT system as in [2] however, since they are derived by eliminating the Lagrange multipliers (see Section 4.3 in our original manuscript).\n\nThis auxiliary LQR approach from [2] is very similar to the approach proposed in PDP and Safe-PDP. However, Safe-PDP claimed that [2] (which uses the implicit function theorem) yields a trajectory derivative that scales quadratically with horizon $T$, which was an important motivation for our work. We will more clearly clarify the difference in approach of our method compared to existing work (e.g., PDP/Safe-PDP/DiffMPC) in the revised manuscript.\n\nRather than simply identifying the block structure from the identities (recently) derived in Gould et al., we believe exploiting the structure of the resulting equations to demonstrate faster computation and substantially more stable trajectory derivatives across a range of problems (especially through inequality constrained problems) is a significant contribution.\n\n> **-How does this method scale with longer and longer trajectories? Need to see X axis as trajectory length and Y axis as solving time.**\n\nSee the global response and attached pdf for additional results to more comprehensively evaluate the scalability w.r.t. trajectory length and problem size.\n\nWe hope you will consider raising your score given we will address your editorial comments and have provided additional scalability results."}, {"review_id": "LNQb0XNQwu", "reviewer": "Reviewer_tsv7", "summary": "There has been recent interest in differentiating through trajectories to obtain first-order derivatives for optimization problems including policy learning, inverse optimal control and model learning. Previous work usually uses a backward recursion which computation scales quadratically with the length of the horizon. In this work, the derivatives are computed exploiting the block-diagonal structure with computational time linearly varying with the length of the horizon. Additionally, it provides the flexibility of parallelizing the computation. This method can also directly compute vector-jacobian products easily. Most of these advantages are verified through training iterations with imitation learning and constrained inverse optimal control tasks on various robotics benchmarks.", "questions": "1) While citing results from Gould et. al such as equation (5) and Proposition 1, the authors are encouraged to provide more explanations about whether the results are a direct application and provide more context in the paper.\n\n2) It is not communicated clearly what is the final take away from sections 4.3 and 4.4. It looks like by assuming H is non-singular, it is possible to do the computation faster using this method. Will this method also result in numerical issues if H is improperly conditioned? If adding the proximal term is the solution, why does a similar type of solution not work for PDP? In other words, add a numerical conditioning to derivatives of c?\n\n3) This paper needs to be self-contained and there is too much dependence on referring to previous papers overall.\n", "limitations": "The authors have discussed the limitations in future work.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "A novel, interesting method to parallelize and accelerate the computation of derivatives is presented by exploiting the block diagonal structure of matrices. This method is numerically better-conditioned as compared to existing algorithms such as PDP and safe-PDP. Further, it is easily amenable for computing vector-jacobian products which provides it with an advantage.", "weaknesses": "The paper is making confusing claims or is making a misclaim.\n\nLine 38 from submission - “Naively applying these identities leads to a quadratic complexity with the length of the trajectory, which is described in prior works [18, 20]” - the citations are PDP and safe PDP.\n\nLine 46 from submission - “Furthermore, we show that the computation of these derivatives is linear with trajectory length, contradicting claims in prior works” - as far as I can tell, the previous backward and forward recursion also scales linearly with trajectory length. The PDP and safe PDP do not explicitly state quadratic complexity with trajectory length anywhere. \n\nExcerpt from PDP: “Third, in the backward pass, unlike differentiable MPC which costs at least a complexity of $O((m+2n)^2 T^2)$ to differentiate a LQR approximation, PDP explicitly solves for first derivative by an auxiliary control system, where thanks to the recursion structure, the memory and computation complexity of PDP is only $O((m+2n)T)$.”\n\nExcerpt from this submission Line 82 - “While the derivative computation in PDP is linear with the number of timesteps, it is inherently a serial calculation, requiring a recursion through time” \n\nIn conclusion, the paper is making contradicting claims about advantage over previous work. Otherwise, the paper is not communicating the claims correctly.\n\nI believe whatever computational benefit is reported is because of the ability to parallelize the computation and the final numerical results are above average but not excellent. The numerical stability problem with PDP is not consistently seen and it is likely that it can be quelled by better numerical conditioning. \n\nAdditive cost functions with quadratic structure is only one type of optimal control problem. The block diagonal structure is dependent on this assumption. There is a discussion in future work about extensions to non-additive cost functions.\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate you taking the time to review our paper and drawing our attention to specific claims within it that appeared unclear.\nIn particular, we apologise for the confusion regarding our comments around quadratic complexity -- we were not intending to make a misclaim here. \n\n> **The paper is making confusing claims or is making a misclaim.**\n\nTo clarify, our comments around quadratic complexity do not relate to the PDP and Safe-PDP methods themselves. They relate to comments in the PDP and Safe-PDP papers about the quadratic complexity of differentiable MPC and CasADi. Reading our submission again, we recognize that this was unclear and apologize for the confusion. The comments that we are referring to are:\n\n* Section 7 in the PDP paper (NeurIPS proceedings version):  *\"unlike differentiable MPC which costs at least a complexity of $O((m+2n)^2 T^2)$ to differentiate a LQR approximation\"*\n* Section 8 of the Safe-PDP paper (NeurIPS proceedings version): *\"Specifically, Safe PDP has a complexity of $O(T)$, while CasADi and Differentiable MPC have at least $O(T^2)$. This is because both CasADi and differentiable MPC are based on the implicit function theorem [75] and need to compute the inverse of a Hessian matrix of the size proportional to $T \\times T$.\".*\n\nThe second comment from the Safe-PDP paper refers to methods using the implicit function theorem (IFT) requiring a matrix inversion of the differential KKT system which is $O(T^3)$ in general. We are showing in our contribution that we can use IFT identities (derived from the same differential KKT system with variable elimination) presented in Gould et al., and recover $O(T)$ complexity. This time complexity does not sacrifice expressiveness; we differentiate through the exact same class of optimal control problems described in the PDP and Safe-PDP paper.\n\nWe will include the specific quotes from the PDP/Safe-PDP paper to clarify our actual claims when we update the manuscript. We are simply refuting that methods based on the IFT require an $O(T^3)$ matrix inversion. In fact, we show that using the IFT yields significant benefits to computation time and numerical stability, especially for problems with inequality constraints.\n\n> **This paper needs to be self-contained and there is too much dependence on referring to previous papers overall.**\n\nWe will improve the self-sufficiency of this paper as described in the global response. We will also better clarify that equation (5) and proposition 1 are taken directly from Gould et al. \n\n> **I believe whatever computational benefit is reported is because of the ability to parallelize the computation and the final numerical results are above average but not excellent.**\n\nPlease see our global response on the additional experiments. We hope this improves our results beyond above average!\n\n> **The numerical stability problem with PDP is not consistently seen and it is likely that it can be quelled by better numerical conditioning.**\n\nWe agree that techniques can be used to improve the stability of PDP, for instance using robust matrix factorizations (e.g., SVD instead of Cholesky/LU). However, we believe that alone cannot significantly improve the performance of Safe-PDP for hard constraints. In fact, in the Safe-PDP implementation, SVD is used for matrix inversion/factorization, whereas we simply use an LU-based solver. We believe our experiments show conclusively that the difference between Safe-PDP and IDOC is significant where hard inequality constraints are present, and doubly so given the simple nature of the optimal control problems we evaluate on.\n\nTo further reinforce our claims, in the global response and attached pdf, we provide a brief discussion on the importance of using hard constraints instead of log-barrier functions.\n\n> **Additive cost functions with quadratic structure is only one type of optimal control problem. The block diagonal structure is dependent on this assumption. There is a discussion in future work about extensions to non-additive cost functions.**\n\nAdditive cost is not the most general formulation, but encompasses many optimal control problems and enables the use of Bellman's principle of optimality. We also point out that PDP and Safe-PDP will not work without an additive cost structure. Recall that our formulation can also extend to include a final cost defined over the whole trajectory, which is not easily achieved under PDP/Safe-PDP. \n\nWe also note that a quadratic cost structure is not required to use our proposed method, only that the cost function is twice differentiable.\n\n> **It looks like by assuming H is non-singular, it is possible to do the computation faster using this method. Will this method also result in numerical issues if H is improperly conditioned?**\n\nYour reasoning is correct, non-singularity of $H$ will allow faster and more stable computation. We accept that poor conditioning in $H$ will cause issues for our method, but these issues will also be problematic for PDP and Safe-PDP since they use (and invert) the same blocks used to build $H$ in their calculations. \n\n> **If adding the proximal term is the solution, why does a similar type of solution not work for PDP? In other words, add a numerical conditioning to derivatives of c?**\n\nA small proximal term is important for the quadrotor experiment, where $H$ is inherently singular. The other experiments on the other hand, avoid using the proximal term altogether. Therefore, the numerical differences cannot be attributed to the addition of the proximal term alone.\n\nWe hope by addressing your major concerns around the claims, you will consider raising your score."}, {"review_id": "KG4SDrm2mE", "reviewer": "Reviewer_fBmg", "summary": "This paper analyzes and develops an efficient method to differentiate through a constrained discrete-time optimal control system, i.e., computing the derivative of an optimal trajectory of a constrained discrete-time optimal control system with respect to the parameters in the system’s cost function, dynamics, and/or constraints, based on the implicit function theorem (IFT). This is a key problem in many learning problems, such as inverse optimal control (inverse reinforcement learning, system identification, or learning from demonstration). The contributions of this paper include (1) the authors analyzed the sparse structures of IFT equations, and showed that the linear time complexity (in the trajectory horizon) can be obtained for solving IFT, instead of quadratic complexity. (2) Based on the analysis, the authors develop an efficient algorithm by parallelization and a vector jacobian auto-diff algorithm. The algorithm is evaluated in comparison the state-of-the-art (PDP [18] and Safe-PDP [20]) for various tasks,  the proposed method is shown at least 2x faster and more numerical stability.  \n", "questions": "Please find my question in the \"weaknesses\" section.", "limitations": "Please find in the \"weaknesses\" section.", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The study of differentiable optimal control is well-motivated, and the problem relates to many problems in learning and robotics. \nThe authors analyzed sparse structure of the IFT equation for the constrained optimal control problem, and showed the linear complexity (in trajectory horizon) of differentiating through optimal control. Besides, they established the connection between solving IFT equation and PDP. Based on the analysis, they develop parallelize and vector Jacobian Products algorithms to accelerate the backward pass, which have been shown effective and stable.\n\nThe algorithm is novel and has shown obvious improvement over the existing method (mainly PDP and Safe-PDP). I believe the algorithm will be useful and of interest of control and learning communities. The paper is well organized and well presented, and I expect the proposed method is important in practice, \n", "weaknesses": "In order to further improve the paper, some claims may need to be further clarified. \n\n- Line 159: since the algorithm requires the identification of a set of active constraints, \\tilde{g}_t, from all inequality constraints, will the use of a threshold cause the numerical issues, eventually leading to bad quality gradient and unstable gradient descent?\n- In Fig. 4 cartpole example, is the instability of Safe PDP without log-barrier functions because of the numerical issues of identifying active and inactive inequality constraints?\n- In my understanding, Safe PDP using log-barrier function has two benefits: first, it avoids the needs of active identification of inequality constraints, and 2) it creates some smoothing effects in the trajectory solution space because the non-differentiability caused by the switch between “active” and “inactive” constraints in the change of parameters. However, as the authors point out, Safe PDP with barrier function is only an approximation method to compute the gradient, thus it is not surprising to have a larger loss, as shown in Fig.4.\n- Currently, most of the examples in experiments consider the loss and stability performance.  In order to support the claim of \"2x\" more efficiency, I think more experiments about algorithm speed test (with respect to PDP) should be done. \n\n\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. We will address your questions as follows:\n\n> **Line 159: since the algorithm requires the identification of a set of active constraints, \\tilde{g}_t, from all inequality constraints, will the use of a threshold cause the numerical issues, eventually leading to bad quality gradient and unstable gradient descent?**\n\nThis is potentially an issue in practice, and this instability is well known when designing numerical solvers which use the so-called ``active-set\" approach. While we didn't experience issues around constraint switching in our experiments, we acknowledge that it could be an issue in practice and may be encountered when scaling these approaches to more difficult problems. \n\n> **In Fig. 4 cartpole example, is the instability of Safe PDP without log-barrier functions because of the numerical issues of identifying active and inactive inequality constraints?**\n\nWe're not sure if the instability of Safe-PDP without barrier functions is caused by identification of active and inactive constraints. While this claim is made in the Safe PDP paper (bottom on page 9 under Problem III), our approach IDOC also needs to identify the active constraints. Of course, we use the same threshold for identifying constraints for a fair comparison. We believe that constraint switching and identification in and of itself is not the root cause of instability for Safe-PDP. We will add this commentary into the manuscript since this is a very useful piece of analysis which arises from the experimental results.\n\n> **In my understanding, Safe PDP using log-barrier function has two benefits: first, it avoids the needs of active identification of inequality constraints**\n\nWe believe that the identification of active/inactive constraints in and of itself is a fast, simple operation (evaluate constraint functions and apply threshold). However, varying numbers of active constraints along the trajectory make it difficult to vectorize the computation of the trajectory derivatives, due to varying-sized blocks. \n\nThis property generally increases computation time compared to using the log-barrier method. Otherwise, the identification of constraints does not appear (at least, in our experiments) to impact gradient quality.\n\n>**it creates some smoothing effects in the trajectory solution space because the non-differentiability caused by the switch between “active” and “inactive” constraints in the change of parameters.**\n\nWe agree that the smooth log-barrier function may yield superior gradient information during learning, since constraints don't need to be satisfied to provide gradient information. This may explain why the log-barrier approach reduces imitation loss faster at the start compared to IDOC. \n\nHowever, as you correctly point out, IDOC provides lower final imitation loss for many of the experiments. This appears to be due to IDOC correctly identifying the hard constraint values. We provide further insight around this in the global response."}, {"review_id": "642HziWhzd", "reviewer": "Reviewer_dzG6", "summary": "This paper considers the optimization problem in the context of discrete-time control of dynamic systems. The paper proposes a method for efficiently and effectively evaluating the Jacobian of saddle points of constrained optimal control problems w.r.t. COC problem parameters. The main technical result is that the computation cost grows linearly with the number of time steps. The proposed approach is evaluated on standard environments like quadrotors and cart-poles.", "questions": "See my comments above.", "limitations": "See my comments above.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The whole paper is well-written and easy to follow. I appreciate the writing in the technical section (Secs. 3-5), which unfolds the core technical method in a nice and coherent way.\n\nThe main technical idea is solid and sound. Fig. 2 turned out to be quite helpful when I was checking the equations. The linear time complexity follows naturally after Proposition 1 is established.", "weaknesses": "I don’t see major fallacies in the main technical result itself, but I still have a few high-level concerns regarding the novelty and usefulness of the proposed technique:\n\n- I feel that Propositions 1, 2, and Fig 2 are direct results of applying existing numerical and mathematical techniques in derivative derivations. From what I understand, no new tools were developed in proving these results. As I am not well-calibrated with NeurIPS’s bar, I will let other reviewers decide how novel the technical method is.\n\n- The writing in the abstract and introduction gives me the impression that the gradient computation is fully parallelizable. I can see how multiplications involving H^(-1) can be parallelized. However, solving the tridiagonal block matrix seems like a sequential procedure to me. I am guessing from reading Lines 234-239 that the implementation used a sequential solver with O(T) time complexity and was the major time bottleneck (correct me if I am wrong). It would be useful to profile the time cost of each step that can and cannot be parallelized.\n\n- The experiments didn’t seem to evaluate the proposed method in large-scale problems. Dynamical systems (and their gradients) like quadrotor can be computed extremely fast (orders of magnitude than real time) because of their few degrees of freedom. The number of time steps (<=50 if I understand correctly from Table 1) does not seem to create a large-sized Hessian (a few hundred by a few hundred?) challenging enough to compute for modern CPUs. Having 2x speed up is still nice, but I had a higher expectation for a parallel (?) linear algorithm given that the baseline is sequential and quadratic.\n\n- On a related note, I didn’t find visualizations of the quadrator/cart-pole/etc environments in the main paper or the supplemental material. Having some images to visualize these tasks would be great.\n\nA very minor comment on writing: Some sentences are quite long and hard to follow, e.g., lines 21-23 and lines 31-34. Splitting them into several small sentences would be better.\n\nSome technical comments:\n\n- Lines 124 and 125: is R^p a typo? It makes more sense to me if we replace R^p with R^d as d is the dimension of \\theta. What is p?\n- Line 150: Why is a transpose operator needed?\n- Proposition 1: same question about p in the definitions of B and C; should it be replaced with d?\n- Proposition 1: Why is C defined as the Jacobian of h w.r.t. \\theta? It makes more sense to me if it is the Jacobian of r w.r.t. \\theta, which is also implied by the dimension n_r.\n- Proposition 2: I feel the linear time complexity O(T) assumes the dimension of \\theta does not grow with the number of time steps, or H^(-1)B won’t be computable in O(T) time. Please correct me if I am wrong.\n- Eqn. (5): it seems more common to flip the sign of lambda and replace -A^T with A^T so that the matrix is symmetric.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed review and kind words around the technical writing.\n\n> **From what I understand, no new tools were developed in proving these results.**\n\nWe agree that no new mathematical tools were developed. However, identifying and exploiting the link between the IFT identities of Gould et al. and optimal control is non-trivial, as evidenced by it being missed by both the learning and control communities. We specifically want to debunk prior statements we observed in the literature that approaches based on IFT identities are inherently quadratic. \n\nWe believe that such assertions need to be refuted formally by demonstrating that there are surprisingly meaningful numerical scalability and stability improvements offered by the novel application of IFT identities to control. We believe that our results (both theoretical and numerical) therefore represent an important novel contribution to both the learning and control communities.\n\n> **However, solving the tridiagonal block matrix seems like a sequential procedure to me. I am guessing from reading Lines 234-239 that the implementation used a sequential solver with O(T) time complexity and was the major time bottleneck (correct me if I am wrong).**\n\nThank you for drawing our attention to the apparent disconnect between our discussion of parallelization and the solution of the tridiagonal block matrix. We will clarify our discussion in the abstract and introduction by highlighting that it is currently easy to parallelize everything except the solution of the tridiagonal block matrix, but that there are a growing number of methods (some of which we cite) that propose parallelizable block tridiagonal solvers.\n\nUnfortunately, the code for these parallelizable block tridiagonal solvers appears unavailable. As a result, we implement a sequential procedure for the block tridiagonal system in Python. Despite this, we still observe significant improvements in computational performance. We envision that future availability and development of custom (parallelizable) block tridiagonal solvers will yield further reductions in computation time.\n\n\n> **It would be useful to profile the time cost of each step that can and cannot be parallelized.**\n\nFor $p=10$k and $T=1$k, our most challenging setting, we observe that the serial computation (block tridiagonal solve using the Thomas algorithm) occupies approx. 45% of total compute time for IDOC full, whereas for IDOC VJP it occupies 18%. For a fixed $p$, this ratio does not change with horizon length $T$. \n\nThe ratio for $p=1$k is 39% and 36% for IDOC full and VJP, respectively. This indicates that the importance of the serial computation diminishes as the problem size (w.r.t. parameters) increases. The serial component is still significant however, and warrants developing parallel block tridiagonal solvers to further improve scalability.\n\n> **The experiments didn’t seem to evaluate the proposed method in large-scale problems.**\n\nWe appreciate you pointing out that our first set of experiments, whilst showing a 2x speedup, could have been more compelling. We have therefore conducted additional large-scale experiments as detailed in the global response and attached pdf in which the dimensions of the parameters, states, and controls are large enough for parallelization to yield benefits over serial computation. \n\nThese additional results (cf. Fig. 1) offer compelling evidence that our IDOC VJP method is approximately 10x faster across all horizon lengths $T$ for a large-scale problem ($p=10000$), with further gains possible as the number of learnable parameters increases.\n\n> **On a related note, I didn’t find visualizations of the quadrator/cart-pole/etc environments in the main paper or the supplemental material. Having some images to visualize these tasks would be great.**\n\nWe will add visualisations (or suitable references) for the benchmark problems to the supplemental material. \n\n> **A very minor comment on writing: Some sentences are quite long and hard to follow, e.g., lines 21-23 and lines 31-34. Splitting them into several small sentences would be better.**\n\nWe agree that some sentences in the submission are too long and could be made clear, thank you for bringing this to our attention. Particular attention will be paid to those you have identified.\n\n### Technical Comments\n\nThank you for carefully reading our work and picking up these typos! We will address these in order.\n\n\n* Yes, this is a typo! $p$ should be $d$.\n* We assume $\\text{D}\\mathcal{L}_\\xi(\\xi)\\in \\mathbb{R}^{(n+m)T + n \\times 1}$ is a column vector. We will make this more explicit.\n* Correct! $p$ should be $d$ here as well.\n* Yes, this should indeed be $r$.\n* As is common in (optimal) control, we assume that $\\theta$ is static across time and is comprised of all relevant parameters across all (time-varying) cost and constraint functions.\n* Agreed, happy to make it symmetric and perform a sign change elsewhere. "}]

Author rebuttal:
# Global Response Overview

We thank the reviewers for providing detailed comments and feedback on our work. We are pleased the reviewers appreciated the computational and numerical performance improvements compared to existing state-of-the-art offered by our use of (structured) implicit differentiation and vector-Jacobian products for optimal control (IDOC and IDOC VJP, respectively). We address a number of common concerns raised by multiple reviewers in a series of posts. Specifically, we address

* Additional results around speed up and numerical stability
* Emphasis around handling problems with inequality constraints
* Making the paper more stand-alone

and leave responses to specific comments from individual reviewers as a reply to their respective posts.

## Additional Numerical Results

We constructed a large, synthetic example to more comprehensively test the computation time saving and numerical stability of IDOC. We sample the blocks, comprising terms in Eq. 4 in the original submission, randomly such that they are 1) large, and 2) poorly conditioned to test the effect of parallelisation and numerical stability, respectively. In summary, our results indicate:

* IDOC, IDOC VJP and PDP **scales linearly** w.r.t. horizon length $T$ and is approx. 10x faster for $p=2000$.
* IDOC VJP scales an **order of magnitude** better than both IDOC full and PDP w.r.t. the number of parameters $p$.
* IDOC is **more numerically stable** than PDP and degrades more gracefully as the numerical conditioning of the blocks deteriorates. 

See the 1 page response pdf for figures (Fig. 1) and more experimental details.

## Importance of Handling Inequality Constrained Problems

We would like to emphasize the significance of the improvement IDOC affords over Safe-PDP for differentiating through optimal control problems with inequality constraints. We acknowledge that the log-barrier approach proposed in Safe-PDP yields similar imitation loss IDOC. However, there are additional benefits to avoiding a log-barrier approach and differentiating through the constrained problem.

* Recovery of the correct constraint function parameters when learning from demonstration data (see Fig. 2 in the response pdf).
* Avoids necessitating initializing the forward pass with a feasible trajectory (otherwise the log-barrier function is undefined).

## Difficulty of Experiments

Some reviewers questioned the difficulty of the experiments. While we agree that scaling up these control problems to more difficult settings (e.g., the DM control suite) is of interest to the community generally, our particular contribution is largely analytical and relates only to the gradient computation. We believe the performance gain of IDOC over PDP/Safe-PDP are already significant, especially when considering the simple nature of the problems we evaluate on. We hope the additional challenging synthetic experiments sufficiently substantiate our claims around numerical stability and parallelization.

## Improving the Stand-Alone Quality of the Paper

We will make the following changes to our paper to make it more stand-alone.

* Describe for each experiment the dynamics equations and cost function parameterization in the appendix.
* Describe the forward solver used (IPOPT) and associated hyperparameters during learning.
* Add more visualizations to the appendix around the experiments.
* Better place the work in the broader RL literature, rather than simply comparing against PDP/Safe-PDP. We believe IDOC is well suited to model-based RL approaches such as the ones proposed in [1], which solve an OC problem (described in Equation 2 in our submission) to produce actions (implicit policy). We hope our work will encourage future work around this direction.

### References

[1] Romero, A., Song, Y. and Scaramuzza, D. (2023). *Actor-Critic Model Predictive Control.* arXiv preprint arXiv:2306.09852.


Meta-review: Thank you for your submission and active engagement throughout the review period. 
After the discussion period, the reviewers and I are in agreement that this is
a reasonable advancement to implicit differentiation for optimal control and
the methodology and experimental results to Pontryagin differentiable programming
clearly demonstrate the advantages.
Decision: Accept (poster)