Title: Efficient Diffusion Policies For Offline Reinforcement Learning
Venue: NeurIPS 2023
Accepted: yes
Keywords: Offline Reinforcement Learning

Diffusion Models
Abstract: Offline reinforcement learning (RL) aims to learn optimal policies from offline datasets, where the parameterization of policies is crucial but often overlooked. Recently, Diffsuion-QL significantly boosts the performance of offline RL by representing a policy with a diffusion model, whose success relies on a parametrized Markov Chain with hundreds of steps for sampling. However, Diffusion-QL suffers from two critical limitations. 1) It is computationally inefficient to forward and backward through the whole Markov chain during training. 2) It is incompatible with maximum likelihood-based RL algorithms (e.g., policy gradient methods) as the likelihood of diffusion models is intractable. Therefore, we propose efficient diffusion policy (EDP) to overcome these two challenges. EDP approximately constructs actions from corrupted ones at training to avoid running the sampling chain. We conduct extensive experiments on the D4RL benchmark. The results show that EDP can reduce the diffusion policy training time from 5 days to 5 hours on gym-locomotion tasks. Moreover, we show that EDP is compatible with various offline RL algorithms (TD3, CRR, and IQL) and achieves new state-of-the-art on D4RL by large margins over previous methods.

Efficient Diffusion Policies for Offline Reinforcement Learning

Abstract

Offline reinforcement learning (RL) aims to learn optimal policies from offline datasets, where the parameterization of policies is crucial but often overlooked. Recently, Diffsuion-QL [37] significantly boosts the performance of offline RL by representing a policy with a diffusion model, whose success relies on a parametrized Markov Chain with hundreds of steps for sampling. However, Diffusion-QL suffers from two critical limitations. 1) It is computationally inefficient to forward and backward through the whole Markov chain during training. 2) It is incompatible with maximum likelihood-based RL algorithms ( e.g. , policy gradient methods) as the likelihood of diffusion models is intractable. Therefore, we propose efficient diffusion policy (EDP) to overcome these two challenges. EDP approximately constructs actions from corrupted ones at training to avoid running the sampling chain. We conduct extensive experiments on the D4RL benchmark. The results show that EDP can reduce the diffusion policy training time from 5 days to 5 hours on gym-locomotion tasks. Moreover, we show that EDP is compatible with various offline RL algorithms (TD3, CRR, and IQL) and achieves new state-ofthe-art on D4RL by large margins over previous methods. Our code is available at https://github.com/sail-sg/edp.

1 Introduction

Offline reinforcement learning (RL) is much desired in real-world applications as it can extract knowledge from previous experiences, thus avoiding costly or risky online interactions. Extending online RL algorithms to the offline domain faces the distributional shift [9] problem. Existing methods mainly focus on addressing this issue by constraining a policy to stay close to the data-collecting policy [6, 39], making conservative updates for Q-networks [17, 15, 40], or combining these two strategies [21, 12]. However, Offline RL can also be viewed as a stateconditional generative modeling problem of actions, where the parameterization of the policy network is important but largely overlooked. Most offline RL works follow the convention of

D-QL   EDP

Locomotion

Kitchen

Adroit

Antmaze

77.6

53.3

56.9

63

85.5

59.4

71.3

73.4

Normalized Score

Pre-SOTA

Ours

parameterizing the policy as a diagonal Gaussian distribution with the learned mean and variance. This scheme might become inferior when the data distribution is complex, especially when offline

data are collected from various sources and present strong multi-modalities [32]. Therefore, more expressive models for the policy are strongly desired.

Recently, Diffusion-QL [37] made a successful attempt by replacing the diagonal Gaussian policy with a diffusion model, significantly boosting the performance of the TD3+BC [6] algorithm. Diffusion models [34, 10] have achieved the new state-of-the-art (SOTA) in image generation tasks [26, 5], demonstrating a superior ability to capture complex data distributions.

Despite the impressive improvement that Diffusion-QL has achieved, it has two critical drawbacks preventing it from practical applications. First , training a diffusion policy with offline RL is computationally inefficient. Consider a parameterized diffusion policy π θ ( a | s ) , Diffusion-QL optimizes it by maximizing the Q value Q ( s , a θ ) of a state s given a policy-generated action a θ ∼ π θ ( a | s ) . However, sampling from a diffusion model relies on a long parameterized Markov chain ( e.g. , 1,000 steps), whose forward inference and gradient backpropagation are unaffordably expensive. Second , diffusion policy is not a generic policy class as it is restricted to TD3-style algorithms. As computing the sample likelihood π θ ( a | s ) is intractable in diffusion models [35], diffusion policy is incompatible with a large family of policy gradient algorithms ( e.g. , V-Trace [22], AWR [28], IQL [15]), which require a tractable and differentiable log-likelihood log π θ ( a | s ) for policy improvement.

In this work, we propose efficient diffusion policy (EDP) to address the above two limitations of diffusion policies. Specifically, we base EDP on the denoising diffusion probabilistic model (DDPM) [10], which learns a noise-prediction network to predict the noise used to corrupt an example. In the forward diffusion process, a corrupted sample follows a predefined Gaussian distribution when the clean example and timestep are given. In turn, given a corrupted sample and predicted noise, we can approximate its clean version by leveraging the reparametrization trick. Based on this observation, to avoid the tedious sampling process, we propose action approximation to build an action from a corrupted one, which can be easily constructed from the dataset. In this way, each training step only needs to pass through the noise-prediction network once, thus substantially reducing the training time. As experimented, by simply adding action approximation, we obtain 2x speed-up without performance loss. Moreover, we apply DPM-Solver [20], a faster ODE-based sampler, to further accelerate both the training and sampling process. Finally, to support likelihood-based RL algorithms, we leverage the evidence lower bound for the likelihood developed in DDPM and approximate the policy likelihood from a constructed Gaussian distribution with variance fixed and mean obtained from action approximation.

We evaluate the efficiency and generality of our method on the popular D4RL benchmarking, as shown in Fig. 1. We first benchmark the efficiency of EDP on gym-locomotion tasks. By replacing the diffusion policy in Diffusion-QL with our EDP, the training time of Diffusion-QL is reduced substantially from five days to five hours (compared to their official code). Meanwhile, we observe slight to clear performance improvements on different tasks as the improved efficiency enables training DDPM with more timesteps than before. Moreover, we plug EDP into three different offline algorithms (including TD3+BC, CRR, and IQL), and the results justify its superiority over standard diagonal Gaussian policies. As a result, EDP set up new state-of-the-art on all four domains in D4RL.

2 Related Work

Offline RL Distributional shift between the learned and behavior policies is offline RL's biggest challenge. Existing research mitigates this problem by making modifications to policy evaluation [24, 17, 15, 21] or policy improvement [9, 39, 6, 33, 38, 42, 41]. For example, conservative Q-learning (CQL) [17] penalizes out-of-distribution actions for having higher Q-values, proving that this is equivalent to optimizing a lower bound of Q-values. Onestep RL [1] conducts policy evaluation on in-distribution data to avoid querying unseen actions. IQL [15] introduces expectile regression [14] to approximate dynamic programming with the Bellman optimality function. TD3+BC explicitly constrains the learned policy by adding a behavior cloning loss to mimic the behavior policy. Instead, CRR and AWR impose an implicit policy regularization by performing policy gradient-style policy updates. Despite their effectiveness, they ignore that the capacity of a policy representation plays a vital role in fitting the data distribution. This paper instead focuses on an orthogonal aspect ( i.e. , policy parameterization) that all the above methods can benefit. Another line of work tries to cast offline RL as a sequence-to-sequence model [3, 11], which is beyond the scope of this work.

