Paper: Efficient Planning with Latent Diffusion (ICLR-2024)

Reviewer comment (Theory):
While I acknowledged the potential strength of hierarchical planning above, since the latent space is weakly regularized, we should not expect there to be too much benefit in leveraging this space for planning as compared to the raw trajectory space.

Author response:
Thank you very much for your question! This is a great question, and the reason behind it is not intuitive. In the default setting, latent action consists of multiple timesteps of state, actions, rewards, and reward-to-go. Using a casual transformer as the encoder and decoder will make the learned latent action representations more predictive. For example, predicting actions based on the state, predicting rewards and reward-to-go based on the state and action. This predictive ability, similar to model-based methods, will make the learned latent action representations more conducive to high-quality planning. At the same time, a weakly regularized latent space does not have much of a negative impact on the quality of the trajectories obtained through planning. This is because the quality of the trajectories is mainly determined by diffusion sampling. LatentDiffuser uses an exact energy-guided diffusion sampling method to sample trajectories. As long as the energy is specified as the expected return value corresponding to high-quality trajectories during sampling, the quality of the sampled trajectories can be guaranteed to some extent.

To verify this point, we also conducted comparative experiments using non-causal transformers. Specifically, we removed the mask part of the causal transformer. This means that during the encoding and decoding process, we allow the model to use the information of the entire subtrajectory to reconstruct any element within that subtrajectory, such as the state using state information. As shown in Figure 6 in the revised version, the experimental results show that LatentDiffuser has a significant performance degradation. Furthermore, we found that using a non-casual transformer is close to the performance when the latent step equals 1. These results are consistent with our previous analysis, and when using a non-casual transformer, the model is also prone to overfitting, causing the learned latent action representations to contain less information, losing a certain degree of "predictability."

We have added the above analysis to Appendix G in the revised version. We hope this analysis can address your concerns.

Reviewer comment (Evaluation):
Could the authors clarify why they expect a weakly regularized latent space to be useful for planning as compared to the raw actions? I see in the appendix that the comparison is made, and raw actions indeed are shown to perform worse. But I would like some intuition for why this should be the case, and if some better hyperparameter tuning could have produced similar results with raw actions.

Author response:
Thank you very much for your question! This is a great question, and the reason behind it is not intuitive. In the default setting, latent action consists of multiple timesteps of state, actions, rewards, and reward-to-go. Using a casual transformer as the encoder and decoder will make the learned latent action representations more predictive. For example, predicting actions based on the state, predicting rewards and reward-to-go based on the state and action. This predictive ability, similar to model-based methods, will make the learned latent action representations more conducive to high-quality planning. At the same time, a weakly regularized latent space does not have much of a negative impact on the quality of the trajectories obtained through planning. This is because the quality of the trajectories is mainly determined by diffusion sampling. LatentDiffuser uses an exact energy-guided diffusion sampling method to sample trajectories. As long as the energy is specified as the expected return value corresponding to high-quality trajectories during sampling, the quality of the sampled trajectories can be guaranteed to some extent.

To verify this point, we also conducted comparative experiments using non-causal transformers. Specifically, we removed the mask part of the causal transformer. This means that during the encoding and decoding process, we allow the model to use the information of the entire subtrajectory to reconstruct any element within that subtrajectory, such as the state using state information. As shown in Figure 6 in the revised version, the experimental results show that LatentDiffuser has a significant performance degradation. Furthermore, we found that using a non-casual transformer is close to the performance when the latent step equals 1. These results are consistent with our previous analysis, and when using a non-casual transformer, the model is also prone to overfitting, causing the learned latent action representations to contain less information, losing a certain degree of "predictability."

We have added the above analysis to Appendix G in the revised version. We hope this analysis can address your concerns.

Reviewer comment (Theory):
The authors use the return to go predictions from the decoder as Q values. Since TD learning is not done to learn optimal Q functions, these Q values correspond to those of the behavior policy. Is there a reason some form of offline Q learning technique such as CQL/BCQ/IQL was not done? Without doing Q learning, should we expect skill stitching to occur?

Author response:
Thank you very much for your question! This is a very good question! 

We primarily considered the following two points as reasons for not using CQL/BCQ/IQL and other TD learning-based offline RL methods:

