Paper: Differentiable Trajectory Optimization as a Policy Class for Reinforcement and Imitation Learning (ICLR-2024)

Reviewer comment (Theory):
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.

Author response:
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.

Reviewer comment (Novelty):
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.

Author response:
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

Reviewer comment (Experiments):
"For DiffTOP, we always use a base policy to generate the action initialization for the trajectory optimization problem" - this makes the comparison between other methods a bit unclear as other methods did not start with an informed initial guess. The experiment results make it seem like DiffTOP cannot be used as a standalone tool for IL, but rather a tool for the purposes of refinement; yet the method section is proposing a standalone IL algorithm.

Author response:
We thank the reviewer for bringing this up. In our experiments, in addition to just comparing to other methods that do not start with an informed initial guess, we also compare to methods that use these informed initial guesses as well. They are “Residual+BC-RNN”, “IBC+Diffusion”, and “Residual+Diffusion'' in Table 1 (for RoboMimic), and “BC + Residual” in Table 2 (for ManiSkill). These methods use informed initial guesses as follow:

For RoboMimic experiments:  
Residual + BC-RNN: We use a pre-trained BC-RNN as the base policy, and learn a residual policy on top of it. The residual policy takes as input the action from the base policy, and outputs a delta action which is added to the base action. This is the most standard and simple way of doing residual learning.    
IBC + Diffusion: A version of IBC that uses the action from a pre-trained Diffusion Policy as the action initialization in the test-time optimization process.   
Residual + Diffusion: Similar to Residual + BC-RNN, but using a pre-trained Diffusion Policy as the base policy.  

For Maniskill experiments:  
BC + residual: Learns a residual policy that takes as input the action from the BC policy and outputs a delta correction.  

As shown in Table 1, DiffTOP outperforms all these 3 methods on RoboMimic, and as shown in Table 2, DiffTOP also outperforms BC+residual on ManiSkill. This demonstrates that even with the same informed initial guess, DiffTOP is still the best method to refine the initial guessed actions. 

We have also added experiments on both RoboMimic and Maniskill to test the performance of DiffTOP without such informed initial guesses; instead, the actions are either initialized to be 0 or sampled from a Gaussian N(0, 1). The results are shown in Table 3 and 4 in Appendix A.2 of the updated paper. We do notice a drop in performance with random or zero-initialized actions, possibly due to the convergence to bad local minima during nonlinear trajectory optimization, as the reviewer mentioned before. We note that this would not be a drawback of applying DiffTOP in practice 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.

Reviewer comment (Novelty):
The introduction of CVAE to generate multimodal samples is reasonable, but there are many other routes that might potentially handle the multimodality problem. For instance, it seems cheaper and more efficient to consider a differentiable stochastic trajectory optimization framework where a small amount of noise can be added to the action at every step, which forces trajectory optimization itself to be a stochastic procedure that will generate multimodal actions. Have the authors considered this formulation?

Author response:
We thank the reviewer for this suggestion! Stochastic trajectory optimization sounds very interesting. However, at this moment, we are not aware of any efficient and robust implementation of a differentiable version of stochastic trajectory optimization. We would definitely like to try it once there is a library for that, which we leave as interesting future work.

Reviewer comment (Experiments):
CVAEs tend to be harder to train compared to recent diffusion-based generative models. Is it possible to train a diffusion model instead of address the multimodality problem? Maybe it is difficult since the author's framework requires latent states?

Author response:
We thank the reviewer for this suggestion. We think it would definitely be possible to use a diffusion model instead of a CVAE to handle the multimodality problem, and given the recent success of diffusion in image generation, it might achieve even better results compared to CVAE. However, as the reviewer mentioned, it might require additional techniques/efforts to correctly combine diffusion with trajectory optimization in the latent space, which is beyond the current scope of the paper. We leave this as interesting future work.

Reviewer comment (Experiments):
How does DiffTOP perform without informed initial guesses in the RoboMimic experiment?

Author response:
We have also added experiments on both RoboMimic and Maniskill to test the performance of DiffTOP without such informed initial guesses; instead, the actions are either initialized to be 0 or sampled from a Gaussian N(0, 1). The results are shown in Table 3 and 4 in Appendix A.2 of the updated paper. We do notice a drop in performance with random or zero-initialized actions, possibly due to the convergence to bad local minima during nonlinear trajectory optimization, as the reviewer mentioned before. We note that this would not be a drawback of applying DiffTOP in practice 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.

Reviewer comment (Evaluation):
No big weaknesses overall. 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.

Author response:
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.

Reviewer comment (Evaluation):
Are there any cases when TD-MPC performs or could theoretically perform better than DiffTOP?

Author response:
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.

Reviewer comment (Evaluation):
What’s the computational cost of using Theseus for differentiable trajectory optimization vs MPPI?

Author response:
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.