% Title: Untitled

% ArXiv ID: 2302.01877




\twocolumn[
\icmltitle{AdaptDiffuser: Diffusion Models as Adaptive Self-evolving Planners}




\icmlsetsymbol{equal}{*}

\begin{icmlauthorlist}
\icmlauthor{Zhixuan Liang}{hku}
\icmlauthor{Yao Mu}{hku}
\icmlauthor{Mingyu Ding}{hku,ucb}
\icmlauthor{Fei Ni}{tju}
\icmlauthor{Masayoshi Tomizuka}{ucb}
\icmlauthor{Ping Luo}{hku,shlab}

\end{icmlauthorlist}

\icmlaffiliation{hku}{Department of Computer Science, The University of Hong Kong, Hong Kong SAR}
\icmlaffiliation{ucb}{University of California, Berkeley, USA}
\icmlaffiliation{tju}{College of Intelligence and Computing, Tianjin University, Tianjin, China}
\icmlaffiliation{shlab}{Shanghai AI Laboratory, Shanghai, China}

\icmlcorrespondingauthor{Ping Luo}{pluo.lhi@gmail.com}

\icmlkeywords{Machine Learning, ICML}

\vskip 0.3in
]



\printAffiliationsAndNotice{}  % leave blank if no need to mention equal contribution

\begin{abstract}
Diffusion models have demonstrated their powerful generative capability in many tasks, with great potential to serve as a paradigm for offline reinforcement learning.
However, the quality of the diffusion model is limited by the insufficient diversity of training data, which hinders the performance of planning and the generalizability to new tasks.
This paper introduces \alias, an evolutionary planning method with diffusion that can self-evolve to improve the diffusion model hence a better planner, not only for seen tasks but can also adapt to unseen tasks.
\alias enables the generation of rich synthetic expert data for goal-conditioned tasks using guidance from reward gradients.
It then selects high-quality data via a discriminator to finetune the diffusion model, which improves the generalization ability to unseen tasks. 
Empirical experiments on two benchmark environments and two carefully designed unseen tasks in KUKA industrial robot arm and Maze2D environments demonstrate the effectiveness of \alias.
For example, \alias not only outperforms the previous art Diffuser~\cite{janner2022planning} by 20.8\% on Maze2D and 7.5\% on MuJoCo locomotion, but also adapts better to new tasks, e.g., KUKA pick-and-place, by 27.9\% without requiring additional expert data.
More visualization results and demo videos could be found on \href{https://adaptdiffuser.github.io/}{our project page}.
\end{abstract}


\begin{figure}[t]
\centering
\begin{subfigure}{0.49\textwidth}
  \centering
  \includegraphics[width=0.9\linewidth]{figs/teaser.pdf}
  \vspace{-2pt}
  \caption{Illustration of \alias.}
  \label{fig:teaser-a}
\end{subfigure}
\begin{subfigure}{0.49\textwidth}
  \centering
  \vspace{8pt}
  \includegraphics[width=0.99\linewidth]{figs/bar.pdf}
  \vspace{-2pt}
  \caption{Performance comparisons on three benchmarks.}
  \label{fig:teaser-b}
\end{subfigure}
\caption{Overall framework and performance comparison of \alias. It enables diffusion models to generate rich synthetic expert data using guidance from reward gradients of either seen or unseen goal-conditioned tasks.
Then, it iteratively selects high-quality data via a discriminator to finetune the diffusion model for self-evolving, leading to improved performance on seen tasks and better generalizability to unseen tasks. }
\vspace{-10pt}
\label{fig:teaser}
\end{figure}


\section{Introduction}
Offline reinforcement learning (RL)~\cite{levine2020offline,prudencio2022survey} aims to learn policies from previously collected offline data without interacting with the live environment.
Traditional offline RL approaches require fitting value functions or computing policy gradients, which are challenging due to limited offline data~\cite{agarwal2020optimistic,kumar2020conservative,wu2019behavior,kidambi2020morel}. 
Recent advances in generative sequence modeling~\cite{chen2021decision, janner2021offline, janner2022planning} provide effective alternatives to conventional RL problems by modeling the joint distribution of sequences of states, actions, rewards and values.
For example, Decision Transformer~\cite{chen2021decision} casts offline RL as a form of conditional sequence modeling, which allows more efficient and stable learning without the need to train policies via traditional RL algorithms like temporal difference learning~\cite{sutton1988learning}.
By treating RL as a sequence modeling problem, it bypasses the need of bootstrapping for long-term credit assignment, avoiding one of the ``deadly triad"~\cite{sutton2018reinforcement} challenges in reinforcement learning.





Therefore, devising an excellent sequence modeling algorithm is essential for the new generation of offline RL. 
The diffusion probability model~\cite{rombach2022high,ramesh2022hierarchical}, with its demonstrated success in generative sequence modeling for natural language processing and computer vision, presents an ideal fit for this endeavor.
It also shows great potential as a paradigm for planning and decision-making.
For example, diffusion-based planning methods~\cite{janner2022planning, ajay2022conditional, wang2022diffusion} train trajectory diffusion models based on offline data and apply flexible constraints on generated trajectories through reward guidance during sampling.
In consequence, diffusion planners show notable performance superiority compared with transformer-based planners like Decision Transformer~\cite{chen2021decision} and Trajectory Transformer~\cite{janner2021offline} on long horizon tasks, while enabling goal-conditioned rather than reward-maximizing control at the same time.



While diffusion-based planners have achieved success in certain areas, their performance is limited by the lack of diversity in their training data.
In decision-making tasks, the cost of collecting a diverse set of offline training data may be high, and this insufficient diversity would impede the ability of the diffusion model to accurately capture the dynamics of the environment and the behavior policy.
As a result, diffusion models tend to perform inferior when expert data is insufficient, and particularly when facing new tasks.
This raises a natural question: can we use the generated heterogeneous data by the reward-guided diffusion model to improve the diffusion model itself since it has powerful generative sequence modeling capability?
As diffusion-based planners can generate quite diverse ``dream" trajectories for multiple tasks which may be different from the original task the training data are sampled from, greatly superior to Decision Transformer \cite{chen2021decision}, enabling the diffusion model to be self-evolutionary makes it a stronger planner, potentially benefiting more decision-making requirements and downstream tasks.



In this paper, we present \alias, a diffusion-based planner for goal-conditioned tasks that can generalize to novel settings and scenarios through self-evolution (see Figure~\ref{fig:teaser}). 
Unlike conventional approaches that rely heavily on specific expert data, \alias uses gradient of reinforcement learning rewards, directly integrated into the sampling process, as guidance to generate diverse and heterogeneous synthetic demonstration data for both existing and unseen tasks. 
The generated demonstration data is then filtered by a discriminator, of which the high-quality ones are used to fine-tune the diffusion model,
resulting in a better planner with significantly improved self-bootstrapping capabilities on previously seen tasks and an enhanced ability of generalizing to new tasks.
As a consequence, \alias not only improves the performance of the diffusion-based planner on existing benchmarks, but also enables it to adapt to unseen tasks without the need for additional expert data.


It's non-trivial to construct and evaluate \alias for both seen and unseen tasks.% on existing benchmarks.
We first conduct empirical experiments on two widely-used benchmarks (MuJoCo~\cite{todorov2012mujoco} and Maze2d) of D4RL~\cite{fu2020d4rl} to verify the self-bootstrapping capability of \alias on seen tasks.
Additionally, we creatively design new pick-and-place tasks based on previous stacking tasks in the KUKA~\cite{schreiber2010fast} industrial robot arm environment, and introduce novel auxiliary tasks (e.g., collecting gold coins) in Maze2D. 
The newly proposed tasks and settings provide an effective evaluation of the generalization capabilities of \alias on unseen tasks.



Our contributions are three-fold:
\textbf{1)} We present \alias, allowing diffusion-based planners to self-evolve for offline RL by generating high-quality heterogeneous data with reward-integrated diffusion model directly and filtering out inappropriate examples with a discriminator. 
\textbf{2)} We apply our self-evolutionary \alias to unseen (zero-shot) tasks without any additional expert data, demonstrating its strong generalization ability and adaptability.
\textbf{3)} Extensive experiments on two widely-used offline RL benchmarks from D4RL as well as our carefully designed unseen tasks in KUKA and Maze2d environments validate the effectiveness of \alias.








\section{Related Works}

\textbf{Offline Reinforcement Learning.}
Offline RL~\cite{levine2020offline,prudencio2022survey} is a popular research field that aims to learn behaviors using only offline data such as those collected from previous experiments or human demonstrations, without the need to interact with the live environment from time to time at the training stage. 