First, even without using TD learning methods, generative approaches based on diffusion models can implement implicit dynamic programming, thereby achieving trajectory stitching. This point has been preliminarily verified through a simple experiment in Appendix A.1 of the Decision Diffuser paper [1]. We also conducted the same experiment in Appendix D.1 of the revised version and obtained similar results. Of course, exploring the trajectory stitching capabilities of generative approaches based on diffusion models is an intriguing and worthwhile problem for deeper investigation.

Secondly, we drew upon the CEP method [2] in the exact energy-guided diffusion sampling section of LatentDiffuser. The authors proposed an offline RL method based on vanilla TD-learning called QGPO in the CEP paper. QGPO requires additional data sampled from the diffusion model during training to train the Q-function, significantly increasing training overhead. Thanks to the introduction of the latent diffusion model in LatentDiffuser, we can precisely use the pretrained decoder as an estimate for the reward-to-go, thereby improving training efficiency.

We hope that the above explanation answers your question. Thanks again for your inquiry!

[1] Ajay, Anurag, et al. "Is Conditional Generative Modeling All You Need for Decision Making?." *ICLR* 2022.

[2] Lu, Cheng, et al. "Contrastive Energy Prediction for Exact Energy-Guided Diffusion Sampling in Offline Reinforcement Learning." *ICML* 2023.

Reviewer comment (Novelty):
I would also encourage the authors to cite a recent paper from Venkatraman et al. [1], which proposes a very similar algorithm.

Author response:
Thank you very much for your suggestions. We apologize for our oversight in this matter. Since this paper was made public on arXiv during the final stages of our work, just before the ICLR submission deadline, we have not discussed it in the related work section. At the end of Appendix C of the revised version, we set aside a separate remark section to emphasize the similarities and differences between this work and LatentDiffuser. The specific content is pasted below.

Recently, Venkatraman et al. [1] proposed a novel algorithm called LDCQ, which is very similar to LatentDiffuser. Specifically, LDCQ also introduces a latent diffusion model to learn a latent action space. Unlike the latent action used by LatentDiffuser, the learned latent action space in LDCQ belongs to the skill space, similar to what is described in Appendix A. Additionally, LDCQ does not perform planning within the learned skill space but instead uses model-free TD-learning methods to choose the optimal skill at each timestep and obtains the final action using a decoder. In summary, it is quite a coincidence that LDCQ and LatentDiffuser belong to two orthogonal approaches to utilizing latent diffusion models in offline RL. The former still adopts the RL framework to model the offline RL problems, while LatentDiffuser approaches the problem from a conditional generative perspective.

Lastly, we thank you once again for your recognition of our work!

[1] Venkatraman, Siddarth, et al. "Reasoning with latent diffusion in offline reinforcement learning." *arXiv preprint arXiv:2309.06599* (2023).

Reviewer comment (Writing):
The paper repeatedly uses terms like the latent action representation and latent action planning, without a carefully derived definition of it. For self consistency, it would be helpful to define these terms more concretely; otherwise, in its current format, the contributions of the paper can be hard to follow.

Author response:
We sincerely apologize for any confusion caused by our unclear description, and we greatly appreciate your suggestions! The term "latent action representation" refers to the learning of representations for latent actions. In our paper, latent actions are considered a latent variable that needs to be generated by encoding via an encoder. On the other hand, "latent action planning" refers to planning within the latent action space rather than the original space. We hope the above explanation can clear up any confusion. In the revised version, we modified the descriptions of latent action representation and latent action planning according to the above explanations.

Reviewer comment (Writing):
The paper draws inspiration from the TAP paper (Jiang et al., 2023) and integrates a diffuser based latent step within this framework, but in its current form it can be quite confusing to follow and difficult to understand.

Author response:
We apologize for the confusion caused due to our content organization. We hope that our response to your concern will help alleviate the problem. As you mentioned, planning based on the diffusion model is not a novel methodology. However, our paper's main contribution is not solely based on using the diffusion model for implicit planning. Specifically, we aim to introduce the concept of latent actions proposed by TAP and perform efficient planning based on the diffusion model in the latent action space. This approach differs from existing methods of implicit planning based on diffusion models.

Our framework also has two advantages over TAP. First, our framework allows algorithms to plan in a continuous latent action space, enhancing the capability of addressing more complex problems. Second, unlike TAP, which separates the learning phase of latent action representation and the planning phase, our framework integrally combines these two phases to enable end-to-end learning.

