Task-Agnostic Pre-training and Task-Guided Fine-tuning for Versatile Diffusion Planner
Authors: Chenyou Fan
Chenjia Bai
Zhao Shan
Haoran He
Yang Zhang
Zhen Wang
Venue: ICML 2025 poster (2025)
Keywords: reinforcement learning
diffusion models
planning
Primary area: reinforcement_learning
TL;DR: A paradigm for learning a versatile diffusion planner from sub-optimal transitions

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). They are costly due to the substantial human efforts required to collect expert data or design reward functions. To address these challenges, we aim to develop a versatile diffusion planner capable of leveraging 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 Sub-Optimal data to learn a Diffusion Planner, 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 quickly 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.

Task Agnostic Pre-training and Task-Guided Fine-tuning for Versatile Diffusion Planner
Chenyou Fan
Chenjia Bai Zhao Shan
Haoran He Yang Zhang
Zhen Wang'
Abstract
1. Introduction
There has been long-standing pursuit to develop agents capable of performing multiple tasks (Reed al., 2022: Lee et al., 2022). Although traditional RL. methods have made significant strides in training agents to master individual tasks (Silver et al., 2016; OpenAl et al., 2019), expanding this capability to handle diverse tasks remains significant challenge due to the diversity task variants and optimiza- tion directions with different rewards. Multi-task RL aims to address this by developing agents via task-conditioned optimization (Yu et al., 2020; Lee al., 2022) parameter- compositional learning (Sun et al., 2022: Lee et al., 2023), However, existing approaches often struggle with policy rep- resentation since they assume shared latent structures across tasks, which causes issues when tasks have significantly dif- ferent dynamics and rewards, limiting their ability to capture the multi-modal distribution of optimal behaviors across di- verse task spaces. Diffusion models (Sohl-Dickstein et al., 2015; Ho al., 2020), originally designed for generative tasks, provide a powerful framework to address these dif- ficulties. Their capacity to capture complex, multi-modal distributions within high-dimensional data spaces (Podell et al., 2024; Ho et al., 2022; Jing et al., 2022) makes them well suited to represent the broad variability encountered in multi-task environments.
Diffusion models have demonstrated their capa- bilities modeling trajectories of multi-tasks. However, existing multi-task planners or policies typically rely on task-specific demonstrations via multi-task imitation, require task-specific re- ward labels to facilitate policy optimization via Reinforcement Learning (RL). They are costly due to the substantial human efforts required to collect expert data or design reward functions. To address these challenges, we aim to develop versatile diffusion planner capable of leverag- ing large-scale inferior data that contains task- agnostic sub-optimal trajectories, with the abil- ity to fast adapt to specific tasks. In this paper, we propose SODP, a two-stage framework that leverages Sub-Optimal data to learn a Diffusion Planner, which is generalizable for various down- stream tasks. Specifically, in the pre-training stage, train 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 cov- erage. Then for downstream tasks, we adopt RL- based fine-tuning with task-specific rewards to quickly 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.
Motivated by this, existing methods have employed diffu- sion models to mimic expert behaviors derived from human demonstrations on various tasks (Pearce et al., 2023; Xu et al., 2023; Chi et al., 2023). However, acquiring task- specific demonstrations often time-consuming and costly, especially in environments requiring specialized domain expertise. Alternative approaches attempt to optimize dif- fusion models with return guidance (He et al., 2023; Liang et al., 2023) or conventional RL paradigm (Wang al., 2023). which demands a large volume of data with reward la- bels for each task. To address the ahove limitations, we won- der whether a generalized diffusion planner can be learned from large amount low- low-quality trajectories without re- ward labels, with the ability adapt quickly various downstream tasks. We only require the inferior data to comprise mixture of sub-optimal state-action pairs from various tasks, which can be easily obtained in the real world.
NNottwwstern Polytechnical University IInstittee of Artificial Intelligence (T) China Telecom 'Shenzhen Research Institute of Northwestern Polytechnical University "Tsinghua University 'Hong Kong University of Science and Technology. Correspon- dence to: Chenjia Bai <baicj eecinatelecom.cn.. Zhen Wang <w <w-zzhe zhen nwpu.cdu.cn>.
Proceedings of the 4dnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Task-Agnosta Pre-training and Task-Guided Fine-tuning for Versatile Diffusion Planner
Offline multitask actions
Online task-speeific actions
Spece
A
Future
Act
>>
Fature sctions
Spm
Task Environment
A
Future actions
Diffusion Model
Task Reward
Fine-tuning with fask reward
(a) Stage 1: Pre-training on multiple task datasets
(b) Stage 2: Fine-tuning on each specific fask
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.
In training, the diffusion planner seeks to model the distribu- tion of diverse trajectories with broad coverage, enabling it to acquire generalizable capabilities and allowing the plan- ner to further concentrate on high-reward regions specific downstream tasks via fast adaptation, An overview of our method is given in Figure 1.
rithm 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 Meta- World (Yu al., 2019), showcasing its superior performance compared to state-of-the-art approaches.
In this paper, we propose 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 capture shared knowledge across tasks. Since the state space may vary between tasks, focusing (e.g., end-effector poses of a robot arm) facilitates task gen- cralization. We frame the pre-training stage as 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., 2022; Bai al., 2025), we adopt RL-based fine-tuning approach to tailor the pre-trained diffusion planner specific down- stream tasks. Specifically, we conduct online interaction based on the pre-trained planner to collect task-specific ex- periences with reward labels, and perform policy gradients to iteratively refine the predicted retion--equunce distribu- tion based on reward feedback of tasks. Through fine-tuning the diffusion planner can gradually adapt toward generat- ing actions with high task-specific rewards and eventually become optimal for the given task. Figure 2 illustrates our method. In pre- training, the model captures diverse behav- ior patterns from training data, encompassing inferior and mediocre actions. After fine-tuning, the model shrinks the action distribution and concentrates on generating optimal
After pretraining
After reward fine-tuning
high
Roward
the common action space
low
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 wide range of behaviors, and the fine-tuned model discards the inferior behaviors coverage high-reward regions.
2. Preliminaries
Multi-task RL We consider the multi-task RL problem involving N tasks, where each task T ~ p(T) is repre- sented by task-specified Markov Decision Process (MDP). Each MDP is defined by tuple (ST AA PT RT PT 7), where S is the state space of task T. A is the global action space, PT (sT az ST x A + ST is the transition function, R (sT af ST x A - R is the reward function, 7 E (0. 1] is the discount factor, and PJ is the initial state distribution. We assume that all tasks action sequences for specific task. Our contributions can be summarized follows. (i) We propose a novel pre- training and fine-tuning paradigm for learning versatile diffusion planner, which leverages sub-optimal transitions to capture the broad action distributions across tasks, and adopt task-spec fine-tuning to transfer the planner to downstream tasks. (ii) We give an efficient fine-tuning algo-
share 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 x? E ST takes an action af E A according to the policy n (aa (aTls?) |sT), and receives a reward rT_ The agent's objective is to determine an optimal pol- icy that maximizes the expected return across all tasks: n arg max, ET E, E EE= 7'r?]
2

---

Task-Agnostic Pre-training and Task-Guided Fine-tuning for Versatile Diffusion Planner
Diffusion Models Diffusion models (Sohl-Dickstein et al., 2015) are type of generative model that first add noise to the data Zo from a unknown distribution (( in K steps through a forward process defined as:
real-world. To overcome this challenge, we train diffusion planner that can learn from offline trajectories transitions (i.e., {(st. ats $t)1 without reward label or task descrip- tions. Specifically, we model the problem as guidance-free generation process (Chi et al., 2023):
g( xk-1) == (Z vI B B,d),
(1)
maxg (s.. D. po st)].
(6)
where 34 is predefined variance schedule. Then, a train- able reverse process is constructed as:
Here, we represent Zo == af = (at, 0r+1..... (++H-1) as an action sequence, where H is the planning horizon and is the timestep sampled from dataset D. We denote st as the historical states at timestep & with length To, i.c., 8, := {s-T.+1.... St-1: st). 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- pre-training objective as:
Po Po(0k=11 (xa)= =- (a No k), Ex),
(2)
where (o (3 k) is the forward process posterior mean as a function of a noise prediction neural network ea (ag,, k) with a learnable parameter 0 (Ho et al., 2020). co(a,,, k) can be trained via surrogate loss as
Lacnoise (0)
[|le co k)]l?]
Lpr train (0) = ( [|| - ca(a 81, A)]l" (7)
(3)
After training, samples can be generated by first drawing Gaussian noise æ and then iteratively denoising æ into noise-free output 26 over K iterations using the trained model co k) by
Following Eq. (4), we can generate action sequences through a series of denoising steps:
1 a --aatc kk,
k-
a
(ookk))oooNNO0 +okN(0,
Vax
VI ōk
1
1 - az
(8) Unlike other models, the dataset 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.
eo (22, ki) + oRN I),
v yok
vI āx
(4)
where ogk == - Br, āg. = I ag and o% = 33..
3. Method
We propose SODP, two-stage framework that leverages large amounts of sub-optimal data to train a diffusion plan- ner 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 incorporat- ing a regularization term to prevent the model from losing acquired skills.
3.2. Reward Fine-tuning for Downstream Tasks
MDP notation. The fine-tuning stage involves two dis- tinet MDPs: one for RL decision process and the other for the diffusion model denoising process. We use the super- script diff (c.g., s ablft to denote the MDP associated with diffusion model denoising process, while no super- script is used for the MDP related to the RL process (e.g., St, at). Additionally, we use k E (K,..., 0) to represent the diffusion timestep and E {l,..., T) represent the trajectory timestep.
3.1. Pre-training with Large-scale Sub-optimal Data
We model the denoising process of our pre-trained diffusion planner as K K-step MDP as follows:
Previous works typically model multi-task RL as a condi- tional generative problem using diffusion models trained on datasets composed of multiple task subsets D = , as:
(st,,, K-K) diff *
diff
K-k-1
O
a
(sshl((s))) pdift diff diff ax
Pahlf( (hill)
(6,,,N00 T));
maxe Dr log (T) (T))
(5)
diff = (6., 8mg))
which requires additional condition y(T to guide diffu- sion model to generate desirable trajectories. For instance, (T) should contain the return of trajectory R(T) and task description Z as prompt. However, the reward label and tra- jectory description may be scarce or costly to obtain in the
((sthil)) 1) = r((a))
if k= K-1 otherwise.
Redii diil dill
(5k .kk
To (a
till difl
dill
po K- K-k-1 | a kk -());)
3

---

Task-Agnost Pre-training and Task-Guided Fine-tuning for Versatile Diffusion Planner
Pretraining
Fine-tuning
Typ Grapp the cnver close
Tat Hammer the wall
Tat Dunk the boal basket
Ta Pull lover danns
dy (ly
Update & for laok Z
Regularize
Froward Process q( (4-1
Policy srudicat
(as)
Soay acciets
actionss
Reverse
Nenay acturss
+
Tr activus
Nony smrvess
polas ilag. 47
Process
actions
log pv(()
Closing
Behavior
Torget policy
Apprnsirate
accuets
actiona
Figure 3. Overview of SODP. We initially pre-train diffusion model using multi-task transition data predict action sequences from historical states. Subsequently, we fine-tune the model on downstream tasks using policy gradient methods, incorporating regularization term mitigate model degradation.
where shifT and d% are the state and action timestep k, Pdiff and pdiff are the initial distribution and transition dy- namics, 5 is the Dirac delta distribution, Rdil is the reward function and (a, ( a.*** K-k 8t) is the pre-trained dif- fusion planner. This formulation allows the state transitions in the MDP to be mapped to the denoising process in the dif- fusion model. The MDP initiates by sampling an initial state s0 Poliff, which corresponds to sample Gaussian noise at at the beginning of the reverse process. At each timestep k. the policy πtlill (ath sdil)) takes action ai based current state <t which corresponds to generate next latent aK a based on current latent a < following Eq. (8). The reward remains zero until noise-free output at is evaluated. Different from previous rext-to-image studies that typically evaluate the final sample using a pre-trained reward model (Black et al., 2024; Fan et al., 2024), we aim to fine-tune the pre- pre-trained diffusion planner to maximize rewards of downstream tasks, which makes constructing reward models for all tasks costly. Therefore, evaluate the generated action sequences in an online RL environment for each specific task T. Specifically, for any given timestep I, we use the planner to generate future ac- tions af = (az, Qs+1,..., (.... and then execute the first T. steps. Then we calculate the discounted cumulative reward from the enviromment to assess the generated sample, expressed as r r(a.) E 20-1, (st, a.). We write r(a.) as shorthand for r(st, a for brevity.
et al., 1999), which estimate the policy gradient and ap- ply stochastic gradient ascent algorithm for updates. The gradient of the objective JJ (0) can be obtained as follows:
V J (0)
K Vo log pa -
rT (T(a)) Σ, (p
(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 lever- age historical sequences, we employ importance sampling, following the approach of proximal policy optimization (PPO) (Schulman al., 2017), and derive the loss function for reward improvement follows:
K
directly
CT (0)
Σ,
Epan (af x (az)
Σ -raaaa))
max (pr(f, (ou), clip (px (6, 1,,1) (12)
-0)]
sc) and is hyperparam- pewa (a; |a!
where (0, Oola)
eter. Then, we can train our model using 4Tmp (0) in an end-to-end manner, which is equivalent to maximizing the objective Eq. (10).
Fine-tuning objective. The objective of fine-tuning our pre- tre-trained diffusion planner is to maximize the expected reward of the generated action sequences for the target down- stream task T, which can be defined as:
Regularization term. Fine-tuning the model solely de- pending 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-Clonn (BC) regulariza- tion term during the fine-tuning process. Concretely, we aim to constrain our policy to closely match a target pol- iey μa, ensuring that does not deviate significantly from μ
JT(()) Σ, E Epp((a))=)) (a.)
(10)
Directly optimizing the objective (T( is intractable since it is infeasible to evaluate the return over all possible ac- tions. Therefore, we utilize policy gradient methods (Sutton

---

Task Task-Agnostic Pre-training and Task-Guided Fine-tuning for Versatile Diffusion Planner
after policy updates. This constraint can be modeled using negative log-likelihood (NLL) loss as:
2021; Huang et al., 2023), and drug design (Schneuing et al.,
2024; Guan et al., 2023). Recent studies have applied them in imitation learning to model human demonstrations and predict future actions (Li et al., 2024; Reuss et al., 2023; He et al., 2024). Other approaches have trained conditional dif- fusion models either as planners (Ajay et al., 2023; Brehmer et al., 2024; Yuan et al., 2024) or policies (Haaneen-Estuur et al., 2023; Kang et al., 2024). However, most of these ef- forts focus 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., 2023) or preference labels (Yu et al., 2024). These methods are limited by their dependence on expert data or explicit task knowledge. In contrast, our method learns broad rattion-sequnne distributions from inferior data to en- hance action priors, enabling effective generalization across range of downstream tasks.
min E Bap [- log
loggo po (a)))
(13)
Following Ho et al. (2020), we can obtain surrogate loss to optimize Eq. (13) as follows:
(0) = =Bp,_[1.K))) [Io((at hi) (co ]]]]
(14) where <(a 1:) represents the ground-truth noise added to a at timestep k.
How to select the target policy? Intuitively, an ideal target policy is the optimal policy that generates samples x² satis- fying C(x²)= 2 C(x)) for all possible æ, where C(æ)) repre- sents a measure of the performance or quality of the sample, such as the accumulated reward for action sequences. Since y is unknown during fine-tuning, we approximate it by sam- pling action sequences a that satisfy C(a)= x C(a"). In practice, we denote a" as the best actions from recent play experience, such as those that yielded the highest n rewards or successfully completed the given task. We then sample a from these proficient actions obtained from online inter- action, nearly equivalent to sampling from μ to regularize the fine-tuning process. We also remark that the BC regu- larizer is not the only way incorporate regularization into Eq. (12). For example, a Kullback-Leibler (KL) divergence between fine-tuned and pre-trained models, or diffusion pre-train loss can be employed to regularize the fine-tuning process, as shown in text-to-image and text-to-speech gen- eration (Fan et al., 2024; Chen et al., 2024). However, we find these regularization may cause the pre-trained planner trap in sub-optimal regions, hindering performance improve- ment. We will further discuss them in experiments.
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., 2022), fairness (Shen et al., 2024), or human preference (Xu et al., 2024), 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., 2023; Clark et al., 2024), 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., 2023). Other methods reformulate the denoising process
MDP and apply policy gradients for fine- tuning (Black et al., 2024; Fan et al., 2024). However, they heavily depend on strong pre-trained models and have proven ineffective in our case. Our goal is to fine-tune less powerful model that has been trained on inferior data.
Combining Eq. (12) with Eq. (14), the loss function for reward fine-tuning in downstream tasks T ~ p(T) is ex- pressed as follows:
Concurrent with our work, DPPO (Ren et al., 2025) also explores reward fine-tuning for refining diffusion planners. However, their approach focuses 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.
CT Cinn--uunig (0)= (0) + (O).
=
(15)
where is a weight coefficient. The overall process of pre-training and fine-tun finentuning using SODP is summarized in Alg. 1. Since our goal is to generate complete trajectories rather than individual segments, we utilize trajectory buffer (Zheng et al., 2022) for estimating the target policy μ.. Further, to ensure the accuracy of the approximation, we generate several proficient trajectories using the pre-trained model the beginning of each iteration.
5. Experiments
In this section, we conduct experiments to evaluate our pro- posed method and address the following questions: (1) How does the performance of SODP compare to other methods that combine offline pre-training with online fine-tuning? (2) How does the performance of SODP compare other multi-task learning approaches?
4. Related Work
Diffusion Models in RL. Diffusson models are a leading class of generative models, achieving state-of-the-an per- formance across a variety of tasks, such as image gener- ation (Ramesh et al., 2021), audio synthesis (Kong et al.,
(3) How does SODP achieve higher rewards during online fine-tuning?
5

---

Task-Agnostic Pre-training and Task-Guided Fine-tuning for Versatile Diffusion Planner
5.1. Experimental Setup
Table I. Average success rate across seeds on Meta-World 50 tasks with random goals (MT50-rand) using sub-optimal data. Each task is evaluated for 50 episodes.
We evaluate SODP in both state-based and image-based environments. We conduct experiments the Meta- World benchmark (Yu et al., 2019) for both state-based and image- based tasks. We also perform image-based experiments on the Adroit benchmark (Rajeswaran et al.. 2018). Image- based experimental results are presented in Appendix C.4.
Method
Meta-World 50 Tasks
RLPD
10. 10.1610.11
IBRL.
25.29 #0.22
Cal-QL
35.09
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., 2023; Hu et al., 2024), we extend all tasks to random- -goal setting, referred to MT50-rand.
MTSAC
42.67=0.11
MTBC
34.53
MTIQL
43.28
MTDQL
17.33
MTDT
12.33
Prompt-DT
18. 18.40.0.16
MIDIFF-P
18.67 #1.32
SEIDIFF-P-NEHHEE
48, 48.9420.95
HarmoDT-R
53.80
HarmoDT-M M
2..20
Harm
57.20
SODP (ours)
60.56_0014
Following previous work (He et al., 2023), we use a sub-optimal offline dataset containing IM transitions for each task. The dataset consists of the first 50% of experi- ences collected from the replay buffer of an SAC (Haarnoja et al., 2018) agent during training. To verify the applicabil- ity 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. (2023).
Datasets.
Transformer architectures with prompt learning facilitate generative planning in multitask offline environments. We extend it with visual extractor in image-based Meta-Worle experiments. (8) HarmoDT (Hu et al., 2024). A D'T-based approach that leverages parameter sharing to exploit task similarities while mitigating the adverse effects of conflict- ing gradients simultaneously. All baselines are trained on the 50M offline transitions.
Baselines. We compare our proposed SODP with several RL baselines that integrate offline pre-train pre-training with online fine-tuning: (1) Cal-QL (Nakamoto et al., 2024). Employs conservative offline initializations to facilitate online fine- tuning. (2) IBRL (Hu al., 2023). Pre-trains a policy using imitation learning and subsequently fine-tunes it RL. policy. (3) RLPD (Ball et al., 2023). Extends SAC by utilizing a sample buffer that incorporates both offline and online data. All baselines, along with SODP, are pre-trained on the same dataset containing 50M transitions and are subsequently fine-tuned on each task with 1M transitions.
5.2. Main Results
We use the average success rate across all tasks as the evalu- ation metric and report the mean and standard deviation of success rates across three seeds. As shown in Table 1, our method achieves 60.56% success rate when learning from inferior data, outperforming all baseline methods. Current offline-online RL methods face challenges in effectively learning from complex, low-quality offline data, resulting in insufficient priors to guide the online fine-tuning. Com- pared the existing state-of-the-ar multi-tast approach, our method demonstrates 5.9% improvement. Notably, when compared to MTDIFF, the current leading method based on diffusion models, our approach shows 24.4% improvement. MTBC performs the worst, as imitation learn- ing heavily depends on data quality, and directly cloning behaviors from sub-optimal data typically results in inferior performance. In contrast, our method models versatile ac- tion distributions from low-quality data and leverages them as priors to guide policy optimization in downstream tasks, leading to improved performance.
Besides, we compare SODP with the following multi-task RL baselines: (1) MTSAC. Extended SAC with one-hot task ID as additional input. (2) MTBC. Extended BC to multi-task learning through network scaling and a task-ID- conditioned actor. (3) MTIQL.. Extended IQL. (Kostrikov et al., 2022) with multi-head critic networks and a task-ID- conditioned actor for multi- task policy learning. (4) MT- DQL. Extended Diffusion-QL. (Wang al., 2023) which is similar to MTIQL. (5) MTDT. Extended Decision Trans- former (DT) (Chen et al., 2021a) to multitask settings by incorporating task ID encoding and state inputs for task- specific learning. (6) Prompt-DT (Xu et al., 2022). An extension of DT, which generates actions by utilizing trajec- tory prompts and reward-to-go signals. (7) MTDIFF (He et al., 2023). uuffussin-bbsed approach that integrates
To further analyze the learning dynamics, we sample 10 tasks and present the learning curves of SODP alongside three offline-online baselines. As shown in Figure 4, SODP demonstrates a higher success rate after fine-tuning, whereas
6

---

Task Task-Agnostic Pre-training and Task-Guided Fine-tuning for Versatile Diffusion Planner
Cal-QL
IBRL
R.PDD
owrs
basketball
box-close
button-pres-topdovr
pantou-press-oodiieee
meme
Suoves
Enx Steps
Steps
Eme Steps
Env Steps
12
Env Steps
handle-poll
bandle-pell-side
peg-insert-side
stick-pull
sweep
Succes Ratc
Eni Steps
Env Stepe
Env Steps
Pri Steps
Eny Steps
Figure 4. Learning dynamics. We sample 10 tasks and present the learning curves of SODP, Cal-QL, IBRL and RLPD across five sceds, X-axis represents environment steps. SODP converges to higher success rates, whereas others struggle with solving the downstream task.
other baseline methods struggle to address the specific task. This performance disparity can be attributed to the inade- quate extraction of meaningful guidance during the offline phase. Due to the presence of numerous low-reward trajec- tories in the sub-optimal dataset, pre-training on such data may misguide the value function. Furthermore, employ- ing a Gaussian policy for behavior cloning fails to capture the multimodal action distributions embedded in the pre- training data. Consequently, this bias can affect the updates during online fine-tuning, potentially trapping the model in sub-optimal plateau or even degrading its overall per- formance. The learning curves of SODP, in comparison to other multi-task baselines, are provided in Appendix C.6. We also conduct additional experiments on offline-online baselines, including both diffusion- based and transformer- based methods, with the results presented in Appendix C.3.
edge. For multi-task RL baselines, HarmoDT remains sta- ble when trained on the augmented dataset, whereas the performance of MTDIFF declines markedly. This may be attributed to the increased presence of inferior data collected during the online interaction phase. The generation process of MTDIFF is guided and subsequently biased by these newly introduced low-reward trajectories, whereas the rela- tively small proportion of low-reward online data compared to the original offline data does not affect the established parameter subspaces in HarmoDT.
Table 2. Average success rate of fine-tuning with 100k steps for offline-online methods and training on augmented sub-optimal data for multi-task RI. methods.
Method
Meta World 50 Tasks
RLPD
7.62#0.10
IBRL
17.08
5.3. Efficiency Validation with Limited Online Samples
Cal-QL
24.60.0.00
To validate the learning efficiency of SODP, we conduct experiments in an extreme scenario where the online fine- tuning steps are reduced to 100k per task. We compare SODP with three offline-online baselines, each using only 100k fine-tuning steps, as well as two leading multi-task RL baselines: MTDIFF and HarmoDT. Since MTDIFF and HarmoDT are purely offline algorithms,, we collect online interaction samples during the SODP fine-tuning process. These samples are incorporated as a supplementary dataset alongside the original data, increasing the dataset size from 50M to 50M 100kx50. 50. We then train both MTDIFF and HarmoDT on this augmented dataset to ensure consistent data usage across our method and the baselines. As shown in Table 2, the performance of our method does not degrade sig- nificantly despite a dramatic reduction in fine-tuning steps, demonstrating the strong efficiency of SODP in learning from limited online transitions. In contrast, all three offline- online baselines experience substantial decline due to insufficient online samples for correcting pre-tr knowl-
MTDIFF-P
27.06
HarmoDT-F
57.37
SODP (ours)
59.26.0.18
5.4. Effectiveness of BC Regularization
To demonstrate the effectiveness of our BC regularization, we conduct an ablation study by fine-tuning the same pre- trained model with our BC regularization and other regular- ization methods. We consider the following variants, each employing a different regularization strategy: (i) SODP w/o regularization. This variant fine-tunes the model directly using Eq. (12) without any regularization. (ii) SODP .kl This variant uses KL regularization term to constrain the divergence between the fine-tuned model and the pre-trained model. (iii) SODP pl. This variant incorporates the original diffusion pre-training loss (PL) into the fine-tuning objec- tive function. The details of these variants are provided in Appendix E, and additional ablation studies on differ-
7

---

Task Task-Agnostic Pre-training and Task-Guided Fine-tuning for Versatile Diffusion Planner
SODP regularization
SODP
_p
SODP (ours)
baseline
baskettball
nuttuu--res-tploowr
hammer
Stooes Rais
0.5
0.4
Environment Stape
Ensironment Staps
Ervinnment Strpp
Ervineeeer Stepe
Figure 5. Learning dynamics for different regularization. Performance declines without any regularization. Both KL and PL regularization confine model to sub-optimal regions, while our BC regularization guides the model toward optimal actions.
Pretrained Model
Fine-tuned tuned Model with SODP jdl
Flre cuned Mede with SCDD_M
Fine tuned Medet with SODP
high
hodertaa
glace- 4tde
Figure 6. Visualization trajectorics using generated actions for different regularization. KL and PL regularization result conservative policies, whereas our BC regularization preserves pre-trained knowledge while discovering new actions that lead to high rewards.
ent fine-tuning methods, such directly fine-tuning with high-quality data, can be found in Appendix C.I.
valuable insights from pre- preitrained knowledge by exploiting discovered high-reward actions (e.g. the central region in plate-slide) while discarding low-reward actions.
Figure 5 demonstrates the effectiveness of our regulariza- tion in achieving higher success rate. Directly fine-tuning the model without regularization leads to the worst per- formance, as the success rate declines due to the model degrading its pre-trained capabilities without constraints. However, adding KL and PL regularization is insufficient, as they cause oscillations near the pre-trained model. We hypothesize that the effectiveness of our BC regularization stems from two factors: (i) enables the model to reuse acquired skills, preventing performance decline; (ii) it fa- cilitates effective exploration of optimal regions by using optimal μ as the target policy. To validate this, we visu- alize trajectories generated by our planner's actions using I-SNE (Van der Maaten & Hinton, 2008). As shown in Figure 6, after fine-tuning with KL and PL regularization, the trajectory distribution remains similar to the pre-training distribution, indicating that the model is reusing learned actions and lacks exploration into new regions. The explo- ration in PL is unstructured as it may lead to worse regions (e.g., the lefer-left region in basketball). In contrast, our method demonstrates superior exploration capabilities to discover new, high-reward regions based on acquired knowl- edge (e.g., the lower-left region in basketball and the bottom region in plate-slide) Meanwhile, the model can derive
5.5. Effectiveness of Pre-training
We investigate the impact of pre-training by comparing the performance of SODP with variant that is fine-tuned from scratch using the same loss function defined in Eq. (15), denoted as SODP.scrathh. Without pre-training, the policy is initially random, which leads inaccurate estimation of the target policy μ in the BC regularization term. To mitigate this issue, we initialize the replay buffer with the same initial rollouts used during the fine-tuning of SODP. These rollouts are generated by the pre-trained model, enabling a more accurate approximation of the target policy μ during early fine-tuning.
We present the learning dynamies of the two methods in Figure 7. Directly fine-tuning the diffusion planner from scratch leads to worse performance. Without pre-training, the planner lacks an action prior to guide its behavior, re- sulting in stagnation as it struggles to explore high-reward regions. Moreover, the learning process becomes unstable, as the limited useful knowledge is easily overwritten by numerous ineffective trials. We further conduct ablation studies examine the effect of varying the pre-training data, with results presented in Appendix C.2.
8

---

Task- Task-Agnostic Pre-training and Task-Guided Fine-tuning for Versatile Diffusion Planner
SODP_scruthh
SODP (ours)
Nurton-press-lupdeeee
hammer
handle-pull-iide
Stooes Rais
0.5
0.
0.4
Envrrnment Stape
Ensironmant Staps
Ervinnment Strpp
Ervinemert Step
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. contrast, pre-training allows the planner to extract useful knowledge, guiding fine-tuning by refining the prior distribution towards more effective behaviors.
Benefits of Multi-Task Pre-Trainn Tre-Training. To explore the ben- efits of pre-training on multi-task dataset compared single-task dataset, conduct experiments by fine-tuning the model on tasks excluded from the pre-training dataset. Specifically, we pre-train a model on the MT-10 dataset (SODP mt10) and fine-tune it on three tasks not included in the pre-training data. We compare SODP_mtID with variant pre-trained exclusively on the basketball dataset (SODP bas). As shown in Table 3. pre-training on multi- task data improves generalizability to unseen tasks multi- task data provide a broader range of action distribution priors compared to single-task data.
dation of China (Grant No.62306242), the Young Elite Scientists Sponsorship Program by CAST (Grant No. 2024QNRC001), and the Yangfan Project of the Shanghai (Grant No.23YF11462200)
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There many potential societal consequences of our work, none which we feel must be specifically highlighted here.
References
Table Average success rate on unseen tasks.
Ajay. A., Du Y., Gupta, A., Tenenbaum, J. B., Jaakkola, T. S., and Agrawal, P. Is conditional generative modeling all you need for decision making? In The Eleventh International Conference on Learning Representations, 2023.
Unseen tasks
SODPLmt10
SODP_bas
drawer-open
34.7 20.06
0.0
plate-slide-side
50.3
0.040.0
handle-pull-side
20.3
0.0
Bai, C., Zhang, Y., Qiu, S., Zhang, Q. Xu, K., and Li, X. Online preference alignment for language models via count-based exploration. In The Thirteenth International Conference on Learning Representations, 2025.
6. Conclusion
We propose SODP, novel framework for training versa- tile diffusion planner using sub-optimal data. By effectively combining pre-training and fine-tuning, we capture broad behavioral patterns drawn from large-scale task tran- sitions and then rapidly adapt them to achieve higher per- formance in specific downstream tasks. During fine-tuning, we introduce 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 em- bodied versatile agents that can effectively learn to solve real-world tasks using inferior data.
Ball, P. J., Smith, L., Kostrikov, L., and Levine, S. Ef- ficient online reinforcement learning with offline data. In International Conference on Machine Learning, pp. 1577-15944 PMLR, 2023.
Black, K., Janner, M., Du, Y., Kostrikov, I., and Levine, S. Training diffusion models with reinforcement learning, The Twelfth International Conference on Learning Representations, 2024.
Brehmer, J., Bose, J., De Haan, P., and Cohen, T.S. Edgi: Equivariant diffusion for planning with embodied agents. Advances in Neural Information Processing Systems, 36. 2024.
Acknowledgements
This work is supported by the National Key Re- search and Development Program China (Grant No. 2024YYF0220000),, the National Natural Science Foun-
Chen, J., Byun, J.-S., Elsner, M., and Perrault, A. Rein- forcement learning for fine-tuning text-to-speech diffu- sion models. arXiv preprint arXiv:2405. /4632, 2024.

---

Task-Agnostic Pre-training and Task-Guided Fine-tuning for Versatile Diffusion Planner
Chen, L., Lu, K., Rajeswaran, A., Lee, K., Grover, A., Laskin, M., Abbeel, P., Srinivas, A., and Mordatch, I. De- cision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34:1 5084-15097, 2021a.
He, H., Bai, C., Pan, L.,, Zhang, W. Zhao, B., and Li, X. Learning an actionable discrete diffusion policy via large- scale actionless video pre-training. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
Chen, N., Zhang, Y. Zen, H., Weiss, R. J., Norouzi, M., Dehak, N., and Chan, W. Wavegrad 2: Iterative re- finement for text-to-speech synthesis. arXiv preprint arXiv:2106. 09660, 2021b.
Ho, J., Jain, A., and Abbeel, Denoising diffusion proba- bilistic models. Advances in neural information process- ing systems, 33:6840-6857 2020.
Ho, J., Chan, W., Saharia, C.,, Whang. J., Gao, R., Gritsenko, A., Kingma, D. P., Poole, B., Norouzi, M., Fleet, D. J., et al. Imagen video: High definition video generation with diffusion models. arXiv preprint arXiv:2210.02300, 2022.
Chi, C., Feng, S., Du, Y., Xu, Z.,, Cousineau, E., Burchfiel, B., and Song, S. Diffusion policy: Visuomotor policy learning via action diffusion. In Proceedings of Robotics: Science and Systems (RSS), 2023.
Hu, H., Mirchandani, S., and Sadigh, D. Imitation bootstrapped reinforcement learning. arXiv preprint arXiv:2311.02 1902/98, 2023.
Clark, K., Vicol, P.. Swersky, K., and Fleet, J. Directly fine-tuning diffusion models on differentiable rewards. In The Twelfth International Conference on Learning Representations, 2024.
Hu, S., Fan, Z., Shen, L.., Zhang, Y., Wang, Y., and Tao, D. Harmodt: Harmony multi-task decision transformer for offline reinforcement learning. In International Confer- ence Machine Learning, 2024.
Fan, Y., Watkins, O., Du, Y., Liu, H., Ryu, M., Boutilier, C.,, Abbeel, P., Ghavamzadeh, M., Lee, K., and Lee, K. Rein- forcement learning for fine-tuning tex--t-immae diffuusinn models. Advances in Neural Information Processing Sys- tems, 36, 2024.
Huang, R., Huang. J., Yang, D., Ren, Liu, L., Li, M., Ye, Text-to-
Z., Liu, J., Yin, X., and Zhao, Z. audio generation with prompt-enhanced diffusion models. In International Conference on Machine Learning, pp. 13916-13932. PMLR, 2023.
Florence, P., Lynch, C.. Zeng, A., Ramirez, O. A., Wahid, A., Downs, L.,, Wong, A., Lee, J., Mordatch, I., and Tompson, J. Implicit behavioral cloning, In Conference on Robot Learning, pp. 168. PMLR, 2022.
Jin, C., Krishnamurthy, A., Simchowitz, M., and Yu, T. Reward-free exploration for reinforcement learning. In International Conference on Machine Learning, pp. 4870- 4879. PMLR, 2020.
Guan, J., Zhou, X., Yang, Y., Bao, Y., Peng, J., Ma, J., Liu, Q., Wang, L., and Gu, Q. DecompDiff: Diffusion models with decomposed priors for structure-based drug design. In Proceedings of the 40th International Confer- ence on Machine Learning, volume 202 of Proceedings of Machine Leurning Research, pp. 11827-1 1846. PMLR, 23-29 Jul 2023.
Jing, B., Corso, G., Chang. J., Barzilay, R., and Jaakkola, T. Torsional diffusion for molecular conformer generation. Advances in Neural Information Processing Systems, 35: 24240-24253, 2022.
Kang, B., Ma, X., Du, C., Pang, T., and Yan, Efficient diffusion policies for offline reinforcement learning. Ad- vances in Neural Information Processing Systems, 36, 2024.
Haarnoja, T., Zhou, A., Abbeel, P., and Levine, S. Soft actor-critic: Off-policy maximum entropy deep reinforce- ment learning with stochastic actor. In International conference on machine learning, pp. 1861-1870. PMLR, 2018.
Kim, C., Park, J., Shin, J., Lee, H., Abbeel, P., and Lee, K. Preference transformer: Modeling human preferences using transformers for RL. In The Eleventh International Conference on Learning Representattonss 2023.
Hansen Hansen-Estruch, P., Kostrikov, I., Janner, M., Kuba, G., and Levine, S. Idql: Implicit gelaarning as actor- critic method with diffusion policies. arXiv preprint arXiv:2304. 10573, 2023.
Kingma, D. Adam: method for stochastic optimization. arXiv preprint arXiv: A 2014.
He, H., Bai, C., Xu, K., Yang, 7.,, Zhang, W., Wang, D., Zhao, B., and Li, X. Diffusson model is an effective planner and data synthesizer for multi-task reinforcement learning. Advances in neural information processing systems, 36:64896 64917, 2023.
Kong, Z., Ping, W., Huang. J., Zhao, K., and Catanzaro, B. Diffwave: A versatile diffusion model for audio synthesis. International Conference on Learning Representations, 2021.
10

---

Task-Agnostic Pre-training and Task-Guided Fine-tuning for Versatile Diffusion Planner
Kostrikov, I., Nair, A., and Levine, S. Offline reinforce- ment learning with implicit q-learning. In International Conference on Learning Representations, 2022.
Pearce, T., Rashid, T., Kanervisto, A., Bignell, D., Sun, M., Georgescu, R., Macua, S. V., Tan, S. 7.,, Momennejad, I., Hofmann, K., and Devlin, S. Imitating human behaviour with diffusion models. In The Eleventh International Conference on Learning Representations, 2023.
Lee, K.-H., Nachum, O., Yang, M. S., Lee, I.., Free- man, D., Guadarrama, S., Fischer, I., Xu, W., Jang, E., Michalewski, H., et al. Multi Multi-game decision transformers. Advances in Neural Information Processing Systems, 35: 27921-27936, 2022.
Perez, E., Strub, F., De Vries, H., Dumoulin, V., and Courville, A. Film: Visual reasoning with general con- ditioning layer. Proceedings of the AAAI conference on artificial intelligence, volume 32, 2018.
Lee, S., Cho, M., and Sung, Y. Parameterizing non- parametrie meta-reinforcement learning tasks via subtask decomposition. Advances in Neural Information Process- ing Systems, 36:43356 43383, 2023.
Podell, D., English, Z., Lacey, K., Blattmann, A., Dockhorn, T., Müller, J., Penna, J., and Rombach, R. SDXL: Im- proving latent diffusion models for high-resolution image synthesis. In The Twelfth International Conference on Learning Representations, 2024.
Li, X., Belagali, V., Shang, J., and Ryoo, M. S. Cross- way diffusion: Improving diffusion-basee visuomotor policy via self-supervised learning. In 2024 IEEE Inter- national Conference on Robotics and Automation (ICRA), pp. 16841-16841 IEEE, 2024.
Prabhudesai, M., Goyal, A., Pathak, D., and Fragkiadaki, K. Aligning text-to-image diffusion models with re- ward backpropagation. arXiv preprint arXiv:2310.03739 2023.
Liang, Z., Mu, Y., Ding, M., Ni, F., Tomizuka, M., and Luo, P. Adaptdiffuser: Diffusion models adaptive self-evolving planners. In International Conference Machine Learning, pp. 20725-20745 PMLR, 2023.
Rajeswaran, A., Kumar, V., Gupta, A., Vezzani, G., Schul- man, J., Todorov, E., and Levine, S. Learning Complex Dexterous Manipulation with Deep Reinforcement Leam- ing and Demonstrations. In Proceedings of Robotics: Science and Systems (RSS), 2018.
Mandlekar, A., Xu, D., Wong, J., Nasiriany, S., Wang, C., Kulkarni, R., Fei-Fei, L., Savarese, S., Zhu, Y., and Martin,Martín, R. What matters in learning from offline human demonstrations for robot manipulation. In 5th Annual Conference on Robot Learning, 2021.
Ramesh, A., Pavlov, M., Goh, G., Gray, S., Voss, C., Rad- ford, A., Chen, M., and Sutskever, I. Zero-shot text-to- image generation. In International conference on ma- chine learning, pp. 8821-8831. Pmlr, 2021.
Nakamoto, M., Zhai, S., Singh, A., Sobol Mark, M., Ma, Y., Finn, C., Kumar, A., and Levine, S. Cal-ql: Calibrated offline rl pre-training for efficient online fine-tuning. Ad- vances in Neural Information Processing Systems, 36, 2024.
Reed, S., Zolna, K., Parisotto, E., Colmenarejo, S. G., Novikov, A., Barth-maron, G. Giménez, M., Sulsky, Y., Kay, J., Springenberg, J. T., Eccles, T., Bruce, J., Razavi, A., Edwards, A., Heess, N., Chen, Y., Hadsell, R., Vinyals, O., Bordbar, M., and de Freitas, N. A gener- alist agent. Transactions on Machine Learning Research, 2022. ISSN 2835-8856.
Nichol, A. Q. and Dhariwal, P. Improved denoising diffu- sion probabilistic models. In International conference on machine learning, pp. 8162-8171. PMLR, 2021.
OpenAI, Akkaya, I., Andrychowiez, M., Chociej, M., Litwin, M., McGrew, B., Petron, A., Paino, A., Plap- pert, M., Powell, G., Ribas, R., Schneider, J., Tezak, N., Tworek, J., Welinder, P., Weng, L., Yuan, Q., Zaremba, W., and Zhang, L. Solving rubik's 's cube with a robot hand. arXiv preprint arXiv: 1900::9000000 07/ 13, 2019.
Ren, A. Z., Lidard, J., Ankile, L. L., Simeonov, A., Agrawal, P., Majumdar, A., Burchfiel, B., Dai, H., and Simchowitz, M. Diffusion policy policy optimization. In The Thir- teenth International Conference on Learning Representa- lions, 2025.
Ostrovski, G., Bellemare, M. G., Oord, A., and Munos, R. Count-based exploration with neural density models. In International conference on machine learning, pp. 2721- 2730. PMLR, 2017.
Reuss, M., Li, M., Jia, X., and Lioutikov, R. Goal con- ditioned imitation learning using score-based diffusion policies. In Robotics: Science and Systems, 2023.
Ouyang, I.., Wu, J., Jiang, X., Almeida, D., Wainwright, C.,, Mishkin, P., Zhang, C., Agarwal, S., Slama, K., Ray, A., et al. Training language models to follow instructions with human feedback. Advances in neural information processing systems, 35:27730-27744 2022.
Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF con- ference on computer vision and pattern recognition, pp. 10684-10695, 2022.

---

Task-Agnostic Pre-training and Task-Guided Fine-tuning for Versatile Diffusion Planner
Schneuing, A., Harris, C., Du, Y., Didi, K., Jamasb, A., Igashov, I., Du, W., Gomes, C., Blundell, T. L.., Lio, P., et al. Structure-based drug design with equivariant diffusion models. Nature Computational Science, 4(12): 899-909, 2024.
Wang, R., Du, S. S., Yang. L.., and Salakhutdinov, R. R. On reward-free reinforcement learning with linear function approximation. Advances in neural information process- 2020.
ing systems,
Wang, Z., Hunt, J. J., and Zhou, M. Diffusion policies as an expressive policy class for offline reinforcement learning. In The Eleventh International Conference Learning Representations, 2023.
Schuhmann, C., Beaumont, R., Vencu, R., Gordon, C., Wightman, R., Cherti, M., Coombes, T., Katta, A., Mullis, C., Wortsman, M., et al. Laion-5b; An open large-scale dataset for training next generation image-text models. Advances in Neural Information Processing Systems, 35: 25278 25294, 2022.
Xu, J., Liu, X., Wu, Y., Tong, Y., Li, Q., Ding, M., Tang, J., and Dong, Y. Imagereward: Learning and evaluating hu- man preferences for text-to-image generation. Advances in Neural Information Processing Systems, 36, 2024.
Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O. Proximal policy optimization algorithms. arXiv preprint arXiv: 170.007777. 2017.
Xu, M., Shen, Y., Zhang, S., Lu, Y., Zhao, D., Tenenbaum. J., and Gan, C. Prompting decision transformer for few- shot policy generalization. In international conference on machine learning, pp.
Sco, Y., Hafner, D., Liu, H., Liu, F., James, S., Lee, K., and Abbeel, P. Masked world models for visual control. In Conference on Robot Learning, pp. 1332-1344. PMLR, 2023.
PMLR, 2022.
Xu, M., Xu, Z., Chi, C., Veloso, M., and Song. S. Xskill: Cross embodiment skill discovery. In Conference on Robot Learning, pp. 3536-3555 PMLR, 2023.
Shen, X., Du. C., Pang, T., Lin, M., Wong, Y., and Kankan- halli, M. Finetuning text-to-image diffusion models for fairness. The Twelfth International Conference on Learning Representations, 2024.
Yu, T., Quillen, D., He, Z., Julian, R., Hausman, K., Finn, C., and Levine, S. Meetawwrld:: A benehmark and evalua- tion for multi-task and meta reinforcement learning. In Conference on Robot Learning (CoRL), 2019.
Silver, D., Huang, A., Maddison, C. J., Guez, A., Sifre, L.., Van Den Driessche, G., Schrittwieser, J., Antonoglou, I., Panneershelvam, V., Lanetot, M., al. Mastering the game of go with deep neural networks and tree search. nature, 529(7587):484 2016.
Yu, T., Kumar, S., Gupta, A., Levine, S., Hausman, K., and Finn, C. Gradient surgery for multi-task learning. Advances in Neural Information Processing Systems, 33: 5824 5836, 2020.
Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N., and Ganguli, S. Deep unsupervised learning using nonequi- librium thermodynamics. In International conference on machine learning, pp. 2256-2265. PMLR, 2015.
Yu, X., Bai, C., Hc, H., Wang, C., and Li, X. Regular- ized conditional diffusion model for multi-task preference alignment. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
Song, J., Meng, C., and Ermon, S. Denoising diffusion implicit models. In International Conference on Learning Representations, 2021.
Yuan, X., Shang, Z., Wang, Z., Wang, C., Shan, Z., Zhu, M., Bai, C., Li, X., Wan, W., and Harada, K. Preference aligned diffusion planner for quadrupedal locomotion control, arXiv preprint 24X 2024.
Sun, I.., Zhang, H., Xu, W., and Tomizuka, M. Paco:
Parameter-composine compostoolo multi-task reinforcement learn- ing. Advances in Neural Information Processing Systems, 05-21495-21507 2022.
Ze, Y., Zhang, G., Zhang, K., Hu, C., Wang, M., and Xu, 3d diffusion policy: Generalizable visuomotor policy learning via simple 3d representations. In Proceedings of Robotics: Science and Systems (RSS), 2024.
Sutton, R. S., McAllester, D., Singh, S., and Mansour, Y. Policy gradient methods for reinforcement learning with function approximation. Advances in neural information processing systems, 12, 1999.
Zheng, Q., Zhang, A., and Grover, A. Online decision trans- former. In international conference on machine learning, pp. 27042-27059. PMLR, 2022.
Van der Maaten, L. and Hinton, G. Visualizing data using (-sne. Journal of machine learning research, 9(11), 2008
Zhou, Q.-Y,, Park, J., and Koltun, V. Open3D: A modern library for 3D data processing. arXiv: 1801009477 2018.
Wang, C., Luo. X., Ross, K., and Li, D. Vr13: A data- driven framework for visual deep reinforcement learning. Advances in Neural Information Processing Systems, 35: 32974-32988, 2022.
12

---

Task Task-Agnostic Pre-training and Task-Guided Fine-tuning for Versatile Diffusion Planner
A. Derivations
A.1. Derivation policy gradient Equation (11)
Assume pa(er (a 0:K |s)rr and Vepa(a, t (s)rr((a)) are continuous (Fan et al., 2024), we have:
Vod" (0) = Vo E, no (as) [a]]]
Σ 0 / (aa)) po(a)lss)da!
E. Vo / | (aa)) (frolatf (a Is) da?
Σ, / T (a.) Vo logopaa (a (s,) po(a, " (s,) da, o
(16)
Σ, / (-((a)) (T())) Ve log
K
(x (s))
II
po k- "1af 81)) pp ((.: da 0.K
k=1
Σ, Ppolas (1 18()
(T(a)) Σ
Va log la
A.2. Derivation loss function in Equation (12)
By using importance sampling approach, we can rewrite Eq. (16) as follows:
(a k-1 aa, 8t)
E, E Bren (a an)
T (Ta)))
Vo log (at st)
(17)
Peu (a la,; si)
Then, we can get a new objective function corresponding to Eq. (17) as:
K at 81) k= Peas (a, aaa 8t)
pa (a, k-
JT (0) max E, (a (a): TT (Ta)))
(18)
Pa ..)) Peau (a. - ...)
denote the probability ratio. Based on PPO (Schulman et al., 2017), we clip Pa: and use the minimum between the clipped and unclipped ratios to derive lower bound of the original objective (18), which serves as our final objective function:
Let Pk (0, 0.1a)
(T max E, Epsny (a (ee) rJ (Taa)) Σ min Pk (o,)) clip (px(0, Oola) 1 1.,.1-))
(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.
13

---

Task-Agnostic Pre-training and Task-Guided Fine-tuning for Versatile Diffusion Planner
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:
Pa(a
2g ~pp |- po po (a))) - Bag-na ~u
log q(a)i ((ak:Ki (()))
po po(a
Bagarp Ppu
E Byla) (09)
log p(a Σ log
aaatia
(20)
K
= =Bag
laga DKila(a ( -a l k-1 la=)))
[log po (o [ (gia K la (I((a)) Ba(a)la))
Following previous work (Ho et al., 2020), the optimization of the bound can be simplified as:
1 (1 - a2)
arg min
g (k) (1 āz)ax ||c(a 1:) co co(a |)]
(21)
where:
(1 oik (1x)(11 Ōx-
o (b)
(22)
1 ak
Here, co ca k) is a noise model that learns to predict the source noise c(a k) which determines af: from a.
B. The Details of SODP
B.1. Theoretical Motivation
The theoretical motivation of SODP is grounded in the provable efficient reward-free exploration in RL (Wang et al., 2020; Jin et al., 2020). This approach employs two-stage learning process for developing an RL policy: a reward-free exploration phase and a reward-basee based optimization phase. The latter stage requires only a limited number of reward labels.
Specifically, in the reward-free stage, the agent optimizes a pure exploration reward, namely, the upper-confidence bound (UCB) term of the value function. Under the linear MDP assumption, we denote the feature of the state-action pair as d0ss, a), and the value function is a linear mapping of o(s, s, a). expressed Q(s, a) abss,a)Te where 0 is weight vector. The covariance matrix for state-action pairs is defined as A = E. Ø(s;, a,) a,) T. Consequently, the UCB-based bonus for reward-free exploration is defined as b, = [Ø(st. ) A TA o(s a,)) Leveraging this bonus, the agent can collect diverse state-action pairs with broad coverage during the exploration stage, thereby providing sufficient amount of information for the subsequent planning phase. In the second stage, the agent performs least-square value iteration (LSVI), a well-studied RL algorithm, by interacting with the environment. The goal of this stage to conduct value and policy updates for a specific task with specific reward function. This two-stage learning paradigm, which combines reward-free exploration and reward-based optimization, can achieve polynomial sample complexity according to theoretical results (Wang et al., 2020). We believe this approach provides solid theoretical foundation for our method's empirical success.
Our algorithm (i.e.. SODP) essentially follows this theoretical motivation but incorporates more practical implementations. Specifically, previous theoretical works have shown that b, defined in linear MDPs is approximately proportional to the reciprocal pseudo-count of the corresponding state-action pair in the dataset. As a result, the reward-free exploration stage in (Wang et al., 2020) similar to count-based exploration, which aims to collect diverse datasets in the first stage (Ostrovski al., 2017). In SODP, we adopt an alternative exploration method (entropy- based exploration) via non-expert SAC policy to provide sufficient exploration in the data collection stage. We also pre-train diffusion planner on this exploratory dataset rather than directly performing reward-based policy optimization. This approach seeks to model the distribution of diverse trajectories with broad coverage and provides a well-initialized starting point for the second stage. In Section 5.3, we empirically demonstrate that SODP requires only a very limited number of online samples to achieve strong performance in the second stage. This finding verifies the theoretical results of (Wang et al., 2020; Jin et al., 2020), which suggest that such a two-stage learning paradigm is sample-efficient. We believe that our practical implementation of this theoretical
14

---

Task-Agnostic Pre-training and Task-Guided Fine-tuning for Versatile Diffusion Planner
framework not only enhances the applicability of our method but also maintains the desirable properties of sample efficiency and robustness.
B.2. Diffusion Policy
We use diffusion policy (Chi al., 2023) to generate future actions. For any given time step L, the model the most recent T, steps of states as input to generate the next Tn action steps. Then, the first T. steps of these generated actions executed in the environment without re-planning, In experiments, use Tр 12, T. 2, To for Meta-World and T = 4, To 2, Ta 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., 2018) to condition on historical states. The implementation based on the code from https: /github. com/CleanDi ffuserTeam/C htmm/CCeaaDifffii parameters. For Adroit, which removes some components in the U-net. 3D-Diffusion-Pon
and we use their default hyper- simplified backbone provided by Simple DP3 (https: //g htt///gittttt /ttggtttbbb/./////g com Yonn e
B.3. Implementation Details
The pseudo-code of SODP is given in Alg. We describe details of pre-training and fine-tuning follows:
For pretraining, cosine schedule for 3g (Nichol & Dhariwal, 2021) and set diffusion steps K 100. We pre-train the model for 5e steps in Meta-Wrold and 3e steps in Adroit.
For fine-tuning, we use DDIM (Song et al., 2021) with 10 sampling steps and rj = We fine-tune each task for le steps in Meta-Worih and 3 steps Adroit. Following DPOK (Fan al., 2024), we perform Psuep {10, 30} gradient steps per episode. We discount factor y 1 for all tasks.
We set Nini E ( (10, 20} for approximating target distribution and À = 1.0 as the BC weight coefficient.
Batch size is set to 256 for both pre-training and fine-tuning.
We use Adam optimizer (Kingma, 2014) with default parameters for both pre-training and fine-tuning. Learning rate is set to 1e-4 for pretraining and le for fine-tuning with exponential decay.
Algorithm 1 SODP: Two-stage framework for learning from sub-optimal data
Input:
diffsuion planner 6, downstream tasks T, multi multi-task sub-optimal data D i=1 Dy, target buffer Burget- replay buffer B, episode length L, pre-train Npr and fine-tune Npr steps
// pre-training model with sub-opt imal data
for 1, ...,, Net do Sample (s, a) ~ D, diffusion time step k Uni.orm((1...., K}), noise NoO,I) Update 0 using the loss function (7)
// fine-tuni ng model for downst ream tasks for Ti E [T...., Tx| do Initialization: 0 + -pr; B, Bunget + Rollout Num proficient trajectories using 0 for r 1,..., Npr do
while not end of the episode do Obtain samples a.:K ~ po 0:K (8)) Execute the first T2 steps and get reward r(a?) B +- BU(s, a 0:K r(a)))
3t + -E4T,,1 t + +tTo
!/ approximate target policy μ
if proficient then
Buaget + Bunget U (a. | E {0, (O, L}}
Compute CT: using batches from B according Eq. (12) Compute L using batches from Busget according Eq. (14) Update 0 using the loss function (15)
15

---

Task-Agnostic Pre-training and Task-Guided Fine-tuning for Versatile Diffusion Planner
C. Extended Results
In this section, we provide our full experimental results:
1. Ablation studies evaluating various fine-tuning strategies.
Analysis of the impact of pre-training dataset quality,
3. Additional experiments for offline-online baselines
4. Evaluation image-basel environments.
Computational efficiency analysis
6. Learning curves of SODP compared to multi-task RL baselines.
C.1. Ablation studies examining other fine-tuning approaches
To demonstrate the effectiveness of our online fine-tuning approach, we compare it with alternative fine-tuning methods: (i SODP _off, which involves fine-tuning using high-quality offline data, and (ii) SODP _off scrathh, which performs direct training with high-quality data without pre-training. Specifically, we fine-tuned the pre-trained models for 100k steps across five tasks, collecting 200 successful episodes (equivalent to 100k steps) for each task. These datasets were then used to independently train five models in an offline setting, utilizing the same loss function as in Eq. (15) (SODP _off). Additionally, to investigate the impact of pre-training on offline fine-tuning, we trained the model directly without pre-training (SODP off scratch).
The experimental results, presented in Table 4, report the success rates averaged over three seeds. Without pre-training, the model lacks the necessary action priors to efficiently identify high-rewart action distributions. Furthermore, directly fine-tuning with high-quality offline data proves insufficient, as static reward labels may fail to provide adequate guidance in dynamic environments, hindering the model's ability to facilitate efficient exploration.
Table Average success rate for different fine-tuning approaches.
Tasks
SODP off
SODP off scratch
SODP
putton-prees-topdoww
58.6740.00
40.67
60.67
hammer
71.33
13.33
73.33+0.03
handle-pull-side side
60.67+00.00
42.67
81.67+0.07
oeg-insert-side
25.3 25.3340.03
0.0.0.0
32.67 +0.06
handle-pull
66.67
31.33440066
75.33
Average success rate
56.53+0.18
25.640.18
64.73
To highlight the importance of modeling the diffusion process as a MDP for reward fine-tuning, we consider an alternative approach that directly applies BC during fine-tuning, using only Eq. (14) as the loss function. As shown in Table 5, directly using BC results in poorer performance, as BC lacks reward labels to effectively guide exploration. While BC during the fine-tuning phase enables access dynamic actions, it is limited to 'imitation' rather than 'evolution," as the model is unable to differentiate between good and bad actions.
Table 5. Average success rate for directly BC during fine-tuning mne-tuning.
Task
Directly BC
SODP
button- puttonnpreesstopdoww
51.3+0.05
60. 60.7
basketball
21.3+0.03
41.2 41.2+0.16
stick-pull
26.7+0.08
50.540.04
16

---

Task-Agnostic Pre-training and Task-Guided Fine-tuning for Versatile Diffusion Planner
C.2. Pre-training using near-optimal data
To evaluate the impact of pre-training data quality on fine-tuning performance, we modified the near-optimal dataset provided by (He et al., 2023) by retaining only the last 50% of the data. This modification ensured that the total number of transitions remained the same as the sub-optimal data used in the main paper, while significantly increasing the proportion of expert trajectories. We refer to this modified dataset as near-optimal data and pre-trained a model on the Meta- Setad 10 tasks. Subsequently, we followed the same fine-tuning procedure outlined in the main paper to fine-tune the model on each task. The experimental results are presented in Table 6. Incorporating more optimal data during the pre- pre-training stage leads to better performance, as the model gains more priors about the optimal action distributions.
Table 6. Average success rate achieved after fine-tuning models pre-trained different datasets.
Tasks
Sub-optimal dataset
Near-optimal dataset
basketball
52.67
80.67
button-press
88.00+0.00 02.0
89.33_0.00
dial-turn
80.67
74.00
drawer-close
100.00+0.00
100. 100.00+0.00
peg-insert-side
62.67++0.0 02
84.67 -0.02
pick-place
36.67+0.00 #0.03
39.33=0.03
push
33.33+0.08
30.67=0.03
reach
68. 10.00
00.03_0.01
sweep-into
60.67+00.0
70.3
window-open
69.33 10.04
100.0 0.000
Average success rate
65.27
80.93_0.16
C.3. Additional experiments for offline-online baselines
We consider more baselines that are diffusion-basee and support offline pre-training to online RL fine-tuning procedures. Specifically, we extend two multi-task cefline-to-onligg training on Meta-World 10 tasks, following (Ren et al., 2025). Specifically, we first pretrain a diffusion- based actor using behavior cloning on an offline dataset and then fine-tune both the actor and critic in an online environment. The results are shown in Table 7, the diffusion actor can be misled by inaccurate value network estimations, causing rapid forgetting of actions learned from offline pretraining. Furthermore, since goals in Meta- Soold randomized, traditional methods struggle to learn how complete them effectively.
methods, DQL (Wang et al., 2023) and IDQL (Hansen-Estruck) et al., 2023), to
We also consider transformer-basee following (Zheng et al., 2022) by pre-training it on the same sub-optimal data and fine-tuning each task online for IM steps on the Meta- 10 tasks. We consider two types of DT: (1) the vanilla DT, which uses a deterministic policy, and (2) the variant proposed in (Zheng et al., 2022), which employs a stochastic policy. The results are presented in Table 7. The stochastic policy improves the success rate by enabling greater exploration in the online environment and discovering more valuable action patterns. However, it still underperforms compared to our method, the diffusion model can generate more diverse actions.
RL baselines. Specifically, we extend the multi-task Decision Transformer (DT)
Table 7. Average success rate of MT-10 tasks for diffusion castus and transformer-based offline-online basclines.
Methods
Success Rate
MT-DQLL MT-IDQL
10.80
Diffusion-basee
13.57 12
AT-OODTdetter
40.07
Transformer-bant based
MT-ODT_stoc
50.47+0.21
SODP (ours)
65.27+0.21
17

---

Task-Agnoste Pre-training and Task-Guided Fine-tuning for Versatile Diffusion Planner
C.4. Experiments in image-based Meta-World and Adroit
The Adroit benchmark includes three dexterous manipulation tasks, requiring a 24-degree-of-freede dexterous hand to solve complex challenges such as in-hand manipulation and tool use. The goals in this environment are also randomized. For Adroit, use images the observation to assess whether our method can scale to high-dimensional input. To collect offline dataset, we train VRL3 (Wang et al., 2022) agent for each task and use the initial 30% experiences (90K transitions) from the converged replay buffer.
The action space for different tasks in Adroit is different and is incompatible with MTDIFF and HarmoDT. Therefore, we compare SODP with following baselines designed for complex environments: (1) BCRNN (Mandlekar et al., 2021). 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 al., 2022). Extended BC with energy-based models (EBM) to train implicit behavioral cloning policies. (3) Diffussio Policy (Chi et al., 2023). A diffusion-based approach that predicts future action sequences based on historical states. (4) DP3 (Ze et al., 2024). A visual imitation learning algorithm that incorporates 3D visual representations into diffusion policies, using 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., 2024).
We scale our method image-based based observations using the Adroit benchmark by employing a point-cloud encoder from DP3 (Ze al., 2024) to process the 3ID scene represented by point clouds. Specifically, we capture depth images directly from the environment and convert them into point clouds using Open3D (Zhou al., 2018). 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 except the imput states are visual features. Following DP3 (Ze et al., 2024). 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 8, our method achieves an 8.2% improvement across all tasks. Since hammer is more challenging than door, our method may need more insightful priors from pre-training to achieve better performance.
Table 8. 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 Door
Algorithm| \
Task
Hammer
Pen
Average
BCRNN
Ozo
Ozo
92s
3.0
BCRNN+3D
8114
0.o
8
5.3
IBC
0.o
0.o
9=2
3.0
IBC+3D
0go
Oao
10a1
3.3
Diffusion Policy
48
50gs
25
31.7
Simple DP3
100
58g4
4626
68.0
DP3
100
62
43
68.3
SODP (ours)
1:
96a1
59
73.9
We further conduct ablation studies on three Adroit environments, fine-tuning the same pre-trained model with and without BC regularization separately. As shown in Table 9, the performance improvement of the model fine-tuned without BC is marginal, whereas the model fine-tuned with BC achieves significantly better results. This indicates that BC also facilitates learning in high-dimensional single-task settings. This improvement can be attributed to the update stride constraint imposed by BC, which ensures that updates not too aggressive, thereby preventing the model from becoming trapped in suboptimal region.
Table Ablation study on the image-based Adroit benchmark. We report average success rate for the pre-trained model, as well as the results of fine-tuning the same model with and without BC regularization.
Adroit Door
Algorithm \ Task
Hammer
Pen
Average
Pre-trained model
20
55-s
28_3
34.3
Fine-tuned model w/o BC
3843
63
30
43.7
Fine-tuned model w BC (ours)
67
96
59
73.9
18

