Entropy-regularized Diffusion Policy with Q-Ensembles for Offline Reinforcement Learning

We appreciate your constructive feedback and the opportunity to improve our work. Your insights are invaluable, and we look forward to incorporating these revisions to strengthen our submission.

Weakness

1. Why do you introduce a discount factor but in all equations use a finite time horizon?

   Note that all considered environments are continuous tasks, and the return is computed using an infinite time horizon. Hence, a discount factor is used. On the other hand, a finite time horizon is used for the SDE steps. Note that the superscript $t$ of $a_i^t$ is the SDE step, and the subscript $i$ is the RL step.

2. Memory Requirement of 64 Q-ensembles.

   For the Antmaze task, the model requires ~3GB and ~5.5GB for $M=4$ and $M=64$, respectively. We will report more details in the revised manuscript.

3. The sensitivity of $\beta$.

   The additional experiments of $\beta=1,2,4$ are provided in General Response Figure 5 and Table 1. These results show that $\beta=4$ performs the best in both antmaze-medium-diverse and antmaze-medium-play environments. The value of $\beta$ determines the amount of pessimism, and of course, it can be tuned for different datasets.

4. Use of "max Q trick".

   We follow the Diffusion-QL to only use max Q-backup in Antmaze tasks. The ablation study of max Q-backup on selected Antmaze tasks is provided in General Response Table 2. One can see that the max Q-backup trick improves the Antmaze results.

5. Experiments on various $T$ values.

   The experiment of $T=3,5,10$ is shown in General Response Table 3. The step $T=5$ gives the best overall performance. Moreover, as shown in General Response Figure 1, when $T=2$, the diffusion model already achieves a good approximation of the original data distribution.

6. Choice of $\eta$.

   We select the value of $\eta$ by checking the most used value for each task group. Compared to Diffusion-QL which carefully tunes parameters for each environment, our settings are more general. Here, $\eta$ is mainly for balancing the two losses but we did not explore different values in the paper.

7. Extend training time.

   We found that extending the training time decreases the performance of the antmaze task in most offline RL methods such as Diffusion-QL. The Q-ensemble technique can alleviate this problem but is still unstable and challenging. We are happy to note and list it as our future work.

8. Non-scalar comparisons.

   We appreciate the suggestion to use non-scalar comparisons for a more comprehensive statistical analysis. We have starred the project and will be happy to illustrate our non-scalar comparisons in the future.

Questions

1. Figures improvement.

   Thank you for the suggestions on figure improvements. Due to the page limitation of the rebuttal PDF, we will try to update the figures (e.g., use IQM plots with error bars, increase the linewidth, and add symbols/markers) directly in the revised manuscript.

2. Notation problems e.g. "$i$" vs. "$t$" and "$a$" vs. "$x$".

   We will consider swapping the notations to align more closely with standard RL terminology, making the presentation clearer.

3. Why different gradient norms.

   Compared to Diffusion-QL which carefully tunes hyperparameters (e.g., the norm value) for each env, our settings are more general: the values are taken close to the average values as in Diffusion-QL and we use the same gradient norm value for the same type of tasks.

4. Missing Reference to SUNRISE.

   Thank you for the suggestion. We will include a discussion on SUNRISE and its relationship to our method, particularly how it compares to using LCB.

5. $Q^\text{LCB}_\psi$ in Algorithm 1.

   Thanks for pointing this out, we will change it to "Update policy $\pi_\phi$ by (15) using $Q^\text{LCB}_\psi$".

6. The sense of "optimality" regarding the reverse time SDE?

   The path for reverse-time SDE can be diverse and tortuous due to the stochasticity of the diffusion term $\mathrm{d}w$, which requires long timesteps to research the real data distribution. Our optimal sampling utilizes the condition posterior trick (see Proposition 3.1) and can generate new samples in very few steps.

7. Mish activation function.

   In this work, we didn't focus on architecture tuning and our network simply follows the Diffusion-QL with three dense layers and the Mish activation function.

8. Time comparison of diffusion policy vs standard approaches.

   We have conducted additional experiments to provide training time and evaluation time comparison between our method and our method with Gaussian policy between different settings. The results are shown in General Response Table 4. While our diffusion policy requires more computational resources during training compared to standard Gaussian policies, the performance gains justify the additional time investment. The evaluation times remain comparable, ensuring practical applicability in real-world scenarios. These results provide a clearer understanding of the computational trade-offs and reinforce the robustness and effectiveness of our proposed method.

9. How does entropy regularization interact with the SDE sampling?

   During training, we first sample action from the SDE and then learn to maximize the approximation of the entropy to increase exploration of the environment in offline settings. In this stage, the MaxEnt is helpful for policy learning. Then, in inference, we can directly sample actions from the SDE and the entropy won't affect the sampling.

Thank you for your detailed review and for recognizing the strengths of our work. We appreciate your positive feedback on the novelty and theoretical robustness of our proposed method.