Policy Parametrization Different RL algorithms may pose different requirements for parameterizing a policy distribution. There are mainly two categories of requirements: 1) The sampling process

is differentiable, such as the deterministic policy in DDPG [18] and TD3 [7]. 2) The log-likelihood of samples is tractable. For example, policy gradient methods [30, 31, 38, 28] optimize a policy based on maximum likelihood estimation (MLE). Therefore, most works represent policy with a diagonal Gaussian distribution with mean and variance parameterized with a multi-layer perceptron (MLP). On the other hand, BCQ [9] and BEAR [16] choose to model policy with a conditional variational autoencoder (CVAE). Recently, Diffusion-QL [37] introduced diffusion models into offline RL and demonstrated that diffusion models are superior at modeling complex action distributions than CV AE and diagonal Gaussian. However, it takes tens to hundreds more time to train a diffusion policy than a diagonal Gaussian one. Moreover, diffusion policy only satisfies the first requirement, which means many other offline RL algorithms can not use it, including the current SOTA IQL.

Our method is motivated to solve the about two limitations. We first propose a more efficient way to train diffusion policies, which reduces training time to the level of a Gaussian policy. Then, we generalize the diffusion policy to be compatible with MLE-based RL methods.

3 Preliminaries

3.1 Offline Reinforcement Learning

A decision-making problem in reinforcement learning is usually represented by a Markov Decision Process (MDP): M = {S , A , P, R, γ } . S and A are the state and action spaces respectively, P ( s ′ | s , a ) measures the transition probability from state s to state s ′ after taking action a while R ( s , a , s ′ ) gives the reward for the corresponding transition, γ ∈ [0 , 1) is the discount factor. A policy π ( a | s ) 2 describes how an agent interacts with the environment. The optimal policy π ∗ ( a | s ) is the one achieves maximal cumulative discounted returns: π ∗ = arg max E π [ ∑ ∞ t =0 γ t r ( s t , a t )] . Reinforcement learning algorithms frequently rely on the definition of value function V ( s ) = E π [ ∑ ∞ t =0 γ t r ( s t , a t ) | s 0 = s ] , and action value (Q) function Q ( s , a ) = E π [ ∑ ∞ t =0 γ t r ( s t , a t ) | s 0 = s , a 0 = a ] , which represents the expected cumulative discounted return of a policy π given the initial state s or state-action pair ( s , a ) .

In the offline RL setting, instead of learning from interactions with the environment, agents focus on learning an optimal policy from a previously collected dataset of transitions: D = { ( s t , a t , s t +1 , r t ) } . Offline RL algorithms for continuous control are usually based on an actor-critic framework that alternates between policy evaluation and policy improvement. During policy evaluation, a parameterized Q network Q ϕ ( s , a ) is optimized based on approximate dynamic programming to minimize the following temporal difference (TD) error L TD ( ϕ ) :

work. Then at the policy improvement step, knowledge in the Q network is distilled into the policy network in various ways. Offline RL methods address the distributional shift [9] problem induced by the offline dataset D by either modifying the policy evaluation step to regularize Q learning or constraining the policy improvement directly. In the following, we will show that our diffusion policy design is compatible with any offline algorithms and can speed up policy evaluation and improvement.

3.2 Diffusion Models

Consider a real data distribution q ( x ) and a sample x 0 ∼ q ( x ) drawn from it. The (forward) diffusion process fixed to a Markov chain gradually adds Gaussian noise to the sample in K steps, producing a sequence of noisy samples x 1 , . . . x K . Note that we use superscript k to denote diffusion timestep to avoid conflicting with the RL timestep. The noise is controlled by a variance schedule β 1 , . . . , β K :

When K → ∞ , x K distributes as an isotropic Gaussian distribution. Diffusion models learn a conditional distribution p θ ( x t -1 | x t ) and generate new samples by reversing the above process:

where p ( x K ) = N ( 0 , I ) under the condition that ∏ K k =1 (1 -β k ) ≈ 0 . The training is performed by maximizing the evidence lower bound (ELBO): E x 0 [log p θ ( x 0 )] ≥ E q [ log p θ ( x 0: K ) q ( x 1: K | x 0 ) ] .

4 Efficient Diffusion Policy

In this section, we detail the design of our efficient diffusion policy (EDP). First, we formulate an RL policy with a diffusion model. Second, we present a novel algorithm that can train a diffusion policy efficiently, termed Reinforcement-Guided Diffusion Policy Learning (RGDPL). Then, we generalize the diffusion policy to work with arbitrary offline RL algorithms and compare our EDP with Diffusion-QL to highlight its superiority in efficiency and generality. Finally, we discuss several methods to sample from the diffusion policy during evaluation.

4.1 Diffusion Policy

Following [37], we use the reverse process of a conditional diffusion model as a parametric policy:

where a K ∼ N ( 0 , I ) . We choose to parameterize π θ based on Denoising Diffusion Probabilistic Models (DDPM) [10], which sets Σ θ ( a k , k ; s ) = β k I to fixed time-dependent constants, and constructs the mean µ θ from a noise prediction model as: µ θ ( a k , k ; s ) = 1 √ α k ( a k -β k √ 1 -¯ α k ϵ θ ( a k , k ; s ) ) , where α k = 1 -β k , ¯ α k = ∏ k s =1 , and ϵ θ is a parametric model.

To obtain an action from DDPM, we need to draw samples from K different Gaussian distributions sequentially, as illustrated in Eqn. (2)-(3). The sampling process can be reformulated as

with the reparametrization trick, where ϵ ∼ N ( 0 , I ) , k is the reverse timestep from K to 0 .

Similar to DDPM, plugging in the conditional Gaussian distributions, the ELBO in Sec. 3.2 can be simplified to the following training objective L diff ( θ ) :

where k follows a uniform distribution over the discrete set { 1 , . . . , K } . It means the expectation is taken over all diffusion steps from clean action to pure noise. Moreover, ϵ ∈ N ( 0 , I ) , and ( a 0 , s ) ∈ D are state-action pairs drawn from the offline dataset. Given a dataset, we can easily and efficiently train a diffusion policy in a behavior-cloning manner as we only need to forward and backward through the network once each iteration. As shown in Diffusion-QL [37], diffusion policies can greatly boost the performance when trained with TD3-based Q learning. However, it still faces two main drawbacks that limit its real-world application: 1) It is inefficient in sampling and training; 2) It is not generalizable to other strong offline reinforcement learning algorithms.

4.2 Reinforcement-Guided Diffusion Policy Learning

To understand how a parametric policy π θ is trained with offline RL algorithms, we start with a typical Q-learning actor-critic framework for continuous control, which iterates between policy evaluation and policy improvement. Policy evaluation learns a Q network by minimizing the TD error L TD ( ϕ ) :

where the next action a ′ ∼ π θ ( ·| s ′ ) . The policy is optimized to maximize the expected Q values :

It is straightforward to optimize this objective when a Gaussian policy is used, but things get much more difficult when a diffusion policy is considered due to its complicated sampling process. Instead, we propose to view the offline RL problem from the perspective of generative modeling, where a diffusion policy can be easily learned in a supervised manner from a given dataset. However, unlike in computer vision, where the training data are usually perfect, offline RL datasets often contain suboptimal state-action pairs. Suppose we have a well-trained Q network Q ϕ , the question becomes how we can efficiently use Q ϕ to guide diffusion policy training procedure. We now show that this can be achieved without sampling actions from diffusion policies.

