DiffTORI: Differentiable Trajectory Optimization for Deep Reinforcement and Imitation Learning

Summary

This paper proposes DiffTOP as a model-based approach to reinforcement learning and behavior cloning. DiffTOP learns a cost function and a dynamics model through differentiable trajectory optimization, and then uses the learned model during inference for online optimization. For model-based RL tasks, DiffTOP is built on a recent work of TD-MPC and resolves the objective mismatch issue that TD-MPC failed to address. For behavior cloning, the DiffTOP can be used with a different loss function tailed to BC, with an additional consideration of learning a multimodal policy which is achieved by the use of CVAE. The proposed DiffTOP framework is evaluated on 15 model-based RL tasks and 35 imitation learning tasks with high-dimensional observation spaces.

Strengths

* The authors demonstrate broad applicability of the proposed DiffTOP framework in both reinforcement learning and behavior cloning settings, with varying problem complexity ranging from simple continuous control (such as cartpole) to complex object manipulation tasks (such as robomimic or maniskill).

Weaknesses

[EDIT] The concerns were resolved throughout the rebuttal. In particular, my initial assessment of the paper regarding the ignorance of TD-MPC2 was inaccurate. Despite the thorough results presented in the paper, I have to say that it is not well contextualized relative to the existing literature. * First and foremost, the idea of leveraging model-based differentiable optimization for reinforcement learning and learning-based control is not new, and the field is rapidly expanding with new methods emerging in multiple domains, including control theory, robotics, and machine learning. For instance, Nikishin et al. [1] proposes an approach to differentiable optimization of trajectories (for an infinite horizon) by leveraging implicit function theorem and deep neural networks. Cheng et al. [2][3] perform differentiable trajectory optimization to learn a feedback policy, although they assume a known model-based function class for the dynamics and the policy. Sacks et al. [4] propose to optimize the inner-loop of an MPC optimization algorithm, though their method is based on MPPI. Given the abundance of similar ideas in the literature, the authors are encouraged to spend time to conduct a more thorough literature review. * Most critically, the authors of TD-MPC recently published a new paper titled TD-MPC2 [5] at ICLR 2024, in which the improvement from TD-MPC includes the mitigation of objective mismatch. Since the motivation is very similar to that of this paper and the authors of this paper completely disregard the new work, the contribution of this paper is highly questionable (even though the algorithmic details are still different between TD-MPC2 and DiffTOP). In particular, the following motivational statement made in Section 4 is no longer true: “Existing model-based RL algorithms such as TD-MPC suffer from the objective mismatch issue … DifftTOP addresses this issue …” [1] Nikishin, Evgenii, Romina Abachi, Rishabh Agarwal, and Pierre-Luc Bacon. "Control-oriented model-based reinforcement learning with implicit differentiation." In Proceedings of the AAAI Conference on Artificial Intelligence, vol. 36, no. 7, pp. 7886-7894. 2022. [2] Cheng, Sheng, Minkyung Kim, Lin Song, Chengyu Yang, Yiquan Jin, Shenlong Wang, and Naira Hovakimyan. "Difftune: Auto-tuning through auto-differentiation." arXiv preprint arXiv:2209.10021 (2022). [3] Cheng, Sheng, Lin Song, Minkyung Kim, Shenlong Wang, and Naira Hovakimyan. "DiffTune $^+ $: Hyperparameter-Free Auto-Tuning using Auto-Differentiation." In Learning for Dynamics and Control Conference, pp. 170-183. PMLR, 2023. [4] Sacks, Jacob, Rwik Rana, Kevin Huang, Alex Spitzer, Guanya Shi, and Byron Boots. "Deep model predictive optimization." arXiv preprint arXiv:2310.04590 (2023). [5] Hansen, Nicklas, Hao Su, and Xiaolong Wang. "TD-MPC2: Scalable, Robust World Models for Continuous Control." In The Twelfth International Conference on Learning Representations (2024).

Questions

* Regarding the use of CVAE for capturing multimodal action distributions for behavioral cloning, it seems that the distribution representation in the latent space is still unimodal Gaussian. Although it should work just fine, I wonder if the authors have considered to leverage a discrete categorical latent variable for explicitly modeling multi-modality even in the latent space to possibly improve the performance (see e.g., [6]) [6] Ivanovic, B. and Pavone, M., 2019. The trajectron: Probabilistic multi-agent trajectory modeling with dynamic spatiotemporal graphs. In Proceedings of the IEEE/CVF international conference on computer vision (pp. 2375-2384).

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

[EDIT] The concerns were resolved throughout the rebuttal. In particular, my initial assessment of the paper regarding the ignorance of TD-MPC2 was inaccurate. * As mentioned above, the contribution of this paper is quite limited especially in light of the recent development of TD-MPC2.