1. Computational cost.

   We acknowledge that diffusion policies require longer training and inference time due to the multi-step sampling process. However, the increased time is manageable and justified by the significant performance gains. A detailed comparison is provided in General Response Table 4. Moreover, techniques like the SDE/ODE solver and diffusion distillation have been proven useful for sample accelerating. We will regard it as our future work and try to apply the faster version to some real-world scenarios.

2. Hyperparameter sensitivity...exploration of hyperparameters like entropy temperature and Q-ensemble size is needed.

   Please let us clarify that the ablation studies on entropy temperature and Q-ensemble size are already provided in Table 3 and Table 4 of the paper. Moreover, we further add two ablation experiments on LCB coefficient $\beta$ and diffusion steps $T$ in General Response Table 1 and Table 3.

3. Challenges and limitations in real-world applications.

   As discussed in the Conclusion section, the main challenges and limitations of our work are the action sample time and the computation cost. And the network is also too simple (only 3 dense layers) for complex real-world applications. Our future work will investigate real-time policy distillation under time and compute constraints and explore other efficient network architectures to address real-world datasets and scenarios.

4. Could the authors provide more intuitive explanations or visualizations to illustrate the workings of the mean-reverting SDE and its impact on action distribution sampling.

   Intuitively, the mean-reverting SDE has a forward process (action to noise) and a reverse process (noise to action). The forward process is only used in training to provide an analytically tractable solution for score network learning. In inference, the reverse process (based on the trained score network) is used as the policy to generate actions conditioned on RL states. A rough visualization of the workings of the mean-reverting SDE is provided in General Response Figure 4. This will also be clarified in the revised manuscript.

Thank you for your rebuttal. However, I believe there may have been a misunderstanding regarding my first question. In Table 1, your method underperforms compared to others (e.g., Diff-QL) in the Adroit and Kitchen environments. I question your explanation attributing this to the narrowness of human demonstrations. I think it's more convincing to support your statement to generate additional data for retraining the policy rather than focusing on tuning $\alpha$? Adroit and Kitchen environments are generally considered more challenging than Gym and AntMaze, which raises further questions about this approach.

Additionally, Table 3 suggests that the auto-selection of entropy temperature is more critical for performance in Adroit and Kitchen than that in Gym and AntMaze. Could you elaborate on why this is the case?

Thank you for your thoughtful review and for recognizing the novelty and potential impact of our work. We appreciate your positive feedback and are happy to provide our point-to-point response below:

1. Influence of multi-modality of the policy...SAC with Q-ensemble and LCB

   In offline RL, the pre-collected datasets are always imbalanced and lead to a complex multi-modal policy rather than a Gaussian. We provided a comparison in Table 2 of the paper, which shows that our model outperforms MSG (Gaussian policy with Q-ensembles) on all AntMaze tasks, meaning that the multi-modal policy is important and improves the Gaussian in offline RL. Also, the ablation study in Table 4 of the paper shows that increasing the ensemble size further improves the performance. Moreover, unfortunately, we can't directly compare our method with SAC on offline RL tasks since SAC and its variants are all online algorithms that need to interact with the environment.

2. Comparison to Other Multi-Modal Policies

   Please allow us to clarify that we have compared the behavior cloning (BC), decision transformer (DT), and Diffusion-QL, that are all non-Gaussian multi-modal policies. We appreciate the reviewer's suggestion and will be happy to implement the mentioned Gaussian mixture and energy-based policies in our future work.

3. Reformulation of the Abstract

   Thank you for pointing out the unclear formulation in the abstract. We will revise the abstract to clarify that the entropy of the policy can be approximated in a tractable way.

4. Computational Limitations and Timings

   We acknowledge that the computational time was only briefly mentioned. We add experiments to compare the training time and evaluation time between different policy types, diffusion steps, and the number of critics. The results are shown in General Response Table 4. These results indicate that while the diffusion policy requires longer training times compared to the Gaussian policy, the increase is manageable and justified by the significant performance gains. The evaluation times, however, remain relatively comparable, suggesting that the practical deployment of the trained models is feasible.

Thank you for your thorough review and insightful feedback. Your comments have been invaluable in guiding our efforts to refine and clarify our work. Below, we address your main concerns in detail:

1. Mean-reverting SDE and VP SDE.

   Our mean-reverting SDE is derived from the famous Ornstein-Uhlenbeck (OU) process [1] which has the following form:

   $$\mathrm{d} x = \theta (\mu - x) \mathrm{d} t + \sigma \mathrm{d}w.$$

   As $t \to \infty$, its marginal distribution $p_t(x)$ converges to a stationary Gaussian towards the mean value $\mu$, which gives the informative name: “mean-reverting”. We assume there is no prior knowledge of the actions and thus set $\mu = 0$ to generate actions from standard Gaussian noise. Then, with $\mu = 0$, the mean-reverting SDE has the same form as VP SDE. However, in [42], no solution of the continuous time SDE was given. The authors start from pertubing data with multiple noise scales and generalize this idea with an infinite number of noise scales which makes the perturbed data distributions evolve according to an SDE. They keep using the solution of DDPM while we use Itô's formula to solve the continuous SDE. Compared to the original VP SDE, our mean-reverting SDE is analytically tractable (Eq.(2)) and thus its score $\nabla_{x} \log p_t(x)$ is easier to learn. More importantly, the solution of the mean-reverting SDE can be used for entropy approximation. We will clarify these points and cite [42] appropriately in the revised manuscript.

