Differentiable Trajectory Optimization as a Policy Class for Reinforcement and Imitation Learning

We thank the reviewer for their time and the valuable comments and suggestions. We address the reviewer’s questions as below.

Is the inability of Theseus to solve a constraint optimization version of the problem a fundamental limitation of how it works, or simply something not implemented yet? Maybe it solves non-linear least-squares problems analytically, and adding constraints would make this impossible?

We thank the reviewer for this question. As we are not the active developers of Theseus, we would not be able to answer this question with perfect accuracy. According to our understanding, Theseus solves the nonlinear least-squares problem numerically; specifically, we used the Levenberg–Marquardt algorithm in our implementation.

The original Theseus paper has the following discussion about adding constraints in their limitation section: “The nonlinear solvers we currently support apply constraints in a soft manner (i.e., using weighted costs). Hard constraints can be handled with methods like augmented Lagrangian or sequential quadratic programs [99, 100], and differentiating through them are active research topics”. Based on this, it seems that adding constraints is not a fundamental limitation of how it works, but rather something that has not been implemented yet.

We thank the reviewer for their time and the valuable comments and suggestions. We address the reviewer’s questions and concerns as below.

How do the authors ensure that a reasonable answer is found in the process of nonlinear trajectory optimization with gradients? If the cost and dynamics was learned to formulate a convex trajectory optimization problem, this would not be a problem as the encoder is forced to learn a representation of that form. However the authors propose to use a fully nonlinear formulation and it's not clear how the training pipeline would be robust against finding bad local minima.

We thank the reviewer for this insightful question. Indeed, there is not much theoretical guarantee on the quality of the converged local minima when solving the nonlinear trajectory optimization problem. One thing we find particularly useful is to provide the nonlinear trajectory optimization problem with a good initialization. This is shown in our ablation study of model-based RL in Figure 4, where we compare the performance of DiffTOP, and DiffTOP without using the action initialization from the learned latent policy (in which case the action is initialized to be 0 in trajectory optimization) on 4 model-based RL tasks . The result shows that without the action initialization from the learned latent policy, the performance drops.

The same results hold in the imitation learning setting, as shown in Table 3 & 4 in Appendix A.2 of the updated paper. Following the reviewer’s questions, we have added experiments to test DiffTOP without informed initial guess in the RoboMimic and ManiSkill experiments in Table 3 & 4. As the results suggest, if the action is initialized to be zero or sampled from a Gaussian of N(0, 1), the performance of DiffTOP is similar or slightly worse to vanilla behavior cloning, possibly due to the convergence to bad local minimal during trajectory optimization. With better action initializations, i.e., using actions from BC-RNN or Diffusion Policy, the performance of DiffTOP improves.

Therefore, a good practice of using DiffTOP is to provide it with a good action initialization. This is often feasible in practice: in model-based RL, as our experiments show, the learned latent policy can be used to provide this action initialization; for imitation learning, one could always first learn a base policy using any behavior cloning algorithm, and then use DiffTOP to further refine the actions from the base policy for better performances.

We have updated Appendix A.2 of the paper to include the above discussion.

The contribution in the IL seems a bit weaker compared to MBRL. In the setting of imitation learning, it seems that the authors are performing a version of inverse RL (which similarly does not assume access to rewards, but learns rewards and dynamics from demonstrations). Making connections to previous methods within inverse RL would be good.

We thank the reviewer for this valuable comment. We agree that when applied to the IL setting, DiffTOP can be viewed as doing inverse RL as well. The differences between DiffTOP and previous IRL methods are as follows.

In DiffTOP, the learned function is not necessarily a “reward” function as those used in a typical RL/MDP setting. It is just a learned “objective function”, such that optimizing it with trajectory optimization would yield actions that minimize the imitation learning loss with respect to the expert actions in the demonstration. We leave exploring the connections with inverse RL for future work.

In contrast to our approach, inverse RL algorithms try to learn a reward in the MDP/RL setting, such that when optimizing the learned reward with RL, the policy’s state-action/state distribution matches the state-action/state distribution of the expert demonstrations. To perform such a distribution match, the loss for learning the reward is usually not the imitation learning loss as used in DiffTOP, but a loss to encourage matching the state-action/state distribution of the expert (e.g., the adversarial loss in GAIL [2]).

We have updated the paper to include the connection to inverse RL at the end of section 5.2.1.

[1] Florence et al, Implicit Behavioral Cloning, CoRL 2021
[2] Ho et al, Generative Adversarial Imitation Learning, NeurIPS 2016

