Maximum Entropy Reinforcement Learning with Diffusion Policy

Thank you for your valuable feedback and suggestions. Here, we aim to address the questions raised in the review.

Q1: Include results on tasks like MetaWorld or DMControl.

We compare MaxEntDP with SAC on DMControl and MyoSuite benchmarks. The results are shown in https://anonymous.4open.science/api/repo/pics-E459/file/DMC_myo.pdf?v=538c572b.

Q2: The importance-weighted sampling technique in the paper is similar to iDEM.

We would like to emphasize that although the expression of our QNE method appears similar to iDEM [1], our method exhibits significantly lower estimation variance and steady performance improvement throughout the training process over the same method replaced with iDEM (Figure 3). In addition, our method does not require gradient computation of the Q-function as in iDEM and thus is more computationally efficient. These advantages demonstrate the superiority of our QNE method on diffusion policy optimization.

Q3: Concern on the assumption of Corollary 3.4.

We use the target Q network to train the diffusion policy. Since we adopt the Exponential Moving Average (EMA) update to smooth the change of the target Q network, it is not difficult to learn a well-trained policy network.

Q4: Keep the hyper-parameter tuning effort the same across baseline algorithms.

For fair comparison, we unify shared hyperparameters (batch size, discount factor, the depth and width of hidden layers, learning rate and replay buffer size) across baseline algorithms. Other hyperparameters strictly follow the settings specified in the original paper or codebase, which are already tuned by their authors.

Q5: One related work model-based diffusion [2] is missing.

We will cite it in revised versions of the paper. The model-based diffusion proposes the Monte Carlo estimation for computing the score function and uses the Monte Carlo score ascent to generate samples following the Boltzmann distribution of a given function. This method is similar to our QNE method, however, QNE has several properties which matter in RL training:

1. We use a parameterized network to approximate the scaled score function, while the model-based diffusion needs to compute the score function using Monte Carlo estimation when generating samples. Therefore, sample generation of diffusion-based diffusion is time-consuming, which will slow the training speed of the RL algorithms.
2. We adopt the ancestral sampling in DDPM to generate samples, that are more diverse than that of Monte Carlo score ascent used in model-based diffusion.
3. We propose to modify the standard Gaussian to the truncated Gaussian in QNE to model the action distribution with a bounded action space. However, model-based diffusion can not address such a bounded distribution.

Thank you for your thoughtful review and constructive feedback. Below, we will address your comments and hope that this clarifies the context of our work.

Q1: Consider more sophisticated RL tasks.

We compare MaxEntDP with SAC on DMContorl and MyoSuite benchmarks. The results are shown in https://anonymous.4open.science/api/repo/pics-E459/file/DMC_myo.pdf?v=538c572b.

Q2: Justification of Eq 59.

We are sorry for omitting the parentheses on $\bigtriangledown_{a_0} w(a_0)$ in Eq 59, which causes trouble in understanding. In fact, we apply the integration by parts formula (line 777) to derive Eq 59, i.e., $\int_{\Omega} v \nabla u \, d\Omega = \int_{\Gamma} u v n \, d\Gamma - \int_{\Omega} u \nabla v \, d\Omega$ (consider $w(a_0)$ as $v$ and $N(a_0)$ as $u$; the term $\int_{\Gamma} u v n \, d\Gamma$ is zero since $w(a_0)$ and $N(a_0)$ decay rapidly at infinity).

Q3: Justification of Eq 65.

In Eq 65, we change the integration variable from $\alpha_t$ to $\sigma(\alpha_t)$ using the equation $\alpha_t=\log \frac{\sigma(\alpha_t)}{1-\sigma(\alpha_t)}$, as illustrated in line 807. Then the integration domain of Eq 65 has been changed to the varying range of the new integration variable $\sigma(\alpha_t)$, which is $(0,1)$.

Q4: The accuracy of the log-likelihood approximation.

Due to the length limit of rebuttal, please refer to our response to the Q2 of reviewer fZmT.

Q5: The paper lacks recent work such as BRO and CrossQ.

We will cite them in future versions. Briefly, the two methods propose some improvements to the SAC algorithm. BRO develops the advanced BroNet architecture, regularization, and optimistic upper-bound Q-value approximation. CrossQ removes the target networks and adopts Batch Normalization for high sample efficiency. Since the improvements of the two methods are also compatible with our MaxEntDP algorithm, it is interesting to see how the performance of MaxEntDP can be enhanced by combining these improvements.

Q6: Why can the noisy actions $p(a_t)$ be substituted with other distributions?

As shown in Eq17, the training target $\epsilon^{*}$ of the noise prediction network depends only on its input $(a_t,\alpha_t)$, which means that we can minimize the L2 loss for each input point. Thus, the noisy action samples from other distributions can also be used for training. This can be seen as a kind of "off-policy" training. Moreover, in the paper we diffuse the action samples in the replay buffer to obtain the surrogate distribution of $p(a_t)$ (line 295). As the training progresses, the actions generated by the diffusion policy will be closer to the target distribution $p(a_0)$, then the diffused distribution of these actions will also be closer to the true distribution of $p(a_t)$, consequently, the training will be more and more "on-policy".

Q7: The reason for using a truncated standard Gaussian in RL tasks with bounded action spaces.

