Diffusion-based Reinforcement Learning via Q-weighted Variational Policy Optimization

We thank the reviewer for your valuable feedback and comments! We itemize the weaknesses you mentioned and answer them.

Q1: While Figure 2 is informative in demonstrating the ideal effect of the entropy term, experimental results on toy examples would provide stronger evidence than an illustration.

A1: We perform a toy experiment on a continuous bandit problem to further present the effect of the entropy term. The results are shown in Figure 1 of the attached pdf file, which is similar to the illustration. We will add it in the final version.

Q2: Since QVPO relies on action samples for both optimization and action selection, this will increase computational demands.

A2: Actually, that is not the weakness but the advantage of QVPO. QVPO can sufficiently utilize the parallel computing ability of GPU via the multiple sampling and selection procedure. We observed that the previous diffusion-based RL methods generally have a low gpu utilization rate since the computational bottleneck of the forward process in the diffusion. In that case, parallelly sampling multiple samples does not add much extra time cost.

Besides, QVPO does not need to use the gradient from Q network to optimize the diffusion policy. This also reduces the time cost to a certain extent compared with other diffusion RL methods like DIPO. Finally, the comprehensive training time and inference time comparison can be referred to in Global Rebuttal.

Q3: The tasks are limited and relatively trivial, making it difficult to tell performance in more complex tasks, such as robot manipulation. In particular, when the action space is larger, it is unclear how the entropy term will scale given its reliance on uniform action distributions.

A3: We need to clarify that these 5 mujoco environments are the standard online RL benchmarks and most of the previous online RL methods (e.g., SAC, DIPO, etc.) only choose them to do the experiments. That is why we verify our algorithm on them.

Besides, to further present the efficiency of QVPO, we also apply it to the recently released HumanoidBench environment, which is a complex continuous control benchmark based on Unitree H1 humanoid with 151 observation dimension and 61 action dimension. The results are shown in Figure 3 of the attached pdf file. It can be observed that model-free RL methods (SAC, PPO) do not work totally, while QVPO (as a model-free method) achieves superior performance compared with advanced model-based methods (Dreamer-v3, TD-MPC2).

The entropy coefficient is set 0.01 here as well. In fact, we found that our algorithm is robust enough with the scaling of $w_{ent}$ in 1e-2 to 1e-3.

Q4(1): Can the author provide an explanation for why the exploration ability of diffusion policies declines with a limited number of diffusion steps?

A4(1): Note that the Gaussian noise is added at each diffusion denoising step. In other words, more diffusion steps lead to more possibilities. Thus, the exploration ability of diffusion policies declines with a limited number of diffusion steps.

Q4(2): Additionally, what happens when the diffusion steps are even fewer, say 5, a common value in other diffusion policy literature? I suppose the entropy term should perform better in such cases.

A4(2): We add the result of the ablation study with 5 diffusion steps in the attached pdf file. It can be found that the performance is poor. Hence, different from offline RL methods, a certain number of diffusion steps is required in online RL setting. This is also verified in previous works like DIPO [R7].

Q5: If the maximum entropy term is used throughout the training, the policy may continue ... potentially hurting performance and convergence. Would a scheduling mechanism improve the final performance?

A5: The entropy term is not constant as you thought. As mentioned in Algorithm 1, line 8. You can find that the actual entropy coefficient $w_{ent}(s)=w_{ent}w_{eq}(s,a_{max})$ is related to $w_{eq}(s,a_{max})=A(s,a_{max}), A(s,a_{max})\ge 0$. When QVPO converges, $A(s,a_{max})\rightarrow 0$ and $w_{ent}(s)\rightarrow 0$ .

Q6: Can the author explain why the Q-weighted VLO loss is superior to QSM regarding the additional errors in policy optimization?

