Title: Diffusion Model is an Effective Planner and Data Synthesizer for Multi-Task Reinforcement Learning
Venue: NeurIPS 2023
Accepted: yes
Keywords: multi-task reinforcement learning

diffusion models

planning

data synthesis
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

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 (MTDIFF), a diffusion-based method that incorporates Transformer backbones and prompt learning for generative planning and data synthesis in multitask offline settings. MTDIFF leverages vast amounts of knowledge available in multi-task data and performs implicit knowledge sharing among tasks. For generative planning, we 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 a single demonstration as a prompt, which enhances the low-quality datasets for even unseen tasks.

1 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), is that the learned model can generalize among different tasks by simply conditioning the model on instructions or prompts [49, 46, 7, 12, 60, 42, 66]. The success of LLMs and vision models inspires us to 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-making 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, 11, 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-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 capabilities. Gato [47] and other generalized agents [29, 65] take transformer-based architecture

𝐴

Multi-task actions

𝐴

𝐴

𝐴

𝐴

𝑅

𝑆

𝐴

𝑆

𝑅

Multi-task transitions

Optimal actions

Diffusion Model

Conditional Generative Process

𝑆

𝑅

R

𝐴

𝑆

𝐴

High fidelity synthetic transitions

Act

Augment

Multi-Task Dataset

(a) Planning

(b) Synthesizing

𝑆prev , 𝑅

𝜏

𝐴

𝑅

𝑆

𝐴

𝑆

𝑅

𝑆

𝑅

R

𝐴

𝑆

𝐴

𝐴

𝐴

𝐴

𝐴

Prompt

Environment

Original Data