However, in practice, offline RL faces a major challenge that standard off-policy RL methods may fail due to the overestimation of values, caused by the distribution deviation between the offline dataset and the policy to learn. Most conventional offline RL methods use action-space constraints or value pessimism~\cite{buckman2020importance} to overcome the challenge~\cite{agarwal2020optimistic, kumar2020conservative, siegel2020keep, wu2019behavior, yang2022regularized}.  
For example, conservative Q-learning (CQL)~\cite{kumar2020conservative} addresses these limitations by learning a conservative Q-function, ensuring the expected value under this Q-function is lower than its true value.




\textbf{Reinforcement Learning as Sequence Modeling.}
Recently, a new paradigm for Reinforcement Learning (RL) has emerged, in which RL is viewed as a generic sequence modeling problem. 
It utilizes transformer-style models to model trajectories of states, actions, rewards and values, and turns its prediction capability into a policy that leads to high rewards. 
As a representative, Decision Transformer (DT)~\cite{chen2021decision} leverages a causally masked transformer to predict the optimal action, which is conditional on an autoregressive model that takes the past state, action, and expected return (reward) into account. It allows the model to consider the long-term consequences of its actions when making a decision. And based on DT, Trajectory Transformer (TT)~\cite{janner2021offline} is proposed to utilize transformer architecture to model distributions over trajectories, repurposes beam search as a planning algorithm, and shows great flexibility across long-horizon dynamics prediction, imitation learning, goal-conditioned RL, and offline RL. Bootstrapped Transformer~\cite{wang2022bootstrapped} further incorporates the idea of bootstrapping into DT and uses the learned model to self-generate more offline data to further improve sequence model training. However, Bootstrapped Transformer could not integrate RL reward into the data synthesizing process directly and can only amplify homogeneous data trivially for its original task, which can boost the performance but cannot enhance the adaptability on another unseen task. Besides, such approaches lack flexibility in adapting to new reward functions and tasks in different environments, as the generated data is not suitable for use in new tasks or environments. 

Diffuser~\cite{janner2022planning} presents a powerful framework for trajectory generation using the diffusion probabilistic model, which allows the application of flexible constraints on generated trajectories through reward guidance during sampling.
The consequent work, Decision Diffuser~\cite{ajay2022conditional} introduces conditional diffusion with reward or constraint guidance for decision-making tasks, further enhancing Diffuser's performance. 
Additionally, Diffusion-QL~\citep{wang2022diffusion}, adds a regularization term to the training loss of the conditional diffusion model, guiding the model to learn optimal actions. 
Nevertheless, the performance of these methods is still limited by the quality of offline expert data, leaving room for improvement in adapting to new tasks or settings.



\textbf{Diffusion Probabilistic Model.}
Diffusion models are a type of generative model that represents the process of generating data as an iterative denoising procedure \cite{sohl2015deep, ho2020denoising}. They have made breakthroughs in multiple tasks such as image generation \cite{song2020denoising}, waveform generation \cite{chen2020wavegrad}, 3D shape generation \cite{zhou2021shape} and text generation \cite{austin2021structured}. These models, which learn the latent structure of the dataset by modeling the way in which data points diffuse through the latent space, are closely related to score matching \cite{hyvarinen2005score} and energy-based models (EBMs) \cite{lecun06atutorial, du2019implicit, nijkamp2019learning, grathwohl2020stein}, as the denoising process can be seen as a form of parameterizing the gradients of the data distribution \cite{song2019generative}. 

Moreover, in the sampling process, diffusion models allow flexible conditioning \cite{dhariwal2021diffusion} and have the ability to generate compositional behaviors \cite{du2020compositional}. It shows that diffusion models own promising potential to generate effective behaviors from diverse datasets and plan under different reward functions including those not encountered during training.



\section{Preliminary}
Reinforcement Learning is generally modeled as a Markov Decision Process (MDP) with a fully observable state space, denoted as $\mathcal{M}=(\mathcal{S}, \mathcal{A}, \mathcal{T}, \mathcal{R}, \gamma)$, where $\mathcal{S}$ is the state space and $\mathcal{A}$ is the action space. 
Besides, $\mathcal{T}$ is the state transition function with the dynamics of this discrete-time system that $\bs_{t+1}=\mathcal{T}(\bs_t, \ba_t)$ at state $\bs_t \in \mathcal{S}$ given the action $\ba_t \in \mathcal{A}$. $\mathcal{R}(\bs_t, \ba_t)$ defines the reward function and $\gamma\in(0,1]$ is the discount factor for future reward.

Considering the offline reinforcement learning as a sequence modeling task, the objective of trajectory optimization is to find the optimal sequence of actions $\ba_{0:T}^*$ that maximizes the expected return with planning horizon $T$, which is the sum of per time-step rewards or costs $R(\bs_t, \ba_t)$:
\begin{equation}
\ba_{0:T}^*=\underset{\ba_{0:T}}{\arg \max } \mathcal{J}(\bs_0, \ba_{0:T})=\underset{\ba_{0:T}}{\arg \max } \sum_{t=0}^T \gamma^t R(\bs_t, \ba_t).
\end{equation}

The sequence data generation methods utilizing diffusion probabilistic models \cite{sohl2015deep, ho2020denoising} pose the generation process as an iterative denoising procedure, denoted by $p_\theta(\btau^{i-1} \mid \btau^i)$ where $\btau$ represents a sequence and $i$ is an indicator of the diffusion timestep.

Then the distribution of sequence data is expanded with the step-wise conditional probabilities of the denoising process,
\vspace{-2pt}
\begin{equation}
p_\theta\left(\btau^0\right)=\int p\left(\btau^N\right) \prod_{i=1}^N p_\theta\left(\btau^{i-1} \mid \btau^i\right) \mathrm{d} \btau^{1: N}
\vspace{-2pt}
\end{equation}
where $p\left(\btau^N\right)$ is a standard normal distribution and $\btau^{0}$ denotes original (noiseless) sequence data. 

The parameters $\theta$ of the diffusion model are optimized by minimizing the evidence lower bound (ELBO) of negative log-likelihood of $p_\theta\left(\btau^0\right)$, similar to the techniques used in variational Bayesian methods.
\begin{equation}
\theta^*=\arg \min _\theta-\mathbb{E}_{\btau^0}\left[\log p_\theta\left(\btau^0\right)\right]
\label{eq:theta_optim}
\end{equation}
What's more, as the denoising process is the reverse of a forward diffusion process which corrupts input data by gradually adding noise and is typically denoted by $q\left(\btau^i \mid \btau^{i-1}\right)$, the reverse process can be parameterized as Gaussian under the condition that the forward process obeys the normal distribution and the variance is small enough \cite{feller2015theory}.
\begin{equation}
p_\theta\left(\boldsymbol{\tau}^{i-1} \mid \boldsymbol{\tau}^i\right)=\mathcal{N}\left(\boldsymbol{\tau}^{i-1} \mid \mu_\theta\left(\boldsymbol{\tau}^i, i\right), \Sigma^i\right)
\label{eq:p1}
\end{equation}
in which $\mu_{\theta}$ and $\Sigma$ are the mean and covariance of the Gaussian distribution respectively.

For model training, with the basis on Eq.~\ref{eq:theta_optim} and \ref{eq:p1}, \cite{ho2020denoising} proposes a simplified surrogate loss:
\begin{equation}
\label{eq:diff_loss}
    \mathcal{L}_{\text{denoise}}(\theta) \coloneqq \mathbb{E}_{i, \btau^{0} \sim q, \epsilon \sim \mathcal{N}}[||\epsilon - \epsilon_{\theta}(\btau^i, i)||^{2}]
\end{equation}
where $i \in \{0, 1, ..., N\}$ is the diffusion timestep, $\epsilon \sim \mathcal{N}(\boldsymbol{0}, \boldsymbol{I})$ is the target noise, and $\btau^{i}$ is the trajectory $\btau^{0}$ corrupted by noise $\epsilon$ for $i$ times. This is equivalent to predicting the mean $\mu_{\theta}$ of $p_\theta\left(\boldsymbol{\tau}^{i-1} \mid \boldsymbol{\tau}^i\right)$ as the function mapping from $\epsilon_{\theta}(\btau^i, i)$ to $\mu_{\theta}(\btau^i, i)$ is a closed-form expression. 



\section{Method}
In this section, we first introduce the basic planning with the diffusion method and its limitations. Then, we propose \alias, a novel self-evolved sequence modeling method for decision-making with the basis of diffusion probabilistic models.
\alias is designed to enhance the performance of diffusion models in existing decision-making tasks, especially the goal-conditioned tasks, and further improve their adaptability in unseen tasks without any expert data to supervise the training process.

