Title: Revisiting Implicit Differentiation for Learning Problems in Optimal Control
Venue: NeurIPS 2023
Accepted: yes
Keywords: implicit differentiation

bi-level optimization; constrained learning and control; safe learning for control
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

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 vectorJacobian 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.

1 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 to underlying parameters in cost functions, system dynamics, and constraints (e.g., state and control limits).

The efficient computation of these trajectory derivatives is important in the (open-loop) solution of optimal control problems for both trajectory optimization and model predictive control (MPC). Such derivatives are also crucial in inverse optimal control (also called inverse reinforcement learning or learning from demonstration), where given expert demonstration trajectories, the objective is to compute parameters of the cost function that best explain these trajectories. Extensions of inverse optimal control that involve inferring parameters of system dynamics in addition to cost functions also subsume system identification, and provide further motivation for the efficient computation of trajectory derivatives. Our proposed method of computing trajectory derivatives enables direct

(e.g., imitation loss)

Control

(e.g., IPOPT, DDP)

Trajectory VJPs

(a) System Figure

(b) Imitation Learning: Rocket Landing

IDOC

104

S-RDP

S-PDP (b)

103

Imitation

102

101

100

0

50

100

15C

Iteration

Limit

三

DeC(≤(0))

E(0)

A

minimization of a loss function defined over optimal trajectories (e.g., imitation loss) using first-order descent techniques and is a direct alternative to existing works [2, 20, 22], including those based on bespoke derivations of Pontryagin's maximum principle (PMP) in discrete time [20, 22].

Analytical trajectory derivatives for COC problems are derived by differentiating through the optimality conditions of the underlying optimization problem and applying Dini's implicit function theorem. Derivatives are recovered as the solution to a system of linear equations commonly referred to as the (differential) KKT system. This framework underlies all existing works [2, 20, 22] as well as our proposed approach. Similarly to Amos et al. [2], we construct the differential KKT system, however we use the identities from Gould et al. [14] that apply variable elimination on the Lagrange multipliers relating to the dynamics and constraints. We show how to exploit the block-sparse structure of the resultant matrix equations to achieve linear complexity in computing the trajectory derivative with trajectory length. Furthermore, we can parallelize the computation, yielding superior scalability and numerical stability on multithreaded systems compared to methods derived from the PMP [20, 22].

Our specific contributions are as follows. First, we derive the analytical trajectory derivatives for a broad class of constrained (discrete-time) optimal control problems with additive cost functions and dynamics described by first-order difference equations. Furthermore, we show that the computation of these derivatives is linear in trajectory length by exploiting sparsity in the resulting matrix expressions. Second, we describe how to parallelize the computation of the trajectory derivatives, yielding lower computation time and superior numerical stability for long trajectories. Third, we show how to directly compute vector-Jacobian products (VJPs) with respect to some outer-level loss over optimal trajectories, yielding further improvements to computation time in the context of bi-level optimization. This setting commonly arises in the direct solution of inverse optimal control problems [17, 33, 34]. Finally, we provide discussion unifying existing methods for computing trajectory derivatives [2, 20, 22]. We dub our method IDOC (Implicit Differentiation for Optimal Control) and validate it across numerous experiments, showing a consistent speedup over existing approaches derived from the PMP. Furthermore, for constrained problems, IDOC provides significantly improved trajectory derivatives, resulting in superior performance for a general learning from demonstration (LfD) task. 1 .

2 Related Work

Differentiating Through Optimal Control Problems. An optimal control (OC) problem consists of (system) dynamics, a cost function, an optimal control policy, and a set of state and/or control constraints. Learning problems in optimal control involve learning some or all of these aspects, with the problem of learning dynamics referred to as system identification [29], the problem of learning the cost function referred to as inverse optimal control (or inverse reinforcement learning or learning from demonstration) [34, 17, 21, 33, 37, 51], and the problem of learning the control policy referred to as reinforcement learning [4]. Recent work [20] has shown that solving these

learning problems can be approached in a unified end-to-end fashion by minimizing task-specific loss functions with respect to (unknown) parameters of the associated OC problem. This unified formulation places great importance on the efficient computation of derivatives of optimal trajectories with respect to said parameters. Methods reliant on computing trajectory derivatives have, however, been mostly avoided (and argued against) in the robotics and control literature due to concerns about computational tractability. For example, bi-level methods of inverse optimal control [34, 33] have mostly been argued against in favor of methods that avoid derivative calculations by instead seeking to satisfy optimality conditions derived from KKT [24, 12] or PMP conditions [17, 32, 23, 21].

Analytical Trajectory Derivatives. Our method is a direct alternative to methods such as DiffMPC [2], PDP [20] and its extension Safe-PDP [22], which differentiate through optimality conditions to derive trajectory derivatives. Common to these methods is identifying that trajectory derivatives can be computed by solving an auxiliary affine-quadratic OC problem, and furthermore, that this can be done efficiently using a matrix Riccati equation. While our approach still differentiates the optimality conditions, we avoid solving matrix Riccati equations and show that this enables easy computation of VJPs, parallelization across timesteps and improved numerical stability. Like Safe-PDP, IDOC computes derivatives through COC problems with smooth inequality constraints, however we show in our experiments that our derivatives are significantly more stable during training. Section 4.3 provides a detailed description of the differences between IDOC and existing methods.

Differentiable Optimization. Our work falls under the area of differentiable optimization, which aims to embed optimization problems within end-to-end learning frameworks. Gould et al. [14] proposed the deep declarative networks framework, and provide identities for differentiating through continuous, inequality-constrained optimization problems. Combinatorial problems [31, 44], as well as specialized algorithms for convex problems [1] have also been addressed. Differentiable optimization has been applied to end-to-end learning of state estimation [49, 42] and motion planning problems [6, 2, 26] in robotics, see Pineda et al. [35] for a recent survey. In addition, numerous machine learning and computer vision tasks such as pose estimation [9, 40], meta-learning [28], sorting [11, 7] and aligning time series [47] have been investigated under this setting.