Summary

The paper introduces DiffTOP, an approach leveraging differentiable trajectory optimization as the policy representation to enhance performance in deep reinforcement learning (RL) and imitation learning (IL). By utilizing the advancements in differentiable trajectory optimization, DiffTOP addresses the "objective mismatch" issue prevalent in prior model-based RL algorithms by optimizing the dynamics model to directly maximize task performance. The method is benchmarked across various robotic manipulation tasks with high-dimensional sensory inputs and demonstrates superior performance over state-of-the-art methods in both model-based RL and IL domains.

Strengths

1. The paper is well-written, structured, and intuitive, making it easy to understand. 2. The experiments are well-executed with various benchmarks. The comparison with baselines is thorough, including state-of-the-art methods like DP3. 3. The paper provides experiments in both imitation learning and reinforcement learning settings, which is commendable. 4. DiffTOP is combined with various existing approaches to demonstrate the generalization of the method.

Weaknesses

1. The work is not well-contextualized. While the paper makes a general claim about a policy class using differentiable trajectory optimization, using such technique is not new in the field, as the authors point out in the related work section. Despite the authors mention some differences in implementation (e.g., whether the dynamics model is learned), the paper presents this differentiable trajectory optimization as a new contribution. It would be beneficial for the introduction and diagrams to highlight the specific differences between the proposed work and other differentiable trajectory optimization approaches, rather than focusing on the differences among EBM Diffusion policy and TD-MPC. Additionally, calling the method Differentiable Trajectory Optimization might be inaccurate since it encompasses a broad range of works and could be misleading, failing to capture the differences. 2. While the execution is commendable, the design choices could be better justified, and more analysis would strengthen the paper. Particularly, as mentioned in the related work section, there are closely related works like [2][19][42], but there is no direct comparison with them in the experiments. More analysis on the design choices, such as using TD-MPC for model-based RL and learning dynamic functions, would also be appreciated. 3. The idea is not particularly novel. The combination of TD-MPC and differentiable trajectory optimization is rather straightforward. 4. It would also be beneficial to have more validation on real robots.

Questions

see the weaknesses section.

Rating

5

Confidence

4

Soundness

3

Presentation

2

Contribution

2

Limitations

yes

Summary

The paper introduces DiffTOP, a novel policy class for reinforcement learning (RL) and imitation learning (IL) that utilizes differentiable trajectory optimization to generate policy actions. DiffTOP leverages recent advancements in differentiable trajectory optimization, allowing end-to-end learning of cost and dynamics functions through gradient computation. The approach addresses the "objective mismatch" problem in model-based RL by optimizing dynamics and reward models to directly maximize task performance. For imitation learning, DiffTOP optimizes actions with a learned cost function at test time, outperforming previous methods. The authors benchmark DiffTOP on 15 model-based RL tasks and 35 imitation learning tasks with high-dimensional inputs like images and point clouds. The results demonstrate that DiffTOP surpasses prior state-of-the-art methods in both domains. The paper includes analysis and ablation studies to provide insights into DiffTOP's learning procedure and performance gains.

Strengths

This paper has several strengths: - The paper presents a robust and technically rigorous study with extensive experiments and ablation studies across a wide range of challenging environments. The proposed method DiffTOP achieves superior results in both RL and IL tasks with high-dimensional sensory observations. - The approach effectively tackles the important "objective mismatch" problem inherent in model-based reinforcement learning. - The ability to compute policy gradients directly with respect to the parameters describing the observation and transition model is a significant advancement, eliminating the need for sample-based estimates. - DiffTOP alleviates the model mismatch problem by learning the model concurrently with optimization, similar to TD-MPC.

Weaknesses

While the paper has many strengths, there are also some potential weaknesses or areas that could be improved: - The paper attempts to condense a large amount of information into limited space, which may affect readability and clarity, particularly for readers less familiar with TD-MPC. - The trajectory optimization solver used (Theseus) does not support constraint optimization, requiring manual unrolling of dynamics instead of presenting it as a constraint to the optimizer.

Questions

- For ManiSkill tasks, have the authors tried more advanced BC baselines (e.g., Diffusion Policy)? It seems the performance of baselines is not very strong.

Rating

8

Confidence

4

Soundness

3

Presentation

2

Contribution

4

Limitations

The authors have discussed a few limitations in the paper.

Summary