A6: As mentioned in lines 43-46, QSM has a doubled approximation error from the alignment process. More concretely, the trained Q value function has an approximation error. QSM utilizes the gradient of this approximated Q function to train the score model. This leads to the doubled approximation error in the score model. That is why the performance of QSM is worse than QVPO as shown in the attached pdf. Besides, the policy improvement stage of DIPO is sampling-based, which implies QVPO is more likely to cross the local optimum.

Q7: In line 300, it is mentioned that action selection is used during inference, which is a standard technique in offline RL. However, is the same technique applied to the baselines? If not, this would be unfair.

A7: We need to clarify that there is a contradiction if we apply this technique to other methods.

As mentioned in offline RL works like IDQL [R6], applying action selection to diffusion policy is to yield a deterministic policy during inference. For MLP or Gaussian policy, we can directly output the deterministic action or the mean of Gaussian policy to achieve this goal. Hence, this trick is not appropriate for classical online RL methods (e.g., SAC, TD3).

For diffusion-based RL, like DIPO, it yields the deterministic policy via fixing the denoising process (i.e., updating with the mean value in each diffusion step). In that case, there exists a contradiction if action selection is also applied during inference.

[R6] Hansen P, et al. Idql: Implicit q-learning as an actor-critic method with diffusion policies.

[R7] Yang L, et al. Policy representation via diffusion probability model for reinforcement learning.

Q1: My main concern is that the design of equivalent Q-weight transformation functions is heuristic and may need some theoretical intuition or backup.

A1: Thank you for raising the concern. Here is the theoretical proof to show the convergence of QVPO with qadv weight transformation function. (We are afraid OpenReview cannot display the equation clearly. You can copy the following equations as well as texts into HackMD for a clear presentation.)

Assume the new diffusion policy after one update can be approximately expressed as

$$$$

where $a^\star\mid s$ is the action that can maximize $Q_{\pi_{old}}(s, a)$ in the state $s$, $\mathcal{N}(a^\star \mid s, \epsilon)$ is the neighborhood of $a^\star\mid s$ with a small radius $\epsilon$, $S_{\mathcal{N}(a^\star\mid s, \epsilon)}$ is the area of this neighborhood, and $p_{data}(s)$ is the sampling distribution of the state. This assumption is straightforward: the training sample's generation probability in the diffusion model will be improved and the improved probability is proportional to its weight.

Now consider the improvement of the RL objective:

$$\begin{aligned} \mathcal{J}(\pi_{new}) - \mathcal{J}(\pi_{old}) & = \mathbb{E}_{s\sim \rho_0}\left[V_{\pi_{new}}(s) - V_{\pi_{old}}(s)\right] \\ & = \mathbb{E}_{s\sim \rho_0}\left[\mathbb{E}_{a\sim \pi_{new}(a\mid s)}\left[Q_{\pi_{new}}(s, a)\right] - V_{\pi_{old}}(s)\right] \\ &=\mathbb{E}_{s\sim \rho_0}\left[\mathbb{E}_{a\sim \pi_{new}(a\mid s)}\left[Q_{\pi_{new}}(s, a)-Q_{\pi_{old}}(s, a)\right]\right] \\ &\quad \quad + \mathbb{E}_{s\sim \rho_0}\left[\mathbb{E}_{a\sim \pi_{new}(a\mid s)}\left[ Q_{\pi_{old}}(s, a)\right] - V_{\pi_{old}}(s)\right] \\ &=\mathbb{E}_{s\sim \rho_0}\left[\mathbb{E}_{a\sim \pi_{new}(a\mid s)}\left[Q_{\pi_{new}}(s, a)-Q_{\pi_{old}}(s, a)\right]\right] + \mathbb{E}_{s,a\sim \rho_0, \pi_{new}(a\mid s)}\left[A_{\pi_{old}}(s,a)\right] \\ \end{aligned}$$

The first term here can be further expanded according to the Bellman equation:

