This paper presents advanced techniques of training diffusion policies for offline reinforcement learning (RL). At the core is a mean-reverting stochastic differential equation (SDE) that transfers a complex action distribution into a standard Gaussian and then samples actions conditioned on the environment state with a corresponding reverse-time SDE, like a typical diffusion policy. We show that such an SDE has a solution that we can use to calculate the log probability of the policy, yielding an entropy regularizer that improves the exploration of offline datasets. To mitigate the impact of inaccurate value functions from out-of-distribution data points, we further propose to learn the lower confidence bound of Q-ensembles for more robust policy improvement. By combining the entropy-regularized diffusion policy with Q-ensembles in offline RL, our method achieves state-of-the-art performance on most tasks in D4RL benchmarks. Code is available at https://github.com/ruoqizzz/Entropy-Regularized-Diffusion-Policy-with-QEnsemble.
Paper
Similar papers
Peer review
Summary
This work introduces a new (fast) SDE-based sampling technique to derive actions from a diffusion based policy.
Strengths
- Compares favorably against most relevant benchmark (Diff-QL) - Good ablation studies; this was helpful to understand the importance of each component. - Strong connection to standard tools in SDEs from diffusion literature
Weaknesses
I'm not sure how big a weakness this is, but there appears to be some disconnect in the theory vs. experiments for discounting, that is infinite-time horizons. Why do you introduce a discount factor but in all equations use a finite time horizon? - High memory requirement: What is the memory cost (i.e. in VRAM) for an ensemble of 64 models? - Inconsistent use of "max Q trick" from [27]. Can you explain why this is used only in a limited set of environments? What happens if it is removed? A performance comparison would be useful. **Experiments** - Given that this is a (mostly) experimental paper, it would be good to have more experiments, in particular some studies on the sensitivity of $\beta$ (is it OK to be equal for all envs? it seems like the magnitude of Q-value variance for the LCB should be MDP-dependent.) - Use of $T=5$ steps: can you show any (simple) experiments illustrating the use of various $T$ values? It seems implicitly dependent on the environment since it somehow controls the expressivity of the policy (please correct me if wrong) - L153, choice of $\eta$ (see above comments) - In figure 3, can the training time be extended? Especially regarding the decrease in performance for antmaze-med-play; if there is always a decrease over time in training this is important to note. (on the flip side, it can be noted as a useful limitation for future work to improve on) - Using non-scalar comparisons would be beneficial to gain a better idea of the statistical comparison of all algorithms. (cf. https://github.com/google-research/rliable in case data is still available, it should be straightforward to generate the plots)
Questions
- it's a bit hard to see variance of green line in Fig 1, Perhaps extracting it / using IQM plots with error bars would be helpful. (I know it is a toy example, but it does a good job illustrating the utility) - may just be me, but I believe using $t$ index in RL terms is more standard than $i$, similarly for $T$ vs. $L$. Perhaps consider swapping this notation? - on a similar notation front, the use of $a$ seems like a poor choice since starting at L76 it is very general (not yet discussing actions). Maybe switch to $x$? - I'm curious why different gradient norms were required; do you have any insight here? Was there parameter/grad divergence in some cases? - Missing reference to SUNRISE? (I know it is UCB instead of LCB, but perhaps still useful to comment on the relationship): https://arxiv.org/pdf/2007.04938 - In algorithm 1, can you please include Q^LCB_psi rather than Q_psi alone (if it is correct)? Otherwise it is confusing where the previously discussed value estimate comes in. - Can you explain a bit more about the sense of "optimality" regarding the reverse time SDE? - Use of Mish activation function seems a bit non-standard. How important is it? - L356: can you give a rough time comparison of the diffusion policy vs. standard approaches? - How does entropy regularization interact with the SDE sampling? Is MaxEnt required for diffusion to be useful? - Can you increase linewidth in Fig 1a/c and Fig 3? Adding symbol/markers would be helpful too for accessibility minor typographical: - L10: "is tractable and that **it** can..." - L27: rephrase "work introduce" - L95-100, maybe point to Fig 1 to see multi-modality of proposed approach
Rating
7
Confidence
3
Soundness
4
Presentation
3
Contribution
4
Limitations
Beyond the weaknesses mentioned above, I believe the limitations were discussed adequately.
Summary
The paper presents Entropy-Regularized Diffusion Policy with Q-Ensembles for offline reinforcement learning. This method addresses Q-value overestimation on out-of-distribution (OOD) data by using a mean-reverting stochastic differential equation (SDE) to transform action distributions into a Gaussian form, combined with entropy regularization for enhanced exploration and Q-ensembles for pessimistic Q-value estimation. The approach achieves state-of-the-art performance across D4RL benchmarks especially on AntMaze.
Strengths
The paper introduces a method that combines entropy regularization with Q-ensembles within the framework of diffusion policies for offline RL, offering a novel solution to tackle Q-value overestimation and limited exploration of OOD samples. The use of a mean-reverting SDE to model action distributions is theoretically robust and aligns well with the goal of transforming actions into a tractable Gaussian form for efficient sampling. The authors provide detailed theoretical contributions, including the tractability of entropy regularization in diffusion policies and the benefits of using Q-ensembles for robust value estimation.
Weaknesses
1. Computational Cost: High computational resources are required, including extensive training epochs on high-capacity GPUs, which limits accessibility and applicability in real-world scenarios with constrained resources. 2. Hyperparameter Sensitivity: Although some experiments with different hyperparameters are provided, a more thorough exploration of hyperparameters like entropy temperature and Q-ensemble size is needed to further validate the method's robustness. 3. Generality: The evaluation is robust but focused on D4RL benchmarks. More experiments on diverse or real-world offline RL tasks would strengthen the claims about the method's generalizability.
Questions
1. Could you please give a more thorough explanation on why the choice of hyperparameters, especially the entropy temperature plays a critical role in the performance. As we know that the adroit and kitchen tasks are more complex than Gym and AntMaze, I suspect on the explanation about narrowness of human demonstrations. 2. What are the potential challenges or limitations in extending this method to real-world offline RL tasks outside the D4RL benchmarks? 3. 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?
Rating
5
Confidence
4
Soundness
3
Presentation
2
Contribution
3
Limitations
The authors have figured out the limitation of the proposed method on high computational cost and large inference time.
Reponse to rebuttal
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?
We sincerely appreciate your timely feedback and follow-up questions. 1. *Why can't we add additional training data?* In offline RL, we typically avoid modifying datasets to ensure fair comparisons across different methods. While adding additional training data is indeed a potential solution for improving performance, in this study, we focused on making algorithmic adjustments to ensure consistency and fairness in evaluation. This approach allows us to directly compare the effectiveness of our method against others without introducing external variables. However, we acknowledge that augmenting data could be beneficial in real-world applications and plan to explore this in future work. 2. *The results in Table 1.* In Table 1, our method underperforms in the Adroit and Kitchen environments compared to Diffusion-QL, mainly due to the fixed entropy temperature $\alpha$ set at 0.01. This fixed $\alpha$ leads the agent to continuously explore the action space throughout the entire training process, even when encountering unseen states. While exploration is generally advantageous, it can be detrimental in environments with limited data variability like Adroit and Kitchen. With sufficient data, the actor is encouraged to explore guided by accurate Q-values estimates; however, in the case of unseen state-action pairs, such exploration may harm performance. Additionally, unlike in antmaze tasks, random actions are more likely to negatively impact performance in more complex environments like Kitchen. Overall, excessive exploration prevents the agent from effectively leveraging learned strategies from human demonstrations, and random actions are more detrimental in Adroit and Kitchen, where precise control is essential. These factors contribute to the lower performance observed compared to Diffusion-QL in these tasks. 3. *Why tuning $\alpha$?* Auto-tuning $\alpha$ is useful because it dynamically adjusts the balance between exploration and exploitation based on the data. Initially, $\alpha$ is set to a non-zero value to encourage exploration of the action space. As training progresses, especially in environments like Adroit and Kitchen, where precise control is crucial, the auto-tuning mechanism reduces $\alpha$ to near zero. Also, with more accurate Q-functions later in training, this shift towards exploitation helps the agent focus on optimal actions, improving performance.
Summary
The paper proposes to use reverse-time SDE as the policy in an actor-critic algorithm. To make it work, entropy regularization is added, for which an entropy approximation scheme is suggested. Furthermore, to improve stability, an ensemble of Q-networks is employed, and the pessimistic lower-confidence bound (LCB) is taken as the value, i.e., $\mathbb{E}[Q] - \beta \sqrt{\mathbb{V}[Q]}$. Evaluations on D4RL show improved performance compared to baselines.
Strengths
Originality: good. The paper proposes a few novel ideas to make diffusion policy work in offline RL. Quality: good. Overall presentation and evaluation are good. Tested only on D4RL, but this is common in offline RL. Clarity: excellent. Significance: excellent. Potentially a new baseline on D4RL, especially AntMaze.
Weaknesses
1) Influence of multi-modality of the policy not sufficiently explored. Is this in the end what makes it better compared to using a Gaussian? Or is it rather the LCB and Q-ensemble? How would SAC with Q-ensemble and LCB-value perform (i.e., not using the diffusion policy)? 2) No comparison to other methods that use multi-modal policies, e.g., Gaussian mixture model policy ([Wasserstein Gradient Flows for Optimizing Gaussian Mixture Policies](https://proceedings.neurips.cc/paper_files/paper/2023/hash/429b5216a4d08850c586fbf809e17877-Abstract-Conference.html)) or energy-based policies ([Maximum Entropy Reinforcement Learning via Energy-Based Normalizing Flow](https://arxiv.org/abs/2405.13629))
Questions
In the abstract, the authors write "we show that the entropy of such a policy is tractable". This sounds like it is analytically tractable. But in fact an approximation is used (Sec. 3.2). I would suggest to reformulate the statement in the abstract to avoid such misunderstanding.
Rating
7
Confidence
4
Soundness
3
Presentation
4
Contribution
3
Limitations
Computational limitation is briefly mentioned in the conclusion. It would be nice to have a more extended discussion of limitations and providing some numbers on the timings.
Summary
The paper proposes that adding entropy regularization to offline RL is beneficial, and using pessimistic Q-value estimation through ensemble methods can provide a better estimate of the Q-value. Figure 1 explicitly shows the benefit of the ensemble Q method. The methods show impressive performance in D4RL.
Strengths
The paper is well-organized and easy to follow. I found the pessimistic ensemble Q trick, which increases estimation accuracy (Figure 1), to be interesting. It is intriguing to see that main entropy to the agent can also benefit offline RL in some tasks.
Weaknesses
I have some trouble overcoming some gaps in the mathematical proofs and found that some mathematical equations may not hold. The empirical performance in Table 1 is not consistent with the training curve in Figure 3. More detailed questions about these two aspects are in the questions section.
Questions
1. A typo in Eq(5). There is an unnecessary right parenthesis. 2. I found Mean-Reverting SDE is a mathematical rewrite of VP SDE, see [42]. VP SDE has form \begin{align*} dx=-\frac{1}{2}\beta(t)xdt+\sqrt{\beta(t)}dw \end{align*} Mean-Reverting SDE has form \begin{align*} dx=-\theta_txdt+\sqrt{2\theta_t}dw \end{align*} 3. These two SDEs are the same in my mind. Then why is it called Mean-Reverting SDE instead of VP SDE in [42], and why is [42] not cited in this context? Can the authors point out any differences between them? 4. Optimal Sampling seems to use the same trick as DDIM[41], using $p(x_{t-1}|x_t,x_0)$ for reverse sampling instead of $p(x_{t-1}|x_t)$ as in DDPM. If that's the case, I believe the authors should point out and properly cite DDIM in the relevant context. 5. It is amazing to see that in Figure 2, optimal sampling can have such great performance when $N=5$. Can the authors also provide its performance in this toy task when $N=1$ and $N=2$? I am curious about their performance with fewer steps. 6. It is hard for me to prove Eq(14). Can the authors elaborate on more details about its proof? 7. Can the authors elaborate more on how Eq(15) is derived from Eq(12) and Eq(13)? I guess the last term in Eq(15) should be $\log p(\hat a_i^1|s_i)$, not $\log p(\hat a_i^1|a_i^T,s_i)$. (There is a typo in Eq(15). It should be $a_i^T$ not $a_t^T$). 8. The paper proposes using automatic entropy tuning as in SAC in line 216. In line 551, the authors mention that the entropy target is the dimension of the action space. I guess it should be the negative dimension of the action space (as SAC does)? For the alpha loss in Eq (45), $\alpha$ depends on $s_i$, while SAC doesn't. Why is the alpha loss designed this way? Is there any difference in making $\alpha$ depend on $s_i$? 9. I found your method's performance in Antmaze-medium-play-v0, Antmaze-medium-diverse-v0, Antmaze-large-play-v0, and Antmaze-large-diverse-v0 in Table 1 to be inconsistent with the training curve in Figure 3. For example, the score of Antmaze-medium-diverse-v0 in the table is 91.6, while the training curve in Figure 2 shows the training curve score is around 40. Why do they differ so much? The same phenomenon occurs in the other three Antmaze environments. Can the authors explain the mismatch? Can the authors also provide training curves for other environments? 10. In Table 1, how is the score recorded? Is it the final round score, online model selection, offline model selection, or moving average? 11. In Eq (34) and (35), should $t$, $t-1$, and $0$ be subscript or superscript? Why is $t-1$ subscript and $t$ superscript? 12. Can the authors elaborate on why Eq(43) holds? 13. In Table 5, the loss type is Likelihood or Noise. What are their mathematical forms? [41] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502, 2020. [42] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456, 2020.
Rating
5
Confidence
4
Soundness
3
Presentation
3
Contribution
3
Limitations
The paper has some gaps in the mathematical proofs. Additionally, more details of the experiments need to be added. I am more than willing to increase my score if my questions are well addressed by the authors.
Thank you for the detailed responses to my questions. I appreciate the effort put into addressing most of my theoretical concerns. It appears that the diffusion training and sampling schedule is derived from IR-SDE. I would suggest including more theoretical components in the revised manuscript to enhance its completeness. However, a few questions regarding the experimental results remain unresolved: 1.The authors mention the use of online model selection to record performance, while the DQL results presented in the main table are based on offline selection. This discrepancy might lead to a mismatch in comparisons. It would be fairer to compare performance using offline model selection. I understand that the BC loss might not function optimally with this method, but do the authors have any results for offline selection performance in Antmaze, Adroit, or Kitchen based on other criteria? 2.Regarding the training curve provided in the rebuttal, specifically Figure 2 in the global response, I noticed that the DQL performance in antmaze-medium-diverse appears low. While it is expected to be around 78 (offline) or 82 (online), the training curve consistently remains below 40. This may not be a significant issue, but I would appreciate it if the authors could provide an explanation for this observation. 3.Lastly, a minor question: Is the parameter $\alpha(s)$ implemented as a neural network? Since my theoretical questions have been resolved, I have raised my score accordingly.
Thank you for your thorough review and your continued engagement with our work. We appreciate the time and effort you have invested. Below, we address your remaining concerns: 1. *Online vs. Offline Model Selection* We acknowledge that a comparison using offline model selection is essential for ensuring fairness. The table below presents a comparison of our method with Diffusion-QL, including both online and offline results. Additionally, we include our method's performance based on offline selection using the BC Loss criterion, selecting the step where the difference between consecutive steps was less than 4e-3. We will conduct further experiments to identify better offline model selection criteria and to evaluate performance across the Antmaze, Adroit, and Kitchen environments. These results will be incorporated into the revised manuscript. | **AntMaze Tasks** | **Diffusion-QL (Offline)** | **Diffusion-QL (Online)** | **Ours (Offline)** | **Ours (Online)** | | ------------------------- | -------------------------- | ------------------------- | ------------------ | ----------------- | | antmaze-umaze-v0 | 93.4 | 96.0 | 99.0 | **100.0** | | antmaze-umaze-diverse-v0 | 66.2 | **84.0** | 67.5 | 79.8 | | antmaze-medium-play-v0 | 77.6 | 79.8 | 84.0 | **91.4** | | antmaze-medium-diverse-v0 | 78.6 | 82.0 | 85.4 | **91.6** | | antmaze-large-play-v0 | 46.6 | 49.0 | 72.6 | **81.2** | | antmaze-large-diverse-v0 | 56.6 | 61.7 | 65.9 | **76.4** | | **Average** | 69.6 | 75.4 | 79.2 | **86.7** | 2. *Training Curve in Antmaze-Medium-Diverse* We obtained the results using [official github code]({https://github.com/Zhendong-Wang/Diffusion-Policies-for-Offline-RL), where the normalized score reflects the online evaluation at each step. The error bars in the graph represent the standard error. We apologize for the minor error that remains in the updated graph; although we corrected the data source, we inadvertently used the 95\% confidence interval for our method's results. As noted in the paper, the training performance of Diffusion-QL is inherently unstable, with performance occasionally dropping to zero. In addition, the best performance for each run occurs at different training steps. We also observed that in 2 out of 5 runs, the score remained at zero from the beginning to the end of training. These factors contribute to the mean and standard error appearing significantly worse than they might under more stable conditions.Using our online model selection approach, the average result we obtained for Diffusion-QL is 39.8 $\pm$ 43.3. The average result of 3 valid runs (without two runs remaining zero) is 66.3 $\pm$ 33.4. 3. *Implementation of $\alpha(s)$ as a Neural Network* Yes, $\alpha(s)$ is implemented as a neural network with a single hidden layer consisting of 32 units.
Rebuttal Acknowledged
I thank the authors for their responses. They answer all of my questions.
Thank you for your positive feedback and for taking the time to review our rebuttal. We are pleased that our responses addressed most of your concerns and we will certainly include the illustrations in the final version of the paper to enhance clarity and completeness.
Decision
Accept (poster)