Diffusion-DICE: In-Sample Diffusion Guidance for Offline Reinforcement Learning

Authors: liyuan  mao
haoran  xu
xianyuan  zhan
weinan  zhang
amy  zhang
u t austin
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 13:19:43

Abstract: One important property of DIstribution Correction Estimation (DICE) methods is that the solution is the optimal stationary distribution ratio between the optimized and data collection policy. In this work, we show that DICE-based methods can be viewed as a transformation from the behavior distribution to the optimal policy distribution. Based on this, we propose a novel approach, Diffusion-DICE, that directly performs this transformation using diffusion models. We find that the optimal policy's score function can be decomposed into two terms: the behavior policy's score function and the gradient of a guidance term which depends on the optimal distribution ratio. The first term can be obtained from a diffusion model trained on the dataset and we propose an in-sample learning objective to learn the second term. Due to the multi-modality contained in the optimal policy distribution, the transformation in Diffusion-DICE may guide towards those local-optimal modes. We thus generate a few candidate actions and carefully select from them to achieve global-optimum. Different from all other diffusion-based offline RL methods, the \textit{guide-then-select} paradigm in Diffusion-DICE only uses in-sample actions for training and brings minimal error exploitation in the value function. We use a didatic toycase example to show how previous diffusion-based methods fail to generate optimal actions due to leveraging these errors and how Diffusion-DICE successfully avoid that. We then conduct extensive experiments on benchmark datasets to show the strong performance of Diffusion-DICE.

Main paper: 
 introduction: 
 We study the problem of offline reinforcement learning (RL), where the goal is to learn effective policies solely from offline data, without any additional online interactions. Offline RL is quite useful for scenarios where arbitrary exploration with untrained policies is costly or dangerous, but sufficient prior data is available, such as robotics [18] or industrial control [56] . Most previous model-free offline RL methods add a pessimism term to off-policy RL algorithms [50, 9, 2] , this pessimism term acts as behavior regularization to avoid extrapolation errors caused by querying the Q-function about values of potential out-of-distribution (OOD) actions produced by the policy [25] . However, explicitly adding the regularization requires careful tuning of the regularization weight because otherwise, the policy will still output actions that are not seen in the dataset. DIstribution Correction Estimation (DICE) methods [35, 41, 34] provide an implicit way of doing so. By applying convex duality, DICE-based methods solve for the optimal stationary distribution ratio between the optimized and data collection policy in an in-sample manner without querying unseen actions [55] .In this paper, we extend the analysis of DICE-based methods: we show that the optimal distribution ratio in DICE-based methods can be extended to a transformation from the behavior distribution to the optimal policy distribution. This different view motivates the use of deep generative models, e.g. diffusion models [42, 15, 43] , to first fit the behavior distribution using their strong expressivity of fitting multi-modal distributions and then directly perform this transformation during sampling. We find that the optimal policy's score function can be decomposed into two terms: the behavior policy's score function and the gradient of a guidance term which depends on the optimal distribution ratio learned by DICE. The first term can be easily obtained from the diffusion model trained on the offline dataset. While it is usually intractable to find a closed-form solution of the second term, we make a subtle mathematical transformation and show its equivalence to solving a convex optimization problem. In this manner, both of these terms can be learned by only dataset samples, providing accurate guidance towards in-distribution while high-value data points. Due to the multi-modality contained in the optimal policy distribution, the transformation may guide towards those local-optimal modes due to stochasticity. We thus generate a few candidate actions and use the value function to select the max from them to go towards the global optimum. We term our method Diffusion-DICE, the guidethen-select procedure in Diffusion-DICE differs from all previous diffusion-based offline RL methods [49, 4, 14, 32] , which are either only guide-based or only select-based. Guide-based methods [32] use predicted values of actions generated by the diffusion behavior policy to guide toward high-return actions. Select-based methods [4, 14] bypass the need for guidance but require sampling a large number of actions from the diffusion behavior policy and using the value function to select the optimal one. All these methods need to query the value function of actions sampled from the diffusion behavior policy, which may produce potential OOD actions and bring overestimation errors in the guiding or selecting process. Our method, however, brings minimal error in the guide-step by using accurate in-sample guidance to generate in-distribution actions with high values, and by doing so, in the select step only a few candidate actions are needed to find the optimal one, which further reduces error exploitation. We use an illustrative toy case to demonstrate the error exploitation in previous methods and how Diffusion-DICE successfully alleviates that. We also verify the effectiveness of Diffusion-DICE in benchmark D4RL offline datasets [8] . Note that Diffusion-DICE also provides a replacement for the Gaussian policy extraction part used in current DICE methods, successfully unleashing the power of DICE-based methods. Diffusion-DICE surpasses both diffusion-based and DICE-based strong baselines, reaching SOTA performance in D4RL benchmarks [8] . We also conduct ablation experiments and validate the superiority of the guide-then-select learning procedure.Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?Answer: [Yes]Justification: Our main claims can be reflected by Section 3 and Section 4.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. 