---

Task-Agnostic Pre-training and Task-Guided Fine-tuning for Versatile Diffusion Planner
We also conduct experiments on image-based Meta-World 10 tasks. Since no existing image-based sub-optimal dataset for Meta-World is available, collect data for the 10 tasks by training separate SAC agents for each task, done in (He et al., 2023), and rendering the environments to obtain image data. We then follow the same procedure as in Adroit to convert the images into point clouds and use the DP3 encoder to extract visual features. For comparison, we consider the following baselines: DP3 and MTDIFF 3D, an extended variant of MTDIFF that employs the same 3D visual encoder used in SODP. The experimental results are presented in Table 10. Our method also outperforms other baselines and achieves higher success rate, demonstrating its generalizability complex input modalities.
Table 10. Average success rate of image- based MT-10 tasks.
Methods
Success rate
DP3
32.6 32.6 +0.23 38. 0 0..882
MTDIFF_3D SODP
47.5+0.18
C.5. Computational efficiency for training and evaluating
We conduct computational efficiency analysis comparing three offline pre-training followed by online fine-tuning baselines and our method. For training time, report the average fine-tuning duration per task over million environment steps, For evaluation time, we report the average inference time per task over 50 episodes. All results are obtained using single NVIDIA RTX 4090 GPU.
Table II. Computational efficiency for olline-online baselines and our method.
Methods
Fine-tuning time
Inference time
RLPD
4.5 5, 5/
37s
IBRL
6 -75h
39s
Cal-QL
8.5 10h
40s
SODP (ours)
11 - 13h
47s
C.6. Learning curves of multi-task baselines
We sample 10 tasks and present the learning curves of SODP alongside two leading baselines, MTDIFF and HarmoDT, across five random seeds. The results shown in Figure 8, where the X-axis represents gradient steps. The first 5 steps (indicated by the black vertical dashed line) correspond to the pre-training phase of SODP, while the subsequent steps represent the fine-tuning phase (marked by the orange vertical dashed line).
MIDIFF
HarmoDT
SODP (ours)
basketball
button-pres--updovee
button button-prest- press- wecpaapreess--eeieiee hopdowa-wal
0.8
0.6
1,0
0.6
Succoss
0.4
04
0.4
0.2
0.0
0.5
0.63
1.5
0.0
0.5
0.64
1.5
0.0
0.6
3.54
1 1,5
0.0
0.5
0.55
1.5
0.0
055
0.6G
1.0
Gradient Steps (M)
Gradient Steps (M)
Gradient Steps (M)
Gradient Stape (M)
Gradient Steps (M)
handle-pull
1.0
hundle-pull-side
peg-insert-side
stick-pull
sweep
0.6
0.8
0.4
0.4
0.2
02-
0.2
0.0
3.56 Gradient Steps (M)
0.5
0.0
0.0
1 1.5
0.0
0.5
0.53
0.0
0.50 Gradicat Steps
Gradient Steps
Gradient Steps (M)
Gradient Stepe (M)
Figure& 8. Learning dynamics. 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 on 50M transitions, followed by fine-tuning with 1M transitions per task. SODP rapidly converges to high success rates, whereas MTDIFF and HarmoDT struggle with some challenging tasks.
SODP rapidly converges to high success rates, surpassing the other two baselines. The pre-training stage equips the planner
19