\subsection{Planning with Task-oriented Diffusion Model}
Following previous work \cite{janner2022planning}, we can re-define the planning trajectory as a special kind of sequence data with actions as an additional dimension of states like:
\begin{equation}
\btau=\left[\begin{array}{llll}
\bs_0 & \bs_1 & ... & \bs_T \\
\ba_0 & \ba_1 & ... & \ba_T
\end{array}\right]
\label{eq:tau}
\end{equation}
Then we can use the diffusion probabilistic model to perform trajectory generation. However, the aim of planning is not to restore the original trajectory but to predict future actions with the highest reward-to-go, the offline reinforcement learning should be formulated as a conditional generative problem with guided diffusion models that have achieved great success on image synthesis \cite{dhariwal2021diffusion}. So, we drive the conditional diffusion process:
\begin{equation}
\label{eq:diff_plan}
    q(\btau^{i+1} | \btau^i), \;\;\;\; p_{\theta}(\btau^{i-1}|\btau^{i}, \boldsymbol{y}(\btau))
\end{equation}
where the new term $\boldsymbol{y}(\btau)$ is some specific information of the given trajectory $\btau$, such as the reward-to-go (return) $\mathcal{J}(\btau^0)$ of the trajectory, the constraints that must be satisfied by the trajectory and so on. On this basis, we can rewrite the optimization objective as,
\begin{equation}
 \theta^*=\arg \min _\theta-\mathbb{E}_{\btau^0}\left[\log p_\theta(\btau^{0} | \boldsymbol{y}(\btau^0))\right]
\label{eq:cond_gen_model}
\end{equation}

\begin{figure}[t]
\centering 
\includegraphics[width=0.95\linewidth]{figs/adaptdiffuser_2.pdf}
\caption{Overall framework of \alias. 
To improve the adaptability of the diffusion model to diverse tasks, rich data with distinct objectives is generated, guided by each task’s reward function. 
During the diffusion denoising process, we utilize a pre-trained denoising U-Net to progressively generate high-quality trajectories.
At each denoising time step, we take the
task-specific reward of a trajectory to adjust the gradient of state and action sequence, thereby creating trajectories that align with specific task objectives. 
Subsequently, the generated synthetic trajectory is evaluated by a discriminator to see if it meets the standards. If yes, it is incorporated into a data pool to fine-tune the diffusion model. 
The procedure iteratively enhances the generalizability of our model for both seen and unseen settings.}
\label{fig:main_frame} 
\end{figure}

Therefore, for tasks aiming to maximize the reward-to-go, we take $\mathcal{O}_{t}$ to denote the optimality of the trajectory at timestep $t$. And $\mathcal{O}_{t}$ obeys Bernoulli distribution with ${p(\mathcal{O}_t=1) = \exp (\gamma^{t} \mathcal{R}(\bs_t, \ba_t))}$. When $p(\mathcal{O}_{1:T} \mid \btau^{i})$ meets specific Lipschitz conditions, the conditional transition probability of the reverse diffusion process can be approximated as \cite{feller2015theory}:
\begin{equation}
\label{eq:guided}
p_\theta(\btau^{i-1} \mid \btau^{i}, \mathcal{O}_{1:T}) \approx \mathcal{N}(\btau^{i-1}; \mu_{\theta} + \alpha\Sigma g, \Sigma)
\end{equation}

\vspace{-20pt}
\begin{align*}
\text{where, } g &= \nabla_{\btau} \log p(\mathcal{O}_{1:T} \mid \btau) |_{\btau = \mu_{\theta}} \\
&= \sum_{t=0}^{T} \gamma^{t} \nabla_{\bs_t,\ba_t} \mathcal{R}(\bs_t, \ba_t) |_{(\bs_t,\ba_t)=\mu_t}
= \nabla_{\btau} \mathcal{J}(\mu_{\theta}).
\end{align*}
\vspace{-.45cm}

Besides, for tasks aiming to satisfy single point conditional constraint (e.g. goal conditioned tasks), the constraint can be simplified by substituting conditional values for the sampled values of all diffusion timesteps $i \in \{0, 1, ..., N\}$.

Although this paradigm has achieved competitive results with previous planning methods which are not based on diffusion models, it only performs conditional guidance during the reverse diffusion process and assumes the unconditional diffusion model is trained perfectly over the forward process. 
However, as depicted in Eq.~\ref{eq:guided}, the quality of generated trajectory $\btau$ depends not only on the guided gradient $g$ but more on the learned means $\mu_{\theta}$ and covariance $\Sigma$ of the unconditional diffusion model. If the learned $\mu_{\theta}$ deviates far from the optimal trajectory, no matter how strong the guidance $g$ is, the final generated result will be highly biased and of low quality. 
Then, learning from Eq.~\ref{eq:diff_loss}, the quality of $\mu_{\theta}$ hinges on the training data, the quality of which, however, is uneven across different tasks, especially on unseen tasks. Previous diffusion-based planning methods have not solved the problem which limits the performance of these methods on both existing and unseen tasks, and thus have poor adaptation ability.


\subsection{Self-evolved Planning with Diffusion}
Therefore, with the aim to improve the adaptability of these planners, we propose \alias, a novel self-evolved decision-making approach based on diffusion probabilistic models, to enhance the quality of the trained means $\mu_{\theta}$ and covariance $\Sigma$ of the forward diffusion process.
\alias relies on self-evolved synthetic data generation to enrich the training dataset which is denoted as $\btau_0$ and synthetic data fine-tuning to boost performance. After that, \alias follows the paradigm depicted in Eq.~\ref{eq:guided} to find the optimal action sequence for the given task with the guidance of reward gradients.

As shown in Figure \ref{fig:main_frame}, to implement \alias, we firstly generate a large number of synthetic demonstration data for unseen tasks which do not exist in the training dataset in order to simulate a wide range of scenarios and behaviors that the diffusion model may encounter in the real world.
This synthetic data is iteratively generated through the sampling process of the original diffusion probabilistic model $\theta_0^*$ with reward guidance, taking the advantage of its great generation ability. We will discuss the details of the synthetic data generation in Section \ref{sec:method_gen} and here we just abbreviate it as a function $\mathcal{G}(\mu_{\theta}, \Sigma, \nabla_{\btau} \mathcal{J}(\mu_{\theta}))$.

Secondly, we design a rule-based discriminator $\mathcal{D}$, with reward and dynamics consistency guidance, to select high-quality data from the generated data pool. Previous sequence modeling methods which predict the rewards $\mathcal{R}(\bs, \ba)$ simultaneously with generated states and actions are unable to solve the dynamics consistency problem that the actual next state with transition model $\bs' = \mathcal{T}(\bs, \ba)$ greatly deviates from the predicted next state. What's more, these deviated trajectories are taken as feasible solutions under previous settings.

To resolve this problem, \alias only takes the state sequence $\bs = \left[\bs_0, \bs_1, ..., \bs_T \right]$ of the generated trajectory and then performs state tracking control using a traditional or neural network-based inverse dynamics model $\mathcal{I}$ to derive real executable actions, denoted as $\widetilde{\ba}_t = \mathcal{I}(\bs_t, \bs_{t+1})$. This step ensures the action that does not violate the robot's dynamic constraints. After that, \alias performs $\widetilde{\ba}_t$ to obtain the revised next state $\widetilde{\bs}_{t+1}=\mathcal{T}\left(\widetilde{\bs}_t, \widetilde{\ba}_t\right)$, and then filters out the trajectories whose revised state $\widetilde{\bs}_{t+1}$ has a too large difference from the generated $\bs_{t+1}$ (measured by MSE $d=||\widetilde{\bs}_{t+1} - \bs_{t+1}||_2$). The remaining trajectories $\widetilde{\bs}$ are then used to predict the reward by $\widetilde{\mathcal{R}}=\mathcal{R}(\widetilde{\bs}, \widetilde{\ba})$ with the new actions $\widetilde{\ba}$ and are selected according to this reward. In this way, we can derive high-quality synthetic data to fine-tune the diffusion probabilistic model. 


We repeat this process multiple times in order to continually improve the model's performance and adapt it to new tasks, ultimately improving its generalization performance. So, it can be formulated as,
\begin{equation}
\begin{split}
&\theta_k^* \quad =\arg \min _\theta-\mathbb{E}_{\hat{\btau}_k}\left[\log p_\theta(\hat{\btau}_k | \boldsymbol{y}(\hat{\btau}_k))\right] \\
 &\btau_{k+1} = \mathcal{G}\left(\mu_{\theta_k^*}, \Sigma, \nabla_{\btau} \mathcal{J}(\mu_{\theta_k^*})\right) \\ 
 &\hat{\btau}_{k+1} = [\hat{\btau_k}, \mathcal{D}(\widetilde{\mathcal{R}}(\btau_{k+1})) ]
\end{split}
\end{equation}
where $k \in \{0, 1, ...\}$ is the number of iteration rounds and the initial dataset $\hat{\btau}_0 = \btau_0$.