2. Optimal Sampling and DDIM.

   Yes, our optimal sampling uses the same conditional posterior distribution$p(x_{t-1} | x_t, x_0)$ as DDIM for reverse sampling and we will cite it in Section 3.1 in the revised manuscript.

3. Performance of Optimal Sampling with $N=1$ and $N=2$.

   We provide new samples generated with fewer steps of the toy task in General Response Figure 1.

4. Proof for Eq.(14).

   Eq.(14) is a multiple conditions Bayes' rule given by

   $$p_\phi(a_i^1 \mid a_i^T, a_i^0, s_i) = \frac{ p_\phi(a_i^T \mid a_i^1, s_i) \; p_\phi(a_i^1 \mid a_i^0, s_i)}{p_\phi(a_i^T \mid a_i^0, s_i)}.$$

   Here we use $p_\phi(a_i^1 \mid a_i^T, s_i)$ instead of $p_\phi(a_i^1 \mid a_i^T, a_i^0, s_i)$ since both $a_i^0$ and $a_i^1$ are computed from the diffusion sampling process and the conditional posterior from $a_i^1$ and $a_i^0$ is a certain Gaussian (with known mean and variance, see Appendix A.4). We will clarify it in the revised manuscript to make the notation and proof clearer.

5. Clarification on Eq.(15).

   Substituting Eq.(13) to Eq.(12) gives the new policy objective and its last term should be $\log (p(\hat{a}_i^1 \mid s_i))$. However, we assume that the terminal state $a_i^T$ must be a standard Gaussian and thus we add the condition $a_i^T$ as $\log (p(\hat{a}_i^1 \mid a_i^T, s_i))$ which can be obtained from Eq.(14) and the results won't be affected. (And thank you for pointing out the typo.)

6. Entropy target should be the negative dimension of the action space.

   Yes, our code uses the negative dimension of the action space as the entropy target, the same as SAC. We will correct it in the paper.

7. The $\alpha$ depends on $s_i$ in Eq.(45) while SAC doesn't. Why is the alpha loss designed this way?

   SAC is an online algorithm that can interact with the environment and collect new samples with various state actions. However, the offline dataset is pre-collected and may be imbalanced across different states. Thus our idea for alpha loss is to assign different entropy values to each state based on available data.

8. Inconsistent between Table 1 and Figure 3.

   We made a mistake in Figure 3, Antmaze-medium-diverse-v0, where we used the wrong data source for the curve, resulting in a major inconsistency. The updated figure is provided in General Response Figure 2. The reason for other minor inconsistencies is that we report the mean of the best results in tables while showing the average values of each step in figures. Moreover, the error bar is a default value in seaborn.lineplot which takes 95% confidence.

9. Training curves for other environments?

   We show some training curves in General Response Figure 3 and will also add other curves to the appendix in the revised manuscript.

10. How is the score recorded?

    Consider that the behavior-cloning (BC) loss is not suitable for exploring OOD samples thus we choose to use online evaluation to select the model. It is worth noting that Diffusion-QL has proven the online-selected model has a similar performance as the offline-selected model. This will be clarified in the revised manuscript.

11. Subscript or superscript "$t$" in Eq.(34) and Eq.(35).

    All actions should use superscript sample steps (i.e., $t$, $t-1$, and $0$). We have now fixed them in the manuscript.

12. Why Eq.(43) holds?

    Eq.(43) holds when action states $a_i^1$ and $a_i^0$ are sequentially sampled from a diffusion process. The term $\pi(a_i^0 \mid a_i^1, \hat{a}_i^0)$ is the conditional posterior sampling from Proposition 3.1. We will elaborate more details in the revised manuscript.

13. Likelihood loss and Noise loss in Table 5.

    The "Noise" loss is similar to the simplified loss in DDPM as shown in Eq.(5). In contrast, the "Likelihood" loss is proposed in IR-SDE [2], which forces the model to learn optimal reverse paths from $x_t$ to $x_{t-1}$. Generally, the Likelihood loss is more stable while the Noise loss tends to generate more stochastic samples. This will be clarified in the revised manuscript.

14. Typos

    Thank you for pointing out these typos and we will fix them in the revised manuscript.

[1] Exact numerical simulation of the Ornstein-Uhlenbeck process and its integral. Physical review E, 1996.

[2] Image restoration with mean-reverting stochastic differential equations. ICML, 2023.