Task-agnostic Pre-training and Task-guided Fine-tuning for Versatile Diffusion Planner

## Abstract

Abstract Diffusion models have demonstrated their capabilities in modeling trajectories of multi-tasks.
However, existing multi-task planners or policies typically rely on task-specific demonstrations via multi-task imitation, or require task-specific reward labels to facilitate policy optimization via Reinforcement Learning (RL). To address these challenges, we aim to develop a versatile diffusion planner that can leverage large-scale inferior data that contains task-agnostic sub-optimal trajectories, with the ability to fast adapt to specific tasks. In this paper, we propose SODP , a two-stage framework that leverages S ub- O ptimal data to learn a D iffusion P lanner, which is generalizable for various downstream tasks.
Specifically, in the pre-training stage, we train a foundation diffusion planner that extracts general planning capabilities by modeling the versatile distribution of multi-task trajectories, which can be sub-optimal and has wide data coverage. Then for downstream tasks, we adopt RL-based fine-tuning with task-specific rewards to fast refine the diffusion planner, which aims to generate action sequences with higher task-specific returns. Experimental results from multi-task domains including Meta-World and Adroit demonstrate that SODP outperforms state-of-the-art methods with only a small amount of data for reward-guided fine-tuning.

## 1 Introduction

There has been a long-standing pursuit to develop agents capable of performing multiple tasks (Reed et al., ; Lee et al., ). While traditional RL methods have made significant strides in training agents to master individual tasks (Silver et al., ; OpenAI et al., ), extending this capability to handle diverse tasks remains a significant challenge due to the diversity of task variants and optimization directions with different rewards.
Multi-task RL aims to address this by developing agents via task-conditioned optimization (Yu et al., ; Lee et al., ) or
parameter-compositional learning (Sun et al., ; Lee et al., ).
However, the inherent diversity in task trajectory distributions makes it challenging to model and accommodate the modeling across different task structures. Diffusion models (Sohl-Dickstein et al., ; Ho et al., ), originally designed for generative tasks, provide a powerful framework for addressing these difficulties. Their capacities to capture complex, multi-modal distributions within high-dimensional data spaces (Podell et al., ; Ho et al., ; Jing et al., ) makes them well-suited for representing the broad variability encountered in multi-task environments.

Motivated by this, existing methods have employed diffusion models to imitate expert behaviors derived from human demonstrations across various tasks (Pearce et al., ; Xu et al., ; Chi et al., ). However, acquiring task-specific demonstrations is often time-consuming and costly, especially in environments requiring specialized domain expertise. Alternative approaches attempt to optimize diffusion models with return guidance (He et al., ; Liang et al., ) or conventional RL paradigm (Wang et al., ), which demands a large volume of data with reward labels for each task.
To address the above limitations, we wonder whether a generalized diffusion planner can be learned from a large amount of low-quality trajectories without reward labels, with the ability to adapt quickly to various downstream tasks. The inferior dataset comprises a mixture of
sub-optimal state-action pairs from various tasks,
which can be easily obtained in the real world. In training, the diffusion planner seeks to model the distribution of diverse trajectories with broad coverage, enabling it to acquire generalizable capabilities and allowing the planner to further concentrate on high-reward regions of specific downstream tasks via fast adaptation. An overview of our method is given in Figure 1.

Figure: Figure 1: The Overall framework. Different colors represent different tasks. The diffusion model is first pre-trained on a mixed dataset drawn from multiple tasks, and is then fine-tuned for each specific task using task-specific rewards.
Refer to caption: x1.png

In this paper, we propose a novel framework to utilize Sub-Optimal data to train a Diffusion Planner (SODP) that can generalize across a wide range of downstream tasks. SODP consists of two stages: pre-training and fine-tuning.
By leveraging a set of trajectories of different tasks for pre-training, we employ action-sequence prediction to capture shared knowledge across tasks. Since the state space may vary between tasks, focusing on the common action space (e.g., end-effector poses of a robot arm) facilitates task generalization.
We frame the pre-training stage as a conditional generative problem that generates future actions based on historical states.
Then, inspired by the remarkable success of RL-based alignment for LLMs (Ouyang et al., ; Glaese et al., ), we adopt an RL-based fine-tuning approach to tailor the pre-trained diffusion planner to specific downstream tasks.
Specifically, we conduct online interaction based on the pre-trained planner to collect task-specific experiences with reward labels, and perform policy gradients to iteratively refine the predicted action-sequence distribution based on reward feedback of tasks. Through fine-tuning, the diffusion planner can gradually adapt toward generating actions with high task-specific rewards and eventually become optimal for the given task.

Figure: Figure 2: Illustration of SODP in Meta-World *button-press-wall* task. We present trajectories generated by the diffusion model after pre-training and fine-tuning of SODP. The pre-trained model captures a wide range of behaviors, and the fine-tuned model discards the inferior behaviors to coverage to high-reward regions.
Refer to caption: x2.png

Figure 2 illustrates our method. In pre-training, the model captures diverse behavior patterns from the training data, encompassing inferior and superior actions. After fine-tuning, the model shrinks the action distribution and concentrates on generating optimal action sequences for a specific task. Our contributions can be summarized as follows. (i) We propose a novel pre-training and fine-tuning paradigm for learning a versatile diffusion planner, which leverages sub-optimal transitions to capture the broad action distributions across tasks, and adopt task-specific fine-tuning to transfer the planner to downstream tasks. (ii) We give an efficient fine-tuning algorithm based on policy gradient for diffusion planners, which progressively shifts the action distribution to concentrate on regions associated with higher task returns. (iii) We conduct extensive experiments using sub-optimal data from state-based Meta-World (Yu et al., ) as well as image-based Adroit (Rajeswaran et al., ), showcasing its superior performance compared to state-of-the-art approaches.

## 2 Preliminaries

#### Multi-task RL

We consider the multi-task RL problem involving $N$ tasks, where each task $\mathcal{T}\sim p(\mathcal{T})$ is represented by a task-specified Markov Decision Process (MDP).
Each MDP is defined by a tuple $(\mathcal{S}^{\mathcal{T}},\mathcal{A},P^{\mathcal{T}},R^{\mathcal{T}},P_{0}^{
\mathcal{T}},\mathcal{\gamma})$, where $\mathcal{S}^{\mathcal{T}}$ is the state space of task $\mathcal{T}$, $\mathcal{A}$ is the global action space, $P^{\mathcal{T}}(s_{t+1}^{\mathcal{T}}|s_{t}^{\mathcal{T}},a_{t}^{\mathcal{T}})
:\mathcal{S}^{\mathcal{T}}\times\mathcal{A}\to\mathcal{S}^{\mathcal{T}}$ is the transition function, $R^{\mathcal{T}}(s_{t}^{\mathcal{T}},a_{t}^{\mathcal{T}}):\mathcal{S}^{\mathcal
{T}}\times\mathcal{A}\to\mathbb{R}$ is the reward function, $\mathcal{\gamma}\in(0,1]$ is the discount factor, and $P_{0}^{\mathcal{T}}$ is the initial state distribution. We assume that all tasks share a common action space, executed by the same agent, while differing in their respective reward functions, state spaces, and transition dynamics.
At each timestep $t$, the agent perceives a state $s_{t}^{\mathcal{T}}\in\mathcal{S}^{\mathcal{T}}$, takes an action $a_{t}^{\mathcal{T}}\in\mathcal{A}$ according to the policy $\pi^{\mathcal{T}}(a_{t}^{\mathcal{T}}|s_{t}^{\mathcal{T}})$, and receives a reward $r_{t}^{\mathcal{T}}$. The agent’s objective is to determine an optimal policy that maximizes the expected return across all tasks: $\pi^{*}=\arg\max_{\pi}\mathbb{E}_{{\mathcal{T}}\sim p(\mathcal{T})}\mathbb{E}_
{a_{t}\sim\pi^{\mathcal{T}}}\big{[}\sum\nolimits_{t=0}^{\infty}\gamma^{t}r_{t}
^{\mathcal{T}}\big{]}$.

#### Diffusion Models

Diffusion models (Sohl-Dickstein et al., ) are a type of generative model that first add noise to the data $\bm{x}_{0}$ from a unknown distribution $q(\bm{x}_{0})$ in $K$ steps through a forward process defined as:

$$ $q(\bm{x}_{k}|\bm{x}_{k-1}):=\mathcal{N}(\bm{x}_{k};\sqrt{1-\beta_{k}}\bm{x}_{k -1},\beta_{k}\bm{I}),$ (1) $$

where $\beta_{k}$ is a predefined variance schedule. Then, a trainable reverse process is constructed as:

$$ $p_{\theta}(\bm{x}_{k-1}|\bm{x}_{k}):=\mathcal{N}(\bm{x}_{k-1};\mu_{\theta}(\bm {x}_{k},k),\Sigma_{k}),$ (2) $$

where $\mu_{\theta}(\bm{x}_{k},k)$ is the forward process posterior mean as a function of a noise prediction neural network $\epsilon_{\theta}(\bm{x}_{k},k)$ with a learnable parameter $\theta$ (Ho et al., ). $\epsilon_{\theta}(\bm{x}_{k},k)$ can be trained via a surrogate loss as

$$ $\mathcal{L}_{\text{denoise}}(\theta):=\mathbb{E}_{k\sim[1,K],x_{0}\sim q, \epsilon\sim\mathcal{N}(0,\bm{I})}\left[\big{\|}\epsilon-\epsilon_{\theta}(\bm {x}_{k},k)\big{\|}^{2}\right].$ (3) $$

After training, samples can be generated by first drawing Gaussian noise $\bm{x}_{K}$ and then iteratively denoising $\bm{x}_{K}$ into a noise-free output $x_{0}$ over $K$ iterations using the trained model $\epsilon_{\theta}(\bm{x}_{k},k)$ by

$$ $\bm{x}_{k-1}=\frac{1}{\sqrt{\alpha_{k}}}\left(\bm{x}_{k}-\frac{1-\alpha_{k}}{ \sqrt{1-\bar{\alpha}_{k}}}\epsilon_{\theta}(\bm{x}_{k},k)\right)+\sigma_{k} \mathcal{N}(0,\bm{I}),$ (4) $$

where $\alpha_{k}:=1-\beta_{k}$, $\bar{\alpha}_{k}:=\prod_{s=1}^{k}\alpha_{s}$ and $\sigma_{k}=\sqrt{\beta_{k}}$.

## 3 Method