---

Task-Agnostic Pre-training and Task-Guided Fine-tuning for Versatile Diffusion Planner
with comprehensive action distribution priors and allows it rapidly transfer and enhance these capabilities across variety of downstream tasks. As result, the pre-training stage significantly accelerates convergence, leading to 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.
D. The Details of Baselines
We describe the details of baselines used for comparison in experiments. For offline-online RL, consider following baselines and use the implementation from https: //t.thub.bb/////// com/ prtommppinneetooooooo dppo:
Cal-QL. The actor network is implemented four-layer MLP and is pre-trained on offline dataset using the SAC approach. During online fine-tuning, it optimizes conservative value functions that are constrained to exceed the value function of the reference policy.
IBRL. The actor network is implemented as three-layer MLP and is pre-trained using imitation learning. During fine-tuning, it leverages an ensemble of N critics, randomly selecting two of them for each update step.
RLPD. The actor network implemented as three-layer MLP and is pre-trained on an offline dataset using the SAC approach. During online fine-tuning, it leverages ensemble of N critics, randomly selecting two of them for each update step.
For multi-task RL, consider following baselines:
MTSAC. The one-hot encoded task IDD incorporated into the original SAC as additional input.
MTBC. The actor network is modeled using 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 latent variable 2. The input to the actor network is then formed by concatenating the original state with this latent variable =
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 3-layer MLP using Mish activation.
MTDQL. Similar MTIQL, multi-head critic network is utilized to predict the Q-value for each task, and the original diffusion actor is extended with an additional task ID input.
MTDT. The task ID is embedded into a latent variable a 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. 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 one-hot task ID. We borrow the official codes from https: /ttt::/ggttttt parameters.
weterrhrheeeMDriir.. MTDiff and use their default hyper-
HarmoDT. Incorporate trainable task-specific masks to address gradient conflict by identifying an optimal harmony subspace 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:// //gi thub. www/chaaaaaaoooccccc.o eahsc/ rmoDT and use their default hyper-parameters.
20