Let's revisit the forward diffusion process in Eqn. 1. A notable property of it is that the distribution of noisy action a k at any step k can be written in closed form: q ( a k | a 0 ) = N ( a k ; √ ¯ α k a 0 , (1 -¯ α k ) I ) . Using the reparametrization trick, we are able to connect a k , a 0 and ϵ by:

Recall that our diffusion policy is parameterized to predict ϵ with ϵ θ ( a k , k ; s ) . By relacing ϵ with ϵ θ ( a k , k ; s ) and rearranging Eqn. (8), we obtain the approximiated action:

In this way, instead of running the reverse diffusion process to sample an action a 0 , we can cheaply construct ˆ a 0 from a state-action pair ( s , a ) in the dataset by first corrupting the action a to a k then performing one-step denoising to it. We will refer to this technique as action approximation in the following. Accordingly, the policy improvement for diffusion policies is modified as follows:

To improve the efficiency of policy evaluation, we propose to replace the DDPM sampling in Eqn. (4) with DPM-Solver [20], which is an ODE-based sampler. The algorithm is defered to the appendix.

4.3 Generalization to Various RL algorithms

There are mainly two types of approaches to realize the objective in Eqn. 7 for policy improvement.

Direct policy optimization : It maximizes Q values and directly backpropagate the gradients from Q network to policy network, i.e. , ∇ θ L π ( θ ) = -∂Q ϕ ( s , a ) ∂ a ∂ a ∂θ . This is only applicable to cases where ∂ a ∂θ is tractable, e.g. , when a deterministic policy a = π θ ( s ) is used or when the sampling process can be reparameterized. Sample algorithms belonging to this category include TD3 [7], TD3+BC [6], and CQL [17]. One can easily verify that both the expensive DDPM sampling in Eqn. (4) and our efficient approximation in Eqn. (9) can be used for direct policy optimization.

Likelihood-based policy optimization : It tries to distill the knowledge from the Q network into the policy network indirectly by performing weighted regression or weighted maximum likelihood:

where f ( Q ϕ ( s , a )) is a monotonically increasing function that assigns a weight to each state-action pair in the dataset. This objective requires the log-likelihood of the policy to be tractable and differentiable. AWR [28], CRR [38], and IQL [15] fall into this category but each has a unique design in terms of the weighting function f . Since the likelihood of samples in Diffusion models is intractable, we propose the following two variants for realizing Eqn. 11.

First, instead of computing the likelihood, we turn to a lower bound for log π θ ( a | s ) introduced in DDPM [10]. By discarding the constant term that does not depend on θ , we can have the objective:

Second, instead of directlying optimizing log π θ ( a | s ) , we propose to replace it with an approximated policy ˆ π θ ( a | s ) ≜ N (ˆ a 0 , I ) , where ˆ a 0 is from Eqn. (9). Then, we get the following objective:

Empirically, we find these two choices perform similarly, but the latter is easier to implement. So we will report results mainly based on the second realization. In our experiments, we consider two offline RL algorithms under this category, i.e. , CRR, and IQL. They use two weighting schemes: f CRR = exp [( Q ϕ ( s, a ) -E a ′ ∼ ˆ π ( a | s ) Q ( s, a ′ ) ) /τ CRR ] and f IQL = exp[( Q ϕ ( s, a ) -V ψ ( s )) /τ IQL ] , where τ refers to the temperature parameter and V ψ ( s ) is an additional value network parameterized by ψ . We defer the details of these two algorithms to Appendix A.

4.4 Comparison to Diffusion-QL

Now we are ready to compare our method and Diffusion-QL comprehensively. Though our EDP shares the same policy parametrization as Diffusion-QL, it differs from Diffusion-QL significantly in the training algorithm. As a result, the computational efficiency and generality of diffusion policies have been improved substantially.

Efficiency The diffusion policy affects both policy evaluation (Eqn. (6)) and policy improvement (Eqn. (7)). First, calculating L TD ( ϕ ) in policy evaluation requires drawing the next action from it. Diffusion-QL uses DDPM sampling while EDP employs a DPM-Solver, which can reduce the sampling steps from 1000 to 15, thus accelerating the training. Second, in policy improvement, Diffusion-QL again applies DDPM for sampling. Then, it calculates the loss function based on sampled actions and backpropagates through the sampling process for network update. This means it needs to forward and backward a neural network for K times each training iteration. As a result, Diffusion-QL can only work with small K , e.g., 5 ∼ 100 . In comparison, our training scheme only passes through the network once an iteration, no matter how big K is. This enables EDP to use a larger K (1000 in our experiments) to train diffusion policy on the more fine-grained scale. The results in Tab. 1 also show a larger K can give better performance.

Generality Diffusion-QL can only work with direct policy optimization, which contains only a small portion of algorithms. Moreover, thanks to their flexibility and high performance, the likelihood-based algorithms are preferred for some tasks ( e.g. , Antmaze). Our method successfully makes diffusion trainable with any RL algorithm.

4.5 Controlled Sampling from Diffusion Policies

Traditionally, a continuous policy is represented with a state-conditional Gaussian distribution. During evaluation time, a policy executes deterministically to reduce variance by outputting the distribution mean as an action. However, with diffusion policies, we can only randomly draw a sample from the underlying distribution without access to its statistics. As a result, the sampling process is noisy, and the evaluation is of high variance. We consider the following method to reduce variance.

Energy-based Action Selection (EAS) Recall that the goal of (offline) RL is to learn a policy that can maximize the cumulative return or values. Though the policy π θ is stochastic, the learned Q ϕ provides a deterministic critic for action evaluation. We can sample a few actions randomly, then use Q ϕ for selection among them to eliminate randomness. EAS first samples N actions from π θ by using any samplers ( i.e. , DPM-Solver), then sample one of them with weights proportional to e Q ( s , a ) . This procedure can be understood as sampling from an improved policy p ( a | s ) ∝ e Q ( s , a ) π θ ( a | s ) . All results will be reported based on EAS. See Appendix. C.4 for the other two methods.

5 Experiments

We conduct extensive experiments on the D4RL benchmark [2] to verify the following assumptions: 1) Our diffusion policy is much more efficient than the previous one regarding training and evaluation costs. 2) Our diffusion policy is a generic policy class that can be learned through direct and likelihood-based policy learning methods. We also provide various ablation studies on the critical components for better understanding.

Baselines Weevaluate our method on four domains in D4RL, including Gym-locomotion, AntMaze, Adroit, and Kitchen. For each domain, we consider extensive baselines to provide a thorough evaluation. The simplest method is the classific behavior cloning (BC) baseline and 10% BC that performs behavior cloning on the best 10% data. TD3+BC [6] combines off-policy reinforcement

learning algorithms with BC. OneStepRL [1] first conducts policy evaluation to obtain the Q-value of the behavior policy from the offline dataset, then use it for policy improvement. AWAC [25], AWR [28], and CRR [38] improve policy improvement by adding advantage-based weights to policy loss functions. CQL [17] and IQL [15] constrain the policy evaluation process by making conservative Q updates or replacing the max operator with expectile regression. We also consider the Decision Transformer (DT) [4] baseline that maps offline RL as a sequence-to-sequence translation problem.

