Title: Diffusion-DICE: In-Sample Diffusion Guidance for Offline Reinforcement Learning
Venue: NeurIPS 2024
Accepted: yes
Keywords: Offline Reinforcement Learning

Diffusion Models
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.

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

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 localoptimal modes. We thus generate a few candidate actions and carefully select from them to approach global-optimum. Different from all other diffusion-based offline RL methods, the 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 avoids that. We then conduct extensive experiments on benchmark datasets to show the strong performance of Diffusion-DICE.

1 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

Truevaluefunction—Learnedvaluefunction

Datasetaction

xDiffusedaction

guide

select

guide

select

guide

guide orselectinpreviousmethods

guide-then-selectinDiffusion-DICE

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.

2 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 Pr ( 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 .

2.1 Distribution Correction Estimation

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:

This 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:

Here 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:

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 ) as

where 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.

2.2 Diffusion Models in Offline Reinforcement Learning

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 t 0 ( 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 ¯ w t is a standard Wiener process in the reverse time.

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 t 0 ( x t | x 0 ) and p 0 t ( 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 ∇ x t 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 ) ,x 0 ∼ q 0 ( x 0 ) ,ϵ ∼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 ∇ a t log π t ( a t | s ) = ∇ a t log π D t ( a t | s ) + ∇ a t 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 j

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.

3 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.

3.1 An Optimal Policy Transformation View of DICE

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:

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:

The proof is given in Appendix B. Intuitively, ∇ a t log π D t ( a t | s ) tells us how to generate actions from π D and ∇ a t log E π D ( a 0 | a t ,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 ∇ a t 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.

3.2 In-sample Guidance Learning for Accurate Policy Transformation

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 )] ,

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. ∇ a t log E π D ( a 0 | a t ,s ) [ w ∗ ( s, a 0 )] can be obtained by solving the following optimization problem:

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 ∇ a t log E π D ( a 0 | a t ,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:

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:

In 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:

We 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.

3.3 Boost Performance with Optimal Action Selection

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 ∇ a t 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 ∇ a t 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 as

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.

Algorithm 1 Diffusion-DICE

1: Initialize value function Q ϕ 1 , V ϕ 2 , diffusion behavior model ϵ ψ , guidance network g θ

2: // Training

3: Pretrain the diffusion behavior model by minimizing E U ( t ) , ( s,a ) ∼D ,ϵ ∼N ( 0 , I ) [ ∥ ϵ ψ ( α t a + σ 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 } i K =1 from π ∗ following Eq.(5), with ∇ a t log π ∗ t ( a t | s ) = ϵ ψ ( a t ,s,t ) -σ t + ∇ a t g θ ( s, a t , t ) 11: Select action according to Eq.(10)

We term this guide-then-select method Diffusion-DICE and present its pseudo-code in Algorithm 1. During the training stage, DiffusionDICE 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,

DatasetActionDistribution

DiffusionBehaviorPolicy

Diffusion-DICE

QGPO

IDQL

1.5

1.5

(guide-then-select)(guide-only)

(select-only)

1.0

0.25

1.0

0.200

0.175

1.5

1

1.0

0.5

0.20

0.5

0.150

0.5

0.5

Guide

0.5

0.0

0.15

0.10

0.5

0.0

0.125

0.100

0.075

-0.5

-1

0.0

0.5

1.0

0.05

1.0

0.050

0.025

1.5

-1.5 -1 -0.5 0 0.5 1 1.5 -1.5 -1 -0.5 0 0.5 1 1.5 -1.5 -1 -0.5 0 0.5 1 1.5

1.0

1.0

1.5

1.5

1.0

0.5

0.0

0.5

1.0

1.5

0.00

1.5

1.5

1.0

-0.5

0.0

0.5

1.0

1.5

0.5

1

0.5

Select

TrueRewardFunction

LearnedRewardFunction

0

0.0

1.5

1.5

-0.5

0.5

1.0

0.75

1.0

1

-1

0.5

0.25

0.50

0.5

0

1.5

-1.5 -1 -0.5 0 0.5 1 1.5 -1.5 -1 -0.5 0 0.5 1 1.5 -1.5 -1 -0.5 0 0.5 1 1.5

1.0

1.0

Diffusion Trajectory

0.0

0.00

0.0

1

1

0.5

0.5

0.25

0.5

2

0.5

0

0.0

1.0

0.50

0.75

1.0

3

-0.5

-1

6

0.5

1.5

1.5

-1.00.50.0

1.0

1.5

1.00

1.5

1.5

1.00.5

0:0

0.5

1.0

1.5

ST I S0 0 S0-I-ST-S'T  S0 0 S0-I-ST-ST T S0 0 S0-T-S'T-

1.0

Diffusion-DICE selects the action from π ∗ ( a | s ) with the highest value. By selecting from a small number of action candidates from π ∗ ( a | s ) , minimal error of the action evaluation model will be exploited. Besides serving as a diffusion-based offline RL algorithm, Diffusion-DICE also introduces a novel policy extraction method from other DICE-based algorithms. By leveraging expressive generative models, Diffusion-DICE can capture the multi-modality in π ∗ and better draw upon the knowledge from d ∗ ( s,a ) d D ( s,a ) and π D .

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.

4 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.

4.1 D4RL Benchmark Datasets:

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.

4.2 Further Experiments on Error Exploitation

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.

5 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.

DICE-based methods 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.

6 Conclusion

In this work, we propose a new diffusion-based offline RL algorithm, Diffusion-DICE. DiffusionDICE 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.

References

A. Ajay, Y. Du, A. Gupta, J. Tenenbaum, T. Jaakkola, and P. Agrawal. Is conditional generative modeling all you need for decision-making? arXiv preprint arXiv:2211.15657 , 2022.

G. An, S. Moon, J.-H. Kim, and H. O. Song. Uncertainty-based offline reinforcement learning with diversified q-ensemble. Proc. of NeurIPS , 2021.

S. Boyd, S. P. Boyd, and L. Vandenberghe. Convex optimization . Cambridge university press, 2004.

H. Chen, C. Lu, C. Ying, H. Su, and J. Zhu. Offline reinforcement learning via high-fidelity generative behavior modeling. arXiv preprint arXiv:2209.14548 , 2022.

L. Chen, K. Lu, A. Rajeswaran, K. Lee, A. Grover, M. Laskin, P. Abbeel, A. Srinivas, and I. Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Proc. of NeurIPS , 2021.

H. Chung, J. Kim, M. T. Mccann, M. L. Klasky, and J. C. Ye. Diffusion posterior sampling for general noisy inverse problems. arXiv preprint arXiv:2209.14687 , 2022.

Z. Ding and C. Jin. Consistency models as a rich and efficient policy class for reinforcement learning. arXiv preprint arXiv:2309.16984 , 2023.

J. Fu, A. Kumar, O. Nachum, G. Tucker, and S. Levine. D4rl: Datasets for deep data-driven reinforcement learning. ArXiv preprint , 2020.

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

S. Fujimoto, H. van Hoof, and D. Meger. Addressing function approximation error in actor-critic methods. In Proc. of ICML , pages 1582-1591, 2018.

S. Fujimoto, D. Meger, and D. Precup. Off-policy deep reinforcement learning without exploration. In Proc. of ICML , pages 2052-2062, 2019.

D. Garg, S. Chakraborty, C. Cundy, J. Song, and S. Ermon. Iq-learn: Inverse soft-q learning for imitation. Advances in Neural Information Processing Systems , 34:4028-4039, 2021.

T. Haarnoja, A. Zhou, P. Abbeel, and S. 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.

P. Hansen-Estruch, I. Kostrikov, M. Janner, J. G. Kuba, and S. Levine. Idql: Implicit q-learning as an actor-critic method with diffusion policies. arXiv preprint arXiv:2304.10573 , 2023.

J. Ho, A. Jain, and P. Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems , 33:6840-6851, 2020.

J. Ho, T. Salimans, A. Gritsenko, W. Chan, M. Norouzi, and D. J. Fleet. Video diffusion models. Advances in Neural Information Processing Systems , 35:8633-8646, 2022.

M. Janner, Y. Du, J. B. Tenenbaum, and S. Levine. Planning with diffusion for flexible behavior synthesis. arXiv preprint arXiv:2205.09991 , 2022.

D. Kalashnikov, J. Varley, Y. Chebotar, B. Swanson, R. Jonschkowski, C. Finn, S. Levine, and K. Hausman. Mt-opt: Continuous multi-task robotic reinforcement learning at scale. ArXiv preprint , 2021.

G.-H. Kim, S. Seo, J. Lee, W. Jeon, H. Hwang, H. Yang, and K.-E. Kim. Demodice: Offline imitation learning with supplementary imperfect demonstrations. In Proc. of ICLR , 2021.

D. Kingma, T. Salimans, B. Poole, and J. Ho. Variational diffusion models. Advances in neural information processing systems , 34:21696-21707, 2021.

D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. In Proc. of ICLR , 2015.

I. Kostrikov, O. Nachum, and J. Tompson. Imitation learning via off-policy distribution matching. In Proc. of ICLR , 2020.

I. Kostrikov, A. Nair, and S. Levine. Offline reinforcement learning with implicit q-learning. ArXiv preprint , 2021.

I. Kostrikov, A. Nair, and S. Levine. Offline reinforcement learning with implicit q-learning. arXiv preprint arXiv:2110.06169 , 2021.

A. Kumar, J. Fu, M. Soh, G. Tucker, and S. Levine. Stabilizing off-policy q-learning via bootstrapping error reduction. Advances in neural information processing systems , 32, 2019.

A. Kumar, A. Zhou, G. Tucker, and S. Levine. Conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems , 33:1179-1191, 2020.

J. Lee, W. Jeon, B. Lee, J. Pineau, and K.-E. Kim. Optidice: Offline policy optimization via stationary distribution correction estimation. In International Conference on Machine Learning , pages 6120-6130. PMLR, 2021.

J. Lee, C. Paduraru, D. J. Mankowitz, N. Heess, D. Precup, K.-E. Kim, and A. Guez. Coptidice: Offline constrained reinforcement learning via stationary distribution correction estimation. arXiv preprint arXiv:2204.08957 , 2022.

J. Li, X. Zhan, H. Xu, X. Zhu, J. Liu, and Y.-Q. Zhang. When data geometry meets deep function: Generalizing offline reinforcement learning. In The Eleventh International Conference on Learning Representations , 2022.

J. Li, X. Hu, H. Xu, J. Liu, X. Zhan, Q.-S. Jia, and Y.-Q. Zhang. Mind the gap: Offline policy optimization for imperfect rewards. arXiv preprint arXiv:2302.01667 , 2023.

C. Lu, Y. Zhou, F. Bao, J. Chen, C. Li, and J. Zhu. Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. Advances in Neural Information Processing Systems , 35:5775-5787, 2022.

C. Lu, H. Chen, J. Chen, H. Su, C. Li, and J. Zhu. Contrastive energy prediction for exact energy-guided diffusion sampling in offline reinforcement learning. In International Conference on Machine Learning , pages 22825-22855. PMLR, 2023.

A. S. Manne. Linear programming and sequential decisions. Management Science , 6(3): 259-267, 1960.

L. Mao, H. Xu, W. Zhang, and X. Zhan. Odice: Revealing the mystery of distribution correction estimation via orthogonal-gradient update. arXiv preprint arXiv:2402.00348 , 2024.

O. Nachum and B. Dai. Reinforcement learning via fenchel-rockafellar duality. arXiv preprint arXiv:2001.01866 , 2020.

O. Nachum, Y. Chow, B. Dai, and L. Li. Dualdice: Behavior-agnostic estimation of discounted stationary distribution corrections. Advances in neural information processing systems , 32, 2019.

O. Nachum, B. Dai, I. Kostrikov, Y. Chow, L. Li, and D. Schuurmans. Algaedice: Policy gradient from arbitrary experience. arXiv preprint arXiv:1912.02074 , 2019.

A. Nair, M. Dalal, A. Gupta, and S. Levine. Accelerating online reinforcement learning with offline datasets. ArXiv preprint , 2020.

A. Rényi. On measures of entropy and information. In Proceedings of the fourth Berkeley symposium on mathematical statistics and probability, volume 1: contributions to the theory of statistics , volume 4, pages 547-562. University of California Press, 1961.

A. Shocher, A. Dravid, Y. Gandelsman, I. Mosseri, M. Rubinstein, and A. A. Efros. Idempotent generative network. arXiv preprint arXiv:2311.01462 , 2023.

H. Sikchi, Q. Zheng, A. Zhang, and S. Niekum. Dual rl: Unification and new methods for reinforcement and imitation learning. arXiv preprint arXiv:2302.08560 , 2023.

J. Sohl-Dickstein, E. Weiss, N. Maheswaranathan, and S. Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learning , pages 2256-2265. PMLR, 2015.

Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456 , 2020.

Y. Song, P. Dhariwal, M. Chen, and I. Sutskever. Consistency models. arXiv preprint arXiv:2303.01469 , 2023.

R. S. Sutton, A. G. Barto, et al. Introduction to reinforcement learning . MIT press Cambridge, 1998.

P. Vincent. A connection between score matching and denoising autoencoders. Neural Comput. , 23(7):1661-1674, 2011. doi: 10.1162/NECO\_A\_00142. URL https://doi.org/10.1162/ NECO_a_00142 .

X. Wang, H. Xu, Y. Zheng, and X. Zhan. Offline multi-agent reinforcement learning with implicit global-to-local value regularization. In Thirty-seventh Conference on Neural Information Processing Systems , 2023.

Z. Wang, A. Novikov, K. Zolna, J. Merel, J. T. Springenberg, S. E. Reed, B. Shahriari, N. Y. Siegel, Ç. Gülçehre, N. Heess, and N. de Freitas. Critic regularized regression. In Proc. of NeurIPS , 2020.

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

Y. Wu, G. Tucker, and O. Nachum. Behavior regularized offline reinforcement learning. ArXiv preprint , 2019.

H. Xu, X. Zhan, J. Li, and H. Yin. Offline reinforcement learning with soft behavior regularization. ArXiv preprint , 2021.

H. Xu, L. Jiang, L. Jianxiong, and X. Zhan. A policy-guided imitation approach for offline reinforcement learning. In Advances in Neural Information Processing Systems , volume 35, pages 4085-4098, 2022.

H. Xu, X. Zhan, H. Yin, and H. Qin. Discriminator-weighted offline imitation learning from suboptimal demonstrations. In International Conference on Machine Learning , pages 2472524742. PMLR, 2022.

H. Xu, X. Zhan, and X. Zhu. Constraints penalized q-learning for safe offline reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 36, pages 8753-8760, 2022.

H. Xu, L. Jiang, J. Li, Z. Yang, Z. Wang, V. W. K. Chan, and X. Zhan. Offline rl with no ood actions: In-sample learning via implicit value regularization. arXiv preprint arXiv:2303.15810 , 2023.

X. Zhan, H. Xu, Y. Zhang, X. Zhu, H. Yin, and Y. Zheng. Deepthermal: Combustion optimization for thermal power generating units using offline reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 36, pages 4680-4688, 2022.

R. Zhang, B. Dai, L. Li, and D. Schuurmans. Gendice: Generalized offline estimation of stationary values. In International Conference on Learning Representations , 2019.

W. Zhou, S. Bajracharya, and D. Held. Latent action space for offline reinforcement learning. In Conference on Robot Learning , 2020.

Z. Zhu, M. Liu, L. Mao, B. Kang, M. Xu, Y. Yu, S. Ermon, and W. Zhang. Madiff: Offline multi-agent learning with diffusion models. arXiv preprint arXiv:2305.17330 , 2023.

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: 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.

2. Limitations

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

Answer: [Yes]

Justification: Limitations are discussed in 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: 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.

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: 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.

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: Code provided in the project link.

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: Experimental settings are given in Section 4 and Appendix D. Details are given in Appendix D.

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: We use sufficient random seeds and provide standard deviation of the normalized scores in Section 4.

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: Details of computation resources are given in Appendix E.

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: This paper doesn't involve human subjects or participants and the datasets in Section 4 are common open-sourced datasets.

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: [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).

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:

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: 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.

13. New Assets

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

Answer: [NA]

Justification:

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:

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:

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 FenchelRockfeller duality or Lagrangian duality to solve it. The regularization term aims at imposing state-action level constraints.[35, 27, 34].

Here D 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 :

Note 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]:

Here 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:

Here f ∗ is a variant of f 's convex conjugate. Note that in the offline RL setting, the inner summation ∑ s ′ 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:

And because of doing so, the optimization objective in 14 can be replaced with:

Also 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.

B Additional Proofs

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:

Proof. Using the definition between d π ( s ) and d π ( s, a ) , we have for any π , d π ( s, a ) = d π ( s ) · π ( a | s ) . Then we have:

Verification of Eq.(6) : The score function of the optimal policy π ∗ ( a | s ) in DICE and the behavior policy π D ( a | s ) satisfy: ∇ a t log π ∗ t ( a t | s ) = ∇ a t log π D t ( a t | s ) + ∇ a t log E a 0 ∼ π D ( a 0 | a t ,s ) [ d ∗ ( s,a 0 ) d D ( s,a 0 ) ] .

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:

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:

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.

Lemma1. 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 y ∗ = log E x ∼ P ( X ) [ f ( x )] .

Proof. Because y is the constant optimizer, this problem can be reformulated as:

We 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:

Because 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:

Theorem 1. ∇ a t log E π D ( a 0 | a t ,s ) [ w ( s, a 0 )] can be obtained by solving the following optimization problem:

The optimal solution θ ∗ satisfies ∇ a t g θ ∗ ( s, a t , t ) = ∇ a t log E π D ( a 0 | a t ,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:

We first take the first-order derivative with respect to each g θ ( s, a t , t ) for this objective and solve its saddle point:

Because 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:

We then take the second-order derivative with respect to each g θ ( s, a t , t ) for this objective:

Because 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:

Proof. 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:

It'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:

Proof. 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:

Because 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:

When 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:

By Lemma 2, we have that f ′ p ( z ) is strictly increasing. This means d ( x · z -f p ( z )) dz has 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.

Leveraging this formulation, one can easily derive the close form solution of 34. In fact, f ∗ p ( x ) has the following formulation:

∗

The final form of f p ( x ) could be derived by directly taking the maximum.

C In-depth Discussion of In-sample Guidance Learning and Other Guidance Methods

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.

Difference between IGL and Contrastive Energy Prediction (CEP) 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 ∇ a t log E p ( a 0 | a t ,s ) [ w ( s, a 0 )] , CEP will first train a model g θ ( s, a t , t ) to fit log E p ( a 0 | a t ,s ) [ w ( s, a 0 )] and then take the gradient of it. Its optimization objective has the following contrastive form:

Note that here K must satisfy K > 1 to ensure the optimal condition is g θ ∗ ( s, a t , t ) = ∇ a t log E p ( a 0 | a t ,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 ∇ a t log E p ( a 0 | a t ,s ) [ w ( s, a 0 )] prefer these a OOD 0 . This is because ∇ a t log E p ( a 0 | a t ,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 ) :

To derive the optimality condition that g θ ( s, ( a t ) ( i ) , t ) = log E p ( ( a 0 ) ( i ) | ( a t ) ( 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:

̸

Note 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 ) .

D Experimental Details

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.

D4RLResultsfor GaussianPolicyandDiffusionPolicy

Scores

GaussianPolicy

100

DiffusionPolicy

80

Nomralized

60

40

20

0

halfcheetah hopper

m

m

m

m-r

m-r

m-r

m-e

m-e

walker2d

m-e

u

u-d

m-p

m-d

I-p

I-d

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.

walker2d-medium-replay-v2

hopper-medium-replay-v2

90

100

80

06

70

09

80

50

70

40

60

30

guide-then-select Qvalue

guide-then-select Q value

20

guide-then-select score

50

guide-then-select score

select-from-behavior Q value

select-from-behavior Q value

10

select-from-behavior score

40

select-from-behavior score

0

200

400

600

800

1000

0

200

400

600

800

1000

walker2d-medium-v2

hopper-medium-v2

90

100

95

80

06

70

85

09

08

75

50

guide-then-select Q value

guide-then-select Q value

guide-then-select score

OL

guide-then-select score

select-from-behavior Q value

select-from-behavior Q value

40

select-from-behavior score

65

select-from-behavior score

0

200

400

600

800

1000

0

200

400

600

800

1000

hopper-medium-v2

walker2d-medium-v2

100

100

100

Nomalized Return

Nomalized Return

Nomalized Return

75

75

75

50

50

50

25

25

25

0.0

0.2

0.4

0.6

0.8

1.0

0.2

0.4

0.6

0.8

1.0

0.2

0.4

0.6

0.8

1.0

Training Steps (1e6)

Training Steps (1e6)

Training Steps (1e6)

hopper-medium-replay-v2

walker2d-medium-replay-v2

100

100

100

Nomalized Return

Nomalized Return

Nomalized Return

75

75

50

50

50

25

25

25

0.2

0.4

0.6

0.8

1.0

0.2

0.4

0.6

0.8

1.0

0.2

0.4

0.6

0.8

1.0

Training Steps (1e6)

Training Steps (1e6)

Training Steps (1e6)

hopper-medium-expert-v2

walker2d-medium-expert-v2

100

100

100

Nomalized Return

Nomalized Return

Nomalized Return

75

75

75

50

50

50

25

25

25

0

0

0

0.0

0.2

0.4

0.6

0.8

1.0

°0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

Training Steps (1e6)

Training Steps (1e6)

Training Steps (1e6)

antmaze-umaze-v2

antmaze-large-play-v2

100

100

100

Nomalized Return

Nomalized Return

Nomalized Return

75

75

75

50

50

50

25

25

25

0.2

0.4

0.6

0.8

1.0

0

0.0

0.2

0.4

0.6

0.8

1.0

0.2

0.4

0.6

0.8

1.0

Training Steps (1e6)

Training Steps (1e6)

Training Steps (1e6)

antmaze-umaze-diverse-v2

antmaze-large-diverse-v2

100

100

100

Nomalized Return

Nomalized Return

Nomalized Return

75

75

75

50

50

50

25

25

25

0

0

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

Training Steps (1e6)

Training Steps (1e6)

Training Steps (1e6)

E Experiments Compute Resources

For all the experiments, we evaluate Diffusion-DICE on either NVIDIA RTX 3080Ti GPUs or NVIDIA RTX 4090 GPUs. The training process of Diffusion-DICE takes about 6 hours on NVIDIA RTX 3080Ti or 4 hours on NVIDIA RTX 4090. The evaluation process of Diffusion-DICE takes about 3 hours on NVIDIA RTX 3080Ti or 2 hours on NVIDIA RTX 4090. Because we parallelize training and testing process, a complete experiment takes about 8 hours on NVIDIA RTX 3080Ti or 4 hours on NVIDIA RTX 4090.

F Broader Impacts

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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nIn 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.\n\nTo 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).\n\nNext, 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.\n\nTo 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. \n\nIn 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. \n\nIn addition, the authors suggest a few techniques to improve the stability of IGLs, such as using piecewise $f$-divergence for the DICE regularization term.\n\nFinally, 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.\n\n\n-----------------------------------------------------\nUpdated the rating from 7 to 8 after the authors' rebuttal.\n\nSTRENGTHS:\nIn my understanding, the paper's contributions are clear, and I also consider that the results are essential for several reasons.\n\nThe paper introduces a novel representation of the diffusion-based optimal policy model and its training method, IGL. In addition, the paper also motivates the solution well. For example, this approach circumvents the drawbacks of previous approaches, which require more than one sample from the behavior policy, and such actions may be out-of-distribution of the environments.\n\nMoreover, the authors provide extensive discussions to help potential readers comprehend the characteristics of previous approaches and the proposed method.\n\nFinally, the paper demonstrates the effectiveness of the proposed method via various experiments, which further supports the authors' claim.\n\nWEAKNESSES:\nOverall, 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.\n\nIn 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.\n\nI recommend revising the paper to address these issues.\n\nQUESTIONS:\nN/A\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-10"}

comments: 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.

STRENGTHS:
- Well-motivated and novel integration of diffusion models with DICE-based methods
- Well-written theoretical justification for the approach
- Strong empirical results that improve upon prior diffusion-policy baselines

WEAKNESSES:
- 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:
Please see the above weaknesses.

LIMITATIONS:
Limitations are discussed in the conclusion.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-11

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe proposed idea is novel and interesting. I like the way the authors connect DICE with diffusion policies, decompose the score functions, and make the guidance score tractable. Theoretically, they have provided careful analysis and derivations to support the claims. Empirically, they have conducted both toy experiments for intuitive understanding, and demonstrating the strong performance on the D4RL benchmark.\n\nWEAKNESSES:\nThere are several weaknesses of the paper I’d like to point out.\n\n1/ 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. \n\n2/ 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.\n\n3/ 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.\n\nReferences:\n\n[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.\n\nQUESTIONS:\nCould 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\n\nLIMITATIONS:\nI think in general this is an interesting work and I don’t see major limitations.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}