\subsection{Reward-guided Synthetic Data Generation}
\label{sec:method_gen}
To improve the performance and adaptability of the diffusion probabilistic model on unseen tasks, we need to generate synthetic trajectory data using the learned diffusion model at the current iteration. We achieve it by defining a series of tasks with different goals and reward functions. 

\textbf{Continuous Reward Function.} For the tasks with continuous reward function, represented by MuJoCo \cite{todorov2012mujoco}, we follow the settings that define a binary random variable indicating the optimality with probability mapped from a continuous value, to convert the reward maximization problem to a continuous optimization problem. We can easily take Eq.~\ref{eq:guided} to generate synthetic results.

\textbf{Sparse Reward Function.} The reward function of tasks as typified by a goal-conditioned problem like Maze2D is a unit step function $\mathcal{J}(\btau) = \boldsymbol{\chi}_{\bs_g }(\btau)$ whose value is equal to 1 if and only if the generated trajectory contains the goal state $\bs_g$. 
The gradient of this reward function is Dirac delta function ~\cite{zhang2021dirac} which is not a classical function and cannot be adopted as guidance. However, if it is considered from the perspective of taking the limit, the constraint can be simplified as replacing all corresponding sampled values with constraints over the diffusion timesteps.

\textbf{Combination.} Many realistic tasks need these two sorts of reward functions simultaneously. For example, if there exists an auxiliary task in Maze2D environment that requires the planner to not only find a way from the start point to the goal point but also collect the gold coin in the maze. This task is more difficult and it's infeasible to add this constraint to the sparse reward term because there is no idea about which timestep the generated trajectory should pass the additional reward point (denoted as $\bs_c$). As a solution, we propose to combine these two sorts of methods and define an auxiliary reward guiding function to satisfy the constraints.
\vspace{-4pt}
\begin{equation}
    \mathcal{J}(\btau) = \sum_{t=0}^{T} ||\bs_t - \bs_c||_p
\label{eq:newj}
\vspace{-4pt}
\end{equation}
where $p$ represents p-norm. Then, with Eq.~\ref{eq:newj} we plug it into Eq.~\ref{eq:guided} as the marginal probability density function and force the last state of the generated trajectory $\btau^0$ to be $\bs_c$.
The generated trajectories that meet the desired criteria of the discriminator are added to the set of training data for the diffusion model learning as synthetic expert data. This process is repeated multiple times until a sufficient amount of synthetic data has been generated. By iteratively generating and selecting high-quality data based on the guidance of expected return and dynamics transition constraints, we can boost the performance and enhance the adaptability of the diffusion probabilistic model.

\section{Experiment}
\subsection{Benchmarks}
\textbf{Maze2D}: Maze2D \cite{fu2020d4rl} environment is a navigation task in which a 2D agent needs to traverse from a randomly designated location to a fixed goal location where a reward of 1 is given. No reward shaping is provided at any other location. The objective of this task is to evaluate the ability of offline RL algorithms to combine previously collected sub-trajectories in order to find the shortest path to the evaluation goal. Three maze layouts are available: ``umaze", ``medium", and ``large". The expert data for this task is generated by selecting random goal locations and using a planner to generate sequences of waypoints that are followed by using a PD controller to perform dynamic tracking. We also provide a method to derive more diverse layouts with ChatGPT in Appendix \ref{appendix:chatgpt}.

\textbf{MuJoCo}: 
MuJoCo \cite{todorov2012mujoco} is a physics engine that allows for real-time simulation of complex mechanical systems. It has  three typical tasks: Hopper, HalfCheetah, and Walker2d. Each task has 4 types of datasets to test the performance of an algorithm: ``medium", ``random", ``medium-replay" and ``medium-expert". The ``medium" dataset is created by training a policy with a certain algorithm and collecting $1$M samples. The ``random" dataset is created by using a randomly initialized policy. The ``medium-replay" dataset includes all samples recorded during training until the policy reaches a certain level of performance. There is also a ``medium-expert" dataset which is a mix of expert demonstrations and sub-optimal data.


\textbf{KUKA Robot}: The KUKA Robot \cite{schreiber2010fast} benchmark is a standardized evaluation tool that is self-designed to measure the capabilities of a robot arm equipped with a suction cup at the end. It consists of two tasks: conditional stacking \cite{janner2022planning} and pick-and-place. More details can be seen in Sec. \ref{sec:pick_place}. By successfully completing these tasks, the KUKA Robot benchmark can accurately assess the performance of the robot arm and assist developers in improving its design.


\subsection{Performance Enhancement on Existing Tasks}
\subsubsection{Experiments on Maze2d Environment}
\vspace{-1pt}
\textbf{Overall Performance.}
Navigation in Maze2D environment takes planners hundreds of steps to reach the goal location. Even the best model-free algorithms have to make great efforts to adequately perform credit assignments and reliably reach the target. 
We plan with \alias using the strategy of sparse reward function to condition on the start and goal location. We compare our method with the best model-free algorithms (IQL \citealt{kostrikov2021offline} and CQL \citealt{kumar2020conservative}), conventional trajectory
optimizer MPPI \cite{williams2015model} and previous diffusion-based approach Diffuser \cite{janner2022planning} in Table \ref{table:maze2d}. This comparison is fair because model-free methods can also identify the location of the goal point which is the only state with a non-zero reward.


\begin{table}[t]
\caption{
    \textbf{Offline Reinforcement Learning Performance in Maze2d Environment.} We show the results of \alias and previous planning methods to validate the bootstrapping effect of our method on a goal-conditioned task.
}
\label{table:maze2d}
\vspace{0.2cm}
\centering
\small
\tabcolsep 3pt
\begin{tabular}{cccccc}
\toprule
\multicolumn{1}{c}{\textbf{Environment}} & \textbf{MPPI} & \textbf{CQL} & \textbf{IQL} & \textbf{Diffuser} & \textbf{AdaptDiffuser} \\
\midrule
U-Maze & 33.2 & 5.7 & 47.4 & 113.9 & \textbf{135.1} \scriptsize{\raisebox{1pt}{$\pm 5.8$}}\\
Medium & 10.2 & 5.0 & 34.9 & 121.5 & \textbf{129.9} \scriptsize{\raisebox{1pt}{$\pm 4.6$}}\\
Large  & 5.1 & 12.5 & 58.6 & 123.0 & \textbf{167.9} \scriptsize{\raisebox{1pt}{$\pm 5.0$}}\\
\midrule
\multicolumn{1}{c}{\textbf{Average}} & 16.2 & 7.7 & 47.0 & 119.5 & \textbf{144.3} \hspace{.58cm} \\
\bottomrule
\end{tabular}
\end{table}


\begin{figure}[t]
    \begin{center}
    \centerline{\includegraphics[width=1.02\columnwidth]{figs/hardcase3.pdf}}
    \vspace{-0.05in}
    \caption{\textbf{Hard Cases of Maze2D with Long Planning Path.} Paths are generated in the Maze2D environment with a specified start 
    \protect{\raisebox{-.05cm}{\includegraphics[height=.3cm]{figs/start.pdf}}}
    and goal
    \protect{\raisebox{-.05cm}{\includegraphics[height=.3cm]{figs/goal.pdf}}}
    condition.}
    \label{fig:maze-hard}
    \end{center}
    \vspace{-20pt}
\end{figure}



\begin{table*}[t]
\caption{\small \textbf{Offline Reinforcement Learning Performance in MuJoCo Environment.} 
We report normalized average returns of D4RL tasks \citep{fu2020d4rl} in the table. And the mean and the standard error are calculated over 3 random seeds.
}
\label{table:mujoco}
\vspace{0.2cm}
\centering
\small
\tabcolsep 4.5pt
\begin{tabular}{ccccccccccccc}
\toprule
\textbf{Dataset} & \textbf{Environment} & \textbf{BC} & \textbf{CQL} & \textbf{IQL} & \textbf{DT} & \textbf{TT} & \textbf{MOPO} & \textbf{MOReL} & \textbf{MBOP} & \textbf{Diffuser} &  \textbf{AdaptDiffuser}\\ 
\midrule
Med-Expert & HalfCheetah & $55.2$ & $91.6$ & $86.7$ & $86.8$ & $95.0$ & $63.3$ & $53.3$ & $\textbf{105.9}$ & $88.9$ & $89.6$ \scriptsize{\raisebox{1pt}{$\pm 0.8$}} \\ 
Med-Expert & Hopper & $52.5$ & $105.4$ & $91.5$ & $107.6$ & $\textbf{110.0}$ & $23.7$ & $108.7$ & $55.1$ & $103.3$ & $\textbf{111.6}$ \scriptsize{\raisebox{1pt}{$\pm 2.0$}} \\ 
Med-Expert & Walker2d & $\textbf{107.5}$ & $\textbf{108.8}$ & $\textbf{109.6}$ & $\textbf{108.1}$ & $101.9$ & $44.6$ & $95.6$ & $70.2$ & $\textbf{106.9}$ & 
$\textbf{108.2}$ \scriptsize{\raisebox{1pt}{$\pm 0.8$}}
\\ 
\midrule
Medium & HalfCheetah & $42.6$ & $44.0$ & $\textbf{47.4}$ & $42.6$ & $\textbf{46.9}$ & $42.3$ & $42.1$ & $44.6$ & $42.8$ & $44.2$ \scriptsize{\raisebox{1pt}{$\pm 0.6$}} \\ 
Medium & Hopper & $52.9$ & $58.5$ & $66.3$ & $67.6$ & $61.1$ & $28.0$ & $\textbf{95.4}$ & $48.8$ & $74.3$ & $\textbf{96.6}$ \scriptsize{\raisebox{1pt}{$\pm 2.7$}}
 \\ 
