Diffusion Policies Creating a Trust Region for Offline Reinforcement Learning

Thank you for your review. We believe there may have been some misunderstandings regarding our paper. We will strive to address your questions thoroughly and hope you will consider re-evaluating our work.

Similarities

Firstly, we have discussed some differences between our method and SRPO in the Global Response. Here we are addressing the reviewer's concerns about "presentation" and "empirical result" similarity.

The presentation. We believe every excellent work, like DQL, SRPO, and IDQL, shares a similar logic in presenting the story in the offline RL setting. The presentation of our work humbly learns from these papers. Thus our presentation flow is not only similar to SRPO but also to DQL and IDQL, as these papers serve as excellent examples for us.

For empirical results. Our method slightly outperforms SRPO on Gym tasks, achieving different SOTA results on various sub-tasks. For instance, SRPO scores higher on walker2d-m-e (114 vs. 110) but lower on hopper-m-e (100 vs. 109). Additionally, our method surpasses SRPO on antmaze-v0 tasks and excels on some antmaze-v2 sub-tasks like antmaze-l-p and antmaze-m-d (Table 6).

Unknown trust region

1. Definition. Formally, by Theorem 1, given any state $s$, we can define a trust region for action by the set {$a | \log p(a|s) \ge \text{threshold}$}, where $\log p(a|s)$ is approximated by the diffusion loss (Eq 4). The magnitude of the threshold can be tuned by the hyperparameter $\alpha$ during the optimization of the final loss (Eq. 5). Thank you for pointing this out; we will make this point clearer in the revision.

2. Trivial solution. For the reviewer's concerns about "move freely" and "how to generate actions...", we have tried to answer them in Global Response. For "trivial solution", such a trivial solution is only possible when dealing with discrete space control questions, and the offline datasets contain information for whole state spaces. However, for continuous control questions, such a trivial solution has a probability measure of 0 since both action and state spaces are continuous, and the size of the offline dataset is finite. Thus, this trivial solution does not apply to our setting. Moreover, even in the discrete case, the trivial solution you mentioned is essentially behavior cloning. When adding the Q function for guidance to improve the policy, $\pi$ will only return the action $a$ with the highest Q value for any given state $s$.

3. Inconsistently. We addressed the reviewer's concerns about "inconsistent" presentation in our global response section. Now, we aim to clarify "why $\pi_\theta$ is instantiated implicitly" in our approach. The reason for the implicit instantiation of $\pi_\theta$ in line 183 is due to the KL loss (Eq 8), which encourages mode seeking (Figure 2). The Gaussian policy has limited expressiveness; therefore, we make $\pi_\theta$ implicit to enhance its expressiveness and ability to seek more modes. It is important to note that Eq 8 is never used in our method; it is included only as a comparison baseline.

Mode seeking behaviors

1. Trivial. The discussion in Section 3 is primarily to demonstrate that our method indeed exhibits mode-seeking behaviors. While SRPO discusses mode-seeking for its own method, our focus here is to discuss the mode-seeking behavior specific to our approach. Therefore, we believe this discussion is necessary and not trivial.

2. Notations. Constrained by the length of the paper, we do not demonstrate all details about the algorithm in Section 3 since it is not our main contribution. However, we did mention in line 191 that more details are deferred to Appendix D. For full details, please refer to Variational Score Distillation [33], Diff-Instruct [23], and Distribution Matching Distillation [37]. These reference indices are consistent with those used in our paper.

3. Reverse KL divergence minimization. This was addressed by SRPO, and since our paper thoroughly compares with SRPO, we believe we have sufficiently addressed the comparison with reverse KL divergence.

4. Theorem 1. Please refer to our global response.

Hard to understand

1. Typos. Thank you for the careful review and valuable suggestions. We will incorporate these corrections in the revised manuscript.

2. Experiments. We have contacted the authors of SRPO by email, and they acknowledged that there is a version mismatch between SRPO and other baselines in Antmaze v0 and v2. Since SRPO does not provide hyperparameters for Antmaze v0, we used the same hyperparameters as for v2 and reported them in the main table. Additionally, as illustrated in line 262 and in Table 6, we also compare our method with SRPO on Antmaze v2, where we use same hyperparameters as those in Antmaze v0 for our method. The SRPO results are taken from the original paper for a fair comparison. And separating table is a good idea. We can separate SRPO from the main table and create an additional table to compare results in the revision, which will only include Antmaze-v2 results if the reviewer believes it is better to present in this way.