3 Constrained Optimal Control Formulation

In this section, we provide an overview of the COC problems we consider in this paper. These problems involve minimizing a cost function with an additive structure, subject to constraints imposed by system dynamics and additional arbitrary constraints such as state and control limits. As a result, they can be formulated as a non-linear program (NLP). We will discuss how to differentiate through these COC problems in Section 4.

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 n → R m be a vector-valued function with vector arguments and let D f ∈ R n × m be the (matrix-valued) derivative where elements are given by

For a scalar-valued function with (multiple) vector-valued inputs f : R n × R m → R evaluated as f ( x, y ) , let the second order derivatives D 2 XY = D X ( D Y f ) ⊤ . See Gould et al. [14] for more details.

3.2 Optimization Formulation for Constrained Optimal Control

We can formulate the (discrete-time) COC problem as finding a cost-minimizing trajectory subject to constraints imposed by (possibly non-linear) system dynamics. In addition, trajectories may be subject to further constraints such as state and control limits.

To begin, let the dynamics governing the COC system at time t be given by x t +1 = f t ( x t , u t ; θ ) , where x t ∈ R n , u t ∈ R m denotes state and control variables, respectively 2 . For time horizon T > 0 ,

let x ≜ ( x 0 , x 1 , . . . , x T ) and u ≜ ( u 0 , u 1 , . . . , u T -1 ) denote the state and control trajectories, respectively. Furthermore, let θ ∈ R d be the parameter vector that parameterizes our COC problem. Our COC problem is then equivalent to the following constrained optimization problem

where f t : R n × R m × R d → R n describe the dynamics, and c t : R n × R m × R d → R , c T : R n × R d → R are the instantaneous and terminal costs, respectively. Furthermore, the COC system may be subject to (vector-valued) inequality constraints g t : R n × R m × R d → R q t and g T : R n × R d → R q T such as control limits and state constraints, for example. Additional equality constraints h t : R n × R m × R d → R s t and h T : R n × R d → R s T can also be included.

The decision variables of Equation 2 are the state and control trajectory ξ ≜ ( x, u ) , whereas parameters θ are assumed to be fixed. Equation 2 can be interpreted as an inequality-constrained optimization problem with an additive cost structure and vector-valued constraints for the initial state x 0 , dynamics, etc. over all timesteps. Let ξ ( θ ) ≜ ( x ⋆ , u ⋆ ) be an optimal solution to Equation 2. We can treat the optimal trajectory ξ ( θ ) as an implicit function of parameters θ , since Equation 2 can be solved to yield an optimal ξ ( θ ) for any (valid) θ .

We can solve Equation 2 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 [30, 19]. Regardless, for the purposes of computing analytical trajectory derivatives using our method (as well as methods that differentiate through optimality conditions [20, 22, 2]), we only need to ensure that our solver returns a vector ξ ( θ ) which is a (local) minimizer to Equation 2. We now describe how to differentiate through these optimal control problems, important in the end-to-end learning context.

4 Trajectory Derivatives using Implicit Differentiation

In this section, we present our identities for computing trajectory derivatives D ξ ( θ ) based on those derived in Gould et al. [14] for general optimization problems by leveraging first-order optimality conditions and the implicit function theorem. We exploit the block structure of the matrices in these identities that arises in optimal control problems to enable efficient computation and furthermore, show that computing trajectory derivatives is linear in the trajectory length T .

Before we present the main analytical result, recall the motivation for computing D ξ ( θ ) . Suppose in the LfD context we have a demonstration trajectory ξ demo (we can extend this to multiple trajectories, but choose not to for notational simplicity). Furthermore, define a loss L ( ξ ( θ ) , ξ demo ) which measures the deviation from predicted trajectory ξ ( θ ) to demonstration trajectory ξ demo . Ultimately, if we wish to minimize the loss L with respect to parameters θ using a first-order decent method, we need to compute D θ L ( ξ ( θ ) , ξ demo ) by applying the chain rule. Specifically, we compute D θ L ( ξ ) = D ξ L ( ξ ) D ξ ( θ ) , which requires trajectory derivative D ξ ( θ ) .

4.1 Preliminaries

First, we first reorder ξ such that ξ = ( x 0 , u 0 , x 1 , u 1 , . . . , x T ) ∈ R n ξ × 1 , 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 ξ t ∈ R n + m represent the subset of variables in ξ associated to time t , with final state ξ T = x T ∈ R n . Next, we stack all constraints defined in Equation 2 into a single vector-valued constraint r comprised of T +2 blocks. Specifically, let r ( ξ ; θ ) ≜ ( r -1 , r 0 , r 1 , . . . , r T ) ∈ R n r , where n r = ( T +1) n + s + q . The first block is given by

r -1 = x 0 , and subsequent blocks are given by

Here, ˜ g t are the subset of active inequality constraints for ξ (detected numerically using a threshold ϵ ) at timestep t . Note, s = ∑ T t =0 s t , q = ∑ T t =0 | ˜ g t | are the total number of additional equality and active inequality constraints (on top of the Tn dynamics and n initial state constraints). Each block represents a group of constraints associated with a particular timestep (except the first block r -1 ). As we will see shortly, grouping decision variables and constraints in this way will admit a favorable block-sparse matrix structure for cost/constraint Jacobians and Hessians required to compute D ξ ( θ ) .

4.2 Analytical Results for Trajectory Derivatives

Proposition 1 (IDOC) . Consider the optimization problem defined in Equation 2. Suppose ξ ( θ ) exists which minimizes Equation 2. Furthermore, assume f t , c t , g t and h t are twice differentiable for all t in a neighborhood of ( θ, ξ ) . If rank ( A ) = n r and furthermore, H is non-singular, then where