Experimental Setup We keep the backbone network architecture the same for all tasks and algorithms, which is a 3-layer MLP (hidden size 256) with Mish [23] activation function following Diffusion-QL [37]. For the noise prediction network ϵ θ ( a k , k ; s ) in diffusion policy, we first encode timestep k with sinusoidal embedding [36], then concatenate it with the noisy action a k and the conditional state s . We use the Adam [13] to optimize both diffusion policy and the Q networks. The models are trained for 2000 epochs on Gym-locomotion and 1000 epochs on the other three domains. Each epoch consists of 1000 iterations of policy updates with batch size 256. For DPM-Solver [20], we use the third-order version and set the model call steps to 15. We reimplement DQL strictly following the official PyTorch code [27] for fair comparison and we refer to DQL (JAX) for all sample efficiency comparisons. We defer the complete list of all hyperparameters to the appendix due to space limits. Throughout this paper, the results are reported by averaging 5 random seeds.

Evaluation Protocal We consider two evaluation metrics in this paper. First, online model selection (OMS), proposed by Diffusion-QL [37], selects the best-performing model throughout the whole training process. However, though OMS can reflect an algorithm's capacity, it is cheating, especially when the training procedure is volatile on some of the tasks. Therefore, we propose another metric to focus on the training stability and quality, which is running average at training (RAT). RAT calculates the running average of evaluation performance for ten consecutive checkpoints during training and reports the last score as the final performance.

5.1 Efficiency and Reproducibility

In this section, we focus on the training and evaluation efficiency of our efficient diffusion policy. We choose the OMS evaluation metric to make a fair comparison with the baseline method Diffusion-QL [37]. We consider four variants of EDP to understand how each component contributes to the high efficiency of our method. 1) EDP is the complete version of our method. It uses the action approximation technique in Eqn. (9) for policy training and uses DPM-Solver for sampling. 2) EDP w/o DPM modifies EDP by replacing DPM-Solver with the original DDPM sampling method in Eqn. 4. 3) EDP w/o AP removes the action approxima-

Training

Sampling

0

×

1

×

2

×

3

×

4

×

5

×

Speedup on Gym-locomotion

DQL

DQL(Jax)

EDP w/o DPM

EDP w/o AP

EDP

tion technique. 4) DQL (JAX) is our Jax implementation of Diffusion-QL.

We first benchmark and compare the training/evaluation speed of the above three variants and Diffusion-QL. We choose walker2d-medium-expert-v2 as the testbed. For training speed, we run each algorithm for 10,000 iterations of policy updates and calculate the corresponding iterations-per-second (IPS). Similarly, we sample 10,000 transitions by interacting with the environment and calculate the corresponding steps-per-second (SPS) for evaluation speed. Based on the visualization in Fig. 2, by taking DQL (JAX) as the baseline, we are able to attribute the performance boost to specific techniques proposed. Specifically, we can observe that action approximation makes 2.3x training and 3.3x sampling faster, while using the DPM-Solver adds an additional 2.3x training speedup. We can observe that DQL (JAX) is 5 × faster than Diffusion-QL, which means our Jax implementation is more computationally efficient than Diffusion-QL's PyTorch code. This demonstrates that both the action approximation technique and DPM-Solver play a critical role in making the training of diffusion policy efficient. However, this technique does not affect the sampling procedure; thus, EDP w/o DPM and DQL (JAX) are on par with each other regarding sampling speed.

To show that EDP does not hurt performance, we compare the normalized scores of all tasks in Tab. 1. The results for Diffusion-QL are directly copied from the paper, where each task is carefully tuned with its own hyperparameters. Instead, DQL (JAX) and EDP use the same hyperparameters for tasks belonging to the same domain. Moreover, since EDP is more computationally efficient, we can train a better score model with large K = 1000 , which is larger than the one (5 ∼ 100) used in Diffusion-QL. Note that training diffusion policies with large K is impossible

in Diffusion-QL, as it needs to forward and backward the neural network K times. Finally, we can observe that EDP and DQL (JAX) are comparable with Diffusion-QL on gym-locomotion tasks but much better on the other three domains. Hence, efficient diffusion policy can boost sample efficiency and improve performance by enabling diffusion training in fine-grained noise prediction.

5.2 Generality and Overall Results

This section aims to evaluate whether EDP is a general policy class. To this end, we train it with direct policy optimization (TD3) and likelihood-based policy optimization (CRR, IQL) methods. Then, we compare them with their feed-forward counterparts and other baseline methods in Table 2. All scores for diffusion policies are reported using the RAT metric, while other scores are directly quoted from their paper. It shows that EDP can beat the standard Gaussian policy parameterized with an MLP on all domains and for all the three algorithms considered. On the gym-locomotion domain, EDP + TD3 gives the best performance (average score 85.5), while likelihood-based policy learning methods are slightly worse. However, on the other three domains (Kitchen, Adroit, and Antmaze), EDP + IQL beats all the other methods by a large margin (more than 10 average scores). Therefore, we conclude that EDP can serve as a plug-in policy class for different RL methods.

5.3 Ablation Study

Evaluation Metrics To reveal that evaluation metric is important, we train EDP with TD3 algorithms on three selected environments: walker2d-medium-expert-v2, hopper-medium-expert-v2, and antmaze-medium-diverse-v0. We then compare the scores for OMS (best) and RAT (average) by plotting the training curves in Fig. 3. On

0.0

0.2

0.4

0.6

0.8

1.0

Normalized Return

Walker2d

Average

Best

Training Steps

0.2

0.4

0.6

0.8

1.0

Hopper

0.0

0.2

0.4

0.6

0.8

1.0

Antmaze

walker2d, the training is stable; thus, both OMS and RAT scores steadily grow and result in close final scores. A similar trend can be observed on the hopper but with a more significant gap between these two metrics. However, these two metrics diverge significantly when the training succeeds and then crashes on antmaze. Therefore, OMS is misleading and can not give a reliable evaluation of algorithms, which explains the necessity of using RAT in Sec. 5.2.

Energy-Based Action Selection

We notice that energy-based action selection (EAS) is a general method and can also be used for arbitrary policies. We apply EAS to normal TD3+BC and find no improvement, which shows EAS is only necessary for diffusion sampling. The results are deferred to the Appendix. Moreover, set the number of actions used in

1

10

100

0.4

0.5

0.6

0.7

0.8

0.9

Normalized Return

halfcheetah

medium

medium-replay

medium-expert

1

10

100

Number of Samples

0.6

0.7

0.8

0.9

1.0

hopper

1

10

100

0.80

0.85

0.90

0.95

1.00

1.05

1.10

walker2d

EAS from 1 to 200, and report the performance on gym-locomotions tasks in Fig. 4. It shows the normalized score monotonically grows as the number of actions increases on 8 out of 9 tasks. In our main experiments, we set the number of actions to 10 by trading off the performance and computation efficiency.

DPM-Solver We are using DPMSolver to speed up the sampling process of diffusion policies. The number of models in DPM-Solver is an important hyper-parameter that affects sampling efficiency and quality. We vary this number from 3 to 30 and compare the performance on gym-locomotion tasks in Fig. 5. We can observe that the performance increases as more

5

10

15

20

25

30

0.4

0.5

0.6

0.7

0.8

0.9

1.0

Normalized Return

halfcheetah

medium

medium-replay

medium-expert

5

10

15

20

25

30

Number of Model Calls

0.2

0.4

0.6

0.8

1.0

hopper

5

10

15

20

25

30

0.0

0.2

0.4

0.6

0.8

1.0

walker2d

steps of model calls are used. The performance gradually plateaus after 15 model calls. Therefore, we use 15 in our main experiments.

6 Conclusion

