Title: Entropy-regularized Diffusion Policy with Q-Ensembles for Offline Reinforcement Learning
Venue: NeurIPS 2024
Accepted: yes
Keywords: Offline Reinforcement Learning; Diffusion Model; Stochastic Differential Equation
Abstract: Diffusion policy has shown a strong ability to express complex action distributions in offline reinforcement learning (RL). However, it suffers from overestimating Q-value functions on out-of-distribution (OOD) data points due to the offline dataset limitation. To address it, this paper proposes a novel entropy-regularized diffusion policy and takes into account the confidence of the Q-value prediction with Q-ensembles. At the core of our diffusion policy is a mean-reverting stochastic differential equation (SDE) that transfers the action distribution into a standard Gaussian form and then samples actions conditioned on the environment state with a corresponding reverse-time process. We show that the entropy of such a policy is tractable and that can be used to increase the exploration of OOD samples in offline RL training. Moreover, we propose using the lower confidence bound of Q-ensembles for pessimistic Q-value function estimation. The proposed approach demonstrates state-of-the-art performance across a range of tasks in the D4RL benchmarks, significantly improving upon existing diffusion-based policies. The code is available at https://github.com/ruoqizzz/entropy-offlineRL.

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

Abstract

Diffusion policy has shown a strong ability to express complex action distributions in offline reinforcement learning (RL). However, it suffers from overestimating Q-value functions on out-of-distribution (OOD) data points due to the offline dataset limitation. To address it, this paper proposes a novel entropy-regularized diffusion policy and takes into account the confidence of the Q-value prediction with Q-ensembles. At the core of our diffusion policy is a mean-reverting stochastic differential equation (SDE) that transfers the action distribution into a standard Gaussian form and then samples actions conditioned on the environment state with a corresponding reverse-time process. We show that the entropy of such a policy is tractable and that can be used to increase the exploration of OOD samples in offline RL training. Moreover, we propose using the lower confidence bound of Q-ensembles for pessimistic Q-value function estimation. The proposed approach demonstrates state-of-the-art performance across a range of tasks in the D4RL benchmarks, significantly improving upon existing diffusion-based policies. The code is available at https://github.com/ruoqizzz/entropy-offlineRL.

1 Introduction

Offline reinforcement learning (RL), also known as batch RL [29] focuses on learning optimal policies from a previously collected dataset without further active interactions with the environment [31]. Although offline RL offers a promising avenue for deploying RL in real-world settings where online exploration is infeasible, a key challenge lies in deriving effective policies from fixed datasets, which usually are diversified and sub-optimal. The direct application of standard policy improvement approaches is hindered by the distribution shift problem [10]. Previous works mainly address this issue by either regularizing the learned policy close to the behavior policy [10, 9] or by making conservative updates for Q-networks [28, 26].

Diffusion models have rapidly become a prominent class of highly expressive policies in offline RL [44, 49]. While this expressiveness is beneficial when modeling complex behaviors, it also means that the model has a higher capacity to overfit the noise or specific idiosyncrasies in the training data. To address this, existing work introduce Q-learning guidance and regard the diffusion loss as a special regularizer adding to the policy improvement process [44, 17, 22]. Such a framework has achieved impressive results on offline RL tasks. However, its performance is limited by pre-collected datasets (or behavior policies) and the learning suffers severe overestimation of Q-value functions on unseen state-action samples [31].

One promising approach is to increase exploration for out-of-distribution (OOD) actions, with the hope that the RL agent can be more robust to diverse Q-values and estimation errors [50]. Previous online RL algorithms achieve this by maximizing the entropy of pre-defined tractable policies such

as Gaussians [35, 13, 15]. Unfortunately, directly computing the log probability of a diffusion policy is almost impossible since its generative process is a stochastic denoising sequence. Moreover, it is worth noting that entropy is seldom used in offline settings because it may lead to a distributional shift issue which may cause overestimation of Q-values on unseen actions in the offline dataset.

Another line of work addresses the overestimation problem by enforcing the Q-values to be more pessimistic [28, 20]. Inspired by this, uncertainty-driven RL algorithms employ an ensemble of Q-networks to provide different Q-value predictions for the same state-action pairs [1, 3]. The variation in these predictions serves as a measure of uncertainty. For state-action pairs exhibiting high predictive variance (e.g., OOD data points), these methods preferentially adopt pessimistic Q-value estimations as policy guidance.

In this paper, we present an entropy-regularized diffusion policy with Q-ensembles for offline RL. At the core of our method is a mean-reverting stochastic differential equation (SDE) [32] which allows us to sample actions from standard Gaussian conditioned on the environment state. We show that such an SDE provides a tractable entropy regularization that can be added in training to increase the exploration of OOD data points. In addition, we approximate the lower confidence bound (LCB) of Q-ensembles to alleviate potential distributional shifts, thereby learning a pessimistic policy to handle high uncertainty scenarios from offline datasets. As illustrated in Figure 1, both entropy regularization and Q-ensembles can improve RL performance on unbalanced offline datasets. The LCB approach further reduces the variance between different trials and provides a better estimation of unseen state-action pairs.

Our model achieves highly competitive performance across a range of offline D4RL benchmark tasks [8] and, in particular, significantly outperforms other diffusion-based approaches in the Antmaze environment. The superior performance demonstrates the effectiveness of the entropy-regularization and Q-ensembles. Overall, the proposed method encourages policy diversity and cautious decisionmaking, enhancing exploration while grounding the policy in the confidence of its value estimates derived from the offline dataset.

2 Background

This section reviews the core concepts of offline RL and then introduces the mean-reverting SDEs and shows how we sample actions from its reverse-time process. Note that there are two types of timesteps for RL and SDE. To clarify that, we use i ∈ { 0 , . . . , N } to denote the RL trajectories' step and t ∈ { 0 , . . . , T } to index diffusion discrete times.

Offline RL. We consider learning a Markov decision process (MDP) defined as M = {S , A , P, R, γ, d 0 } , where S and A are the state and action spaces, respectively. The state transition probability is denoted P ( s i +1 | s i , a i ) and R : S × A → R represents a reward function,