Medium & Walker2d & $75.3$ & $72.5$ & $78.3$ & $74.0$ & $79.0$ & $17.8$ & $77.8$ & $41.0$ & $79.6$ & $\textbf{84.4}$ \scriptsize{\raisebox{1pt}{$\pm 2.6$}} \\ 
\midrule
Med-Replay & HalfCheetah & $36.6$ & $\textbf{45.5}$ & $\textbf{44.2}$ & $36.6$ & $41.9$ & $53.1$ & $40.2$ & $42.3$ & $37.7$ & $38.3$ \scriptsize{\raisebox{1pt}{$\pm 0.9$}}
\\ 
Med-Replay & Hopper & $18.1$ & $95.0$ & $94.7$ & $82.7$ & $91.5$ & $67.5$ & $\textbf{93.6}$ & $12.4$ & $\textbf{93.6}$ &
$\textbf{92.2}$ \scriptsize{\raisebox{1pt}{$\pm 1.5$}} \\ 
Med-Replay & Walker2d & $26.0$ & $77.2$ & $73.9$ & $66.6$ & $82.6$ & $39.0$ & $49.8$ & $9.7$ & $70.6$ & $\textbf{84.7}$ \scriptsize{\raisebox{1pt}{$\pm 3.1$}} \\ 
\midrule
\multicolumn{2}{c}{\textbf{Average}} & 51.9 & 77.6 & 77.0 & 74.7 & 78.9 & 42.1 & 72.9 & 47.8 & 77.5 & \textbf{83.4} \hspace{.58cm} \\
\bottomrule
\end{tabular}
\vspace{-4pt}
\end{table*}

As shown in Table \ref{table:maze2d}, scores achieved by \alias are over 125 in all maze sizes and are 20 points higher than those of Diffuser in average, indicating our method's strong effectiveness in goal-conditioned tasks.

\textbf{Visualization of Hard Cases.}
In order to more intuitively reflect the improvement of our method compared with previous Diffuser \cite{janner2022planning}, we select one difficult planning example of Maze2D-Medium and one of Maze2D-Large respectively for visualization, as shown in Figure \ref{fig:maze-hard}. Among the Maze2D planning paths with sparse rewards, the example with the longest path to be planned is the hardest one. Therefore, in Maze2D-Medium (Fig. \ref{fig:maze-hard} (a) (b)), we designate the start point as (1, 1) with goal point (6, 6), while in Maze2D-Large (Fig. \ref{fig:maze-hard} (c) (d)), we specify the start point as (1, 7) with goal point (9, 7) in the figure.

It can be observed from Fig. \ref{fig:maze-hard} that in Hard Case 1, \alias generates a shorter and smoother path than that generated by Diffuser. So, AdaptDiffuser achieves a larger reward. And in Hard Case 2, previous Diffuser method even fails to plan while our \alias derives a feasible path.

\subsubsection{Experiments on MuJoCo Environment}
\label{sec:mujoco}
MuJoCo tasks are employed to test the performance enhancement of our \alias learned from heterogeneous data of varying quality using the publicly available D4RL datasets \citep{fu2020d4rl}. We evaluate our approach with a number of existing algorithms that cover a variety of data-driven methodologies, including model-free RL algorithms like CQL \cite{kumar2020conservative} and IQL \cite{kostrikov2021offline}; return-conditioning approaches like Decision Transformer (DT) \cite{chen2021decision}; and model-based RL algorithms like Trajectory Transformer (TT) \cite{janner2021offline}, MOPO \cite{yu2020mopo}, MOReL \cite{kidambi2020morel}, and MBOP \cite{argenson2020model}. The results are shown in Table \ref{table:mujoco}. Besides, it is also worth noting that in the MuJoCo environment, the state sequence $\widetilde{\bs}$ derived by taking the generated actions $\ba$ is very close to the generated state sequence $\bs$, so we directly use $\widetilde{\mathcal{R}}(\bs, \ba)=\mathcal{R}(\bs, \ba)$ in this dataset.

\begin{figure}[t]
    \begin{center}
\centerline{\includegraphics[width=0.99\columnwidth]{figs/maze_adapt.pdf}}
    \vspace{-0.05in}
    \caption{\textbf{Maze2d Navigation with Gold Coin Picking Task.} Subfigures (a) (b) show the optimal path when there are no gold coins in the Maze. (The generated routes walk at the bottom of the Maze.) 
    And subfigures (c) (d) add additional reward \protect{\raisebox{-.10cm}{\includegraphics[height=.4cm]{figs/coins.pdf}}} in (4,2) position of the Maze. The planners generate new paths that pass through the gold coin as shown in subfigures (c) (d). (The newly generated routes walk in the middle of the maze.)
    }
    \label{fig:maze_adapt}
    \end{center}
    \vspace{-0.3in}
\end{figure}

Observed from the table, our method \alias is either competitive or outperforms most of the offline RL baselines across all three different locomotion settings. And more importantly, compared with Diffuser \cite{janner2022planning}, our method achieves higher reward in almost all the datasets and improves the performance greatly, especially in ``Hopper-Medium" and ``Walker2d-Medium" environments. 
We analyze that this is because the quality of the original data in the ``Medium dataset" is poor, so \alias has an evident effect on improving the quality of the training dataset, thus significantly enhancing the performance of the planner based on the diffusion probabilistic model. 
The results of the ``Medium-Expert" dataset verify this analysis because the quality of original data in the ``Medium-Expert" dataset (especially the Halfcheetah environment) has been good enough, making the generation of new data only has a little gain on the model performance.

\subsection{Adaptation Ability on Unseen Tasks}
\subsubsection{Maze2d with Gold Coin Picking Task}
On top of existing Maze2D settings, we carefully design a new task that requires the agent to navigate as well as pick all gold coins in the maze. We show an example with an additional reward in (4, 2) in Figure \ref{fig:maze_adapt}.

We can see that when there is no additional reward, both Diffuser \cite{janner2022planning} and our method \alias choose the shorter path at the bottom of the figure to reach the goal point. But, when additional reward is added in the (4, 2) position of the maze, both planners change to the path walking in the middle of the figure under the guidance of rewards. However, at this time, the path generated by Diffuser causes the agent to collide with the wall, while \alias generates a smoother collision-free path, reflecting the superiority of our method.

\begin{table}[t]
\caption{
    \textbf{Adaptation Performance on Pick-and-Place Task}
}
\vspace{0.2cm}
\label{table:kukapick}
\tabcolsep 8pt
\centering
\small
\begin{tabular}{ccc}
\toprule
\multicolumn{1}{c}{\textbf{Environment}} & \textbf{Diffuser} & \textbf{AdaptDiffuser} \\
\midrule
Pick and Place setup 1 & 28.16 \scriptsize{\raisebox{1pt}{$\pm 2.0$}} & \textbf{36.03} \scriptsize{\raisebox{1pt}{$\pm 2.1$}} \\
Pick and Place setup 2 & 35.25 \scriptsize{\raisebox{1pt}{$\pm 1.4$}} & \textbf{39.00} \scriptsize{\raisebox{1pt}{$\pm 1.3$}} \\
\midrule
\multicolumn{1}{c}{\textbf{Average}} & 31.71 \hspace{.58cm} & \textbf{37.52} \hspace{.58cm} \\
\bottomrule
\end{tabular}
\vspace{-5pt}
\end{table}

\begin{figure}[t]
    \begin{center}
    \centerline{\includegraphics[width=0.9\columnwidth]{figs/pick_vis.pdf}}
    \vspace{-0.05in}
    \caption{\textbf{Visualization of KUKA Pick-and-Place Task.} We require the KUKA Arm to move the blocks from their random initialized positions on the right side of the table to the left and arrange them in the order of yellow, blue, green, and red (from near to far).}
    \label{fig:pick_vis}
    \end{center}
\vspace{-15pt}
\end{figure}