Responses to Questions

For Questions 2 and 3, we believe they have been covered in the global response. For Question 1, in offline RL, typically after training and during interaction with the environment, for any given state, we want the agent to generate the action with the highest Q value. Therefore, diversity, in the sense of generating a batch of diverse actions for a given state, is not essential. As long as the action has a high Q value, the diversity of the generated actions does not contribute to increasing the final cumulative reward. A good example is the deterministic policy, which, for a given state, generates only one action. Although it lacks diversity, it can still achieve a high Q value.

We appreciate the reviewer's recognition of our method's novelty, motivation, and empirical performance. Thank you for the careful review.

Responses to Weakness 1.

Thank you for the good suggestion. You are absolutely right that some modes are also missed by $\mathcal{L}_{\text{KL}}$. We will modify the caption accordingly in the revision to accurately reflect this.

Responses to Weakness 2.

Thank you for pointing this out and for your valuable suggestions. We will continue polishing our paper to enhance clarity. We apologize for the oversight in Line 183 with the word "instantiates." We will correct this and thoroughly review the entire paper to eliminate any other typographical errors. Additionally, we will correct the inconsistent wording and change "cooperative policy" to "dual policy" to provide clearer communication to the readers.

Responses to Question 1.

Thank you for the insightful observation. We hypothesize that this behavior is due to neural networks' tendency to learn continuous functions, which can lead to the connection of some modes with artificial dots. This phenomenon also occurs when using diffusion models to mimic distributions with isolated density modes. For example, in Figure 2 of DQL [1], there are similar artifacts in the interpolation of modes when we use pure diffusion policy to do behaviour cloning.

Yes, we have also designed an RL algorithm based on $\mathcal{L}_{\text{KL}}$, as shown in Figure 3 and Appendix G.2. We suspect that interpolation may cause some out-of-distribution action generation, which could explain the performance drop observed in the learning curve in Figure 8.

[1] Wang, Z., Hunt, J. J., & Zhou, M. (2022). Diffusion policies as an expressive policy class for offline reinforcement learning. arXiv preprint arXiv:2208.06193.

Responses to Question 2.

Thank you for the question. We are happy to make some clarification on this point.

1. We don't expect the Gaussian policy to capture all the information from the diffusion model due to the unimodal nature of Gaussian policies. However, the Gaussian policy is sufficient to capture the optimal mode, which has the highest Q value and is usually unimodal. For example, in Figure 2 of DQL [1], even though the diffusion policy initially captures all behavior modes, it eventually focuses only on one optimal mode after policy improvement (guided by Q value). Thus, we believe that in an offline RL setting, the Gaussian policy is able to capture one optimal behavior mode when constrained by the diffusion policy.

2. We chose the Gaussian policy because it is widely used in RL settings, such as in IQL [2] and SAC [3], making it a good starting point to demonstrate the performance of our algorithm. However, our algorithm can accommodate any one-step policy. One possible extension is to use an implicit policy parameterized by a neural network to enhance the expressiveness of the policy and cover multimodal distributions if necessary.

[2] Kostrikov, I., Nair, A., & Levine, S. (2021). Offline reinforcement learning with implicit q-learning. arXiv preprint arXiv:2110.06169.

[3] Haarnoja, T., Zhou, A., Abbeel, P., & Levine, S. (2018, July). Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning (pp. 1861-1870). PMLR.

Responses to Limitations.

Thanks for the discussion. We evaluate our algorithm using widely used benchmarks for fair comparison with other offline RL algorithms. We acknowledge that the Gaussian policy may not have sufficient expressiveness for complex tasks. However, our algorithm framework can accommodate any one-step policy. Therefore, any policy, such as an implicit policy defined by passing random noise through a deep neural network, can potentially be used in our algorithm to handle more complex tasks. This represents an interesting future direction to explore how these diffusion-based RL algorithms perform in more complex settings.

