Revisiting Implicit Differentiation for Learning Problems in Optimal Control

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.

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. The 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: - 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. - 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. - 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. - 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. 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. Some technical comments: - 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? - Line 150: Why is a transpose operator needed? - Proposition 1: same question about p in the definitions of B and C; should it be replaced with d? - 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. - 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. - 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.

Questions

See my comments above.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

See my comments above.

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.

Strengths

The study of differentiable optimal control is well-motivated, and the problem relates to many problems in learning and robotics. The 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. The 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,

Weaknesses

In order to further improve the paper, some claims may need to be further clarified. - 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? - 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? - 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. - 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.

Questions

Please find my question in the "weaknesses" section.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

Please find in the "weaknesses" section.

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.

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. Line 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. Line 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. Excerpt 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)$.” Excerpt 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” In conclusion, the paper is making contradicting claims about advantage over previous work. Otherwise, the paper is not communicating the claims correctly. 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. The numerical stability problem with PDP is not consistently seen and it is likely that it can be quelled by better numerical conditioning. 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.

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. 2) 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? 3) This paper needs to be self-contained and there is too much dependence on referring to previous papers overall.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

The authors have discussed the limitations in future work.

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

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

2 fair

Contribution

3 good

Limitations

Yes

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.

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

Questions

- How does the gradient stability of the proposed method compare to PDP and its related methods as horizon increases? - How does the proposed method fair with much longer horizons than those considered in the experiments? - What is the solver used in the forward pass and how are the Lagrange multipliers found?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

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.

Summary

The paper introduces a method for calculating analytic trajectory gradients in constrained optimal control problems using implicit differentiation, with following contributions: * Shows that computation of these derivatives can be linear in trajectory time-steps, utilizing the structure in the matrices in the gradient computation. * Shows how to parallelize the gradient computation for reduced overall computation time and better numerical stability. * Shows direct computation of vector-jacobian products for find optimal trajectories with some outer loss. The 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.

Strengths

Strenghts: The contribution of the method can be practically useful for trajectory gradient methods allowing for faster computation. Originality - 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. 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. Clarity - 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. 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. 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: 1. The paper lacks comparison on other experimental setting such as sysID and control from the baseline paper [18]. 2. 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. Minor Weaknesses: 1. Improve the legends in figure 3, and enlarge figure 1(a). In general, all the figures needs improvement.

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. 1. 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? 2. 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?

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

2 fair

Limitations

Limitations: * No clear experiments about how the method scales with more challenging and complex environments such as DM control suit, etc.

Acknowledgement

I have read the author response and other reviews. The authors have given a fair response to my questions. The contribution of the paper is better appreciated after this. Many of the differentiable trajectory optimization methods do run into numerical issues. However, I believe the paper can still be improved upon and presented better. Further, it helps to include a further discussion on ILQR/DDP literature from robotics and control. e.g. [1]. At this time, I increase my score to 5. [1] Roulet, Vincent, et al. "Iterative linear quadratic optimization for nonlinear control: Differentiable programming algorithmic templates." arXiv preprint arXiv:2207.06362 (2022).

Thank you for the response

After reading the rebuttal, I remain supportive of the paper and will raise my score to weak accept. I think there is still room for improving the experiments, but I am OK with the current paper with all the promised changes incorporated.