We introduce the recently proposed exact energy-guided diffusion sampling method CEP to achieve efficient planning in the latent action space. However, CEP uses diffusion models to parameterize a policy rather than directly generating entire trajectories, so no planning step is involved. Therefore, we establish the theoretical equivalence between planning and exact energy-guided diffusion sampling, deeply integrating CEP into the planning process based on the latent action space. We also improved the CEP algorithm. Specifically, CEP requires additional TD-learning for a Q-function to sample high-quality trajectories. In LatentDiffuser, due to the definition of latent actions, we can directly use the reward-to-go estimates output by the latent diffusion model to replace the Q-function in CEP, thereby reducing the training cost. 

In addition, due to the flexibility of the 'latent action' definition, our framework can also efficiently plan in the raw action and skill spaces (see Appendix F).

In summary, LatentDiffuser is not a simple rearrangement of existing technologies. Although LatentDiffuser is based on existing technologies, neither the latent diffusion model nor the exact energy-guided sampling method (CEP) has been fully applied to the planning paradigm in offline RL. Each tool has its necessity and rationality for implementing efficient planning, and we have also adapted and improved these tools in a task-oriented manner. Our main contribution is precisely manifested in these tools' organic integration, deep adaptation, and improvement.

Reviewer comment (Novelty):
The model seems overly complicated to understand and it is not clear what the contributions of the paper really are. The reviewer understands the relations with the TAP paper and how the work is trying to build from it, but questions whether the planning approach based on diffusion contains any novel approach.

Author response:
We apologize for the confusion caused due to our content organization. We hope that our response to your concern will help alleviate the problem. As you mentioned, planning based on the diffusion model is not a novel methodology. However, our paper's main contribution is not solely based on using the diffusion model for implicit planning. Specifically, we aim to introduce the concept of latent actions proposed by TAP and perform efficient planning based on the diffusion model in the latent action space. This approach differs from existing methods of implicit planning based on diffusion models.

Our framework also has two advantages over TAP. First, our framework allows algorithms to plan in a continuous latent action space, enhancing the capability of addressing more complex problems. Second, unlike TAP, which separates the learning phase of latent action representation and the planning phase, our framework integrally combines these two phases to enable end-to-end learning.

We introduce the recently proposed exact energy-guided diffusion sampling method CEP to achieve efficient planning in the latent action space. However, CEP uses diffusion models to parameterize a policy rather than directly generating entire trajectories, so no planning step is involved. Therefore, we establish the theoretical equivalence between planning and exact energy-guided diffusion sampling, deeply integrating CEP into the planning process based on the latent action space. We also improved the CEP algorithm. Specifically, CEP requires additional TD-learning for a Q-function to sample high-quality trajectories. In LatentDiffuser, due to the definition of latent actions, we can directly use the reward-to-go estimates output by the latent diffusion model to replace the Q-function in CEP, thereby reducing the training cost. 

In addition, due to the flexibility of the 'latent action' definition, our framework can also efficiently plan in the raw action and skill spaces (see Appendix F).

In summary, LatentDiffuser is not a simple rearrangement of existing technologies. Although LatentDiffuser is based on existing technologies, neither the latent diffusion model nor the exact energy-guided sampling method (CEP) has been fully applied to the planning paradigm in offline RL. Each tool has its necessity and rationality for implementing efficient planning, and we have also adapted and improved these tools in a task-oriented manner. Our main contribution is precisely manifested in these tools' organic integration, deep adaptation, and improvement.

Reviewer comment (Theory):
There is uncertainty about whether equation 4 is the objective that needs to be optimized, with the first term generating the planning trajectory based on the latents from the diffuser.

Author response:
Sure, your understanding is entirely correct. I would like to explain Equation 4 further, hoping to provide you with a more intuitive understanding of our paper on a more specific level. 

The reconstruction and entropy terms for the above objective function are easily estimated for any explicit encoder. The challenging aspect pertains to training the cross entropy term, which involves the score-based prior. We adopt a more straightforward yet efficacious approach by training a VAE $\{q_{\boldsymbol{\phi}},p_{\boldsymbol{\psi}}\}$ and a score-based diffusion model $\{q_{\boldsymbol{\theta}}\}$ consecutively based on the offline dataset $\mathcal{D}_{\tau}$.

Reviewer comment (Experiments):
Experimental results are compared with few baselines using the D4RL setup, and the reviewer believes all the results from sections 5.1 - 5.3 are using the same setup, although studying different aspects of the proposed algorithm. It would have been useful to include much simpler proof-of-concept experiments, especially focusing on the planning part and how the diffuser plays a role there.