γ ∈ (0 , 1] is the discount factor, and d 0 is the initial state distribution. The goal of RL is to maximize the cumulative discounted reward ∑ ∞ i =0 γ i E a i ∼ π ( s i ) [ r ( s i , a i ) ] with a learned policy π . In contrast to online RL which requires continuous interactions with the environment, offline RL directly learns the policy from the static dataset D = { ( s i , a i , r i , s i +1 ) } N D i =1 . In the offline setting, two primary challenges are frequently encountered: over-conservatism and a limited capacity to effectively utilize diversified datasets [31]. To address the issue of limited capacity, diffusion models have recently been employed to learn complex behavior policies from datasets [44].

Mean-Reverting SDE. Assume that we have a random variable a 0 sampled from an unknown distribution p 0 ( a ) . The mean-reverting SDE [32] is a diffusion process { a t } t ∈ [0 ,T ] that gradually injects noise to a 0 :

where w is the standard Wiener process, θ t and σ t are predefined positive parameters that characterize the speed of mean reversion and the stochastic volatility, respectively. Compared to IR-SDE [32], we set the mean to 0 to let the process drift to pure noise to fit the RL environment. The mean can however be tuned to high-reward actions in the offline dataset or prior knowledge. By setting σ 2 t = 2 θ t for all diffusion steps, the solution to the forward SDE ( τ < t ) is given by

where ¯ θ τ : t := ∫ t τ θ z d z are known coefficients [32]. In the limit t →∞ , the marginal distribution p t ( a ) = p ( a t | a 0 ) converges to a standard Gaussian N (0 , I ) . This gives the forward process its informative name, i.e. ' mean-reverting '. Then, Anderson [2] states that we can generate new samples from Gaussian noises by reversing the SDE (1) as

where a T ∼ N (0 , I ) and ¯ w is the reverse-time Wiener process. This reverse-time SDE provides a strong ability to fit complex distributions, such as the policy distribution represented in the dataset D . Moreover, the ground truth score ∇ a log p t ( a ) is acquirable in training. We can thus combine it with the reparameterization trick

to train a time-dependent neural network ϵ ϕ using noise matching on randomly sampled timesteps:

where ϵ t ∼ N (0 , I ) is a Gaussian noise and { a t } T t =0 denotes the discretization of the diffusion process. See Appendix A.1 for more details about the solution, reverse process, and loss function.

Sample Actions with SDE. Most existing RL algorithms employ unimodal Gaussian policies with learned mean and variance. However, this approach encounters a challenge when applied to offline datasets, which are typically collected by a mixture of policies and therefore hard to represent by a simple Gaussian model. Thus we prefer to represent the policy with an expressive model such as the reverse-time SDE. More specifically, the forward SDE provides theoretical guidance to train the neural network, and the reverse-time SDE (3) generates actions from Gaussian noise conditioned on the current environment state as a typical score-based generative process [43].

3 Method

We present the three core components of our method: 1) an efficient sampling strategy based on the mean-reverting SDE; 2) an entropy regularization term that enhances action space exploration; and 3) a pessimistic evaluation with Q-ensembles that avoids overestimation of unseen actions.

3.1 Optimal Sampling with Mean-Reverting SDE

We have shown how to sample actions with reverse-time SDEs in Section 2. However, generating data from the standard mean-reverting SDE [32] requires many diffusion steps and is sensitive to the noise scheduler [36]. To improve sample efficiency, we propose generating actions from the posterior distribution p ( a t -1 | a t ) conditioned on a 0 . This approach ensures fast convergence of the generative process while preserving its stochasticity.

Proposition 3.1. Given an initial variable a 0 , for any diffusion state a t at time t ∈ [1 , T ] , the posterior of the mean-reverting SDE (1) conditioned on a 0 is

which is a Gaussian with mean and variance given by:

where θ ′ i := ∫ i i -1 θ t dt and ¯ θ t is to substitute ¯ θ 0: t for clear notation.

The proof is provided in Appendix A.2. Moreover, thanks to the reparameterization trick [25], we can approximate the variable a 0 by reformulating Eq. (4) to

where ϵ ϕ is the learned noise prediction network. Then we combine Eq. (8) with Eq. (6) to iteratively construct the sampling process. In addition, it can be proved that the distribution mean

˜ µ t ( a t , a 0 ) is the optimal reverse path from a t to a t -1 (see Appendix A.3). Here, we illustrate a simple example of data reconstruction with different diffusion steps in Figure 2. Additional results are shown are provided in Appendix B.1. The proposed optimal sampling process from Proposition 3.1 requires only 5 steps, versus over 30 for the standard process. It clearly shows that the proposed optimal sampling is more efficient than the standard reverse-time SDE process.

Notation Note Recall that we have

two distinct types of timesteps for RL and SDE denoted by i and t , respectively. To clarify the notation, in the following sections, we use a t i to represent the intermediate variable of an action taken at RL trajectory step i with SDE timestep t , as a t i = a t at state s i . Therefore, the action to take for state s i is the final sampled action a i denoted by a 0 i . Hence, the policy is given by

While we cannot sample directly from this distribution we can efficiently sample the SDE's reverse joint distribution as

where p ( a T ) = N (0 , I ) is Gaussian noise and the generative process is conditioned on the environment state s i . So to take an action from π ϕ ( a 0 i | s i ) , we sample from the joint distribution using Eq. (6) and Eq. (8) and finally pick out a 0 as our sampled action. The visualization of is out method is provided in Appendix A.5.

3.2 Diffusion Policy with Entropy Regularization

The simplest strategy of learning a diffusion policy is to inject Q-value function guidance to the noise matching loss (5), in the hope that the reverse-time SDE (3) would learn to sample actions with higher values. This can be easily achieved by minimizing the following objective:

where Q ψ is the state-action value function approximated by a neural network, see Section 3.3.

Data Distribution

Reconstruction (N=5)

Reconstruction (N=10)

Reconstruction (N=30)

Reconstruction (N=50)

Reverse-time SDE

Optimal Sampling