\subsubsection{KUKA Pick and Place Task}
\label{sec:pick_place}
\textbf{Task Specification.}
There are two tasks in the KUKA robot arm environment. One is the conditional stacking task, as defined in \cite{janner2022planning}, where the robot must correctly stack blocks in a predetermined order on a designated location, using blocks that have been randomly placed. 
And the other is the pick-and-place task designed by us, which aims to place the randomly initialized blocks in their own target locations in a predetermined order. The reward functions of both tasks are defined as one upon successful placements and zero otherwise.


To test the adaptation capability of \alias and other baselines, we only provide expert trajectory data for the conditional stacking task, which is generated by PDDLStream \cite{garrett2020pddlstream}, but we require the planner to generalize to pick-and-put task without any expert data. The performance of the pick-and-place task is supposed to be a good measure of the planner's adaptability.





\textbf{Adaptation Performance.}
In KUKA pick-and-place task, we define the guidance of the conditional diffusion model as the gradient of the reward function about the distance between the current location and the target location. Then, the adaptation performance is displayed in Table \ref{table:kukapick}.

There are two setups in KUKA benchmark. In setup 1, the four blocks are initialized randomly on the floor, while in setup 2, the four blocks are stacked at a random location at the beginning.
As shown in Table \ref{table:kukapick}, \alias outperforms Diffuser greatly on both setups while achieving higher performance at setup 2 because all of the blocks start from the same horizontal position. We visualize a successful case of the KUKA pick-and-place task in Figure \ref{fig:pick_vis}, and more visualization results can be seen in Appendix \ref{appedix:vis_kuka}.

\subsection{Ablation Study}

\subsubsection{Ablation on Iterative Phases}
In order to verify the lifting effect of iterative data generation of our method \alias to improve the performance of the planner, we conduct an ablation experiment on the number of iterative phases of \alias in the MuJoCo environment of D4RL.
\begin{table}[t]
\caption{
\textbf{Ablation on Iterative Phases.} The mean and the standard error are calculated over 3 random seeds.
}
\label{table:mujoco_iterphases}
\vspace{0.2cm}
\centering
\footnotesize
\begin{tabular}{cccc}
\toprule
\textbf{Dataset} & \textbf{Environment} & \textbf{$1^{\text{st}}$ Phase} & \textbf{$2^{\text{nd}}$ Phase} \\ 
\midrule
Medium-Expert & HalfCheetah & $\textbf{89.3}$ \scriptsize{\raisebox{1pt}{$\pm 0.6$}} & $\textbf{89.6}$ \scriptsize{\raisebox{1pt}{$\pm 0.8$}} \\ 
Medium-Expert & Hopper & $\textbf{110.7}$ \scriptsize{\raisebox{1pt}{$\pm 3.2$}} & $\textbf{111.6}$ \scriptsize{\raisebox{1pt}{$\pm 2.0$}}  \\ 
Medium-Expert & Walker2d & $\textbf{107.7}$ \scriptsize{\raisebox{1pt}{$\pm 0.9$}} & $\textbf{108.2}$ \scriptsize{\raisebox{1pt}{$\pm 0.8$}}\\ 
\midrule
Medium & HalfCheetah & $\textbf{43.8}$ \scriptsize{\raisebox{1pt}{$\pm 0.5$}} & $\textbf{44.2}$ \scriptsize{\raisebox{1pt}{$\pm 0.6$}} \\ 
Medium & Hopper & $95.4$ \scriptsize{\raisebox{1pt}{$\pm 3.4$}} & $\textbf{96.6}$ \scriptsize{\raisebox{1pt}{$\pm 2.7$}} \\ 
Medium & Walker2d & $83.2$ \scriptsize{\raisebox{1pt}{$\pm 3.5$}} & $\textbf{84.4}$ \scriptsize{\raisebox{1pt}{$\pm 2.6$}} \\ 
\midrule
\multicolumn{2}{c}{\textbf{Average}} & 88.4 \hspace{.58cm} & \textbf{89.1} \hspace{.58cm} \\
\bottomrule
\end{tabular}
\vspace{-10pt}
\end{table}

As shown in Table \ref{table:mujoco_iterphases}, with ``Medium" dataset, due to the low quality of the original dataset, although the data generated in the first phase has greatly supplemented the training dataset and greatly improved the performance (referring to Sec \ref{sec:mujoco}), the performance achieved after the second phase is still significantly improved compared with that of the first phase. However, for ``Medium-Expert" dataset, because the expert data of the dataset has covered most of the environment, and the newly generated data is only more suitable for the planner to learn. So, after a certain improvement in the first phase, the subsequent growth is not obvious. The above experiments verify the effectiveness of \alias for the multi-phase iterative paradigm, and also show that the boosting effect is no longer obvious after the algorithm performance reaches a certain level.


\subsubsection{Ablation on Insufficient Data \& Training}
\label{ablation:limit_data}
To demonstrate the superiority of our method over previous diffusion-based work Diffuser~\cite{janner2022planning} when the expert data is limited and the training is insufficient,
we conducted experiments on the Maze2d-Large dataset using different percentages of expert data (e.g. 20\%, 50\%) with only 25\% training steps to train our model.
The results are shown in Table~\ref{tab:ablation_amount}. 
The setting 100\%$\mathcal{D}$ denotes the full training setting.
We can see our \alias, which uses only 50\% data and 25\% training steps, beats the fully trained Diffuser.
\alias can achieve good performance with a small amount of expert data and training steps.

\begin{table}[t]
\caption{\textbf{Ablation study on different amounts of expert data.}}
\vspace{0.2cm}
\centering
\footnotesize
\tabcolsep 8pt
\begin{tabular}{ccccc}
\toprule
   \textbf{Amount of Data} & \textbf{20\% $\mathcal{D}$}   & \textbf{50\% $\mathcal{D}$} & \textbf{100\%$\mathcal{D}$} \\
\midrule
    Diffuser      & 105.0  & 107.9  & 123.0  \\
    AdaptDiffuser & \textbf{112.5}  & \textbf{123.8}  & \textbf{167.9}  \\      
\bottomrule
\end{tabular}
\label{tab:ablation_amount}
\vspace{-5pt}
\end{table}

\subsubsection{Model Size and Running Time}
We show the model size of \alias measured by the number of parameters in Table \ref{table:model_size} here. And we also analyze the testing time and training time performance in Appendix \ref{appendix:time}. From the analysis, we can see that the inference time of \alias is almost equal to that of Diffuser \citep{janner2022planning}.

\begin{table}[t]
\caption{\small \textbf{Model Size of \alias.}}
\label{table:model_size}
\vspace{0.2cm}
\centering
\small
\begin{tabular}{cc}
\toprule
\textbf{Environment} & \textbf{Total Parameters (Model Size)} \\
\midrule
MuJoCo & 3.96 M \\
Maze2D & 3.68 M \\
KUKA Robot & 64.9 M \\
\bottomrule
\end{tabular}
\vspace{-5pt}
\end{table}

\section{Conclusion}

We present \alias, a method for improving the performance of diffusion-based planners in offline reinforcement learning through self-evolution. By generating diverse, high-quality and heterogeneous expert data using a reward-guided diffusion model and filtering out infeasible data using a rule-based discriminator, \alias is able to enhance the performance of diffusion models in existing decision-making tasks, especially the goal-conditioned tasks, and further improve the adaptability in unseen tasks without any expert data. Our experiments on two widely-used offline RL benchmarks and our carefully designed unseen tasks in KUKA and Maze2D environments validate the effectiveness of \alias. 

\textbf{Discussion of Limitation.} Our method achieves better performance by generating high-quality synthetic data but increases the amount of computation required in training with almost no increase in inference time. Besides, although AdaptDiffuser has proven its effectiveness in several scenarios (e.g. MuJoCo, Maze2d, KUKA), it still faces challenges in high-dimensional observation space tasks. More detailed discussions are given in Appendix \ref{appendix:discuss}.

\textbf{Future works.} 
Further improving the sampling speed and exploring tasks with high-dimensional input are potential areas for future works. 
And with the help of ChatGPT \cite{ouyang2022training}, we can use prompts to directly generate diverse maze settings to assist synthetic data generation which is also a promising direction. We provide some examples in Appendix \ref{appendix:chatgpt}.



\section*{Acknowledgements}

This paper is partially supported by the National Key R\&D Program of China No.2022ZD0161000 and the General Research Fund of Hong Kong No.17200622.