preliminaries: 
 We consider the RL problem presented as a Markov decision process [45] , which is specified by a tuple M = ⟨S, A, P, d 0 , r, γ⟩. Here S and A are state and action space, P(s ′ |s, a) and d 0 denote transition dynamics and initial state distribution, r(s, a) and γ represent reward function and discount factor, respectively. The goal of RL is to find a policy π(a|s) which maximizes expected return E[ ∞ t=0 γ t • r(s t , a t )]. Another equivalent LP form of expected return is E (s,a)∼d π [r(s, a)], where d π (s, a) := (1 -γ) ∞ t=0 γ t P r(s t = s, a t = a) is the normalized discounted stationary distribution of π [36] . For simplicity, we refer to d π (s, a) as the stationary distribution of π. Offline RL considers the setting where interaction with the environment is prohibited, and one need to learn the optimal π from a static dataset D = {s i , a i , r i , s ′ i } N i=1 . We denote the empirical behavior policy of D as π D .DICE methods [41] incorporate the LP form of expected return J (π) = E (s,a)∼d π [r(s, a)] with a regularizer D f (d π ||d D ) = E (s,a)∼d D [f ( d π (s,a) d D (s,a) )], where D f is the f -divergence induced by a convex function f [3] . More specifically, DICE methods try to find an optimal policy π * that satisfies:EQUATIONThis objective is generally intractable due to the dependency on d π (s, a), especially under the offline setting. However, by imposing the Bellman-flow constraint [33] a∈A d(s, a) = (1 -γ)d 0 (s) + γ (s ′ ,a ′ ) d(s ′ , a ′ )p(s|s ′ , a ′ ) on states and applying Lagrangian duality and convex conjugate, its dual problem has the following tractable form:EQUATIONHere f * is a variant of f 's convex conjugate and T V (s, a) = r(s, a)+γE s ′ ∼p(•|s,a) [V (s ′ )] represents the Bellman operator on V . In practice, one often uses a prevalent semi-gradient technique in RL that estimates T V (s, a) with Q(s, a) and replaces the initial state distribution d 0 with dataset distribution d D to stabilize learning [41, 34] . In addition, because D usually cannot cover all possible s ′ for a specific (s, a), DICE methods only use a single sample of the next state s ′ . The update of Q(s, a) and V (s) in DICE methods are as follows and we refer to a detailed derivation in Appendix A:EQUATION)Note that learning objectives of DICE-methods can be calculated solely with a (s, a, s ′ ) sample from D, which is totally in-sample. One important property of DICE methods is that Q * and V * have a relationship with the optimal stationary distribution ratio w * (s, a) asEQUATIONwhere d * is the stationary distribution of π * . To get π * from w * , previous policy extraction methods in DICE methods include weighted behavior cloning [34] , information projection [27] or policy gradient [37] . All these methods parametrize an unimodal Gaussian policy in order to compute log π(a|s) [13] , which greatly limits its expressivity.Diffusion models [42, 15, 43] are generative models based on a Markovian noising and denoising process. Given a random variable x 0 and its corresponding probability distribution q 0 (x 0 ), the diffusion model defines a forward process that gradually adds Gaussian noise to the samples from x 0 to x T (T > 0). Kingma et al. [20] shows there exists a stochastic process that has the same transition distribution q t0 (x t |x 0 ) and Song et al. [43] shows that under some conditions, this process has an equivalent reverse process from T to 0. The forward process and the equivalent reverse process can be characterized as follows, where wt is a standard Wiener process in the reverse time.EQUATION)Here we slightly abuse the subscript t to denote the diffusion timestep. α t , σ t are the noise schedule and f (t), g(t) can be derived from α t , σ t [31] . For simplicity, we denote q t0 (x t |x 0 ) and p 0t (x 0 |x t ) as q(x t |x 0 ) and p(x 0 |x t ), respectively. To sample from q 0 (x 0 ) by following the reverse stochastic differential equation (SDE), the score function ∇ xt log q t (x t ) is required. Typically, diffusion models use denoising score matching to train a neural network ϵ θ (x t , t) that estimates the score function [46, 15, 43] , by minimizing E t∼U (0,T ),x0∼q0(x0),ϵ∼N (0,I) [∥ϵ θ (x t , t) -ϵ∥ 2 ], where x t = α t x 0 + σ t ϵ.As we mainly focus on diffusion policy in RL, this objective is usually impractical because q 0 (x 0 ) is expected to be the optimal policy π * (a|s). A more detailed discussion is given in Appendix A.To make diffusion models compatible with RL, there are generally two approaches: guide-based and select-based. Guide-based methods [32, 17] incorporate the behavior policy's score function with an additional guidance term. Specifically, they learn a time-dependent guidance term J t and use it to drift the generated actions towards high-value regions. The learning objective of J t can be generally formalized with L(J t (a t ), w(s,{a i 0 } K i=1 ))), where w(s, {a i 0 } K i=1 ) are critic-computed values on K diffusion behavior samples. L can be a contrastive objective [32] or mean-square-error objective [17] .After training, the augmented score function ∇ at log π t (a t |s) = ∇ at log π D t (a t |s) + ∇ at J t (a t ) is used to characterize the learned policy.Select-based methods [4, 14] utilize the observation that for some RL algorithms, the actor induced through critic learning manifests as a reweighted behavior policy. To sample from the optimized policy, these methods first sample multiple candidates {a i 0 } N i=1 from the diffusion behavior policy and then resample from them using critic-computed values w(s, {a i 0 } N i=1 ). More precisely, the sampling procedure follows the categorical distribution P r[a = a j 0 |s] = w(s,a j 0 )N i=1 w(s,a i 0 ) . Error exploitation As we can see, both guide-based and select-based methods need the information of w(s, {a i 0 } N i=1 ) to get the optimal action. However, this term may bring two sources of errors. One is the diffusion model's approximation error in modeling complicated policy distribution, and the other is the critic's error in evaluating unseen actions. Although trained on offline data, the diffusion model may still generate OOD actions (especially with frequent sampling) and the learned critic can make erroneous predictions on these OOD actions, causing the evaluated value on these actions to be over-estimated due to the learning nature of value functions [11, 25] . As a result, the generation of high-quality actions in existing methods is greatly affected due to this error exploitation, which we will also show empirically in the next section.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes]Justification: Limitations are discussed in Section 6Guidelines:• 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. 