This combination regards diffusion loss as a behavior-cloning term that learns the overall action distribution from offline datasets. However, the training is limited to existing data samples and the Q-learning term is sensitive to unseen actions. To address it, we propose to add an additional entropy term H = E s i ∼D [ -log π ϕ ( · | s i )] to increase the exploration of the action space during training and rewrite the policy loss (11) to

where α is a hyperparameter that determines the relative importance of the entropy term versus Q-values, and λ = η / E ( s,a ) ∼D [ | Q ψ ( s, a ) | ] to normalize the scale of the Q-values and balance loss terms. Iteratively generating the action a 0 i though a reverse diffusion process is computationally costly but, with an estimated noise ϵ ϕ from diffusion term (5), we can thus directly use it to approximate a 0 i based on Eq. (8) for more efficient training.

Entropy Approximation. It is worth noting that the log probability of the policy log( π ϕ ( a 0 i | s i )) is in general intractable in the diffusion process. However, we found that the log probability of the joint distribution in Eq. (10) is tractable when conditioned on the sampled action a 0 i . Proposition 3.1 further shows that the conditional posterior from a 1 i to a 0 i is Gaussian, meaning that

where C is a constant and a 1 i can be approximated using Eq. (2) similar to a 0 i . The proof is provided in Appendix A.4. Then we can focus on the conditional reverse marginal distribution p ϕ ( a 1 i | a T i , s i ) that determines the exploration of actions and is acquirable via Bayes' rule:

Since all terms in Eq. (14) can be computed with Eq. (2), we can rewrite the policy objective as

where ˆ a 0 i and ˆ a 1 i are approximate values calculated based on samples from the diffusion term. Note that the temperature α usually plays an important role in the maximum entropy RL framework and we thus provide a detailed analysis in Section 4.3.

3.3 Pessimistic Evaluation via Q-ensembles

Entropy regularization encourages diffusion policies to explore the action space, reducing the risk of overfitting pre-collected data. However, in offline RL, since the agent cannot collect new data during training, this exploration can lead to inaccuracies in value estimation for unseen state-action pairs [3, 11]. Instead of staying close to the behavior policy and being overly conservative, considering the uncertainty in the value function is an alternative approach.

In this work, we consider a pessimistic variant of a value-based method to manage the uncertainty and risks, i.e., the lower confidence bounds (LCB) with Q-ensembles. More specifically, we use an ensemble of Q-functions with independent targets to obtain an accurate LCB of Q-values. Each Q-function is updated based on its own Bellman target without sharing targets among ensemble members [11], as follows:

where ψ m , ¯ ψ m are the parameters of the Q network and Q-target network for the m th Q-function. Then, the pessimistic LCB values are derived by subtracting the standard deviation from the mean of the Q-value ensemble,

where β ≥ 0 is a hyperparameter determining the amount of pessimism, V [ Q ψ m ] is the variance of the ensembles, and m ∈ { 1 , . . . , M } where M the number of ensembles. Then, Q LCB ψ is used in the policy improvement step to balance entropy regularization and ensure robust performance. Finally, we use Q LCB ψ as the Q ψ to (15). We summarize our method in Algorithm 1.

4 Experiment

In this section, we evaluate our methods on standard D4RL offline benchmark tasks [8] and provide a detailed analysis of entropy regularization, Q-ensembles, and training stability.

4.1 Setup

Datasets We evaluate our approach on four D4RL benchmark domains: Gym, AntMaze, Adroit, and Kitchen. In Gym, we examine three robots (halfcheetah, hopper, walker2d) across sub-optimal (medium), near-optimal (medium-expert), and diverse (medium-replay) datasets. The AntMaze domain challenges a quadrupedal ant robot to navigate mazes of varying complexities. The Adroit domain focuses on highdimensional robotic hand manipulation, using datasets from human demonstrations and robot-imitated human actions. Lastly, the Kitchen domain explores different tasks within a simulated kitchen. These domains collectively provide a comprehensive framework for assessing RL algorithms across diverse scenarios.

Algorithm 1 Diffusion Policy with Q-Ensembles

Initialize parameters for π ϕ , π ¯ ϕ , { Q ψ m , Q ¯ ψ m } M m =1 . for each iteration do Sample mini-batch { ( s i , a i , r i , s i +1 ) } from D . # Ensemble-Q learning Generate a 0 i +1 ∼ π ¯ ϕ ( a i +1 | s i ) with (6) and (8). Update Q-networks { Q ψ m } M m =1 by (16). # Diffusion policy learning Sample { a t i } t ∈ [0 ,T ] from a i with (4). Predict noise and approximate a 0 i , a 1 i with (8). Update policy π ϕ by (15) using Q LCB ψ from (17). # Update target networks ¯ ϕ ← ηϕ +(1 -η ) ¯ ϕ , ¯ ψ m ← ηψ m +(1 -η ) ¯ ψ m , m ∈ { 1 , . . . , M } end for

Implementation Details Following Diffusion-QL [44], we keep the network structure the same for all tasks with three MLP layers (hidden size 256, Mish activation [34]), and train models for 2000 epochs for Gym and 1000 epochs for others. Each epoch consists of 1000 training steps with a batch size of 256. We use Adam [24] to optimize both SDE and the Q-ensembles. Each model is evaluated by 10 trajectories for Gym tasks and 100 trajectories for others. In addition, our model is trained on an A100 GPU with 40GB memory for about 8 hours per task, and results are averaged over five random seeds.

Hyperparameters We keep key hyperparameters consistent: Q-ensemble size 64, LCB coefficient β = 4 . 0 . The entropy temperature α = 0 . 01 for Gym and AntMaze tasks and automated for Adroit and Kitchen tasks. The SDE sampling step is set to T = 5 for Gym and Antmaze tasks, T = 10 for Adroid and Kitchen tasks. For 'medium' and 'large' datasets of AntMaze, we use max Q-backup following Wang et al. [44] and Kumar et al. [28]. We also introduce the maximum likelihood loss for SDE training as proposed by Luo et al. [32]. More details are in Appendix B.2.

4.2 Comparison with other Methods