$$\mathbb{E}_{s\sim \rho_0}\left[\mathbb{E}_{a\sim \pi_{new}(a\mid s)}\left[Q_{\pi_{new}}(s, a)-Q_{\pi_{old}}(s, a)\right]\right] = \gamma \mathbb{E}_{s\sim d^1_{\pi_{new}}}\left[V_{\pi_{new}}(s)-V_{\pi_{old}}(s)\right]$$

where $d^1_{\pi_{new}}$ denotes the probability distribution of state in time step $t=1$ with policy $\pi_{new}$. Repeating the above operation, we will obtain:

$$\begin{aligned} \mathcal{J}(\pi_{new}) - \mathcal{J}(\pi_{old}) &= \sum_{t=0}^\infty \gamma^t\mathbb{E}_{s,a \sim d_{\pi_{new}}^t, \pi_{new}}\left[A_{\pi_{old}}(s,a)\right] \\ &= \frac{1}{1-\gamma} \mathbb{E}_{s \sim d_{\pi_{new}}}\left[\mathbb{E}_{a\sim \pi_{new}(\cdot\mid s)}\left[A_{\pi_{old}}(s,a)\right]\right] \\ &\approx \frac{1}{1-\gamma} \mathbb{E}_{s \sim d_{\pi_{new}}}\Bigg[(1-p_{data}(s)A_{\pi_{old}}(s,a^\star)\eta)\mathbb{E}_{a\sim \pi_{old}(\cdot\mid s)}\left[A_{\pi_{old}}(s,a)\right] \\& \quad \quad + p_{data}(s)A^2_{\pi_{old}}(s,a^\star) \eta \frac{\mathbb{I}_{a\in \mathcal{N}(a^\star\mid s, \epsilon)}(a)}{S_{\mathcal{N}(a^\star\mid s, \epsilon)}}\Bigg] \\ & = \frac{1}{1-\gamma} \mathbb{E}_{s \sim d_{\pi_{new}}}\left[p_{data}(s)A^2_{\pi_{old}}(s, a^\star) \eta \frac{\mathbb{I}_{a\in \mathcal{N}(a^\star\mid s, \epsilon)}(a)}{S_{\mathcal{N}(a^\star\mid s, \epsilon)}}\right] \ge 0 \end{aligned}$$

Q2: Have the authors considered tuning rewards (for example, adding a fixed constant to all rewards to make all rewards positive) to guarantee Q values are all positive?

A2: Thank you for your question. Yes, we did conduct such a test by adding a fixed constant to all rewards at the very beginning. However, it does not work well after the policy has been improved to a certain extent. This is because the relative difference of Q value will be reduced with this operation. For example, if we have $Q(s,a_1)=1, Q(s,a_2)=2$, and $Q^{'}(s,a_1)=11, Q^{'}(s,a_2)=12$ after adding a fixed constant $10$, the raw relative difference is $\frac{2-1}{2}=0.5$ and the modified relative difference is $\frac{12-11}{12}=0.083$. Besides, in many real applications, how to set this fixed constant to ensure all $Q$ values positive is hard.

Q3: There is a typo in line 208 where the right parenthesis should be "]".

A3: Thank you for pointing out this. We will correct this typo in the final version.

Q4: For Eq(9), how do you estimate $\max_a Q(s,a)$ in practice?

A4: As mentioned in lines 219-220, we can obtain the approximate $\max_a Q(s,a)$ via enough samples from the diffusion. In practice, we choose around 64 samples.

Q5: Even though critic training is not the main point of the paper, I suggest adding a section in the Appendix to explicitly describe it for completeness.

A5: Thank you for your valuable suggestion. The critic training part is standard and the same as TD3 and DiPO. We will add the critic training part to the appendix in the final version.

Q6: In Figure 5, some $K_b, K_t$ combinations may lead to extreme training collapse. Did the authors experiment with higher values of $K_b, K_t$ while maintaining $K_t<K_b$?

A6: Thank you for your constructive suggestion. Firstly, we need to clarify that the fluctuation of the reward curve is a common phenomenon in RL and cannot be viewed as a training collapse, since it recovers quickly to the previous level. In Figure 5, we only plot the result of one run for each case without a window smoothing. That is why the reward curves look unstable.