diffusion-dice: 
 In this section, we introduce our method, Diffusion-DICE. We start from an extension of DICE-based method that considers DICE as a transformation from the behavior distribution to the optimal policy distribution, which motivates us to use diffusion models to perform such transformation. We then propose a guide-then-select procedure to achieve the best action, i.e., we propose in-sample guidance learning for accurate policy transformation and use the critic to do optimal action selection to boost the performance. We also propose a piecewise f -divergence to stabilize the gradient during learning. We give an illustration of the guide-then-select paradigm and use a toycase to showcase the error exploitation in previous methods and how Diffusion-DICE successfully alleviates that.As mentioned before, DIstribution Correction Estimation (DICE) methods, an important line of work in offline RL and IL provides us with an elegant way to estimate the optimal stationary distribution ratio between d * (s, a)) and d D (s, a) [28, 41, 34] . We show that this ratio also directly indicates a proportional relationship between the optimal in-support policy π * (a|s) and the behavior policy π D (a|s). This proportional relationship enables us to transform π D into π * .Our key observation is that the definition of d π (s, a) inherently reveals a bijection between π(a|s) and d π (s, a). Given a relationship between d * (s, a) and d D (s, a), we can use this bijection to derive a relationship between π * (a|s) and π D (a|s). We formalize the bijection and derived relationship as:π(a|s) = d π (s, a) a∈A d π (s, a)da , π * (a|s) = d * (s,a) d D (s,a) a∈A d * (s,a) d D (s,a) π D (a|s)da π D (a|s).The proof is given in Appendix B. Since the denominator in the second equation involves an integral over a ∈ A and is unrelated to the chosen action, the relationship indicates π * (a|s) ∝ d * (s,a) d D (s,a) π D (a|s). This means that the transformation between d D (s, a) and d * (s, a) can be extended to the transformation between π D (a|s) and π * (a|s). This transformation motivates the use of deep generative models, e.g. diffusion models, to first fit the behavior distribution using their strong expressiveness and then directly perform this transformation during sampling. More specifically, the score function of the optimal policy and the behavior policy satisfy the following relationship:EQUATION)The proof is given in Appendix B. Intuitively, ∇ at log π D t (a t |s) tells us how to generate actions from π D and ∇ at log E π D (a0|at,s) [w * (s, a 0 )] tells us how to transform from π D to in-support π * during the reverse diffusion process. As mentioned before, when representing π * with a diffusion model, all we need is its score function ∇ at log π * t (a t |s). Equivalently, we focus on the right-hand side of Eq. (6) . The first term is just the score function of π D , which is fairly easy to obtain from the offline dataset [14, 4] . To perform the transformation, we still require the second term. Fortunately, DICE allows us to acquire the inner ratio term in an in-sample manner directly from the offline dataset, and we will show in the next section how to exactly compute the whole second term using offline dataset.Although DICE provides us with the optimal stationary distribution ratio, which is the cornerstone of the transformation, the second term on the right-hand side of Eq.( 6) is still intractable due to the conditional expectation. Our key observation is that for arbitrary non-negative function f (x), the optimizer of the following convex problem is unique and takes the desired form of log-expectation. Lemma 1. Given a random variable X and its corresponding distribution P (X), for any nonnegative function f (x), the following problem is convex and its optimizer is given by y* = log E x∼P (X) [f (x)], min y E x∼P (X) [f (x) • e -y + y].It is evident that the optimizer can be used to derive the second term in Eq.( 6). In fact, we show the second term can be obtained directly from the offline dataset, if we optimize the following objective.Here g θ is a guidance network parameterized by θ and we denote d * (s,a) d D (s,a) as w * (s, a) for shorthand. Theorem 1. ∇ at log E π D (a0|at,s) [w * (s, a 0 )] can be obtained by solving the following optimization problem:EQUATIONas the optimal solution θ * satisfies ∇ at g θ * (s,a t , t) = ∇ at log E π D (a0|at,s) [w(s, a 0 )].This objective has two appealing properties. Firstly, compared to guide-based and select-based methods, we don't need to use multiple actions generated by the behavior diffusion model, we only need one sample from π D (a|s) to estimate the second expectation, enabling an unbiased estimator of this objective using only offline dataset possible. Secondly, because w(s, a) must be non-negative to ensure a valid transformation, this objective is convex with respect to g θ (s, a t , t). This indicates a guarantee of convergence to g θ * (s, a t , t) under mild assumptions (see proof in Appendix B).Directly inducing ∇ at log E π D (a0|at,s) [w * (s, a 0 )] from the dataset is crucial in the offline setting. Firstly, it avoids the evaluation of OOD actions and solely depends on reliable value predictions of in-sample actions. Consequently, the gradient exploits minimal error in the critic, enabling it to more accurately transform π D into in-support π * . Moreover, because all values of w * (s, a) used in our method are based on in-sample actions, applying this gradient in the reverse process will guide the samples towards in-sample actions with high w * (s, a) value.We refer to this method as In-sample Guidance Learning (IGL). Note that previous methods either have a biased estimate of the gradient [17, 6] , or have to rely on value predictions for diffusion generated actions, which could be OOD [32] . We refer to Appendix C for an in-depth discussion of IGL against other guidance methods.Stablizing gradient using piecewise f -divergence. In practice, there exists one issue when computing the guidance term. Note that in Eq.( 4), w * (s, a) could become zero, depends on the choice of f . Given a specific a t , if w * (s, a 0 ) = 0 for actions under π D (a 0 |a t , s), the second gradient term will not be well-defined due to taking the logarithm of 0. In practice, this can result in an unstable gradient. To avoid such issue, (f ′ ) -1 (x) should be positive for any given x. Also, to facilitate the optimization process, f * (x) should possess a closed-form solution and good numerical properties. Unfortunately, none of the commonly used f -divergence can accommodate both. However, considering the following two f -divergences:Divergence f (x) f * (x) (f ′ ) -1 (x) Reverse KL x log x e x-1 e x-1 Pearson χ 2 (x -1) 2 max(-1, x 2 4 + x) max(0, x 2 + 1)It is obvious that Reverse KL divergence possesses the positive property but exhibits numerical instability given a large x due to the exponential function in f * (x), while Pearson χ 2 divergence avoids the instability in the exponential function, its (f ′ ) -1 (x) value is negative when x < -2.To take advantage of both divergences while avoiding their drawbacks, we propose the following piecewise f -divergence that has properties similar to Pearson χ 2 when x is large while has properties similar to Reverse KL when x is small:EQUATIONIn Appendix B, we prove that f (x) is a well-defined f -divergence. We also prove that f (x) has the following proposition: Proposition 1. Given the piecewise f -divergence in Eq.( 8), (f ′ ) -1 (x) and f * (x) has the following formulation:EQUATIONWe can see that given any x, the exponential function in (f ′ ) -1 (x) ensures a strictly positive value. Meanwhile, f * (x) is at most a quadratic polynomial, which ensures numerical stability in practice. Note that f (x) enables a stable policy transformation and can also be applied to other DICE-based methods when the optimal stationary distribution ratio needs to be positive.After transforming π D into in-support π * , we now introduce the select-step to boost the performance during evaluation. One issue is the multi-modality in the optimal policy, this partially arises from the policy constraint used to regularize its output [9, 25, 55, 34] . Under policy constraint, the regularized optimal policy is unlikely to be deterministic and may be multi-modal.More specifically, in our method, as we strictly follow the score function ∇ at log π * t (a t |s) of insupport optimal policy π * to generate high-value actions during the reverse diffusion process, any mode that has non-zero probability under π * (a|s) could be generated, which lead to a sub-optimal choice. Similar issues have been discussed in previous works [14, 13] . In Diffusion-DICE, it's natural to leverage the optimal critic Q * derived from DICE in Eq.( 3) to identify the best action. Given a specific state s, we first follow ∇ at log π * t (a t |s) in the reverse diffusion process to generate a few actions from π * (a|s). Then we evaluate these actions with Q * and select the action with the highest Q * (s, a) value as the policy's output asEQUATION) Algorithm 1 Diffusion-DICE 1: Initialize value function Q ϕ 1 , V ϕ 2 , diffusion behav- ior model ϵ ψ , guidance network g θ 2: // Training 3: Pretrain the diffusion behavior model by minimizing E U (t),(s,a)∼D,ϵ∼N (0,I) [∥ϵ ψ (αta + σtϵ, s, t) -ϵ∥ 2 ] 4: for t = 1, 2, • • • , N do 5: Sample transitions (s, a, r, s ′ ) ∼ D 6:Update Q ϕ 1 and V ϕ 2 by minimizing Eq.(3) 7:Update g θ by minimizing Eq.( 7) 8: end for 9: // Evaluation 10: Sample {a i 0 } K i=1 from π * following Eq.( 5), with ∇a t log π * t (at|s) = ϵ ψ (a t ,s,t) -σ t + ∇a t g θ (s, at, t) 11: Select action according to Eq. (10) Different from previous guide-only or selectonly methods, we base the select stage on the guide stage. As mentioned before, using IGL accurately guides the generated candidate actions towards in-support actions with high value. With the assistance of the guide stage, we can sample candidates from π * (a|s), which has a high probability of being high-quality. This means only a small number of candidates are required to be sampled, which reduces the probability of sampling OOD actions. This leads to minimal error exploitation while still attaining high returns. We give an illustration of the guide-then-select paradigm in Figure 1 .We term this guide-then-select method Diffusion-DICE and present its pseudo-code in Algorithm 1. During the training stage, Diffusion-DICE estimates the optimal stationary distribution ratio using DICE with our piecewise f -divergence and calculates the guidance with IGL, both in an in-sample manner. During the testing stage, Toycase validation. We use a toycase to validate that the guide-then-select paradigm used in Diffusion-DICE indeed brings minimal error exploitation that other diffusion-based methods suffer from. Here we aim to solve a 2-D bandit problem given a fixed action dataset. The action space is continuous and actions in the offline dataset follow a bivariate standard normal distribution constrained within an annular region. We show the dataset distribution π D , the learned diffusion behavior policy πD , the ground truth reward R and the predicted reward R in Figure 2 . Note that the true optimal reward occurs on the outer circle. However, due to limited data coverage, the learned reward function exploits overestimation error on unseen regions, e.g., actions inside the inner circle have erroneous high values. What's worse, due to fitting errors, the diffusion behavior policy may generate such overestimated actions. We take Diffusion-DICE with a guide-based method, QGPO [32] , and a select-based method, IDQL [14] for comparison.More specifically, we visualize the sample generation process in different methods. For guide-based method QGPO, as it utilizes actions generated by πD and the value of R on them to obtain guidance, the guidance exploits overestimation errors of OOD actions in the center and drifts the generated actions towards them. Note that adding a select step in QGPO will not help as the guide-step is already incorrect. IDQL, because it is free of guidance, requires a large number of generated candidate actions to ensure the coverage of optimal action. IDQL utilizes R to select from those suboptimal or out-of-distribution actions, which however, are the sources of error exploitation in R.For Diffusion-DICE, because the first guide-step accurately guides the generated actions towards in-support while high-value regions (i.e., actions around the outer circle), in the select step, the learned value function R could make a correct evaluation and successfully select the true best actions in the outer circle. Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [Yes] Justification: Seed Appendix B 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. 