and Lagrange multipliers λ ∈ R n r satisfies λ ⊤ A = D ξ J ( ξ ; θ ) .

Proof. This is a direct application of Proposition 4.5 in Gould et al. [14].

Remark 1. Matrix H is block diagonal with T +1 blocks.

To see this, ˆ H = D 2 ξξ J ( ξ ; θ ) has a block diagonal structure with blocks given by D 2 ξ t ξ t c t ( ξ t ; θ ) , which follows from the assumption of additive costs in the COC problem described in Equation 2. In addition, constraints g t , h t depend only on ξ t and furthermore, the dynamics constraint is first order in ξ t +1 . Therefore, D 2 ξξ r t ( ξ ; θ ) i is only non-zero in the block relating to ξ t for all t and i . We plot the block-sparsity structure of H in Figure 2a.

Remark 2. Matrix A is a block-banded matrix that is two blocks wide.

This is shown by noting that the only non-zero blocks in A correspond to D ξ t r t ( ξ ; θ ) and D ξ t +1 r t ( ξ ; θ ) . The former relates to g t , h t and the -f t ( ξ t ; θ ) component of the dynamics constraint. The latter only relates to the x t +1 component of the dynamics constraint. Finally, the initial condition block r -1 only depends on ξ 0 , hence the only non-zero block is given by D ξ 0 r -1 ( ξ 0 ; θ ) . We plot the block-sparsity structure of A in Figure 2b.

Proposition 2. Evaluating Equation 4 has O ( T ) time complexity.

Proof. From Remarks 1 and 2, H and A have a block diagonal and block-banded structure with T +1 and 2 T +2 blocks, respectively. Therefore, we can evaluate H -1 A ⊤ in O ( T ) time and furthermore, H -1 A ⊤ has the same structure as A . It follows that we can also evaluate AH -1 B -C in O ( T ) time after partitioning B and C into blocks based on groupings of ξ and r .

Next, observe that AH -1 A ⊤ yields a block tridiagonal matrix with T + 2 blocks on the main diagonal (see Figure 2c for a visualization of the block-sparse structure). As a result, we can evaluate ( AH -1 A ⊤ ) -1 ( AH -1 B ) by solving the linear system ( AH -1 A ⊤ ) Y = AH -1 B for Y in O ( T )

1

2

0

0

1

0

H

0

52

3

0

0

0

H2

0

0

3

0

0

0

H3

(b) Matrix

A

1

2

3

r1

I

0

0

0

0

0

0

ro

-Deoro

0

0

I

0

0

0

r1

0

-Dr1

0

I

0

0

r2

0

0

-De2r2

I

r3

0

0

0

Dears

r-1

ro r_1

ro r1

r2

0

0

0

0

r1

0

0

0

r3

0

0

0

(c) Matrix

AH

r2

0

0

0

-

1

A

r3

0

0

0

⊤

time using any linear time (w.r.t. number of diagonal blocks) block tridiagonal linear solver. A simple example is the block tridiagonal matrix algorithm, easily derived from block Gaussian elimination.

Finally, H -1 B can easily be evaluated in O ( T ) time by solving T +1 linear systems comprised of the blocks of H and B . We conclude that evaluating Equation 4 requires O ( T ) operations.

4.3 Comparison to DiffMPC and PDP

All existing methods for computing analytical trajectory derivatives involve solving the linear system

where blocks H,A,B,C are defined in Section 4.2 and D ξ is the desired trajectory derivative. We call Equation 5 the (differential) KKT system and the matrix K the (differential) KKT matrix . Note, D ξ is unique if K is non-singular 3 . To better contextualize IDOC, we now briefly compare and contrast how previous methods such as DiffMPC [2] and PDP [20, 22] solve this KKT system.

DiffMPC. DiffMPC [2] proposes a method for differentiating through Linear Quadratic Regulator (LQR) problems, which are OC problems where the cost function is (convex) quadratic and the dynamics are affine. The authors show that solving the LQR problem using matrix Riccati equations can be viewed as an efficient method for solving a KKT system which encodes the optimality conditions. In addition, they show that computing trajectory derivatives involves solving a similar KKT system, motivating efficient computation by solving an auxiliary LQR problem in the backward pass. The matrix equation interpretation of the backward pass allows the efficient computation of VJPs for some downstream loss L ( ξ ) in a bi-level optimization context.

DiffMPC extends to handling non-convex OC problems with box constraints on control inputs via an iterative LQR-based approach. Specificially, such problems are handled by first computing VJPs w.r.t. the parameters of an LQR approximation to the non-convex problem around the optimal trajectory ξ, λ . Next, the parameters of the approximation are differentiated w.r.t. the underlying parameters θ and combined using the chain rule. However, differentiating the quadratic cost approximation and dynamics requires evaluating costly higher-order derivatives, e.g., d dθ d 2 c ( ξ ) dξ 2 , which are 3D tensors. These tensors are dense in general, although problem specific sparsity structures may exist.

PDP/Safe-PDP. PDP [20] and its extension, Safe-PDP [22], take a similar approach to DiffMPC in deriving trajectory derivatives. PDP derives trajectory derivatives by starting with PMP, which is well-known in the control community and applies to non-convex OC problems. Furthermore, Jin et al. [20] shows that the PMP and KKT conditions are equivalent in the discrete-time COC setting. Trajectory derivatives are obtained by differentiating the PMP conditions, yielding a new set of PMP conditions for an auxiliary LQR problem, which can be solved using matrix Riccati equations. This

approach is fundamentally identical to DiffMPC, with only superficial differences for LQR problems (solving a differential KKT system versus differentiating PMP conditions) 4 .