We have conducted experiments with higher $K_b$ and have presented the results in the attached PDF file. It can be found that a high value of $K_b$ will limit the exploration of QVPO and tend to be stuck at a local optimum.

We thank the reviewer for his/her careful reading and valuable suggestions. Below we will answer your concerns point-by-point.

Q1: The experiments are limited to D4RL datasets, which do not cover the full spectrum of possible RL environments. Can the proposed algorithm be adapted for tasks with discrete action spaces or sparse rewards, and if so, how?

A1: Thanks for your question. We need to clarify two fundamental points:

(1) We never use D4RL datasets. QVPO is an online RL method instead of an offline one. Only offline RL requires datasets such as D4RL. We perform our experiments in 5 mujoco online environments rather than D4RL datasets. These mujoco tasks are quite standard benchmarks in existing online RL methods. Besides, we also added the experiments on the recently released HumanoidBench environment in the attached pdf file. It is a very complex environment and QVPO still achieves the SOTA performance.

(2) QVPO is designed to handle complex continuous control tasks rather than for discrete action spaces or sparse rewards. For discrete action spaces, we can directly sample in the discrete probability simplex output by the neural network and achieve multimodal policy. However, it is not trivial to achieve multimodal policy in continuous action distribution, since we cannot output the corresponding probability density function via neural networks. In that case, we consider utilizing diffusion model to enhance the multimodality and exploration ability of RL agent in continuous control tasks. Besides, QVPO does not conflict with most existing RL methods for sparse rewards like HER [R5]. Thus, we can incorporate HER into QVPO in environments with sparse rewards.

Q2: While the paper includes ablation studies, they are not exhaustive. Additional ablation studies are needed to isolate the contributions of individual components, such as the Q-weighted variational loss and the entropy regularization term. It is unclear how much each component independently contributes to the overall performance improvements.

A2: Thanks for raising the concern. We need to clarify that the main body of QVPO is the Q-weighted variational loss and QVPO must use this loss to train the diffusion policy. Thus, there is no way to do ablation study on the Q-weighted variational loss. Besides, the ablation study on the entropy regularization term is shown in Figure 4 of the original paper.

Q3: The proposed method for reducing diffusion policy variance via action selection aims to improve sample efficiency. However, this approach may sacrifice the exploration capabilities of the diffusion policy, thus limiting its multimodal potential. This trade-off is not discussed in the paper, and no ablation study is provided to evaluate the impact.

A3: We did conduct such ablation study with the trade-off analysis in our paper. The action selection number $K_b$ in efficient diffusion policy is actually the trade-off between exploration and exploitation. As shown in Figure 5 and mentioned in lines 320-330, the diffusion policy converges slowly when $K_b=1$ (i.e., without efficient diffusion policy). This means QVPO without efficient diffusion policy cannot exploit the information of Q value very well. Besides, we also add the experiment with $K_b=20$ to show that too high action selection number will lead to a limited exploration. In that case, setting $K_b=4$ can balance exploration and exploitation well.

Q4: How is the diffusion policy initialized, and what impact does this initialization have on the performance of the algorithm? If the policy is trained from scratch, how is the Q-function trained at the beginning stage?

A4: Thanks for your question. The diffusion policy and Q-function are both trained from scratch without any special initilization trick.

Q5: The efficient diffusion policy method reduces policy variance but may limit exploration. How do you balance exploration and exploitation, and can you provide an ablation study to verify the effects of this trade-off?

A5: The answer can be referred to in A3.

Q6: From Figure 3, QVPO appears to converge faster. Can you provide an explanation or experiments to support this observation?

A6: Thanks for your question. Compared with DIPO and QSM, QVPO can explore broader action space and achieve multimodal training samples with the sampling-based update and entropy regularization term. We think that is why QVPO converges faster.

