Paper: Contrastive Diffuser: Planning Towards High Return States via Contrastive Learning (ICLR-2024)

Reviewer comment (Reproducibility):
The training process of the model remains unclear. While Equation 14 suggests end-to-end training, it is unclear where the contrastive loss is integrated. If added to the diffusion probabilistic model, which aims to reconstruct the un-corrupted trajectories, will this added loss diverge the learning, making the training unstable?

Author response:
The pseudo code of CDiffuser's training process is available in Appendix A.1, and the hyper parameters of  each environment are shown in  Appendix A.3. The contrastive loss is added to the diffusion probabilistic model.  To evaluate how contrastive loss impacts the training process, we visualized the loss during the training process with various hyperparameter values, as is shown in Figure 8 (Appendix A.2). As Figure 8 illustrates,  the added contrastive loss actually has no significant impact on the stability of training.

Reviewer comment (Evaluation):
The claim of 'significant improvements in medium and medium-replay datasets' seems overstated. The improvements are noticeable in only one task from each dataset compared to DD.

Author response:
Thank you for your comment. Here, the significant improvements refers to the comparison of CDiffuser and Diffuser, as CDiffuser takes Diffuser as the backbone. As is shown in Table 1, compared with Diffuser, our approach CDiffuser demonstrates clear advantages across all the environments. Nevertheless, we have modified the description of significant improvements in Section 4.2 to avoid any ambiguous expressions.

Reviewer comment (Experiments):
Extending experiments to more complex control tasks or scenarios with high-dimensional state/action spaces would substantiate the method's effectiveness.

Author response:
Thank you for your comment.  Following Diffuser, to evaluate the performance of CDiffuser on tasks beyond simple periodic tasks such as hopper, walker2d and halfcheetah, we conduct additional experiments on another three RL environments maze2d-umaze-v1, maze2d-medium-v1, maze2d-large-v1 in Section 4.1. which is not periodic  and tests the model's ability for long-horizon planning. Same to Diffuser, we compare CDiffuser with CQL, IQL, and Diffuser. The results are available in Table 1, and are summarized as follows:
|Environment|CQL|IQL|Diffuser|CDiffuser|
|:----|:----|:----|:----|:----|
|Maze2D-UMaze|5.7|47.4|113.9|**142.9±2.2**|
|Maze2D-Medium|5.0|34.9|121.5|**140.0±0.7**|
|Maze2D-Large|12.5|58.6|123.0|**131.5±3.2**|

As can be observed, CDiffuser achieves better results than baselines across all three environments, especially on maze2d-umaze-v1, where CDiffuser demonstrated a clear advantage of 25.4% over Diffuser. The results in the Maze2d environments indicate that CDiffuser is also effective in scenarios that require future planning.

Reviewer comment (Experiments):
A comparative test incorporating DD + Contrastive Learning would add effectiveness to the proposed method.