For non-convex problems, PDP does not evaluate higher-order derivatives, unlike DiffMPC. However, it is not clear how to compute VJPs under approaches derived using the PDP framework. SafePDP extended PDP to handle arbitrary, smooth inequality constraints using the constrained PMP conditions, which extends the capability of DiffMPC to handle box constraints. Trajectory derivatives are computed by solving an auxiliary equality constrained LQR problem. Experiments in both Jin et al. [22] and Section 6 show that Safe-PDP yields unreliable derivatives for COC problems. While Jin et al. [22] conjected that instability was due to the set of active constraints changing between iterations, we find that IDOC still learns reliably in the presence of constraint switching. We instead observed that the poor gradient quality arises naturally from the equality constrained LQR solver used for the backward pass [25] not matching the solution obtained by directly solving the KKT system.

IDOC. In contrast, we solve Equation 5 for D ξ by first applying variable elimination on D λ , yielding Proposition 1. We will discuss in Section 5 how the equations resulting from this approach admit parallelization and favorable numerical stability compared to DiffMPC and PDP's auxiliary LQR approach. Furthermore, VJPs can be easily computed (unlike PDP) without evaluating higherorder derivatives (unlike DiffMPC). IDOC handles arbitrary smooth constraints like Safe-PDP, and we show in Section 6 that our derivatives for COC problems are reliable during training. However, we require the additional assumption that H is non-singular (which holds if and only if all blocks in H are non-singular), which is not required for K to be non-singular. An oftentimes effective solution when H is singular, initially proposed by Russell et al. [39], is to set H = H + δ 2 I for small δ , which is analogous to adding a proximal term to the cost function in Equation 2. We will now describe how to evaluate Equation 4 in linear time by exploiting the block-sparse structure of the matrix equation.

5 Algorithmic Implications of IDOC

5.1 Parallelization and Numerical Stability

Parallelization. To evaluate Equation 4, we can leverage the block diagonal structure of H and block-banded structure of A to compute H -1 A ⊤ and H -1 B in parallel across all timesteps. Specifically, we evaluate the matrix product block-wise, e.g., H -1 t B t for all t in parallel (similarly with A ). Following this, we can then compute A ( H -1 A ⊤ ) in parallel also using the same argument.

In addition, we can solve the block tridiagonal linear systems involving AH -1 A ⊤ by using a specialized parallel solver [5, 36, 27, 41, 15, 18]. Unfortunately, none of these methods have open source code available, and so in our experiments, we implement the simple block tridiagonal matrix algorithm. Implementing a robust parallel solver is an important direction for future work, and will further improve the scalability and numerical stability of IDOC.

Numerical Stability. In addition to parallelization, another benefit of avoiding Riccati-style recursions for computing derivatives is improved numerical stability. For long trajectories and/or poorly conditioned COC problems (e.g., stiff dynamics), IDOC reduces the rounding errors that accumulate in recursive approaches. We show superior numerical stability in our experiments compared to PDP in Section 6, despite using the simple recursive block tridiagonal matrix algorithm for solving AH -1 A ⊤ . 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 D ξ ( θ ) given in Equation 4 is that we can now directly compute VJPs given some outer loss L ( ξ ) ∈ R over the optimal trajectory. Let v ≜ D ξ L ( ξ ) ⊤ ∈ R n ξ × 1 be the gradient of the loss w.r.t. trajectory ξ ( θ ) . The desired gradient D θ L ( ξ ( θ )) is then given by D θ L ( ξ ) = v ⊤ D ξ ( θ ) using the chain rule. The resultant expression is

(a) Cartpole

(b) 6-DoF Quadrotor

(c) Robotarm

(d) 6-DoF Rocket

103

IDOC

S-PDP

loss

102

S-PDP(O)

Imitation

101

10°

10-1

10-21

0

50

100150200250300

Iteration

一

IDOC

S.PDP

loss

101

S-PDP(b)

hitation

100

E

10-1

10-2

0

250

5007501000 1250150C

Iteration

103

IDOC

S-PDP

loss

102

S-PDP (b)

Imitation

101

10°

10-1

10-2

0

250500 7501000 12501500

Iteration

IDOC

103

S-PDP

loss

S-PDP (b)

Imitation

102

101

10°

0

60

120

180

240

30C

Iteration

The simple observation here is that we do not need to construct D θ ξ ⋆ ( θ ) 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 D ξ ( θ ) and then multiplying with v .

To see this, we follow the example in Gould et al. [14] and assume the blocks in H have been factored. Then for a single block (ignoring constraints for simplicity), evaluating v ⊤ ( H -1 t B t ) is O (( n + m ) 2 p ) while evaluating ( v ⊤ H -1 t ) B t is O (( n + m ) 2 +( n + m ) p ) . Evaluating the VJP directly significantly reduces computation time compared to constructing the full trajectory derivative for problems with higher numbers of state/control variables and tunable parameters.

6 Experiments

While our contributions are largely analytical, we have implemented the identities in Equation 4 to verify our claims around numerical stability and computational efficiency on a number of simulated COC environments. We will show that IDOC is able to compute trajectory derivatives significantly faster compared to its direct alternative PDP [20] and Safe-PDP [22] with superior numerical stability.

6.1 Experimental Setup

We evaluate IDOC against PDP [20] and Safe-PDP [22] in an LfD setting across four simulation environments proposed in Jin et al. [22], as well as a synthetic experiment. The simulation environments showcase the gradient quality for a realistic learning task, whereas the synthetic experiment is designed to measure numerical stability and computation times. For Safe-PDP, we evaluate against the method proposed for inequality constrained problems (Safe-PDP), as well as using an approximate log-barrier problem in place of the full problem (Safe-PDP (b)). The IPOPT solver [45] is used in the forward pass to solve the COC problem, and Lagrange multipliers λ are extracted from the solver output. More generally however, note that the method proposed in Gould et al. [14] can be used to recover λ if another solver is used where λ is not provided.