Q7: How does the computational complexity of QVPO compare with existing diffusion-based online RL methods in terms of training time and resource utilization?

A7: The training and inference time of QVPO and other online RL methods are shown in Global Rebuttal. The GPU resource utilization rate of QVPO/DIPO is 55%/28% and memory consumption of QVPO/DIPO is 598MB/1354MB. In that case, QVPO can sufficiently make use of the parallel computing ability of GPU.

[R5] Andrychowicz M, Wolski F, Ray A, et al. Hindsight experience replay[J]. Advances in neural information processing systems, 2017, 30.

Q1: You make claims about DIPO and QSM that don't seem to be proven anywhere. You really should compare to QSM.

A1(DIPO): We did compare QVPO with DIPO in Figure 4 (original paper). Besides, as mentioned in lines 110-114, DIPO creates a dedicated buffer for diffusion policy, updates the state-action pairs in this buffer via the gradient of Q network, and finally uses the state-action pair from this buffer to train the diffusion policy. In that case, the training samples of DIPO will be limited to the vicinity of previously explored actions, thus leading to limited exploration.

A1(QSM): The code of QSM was not yet released when this paper was submitted. So we did not compare QVPO with QSM before. We have added the QSM results of 5 Mujoco Benchmarks in Figure 2 of the attached pdf file, following the recently released official implementation. We can see the performance of QSM is much worse than QVPO and DIPO, and even worse than SAC and TD3 in the Walker and Humanoid environments, which is also shown in Figure 3 of the QSM original paper. As we mentioned in lines 43-46, QSM has a doubled approximation error from the alignment process. The trained Q value function has an approximation error. QSM utilizes the gradient of this approximated Q function to train the score model.

Q2: You take credit ... "efficient diffusion policy" trick is used extensively in other work and the weighted VLO loss is one of the methods presented in EDP (Kang et al.) ... This should be rewritten and these works should be cited. ... This paper still has reasonable novelty but that similarity should be better acknowledged.

A2: We will add more illustrations and cite these works properly in the final version. Note that these works are all offline RL. They used "efficient diffusion policy" to yield deterministic policy during inference, while our QVPO utilizes "efficient diffusion policy" in the online interaction with the environment to obtain more efficient transition data. Our motivation is different from these existing works.

Besides, the weighted VLO loss of QVPO is different from EDP. In the weighted VLO loss, QVPO trains diffusion policy with the weighted sample from the current policy, while EDP trains the diffusion policy with the weighted sample from policy in the offline dataset (replay buffer).

Q3-1: The comparison to other RL algorithms is unfair. Need to either tune all relevant hyperparameters for other algorithms or fix yours.

A3-1: Most hyperparameters in our experiments are the same, except $K_t$ in Hopper and $w_{ent}$ in HalfCheetah, as shown in Table 3. During the rebuttal, we have fixed all hyperparameters for QVPO and the new results are shown in the attached pdf file for 5 Mujoco benchmarks. QVPO shows similar performance and still outperforms all baselines with a clear margin.

Q3-2: Is $\omega_{ent}$ difficult to tune?

A3-2: No, quite easy to set this value. Just make the scaling within 1e-3 to 1e-2.

Q4: Claim diffusion policy is a good fit for RL due to its ability to model multimodal action distributions, but don't conduct experiments on such environment.

A4: Firstly, existing online RL methods almost all verify their performance on these robot locomotion tasks. In that case, we can compare QVPO with previous methods in a relatively fair manner with these standard environments.

Second, the multimodality of policy is important because it can avoid the policy stuck at the local optimum. Moreover, our tested tasks such as Humanoid-v3 cannot be solved very well with a MLP policy.

Third, to further test QVPO for more complex tasks, we applied QVPO to the recently released HumanoidBench environment. The benchmark is based on Unitree H1 humanoid robot with 151 state dimension and 61 action dimension. Please refer to Global Response for more details.