The paper presents a method that uses Differentiable Trajectory Optimization as a policy representation. The proposed method extends the work of Temporal difference learning for model predictive control (TD-MPC) by incorporating a policy-gradient loss for which analytical backpropagation is possible thanks to the differentiable properties of the used trajectory optimization scheme. The paper evaluates the effectiveness of trajectory optimization as a policy representation for both model-based RL and Imitation Learning, benchmarking the approach in the DeepMind control tasks, the MetaWorld benchmark, the Robomimic benchmark, and the ManiSkill benchmark. Furthermore, the work presents comparisons against other policy representations like traditional feed-forward policies, Energy-based methods (EBMs), and Diffusion Policies.

Strengths

- The paper is well-written, the work is interesting and it presents a simple yet effective extension of TD-MPC. - The method enables the use of trajectory optimization using high-dimensional observations and it is also able to deal with multimodality in the solution space. - The method is evaluated thoroughly in different setups (Model-based RL, IL) and against other classes of policy representations, outperforming most baselines in terms of sample efficiency and final reward.

Weaknesses

[Minor] - The main weakness of the approach is the high computational cost of solving trajectory optimization at test time. The appendix reports 0.052 seconds to infer the action for one timestep (20 Hz), despite the fact that the time horizon is considerably short (1 to 3 steps). Such inference speed might hinder the deployment in certain real-world scenarios.

Questions

- Is there any specific reason why the Levenberg-Marquardt solver was used? Using a second-order solver might lead to fewer iterations to get to a solution, but the time per iteration might also be higher. A simpler solver (SGD) could do the job and be faster, improving the overall training time and inference speed. - Are there any insights on why increasing the time horizon for the trajectory-optimization policy leads to the slight decrease in performance reported in Appendix A2.3? Typically, a longer time horizon leads to trajectories of better quality. Does the trajectory optimization problem reach convergence for the longer time horizons that were tested?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes. Limitations of the method are mentioned and I agree with the author's claims that there are no direct potential societal impacts.

Response to Authors and Acknowledgement of Inaccurate Feedback in Initial Review

I sincerely thank the authors for their time and effort to prepare the rebuttal. Please find below my response. * > We thank the reviewer for bringing up these related works. We agree that the field of using differentiable optimization for end-to-end learning is rapidly expanding, and we have already discussed and cited many relevant works in the “Differentiable optimization” paragraph in the related work section. We are sorry for the omission of these 4 related works and we have updated section 2 of our paper to include a thorough discussion of them, detailed as follows: * Thank you for conducting the thorough literature review. It should further clarify the uniqueness of DiffTOP when compared against prior approaches that share similarity. * > We thank the reviewer for raising this issue. We respectfully disagree with this assessment. TD-MPC2 [5] does not include any improvements over TD-MPC [6] that mitigate the objective mismatch issue. TD-MPC2 inherits the same training objective from TD-MPC, as evidenced by Equation (3) in TD-MPC2 and Equations (8), (9), and (10) in TD-MPC. This means that TD-MPC2 still suffers from the objective mismatch issue, where the dynamics model is optimized to predict future states in latent space (via the latent state consistency loss), which is not necessarily aligned with the goal of achieving high task return when using the dynamics model for planning. The primary improvements in TD-MPC2 over TD-MPC are related to exploring different architectural variations and new algorithmic design choices for more stable training in multi-task settings, rather than addressing the objective mismatch issue. Therefore, we believe our contribution and motivation remains valid. We have updated section 2 of our paper to include a discussion of TD-MPC2. * Thank you for your response. I appreciate the level of details provided in the exposition and apologize for mistakenly making an inaccurate claim in my original review; I had been confused in part by the following statement made in Section 3.1 of TD-MPC2 paper, which I cite below for transparency: * > However, accurately predicting raw future observations (e.g., images or proprioceptive features) over long time horizons is a difficult problem, and does not necessarily lead to effective control (Lambert et al., 2020). Rather than explicitly modeling dynamics using reconstruction, TD-MPC2 aims to learn a maximally useful model: a model that accurately predicts outcomes (returns) conditioned on a sequence of actions. Here, Lambert et al. (2020) is the original objective mismatch paper. It implicitly states that TD-MPC2 takes care of the objective mismatch issue. However, based on your response (and upon closer look at Appendix A of TD-MPC2), I agree that the model objective of TD-MPC2 is largely the same as that of TD-MPC. Thus, the results presented in Figure 3 of the present paper already confirms the superiority of DiffTOP over TD-MPC/TD-MPC2 in terms of directly optimizing the task performance. In the revised paper, the authors may want to explicitly counter-argue the above statement made by the TD-MPC2 authors, rather than just stating that TD-MPC suffers from the objective mismatch issue (e.g., in Section 4.1), since these two statements are indeed conflicting with each other and may cause confusion for other people as well. Again, I appreciate the authors for taking the review comments seriously. I have updated my review and revised the score.