---

Task-Agnostic Pre-training and Task-Guided Fine-tuning for Versatile Diffusion Planner
E. Variants of SODP
In Eq. (15), we introduce a BC regularization term to preserve the pre-trained knowledge and demonstrate effectiveness compared to existing regularization approaches. We detailed other regularization methods below.
SODP w/o regularization. This variant is similar to DDPO (Black et al., 2024) and DPPO (Ren et al., 2025), which fine-tunes the model directly using Eq. (12) without any regularization.
SODP kl. This variant is similar to DPOK (Fan al., 2024), with the addition of KL regularization term constrain the divergence between the fine-tuned model and the pre-trained model. Specifically, the regularization term Lxl. used in DPOK is expressed as:
K
(0)
KL (po (2 (x) (xX-1 (1
(23)
SODP pl. This variant is similar to DLPO (Chen et al., 2024), incorporating the original diffusion pre-training loss (PL. into the fine-tuning objective to prevent the model from deviation. Specifically, the regularization term CpL. used in DLPO is expressed as:
(0) = [[x-(]_)) [Ie(zx: k) c0(2 cx(xxx)))] A)]?
(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 μ Prre-train and DLPO selects μr The rationale behind their selection is based the assumption that 9 Apre-train f" This is reasonable for methods like DPOK and DLPO, which utilize pre-trained models such as Stable Diffusion (Rombach et al., 2022) and WaveGrad2 (Chen et al., 2021b). 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.
However, this assumption does not apply to our pre-trained planner, as the 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 sub-optimal region. While PL regularization allows some slight exploration, it is uncontrolled. Consequently, as shown in Section 5.4, we observe that the KI. regularization almost remains unchanged and the PL regularization slightly increases the performance in basketball but decreases in other tasks.
F. Comparison to DPPO
We summarize some similarities and differences between our work and the concurrent work DPPO (Ren al., 2025) as follows:
Both DPPO and our approach formulate the diffusion policy denoising process an MDP and use policy gradients fine-tune the model for higher environment rewards.
DPPO demonstrated that reward-based based RL fine-tuning promotes effective exploration, which is consistent with observations.
While DPPO requires task-specific expert demonstrations for pre-training, our method pre-trains 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 current regularization methods and propose a novel BC regularization term. By employing our regularizer, the pre-traing pre-trained model achieves higher success rates after fine-tuning.
Unlike DPPO, we don't employ advantage estimator.
21

---

Task-Agnostic Pre-training and Task-Guided Fine-tuning for Versatile Diffusion Planner
G. Detailed Performance of Offline-Online Methods
We report the success rates of both pre-training and fine-tuning for three offline-online baselines, as well as our method. Specifically, we evaluate each pre-trained model over 50 episodes across three random seeds and present the results in Table 12. The pre-training performance of SODP surpasses that of the other baselines, indicating the advantage of diffusion models in capturing multi-modal action distributions. Furthermore, when fine-tuned with reinforcement learning- based rewards, SODP shows greater performance improvement compared to traditional offline-online methods, highlighting the effectiveness of our fine-tuning strategy in generalizing more efficiently to downstream tasks.
Table 12. Performance of the model after pre-training and fine-tuning for offline-online baselines and our method.
Methods
After pre-training
After fine-tuning
RLPD IBRL
6.40
10.164..11
14.68
25.2910.22
Cal-QL
15.8240.06
35.0940.12
SODP (ours)
31.76+0.11
60.56_0.14
H. 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 Table 13.
22

---

Task-Agnostic Pre-training and Task-Guided Fine-tuning for Versatile Diffusion Planner
Table 13. Evaluated return of SODP pre-trained model and fine-tuned model for cach task in M750-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 + 100.7
2347.1 = 580.8
bin-picking-v2
96.8 + 23.9
602.7 72.8
button-pree press-topdowD-y
1405 20.3
1679 25.9
button-press-v2
1397 15.6
2
button-press-wall--ee
1375 10.58
2524.7 18.1
noffee-button-ve
293.2 12.1
451.5 14.4
coffee-pul pullee22
39.5 6.2
117.9 23.3 23.3
coffee-pusie puitee-
33.8 6.1
2
dial-turn-v2
1217.7 = 239.3
1557.3 226.7
disassemble-v2
237 : 117.2
5 164.6
door-close-v2
3347.7 = 124.9
4116.3 = 118.6
door-lock-v2
1042.3 94.9
2491 79.5
door-open-v2
2036.3 + 79.3
2460.3 57.7
door-unlock-v2
1335 + 46.9
2257.7 = 323.0
hand-insert-v2
85.9 = 56.7
449.5 54.9
drawer-close-v2
2468.3 + 167.2
3953.7 214.3
drawer-open-v2
1656 45.6
2489.7 = 188.4
faucet-open-v2
2728.7 = 424.1
4094.7 = 290.3
faucet-close-v2
2156.7 + 113.6
3772 70.1
handle-press-side--22
1919.7 + 449.5
3478.3 98.0
handle-press-v2
2216.3 + 182.0
3415.7 = 221.6
handle-pulle
sull-s
1351.7 = 119.0
2665.7 = 243.9
handle-pull-v2
1510.7 + 111.6
2734 + 64.3
lever-pull-v2
650.7 + 32.5
1068.8 + 110.4
peg-insert-side-v2
300.3 + 122.2
1969.7 + 237.6
pick-place-wall-v22
596.7 = 10.6
1175.7 150.1
pick-out-of-hole-22
38.5 6.3
106.7 7.9
reach-v2
2664.3 77.5
3083.7 + 149.7
push-back-v2
55.8 + 26.7
350.9 = 30.3
push-v2
46.8 + 35.9
148.9 = 43.9
pick-place-v2
3.9 + 0.2
5.9 + 1.8
plate-slide-v2
1268.7 75.8
2862 234.5
plate-slide-side-22
826.4 = 54.3
1929.7 = 104.5
plate-slide-back-22
795.8 62.9
1587.7 = 125.0
plate-shide-buckk-sideew22
626.7 36.9
1541.3 + 78.5
soccer-v2
863.8 159.5
1234.2 + 225.8
push-wall-v2
175.2 = 35.0
471.7 + 73.9
shelf-place-v2
260.4 111.5
785.1 81.5
sweep-into-v2
621.0 132.9
1282.3 + 135.9
sweep-v2
442.3 83.0
= 1081.7=58.. 58.0
window-open-v2
1342.3 60.1
2474.3 266.4
window-close-v2
1087.7 78.9
1816.7 166.3
assembly-v2
282.5 + 4.3
446.1 # 26.9
button-
prtton-ppess-topooonnwwll-->>
1374 16.1
1702.7 71.5
hammer-v2
1678.3 52.5
1907.7 25.4
peg-unplug-side-v2
34.2 2.9
52.9
reach-wall-v2
3373.7 41.7
3839.7 - 69.3
stick-push-v2
412.9 = 95.8
833.5 99.1
stick-pull-v2
1977 = 155.9
3116.3 58.1
box-close-v2
692.3 = 22.8
1300.1 53.2
23


Reviews:
[{"review_id": "p0zNkxIL7P", "reviewer": "Reviewer_aKBX", "summary": "The paper presents SODP (Sub-Optimal Diffusion Planner), a two-stage framework that leverages task-agnostic pre-training on sub-optimal trajectory data followed by reward-guided fine-tuning for multi-task reinforcement learning (RL). The approach aims to train a generalizable diffusion-based planner that can extract broad action distributions from large-scale, low-quality offline data and adapt efficiently to new tasks using policy gradient fine-tuning with behavior cloning (BC) regularization. The pre-training phase learns a guidance-free generative model that predicts future actions given past states, without requiring task-specific rewards or expert demonstrations. In the fine-tuning phase, the planner is refined using policy gradients to maximize task-specific returns while maintaining pre-trained knowledge via BC regularization. Empirical evaluations on Meta-World and Adroit show that SODP outperforms state-of-the-art multi-task RL and diffusion-based methods, demonstrating strong sample efficiency, generalization to unseen tasks, and robustness with limited fine-tuning steps. The paper claims that pre-training on diverse sub-optimal data enables the planner to acquire broad action priors, which accelerate task-specific adaptation and achieve higher rewards compared to traditional RL baselines.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer aKBX,\n\nWe sincerely appreciate your precious time and constructive comments. In the following, we would like to answer your concerns separately.\n\n**W1**: ...there are areas where it could improve...: \n\n**W1.1**: ... analyzing ... performance degradation over long adaptation periods\n\n**R1.1**: As shown in Tables 1 and 2, increasing the fine-tuning steps from 100k to 1M does not degrade performance. After 1M steps, the learning curve converges and remains stable.\n\n**W1.2**: BC regularization prevents catastrophic forgetting is not explicitly tested ...\n\n**R1.2**: We fine-tune two models on the same *button-press-topdown* task using the same pre-trained model, with and without BC regularization, respectively. We then evaluate them on other tasks from the pre-training dataset. As shown below, the performance of the model fine-tuned with BC regularization does not decline significantly, whereas the model without BC exhibits a noticeable performance drop, indicating that it forgets the acquired skills.\n||Pre-trained model|w/o BC|w/ BC|\n|-|-|-|-|\n|*button-press-topdown-wall*| 52.67 ± 0.01|21.33 ± 0.05|65.33 ± 0.05\n|*button-press-wall*| 45.33 ± 0.01|20.00 ± 0.07|52.67 ± 0.03\n|*door-close*| 86.00 ± 0.05|47.33 ± 0.06|84.67 ± 0.04\n\n**W1.3**: ...transfer efficiency between tasks\n\n**R1.3**: As shown in R1.2, the model fine-tuned on the *button-press-topdown* task also improves its performance on *button-press-topdown-wall*  and *button-press-wall*, as these tasks share similar goals, allowing the learned skills to be transferred.\n\n**W2**：...theoretical motivation for pre-training on sub-optimal data ...\n\n**R2**：Thank you for your question. The theoretical motivation of our method is grounded in the provable efficient reward-free exploration in RL [1,2]. This approach employs a two-stage learning process for developing an RL policy: a reward-free exploration phase and a reward-based optimization phase.\n\nSpecifically, in the reward-free stage, the agent optimizes a purely exploration reward, namely, the upper-confidence bound (UCB) term of the value function. Leveraging a UCB-based bonus, the agent is able to collect diverse state-action pairs with broad coverage, thereby providing sufficient information for the subsequent planning phase. In the second stage, the agent performs least-square value iteration (LSVI) by interacting with the environment. The goal of this stage is to conduct value and policy updates for a specific task with a specific reward function. This two-stage learning paradigm, which combines reward-free exploration and reward-based optimization, can achieve polynomial sample complexity according to theoretical results [1]. We believe this approach provides a solid theoretical foundation for our method's empirical success.\n\nOur algorithm essentially follows this theoretical motivation but incorporates more practical implementations. The reward-free exploration stage in [1] is similar to count-based exploration, which aims to collect diverse datasets in the first stage [3]. In SODP, we adopt an alternative exploration method (entropy-based exploration) via a non-expert SAC policy to provide sufficient exploration in the data collection stage. We also pre-train a diffusion planner on this exploratory dataset rather than directly performing reward-based policy optimization. We empirically demonstrate that SODP requires few online samples to achieve strong performance in the second stage. This finding verifies the theoretical results of [1,2], which suggest that such a two-stage learning paradigm is sample-efficient.\n\nConcerning the BC-regularization, we would like to provide further clarification. We follow the principled RL-tuning formulation used in LLMs fine-tuning [4], which is defined as:\n\n$L=\\mathbb{E} _ {s,a}[R(s,a)]-\\beta D _ {\\rm KL}[\\pi_\\theta(a|s)\\|\\pi _ {\\rm ref}(a|s)].$\n\nHere, the first term aims to maximize the reward during fine-tuning, while the second term ensures that the learned policy $\\pi _ \\theta$ remains close to the reference policy $\\pi _ {\\rm ref}$. In the context of SODP, the $\\pi _ {\\rm ref}$ is the diffusion planner trained on a large-scale non-expert dataset in the first stage, and $\\pi_\\theta$ is the policy learned for a specific task in the second stage. However, directly calculating the KL-divergence term for diffusion models is not feasible due to their complex nature. Therefore, we adopt BC-regularization as an alternative to achieve similar effects, ensuring that the learned policy retains the desirable properties of the reference policy while adapting to the specific task.\n\n**W3**：... lacks computational efficiency analysis.\n\n**R3**: Please refer to R2 for Reviewer 4vpc.\n\n[1] On reward-free reinforcement learning with linear function approximation.\n\n[2] Reward-free exploration for reinforcement learning.\n\n[3] Count-based exploration with neural density models.\n\n[4] Training language models to follow instructions with human feedback.", "claims_and_evidence": "The claims in the submission are largely supported by clear empirical evidence, particularly those regarding SODP’s superior performance over state-of-the-art multi-task RL and diffusion-based methods, as demonstrated through success rates in Meta-World, learning curves, and fine-tuning efficiency studies. The claim that BC regularization improves fine-tuning stability is also well-supported by ablation studies comparing different regularization techniques. However, some claims lack sufficient justification. Specifically, the assertion that task-agnostic pre-training on sub-optimal data leads to strong generalization across tasks is not rigorously analyzed beyond empirical results, and the paper does not provide a theoretical explanation or deeper analysis of learned priors. Additionally, while SODP is claimed to be efficient for real-world applications, the paper does not discuss the computational overhead of diffusion-based inference compared to standard policy networks. Lastly, the claim that BC regularization prevents catastrophic forgetting is not explicitly tested through a retention or forgetting analysis. Strengthening these aspects with additional theoretical discussion, computational cost analysis, and empirical studies on task transfer and model forgetting would make the paper’s claims more convincing.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria are well-aligned with the problem of multi-task reinforcement learning using diffusion models. The use of Meta-World and Adroit as benchmark environments is appropriate, as they provide diverse manipulation tasks that test generalization and fine-tuning efficiency across different task distributions. The evaluation methodology effectively measures success rates, fine-tuning efficiency, and generalization to unseen tasks, which are key factors in assessing the effectiveness of a multi-task planner. Additionally, the comparison against state-of-the-art RL baselines, including diffusion-based planners (MTDIFF, HarmoDT) and offline-online RL methods (RLPD, IBRL, Cal-QL), ensures a fair assessment of SODP’s performance. However, the evaluation could be further strengthened by including runtime comparisons to assess the computational feasibility of diffusion-based planning versus traditional policy networks. Additionally, while success rate is a relevant metric, analyzing the stability of fine-tuning (e.g., performance degradation over long adaptation periods) and transfer efficiency between tasks would provide deeper insights. Overall, the experimental setup is well-structured, but adding computational efficiency analysis and long-term stability studies would enhance the evaluation.", "theoretical_claims": "The paper primarily focuses on an empirical approach rather than providing extensive theoretical claims, but it does include a derivation of policy gradient updates for fine-tuning the diffusion planner in Appendix A. The derivation follows standard policy gradient principles and importance sampling, aligning with established reinforcement learning literature (e.g., PPO and policy optimization methods). The formulation of reward fine-tuning as a K-step Markov Decision Process (MDP) for the diffusion denoising process appears correct in structure, mapping the diffusion process to reinforcement learning updates in a reasonable way. However, the paper does not provide formal theoretical guarantees on why pre-training with sub-optimal data leads to better generalization, which remains an empirical claim rather than a proven result. Additionally, while the policy gradient derivation follows standard RL formulations, it would be useful to include a more detailed discussion of convergence properties or potential biases introduced by the BC regularization term. There were no critical errors in the derivations, but a more formal analysis of pre-training’s impact on generalization and stability would strengthen the theoretical foundation.", "experimental_designs_or_analyses": "I reviewed the soundness and validity of the experimental design and analyses presented in the paper. The benchmark selection (Meta-World, Adroit) and success rate evaluation are appropriate for assessing multi-task reinforcement learning performance. The experimental design includes comparisons against strong baselines, such as diffusion-based RL methods (MTDIFF, HarmoDT) and offline-online RL baselines (RLPD, IBRL, Cal-QL), ensuring a fair evaluation. The ablation studies on BC regularization vs. KL and PL regularization, as well as the impact of pre-training on unseen tasks, further strengthen the empirical findings. However, some analyses are missing, such as a computational efficiency comparison and an analysis of forgetting during fine-tuning. While the learning curves and task success rates support the claims, a more detailed breakdown of failure cases and sensitivity analyses would enhance robustness. Overall, the experimental design is strong, but adding efficiency analysis and long-term stability studies would further validate the findings.", "relation_to_broader_scientific_literature": "The key contributions of the paper build on prior work in diffusion models for reinforcement learning, multi-task RL, and offline-to-online fine-tuning. Diffusion-based models have been used in trajectory generation and planning (e.g. MTDIFF), but these approaches often rely on expert demonstrations or task-conditioned optimization. The paper extends this line of research by proposing a task-agnostic pre-training strategy using sub-optimal data, similar in spirit to unsupervised pre-training in language models, but applied to RL. The fine-tuning approach builds on policy gradient-based RL fine-tuning of generative models, but introduces a behavior cloning (BC) regularization term to preserve pre-trained knowledge while optimizing for task-specific rewards. This aligns with prior findings that pre-trained policies can be refined through RL but may suffer from forgetting or policy collapse without regularization (e.g., RLHF for LLMs, reward-guided fine-tuning of generative models). The work also connects to multi-task RL approaches that aim to learn task-invariant representations (e.g., MTSAC, HarmoDT), but differs by leveraging diffusion models to model broad action distributions without explicit task conditioning. While prior work has explored multi-task pre-training for policy learning, this paper specifically demonstrates that sub-optimal data can serve as a useful training signal for generalization. The findings contribute to the broader literature by showing that diffusion models can learn meaningful priors from noisy data and be efficiently fine-tuned for high-reward adaptation, a paradigm that could be extended to real-world robotic control and hierarchical RL settings.", "essential_references_not_discussed": "The paper includes references and comparison to most of the relevant works in this domain.", "other_strengths_and_weaknesses": "Strengths\n- The paper demonstrates strong originality in its approach to training a diffusion-based planner using sub-optimal data, rather than relying on expert demonstrations or reward-conditioned learning, which is a notable departure from prior diffusion RL methods\n- The proposed two-stage framework (task-agnostic pre-training + reward-guided fine-tuning) is an interesting adaptation of pretraining-finetuning paradigms used in language models and computer vision but applied to multi-task reinforcement learning.\n- The integration of behavior cloning (BC) regularization during fine-tuning is also an important contribution, as it effectively prevents policy degradation while allowing adaptation to high-reward behaviors.\n\nWeakness\n- While the paper is methodologically strong, there are areas where it could improve in clarity and broader significance.\n-  Some explanations, particularly in theoretical motivation for pre-training on sub-optimal data, are not well justified beyond empirical results.\n- While the experimental results are compelling, the paper lacks computational efficiency analysis.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "N/A"}}, {"review_id": "lnOqzGT7bv", "reviewer": "Reviewer_w6L1", "summary": "The paper proposes a two-stage framework that pre-trains a diffusion planner using a large number of suboptimal demonstrations and performs task-specific finetuning through reinforcement learning. The diffusion planner is first pre-trained on multi-task offline data of low quality to predict action sequences given historical states. To further optimize performance for particular tasks, the planner is finetuned with policy gradients using task rewards in an online manner. The proposed method achieves strong performance on MetaWorld tasks, outperforming several offline-to-online and multi-task RL baselines.\n\n## Update after rebuttal\nI appreciate the authors' response and effort in the additional experiments. Most of my concerns and questions have been addressed. Please include the results above in the final version. I will change my score to 3.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer w6L1,\n\nWe sincerely appreciate your precious time and constructive comments. In the following, we would like to answer your concerns separately.\n\n**W1**: Different sets of MetaWorld tasks are selected for different ablation studies ... provide some context information or rationale behind the choice of particular tasks ...\n\n**R1**: Thank you for your question. Our main experiments were conducted on Meta-World 50 tasks, demonstrating the effectiveness of our method compared to baseline approaches. For ablation studies, such as investigating the impact of different pre-training datasets and evaluating performance on image-based Meta-World, we conducted experiments on Meta-World 10 tasks. This subset is a well-established benchmark that includes tasks of varying difficulty levels and is widely used in multi-task learning research [1,2]. We find it sufficient to assess the effect of the pre-training dataset and the model’s generalization ability to high-dimensional image observations. For other ablation studies, we primarily selected tasks from those used to visualize the learning curves. These tasks were chosen to cover a range of difficulty levels, ensuring a comprehensive comparison of our method's performance.\n\n**W2**: Table 1 should include more baselines that (1) are diffusion-based and (2) support offline re-training to online RL fine-tuning procedures for better comparison\n\n**R2**: We extend two diffusion-based methods, DQL [3] and IDQL [4], to multi-task offline-to-online training on Meta-World 10 tasks, following [5]. Specifically, we first pretrain a diffusion-based actor using behavior cloning on an offline dataset and then fine-tune both the actor and critic in an online environment. However, as the results below show, the diffusion actor can still be misled by inaccurate value network estimations, causing rapid forgetting of actions learned from offline pretraining. Furthermore, since goals in Meta-World are randomized, traditional methods struggle to learn how to complete them effectively.\n\n|Methods|Success rate on Meta-World 10 tasks|\n|-|-|\n|MT-DQL|10.80 ± 0.11| \n|MT-IDQL|13.57 ± 0.13| \n|SODP (ours)|65.27 ± 0.21| \n\n**W3**: ... provide ablations on additional environments, such as image-based Adroit ...\n\n**R3**: Thank you for your suggestion. We conduct ablation studies on three Adroit environments, fine-tuning the same pre-trained model with and without BC regularization separately. As shown below, the performance improvement of the model fine-tuned without BC is marginal, whereas the model fine-tuned with BC achieves significantly better results. This indicates that BC can also facilitate learning in high-dimensional single-task settings. This improvement can be attributed to the update stride constraint imposed by BC, which ensures that updates are not too aggressive, thereby preventing the model from becoming trapped in a suboptimal region [6,7].\n\n||Pre-trained model |w/o BC| w/ BC|\n|-|-|-|-|\n|Hammer|20 ± 8|38 ± 3|67 ± 6|\n|Door|55 ± 8|63 ± 3|96 ± 1|\n|Pen|28 ± 3|30 ± 5|59 ± 4|\n\n**S1**: For offline-to-online settings ... show the performance before and after online finetuning\n\n**R4**: Thank you for your suggestion. We evaluate the performance of the pre-trained model over 50 episodes across three seeds and report the results below. The pre-training performance of SODP surpasses that of other baselines, demonstrating the advantage of diffusion models in capturing multi-modal action distributions. Additionally, with RL-based reward fine-tuning, the performance improvement exceeds that of traditional offline-online methods, highlighting the effectiveness of our fine-tuning strategy in generalizing to downstream tasks more efficiently.\n\n|Methods|After pre-training| After fine-tuning|\n|-|-|-|\n|RLPD|6.4 ± 0.03|10.16 ± 0.11|\n|IBRL|14.68 ± 0.06|25.29 ± 0.22|\n|Cal-QL|15.82 ± 0.06|35.09 ± 0.12|\n|SODP (ours)|31.76 ± 0.11|60.56 ± 0.14| \n\n**Q1**: In MetaWorld, are dense rewards or sparse success signals used for RL finetuning?\n\n**R5**: We use dense rewards for RL finetuning because it is necessary to provide the reward signal $r(a _ 0)$ for each samling process to guide the update of $p _ \\theta$, as shown in Eq.(12). However, as shown in Table 2, we achieve good performance with only 100k online steps, which is 10% of the unlabeled data used for pre-training.\n\n[1] Hard Tasks First: Multi-Task Reinforcement Learning Through Task Scheduling. ICML 2024.\n\n[2] Decomposed Prompt Decision Transformer for Efficient Unseen Task Generalization. NeurIPS 2024.\n\n[3] Idql: Implicit q-learning as an actor-critic method with diffusion policies. arXiv:2304.10573\n\n[4] Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning. ICLR 2023.\n\n[5] Diffusion Policy Policy Optimization. arXiv:2409.00588.\n\n[6] Trust Region Policy Optimization. ICML 2015.\n\n[7] Proximal Policy Optimization Algorithms. arXiv:1707.06347.", "claims_and_evidence": "Please see “Other strengths and weaknesses” and “Other Comments Or Suggestions” section below.", "methods_and_evaluation_criteria": "Please see “Other strengths and weaknesses” and “Other Comments Or Suggestions” section below.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "Please see “Other strengths and weaknesses” and “Other Comments Or Suggestions” section below.", "relation_to_broader_scientific_literature": "N/A", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Strengths:\n- The proposed method effectively leverages a large suboptimal dataset to learn a diverse behavior prior, which can facilitate the second stage of task-specific finetuning\n- Extensive ablation studies demonstrate the effectiveness of design choices such as multi-task pretraining and behavior cloning regularization, especially with limited finetuning steps.\n- The paper is well-written and easy to read\n\nWeaknesses:\n- Different sets of MetaWorld tasks are selected for different ablation studies. It would be helpful to provide some context information or rationale behind the choice of particular tasks, or rigorously keep the same choice of tasks throughout ablation studies.\n- Table 1 should include more baselines that (1) are diffusion-based and (2) support offline (pre-)training to online RL fine-tuning procedures for better comparison. The bulk of current baselines are trained offline on suboptimal data without finetuning, and the selected offline-to-online baselines do not use diffusion backbones, which render comparisons unfair.\n- It appears that BC regularization is the essential component of the proposed method. It would be helpful to provide ablations on additional environments, such as image-based Adroit, to strengthen the conclusion further", "other_comments_or_suggestions": "1. For offline-to-online settings, it would be helpful to show the performance before and after online finetuning (e.g. in Table 1)", "questions_for_authors": "1. In MetaWorld, are dense rewards or sparse success signals used for RL finetuning?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "YKGZBwEBRb", "reviewer": "Reviewer_aUqS", "summary": "This paper proposes SODP, a two-stage framework for training versatile diffusion planners using sub-optimal data. SODP first pre-trains a foundation model on mixed sub-optimal trajectories without requiring expert demonstrations or reward labels, then employs RL-based fine-tuning with task-specific rewards and behavior cloning regularization. Experiments on Meta-World and Adroit benchmarks demonstrate superior performance over state-of-the-art methods with limited fine-tuning data.\n\n## update after rebuttal\nAs the rebuttal has addressed most of my concerns. I'd like to keep my score.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer aUqS,\n\nWe sincerely appreciate your precious time and constructive comments. In the following, we would like to answer your concerns separately.\n\n**C1**: The computational costs of pre-training (500k steps) are omitted\n\n**R1**: Thank you for your question. The data used for pre-training are sub-optimal trajectories. Compared to scarce expert demonstrations, collecting sub-optimal data is more cost-effective. Thus, the pre-training stage is data-efficient, allowing us to use relatively inexpensive training data to obtain a reasonable initial policy.\nThe pre-training stage is a supervised learning process in which the diffusion model is trained offline. This setup enables the use of data parallelism to efficiently fit the action distribution. In contrast, online interaction is considerably more expensive because data must be collected and used iteratively to optimize the model. Consequently, our primary focus is on the RL fine-tuning process, aiming to improve the efficiency of data utilization and reduce the costs associated with online data collection.\nOverall, large-scale offline data can be leveraged for pre-training, making total computational cost less of a concern. The primary challenge in RL is not the overall computational expense but rather the efficiency of data utilization, as the cost of interaction in the fine-tuning stage is significantly higher.\n\n**C2**: No experiments on entirely novel tasks to evaluate true generalization\n\n**R2**: We fine-tune a model on the *hammer* task and evaluate it on tasks that were not included in the pre-training dataset. The results are presented below.\n||Pre-trained model | Fine-tuned model on *hammer*|\n|-|-|-|\n|*drawer-close*|54.67 ± 0.06|75.33 ± 0.01|\n|*faucet-open*|19.33 ± 0.01|32.67 ± 0.04|\n\n**C3**: I think Appendix C is very important to support the claims. The authors had better merge them with the results in main paper.\n\n**R3**: Thank you for your suggestion. Appendix C indeed provides further evidence supporting the advantages of our methods. We conduct additional ablation studies on various fine-tuning strategies, such as fine-tuning with high-quality offline data and directly applying behavior cloning on online transitions. Additionally, we fine-tune models pre-trained on different datasets to validate the scalability of our fine-tuning approach. We will merge and discuss these ablation studies in the revised version to enhance clarity.\n\n**W1**: Insufficient analysis of scaling with task complexity.\n\n**R4**: As shown in Figure 8, all methods achieve a high success rate in the simple movement task *handle-pull-side*, which requires the arm to pull a handle upward. However, in more complex two-stage tasks such as *basketball* and *hammer*, which demand precise arm control to interact with small objects, both MTDIFF and HarmoDT exhibit a success rate below 10%.\n\nDespite this challenge, our method consistently outperforms the baselines, achieving higher performance (e.g., 80% success rate in the *hammer* task). This improvement can be attributed to the broad action prior learned during pre-training, which enables effective knowledge transfer across tasks. This is particularly advantageous for challenging tasks where task-specific feedback is limited, requiring the model to leverage skills acquired from other tasks. As shown in Figure 7, fine-tuning from scratch results in a substantial decline in success rate for the *hammer* task, dropping from 80% to 20%.\n\n**Q1**: How does SODP's computational efficiency compare to baselines in terms of training time and inference overhead?\n\n**R5**: Please refer to R2 for Reviewer 4vpc.\n\n**Q2**: the performance of MTDIFF declines markedly. How could it happen? Is it because we do not use a good way to distill the new information into the model?\n\n**R6**: Thank you for your question. We hypothesize that this is due to the 'mode shift' problem identified in previous theoretical work [1,2], which suggests that diffusion models struggle to fit the target distribution when distant modes exist. The original dataset was collected using an SAC agent with a Gaussian policy, while the newly introduced data was collected using our pre-trained model with a diffusion policy. As a result, the data distributions obtained by these policies differ significantly. When combined into a single dataset, the overall distribution contains two distant modes, which may severely hinder the model's ability to accurately capture such complex distributions.\n\n**Q3**: Are there specific task categories where SODP struggles that would help identify the advantage of this paper?\n\n**R7**: As shown in Fig. 8, other baselines often struggle with tasks that require the arm to pick and move a small object (e.g., *basketball*, *sweep*), whereas our method achieves higher performance.\n\n[1] On the Generalization Properties of Diffusion Models. NeurIPS 2023.\n\n[2] Statistical Efficiency of Score Matching: The View from Isoperimetry. arXiv:2210.00726.", "claims_and_evidence": "The main claims are supported by convincing evidence. SODP achieves 60.56% success rate on Meta-World benchmarks compared to 57.20% for the closest competitor. Ablation studies (especially Appendix C) demonstrate the effectiveness of BC regularization, importance of pre-training, and benefits of the two-stage approach. The efficiency claim is supported by strong performance with just 100k fine-tuning steps.", "methods_and_evaluation_criteria": "The methods are appropriate, using established benchmarks (Meta-World and Adroit) and standard evaluation metrics (success rate). The approach comprehensively compares against both offline-online RL methods and multi-task RL baselines. The experimental design for testing data efficiency is sensible.", "theoretical_claims": "The paper provides theoretical derivations for the policy gradient, loss function, and behavior cloning loss. These derivations appear sound but don't offer novel theoretical insights beyond established techniques from diffusion models and policy optimization literature.", "experimental_designs_or_analyses": "The experimental designs are generally sound, with comprehensive comparisons and ablation studies. Visualizations effectively demonstrate how the method improves exploration. Limitations include:\n- The computational costs of pre-training (500k steps) are omitted\n- No experiments on entirely novel tasks to evaluate true generalization", "relation_to_broader_scientific_literature": "The paper builds upon recent works using diffusion models for trajectory planning but innovates by learning from sub-optimal data rather than expert demonstrations. It effectively relates to multi-task RL methods and diffusion model fine-tuning approaches from other domains.", "essential_references_not_discussed": "None.", "other_strengths_and_weaknesses": "Minor Strengths: \n1. Novel use of sub-optimal data\n2. Extension to image-based observations.\n\nMinor Weaknesses: \n1. Insufficient analysis of scaling with task complexity.", "other_comments_or_suggestions": "None.", "questions_for_authors": "1. How does SODP's computational efficiency compare to baselines in terms of training time and inference overhead?\n2. From Tab.2, it's interesting to find that with online data collection, the baseline HarmoDT remains stable when trained on the augmented dataset, whereas the performance of MTDIFF declines markedly. How could it happen? Is it because we do not use a good way to distill the new information into the model?\n3. Are there specific task categories where SODP struggles that would help identify the advantage of this paper?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "idSv2Mp21e", "reviewer": "Reviewer_4vpc", "summary": "This paper proposes a two-stage framework for training versatile diffusion policies.  The core idea is to learn generalized diffusion policies from a large amount of low-quality multi-task trajectories, followed by task-specific online RL fine-tuning that quickly tailor the pre-trained diffusion policy to the downstream task.  In particular, this paper uses a learning objective, similar to PPO, for fine-tuning the diffusion policy with online interaction.  To enhance the learning stability, this paper further introduces a behavior cloning loss that bias the policy toward the distribution of high-reward actions.  This paper evaluates the proposed method on MetaWorld 50 tasks, outperforming other offline/online RL baselines.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer 4vpc,\n\nWe sincerely appreciate your precious time and constructive comments. In the following, we would like to answer your concerns separately.\n\n**W1**: would offline+online RL transformers work better?\n\n**R1**: We extend the multi-task Decision Transformer (DT) following [1] by pre-training it on the same sub-optimal data and fine-tuning each task online for 1M steps on the Meta-World 10 tasks. We consider two types of DT: (1) the vanilla DT, which uses a deterministic policy, and (2) the variant proposed in [1], which employs a stochastic policy. The results are presented below.\n|Methods|Success rate after pre-training |Success rate after fine-tuning|\n|-|-|-|\n|MT-ODT_deter|24.83 ± 0.10|40.07 ± 0.15| \n|MT-ODT_stoc|37.20 ± 0.12|50.47 ± 0.21| \n|SODP (ours)|31.60 ± 0.14|65.27 ± 0.21| \n\n**W2**: ...gradients are back-propagated across the whole denoising rollout... analysis on training and testing efficiency...\n\n**R2**: Similar to previous work [2], we do not require gradients to be back-propagated through the entire denoising chain at once. From Eq. (11) and (12), each step only requires the final reward $r(a^0)$ to compute the gradient. This allows us to back-propagate the gradient for each step separately, which improves memory efficiency. Furthermore, we employ DDIM sampling during the fine-tuning process, requiring only 10 denoising steps, which reduces both training and inference time. We present the computational efficiency results below. All results are obtained using a single 4090 GPU.\n\nPrevious works on fine-tuning text-to-image diffusion models [3] propose a more efficient approach that directly back-propagates the reward gradient through the denoising chain. This method bypasses the need to model the denoising process as a MDP and reduces memory consumption by truncating backpropagation. These reward-backpropagation methods rely on differentiable rewards. However, such differentiable rewards are impractical in RL because reward functions vary across different environments and cannot be easily predicted by a differentiable reward model. Therefore, we use the raw, non-differentiable environment reward and employ PPO for fine-tuning.\n\nRecent studies [4,5] have improved RL-based diffusion model fine-tuning by simplifying the original PPO or introducing dense reward signals at each denoising step, leading to greater effectiveness and better performance. We leave the exploration of optimizing SODP for future work.\n\n|Methods|Fine-tuning time for each task over 1M environment steps| Inference time for each task over 50 episodes|\n|-|-|-|\n|RLPD|4.5 - 5.5 h|~ 37s|\n|IBRL|6 - 7.5 h|~ 39s|\n|Cal-QL|8.5 - 10 h|~ 40s|\n|SODP (ours)|11 - 13 h|~ 47s|\n\n**Q1**: ... the performance of pre-training a Decision transformer with the same sub-optimal trajectories and fine-tuning it with online RL\n\n**R3**: See R1. \n\n**Q2**: What is the training and testing efficiency?\n\n**R4**: See R2.\n\n**Q3**: I'm wondering if the behavior cloning regularization is needed, if you consider a preference-based RL objectives?\n\n**R5**: Preference serves as a weak label when the reward function is hard to obtain and must be converted into an explicit or implicit reward signal [6,7] for model fine-tuning. In our setting, the model interacts with online environments and can readily obtain rewards to evaluate the performance of given actions. Therefore, we can directly fine-tune the model using these ground-truth reward labels without the additional steps of constructing and translating preference labels. The recent success of DeepSeek-R1 [8] further demonstrates the effectiveness of using rewards to improve LLM performance.\nWhether through preferences or rewards, the objective of fine-tuning is a KL-constrained reward maximization problem [6,7], defined as: $L=\\mathbb{E} _ {s,a}[R(s,a)]-\\beta D _ {\\rm KL}[\\pi_\\theta(a|s)\\|\\pi_{\\rm ref}(a|s)].$\nPreference-based methods such as DPO offer an alternative approach to injecting update gradients into the model but still require constraints to regulate the update step size. However, directly computing the KL-divergence term for diffusion models is not feasible due to their complexity. Therefore, we employ BC-regularization as an alternative to achieve similar effects, ensuring that the learned policy preserves the desirable properties of the reference policy while adapting to the specific task.\n\n[1] Online decision transformer. \n\n[2] DPOK: Reinforcement Learning for Fine-tuning Text-to-Image Diffusion Models.\n\n[3] Directly fine-tuning diffusion models on differentiable rewards.\n\n[4] A Simple and Effective Reinforcement Learning Method for Text-to-Image Diffusion Fine-tuning.\n\n[5] Aligning Few-Step Diffusion Models with Dense Reward Difference Learning.\n\n[6] Training language models to follow instructions with human feedback.\n\n[7] Direct Preference Optimization: Your Language Model is Secretly a Reward Model.\n\n[8] DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning.", "claims_and_evidence": "Yes, the claims are supported by the experimental results presented in Table 1 and 2.  Moreover, Table 3 demonstrates the efficacy of multi-task pre-training.", "methods_and_evaluation_criteria": "This paper tests baselines and the proposed method with success rate, which is a commonly adopted metric for evaluating robot policies.", "theoretical_claims": "The theoretical proof (A.1, A.2 and A.3) in the Appendix looks correct to me.", "experimental_designs_or_analyses": "The experimental designs are reasonable", "relation_to_broader_scientific_literature": "This paper contributes to the field of fine-tuning diffusion generative models, a crucial aspect of image and video generation, as well as chemical and material discovery", "essential_references_not_discussed": "I believe the references are complete", "other_strengths_and_weaknesses": "Strengths:\n\n1. One of the core ideas of this paper is to pre-train a generalized diffusion policy with scaled-up sub-optimal training data.  The idea is practically reasonable, since it is much cheaper to collect non-expert trajectories than expert ones.  This paper demonstrates that, by incorporating the knowledge from large-scale sub-optimal trajectories, the pre-trained policy adapts to downstream tasks more effectively and rapidly compared to the policy trained from scratch.\n\n2. Another key insight of this paper to enhance diffusion policies with online RL fine-tuning, which is also technically sound.  It is commonly known that policies learned with behavior cloning have issues with distributional drift, where models perform poorly on unseen scenario, and reinforcement learning could fix these issues.  This paper proposes to enhance diffusion policies, which are often pre-trained with imitation learning, by online RL fine-tuning.  With a simple PPO-like objective and behavior cloning-based regularization, the performance of the diffusion policy is improved by a large margin.\n\n3. This paper is well-written and easy to follow.  I can easily understand the high-level idea and low-level details from the paper.\n\n4. This paper presents good visualization (e.g. Figure 2 and 6), helping readers understand the effect of each component.\n\n--- \n\nWeaknesses:\n\n1. This paper lacks the analysis on the effect of separate components.  I believe this paper has three main ideas: 1) pre-training models with large-scale sub-optimal trajectories, 2) fine-tuning models with online RL, and 3) diffusion modeling.  This paper did a good job on analyzing the effect of large-scale pretraining and RL fine-tuning, but does not ablate the effectiveness of diffusion modeling.  \n\nMore specifically, Table compares to offline RL transformers (HarmoDT), offline RL diffusion models (MTDIFF), and offline+online RL baselines (which often use MLP Gaussian policies).  This paper proposed an offline+online RL diffusion model.  From the comparison to offline RL diffusion models, we know the effectiveness of oneline RL.  From the comparison to offline+online RL baselines, we know diffusion models are better than MLP Gaussian policies.  However, the comparison to offline RL transformers involves two factors (offline RL vs. offline+online RL and transformers vs. diffusion models).  It's unclear for me if the superior performance is mainly resulted from online RL fine-tuning, or the combination of diffusion modeling and online RL.  In other words, would offline+online RL transformers work better?\n\n2. This paper does not present the computational limitation of the proposed method.  From Eqn 12, I suppose gradients are back-propagated across the whole denoising rollout, which may require longer training time and larger memory usage.  It'd be great if the authors could provide analysis on training and testing efficiency, compared to others.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "1. What is the performance of pre-training a Decision transformer with the same sub-optimal trajectories and fine-tuning it with online RL?\n\n2. What is the training and testing efficiency?\n\n3. I'm wondering if the behavior cloning regularization is needed, if you consider a preference-based RL objectives [1]?\n\n---\n\nReference:\n\n[1] Direct Preference Optimization: Your Language Model is Secretly a Reward Model. Rafailov et al.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "N/A"}}]

Author rebuttal:

Meta-review: The paper introduces SODP, a novel two-stage framework for learning versatile diffusion planners by leveraging task-agnostic sub-optimal data for pre-training followed by task-guided fine-tuning. Reviewers appreciate the framework for its originality, the practical significance of using sub-optimal data, strong empirical performance on challenging benchmarks, sample efficiency in fine-tuning, and the effectiveness of behavior cloning regularization. However, initial concerns were raised regarding the theoretical grounding for generalization from sub-optimal data, computational efficiency, and the need for more comprehensive comparisons and ablations. Overall, the paper is recommended for acceptance, as the authors' rebuttal effectively addressed most concerns and solidified the work's contribution.
Decision: Accept (poster)