Diffusion Model is an Effective Planner and Data Synthesizer for Multi-Task Reinforcement Learning
Authors: Haoran He
Chenjia Bai
Kang Xu
Zhuoran Yang
Weinan Zhang
Dong Wang
Bin Zhao
Xuelong Li
Venue: NeurIPS 2023 poster (2023)
Keywords: multi-task reinforcement learning
diffusion models
planning
data synthesis
TL;DR: A  diffusion-based method to facilitate multi-task reinforcement learning

Abstract: Diffusion models have demonstrated highly-expressive generative capabilities in vision and NLP. Recent studies in reinforcement learning (RL) have shown that diffusion models are also powerful in modeling complex policies or trajectories in offline datasets. However, these works have been limited to single-task settings where a generalist agent capable of addressing multi-task predicaments is absent. In this paper, we aim to investigate the effectiveness of a single diffusion model in modeling large-scale multi-task offline data, which can be challenging due to diverse and multimodal data distribution. Specifically, we propose Multi-Task Diffusion Model (\textsc{MTDiff}), a diffusion-based method that incorporates Transformer backbones and prompt learning for generative planning and data synthesis in multi-task offline settings. \textsc{MTDiff} leverages vast amounts of knowledge available in multi-task data and performs implicit knowledge sharing among tasks. For generative planning, we find \textsc{MTDiff} outperforms state-of-the-art algorithms across 50 tasks on Meta-World and 8 maps on Maze2D. For data synthesis, \textsc{MTDiff} generates high-quality data for testing tasks given a single demonstration as a prompt, which enhances the low-quality datasets for even unseen tasks.