We propose SODP, a two-stage framework that leverages large amounts of sub-optimal data to train a diffusion planner that can generalize to downstream tasks. The process is depicted in Figure 3.
In the pre-training stage, we train a guidance-free diffusion model to predict future actions based on historical states, using an mixture offline dataset cross tasks without reward labels.
In the fine-tuning stage, we refine the pre-trained model using policy gradient to maximize the task-specific rewards, additionally incorporating a regularization term to prevent the model from losing acquired skills.

Figure: Figure 3: Overview of SODP. We initially pre-train a diffusion model using multi-task transition data to predict action sequences from historical states. Subsequently, we fine-tune the model on downstream tasks using policy gradient methods, incorporating a regularization term to mitigate model degradation.
Refer to caption: x3.png

### 3.1 Pre-training with Large-scale Sub-optimal Data

Previous works (He et al., ) typically model multi-task RL as a conditional generative problem using diffusion models trained on datasets composed of multiple task subsets $\mathcal{D}=\cup_{i=1}^{N}\mathcal{D}_{i}$, as:

$$ $\max\nolimits_{\theta}\mathbb{E}_{\tau\sim\cup_{i}\mathcal{D}_{i}}\big{[}\log p _{\theta}(\bm{x}_{0}(\tau)\>\big{|}\>\bm{y}(\tau)\big{]},$ (5) $$

which requires additional condition $\bm{y}(\tau)$ to guide diffusion model to generate desirable trajectories. For instance, $\bm{y}(\tau)$ should contain the return of trajectory $R(\tau)$ and task description $Z$ as prompt.
However, the reward label and trajectory description may be scarce or costly to obtain in the real-world.
To overcome this challenge, we train a diffusion planner that can learn from offline trajectories transitions (i.e., $\{(s_{t},a_{t},s_{t+1})\}$) without reward label or task descriptions.
Specifically, we model the problem as a guidance-free generation process  (Chi et al., ):

$$ $\max\nolimits_{\theta}\mathbb{E}_{(\bm{s}_{t},\bm{a}_{t})\sim\cup_{i}\mathcal{ D}_{i}}\big{[}\log p_{\theta}(\bm{a}_{t}^{0}\>\big{|}\>\bm{s}_{t})\big{]}.$ (6) $$

Here, we represent $\bm{x}_{0}:=\bm{a}_{t}^{0}=(a_{t},a_{t+1},...,a_{t+H-1})$ as an action sequence, where $H$ is the planning horizon and $t$ is the timestep sampled from dataset $\mathcal{D}$. As previous work (Chi et al., ), we denote $\bm{s}_{t}$ as the historical states at timestep $t$ with length $T_{o}$, i.e., $\bm{s}_{t}:=\{s_{t-T_{o}+1},\dots s_{t-1},s_{t}\}$. The formulation in Eq. (6) enables the model to learn the broad action-sequence distribution of multi-tasks depending on previous observations, without requiring additional guidance.
To train our planning model, we modify Eq. (3) to obtain our pre-training objective as follows:

$$ $\mathcal{L}_{\text{pre-train}}(\theta)=\mathbb{E}_{k\sim[1,K],(\bm{s}_{t},\bm{ a}^{0}_{t})\sim D,\epsilon\sim\mathcal{N}(0,\mathbf{I})}\left[\left\|\epsilon- \epsilon_{\theta}(\bm{a}_{t}^{k},\bm{s}_{t},k)\right\|^{2}\right].$ (7) $$

Following Eq. (4), we can generate action sequences through a series of denoising steps:

$$ $\bm{a}^{k-1}_{t}=\frac{1}{\sqrt{\alpha_{k}}}\left(\bm{a}^{k}_{t}-\frac{1- \alpha_{k}}{\sqrt{1-\bar{\alpha}_{k}}}\epsilon_{\theta}(\bm{a}_{t}^{k},\bm{s}_ {t},k)\right)+\sigma_{k}\mathcal{N}(0,\bm{I}).$ (8) $$

Unlike other models, the dataset $\mathcal{D}$ we used for the pre-training stage is not restricted to expert trajectories. As shown in Figure 2, we aim to train a foundation model that captures diverse behaviors and learns general capabilities from inferior trajectories, enabling the planner to enhance its representation and action priors through pre-training before learning on downstream tasks.

### 3.2 Reward Fine-tuning for Downstream Tasks

#### MDP notation.

The fine-tuning stage involves two distinct MDPs: one for RL decision process and the other for the diffusion model denoising process. We use the superscript diff (e.g., $s_{k}^{\text{diff}}$, $a_{k}^{\text{diff}}$) to denote the MDP associated with diffusion model denoising process, while no superscript is used for the MDP related to the RL process (e.g., $s_{t}$, $a_{t}$). Additionally, we use $k\in\{K,\dots,0\}$ to represent the diffusion timestep and $t\in\{1,\dots,T\}$ to represent the trajectory timestep.

We model the denoising process of our pre-trained diffusion planner as a $K$-step MDP as follows:

$$ $\displaystyle s_{k}^{\text{diff}}=(s_{t},\bm{a}_{t}^{K-k}),\quad\;\;a_{k}^{ \text{diff}}=\bm{a}_{t}^{K-k-1},\quad\;\;P_{0}^{\text{diff}}(s_{0}^{\text{diff }})=(\delta_{s_{t}},\mathcal{N}(0,I)),$ $\displaystyle P^{\text{diff}}(s_{k+1}^{\text{diff}}\mid s_{k}^{\text{diff}},a_ {k}^{\text{diff}})=(\delta_{s_{t}},\delta_{a_{k}^{\text{diff}}}),\quad\;\;R^{ \text{diff}}(s_{k}^{\text{diff}},a_{k}^{\text{diff}})=\begin{cases}r(s_{k+1}^{ \text{diff}})=r(\bm{a}_{t}^{0})&\text{if }\;k=K-1,\\ 0&\text{otherwise}.\end{cases},$ $\displaystyle\pi_{\theta}^{\text{diff}}(a_{k}^{\text{diff}}\mid s_{k}^{\text{ diff}})=p_{\theta}(\bm{a}_{t}^{K-k-1}\mid\bm{a}_{t}^{K-k},\bm{s}_{t}),$ (9) $$

where $s_{k}^{\text{diff}}$ and $a_{k}^{\text{diff}}$ are the state and action at timestep $k$, $P_{0}^{\text{diff}}$ and $P^{\text{diff}}$ are the initial distribution and transition dynamics, $\delta$ is the Dirac delta distribution, $R^{\text{diff}}$ is the reward function and $p_{\theta}(\bm{a}_{t}^{K-k-1}\mid\bm{a}_{t}^{K-k},\bm{s}_{t})$ is the pre-trained diffusion planner. This formulation allows the state transitions in the MDP to be mapped to the denoising process in the diffusion model. The MDP initiates by sampling an initial state $s_{0}^{\text{diff}}\sim P_{0}^{\text{diff}}$, which corresponds to sample Gaussian noise $\bm{a}^{K}_{t}$ at the beginning of the reverse process. At each timestep $k$, the policy $\pi_{\theta}^{\text{diff}}(a_{k}^{\text{diff}}\mid s_{k}^{\text{diff}})$ takes an action $a_{k}^{\text{diff}}$ based on current state $s_{k}^{\text{diff}}$, which corresponds to generate next latent $\bm{a}_{t}^{K-k-1}$ based on current latent $\bm{a}_{t}^{K-k}$ following Eq. (8). The reward remains zero until a noise-free output $\bm{a_{0}^{t}}$ is evaluated.
Different from previous text-to-image studies that typically evaluate the final sample using a pre-trained reward model (Black et al., ; Fan et al., ), we aim to fine-tune the pre-trained diffusion planner to maximize rewards of downstream tasks,
which makes constructing reward models for all tasks costly. Therefore, we directly evaluate the generated action sequences in an online RL environment for each specific task $\mathcal{T}$.
Specifically, for any given timestep $t$, we use the planner to generate future actions $\bm{a}_{t}^{0}=(a_{t},a_{t+1},\dots,a_{t+H-1})$ and then execute the first $T_{a}$ steps. Then we calculate the discounted cumulative reward from the environment to assess the generated sample, expressed as $r(\bm{a}_{t}^{0})=\sum\nolimits_{t}^{T_{a}}\gamma^{t-1}r^{\mathcal{T}}(s_{t},a
_{t})$. We write $r(\bm{a}_{t}^{0})$ as shorthand for $r(s_{t},\bm{a}_{t}^{0})$ for brevity.

#### Finetuning objective.

The objective of fine-tuning our pre-trained diffusion planner is to maximize the expected reward of the generated action sequences for the target downstream task $\mathcal{T}$, which can be defined as:

$$ $\displaystyle J^{\mathcal{T}}(\theta)=\sum\nolimits_{t}\mathbb{E}_{p_{\theta}( \bm{a}^{0}_{t}|\bm{s}_{t})}[r^{\mathcal{T}}(\bm{a}^{0}_{t})].$ (10) $$

Directly optimizing the objective $J^{\mathcal{T}}(\theta)$ is intractable since it is infeasible to evaluate the return over all possible actions.
Therefore, we utilize policy gradient methods (Sutton et al., ), which estimate the policy gradient and apply a stochastic gradient ascent algorithm for updates. The gradient of the objective $J^{\mathcal{T}}(\theta)$ can be obtained as follows:

$$ $\nabla_{\theta}J^{\mathcal{T}}(\theta)=\sum\nolimits_{t}\mathbb{E}_{p_{\theta} (\bm{a}_{t}^{0:K}|\bm{s}_{t})}\left[r^{\mathcal{T}}(\bm{a}^{0}_{t})\sum_{k=1}^ {K}\nabla_{\theta}\log p_{\theta}(\bm{a}_{t}^{k-1}|\bm{a}_{t}^{k},\bm{s}_{t}) \right].$ (11) $$

However, optimizing with Eq. (11) can be computationally intensive, as it requires generating new samples after each optimization step. To enhance sample efficiency and leverage historical sequences, we employ importance sampling, following the approach of proximal policy optimization (PPO) (Schulman et al., ), and derive the loss function for reward improvement as follows:

$$ $\mathcal{L}_{\text{Imp}}^{\mathcal{T}}(\theta)=\sum\nolimits_{t}\mathbb{E}_{p_ {\theta_{\text{old}}}(\bm{a}_{t}^{0:K}|\bm{s}_{t})}\left[\sum_{k=1}^{K}-r^{ \mathcal{T}}(\bm{a}^{0}_{t})\text{max}\Big{(}\rho_{k}(\theta,\theta_{\text{old }}),\text{clip}\left(\rho_{k}(\theta,\theta_{\text{old}}),1+\epsilon,1- \epsilon\right)\Big{)}\right],$ (12) $$

where $\rho_{k}(\theta,\theta_{\text{old}})=\frac{p_{\theta}(\bm{a}_{t}^{k-1}|\bm{a}_
{t}^{k},\bm{s}_{t})}{p_{\theta_{\text{old}}}(\bm{a}_{t}^{k-1}|\bm{a}_{t}^{k},
\bm{s}_{t})}$ and $\epsilon$ is a hyperparameter. Then, we can train our model using $\mathcal{L}_{\text{Imp}}^{\mathcal{T}}(\theta)$ in an end-to-end manner, which is equivalent to maximizing the objective in Eq. (10).

#### Regularization term.

However, fine-tuning the model solely depending on the reward is insufficient since the model may step too far, which can lead to performance collapse and instability during reward maximization.
To address this problem, we introduce a Behavior-Clone (BC) regularization term during the fine-tuning process. Concretely, we aim to constrain our policy $\theta$ to closely match a target policy $\mu$, ensuring that $\theta$ does not deviate significantly from $\mu$ after policy updates. This constraint can be modeled using a negative log-likelihood (NLL) loss as:

$$ $\mathop{\min}_{\theta}\mathbb{E}_{\bm{a}_{\mu}^{0}\sim p_{\mu}}\left[-\log p_{ \theta}(\bm{a}_{\mu}^{0})\right].$ (13) $$

Following Ho et al. (), we can obtain a surrogate loss to optimize Eq. (13) as follows:

$$ $\mathcal{L}_{\text{BC}}(\theta)=\mathbb{E}_{k\sim[1,K],\bm{a}^{k}_{\mu}\sim p_ {\mu}}\left[\left\|\epsilon(\bm{a}_{\mu}^{k},k)-\epsilon_{\theta}(\bm{a}_{\mu} ^{k},k)\right\|^{2}\right],$ (14) $$

where $\epsilon(\bm{a}_{\mu}^{k},k)$ represents the ground-truth noise added to $\bm{a}_{\mu}^{k}$ at timestep $k$, which can be calculated as
$\epsilon(\bm{a}_{\mu}^{k},k)=\frac{\bm{a}_{\mu}^{k}-\sqrt{\bar{\alpha}_{k}}
\cdot\bm{a}_{\mu}^{0}}{\sqrt{1-\bar{\alpha}_{k}}}.$

How to select the target policy? Intuitively, an ideal target policy is the optimal policy that generates samples $x^{*}$ satisfying $\mathcal{C}(\bm{x}^{*})\geq\mathcal{C}(\bm{x})$ for all possible $\bm{x}$, where $\mathcal{C}(\bm{x})$ represents a measure of the performance or quality of the sample, such as the accumulated reward for action sequences. Since $\mu$ is unknown during fine-tuning, we approximate it by sampling action sequences $\bm{a}$ that satisfy $\mathcal{C}(\bm{a})\approx\mathcal{C}(\bm{a}^{*})$.
In practice, we denote $\bm{a}^{*}$ as the best actions from recent play experience, such as those that yielded the top $n$ highest rewards or successfully completed the given task. We then sample $\bm{a}^{k}$ from these proficient action sequences obtained during online interaction, nearly equivalent to sampling from $\mu$ to regularize the fine-tuning process. We also remark that the BC regularizer is not the only way to incorporate regularization into Eq. (12). For example, a Kullback–Leibler (KL) divergence between the fine-tuned and pre-trained models, or a diffusion pre-train loss can be employed to regularize the fine-tuning process, as shown in text-to-image and text-to-speech generation (Fan et al., ; Chen et al., ).
However, we find these regularization may cause the pre-trained planner trap in sub-optimal regions, hindering performance improvement. We will further discuss them in experiments.

Combining Eq. (12) with Eq. (14), the loss function for reward fine-tuning in downstream tasks $\mathcal{T}\sim p(\mathcal{T})$ is expressed as follows:

$$ $\mathcal{L}^{\mathcal{T}}_{\text{fine-tuning}}(\theta)=\mathcal{L}_{\text{Imp} }^{\mathcal{T}}(\theta)+\lambda\mathcal{L}_{\text{BC}}(\theta),$ (15) $$

where $\lambda$ is a weight coefficient. The overall process of pre-training and fine-tuning using SODP is summarized in Alg. 1 in the appendix. Since our goal is to generate complete trajectories rather than individual segments, we utilize a trajectory-level buffer (Zheng et al., ) for estimating the target policy $\mu$. Further, to ensure the accuracy of the approximation, we generate several proficient trajectories using the pre-trained model at the beginning of each iteration.

## 4 Related Work

Diffusion Models in RL.
Diffusion models are a leading class of generative models, achieving state-of-the-art performance across a variety of tasks, such as image generation (Ramesh et al., ), audio synthesis (Kong et al., ; Huang et al., ), and drug design (Schneuing et al., ; Guan et al., ). Recent studies have applied them in imitation learning to model human demonstrations and predict future actions (Li et al., ; Reuss et al., ). Other approaches have trained conditional diffusion models either as planners (Ajay et al., ; Brehmer et al., ) or policies (Hansen-Estruch et al., ; Kang et al., ). However, most of these efforts focus on single-task settings. While some recent works aim to extend diffusion models to multi-task scenarios, they often rely on additional conditions, such as prompts (He et al., ) or preference labels (Yu et al., ). These methods are limited by their dependence on expert data or explicit task knowledge. In contrast, our method learns broad action-sequence distributions from inferior data to enhance action priors, enabling effective generalization across a range of downstream tasks.

Fine-tuning Diffusion Models.
Despite the impressive success of diffusion models, they often face challenges in aligning with specific downstream objectives, such as image aesthetics (Schuhmann et al., ), fairness (Shen et al., ), or human preference (Xu et al., ), primarily due to their training on unsupervised data. Some methods have been proposed to address this issue by directly fine-tuning models using downstream objectives (Prabhudesai et al., ; Clark et al., ), but they rely on differentiable reward models, which are impractical in RL since accurately modeling rewards with neural networks is quite costly (Kim et al., ). Other methods reformulate the denoising process as an MDP and apply policy gradients for fine-tuning (Black et al., ; Fan et al., ). However, they heavily depend on strong pre-trained models and have proven ineffective in our case. Our goal is to fine-tune a less powerful model that has been trained on inferior data.

Concurrent with our work, DPPO (Ren et al., ) also explores reward fine-tuning for refining RL diffusion planners. However, their approach focuses exclusively on single-task settings and allows access to expert demonstrations. In contrast, we train our model on multi-task data without the need for superior demonstrations. Additionally, we analyze the limitations of current regularization methods for versatile RL diffusion models and propose a new regularizer that improves the performance of sub-optimal pre-trained models.

## 5 Experiments

In this section, we conduct experiments to evaluate our proposed method and address the following questions:
(1) How does SODP’s performance compare to current methods?
(2) Can SODP scale to high-dimensional observation inputs?
(3) How does SODP achieve higher rewards during online fine-tuning?

### 5.1 Experimental Setup

We evaluate SODP in both state-based and image-based environments. For state-based environments, we conduct experiments using the Meta-World benchmark (Yu et al., ). For image-based environments, we perform experiments using the Adroit benchmark (Rajeswaran et al., ).

Meta-World. The Meta-World benchmark comprises 50 distinct manipulation tasks, each requiring a Sawyer robot to interact with various objects. These tasks are designed to assess the robot’s ability to handle different scenarios, such as grasping, pushing, pulling, and manipulating objects of varying shapes, sizes, and complexities. While the state space and reward functions differ across tasks, the action space remains consistent. Following recent studies (He et al., ; Hu et al., ), we extend all tasks to a random-goal setting, referred to as MT50-rand.

Adroit. The Adroit benchmark includes three dexterous manipulation tasks, requiring a 24-degree-of-freedom dexterous hand to solve complex challenges such as in-hand manipulation and tool use. The goals in this environment are also randomized. For Adroit, we use images as the observation to assess whether our method can scale to high-dimensional input.

Datasets. Following previous work (He et al., ), for Meta-World, we use a sub-optimal dataset comprising the first 50% of experiences (50M transitions) obtained from the replay buffer of a SAC (Haarnoja et al., ) agent during training. To verify the applicability of our method to tasks of varying difficulty levels, we divide the entire dataset into four subsets based on the task categories presented in Seo et al. (). All baselines and our pre-training stage are trained on sub-optimal datasets. We then use 1M online transitions per task for fine-tuning.
For Adroit, we train a VRL3 (Wang et al., ) agent for each task and use the initial 30% experiences (90K transitions) from the converged replay buffer. All baselines are trained on expert demonstrations and our pre-training stage is trained on sub-optimal transitions.

Baselines.
For Meta-World, we compare our proposed SODP with the following baselines:
(1) MTBC.Extended BC to multi-task offline policy learning through network scaling and a task-ID-conditioned actor.
(2) MTIQL.Extended IQL (Kostrikov et al., ) with multi-head critic networks and a task-ID-conditioned actor for multi-task policy learning.
(3) MTDT.Extended Decision Transformer (DT) (Chen et al., ) to multitask settings by incorporating task ID encoding and state inputs for task-specific learning.
(4) Prompt-DT(Xu et al., ). An extension of DT, which is designed to leverage multi-task data to generalize policies for previously unseen tasks. It generates actions by utilizing trajectory prompts and reward-to-go signals.
(5) MTDIFF(He et al., ). A diffusion-based approach that integrates Transformer architectures with prompt learning to facilitate generative planning in multitask offline environments.
(6) HarmoDT(Hu et al., ). A DT-based approach that leverages parameter sharing to exploit task similarities while mitigating the adverse effects of conflicting gradients simultaneously.
The results for these baselines are directly replicated from those reported in HarmoDT (Hu et al., ).

Learning policies for Adroit is more challenging due to the significantly larger state and action spaces. The baselines used in Meta-World struggled to handle this high-dimensional data structure. Therefore, we compare SODP with following baselines designed for complex environments:
(1) BCRNN(Mandlekar et al., ). A variant of BC that employs a Recurrent Neural Network (RNN) as the policy network, predicting the sequence of actions based on the sequence of states as input.
(2) IBC(Florence et al., ). Extended BC with energy-based models (EBM) to train implicit behavioral cloning policies.
(3) Diffusion Policy(Chi et al., ). A diffusion-based approach that predicts future action sequences based on historical states.
(4) DP3(Ze et al., ). A visual imitation learning algorithm that incorporates 3D visual representations into diffusion policies, using a point clouds encoder to process visual observations into visual features. The results for these baselines are directly replicated from those reported in DP3 (Ze et al., ).

### 5.2 Results

We use the average success rate across all tasks as the evaluation metric and report the mean and standard deviation of success rates across three seeds. All baselines are trained on sub-optimal data. As shown in Table 1, our method achieves over a 60% success rate when learning from inferior data, outperforming all baseline methods. Compared to the existing state-of-the-art approach,

**Table 1: Average success rate across 3 seeds on Meta-World 50 tasks with random goals (MT50-rand), using sub-optimal data. Each task is evaluated for 50 episodes.**
| Method | Meta-World 50 Tasks |
| --- | --- |
| MTBC | $34.53_{\pm 1.25}$ |
| MTIQL | $43.28_{\pm 0.90}$ |
| MTDT | $42.33_{\pm 1.89}$ |
| Prompt-DT | $48.40_{\pm 0.16}$ |
| MTDIFF-P | $48.67_{\pm 1.32}$ |
| MTDIFF-P-ONEHOT | $48.94_{\pm 0.95}$ |
| HarmoDT-R | $53.80_{\pm 1.07}$ |
| HarmoDT-M | $57.20_{\pm 0.73}$ |
| HarmoDT-F | $57.20_{\pm 0.68}$ |
| SODP (ours) | $\textbf{60.56}_{\pm 0.14}$ |

our method demonstrates a 5.9% improvement. Notably, when compared to MTDIFF, the current leading method based on diffusion models, our approach shows a 24.4% improvement. MTBC performs the worst, as imitation learning heavily depends on data quality, and directly cloning behaviors from sub-optimal data typically results in inferior performance. In contrast, our method models versatile action distributions from low-quality data and leverages them as priors to guide policy optimization in downstream tasks, leading to improved performance.

To further analyze the learning dynamics, we sample 10 tasks and present their learning curves of SODP alongside two leading baselines, MTDIFF and HarmoDT, across five seeds. As shown in Figure 4, SODP rapidly converges to high success rates, surpassing the other two baselines. The pre-training stage equips the planner with comprehensive action distribution priors and allows it to rapidly transfer and enhance these capabilities across a variety of downstream tasks. As a result, the pre-training stage significantly accelerates convergence, leading to more efficient learning in the fine-tuning stage.
The two baseline approaches struggle to address complex and challenging tasks such as basketball and hammer. In contrast, our method effectively guides the model to generate proficient actions, demonstrating the benefits of fine-tuning with policy gradient concerning return maximization. Moreover, while HarmoDT exhibits instability across different random seeds, our method demonstrates robustness against randomness.

Figure: Figure 4: Learning efficiency. We sample 10 tasks and present the learning curves of SODP, MTDIFF, and HarmoDT across five seeds. X-axis represents gradient steps. We pre-train the planner for $5e^{5}$ steps, followed by fine-tuning with a smaller number of steps. SODP rapidly converges to high success rates, whereas MTDIFF and HarmoDT struggle with some challenging tasks.
Refer to caption: x4.png

Does SODP generalize to high-dimensional observations? We scale our method to image-based

**Table 2: Average success rate across 3 seeds on Adroit 3 tasks. IBC and BCRNN are extended by incorporating the DP3 point cloud encoder, resulting in IBD+3D and BCRNN+3D.**
|  | Adroit |  |  |  |
| --- | --- | --- | --- | --- |
| Algorithm $\backslash$ Task | Hammer | Door | Pen | Average |
| BCRNN | $0_{\pm 0}$ | $0_{\pm 0}$ | $9_{\pm 3}$ | 3.0 |
| BCRNN+3D | $8_{\pm 14}$ | $0_{\pm 0}$ | $8_{\pm 1}$ | 5,3 |
| IBC | $0_{\pm 0}$ | $0_{\pm 0}$ | $9_{\pm 2}$ | 3.0 |
| IBC+3D | $0_{\pm 0}$ | $0_{\pm 0}$ | $10_{\pm 1}$ | 3.3 |
| Diffusion Policy | $48_{\pm 17}$ | $50_{\pm 5}$ | $25_{\pm 4}$ | 31.7 |
| Simple DP3 | $100_{\pm 0}$ | $58_{\pm 4}$ | $46_{\pm 5}$ | 68.0 |
| DP3 | $100_{\pm 0}$ | $62_{\pm 4}$ | $43_{\pm 6}$ | 68.3 |
| SODP (ours) | $67_{\pm 6}$ | $96_{\pm 1}$ | $59_{\pm 4}$ | 73.9 |

observations using the Adroit benchmark by employing a point-cloud encoder from DP3 (Ze et al., ) to process the 3D scene represented by point clouds. Specifically, we capture depth images directly from the environment and convert them into point clouds using Open3D (Zhou et al., ). These point clouds are then processed by the DP3 Encoder, which maps them into visual features. We then train our diffusion planner following the same procedure in Algorithm 1 except the input states are visual features. Following DP3 (Ze et al., ), We compute the average of the highest 5 evaluation success rates during training and report the mean and std across 3 seeds. As shown in Table 2, our method outperforms the baseline approaches and achieves an 8.2% improvement, demonstrating its capability to scale up to high-dimensional inputs like images and point-clouds.

### 5.3 Effectiveness of BC Regularization

To demonstrate the effectiveness of our BC regularization, we conduct an ablation study on fine-tuning same pre-trained model with our BC regularization and other variants. We consider following variants:

- •
SODP w/o regularization. This variant is similar to DDPO (Black et al., ) and DPPO (Ren et al., ), which fine-tunes the model directly using Eq. (12) without any regularization.
- •
SODP_kl. This variant is similar to DPOK (Fan et al., ), with the addition of a KL regularization term to constrain the divergence between the fine-tuned model and the pre-trained model.
- •
SODP_pl. This variant is similar to DLPO (Chen et al., ), incorporating the original diffusion pre-training loss (PL) into the fine-tuning objective to prevent the model from deviation.

The details of these variants are presented in Appendix D. Figure 5 demonstrates the effectiveness of our regularization in achieving a higher success rate. We observe that directly fine-tuning the model without any regularization results in the worst performance, with a decline in success rate, as the model may degrade the capabilities acquired from pre-training due to the lack of constraints.
However, adding KL and PL is insufficient, as they cause oscillations near the pre-trained model. This aligns with the original intent of these regularizers, which is to prevent excessive deviation.
This is reasonable for methods like DPOK and DLPO, which utilize pre-trained models such as Stable Diffusion (Rombach et al., ) and WaveGrad2 (Chen et al., ). These models already exhibit strong generative capabilities without fine-tuning, and the goal is to make slight adjustments to align them with more fine-grained attributes, such as aesthetic scores and human preferences.

In contrast, our model is pre-trained on sub-optimal data and lacks the ability to solve complex tasks. We expect it to develop new skills for completing these tasks through fine-tuning. However, directly applying KL regularization to the pre-trained model leads to conservative policies that heavily rely on the existing capability, thereby confining the model to a sub-optimal region. While PL regularization allows some slight exploration, it is uncontrolled and random. Consequently, we observe that the KL regularization almost remains unchanged and the PL regularization slightly increases the performance in basketball but decreases in other tasks.

Figure: Figure 5: Learning efficiency for different regularization. X-axis represents environment steps. Performance declines significantly without any regularization. Both KL and PL regularization confine the model to sub-optimal regions. In contrast, our BC regularization effectively guides the model away from these sub-optimal areas, facilitating the attainment of optimal actions.
Refer to caption: x5.png

#### Visualization.

We hypothesize that the effectiveness of our BC regularization lies in two aspects: (i) it ensures that our model can reuse the skills it has acquired, thereby preventing a decline in performance; (ii) It guides our model to effectively explore optimal regions due to the utilization of optimal $\mu$ as the target policy. To demonstrate this, we visualize trajectories of using the actions generated by our planner using t-SNE (Van der Maaten & Hinton, ). As shown in Figure 6, the trajectory distribution after fine-tuning with KL regularization closely resembles the original pre-training distribution, indicating that the model is reusing learned actions and lacks exploration into new regions. The exploration in PL is unstructured as it may lead to worse regions (e.g., the upper-left region in basketball). In contrast, our method demonstrates superior exploration capabilities to discover new, high-reward regions based on acquired knowledge (e.g., the lower-left region in basketball and the bottom region in plate-slide). Meanwhile, the model can derive valuable insights from pre-trained knowledge by exploiting discovered high-reward actions (e.g. the central region in plate-slide) while discarding low-reward actions.

Figure: Figure 6: Visualization of trajectories using generated actions for different regularization. KL and PL regularization results in conservative policies with distributions closely resembling the original. Our BC regularization retains pre-trained knowledge while effectively discovering new actions that can lead to high rewards.
Refer to caption: x6.png

### 5.4 Effectiveness of Pre-training

We investigate the impact of pre-training. We compare the performance of SODP with a version trained from scratch (SODP_scratch). For SODP_scratch, we use the same rollouts generated by the pre-trained model to approximate the target policy and initialize the replay buffer.

Figure 7 shows that fine-tuning the planner from scratch results in worse performance. Without pre-training, the planner lacks an action prior to guide its behavior, leading to stagnation as it struggles to move towards high reward regions. Additionally, it becomes unstable, as the limited useful knowledge is easily disrupted by a large number of ineffective trials.

Figure: Figure 7: Effectiveness of pre-training. X-axis represents environment steps. Fine-tuning from scratch struggles to identify high-reward actions due to the lack of representation prior. In contrast, pre-training allows the planner to extract useful knowledge, guiding fine-tuning by refining the prior distribution towards more effective behaviors.
Refer to caption: x7.png

## 6 Conclusion

We propose SODP, a novel framework for training a versatile diffusion planner using sub-optimal data. By effectively combining pre-training and fine-tuning, we capture broad behavioral patterns drawn from large-scale multi-task transitions and then rapidly adapt them to achieve higher performance in specific downstream tasks. During fine-tuning, we introduce a BC regularization method, which preserves the pre-trained model’s capabilities while guiding effective exploration. Experiments demonstrate that SODP achieves superior performance across a wide range of challenging manipulation tasks. In future work, we aim to develop embodied versatile agents that can effectively learn to solve real-world tasks using inferior data.

## Appendix A Derivations

### A.1 Derivation of policy gradient in Equation  ( 11 )

Assume $p_{\theta}(\bm{a}_{t}^{0:K}|\bm{s}_{t})r^{\mathcal{T}}(\bm{a}_{t}^{0})$ and $\nabla_{\theta}p_{\theta}(\bm{a}_{t}^{0:K}|\bm{s}_{t})r^{\mathcal{T}}(\bm{a}_{
t}^{0})$ are continuous (Fan et al., ), we have:

$$ $\displaystyle\nabla_{\theta}J^{\mathcal{T}}(\theta)$ $\displaystyle=\nabla_{\theta}\sum\nolimits_{t}\mathbb{E}_{p_{\theta}(\bm{a}_{t }^{0}|\bm{s}_{t})}\left[r^{\mathcal{T}}(\bm{a}_{t}^{0})\right]$ (16) $\displaystyle=\sum\nolimits_{t}\left[\nabla_{\theta}\int r^{\mathcal{T}}(\bm{a }_{t}^{0})\cdot p_{\theta}(\bm{a}_{t}^{0}|\bm{s}_{t})d\bm{a}_{t}^{0}\right]$ $\displaystyle=\sum\nolimits_{t}\left[\nabla_{\theta}\int r^{\mathcal{T}}(\bm{a }_{t}^{0})\cdot\left(\int p_{\theta}(\bm{a}_{t}^{0:K}|\bm{s}_{t})d\bm{a}^{1:K} _{t}\right)d\bm{a}_{t}^{0}\right]$ $\displaystyle=\sum\nolimits_{t}\left[\int r^{\mathcal{T}}(\bm{a}_{t}^{0})\cdot \nabla_{\theta}\log p_{\theta}(\bm{a}_{t}^{0:K}|\bm{s}_{t})\cdot p_{\theta}( \bm{a}_{t}^{0:K}|\bm{s}_{t})\;d\bm{a}_{t}^{0:K}\right]$ $\displaystyle=\sum\nolimits_{t}\left[\int r^{\mathcal{T}}(\bm{a}_{t}^{0})\cdot \nabla_{\theta}\log\left(p_{K}(\bm{a}^{K}_{t}|\bm{s}_{t})\prod_{k=1}^{K}p_{ \theta}(\bm{a}^{k-1}_{t}|\bm{a}^{k}_{t},\bm{s}_{t})\right)\cdot p_{\theta}(\bm {a}_{t}^{0:K}|\bm{s}_{t})\;d\bm{a}_{t}^{0:K}\right]$ $\displaystyle=\sum\nolimits_{t}{\mathbb{E}_{p_{\theta}(\bm{a}_{t}^{0:K}|\bm{s} _{t})}}\left[r^{\mathcal{T}}(\bm{a}_{t}^{0})\sum_{k=1}^{K}\nabla_{\theta}\log p _{\theta}(\bm{a}^{k-1}_{t}|\bm{a}^{k}_{t},\bm{s}_{t})\right].$ $$

### A.2 Derivation of loss function in Equation ( 12 )

By using importance sampling approach, we can rewrite Eq. (16) as follows:

$$ $\sum\nolimits_{t}\mathbb{E}_{p_{\theta_{\text{old}}}(\bm{a}_{t}^{0:K}|\bm{s}_{ t})}\left[r^{\mathcal{T}}(\bm{a}^{0}_{t})\sum_{k=1}^{K}\frac{p_{\theta}(\bm{a} _{t}^{k-1}|\bm{a}_{t}^{k},\bm{s}_{t})}{p_{\theta_{\text{old}}}(\bm{a}_{t}^{k-1 }|\bm{a}_{t}^{k},\bm{s}_{t})}\nabla_{\theta}\log p_{\theta}(\bm{a}^{k-1}_{t}| \bm{a}^{k}_{t},\bm{s}_{t})\right]$ (17) $$

Then, we can get a new objective function corresponding to Eq. (17) as:

$$ $J^{\mathcal{T}}_{\theta_{\text{old}}}(\theta)=\mathop{\max}_{\theta}\sum \nolimits_{t}\mathbb{E}_{p_{\theta_{\text{old}}}(\bm{a}_{t}^{0:K}|\bm{s}_{t})} \left[r^{\mathcal{T}}(\bm{a}^{0}_{t})\sum_{k=1}^{K}\frac{p_{\theta}(\bm{a}_{t} ^{k-1}|\bm{a}_{t}^{k},\bm{s}_{t})}{p_{\theta_{\text{old}}}(\bm{a}_{t}^{k-1}| \bm{a}_{t}^{k},\bm{s}_{t})}\right]$ (18) $$

Let $\rho_{k}(\theta,\theta_{\text{old}})=\frac{p_{\theta}(\bm{a}_{t}^{k-1}|\bm{a}_
{t}^{k},\bm{s}_{t})}{p_{\theta_{\text{old}}}(\bm{a}_{t}^{k-1}|\bm{a}_{t}^{k},
\bm{s}_{t})}$ denote the probability ratio.
Based on PPO (Schulman et al., ), we clip $\rho_{k}$ and use the minimum between the clipped and unclipped ratios to derive a lower bound of the original objective (18), which serves as our final objective function:

$$ $J^{\mathcal{T}}_{\text{clip}}(\theta)=\mathop{\max}_{\theta}\sum\nolimits_{t} \mathbb{E}_{p_{\theta_{\text{old}}}(\bm{a}_{t}^{0:K}|\bm{s}_{t})}\left[r^{ \mathcal{T}}(\bm{a}^{0}_{t})\sum_{k=1}^{K}\text{min}\Big{(}\rho_{k}(\theta, \theta_{\text{old}}),\text{clip}\left(\rho_{k}(\theta,\theta_{\text{old}}),1+ \epsilon,1-\epsilon\right)\Big{)}\right]$ (19) $$

To refine our pre-trained planner, we employ the negative of objective (19) as the loss function to facilitate reward maximization during fine-tuning.

### A.3 Derivation of loss function in Equation ( 14 )

Directly computing and minimizing the NLL is difficult. However, we can derive an upper bound of Eq. (14) as follows:

$$ $\displaystyle\mathbb{E}_{\bm{a}_{\mu}^{0}\sim p_{\mu}}\left[-\log p_{\theta}( \bm{a}_{\mu}^{0})\right]$ $\displaystyle\leq\mathbb{E}_{\bm{a}_{\mu}^{0}\sim p_{\mu}}\left[\mathbb{E}_{q( \bm{a}^{1:K}_{\mu}|\bm{a}^{0}_{\mu})}\left[-\log\frac{p_{\theta}(\bm{a}_{\mu}^ {0:K})}{q(\bm{a}^{1:K}_{\mu}|\bm{a}_{\mu}^{0})}\right]\right]$ (20) $\displaystyle=\mathbb{E}_{\bm{a}_{\mu}^{0}\sim p_{\mu}}\left[\mathbb{E}_{q(\bm {a}^{1:K}_{\mu}|\bm{a}^{0}_{\mu})}\left[-\log p(\bm{a}^{K}_{\mu})-\sum_{k=1}^{ K}\log\frac{p_{\theta}(\bm{a}_{\mu}^{k-1}|\bm{a}_{\mu}^{k})}{q(\bm{a}_{\mu}^{k }|\bm{a}_{\mu}^{k-1})}\right]\right]$ $\displaystyle=\mathbb{E}_{\bm{a}_{\mu}^{0}\sim p_{\mu}}\left[\sum_{k=2}^{K} \mathbb{E}_{q(\bm{a}^{k}_{\mu}|\bm{a}^{0}_{\mu})}D_{\text{KL}}[q(\bm{a}^{k-1}_ {\mu}|\bm{a}^{k}_{\mu},\bm{a}^{0}_{\mu})||p(\bm{a}^{k-1}_{\mu}|\bm{a}^{k}_{\mu })]+\right.$ $\displaystyle\left.D_{\text{KL}}(q(\bm{a}^{K}_{\mu}|\bm{a}^{0}_{\mu})||p(\bm{a }^{K}_{\mu}))-\mathbb{E}_{q(\bm{a}^{1}_{\mu}|\bm{a}^{0}_{\mu})}\left[\log p_{ \theta}(\bm{a}_{\mu}^{0}|\bm{a}_{\mu}^{1})\right]\right]$ $$

Following previous work (Ho et al., ), the optimization of the bound can be simplified as:

$$ $\arg\min_{\theta}\frac{1}{2\sigma_{q}^{2}(k)}\frac{(1-\alpha_{k})^{2}}{(1-\bar {\alpha}_{k})\alpha_{k}}\left\|\epsilon(\bm{a}_{\mu}^{k},k)-\epsilon_{\theta}( \bm{a}_{\mu}^{k},k)\right\|^{2}$ (21) $$

where:

$$ $\sigma_{q}^{2}(k)=\frac{(1-\alpha_{k})(1-\bar{\alpha}_{k-1})}{1-\bar{\alpha}_{ k}}$ (22) $$

Here, $\epsilon_{\theta}(\bm{a}_{\mu}^{k},k)$ is a noise model that learns to predict the source noise $\epsilon(\bm{a}_{\mu}^{k},k)$ which determines $\bm{a}_{\mu}^{k}$ from $\bm{a}_{\mu}^{0}$.

## Appendix B The Details of SODP

### B.1 Diffusion Policy

We use diffusion policy (Chi et al., ) to generate future actions. For any given time step $t$, the model uses the most recent $T_{o}$ steps of states as input to generate the next $T_{p}$ action steps. Then, the first $T_{a}$ steps of these generated actions are executed in the environment without re-planning. In our experiments, we use $T_{p}=12,T_{o}=2,T_{a}=8$ for Meta-World and $T_{p}=4,T_{o}=2,T_{a}=3$ for Adroit.

We employ a CNN-based diffusion policy as our noise model, utilizing a U-net architecture that incorporates Feature-wise Linear Modulation (FiLM) (Perez et al., ) to condition on historical states. The implementation is based on the code from [https://github.com/CleanDiffuserTeam/CleanDiffuser](https://github.com/CleanDiffuserTeam/CleanDiffuser), and we use their default hyper-parameters. For Adroit, we use a simplified backbone provided by Simple DP3 ([https://github.com/YanjieZe/3D-Diffusion-Policy](https://github.com/YanjieZe/3D-Diffusion-Policy)), which removes some components in the U-net.

### B.2 Implementation Details

The pseudo-code of SODP is given in Alg. 1. We describe details of pre-training and fine-tuning as follows:

- •
For pretraining, we use cosine schedule for $\beta_{k}$ (Nichol & Dhariwal, ) and set diffusion steps $K=100$. We pre-train the model for $5e^{5}$ steps in Meta-Wrold and $3e^{3}$ steps in Adroit.
- •
For fine-tuning, we use DDIM (Song et al., ) with $10$ sampling steps and $\eta=1$. We fine-tune each task for $1e^{6}$ steps in Meta-World and $3e^{3}$ steps in Adroit. Following DPOK (Fan et al., ), we perform $p_{\text{step}}\in\{10,30\}$ gradient steps per episode. We set discount factor $\gamma=1$ for all tasks.
- •
We set $N_{\text{init}}\in\{10,20\}$ for approximating target distribution.
- •
Batch size is set to 256 for both pre-training and fine-tuning.
- •
We use Adam optimizer (Kingma, ) with default parameters for both pre-training and fine-tuning. Learning rate is set to $1e^{-4}$ for pretraining and $1e^{-5}$ for fine-tuning with exponential decay.

Figure: Algorithm 1 SODP: Two-stage framework for learning from sub-optimal data

## Appendix C The Details of Baselines

We describe the details of baselines used for comparison in our experiments. For Meta-World, we consider following baselines:

- •
MTBC.
The actor network is modeled using a 3-layer MLP with Mish activation. In training and inference, the scalar task ID is processed through a separate 3-layer MLP with Mish activation to produce a latent variable $z$. The input to the actor network is then formed by concatenating the original state with this latent variable $z$
- •
MTIQL.
Similar to MTBC, the actor network incorporates the task ID through a task-aware embedding. A multi-head critic network is employed to estimate the $Q$-values for each task, with each head being parameterized by a 3-layer MLP using Mish activation.
- •
MTDT.
The task ID is embedded into a latent variable $z$ of size 12. This latent variable is then concatenated with the raw state to form the input tokens.
- •
Prompt-DT.
Actions are generated based on trajectory prompts and the reward-to-go. A GPT-2 transformer model is utilized as the noise network.
- •
MTDIFF.
Actions are generated by a GPT-based diffusion model that incorporates prompt learning to capture task knowledge. MTDIFF considers a variant: MTDIFF-ONEHOT, which replaces the prompt with a one-hot task ID. We borrow the official codes from [https://github.com/tinnerhrhe/MTDiff](https://github.com/tinnerhrhe/MTDiff) and use their default hyper-parameters.
- •
HarmoDT.
Incorporate trainable task-specific masks to address gradient conflict by identifying an optimal harmony subspace of parameters for each task. There are three variants of HarmoDT: HarmoDT-R, which keeps task masks unchanged; HarmoDT-F and HarmoDT-M utilize different methods to weight masks.
We borrow the official codes from [https://github.com/charleshsc/HarmoDT](https://github.com/charleshsc/HarmoDT) and use their default hyper-parameters.

For Adroit, we consider following baselines:

- •
BCRNN.
A variant of BC that models the policy network as an RNN. The network is trained on temporal sequences of length $H$, denoted as $(s_{t},a_{t},...,s_{t+H},a_{t+H})$, to predict action sequences based on historical states.
- •
IBC.
BC is represented as a conditional energy-based modeling problem, where implicit policies are trained to imitate expert demonstrations.
- •
Diffusion Policy.
The generation of robot behaviors is formulated as a conditional denoising diffusion process, where the diffusion model predicts action sequences based on given observations as conditions.
- •
DP3.
Diffusion Policy is extended by incorporating 3D visual representations. The 3D scenes from the environment are represented as point clouds, which are then cropped and downsampled to reduce redundant information. These processed point clouds are passed through an MLP to generate visual representations, which serve as conditions for the diffusion models.

## Appendix D Variants of SODP

In Eq. (15), we introduce a BC regularization term to preserve the pre-trained knowledge and demonstrate its effectiveness compared to two existing regularization approaches presented in DPOK(Fan et al., ) and DLPO (Chen et al., ). Specifically, the regularization term $\mathcal{L}_{\text{KL}}$ used in DPOK is expressed as:

$$ $\mathcal{L}_{\text{KL}}(\theta)=\sum_{k=1}^{K}\text{KL}(p_{\theta}(x_{k-1}|x_{ k})||p_{\text{pre}}(x_{k-1}|x_{k})).$ (23) $$

And the regularization term $\mathcal{L}_{\text{PL}}$ used in DLPO is expressed as:

$$ $\mathcal{L}_{\text{PL}}(\theta)=\mathbb{E}_{k\sim[1,K],p_{\theta}(x_{1:K})} \left[\left\|\epsilon(x_{k},k)-\epsilon_{\theta}(x_{k},k)\right\|^{2}\right].$ (24) $$

These methods can be considered as different approaches to selecting the target policy in line with our analysis and can be seen as variants of Eq. (14), where DPOK selects $\mu=\theta_{\text{pre-train}}$ and DLPO selects $\mu=\theta$. The rationale behind their selection is based on the assumption that $\theta\approx\theta_{\text{pre-train}}\approx\theta^{*}$. This assumption is reasonable in text-to-image or text-to-speech tasks, as the pre-trained models they used are already strong and perform exceptionally well even without fine-tuning.
However, this assumption does not apply to our pre-trained planner, as the model is trained on sub-optimal data. As a result, as shown in Section 5.3, these regularization methods may lead the pre-trained planner to be stuck in inferior regions, limiting its ability to improve performance.

## Appendix E Comparison to DPPO

We summarize some similarities and differences between our work and the concurrent work DPPO (Ren et al., ) as follows:

- •
Both DPPO and our approach formulate the diffusion policy denoising process as an MDP and use policy gradients to fine-tune the model for higher environment rewards.
- •
DPPO demonstrated that reward-based RL fine-tuning promotes effective exploration, which is consistent with our observations.
- •
While DPPO requires task-specific expert demonstrations for pre-training, our method pre-trains a foundation model capable of capturing useful behavior patterns from multi-task inferior data.
- •
We show that directly fine-tuning the pre-trained planner without any regularization, as done in DPPO, fails in the multi-task setting. We further analyze the limitations of current regularization methods and propose a novel BC regularization term. By employing our regularizer, the pre-trained model achieves higher success rates after fine-tuning.
- •
Unlike DPPO, we don’t employ advantage estimator.

## Appendix F Single-Task Performance

We evaluate the performance for each task for 50 episodes. We report the average evaluated return of pre-trained and fine-tuned models in Table 3.

**Table 3: Evaluated return of SODP pre-trained model and fine-tuned model for each task in MT50-rand. We report the mean and standard deviation for 50 episodes for each task.**
| Tasks | Return of pre-trained model | Return of fine-tuned model |
| --- | --- | --- |
| basketball-v2 | $133.5\pm 100.7$ | $2347.1\pm 580.8$ |
| bin-picking-v2 | $96.8\pm 23.9$ | $602.7\pm 72.8$ |
| button-press-topdown-v2 | $1405\pm 20.3$ | $1679\pm 25.9$ |
| button-press-v2 | $1397\pm 15.6$ | $2452.7\pm 89.3$ |
| button-press-wall-v2 | $1375\pm 10.58$ | $2524.7\pm 18.1$ |
| coffee-button-v2 | $293.2\pm 12.1$ | $451.5\pm 14.4$ |
| coffee-pull-v2 | $39.5\pm 6.2$ | $117.9\pm 23.3$ |
| coffee-push-v2 | $33.8\pm 6.1$ | $273.3\pm 36.6$ |
| dial-turn-v2 | $1217.7\pm 239.3$ | $1557.3\pm 226.7$ |
| disassemble-v2 | $237\pm 117.2$ | $502\pm 164.6$ |
| door-close-v2 | $3347.7\pm 124.9$ | $4116.3\pm 118.6$ |
| door-lock-v2 | $1042.3\pm 94.9$ | $2491\pm 79.5$ |
| door-open-v2 | $2036.3\pm 79.3$ | $2460.3\pm 57.7$ |
| door-unlock-v2 | $1335\pm 46.9$ | $2257.7\pm 323.0$ |
| hand-insert-v2 | $85.9\pm 56.7$ | $449.5\pm 54.9$ |
| drawer-close-v2 | $2468.3\pm 167.2$ | $3953.7\pm 214.3$ |
| drawer-open-v2 | $1656\pm 45.6$ | $2489.7\pm 188.4$ |
| faucet-open-v2 | $2728.7\pm 424.1$ | $4094.7\pm 290.3$ |
| faucet-close-v2 | $2156.7\pm 113.6$ | $3772\pm 70.1$ |
| handle-press-side-v2 | $1919.7\pm 449.5$ | $3478.3\pm 98.0$ |
| handle-press-v2 | $2216.3\pm 182.0$ | $3415.7\pm 221.6$ |
| handle-pull-side-v2 | $1351.7\pm 119.0$ | $2665.7\pm 243.9$ |
| handle-pull-v2 | $1510.7\pm 111.6$ | $2734\pm 64.3$ |
| lever-pull-v2 | $650.7\pm 32.5$ | $1068.8\pm 110.4$ |
| peg-insert-side-v2 | $300.3\pm 122.2$ | $1969.7\pm 237.6$ |
| pick-place-wall-v2 | $596.7\pm 10.6$ | $1175.7\pm 150.1$ |
| pick-out-of-hole-v2 | $38.5\pm 6.3$ | $106.7\pm 7.9$ |
| reach-v2 | $2664.3\pm 77.5$ | $3083.7\pm 149.7$ |
| push-back-v2 | $55.8\pm 26.7$ | $350.9\pm 30.3$ |
| push-v2 | $46.8\pm 35.9$ | $148.9\pm 43.9$ |
| pick-place-v2 | $3.9\pm 0.2$ | $5.9\pm 1.8$ |
| plate-slide-v2 | $1268.7\pm 75.8$ | $2862\pm 234.5$ |
| plate-slide-side-v2 | $826.4\pm 54.3$ | $1929.7\pm 104.5$ |
| plate-slide-back-v2 | $795.8\pm 62.9$ | $1587.7\pm 125.0$ |
| plate-slide-back-side-v2 | $626.7\pm 36.9$ | $1541.3\pm 78.5$ |
| soccer-v2 | $863.8\pm 159.5$ | $1234.2\pm 225.8$ |
| push-wall-v2 | $175.2\pm 35.0$ | $471.7\pm 73.9$ |
| shelf-place-v2 | $260.4\pm 111.5$ | $785.1\pm 81.5$ |
| sweep-into-v2 | $621.0\pm 132.9$ | $1282.3\pm 135.9$ |
| sweep-v2 | $442.3\pm 83.0$ | $1081.7\pm 58.0$ |
| window-open-v2 | $1342.3\pm 60.1$ | $2474.3\pm 266.4$ |
| window-close-v2 | $1087.7\pm 78.9$ | $1816.7\pm 166.3$ |
| assembly-v2 | $282.5\pm 4.3$ | $446.1\pm 26.9$ |
| button-press-topdown-wall-v2 | $1374\pm 16.1$ | $1702.7\pm 71.5$ |
| hammer-v2 | $1678.3\pm 52.5$ | $1907.7\pm 25.4$ |
| peg-unplug-side-v2 | $34.2\pm 2.9$ | $52.9\pm 4.6$ |
| reach-wall-v2 | $3373.7\pm 41.7$ | $3839.7\pm 69.3$ |
| stick-push-v2 | $412.9\pm 95.8$ | $833.5\pm 99.1$ |
| stick-pull-v2 | $1977\pm 155.9$ | $3116.3\pm 58.1$ |
| box-close-v2 | $692.3\pm 22.8$ | $1300.1\pm 53.2$ |

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
questions: 1. Is the 3D visual feature process not compatible with MTDIFF? Have the authors tried combining the 3D visual feature process with MTDIFF on the Adroit task?
2. What does the vertical orange dashed line indicating in Figure 4?
soundness: 3
strengths: 1. The paper is well-written and includes illustrative figures that aid comprehension.

2. It reports that SODP not only surpasses baseline methods in performance but also shows rapid convergence and robustness across diverse tasks and input modalities. Additionally, the included ablation study offers valuable insights into the design and effectiveness of the proposed methods.
summary: This paper introduces a novel framework called SODP, designed to train a multi-task diffusion planner using sub-optimal data. SODP aims to reduce dependency on task-specific labeled data through a two-stage training process. Initially, a diffusion model is pre-trained without the need for task-specific rewards or demonstrations. Subsequently, the model undergoes a second stage of online interactive reinforcement learning-based fine-tuning to rapidly refine its capabilities. The effectiveness of this methodology is validated on two multi-task domains: Meta-World and Adroit. Experimental results show that SODP outperforms the considered baseline methods, highlighting its potential in multi-task learning environments.
weaknesses: My primary concerns arise from the evaluation protocols, specifically:

1.	Was the fine-tuning stage also applied to the baseline methods? If not, this might represent an unfair comparison, as all baselines are trained on sub-optimal data while SODP gains an advantage from an additional online fine-tuning stage.
2.	Regarding the Adroit task, was the 3D visual feature process (from DP3) also employed for the Diffusion Policy in Table 2? Moreover, what is the performance of SODP when this process is not utilized?

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
questions: * For different tasks with different state spaces, how to train these different tasks?
* After fine-tuning the pre-trained model on the specific task, can you evaluate the performance of the fine-tuned model on the other tasks? I'm very concerned that the catastrophic forgetting of the model has been addressed.
* Can you explain how to select the target policy in details?
soundness: 3
strengths: * It leverages sub-optimal, task-agnostic data for pre-training, which contrasts with traditional reliance on task-specific demonstrations or expert data.
* The proposed behavior-cloning (BC) regularization in the fine-tuning stage offers a creative solution to the problem of model drift, allowing the model to maintain useful pre-trained capabilities while adapting to new tasks.
summary: SODP is a two-stage diffusion-based framework designed to enhance multi-task planning in reinforcement learning. This framework comprises a pre-training phase that leverages large-scale, sub-optimal, task-agnostic data to learn a broad range of behaviors and a fine-tuning phase that uses task-specific rewards for adaptation. The pre-trained model captures general action patterns, allowing it to rapidly adapt to diverse tasks through fine-tuning, guided by reinforcement learning techniques. SODP demonstrates improved generalization in handling diverse tasks with minimal task-specific data and outperforms existing methods on benchmarks like Meta-World and Adroit.
weaknesses: * Action and State Space Limitations: SODP requires tasks to share the same action space, limiting adaptability to tasks with different action spaces. It’s unclear how the method handles varying state space dimensions.

* Typographical Errors: There are minor errors, such as $a_0^t$ in Line 215, which should be $a^0_t$.

* Dependence on Online Fine-tuning: SODP relies on an online environment for fine-tuning, which may limit practical applications. The paper also lacks details on the number of fine-tuning steps and associated costs.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 3
questions: See weakness. 
I believe that more experiments are needed to illustrate the soundness and Contribution of the paper's SETTINGS and Designs.
soundness: 2
strengths: Pros:

- The paper is well-written with clear organization and easy to follow.

- The fine-tuning stage presents a systematic design, particularly in its BC regularization mechanism that effectively balances between preserving pre-trained knowledge and exploring new high-reward behaviors, supported by comprehensive ablation studies comparing different regularization strategies.
summary: This paper introduces SODP, a two-stage framework for training versatile diffusion planners for robotic manipulation tasks. The framework first pre-trains a diffusion model using large-scale sub-optimal data from multiple tasks without reward labels, then fine-tunes it for specific downstream tasks using policy gradient optimization with a novel BC regularization.
weaknesses: Cons:
1. While the paper claims using sub-optimal data for pre-training as an innovation, there appears to be no specialized design or methodological advancement in handling such data. The pre-training stage follows standard diffusion model training procedures, identical to those used with expert demonstrations.  
2. A significant concern arises regarding the fairness of experimental comparisons. First, the learning curves in Figure 4 only display fine-tuning steps while omitting SODP's substantial pre-training phase (5e5 steps), potentially understating the total computational requirements. Second, the paper fails to clarify whether baseline methods undergo pre-training or have access to online data collection during training. This is particularly important as SODP benefits from continuous online data collection to build its target policy, which may give it an unfair advantage if baselines are limited to offline data.
3. A fundamental limitation is that the paper fails to isolate whether the performance gains stem from sub-optimal data pre-training or from the online data collection during fine-tuning. Critical control experiments are missing, particularly (1) fine-tuning with high-quality offline data and (2) direct training with high-quality data without pre-training. Without these comparisons, it remains unclear if the proposed two-stage framework and the use of sub-optimal data are truly necessary, as the improvements might simply result from the continuous collection of high-quality samples during online fine-tuning.
4. A direct empirical comparison between optimal and sub-optimal pre-training data is missing. Specifically, the paper should compare the current approach (using the first 50% of SAC training data) with a variant using high-reward trajectories (e.g., the last 30% of converged SAC data) for pre-training. This would help validate whether sub-optimal data is truly sufficient or if optimal data would lead to substantially better performance.
5. Why not directly BC in finetune phase. Since the method already collects high-reward trajectories during online interaction for BC regularization, why not directly fine-tune the pre-trained model using imitation learning on these trajectories? This simpler approach might achieve similar performance without the complexity of MDP formulation and policy gradient optimization. Without this comparison, it's unclear whether the proposed fine-tuning framework provides substantial benefits over straightforward behavior cloning of high-reward samples.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: 1. Is the weight coefficient λ in Eq. 15 sensitive to changes? I could not find any implementation details or discussion regarding this hyperparameter. A sensitivity analysis for this hyperparameter should be provided to show how different values of λ affect the performance of their method.
2. In line 332, the statement “state space differs across tasks” could be clarified further. Does this mean that the dimensions of the state space differ across tasks or that the physical meanings of each index in the state space vary between tasks? If it refers to dimensional differences, how is this managed during the pre-training stage with multi-task data?
3. Could you clarify the meaning of line 968: “We show that directly fine-tuning the pre-trained planner without any regularization, as done in DPPO, fails in the multi-task setting”? As I understand it, fine-tuning task-specific rewards would shift to a single-task setting in the fine-tuning stage. Do you mean that the diffusion planner is actually fine-tuned on multiple tasks simultaneously?
soundness: 2
strengths: 1. The approach of using sub-optimal, multi-task trajectories for pre-training to enhance downstream task performance after fine-tuning is interesting.
2. The paper offers a comprehensive discussion on various regularization techniques during fine-tuning, and the experiments show that the proposed BC-regularization method effectively balances reward maximization with preventing performance collapse.
summary: This paper proposes a pre-training and fine-tuning method, SODP, for diffusion planner. The model is pre-trained on sub-optimal and multi-task data offline. Unlike reinforcement learning approaches, the pre-training stage does not require rewards. Additionally, the model is designed to operate without task descriptions or trajectory returns as conditions. The authors aim to show that generalizable capabilities can be acquired from multi-task sub-optimal data in pre-training, providing helpful knowledge for task-specific fine-tuning in an online reinforcement learning setting.
weaknesses: 1. Clarification is needed on whether both multi-task and sub-optimal data in pre-training equally contribute to improved fine-tuning performance for specific tasks.
    1. There is no experiment using only single-task sub-optimal data for pre-training. Although previous work [1] suggests that multi-task diffusion planners outperform single-task setups, the setting in this paper is different, involving both pre-training and fine-tuning. It remains unclear whether multi-task data specifically benefits downstream RL fine-tuning or if the model could already gain useful patterns from sub-optimal data by training on the same task as the target downstream task. If this is the case, unrelated tasks in pre-training may be redundant. Thus, experiments or analyses are needed to separate the contributions of multi-task data versus sub-optimal data, at least for a few representative tasks.
    2. The choice of using 50% and 30% of experiences collected from RL agents' replay buffers in Meta-World and Adroit needs to be explained. Additionally, results from varying the ratios of experience while keeping the total number of pre-training transitions fixed would be valuable. This analysis could clarify if there is a trade-off between fine-tuning performance and the quality of pre-training data (i.e. if using lower-quality data for pre-training negatively impacts fine-tuning). A comparison between fine-tuning results using expert data (100% experience) versus sub-optimal data with the exact transition count would be especially insightful.
2. Key baselines are missing, and the baselines used in state-based and image-based environments are inconsistent. 
    1. Since online interactions with the environment are allowed in the fine-tuning stage, online RL methods should be included for a fair comparison. Currently, only BC and offline RL are used as baselines. 
    2. Both BC and offline RL methods are compared in Meta-World, whereas offline RL methods are excluded in Adroit. 
    3. Given that diffusion-based policies in BC, offline RL, and online RL [2-7] have shown strong performance compared to non-diffusion methods in prior work, including offline and online diffusion-based RL baselines would offer a more complete comparison for evaluating the fine-tuning approach with online RL for the diffusion model.
    4. Overall, the most critical baselines to prioritize, in both image and state-based environments, would be online RL (non-diffusion, such as PPO) and diffusion-based offline and online RL methods [2-7]. Including these baselines would allow for a fairer comparison by incorporating methods that offer strong modeling capabilities with diffusion and interaction with the environment.
3. The problem setting for sub-optimal data collection may lack precision, as in line 52, “sub-optimal data can be easily obtained in the real world.” Concrete examples or scenarios of how sub-optimal data might be collected in real-world settings would be helpful. Discussing potential challenges or limitations of these collection methods would provide readers with a clearer understanding of the motivation’s validity.
    *  Currently, the experimental setup collects sub-optimal data by training RL agents and treating part of the replay buffer as "inferior data." However, without rewards, identifying data that qualifies as sub-optimal could be challenging or even infeasible, especially since, as noted in line 153, “reward labels may be scarce or costly to obtain.” If rewards were readily accessible in the offline data or during interaction with the environment, then a diffusion planner could be directly trained using offline or online RL methods [2-7], and pre-training on reward-free sub-optimal data might no longer be necessary. This raises questions about the practicality of the motivation and problem setting proposed in the paper. It would be more meaningful if data with mixed or random levels of inferiority—without requiring a curated collection mechanism—could still enhance downstream fine-tuning during pre-training.
4. Several descriptions are either inconsistent or need further clarification:
    1. The claim in line 362, “The baselines used in Meta-World struggled to handle this high-dimensional data structure,” lacks supporting analysis or experimental evidence. Where in the paper is this conclusion demonstrated? 
    2. Are all baselines for Meta-World evaluated in a multi-task setting, while those used in Adroit are in a single-task setting? Additionally, what does "Simple DP3" in Table 2 refer to? This term is not explained anywhere in the paper.
    3. There are inconsistencies in the experimental setup descriptions. In Sec. 5.1 (EXPERIMENTAL SETUP), line 344 states, “All baselines and our pre-training stage are trained on the same dataset” for Meta-World, while line 346 states, “All baselines are trained on expert demonstrations and our pre-training stage is trained on sub-optimal transitions” for Adroit. However, in Sec. 5.2 (RESULTS), line 376 says, “All baselines are trained on sub-optimal data.” Could you clarify which statement is correct? Do lines 344 and 346 mean that all baselines in Meta-World are trained on the same sub-optimal data used for SODP pre-training, while expert data are used in Adroit baselines? If so, what is the reason behind using different data sources for each environment?
    4. This could introduce unfairness if it is accurate that baselines are trained with sub-optimal data. For SODP fine-tuning, the model can interact with the environment to optimize actions guided by rewards, whereas the BC and offline-RL methods do not have this privilege. Additionally, while offline RL methods can leverage low-reward transitions, BC methods typically require expert-level data to imitate expert policies effectively. Sub-optimal data, especially when only 50% or 30% of the RL agents' replay buffer is used, likely does not approximate expert-level quality. It would be more convincing if offline-RL and BC baselines were trained with expert data or even allowed to interact with the environment, as in methods like GAIL [8].
5. Concerns regarding experimental results and setup:
    1. There is an imbalance in the number of tasks between the state-based and image-based environments: 50 tasks are used in Meta-World, while only 3 are used in Adroit. Why wasn’t image observation utilized in Meta-World to create an image-based environment there as well? Including more image-based tasks would provide a more comprehensive evaluation of SODP’s image-based performance and could clarify whether the performance drop observed in tasks like “Hammer” is an isolated case, as shown in Table 2. If SODP generally struggles with challenging image-based tasks, such as “Hammer,” this would indicate limited generalization ability.
    2. The claim in line 432—“training on inferior data is more difficult, and the baselines trained on expert data perform better in this case”—is not entirely convincing, as SODP has access to interact with the environment and obtain rewards in the fine-tuning stage, which should be advantageous even if the baselines are trained with expert data.
    3. Does the same trend appear in state-based environments, where SODP performs worse in challenging tasks, as observed in the image-based setting? Could you provide task-specific performance results from Table 1 instead of just the average success rate? This would allow for a more precise identification of any special cases.
    4. In Fig. 7, it appears unreasonable that SODP_scratch could not succeed despite being able to interact with the environment and optimize with rewards. Additional experiments with a diffusion planner trained from scratch using online RL methods [2] should be included to see if online RL methods also fail entirely. Training with online RL from scratch would be a fairer comparison for SODP_scratch.
    5. Line 511 states, “We use the same rollouts generated by the pre-trained model to approximate the target policy and initialize the replay buffer.” Could you clarify the methodology used here and discuss its validity for testing the performance of SODP_scratch?

**References**

[1] He et al. "Diffusion model is an effective planner and data synthesizer for multi-task reinforcement learning." In NeurIPS 2024.

[2] Yang et al. "Policy representation via diffusion probability model for reinforcement learning." arXiv 2023.

[3] Kang et al. "Efficient diffusion policies for offline reinforcement learning." In NeurIPS 2024.

[4] Chen et al. "Offline reinforcement learning via high-fidelity generative behavior modeling." In ICLR 2023.

[5] Hansen-Estruch et al. "Idql: Implicit q-learning as an actor-critic method with diffusion policies." arXiv 2023.

[6] Wang et al. "Diffusion policies as an expressive policy class for offline reinforcement learning." In ICLR 2023.

[7] Psenka et al. "Learning a diffusion model policy from rewards via q-score matching." In ICML 2024.

[8] Ho et al.. "Generative adversarial imitation learning." In NeurIPS 2016.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 3
questions: - In equation (11), I assume $\sum_{k=1}^K \triangledown \log p_\theta (a_t^{k-1} | a^{k-1}_t, s_t)$ is modeling $ \triangledown p(a_t | s_t) $ in policy gradient for the RL MDP. The probability of a specific sample is known to be intractable in the diffusion process. Can you provide proof for calculating the gradient of probability of a specific sample in the diffusion process? Note, DDPO is applying the policy gradient to each step of the diffusion reverse process, where the probability of each reverse step is defined as an isomorphic gaussian distribution, whose probabilities are well-defined.
- In equation (12), is "generating new samples" (page 5, line 238) referring to samples in the RL MDP or the diffusion MDP? If it is in the RL MDP, how does applying the surrogate loss to the diffusion MDP help solving the RL MDP?
soundness: 2
strengths: - Introduction of BC Surrogate Loss: The paper introduces a novel Behavior Cloning (BC) surrogate loss that effectively regularizes the fine-tuning process, ensuring stability and preserving pre-trained competencies.
- Effective Use of Sub-Optimal Data: SODP demonstrates an innovative approach for leveraging sub-optimal data, overcoming the need for high-quality labels while still achieving strong performance across tasks.
- Strong Empirical Results: The experimental results on Meta-World and Adroit benchmarks show clear advantages over state-of-the-art methods, highlighting the practical effectiveness of the proposed framework.
summary: The paper presents a novel framework, SODP (Sub-Optimal Diffusion Planner), designed for optimizing reinforcement learning (RL) models by integrating Proximal Policy Optimization (PPO) with a Behavior Cloning (BC) loss function. This approach leverages PPO to maximize task-specific returns during fine-tuning, while incorporating BC regularization to preserve competencies gained during the pre-training stage. The framework addresses the challenge of utilizing sub-optimal data, enabling effective policy adaptation across various downstream tasks.

The authors validate their methodology through extensive empirical evaluations in multi-task environments, demonstrating that SODP can efficiently adapt and learn from low-quality datasets. The experimental results show that the proposed framework consistently outperforms state-of-the-art methods, highlighting its ability to produce competitive outcomes despite the limitations of pre-training on sub-optimal data.
weaknesses: - Insufficient Explanation of Key Equations: The derivation of Equations (11) and (12) is not fully explained. Specifically, the summation of the log probabilities in the reverse diffusion process in Equation (11) lacks clarity, as it differs from the standard $p(a|s)$ formulation commonly used in reinforcement learning. This discrepancy requires further explanation to justify its relevance and correctness.
- Misapplication of Trust Region Loss: In Equation (12), the trust region loss is applied to the diffusion process rather than directly addressing the reinforcement learning process. This choice appears to overlook the key challenge of improving sample efficiency in reinforcement learning, which could undermine the effectiveness of the method in dealing with limited data.
- Over-reliance on BC Regularization: I suspect behavior cloning the better samples from the replay buffer starting from a diffusion model pre-trained on sub-optimal data is the key to the performance gain over baselines in the proposed method. Please consider an experiment with only the BC regularization term and no fine-tuning term. This is supported by the fact that SODP w/o regularization decrease in performance with online samples.

Meta-review: The paper introduces SODP, a two-stage framework for robotic manipulation tasks, which utilizes suboptimal data to train a Diffusion Planner. The approach pre-trains a diffusion model on sub-optimal multi-task data and fine-tunes it using policy gradient optimization with a novel BC regularization. The framework leverages sub-optimal data to avoid reliance on high-quality labels while achieving strong performance during fine-tuning.

Reasons to accept
- The paper is well-structured and easy to follow.
- The proposed BC regularization effectively balances retaining pre-trained knowledge and exploring new high-reward behaviors, supported by thorough ablations.
- The paper presents a novel method to leverage sub-optimal data.
- The proposed method achieves strong empirical performance on Meta-World and Adroit

Reasons to reject
- The role of pre-training with sub-optimal data is not rigorously isolated, and critical experiments (e.g., pre-training with high-reward data or direct fine-tuning without pre-training) are missing.
- Experimental results omit the computational costs of pre-training (e.g., 500k steps) and lack clarity on whether baselines also benefit from online data collection, potentially biasing comparisons.
- The absence of comparisons with straightforward methods, such as direct behavior cloning during fine-tuning, weakens the argument for the necessity of the proposed framework.
- The performance gains might primarily stem from BC regularization rather than the fine-tuning framework, which undermines the novelty of the approach.

While this paper studies a promising research direction and presents an interesting approach, I believe its weaknesses outweigh its strengths. Consequently, I recommend rejecting the paper.

Decision: Reject