Imitation Learning/LfD. The LfD setting involves recovering the model parameters θ that minimizes the mean-squared imitation error to a set of N demonstration trajectories, given by

200

400

600

800

1000

Horizon (

T

)

10

-

2

10

-

1

10

0

10

1

Computation Time (s)

PDP

IDOC Full

IDOC VJP

(a) Horizon length scalability.

0

2000

4000

6000

8000

10000

Num. Parameters (

d

)

10

-

2

10

-

1

10

0

10

1

PDP

IDOC Full

IDOC VJP

(b) Model size scalability.

1

2

3

4

5

6

7

8

log(

(

H

))

10 2

10 3

10 4

10 5

10 6

10 7

10 8

D

(

)

64

D

(

)

32

1

PDP

IDOC Full

IDOC VJP

(c) Numerical stability.

L ( ξ ( θ ) , ξ demo ) ≜ 1 N ∑ i ∥ ξ ( θ ) i -ξ demo i ∥ 2 . We include all parameters in the cost, dynamics and constraint functions in θ and furthermore, evaluate performance both with (S-PDP) and without (PDP) inequality constraints. We perform experiments in four standard simulated environments: cartpole, 6-DoF quadrotor 5 , 2-link robot arm and 6-DoF rocket landing. For a detailed description of the imitation learning problem as well as each COC task, see the appendix. In addition, we provide timings for all methods within the simulation environments in the appendix.

Synthetic Benchmark. The simulation experiments described above are not sufficiently large-scale to adequately measure the benefits of parallelization and numerical stability afforded by IDOC over PDP. To demonstrate these benefits, we constructed a large-scale, synthetic experiment where the blocks required to construct H,A,B and C (as discussed in Section 4) are generated randomly. Computation time against the horizon length T , as well as the number of parameters d is reported. Numerical stability is measured by comparing mean-absolute error between trajectory derivatives evaluated under 32-bit and 64-bit precision for varying condition numbers over the blocks of H . State and control dimensions are fixed at n = 50 and m = 10 . Experiments are run on an AMD Ryzen 9 7900X 12-core processor, 128Gb RAM and Ubuntu 20.04. See the appendix for more details.

6.2 Imitation Learning/LfD Results

In this section, we evaluate the effectiveness of gradients produced from IDOC against PDP and Safe-PDP [20, 22] for the imitation learning/LfD task. In the equality constrained setting, IDOC and PDP yield identical (up to machine precision) trajectory derivatives and imitation loss throughout learning; see the appendix for more detailed results and analysis. However, as shown in Figure 3, when inequality constraints are introduced, the derivatives produced by Safe-PDP and IDOC differ significantly. Safe-PDP fails to reduce the imitation loss due to unreliable gradients, whereas IDOC successively decreases the imitation loss. While Safe-PDP (b), which differentiates through a log-barrier problem, also provides stable learning, it ultimately yields higher imitation loss compared to IDOC. This is because

the barrier problem is an approximation to the true COC problem. Further discussion around using log-barrier methods and COC problems is provided in Section 7.1.

6.3 Synthetic Benchmark Results

Figure 4 illustrates the results of the synthetic benchmark. We observed that IDOC and PDP have similar computation time and scalability with problem size when computing full trajectory derivatives. However, computing VJPs with IDOC significantly reduces computation time and improves scalability

by an order of magnitude with model size (i.e., d = | θ | ). Numerical stability experiments show that IDOC (full and VJP) yields lower round-off errors and improved stability compared to PDP. Using more sophisticated block tridiagonal solvers will further improve stability.

7 Discussion and Future Work

7.1 Differentiating through Log-Barrier Methods

