# Learning from Random Demonstrations: Offline Reinforcement Learning with
Importance-Sampled Diffusion Models

Zeyu Fang
George Washington University
<EMAIL_ADDRESS>
&Tian Lan
George Washington University
<EMAIL_ADDRESS>

###### Abstract

Generative models such as diffusion have been employed as world models in
offline reinforcement learning to generate synthetic data for more effective
learning. Existing work either generates diffusion models one-time prior to
training or requires additional interaction data to update it. In this paper,
we propose a novel approach for offline reinforcement learning with closed-
loop policy evaluation and world-model adaptation. It iteratively leverages a
guided diffusion world model to directly evaluate the offline target policy
with actions drawn from it, and then performs an importance-sampled world
model update to adaptively align the world model with the updated policy. We
analyzed the performance of the proposed method and provided an upper bound on
the return gap between our method and the real environment under an optimal
policy. The result sheds light on various factors affecting learning
performance. Evaluations in the D4RL environment show significant improvement
over state-of-the-art baselines, especially when only random or medium-
expertise demonstrations are available – thus requiring improved alignment
between the world model and offline policy evaluation.

##  1 Introduction

Offline Reinforcement Learning (RL) methods have received much recent
attention Levine et al. (2020); Prudencio et al. (2023), due to their
abilities to train policies based on offline datasets (also known as
demonstrations) that are collected using a behavior policy, rather than
through expensive (and sometimes dangerous) online interactions Kiran et al.
(2021). The available datasets can only include limited and fixed
transitions/trajectories. Thus, as the learned policy gradually deviates from
the behavior policy used for collecting data, the learned policy and the
estimated value function would overestimate out-of-distribution actions in
unseen dynamics. This is known as the distribution shift Kumar et al. (2019),
one of the key challenges for offline RL.

To this end, solutions have been proposed to augment the offline training
dataset by building world models Yu et al. (2021); Rigter et al. (2022);
Matsushima et al. (2020), besides the approach of policy regularization Kumar
et al. (2019); Rashidinejad et al. (2021). The idea is to learn a synthetic
model, known as the world model, representing the transition dynamics of the
underlying Markov Decision Process (MDP). Once trained with the offline
dataset, the world model can be leveraged to generate additional synthetic
trajectories for offline RL. Existing work include world models generated by
VAE Ha and Schmidhuber (2018); Hafner et al. (2023); Ozair et al. (2021), GAN
Eysenbach et al. (2022), Transformers Janner et al. (2021), and more recently
Diffusion Ding et al. (2024); Lu et al. (2023). However, most existing work
either obtain a world model one-time prior to policy training Yu et al.
(2020); Kidambi et al. (2020); Janner et al. (2019) or require additional
online interaction data to adapt the world model Kaiser et al. (2019); Hafner
et al. (2019), neither effectively mitigating distribution shift with only
offline data. Further, the return gap between world models (e.g., diffusion
models) and the real environment of such offline RL algorithms remains to be
analyzed.

Figure 1: The existing algorithms as shown in (a) use a fixed world model in
offline RL (or require additional data from the environment to update). Our
proposed approach, as illustrated in (b1) to (b3), adapts the diffusion model
from $\theta_{k}$ to $\theta_{k+1}$ with importance-sampling regarding the
updated policy $\pi_{k+1}$, and then evaluates the current policy $\pi_{k+1}$
with guided diffusion to obtain $\pi_{k+2}$. Each point refers to a sample in
the offline dataset, while darker points are given higher weight in loss
calculation for diffusion world model update. The performance of the proposed
method is analyzed in this paper.

This paper proposes a novel approach for offline RL with closed-loop policy
evaluation and world-model adaptation. Our approach encompasses two
collaborative components: (i) A guided diffusion world model to directly
evaluate the target policy by generating synthetic trajectories with actions
drawn from it; and (ii) An importance-sampled world model update to align the
world model with the target policy. The two components work in a closed-loop
operation throughout training, as illustrated in Figure 1. Existing work often
considers diffusion world models as a data synthesizer Lu et al. (2023) or
planner Janner et al. (2022); Ding et al. (2024). It generates additional
synthetic trajectories, following data collected using a behavior policy, to
support policy updates as shown in Figure 1(a). Our proposed approach, on the
other hand, continually adapts the diffusion model using importance sampling
with respect to the distribution shift between the current target policy
$\pi_{k}$ and the behavior policy $\pi_{b}$. The updated diffusion model is
then used to evaluate $\pi_{t}$ for policy improvement, with a sequence of
actions drawn from $\pi_{k}$ used as guidance/input to the world model. We
note that this guided policy evaluation requires iterating between next-state
generation using the guided diffusion world model and next-action sampling
from current target policy. It is different from previous work using diffusion
on trajectory synthesize Lu et al. (2023), planning Janner et al. (2022); Ding
et al. (2024), or policy representation Wang et al. (2022); Chen et al. (2022)
in RL.

We analyze the performance of the proposed learning algorithm, namely Adaptive
Diffusion World-Model for Policy Evaluation (ADEPT). It addresses the
distribution shift problem in offline RL, especially when only random and
medium-expertise demonstrations are available. We provide the bound of value
discrepancy between actual environment and our diffusion model under the same
policy, and show that the monotonic improvement can be guaranteed when the
one-step policy update under the model is larger than this bound. We further
decompose the bound into three factors: state transition error, reward
prediction error and policy shift, and discuss how ADEPT lowers these factors
to narrow the bound. To our best of knowledge, this is the first analysis for
offline RL with diffusion world models. We note that ADEPT works best with
datasets consisting of mainly random and medium demonstrations, where the
distribution shift becomes more severe as target policy moves toward optimum.
The proposed algorithm is evaluated on D4RL benchmark Fu et al. (2020) in
three MuJoCo environments. The results show that ADEPT improves the baseline
IQL method Kostrikov et al. (2021) with an average of 211.8% on random, 15.6%
on medium and 19.4% on medium-replay, and outperforms other SOTA offline RL
methods including diffusion-based algorithms.

The contributions of this work can be summarized as follows:

  1. 1.

We propose an offline RL algorithm, ADEPT, for a closed-loop operation of
policy evaluation with guided diffusion and world-model adaptation with
importance sampling.

  2. 2.

We provide theoretical proof of bounding the return gap between ADEPT (using
offline RL and diffusion world models) and actual environment under an optimal
policy.

  3. 3.

We evaluate our method on the D4RL benchmarks and demonstrate significant
improvement over state-of-the-art baselines, especially on random and medium
datasets.

##  2 Related Works

#### Offline RL.

Offline RL faces the distribution shift problem Kumar et al. (2020) due to
data collected using a specific behavior policy. Various methods have been
proposed to regularize an offline RL policy and address this issue. In
particular, MOReL Kidambi et al. (2020), MOPO Yu et al. (2020), VI-LCB
Rashidinejad et al. (2021) and COMBO Yu et al. (2021) develop different ways
to penalize reward or value function in unseen state and action pairs to
address the out-of-distribution issues. Other model-free methods, including
BCQ Fujimoto et al. (2019), IQL Kostrikov et al. (2021), CQL Kumar et al.
(2020), and TD3+BC Fujimoto and Gu (2021), develop different conservatism
mechanisms that are defined on action or value function, forcing the policy to
act more conservatively. Our proposed ADEPT framework can be combined with any
of these offline RL algorithms.

#### World Models for Offline RL.

The use of world models to generate synthetic data for offline RL was first
proposed in Ha and Schmidhuber (2018), utilizing VAE and RNN for predicting
state transitions. Following this approach, various world models with advanced
capabilities of fitting desired distributions have been proposed, including
convolutional U-networks Kaiser et al. (2019), vector-quantized autoencoders
Ozair et al. (2021), generative adversarial networks Eysenbach et al. (2022),
energy-based models Boney et al. (2020), Transformers Janner et al. (2021),
and diffusion Ding et al. (2024); Lu et al. (2023). They mainly use world
models for trajectory synthesis, with limited adaptability of the world models
(except for using additional online data Lu et al. (2023)).

#### The Use of Diffusion Models in RL.

Diffusion is a state-of-the-art technique for generating synthetic samples of
images and text data Ho et al. (2020). It was first introduced as planner or
policy representation in offline RL Janner et al. (2022), where the diffusion
model directly generates trajectories that are used for execution. This is
further extended to conditional actions Ajay et al. (2022), meta-RL Ni et al.
(2023), hierarchical tasks Li et al. (2023), multi-task problems He et al.
(2023), multi-agent tasks Zhu et al. (2023) and safe planning Xiao et al.
(2023). Diffusion models are also employed for policy expression Wang et al.
(2022); Chen et al. (2022), imitation learning Hegde et al. (2024) and reward
modeling Nuti et al. (2023). Diffusion models are adopted as a data
synthesizer to generate additional synthetic data based on offline datasets
before policy training Lu et al. (2023). Later, a conditional diffusion world
model was proposed to generate trajectories from current state and action, to
support offline value-based RL Ding et al. (2024). Different from these
existing works, we propose ADEPT for closed-loop policy evaluation and world-
model adaptation, and provide the theoretical analysis of the return gap
between ADEPT and actual environment.

##  3 Preliminary

#### Offline RL using World Models