experiments: 
 In this section, we present empirical evaluations of Diffusion-DICE. To validate Diffusion-DICE's ability to transform the behavior policy into the optimal policy while maintaining minimal error exploitation, we conduct experiments on two fronts. On one hand, we evaluate Diffusion-DICE on the D4RL offline RL benchmark and compare it with other strong diffusion-based and DICE-based methods. On the other hand, we use alternative criteria to demonstrate that the guide-then-select paradigm results in minimal error exploitation. Experimental details are shown in Appendix D.We first evaluate Diffusion-DICE on the D4RL benchmark [8] and compare it with several related algorithms. For the evaluation tasks, we select MuJoCo locomotion tasks and AntMaze navigation tasks. While MuJoCo locomotion tasks are popular in offline RL, AntMaze navigation tasks are more challenging due to their stronger need for trajectory stitching. For baseline algorithms, we selected state-of-the-art methods not only from traditional methods that use Gaussian-policy (including DICEbased methods) but also from diffusion-based methods. Gaussian-policy-based baseline includes CQL [26] , in-sample based methods IQL [24] , SQL [55] and DICE-based method O-DICE [34] . Notably, O-DICE is a recently proposed DICE-based algorithm that stands out among various DICE-based offline RL methods. Diffusion-policy-based baseline includes Diffusion-QL [49] SfBC [4] , QGPO [32] and IDQL [14] . We also compare Diffusion-DICE with its Gaussian-policy counterpart to show the benefit of using Diffusion-policy in Appendix D. While SfBC and IDQL simply sample from behavior policy candidates and select according to the action evaluation model, Diffusion-QL and QGPO will guide generated actions towards high-value ones. It is worth noting that Diffusion-QL will also resample from generated actions after the guidance.The results show that Diffusion-DICE outperforms all other baseline algorithms, including previous SOTA diffusion-based methods, especially on MuJoCo medium, medium-replay datasets, and AntMaze datasets. The consistently better performance compared with Diffusion-QL and QGPO demonstrates the essentiality of in-sample guidance learning. Compared with SfBC and IDQL, Diffusion-DICE also shows superior performance, even with fewer action candidates. This is because the guide stage provides the in-support optimal policy for the select stage to sample from, which underscores the necessity of sampling carefully from an in-support optimal action distribution, rather than from the behavior distribution. We refer to the comparison of candidate numbers under different environments in Appendix D. Furthermore, the substantial performance gap between Diffusion-DICE and O-DICE reflects the multi-modality of DICE's optimal policy and firmly positions Diffusion-DICE as a superior policy extraction method for other DICE-based algorithms.We then continue to verify that the guide-then-select paradigm used in Diffusion-DICE indeed exploits minimal error. This is obvious for the guide stage because we only leverage in-sample actions for both critic training and guidance learning. For the select stage, additional evidence is needed, as Diffusion-DICE also selects actions from generated candidates. Our key observation is that if the action value remains high during evaluation but results in a low return trajectory, these action values must suffer from overestimation error. Based on this, we choose to compare the average state-action value and the average return between actions selected by our guide-then-select paradigm and simply select from the behavior policy. Given identical critic and behavior policy, we run Diffusion-DICE and its guidancefree variant for 10 episodes, comparing their average Q(s, a) values and their normalized scores. The results are shown in Figure 3 . We also compare the learning curves of their average Q(s, a) values over time in Appendix D. It's clear from the results that our guide-then-select paradigm achieves better performance while having less overestimated Q(s, a) values. This result validates the minimal error exploitation in Diffusion-DICE.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: Experimental details are given in Appendix D. 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. 
related work: 
 Offline RL To tackle the distributional shift problem, most model-free offline RL methods augment existing off-policy RL methods with a behavior regularization term. Behavior regularization can appear explicitly as divergence penalties [50, 25, 51, 9] , implicitly through weighted behavior cloning [48, 38, 52] , or more directly through careful parameterization of the policy [11, 58] . Another way to apply action-level regularization is via modification of value learning objective to incorporate some form of regularization, to encourage staying near the behavioral distribution and being pessimistic about OOD state-action pairs [26, 24, 55, 47] . There are also several works incorporating action-level regularization through the use of uncertainty [2] or distance function [29] .All these methods are based on the actor-critic framework and use unimodal Gaussian policy. However, several works indicate their limited ability to model multi-modal policy distribution [49, 14, 4] , thereby leading to suboptimal performance. To remedy this, it's natural to employ powerful generative models to represent the policy. DT [5] uses transformer as the policy, Diffusion-QL [49] , SfBC [4] , QGPO [32] and IDQL [14] leverage diffusion models to represent policy. Other generative models like CVAE and consistency model have also been used as policy in offline RL [58, 54, 7] . Another line of methods utilizes diffusion models for trajectory-level planning by generating high-return trajectories and taking the corresponding action of the current state [17, 1] . While generative models are widely used in offline RL, few methods consider the existing errors in these models and whether the generation process exploits these errors.The core of DICE-based methods revolves around the ratio of the stationary distribution between two policies. This ratio can serve as the estimation target in off-policy evaluation [36, 57] or as part of a state-action level constraint term in RL [37, 41, 34] . In offline IL [53] , it can connect the optimal stationary distribution of a regularized MDP with the dataset distribution [22, 12, 19] . In constrained RL, this ratio can induce the discounted sum of cost without an additional function approximator [28] . Under the imperfect rewards setting, this ratio can reflect the gap between the given rewards and the underlying perfect rewards [30] . All of these DICE methods consider this ratio as a transformation from one stationary distribution to another. In this paper, however, we extend this ratio as a transformation from the behavior policy to the optimal policy and propose a novel way of doing so by using diffusion models, which also serve as a replacement for the Gaussian-based policy extraction in DICE.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 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). 