Diffusion policy has emerged as an expressive policy class for offline reinforcement learning. Despite its effectiveness, diffusion policy is limited by two drawbacks, hindering it from wider applications. First, training a diffusion policy requires to forward and backward through a long parameterized Markov chain, which is computationally expensive. Second, the diffusion policy is a restricted policy class that can not work with likelihood-based RL algorithms, which are preferred in many scenarios. We propose efficient diffusion policy (EDP) to address these limitations and make diffusion policies faster, better, and more general. EDP relies on an action approximation to construct actions from corrupted ones, thus avoiding running the Markov chain for action sampling at training. Our benchmarking shows that EDP achieves 25 × speedup over Diffusion-QL at training time on the gym-locomotion tasks in D4RL. We conducted extensive experiments by training EDP with various offline RL algorithms, including TD3, CRR, and IQL, the results clearly justify the superiority of diffusion policies over Gaussian policies. As a result, EDP set new state-of-the-art on all four domains in D4RL.

References

David Brandfonbrener, Will Whitney, Rajesh Ranganath, and Joan Bruna. Offline rl without off-policy evaluation. Advances in Neural Information Processing Systems , 34:4933-4946, 2021.

David Brandfonbrener, Will Whitney, Rajesh Ranganath, and Joan Bruna. Offline rl without off-policy evaluation. Advances in Neural Information Processing Systems , 34:4933-4946, 2021.

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

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

Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in Neural Information Processing Systems , 34:8780-8794, 2021.

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

Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor-critic methods. In International conference on machine learning , pages 1587-1596. PMLR, 2018.

Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor-critic methods. In International conference on machine learning , pages 1587-1596. PMLR, 2018.

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

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems , 33:6840-6851, 2020.

Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems , 34:12731286, 2021.

Bingyi Kang, Xiao Ma, Yirui Wang, Yang Yue, and Shuicheng Yan. Improving and benchmarking offline reinforcement learning algorithms. arXiv preprint arXiv:2306.00972 , 2023.

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

Roger Koenker and Kevin F Hallock. Quantile regression. Journal of economic perspectives , 15(4):143-156, 2001.

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

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

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

Timothy P Lillicrap, Jonathan J Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. arXiv preprint arXiv:1509.02971 , 2015.

Timothy P Lillicrap, Jonathan J Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. arXiv preprint arXiv:1509.02971 , 2015.

Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. arXiv preprint arXiv:2206.00927 , 2022.

Xiao Ma, Bingyi Kang, Zhongwen Xu, Min Lin, and Shuicheng Yan. Mutual information regularized offline reinforcement learning. arXiv preprint arXiv:2210.07484 , 2022.

Michael Mathieu, Sherjil Ozair, Srivatsan Srinivasan, Caglar Gulcehre, Shangtong Zhang, Ray Jiang, Tom Le Paine, Konrad Zolna, Richard Powell, Julian Schrittwieser, et al. Starcraft ii unplugged: Large scale offline reinforcement learning. In Deep RL Workshop NeurIPS 2021 , 2021.

Diganta Misra. Mish: A self regularized non-monotonic neural activation function. arXiv preprint arXiv:1908.08681 , 4(2):10-48550, 2019.

Ofir Nachum, Mohammad Norouzi, Kelvin Xu, and Dale Schuurmans. Bridging the gap between value and policy based reinforcement learning. Advances in neural information processing systems , 2017.

Ashvin Nair, Abhishek Gupta, Murtaza Dalal, and Sergey Levine. Awac: Accelerating online reinforcement learning with offline datasets. arXiv preprint arXiv:2006.09359 , 2020.

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

Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems 32 , pages 8024-8035. Curran Associates, Inc., 2019.

Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. Advantage-weighted regression: Simple and scalable off-policy reinforcement learning. arXiv preprint arXiv:1910.00177 , 2019.

Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. Advantage-weighted regression: Simple and scalable off-policy reinforcement learning. arXiv preprint arXiv:1910.00177 , 2019.

John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In International conference on machine learning , pages 1889-1897. PMLR, 2015.

John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347 , 2017.

Nur Muhammad Mahi Shafiullah, Zichen Jeff Cui, Ariuntuya Altanzaya, and Lerrel Pinto. Behavior transformers: Cloning k modes with one stone. arXiv preprint arXiv:2206.11251 , 2022.

Noah Siegel, Jost Tobias Springenberg, Felix Berkenkamp, Abbas Abdolmaleki, Michael Neunert, Thomas Lampe, Roland Hafner, Nicolas Heess, and Martin Riedmiller. Keep doing what worked: Behavior modelling priors for offline reinforcement learning. In International Conference on Learning Representations , 2020.

Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning , pages 2256-2265. PMLR, 2015.

Yang Song, Conor Durkan, Iain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. Advances in Neural Information Processing Systems , 34:14151428, 2021.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

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

Ziyu Wang, Alexander Novikov, Konrad Zolna, Josh S Merel, Jost Tobias Springenberg, Scott E Reed, Bobak Shahriari, Noah Siegel, Caglar Gulcehre, Nicolas Heess, et al. Critic regularized regression. Advances in Neural Information Processing Systems , 33:7768-7778, 2020.

Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv preprint arXiv:1911.11361 , 2019.

Tianhe Yu, Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine, and Chelsea Finn. Combo: Conservative offline model-based policy optimization. Advances in neural information processing systems , 34:28954-28967, 2021.

Yang Yue, Bingyi Kang, Xiao Ma, Gao Huang, Shiji Song, and Shuicheng Yan. Offline prioritized experience replay. arXiv preprint arXiv:2306.05412 , 2023.

Yang Yue, Bingyi Kang, Xiao Ma, Zhongwen Xu, Gao Huang, and Shuicheng Yan. Boosting offline reinforcement learning via data rebalancing. arXiv preprint arXiv:2210.09241 , 2022.

Broader Impact

In this paper, we propose an efficient yet powerful policy class for offline reinforcement learning. We show that this method is superior to most existing methods on simulated robotic tasks. However, some war robots or weapon robots might employ our EDP to learn strategic agents considering the generalization ability of EDP. Depending on the specific application scenarios, it might be harmful to domestic privacy and safety.

A Reinforcement Learning Algorithms

In this section, we introduce the details of the RL algorithms experimented.

A.1 TD3+BC

TD3 [8] is a popular off-policy RL algorithm for continuous control. TD3 improves DDPG [19] by addressing the value overestimation issue. Specifically, TD3 adopts a double Q-learning paradigm that computes the TD-target ˆ Q ( s , a ) as

where π ( s ′ ) is a deterministic policy, Q 1 and Q 2 are two independent value networks. Specifically, TD3 takes only Q 1 for policy improvement

Built on top of TD3, TD3+BC simply adds an additional behavior cloning term for its policy improvement

where α is a hyper-parameter that balances these two terms. However, as the scale of Q-values are different from the behavior cloning loss, TD3+BC normalizes it with 1 N N ∑ i =1 | Q ( s i , ˆ a i ) | for numerical stability.

In the context of EDP, we observe that 1) the behavior cloning term can be naturally achieved as a diffusion loss as defined in Eqn. 5; 2) the sampled action ˆ a is replaced by action approximation as in Eqn. 9 for efficient policy improvement. Different from the original TD3+BC that uses only Q 1 for policy improvement, we sample from Q 1 and Q 2 with equal probability for each policy improvement step.

A.2 Critic Regularized Regression

CRR follows Advantage Weighted Regression (AWR) [29], which is a simple yet effective off-policy RL method, for policy improvement. Specifically, AWR considers a constrained policy improvement step