I would like to thank the authors for the detailed response, I agree with the points in the rebuttal and will keep the score as is.

I have questions regarding the additional PG loss for the RL case. Does it make sense to optimize the model parameters w.r.t. the Q-function? Wouldn't this loss let the dynamics model prefer to hallucinate to obtain a high q-function? Learning the dynamics and reward model is a supervised learning problem while learning the q-function is not. Therefore, I would expect that the supervised learning problem is much easier and imagine that the additional information flowing through this additional PG loss is not that important. What is the intuition of the authors for the additional PG loss? For imitation learning, the motivation for the additional bc loss is much clearer.

We thank the reviewer for this insightful question. Optimizing the model parameters with the additional PG loss through differentiable trajectory optimization is one of the core contributions of our paper (besides the imitation learning modification as the reviewer noted). The intuition behind this is as follows:

As the reviewer mentioned, since learning the dynamics and reward model is a supervised learning problem that have been well-studied, existing model-based RL algorithms adopt this paradigm: they first learn a latent dynamics and reward model via supervised learning, and during inference time, they perform planning (such as MPPI) using the learned dynamics and reward model to generate the actions to take. However, recent studies [1, 2, 3] have shown that there is a fundamental “objective mismatch issue” with such a paradigm, i.e., models that achieve better training performance (e.g., lower MSE) in learning a dynamics/reward model are not necessarily better for control. These papers explain this phenomenon in more detail, but the intuition here is as follows: a learned dynamics model will typically not perfectly drive the dynamics loss to 0. Two different learned dynamics models will have different types of errors, and the model with the lower MSE loss is not necessarily the model that will be most useful for optimizing actions to maximize reward; it depends on what states and actions are causing those errors and what types of errors they are. For example, errors in predicting the dynamics for actions that lead to low rewards are not very important.

DiffTOP addresses this issue: by computing the PG loss on the optimized actions from trajectory optimization and differentiating through the trajectory optimization process, the dynamics and reward functions are both optimized directly to maximize the task performance. As there will inevitably be errors in the learned dynamics model, the additional PG loss regularizes the dynamics model to have low error in states and actions that lead to high rewards, i.e., those that are important for maximizing the task performance, instead of driving the dynamics model to hallucinate. Besides, we keep the dynamics prediction loss term when training the dynamics model, which should also prevent hallucination. In our experiments, DiffTOP with the added PG loss greatly outperformed the TD-MPC baseline that does not have this loss term, demonstrating its effectiveness.

[1] Lamber et al, Objective mismatch in model-based reinforcement learning, Learning for Dynamics and Control (L4DC), 2020
[2] Eysenbach et al, Mismatched no more: Joint model-policy optimization for model-based rl, NeurIPS 2022
[3] Ghugare et al, Simplifying model-based rl: learning representations, latent-space models, and policies with one objective, ICLR 2022

We thank the reviewer for their time and the valuable comments and suggestions. We address the reviewer’s questions and concerns as below.

For example, the algorithm is still very unclear to me, which model parameters exist and are optimized together.

We have updated section 4.2 and 4.3 of the paper to make the model parameters more clear, and hope the following clarifications help with addressing this concern.

At the beginning of Section 4.2, we have added the following sentence to more explicitly state the model parameters that are optimized: “We use $\theta$ to denote all learnable model parameters to be optimized in DiffTOP, including the parameters of the encoder $h_\theta$, the latent dynamics model $d_\theta$, the reward predictor $R_\theta$, and the Q value predictor $Q_\theta$” . Equation 6 in the paper illustrates the loss that is used to optimize these parameters.

For imitation learning, we have added the following sentence at the end of section 4.3 to more explicitly show the model parameters that are optimized in this setting: “Similarly, We use $\theta$ to denote all learnable model parameters to be optimized in DiffTOP, which includes the parameters of the encoder $h_\theta$, the latent dynamics model $d_\theta$, and the cost function $f_\theta$ in the imitation learning setting”. Equation 8 in the paper shows the loss that is used to optimize these parameters.

It would be great if the authors could introduce their algorithm more slowly and rely on less knowledge from the TD-MPC paper. While mentioning the relation to TD-MPC is important, requiring excellent knowledge of TD-MPC to read the paper is very limiting.