We consider an unknown Markov Decision Process (MDP), referred to as the
environment. Supposing the MDP is fully-observable with discrete time, it
could be defined by the tuple $M=(S,A,P,R,\mu_{0},\gamma)$. $S$ and $A$ are
the state and action spaces, respectively. $P(s_{t+1}|s_{t},a)$ is the
transition probability and $R:S\times A\to\mathbb{R}$ is the reward function.
$\mu_{0}$ is the initial state distribution and $\gamma$ is the discount
factor. We consider an agent acts within the environment based on a policy
$\pi(a|s)$ repeatedly. In each time step, the agent receives a state $s_{t}$
and samples an action via its policy $a_{t}\sim\pi(\cdot|s_{t})$. The
environment transits into a new state $s_{t+1}\sim P(\cdot|s_{t},a_{t})$ and
returns a reward $r_{t}=R(s_{t},a_{t})$. After a whole episode of
interactions, a trajectory
$\tau=(s_{0},a_{0},r_{0},s_{1},\dots,s_{T},a_{T},r_{T})$ will be generated,
which contains states, actions and rewards of maximum length $T$. Based on
that, the goal of RL is to learn an optimal policy $\pi^{*}$ to maximize the
expectation of cumulative rewards from this MDP:
$\pi^{*}=\arg\max_{\pi}\mathbb{E}_{s,a\sim\pi}(\sum_{t=0}^{T}\gamma^{t}r_{t})$.

Specifically in offline model-based RL, only a dataset of trajectories
$\mathcal{D}$ is available. Therefore, a prediction model of the environment
is introduced, denoted as world model $\hat{M}$, to improve sample efficiency
for further learning and planning. Commonly the world model learns a single-
step transition approximating the real dynamics $M$ of the environment in a
supervised method based on $\mathcal{D}$. Hence, once a world model has been
trained, it could replace the real environment to generate synthetic
trajectories. Similar to standard RL, an initial state $s_{0}$ is sampled
first from datasets, and based on that the interactions start. After certain
length of steps $H$, referred as horizon, a synthetic trajectory
$\hat{\tau}=(\hat{s}_{0},\hat{a}_{0},\hat{r}_{0},...,\hat{s}_{H},\hat{a}_{H},\hat{r}_{H})$
is generated, in which
$\hat{s}_{t+1}\sim\hat{P}(\cdot|\hat{s}_{t},\hat{a}_{t})$ and
$\hat{r}_{t}=\hat{R}(s_{t},a_{t})$. These imaginary trajectories are added
into the experience buffer for RL algorithms to optimize its policy.

#### Diffusion Model

Diffusion model is one of the SOTA deep generative models for images and
videos. The purpose of diffusion model is to learn an underlying data
distribution $q(\bm{x_{0}})$ from a dataset
$\mathcal{D}=\left\\{\bm{x_{i}}\right\\}$. In Denoising Diffusion
Probabilistic Model (DDPM) Nichol and Dhariwal (2021), the synthetic data
generation is conducted by denoising real data $x_{0}$ from noises
$\mathcal{N}(\bm{0},\bm{I})$ with $K$ steps. To train such a denoising step,
we first formulate a predefined forward process
$q\left(\bm{x_{k}}|\bm{x_{k}-1}\right)=\mathcal{N}\left(\sqrt{\alpha_{k-1}}\bm{x_{k-1}},\sqrt{1-\alpha_{k-1}}\bm{I}\right)$
to add noise on real data in the dataset step by step. Finally the
distribution would be similar to Gaussian noises. Here,
$\alpha_{k}=1-\beta_{k}$, and $\beta_{1:K}$ is a certain various schedule. The
diffusion model define another parameterized reverse process
$p_{\theta}\left(\bm{x_{k-1}}|\bm{x_{k}}\right)=\mathcal{N}\left(\mu_{\theta}(\bm{x_{k}}),\bm{\Sigma_{k}}\right)$
to denoise the real data from the Gaussian noise
$\mathcal{N}\left(0,\bm{I}\right)$. By defining
$\bar{\alpha}_{k}=\prod_{i=1}^{k}\alpha_{i}$, $\mu_{\theta}$ and
$\bm{\Sigma_{k}}$ can be rewritten as follows:

$\displaystyle\mu_{\theta}(\bm{x_{k}})=\frac{1}{\sqrt{a_{k}}}\left(\bm{x_{k}}-\frac{\beta_{k}}{\sqrt{1-\bar{\alpha}_{k}}}\bm{\epsilon}_{\theta}(\bm{x_{k}},k)\right)\
\ {\rm and}\ \
\bm{\Sigma_{k}}=\beta_{k}\frac{1-\bar{\alpha}_{k-1}}{1-\bar{\alpha}_{k}}\bm{I}.$
(1)

Here, $\bm{\epsilon}_{\theta}$ is the parameterized noise prediction model to
be trained. Therefore, the loss function of the diffusion model is defined as:
$L(\theta)=\mathbb{E}_{k\sim[1,K],x_{0}\sim q,\bm{\epsilon}\sim
N(\bm{0},\bm{I})}(\lVert\bm{\epsilon}-\bm{\epsilon}_{\theta}(\bm{x}_{k},k)\rVert^{2})$,
where $\bm{\epsilon}$ is the real noise added in each step. When generating
the synthetic data $\hat{x_{0}}$, beginning with $\hat{x}_{K}$ sampled from
$\mathcal{N}(\bm{0},\bm{I})$, predicts noise possibly added in each step until
reaching $\hat{x}_{0}$.

In this work, the diffusion model is guided by current state and action to
predict the next state. Such a setting leads to modeling a conditional
distribution $q(\bm{x}|\bm{y})$, in which $y$ is the attribute or label of the
generated samples. To handle this, two kinds of conditional diffusion models
have been proposed: classifier-guided and classifier-free methods. In
classifier-guided methods, an extra parameterized classifier
$p_{\phi}(\bm{y}|\bm{x_{k}})$ is trained on noisy sample of $\bm{x}$ and label
$\bm{y}$. On the contrary, classifier-free methods don’t train a separate
classifier, but regard the condition as an extra input of the noise prediction
model. In that case, the noise prediction equation is replaced by:

$\bm{\epsilon}_{\theta}(\bm{x_{k}},k|\bm{y})=\bm{\epsilon}_{\theta}(\bm{x_{k}},\varnothing,k)+\omega(\bm{\epsilon}_{\theta}(\bm{x_{k}},\bm{y},k)-\bm{\epsilon}_{\theta}(\bm{x_{k}},\varnothing,k)),$
(2)

in which $\varnothing$ is the default value representing no conditions. In
ADEPT, we choose the classifier-free method to train the diffusion model.

##  4 Methodology

Figure 2: An overview of our ADEPT algorithm. It iteratively leverages a
guided diffusion world model to directly evaluate the target policy with
actions drawn from it, and then performs an importance-sampled world model
update to adaptively align the world model with the updated policy. An upper
bound on the return gap is analyzed in Section 4.3.

In this section, we illustrate details of ADEPT. As shown in Figure 2, the two
key components in ADEPT: policy evaluation on guided diffusion world model,
and importance-sampled world model update, work in a closed-loop operation
throughout the whole training process. The explanations of these two
components are covered in 4.1 and 4.2, respectively. In 4.3, we provide
theoretical derivation to bound the return gap of the proposed method.

###  4.1 Policy Evaluation using Guided Diffusion Model

Before the policy iteration, we first utilize the offline dataset to
initialize guided diffusion world model $\mathcal{M}_{\theta,\eta}$,
consisting of noise prediction model $\bm{\epsilon}_{\theta}$ and a multi-
layer perceptron $r_{\eta}$, to simulate the conditional distribution of the
one-step transition $P(s_{t+1}|s_{t},a_{t})$ and reward function
$R(s_{t},a_{t})$. The offline dataset is normalized by linearly mapping each
dimension of the state space to have $0$ mean and $1$ standard deviation.
While training the diffusion model, we sample a minibatch of tuples
$(s_{t},a_{t},r_{t},s_{t+1})$ consisting of the state, action, reward and next
state from the normalized offline datasets in each iteration. As mentioned in
3, we follow DDPM and adopt classifier-free method to learn the conditional
probabilities of the state transition based on current state $s_{t}$ and
action $a_{t}$, _i.e._ , we replace $x$ with $s_{t+1}$ and $y$ with
$(s_{t},a_{t})$ in equation (2). Since the reward function plays a significant
role in RL training, the model should emphasize more on its accuracy and
decouple reward function from state transition dynamics. Therefore, we
separately train $r_{\eta}(s_{t},a_{t},s_{t+1})$ to predict the reward
function and the terminal signal. Introducing $s_{t+1}$ as an extra input
significantly improves the accuracy of terminal signal prediction, and in
evaluation the next state is sampled and denoised by $\bm{\epsilon}_{\theta}$.

Once $\mathcal{M}_{\theta,\eta}$ is initialized, we could use it to interact
with the current policy $\pi_{\phi}$, generating data to evaluate the current
policy. At the beginning of each policy evaluating iteration,
$\mathcal{M}_{\theta,\eta}$ randomly samples a state from $\mathcal{D}$ to be
the start state $\hat{s}_{0}$, even though it may appear as a middle state in
the real trajectory. Based on the current state $\hat{s}_{t}$, an action
$\hat{a}_{t}\sim\pi_{\phi}(\hat{s}_{t})$ is sampled given target policy.
Conditioned by $\hat{s}_{t}$ and $\hat{a}_{t}$, the diffusion model generates
the next state $\hat{s}_{t+1}$ after $K$ denoising steps via
$\bm{\epsilon}_{\theta}$. Next, $r_{\eta}$ generates the reward $\hat{r}_{t}$
and terminal signal $d_{t}$, given $\hat{s_{t}},a_{t}$, and $\hat{s}_{t+1}$,
in which $\hat{s}_{t+1}$ and $\hat{r_{t}}$ are clipped into the range of the
offline dataset. Such an iteration continues till when $d_{t}$ is true or $t$
reaches horizon $H$. The generated trajectory $\hat{\tau}$ is added into the
replay buffer $\hat{\mathcal{D}}$, combined with $\mathcal{D}$ to improve the
policy via offline RL algorithms.