where A ( s , a ) is the advantage function, µ is a behavior policy that is used to generate trajectories during off-policy RL, d µ ( s ) is the state distribution induced by µ , D KL is the KL divergence, and ϵ is a threshold parameter. This constrained optimization problem can be solved in closed form, which gives an optimal policy of

with Z ( s ) being the partition function and β is a hyper-parameter. For policy improvement, AWR simply distills the one-step improved optimal policy to the learning policy π by minimizing the

KL-divergence

CRR performs policy improvement in the same way as AWR, which simply replaces the sampling distribution with a fixed dataset

As a result, this naturally imposes an implicit constraint on its policy improvement step.

However, as computing the log-likelihood log π ( a | s ) is intractable in diffusion models, in practical implementations, we use Eqn. 13 instead. In addition, we compute the advantage by

where ˆ a i ∼ N (ˆ a 0 , σ ) is a sampled action with the mean of approximated action and an additional standard deviation. In our experiment, we found using a fixed standard deviation, an identity matrix of σ = I , β = 1 , and sample size N = 10 generally produce a good performance.

A.3 Implicit Q Learning

Similar to CRR, IQL also adopts the AWR-style policy improvement that naturally imposes a constraint to encourage the learning policy to stay close to the behavior policy that generates the dataset. Different from CRR query novel and potentially out-of-distribution (OOD) actions when computing advantages, IQL aims to completely stay in-distribution with only dataset actions, while maintaining the ability to perform effective multi-step dynamic programming during policy evaluation. IQL achieves this by introducing an additional value function V ( s ) and performing expectile regression. Specifically, the policy evaluation in IQL is implemented as

where L τ 2 is the expectile regression loss defined as L τ 2 ( x ) = | τ -✶ ( x < 0) | x 2 with hyper-paramter τ ∈ (0 , 1) . The intuition behind IQL is that with a larger τ , we will be able to better approximate the max operator. As a result, IQL approximates the Bellman's optimality equation without querying OOD actions.

In practical implementations, we also adopt double Q learning for IQL, where we replace the Q ( s , a ) in Eqn. 22 with min( Q 1 ( s , a ) , Q 2 ( s , a )) , and then use the updated value network to train both Q value networks. For IQL, we follow the policy improvement steps of CRR, as described in Eqn. 20 and Eqn. 21. The key difference is that instead of sampling actions to compute

parameters, we use a temperature of β = 1 , a fixed standard deviation σ = I , and for expectile ratio τ , we use τ = 0 . 9 for antmaze-v0 environments and τ = 0 . 7 for other tasks.

B Reinforcement Guided Diffusion Policy Details

The overall algorithm for our Reinforcement Guided Diffusion Policy Learning is given in Alg. 1. The detailed algorithm for energy-based action selection is given in Alg. 2.

C Environmental Details

C.1 Hyper-Parameters

We detail our hyperparameters in Tab. 3.

Algorithm 1: Reinforcement Guided Diffusion Policy

Algorithm 2: Energy-based Action Selection

C.2 More Results

We first expand Tab. 1 by providing detailed numbers for each of the tasks used in Tab. 4.

We report the performance of EDP trained with TD3, CRR, and IQL in Tab. 5, where we directly compare the scores of different evaluation metrics, i.e. , OMS and RAT. We can observe that there are huge gaps between OMS and RAT for all domains and all algorithms. However, IQL and CRR are relatively more stable than TD 3 . For example, on the antmaze domain, TD3 achieves a best score of 80.5, while the average score is just 29.8. In comparison, the best and average scores of IQL are 89.2 and 73.4, respectively.

C.3 More Results on EAS

We compare normal TD3+BC, TD3+BC with EAS for evaluation and TD3 + EDP in Tab. 6.

C.4 More Experiments on Controlled Sampling

As in Sec. 4.5, we discussed reducing variance with EAS. We now detail another two methods experimented as below.

Policy scaling Instead of sampling from the policy directly, we can sample from a sharper policy π τ θ ( a | s ) , where τ > 1 . The scaled policy π τ θ shares the same distribution modes as π θ but with reduced variance. Since diffusion policies are modeling the scores log π θ ( a t | s ) , policy scaling can be easily achieved by scaling the output of noise-prediction network by the factor τ . We conduct experiments on the gym-locomotion tasks in D4RL by varying τ from 0 . 5 to 2 . 0 , as shown in Fig. 6, the results show the best performance is achieved when τ = 1 . 0 . It means sampling from a scaled policy does not work.

Deterministic Sampling This method is based on the observation that the sampling process of the DPM-Solver is deterministic except for the first step. The first step is to sample from an isotropic Gaussian distribution. We modify it to use the mean, thus avoiding stochasticity. Consider a initial noise a K ∼ N ( 0 , I ) , we rescale a K by a noise scale factor. We show how this factor affects the final performance by varying it from 0.0 to 1.0. As illustrated in Fig. 7, the best performance is achieved at zero noise in most cases, and a normal a K performs worst. This means reducing the variance of the initial noise is able to improve the performance of a diffusion policy. However, the best performance achieved in this way still falls behind EAS.

C.5 Computational Cost Comparison between EDP and Feed-Forward Policy networks

We benchmark the training speed of TD3+BC with EDP on walker2d-medium-expert-v2, by training each agent for 10,000 iterations of policy updates. The training speed for TD3+BC is 689 iterationsper-second (IPS), while EDP is 412 IPS. In other words, it takes around 3 hours to train an agent with the feed-forward network, while EDP needs around 5 hours. We also conducted experiments by double the network used in TD3+BC; unfortunately, there was no performance gain on the locomotion tasks (75.3 to 75.6). Moreover, both feed-forward policy and diffusion policy utilize a 3-layer MLP (hidden size 256) as the backbone network. Therefore, the network capacity should be comparable.

1.0

1.5

2.0

0.0

0.2

0.4

0.6

0.8

Normalized Return

halfcheetah

medium

medium-replay

medium-expert

1.0

1.5

2.0

Score Scale

0.0

0.2

0.4

0.6

0.8

1.0

hopper

1.0

1.5

2.0

0.0

0.2

0.4

0.6

0.8

1.0

walker2d

C.6 The effect of action approximation

We compare with and without action approximation on the following three environments by using OMS metric (Tab. 4). In Tab. 7, the DDPM column will forward and backward a policy network 100 times at training time, but action approximation only needs once. We can observe that action approximation will slightly harm the performance, when the same number of diffusion steps is used. However, it supports training diffusion policies with larger K (e.g. 1000), while Diffusion-QL does

0.00

0.25

0.50

0.75

1.00

0.4

0.5

0.6

0.7

0.8

0.9

Normalized Return

halfcheetah

medium

medium-replay

medium-expert

0.00

0.25

0.50

0.75

1.00

0.6

0.7

0.8

0.9

1.0

1.1

hopper

0.00

0.25

0.50

0.75

1.00

0.80

0.85

0.90

0.95

1.00

1.05

1.10

walker2d

not. Increasing K is able to avoid performance drop as evidenced by the last column of Tab. 2 and Tab. 4.

D Negative Societal Impacts