conclusion: 
 In this work, we propose a new diffusion-based offline RL algorithm, Diffusion-DICE. Diffusion-DICE uses a guide-then-select paradigm to select the best in-support actions while achieving minimal error exploitation in the value function. Diffusion-DICE also serves as a replacement for the Gaussian-based policy extraction part in current DICE methods, successfully unleashing the power of DICE-based methods. Through toycase illustration and extensive experiments, we show that Diffusion-DICE outperforms prior SOTA methods on a variety of datasets, especially those with multi-modal complex behavior distribution. One limitation of Diffusion-DICE is the sampling process of diffusion models is costly and slow. Another limitation is the training of diffusion models may suffer in low-data regimes. One future work is using more advanced generative models [44, 40] as a better choice.NeurIPS Paper Checklist 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes]Justification: Potential positive impacts and negative impacts are discussed in Appendix F.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). 
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: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. 
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: All assets we use are properly credited, and the corresponding license are given in Appendix D.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. 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. 
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: 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.A A More Detailed Discussion of DICE and Diffusion Model in RL Derivation of learning objectives in DICE DICE algorithms consider the following regularized RL problem as a convex programming problems with Bellman-flow constraints and apply Fenchel-Rockfeller duality or Lagrangian duality to solve it. The regularization term aims at imposing state-action level constraints. [35, 27, 34] .EQUATIONHereD f (d π (s, a)∥d D (s, a)) is the f -divergence which is defined with D f (P ∥Q) = E ω∈Q f P (ω) Q(ω). Directly solving π * is impossible because it's intractable to calculate d π (s, a). However, one can change the optimization variable from π to d π because of the bijection existing between them. Then with the assistance of Bellman-flow constraints, we can obtain an optimization problem with respect to d:EQUATIONNote that the feasible region has to be {d : ∀s ∈ S, a ∈ A, d(s, a) ≥ 0} because d should be non-negative to ensure a valid corresponding policy. After applying Lagrangian duality, we can get the following optimization target following [27] :EQUATIONHere we denote ω(s, a) as d(s,a) d D (s,a) for simplicity. By incorporating the non-negative constraint of d and again solving the constraint problem with Lagrangian duality, we can derive the optimal solution w * (s, a) for the inner problem and thus reduce the bilevel optimization problem to the following optimization problem:EQUATIONHere f * is a variant of f 's convex conjugate. Note that in the offline RL setting, the inner summations ′ p(s ′ |s, a)V (s ′) is usually intractable because of limited data samples. To handle this issue and increase training stability, DICE methods usually use additional network Q(s, a) to fit r(s, a) + γ s ′ p(s ′ |s, a)V (s ′ ), by optimizing the following MSE objective:EQUATIONAnd because of doing so, the optimization objective in 14 can be replaced with:EQUATIONAlso note that to increase the diversity of samples, one often extends the distribution of initial state d 0 to d D by treating every state in a trajectory as initial state [22] .Discussion of diffusion policy and other methods in offline RL In offline RL, as discussed before, most diffusion-based algorithms utilize diffusion models to represent policy π. This policy can either be the behavior policy [14, 4] , the optimal policy [49, 32] . When representing policy with diffusion model, the ultimate goal is to model the optimal policy π * with diffusion model. Different from other tasks in image generation or video generation [43, 16] , RL tasks will not provide data samples under π * , while the agent has to learn such optimal policy. This makes directly leveraging the diffusion model to fit the optimal policy impossible. However, besides these methods that are based on diffusion policy, some methods utilize a diffusion model to fit the entire trajectory's distribution [17, 1] . These methods try to generate the optimal trajectory based on the current state. After generating the entire trajectory, these methods simply take the action corresponding to the current state. It's worth noting that although these methods avoid modeling π * , extra effort is required to generate high-quality trajectories. 
 
 Appendix: Optimal policy transformation with stationary distribution ratio: Given that π D is the behavior policy, d * (s,a) d D (s,a) and π * are the optimal stationary distribution correction and its corresponding optimal policy. π * and π D satisfy the following proposition:EQUATIONProof. Using the definition between d π (s) and d π (s, a), we have for any π, d π (s, a) = d π (s)•π(a|s).Then we have:EQUATIONVerification of Eq.( 6): The score function of the optimal policy π * (a|s) in DICE and the behavior policy π D (a|s) satisfy:∇ at log π * t (a t |s) = ∇ at log π D t (a t |s) + ∇ at log E a0∼π D (a0|at,s) [ d * (s,a0) d D (s,a0)]. We mainly follow the proof in Lu et al. [32] to verify this equation. Given the relationship between π * (a|s) and π * (a|s) in Eq.( 18) we have: π * t (a t |s) = p(a t |a 0 )π * 0 (a 0 |s)da 0 = p(a t |a 0 , s)π * (a 0 |s)da 0 = p(a t |a 0 , s) d * (s,a0) d D (s,a0) a∈A d * (s,a) d D (s,a) π D (a|s)da π D (a 0 |s)da 0 = p(a t |a 0 , s)π D 0 (a 0 |s) d * (s,a0) d D (s,Note that the second equation holds because p(a t |a 0 ) represents the forward transition probability in the diffusion process, which is unrelated to s and π * 0 is exactly the optimal policy π * . The forth equation holds for the similar reasons. We exchange the condition variable to derive the final relationship between π * t (a|s) and π * t (a|s). Having this relationship, we can directly calculate their score functions as follows: ∇ at log π * t (Note that the second equation holds because the integral a∈A d * (s,a) d D (s,a) π D (a|s)da is unrelevant to a 0 , a t and because of so, its gradient with respect to a t becomes 0.Lemma 1. Given a random variable X and its corresponding distribution P (X), for any non-negative function f (x), the following problem is convex and its optimizer is given by yEQUATIONProof. Because y is the constant optimizer, this problem can be reformulated as:EQUATIONWe first verify that this problem is actually a convex problem with respect to y. Because f (x) is a non-negative function, we can take the second derivative of this objective and get:EQUATIONBecause f (x) is a non-negative function, the second derivative is also non-negative, which means this problem is a convex problem with respect to y. For a convex problem, the optimizer can be induced from the first-order condition:EQUATIONTheorem 1. ∇ at log E π D (a0|at,s) [w(s, a 0 )] can be obtained by solving the following optimization problem:EQUATIONThe optimal solution θ * satisfies ∇ at g θ * (s, a t , t) = ∇ at log E π D (a0|at,s) [w(s, a 0 )].Proof. Because the forward diffusion process simply adds noise to a 0 and is unrelated to s, taking s as a condition doesn't affect p(a t |a 0 ). This means p(a t |a 0 , t) = p(a t |a 0 ). By changing the conditional variables from (a 0 ) (1:K) to (a t ) (1:K) we have:E t∼U (0,T ) E π D (a0|s) E p(at|a0) w(s, a 0 )e -g θ (s,at,t) + g θ (s, a t , t)= E t∼U (0,T ) E π D (a0|s) E p(at|a0,s) w(s, a 0 )e -g θ (s,at,t) + g θ (s, a t , t)= E t∼U (0,T ) E π D (at,a0|s) w(s, a 0 )e -g θ (s,at,t) + g θ (s, a t , t)= E t∼U (0,T ) E π D (at|s) E π D (a0|at,s) w(s, a 0 )e -g θ (s,at,t) + g θ (s, a t , t)= E t∼U (0,T ) E π D (at|s) E π D (a0|at,s) w(s, a 0 ) e -g θ (s,at,t) + g θ (s, a t , t)We first take the first-order derivative with respect to each g θ (s, a t , t) for this objective and solve its saddle point:∇ g θ (s,at,t) E t∼U (0,T ) E π D (at|s) E π D (a0|at,s) w(s, a 0 ) e -g θ (s,at,t) + g θ (s, a t , t)EQUATIONBecause any possible intermediate action a t has a positive value of π D (a t |s), the saddle point g θ * (s, a t , t) must satisfy:-E π D (a0|at,s) w(s, a 0 ) e -g θ * (s,at,t) + 1 = 0EQUATIONand then ∇ at g θ * (s, a t , t) = ∇ at log E π D (a0|at,s) [w(s, a 0 )]We then take the second-order derivative with respect to each g θ (s, a t , t) for this objective:∇ 2 g θ (s,at,t) E t∼U (0,T ) E π D (at|s) E π D (a0|at,s) w(s, a 0 ) e -g θ (s,at,t) + g θ (s, a t , t)EQUATIONBecause w(s, a 0 ) has to be non-negative for any a 0 to ensure a valid transformation from π D to π * , it's obvious that the second-order derivative is always non-negative. This means this objective is convex with respect to g θ (s, a t , t) and has convergence guarantee given unlimited model capacity.Lemma 2. The following f p (x) is a well-defined f -divergence:EQUATIONProof. By the definition of f -divergence in Rényi [39] , f p (x) should be a convex function from [0, +∞) → (-∞, +∞]. f p (x) should be finite for all x > 0 and f p (1) = 0, f p (0) = lim x→0 + f p (x).We first verify that f p (x) is convex. By taking first-order and second-order derivatives on f p (x) we have:EQUATIONIt's obvious that f ′ p (x) is continuous and f ′′ p (x) is strictly positive. According to the second-order condition for convex functions, f p (x) is convex.We then move on to verify the rest of the properties. It's obvious that given a specific x > 0, f p (x) is finite and f p (1) = 0. In mathematics, the value of x log x when x = 0 is defined by its right limit, which means f p (0) = lim x→0 + f p (x) = 1.Proposition 1: Given the piece-wise f -divergence in eq (30) , (f ′ p ) -1 (x) and f * p (x) has the following formulation:EQUATIONProof. We begin with the proof of (f ′ p ) -1 (x). Given Lemma 2, the form of (f ′ p ) -1 (x) directly comes from the continuity and monotonic increasing property of f ′ p (x). One can easily verify (f ′ p ) -1 (x)'s form by calculating the reverse function of f ′ p (x). We then continue with the proof of f * p (x). By the definition of Fenchel conjugate we have:EQUATIONBecause here f p (x) is a piecewise function, the supremum could be achieved in either [0, 1) or [1, +∞) . This means we should first take supremum over two intervals respectively and then choose the greater value as f * p (x). More specifically, we perform the following decomposition:EQUATIONWhen constrained to either [0, 1) or [1, +∞), f p (z) is no longer piecewise, and the inner supremum becomes the Fenchel conjugate, albeit under a narrower interval. To calculate Fenchel conjugate under a limited interval, one should observe that for any given x, the derivative of the inner objective satisfies:EQUATIONBy Lemma 2, we have that f ′ p (z) is strictly increasing. This meansd(x•z-fp(z)) dzhas only three possible behaviors on any interval: always positive, always negative, or transitioning from positive to negative with a zero crossing within the interval. Then for any interval[a, b], sup z∈[a,b] x • z -f p (z)has the following formulation. Note that similar results can be extended to open intervals.EQUATIONLeveraging this formulation, one can easily derive the close form solution of 34. In fact, f * p (x) has the following formulation:EQUATIONThe final form of f * p (x) could be derived by directly taking the maximum.Difference between IGL and other inexact guidance methods One important property of IGL is that it can induce the exact optimal policy, rather than a policy similar to the optimal policy, if the optimal policy can be represented as a weighted behavior policy. Some of the previous methods use inexact guidance terms and thus have no guarantee of the induced policy [17, 16, 6 ]. Among them, Janner et al. [17] uses a mean-square-error (MSE) objective to train the guidance model. Ho et al. [16] , Chung et al. [6] leverage training-free guidance by reusing the pre-trained diffusion model in the data prediction formulation to characterize intermediate guidance.Besides the inexact guidance methods mentioned above, Contrastive Energy Prediction (CEP) [32] , is a guidance learning method that could induce exactly the desired policy if it can be represented as a weighted behavior policy. However, CEP is still prone to exploit overestimation error in the critic. We provide additional explanations for the differences between Contrastive Energy Prediction and In-sample Guidance Learning, as well as the reasons for these differences. Intuitively, the difference mainly stems from CEP's inherent property as a contrastive loss.In order to obtain ∇ at log E p(a0|at,s) [w(s, a 0 )], CEP will first train a model g θ (s, a t , t) to fit log E p(a0|at,s) [w(s, a 0 )] and then take the gradient of it. Its optimization objective has the following contrastive form:EQUATIONNote that here K must satisfy K > 1 to ensure the optimal condition is g θ * (s, a t , t) = ∇ at log E p(a0|at,s) [w(s, a 0 )]. This means under offline setting, additional fake actions are required on every state in the dataset. CEP first uses diffusion model to pre-train a behavior policy πD from the dataset and then generates multiple fake actions (a 0 ) (1:K) ∼ πD (a 0 |s) for each state in the dataset. The second expectation in the objective will be calculated using these fake actions. It's clear that when πD mistakenly generates OOD actions, CEP still uses action evaluation model to predict w(s, a OOD 0 ) on these OOD actions. The introduction of w(s, a OOD 0 ) can make ∇ at log E p(a0|at,s) [w(s, a 0 )] prefer these a OOD 0 . This is because ∇ at log E p(a0|at,s) [w(s, a 0 )] will guide the samples towards actions with high w(s, a 0 ) value. These high w(s, a 0 ) value actions can be OOD actions whose w(s, a 0 ) is overestimated. The catastrophic result is that CEP may guide action samples towards OOD actions with erroneously high w(s, a 0 ) value rather than real high-quality actions.Finally and most importantly, we will explain from a probabilistic perspective why the contrastive loss in CEP can't be estimated with one sample from π D (a|s). The derivation of the optimality condition also depends on changing the conditional variables from (a 0 ) (1:K) to (a t ) (1:K) :E P (t) E π D (a0) (1:K) |s E p (at) (1:K) |(a0) (1:K) ,s - K i=1 w s, (a 0 ) (i) log e g θ s,(at) (i) ,t K j=1 e g θ s,(at) (j) ,t = E P (t) E p((a0) (1:K) ,(at) (1:K) |s) - K i=1w s, (a 0 ) (i) log e g θ s,(at) (i) ,t K j=1 e g θ s,(at) (j) ,tEQUATIONTo derive the optimality condition that g θ s, (a t) (i) , t = log E p (a0) (i) |(at) (i) ,s [w s, (a 0 ) (i) ]for each i, one have to decompose the joint probability p (a 0 ) (1:K) |(a t ) (1:K) , s into the product of each (a 0 ) i 's conditional probability:= E P (t) E p((at) (1:K) |s) E Π K i=1 p (a0) (i) |(at) (1:K) ,s - K i=1w s, (a 0 ) (i) log e g θ s,(at) (i) ,t K j=1 e g θ s,(at) (j) ,t= E P (t) E p((at) (1:K) |s) - K i=1 E p (a0) (i) |(at) (i) ,sw s, (a 0 ) (i) log e g θ s,(at) (i) ,t K j=1 e g θ s,(at) (j) ,tNote that here we absorb the expectation over Π K i=1 p (a 0 ) (i) |(a t ) (1:K) , s into the sum over i and replace it with p (a 0 ) (i) |(a t ) (i) , s . This is because each reverse process p (a 0 ) (i) |(a t ) (i) , s is independent of other (a t ) (j̸ =i) .The requirement of decomposing joint probability p (a 0 ) (1:K) |(a t ) (1:K) , s means that (a 0 ) (1:K) should be mutually independent from each other, once (a t ) (1:K) and s are determined. Obviously, if we only have one sample of a 0 , (a 0 ) (1:K) must be identical and not independent. On the other hand, because the objective of In-sample Guidance Learning doesn't contain expectation over any joint distribution, it can be estimated with only one sample from π D (a|s).Additional details for the practical algorithm In this part, we further provide the implementation details based on the given pseudo-code. Firstly, the score network ϵ θ we used is based on a U-net architecture, which is fairly common in diffusion-based RL algorithms [1, 59, 32] . For the DICE part, we use the double Q-learning trick [10] for the training of Q. For sampling, we use DPM-solver [31] to accelerate the sampling of the diffusion model, while the score function is given by IGL. The implementation of our score model and sampling process are based on DPM-solver, which uses MIT license. Diffusion-DICE contains 2 hyper-parameters in total, which are α in the DICE part and K during sampling.Toy case experimental details As shown in Fig. 2 , the behavior policy follows a standard Gaussian distribution constrained in an annular region. We sample 1500 times from the ground truth behavior policy to construct the offline dataset. The reward model R is represented with a 3-layer MLP and because we don't have to model complicated policy distribution, the score network ϵ θ in the toy case is simply multi-layer MLP. For better visualization, we sample 500 candidates during the guide stage. For QGPO, because it doesn't have the select stage, the sample distribution remains the same. For the other 2 methods, we select actions from candidates generated either from diffusion behavior policy (IDQL), or from guided policy (Diffusion-DICE) for 64 times. The reverse diffusion trajectories contain 16 action outputs. Because these 3 methods stand for different paradigms, we tuned different hyper-parameter sets for all of them to achieve the best performance.D4RL experimental details For all tasks, we ran Diffusion-DICE for 10 6 steps and reported the final performance. In MuJoCo locomotion tasks, we computed the average mean returns over 10 evaluations across 5 different seeds. For AntMaze tasks, we calculated the average over 50 evaluations, also across 5 seeds. Because the inference of diffusion model is relatively slow, we conduct the evaluation every 4 • 10 4 training steps. Following previous research, we standardized the returns by dividing the difference in returns between the best and worst trajectories in MuJoCo tasks. In AntMaze tasks, we subtracted 1 from the rewards following previous methods [23, 34] . Note that for the D4RL benchmark and its corresponding datasets, Apache-2.0 license is used.For the network, we use a 3-layer MLP with 256 hidden units to represent Q and V . For the guidance network g θ , we use a slightly more complicated 4-layer MLP with 256 hidden units. Note that the time embedding in the guidance network is the Gaussian Fourier projection. We use Adam optimizer [21] to update all the networks, with a learning rate of 3 • 10 -4 . The target network of Q used for double Q-learning trick is soft updated with a weight of 5 • 10 -3 .For the baseline algorithms, we report the performance of IDQL under any number of hyperparameters. Other baseline methods' results are sourced from their original papers. For all tasks, we list the chosen hyper-parameters in Table 2 . As mentioned before, Diffusion-DICE requires significantly fewer action candidates during the select stage and we present the comparison of K between Diffusion-DICE and IDQL in Table 3 .To further demonstrate the superior performance of Diffusion-DICE compared to traditional DICE algorithms, we conduct ablation studies by replacing the policy in Diffusion-DICE with a simple Gaussian policy, while preserving the piecewise f -divergence. The results are given in Figure 4 . We also perform hyper-parameter study for candidate number K. To compare the performance of Diffusion-DICE under different K, we select some of the datasets in D4RL benchmark and test its performance following the same protocol before. Because the optimal K for MuJoCo locomotion tasks and AntMaze navigation tasks have significant differences in magnitude, we chose to test 5 neighboring orders of magnitude around the optimal K value. The results are given in Table 4 .Comparison of average Q(s, a) curves For a detailed comparison between the guide-then-select paradigm and the select-from-behavior paradigm regarding error exploitation, we present their average Q(s, a) curves normalized returns in Fig. 5 . It's evident from the results that the average Q(s, a) curves of the select-from-behavior paradigm almost remain above the guide-then-select paradigm, while the latter exhibits better performance.Learning curves on D4RL benchmark We present the learning curves on D4RL benchmark in Figure 6 and Figure 7 . Besides the positive social impacts that Diffusion-DICE can help solve various practical offline RL tasks, for example, robotics, health care, industrial control tasks, some potential negative impacts also exists. Part of the Diffusion-DICE algorithm, the In-sample Guidance Learning, can also be applied to other generative tasks like image generation. This could be used to generate fake images, which could possibly mislead the public after being spread online. 


Decision: The reviewers found the proposed approach novel and interesting. The method was found theoretically sound and the experimental results convincing.  Questions mainly concerned the clarification of the presentation of the mehtod, and the need for better positioning with respect to existing offline RL methods. 

The rebuttal focused on clarifying the contributions of this work with respect to related approaches, and provided aditional results on the kitchen and adroit environments, demonstrating the benefit of Diffusion-DICE as compared to the chosen panel of baselines. These additional clarifications were appreciated by the reviewers during the discussion phase, and a consensus was reached towards the acceptance of this work to NeurIPS 2024.

Reviews: review: paper_summary: paper_summary: This paper introduces Diffusion-DICE for offline reinforcement learning. Diffusion-DICE motivates from the transformation between the behaviour distribution and the optimal distribution, which inspires the use of generative models for behaviour distribution modelling. Next, Diffusion-DICE decomposes the policy score function into two components, one from the behaviour distribution, another from the guidance, i.e., the transformation. Lastly, Diffusion-DICE employs a guide-then-select paradigm, which uses only in-sample actions for training to avoid out-of-distribution issues. In their experiments, Diffusion-DICE has achieved strong performance compared with baselines on the D4RL benchmark.
strength_weakness: strength_weakness: There are several weaknesses of the paper I’d like to point out.

1/ The biggest issue is the presentation. Although I do like the idea of the work and recognise its contributions, I found the paper very hard to follow and needed to read through the paper to understand the introduction. The way authors presented the guide-then-select is confusing. I’d suggest authors provide a bit more background and carefully define the “guidance term”, and how it relates to the RL before using it in both abstract and introduction. 

2/ Achieving in-sample learning of offline diffusion RL is not new. Efficient Diffusion Policy (EDP) [1] has introduced an IQL-based variant which naturally allows training Q-values using only in-sample data, without querying out-of-distribution actions during policy evaluation. I’d suggest the authors carefully check the claims and avoid over claiming.

3/ The D4RL experiments are only conducted on the locomotion tasks and the antmaze tasks. The commonly tested kitchen and adroit tasks are missing, which weakens the claim of the paper.

References:

[1] Kang, B., Ma, X., Du, C., Pang, T., & Yan, S. (2024). Efficient diffusion policies for offline reinforcement learning. Advances in Neural Information Processing Systems, 36.
questions: questions: Could you provide a bit more discussion about the differences of using guidance for inference with the Diffusion-QL style inference, which directly guides the sampling process towards actions with high returns? It would be interesting to understand the pros and cons of these two different paradigms
limitations: limitations: I think in general this is an interesting work and I don’t see major limitations.
paper_summary: paper_summary: The paper introduces a novel offline reinforcement learning approach that leverages diffusion models integrated with DICE-based methods. The proposed guide-then-select paradigm aims to minimize error exploitation. The resultant algorithm achieves state-of-the-art performance on D4RL benchmark tasks.
strength_weakness: strength_weakness: - Missing comparison in Table 1 of more optimal Gaussian policy methods, e.g. EDAC [1]
- Selection of D4RL datasets is limited, e.g. what about expert/random datasets? Further environments like Adroit would also be interesting
- Discussion of hyperparameter choice in the Appendix is important and should be included in the experimental section.
- A comparison of the inference speed of Diffusion-DICE and prior baselines would be valuable.

Minor:
- Line 70: “M” -> “M=”
- Line 73: LP abbreviation not explained
- There is concurrent related work [2] which also performs an analogous transformation between the behavior distribution to an online policy with diffusion models for synthetic data generation.

[1] Uncertainty-Based Offline Reinforcement Learning with Diversified Q-Ensemble. Gaon An, Seungyong Moon, Jang-Hyun Kim, Hyun Oh Song. NeurIPS, 2021.

[2] Policy-Guided Diffusion. Matthew Thomas Jackson, Michael Tryfan Matthews, Cong Lu, Benjamin Ellis, Shimon Whiteson, Jakob Foerster. RL Conference, 2024.
questions: questions: Please see the above weaknesses.
limitations: limitations: Limitations are discussed in the conclusion.
paper_summary: paper_summary: The paper mainly addresses the problem of obtaining the optimal policies in the distribution correction estimation (DICE) setting, which is one popular offline RL approach. Note that the offline RL assumes that we can’t interact with an environment and only have access to a dataset—sets of (state, action, and next state) tuples collected by some behavior policy. For the given history, DICE estimates the optimal value functions $V^*$ and $Q^*$ and obtains the optimal stationary distribution ratio $w^*$, which is the rate between the optimal policy’s state occupancy distribution $d^*$ and the one of the behavior policy $d^D$. Even if it is possible to get the optimal value functions and the ratio via DICE, obtaining the corresponding optimal policy $\pi^*$ from them is still challenging.

In this context, the paper proposes a novel training method for diffusion-based models to learn the optimal policy from the DICE-driven optimal value functions and the ratio. Unlike common density modeling, learning the diffusion-based models from the optimal value functions is not straightforward. Such a problem is closer to variational inference than density modeling; specifically, we don’t have samples from the target distribution to diffuse. Moreover, unlike a typical variational inference problem, sampling from the models is not available due to the nature of the offline RLs.

To do that the paper, the paper first introduces that the optimal policy can be represented by a product of the optimal stationary distribution ratio $w^*$ and the behavior policy $\pi^D$. Since the product of two distributions is unnormalized, the complete expression includes the normalization constant, which is the expected value of $w^*$ under the behavior policy (Line 155).

Next, the paper assumes that the behavior policy can be represented by a diffusion-based model with the forward process, which will also be used to perturb the policy of interest. Then, the authors show that the optimal policy at the perturbation level $t$ can be represented by using the behavior’s perturbed distribution at the same $t$ and the DICE ratio $w^*$ (Equation 6). In particular, this representation includes the logarithm of the expectation of the DICE ratio $w^*$ under the posterior distribution of the clean action $a_0$ for a given perturbed action $a_t$. While this novel representation doesn’t require sampling from the model policy anymore, computing the logarithm of the Monte Carlo estimate of any expectation is biased in general.

To circumvent this, the authors propose using the tangent transform, one of the variational approximation methods; thus, the quantity inside the logarithm is represented by an optimization problem, as in Equation 7. Interestingly, the new training objective only requires sampling from the behavior policy, which is suitable for the offline RLs. 

In summary, the paper introduces a new representation of the diffusion-based optimal policy model by using the diffusion-based behavior policy and time-dependent learnable terms that will be learned by the convex problem described in Equation 7. The authors refer to this approach as In-sample Guidance Learning (IGL). While there have been a few previous approaches to learning diffusion-based optimal policy models, the authors point out that the IGL shows some benefits. For example, IGL only requires a single sample from the behavior policy, which can be obtained from the history, while some previous approaches require more than one sample from the behavior, which is not favorable in the offline RLs. 

In addition, the authors suggest a few techniques to improve the stability of IGLs, such as using piecewise $f$-divergence for the DICE regularization term.

Finally, the paper discusses some failure cases of previous approaches and how the proposed method bypasses them. It also demonstrates the efficacy of the proposed method via experiments on several benchmark datasets.


-----------------------------------------------------
Updated the rating from 7 to 8 after the authors' rebuttal.
strength_weakness: strength_weakness: Overall, the paper presents a novel method for learning optimal policy in DICE, which is a valuable contribution to the field. However, improvements in the presentation would greatly enhance the clarity and comprehensibility of the manuscript.

In particular, several equations within the paper omit the definitions of variables, which can lead to confusion—for example, $a_0$ in Line 119. In addition, some variables overlap while they are independent. For instance, in Line 155, the variable of the integration overlaps with the $a$ in the numerator.

I recommend revising the paper to address these issues.
questions: questions: N/A
limitations: limitations: N/A

Comments: title: Kind Reminder: Discussion Period Ending
comment: Thank you for your clarifications, I will raise my score.
Dear Reviewer uMN3,

We sincerely apologize for any inconvenience this reminder may cause. We just wanted to kindly remind you that the discussion period will conclude tomorrow.

As the discussion period is nearing its end, we wanted to check if you have any remaining questions or concerns. **We would be more than happy to address further inquiries you may have.** We understand how busy you must be during this time, and we truly appreciate the effort and time you've dedicated to the rebuttal process.

Thank you very much, and we look forward to your response.

Best regards,

The Authors of Paper 16567
I do appreciate the authors' efforts and detailed explanations. I feel most of my concerns are addressed. I still feel this is a good paper, although certain efforts are still needed for a better presentation. I'll keep my original score and vote for an acceptance.