We would like to point that MaxEntDP and SAC model the bouded policy distribution in different ways. SAC outputs an unbounded distribution and transforms it to a bounded one by applying a tanh function, while MaxEntDP directly models a bounded distribution following the common practices in the image generation domain. To learn this bounded distribution, we need to change the target distribution to $p(a_0)\propto\exp(\frac{1}{\beta}Q(a_0))I_{a_0\in[lb,ub]}$, where $I_{a_0\in [lb,ub]}$ is an indicator function to check if $a_0$ lies within the bound. Similar to lemma 3.1, we can obtain that the reverse transition distribution has be changed to $p(a_0|a_t)\propto\exp(\frac{1}{\beta}Q(a_0)) I_{a_0\in[lb,ub]} \mathcal{N}(a_0|\frac{1}{\sqrt{\sigma(\alpha_t)}} a_t,\frac{\sigma(-\alpha_t)}{\sigma(\alpha_t)}I)$. Taking the indicator function and the Gaussian distribution together, we can derive a truncated Gaussian distribution with a bound of $[lb,ub]$. Then, the conditional distribution $p(a_0|a_t)$ can be seen as a truncated Gaussian distribution of $a_0$ weighted by the exponential of the Q-function. Therefore, for diffusion policy training with a bounded action space, we only need to modify the standard Gaussian distribution in the QNE method to a truncated Gaussian distribution.

Q8: Performance comparison with other Diffusion-Based methods.

We would like to emphasize that none of the competing diffusion-based methods performs consistently well on all tasks (DACER struggles on Walker2d, and others underperform on HalfCheetah). However, our MaxEntDP displays consistent sample efficiency and stability across all tasks.

Q9: The meaning of "training step" in the learning curves; K-L in line 119.

We use "training step" because both DACER and QVPO use it as the x-axis of the learning curve, which means the training steps of the actor/critic network. Since we use a UTD of 1, it is the same as the number of environment interactions. And we will modify the 'K-L' in line 119 in revised versions.

Thank you for your valuable feedback and suggestions. Below we will address each concern raised in the review.

Q1: Can MaxEntDP learn different solutions in high-dimensional multi-goal tasks from DDiffPG? And provide evidence for improved exploration.

We tested MaxEntDP and SAC on four versions of the AntMaze environment in DDiffPG (using dense rewards) and visualized the generated trajectories in Figure 1 of https://anonymous.4open.science/api/repo/pics-E459/file/antmaze.pdf?v=eebbbe1f. The results confirm that MaxEntDP can learn diverse behavior modes even in challenging high-dimensional RL tasks, while SAC fails to learn different solutions. In addition, we visualized state coverage for MaxEntDP and SAC (see Figure 2 in the above link). The results show that MaxEntDP explores multiple behavior modes and exhibits broader state coverage. This highlights the advantage of using a diffusion policy for efficient exploration.

Q2: No evidence supporting numerical integration as an effective approximation.

According to the Law of Large Numbers, numerical integration accuracy improves with larger diffusion steps $T$ and sample numbers $N$. To exhibit the accuracy of different $T$ and $N$, we conducted experiments on a 2D toy example (a mixture of four Gaussians) and presented the results in https://anonymous.4open.science/api/repo/pics-E459/file/probability.pdf?v=86f81805. As shown in the figure, our setting in the paper ($T=20, N=50$) provides an effective probability approximation. Moreover, when fewer samples ($T=20, N=10$) are used, despite some estimation errors, our method still assigns higher values to high-probability regions, which can be considered as an intrinsic curiosity reward to promote exploration on the action region with low policy probability.

Q3: The claim that MaxEntDP achieves the optimal MaxEnt policy with sufficient model capacity.

We do not make such a claim in the paper. Instead, we argue that incorporating a diffusion policy into MaxEnt RL improves exploration and moves the policy closer to the optimal solution, as proposed in the abstract. The 2D toy example in the paper reveals that MaxEntDP can explore the state-action space efficiently and finally learn a multimodal policy. The Mujoco experiments show performance improvements over other generative models and diffusion-based approaches, supporting our main claim.

Q4: Concern about Monte Carlo sampling.

Your concern is reasonable since a higher estimation accuracy requires more Monte Carlo samples. However, in our experiments, we find that a small sample number of 1000 (for both diffusion policy optimization and action probability estimation) is enough to obtain a good performance. Additionally, as the samples can be processed parallelly and GPU throughput continues to improve, Monte Carlo sampling is not a bottleneck.

Q5: Necessity of action selection and its application to all baselines.

Action selection is crucial because our method approximates the exponential of the Q-function, generating both high- and low-return actions. While beneficial for exploration, this can reduce test-time performance. To mitigate this, we apply action selection (only in testing) to pick the action with the highest Q-value. Similar techniques are used in SAC (take the Gaussian mean) and EBFlow. To ensure fairness, we applied action selection to all diffusion-based baselines (with a candidate number of 10) and reported results in https://anonymous.4open.science/api/repo/pics-E459/file/comparison_diffusion_selection.pdf?v=a952c896. Our MaxEntDP continued to demonstrate high sample efficiency and stability across all tasks.

Q6: Add experiments in more challenging environments.

We compared MaxEntDP with SAC on the DMControl and MyoSuite benchmarks. Results are shown in https://anonymous.4open.science/api/repo/pics-E459/file/DMC_myo.pdf?v=538c572b.

Q7: Include a comparison with BRO.

BRO improves SAC with advanced network architectures, regularization, and optimistic Q-value estimation. Since BRO's enhancements can be applied to any baseline (non-diffusion and diffusion-based) and MaxEntDP, a direct comparison would be unfair. However, integrating BRO's improvements into MaxEntDP is an interesting future direction.

Q8: Performance gains of MaxEntDP over existing diffusion-based baselines.

None of the competing diffusion-based methods consistently outperforms across all tasks—e.g., DACER struggles on Walker2d, and others underperform on HalfCheetah. In contrast, MaxEntDP exhibits consistent sample efficiency and stability across all tasks.

Q9: Missing reference to DDiffPG.

Thanks for your valuable suggestion. We will cite it in future revisions. While DDiffPG explicitly distinguishes different behavior modes and learns a Q-function for each mode, MaxEntDP employs a single Q-function, making it simpler and more computationally efficient.