In this paper, we propose an efficient yet powerful policy class for offline reinforcement learning. We show that this method is superior to most existing methods on simulated robotic tasks. However, some war robots or weapon robots might employ our EDP to learn strategic agents considering the generalization ability of EDP. Depending on the specific application scenarios, it might be harmful to domestic privacy and safety.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper focuses on the improvement of computation efficiency of Diffusion-QL by adopting the property of marginal distribution in the diffusion model and the variance control scheme proposed by DPM-Solver. Besides, this paper extends the scope of compatibility with other offline RL methods, from value-based to policy gradient methods. This is achieved by directly approximating the clean examples from corrupted examples at arbitrary diffusion time steps.\n\nSTRENGTHS:\n- This paper is well-written, and the proposed method is easy to understand.\n- The evaluation of value at risk seems interesting.\n- The technique seems sound, and the results seem strong.\n\nWEAKNESSES:\n- The discussion of limitation seems to be neglected.\n- More clear discussion is needed to highlight the contribution of this paper and show the novelty.\n- The proposed algorithm seems to significantly increase the inference computation complexity by introducing 10000 iterations for each action.\n- More experiments are needed to support the contributions sufficiently.\n- Some typos are needed to fix through further proofreading.\n\nQUESTIONS:\nSee my detailed comments below.\n\n==Major concerns==\n\n- In lines 233-235 of Section 4.5: according to Equation (9), the action approximation is the mean value of action Gaussian distribution, so why adopting $\\hat{a}^0$ can not reduce high variance?\n- In Section 4.4, the improvement of sample efficiency (from $ K=1000$ to $K\\approx 15$) comes from DPM-Solver rather than the action approximation. Firstly, if the author adopts action approximation, whether the $a_K$ comes from normal Gaussian distribution? If so, the noise prediction function $\\epsilon_theta (a^K,K;s)$ should predict the action based on the state and the Gaussian noise $a^K$. Intuitively, generating proper actions based on states seems more difficult, i.e., that classical policy function.\n- As for the log-likelihood, why can we not obtain the exact log-likelihood of $\\pi_{\\theta}(a|s)$ through $log~p(a^K)+\\sum_{k=1}^{K}log~p_{\\theta}(a^{k-1}| a^{k}, s)$?\n- Does the action approximation improve the computation efficiency of diffusion policy? \n- I hope the authors explain how they use the action approximation during training.\n- During the sampling stage, the DPM-Solver speeds up generation rather than action approximation. If action approximation just works on Equation (10), I think it is very similar to the “Target Policy Smoothing Regularization” technique in Section 5.3 of the TD3 paper. I want to know why the authors use the predicted $\\hat{a}^0$ rather than $a^0+\\epsilon$.\n\n\n==Minor concerns==\n- How does the author avoid out-of-distribution actions' effects on the Q function?\n- DPM-Solver directly makes me understand the sample efficiency of EDP, but I can not understand why the training efficiency can also be improved? So I hope more description of the reason.\n- In line 318-320 of Section 5.1: during training the diffusion model, we sample a step $t\\sim [1, K]$, then train the diffusion model according to Equation (5). If the $K$ is large, we can still train the diffusion model. So I hope the authors explain this claim more clearly.\n\n\n==Typos==\n- Line 188: Eqn. 7 -> Eqn. (7).\n- Line 176: Eqn. 1 -> Eqn. (1).\n- Line 201: Eqn. 11 -> Eqn. (11).\n- Line 292: Eqn. 4 -> Eqn. (4).\n- Line 518 of Appendix: Eqn. 5 -> Eqn. (5).\n- Line 519 of Appendix: Eqn. 9 -> Eqn. (9).\nSimilar typos can also be found in many places in Appendix.\n- Line 565 of Appendix: Tab. ??.\n- Line 575 of Appendix: Tab. ??.\n\nLIMITATIONS:\nThe authors do not discuss the limitations and discuss broader impacts.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose a method to efficiently train diffusion based policies in the offline-RL setting. The authors suggest three main tricks to enable this: 1) Removing the need to backpropagate through the diffusion sampling chain to update the policy by using what the authors call action approximation; 2) Replacement of intractable policy log likelihoods in reinforcement learning objectives with the ELBO, which allows using diffusion policies generally with many popular offline RL algorithms; 3) Using a fast DPM-Solver for action sampling during policy execution. Altogether, these tricks allow for relatively fast training of expressive multimodal policies which achieve state of the art results across the D4RL benchmark suite.\n\nSTRENGTHS:\n* The paper is generally well written and easy to follow.\n* The use of action approximation facilitates training diffusion policy models with much larger diffusion noising time K, with the authors using K=1000. The RL community has recently taken great interest in the prospect of using the expressiveness of diffusion models for policies, and this simple trick seems to make diffusion policy training practical without any noticeable drawbacks.\n* The authors also demonstrate that we can use diffusion policies with other common offline RL algorithms such as TD3+BC, CRR and IQL by simply replacing the log likelihood with the diffusion ELBO. This is a valuable contribution to the field, as it could open the door to use diffusion policies to work with many general RL algorithms, a strict improvement over simple gaussian policies.\n* The authors show strong empirical performance across the D4RL benchmark suite. EDP policies are especially strong compared to the FF counterparts on the more challenging antmaze and kitchen suite, which involve high multimodality due to undirected demonstrations.\n\nWEAKNESSES:\n* The main proposal to efficiently train diffusion policies hinges on action approximation, which uses a much higher variance estimate of $\\hat{a}^0$. There are no ablations to show how this affects the final diffusion policy as compared to one trained without action approximation. There is a comparison against Diffusion-QL in the paper, but with a very large difference in diffusion timesteps K, so the effect of action approximation itself is difficult to gauge. I have listed some questions related to this point below, which if clarified would be helpful.\n* The use of DPM-Solver is highlighted as a major contribution is speeding up the sampling process compared to DDPM sampling in Diffusion-QL. It is a well known fact that faster diffusion sampling methods than DDPM exist which can be used for sampling, and Diffusion-QL could have used this as well. While it is useful that the authors show the time savings for using DPM-Solver over DDPM sampling, this is not a novel contribution.\n* There isn’t discussion about important hyperparameters related to the method. A very important hyperparameter that could be discussed more the diffusion timesteps K (I have noted this as a question below as well). Are EDPs more brittle to train than their FF counterparts in algorithms like IQL with regards to changing hyperparameters?\n\nQUESTIONS:\n* Generally, diffusion models generate higher quality samples with larger diffusion timesteps K, and the authors note in Tab. (1) that training EDP with K=1000 matched or beat Diffusion-QL which was trained with much smaller K~5 to 100. The authors cite this ability to train diffusion models with higher K without having to backprop through the chain as one of the main advantages of EDP. I am curious how much of an impact K actually has in these offline RL benchmark tasks, given that Diffusion-QL has similar scores. I am interested to see an ablation study (at least in some environments) where EDP is trained with different values of K starting at K=1, to see if generally increasing K actually significantly improves policies in these domains.\n* Energy-based Action Selection (EAS) seems necessary to come close to the performance of DQL using the proposed EDP policies. The appendix shows that other ways to sample actions from the diffusion model result in much worse performance. DQL does not to my knowledge require EAS to produce good action samples from its diffusion policy. Is this due to the effect of action approximation resulting in training a worse diffusion model?\n* As an extension to the above point, have the authors considered using sampling guidance techniques like Classifer-Free Guidance (CFG) to reduce variance of the final policy?\n* The authors allow for training diffusion policies with popular offline RL algorithms which require policy log likelihoods by replacing it with the ELBO. Could this be done with lower bounds of simpler deep generative models, such as the ELBO of a Variational Autoencoder as well? I am not aware of prior work that has done this. It could be that diffusion policies are overkill for the expressiveness needed to learn high value policies in the tasks under consideration, in which case these same tricks could be applied to train simpler VAE policies for example.\n\nLIMITATIONS:\nThe authors have not addressed technical limitations of their work. They have briefly addressed societal impact of better reinforcement learning methods in the appendix.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes to learn a policy for several offline RL tasks in the D4RL benchmark by parameterizing a policy with a diffusion model. The authors claim that their approach is computationally efficient and more compatible with several other RL approaches like maximum likelihood based approaches when compared to Diffusion-QL which is the primary baseline in their experimental studies.\n\nSTRENGTHS:\n1) The paper is well written and easy to read\n2) The authors give a clear and concise preface to the method and results in the introduction part of the paper which sets the flow for the rest of the paper.\n3) The authors compare their approach with the Diffusion-QL method which introduced the idea of formulating policies as a diffusion model. They explain the issues with the Diffusion-QL approach and try to overcome those shortcomings through EDP\n4) Their approach is compatible with max-likelihood based approaches such as IQL (which is one of the SoTA in Offline RL)\n5) The run extensive experiments and ablations on the D4RL benchmark and the results reported are consistently better than Diffusion-QL\n\nWEAKNESSES:\n1) Line 206: The authors claim that both policy approximations are empirically similar when calculating the objective though there is no evidence provided for the same. Would be good to know their intuition on that as well.\n2) There is no discussion around why EDP + IQL works significantly better in some tasks like kitchen, adroit whereas it under-performs in the locomotion tasks (where EDP + TD3) performs better. Would be good to know why one works better than another in these cases. Is there any investigation done in that direction?\n3) Barely no discussion around limitations of the approach. \n4) Would the approach be robust to data from multi-modalities? There is a mention of why that would be an issue with gaussian policies but there is no details on that in the latter parts of the paper.\n5) Would the method work on more complex offline-RL task like CARLA? (why was that excluded from the investigation? I understand that diffusion-QL does not test their method on this task but was it ever tried out?)\n\nQUESTIONS:\nPlease refer to some of the questions mentioned in the weakness section;\n\nMinor comments:\n\n1) Is the speed only reason for using the DPM-Solver? Are there other alternatives that were considered\n\n2) Please rephrase line 274 (Do not like phrasing the evaluation protocol of another peer-reviewed paper to be labelled as cheating)\n\nLIMITATIONS:\nI do not see any discussion around limitations of the approach or other assumptions made.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe focus of this paper is to enhance the diffusion policies introduced in Diffusion-QL for offline reinforcement learning. The authors address the challenges of training and sampling efficiency by incorporating action approximation and employing an advanced ODE solver for diffusion policies. They conducted extensive experiments to demonstrate the effectiveness of their proposed approach, known as Efficient Diffusion Policy (EDP).\n\nSTRENGTHS:\n1.\tThis paper is well structured and easy to follow.  \n2.\tThe conducted experiments are of enough quantity and quality.\n\nWEAKNESSES:\n> I think some claims need to be more careful or accurate. \n\n“… reduce the diffusion policy training time from 5 days …”\n“However, it takes tens to hundreds more time to train a diffusion policy than a diagonal Gaussian one”\t\n\nI tried Diffusion-QL before and I don’t think it takes about 5 days for training in their default setting (K=5). In my case, it is only about 20 hours, so I am not sure what setting is inferred here to be 5 days. \n\nMoreover, “In comparison, our training scheme only passes through the network once an iteration, no matter how big K is.” This statement is in correct. According to Algorithm 1 and the description in the paper, “To improve the efficiency of policy evaluation, we propose to replace the DDPM sampling in Eqn. (4) with DPM-Solver”, the policy evaluation part still needs the action samples with the full reverse process, which takes a number of steps even with DPM-Solver. \n\n> The novelty is kind of limited. In short words, EDP tried two things: 1. replace the real action samples generated by the full reverse process in policy improvement step with one-time real action prediction. 2. replace the original DDPM solver with DPM-solver. \n\n> The Energy-based Action Selection (EAS) is not a new thing and has been studied by other works, such as “Offline Reinforcement Learning via High-Fidelity Generative Behavior Modeling (ICLR 2023)”.\n\n> Following diffusion model conventions, I use $x$ to represent $a$ here. This paper uses $\\hat{x_{0, t}}$ to replace $x_0$ in policy improvement step. Note using $\\hat{x_{0, t}}$ to replace $x_0$ is not a common thing in diffusion models, since the $\\hat{x_{0, t}}$ has different meaning compared to $x_0$. $\\hat{x_{0, t}}$ usually represents the mean of all $\\hat{x_0}$ that could goes back to the true $x_0$ distribution from time step $t$. Hence, $\\hat{x_{0, t}}$ could be with much noise and averaged when $t$ is large. Not sure why here the $\\hat{x_{0, t}}$ could replace the $x_0$ in Eq (7). If so, why this procedure could not be applied onto Eq (6) as well?\n\nQUESTIONS:\nSee weakness.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposed EDP to address the existing limitation of diffusion policy in offline RL. EDP relies on an action approximation to construct actions from corrupted ones, thus avoiding running the Markov chain for action sampling at training. shows that EDP achieves 25+ speedup over Diffusion-QL at training time on the gym-locomotion tasks in D4RL. Extensive experiments by training EDP with various offline RL algorithms, including TD3, CRR, and IQL justify the superiority of the proposed method.\n\nSTRENGTHS:\n* The proposed method EDP focuses on improving the Diffusion-QL baseline from two aspects including computational cost and limited applicability to different RL algorithms, which is the major contribution of this paper. It provides a new way to apply a complicated generative model in RL.\n* Extensive evaluation of datasets of different domains is performed. EDP can outperform several widely-used algorithms by replacing the originally-used Gaussian policy with the refined diffusion policy. And both the training and sampling speed of EDP are way faster than that of Diffusion-QL. Besides, it seems EDP is easier to tune compared with Diffusion-QL as it uses the same hyperparameters for tasks belonging to the same domain and could achieve satisfactory performance.\n* The unreasonableness of evaluation protocol OMS is pointed out and RAT is adopted to justify the performance of the proposed method and baselines.\n\nWEAKNESSES:\n* The compared baselines  in Table 2 and Figure 1 are a little bit weak, which cannot demonstrate the superiority of EDP.  As the authors claimed that the proposed method outperforms pre-sota methods, more powerful model-free baselines should be included such as EDAC[1], RORL[2] on Locomotion tasks, X-QL[3], InAC[4], BPPO[5] on other domains. Otherwise, I believed the superiority of the proposed method is overclaimed to some extent.\n\n  [1] An, Gaon, et al. \"Uncertainty-based offline reinforcement learning with diversified q-ensemble.\" *Advances in neural information processing systems* 34 (2021): 7436-7447.\n\n  [2] Yang, Rui, et al. \"Rorl: Robust offline reinforcement learning via conservative smoothing.\" *Advances in Neural Information Processing Systems* 35 (2022): 23851-23866.\n\n  [3] Garg, Divyansh, et al. \"Extreme Q-Learning: MaxEnt RL without Entropy.\" *The Eleventh International Conference on Learning Representations*. 2022.\n\n  [4] Xiao, Chenjun, et al. \"The In-Sample Softmax for Offline Reinforcement Learning.\" *The Eleventh International Conference on Learning Representations*. 2022.\n\n  [5] Zhuang, Zifeng, et al. \"Behavior Proximal Policy Optimization.\" *The Eleventh International Conference on Learning Representations*. 2022.\n\nQUESTIONS:\n* It seems there exist mismatch of results between pre-sota  methods from Figure 1 and Table 2?\n* Line 227 points out that Tab 1 reflects the benefits of increasing K. However, it seems Tab 1 doesn't include this information?\n* The GPU memory needed by Diffusion-QL is quite large. I'm curious whether EDP has the same problem?\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-10"}