Author response:
Thank you very much for your suggestion. Since existing generative methods for solving offline RL problems primarily use Gym locomotion as a proof-of-concept benchmark, our original paper version also followed the tradition. It is somewhat challenging to develop a suitable proof-of-concept task to validate the implicit planning capabilities of the diffusion model in a short period, as planning is a very broad concept. To address this, we plan to construct a simple experiment from the perspective of trajectory stitching.

Concretely, we designed an experiment similar to the Decision Diffuser to demonstrate the importance of planning. Most tasks' offline datasets contain a large number of suboptimal trajectories. To learn better policies rather than just simple behavior cloning, trajectory stitching is one of the essential abilities algorithms must possess. To validate whether the LatentDiffuser can achieve trajectory stitching through implicit planning, we adopted the same experimental setup as the Decision Diffuser.

In the maze-2D-open environment, the objective is to navigate towards the target area on the right side, with the reward being the negative distance to this target area. The training dataset comprises $500$ trajectories originating from the left side and terminating at the bottom and $500$ trajectories starting from the bottom and ending at the right side. Each trajectory is constrained to a maximum length of $50$. At test time, the agent begins on the left side and aims to reach the right side as efficiently as possible. As demonstrated in Figure 3 in the revised version and consistent with the findings of the Decision Diffuser, the LatentDiffuser can effectively stitch trajectories from the training dataset to produce trajectories that traverse from the left side to the right side in (near) straight lines.

We hope that this simple experiment on trajectory stitching can help you better understand the implicit planning capabilities of diffusion models and the importance of planning. We have added the aforementioned experiment to Appendix D.1.

Reviewer comment (Evaluation):
The experimental results need an explanation of performance in comparison to the baselines.

Author response:
Thank you very much for your suggestion, which helps enhance the manuscript’s completeness. We have added the following analysis of the experimental results in Appendix D:

This section will then delve into a more detailed analysis of the performance differences among different baselines across various tasks. To provide an intuitive comparison of different algorithms, we classify them from three perspectives — planning, hierarchy, and generative — according to the classification method shown in Table 5. Firstly, the Gym locomotion task has a long horizon, dense rewards, and low action dimensions, making it a baseline test task for offline RL. The results from Table 4 show that generative methods based on diffusion models generally perform better. The community currently attributes this to diffusion models' more powerful representation capabilities in modeling more complex policies or environmental models. However, LatentDiffuser does not demonstrate its advantages well in the low-dimensional action space. Although LatentDiffuser approaches the SOTA performance on this task, it is mainly due to a better diffusion sampling method, which is supported by the solid performance of the QGPO method. Due to dense rewards, planning and hierarchy-based methods, such as TAP and HDMI, have not achieved the best results.

Secondly, the Adroit task is characterized by a high-dimensional action space. This leads to the best performance for TAP and LatentDiffuser (see Table 2), two methods based on latent action, which experimentally verify the effectiveness of latent action. Additionally, generative methods based on diffusion models generally exhibit better performance. However, due to the shorter horizon of the Adroit task, the HDMI method, which is based on planning and hierarchy, does not achieve the best performance.

Lastly, the AntMaze task has a longer horizon and very sparse rewards. This allows latent action ample room for improvement (see Table 2). Moreover, methods based on planning and hierarchy also achieve good results, such as HDMI. In this task, non-generative methods based on planning and hierarchy, such as ComPILE and GoFAR, approach the performance of generative methods without planning and hierarchy (D-QL).

We hope that the above content allows you to gain a deeper understanding of the experimental results.

Reviewer comment (Writing):
The paper requires more proofreading as it contains typos (such as “emploies”).

Author response:
Thank you for your thorough review! We have polished the entire text and corrected the typos in the revised version.

Reviewer comment (Novelty):
While I think this is not a very big deal, the paper relies mostly on existing methodologies.

Author response:
Thank you very much for recognizing our work! As you mentioned, we also believe that the appropriate organic combination of existing tools to solve current open problems is a behavior that the community should encourage. Although LatentDiffuser is based on existing technologies, neither the latent diffusion model nor the exact energy-guided sampling method (CEP) has been fully applied to the planning paradigm in offline RL. We think it is unreasonable to design new tools or methods directly without verifying the effectiveness of existing tools.