Diffusion Model is an Effective Planner and Data Synthesizer for Multi-Task Reinforcement Learning
Haoran He 1.2+
Chenjia Bai
Kang Xu 2.3
Zhuoran Yang*
Weinan Zhang'
Dong Wang?
Bin Zhao
Xuclong Li
Shanghai Jiao Tong University
Shanghai Artificial Intelligence Laboratory
FFudan University
"Yale University
NNottwweserrn Polytechnical University
Abstract
Diffusion models have demonstrated highly-expressive generative capabilities in vision and NL.. Recent studies in reinforcement learning (RL.) have shown that diffusion models are also powerful in modeling complex policies or trajectories in offline datasets. However, these works have been limited to single-task settings where generalist agent capable of addressing multi-task predicaments is absent. this paper, we aim to investigate the effectiveness of a single diffusion model in modeling large-scale multi-task offline data, which can be challenging due diverse and multimodal data distribution. Specifically, we propose Multi-Tash Dif- fusion Model (MTDIFF), fiffusion-basee method that incorporates Transformer backbones and prompt learning for generative planning and data synthesis multi- task offline settings. MTDIFF leverages vast amounts of knowledge available in multi-task data and performs implicit knowledge sharing among tasks. For gen- erative planning, find MTDIFF outperforms state-of-the-art algorithms across 50 tasks on Meta-World and 8 maps on Maze2D. For data synthesis, MTDIFF generates high-quality data for testing tasks given single demonstration a prompt, which enhances the low-quality datasets for even unseen tasks.
Introduction
The high-capacity generative models trained on large, diverse datasets have demonstrated remarkable success across vision and language tasks. An impressive and even preternatural ability of these models, e.g. large language models (LLMs), that the learned model can generalize among different tasks by
simply conditioning the model on instructions prompts 49 46
60 42 66). The success of 12
LLMs and vision models inspires us utilize the recent generative model to learn from large-scale offline datasets that include multiple tasks for generalized decision-making in reinforcement learning (RL). Thus far, recent attempts in offline decision -maiing take advantage of the generative capability of diffusion models [54] 20 to improve long-term planning [21 2 or enhance the expressiveness of policies [63] 8. However, these works are limited to small-scale datasets and single-task settings where broad generalization and general-purpose policies are not expected. In multi-task offline RL which considers learning a single model to solve multi-task problems, the dataset often contains noisy, multimodal, and long-hol mag-horizon trajectories collected by various policies across tasks and with various qualities, which makes it more challenging to learn policies with broad generalization and transferable take transformer-basee architecture
capabilities. Gato 47 and other generalized agents 29 65
"The work was conducted during the internship of Haoran He Shanghai Artificial Intelligence Laboratory. Corresponding authors: Chenjia Bai (baichenjia Saaichenjia@pjlab.ooog ojl.b.org.cn) Xuelong Li
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Multi-Task Dotaset
Prompt
Multi-task actions
Optimal actions
Planning
pree R2
Act
Environment
Multi-task transitions
High fidelity syuthetic transitions
(b) Synthesizing
S
A
Augment
A
Diffiusion Model
Conditional Generative Process
Original Data
Overall architecture of MTDIFF. Different colors represent different tasks. S, A and R denote the state sequence, action sequence, and reward sequence from multi-task data, respectively. Sprey and R, represent historical states and normalized return.
Figure I:
62 via sequence modeling to solve multi-task problems, while they are highly dependent on the optimality of the datasets and are expensive to train due to the huge number of parameters.
To address the above challenges, we propose a novel diffusion model further explore its generaliz- ability in multi-task setting. We formulate the learning process from multi-task data as denoising problem, which benefits the modeling of multimodal data. Meanwhile, we develop relatively lightweight architecture by using a GPT backbone [45) to model sequential trajectories, which has less computation burden and improved sequential modeling capability than previous U-Net-based 48) diffusion models (21 33 To disambiguate tasks during training and inference, instead of providing e.g. one-hot task identifiers, we leverage demonstrations as prompr conditioning, which exploits the few-shot abilities of agents [50 64 69].
We name our method the Multi-Task Diffusion Model (MTDIFF). As shown in Figure 1 investigate two variants of MTDIFF for planning and data synthesis to further exploit the utility of diffusion models, denoted as MTDIEF-P and MTDIFF-S, respectively. (a) For planning, MTDIFF-P learns a prompt embedding to extract the task-relevant representation, and then concatenates the embedding with the trajectory's normalized return and historical states as the conditions of the model. During training, MTDIFF-P learns predict the corresponding future action sequence given the conditions, and we call this process generative planning (78). During inference, given few-shot prompts and the desired return, MTDIFF-P tries denoise out the optimal action sequence starting from the current state. Surprisingly. MTDIFF-P can adapt to unseen tasks given well-constructed prompts that contain task information.
(b) By slightly changing the inputs and training strategy, we can unlock the abilities of our diffusion model for data synthesis. Our insight is that the diffusion model, which compresses informative multi-task knowledge well, more effective and generalist than previous methods that only utilize single-task data for augmentation [38] 52, (8). Specifically, MTDIFF-S learns to estimate the joint conditional distribution of the full transitions that contain states, actions, and rewards based on the task-orienter prompt. Different from MTDIFF-P, MTDIFF-S learns to synthesize data from the underlying dynamic environments for each task. Thus MTDIFF-S only needs prompt conditioning to identify tasks. We empirically find that MTDIFF-S synthesizes high-fidelity data for multiple tasks, including both seen and unseen ones, which can be further utilized for data augmentation to expand the offline dataset and enhance policy performance.
To summarize, MTDIFF is a ceffuusion-based method that leverages the multimodal generative ability of diffusion models, the sequential modeling capability of GPT architecture, and the few-shot generalizability of prompt learning for multi-task RL. To the best of our knowledge, we are the first to achieve both effective planning and data synthesis for multi-task RL via diffusion models. Our contributions include: we propose MTDIFF, novel GPT-based diffusion model that illustrates the supreme effectiveness in multi-task trajectory modeling for both planning and data synthesis; (ii) we incorporate prompt learning into the diffusion framework learn to generalize across different tasks and even adapt to unseen tasks; (iii) our experiments on Meta-World and Maze2D benchmarks demonstrate that MIFE an effective planner to solve the multi-task problem, and also powerful data synthesizer to augment offline datasets in the seen unseen tasks.

---

2 Preliminaries
2.1 Reinforcement Learning
MDP and Multi-task MDP. A Markov Decision Process (MDP) is defined by tuple (S,A,P, P, 7, where S is the state space, A is the action space, P S x A - is the transition function, R S x A x S + R is the reward function for any transition, 7 E (0, is discount factor, and μ is the initial state distribution. At each timestep i, the agent chooses an action a by following the policy scalar reward rt- In single-task RL, the goal is to learn policy π = arg max; Ea,~ Ei= r'rs] by maximizing the expected cumulative reward of the corresponding task.
S - AA. Then the agent obtains the next state $1+1 and receives
In multi-task setting, different tasks can have different reward functions, state spaces and transition functions. We consider all tasks to share the same action space with the same embodied agent. Given a specific task ~ (T)), a task sask-specified MDP can be defined as (ST .A, PT, R ,,4" y,TT, :7). Instead of solving single MDP, the goal of multi-task RL is to find an optimal policy that maximizes expected return over all the tasks: π = arg max; (TTp(())B,,aar)
Multi-Task Offline Decision-Making. In offline decision-making, the policy is learned from static dataset of transitions {18,, 8/,, T3 }) collected by an unknown behavior policy T3 (30 In the multi-task offline RL setting, the dataset D is partitioned into per-task subsets as D = Uا Dr, where D, consists of experiences from task T. The key issue of RL the offline setting is the distribution shift problem caused by lemporal-difference (TD) learning. In our work, we extend the idea of Decision Diffuser [2] by considering multi-task policy learning as a conditional generative process without fitting value function. The insight is to take advantage of the powerful distribution modeling ability of diffusion models for multi-task data, avoiding facing the risk of distribution shift.
Offline RL learns policies from a static dataset, which makes the quality and diversity of the dataset significant [15]. One can perform data perturbation (28 to up-sample the offline dataset. Alternatively, we synthesize new transitions (s, a, 8', r)) by capturing the underlying MDP of a given task via diffusion models, which expands the original dataset and leads to significant policy improvement.
2.2 Diffusion Models
We employ diffusion models to learn from multi-task data D = D; in this paper. With the sampled trajectory from D, we denote (x((T) as the k-step denoised output of the diffusion model, and y()) is the condition which represents the task attributes and the trajectory's 's optimality (e.g.. returns). A forward diffusion chain gradually adds noise to the data ZD(T) ~ q(x(T)) in steps with pre-defined variance schedule 3r. which can be expressed as
= = (N(xx(T); v1 - Be (T),
In this paper, we adopt Variance Preserving (VP) beta schedule [67 and define Bk
1 - exp (-8min (+) 0.5(3max - Bmin) 2 where Bmax 10 and Bmi = 0.1 are
constants, A trainable reverse diffusion chain, constructed as ( (T) (T)) = N (T)
(o (I=(()), v(T), Ex), can be optimized by simplified surrogate loss 20)
Ldenoiee == E (1.)' K)), (F) [||e= <0(xx(T), b(T),k)))
(2)
where Eg parameterized by deep neural network is trained to predict the noise N I) added to
the dataset sample (T) to produce (T). By setting ag := 1 - Bx and ã := II Cts, we obtain
Be
3h xk-1(()
(T((T)
(x) (T), k)) + + /Bko, o NO, 1), for k = (K, 1}.
VI as
Classifier-free guidance [9 aims to learn the conditional distribution ((=(T)lu(T)) without sep- arately training classifier. In the training stage, this method needs to learn both a conditional <0()))(T), (T)), k) and an unconditional (x (T), 2, model, where (()) is dropped out. Then the perturbed noise eo (0(xx((), 20, k) a ( (T), (T), k) - ca ((xx((). 23, k)) used to generate samples latter, where can be recognized as the guidance scale.
3

---

Reverse Denoising, Process
counte--lockwwicc
Bypass a and prese button
Slide plate cobinet
An example of the denoising process of MTDIFF. We choose 4 tasks for visualization and æx (T) is sampled from the Gaussian noise for each task. Since different tasks require different manipulation skills, the corresponding action sequences dispersed in the embedding space. Our model learns such properties and generates task-specific sequences based on task-relevant prompts.
Figure 2:
Methodology
3.1 Diffusion Formulation
To capture the multimodal distribution of the trajectories sampled from multiple MDPs, we formulate the multi-task trajectory modeling conditional generative problem via diffusion models:
max; E,,u, (o y(r)),
where (T) the generated desired sequence and y(T) is the condition. (7( will then be used for generative planning or data synthesis through conditional reverse denoising process po for specific
tasks. Maximizing Eq. (3) can be approximated by maximizing variational lower bound [20].
In terms of different inputs and outputs in generative planning and data synthesis, x(T) can be represented in different formats. We consider two choices to formulate xx()) in MTDIIF-P and MTDIFF-S, respectively. (i) For MTDIFF-P, z(T) represents the action sequence for planning. We model the action sequence defined as:
x (T) = (az, ac+1. ... a1+H-1))+
(4)
with the context condition as
yd (=)
8'(r):= (Z, (1
(5)
where t, H, R(T) and L denote the time visited in trajectory т, the length of the input sequence æ. the normalized cumulative return under t and the length of the observed state history, respectively. Z the task-relevant information as prompt. We use v'(r) as an ordinary condition that is injected into the model during both training and testing, while considering
R(T) the classifier-free guidance to obtain the optimal action sequence for a given task. (ii) For data synthesis in MTDIFF-S, the inputs and outputs become the transition sequence that contains states, actions, and rewards, and then the outputs are utilized for data augmentation. We define the transition sequence as:
[8t 8t41
8t4 ff 1-1
(T) =
ag a++1
(6)
It T++1
with the condition:
y (T( = [2),
(7)
where y (7) takes the same conditional approach as y (T)). Figure illustrates the reverse denoising
process of MTDIFF-P learned on multi-task datasets collected in Mera-World [73]. The result demonstrates that our diffusion model successfully distinguishes different tasks and finally generates the desired TT(T). We illustrate the data distribution of zo(T) in a 2D space with dimensional reduction via T-SNE [6], as well as rendered states after executing the action sequence. The result shows that, with different task-specific prompts as conditions, the generated planning sequence for a specific task will be separate from sequences other tasks, which verifies that MID can leam the distribution of multimodal trajectories based on (T).

---

3.2 Prompt, Training and Sampling
In multi-task RL and LLM-driven decision-making, existing works use one-hot task identifiers [56] [4 or language descriptions [ 6 as conditions in multi-task training. Nevertheless, argue that the one-hot encoding for cach task [22] [4] suffices for learning a repertoire of training tasks while cannot generalize novel tasks since it does not leverage semantic similarity between tasks. In addition, the language descriptions [SI o 4T 40] of tasks require large amounts of human labor annotate and encounter challenges related to ambiguity [2].. In MTDIFF, we use expert demonstrations consisting of a few trajectory segments to construct more expressive prompts in multi- task settings. The incorporation of prompt learning improves the model's ability for generalization and extracting task-relevant information to facilitate both generative planning and data synthesis. We remark that similar method has also been used PromptDT [69] Nonetheless, how such a prompt contributes within diffusion-baser framework remains to be investigated.
Specifically, we formulate the task-specific label Z as trajectory prompts that contain states and actions:
S 4++
Z
(8)
where each element with star-script is associated with trajectory prompt, and J is the number of environment steps for identifying tasks. With the prompts as conditions, MTDIFF can specify the task by implicitly capturing the transition model and the reward function stored in the prompts for better generalization to unseen tasks without additional parameter-tuning.
In terms of decision-making in MTDIFF-P, we aim devise the optimal behaviors that maximize return. Our approach is to utilize the diffusion model for action planning via classifier guidance [19]. Formally, an optimal action sequence x is sampled by starting with Gaussian noise xx (T) and refining x (T) into (T) at each intermediate timestep with the perturbed noise:
eo (T)(()) (T), 8, k) + (< (x) w'))
ea (=)())) (T),
(9)
where y' (r) is defined in Eq. (5). R(T) is the normalized return of and a is hyper-parameter that seeks to augment and extract the best portions of trajectories in the dataset with high return. During training, we follow DDPM (20] well as glassifier-fre guidance [99] to train the reverse diffusion process po. parameterized through the noise model cg, with the following loss:
(0)
[|| co (o0 (T), R) 320 |]]
(10)
Note that with probability p sampled from Bernoulli distribution, we ignore the conditioning
return R(T). During inference, we adopt the low-temperature sampling technique 2 to produce high-likelihood sequences. We sample (7) ~ (1 1. ( Rmax (T), k 1), -1x
where the variance is reduced by B [0, 1) for generating action sequences with higher optimality.
For MTDIFF-S, since the model aims to synthesize diverse trajectories for data augmentation, which does not need to take any guidance like R(T), we have the following loss:
(*) : (o00
c) (T),
(11)
We sample x (T) ~ ~N (F), 1), 1) The evaluation process is given in Fig. 3
3.3 Architecture Design
Notably, the emergence of Transformer [62] and its applications on generative modeling 44
provides a promising solution to capture interactions between modalities of different tasks. Naturally. instead of U-Net [48) which commonly used in previous single-task diffusion RL works 21 33). we parameterize with novel transformer architecture. We adopt GPT2 45) architecture for implementation, which excels in sequential modeling and offers a favorable balance between performance and computational efficiency. Our key insight is to train the diffusion model in a unified manner to model multi-task data, treating different inputs as tokens a unified architecture, which expected to enhance the efficiency of diverse information exploitation.
5

---

Mrlirr :
Training
Evaluation
Predicted
NrDur-F
Feweabal Pronge
Pa
Arbors
Prediction Head
Pa
Tranations
GPT2 Transformer
Pazz fooses
Layer Norm
Milless
Positional Encoding
Prompt
MLP
MLP P
MLP
History
MT
Inputs
Actions
@
Transition:
Figure 3: Model architecture of MTDIFF, which treats different inputs as tokens in a unified architecture. The two key designs are (i) the trainable GPT2 Transformer which enhances sequential modeling, and (ii) the MLPs and prediction head which enable efficient training.
As shown in Figure 3 first, different raw inputs x are embedded into embeddings h of the same size via separate MLPs f. which can be expressed as:
hp (xppprtppt ..Try (r (imestoo)) >> for prompt and diffusion timestep
hrr Jry ( > for MTDIFF-S
hy = fa(a 1, = Su (I (ishistory,, ), RR fn( ro.tum))))
for MTDIFF-P
Then, the embeddings hp and hTT are prepended as follows to formulate input tokens for MTDIFF-P and MTDIFF-S, respectively:
h Pooeens IN(NNN [hr, hers, h t Mi (N (hr, hri, hyr]+ EPos),
where pos is the positional embedding, and LN denotes layer normalization 33 for stabilizing train- ing. In our implementation, we strengthen the condition of the stacked inputs through multiplication with the diffusion timestep hry and addition with the return h GPT2 is decoder-only transformer that incorporates self-attention mechanism capture dependencies between different positions in the input sequence. We employ the GPT2 architecture as trainable backbone in MTDIFF to handle sequential inputs. It outputs updated representation as:
out transformer(("
hout transformer(Hicteen
Finally, given the output representation, use prediction head consisting of fully connected layers to predict the corresponding noise at diffusion timestep k. Notice that the predicted noise shares the same dimensional space the original inputs, which differs from the representation size d. This noise is used in the reverse denoising process po during inference. We summarize the details of the training process, architecture and hyperparameters used in MTDIFF in Appendix A
4 Related Work
Diffussio Models in RL. Diffusion models have emerged as a powerful family of deep generative models with a sccrd-bbeakkin performance in many applications across vision, language and com- 31 32). Recent works in RL have demonstrated the capability of diffusion models learn the multimodal distribution of offline policies [63] 43
binatorial optimization 44 46 I7
human problem as a conditional gener- ative process (2 and learn to generate the trajectories satisfying conditioned constraints. However, these works are limited to the single-task settings, while we further study the trajectory modeling and generalization problems of diffusion models in multi-task settings.
behaviors 8). Other works formulate the sequential decision-makinn
Multi-Task RL and Few-Shot RL. Multi-task RL aims to learn a shared policy for a diverse set of tasks. The main challenge of multi-task RL is the conflicting gradients among different tasks, and

---

previous online RL works address this problem via gradient surgery [74]. conflict-averse learning [4], and parameter composition [70 [56]. Instead, MTDIFF addresses such a problem in an offline setting through conditional generative process via a novel transformer architecture. Previous Decision-Tranee Transformer (DT))-aaeed methods [47 29 72 which consider handling multi-task problems, mainly rely on expert trajectories and entail substantial training expenses. Scaled-QL. (26) adopts separate networks for different tasks and is hard to generalize to new tasks. Instead of focusing on the performance of training tasks in multi-task RL, few-shot RL aims to improve the generalizability in novel tasks based on the learned multi-task knowledge. Nevertheless, these methods need additional context encoders 777 [99 gradient descents in the finetuning stage [57] [58 29 In contrast, use prompts for few-shot generalization without additional sraueterr-tning
Data Augmentation for RL. Data augmentation
13 39 has been verified to be effective in RL. Previous methods incorporate various data augmentations (e.g. adding noise, random translation) on observations for visual-based RI. [71 28, 52 27), which ensure the agents learn on multiple views of the same observation. Differently, we focus on data augmentation via synthesizing new experiences
rather than perturbing the origin one. Recent works 776 [0] consider augmenting the observations of robotic control using a text-guided diffusion model whilst maintaining the same action, which differs from our approach that can synthesize novel action and reward labels. The recently proposed SynthER (88 is closely related to our method by generating transitions of trained tasks via a diffusion model. However, SynthER is studied in single-task settings, while we investigate whether diffusion model can accommodate all knowledge multi- talti-task datasets and augment the data for novel tasks.
5 Experiments
In this section, we conduct extensive experiments to answer the following questions: (1) How does MTDIFT-P compare to other offline and online baselines the multi-task regime? (2) Does MTDIFF- s synthesize high-fidelity data and bring policy improvement?
(3) How is MTDIFF-S compared with other augmentation methods for task RL? (4) Does the synthetic data of MTDIFF-S match the original data distribution? (5) Can both MTDIFF-P and MTDIFF-S generalize to unseen tasks?
5.1 Environments and Baselines
Meta-World Tasks The Meta-World benchmark [3] contains 50 qualitatively-disticc manipu- lation tasks. The tasks share similar dynamics and require Sawyer robot interact with various objects with different shapes. joints, and connectivity. In this setup, the state space and reward functions of different tasks are different since the robot is manipulating different objects with dif- ferent objectives. At each timestep, the Sawyer robot receives a 4-dimensional fine-grained action, representing the 3D position movements of the end effector and the variation of gripper openness. The original Meta-World environment is configured with a fixed goal, which is more restrictive and less realistic in robotic learning. Following recent works 170, (6). we extend all the tasks to random-goal setting and refer to it as MT50-rand, We use the average success rate of all tasks as the evaluation metric.
By training a SAC [88 agent for each task in isolation, we utilize the experience collected in the replay buffer as our offline dataset. Similar to 26), we consider two different dataset compositions: (i) Near-optimal dataset consisting of the experience (100M transitions) from random to expert (convergence) in SAC-Replay, and (ii) Sub-optimal dataset consisting of the initial 50% of the trajectories (50M transitions) from the replay buffer for each task, where the proportion of expert data decreases lot. We summarize more details about the dataset Appendix G
is navigation task that requires agent to traverse from a randomly designated location to a fixed goal in the 2D map. The reward is if succeed and 0 otherwise. Maze2D can evaluate the ability of RL algorithms to stitch together previously collected sub-tra sub-trajectories, which helps the agent find the shortest path to evaluation goals. We use the agent's 's scores as the evaluation metrie. The offline dataset is collected by selecting random goal locations and using planner to generate sequences of waypoints by following a PD controller.
Maze2D Tasks Maze2D [5|
Baselines We compare our proposed MTDIFF (MTDIFF-P and MTDIFF-S) with the following basclines. Each baseline has the same batch size and training steps MTDIFF. For MTDIFF-P,
7

---

have following baselines: (i) PromptDT. PromptDT [69] built on Decision-Transforee (DT) [ aims to learn from multi-task data and generalize the policy to unseen tasks. PromptDT generates actions based on the trajectory prompts and reward-to-go. We use the same GPT2-network as in MTDIFF-P. The main difference between our method and PromptDDT is that we employ diffusion models for generative planning. (ii) MTDT. We extend the DT architecture ( to learn from multi- task data. Specifically, MTDT concatenates an embedding 2 and a state as the input tokens, where is the encoding of task ID. In evaluation, the reward-to-go and task ID are fed into the Transformer to provide task-specific information. MTDT also uses the same GPT2-network as in MTDIFF-P. Compared to MTDT, our model incorporates prompt and diffusion framework to learn from the multi-task data. (iii) MTCQL. Following scaled- QL [26), we extend CQL (25] with multi-head critic networks and task-ID conditioned actor for multi-task policy learning. (iv) MTTQL. We extend IQL. (24) for multi-task learning using a similar revision of MTCQL. The TD-based baselines (i.e., MTCQL and MTIQL)) are used to demonstrate the effectiveness of conditional generative modeling for multi-task planning. (v) MTBC. We extend Behavior cloning (BC) to multi-task offlin policy learning via network scaling and task-ID conditioned actor that is similar MTCQL and MTIQL..
As for MTDIFF-S, we compare it with two baselines that perform direct data augmentation in offline RL. (i) RAD. We adopt the random amplitude scaling (28) that multiplies random variable to states, i... s' = s x 2, where & ~ Uniform(a. .3.. This augmentation technique has been verified to be effective for state-based RL. (ii) S4RL. We adopt the adversarial state training [52] by taking gradients with respect to the value function to obtain a new state, i.c. s' < s + eV <VeJQ(T(s)) where Jo is the policy evaluation update performed via Q function, and is the size of gradient steps. We summarize the details of all the baselines in Appendix B
5.2 Result Comparison for Planning
How does MTDIFF-P compare to baselines in the multi-task regime? For a fair comparison, we add MTDIFF-P-
as a variant of MTDIFF-P by replacing the prompt with a one-hot task-ID, which is used in the base- lines except for PromptDT. The first action generated by MTDIFF-P is used to interact with the cnvironment. According to Tab, ! have the following key observations. (i) Our method achieves better performance than base- lines in both near-optimal and sub-optimal set- tings. For near-optimal datasets, MTDIFF-P and P-DEEET--ONEHOT achieve about 60% success rate, significantly outperforming other methods and performing comparably with MTBC. However, the performance of MTBC decreases lot in sub-optimal datasets. BC is hard to handle the conflict behaviors in experi- ences sampled by a mixture of policies with different returns, which has also been verified in previous offline imitation methods [35 68 In contrast, both MTDIFF-P and MTDIFF-ONEHE perform the best sub-optimal datasets. (ii) We compare MTDIFF-P with two SOTA multi-task online RL methods, CARE (53] and PaCo (56), which are trained for 100M steps in MT50-rand. MTDIFF-P outperforms both of them given the near-optimal dataset, demonstrating the potential of solving the multi-task RL problems in an offline setting. (iii) MTDT is limited in distinguishing different tasks with task ID while PromptDT performs better, which demonstrates the effect of prompts in multi-task settings. (iv) As for TD- basclines, we find MTCQL almost fails while MTIQL performs well. We hypothesize that since MTCQL penalizes the OOD actions for each task, it will hinder other tasks' learning since different tasks can choose remarkably different actions when facing similar states. In contrast, IQL learns a value function without querying the values of OOD actions.
Table I: Average success rate across 3 seeds on Meta- V2 MT50 with random goals (MT50-rand), Each task is evaluated for 50 episodes.
Methods
Near-optimal
Sub-optimal
we
CARE (Online)
50.8
PaCo
36 (Online)
57.3 1.3
MTDT
20.99 1 2.66
20.63 : 2.21
PromptDT
45.68 + 1.84
39.76 + 2.79
MTBC
60.39 1 0.86
34.53 :: 1.25
MTCQL MTIQL
56.21 1.39
43.28 0.90
MTDIFF-P (ours)
39.53 1.12
48.67 + 1.32
MIDUN-P-ONEHOT (ours)
61.32 + 0.89
48.94 + 0.95
We remark that MTDIFF-P based on GPT network outperforms that with U-Net in similar model size, and detailed results are given in Appendix C Overall, MTDIFF-P is an effective planner in multi-task settings including sub-optimal datasets where we must stitch useful segments of suboptimal trajectories, and near-optimal datasets where we mimic the best behaviors. Meanwhile, we argue that although ATDUFF-P-ONEHEE performs well, it cannot generalize to unseen tasks without prompts.
8

---

K
E E
Scores
(a) MTDIEF-P
(b) PrompiDT
(c) MTDIFF-P
(d) PrompiDTt(aaed)
(1) Seen map
(2) Unseen map
Figure 4: Seen and unseen maps of Maze2D with long planning path.
Figure 5: Average scores ob- tained in Maze2D maps.
Goal position is denoted as P.
Does MTDIFF-P generalize to unseen tasks? We further carry out experiments on Maze2D to evaluate the generalizability of MTDIFF-P. We select PromptDT as our baseline, as it has demonstrated both competitive performances on training tasks and adaptation ability for unseen tasks (69] We use 8 different maps for training and one new map for adaptation evaluation. The setup details are given in Appendix D We evaluate these two methods map. The average scores obtained in 8 training maps are referred to Figure 5 To further illustrate the advantages of our method compared to PromptDT, we select one difficult training map and the unseen map for visualization, as shown in Figure 4 According to the visualized path, we find that (1) for maps in training, MTDIFF-P generates shorter and smoother path, and (2) for unseen maps, PromptDT fails to obtain reasonable path while MTDIFF-P succeed, which verifies that MTDIFF-P perform few-shot adaptation based on trajectory prompts and the designed architecture.
both seen maps and an unseen
5.3 Results for Augmentation via Data Synthesis
Does MTDIFF-S synthesize high-fidelity data and bring policy improvement? We train MTDIFF- near-optimal datasets from 45 tasks evaluate its generalizability. We select 3 training tasks and unseen tasks, and measure the policy improvement of offline RL training (i.e., TD3-BC [16] with data augmentation. For each evaluated task, MTDIFF-S synthesizes 2M transitions to expand the original IM dataset. From the results summarized in Table MTDIFF-S can boost the offline performance for all tasks and significantly increases performance by about 180%, 131 %6 and 161 % for box-close, hand-insert, and coffee-push, respectively.
MTDiff single
MTDiffss
SynthER
Original
100
80
Success Rate (%)
Success Rate (%)
40
%
20
Figure Average success rate across 3 seeds on coffee-push. Dataset is aug- mented by MTDIFFFS trained under differ- ent number of tasks.
coffee-push disassemble
dial-tum
Figure 6: Results of multi-task and single- task aug- mentation
How does MTDifT-s perform and benefit from multi-task training? From Table 2 we find MTDIFF-S achieves superior policy improvement in seen tasks compared with previous SOTA augmentation methods (i.e., S4RL and RAD) that are developed in single-task RL. We hypothesize that, by absorbing vast knowledge of multi-task data training. MTDIFF-S can perform implicit data sharing [75] by integrating other tasks' knowledge into data synthesis of the current task. To verify this hypothesis, we select three tasks (i.e., coffee-push, disassemble and dial-turn) to re-train MTDIFF-S on the corresponding single-task dataset. We denote this variant as MTDIFF-S-SINCE We implement SynthER (3 to further confirm the benefits multi-task training. SynthER also utilizes diffusion models to enhance offline datasets for policy improvement, but it doesn't it take into

---

account the aspect of learning from multi-task data. We observe that MTDIFF-S outperforms both two methods, as shown in Figure 6 What's more, we re-train MTDIFF-S on 10, 20, and 30 tasks respectively, in order to obtain the relationship between the performance gains and the number of training tasks. Our findings, as outlined in Fig. [ provide compelling evidence in support of our hypothesis: MTDIFF-S exhibits progressively superior data synthesis performance with increasing task diversity.
Does MTDIFF-S generalize to unseen tasks? We answer this question by conducting offline RL training on the augmented datasets of 3 unseen tasks. According to Table 2 MTDIFF-S is well- generalized and obtains significant improvement compared to the success rate of original datasets. MTDIFF-S boost the policy performance by 131%, 180% and 32% for hand-insert, box-close and bin-picking, respectively. We remark that S4RL performs the best on the two unseen tasks, i.e., box-close and bin-picking, since it utilizes the entire datasets to train Q-functions and obtains the augmented states. Nevertheless, we use much less information (i.e., a single trajectory prompts) for augmentation.
Table 2: Average success rate across 3 seeds on Meta- Sita-World-Ve V2 single task with random goals. Each selected task is evaluated for 500 episodes.
Tasks
Original
S4RL
RAD
MTDIFF-SOOURS
box-close
23.40 7.11
+ 3.51 211%]
71.20 L.8
203%)
08.3118.61
180%
Unseen Tasks
hand-insert
30.60
= 9.77
60.20 1.57 (1 96%)
43.79 3.44
13%)
70.87 + 4.59 (1 131%)
bin-picking
42.13
14.33
72:20
+ 4.17(1 71%)
43.27 2 1.38
2%)
36.13 7.03( 32%]
sweep-into
91.8 1.14
90.53 3.52() 1%)]
58.06 + 9.86 4961
02.87
198
Seen Tasks
coffec-push
28.60 + 1/1.55
28.73 8.11 0.15%
33. 33.19 + 2686(1 16%)
74.67 + 6.79 161%)
disussemble
60.20 + 16.29
52.20 + 5.68(] 12%)
60.93 + 20. 20. 1%)
69.00 + 4.72( 14.6%)
Does the synthetic data of MTDIFF-S match the original data distribution? We select 4 tasks and use T-SNE [61] to visualize the distribution of original data and synthetic data. We find the synthetic data overlap and expand the original data distribution while also keeping consistency with the underlying MDP. The visualization results and further analyses are given in Appendix E
6 Conclusion
We propose MTDIFF, diffusion based effective planner and data synthesizer for multi-task RL. With the trajectory prompt and unified GPT-based architecture, MTDIFF can model multi-task data and generalize unseen tasks. We show that in the MT50-rand benchmark containing fine-grained manipulation tasks, MTDIFF-P generates desirable behavior for each task via few-shot prompts. By compressing multi-task knowledge in single model, we demonstrate that MTDIFF-S greatly boosts policy performance by augmenting original offline datasets. In future research, we aim to develop practical multi-task algorithm for real robots to trade off the sample speed and generative quality. We further discuss the limitations and broader impacts of MTDIFF in Appendix
Acknowledgments
This work is supported by the National Natural Science Foundation of China (Grant No.62306242266276666). the National Key R&D Program of China (Grant No.20222D01601000 Shanghai Municipal Science and Technology Major Project (202IHHZDZZ0102) and Shanghai Artificial Intelligence Laboratory.
References
[1] Michael Ahn, Anthony Brohan, Noah Brown, Yevgen Chebotar, Omar Cortes, Byron David, Chelsea Finn, Keerthana Gopalakrishnan, Karol Hausman, Alex Herzog, et al. Do i can, not as i say: Grounding language in robotic affordances. arXiv preprint arXiv:2204.01691, 2022.
[2] Anurag Ajay, Yilun Du, Abhi Gupta, Joshua B. Tenenbaum, Tommi S. Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision making? In The Eleventh International Conference on Learning Representations, 2023.
[3] Lei Jimmy Ba, Jamie Ryan Kiros, and Geoffrey F. Hinton. Layer normalization. CoRR, Abs/1607.0644 2016. URL httt:///axiii org/ aws//b/... 06450
10

---

[4] Shaojie Bai, J Zico Kolter, and Vladlen Koltun. An empirical evaluation of generic convolutional and recurrent networks for sequence modeling. arXiv preprint 1rXiv:1803.0127/ 2018.
[5] Fan Bao, Shen Nie, Kaiwen Xue, Chongxuan Li, Shi Pu, Yaole Wang. Gang Yue, Yue Cao, Hang Su, and Jun Zhu, One transformer fits all distributions in multi-modal diffusion at scale. arXiv preprint arXiv:2303. 06555, 2023.
[6] Authony Brohan, Noah Brown, Justice Carbajal, Yevgen Chebotar, Joseph Dabis, Chelsea Finn, Keerthana Gopalakrishnan, Karol Hausman, Alex Herzog, Jasmine Hsu, et al. Rt-1: Robotics transformer for real-world control at scale. arXiv preprint arXiv:2212.06681 2022.
[7] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
[8] Huayu Chen, Cheng Lu, Chengyang Ying. Hang Su, and Jun Zhu. Offline reinforcement learning via high-fidelity generative behavior modeling. In The Eleventh International Conference Learning Representations, 2023.
[9] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors. Advances in Neural Information Processing Systems, 2021. URL httpps://openreviees net/ forum?id=a// APMM499d
[10] Zoey Chen, Sho Kiami, Abhishek Gupta, and Vikash Kumar. Genaug: Retargeting behaviors to unseen situations via generative augmentation. arXiv preprint arXiv:2302.06671, 2023.
[11] Cheng Chi, Siyuan Feng, Yilun Du, Zhenjia Xu, Eric Cousineau, Benjamin Burchfiel, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion. arXiv preprint arXiv:2303.04137, 2023.
[12] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmano, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311 2022.
[13] Karl Cobbe, Oleg Klimov, Chris Hesse, Taehoon Kim, and John Schulman. Quantifying generalization in reinforcement learning. In International Conference on Machine Learning, pages 1282-1289. PMLR, 2019.
[14] Frederik Ebert, Yanlai Yang, Karl Schmeckpeper, Bernadette Bucher, Georgios Georgakis, Kostas Daniilidis, Chelsea Finn, and Sergey Levine Bridge data: Boosting generalization of robotic skills with cross-domain datasets. CoRR, abs/2109.13396,, 2021. URL https
//arxiv.orz wwtp/og///bbbbaee www.ab//9 13396
[15] Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4RL: datasets for deep data-driven reinforcement learning. CoRR, abs/2004.07219, 2020. URL https
7/arxiv. org/ abwaaa////...// 07219
[16] Scott Fujimoto and Shixiang Shane Gu. A minimalist approach to offline reinforcement learning.
In Thirty-Fifth Conference on Neural Information Processing Systems, 2021.
[17] Shansan Gong, Mukai Li, Jiangtao Feng, Zhiyong Wu, and Lingpeng Kong. Diffuseq: Sequence to sequence text generation with diffusion models. In The Eleventh International Conference on Learning Representations, 2023. URL httpp:://opereviiit net//formmid==j0j--
rLVXsj
[18] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off- policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning, pages 1861-1870. PMLR, 2018.
[19] Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. In NeurlPS 2021 Workshop on Deep Generative Models and Downstream Applications, 2021. WtR.lhtttp:///peeeeiei het:/forumPdd-qpBBAkkk.Y...

---

[20] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems,
2020.
[21] Michael Janner, Yilun Du, Joshua B. Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. In International Conference Machine Learning, 2022.
[22] Dmitry Kalashnikov, Jacob Varley, Yevgen Chebotar, Benjamin Swanson, Rico Jonschkowski, Chelsea Finn, Sergey Levine, and Karol Hausman. Mt-opt: Continuous multi task robotic reinforcement learning scale. CoRR, abs/2104.08212, 2021. URL atttp:///ariiv org/
abs/2104.00
[23] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
[24] Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offiin reinforcement learning with implicit q-learning. In International Conference on Learning Representations, 2022. URL https
[25] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative for offline reinforcement learning. CoRR, abs/2006.04779, 2020. URL /tttps://arxiv.org. abs/2006.04 04779
[26] Aviral Kumar, Rishabh Agarwal, Xinyang Geng, George Tucker, and Sergey Levine. Offline q-learning on diverse multi-task data both scales and generalizes. In The Eleventh International Conference on Learning Representations, 2023.
[27] Michael Laskin, Aravind Srinivas, and Pieter Abheel. Curl: Contrastive unsupervised represen- tations for reinforcement learning. In International Conference on Machine Learning, pages 5639-5650. PMLR, 2020.
[28] Misha Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Reinforcement learning with augmented data. Advances in neural information processing systems, 13:19884-19895, 2020.
[29] Kuang- Huei Lee, Ofir Nachum. Mengjiao Sherry Yang. Lisa Lee, Daniel Freeman, Sergio Guadarrama, Ian Fischer, Winnie Xu, Eric Jang, Henryk Michalewski, et al. Multi-game decision transformers. Advances Neural Information Processing Systems, 35:27921-27936. 2022.
[30] Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. CoRR, abs/2005.01643, 2020. URL attps://arxics org/ abs/ 01643
[31] Xiang Lisa Li, John Thickstun, Ishaan Gulrajani, Percy Liang, and Tatsunori Hashimoto. Diffusion-LM improves controllable text generation. In Alice H. Oh, Alekh Agarwal. Danielle Belgrave, and Kyunghyun Cho, editors, Advances Neural Information Processing Systems, 2022. URL httpstts:::///// /https://openreiieee wett://rrumlddd3311Ess...
[32] Yang Li, Jinpei Guo, Runzhong Wang, and Junchi Yan. From distribution learning in training gradient search in testing for combinatorial optimization. In Advances in Neural Information Processing Systems, 2023.
[33] Zhixuan Liang, Yao Mu, Mingyu Ding, Fei Ni, Masayoshi Tomizuka, and Ping Luo. Adaptdif- fuser: Diffusion models as adaptive self-evolving planners. ArXiv, abs/2302.01 2023.
[34] Bo Liu, Xingchao Liu, Xiaojie Jin, Peter Stone, and Qiang Liu. Confliet-averse gradient descent for multi-task learning. Advances in Neural Information Processing Systems, 34: 14888888-18890
2021.
[35] Minghuan Liu, Hanye Zhao, Zhengyu Yang, Jian Shen, Weinan Zhang, Li Zhao, and Yan Liu. Curriculum offline imitating learning. In A. Beygelzimer, Y. Dauphin, P. Liang, and
J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021. URL https://openrevies.. wtt.://iww....eoooooooo/oonnnnaa
12

---

[36] Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm- solver+t: Fast solver for guided sampling of diffusion probabilistic models. arXiv preprint arXiv:2211.01095 2022.
[37] Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. DoM-solver: fast ODE solver for diffusion probabilistic model sampling in around 10 steps. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[38] Cong Lu, Philip Ball, and Jack Parker-Holder. Synthetic experience replay. In Workshop on Reincarnating Reinforcement Learning at ICLR 2023, 2023. URL ltttps///pperreiiii Sattorummidd09933Tyyae
[39] Guozheng Ma, Zhen Wang, Zhecheng Yuan, Xueqian Wang, Bo Yuan, and Dacheng Tao. A comprehensive survey of data augmentation in visual reinforcement learning. arXiv preprint arXiv:2210004561, 2022.
[40] Oier Mees, Lukas Hermann, and Wolfram Burgard. What matters in language conditioned robotic imitation learning over unstructured data. IEEE Robotics and Automation Letters, 7(4): 11205-11212, 2022. doi: 10.11099LRA2222.339900
[41] Oier Mees, Lukas Hermann, Erick Rosete-Beas, and Wolfram Burgard. Calvin: A benchmark for language-conditioned policy learning for long-horizon robot manipulation tasks. IEEE Robotics and Automation Letlers, 7(3):7527-7334, 2022. doi: 10.
[42] OpenAL. Gpt-4 technical report, 2023.
[43] Tim Pearce, Tabish Rashid, Anssi Kanervisto, Dave Bignell, Mingfei Sun, Raluca Georgesou, Sergio Valcarcel Macua, Shan Zheng Tan, Ida Momennejad, Katja Hofmann, and Sam Devlin. Imitating human behaviour with diffusion models. In The Eleventh International Conference on Learning Representations, 2023.
[44] William Peebles and Saining Xie. Scalable diffusion models with transformers. arXiv preprint arXiv:2212.09748, 2022.
[45] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, al. Language models are unsupervised multitask learners. OpenAl blog, 188)9, 2019.
[46] Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents, 2022.
[47] Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Barth-Maron Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, et al. A generalist agent, arXiv preprint arXiv:2205. 06 2022.
[48] Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical Image Computing and Computer-Assistee Iitervention-MICCA 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18, pages 234-241. Springer, 2015.
[49] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, S. Sara Mahdavi, Rapha Gontijo Lopes, Tim Salimans, Jonathan Ho, David J Fleet, and Mohammad Norouzi. Photorealistic text-to-image diffusion models with deep language understanding, 2022.
[50] Victor Sanh, Albert Webson, Colin Raffel, Stephen Bach, Lintang Sutawika, Zaid Alyafeai, Antoine Chaffin, Arnaud Stiegler, Arun Raja, Manan Dey, M Saiful Bari, Canwen Xu, Urmish Thakker, Shanya Sharma Sharma, Eliza Szezechla, Taewoon Kim, Gunjan Chhablani, Nihal Nayak, Debajyoti Datta, Jonathan Chang, Mike Tian-Jian Jiang, Han Wang, Matteo Manica, Sheng Shen, Zheng Xin Yong, Harshit Pandey, Rachel Bawden, Thomas Wang, Trishala Neeraj, Jos Rozen, Abheesht Sharma, Andrea Santilli, Thibault Fevry, Jason Alan Fries, Ryan Teehan, Teven Le Scao, Stella Biderman, Leo Gao, Thomas Wolf, and Alexander M Rush. Multitask prompted training enables zero-shot task generalization. In International Conference on Learning Representations, 2022.
13

---

[51] Mohit Shridhar, Jesse Thomason, Daniel Gordon, Yonatan Bisk, Winson Han, Roozbeh Mot- taghi, Luke Zettlemoyer, and Dieter Fox. ALFRED: A Benchmark for Interpreting Grounded Instructions for Everyday Tasks. In The IEEF Conference on Computer Vision and Puttern Recognition (CVPR), 2020. URL SeIl.httss:.. /tttps://arivv. owgp wwtg/bg//////o//o aww/aaa///1/ 01734
[52] Samarth Sinha, Ajay Mandlekar, and Animesh Garg. S4rl: Surprisingly simple self-suueriision for offline reinforcement learning robotics. In Conference on Robot Learning, pages 907-917. PMLR, 2022.
[53] Shagun Sodhani, Amy Zhang, and Joelle Pineau. Multi-task reinforcement learning with context- based representations. In International Conference on Machine Learning, pages 9767-9779. PMLR, 2021.
[54] Jascha Sohl- Dickstein, Eric Weiss, Niru Maheswaranathan. and Surya Ganguli. Deep unsu- pervised learning using nonequilibrium thermodynamics. In Francis Bach and David Blei, editors, Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pages 2256 2265, Lille, France, 07-09 Jul 2015. PMLR. URL htttt:://proceeeeeoiio
mlr wwwppp/wwssseeeeepeepeptppepepse wwtp:////oooooiossssoaaaoa httteemiin..mm111111
[55] Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models, 2023.
[56] Lingfeng Sun, Haichao Zhang, Wei Xu, and Masayoshi Tomizuka. Paco: Parameter- compositional multi-task reinforcement learning. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[57] Yanchao Sun, Shuang Ma, Ratnesh Madaan, Rogerio Bonatti, Furong Huang, and Ashish Kapoor. SMART: Self-supervised multi-task pretraining with control transformers. In The Eleventh International Conference on Learning Representations, 2023.
[58) Adrien Ali Taiga, Rishabh Agarwal, Jesse Farebrother, Aaron Courville, and Marc G Belle- mare. Investigating multi-task pretraining and generalization in reinforcement learning. In The Eleventh International Conference on Learning Representations, 2023. URL https
/ /ooenreview
[59] Denis Tarasov, Alexander Nikulin, Dmitry Akimov, Vladislav Kurenkov, and Sergey Kolesnikov. CORL: Research-oriented deep offline reinforcement learning library. In 3rd Offline RL Workshop: Offiiee RL as "Launchpad", 2022. URL https://opeereevier http:/pperreiiee.eeefffrooorr net/forum?
id=SyAS49BBcv
[60] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. (tt.::2322.199911. 13911, 2023.
[61] Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine learning research, 9(11), 2008.
[62] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[63] Zhendong Wang, Jonathan J Hunt, and Mingyuan Zhou. Diffusion policies an expressive policy class for offline reinforcement learning. In The Eleventh International Conference on Learning Representations, 2023.
[64] Jason Wei, Maarten Bosma, Vincent Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, Andrew M. Dai, and Quoe V L.e. Finetuned language models are zero-shot learners. In International Conference on Learning Representations, 2022.
[65] Ying Wen, Ziyu Wan, Ming Zhou, Shufang Hou, Zhe Cao, Chenyang Le, Jingxiao Chen, Zheng Tian, Weinan Zhang, and Jun Wang. On realization of intelligent decision-makinn makinn the real world: A foundation decision model perspective. arXiv preprint arXiv:2212.12669 2022.
14

---

[66] Chenfei Wu, Shengming Yin, Weizhen Qi, Xiaodong Wang, Zecheng Tang, and Nan Duan. Visual chatgpt: Talking, drawing and editing with visual foundation models. arXiv preprint arXiv:2303.04671 2023.
[67] Zhisheng Xiao, Karsten Kreis, and Arash Vahdat. Tackling the generative learning trilemma with denoising diffusion gans. In International Conference on Learning Representations, 2022.
[68] Haoran Xu, Xianyuan Zhan, Honglei Yin, and Huiling Qin. Discriminator-weise imitation learning from suboptimal demonstrations. In Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 24725-24742. PMLR, 17-23 Jul 2022.
offline
[69] Mengdi Xu, Yikang Shen, Shun Zhang, Yuchen Lu, Ding Zhao, Joshua Tenenbaum, and Chuang Gan. Prompting decision transformer for few-shot policy generalization. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari. Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Muchine Learning Research, pages 24631-24645 PMLR, 17-23 Jul 2022. URL. https //proceedings.nu mlr. prwss/ v12./v6222/// wttp//v2222x2222ee html
[70] Ruihan Yang, Huazhe Xu, Yi Wu, and Xiaolong Wang. Multi-task reinforcement learning with soft modularization. Advances in Neural Information Processing Systems, 33:4767-4777, 2020.
[71] Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image augmentation is all you need: Regu- larizing deep reinforcement learning from pixels. In International Conference on Learning Representations, 2021. URL
net.//nee/eeeeee orrm? 1d=GG6-6sTvvatt
[72] Albert Yu and Ray Mooney, Using both demonstrations and language instructions to efficiently learn robotic tasks. In The Eleventh International Conference on Learning Representations, 2023. URL
[73] Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on Robot Learning (CoRL), 2019. URLhttpss //arxiv. orgg /bs//
10897
[74] Tianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine, Karol Hausman, and Chelsea Finn. Gradient surgery for multi-task learning. In H. Larochelle, M. Ranzato, R. Hadsell, M.R. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 5824 5836. Curran Associates, Inc., 2020. httphttsss::pooceeeieiiii neerddddddddiinnnnss.eee.eee.ere cc/ Pap
paper_ files/
[75] Tianhe Yu, Aviral Kumar, Yevgen Chebotar, Karol Hausman, Sergey Levine, and Chelsea Finn. Conservative data sharing for multi-task offline reinforcement learning. Advances in Neural Information Processing Systems, 34:11501-11516, 2021.
[76] Tianhe Yu, Ted Xiao, Austin Stone, Jonathan Tompson, Anthony Brohan, Su Wang, Jaspiar Singh, Clayton Tan, Jodilyn Peralta, Brian Ichter, et al. Scaling robot learning with semantically imagined experience. arXiv preprint 2rxiv:2302.111500 2023.
[77] Haoqi Yuan and Zongqing Lu. Robust task representations for offline reiffreeinor learning via contrastive learning. In International Conference on Machine Learning, pages 25747-25759. PMLR, 2022.
[78] Haichao Zhang, Wei Xu, and Haonan Yu. Generative planning for temporally coordinated exploration in reinforcement learning. In International Conference on Learning Representations, 2022. URL httpss://openrreieee http:://pppaooooeeeeeeeeeeeeeeee
[79] Luisa Zintgraf. Kyriacos Shiarlis, Maximilian Igl, Sebastian Schulze, Yarin Gal, Katja Hofmann, and Shimon Whiteson. Varibad: A very good method for bayes-adaptive deep rl via meta- learning. In International Conference on Learning Representations, 2020.
15

---

A The Details of MTDIFF
In this section, we give the pseudocodes of MTDIFF-P and MTDIFF-S Alg. and Alg. 2 respec- tively. Then we describe various details of the training process, architecture and hyperparameters:
We set the per-task batch size as 8, so the total batch size is 400. We train our model using Adam optimizer 23
with 2e learning rate for 2 train steps.
We train MTDIFT on NVIDIA GeForce RTX 3080 for around 50 hours.
We represent the noise model as the transformer- based architecture described in Section 3.3. MLP fr which processes prompt is 3-layered MLP (prepended by layer norm 3 and with Mish activation). MLP fr; which processes diffusion timestep fR is 2-layered MLP (prepended by Sinusoidal embedding and with Mish activation). JR which processes conditioned Return and fH which processes state history are 3-layered MLPs with Mish activation. fA which processes actions, JTR which process transitions and prediction head are 2-layered MLPs with Mish activation. The GPT2 transformer is configured as hidden layers and 4 attention heads, The code of GPT2 is borrowed from https://gi thub. com/ kzl/ wttp//ecsiieeeaee sn transformer
We choose the probability p of removing the conditioning information be 0.25.
In MTDIFF-P, we choose the state history length L 2 for Meta-World and L 5 for Maze2D.
We choose the trajectory prompt length J = 20.
We use = 200 for diffusion steps.
We set guidance scale a 1.2 for extracting near-optimal behavior.
We choose B 0.5 for low temperature sampling.
Algorithm MTDIFF-P Training and Evaluation
Training Process
Initialize: training tasks Trrain training iterations N, multi-task dataset D, per-task batch size M, multi-task trajectory prompts Z
1: for = to N do
2: for Each task Ti E Titrain do
3:
Sample action sequences (T:) of length H and corresponding state history Sprey of length from D, with batch size M Compute normalized return R(T,) under
4:
5: Sample trajectory prompts T. of length from Z; with batch size end for
6:
7: Get batch B = (T)(T;), Sprev R(T,)) T [=]
8: Randomly sample a diffusion timestep (1, K) and obtain noisy sequences x (T,)
9: Omit the R(T) condition with probability 3 ~ Bern(p)
10: Compute (0) and update MTDIFE-P model
11: end for
# Evaluation Process
1: Given task, reset the environment and desired return Rimax(T)
2: Obtain the initial state history ho. few-shot prompts
3: Set low-temperature sampling scale 8, classifier-free guidance scale a
4: for L = to mmxx do
5: Initialize xxx (T) (N 3T)
6: Sample 7" ~ Z, and formulate ( [h,, for k K to do
7:
8:
? ē= Eg c0 ( k)k) + ( (T), Rmax ( (#x-1, 1)-1) + Denoise(="(T)...
9:
x (T) (
11:
end for
Execute the first action from x the current action interact with the environment
13: Obtain the next state, and update hi
end for
16

---

Algorithm 2 MTDIFF-S Training and Data Synthesis
# Training Process
Initialize: training tasks T training iterations N, multi-task dataset D, per-task batch size M, multi-task trajectory prompts Z
for 1 to N do
2: for Each task Ti E T do
Sample transition sequences (i) of length H from D, with batch size M
4: Sample trajectory prompts of of length J from Z₃ with batch size M
5: end for
6: Get batch B = ( 7 J==1 17'"'""
7: Randomly sample a diffusion timestep k ~ u(1, K) and obtain noisy sequences a: (T:)
8: Compute (0) and update MTDIFF-S model
9: end for
# Data Synthesis Process
Initialize: synthetic dataset D = Ø, synthesizing times M
1: Given task, obtain few-shot prompts Z
2: for m
M do
3: Initialize (T) I)
4: Sample -" ~ Z, and formulate (T) = [T*]
5: for k K to do
6: ? = Eg [ (T), kT)
7:
Ex-1) <- Denoiso ( (T).
8:
x (T) N (1 1
9:
end for
10: Update D = D U xT(TT)
11:
end for
B The Details of Baselines
In this section, we describe the implementation details of the baselines:
PromptDT uses the same prompts and GPT2 transformer in MTDIFF-P for taining. We borrow the code from
com wwom///x33//pooopttooo for implementation.
MTDT embeds the tskalD which indicates the task to a embedding with size 12, then is concatenated with the raw state. With such conditioning, we broaden the original state space to equip DT with the ability to identify tasks in this multi-task setting. We keep other hyperparam- eters and implementation details the same as the official version btctps://gitan com/ com/kz1/ decision- www.eat-taaasfofrrreotoom
MTIQL uses multi-head critic network to predict the Q value for each task, and each head is parameterized with a 3-layered MLP (with Mish activation). The actor-network is parameterized with 3-layered MLP (with Mish activation). During training and inference, the scalar task is embedded via 3-layered MLP (with Mish activation) into latent variable 2, and the input of the actor becomes the concatenation of the original state and We build MTIQL based on the code https://githuu: com/ cinkoff-ai/Con
59
MTCQL is applied with a similar revision in MTIQL. The main difference is that MTCQL is based on CQL [25] algorithm instead of the IQL algorithm (24). We build MTCQL based on the
code https: / thub. http://inofffffiieeee
[59].
MTBC uses similar taskID-cogmnn taskiD-connitional actor in MTIQL and MTCQL. For training and inference, the scalar taskiD is embedded via a 3-layered MI.P (with Mish activation) into latent variable and the input of the actor becomes the concatenation of the original state and 2. The actor is parameterized with 3-layered MLP and outputs predicted actions.
RAD adopts the random amplitude scaling (28 that multiplies random variable to states during training, i.c.. s' where Uniform/a, B). We choose 0.8 and 3 1.2.
17

---

S4RL adopts the adversarial state training 152 by taking gradients with respect to the value function obtain new state, i.e. + <V (J)(((()) where Ng the policy evaluation update performed via Q function, and is the size of gradient steps. We choose = 0.01.
C Ablation Study on Model Architecture
The architecture described in $3.3 handles input types of different modalities as tokens that share similar formats, actively capturing interactions between modalities. The incorporation of transformer also helpful for sequential modeling. To ablate the effectiveness our architecture design, train MTDIFF-P using U-Net with similar model size to ours on the near-optimal dataset. We use Temporal Convolutional Network (TCN)
encode the prompt into an embedding 2p. and inject in the U-Net layers condition. We follow the conditional approach in 2 and borrow the code for com/ annerm/ diffuser [21] The results summarized in
temporal U-Net from Table show that our model architecture outperforms U-Net to learn from multi-task datasets.
Table 3: Average success rate across 3 different seeds of MTDIFF-P and MTDIFF-P (U-Net) on MT50-rand.
Methods MTDIFF-P
Success rate near-optimal dataset (%)
Success rate sub-optimal dataset (%)
59.53 1.12
48.67 1.32
MTDIFFF (U-Net)
55.67 1.27
47.42
D Environmental Details of Maze2D
map
map
map
(d) map
H E
+
(e) map 5
(1) map 6
gg) map 7
(h) map 8
Figure 8: 2D visualization eight training maps designed in Maze2D.
We design eight different training maps for multi-task training, which shown Fig. g Different tasks have different reward functions and transition functions. For generalizability evaluation, we have designed one new unseen map. Although in Fig. 4, MTDIFF is able generalize on new maps while PromptDT fails, we should acknowledged that MTDIFF may fail at some difficult unseen cases, as shown in Fig. [ The reasons may lie folds: One is the inherent difficulty of the case itself, and the other is that the case's deviation from the distribution of the training cases surpasses the upper threshold of generalizability of MTDIFF. We also provide another map where MTDIFF succeeds while PromptDT fails in Fig. I0
18

---

We collect 35k episodes together to train our model and PromptDT. The episodic length is set as 600 for training and 200 for evaluation. After training with 512 batch size for 2e5 gradient steps, we evaluate these methods seen and unseen maps.
(a) PromptDT
(b) MTD1FF
(c) PromptDT
(d) MTDIFF
(1) Failure case 1
(2) Failure case 2
Figure 9: 2D visualization of 2 difficult cases where both PromptDT and MTDIFF both fail. These cases are both unseen during training. Goal position is denoted F.
(a) MTDIFF-P
(b) PrompiDT (failed)
Figure 10: Unseen maps of Maze2D with long planning path. MTDIFF-P reach the goal while PromptDT fails. Goal position denoted F.
E
Data Analysis
E.1
Distribution Visualization
We find the synthetic data is high-fidelity, covering broadening the original data distribution, which makes the offline RL method performs better in the augmented dataset. The distribution visualization is shown in Fig.
(a) Coffee-push
(b) Disassemble
(c) hand-insert
(d) box-close
Figure 11: 2D visualization of sampled synthetic data and original data via T-SNE 61). The data is sampled from tasks coffee-push, disassemble, hand-insert and box-close respectively.
19

---

E.2 Statistical Analysis
Following SynthER [38], we measure the Dynamics Error (MSE between the augmented next state and true next state), and 1.2 Distance from Dataset (minimum 1.2 distance of each datapoint from the dataset) for the augmented data of each method (i.e., MTDIFF-S, S4RI. and RAD), as shown in T Since S4RI performs data augmentation by adding data within the e-ball of the original data points, it has the smallest dynamics error with a small e. RAD performs random amplitude scaling and causes the largest dynamies error. We remark that S4RL performs local data augmentation around the original data and can be limited in expanding the data coverage of offline datasets. In contrast, our method generates data via diffusion model without explicit constraints to the original data points, which also has small dynamics error and significantly improves the data coverage, benefiting the offline RL training.
Table 4: Comparing L2 distance from the training dataset and dynamics error under each method.
Methods
Dynamics Error
L2 Distance from Dataset
MTDIFF-S
0.0174
0.4552
S4RL
0.0001
0.4024
RAD
0.0641
0.4617
F Limitations and Discussions
In this section, we will discuss the limitations and broader impacts of our proposed method MTDIFF.
Limitation. Diffusion models are bottlenecked by their slow sampling speed, which caps the potential of MTDIFF for real-time control. How to trade off the sampling speed and generative quality remains to be crucial research topic. For concrete example in MetaWorld, it takes average 1.9s in wall-clock time to generate one action sequence for planning (hardware being a 3090 GPU). We can improve the inference speed by leveraging recent sampler called DPM-solver 37 36 to decrease the diffusion steps required to 0.2x without any loss performance, and using larger batch size (leveraging the parallel computing power of GPUs) to evaluate multiple environments at once. Thus the evaluation run-time roughly matches the run-time of non-diffusion algorithms (diffusion step is 1). In addition, consistency models (55) are recently proposed to support one-step and few-step generation, while the upper performance of what such models can achieve is still vague.
Broader Impacts. As far as we know, MTDIFF is the first proposed ceffusion-based approach for multi-task reinforcement learning. It could be applied to multi-task decision decision-making, and also could be used to synthesize more data to boost policy improvement. MTDIFF provides solution for achieving generalization in reinforcement learning.
G Dataset collection
Meta-World. We train Soft Actor-Critic (SAC) 18 policy in isolation for each task from scratch until convergence. Then we collect 1M transitions from the SAC replay buffer for each task, consisting of recording samples in the replay buffer observed during training until the policy reaches the convergence of performance. For this benchmark, we have two different dataset compositions:
Near-optimal dataset consisting of the experience (100M transitions) from random to expert (convergence) in SAC-Replay.
Sub-optimal dataset consisting of the initial 50% of the trajectories (50M transitions) of the near-optimal dataset for each task, where the proportion of expert data decreases a lot.
To visualize the optimality of each dataset clearly, we plot the univariate distribution of return each kind of dataset in Fig. 12 Our dataset is available at https: //tit./ttttt/ 1y/3MWf40w
20

---

Dendlly
Denaity
Normalloed Betum
Normalized Return
(a) Near-optimal dataset
(b) Sub-optimal dataset
Figure 12: Density visualization of the normalized return in the dataset.
Maze2D The offline dataset collected by selecting random goal locations and using a planner to generate sequences of waypoints by following a PD controller. We borrow the code from https
//github. comm/aaaaaaFFuuuttiiicci Foundation/D4Rl collect 35k episodes in total.
generate datasets for 8 training maps. We [5
H
Differences Between PromptDT and MTDIFF-P
The remarkable superiority of MTDIFF-P over PromptDT emerges from our clegant incorporation of transformer architecture and trajectory prompt within the diffusion model framework, effectivelyy modeling the multi-task trajectory distribution. PromptDT is built on Decision Transformer and is trained in autoregressive manner, which is limited to predicting actions step by step. However, MTDIFF-P leverages the potency of sequence modeling, empowering to perform trajectory genera- tion adeptly. MTDIFF-P has demonstrated SOTA performance in both multi-tasi maltiitask decision-making and data synthesis empirical experiments, while PromptDDT fails to contribute to data synthesis. Tech- nically, MTDIFF- extends Decision Diffuser 2 into the multi-task scenario, utilizing classifier-free guidance for generative planning to yield high expected returns. To further verify our claim, we train model the publicly available PromptDTT datasets [69]. i.e., Cheetah-vel and Ant-dir. These chosen environments have been judiciously selected due to their inherent diversity of tasks, serving robust test validate the capability multi-task learning. We report the scores (mean and std for 3 seeds) in Table 5
Table 5: Average scores obtained by MTDIFF-P and PromptDT across 3 seeds. We observed that MTDIFF-P outperforms PromptDT largely, demonstrating its high efficacy and potency.
Methods
Cheetah-vel
Ant-dir 602.17 + 1.68
MTDIFF-P -29.09 + 0.31
PromptDT -34.43 2.33
409.81 9.69
I Single-Task Performance
We train one MTDIFF-P model on MT50-rrand and evaluate the performance for each task for 50 episodes. We report the average evaluated return in Table 6
21

---

Table 6: Evaluated return of MTDIFF-P for each task in MT500raand We report the mean and standard deviation for 50 episodes for each task.
Tasks
Return on near-optimal dataset
Return on sub-optimal dataset
baskethall-v2
2735.7 1927.7
2762.7 1928.8
bin-picking-22
733.7 + 1211.8
59.6 32.3
button-press-ton pautoon-peex-ttodoon--22
1491.6 390.9
1395.8 + 332.6
button-press-v2
2419.2 413.7
2730.4 514.3
putton-press-walll-22
1474.718887
2613.1 + 906.8
coffee-button-v2
3157.9 1274.3
1649.4 1 869.5
coffee-pull-v2
437.5 1: 597.3
98.8 115.4
coffce-push-v2
71.7
dial-turn-v2
2848.1 996.3
2244.5 881.5
disassemble-v2
369.9 265.4
3
door-closc-v2
4 377.6
4
door-lock-v2
3082.4 1041.4
door-open-v2
3458.1 1 960.5
2457.0 833.0
door-unlock-v2
2082.6 1 1568.4
3078.8 1158.1
hand-insert-v2
927.9 1 1527.6
288.8 1 932.1
drawer-close-v2
4824.5 8.8
1
drawer-open-v2
3530.4 1512.6
2218.8 529.3
faucet-open-v2
198.1
4245.6 = 575.2
faucet-close-v2
1
4624.1 = 107.5
handle-pres prendl--rres-ssde--e
3423.3 1075.6
3995.8 1130.0
handle- press-v2
handle-press-v2
3397.9 1752.1
3125.4 1724.7
handle-pull-side--2
314155 1775.4
1474.2 1392.9
handle-pull-v2
2845.7 1970.6
2856.8 1699.8
lever-pull-v2
3383.5 1299.2
3125.6 1433.1
peg-insert-side--2
915.1 1624.8
497.8 + 1015.3
pick-place-wall-22
649.8 + 1288.6
288.2 : 988.1
pick-out- off--out-of-hole-ve
1792.8 1959.9
700.5 1190.3
reach-v2
1 1 645.4
166.48842.99
push-back-v2
97.2 611.3
566.3 1042.9
push-v2
64.6
pick-place-12
166.0 759.2
7.8
plate-slide-v2
4096.4 1202.5
4 495.8
plate-slide-side--2
+
+08900+0444
plate-slide-backk-22
4 373.0
3963.4 927.3
platesslidebacc--side-
3872.0 1 1151.6
4186.5 772.6
soccer-v2
443.5 1 785.5
480.4 994.8
push-wall-v2
873.9 1 1718.3
705.2 1535.7
shelf-place-v2
204.5 1: 714.3
366.1 919.0
sweep-into-v2
1297.9 1: 1661.5
506.2 1 1368.5
sweep-v2
1397.1 + 1922.3
599.1 1359.2
window-open-v2
1453.6 1101.4
3095.5 + 751.1
window-close-v2
2963.9 875.3
3177.6 737.8
asscmbly-v2
2470.7 + 1758.6
663.7 28.7
buttonrpress-lopooon-walllll
1270.9 214.5
1199.9 1: 203.7
hammer-v2
868.7 1: 983.4
1024.4 1024.3
peg-unplug-side-v2
1439.1 4: 1817.8
136.4 508.6
reach-wall-v2
4249.9 + 536.0
1191.4 112.9
stick-push-v2
1288.5 + 1587.3
790.1 1097.6
stick-pull-v2
601.4 4 1346.3
287.1 + 922.8
box-close-vy
2683.7 + 1823.4
2273.3 + 1707.3
22


Reviews:
[{"review_id": "qkOywTeza6", "reviewer": "Reviewer_8zHR", "summary": "The author propose a diffusion model MTDIFF for multi-task RL. The goal is to leverage diffusion process and transformer backbone to have a sota generative model for RL. The author demonstrates its effectiveness with generative planning on meta world and data augmentation.", "questions": "With generative planning, can the authors show the visualization of how model get from random noise to optimal trajectory? Like a GIF would be nice.", "limitations": "N/A", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "Combing transformer and diffusion together is a straightforward idea, but the paper excuted it very well. Also leveraging the power of prompting in generative planning is a cool idea. The analysis on how MTDIFF benefit from multi-tasking in data synthesis is interesting.", "weaknesses": "The paper also seems unpolished. On table 1, row MTIQL is empty. Also the figure 5 should be \"We evaluate these two methods on both seen maps and an unseen map. The average scores obtained in 8 training maps are referred to Figure 5\", but I don't know which is PromptDT which is yours.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable comments and a positive assessment of our work! We are glad that they found that our paper provides an interesting analysis and executes a straightforward idea very well. To address your concerns, we would polish our paper in the next version to make it clear, and our detailed response follows:\n\n>1. On table 1, row MTIQL is empty.\n\nIn Table 1, the row of **MTCQL** remains empty as MTCQL almost fails on the Metaworld MT50-v2 benchmark. The corresponding analysis is in lines 307-309 in our paper. We hypothesize that the failure of MTCQL emanates from penalizing the O.O.D. actions across diverse tasks explicitly, thereby exacerbating the challenge of distribution shift within the ambit of multi-task training. Through our empirical experiments, we observed that MTCQL would firstly achieve around 10% success rate at the early training stage, and then drop down to almost 0% success rate. \n\n>2. Also the figure 5 should be \"We evaluate these two methods on both seen maps and an unseen map. The average scores obtained in 8 training maps are referred to Figure 5\", but I don't know which is PromptDT which is yours.\".\n\nWe are sorry that we mistook the x-label of Fig. 5. This figure quantifies the average scores obtained in 8 training maps of MTDiff-p and PromptDT, respectively. The final corrected version of Figure 5 is referred to Fig. 1 in the global PDF (with standard deviation added).\n\n>3. With generative planning, can the authors show the visualization of how model get from random noise to optimal trajectory? Like a GIF would be nice.\n\nThank you for this good suggestion! We present a visual depiction of the model's performance improvement as the number of denoised steps increases on Maze2D, and the result can be seen in Fig. 3 in the global PDF. \n\nWe appreciate the reviewer’s feedback again which helped us to improve the quality of the paper. We also hope that our response have sufficiently addressed the reviewer’s concerns."}, {"review_id": "fRK2hPqRFK", "reviewer": "Reviewer_EBNs", "summary": "This paper extends diffusion-based planners to multi-task settings by combining prompt learning. Specifically, a few segments from expert demonstrations are used as task prompts to distinguish different tasks and guide the diffusion model to generate task-specific trajectories. A classifier-free guidance approach is used to sample trajectories that yield high expected returns. Unlike previous diffusion-based planners, the authors utilize a Transformer as the denoising network. The experiments conducted demonstrate a marked improvement in performance.", "questions": "Please see weaknesses.", "limitations": "Yes.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper is easy to follow. I feel quite enjoyable while reading this paper. ", "weaknesses": "- The experiment section needs further revision:\n  - Figure 5 x-label is not correct.\n- What are the quantitative results on Maze2D unseen map\n- Figure 6, what is the criterion for selecting these 2 tasks? This experiment compares the model training from scratch and the model with pre-training. The results are not that surprising since pre-training provides a better initialization in general.\n- line 334: The subtitle is confusing. This paragraph is mainly about figure 6, not comparing other augmentation methods.\n- In Table 2, can the authors add results with 2M random data expanded for S4RL and RAD? Just want to remove the effect of different data sizes.\n- In Table 2, what are the original results? Is it the maximal return of the dataset?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thorough and detailed review and welcome suggestions for improvement.  Here, we address your concerns as follows:\n\n>1. The experiment section needs further revision: Figure 5 x-label is not correct.\n\nWe apologize for mistaking the x-label of Fig. 5. This figure quantifies the average scores obtained in 8 training maps of MTDiff-p and PromptDT, respectively. The final corrected version of Figure 5 is referred to Fig. 1 in the global PDF (with standard deviation added).\n\n>2. What are the quantitative results on Maze2D unseen map？\n\nThank you for this question! We have added another unseen map for thorough comparison, which is rendered in Fig. 4 within the global PDF. Subsequently, we evaluate *MTDiff-p* and PromptDT on the 2 unseen maps for 100 episodes (150 steps for an episode), and report average scores as follows: \n\n|   Method   | Scores  |\n| :--------: | :-----: |\n| *MTDiff-p* | $23.37$ |\n|  PromptDT  | $17.03$ |\n\n>3. Figure 6, what is the criterion for selecting these 2 tasks? This experiment compares the model training from scratch and the model with pre-training. The results are not that surprising since pre-training provides a better initialization in general.\n\nWe would like to provide further clarification regarding Figure 6. It is essential to note that the figure does not illustrate the model acquired through pre-training; rather, it serves as a comparative representation of policies learned with distinct synthetic datasets generated by single-task and multi-task models. \n\nConcretely, we train *MTDiff-s* on the multi-task dataset encompassing 45 tasks and *MTDiff-s-single* on the single-task dataset. The models, *MTDiff-s* and *MTDiff-s-single*, are then leveraged to generate task-specific data (i.e., for *coffee-push* and *disassemble* in Fig. 6) for data augmentation in offline RL learning. We compare the performance trained on these two kinds of augmented data, and find *MTDiff-s* shows a significant advantage compared to *MTDiff-s-single*. The reason would be that *MTDiff-s* can conduct implicit knowledge sharing [1] among tasks and transfer knowledge from other tasks to expand the data coverage of the target task, which is also appreciated by Reviewer 8zHR and UScg. As a result, both methods in Fig. 6 are trained from scratch with augmented data without pre-training.\n\nThe reason why we select these 2 tasks is that they are difficult and the policies perform relatively poorly learned with their original dataset. They are suitable for evaluating policy improvement after data augmentation to validate our hypothesis. We also added new experiments on another new task *dia-turn* and observed similar results, which is referred to Fig. 2 in the global PDF.\n\n>3. The subtitle is confusing.\n\nThis paragraph highlights the superior performance of *MTDiff-s* in comparison to the baseline methods and offers a comprehensive analysis of the advantageous aspects derived from its multi-task training paradigm. To enhance clarity, we propose revising the subtitle to read:  \"How does *MTDiff-s* perform and benefit from multi-task training?\" \n\n>4. In Table 2, can the authors add results with 2M random data expanded for S4RL and RAD?\n\nThanks for the question. We would like to clarify that both S4RL and RAD have already used additional augmented data in their training process, so it is no need to add more samples to the dataset. Specifically, during each training step of S4RL or RAD, a batch of transitions are sampled and the algorithm augments the transitions to generate new ones. As a result, the size of newly added training samples is batch_size $\\times$ training_steps (e.g., 400$\\times$1M), which is much larger than the 2M augmented data in *MTDiff-s*.\n\n>5. In Table 2, what are the original results? Is it the maximal return of the dataset?.\n\nThe original result is the success rate attained by the TD3-BC agent following training on the original 1 million (1M) dataset.\n\n[1] Tianhe Yu, et al. Conservative data sharing for multi-task offline reinforcement learning. NeurIPS 2021."}, {"review_id": "X6YqP2vEaO", "reviewer": "Reviewer_UScg", "summary": "- The paper investigates the effectiveness of learning a diffusion model for modelling multi-task offline data. To do so, the paper introduces two variants of using a learned diffusion model: (a) by planning over a sequence of actions, (b) generating data, and using the generated data for offline policy optimisation/improvement. \n\n- The paper compares the proposed method to various different baselines, and shows decent improvements over all the baselines. \n\n- The paper shows the ability of the proposed model to generate useful data, by augmenting kow-quality datasets. \n\n- The paper can be seen as an extension of Decision Diffuser for multi-task scenario.\n\n", "questions": "- For Table H in appendix, it will be useful to report results for other baselines too. \n- It will be useful to see how the performance changes as a result of varying the generated data. For example: MTDIFF-S synthesizes 2M transitions to expand the original 1M dataset. It will be helpful to report reports as to how the baselines perform by varying amount of data generated or augmented. \n- It will also be helpful to see how the performance changes as the MTDIFF-S  is trained on less amount of tasks (i.e., instead of 45 tasks it's trained on 5/10/15/20/30 tasks, as the working hypothesis is MTDIFF-S can perform implicit data sharing, so will be useful to verify this by decreasing the number of tasks (figure 6 only shows the comparison for 2 tasks).", "limitations": "See Weaknesses and Questions. ", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is very well written.\n- The paper does a good job comparing to various strong baselines, as well as showing the usefulness of the generated synthetic data. \n- The proposed method is shown to be effective planner for solving multi-task problem on Maze2D and Meta-World benchmarks. ", "weaknesses": "No such weakness as such.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed feedback and for a positive assessment of our work! We carefully address your concerns as follow:\n\n>1. For Table H in appendix, it will be useful to report results for other baselines too.\".\n\nThanks for this good suggestion! Considering the space limitation during this rebuttal period, we would report the results in the next version of our paper. \n\n>2. It will be useful to see how the performance changes as a result of varying the generated data.\n\nThank you for this constructive suggestion! We re-run the experiments on coffee-push and disassemble tasks, expanding the original 1M dataset with 0.5M, 1M and 2M generated data for each task respectively. Then we continue to use TD3-BC as the downstream algorithm to measure the policy improvement. The result shows the performance increases along with more augmented data. Meanwhile, we find our method with only 0.5M generated data even outperforms RAD and S4RL which use 2M and 10M augmented data, respectively.\n\n\n| Tasks       | 2M Generated    |  1M Generated  | 0.5M Generated | Origin          |\n| ----------- | --------------- | :------------: | :------------: | --------------- |\n| coffee-push | $74.67\\pm 6.79$ | $70.66\\pm7.78$ | $65.43\\pm9.91$ | $28.60\\pm14.55$ |\n| disassemble | $69.00\\pm 4.72$ | $63.6\\pm4.62$  | $60.83\\pm4.40$ | $60.20\\pm16.29$ |\n\n>3. It will also be helpful to see how the performance changes as the MTDIFF-s is trained on less amount of tasks.\n\nThank you for this constructive suggestion! We re-train *MTDiff-s* on 30/20/10 tasks respectively,  and then measure the policy improvement on the coffee-push task across 3 seeds to verify our hypothesis that multi-task training enables implicit data sharing in *MTDiff-s*.\n\nOur findings, as outlined below, provide compelling evidence in support of our hypothesis: *MTDiff-s* exhibits progressively superior data synthesis performance with increasing task diversity.\n\n| Number of Taks | coffee-push     |\n| -------------- | --------------- |\n| 45             | $74.67\\pm 6.79$ |\n| 30             | $68.33\\pm9.00$  |\n| 20             | $66.53\\pm7.12$  |\n| 10             | $65.06\\pm 9.39$ |\n| 1              | $63.41\\pm8.99$  |\n"}, {"review_id": "udFMypDAQx", "reviewer": "Reviewer_nEaR", "summary": "The paper studies the use of diffusion models in offline multi-task reinforcement learning for planning and synthetic data generation. Both approaches use prompting to encode task-specific conditions for the generative model along with a transformer backbone. In the multitask setting, the approach outperforms prior SOTA baselines on MT50. In the synthetic data setting, the approach improves downstream TD3+BC performance.\n", "questions": "- How would the prompting perform if the tasks were not readily discernible from the initial transitions? E.g. multi-task settings with sparse task-specific reward \n- What is the minimum data required to make the algorithm work?\n- What is the speed of planning and sampling transitions with the transformer model? How does this compare to related methods?\n- How does MTDIFF-s-single compare to Synthetic Experience Replay? I.e. what is the benefit of modeling full trajectories instead of transitions?\n", "limitations": "Discussion of limitations hidden in Appendix, should be moved to the main paper.\n", "rating": 6, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- Clear and well-written presentation of the method\n- A thorough set of baselines and strong evaluation on the challenging MT50 benchmark for MTDIFF-p\n- Elegant method of removing the requirement of one-hot task encoding by transition prompting\n- Strong results showing positive multi-task transfer in both planning and synthetic data generation, the algorithm neatly extends single-task equivalents in both areas\n", "weaknesses": "- Only custom settings with large amounts of data are considered, it would be useful to understand what the minimum data needed to be effective is. Additionally, it would be useful to evaluate on pre-existing offline datasets for a more representative comparison.\n- In Figure 6, it would be useful to compare MTDIFF-s-single with the existing single task baseline, Synthetic Experience Replay. [1] \n- More thorough analysis of data quality would be valuable rather than just downstream performance, e.g. error of the transitions. Furthermore, downstream RL performance is only computed with one RL algorithm.\n- The two parts of the paper - multitask planning and synthetic data generation - are not necessarily connected. For example, the synthetic data (MTDIFF-s) is used for an entirely different RL algorithm, TD3+BC, and does not contribute to the performance of the planning algorithm (MTDIFF-p).\n\nMinor:\n- Line 242: Typo ‘fine-grind’\n- Line 248: the author’s definition of ‘near-optimal’ is the same as D4RL ‘full-replay’ and ‘sub-optimal’ is the same as D4RL ‘medium-replay’. These descriptions may be more clear to offline RL practitioners\n- Figure 4: choosing one sampled trajectory for each approach is likely cherry-picking, I would suggest rendering N samples\n- Figure 5: missing standard deviation\n- Figure 6: missing the baseline without synthetic data\n- Line 15: Unclear what ‘high-quality’ and ‘low-quality’ mean as it seems MTDIFF-s models the original distribution\n\n[1] Synthetic Experience Replay. Cong Lu, Philip J. Ball, Yee Whye Teh, Jack Parker-Holder.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ">1. ... custom settings with large amounts of data ... minimum data needed to be effective is.\n\nWe downsample the \"near-optimal dataset\" to 0.1$\\times$/ 0.2$\\times$/ 0.3$\\times$ the size via random selection. We observed that the performance of *MTDiff-p* decreases with dataset reduced, dropping at $36.06\\%\\pm2.36\\%$ success rate for 0.1$\\times$ size data, which can be caused by missing parts of optimal trajectories. To validate that, we re-train *MTDiff-p* on 0.1$\\times$ size data containing *only* expert trajectories. In this case, *MTDiff-p* can even obtain $85.05\\%\\pm1.54\\%$ success rate. It is reasonable that the most critical factor for goal-conditional planning is the optimality of the dataset instead of data quantity, which has also been verified in [1]. In our paper, we consider the more challenging setting, i.e., using data with low return accounting for majority of the dataset.\n\n>2. evaluate on pre-existing offline datasets for a more representative comparison.\n\nOur method aims to learn a diffusion model from multi-task offline data. However, pre-existing datasets like D4RL only contain single-task datasets in each domain. As a result, we collect multi-task datasets for evaluation. Meta-world is chosen since it contains 50 various and difficult tasks, and it is well-known for its efficiency. For reproductivity, we give the code in supplementary materials and the link to our dataset in Appendix. We remark that we added new experiments on Mujoco and *MTDiff-p* still achieves the SOTA performance, where the details can be seen in our response to Reviewer *gBBv*.\n\n>3. In Figure 6, it would be useful to compare MTDIFF-s-single with the existing single task baseline.\n\nWe remark that SynthER [2] and MTDiff are almost concurrent work. Fig. 6 is used to support our hypothesis that *MTDiff-s* benefits from multi-task training instead of how *MTDiff-s-single* outperforms other methods. To compare *MTDiff-s-single* with SynthER, we borrowed the open-sourced code and trained SynthER on coffee-push, disassemble, and dial-turn tasks. We added the result to Fig. 6, and the revised vision is referred to Fig. 2 in **Global PDF**. We find that SynthER still underperforms *MTDiff-s*.\n\n>4. More thorough analysis of data quality ... downstream RL performance is only computed with one RL algorithm.\n\nWe refer to **Global PDF** for the new experiments and explanations. \n\n>5. ...multitask planning and synthetic data generation are not necessarily connected...\n\nThanks for your question. Indeed, MTDIFF-s can augment the multi-task dataset to benefit MTDIFF-p, while it can be biased in comparison. First, MTDIFF-p is a generative planning method and its performance is closely connected to the data optimality rather than the data coverage (see [1] and R1). MTDIFF-s is designed to improve the data coverage and it is more suitable to use offline RL for evaluation. Second, for a specific task, even MTDIFF-p's performance improves with data augmentation by MTDIFF-s, it can be caused by the augmented data of other tasks since MTDIFF-p can conduct implicit knowledge sharing among tasks. As a result, we conduct single-task augmentation and use offline RL for evaluation.\n\n\n>6. Minor error\n* We will fix the typo in the next version. \n* We remark that definition of \"near-optimal\" and \"sub-optimal\" comes from paper [3], and we would mention replay and medium-replay based on your suggestion.\n* We show more rendered cases for Maze2D in Appendix D, however, we added another example (see Fig. 4 in **Global PDF**) to demonstrate the effectiveness of our method.\n* The corrective Fig. 5 is referred to Fig. 1 in **Global PDF**.\n* We added the results of the baseline without synthetic data into Fig. 6, which is referred to Fig. 2 in **Global PDF**.\n* A good quality means the generated data has (1) low dynamics errors to follow the true dynamics of the original task, and (2) extended data coverage to enlarge the dataset. Fig. 4 in appendix shows the synthetic data expands the coverage of the original data, thus boosting the performance of offline RL.\n\n>6. How would the prompting perform if the tasks were not readily discernible from the initial transitions?\n\nThanks for the question! It is possible that the tasks cannot be identified from their initial transitions. Nevertheless, the prompt contains both the initial transitions and the task-specific label $Z$, where $Z=(s^*_i,a^*_i,\\ldots,s^*\\_{i+J-1},a^*\\_{i+J-1})$ is the trajectory prompt sampled from an expert trajectory to provide task-identified information. $Z$ is injected into the model as a condition during both training and testing.\n\n>7. ...speed of planning and sampling ...? How does this compare to related methods?\n\nIt is known that the sampling speed of diffusion model according to the diffusion process can be quite slow, which we have discussed in Appendix F. For a concrete example in MetaWorld, it takes on average 1.9s in wall-clock time to generate one action sequence for planning (hardware being a 3090 GPU). We can improve the inference speed by leveraging a recent sampler called DPM-solver [4] to decrease the diffusion steps required to 0.2$\\times$ without any loss in performance, and using a larger batch size (leveraging the parallel computing power of GPUs) to evaluate multiple environments at once. Thus the evaluation run-time roughly matches the runtime of non-diffusion algorithms (diffusion step is 1). We will add an section to discuss the sampling speed of our method in Appendix in the next version.\n\n>8. Discussion of limitations ... moved to main paper.\n\nWe would move it into the Conclusion section in the next version.\n\n[1] Rethinking Goal-Conditioned Supervised Learning and Its Connection to Offline RL. ICLR 2022\n\n[2] Synthetic Experience Replay. ICLRW 2023\n\n[3] Offline Q-learning on Diverse Multi-Task Data Both Scales And Generalizes. ICLR 2023\n\n[4] DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps. NeurIPS 2022"}, {"review_id": "R07Q7lELvd", "reviewer": "Reviewer_gBBv", "summary": "This paper proposed Multi-Task Diffusion Model, a diffusion-based method that incorporates Transformer backbones and prompt learning for generative planning and data synthesis in multitask offline settings. The performance of the proposed model on Meta-World and Maze2D benchmarks was shown.", "questions": "# PromptDT showed performance on MuJoCo. How is the performance of the proposed model on MuJoCo?\n# This paper seemed to replace transformer in PromptDT with diffusion model. It’s better to explain the difference between the proposed model and PromptDT.", "limitations": "The related work was introduced in Section 4. It would be clearer to introduce related works before Section 2 Preliminaries.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "This paper is the first to achieve both effective planning and data synthesis for multi-task RL via diffusion model and GPT.\nThis paper is well written and easy to follow.\n", "weaknesses": "# The difference between the proposed model and PromptDT is not well explained.\n# Additional environment and ablation experiments may be more convincing.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable comments and a positive assessment of our work! We are glad that you find our paper easy to follow, well-written, and the first to achieve both effective planning and data synthesis for multi-task RL. To address your concerns, we have added additional experiments on MuJoCo to validate the efficacy of our method. Our detailed response follows:\n\n>1. The difference between the proposed model and PromptDT is not well explained.\n\nThank you for this good question! We would enrich the appendix with a dedicated section discussing their differences in the next version of our paper. The remarkable superiority of *MTDiff-p* over PromptDT emerges from our elegant incorporation of transformer architecture and trajectory prompt within the **diffusion model framework**, effectively modeling the multi-task trajectory distribution. PromptDT is built on Decision Transformer and it is trained in an autoregressive manner, which is limited to predicting actions step by step. However, *MTDiff-p* leverages the potency of sequence modeling, empowering it to adeptly perform trajectory generation. *MTDiff-p* has demonstrated SOTA performance in both multi-task decision-making and data synthesis empirical experiments, while PromptDT fails to contribute to data synthesis. Technically, *MTDiff-p* extends Decision Diffuser [1] into the multi-task scenario, utilizing classifier-free guidance for generative planning to yield high expected returns, which is also recognized by Reviewer EBNs. \n\n[1] Anurag Ajay, et al. Is conditional generative modeling all you need for decision making? International Conference on Learning Representations, 2023.\n\n>2. PromptDT showed performance on MuJoCo. How is the performance of the proposed model on MuJoCo?\n\nWe have added new experiments to compare our method with PromptDT on the MuJoCo benchmark. We trained our model on the publicly available PromptDT datasets, i.e., *Cheetah-vel* and *Ant-dir*. These chosen environments have been judiciously selected due to their inherent diversity of tasks, serving as a robust test to validate the capability of multi-task learning. We report the scores (mean and std for 3 seeds) as follows:\n\n| Methods  | Cheetah-vel      | Ant-dir          |\n| -------- | ---------------- | ---------------- |\n| MTDiff-p | $-29.09\\pm 0.31$ | $602.17\\pm 1.68$ |\n| PromptDT | $-34.43\\pm 2.33$ | $409.81\\pm 9.69$ |\n\nWe observed that *MTDiff-p* outperforms PromptDT largely, demonstrating its high efficacy and potency."}]

Author rebuttal:
### General Response

We thank all of the reviewers for their time and insightful comments. Furthermore, we are very glad to find that reviewers generally recognized our key contributions and clear presentation of our paper:

#### Contributions:

* **Method:** "This paper is the first to achieve both effective planning and data synthesis for multi-task RL via diffusion model and GPT" [gBBv]. The proposed method is shown to be effective and positive, which extends previous single-task models to solving more complex multi-task problems [UScg, nEaR]. "Leveraging the power of prompting instead of one-hot task encoding is a cool and elegant idea" [8zHR, nEaR]. "The analysis of how MTDIFF benefits from multi-tasking in data synthesis is interesting" [8zHR]. 
* **Experiment:** The paper does a good job comparing various strong baselines, and shows decent improvements over all the baselines [UScg]. The experiments conducted demonstrate a marked improvement in performance [EBNs]. The paper shows a thorough set of baselines and the approach outperforms prior SOTA baselines [nEaR].
* **Presentation:** The paper is well written [UScg, gBBv]. The paper shows clear and well-written presentation of the method [nEaR]. The paper is easy to follow [EBNs, gBBv].

Meanwhile, we thank all the reviewers for their helpful and constructive feedback to improve the quality of our work. In addition to the pointwise responses below, we would carefully update our paper in the next version to incorporate the valuable suggestions of the reviewers:

* [**UScg**] We would add the experimental results about varying the generated data of *MTDiff-s* and varying the trained tasks of *MTDiff-s* (see the response to Reviewer UScg).
* [**UScg**] We would report performance for baselines on each single task in MT50-v2.
* [**nEaR**] We would add one more unseen Maze2D map where *MTDiff-p* succeeds while PromptDT fails (see Fig. 4 in global PDF).
* [**nEaR**] We would update Figure 6 to compare with more baselines (see Fig. 2 in global PDF).
* [**nEaR**] We would polish some of our expressions (e.g. definition) and add a section to discuss the sampling speed of our model (see the response to Reviewer nEaR).
* [**nEaR**] We would add experimental results about data quality analysis and IQL performance improvement (see Table 1 and Table 2 in the global PDF).
* [**EBNs**] We would revise the subtitle the line 334 (see the response to Reviewer EBNs).
* [**EBNs**] We would add quantitative results on Maze2D unseen map (see the response to Reviewer EBNs).
* [**8zHR, EBNs, nEaR**] We would update Figure 5 (see Fig. 1 in global PDF).
* [**gBBv**] We would add a section in the appendix to demonstrate the difference between our model and PromptDT, and add the experimental results on MuJoCo benchmark (see the response to Reviewer gBBv).

We hope to have addressed all the raised concerns and would be happy to respond to further questions and suggestions.

Meta-review: The paper proposes Multi-Task Diffusion Model (MTDiff), a method that integrates Transformers and diffusion models for generative planning and data synthesis in multi-task offline reinforcement learning settings. It demonstrates the efficacy of MTDiff in a variety of benchmarks, including Meta-World and Maze2D, for both planning and synthetic data generation tasks.

Reviewers generally commend the well-executed integration of Transformers and diffusion models, along with the innovative use of prompting in generative planning. The paper also stands out for its clear writing and strong evaluation against various baselines. However, there are areas for improvement such as the unpolished initial version, as pointed out in the figure and table issues. A more thorough analysis on data quality would also add value, and there is a disconnect between the paper's two main components - multi-task planning and synthetic data generation. Some of these weaknesses have been addressed during the rebuttal stage. Given these considerations, the paper is recommended for acceptance with the understanding that the authors will make the necessary revisions.
Decision: Accept (poster)