[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 to further explore its generalizability in a multi-task setting. We formulate the learning process from multi-task data as a denoising problem, which benefits the modeling of multimodal data. Meanwhile, we develop a 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 prompt 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, we investigate two variants of MTDIFF for planning and data synthesis to further exploit the utility of diffusion models, denoted as MTDIFF-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 to 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 to 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, is more effective and generalist than previous methods that only utilize single-task data for augmentation [38, 52, 28]. Specifically, MTDIFF-S learns to estimate the joint conditional distribution of the full transitions that contain states, actions, and rewards based on the task-oriented 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 diffusion-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: (i) we propose MTDIFF, a 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 to learn to generalize across different tasks and even adapt to unseen tasks; (iii) our experiments on Meta-World and Maze2D benchmarks demonstrate that MTDIFF is an effective planner to solve the multi-task problem, and also a powerful data synthesizer to augment offline datasets in the seen or unseen tasks.

2 Preliminaries

2.1 Reinforcement Learning

MDP and Multi-task MDP. A Markov Decision Process (MDP) is defined by a tuple ( S , A , P , R , µ, γ ) , where S is the state space, A is the action space, P : S × A → S is the transition function, R : S × A × S → R is the reward function for any transition, γ ∈ (0 , 1] is a discount factor, and µ is the initial state distribution. At each timestep t , the agent chooses an action a t by following the policy π : S → ∆ A . Then the agent obtains the next state s t +1 and receives a scalar reward r t . In single-task RL, the goal is to learn a policy π ∗ = arg max π E a t ∼ π [ ∑ ∞ t =0 γ t r t ] by maximizing the expected cumulative reward of the corresponding task.

In a 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 ∼ p ( T ) , a task-specified MDP can be defined as ( S T , A , P T , R T , µ T , γ ) . Instead of solving a single MDP, the goal of multi-task RL is to find an optimal policy that maximizes expected return over all the tasks: π ∗ = arg max π E T ∼ p ( T ) E a t ∼ π T [ ∑ ∞ t =0 γ t r T t ] .

Multi-Task Offline Decision-Making. In offline decision-making, the policy is learned from a static dataset of transitions { ( s j , a j , s ′ j , r j ) } N j =1 collected by an unknown behavior policy π β [30]. In the multi-task offline RL setting, the dataset D is partitioned into per-task subsets as D = ∪ N i =1 D i , where D i consists of experiences from task T i . The key issue of RL in the offline setting is the distribution shift problem caused by temporal-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 a 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, s ′ , 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 = ∪ N i =1 D i in this paper. With τ the sampled trajectory from D , we denote x k ( τ ) as the k -step denoised output of the diffusion model, and y ( τ ) is the condition which represents the task attributes and the trajectory's optimality (e.g., returns). A forward diffusion chain gradually adds noise to the data x 0 ( τ ) ∼ q ( x ( τ )) in K steps with a pre-defined variance schedule β k , which can be expressed as

In this paper, we adopt Variance Preserving (VP) beta schedule [67] and define β k = 1 -exp ( -β min ( 1 K ) -0 . 5( β max -β min ) 2 k -1 K 2 ) , where β max = 10 and β min = 0 . 1 are constants. A trainable reverse diffusion chain, constructed as p θ ( x k -1 ( τ ) | x k ( τ ) , y ( τ )) := N ( x k -1 ( τ ) | µ θ ( x k ( τ ) , y ( τ ) , k ) , Σ k ) , can be optimized by a simplified surrogate loss [20]:

where ϵ θ parameterized by a deep neural network is trained to predict the noise ϵ ∼ N ( 0 , I ) added to the dataset sample x 0 ( τ ) to produce x k ( τ ) . By setting α k := 1 -β k and ¯ α k := ∏ k s =1 α s , we obtain

Classifier-free guidance [19] aims to learn the conditional distribution q ( x ( τ ) | y ( τ )) without separately training a classifier. In the training stage, this method needs to learn both a conditional ϵ θ ( x k ( τ ) , y ( τ ) , k ) and an unconditional ϵ θ ( x k ( τ ) , ∅ , k ) model, where y ( τ ) is dropped out. Then the perturbed noise ϵ θ ( x k ( τ ) , ∅ , k ) + α ( ϵ θ ( x k ( τ ) , y ( τ ) , k ) -ϵ θ ( x k ( τ ) , ∅ , k )) is used to generate samples latter, where α can be recognized as the guidance scale.

Bypass a wall and press a button

Slide a plate into a cabinet

Close a door

Rotate the faucet counter-clockwise

Reverse Denoising Process

3 Methodology

3.1 Diffusion Formulation

To capture the multimodal distribution of the trajectories sampled from multiple MDPs, we formulate the multi-task trajectory modeling as a conditional generative problem via diffusion models:

where x 0 ( τ ) is the generated desired sequence and y ( τ ) is the condition. x 0 ( τ ) will then be used for generative planning or data synthesis through conditional reverse denoising process p θ for specific tasks. Maximizing Eq. (3) can be approximated by maximizing a variational lower bound [20].

In terms of different inputs and outputs in generative planning and data synthesis, x ( τ ) can be represented in different formats. We consider two choices to formulate x ( τ ) in MTDIFF-P and MTDIFF-S, respectively. (i) For MTDIFF-P , x ( τ ) represents the action sequence for planning. We model the action sequence defined as:

with the context condition as

where t , H , R ( τ ) and L denote the time visited in trajectory τ , the length of the input sequence x , the normalized cumulative return under τ and the length of the observed state history, respectively. Z is the task-relevant information as prompt . We use y ′ ( τ ) as an ordinary condition that is injected into the model during both training and testing, while considering R ( τ ) as 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:

with the condition:

where y s ( τ ) takes the same conditional approach as y ′ ( τ ) . Figure 2 illustrates the reverse denoising process of MTDIFF-P learned on multi-task datasets collected in Meta-World [73]. The result demonstrates that our diffusion model successfully distinguishes different tasks and finally generates the desired x 0 ( τ ) . We illustrate the data distribution of x 0 ( τ ) in a 2D space with dimensional reduction via T-SNE [61], as well as the 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 of other tasks, which verifies that MTDIFF can learn the distribution of multimodal trajectories based on y ( τ ) .

3.2 Prompt, Training and Sampling

In multi-task RL and LLM-driven decision-making, existing works use one-hot task identifiers [56, 74] or language descriptions [1, 6] as conditions in multi-task training. Nevertheless, we argue that the one-hot encoding for each task [22, 14] suffices for learning a repertoire of training tasks while cannot generalize to novel tasks since it does not leverage semantic similarity between tasks. In addition, the language descriptions [51, 1, 6, 41, 40] of tasks require large amounts of human labor to annotate and encounter challenges related to ambiguity [72]. In MTDIFF, we use expert demonstrations consisting of a few trajectory segments to construct more expressive prompts in multitask 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 a similar method has also been used in PromptDT [69]. Nonetheless, how such a prompt contributes within a diffusion-based framework remains to be investigated.

Specifically, we formulate the task-specific label Z as trajectory prompts that contain states and actions:

where each element with star-script is associated with a 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 to devise the optimal behaviors that maximize return. Our approach is to utilize the diffusion model for action planning via classifier-free guidance [19]. Formally, an optimal action sequence x p 0 ( τ ) is sampled by starting with Gaussian noise x K ( τ ) and refining x p k ( τ ) into x p k -1 ( τ ) at each intermediate timestep with the perturbed noise:

where y ′ ( τ ) is defined in Eq. (5). R ( τ ) is the normalized return of τ , and α is a 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] as well as classifier-free guidance [19] to train the reverse diffusion process p θ , parameterized through the noise model ϵ θ , with the following loss:

Note that with probability p sampled from a Bernoulli distribution, we ignore the conditioning return R ( τ ) . During inference, we adopt the low-temperature sampling technique [2] to produce high-likelihood sequences. We sample x p k -1 ( τ ) ∼ N ( µ θ ( x p k -1 , y ′ ( τ ) , R max ( τ ) , k -1) , β Σ k -1 ) , where the variance is reduced by β ∈ [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 ( τ ) , we have the following loss:

We sample x s k -1 ( τ ) ∼ N ( µ θ ( x s k -1 , y s ( τ ) , k -1) , Σ k -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, 5, 47, 6] provides a promising solution to capture interactions between modalities of different tasks. Naturally, instead of U-Net [48] which is commonly used in previous single-task diffusion RL works [21, 2, 33], we parameterize ϵ θ with a 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 in a unified architecture, which is expected to enhance the efficiency of diverse information exploitation.

!

∗

"

∗

"

∗

!

∗

!

∗

"

∗

!

k

MLP

1

1

2

2

$

$

!

"

"

"

! - # + 1

!

!

! + 1

! + '

!

"

%

!

!

!

"

! + '

!

! + '

Actions

Prediction Head

Predicted noise

Positional Encoding

GPT2 Transformer

:

:

:

:

:

Inputs

Timestep

Transitions

History

Prompt

+

MLP

MLP

Actions*

Transitions*

!

!

!

!

Layer Norm

!

!

!

"

Few-shot Prompt

(

!"#

)

$%&'(:$

!

Training

Evaluation

#

:

Return

+

Denoising Process

MTD

IFF

-

S

MTD

IFF

-

P

MTD

IFF

-

S

MTD

IFF

-

P

MTD

IFF-S

MTD

IFF-P

As shown in Figure 3, first, different raw inputs x are embedded into embeddings h of the same size d via separate MLPs f , which can be expressed as:

Then, the embeddings h P and h Ti are prepended as follows to formulate input tokens for MTDIFF-P and MTDIFF-S, respectively:

where E pos is the positional embedding, and LN denotes layer normalization [3] for stabilizing training. In our implementation, we strengthen the condition of the stacked inputs through multiplication with the diffusion timestep h Ti and addition with the return h p R . GPT2 is a decoder-only transformer that incorporates a self-attention mechanism to capture dependencies between different positions in the input sequence. We employ the GPT2 architecture as a trainable backbone in MTDIFF to handle sequential inputs. It outputs an updated representation as:

Finally, given the output representation, we use a 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 as the original inputs, which differs from the representation size d . This noise is used in the reverse denoising process p θ during inference. We summarize the details of the training process, architecture and hyperparameters used in MTDIFF in Appendix A.

4 Related Work

Diffusion Models in RL. Diffusion models have emerged as a powerful family of deep generative models with a record-breaking performance in many applications across vision, language and combinatorial optimization [49, 46, 17, 31, 32]. Recent works in RL have demonstrated the capability of diffusion models to learn the multimodal distribution of offline policies [63, 43, 11] or human behaviors [8]. Other works formulate the sequential decision-making problem as a conditional generative 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.

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 [34], and parameter composition [70, 56]. Instead, MTDIFF addresses such a problem in an offline setting through a conditional generative process via a novel transformer architecture. Previous Decision-Transformer (DT)-based 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 [77, 79] or gradient descents in the finetuning stage [57, 58, 29]. In contrast, we use prompts for few-shot generalization without additional parameter-tuning.

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 RL [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 [76, 10] 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 [38] 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 a diffusion model can accommodate all knowledge of multi-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 MTDIFF-P compare to other offline and online baselines in the multi-task regime? (2) Does MTDIFFS synthesize high-fidelity data and bring policy improvement? (3) How is MTDIFF-S compared with other augmentation methods for single-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 [73] contains 50 qualitatively-distinct manipulation tasks. The tasks share similar dynamics and require a Sawyer robot to 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 different 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 [70, 56], we extend all the tasks to a 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 [18] 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 a lot. We summarize more details about the dataset in Appendix G.

Maze2D Tasks Maze2D [15] is a navigation task that requires an agent to traverse from a randomly designated location to a fixed goal in the 2D map. The reward is 1 if succeed and 0 otherwise. Maze2D can evaluate the ability of RL algorithms to stitch together previously collected sub-trajectories, which helps the agent find the shortest path to evaluation goals. We use the agent's scores as the evaluation metric. The offline dataset is collected by selecting random goal locations and using a planner to generate sequences of waypoints by following a PD controller.

Baselines We compare our proposed MTDIFF (MTDIFF-P and MTDIFF-S) with the following baselines. Each baseline has the same batch size and training steps as MTDIFF. For MTDIFF-P , we

have following baselines: (i) PromptDT. PromptDT [69] built on Decision-Transformer (DT) [9] 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 PromptDT is that we employ diffusion models for generative planning. (ii) MTDT. We extend the DT architecture [9] to learn from multitask data. Specifically, MTDT concatenates an embedding z and a state s as the input tokens, where z 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 a task-ID conditioned actor for multi-task policy learning. (iv) MTIQL. 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 offline policy learning via network scaling and a task-ID conditioned actor that is similar to 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 a random variable to states, i.e., s ′ = s × z , where z ∼ Uniform[ α, β ] . 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.e. s ′ ← s + ϵ ∇ s J Q ( π ( s )) , where J Q is the policy evaluation update performed via a 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,

Table 1: Average success rate across 3 seeds on MetaWorld-V2 MT50 with random goals (MT50-rand). Each task is evaluated for 50 episodes.

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-ONEHOT perform the best in 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-based baselines, 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.

we add MTDIFF-P-ONEHOT as a variant of MTDIFF-P by replacing the prompt with a one-hot task-ID, which is used in the baselines except for PromptDT. The first action generated by MTDIFF-P is used to interact with the environment. According to Tab. 1, we have the following key observations. (i) Our method achieves better performance than baselines in both near-optimal and sub-optimal settings. For near-optimal datasets, MTDIFF-P and MTDIFF-P-ONEHOT achieve about 60% success rate, significantly outperforming other methods and performing comparably with MTBC. However, the performance of MTBC decreases a lot in sub-optimal datasets. BC is hard to handle the conflict behaviors in experi-

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 MTDIFF-P-ONEHOT performs well, it cannot generalize to unseen tasks without prompts.

(a) MTDIFF-P

(b) PromptDT

(c) MTDIFF-P

(d) PromptDT(failed)

(1) Seen map

(2) Unseen map

MTDiff-p

PromptDT

0

20

40

60

80

100

Scores

Diffusion Model is an Effective Planner and Data Synthesizer for Multi-Task RL Anonymous Author(s) Affiliation Address email 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 on both seen maps and an unseen 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 seen maps in training, MTDIFF-P generates a shorter and smoother path, and (2) for unseen maps, PromptDT fails to obtain a reasonable path while MTDIFF-P succeed, which verifies that MTDIFF-P can perform few-shot adaptation based on trajectory prompts and the designed architecture.

5.3 Results for Augmentation via Data Synthesis

Does MTDIFF-S synthesize high-fidelity data and bring policy improvement? We train MTDIFFS on near-optimal datasets from 45 tasks to evaluate its generalizability. We select 3 training tasks and 3 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 1M dataset. From the results summarized in Table 2, MTDIFF-S can boost the offline performance for all tasks and significantly increases performance by about 180%, 131%, and 161% for box-close , hand-insert , and coffee-push , respectively.

Submitted to 37th Conference on Neural Information Processing Systems (NeurIPS 2023). Do not distribute.

coffee-push

disassemble

dial-turn

0

10

20

30

40

50

60

70

80

Success Rate (%)

MTDiff-s-single

MTDiff-s

SynthER

Original



























How does MTDiff-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 in 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-SINGLE. We implement SynthER [38] to further confirm the benefits of multi-task training. SynthER also utilizes diffusion models to enhance offline datasets for policy improvement, but it doesn't 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. 7, 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 wellgeneralized 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 as prompts) for augmentation.

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, a 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 to 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 a single model, we demonstrate that MTDIFF-S greatly boosts policy performance by augmenting original offline datasets. In future research, we aim to develop a 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 F.

Acknowledgments

This work is supported by the National Natural Science Foundation of China (Grant No.62306242&62076161), the National Key R&D Program of China (Grant No.2022ZD0160100), Shanghai Municipal Science and Technology Major Project (2021SHZDZX0102) and Shanghai Artificial Intelligence Laboratory.

References

Michael Ahn, Anthony Brohan, Noah Brown, Yevgen Chebotar, Omar Cortes, Byron David, Chelsea Finn, Keerthana Gopalakrishnan, Karol Hausman, Alex Herzog, et al. Do as i can, not as i say: Grounding language in robotic affordances. arXiv preprint arXiv:2204.01691 , 2022.

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.

Lei Jimmy Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization. CoRR , abs/1607.06450, 2016. URL http://arxiv.org/abs/1607.06450 .

Shaojie Bai, J Zico Kolter, and Vladlen Koltun. An empirical evaluation of generic convolutional and recurrent networks for sequence modeling. arXiv preprint arXiv:1803.01271 , 2018.

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.

Anthony 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.06817 , 2022.

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.

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 on Learning Representations , 2023.

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 https://openreview.net/ forum?id=a7APmM4B9d .

Zoey Chen, Sho Kiami, Abhishek Gupta, and Vikash Kumar. Genaug: Retargeting behaviors to unseen situations via generative augmentation. arXiv preprint arXiv:2302.06671 , 2023.

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.

Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311 , 2022.

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.

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.org/abs/2109.13396 .

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: //arxiv.org/abs/2004.07219 .

Scott Fujimoto and Shixiang Shane Gu. A minimalist approach to offline reinforcement learning. In Thirty-Fifth Conference on Neural Information Processing Systems , 2021.

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 https://openreview.net/forum?id=jQj-_ rLVXsj .

Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Offpolicy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning , pages 1861-1870. PMLR, 2018.

Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. In NeurIPS 2021 Workshop on Deep Generative Models and Downstream Applications , 2021. URL https://openreview. net/forum?id=qw8AKxfYbI .

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems , 33:6840-6851, 2020.

Michael Janner, Yilun Du, Joshua B. Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. In International Conference on Machine Learning , 2022.

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 at scale. CoRR , abs/2104.08212, 2021. URL https://arxiv.org/ abs/2104.08212 .

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learning. In International Conference on Learning Representations , 2022. URL https: //openreview.net/forum?id=68n2s9ZJWF8 .

Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. CoRR , abs/2006.04779, 2020. URL https://arxiv.org/ abs/2006.04779 .

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.

Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Curl: Contrastive unsupervised representations for reinforcement learning. In International Conference on Machine Learning , pages 5639-5650. PMLR, 2020.

Misha Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Reinforcement learning with augmented data. Advances in neural information processing systems , 33:19884-19895, 2020.

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 in Neural Information Processing Systems , 35:27921-27936, 2022.

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 https://arxiv.org/abs/2005.01643 .

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 in Neural Information Processing Systems , 2022. URL https://openreview.net/forum?id=3s9IrEsjLyk .

Yang Li, Jinpei Guo, Runzhong Wang, and Junchi Yan. From distribution learning in training to gradient search in testing for combinatorial optimization. In Advances in Neural Information Processing Systems , 2023.

Zhixuan Liang, Yao Mu, Mingyu Ding, Fei Ni, Masayoshi Tomizuka, and Ping Luo. Adaptdiffuser: Diffusion models as adaptive self-evolving planners. ArXiv , abs/2302.01877, 2023.

Bo Liu, Xingchao Liu, Xiaojie Jin, Peter Stone, and Qiang Liu. Conflict-averse gradient descent for multi-task learning. Advances in Neural Information Processing Systems , 34:18878-18890, 2021.

Minghuan Liu, Hanye Zhao, Zhengyu Yang, Jian Shen, Weinan Zhang, Li Zhao, and Tie-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://openreview.net/forum?id=q6Kknb68dQf .

Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpmsolver++: Fast solver for guided sampling of diffusion probabilistic models. arXiv preprint arXiv:2211.01095 , 2022.

Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. DPM-solver: A 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.

Cong Lu, Philip J. Ball, and Jack Parker-Holder. Synthetic experience replay. In Workshop on Reincarnating Reinforcement Learning at ICLR 2023 , 2023. URL https://openreview. net/forum?id=0a9p3Ty2k_ .

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:2210.04561 , 2022.

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.1109/LRA.2022.3196123.

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 Letters , 7(3):7327-7334, 2022. doi: 10.1109/LRA.2022.3180108.

OpenAI. Gpt-4 technical report, 2023.

Tim Pearce, Tabish Rashid, Anssi Kanervisto, Dave Bignell, Mingfei Sun, Raluca Georgescu, 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.

William Peebles and Saining Xie. Scalable diffusion models with transformers. arXiv preprint arXiv:2212.09748 , 2022.

Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog , 1(8):9, 2019.

Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents, 2022.

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.06175 , 2022.

Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical Image Computing and Computer-Assisted Intervention-MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18 , pages 234-241. Springer, 2015.

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.

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 Szczechla, 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.

Mohit Shridhar, Jesse Thomason, Daniel Gordon, Yonatan Bisk, Winson Han, Roozbeh Mottaghi, Luke Zettlemoyer, and Dieter Fox. ALFRED: A Benchmark for Interpreting Grounded Instructions for Everyday Tasks. In The IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , 2020. URL https://arxiv.org/abs/1912.01734 .

Samarth Sinha, Ajay Mandlekar, and Animesh Garg. S4rl: Surprisingly simple self-supervision for offline reinforcement learning in robotics. In Conference on Robot Learning , pages 907-917. PMLR, 2022.

Shagun Sodhani, Amy Zhang, and Joelle Pineau. Multi-task reinforcement learning with contextbased representations. In International Conference on Machine Learning , pages 9767-9779. PMLR, 2021.

Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised 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 https://proceedings.mlr.press/v37/sohl-dickstein15.html .

Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models, 2023.

Lingfeng Sun, Haichao Zhang, Wei Xu, and Masayoshi Tomizuka. Paco: Parametercompositional multi-task reinforcement learning. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022.

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.

Adrien Ali Taiga, Rishabh Agarwal, Jesse Farebrother, Aaron Courville, and Marc G Bellemare. Investigating multi-task pretraining and generalization in reinforcement learning. In The Eleventh International Conference on Learning Representations , 2023. URL https: //openreview.net/forum?id=sSt9fROSZRO .

Denis Tarasov, Alexander Nikulin, Dmitry Akimov, Vladislav Kurenkov, and Sergey Kolesnikov. CORL: Research-oriented deep offline reinforcement learning library. In 3rd Offline RL Workshop: Offline RL as a 'Launchpad' , 2022. URL https://openreview.net/forum? id=SyAS49bBcv .

Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971 , 2023.

Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine learning research , 9(11), 2008.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

Zhendong Wang, Jonathan J Hunt, and Mingyuan Zhou. Diffusion policies as an expressive policy class for offline reinforcement learning. In The Eleventh International Conference on Learning Representations , 2023.

Jason Wei, Maarten Bosma, Vincent Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, Andrew M. Dai, and Quoc V Le. Finetuned language models are zero-shot learners. In International Conference on Learning Representations , 2022.

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-making in the real world: A foundation decision model perspective. arXiv preprint arXiv:2212.12669 , 2022.

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.

Zhisheng Xiao, Karsten Kreis, and Arash Vahdat. Tackling the generative learning trilemma with denoising diffusion gans. In International Conference on Learning Representations , 2022.

Haoran Xu, Xianyuan Zhan, Honglei Yin, and Huiling Qin. Discriminator-weighted offline 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.

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 Machine Learning Research , pages 24631-24645. PMLR, 17-23 Jul 2022. URL https: //proceedings.mlr.press/v162/xu22g.html .

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.

Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. In International Conference on Learning Representations , 2021. URL https://openreview.net/forum?id=GY6-6sTvGaf .

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 https://openreview.net/forum?id=4u42KCQxCn8 .

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. URL https://arxiv.org/abs/ 1910.10897 .

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.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems , volume 33, pages 5824-5836. Curran Associates, Inc., 2020. URL https://proceedings.neurips.cc/ paper_files/paper/2020/file/3fe78a8acf5fda99de95303940a2420c-Paper.pdf .

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.

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 arXiv:2302.11550 , 2023.

Haoqi Yuan and Zongqing Lu. Robust task representations for offline meta-reinforcement learning via contrastive learning. In International Conference on Machine Learning , pages 25747-25759. PMLR, 2022.

Haichao Zhang, Wei Xu, and Haonan Yu. Generative planning for temporally coordinated exploration in reinforcement learning. In International Conference on Learning Representations , 2022. URL https://openreview.net/forum?id=YZHES8wIdE .

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 metalearning. In International Conference on Learning Representations , 2020.

A The Details of MTDIFF

In this section, we give the pseudocodes of MTDIFF-P and MTDIFF-S in Alg. 1 and Alg. 2, respectively. 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 2 e -4 learning rate for 2 e 6 train steps.

We train MTDIFF 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 f P which processes prompt is a 3-layered MLP (prepended by a layer norm [3] and with Mish activation). MLP f Ti which processes diffusion timestep f R is a 2-layered MLP (prepended by a Sinusoidal embedding and with Mish activation). f R which processes conditioned Return and f H which processes state history are 3-layered MLPs with Mish activation. f A which processes actions, f TR which process transitions and prediction head are 2-layered MLPs with Mish activation. The GPT2 transformer is configured as 6 hidden layers and 4 attention heads. The code of GPT2 is borrowed from https://github.com/kzl/decision-transformer .

We choose the probability p of removing the conditioning information to 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 K = 200 for diffusion steps.

We set guidance scale α = 1 . 2 for extracting near-optimal behavior.

We choose β = 0 . 5 for low temperature sampling.

Algorithm 1 MTDIFF-P Training and Evaluation # Training Process Initialize : training tasks T train , training iterations N , multi-task dataset D , per-task batch size M , multi-task trajectory prompts Z 1: for n = 1 to N do 2: for Each task T i ∈ T train do 3: Sample action sequences x p 0 ( τ i ) of length H and corresponding state history S prev i of length L from D i with batch size M 4: Compute normalized return R ( τ i ) under τ i 5: Sample trajectory prompts τ ∗ i of length J from Z i with batch size M 6: end for 7: Get a batch B = { x p 0 ( τ i ) , τ ∗ i , S prev i , R ( τ i ) } |T train | i =1 8: Randomly sample a diffusion timestep k ∼ U (1 , K ) and obtain noisy sequences x p k ( τ i ) 9: Omit the R ( τ ) condition with probability β ∼ Bern( p ) 10: Compute L p ( θ ) and update MTDIFF-P model 11: end for # Evaluation Process 1: Given a task, reset the environment and set desired return R max ( τ ) 2: Obtain the initial state history h 0 , few-shot prompts Z 3: Set low-temperature sampling scale β , classifier-free guidance scale α 4: for t = 0 to t max do 5: Initialize x p K ( τ ) ∼ N ( 0 , β I ) 6: Sample τ ∗ ∼ Z , and formulate y ′ ( τ ) = [ h t , τ ∗ ] 7: for k = K to 1 do 8: ˆ ϵ = ϵ θ ( x p k ( τ ) , y ′ ( τ ) , ∅ , k ) + α ( ϵ θ ( x p k ( τ ) , y ′ ( τ ) , R max ( τ ) , k ) -ϵ θ ( x p k ( τ ) , y ′ ( τ ) , ∅ , k ) ) 9: ( µ k -1 , Σ k -1 ) ← Denoise( x p k ( τ ) , ˆ ϵ ) 10: x p k -1 ( τ ) ∼ N ( µ k -1 , β Σ k -1 ) 11: end for 12: Execute the first action from x p 0 ( τ ) as the current action to interact with the environment 13: Obtain the next state, and update h t 14: end for

Algorithm 2 MTDIFF-S Training and Data Synthesis

# Training Process

Initialize : training tasks T train , training iterations N , multi-task dataset D , per-task batch size M , multi-task trajectory prompts Z

1: for n = 1 to N do 2: for Each task T i ∈ T train do 3: Sample transition sequences x s 0 ( τ i ) of length H from D i with batch size M 4: Sample trajectory prompts τ ∗ i of length J from Z i with batch size M 5: end for 6: Get a batch B = { x s 0 ( τ i ) , τ ∗ i } |T train | i =1 7: Randomly sample a diffusion timestep k ∼ U (1 , K ) and obtain noisy sequences x s k ( τ i ) 8: Compute L s ( θ ) and update MTDIFF-S model 9: end for

# Data Synthesis Process

Initialize : synthetic dataset D = ∅ , synthesizing times M

1: Given a task, obtain few-shot prompts Z 2: for m = 1 to M do 3: Initialize x s K ( τ ) ∼ N ( 0 , I ) 4: Sample τ ∗ ∼ Z , and formulate y s k ( τ ) = [ τ ∗ ] 5: for k = K to 1 do 6: ˆ ϵ = ϵ θ ( x s k ( τ ) , y s ( τ ) , k ) 7: ( µ k -1 , Σ k -1 ) ← Denoise( x s k ( τ ) , ˆ ϵ ) 8: x s k -1 ( τ ) ∼ N ( µ k -1 , Σ k -1 ) 9: end for 10: Update D = D∪ x s 0 ( τ ) 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 https://github.com/mxu34/prompt-dt for implementation.

MTDT embeds the tskaID which indicates the task to a embedding z with size 12, then z 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 hyperparameters and implementation details the same as the official version https://github.com/kzl/ decision-transformer/ .

MTIQL uses a 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 a 3-layered MLP (with Mish activation). During training and inference, the scalar taskID is embedded via 3-layered MLP (with Mish activation) into latent variable z , and the input of the actor becomes the concatenation of the original state and z . We build MTIQL based on the code https://github.com/tinkoff-ai/CORL [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://github.com/tinkoff-ai/CORL [59].

MTBC uses a similar taskID-cognitional actor in MTIQL and MTCQL. For training and inference, the scalar taskID is embedded via a 3-layered MLP (with Mish activation) into latent variable z , and the input of the actor becomes the concatenation of the original state and z . The actor is parameterized with a 3-layered MLP and outputs predicted actions.

RAD adopts the random amplitude scaling [28] that multiplies a random variable to states during training, i.e., s ′ = s × z , where z ∼ Uniform[ α, β ] . We choose α = 0 . 8 and β = 1 . 2 .

S4RL adopts the adversarial state training [52] by taking gradients with respect to the value function to obtain a new state, i.e. s ′ ← s + ϵ ∇ s J Q ( π ( s )) , where J Q is the policy evaluation update performed via a 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 a transformer is also helpful for sequential modeling. To ablate the effectiveness of our architecture design, we train MTDIFF-P using U-Net with a similar model size to ours on the near-optimal dataset. We use a Temporal Convolutional Network (TCN) [4] to encode the prompt into an embedding z p , and inject it in the U-Net layers as a condition. We follow the conditional approach in [2] and borrow the code for temporal U-Net from https://github.com/jannerm/diffuser [21]. The results summarized in Table 3 show that our model architecture outperforms U-Net to learn from multi-task datasets.

D Environmental Details of Maze2D

(a) map 1

(b) map 2

(c) map 3

(d) map 4

(e) map 5

(f) map 6

(g) map 7

(h) map 8

3

田

2

+

We design eight different training maps for multi-task training, which are shown at Fig. 8. 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 to generalize on new maps while PromptDT fails, we should acknowledged that MTDIFF may fail at some difficult unseen cases, as shown in Fig. 9. The reasons may lie in 2 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. 10.

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 on seen and unseen maps.

(a) PromptDT

(b) MTDIFF

(c) PromptDT

(d) MTDIFF

(1) Failure case 1

(2) Failure case 2

(a) MTDIFF-P

(b) PromptDT (failed)

E Data Analysis

E.1 Distribution Visualization

We find the synthetic data is high-fidelity, covering or even broadening the original data distribution, which makes the offline RL method performs better in the augmented dataset. The distribution visualization is shown in Fig. 11.

60

40

20

0

20

40

60

80

60

40

20

0

20

40

60

80

Original Data

Synthetic Data

(a) Coffee-push

60

40

20

0

20

40

60

80

60

40

20

0

20

40

60

(b) Disassemble

60

40

20

0

20

40

60

80

60

40

20

0

20

40

60

(c) hand-insert

75

50

25

0

25

50

75

40

20

0

20

40

(d) box-close

E.2 Statistical Analysis

Following SynthER [38], we measure the Dynamics Error (MSE between the augmented next state and true next state), and L2 Distance from Dataset (minimum L2 distance of each datapoint from the dataset) for the augmented data of each method (i.e., MTDIFF-S, S4RL and RAD), as shown in Table 4. Since S4RL performs data augmentation by adding data within the ϵ -ball of the original data points, it has the smallest dynamics error with a small ϵ . RAD performs random amplitude scaling and causes the largest dynamics 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.

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 a crucial research topic. 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 [37, 36] to decrease the diffusion steps required to 0 . 2 × 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 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 diffusion-based approach for multi-task reinforcement learning. It could be applied to multi-task decision-making, and also could be used to synthesize more data to boost policy improvement. MTDIFF provides a 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 in each kind of dataset in Fig. 12. Our dataset is available at https://bit.ly/3MWf40w .

0.0

0.2

0.4

0.6

0.8

1.0

Normalized Return

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

Density

(a) Near-optimal dataset

0.0

0.2

0.4

0.6

0.8

1.0

Normalized Return

0

1

2

3

4

5

6

Density

(b) Sub-optimal dataset

Maze2D The offline dataset is 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.com/Farama-Foundation/D4RL [15] to generate datasets for 8 training maps. We collect 35k episodes in total.

H Differences Between PromptDT and MTDIFF-P

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 perform trajectory generation adeptly. 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 [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 our model on the publicly available PromptDT datasets [69], 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) in Table 5.

I Single-Task Performance

We train one MTDIFF-P model on MT50-rand and evaluate the performance for each task for 50 episodes. We report the average evaluated return in Table 6.

Peer reviews:
comments: 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.

STRENGTHS:
This paper is the first to achieve both effective planning and data synthesis for multi-task RL via diffusion model and GPT.
This paper is well written and easy to follow.

WEAKNESSES:
# The difference between the proposed model and PromptDT is not well explained.
# Additional environment and ablation experiments may be more convincing.

QUESTIONS:
# PromptDT showed performance on MuJoCo. How is the performance of the proposed model on MuJoCo?
# 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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-06-29

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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\nSTRENGTHS:\n- 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\nWEAKNESSES:\n- 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\nQUESTIONS:\n- 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\nLIMITATIONS:\nDiscussion of limitations hidden in Appendix, should be moved to the main paper.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\n- 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\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\nNo such weakness as such.\n\nQUESTIONS:\n- 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).\n\nLIMITATIONS:\nSee Weaknesses and Questions.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

comments: 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.

STRENGTHS:
This paper is easy to follow. I feel quite enjoyable while reading this paper.

WEAKNESSES:
- The experiment section needs further revision:
  - Figure 5 x-label is not correct.
- What are the quantitative results on Maze2D unseen map
- 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.
- line 334: The subtitle is confusing. This paragraph is mainly about figure 6, not comparing other augmentation methods.
- 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.
- In Table 2, what are the original results? Is it the maximal return of the dataset?

QUESTIONS:
Please see weaknesses.

LIMITATIONS:
Yes.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-07

comments: 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.

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.

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
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-08