We thank the reviewer for bringing this to our attention. Indeed, the model-based RL part of DiffTOP builds upon TD-MPC, thus requiring some knowledge of TD-MPC for understanding the paper. We do acknowledge that with the page limit, we are not able to cover every detail of the TD-MPC paper. However, we would like to clarify that we do have a section in our paper, i.e., section 3.2, model-based RL preliminaries, to explain the notations and external knowledge from TD-MPC that is needed for understanding our method, in order to make the paper self-contained. Following the reviewer’s suggestion, we have updated section 3.2 with more explanations of TD-MPC to add more buildup for introducing our algorithm, and for a better understanding of the paper.

Please let us know if you have any more suggestions on how to improve the readability of the paper – we would be more than happy to further incorporate them.

We thank the reviewer for their time and the valuable comments and suggestions. We address the reviewer’s questions as below.

It would be good to see the comparison of the computational efficiency of given algorithms and plots for return vs wall-clock time in addition to the number of samples / Could you share the training plots of returns vs wall-clock time at least for some of the environments? It would be useful to compare not only the sample but also the computational efficiency of the different algorithms.

We thank the reviewer for this suggestion. We have added plots for return vs wall-clock time of TD-MPC and DiffTOP on some of the model-based RL tasks, shown in Figure 7 in the updated paper, along with a discussion of the results in Appendix A.1.

To provide a short summary here for the reviewer’s convenience:
In terms of computational efficiency (e.g. wall-clock time needed to achieve a given level of performance), the results are environment-dependent. Compared to TD-MPC, DiffTOP is more computationally efficient on 3 environments, has similar computational efficiency on 3 environments, and worse computational efficiency on 3 environments (please see Figure 7 of the updated paper). The main advantage of our method is not computational efficiency but sample efficiency, as shown by our superior performance in Figure 3.

Compared to TD-MPC, it takes more wall-clock time for DiffTOP to finish the training of 1M environment steps. The major reason is that solving and back-propagating through the trajectory optimization problem is slower than using MPPI as in TD-MPC. As a reference, it takes 0.052 seconds to use Theseus to solve and differentiate through the trajectory optimization problem, and 0.0092 seconds to use MPPI as in TD-MPC.

The research community is actively developing better and faster algorithms/software libraries for differentiable trajectory optimization, which could improve the computation efficiency of DiffTOP. For example, in all our experiments, we used the default CPU-based solver in Theseus. Theseus also provides a more advanced solver named BaSpaCho, which is a batched sparse Cholesky solver with GPU support. When we switch from the default CPU-based solver to BaSpaCho, the time cost of solving the trajectory optimization problem is reduced by 22% from 0.052 seconds to 0.041 seconds.

Are there any cases when TD-MPC performs or could theoretically perform better than DiffTOP?

We thank the reviewer for this interesting question. Empirically, as shown in Figure 3 in the paper, DiffTOP always performs better than or on par with TD-MPC. Theoretically, if the latent dynamics model and the latent reward model is perfectly learned without any error, i.e., if they exactly match the ground-truth dynamics and reward model, and if enough samples are used in MPPI, TD-MPC should be able to achieve the maximal possible performance and DiffTOP would not be able to further outperform it. However, in practice, since there will always be errors in the learned latent dynamics and reward model, DiffTOP outperforms TD-MPC by directly learning a latent space that is useful for control, instead of just minimizing the dynamics or reward prediction error.

What’s the computational cost of using Theseus for differentiable trajectory optimization vs MPPI?

We thank the reviewer for this question. We have performed an analysis on the computation cost, and updated the paper to have a discussion about this in Appendix A.1. The results are summarized as follows:

For inferring the action at one time step, the computation cost of using the default CPU-based solver in Theseus for differentiable trajectory optimization is 0.052 seconds. For MPPI, the time cost is 0.0092 seconds. However, as mentioned before, the computation cost of DiffTOP can be further reduced with better algorithms/solvers. For example, if we switch from the default CPU-based solver in Theseus to a more advanced batched sparse Cholesky solver with GPU support, the computation cost of solving the trajectory optimization problem is reduced by 22% from 0.052 second to 0.041 second. As the community is actively developing better libraries/algorithms for differentiable trajectory optimization, we believe the computational efficiency of DiffTOP would further improve in the future.

The paper was accepted in July. We were slightly late in pushing to Arxiv. But we would really appreciate it if you can include it because it is indeed the closest related work to yours

Thank you for your interest in our paper, and thanks for drawing our attention to your paper.
As your paper is released on arxiv on Oct 23, 2023, which is past the ICLR submission date, we are not aware of it at the time of submission. It looks very relevant indeed, we will read it in more detail, and include it in the related work in the updated version.