Author response:
Thank you for your suggestion! Actually, we once had intentions of applying the contrast mechanism on DD, however, we encountered the same problem with other researchers: we are unable to reproduce the results reported in DD (https://github.com/anuragajay/decision-diffuser/issues/1) . Nevertheless, here are the results of our early experiments:

|Environment|DD|DD + Contrastive Learning|
|:----|:----|:----|
|Hopper-Med-Expert|66.7±1.4|**83.0 ± 0.9**|

The results show that incorporating contrastive learning into DD also brings much improvement.

Reviewer comment (Presentation):
Figure 6 requires more explanation, particularly regarding the methodology for generating and comparing states in each showcased scenario.

Author response:
Thank you for your suggestion. We have adjusted Figure 6 and the descriptions corresponding to it in Section 4.4.

Reviewer comment (Presentation):
The distinctions among the three models presented in Figure 5 are not very clear to me.

Author response:
Thank you for your suggestion. We have replotted Figure 5. Colored scatters are the states collected during the models' interaction with the environment, while the gray scatters are the states in the dataset.  As can be observed, compared with Decision Diffuser and Diffuser, CDiffuser achieves higher rewards in both in-distribution areas(circled with blue) and out-of-distribution areas(circled with red).

Reviewer comment (Evaluation):
The performance improvement may not be significant. According to Table 1, the performance is highly comparable to DD.

Author response:
Thank you for your comment. Here, the significant improvements refers to the comparison of CDiffuser and Diffuser, as CDiffuser takes Diffuser as the backbone. As is shown in Table 1, compared with Diffuser, our approach CDiffuser demonstrates clear advantages across all the environments. Nevertheless, we have modified the description of significant improvements in Section 4.2 to avoid any ambiguous expressions.

Reviewer comment (Experiments):
The benchmark only uses Gym.

Author response:
Thank you for your comment.  Following Diffuser, to evaluate the performance of CDiffuser on tasks beyond simple periodic tasks such as hopper, walker2d and halfcheetah, we conduct additional experiments on other three RL environments maze2d-umaze-v1, maze2d-medium-v1, maze2d-large-v1 in Section 4.1. which is not periodic and tests the model's ability for long-horizon planning. Same to Diffuser, we compare CDiffuser with CQL, IQL and Diffuser. The results are available at Table 1, and are summarized as follows:

|Environment|CQL|IQL|Diffuser|CDiffuser|
|:----|:----|:----|:----|:----|
|Maze2D-UMaze|5.7|47.4|113.9|**142.9±2.2**|
|Maze2D-Medium|5.0|34.9|121.5|**140.0±0.7**|
|Maze2D-Large|12.5|58.6|123.0|**131.5±3.2**|


We can observe, CDiffuser achieves better results than baselines across all three environments, especially on maze2d-umaze-v1, where CDiffuser demonstrated a clear advantage of 25.4% over Diffuser. The results in the Maze2d environments indicate that CDiffuser is also effective in scenarios that require future planning.

Reviewer comment (Experiments):
The method is using one step generation from EDP. However, Table 1 and ablation study do not include the comparison against this method.

Author response:
Thank you for your comment. CDiffuser and EDP leverage diffusion models differently. Specifically, following Diffuser and Decision Diffuser, CDiffuser models RL problems as sequence generation problems and leverages diffusion models to generate subsequent  trajectories at each time step. However, methods like EDP and Diffusion-QL leverage diffusion models to act as the policy network, which only predicts action for the current time. In conclusion,  CDiffuser and EDP adopt completely different pipelines.  Considering these, we compare CDiffuser with sequence generation models such as Diffuser and Decision Diffuser rather than EDP.

Reviewer comment (Theory):
The method highly relies on EDP to make the contrastive loss differentiable through the generated states, from my understanding. However, this could be hard to generalize to other diffusion-based methods.

Author response:
Thank you for noticing the extension of CDiffuser on other RL methods. We would like to clearify that Firstly, CDiffuser is independent of EDP. As is mentioned in the reply for Q3, CDiffuser and EDP adopted completely different pipelines. Secondly, transplanting the CDiffuser framework to other RL methods is quite convenient. For methods that predict the future states such as Diffuser and Decision Diffuser, the CDiffuser framework can be naturally applied to them as  our contrastive loss $\mathcal{L}_{c}$ in Equation (13) only requires states to contrast. For other methods that only predict actions,  such as Diffusion-QL and EDP, we just need to make little modifications to enable them to predict states while predicting actions, and then our framework can be applied to these methods.  Moreover, the CDiffuser framework is not limited to contrast over states. For example, we can conduct contrast over actions. We leave that as our future work.

Reviewer comment (Experiments):
The original diffuser separately trains the auxiliary return prediction model on all data. This modification is not discussed and experimentally validated.

Author response:
We separately train the auxiliary return prediction model on all data as is described in Equation (12) and Equation (14), just the same as Diffuser does. The parameters of return predictor $\mathcal{J}$ and diffusion model are independent, therefore updating the parameters of $\mathcal{J}$ separately is equal to updating $\mathcal{J}$'s parameters with $\mathcal{L}$.

Reviewer comment (Theory):
The contrastive loss Equation (9) seems to not be a common form. Usually the denominator considers all the samples, for example in [2,3]. This is a concern on the correctness of this implementation and a justification is needed.

Author response:
Thank you for your question. **Firstly**, sampling a batch as negative samples is commonly adopted in many previous works[1, 4, 5]. The underlying reason is that Equation (9) is iterated for a huge number of steps (1e6 for CDiffuser), which makes Equation (9) equivalent to considering all the samples in the denominator.
**Secondly**, the contrastive loss in Equation (9) is derived from InfoNCE in [2,3]. We made a few modifications to adapt it to the scenario of a large number of samples. For the negative samples, we randomly sample a batch of size $\kappa$ for calculating the distance of the sample $\hat{s}$ from the original sample space. The reason for this modification is that even if we use the anchor function to classify a small portion of samples as negative (For example, 30%), calculating the distance for all samples is impractical due to the extremely large number of samples (190,000 in the case of med-expert). However, this challenge can be approximately addressed by randomly sampling a batch of size $\kappa$ and iterating a huge number of steps externally (1e6 for CDiffuser), as is the denominator of Equation (9). For the positive samples, we take multiple positive samples for each sample and average the loss values. This extends InfoNCE in [2,3] to the case of multiple positive samples.  The underlying reason is that we want the model learn a distribution close to the distribution of positive samples, not just reproducing a single positive sample. It is worth noting that the numerator of Equation (9) should have a coefficient of $
\frac{1}{\kappa}$, however, this will transform into a coefficient during the process of gradient descent. Therefore, we merge it into the weight of the contrastive learning loss, $\lambda_c$, and omit it from equation (9). 

[1]Wang, Hao, et al. "Knowledge-Adaptive Contrastive Learning for Recommendation." Proceedings of the Sixteenth ACM International Conference on Web Search and Data Mining. 2023.\
[2] Oord, Aaron van den, Yazhe Li, and Oriol Vinyals. "Representation learning with contrastive predictive coding." arXiv preprint arXiv:1807.03748 (2018).\
[3] Khosla, Prannay, et al. "Supervised contrastive learning." Advances in neural information processing systems 33 (2020): 18661-18673.\
[4] Chen, Ting, et al. "A simple framework for contrastive learning of visual representations." International conference on machine learning. PMLR, 2020.\
[5] Wonsung Lee, Jaeyoon Chun, Youngmin Lee, Kyoungsoo Park, and Sungrae Park. 2022. Contrastive Learning for Knowledge Tracing. In Proceedings of the ACM Web Conference 2022 (WWW '22). Association for Computing Machinery, New York, NY, USA, 2330–2338. https://doi.org/10.1145/3485447.3512105

Reviewer comment (Evaluation):
While being best and second best in the med-replay datasets, it could be argued whether being so close to the other results can be called significant improvements, and highlighting the second best is potentially done to have the results in the best possible light. However, the authors put their results in ample perspective and give reasonable hypotheses about the impact of expert examples.

Author response:
Thank you for your comment. Here, the significant improvements refers to the comparison of CDiffuser and Diffuser, as CDiffuser takes Diffuser as the backbone. As is shown in Table 1, compared with Diffuser, our approach CDiffuser demonstrates clear advantages across all the 12 settings. Nevertheless, we have modified the description of significant improvements in Section 4.2 to avoid any ambiguous expressions.

Reviewer comment (Presentation):
In Figure 5 it is hard to see what is supposed to be in and out of distribution. Maybe some circles could help make the points from section 4.4.

Author response:
Thank you for your suggestion. We have replotted Figure 5. Colored scatters are the states collected during the models' interaction with the environment, while the gray scatters are the states in the dataset.  As can be observed, compared with Decision Diffuser and Diffuser, CDiffuser achieves higher rewards in both in-distribution areas(circled with blue) and out-of-distribution areas(circled with red).

Reviewer comment (Presentation):
All three figures also look very alike. While the idea of comparison is clear, the overall value of this is uncertain.

Author response:
Thank you for your suggestion. We have replotted Figure 5. Colored scatters are the states collected during the models' interaction with the environment, while the gray scatters are the states in the dataset.  As can be observed, compared with Decision Diffuser and Diffuser, CDiffuser achieves higher rewards in both in-distribution areas(circled with blue) and out-of-distribution areas(circled with red).

Reviewer comment (Presentation):
The nuanced color changes in the figures are hard to see, and some people are color blind.

Author response:
Thank you for your suggestion. We have replotted Figure 5. Colored scatters are the states collected during the models' interaction with the environment, while the gray scatters are the states in the dataset.  As can be observed, compared with Decision Diffuser and Diffuser, CDiffuser achieves higher rewards in both in-distribution areas(circled with blue) and out-of-distribution areas(circled with red).