Thank you for your review. We appreciate the time and effort you have dedicated to evaluating our work. It seems there may have been some misunderstandings regarding the main points of our paper. We will do our best to clarify these issues and hope you will consider re-evaluating our work.

Responses to Weakness 1,2,3

We try to address the reviewer's concerns in the Global Response, where we introduce the main goal, contribution, and presentation logic of our paper. We believe this global response adequately addresses the reviewer's concerns regarding the weaknesses. If the reviewer still has any questions about these aspects, we are happy to discuss them further.

Response to Question 1

1.From Toy Example: Let us illustrate this with toy examples. In Figure 2, the first column shows the offline action data distribution, and the second column shows the trust region loss for different actions. We can observe that actions lying in the high-density region of the raw data manifold have less trust region loss. This means that actions similar to in-sample data are trusted and have lower loss. In contrast, actions outside the in-sample data manifold, as shown in the second and third columns, have large trust region loss. This visualization demonstrates that for a given fixed diffusion model $\mu_\phi$, the loss $E[||\mu_\phi(a_\theta+\sigma_t\epsilon,t|s)-a_\theta||^2]$ can be used to check whether the action $a_\theta$ generated by the one-step policy $\pi_\theta$ is similar to or different from in-sample actions. If the generated data is similar to in-sample data, it is trusted and has a lower trust region loss.

2.From Theoretical Perspective: From Theorem 1, the trust region can be defined by the conditional log likelihood. Specifically, for any given state $s$, a trust region of action is a set {$a|\log p(a|s)\geq\text{threshold}$} where the conditional log likelihood is approximated by conditional diffusion loss (Eq 4). The magnitude of the threshold can be tuned by the hyperparameter $\alpha$ during the optimization of the final loss (Eq. 5).

3.Further Discussion: The discussion about the "trust region" is also covered from Line 117 to Line 122. If the reviewer has any further questions about understanding this loss, we are happy to discuss it in more detail.

Response to Question 2

As our goal is not to design a new weight schedule for the diffusion model, we are using the state-of-the-art diffusion weight schedule, EDM [3], which is discussed in Section 4.3 and Appendix C for completeness. However, other diffusion training schedules, such as VP and VE [4], can also be accommodated within our algorithm framework.

[3] Karras, T., Aittala, M., Aila, T., & Laine, S. (2022). Elucidating the design space of diffusion-based generative models. Advances in neural information processing systems, 35, 26565-26577.

[4] Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., & Poole, B. (2020). Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456.

Response to Question 3

In brief, we do not use a diffusion policy to generate actions, thereby avoiding iterative denoising. A comprehensive explanation of this point is discussed in the global response. If the reviewer has any further questions about this point, we are happy to discuss them.

Response to Limitation 1

We are confident that by considering the global response and gaining a more comprehensive understanding of our paper, the logic and presentation will become clearer and more accessible for readers. We kindly request that the reviewer re-evaluate our paper in light of these clarifications.

Response to Limitation 2

As discussed in the global response, our paper is not the same as DQL [1]. On the contrary, we introduce a completely different training scheme to address the drawbacks in DQL. Additionally, [2] is used to support Theorem 1, and we believe that our use of this reference does not diminish the novelty of our new offline RL method design in any way.

We appreciate the reviewer's recognition of novelty, presentation and performance improvement of our paper. Below, we provide detailed clarifications and answers to solve the remaining concerns.

Responses to Weakness 1.

We first want to emphasize that one major contribution of DTQL is for speeding up Diffusion-QL. We proposed a dual policy training approach and achieved dramatically improvement in terms of both training and inference efficiency.

For D4RL benchmark performance, we provide more clarifications below.

1. Compared with DQL, our major baseline, our method demonstrates dramtically speeding up, superior average score and greater stability. We expand the explanation for stability here. As shown in Figure 8 of [2], the training curve of DQL is less stable due to the challenges in computing gradients of the Q-value function while backpropagating through all diffusion timesteps, which may result in a vanishing gradient problem. In contrast, we propose a dual policy trainig approach, and train a one-step Gaussian policy with trust region defined by a diffusion policy. This avoids the gradient vanishing issue and mitigates training instability, as shown in the learning curve in Appendix H.
2. DTQL outperforms IDQL in Gym and slightly underperforms it in Antmaze. However, IDQL is 10x slower than us in inference.
3. Regarding SRPO, our method not only scores higher on Gym tasks but also surpasses SRPO on Antmaze-v0 as shown in Table 1 (from 30.1 to 73.6) and Antmaze-v2 as shown in Appendix G.3.