\begin{thebibliography}{49}
\providecommand{\natexlab}[1]{#1}
\providecommand{\url}[1]{\texttt{#1}}
\expandafter\ifx\csname urlstyle\endcsname\relax
  \providecommand{\doi}[1]{doi: #1}\else
  \providecommand{\doi}{doi: \begingroup \urlstyle{rm}\Url}\fi

\bibitem[Agarwal et~al.(2020)Agarwal, Schuurmans, and
  Norouzi]{agarwal2020optimistic}
Agarwal, R., Schuurmans, D., and Norouzi, M.
\newblock An optimistic perspective on offline reinforcement learning.
\newblock In \emph{International Conference on Machine Learning}, pp.\
  104--114. PMLR, 2020.

\bibitem[Ajay et~al.(2023)Ajay, Du, Gupta, Tenenbaum, Jaakkola, and
  Agrawal]{ajay2022conditional}
Ajay, A., Du, Y., Gupta, A., Tenenbaum, J., Jaakkola, T., and Agrawal, P.
\newblock Is conditional generative modeling all you need for decision-making?
\newblock In \emph{International Conference on Learning Representations}, 2023.

\bibitem[Argenson \& Dulac-Arnold(2021)Argenson and
  Dulac-Arnold]{argenson2020model}
Argenson, A. and Dulac-Arnold, G.
\newblock Model-based offline planning.
\newblock In \emph{International Conference on Learning Representations}, 2021.

\bibitem[Austin et~al.(2021)Austin, Johnson, Ho, Tarlow, and van~den
  Berg]{austin2021structured}
Austin, J., Johnson, D.~D., Ho, J., Tarlow, D., and van~den Berg, R.
\newblock Structured denoising diffusion models in discrete state-spaces.
\newblock In \emph{Advances in Neural Information Processing Systems}, 2021.

\bibitem[Buckman et~al.(2021)Buckman, Gelada, and
  Bellemare]{buckman2020importance}
Buckman, J., Gelada, C., and Bellemare, M.~G.
\newblock The importance of pessimism in fixed-dataset policy optimization.
\newblock In \emph{International Conference on Learning Representations}, 2021.

\bibitem[Chen et~al.(2021{\natexlab{a}})Chen, Lu, Rajeswaran, Lee, Grover,
  Laskin, Abbeel, Srinivas, and Mordatch]{chen2021decision}
Chen, L., Lu, K., Rajeswaran, A., Lee, K., Grover, A., Laskin, M., Abbeel, P.,
  Srinivas, A., and Mordatch, I.
\newblock Decision transformer: Reinforcement learning via sequence modeling.
\newblock \emph{Advances in neural information processing systems},
  34:\penalty0 15084--15097, 2021{\natexlab{a}}.

\bibitem[Chen et~al.(2021{\natexlab{b}})Chen, Zhang, Zen, Weiss, Norouzi, and
  Chan]{chen2020wavegrad}
Chen, N., Zhang, Y., Zen, H., Weiss, R.~J., Norouzi, M., and Chan, W.
\newblock Wavegrad: Estimating gradients for waveform generation.
\newblock In \emph{International Conference on Learning Representations},
  2021{\natexlab{b}}.

\bibitem[Dhariwal \& Nichol(2021)Dhariwal and Nichol]{dhariwal2021diffusion}
Dhariwal, P. and Nichol, A.~Q.
\newblock Diffusion models beat {GAN}s on image synthesis.
\newblock In \emph{Advances in Neural Information Processing Systems}, 2021.

\bibitem[Du \& Mordatch(2019)Du and Mordatch]{du2019implicit}
Du, Y. and Mordatch, I.
\newblock Implicit generation and generalization in energy-based models.
\newblock In \emph{Advances in Neural Information Processing Systems}, 2019.

\bibitem[Du et~al.(2020)Du, Li, and Mordatch]{du2020compositional}
Du, Y., Li, S., and Mordatch, I.
\newblock Compositional visual generation with energy based models.
\newblock In \emph{Advances in Neural Information Processing Systems}, 2020.

\bibitem[Fan et~al.(2022)Fan, Wang, Jiang, Mandlekar, Yang, Zhu, Tang, Huang,
  Zhu, and Anandkumar]{fan2022minedojo}
Fan, L., Wang, G., Jiang, Y., Mandlekar, A., Yang, Y., Zhu, H., Tang, A.,
  Huang, D.-A., Zhu, Y., and Anandkumar, A.
\newblock Minedojo: Building open-ended embodied agents with internet-scale
  knowledge.
\newblock In \emph{Thirty-sixth Conference on Neural Information Processing
  Systems Datasets and Benchmarks Track}, 2022.

\bibitem[Feller(2015)]{feller2015theory}
Feller, W.
\newblock On the theory of stochastic processes, with particular reference to
  applications.
\newblock In \emph{Selected Papers I}, pp.\  769--798. Springer, 2015.

\bibitem[Fu et~al.(2020)Fu, Kumar, Nachum, Tucker, and Levine]{fu2020d4rl}
Fu, J., Kumar, A., Nachum, O., Tucker, G., and Levine, S.
\newblock D4rl: Datasets for deep data-driven reinforcement learning.
\newblock \emph{arXiv preprint arXiv:2004.07219}, 2020.

\bibitem[Garrett et~al.(2020)Garrett, Lozano-P{\'e}rez, and
  Kaelbling]{garrett2020pddlstream}
Garrett, C.~R., Lozano-P{\'e}rez, T., and Kaelbling, L.~P.
\newblock Pddlstream: Integrating symbolic planners and blackbox samplers via
  optimistic adaptive planning.
\newblock In \emph{Proceedings of the International Conference on Automated
  Planning and Scheduling}, volume~30, pp.\  440--448, 2020.

\bibitem[Grathwohl et~al.(2020)Grathwohl, Wang, Jacobsen, Duvenaud, and
  Zemel]{grathwohl2020stein}
Grathwohl, W., Wang, K.-C., Jacobsen, J.-H., Duvenaud, D., and Zemel, R.
\newblock Learning the stein discrepancy for training and evaluating
  energy-based models without sampling.
\newblock In \emph{International Conference on Machine Learning}, 2020.

\bibitem[Ho et~al.(2020)Ho, Jain, and Abbeel]{ho2020denoising}
Ho, J., Jain, A., and Abbeel, P.
\newblock Denoising diffusion probabilistic models.
\newblock \emph{Advances in Neural Information Processing Systems},
  33:\penalty0 6840--6851, 2020.

\bibitem[Hyv{{\"a}}rinen(2005)]{hyvarinen2005score}
Hyv{{\"a}}rinen, A.
\newblock Estimation of non-normalized statistical models by score matching.
\newblock \emph{Journal of Machine Learning Research}, 2005.

\bibitem[Janner et~al.(2021)Janner, Li, and Levine]{janner2021offline}
Janner, M., Li, Q., and Levine, S.
\newblock Offline reinforcement learning as one big sequence modeling problem.
\newblock \emph{Advances in neural information processing systems},
  34:\penalty0 1273--1286, 2021.

\bibitem[Janner et~al.(2022)Janner, Du, Tenenbaum, and
  Levine]{janner2022planning}
Janner, M., Du, Y., Tenenbaum, J., and Levine, S.
\newblock Planning with diffusion for flexible behavior synthesis.
\newblock In \emph{International Conference on Machine Learning}, pp.\
  9902--9915. PMLR, 2022.

\bibitem[Kidambi et~al.(2020)Kidambi, Rajeswaran, Netrapalli, and
  Joachims]{kidambi2020morel}
Kidambi, R., Rajeswaran, A., Netrapalli, P., and Joachims, T.
\newblock Morel: Model-based offline reinforcement learning.
\newblock \emph{Advances in neural information processing systems},
  33:\penalty0 21810--21823, 2020.

\bibitem[Kingma \& Ba(2015)Kingma and Ba]{kingma2014adam}
Kingma, D.~P. and Ba, J.
\newblock Adam: A method for stochastic optimization.
\newblock In \emph{International Conference on Learning Representations}, 2015.

\bibitem[Kostrikov et~al.(2022)Kostrikov, Nair, and
  Levine]{kostrikov2021offline}
Kostrikov, I., Nair, A., and Levine, S.
\newblock Offline reinforcement learning with implicit q-learning.
\newblock In \emph{International Conference on Learning Representations}, 2022.

\bibitem[Kumar et~al.(2020)Kumar, Zhou, Tucker, and
  Levine]{kumar2020conservative}
Kumar, A., Zhou, A., Tucker, G., and Levine, S.
\newblock Conservative q-learning for offline reinforcement learning.
\newblock \emph{Advances in Neural Information Processing Systems},
  33:\penalty0 1179--1191, 2020.

\bibitem[LeCun et~al.(2006)LeCun, Chopra, Hadsell, Huang, and
  et~al.]{lecun06atutorial}
LeCun, Y., Chopra, S., Hadsell, R., Huang, F.~J., and et~al.
\newblock A tutorial on energy-based learning.
\newblock In \emph{Predicting Structured Data}. MIT Press, 2006.

\bibitem[Levine et~al.(2020)Levine, Kumar, Tucker, and Fu]{levine2020offline}
Levine, S., Kumar, A., Tucker, G., and Fu, J.
\newblock Offline reinforcement learning: Tutorial, review, and perspectives on
  open problems.
\newblock \emph{arXiv preprint arXiv:2005.01643}, 2020.

\bibitem[Misra(2020)]{misra2019mish}
Misra, D.
\newblock Mish: A self regularized non-monotonic activation function.
\newblock In \emph{31st British Machine Vision Conference 2020, {BMVC} 2020,
  Virtual Event, UK, September 7-10, 2020}, 2020.

\bibitem[Nijkamp et~al.(2019)Nijkamp, Hill, Zhu, and Wu]{nijkamp2019learning}
Nijkamp, E., Hill, M., Zhu, S.-C., and Wu, Y.~N.
\newblock Learning non-convergent non-persistent short-run {MCMC} toward
  energy-based model.
\newblock In \emph{Advances in Neural Information Processing Systems}, 2019.

\bibitem[Ouyang et~al.(2022)Ouyang, Wu, Jiang, Almeida, Wainwright, Mishkin,
  Zhang, Agarwal, Slama, Ray, et~al.]{ouyang2022training}
Ouyang, L., Wu, J., Jiang, X., Almeida, D., Wainwright, C., Mishkin, P., Zhang,
  C., Agarwal, S., Slama, K., Ray, A., et~al.
\newblock Training language models to follow instructions with human feedback.
\newblock \emph{Advances in Neural Information Processing Systems},
  35:\penalty0 27730--27744, 2022.

\bibitem[Prudencio et~al.(2022)Prudencio, Maximo, and
  Colombini]{prudencio2022survey}
Prudencio, R.~F., Maximo, M.~R., and Colombini, E.~L.
\newblock A survey on offline reinforcement learning: Taxonomy, review, and
  open problems.
\newblock \emph{arXiv preprint arXiv:2203.01387}, 2022.

\bibitem[Ramesh et~al.(2022)Ramesh, Dhariwal, Nichol, Chu, and
  Chen]{ramesh2022hierarchical}
Ramesh, A., Dhariwal, P., Nichol, A., Chu, C., and Chen, M.
\newblock Hierarchical text-conditional image generation with clip latents.
\newblock \emph{arXiv preprint arXiv:2204.06125}, 2022.

\bibitem[Rombach et~al.(2022)Rombach, Blattmann, Lorenz, Esser, and
  Ommer]{rombach2022high}
Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B.
\newblock High-resolution image synthesis with latent diffusion models.
\newblock In \emph{Proceedings of the IEEE/CVF Conference on Computer Vision
  and Pattern Recognition}, pp.\  10684--10695, 2022.

\bibitem[Ronneberger et~al.(2015)Ronneberger, Fischer, and
  Brox]{ronneberger2015u}
Ronneberger, O., Fischer, P., and Brox, T.
\newblock U-net: Convolutional networks for biomedical image segmentation.
\newblock In \emph{Medical Image Computing and Computer-Assisted
  Intervention--MICCAI 2015: 18th International Conference, Munich, Germany,
  October 5-9, 2015, Proceedings, Part III 18}, pp.\  234--241. Springer, 2015.

\bibitem[Schreiber et~al.(2010)Schreiber, Stemmer, and
  Bischoff]{schreiber2010fast}
Schreiber, G., Stemmer, A., and Bischoff, R.
\newblock The fast research interface for the kuka lightweight robot.
\newblock In \emph{IEEE workshop on innovative robot control architectures for
  demanding (Research) applications how to modify and enhance commercial
  controllers (ICRA 2010)}, pp.\  15--21. Citeseer, 2010.

\bibitem[Siegel et~al.(2020)Siegel, Springenberg, Berkenkamp, Abdolmaleki,
  Neunert, Lampe, Hafner, Heess, and Riedmiller]{siegel2020keep}
Siegel, N., Springenberg, J.~T., Berkenkamp, F., Abdolmaleki, A., Neunert, M.,
  Lampe, T., Hafner, R., Heess, N., and Riedmiller, M.
\newblock Keep doing what worked: Behavior modelling priors for offline
  reinforcement learning.
\newblock In \emph{International Conference on Learning Representations}, 2020.

\bibitem[Sohl-Dickstein et~al.(2015)Sohl-Dickstein, Weiss, Maheswaranathan, and
  Ganguli]{sohl2015deep}
Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N., and Ganguli, S.
\newblock Deep unsupervised learning using nonequilibrium thermodynamics.
\newblock In \emph{International Conference on Machine Learning}, pp.\
  2256--2265. PMLR, 2015.

\bibitem[Song et~al.(2021)Song, Meng, and Ermon]{song2020denoising}
Song, J., Meng, C., and Ermon, S.
\newblock Denoising diffusion implicit models.
\newblock In \emph{International Conference on Learning Representations}, 2021.

\bibitem[Song \& Ermon(2019)Song and Ermon]{song2019generative}
Song, Y. and Ermon, S.
\newblock Generative modeling by estimating gradients of the data distribution.
\newblock In \emph{Advances in Neural Information Processing Systems}, 2019.

\bibitem[Sutton(1988)]{sutton1988learning}
Sutton, R.~S.
\newblock Learning to predict by the methods of temporal differences.
\newblock \emph{Machine learning}, 3\penalty0 (1):\penalty0 9--44, 1988.

\bibitem[Sutton \& Barto(2018)Sutton and Barto]{sutton2018reinforcement}
Sutton, R.~S. and Barto, A.~G.
\newblock \emph{Reinforcement learning: An introduction}.
\newblock MIT press, 2018.

\bibitem[Todorov et~al.(2012)Todorov, Erez, and Tassa]{todorov2012mujoco}
Todorov, E., Erez, T., and Tassa, Y.
\newblock Mujoco: A physics engine for model-based control.
\newblock In \emph{2012 IEEE/RSJ international conference on intelligent robots
  and systems}, pp.\  5026--5033. IEEE, 2012.

\bibitem[Wang et~al.(2022)Wang, Zhao, Luo, Ren, Zhang, and
  Li]{wang2022bootstrapped}
Wang, K., Zhao, H., Luo, X., Ren, K., Zhang, W., and Li, D.
\newblock Bootstrapped transformer for offline reinforcement learning.
\newblock In \emph{Advances in Neural Information Processing Systems}, 2022.

\bibitem[Wang et~al.(2023)Wang, Hunt, and Zhou]{wang2022diffusion}
Wang, Z., Hunt, J.~J., and Zhou, M.
\newblock Diffusion policies as an expressive policy class for offline
  reinforcement learning.
\newblock In \emph{International Conference on Learning Representations}, 2023.

\bibitem[Williams et~al.(2015)Williams, Aldrich, and
  Theodorou]{williams2015model}
Williams, G., Aldrich, A., and Theodorou, E.
\newblock Model predictive path integral control using covariance variable
  importance sampling.
\newblock \emph{arXiv preprint arXiv:1509.01149}, 2015.

\bibitem[Wu \& He(2018)Wu and He]{wu2018group}
Wu, Y. and He, K.
\newblock Group normalization.
\newblock In \emph{Proceedings of the European conference on computer vision
  (ECCV)}, pp.\  3--19, 2018.

\bibitem[Wu et~al.(2019)Wu, Tucker, and Nachum]{wu2019behavior}
Wu, Y., Tucker, G., and Nachum, O.
\newblock Behavior regularized offline reinforcement learning.
\newblock \emph{arXiv preprint arXiv:1911.11361}, 2019.

\bibitem[Yang et~al.(2022)Yang, Wang, Zheng, Feng, and
  Zhou]{yang2022regularized}
Yang, S., Wang, Z., Zheng, H., Feng, Y., and Zhou, M.
\newblock A regularized implicit policy for offline reinforcement learning.
\newblock \emph{arXiv preprint arXiv:2202.09673}, 2022.

\bibitem[Yu et~al.(2020)Yu, Thomas, Yu, Ermon, Zou, Levine, Finn, and
  Ma]{yu2020mopo}
Yu, T., Thomas, G., Yu, L., Ermon, S., Zou, J.~Y., Levine, S., Finn, C., and
  Ma, T.
\newblock Mopo: Model-based offline policy optimization.
\newblock \emph{Advances in Neural Information Processing Systems},
  33:\penalty0 14129--14142, 2020.

\bibitem[Zhang(2021)]{zhang2021dirac}
Zhang, L.
\newblock Dirac delta function of matrix argument.
\newblock \emph{International Journal of Theoretical Physics}, 60\penalty0
  (7):\penalty0 2445--2472, 2021.

\bibitem[Zhou et~al.(2021)Zhou, Du, and Wu]{zhou2021shape}
Zhou, L., Du, Y., and Wu, J.
\newblock {3D} shape generation and completion through point-voxel diffusion.
\newblock In \emph{International Conference on Computer Vision}, 2021.

\end{thebibliography}