In the imitation learning setting, we assume that demonstration data ξ demo is an optimal solution to a COC problem with (unknown) parameters θ ⋆ . Furthermore, we assume that a subset of timesteps A ⊆ { 1 , . . . , 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 θ ⋆ , log g t ( x demo t , u demo t ; θ ⋆ ) are undefined for t ∈ A . Therefore to recover ξ ( θ ) = ξ demo , we must have g t ( x t ( θ ) , u t ( θ )) < 0 for t ∈ A , i.e., we cannot recover the true constraint function through minimizing the imitation loss. We verify this using the robot arm environment, and present the results in Figure 5, plotting the error between the estimated and true constraint value. We observe that IDOC recovers the true constraint value more closely compared to Safe-PDP (b).

Generality of IFT. The generality of the differentiable optimization framework, and the matrix equation formulation for trajectory derivatives given in Equation 4 yield additional conceptual benefits which may help us tackle even broader classes of COC problems. For example, we can relax the additive cost assumption and add a final cost defined over the full trajectory such as

where Q = CC ⊤ and C ∈ R n ξ × k for k ≪ n ξ is full rank. For IDOC, we can simply use the matrix inversion lemma to invert H in O ( T ) . However, this is more complicated for methods derived from the PMP which rely on the specific additive cost structure of the underlying COC problem.

Dynamic Games. A generalization of our work is to apply the differentiable optimization framework to handle learning problems that arise in dynamic games [3] and have recently begun to employ PDP-based approaches [10]. Being able to differentiate through solution and equilibrium concepts that arise in dynamic games enables the solution of problems ranging from minimax robust control [3] to inverse dynamic games [10, 33], and is a promising direction for future work.

Fast forward passes. While we have proposed a more efficient way of computing trajectory derivatives (i.e., the backward pass), we observe that the forward pass to solve the COC problem bottlenecks the learning process. Significant effort must be placed on developing fast solvers for COC problems for hardware accelerators to allow learning to scale to larger problem sizes.

Combining MPC and Deep Learning. A recent application for computing trajectory derivatives is combining ideas from deep learning (including reinforcement learning) with MPC [38, 50, 46, 48]. Common to these approaches is using a learned model to select the parameters for the MPC controller. By allowing for robust differentiation through a broader class COC problems compared to previous approaches, we hope that IDOC will allow for future development in this avenue of research.

8 Conclusion

In this paper, we present IDOC, a novel approach to differentiating through COC problems. Trajectory derivatives are evaluated by differentiating KKT conditions and using the implicit function theorem. Contrary to prior works, we do not solve an auxiliary LQR problem to efficiently solve the (differential) KKT system for the trajectory derivative. Instead, we apply variable elimination on the KKT system and solve the resultant matrix equations directly. We show that linear time evaluation is possible by appropriately considering equation structure. In fact, we show that IDOC is faster and more numerically stable in practice compared to methods derived from the PMP. We hope that our discussion connecting the fields of inverse optimal control and differentiable optimization will lead to future work which enables differentiability of a broader class of COC problems.

References

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.

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.

T. Ba¸ sar and G. J. Olsder. Dynamic noncooperative game theory . Academic Press, London/New York, 1982. ISBN 978-0-08-095666-4.

D. Bertsekas. Reinforcement learning and optimal control . Athena Scientific, 2019.

R. Bevilacqua, B. Codenotti, and F. Romani. Parallel solution of block tridiagonal linear systems. Linear Algebra and its Applications , 104:39-57, 1988.

M. Bhardwaj, B. Boots, and M. Mukadam. Differentiable Gaussian process motion planning. Proc. of the IEEE International Conference on Robotics and Automation (ICRA) , 2020.

M. Blondel, O. Teboul, Q. Berthet, and J. Djolonga. Fast differentiable sorting and ranking. In Proc. of the International Conference on Machine Learning (ICML) , 2020.

S. Boyd and L. Vandenberghe. Convex Optimization . Cambridge University Press, 2004. doi: 10.1017/CBO9780511804441.

D. Campbell, L. Liu, and S. Gould. Solving the Blind Perspective-n-Point Problem End-To-End With Robust Differentiable Geometric Optimization. In Proc. of the European Conference on Computer Vision (ECCV) , 2020.

K. Cao and L. Xie. Game-Theoretic Inverse Reinforcement Learning: A Differential Pontryagin's Maximum Principle Approach. IEEE Transactions on Neural Networks and Learning Systems , pages 1-8, 2022. doi: 10.1109/TNNLS.2022.3148376.

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.

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):14741488, 2017. doi: 10.1177/0278364917745980.

P. E. Gill, W. Murray, and M. A. Saunders. SNOPT: An SQP Algorithm for Large-Scale Constrained Optimization. SIAM Review , page 99-131, 2005.

S. Gould, R. Hartley, and D. Campbell. Deep Declarative Networks. IEEE Transactions on Pattern Analysis and Machine Intelligence , 44:3988-4004, Aug 2022. doi: 10.1109/TPAMI. 2021.3059462.

I. N. Hajj and S. Skelboe. A multilevel parallel solver for block tridiagonal and banded linear systems. Parallel Computing , 15:21-45, 1990.

C. R. Harris, K. J. Millman, S. J. van der Walt, R. Gommers, P. Virtanen, D. Cournapeau, E. Wieser, J. Taylor, S. Berg, N. J. Smith, R. Kern, M. Picus, S. Hoyer, M. H. van Kerkwijk, M. Brett, A. Haldane, J. F. del Río, M. Wiebe, P. Peterson, P. Gérard-Marchant, K. Sheppard, T. Reddy, W. Weckesser, H. Abbasi, C. Gohlke, and T. E. Oliphant. Array programming with NumPy. Nature , 585:357-362, 2020.

K. Hatz, J. P. Schlöder, and H. G. Bock. Estimating parameters in optimal control problems. SIAM Journal on Scientific Computing , 34(3):A1707-A1728, 2012. doi: 10.1137/110823390.

S. Hirshman, K. Perumalla, V. Lynch, and R. Sanchez. Bcyclic: A parallel block tridiagonal matrix cyclic solver. Journal of Computational Physics , 229:6392-6404, 2010.

T. A. Howell, B. E. Jackson, and Z. Manchester. ALTRO: A Fast Solver for Constrained Trajectory Optimization. In Proc. of the IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS) , pages 7674-7679, 2019.

W. Jin, Z. Wang, Z. Yang, and S. Mou. Pontryagin Differentiable Programming: An End-to-End Learning and Control Framework. Advances in Neural Information Processing Systems , 33: 7979-7992, 2020.

W. Jin, D. Kuli´ c, 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.1177/0278364921996384.

W. Jin, S. Mou, and G. J. Pappas. Safe Pontryagin Differentiable Programming. Advances in Neural Information Processing Systems , 34:16034-16050, 2021.

M. Johnson, N. Aghasadeghi, and T. Bretl. Inverse optimal control for deterministic continuoustime nonlinear systems. In Decision and Control (CDC), 2013 IEEE 52nd Annual Conference on , pages 2906-2913, Dec 2013.

A. Keshavarz, Y. Wang, and S. Boyd. Imputing a convex objective function. In Intelligent Control (ISIC), 2011 IEEE International Symposium on , pages 613-619. IEEE, 2011.

F. Laine and C. Tomlin. Efficient computation of feedback control for equality-constrained lqr. In Proc. of the IEEE International Conference on Robotics and Automation (ICRA) , 2019. doi: 10.1109/ICRA.2019.8793566.

B. Landry, Z. Manchester, and M. Pavone. A differentiable augmented lagrangian method for bilevel nonlinear optimization. In Robotics: Science and Systems , 2019.

E. László, M. Giles, and J. Appleyard. Manycore algorithms for batch scalar and block tridiagonal solvers. ACM Transactions on Mathematical Software , 42(4), 2016.

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 Pattern Recognition (CVPR) , 2019.