We compare our method with extensive baselines for each domain to provide a thorough evaluation and to understand the contributions of different components in our approach. The most fundamental among these are the behavior cloning (BC) method, BCQ [10] and BEAR [27] which restrict the policy to dataset behavior, highlighting the need for policy regularization and exploration. We also assess against Diffusion-QL [44] which integrates a diffusion model for policy regularization guided by Q-values. This comparison isolates the benefits of our enhanced sampling process and Q-ensemble integration. Our comparison includes CQL [28] and IQL [26], known for conservative Q-value updates. Additionally, we consider EDP [21], a variant of IQL with an efficient diffusion policy, and IDQL [17], which combines IQL as a critic with behavior cloning diffusion policy reweighted by learned Q-values. These comparisons evaluate the effectiveness of integrating diffusion policies with conservative value estimation. Finally, we include MSG [11], which combines independent Q-ensembles with CQL, and DT [4], treating offline RL as a sequence-to-sequence translation problem. These baselines help assess the robustness and generalizability of our method across different approaches. The performance comparison between baselines and ours is reported in Table 1 (Gym, Adroit and Kitchen) and Table 2 (AntMaze). Detailed results are discussed below.

Gym tasks Most approaches perform well on Gym 'medium-expert' and 'medium-replay' tasks with high-quality data but drop severely on 'medium' tasks with suboptimal trajectories. DiffusionQL [44] achieves a better performance through a highly expressive diffusion policy. Our method further improves performance across all three 'medium' tasks. The results illustrate the efficacy of combining diffusion policy with entropy regularization and Q-ensembles in preventing overfitting to suboptimal behaviors. By maintaining policy stochasticity, our algorithm encourages the exploration of action spaces, potentially discovering better strategies than those in the dataset.

Adroit and Kitchen Most offline approaches cannot achieve expert performance on these tasks due to the narrowness of human demonstrations in Adroit and the indirect, multitask data in Kitchen [44]. Our method outperforms all other approaches in the Kitchen tasks which suggests its ability to 'stitching' the dataset and generalization. In addition, we fix the entropy coefficient α to be the same as other tasks for a robust setting. Even so, our method still achieves a competitive performanc in Adroit tasks. This fixed α 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. Additionally, unlike in antmaze tasks, random actions are more likely to negatively impact performance in tasks where precise control is essential like Adroit. Moreover, it's worth noting that slightly tuning α leads to a SOTA performance, as illustrated in Table 3.

AntMaze AntMaze tasks are more challenging, requiring point-to-point navigation with sparse rewards from sub-optimal trajectories [8]. As shown in Table 2, traditional behavior cloning methods (BC and DT) get 0 rewards on AntMaze medium and large environments. Our method shows excellent performance on all the tasks in AntMaze even with large complex maze settings and outperforms other methods by a margin. The result is not surprising because the entropy regularization incentivizes the policy to explore various sub-optimal trajectories within the dataset and stitch them to find a path toward the goal. In tasks with sparse rewards, this can be crucial because it prevents premature convergence to suboptimal deterministic policies. Additionally, employing the LCB of Q-ensembles effectively reduces the risk of taking low-value actions, enabling the development of robust policies.

In general, employing consistent hyperparameters for each domain, along with fixed entropy temperature α , LCB coefficient β , and ensemble size M across all tasks, our method not only achieves substantial overall performance but also outperforms prior works in the challenging AntMaze tasks. By the comparison with MSG[11] (Q-ensemble alone) and Diffusion-QL (Diffusion alone), our method further improves results demonstrating its effectiveness in handling complex environments with sparse rewards by effectively combining suboptimal trajectories to find better solutions via action space exploration.

4.3 Analysis and Discussion

We first study the core components of our method: entropy regularization and Q-ensemble. Then we show that adding both significantly improves the training robustness of diffusion-based policies.

Entropy Regularization The core idea of applying entropy regularization in offline RL is to increase the exploration of new actions such that the estimation of Q-functions is more accurate, especially for datasets with unbalanced action distribution such as the toy example in Figure 1. Here we report the results of training the diffusion policy with different entropy temperatures in Table 3. It is observed that our method with positive entropy coefficients performs better than that without the entropy term. In addition, we can extend our model with an automatic entropy adjustment

similar to the work in [16]. This approach is marked as ' auto ' in Table 3. The results show that auto-tuning the entropy temperature further improves the performance in the Adroit and Kitchen domains. Please refer to Appendix B.3 for more details.

Q-Ensembles We evaluate our method under different numbers of Q networks M ∈ { 2 , 4 , 64 } in the AntMaze environment to explore the effectiveness of Q-ensembles. The results with average performance within 5 different seeds are provided in Table 4. The key observations are 1) As the M increases, the model gets better performance and the training process becomes more stable; 2) The standard deviation in the results decreases as M increases, suggesting larger ensembles not only perform better on

Table 4: Ablation experiments of our entropybased diffusion policy with different ensemble sizes on selected AntMaze tasks.

average but also provide more reliable and consistent results. 3) While increasing M from 2 to 4 shows a substantial improvement, the performance gains decrease with an even larger size. It is worth noting that other offline RL approaches like Diffusion-QL [44] also adopt two Q networks for training robustness. See Appendix B for more detailed results.

LCB coefficients β We evaluate our method with β values of 1, 2, and 4 on AntMaze-medium environments Figure 3 demonstrates that adjusting the LCB coefficient improves performance, partic-

0.0

0.2

0.4

0.6

0.8

Normalized Score

1

2

4

LCB

0.0

0.2

0.4

0.6

0.8

Normalized Score

1

2

4

LCB

ularly for higher values, which helps in managing the exploration-exploitation trade-off effectively. In addition, the numerical results are provided in Appendix Table 8.

Training Stability and Computational Time Empirically we observe that the training of diffusion policies is always unstable, particularly for sparse-reward environments such as AntMaze medium and large tasks. Our method alleviates this problem by incorporating the entropy regularization and Q-ensembles as stated in the introduction. Here, we further show the comparison of training Diffusion-QL and our method on four AntMaze tasks as illustrated in Figure 4, maintaining the same number of diffusion steps T = 5 for both. It is observed that the performance of Diffusion-QL even drops down as the training step increases, while our method is substantially more stable and achieves higher results throughout all the training processes. We also included a detailed comparison of training and evaluation times for Gaussian and diffusion policies with Q-ensembles in Table 5. Increasing M from 2 to 64 almost does not influence the evaluation time. The diffusion step T has more impact on both training and evaluation time which is a common problem in diffusion models.