Q5(1): The ablation studies are a bit incomplete in my opinion. It would be good to see results from all five environments.

A5(1): Due to the time and resource limit, we cannot complete these additional experiments. We will add them in our final version. Actually, the results for other 4 environments are similar but we did not record them.

Q5(2): It would be good to run an ablation where you omit the policy loss weighting ... I assume this ablation will go well ... $K_b=1$ still learns a policy but this will strengthen the paper.

A5(2): The requested ablation study $K_b=1$ is shown in Figure 5 of the original paper. Offline diffusion RL methods like EDP train the diffusion policy with samples from offline datasets or replay buffer, while QVPO trains the diffusion policy with selected samples from the diffusion policy itself.

Q6: You mention qcut. Either show an experiment proving that or remove it.

A6: In the beginning, qcut is mentioned to present the original motivation of QVPO. We will remove qcut part in the final version.

Q7: The training and inference time of QVPO.

A7: Please refer to the Global Response.

Q8: The second and third sentence of the abstract seem to contradict one another.

A8: Thanks and we will modify the two sentences.

Q9: I'm very confused by Figure 2. How did you generate this? What does 'explorable area' mean?

A9: Figure 2 is exported by powerpoint to visualize the effect of diffusion entropy regularization term. The 'explorable area' means the area with a certain generation probability of diffusion policy.

Q10: Figure 4 is a nice result but it would be helpful to also see steps=100 & ent.

Q10: The entropy term is introduced for the diffusion model with few denoising steps. QVPO with steps=100 with entropy term will not improve the result. We will add this result in the final version.

Q11: Suggestions to improve the presentation (some minor format issues) ...

A11: We will modify them.

Thanks for your valuable suggestions as well. We are glad that our responses help you solve most of your concerns. Here are some further illustrations for your concerns.

Why results of DSAC-T are so much stronger than ours?

We want to clarify that the difference in performance between DSAC-T [R8] and QVPO merely comes from different implementation settings. As shown in Figure 2 of DSAC-T [R8], the x-axis is not the online interaction steps as in our paper (i.e., training epoch) but the training iteration. As shown in Table 3 of DSAC-T [R8], DSAC-T performs one training iteration with 20 online interaction steps (sample batch size in original paper) for off-policy RL methods. In that case, to fairly compare QVPO with DSAC-T, you need to multiply all the coordinates on the x-axis by 20. With that operation, you will obtain results similar to ours on SAC. Actually, it is meaningful to use interaction steps (ours) rather than training iterations (DSAC-T) as the x-axis, since the former can show sample efficiency of different algorithms better in the online RL paradigm.

We want to highlight that our setting is the same as the most existing works, including SAC [R8], DIPO [R9], QSM [R10], TD-MPC2 [R1], Dreamer-v3 [R2], etc. Our results on all other baselines including SAC are also consistent with these works. Please refer to Figure 1 of SAC [R8], Figure 6 of DIPO [R9], Figure 3 of QSM [R10], Figure 12 of TD-MPC2 [R1], etc.

Why PPO performs poorly on Humanoid and HalfCheetah? PPO does perform poorly with limited interactions since it is an on-policy RL method. Many previous works such as DIPO [R7], SAC [R9], Dreamer-v3 [R2] also show similar results to ours. Please refer to Figure 6 in DIPO [R7], Figure 1 in SAC [R9], and Figure 14 in Dreamer-v3 [R9].

[R8] Duan J, Wang W, Xiao L, et al. DSAC-T: Distributional soft actor-critic with three refinements[J]. arXiv preprint arXiv:2310.05858, 2023.

[R9] Haarnoja T, Zhou A, Hartikainen K, et al. Soft actor-critic algorithms and applications[J]. arXiv preprint arXiv:1812.05905, 2018.

[R10] Psenka M, Escontrela A, Abbeel P, et al. Learning a diffusion model policy from rewards via q-score matching[J]. arXiv preprint arXiv:2312.11752, 2023.