###  4.2 Importance-Sampled World Model Update

Once the policy is updated, there is a policy shift between $\pi_{\phi}$ and
the behavior policy $\pi_{\mathcal{D}}$ that collects $\mathcal{D}$. The
distribution estimation from $\mathcal{M}_{\theta,\eta}$ could lose accuracy
under the new policy. To handle this, we adopt the importance-sampling
technique to update $\mathcal{M}_{\theta,\eta}$ with offline dataset, guiding
$\hat{\mathcal{D}}$ towards the accurate distribution under the current
policy. This is achieved by re-weighting the loss function of multiple samples
to reduce the discrepancy between $\pi_{\phi}$ and $\pi_{\mathcal{D}}$. Even
if $\pi_{\mathcal{D}}$ is not available, it’s not hard to estimate the
behavior policy from the offline dataset via behavior cloning (BC) Nair et al.
(2018). For each transition
$\\{(s^{i}_{t_{i}},a^{i}_{t_{i}},r^{i}_{t_{i}},s^{i}_{t_{i}+1})\\}_{i=1}^{N}$
in the training batch, given their loss as $l_{i}(\theta,\eta)$, the
importance weight as $\omega_{i}$, the total loss $\mathcal{L}(\theta,\eta)$
of the whole training batch is calculated as:

$\mathcal{L}(\theta,\eta)=\frac{1}{N}\sum_{i=1}^{N}\omega_{i}l_{i}(\theta,\eta)=\frac{1}{N}\sum_{i=1}^{N}\frac{\pi_{\phi}(a^{i}_{t_{i}}|s^{i}_{t_{i}})}{\pi_{\mathcal{D}}(a^{i}_{t_{i}}|s^{i}_{t_{i}})}l_{i}(\theta,\eta).$
(3)

Generally, the state-action pairs that have more probabilities under the
current policy are associated with a larger weight in loss calculation.

Hyperparameters: denoising steps $K$, schedule $\beta_{1:K}$, evaluation steps
per epoch $N_{e}$, batchsize for policy improvement $B_{p}$, Batchsize for
world model update $B_{m}$, Horizon $H$

Input: offline dataset $\mathcal{D}$, diffusion world model
$\mathcal{M}_{\theta,\eta}$

Initialize target policy $\pi_{\phi}$, replay buffer
$\mathcal{\hat{D}}=\varnothing$; Normalized the dataset $\mathcal{D}$

Initialize $\mathcal{M}_{\theta,\eta}$ with $\mathcal{D}$ and
$\pi_{\mathcal{D}}$ till convergence:
$IWU(\pi_{\mathcal{D}},\mathcal{D},\theta,\eta)$

while  _not converged_ do

for  _$j=0\to N_{e}$_ do

$\mathcal{\hat{D}}\leftarrow\mathcal{\hat{D}}\cup
PE(\mathcal{M}_{\theta,\eta},\pi_{\phi})$

Sample
$\mathcal{B}=\\{(s^{i}_{t_{i}},a^{i}_{t_{i}},r^{i}_{t_{i}},s^{i}_{t_{i}+1})\\}_{i=1}^{B_{p}}\sim\mathcal{D}\cup\mathcal{\hat{D}}$

Update $\phi$ with $\mathcal{B}$ via offline RL methods

$IWU(\pi_{\phi},\mathcal{D},\theta,\eta)$

Subroutine: Importance-Sampled World-Model Update $(IWU)$:

Sample batch
$\\{(s^{i}_{t_{i}},a^{i}_{t_{i}},r^{i}_{t_{i}},s^{i}_{t_{i}+1})\\}_{i=1}^{B_{m}}\sim\mathcal{D}$

for  _$i=0\to B_{m}$_ do

$k\sim Uniform(\\{1,2,\dots,K\\})$;
$\bm{\epsilon}\sim\mathcal{N}(\bm{0},\bm{I})$

$s_{noise}=\sqrt{\bar{\alpha}_{k}}s^{i}_{t_{i}+1}+\sqrt{1-\bar{\alpha}_{k}}\bm{\epsilon}$

Get Importance-sampling weight $\omega_{i}$ under $\pi$ via Equation (3)

$l_{i}=\lVert\bm{\epsilon}-\bm{\epsilon}_{\theta}(s_{noise},(s^{i}_{t_{i}},a^{i}_{t_{i}}),k)\rVert^{2}+\lVert
r_{t}-r_{\eta}(s^{i}_{t_{i}},a^{i}_{t_{i}},s^{i}_{t_{i}+1})\rVert^{2}$

Calculate $\mathcal{L}(\theta,\eta)$ via Equation (3) and take gradient step
on it.

Subroutine: Policy Evaluation using Guided Diffusion Model ($PE$):

Randomly select a state from $\mathcal{D}$ as $\hat{s}_{0}$.

for  _$t=0\to H$_ do

$x_{K}\sim\mathcal{N}(\bm{0},\bm{I});\hat{a}_{t}\sim\pi(\hat{s}_{t})$

for  _$k=K-1\to 0$_ do

$x_{k}\sim\mathcal{N}(\mu_{\theta}(x_{k+1}|\hat{s_{t}},\hat{a}_{t}),\bm{\Sigma_{k}})$,
in which $\mu_{\theta}$ and $\bm{\Sigma_{k}}$ are defined in Eq. (1).

$\hat{s}_{t+1}$ =
$x_{0};\hat{r}_{t},d_{t}=r_{\eta}(\hat{s}_{t},\hat{a}_{t},\hat{s}_{t+1})$

If $d_{t}==True$ then break

Return the trajectory
$\hat{\tau}=(\hat{s}_{0},\hat{a}_{0},\hat{r}_{0},\hat{s}_{1},\dots\hat{s}_{t},\hat{a}_{r},\hat{r}_{t},\hat{s}_{t+1})$

Algorithm 1 Our proposed ADEPT algorithm

The complete training procedure of ADEPT is illustrated in Algorithm 1. Two
subroutines: policy evaluation and world-model update alternate iteratively
until convergence. In this work we select the state-of-the-art algorithm IQL
Kostrikov et al. (2021) and traditional method SAC Haarnoja et al. (2018) as
the offline RL methods to show the performance. The selection of
hyperparameters is a tradeoff between performance and efficiency, which will
be discussed and analyzed in the appendix.

###  4.3 Return Gap Analysis

To show that the expected return improvement could be guaranteed when adopting
diffusion model as the world model to train the policy, we wish to provide a
lower bound $C$ of the returns of the policy:

$J(\pi)=\mathbb{E}_{s_{t},a_{t}\sim\pi}\sum_{t=0}^{T}\gamma^{t}r_{t}.$ (4)

By optimizing the lower bound, the performance could be guaranteed non-
decreasing:

$J(\pi)\geq\hat{J}(\pi)-C.$ (5)

Here, $J(\pi)$ and $\hat{J}(\pi)$ are the returns of the policy under the real
environment and the world model used for training and evaluating, which is the
diffusion model in this work. With such a condition, we could guarantee
improvement under the actual environment if the returns under the simulating
model are promoted by at least $C$. Furthermore, $C$ is expected to be
expressed in terms of the error quantities of the model. We denote the reward
prediction error as $\hat{\varepsilon}_{r}$, the model transition error as
$\hat{\varepsilon}_{m}$, and policy distribution shift error as
$\hat{\varepsilon}_{\pi}$. Their detailed definitions are presented as
follows:

######  Definition 4.1.

We define $\hat{\varepsilon}_{r}$ to be the maximal expectation of total-
variation distance (TV-distance) of the probabilities between predicted reward
and true reward under the target policy $\pi_{\phi}$.

$\max_{t}\mathbb{E}_{s_{t},a_{t}\sim\pi_{\phi}}D_{TV}(P(r_{t}|s_{t},a_{t})\|P_{\eta}(r_{t}|s_{t},a_{t}))\leq\hat{\varepsilon}_{r}.$
(6)

Although in practice the reward prediction model also takes $s_{t+1}$ as an
input, it can still be considered only conditioned by $s_{t}$ and $a_{t}$,
since $s_{t+1}$ is generated from denoising steps.

######  Definition 4.2.

Similar to the Definition 4.1, $\hat{\varepsilon}_{m}$ is defined as the
maximal expected TV-distance of the probabilities between predicted next state
and true value under $\pi_{\mathcal{D}}$.