L. Ljung. System identification: theory for the user . Prentice Hall PTR, Upper Saddle River, NJ, 1999.

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.

A. Mensch and M. Blondel. Differentiable dynamic programming for structured prediction and attention. In Proc. of the International Conference on Machine Learning (ICML) , 2018.

T. L. Molloy, J. J. Ford, and T. Perez. Finite-horizon inverse optimal control for discrete-time nonlinear systems. Automatica , 87:442 - 446, 2018. doi: 10.1016/j.automatica.2017.09.023.

T. L. Molloy, J. Inga Charaja, S. Hohmann, and T. Perez. Inverse Optimal Control and Inverse Noncooperative Dynamic Game Theory: A Minimum-Principle Approach . Springer International Publishing, 2022.

K. Mombaur, A. Truong, and J.-P. Laumond. From human to humanoid locomotion-an inverse optimal control approach. Autonomous robots , 28(3):369-383, 2010.

L. Pineda, T. Fan, M. Monge, S. Venkataraman, P. Sodhi, R. T. Chen, J. Ortiz, D. DeTone, A. Wang, S. Anderson, J. Dong, B. Amos, and M. Mukadam. Theseus: A Library for Differentiable Nonlinear Optimization. Advances in Neural Information Processing Systems , 2022.

E. Polizzi and A. H. Sameh. A parallel hybrid banded system solver: the spike algorithm. Parallel Computing , 32(2):177-194, 2006.

N. D. Ratliff, J. A. Bagnell, and M. A. Zinkevich. Maximum margin planning. In Proceedings of the 23rd international conference on Machine learning , pages 729-736. ACM, 2006.

A. Romero, Y. Song, and D. Scaramuzza. Actor-critic model predictive control, 2023.

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.

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.

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.

P. Sodhi, E. Dexheimer, M. Mukadam, S. Anderson, and M. Kaess. LEO: Learning energy-based models in factor graph optimization. In Proc. of the Conference on Robot Learning (CoRL) , 2021.

R. Tedrake. Underactuated Robotics . 2023. URL https://underactuated.csail.mit. edu .

M. Vlastelica, A. Paulus, V. Musil, G. Martius, and M. Rolínek. Differentiation of blackbox combinatorial solvers. In Proc. of the International Conference on Learning Representations (ICLR) , 2020. *Equal Contribution.

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.

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.

M. Xu, S. Garg, M. Milford, and S. Gould. Deep declarative dynamic time warping for endto-end learning of alignment paths. In Proc. of the International Conference on Learning Representations (ICLR) , 2023.

F. Yang, C. Wang, C. Cadena, and M. Hutter. iPlanner: Imperative Path Planning. In Proceedings of Robotics: Science and Systems , Daegu, Republic of Korea, July 2023. doi: 10.15607/RSS. 2023.XIX.064.

B. Yi, M. Lee, A. Kloss, R. Martín-Martín, and J. Bohg. Differentiable factor graph optimization for learning smoothers. In Proc. of the IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS) , 2021.

M. Zanon and S. Gros. Safe Reinforcement Learning Using Robust MPC. IEEE Transactions on Automatic Control , 66:3638-3652, 2021.

B. D. Ziebart, A. L. Maas, J. A. Bagnell, and A. K. Dey. Maximum entropy inverse reinforcement learning. In AAAI Conference on Artificial Intelligence , pages 1433-1438, 2008.

A Experimental Setup for Imitation Learning/LfD Task

We use the code and simulation environments provided by [22] and follow a similar experimental setup, described in more detail in this section. The LfD experimental setting without inequality constraints is referred to in [20] as the inverse reinforcement learning (IRL) task, whereas including inequality constraints is referred to in [22] as the constrained inverse optimal control (CIOC) task.

Demonstration Trajectories. Up to five demonstration trajectories are used in the LfD setting with no inequality constraints. Each demonstration trajectory is generated by solving a COC problem using the same underlying parameters θ , 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, θ is initialized by adding uniform noise to the true value. For all methods, we use the gradient descent with the same learning rate for a given environment.

Environment Specifications. Table 1 provides summary specifications of the simulation environments used in our experiments, which are consistent with prior works [20, 22]. Recall that n , m denote the number of states and controls, respectively. The number of parameters is denoted d = | θ | and T is the horizon length.

Additional Environment Parameters. Additional specifications are presented in Table 2. These include log-barrier parameter γ used for Safe-PDP (b), time discretization for Forward Euler to discretize continuous time dynamics, the learning rate for minimizing the imitation loss and finally the number of demonstration trajectories per environment. (E) refers to experiments without constraints, whereas (I) means inequality constraints are present.

A.1 Cartpole

The cartpole environment relates to the swing-up and balance task, where a simple pendulum (i.e., massless pole and point mass on the end of the pole) is swinging on a cart and the objective is to balance the pendulum above the cart. Only a horizontal force can be applied to drive the cart.

Dynamics. The state variable x = [ y, q, ˙ y, ˙ q ] ⊤ ∈ R 4 is comprised of the horizontal position of the cart y , counter-clockwise angle of the pendulum θ (from the hanging position) and their respective velocities. Control u ∈ R relates to the horizontal force applied to the cart. Dynamics parameters θ dyn = [ m c , m l p, ℓ ] ⊤ ∈ R 3 relate to the the mass of the cart m c , mass of the point mass at the end of the pole m p and length of the pole ℓ . Gravity is set at g = 10 . The continuous time dynamics are given by

with a detailed derivation provided in [43, Ch. 3].

Cost Functions. The cost function is a quadratic cost over states and controls. Specifically,

with x goal = [0 , 0 , -π, 0] ⊤ . We specify that W ≜ diag ( w ) , where w = [ w y , w q , w ˙ y , w ˙ q ] ≥ 0 and furthermore, w u ≥ 0 . We have that θ cost = [ w ⊤ , w ⊤ u ] ⊤ ∈ R 5 .