[2] Lu, C., Chen, H., Chen, J., Su, H., Li, C., & Zhu, J. (2023, July). Contrastive energy prediction for exact energy-guided diffusion sampling in offline reinforcement learning. In International Conference on Machine Learning (pp. 22825-22855). PMLR.

Responses to Weakness 2.

Thank you for your suggestion. We have updated the ablation study to include results across 5 random seeds in the PDF of global response, along with the standard error. We would like to make the following clarifications:

1. The blue line in Figure 5 (a), only used for ablation study, is quite noisy because it does not include the entropy term. After adding the entropy term, the orange line for the reward curve of the proposed algorithm becomes much more stable.
2. We acknowledge that the orange line in the antmaze-large-diverse-v0 appears noisy. However, this level of variations is common in the offline RL setting and is, in fact, found to be more stable than DQL.

Responses to Question 1.

Thanks for giving us opportunity to explain about this, and we will also make more discussion about it in the revision. In sum, KL regularization in Figure 3 is not the same as reverse KL used in SRPO. In Figure 3, we are using the loss from Diff-Instruct [3], Variational Score Distillation [4] and Distribution Score Matching [5]. The detailed algorithm is described in Appendix D. An analysis of the gradients shows distinct differences between the Score Distillation Sampling (SDS) loss in SRPO [1] and our loss:

• As discussed beginning at line 220, for SRPO [1], we have:

$$\nabla_\theta L_{\text{SDS}} =E_{t,s,\epsilon}\left[w(t)(\epsilon_\phi(z_t,t|s) - \epsilon)\frac{\partial z_t}{\partial \theta}\right]$$

• In contrast, the loss we used in Figure 3 is:

$$\nabla_\theta L_{\text{KL}} = E_{t,s,\epsilon}\left[w(t)(\epsilon_\phi(z_t,t|s) - \epsilon_{\text{fake}}(z_t,t|s))\frac{\partial z_t}{\partial \theta}\right]$$

where $\epsilon_{\text{fake}}$ is the score function we learned from the current policy. Thus, the loss we used in Figure 3 provides an updating direction that can be rather “fine” and “sharp” due to the difference between pretrained $\epsilon_\phi$ and $\epsilon_{\text{fake}}$. This approach is encouraging for covering more modes (as shown in Figure 2) and differs from the reverse KL loss (SDS loss) in SRPO. A more detailed comparison has been discussed in [4].

[3] Luo, W., Hu, T., Zhang, S., Sun, J., Li, Z., & Zhang, Z. (2024). Diff-instruct: A universal approach for transferring knowledge from pre-trained diffusion models. Advances in Neural Information Processing Systems, 36.

[4] Wang, Z., Lu, C., Wang, Y., Bao, F., Li, C., Su, H., & Zhu, J. (2024). Prolificdreamer: High-fidelity and diverse text-to-3d generation with variational score distillation. Advances in Neural Information Processing Systems, 36.

[5] Yin, T., Gharbi, M., Zhang, R., Shechtman, E., Durand, F., Freeman, W. T., & Park, T. (2024). One-step diffusion with distribution matching distillation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 6613-6623).

Responses to Question 2.

Thank you for your observation. We will clarify this point in the revision.

1. SRPO uses a deterministic policy, whereas our method employs a stochastic policy. Consequently, after the network forward pass, we need to resample to generate an action, which adds to the inference time.
2. Additionally, we utilize the (stochastic) max Q trick in inference, as in DQL, which involves generating $N$ candidate actions for a given state and then selecting an action randomly with a weight proportional to $\exp{Q(s,a)}$.

Therefore, these two factors—stochastic resampling and max Q —contribute to making our inference slightly slower than the deterministic policy used in SRPO. We will provide a more detailed explanation of this point in the revised version.