5 Related Work

Generative Diffusion Models and Meanreverting SDEs Recent advancements have integrated diffusion models [18, 42, 41, 40] and SDEs [43, 32, 45, 39] for realistic generative modeling. The development of Denoising Diffusion Probabilistic Models [18] showcases the ability of diffusion models to generate high-fidelity images through iterative reverse diffusion processes guided by deep neural networks, achieving state-of-the-art performance in generative tasks. In [33, 45, 39], meanreverting SDEs are applied to speech processing and image restoration tasks. These SDEs, similar to (1) but with different parameters, ensure our policy adapts across various distributions without bias. The gen-

0.0

0.2

0.4

0.6

0.8

Normalized Score

Our

Diffusion-QL

0.0

0.2

0.4

0.6

0.8

Normalized Score

0

5

10

15

Training Progress (#50k steps)

0.0

0.2

0.4

0.6

Normalized Score

antmaze-large-play-v0

0

5

10

15

Training Progress (#50k steps)

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

Normalized Score

antmaze-large-diverse-v0

eral applicability of our method is demonstrated in 4 D4RL benchmark domains. The comparison between our SDE and [43] are provided in Appendix A.6.

Diffusion Models in Offline RL Diffusion models in offline RL have gained growing attention for their potent modeling capabilities. In Janner et al. [19], diffusion models are introduced as trajectory planners trained with offline datasets for guided sampling, significantly mitigating compounding errors in model-based planning [46]. Diffusion models are also used as data synthesizers [5, 48], generating augmented training data to enhance offline RL robustness. Additionally, diffusion models approximate behavior policies [44, 21, 17], integrating Q-learning for policy improvement, though this can lead to overly conservative policies.

Entropy Regularization In online RL, maximum entropy strategies encourage exploration by maximizing rewards while maintaining high entropy [14, 15]. This approach develops diverse skills [7] and adapts to unseen goals [38]. However, its application in offline RL is challenging due to the multi-modal nature of datasets from various policies and expert demonstrations.

Uncertainty Measurement Balancing exploration and exploitation is crucial when data is limited. Online RL methods like bootstrapped DQN [37] and Thompson sampling [30] estimate uncertainty for exploration guidance. In offline RL, handling uncertainty is critical due to the lack of environment interaction. Model-based methods like MOPO [47] and MORel [23] measure and penalize uncertain model dynamics. Similarly, model-free methods like EDAC [1] and MSG [11] use Q-network ensembles to obtain pessimistic value estimations for policy guidance.

6 Conclusion

In this work, we present an entropy-regularized diffusion policy for offline RL, introducing meanreverting SDEs as the base framework to provide tractable entropy. Our theoretical contributions include deriving an approximated entropy for a diffusion model, enabling its integration as an entropy regularization component within the policy loss function. We also propose an optimal sampling process, ensuring the fast convergence of action generation from diffusion policy. Additionally, we enhance our method by incorporating Q-ensembles to handle the data uncertainty. Our experimental results show that combining entropy regularization with the LCB approach leads to a more robust policy, achieving state-of-the-art performance across offline RL benchmarks, particularly in AntMaze tasks with sparse rewards and suboptimal trajectories.

Future Work While the proposed method performs well on most D4RL tasks, the diffusion policy requires longer time when executed on compute- and power-constrained devices. Our future work will investigate real-time policy distillation under time and compute constraints to address this challenge.

References

Gaon An, Seungyong Moon, Jang-Hyun Kim, and Hyun Oh Song. Uncertainty-based offline reinforcement learning with diversified q-ensemble. Advances in neural information processing systems , 34:7436-7447, 2021.

Brian DO Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications , 12(3):313-326, 1982.

Chenjia Bai, Lingxiao Wang, Zhuoran Yang, Zhihong Deng, Animesh Garg, Peng Liu, and Zhaoran Wang. Pessimistic bootstrapping for uncertainty-driven offline reinforcement learning. arXiv preprint arXiv:2202.11566 , 2022.

Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems , 34:15084-15097, 2021.

Zoey Chen, Sho Kiami, Abhishek Gupta, and Vikash Kumar. Genaug: Retargeting behaviors to unseen situations via generative augmentation. arXiv preprint arXiv:2302.06671 , 2023.

Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in neural information processing systems , 34:8780-8794, 2021.

Benjamin Eysenbach, Abhishek Gupta, Julian Ibarz, and Sergey Levine. Diversity is all you need: Learning skills without a reward function. arXiv preprint arXiv:1802.06070 , 2018.

Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219 , 2020.

Scott Fujimoto and Shixiang Shane Gu. A minimalist approach to offline reinforcement learning. Advances in neural information processing systems , 34:20132-20145, 2021.

Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In International conference on machine learning , pages 2052-2062. PMLR, 2019.

Kamyar Ghasemipour, Shixiang Shane Gu, and Ofir Nachum. Why so pessimistic? estimating uncertainties for offline rl through ensembles, and why their independence matters. Advances in Neural Information Processing Systems , 35:18267-18281, 2022.

Daniel T Gillespie. Exact numerical simulation of the ornstein-uhlenbeck process and its integral. Physical review E , 54(2):2084, 1996.

Tuomas Haarnoja, Haoran Tang, Pieter Abbeel, and Sergey Levine. Reinforcement learning with deep energy-based policies. In International conference on machine learning , pages 1352-1361. PMLR, 2017.

Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning , pages 1861-1870. PMLR, 2018.

Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning , pages 1861-1870. PMLR, 2018.

Tuomas Haarnoja, Aurick Zhou, Kristian Hartikainen, George Tucker, Sehoon Ha, Jie Tan, Vikash Kumar, Henry Zhu, Abhishek Gupta, Pieter Abbeel, et al. Soft actor-critic algorithms and applications. arXiv preprint arXiv:1812.05905 , 2018.

Philippe Hansen-Estruch, Ilya Kostrikov, Michael Janner, Jakub Grudzien Kuba, and Sergey Levine. Idql: Implicit q-learning as an actor-critic method with diffusion policies. arXiv preprint arXiv:2304.10573 , 2023.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems , 33:6840-6851, 2020.

Michael Janner, Yilun Du, Joshua B Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. arXiv preprint arXiv:2205.09991 , 2022.

Ying Jin, Zhuoran Yang, and Zhaoran Wang. Is pessimism provably efficient for offline rl? In International Conference on Machine Learning , pages 5084-5096. PMLR, 2021.

Bingyi Kang, Xiao Ma, Chao Du, Tianyu Pang, and Shuicheng Yan. Efficient diffusion policies for offline reinforcement learning. arXiv preprint arXiv:2305.20081 , 2023.

Bingyi Kang, Xiao Ma, Chao Du, Tianyu Pang, and Shuicheng Yan. Efficient diffusion policies for offline reinforcement learning. arXiv preprint arXiv:2305.20081 , 2023.

Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. Morel: Model-based offline reinforcement learning. Advances in neural information processing systems , 33:21810-21823, 2020.

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114 , 2013.

Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learning. arXiv preprint arXiv:2110.06169 , 2021.

Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off-policy q-learning via bootstrapping error reduction. Advances in Neural Information Processing Systems , 32, 2019.

Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems , 33:1179-1191, 2020.

Sascha Lange, Thomas Gabel, and Martin Riedmiller. Batch reinforcement learning. In Reinforcement learning: State-of-the-art , pages 45-73. Springer, 2012.

Tor Lattimore and Csaba Szepesvári. Bandit algorithms . Cambridge University Press, 2020.

Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643 , 2020.

Ziwei Luo, Fredrik K Gustafsson, Zheng Zhao, Jens Sjölund, and Thomas B Schön. Image restoration with mean-reverting stochastic differential equations. International Conference on Machine Learning , 2023.

Ziwei Luo, Fredrik K Gustafsson, Zheng Zhao, Jens Sjölund, and Thomas B Schön. Refusion: Enabling large-size realistic image restoration with latent-space diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 1680-1691, 2023.

Diganta Misra. Mish: A self regularized non-monotonic activation function. arXiv preprint arXiv:1908.08681 , 2019.

Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. In International conference on machine learning , pages 1928-1937. PMLR, 2016.

Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning , pages 8162-8171. PMLR, 2021.

Ian Osband, Charles Blundell, Alexander Pritzel, and Benjamin Van Roy. Deep exploration via bootstrapped DQN. Advances in neural information processing systems , 29, 2016.

Sinno Jialin Pan and Qiang Yang. A survey on transfer learning. IEEE Transactions on knowledge and data engineering , 22(10):1345-1359, 2009.

Julius Richter, Simon Welker, Jean-Marie Lemercier, Bunlong Lay, and Timo Gerkmann. Speech enhancement and dereverberation with diffusion-based generative models. IEEE/ACM Transactions on Audio, Speech, and Language Processing , 2023.

Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 10684-10695, 2022.

Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, et al. Photorealistic text-toimage diffusion models with deep language understanding. Advances in neural information processing systems , 35:36479-36494, 2022.

Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502 , 2020.

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.

Zhendong Wang, Jonathan J Hunt, and Mingyuan Zhou. Diffusion policies as an expressive policy class for offline reinforcement learning. arXiv preprint arXiv:2208.06193 , 2022.

Simon Welker, Julius Richter, and Timo Gerkmann. Speech enhancement with score-based generative models in the complex STFT domain. In Proc. Interspeech 2022 , pages 2928-2932, 2022. doi: 10.21437/ Interspeech.2022-10653.

Chenjun Xiao, Yifan Wu, Chen Ma, Dale Schuurmans, and Martin Müller. Learning to combat compounding-error in model-based reinforcement learning. arXiv preprint arXiv:1912.11206 , 2019.

Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Model-based offline policy optimization. Advances in Neural Information Processing Systems , 33:14129-14142, 2020.

Tianhe Yu, Ted Xiao, Austin Stone, Jonathan Tompson, Anthony Brohan, Su Wang, Jaspiar Singh, Clayton Tan, Jodilyn Peralta, Brian Ichter, et al. Scaling robot learning with semantically imagined experience. arXiv preprint arXiv:2302.11550 , 2023.

Zhengbang Zhu, Hanye Zhao, Haoran He, Yichao Zhong, Shenyu Zhang, Yong Yu, and Weinan Zhang. Diffusion models for reinforcement learning: A survey. arXiv preprint arXiv:2311.01223 , 2023.

Brian D Ziebart. Modeling purposeful adaptive behavior with the principle of maximum causal entropy . Carnegie Mellon University, 2010.

Appendix

A Proof

A.1 Solution to the Forward SDE

Given the forward Stochastic Differential Equation (SDE) represented by

where θ t and σ t are time-dependent positive functions, and w denotes a standard Wiener process. We consider the special case where σ 2 t = 2 θ t for all t . The solution for the transition probability from time τ to t ( τ < t ) is given by

Proof. The proof is in general similar to that in IR-SDE [32]. To solve Equation (18), we introduce the transformation

and apply Itô's formula to obtain

Integrating from τ to t , we get

we can analytically compute the two integrals as θ t and σ t are scalars and then obtain

Rearranging terms and dividing by e ¯ θ t , we obtain

The integral term is actually a Gaussian random variable with mean zero and variance

under the condition σ 2 t = 2 θ t . Thus, the transition probability is

This completes the proof.

Loss function From (19), the marginal distribution of p ( a ( t )) can be written as

where we substitute ¯ θ 0: t with ¯ θ t for clear notation. During training, the initial diffusion state a 0 is given and thus we can obtain the ground truth score ∇ a log p t ( a ) based on the marginal distribution:

which can be approximated using a neural network and optimized with score-matching loss. Moreover, the marginal distribution (27) gives the reparameterization of the state:

where ϵ t is a standard Gaussian noise ϵ t ∼ N (0 , I ) . By substituting (29) into (28), the score function can be re-written in terms of the noise as

Then we follow the practical settings in diffusion models [18, 6] to estimate the noise with a timedependent neural network ϵ ϕ and optimize it with a simplified noise matching loss:

where t is a randomly sampled timestep and { a t } T t =0 denotes the discretization of the diffusion process. And this loss (31) is the same as (5) in the main paper.

A.2 Sampling from the Posterior

Proposition 3.1. Given an initial variable a 0 , for any diffusion state a t at time t ∈ [1 , T ] , the posterior of the mean-reverting SDE (1) conditioned on a 0 is

which is a Gaussian with mean and variance given by:

where θ ′ i := ∫ i i -1 θ t dt and ¯ θ t is to substitute ¯ θ 0: t for clear notation.

Proof. The posterior of SDE can be derived from Bayes' rule,

Recall that the transition distribution p ( a t | a t -1 ) and p ( a t | a 0 ) can be known with the solution to the forward SDE. Since all the distributions are Gaussian, the posterior will also be a Gaussian.

where C ( a t , a 0 ) is some function not involving ( a t -1 ) 2 . With the standard Gaussian density function, the mean and the variance can be computed:

Thus we complete the proof.

A.3 Optimal Reverse Path

In addition, we can prove the distribution mean ˜ µ t ( a t , a 0 ) is the optimal reverse path from a t to a t -1 .

Proof. As stated in Proposition 3.1, the posterior is a Gaussian distribution and can be derived by Bayes' rule. Thus it is natural to find the optimal reverse path by minimizing the negative log-likelihood according to

From (34), we have

Then we can directly solve (37) by computing the gradient of the negative log-likelihood and setting it to 0 :

Since (39) is linear, we get

This completes the proof. Note that the second-order derivative is a positive constant, and thus a t -1 ∗ is the optimal point. And we find that this optimal reverse path is the same as our posterior distribution mean as shown in Proposition 3.1.

A.4 Entropy Approximation

Proposition 3.1 further shows that the conditional posterior from a 1 i to a 0 i is Gaussian, meaning that

where C is a constant and a 1 i can be approximated using

Proof. Let's consider sequentially sampled action states a 0 i , a 1 i from our optimal sampling strategy. Then we have

where ˆ a 0 i is the approximated action's initial state from Eq. (8). Proposition 3.1 shows that the conditional posterior from a 1 i to a 0 i is a Gaussian with certain mean and variance , meaning that the term π ( a 0 i | a 1 i , ˆ a 0 i ) is a computable constant and thus we can write

A.5 Visualization

For a more intuitive explanation of our approach, Figure 5 outlines the forward and reverse processes of the mean-reverting SDE used for action prediction.

Forward:Mean-revertingSDE

Data

Noise

p(a/at-1)

0

2

T

a

D

a

D

a

Reverse:optimal sampling for action prediction

A.6 Comparison to VP SDE

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

As t →∞ , its marginal distribution p t ( x ) converges to a stationary Gaussian with mean value µ , which explains the name: 'mean-reverting'. We assume that there is no prior knowledge of the actions and thus set µ = 0 to generate actions from standard Gaussian noise. Then, with µ = 0 , the mean-reverting SDE has the same form as VP SDE. However, in [43], no solution of the continuoustime SDE was given. The authors start from perturbing 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-time SDE. Compared to the original VP SDE, our mean-reverting SDE is analytically tractable, see (2) and thus its score ∇ x log p t ( x ) is easier to learn. More importantly, the solution of the mean-reverting SDE can be used for entropy approximation.

B Additional Experiments Details

B.1 More experiments on proposed optimal sampling with different sample step

Weadded the figures of data generation with fewer steps ( T = 1 and T = 2 ) for the toy task in Section 3.1. The results show that the optimal sampling strategy significantly outperforms the reverse-time SDE in all steps, further demonstrating the efficiency and effectiveness of our method.

B.2 Hyperparameters

As stated in Section 4.1, we keep our key hyperparameters, entropy weight α = 0 . 01 , ensemble size M = 64 , LCB coefficient β = 4 and diffusion steps T = 5 for all tasks in different domains. As for others related to our algorithm, we consider the policy learning rate, Q-learning weight η , and whether to use max Q backup. For implementation details, we consider the gradient clip norm, diffusion loss type, and whether to clip action at every diffusion step. We keep the hyperparameter same for tasks in the same domain except for the AntMaze domain. We use max Q backup [28] for complete tasks. The hyperparameter settings are shown in Table 6.

B.3 Automating Entropy Adjustment

It is possible to consider automating entropy adjustment similar to [15] but α depends on the state since the offline dataset is pre-collected and may be imbalanced across different states shown in

Data Distribution

Reconstruction (T=1)

Reconstruction (T=2)

Reconstruction (T=5)

Reconstruction (T=10)

Reconstruction (T=30)

Reconstruction (T=50)

Reverse-time SDE

Optimal Sampling

where α ( s i ) is implemented as a neural network with a single hidden layer consisting of 32 units and ¯ H is a desired minimum expected entropy which represents the desired level of exploration which can be set as a function of action space dimension without the need for extensive hyperparameter tuning across different tasks.

B.4 More Analysis for Q-Ensembles

Here, we provide more detailed experiments for analyzing the effect of ensemble sizes M as we discussed in Section 4.3. More specifically, the results of different ensemble sizes are reported in Table 7 and Figure 8, in which we also provide the variance that further shows the robustness of our method.

pen-human-v1

pen-cloned-v1

0

2

4

6

8

10

12

14

Return (Quality)

kitchen-complete-v0

kitchen-partial-v0

kitchen-mixed-v0

20

40

60

80

100

Return (Quality)

antmaze-large-play-v0

antmaze-large-diverse-v0

0.0

0.2

0.4

0.6

0.8

1.0

Return (Quality)

0

5

10

15

Training Progress (#50k steps)

0.0

0.2

0.4

0.6

0.8

Normalized Score

M

-1

M

-4

M

-64

0

5

10

15

Training Progress (#50k steps)

0.0

0.2

0.4

0.6

0.8

Normalized Score

0

5

10

15

Training Progress (#50k steps)

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

Normalized Score

antmaze-large-play-v0

0

5

10

15

Training Progress (#50k steps)

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

Normalized Score

antmaze-large-diverse-v0

B.5 Ablation study on LCB coefficients β

To explore the impact of different LCB coefficients β . We add an experiment of our method with β values of 1, 2, and 4 on AntMaze-medium environments. Figure 3 demonstrates that adjusting the LCB coefficient improves performance, particularly for higher values, which helps in managing the exploration-exploitation trade-off effectively. In addition, the numerical results are provided in Table 8.

B.6 Ablation study on Diffusion step T

We evaluated the impact of varying the number of diffusion steps on a range of tasks, including AntMaze, Gym, and Kitchen in Table 9. Our findings indicate that while increasing the number of steps generally improves performance, five steps provide the best balance across different tasks and

between performance and computational time in Gym and Antmaze tasks. For more complex tasks as in Kitchen and Pen, we choose T = 10 .

B.7 Ablation study on Max Q-back

We conducted experiments with and without max Q-backup on AntMaze tasks in Table 10. The inclusion of max Q-backup significantly enhances performance, particularly in more complex environments (e.g., Antmaze-large).

B.8 Offline vs Online Model Selection

We use the online experience to evaluation our model during training. Table 11 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.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The abstract and introduction clearly describe the introduction of meanreverting SDEs for entropy-regularized diffusion policies, the integration of Q-ensembles for robust policy improvement, and the empirical validation through performance on D4RL benchmarks, which are thoroughly discussed and validated in the subsequent sections of the paper.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: See Section 6.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: See Appendix A for more details.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: The code is provided in supplementary materials and additional experiments Details can be found in Appendix B.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: The code with a README file is provided.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: The experiment setting and main hyperparameters are shown in Section 4 and additional experiments Details can be found in Appendix B.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: Our results are reported averaging by 5 random seeds.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: The computer resources are shown in Section 4.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: The dataset used in the paper is public without privacy-related data.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: This paper presents work whose goal is to advance the field of machine learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: The paper poses no such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We use the dataset D4RL for training offline RL agents and cite the original paper in Section 4.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: The paper does not release new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work introduces a new (fast) SDE-based sampling technique to derive actions from a diffusion based policy.\n\nSTRENGTHS:\n- Compares favorably against most relevant benchmark (Diff-QL)\n- Good ablation studies; this was helpful to understand the importance of each component.\n- Strong connection to standard tools in SDEs from diffusion literature\n\nWEAKNESSES:\nI'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?\n\n- High memory requirement: What is the memory cost (i.e. in VRAM) for an ensemble of 64 models?\n- 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.\n\n\n**Experiments**\n- 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.)\n- 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)\n- L153, choice of $\\eta$ (see above comments)\n- 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)\n- 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)\n\nQUESTIONS:\n- 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)\n- 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?\n- 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$?\n- I'm curious why different gradient norms were required; do you have any insight here? Was there parameter/grad divergence in some cases?\n- 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\n- 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.\n- Can you explain a bit more about the sense of \"optimality\" regarding the reverse time SDE?\n- Use of Mish activation function seems a bit non-standard. How important is it?\n- L356: can you give a rough time comparison of the diffusion policy vs. standard approaches?\n- How does entropy regularization interact with the SDE sampling? Is MaxEnt required for diffusion to be useful?\n- Can you increase linewidth in Fig 1a/c and Fig 3? Adding symbol/markers would be helpful too for accessibility\n\nminor typographical:\n- L10: \"is tractable and that **it** can...\"\n- L27: rephrase \"work introduce\"\n- L95-100, maybe point to Fig 1 to see multi-modality of proposed approach\n\nLIMITATIONS:\nBeyond the weaknesses mentioned above, I believe the limitations were discussed adequately.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-06-24"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\n1. 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.\n\n2. 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.\n\n3. 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.\n\nQUESTIONS:\n1. 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.\n\n2. What are the potential challenges or limitations in extending this method to real-world offline RL tasks outside the D4RL benchmarks?\n\n3. 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?\n\nLIMITATIONS:\nThe authors have figured out the limitation of the proposed method on high computational cost and large inference time.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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]}$.\n\nEvaluations on D4RL show improved performance compared to baselines.\n\nSTRENGTHS:\nOriginality: good. The paper proposes a few novel ideas to make diffusion policy work in offline RL.\nQuality: good. Overall presentation and evaluation are good. Tested only on D4RL, but this is common in offline RL.\nClarity: excellent.\nSignificance: excellent. Potentially a new baseline on D4RL, especially AntMaze.\n\nWEAKNESSES:\n1) 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)?\n2) 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))\n\nQUESTIONS:\nIn 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.\n\nLIMITATIONS:\nComputational 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nI have some trouble overcoming some gaps in the mathematical proofs and found that some mathematical equations may not hold.\n\nThe empirical performance in Table 1 is not consistent with the training curve in Figure 3.\n\nMore detailed questions about these two aspects are in the questions section.\n\nQUESTIONS:\n1. A typo in Eq(5). There is an unnecessary right parenthesis.\n\n2. I found Mean-Reverting SDE is a mathematical rewrite of VP SDE, see [42]. VP SDE has form \n\\begin{align*}\n    dx=-\\frac{1}{2}\\beta(t)xdt+\\sqrt{\\beta(t)}dw\n\\end{align*}\nMean-Reverting SDE has form \n\\begin{align*}\n    dx=-\\theta_txdt+\\sqrt{2\\theta_t}dw\n\\end{align*}\n\n3. 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?\n\n4. 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.\n\n5. 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.\n\n6. It is hard for me to prove Eq(14). Can the authors elaborate on more details about its proof?\n\n7. 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$).\n\n8. 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$?\n\n9. 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?\n\n10. In Table 1, how is the score recorded? Is it the final round score, online model selection, offline model selection, or moving average?\n\n11. In Eq (34) and (35), should $t$, $t-1$, and $0$ be subscript or superscript? Why is $t-1$ subscript and $t$ superscript?\n\n12. Can the authors elaborate on why Eq(43) holds?\n\n13. In Table 5, the loss type is Likelihood or Noise. What are their mathematical forms?\n\n\n[41] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502, 2020.\n\n[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.\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}