Constraints. For the setting with constraints, we apply box constraints for the cart position and controls. Specifically, we enforce | y | ≤ y max and | u | ≤ u max. We let constraint parameters θ constr = [ y max , u max ] ⊤ . The final parameter vector is given by θ = [ θ ⊤ dyn , θ ⊤ cost , θ ⊤ constr ] ⊤ .

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 ⊤ , ˙ p ⊤ , q ⊤ , ω ⊤ ] ⊤ ∈ R 13 , where p ∈ R 3 represents the position of the center-of-mass of the quadrotor, ˙ p ∈ R 3 represents linear velocity, q ∈ R 4 is a unit quaternion representation of the quadrotor attitude and ω ∈ R 3 represents angular velocity. For the inequality constrained setting, non-linear (squared-norm) constraints on quadrotor position, given by ∥ p ∥ 2 2 ≤ r max and thrust limits are applied.

A.3 2-Link Robot Arm

See Section E in the appendix in [20] for a reference to the full description of the robot arm problem. The objective of this OC problem is to move the robot arm to a goal configuration. The state is given by the orientation of the base and elbow joint, as well as respective velocities [ q 1 , q 2 , ˙ q 1 , ˙ q 2 ] ⊤ ∈ R 4 . The control inputs u = [ u q , u 2 ] ⊤ ∈ R 2 relate to applying torques at each joint. Constraints for the inequality constrained setting relate to box constraints on the joint positions, as well as torque limits. The cost function is similar to cartpole, i.e., a quadratic penalty to a goal state and over the controls.

A.4 6-DoF Rocket Landing

See Section I in the appendix in [20] for the full definition of the 6-DoF rocket landing problem. The objective of this task is to land a rocket modeled as a rigid body (softly) onto a 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.

and

(a) Cartpole

(b) 6-DoF Quadrotor

(c) Robotarm

(d) 6-DoF Rocket

3

1e-10

1e-10

3

1e-10

Rel. diff. imitation loss (%)

Rel. diff. imitation loss (%)

2

Rel. diff. imitation loss (%)

Rel. diff. imitation loss (%)

N

1

1

D

-1

-8

2

2

10

2000

4000

0009

8000

10000

2000

4000

6000

8000

10000

2000

4000

6000

8000

10000

200

400

600

800

1000

Iteration

Iteration

Iteration

Iteration

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 (i.i.d.) from a standard Gaussian. Hessian blocks are made symmetric using H sym = ( H + H ⊤ ) / 2 , where H 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 v = D ξ L ( ξ ) ⊤ are also sampled i.i.d. from a standard Gaussian.

C Additional Results for the Imitation Learning/LfD Tasks

In this section, we provide additional detailed analysis around the imitation learning/LfD 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 are 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/LfD tasks evaluated in the simulation environments. Given the relatively small size for the COC problems, there is very limited advantage with respect to compute time for parallel evaluation of trajectory derivatives. However, as a proxy to parallelization, IDOC can be vectorized across timesteps using the numerical linear algebra library Numpy [16]. All experiments in this section are run on a single thread of an AMDRyzen 9 7900X 4.7Ghz desktop CPU.

The experimental setting where inequality constraints are present is 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 C presented in Equation 4 are still block structured, the blocks may not be of uniform size in this setting. We present in Figure 7a and 7b the mean computation time (along with the upper/lower bound) per iteration across five trials with 10k iterations for each trial for all environments except for rocket (1k). We observe around a 2 × speedup over PDP by computing our IFT/DDN gradients, even before direct computation of vector-Jacobian products. Computing VJPs directly yields a 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 of 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 to the necessary computational overheads required for identifying and batching blocks.

Note that Safe-PDP [22] (for hard inequality constraints) and PDP [20] (used for differentiating through log-barrier problems) derivatives are implemented using a custom LQR and equality constrained LQR solver, respectively. These solvers have markedly different implementations, which explains differences in computation time between Safe-PDP and PDP.

(a) Hard Constraints

sw)

2.5

IDOC VJP

IDOC

w)

2.5

IDOC VJP

IDOC

time

2.0

Safe-PDP

time

2.0

PDP

Computationt

1.5

Computationt

1.5

1.0

1.0

0.5

0.5

0.0

0.0

CartpoleQuadrotorRobotarm

Rocket

CartpoleQuadrotorRobotarm

Rocket

Environment

Environment

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nI 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.\n\nQUESTIONS:\nSee my comments above.\n\nLIMITATIONS:\nSee my comments above.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\nThe 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\nWEAKNESSES:\nIn 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\nQUESTIONS:\nPlease find my question in the \"weaknesses\" section.\n\nLIMITATIONS:\nPlease find in the \"weaknesses\" section.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThere 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.\n\nSTRENGTHS:\nA 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.\n\nWEAKNESSES:\nThe 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\nQUESTIONS:\n1) 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\nLIMITATIONS:\nThe authors have discussed the limitations in future work.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

comments: 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.

STRENGTHS:
- Good paper structure, reasonable flow.
- 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.
- Good to move derivative w.r.t to trajectory in Section 4 as a motivation for this work. (i.e. move to Section 3)
- 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.
- Assumes theta is a vector, does not apply for non-linear approximations such as DNNs
- 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.

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.
- Confused, why prior work fail to identify the block structure? It seems to directly follow from prior work (Gould et al.)

LIMITATIONS:
Yes
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-07-10

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nDifferentiating 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nThere 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-18"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nStrenghts:\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.\n\nWEAKNESSES:\nMajor 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.\n\nQUESTIONS:\nI 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?\n\nLIMITATIONS:\nLimitations:\n* No clear experiments about how the method scales with more challenging and complex environments such as DM control suit, etc.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-08-01"}