Furthermore, we would like to add that by combining the latent diffusion model and CEP with the TAP method, we can not only implement the planning process in TAP more holistically and highly performing but also improve the CEP algorithm. Specifically, CEP requires additional TD-learning for a Q-function to sample high-quality trajectories. In LatentDiffuser, due to the definition of latent actions, we can directly use the reward-to-go estimates output by the latent diffusion model to replace the Q-function in CEP, thereby reducing the training cost. Due to the flexibility of the 'latent action' definition, our framework can also efficiently plan in the raw action and skill spaces (see Appendix F).

In summary, LatentDiffuser is not a simple rearrangement of existing technologies. Each tool has its necessity and rationality for implementing efficient planning, and we have also adapted and improved these tools in a task-oriented manner.

Reviewer comment (Writing):
The overall notation, in particular where x and y get used, together with tau, and state tuples, can be a bit confusing.

Author response:
We apologize for any confusion caused by the misuse of symbols. In the original version, we used boldface x to represent the noised trajectory and the latent action. In the revised version, we have amended the symbols in Section 3, using $\tau_{k}$ to represent the noised trajectory directly. We hope that these modifications will make the manuscript more readable.

Reviewer comment (Experiments):
Can the authors comment on how the sequence length is determined? Is the sequence terminated externally?

Author response:
Thank you for your question. We are unsure what "sequence length" means in this context. We will explain all parts of the paper that involve sequences.

1. If you are referring to the horizon or episode length, the length is determined by the simulator, and when it is terminated is also decided by the simulator. Termination may occur due to reaching a termination condition (task completion or failure) or exceeding the maximum length and being truncated.
2. If you are referring to the timesteps corresponding to the latent action or latent steps, the length is a fixed hyperparameter. In all experiments, the default setting is 3. We also conducted ablation studies on this hyperparameter in Appendix G.
3. If you are referring to the length of the sub-trajectory planned during the planning phase (corresponding to the length of the latent action sequence in LatentDiffuser's energy-guided sampling), the length is also a fixed hyperparameter. In all experiments, the default setting is 40. We also conducted ablation studies on this hyperparameter in Appendix G.
4. If you are referring to the diffusion steps, the length is also a fixed hyperparameter. In all experiments, the default setting is 100. We also conducted ablation studies on this hyperparameter in Appendix G.

Reviewer comment (Theory):
A reference for the proof for theorem 1 is missing.

Author response:
We apologize for our oversight. In the revised version, we have added references to the proof process.

Reviewer comment (Presentation):
Is the text inconsistent with Fig 2? It seems like the initial state in Fig 2 is labelled s0, while in the text there is only ever reference to s1.

Author response:
We apologize. This was indeed an oversight on our part. In the revised version, we have updated Figure 2 in the main text and Figure 5 in the Appendix so that the indices of the sequences start from 1.

Reviewer comment (Writing):
On page 5 typo "emploies".

Author response:
Thank you for your thorough review! We have polished the entire text and corrected the typos in the revised version.

Reviewer comment (Evaluation):
The paper could be stronger if the author could provide some visualization/analysis what the latent diffusers learned

Author response:
We greatly appreciate your taking the time to carefully review our paper and provide valuable suggestions and constructive feedback. This has significantly helped us improve the rigor and completeness of our paper. Additionally, we are grateful for your recognition of our work! Below, we have organized your questions and will address them individually.

`Q1: The paper could be stronger if the author could provide some visualization/analysis what the latent diffusers learned.`

Thank you very much for your suggestion. We have visualized the learned latent action space in Appendix H in the revised version. The specific analysis is as follows.

To gain a more intuitive understanding of the latent action space learned by LatentDiffuser, this section presents a visualization of the latent actions and the corresponding trajectories obtained by decoding them. Specifically, we use the fully trained LatentDiffuser in the Hopper task to sample $5$ trajectories and apply the t-SNE method to reduce the dimensionality of the latent actions associated with these $5$ trajectories for visualization, as shown in Figure 8.
In the visualization of the trajectories, a random trajectory is selected. To facilitate presentation (due to the similarity in the robot shape at adjacent timesteps), we downsampled this trajectory by taking one latent action for every $5$ latent action and overlaid the trajectory images obtained by decoding the adjacent $3$ latent actions with different opacities to achieve the trajectory displayed in Figure 8. It can be seen from the figure that the latent action space learned by LatentDiffuser is a more compact action space, which to some extent, has learned a specific type of macro-action or "skill."

We hope these visualization results give you a more intuitive understanding of our algorithm. Once again, we appreciate your recognition of our work!