$\max_{t}\mathbb{E}_{s_{t},a_{t}\sim\pi_{\phi}}D_{TV}(P(s_{t+1}|s_{t},a_{t}\|P_{\theta}(s_{t+1}|s_{t},a_{t}))\leq\hat{\varepsilon}_{m},$
(7)

In practice, $\hat{\varepsilon}_{r}$ and $\hat{\varepsilon}_{m}$ could be
estimated by measuring the mean square error (MSE) between the predicted
states or rewards and their values in the offline dataset in practice. To be
noticed that $s_{t}$ and $a_{t}$ are sampled from the distribution under
$\pi_{\phi}$ in both equations, but the diffusion model can only be trained
from the offline dataset collected by $\pi_{\mathcal{D}}$. Therefore, it’s
necessary to take importance-sampling technique to keep the diffusion model
fitting the new distribution under $\pi_{\phi}$.

######  Definition 4.3.

$\hat{\varepsilon}_{\pi}$ is denoted by the maximal TV-distance of the policy
between iterations.

$\max_{s}D_{TV}(\pi(a|s)\|\pi_{\phi}(a|s))\leq\hat{\varepsilon}_{\pi}.$ (8)

This error measures how the evaluating policy $\pi$ has shifted from the
target policy $\pi_{\phi}$ that is evaluated under diffusion world model and
collects the dataset for policy improvement.

######  Theorem 4.4.

Given $\hat{\varepsilon}_{r}$, $\hat{\varepsilon}_{m}$ and
$\hat{\varepsilon}_{\pi}$, the bound $C$ between the true return and the IS-
diffusion model-based return can be expressed as follows:

$J(\pi)\geq\hat{J}_{\theta,\eta}(\pi)-2r_{max}(\frac{\hat{\varepsilon}_{r}+2\hat{\varepsilon}_{\pi}}{1-\gamma}+\frac{\gamma(2\hat{\varepsilon}_{\pi}+\hat{\varepsilon}_{m})}{(1-\gamma)^{2}}).$
(9)

###### Proof.

We give an outline of our proof next and display the whole proof in the
Appendix. To prove this, we first divide the discrepancy between $J(\pi)$ and
$\hat{J}_{\theta,\eta}(\pi)$ into $\mathcal{L}_{1}=J(\pi)-J(\pi_{\phi})$,
$\mathcal{L}_{2}=J(\pi_{\phi})-\hat{J}_{\theta,\eta}(\pi_{\phi})$ and
$\mathcal{L}_{3}=\hat{J}_{\theta,\eta}(\pi_{\phi})-\hat{J}_{\theta,\eta}(\pi)$.
Then we will prove several lemmas to show that if the reward is finite,
$\mathcal{L}_{1}$ and $\mathcal{L}_{3}$ could be bounded by
$\hat{\varepsilon}_{\pi}$, while $\mathcal{L}_{2}$ could be bounded by
$\hat{\varepsilon}_{m}$ and $\hat{\varepsilon}_{r}$. Finally we combined these
three bounds together to complete the theorem. ∎

######  Lemma 4.5.

We define $r_{max}=\max_{s,a}r(s,a)$. For any $\pi,\hat{\pi}$ under any two
environmental model $\mathcal{M}$ and $\hat{\mathcal{M}}$, satisfying

$\max_{t}\mathbb{E}_{s_{t},a_{t}\sim\pi}D_{TV}(P(r_{t}|s_{t},a_{t})\|\hat{P}(r_{t}|s_{t},a_{t}))\leq\varepsilon_{r},$
(10)

we have

$\left|J(\pi)-\hat{J}(\hat{\pi})\right|\leq
r_{max}\sum_{t}\gamma^{t}(2\varepsilon_{r}+2D_{TV}(P_{t}(s,a)\|\hat{P}_{t}(s,a))).$
(11)

Here, $P_{t}(s,a)$ denotes the joint distribution of state and action at time
step $t$ under $\mathcal{M}$ and $\pi$, while $\hat{P}_{t}(s,a)$ is under
$\hat{\mathcal{M}}$ and $\hat{\pi}$. This lemma decomposes the discrepancy
into the form of reward error and the state-action joint distribution shift.
Next, we further analyze the bound of the latter.

######  Lemma 4.6.

For $P_{t}$ under $\pi$ and $\mathcal{M}$, and $\hat{P}_{t}$ under $\hat{\pi}$
and $\hat{\mathcal{M}}$, there existing $\delta$ _s.t._

$\max_{t}\mathbb{E}_{s,a\sim\pi}D_{TV}(P_{t+1}(s^{\prime},a^{\prime}|s_{t}=s,a_{t}=a)\|\hat{P}_{t+1}(s^{\prime},a^{\prime}|s_{t}=s,a_{t}=a))\leq\delta.$
(12)

Then we have

$D_{TV}(P_{t}(s,a)\|\hat{P}_{t}(s,a))\leq
t\delta+D_{TV}(P_{0}(s,a)\|\hat{P}_{0}(s,a)).$ (13)

This lemma provides the upper bound of the state-action joint distribution
shift at time step $t$, when the initial state distribution is shared and the
state transition distribution under different policies and models is bounded.
Next step is to decompose the joint distribution into the form of
$\varepsilon_{m}$ and $\varepsilon_{\pi}$.

######  Lemma 4.7.

For any 2 different joint distributions of state and action pairs $P$ and
$\hat{P}$, we have

$D_{TV}(P(s,a)\|\hat{P}(s,a))\leq
D_{TV}(P(s)\|\hat{P}(s))+\max_{s}D_{TV}(\pi(a|s)\|\hat{\pi}(a|s)).$ (14)

To be noticed that this lemma is also correct for conditional distribution of
$s$ and $a$. By combining Lemma 4.6 and Lemma 4.7, the joint distribution
shift can be expressed with $\varepsilon_{r}$ and $\varepsilon_{m}$.
Therefore, when applied with the error defined in diffusion model, the final
bound $C$ could be concluded.

With this theorem, the monotonic improvement of the true return $J(\pi)$ is
guaranteed when the returns under IS-diffusion model
$\hat{J}_{\theta,\eta}(\pi)$ is improved by more than $C$. Besides, by
lowering $\hat{\varepsilon}_{r}$, $\hat{\varepsilon}_{m}$, and
$\hat{\varepsilon}_{\pi}$, this bound could be further narrowed. These errors
are mainly caused by the following reasons:

  * •

The compounding error: Though the diffusion model has lower prediction error
compared with conventional world models, as a one-step transition model, the
variation in each step accumulates as the compounding error, which
significantly increases $\hat{\varepsilon}_{m}$.

  * •

Out of distribution: While using diffusion world model for policy evaluation,
the action derived from the policy can drive the state out of distribution. In
that case, the generated state and reward become unstable or even illegal,
causing a high value of $\hat{\varepsilon}_{m}$ and $\hat{\varepsilon}_{r}$.

  * •

Policy shifting: While the target policy keeps getting improved with both the
offline dataset and generated replay buffer, it’s drifting away from the
original behavior policy that collects the data, leading to a high value of
$\hat{\varepsilon}_{\pi}$.

To handle these, in our method $H$ is set to be small enough to limit the
accumulating variations. Since the initial state is randomly chosen, the whole
dataset can be covered even with small horizons. The generated states and
rewards are clipped within the range of the dataset to avoid out-of-
distribution. Finally to solve the distribution shift problem, inspired by the
model adaption to evolving policy method in online RL Wang et al. (2023), we
adopt importance-sampling in offline RL to continuously align the diffusion
model $\mathcal{M}_{\theta,\eta}$ with the new distribution under current
policy $\pi_{\phi}$ to lower $\hat{\varepsilon}_{\pi}$.

##  5 Experiment

Our experiments are designed to evaluate: 1\. The performance of ADEPT with
adaptive diffusion world model and offline RL updates, compared with other
SOTA algorithms including diffusion-based methods. 2\. The effectiveness of
the proposed importance sampling and policy improvement in ADEPT. We train and
test our method on multiple environments and datasets in D4RL Fu et al. (2020)
to show the quantitative results, and further analyze our method with ablation
study.

###  5.1 Numerical Evaluation

In this section, we evaluate our proposed ADEPT with 2 widely-used offline RL
methods: IQLKostrikov et al. (2021) and SACHaarnoja et al. (2018) over the
MuJoCo locomotion environments (e.g., halfcheetah, walker2d, and hopper) and
on 3 different datasets (random, medium and medium-replay). We select a number
of SOTA algorithms as baselines, including model-free methods TD3+BC Fujimoto
and Gu (2021), CQL Kumar et al. (2020), model-based methods such as MOPO Yu et
al. (2020), COMBO Yu et al. (2021), and diffusion-based methods as SyntheER Lu
et al. (2023) and Diffuser Janner et al. (2022), as well as behavior cloning.
All experiments are conducted with the same training hyperparameters. The
comparison is summarized in Table 1.

Environment | Dataset | SAC | TD3+BC | CQL | IQL | MOPO | COMBO | SynthER | Diffuser | ADEPT+ (Ours)
---|---|---|---|---|---|---|---|---|---|---
IQL | SAC
halfcheetah | rnd | 30.5 | 11.3 | 35.4 | 12.5 | 35.4 | 38.8 | 17.2 | 3.6 | $34.5\pm 1.1$ | $\bm{39.5\pm 0.8}$
walker2d | 4.1 | 0.6 | 7.0 | 5.4 | 13.6 | 7.0 | 4.2 | 3.5 | $10.3\pm 2.2$ | $\bm{23.6\pm 0.9}$
hopper | 11.3 | 8.6 | 10.8 | 7.5 | 11.7 | 17.9 | 7.7 | 6.3 | $34.4\pm 0.9$ | $\bm{37.5\pm 5.7}$
halfcheetah | med | -4.3 | 48.1 | 44.4 | 47.4 | 42.3 | 54.2 | 49.6 | 42.8 | $55.2\pm 0.5$ | $\bm{56.6\pm 2.3}$
walker2d | 0.9 | 82.7 | 79.2 | 78.3 | 17.8 | 81.9 | 84.7 | 79.6 | $\bm{97.2\pm 6.5}$ | $70.9\pm 4.7$
hopper | 0.8 | 60.4 | 58.0 | 66.3 | 28.0 | 97.2 | 72.0 | 74.3 | $69.6\pm 3.5$ | $32.8\pm 3.9$
halfcheetah | mix | 29.9 | 44.8 | 46.2 | 44.2 | 53.1 | 55.1 | 46.6 | 37.7 | $51.2\pm 1.2$ | $\bm{59.3\pm 4.1}$
walker2d | 21.0 | 85.6 | 26.7 | 94.7 | 39.0 | 56.0 | 83.3 | 70.6 | $\bm{105.3\pm 1.4}$ | $85.0\pm 2.2$
hopper | 70.4 | 64.4 | 48.6 | 73.9 | 67.5 | 89.5 | 103.2 | 93.6 | $97.6\pm 3.7$ | $\bm{104.0\pm 2.7}$
Average | 18.3 | 45.2 | 39.6 | 47.8 | 55.3 | 34.3 | 52.1 | 45.8 | $\bm{61.7\pm 2.3}$ | $56.6\pm 3.0$

Table 1: The evaluation of ADEPT compared with other SOTA algorithms including
offline RL and diffusion-based algorithms, on D4RL MuJoCo environments with
random(rnd), medium(med) and medium-replay(mix) datasets. We show the mean of
standard deviation of the performance over 5 different seeds. ADEPT
significantly outperforms the baselines and achieves very stable results,
especially on random datasets due to its closed-loop operation supporting
iterative diffusion world model adaptation and offline policy improvement.

From our experimental result, the proposed ADEPT algorithms outperform the
existing SOTA offline RL algorithms in most of the environments, especially in
random dataset. This is consistent with our hypothesis that adapting the
diffusion world model in offline RL is more critical, when there is a lack of
expert demonstrations and the distribution shift becomes more severe as target
policy moves toward optimum. Compared to the original SAC method which
performs poorly due to lack of policy regularization, the diffusion model
generated data has significantly improved its performance, with an average of
119.2% and 104.7% gain on random and medium-replay datasets, respectively.

This is because adaptive diffusion model guided policy improvement could be
viewed as a conservative regularization method. For out-of-distribution
situations, the diffusion model generates more relevant state transition
results and corresponding reward function, which mitigates overestimation. For
IQL that has already regularized its value function, the adaptive diffusion
world model also provides synthetic data to complement the offline dataset,
leading to an average of 211.8% and 19.4% performance gain over IQL on random
and medium-replay datasets, respectively.

Compared with other diffusion-based offline RL algorithms, our method still
gets significant advances in random and medium-replay datasets. Such a result
shows that with a lack of expert demonstration and plenty of random data, a
closed-loop iterative algorithm for diffusion world model adaptation and
offline policy improvement could become more advantageous than using diffusion
models as a policy representation like Diffuser Janner et al. (2022), or
generating synthetic dataset one time before training as SynthER Lu et al.
(2023). We also note that such a method may not be necessary when the dataset
is collected by an expert policy, as less adaptation to the world model is
needed.

###  5.2 Ablation Study

In the ablation study, we intend to validate the necessity of importance
sampling and the effectiveness of each part in the close-loop ADEPT algorithm.
To accomplish this, we compare our methods under different settings: (1)
offline RL with no generated data used, _i.e._ , standard IQL and SAC; (2)
offline RL with diffusion model trained one-time before training, so no
importance sampling technique, _i.e._ , Diff+IQL and Diff+SAC; (3) replacing
policy improvement with behavior cloning to demonstrate the impact of offline
RL. These three settings show the importance of adaptive diffusion world model
with importance sampling, and offline reinforcement learning, respectively. We
also perform additional experiments to choose the hyperparameters and validate
the effectiveness of subordinate techniques like data normalization and
clipping. The results are shown in the Appendix.

Figure 3: The training curves of different methods on halfcheetah environment
for our ablation study. In particular, Diff+IQL and Diff+SAC generate
synthetic data while removing the use of importance sampling, using a
diffusion model trained one-time before training. IQL and SAC further remove
the use of any synthetic data. All other settings remain the same with ADEPT
methods. The results demonstrate the critical role of diffusion world model
adaptation in offline RL.

Figure 3 shows the training curves of different algorithms in the halfcheetah
environment in our ablation study. With no synthetic data, the original SAC
algorithm fails in medium and medium-replay datasets (even lower than behavior
cloning), while IQL fails in the random datasets. The diffusion world model
generated samples help improve the performance of these two methods in all
kinds of datasets. Besides, adopting importance sampling for diffusion model
adaptation improves the performance of SAC substantially in all three datasets
while only improving IQL significantly in random dataset. That is consistent
with the hypothesis that samples in the random dataset are more scattered, so
the importance sampling method for diffusion world model adaptation could have
a higher impact on the generated distribution, aligning the world model with
the target policy. The improvement on SAC shows that importance sampling
reduces the distribution shift and helps SAC to reach comparable performance
as IQL, even with no extra regularization methods used.

##  6 Conclusion

This paper proposes ADEPT to enable (i) a guided diffusion world model to
directly evaluate the target policy in offline reinforcement learning and (ii)
an importance-sampled world model update to adaptively align the world model
with the updated policy, in a closed-loop operation. Our analysis of the
algorithm provides an upper bound on the return gap and illuminates key
factors affecting the learning performance. Evaluations on D4RL show
significantly improved results over state-of-the-art baselines. On random and
medium-replay datasets, we obtain 119.2% and 104.7% performance improvements
for SAC, and 211.8% and 19.4% for IQL on average. Our work provides important
insights into the use of world model with offline reinforcement learning.
However, the theoretical analysis and experiments are based on simple
environments such as MuJoCo tasks. In practice, the three error sources could
be hard to reduce, especially in complicated partial-observable environments,
which is the limitation of this paper and remains to be further researched.

## References

  * Levine et al. (2020) Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu.  Offline reinforcement learning: Tutorial, review, and perspectives on open problems.  _arXiv preprint arXiv:2005.01643_ , 2020.
  * Prudencio et al. (2023) Rafael Figueiredo Prudencio, Marcos ROA Maximo, and Esther Luna Colombini.  A survey on offline reinforcement learning: Taxonomy, review, and open problems.  _IEEE Transactions on Neural Networks and Learning Systems_ , 2023.
  * Kiran et al. (2021) B Ravi Kiran, Ibrahim Sobh, Victor Talpaert, Patrick Mannion, Ahmad A Al Sallab, Senthil Yogamani, and Patrick Pérez.  Deep reinforcement learning for autonomous driving: A survey.  _IEEE Transactions on Intelligent Transportation Systems_ , 23(6):4909–4926, 2021.
  * Kumar et al. (2019) Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine.  Stabilizing off-policy q-learning via bootstrapping error reduction.  _Advances in neural information processing systems_ , 32, 2019.
  * Yu et al. (2021) 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.
  * Rigter et al. (2022) Marc Rigter, Bruno Lacerda, and Nick Hawes.  Rambo-rl: Robust adversarial model-based offline reinforcement learning.  _Advances in neural information processing systems_ , 35:16082–16097, 2022.
  * Matsushima et al. (2020) Tatsuya Matsushima, Hiroki Furuta, Yutaka Matsuo, Ofir Nachum, and Shixiang Gu.  Deployment-efficient reinforcement learning via model-based offline optimization.  _arXiv preprint arXiv:2006.03647_ , 2020.
  * Rashidinejad et al. (2021) Paria Rashidinejad, Banghua Zhu, Cong Ma, Jiantao Jiao, and Stuart Russell.  Bridging offline reinforcement learning and imitation learning: A tale of pessimism.  _Advances in Neural Information Processing Systems_ , 34:11702–11716, 2021.
  * Ha and Schmidhuber (2018) David Ha and Jürgen Schmidhuber.  World models.  _arXiv preprint arXiv:1803.10122_ , 2018.
  * Hafner et al. (2023) Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap.  Mastering diverse domains through world models.  _arXiv preprint arXiv:2301.04104_ , 2023.
  * Ozair et al. (2021) Sherjil Ozair, Yazhe Li, Ali Razavi, Ioannis Antonoglou, Aaron Van Den Oord, and Oriol Vinyals.  Vector quantized models for planning.  In _international conference on machine learning_ , pages 8302–8313. PMLR, 2021.
  * Eysenbach et al. (2022) Benjamin Eysenbach, Alexander Khazatsky, Sergey Levine, and Russ R Salakhutdinov.  Mismatched no more: Joint model-policy optimization for model-based rl.  _Advances in Neural Information Processing Systems_ , 35:23230–23243, 2022.
  * Janner et al. (2021) Michael Janner, Qiyang Li, and Sergey Levine.  Offline reinforcement learning as one big sequence modeling problem.  _Advances in neural information processing systems_ , 34:1273–1286, 2021.
  * Ding et al. (2024) Zihan Ding, Amy Zhang, Yuandong Tian, and Qinqing Zheng.  Diffusion world model.  _arXiv preprint arXiv:2402.03570_ , 2024.
  * Lu et al. (2023) Cong Lu, Philip J Ball, and Jack Parker-Holder.  Synthetic experience replay.  _arXiv preprint arXiv:2303.06614_ , 2023.
  * Yu et al. (2020) Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma.  Mopo: Model-based offline policy optimization.  _Advances in Neural Information Processing Systems_ , 33:14129–14142, 2020.
  * Kidambi et al. (2020) Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims.  Morel: Model-based offline reinforcement learning.  _Advances in neural information processing systems_ , 33:21810–21823, 2020.
  * Janner et al. (2019) Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine.  When to trust your model: Model-based policy optimization.  _Advances in neural information processing systems_ , 32, 2019.
  * Kaiser et al. (2019) Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, et al.  Model-based reinforcement learning for atari.  _arXiv preprint arXiv:1903.00374_ , 2019.
  * Hafner et al. (2019) Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi.  Dream to control: Learning behaviors by latent imagination.  _arXiv preprint arXiv:1912.01603_ , 2019.
  * Janner et al. (2022) Michael Janner, Yilun Du, Joshua B Tenenbaum, and Sergey Levine.  Planning with diffusion for flexible behavior synthesis.  _arXiv preprint arXiv:2205.09991_ , 2022.
  * Wang et al. (2022) 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.
  * Chen et al. (2022) Huayu Chen, Cheng Lu, Chengyang Ying, Hang Su, and Jun Zhu.  Offline reinforcement learning via high-fidelity generative behavior modeling.  _arXiv preprint arXiv:2209.14548_ , 2022.
  * Fu et al. (2020) Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine.  D4rl: Datasets for deep data-driven reinforcement learning.  _arXiv preprint arXiv:2004.07219_ , 2020.
  * Kostrikov et al. (2021) Ilya Kostrikov, Ashvin Nair, and Sergey Levine.  Offline reinforcement learning with implicit q-learning.  _arXiv preprint arXiv:2110.06169_ , 2021.
  * Kumar et al. (2020) Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine.  Conservative q-learning for offline reinforcement learning.  _Advances in Neural Information Processing Systems_ , 33:1179–1191, 2020.
  * Fujimoto et al. (2019) Scott Fujimoto, David Meger, and Doina Precup.  Off-policy deep reinforcement learning without exploration.  In _International conference on machine learning_ , pages 2052–2062. PMLR, 2019.
  * Fujimoto and Gu (2021) Scott Fujimoto and Shixiang Shane Gu.  A minimalist approach to offline reinforcement learning.  _Advances in neural information processing systems_ , 34:20132–20145, 2021.
  * Boney et al. (2020) Rinu Boney, Juho Kannala, and Alexander Ilin.  Regularizing model-based planning with energy-based models.  In _Conference on Robot Learning_ , pages 182–191. PMLR, 2020.
  * Ho et al. (2020) Jonathan Ho, Ajay Jain, and Pieter Abbeel.  Denoising diffusion probabilistic models.  _Advances in neural information processing systems_ , 33:6840–6851, 2020.
  * Ajay et al. (2022) Anurag Ajay, Yilun Du, Abhi Gupta, Joshua Tenenbaum, Tommi Jaakkola, and Pulkit Agrawal.  Is conditional generative modeling all you need for decision-making?  _arXiv preprint arXiv:2211.15657_ , 2022.
  * Ni et al. (2023) Fei Ni, Jianye Hao, Yao Mu, Yifu Yuan, Yan Zheng, Bin Wang, and Zhixuan Liang.  Metadiffuser: Diffusion model as conditional planner for offline meta-rl.  _arXiv preprint arXiv:2305.19923_ , 2023.
  * Li et al. (2023) Wenhao Li, Xiangfeng Wang, Bo Jin, and Hongyuan Zha.  Hierarchical diffusion for offline decision making.  In _International Conference on Machine Learning_ , pages 20035–20064. PMLR, 2023.
  * He et al. (2023) Haoran He, Chenjia Bai, Kang Xu, Zhuoran Yang, Weinan Zhang, Dong Wang, Bin Zhao, and Xuelong Li.  Diffusion model is an effective planner and data synthesizer for multi-task reinforcement learning.  _arXiv preprint arXiv:2305.18459_ , 2023.
  * Zhu et al. (2023) Zhengbang Zhu, Minghuan Liu, Liyuan Mao, Bingyi Kang, Minkai Xu, Yong Yu, Stefano Ermon, and Weinan Zhang.  Madiff: Offline multi-agent learning with diffusion models.  _arXiv preprint arXiv:2305.17330_ , 2023.
  * Xiao et al. (2023) Wei Xiao, Tsun-Hsuan Wang, Chuang Gan, and Daniela Rus.  Safediffuser: Safe planning with diffusion probabilistic models.  _arXiv preprint arXiv:2306.00148_ , 2023.
  * Hegde et al. (2024) Shashank Hegde, Sumeet Batra, KR Zentner, and Gaurav Sukhatme.  Generating behaviorally diverse policies with latent diffusion models.  _Advances in Neural Information Processing Systems_ , 36, 2024.
  * Nuti et al. (2023) Felipe Nuti, Tim Franzmeyer, and João F Henriques.  Extracting reward functions from diffusion models.  _arXiv preprint arXiv:2306.01804_ , 2023.
  * Nichol and Dhariwal (2021) Alexander Quinn Nichol and Prafulla Dhariwal.  Improved denoising diffusion probabilistic models.  In _International conference on machine learning_ , pages 8162–8171. PMLR, 2021.
  * Nair et al. (2018) Ashvin Nair, Bob McGrew, Marcin Andrychowicz, Wojciech Zaremba, and Pieter Abbeel.  Overcoming exploration in reinforcement learning with demonstrations.  In _2018 IEEE international conference on robotics and automation (ICRA)_ , pages 6292–6299. IEEE, 2018.
  * Haarnoja et al. (2018) Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine.  Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor.  In _International conference on machine learning_ , pages 1861–1870. PMLR, 2018.
  * Wang et al. (2023) Xiyao Wang, Wichayaporn Wongkamjan, Ruonan Jia, and Furong Huang.  Live in the moment: Learning dynamics model adapted to evolving policy.  In _International Conference on Machine Learning_ , pages 36470–36493. PMLR, 2023.
  * Zhang et al. (2024a) Zuyuan Zhang, Hanhan Zhou, Mahdi Imani, Taeyoung Lee, and Tian Lan.  Collaborative ai teaming in unknown environments via active goal deduction.  _arXiv preprint arXiv:2403.15341_ , 2024a.
  * Zou et al. (2024) Yifei Zou, Zuyuan Zhang, Congwei Zhang, Yanwei Zheng, Dongxiao Yu, and Jiguo Yu.  A distributed abstract mac layer for cooperative learning on internet of vehicles.  _IEEE Transactions on Intelligent Transportation Systems_ , 2024.
  * Gao et al. (2024) Mengtong Gao, Yifei Zou, Zuyuan Zhang, Xiuzhen Cheng, and Dongxiao Yu.  Cooperative backdoor attack in decentralized reinforcement learning with theoretical guarantee.  _arXiv preprint arXiv:2405.15245_ , 2024.
  * Zhang et al. (2024b) Zuyuan Zhang, Mahdi Imani, and Tian Lan.  Modeling other players with bayesian beliefs for games with incomplete information.  _arXiv preprint arXiv:2405.14122_ , 2024b.
  * Wáng et al. (2023) Yì Xiáng J Wáng, Zhi-Hui Lu, Jason CS Leung, Ze-Yu Fang, and Timothy CY Kwok.  Osteoporotic-like vertebral fracture with less than 20% height loss is associated with increased further vertebral fracture risk in older women: the mros and msos (hong kong) year-18 follow-up radiograph results.  _Quantitative Imaging in Medicine and Surgery_ , 13(2):1115, 2023.
  * Fang et al. (2022) Zeyu Fang, Jian Zhao, Mingyu Yang, Wengang Zhou, Zhenbo Lu, and Houqiang Li.  Coordinate-aligned multi-camera collaboration for active multi-object tracking.  _arXiv preprint arXiv:2202.10881_ , 2022.
  * Fang et al. (2023) Zeyu Fang, Jian Zhao, Wengang Zhou, and Houqiang Li.  Implementing first-person shooter game ai in wild-scav with rule-enhanced deep reinforcement learning.  In _2023 IEEE Conference on Games (CoG)_ , pages 1–8. IEEE, 2023.
  * Zhou et al. (2022) Hanhan Zhou, Tian Lan, and Vaneet Aggarwal.  Pac: Assisted value factorization with counterfactual predictions in multi-agent reinforcement learning.  _Advances in Neural Information Processing Systems_ , 35:15757–15769, 2022.
  * Zhou et al. (2023) Hanhan Zhou, Tian Lan, and Vaneet Aggarwal.  Value functions factorization with latent state information sharing in decentralized multi-agent policy gradients.  _IEEE Transactions on Emerging Topics in Computational Intelligence_ , 2023.
  * Mei et al. (2023) Yongsheng Mei, Hanhan Zhou, Tian Lan, Guru Venkataramani, and Peng Wei.  Mac-po: Multi-agent experience replay via collective priority optimization.  _arXiv preprint arXiv:2302.10418_ , 2023.
  * Chen et al. (2023) Jingdi Chen, Hanhan Zhou, Yongsheng Mei, Gina Adam, Nathaniel D Bastian, and Tian Lan.  Real-time network intrusion detection via decision transformers.  _arXiv preprint arXiv:2312.07696_ , 2023.
  * Chen et al. (2021) Jingdi Chen, Yimeng Wang, and Tian Lan.  Bringing fairness to actor-critic reinforcement learning for network utility optimization.  In _IEEE INFOCOM 2021-IEEE Conference on Computer Communications_ , pages 1–10. IEEE, 2021.

##  Appendix A Proof of the Theorem

In this appendix, we show the complete proof for the return gap bound proposed
in the main paper.

######  Lemma A.1.

For any 2 different joint distribution of state and action pairs $P$ and
$\hat{P}$ under two policies $\pi$ and $\hat{pi}$, we have

$D_{TV}(P(s,a)\|\hat{P}(s,a))\leq
D_{TV}(P(s)\|\hat{P}(s))+\max_{s}D_{TV}(\pi(a|s)\|\hat{\pi}(a|s)).$

###### Proof.

$\displaystyle D_{TV}(P(s,a)\|\hat{P}(s,a))$
$\displaystyle=\frac{1}{2}\sum_{s,a}|P(s,a)-\hat{P}(s,a)|$
$\displaystyle=\frac{1}{2}\sum_{s,a}|P(s)\pi(a|s)-\hat{P}(s,a)\hat{\pi}(a|s)|$
$\displaystyle=\frac{1}{2}\sum_{s,a}|P(s)\pi(a|s)-P(s)\hat{\pi}(a|s)+P(s)\hat{\pi}(a|s)-\hat{P}(s,a)\hat{\pi}(a|s)|$
$\displaystyle\leq\frac{1}{2}\sum_{s,a}P(s)|\pi(a|s)-\hat{\pi}(a|s)|+\frac{1}{2}\sum_{s,a}|P(s)-\hat{P}(s)|\hat{\pi}(a|s)$
$\displaystyle=\mathbb{E}_{s\sim
P}\left[D_{TV}(\pi(a|s)\|\hat{\pi}(a\|s))\right]+\frac{1}{2}\sum_{s}|P(s)-\hat{P}(s)|$
$\displaystyle=\mathbb{E}_{s\sim
P}\left[D_{TV}(\pi(a|s)\|\hat{\pi}(a\|s))\right]+D_{TV}(\sum_{s}|P(s)-\hat{P}(s))$
$\displaystyle\leq
D_{TV}(\sum_{s}|P(s)-\hat{P}(s))+\max_{s}D_{TV}(\pi(a|s)\|\hat{\pi}(a|s))$

∎

To be noticed that this equation can be extended to conditional probabilities:

$D_{TV}(P(s^{\prime},a^{\prime}|s,a)\|\hat{P}(s^{\prime},a^{\prime}|s,a))\leq
D_{TV}(P(s^{\prime}|s,a)\|\hat{P}(s^{\prime}|s,a))+\max_{s^{\prime}}D_{TV}(\pi(a^{\prime}|s^{\prime})\|\hat{\pi}(a^{\prime}|s^{\prime})).$

######  Lemma A.2.

For $P_{t}$ under $\pi$ and $\mathcal{M}$, and $\hat{P}_{t}$ under $\hat{\pi}$
and $\hat{\mathcal{M}}$ at time step $t$, given
$P_{0}(s)=\hat{P}_{0}(s),\forall s$, and there existing $\delta$ _s.t._

$\max_{t}\mathbb{E}_{s,a\sim\pi}D_{TV}(P_{t+1}(s^{\prime},a^{\prime}|s_{t}=s,a_{t}=a)\|\hat{P}_{t+1}(s^{\prime},a^{\prime}|s_{t}=s,a_{t}=a))\leq\delta.$
(15)

Then we have

$D_{TV}(P_{t}(s,a)\|\hat{P}_{t}(s,a))\leq
t\delta+D_{TV}(P_{0}(s,a)\|\hat{P}_{0}(s,a)).$ (16)

###### Proof.

$\displaystyle\left|P_{t}(s^{\prime},a^{\prime})-\hat{P}_{t}(s^{\prime},a^{\prime})\right|=\left|\sum_{s,a}\left(P_{t-1}(s,a)P(s^{\prime},a^{\prime}|s,a)-\hat{P}_{t-1}(s,a)\hat{P}(s^{\prime},a^{\prime}|s,a)\right)\right|$
$\displaystyle\leq\sum_{s,a}\left|\left(P_{t-1}(s,a)P(s^{\prime},a^{\prime}|s,a)-\hat{P}_{t-1}(s,a)\hat{P}(s^{\prime},a^{\prime}|s,a)\right)\right|$
$\displaystyle=\sum_{s,a}\left|\left(P_{t-1}(s,a)\left(P(s^{\prime},a^{\prime}|s,a)-\hat{P}(s^{\prime},a^{\prime}|s,a)\right)\right)+\hat{P}(s^{\prime},a^{\prime}|s.a)\left(P_{t-1}(s,a)-\hat{P}_{t-1}(s,a)\right)\right|$
$\displaystyle\leq\sum_{s,a}\left(P_{t-1}(s,a)\left|P(s^{\prime},a^{\prime}|s,a)-\hat{P}(s^{\prime},a^{\prime}|s.a)\right|+\hat{P}(s^{\prime},a^{\prime}|s,a)\left|P_{t-1}(s,a)-\hat{P}_{t-1}(s,a)\right|\right)$
$\displaystyle=\mathbb{E}_{s,a\sim
P_{t-1}}\left|P(s^{\prime},a^{\prime}|s,a)-\hat{P}(s^{\prime},a^{\prime}|s,a)\right|+\sum_{s,a}\hat{P}(s^{\prime},a^{\prime}|s,a)\left|P_{t-1}(s,a)-\hat{P}_{t-1}(s,a)\right|$

Therefore, we have:

$\displaystyle
D_{TV}(P_{t}(s,a)\|\hat{P}_{t}(s,a))=\frac{1}{2}\sum_{s^{\prime},a^{\prime}}\left|P_{t}(s^{\prime},a^{\prime})-\hat{P}_{t}(s^{\prime},a^{\prime})\right|$
$\displaystyle\leq\frac{1}{2}\sum_{s^{\prime},a^{\prime}}\left(\mathbb{E}_{s,a\sim
P_{t-1}}\left|P(s^{\prime},a^{\prime}|s,a)-\hat{P}(s^{\prime},a^{\prime}|s,a)\right|+\sum_{s,a}\hat{P}(s^{\prime},a^{\prime}|s,a)\left|P_{t-1}(s,a)-\hat{P}_{t-1}(s,a)\right|\right)$
$\displaystyle=\mathbb{E}_{s,a\sim
P_{t-1}}D_{TV}(P(s^{\prime},a^{\prime}|s,a)\|\hat{P}(s^{\prime},a^{\prime}|s,a))+\frac{1}{2}\sum_{s^{\prime},a^{\prime}}\sum_{s,a}\hat{P}(s^{\prime},a^{\prime}|s,a)\left|P_{t-1}(s,a)-\hat{P}_{t-1}(s,a)\right|$
$\displaystyle\leq\delta+\frac{1}{2}\sum_{s,a}\left|P_{t-1}(s,a)-\hat{P}_{t-1}(s,a)\right|=\delta+D_{TV}\left(P_{t-1}(s,a)\|\hat{P}_{t-1}(s,a)\right)$
$\displaystyle\leq t\delta+D_{TV}(P_{0}(s,a)\|\hat{P}_{0}(s,a))$

Furthermore, if the initial state distribution is the same for the two
environmental model, _i.e.,_ $P_{0}(s)=\hat{P}_{0}(s),\forall s$, then by
applying Lemma A.1 we have:

$\displaystyle D_{TV}(P_{t}(s,a)\|\hat{P}_{t}(s,a))$ $\displaystyle\leq
t\delta+D_{TV}(P_{0}(s)\|\hat{P}_{0}(s))+\max_{s}D_{TV}(\pi(a|s)\|\hat{\pi}(a|s))$
$\displaystyle\leq t\delta+\max_{s}D_{TV}(\pi(a|s)\|\hat{\pi}(a|s))$

∎

######  Lemma A.3.

We define $r_{max}=\max_{s,a}r(s,a)$. For any $\pi,\hat{\pi}$ under any two
environmental model $\mathcal{M}$ and $\hat{\mathcal{M}}$, satisfying

$\max_{t}\mathbb{E}_{s_{t},a_{t}\sim\pi}D_{TV}(P(r_{t}|s_{t},a_{t})\|\hat{P}(r_{t}|s_{t},a_{t}))\leq\varepsilon_{r},$
(17)

we have

$\left|J(\pi)-\hat{J}(\hat{\pi})\right|\leq
r_{max}\sum_{t}\gamma^{t}(2\varepsilon_{r}+2D_{TV}(P_{t}(s,a)\|\hat{P}_{t}(s,a))).$
(18)

###### Proof.

$\displaystyle\left|J(\pi)-\hat{J}(\hat{\pi})\right|=\left|\sum_{t}\gamma^{t}\sum_{s,a}\left(P_{t}(s,a)r(s,a)-\hat{P}_{t}(s,a)\hat{r}(s,a)\right)\right|$
$\displaystyle=\left|\sum_{t}\gamma^{t}\left(\sum_{s,a}\left[P_{t}(s,a)\left(r(s,a)-\hat{r}(s,a)\right)+\hat{r}(s,a)\left(P_{t}(s,a)-\hat{P}_{t}(s,a)\right)\right]\right)\right|$
$\displaystyle\leq\sum_{t}\gamma^{t}\left(\left|\sum_{s,a,r}rP_{t}(s,a)\left[P_{(}r|s,a)-\hat{P}(r|s,a)\right]\right|+r_{max}\sum_{s,a}\left|P_{t}(s,a)-\hat{P}_{t}(s,a)\right|\right)$
$\displaystyle\leq\sum_{t}\gamma^{t}r_{max}\left(\mathbb{E}_{s,a\sim\pi}\sum_{r}\left|P(r|s,a)-\hat{P}(r|s,a)\right|\right)$
$\displaystyle\leq
r_{max}\sum_{t}\gamma^{t}\left(2\varepsilon_{r}+2D_{TV}(P_{t}(s,a)\|\hat{P}_{t}(s,a))\right)$

∎

######  Theorem A.4.

Given $\hat{\varepsilon}_{r}$, $\hat{\varepsilon}_{m}$ and
$\hat{\varepsilon}_{\pi}$ with the Definition 4.1, 4.2 and 4.3, the bound $C$
between the true return $J$ and the IS-diffusion model-based return
$\hat{J}_{\theta,\eta}$ under a same evaluating policy $\pi$ can be expressed
as follows:

$J(\pi)\geq\hat{J}_{\theta,\eta}(\pi)-2r_{max}(\frac{\hat{\varepsilon}_{r}+2\hat{\varepsilon}_{\pi}}{1-\gamma}+\frac{\gamma(2\hat{\varepsilon}_{\pi}+\hat{\varepsilon}_{m})}{(1-\gamma)^{2}}).$
(19)

###### Proof.

We denote $\pi_{\mathcal{D}}$ as the policy collecting the trajectories in the
diffusion world model. The return gap could be separated into:

$|J(\pi)-\hat{J}_{\theta,\eta}(\pi)|\leq|J(\pi)-J(\pi_{\mathcal{D}})|+|J(\pi_{\mathcal{D}})-\hat{J}_{\theta,\eta}(\pi_{\mathcal{D}})|+|\hat{J}_{\theta,\eta}(\pi_{\mathcal{D}})-\hat{J}_{\theta,\eta}(\pi)|$

By applying Lemma A.3, we have:

$\left|J(\pi)-J(\pi_{\mathcal{D}})\right|\leq
r_{max}\sum_{t}\gamma^{t}\left(2\hat{\varepsilon}_{r}+2D_{TV}(P_{t}(s,a)\|\hat{P_{t}}(s,a))\right)$

Considering the condition in Lemma A.2, according to Lemma A.1 it can be
bounded by:

$\displaystyle\max_{t}\mathbb{E}_{s,a\sim\pi}D_{TV}(P_{t+1}(s^{\prime},a^{\prime}|s,a)\|\hat{P}_{t+1}(s^{\prime},a^{\prime}|s,a))$
$\displaystyle\leq\max_{t}\mathbb{E}_{s,a\sim\pi}D_{TV}(P(s^{\prime}|s,a)\|\hat{P}(s^{\prime}|s,a))+\max_{s}D_{TV}(\pi(a|s)\|\hat{\pi}(a|s))$

Since in this term the environmental model is the same while policy is
variant, the state transition error and reward error are zero, and we have:

$\max_{t}\mathbb{E}_{s,a\sim\pi}D_{TV}(P_{t+1}(s^{\prime},a^{\prime}|s,a)\|\hat{P}_{t+1}(s^{\prime},a^{\prime}|s,a))\leq\hat{\varepsilon}_{\pi}$

Next, we replace $\delta$ with $\hat{\varepsilon}_{\pi}$ in Lemma A.2 and get:

$\displaystyle\left|J(\pi)-J(\pi_{\mathcal{D}})\right|$ $\displaystyle\leq
2r_{max}\sum_{t}\gamma^{t}D_{TV}(P_{t}(s,a)\|\hat{P_{t}}(s,a))$
$\displaystyle\leq 2r_{max}\sum_{t}\gamma^{t}(t+1)\hat{\varepsilon}_{\pi}$

The third term could be analysed similarly, and we get:

$\displaystyle\left|\hat{J}_{\theta,\eta}(\pi_{\mathcal{D}})-\hat{J}_{\theta,\eta}(\pi)\right|\leq
2r_{max}\sum_{t}\gamma^{t}(t+1)\hat{\varepsilon}_{\pi}$

For the second term referring to the return gap of the behavior policy under
the real environment and model, there’s no policy shift but having the state
transition error:

$\max_{t}\mathbb{E}_{s,a\sim\pi}D_{TV}(P_{t+1}(s^{\prime},a^{\prime}|s,a)\|\hat{P}_{t+1}(s^{\prime},a^{\prime}|s,a))\leq\hat{\varepsilon}_{m}$

By replacing $\delta$ with $\hat{\varepsilon}_{m}$ in Lemma A.2 we get:

$\left|J(\pi_{\mathcal{D}})-\hat{J}_{\theta,\eta}(\pi_{\mathcal{D}})\right|\leq
2r_{max}\sum_{t}\gamma^{t}\left(\hat{\varepsilon}_{r}+(t+1)\hat{\varepsilon}_{m}\right)$

Finally, we summed the bounds of all the three terms and get:

$|J(\pi)-\hat{J}_{\theta,\eta}(\pi)|\leq
2r_{max}\sum_{t}\gamma^{t}\left(\hat{\varepsilon}_{r}+(t+1)\hat{\varepsilon}_{m}+2(t+1)\hat{\varepsilon}_{\pi}\right)=2r_{max}(\frac{\hat{\varepsilon}_{r}+2\hat{\varepsilon}_{\pi}}{1-\gamma}+\frac{\gamma(2\hat{\varepsilon}_{\pi}+\hat{\varepsilon}_{m})}{(1-\gamma)^{2}})$

∎

##  Appendix B Experiment Details

We use D4RL Fu et al. [2020] datasets for evaluation, their code could be
found at https://github.com/Farama-Foundation/D4RL. Their datasets are
licensed under the Creative Commons Attribution 4.0 License (CC BY), and their
code is licensed under the Apache 2.0 License.

###  B.1 Baselines

We select a number of SOTA baselines algorithms, including model-free methods
IQLKostrikov et al. [2021], SAC Haarnoja et al. [2018], TD3+BC Fujimoto and Gu
[2021], CQL Kumar et al. [2020], model-based methods such as MOPO Yu et al.
[2020], COMBO Yu et al. [2021], and diffusion-based methods as SyntheER Lu et
al. [2023] and Diffuser Janner et al. [2022]. We run the IQL and SAC code for
evaluation to get the result, while other results on D4RL dataset are obtained
from the original paper of each method. Specially, Diffuser doesn’t report
their results in random dataset, therefore we run its code from
https://github.com/jannerm/diffuser for evaluation. In addition, we cite to
the following works. Zhang et al. [2024a], Zou et al. [2024], Gao et al.
[2024], Zhang et al. [2024b], Wáng et al. [2023], Fang et al. [2022, 2023],
Zhou et al. [2022, 2023], Mei et al. [2023], Chen et al. [2023, 2021]

###  B.2 Computational Resources and Costs

All of the experiments in this paper are conducted on a server with an AMD
EPYC 7513 32-Core Processor CPU and an NVIDIA RTX A6000 GPU. The training of
the diffusion model costs approximately 1 hour for 1M gradient steps. The
offline training with importance-sampling costs nearly 5 hours for 1M gradient
steps.

###  B.3 Hyperparameter Settings

We show the hyperparameters used in the training process in Table 2 and 3.
These hyperparameters are shared in all of the environments.

Parameter |  Value
---|---
denoising steps |  10
$s$ for cosine schedule |  $1\times 10^{-4}$
train batchsize |  256
importance-sampling batchsize |  1024
learning rate |  $3\times 10^{-4}$
optimizer |  Adam
hidden dimension |  256
error model depth |  3
model training steps |  $1\times 10^{6}$
Horizon |  10
condition guidance weight |  0.1
Table 2: Diffusion Training Hyperparameters Parameter |  Value
---|---
$\gamma$ |  0.99
learning rate |  $3\times 10^{-4}$
train batch |  256
replay buffer size |  $1\times 10^{6}$
evaluation steps per epoch |  $1000$
gradient steps per epoch |  $1000$
training epochs |  $1000$
optimizer |  Adam
hidden dimension |  256
soft target updata rate (SAC) |  $5\times 10^{-3}$
$\alpha$ (IQL) |  0.005
$\tau$ (IQL) |  0.7
$\beta$ (IQL) |  3.0
Table 3: Offline RL Hyperparameters

##  Appendix C Additional Ablation Studies

In ADEPT, high values of the denoising steps $K$ and Importance-sampling
batchszie $B_{m}$ generally has higher accuracy on the state prediction and
better performance in model adjustment since it could cover more samples in
the datasets. However, increasing $K$ will extends the training procedure
since the diffusion model executes more denoising steps in state generation.
Also, a large $H$ decelerates the training of the RL agent, since there are
more policy inference steps in each updating.

Environment | baseline | $K=5$ | $K=20$ | $B_{m}=256$ | $B_{m}=2048$ | no clipping
---|---|---|---|---|---|---
halfcheetah-random | 39.5 | 30.0 | 33.2 | 42.1 | 40.5 | 18.52
walker2d-random | 23.6 | 24.0 | 21.5 | 8.2 | 23.1 | 13.32
hopper-random | 37.5 | 33.0 | 34.7 | 32.8 | 34.4 | 32.7
halfcheetah-medium | 56.6 | 55.4 | 60.4 | 56.1 | 55.2 | 2.9
walker2d-medium | 70.9 | 52.7 | 73.7 | 22.6 | 66.0 | 21.7
hopper-medium | 32.8 | 4.4 | 30.1 | 4.6 | 31.7 | 2.7
halfcheetah-medium-replay | 59.3 | 49.5 | 57.1 | 53.6 | 59.3 | 46
walker2d-medium-replay | 85.5 | 86.4 | 88.6 | 22.1 | 82.4 | 16.4
hopper-medium-replay | 104.0 | 77.8 | 104.6 | 98.7 | 99.1 | 41.0
Table 4: Ablation study for denoising steps

We conducted additional ablation study on the value of $K$ and $B_{m}$, and
the use of clipping technique. The results are shown in the Table 4. Based on
these results we can conclude that for simple environments like MuJoCo tasks,
improving $K$ from 10 to 20 doesn’t change much on performance, while the
evaluation time is nearly doubled. If $K$ is too low, the estimation error can
be large. The selection of $B_{m}$ could have similar influence on the
performance as $K$. Therefore, we select $K$ to be 10 and $B_{m}$ to be 1024
in our method. Besides, clipping of the generated state and reward have a
significant affects on performance improvement while avoiding out-of-
distribution outputs.