SlowGuess/NeurIPS_2023_OCR huggingface dataset:
where  $E^{\mathrm{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_R^p$ . 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:
$$
h _ {\text {o u t}} ^ {p} = \operatorname {t r a n s f o r m e r} \left(h _ {\text {t o k e n s}} ^ {p}\right), \quad h _ {\text {o u t}} ^ {s} = \operatorname {t r a n s f o r m e r} \left(h _ {\text {t o k e n s}} ^ {s}\right).
$$
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_{\theta}$  during inference. We summarize the details of the training process, architecture and hyperparameters used in MTDIFF in Appendix A.
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.
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 MTDIFF-S 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 multi-task 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 \times z$ , where  $z \sim \mathrm{Uniform}[\alpha, \beta]$ . 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' \gets s + \epsilon \nabla_s \mathbb{J}_Q(\pi(s))$ , where  $\mathbb{J}_Q$  is the policy evaluation update performed via a  $Q$  function, and  $\epsilon$  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-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
Table 1: Average success rate across 3 seeds on MetaWorld-V2 MT50 with random goals (MT50-rand). Each task is evaluated for 50 episodes.
Methods	Near-optimal	Sub-optimal
CARE [53] (Online)	50.8 ± 1.0	-
PaCo [56] (Online)	57.3 ± 1.3	-
MTDT	20.99 ± 2.66	20.63 ± 2.21
PromptDT	45.68 ± 1.84	39.76 ± 2.79
MTBC	60.39 ± 0.86	34.53 ± 1.25
MTCQL	-	-
MTIQL	56.21 ± 1.39	43.28 ± 0.90
MTDIFF-p (ours)	59.53 ± 1.12	48.67 ± 1.32
MTDIFF-P-ONEHOT (ours)	61.32 ± 0.89	48.94 ± 0.95
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 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
(1) Seen map
(b) PromptDT
(c) MTDIFF-P
(2) Unseen map
(d) PromptDT(failed)
Figure 4: Seen and unseen maps of Maze2D with long planning path. Goal position is denoted as  $\blacksquare$ .
Figure 5: Average scores obtained in 8 Maze2D maps.
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 MTDIFF-S 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.
Figure 6: Results of multi-task and single- task augmentation
Figure 7: Average success rate across 3 seeds on coffee-push. Dataset is augmented by MTDIFF-S trained under different number of tasks.
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.
Table 2: Average success rate across 3 seeds on Meta-World-V2 single task with random goals. Each selected task is evaluated for 500 episodes.
Tasks	Original	S4RL	RAD	MTDIFF-S(OURS)
Unseen Tasks	box-close	23.46 ± 7.11	73.13 ± 3.51 (↑ 211%)	71.20 ± 4.84 (↑ 203%)	65.73 ± 8.36 (↑ 180%)
hand-insert	30.60 ± 9.77	60.20 ± 1.57 (↑ 96%)	43.79 ± 3.44 (↑ 43%)	70.87 ± 3.59 (↑ 131%)
bin-picking	42.13 ± 14.33	72.20 ± 4.17 (↑ 71%)	43.27 ± 4.38 (↑ 2%)	55.73 ± 7.63 (↑ 32%)
Seen Tasks	sweep-into	91.8 ± 1.14	90.53 ± 3.52 (↓ 1%)	88.06 ± 9.86 (↓ 4%)	92.87 ± 1.11 (↑ 1%)
coffee-push	28.60 ± 14.55	28.73 ± 8.44 (↑ 0.4%)	33.19 ± 2.86 (↑ 16%)	74.67 ± 6.79 (↑ 161%)
disassemble	60.20 ± 16.29	52.20 ± 5.68 (↓ 12%)	60.93 ± 20.80 (↑ 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.
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.
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.
[1] 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.
[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 E. Hinton. Layer normalization. CoRR, abs/1607.06450, 2016. URL http://arxiv.org/abs/1607.06450.
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  $2e^{-4}$  learning rate for  $2e^{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/kz1/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  $\alpha = 1.2$  for extracting near-optimal behavior.
- We choose  $\beta = 0.5$  for low temperature sampling.
Algorithm 1 MTDIFF-P Training and Evaluation
# Training Process   
Initialize: training tasks  $\mathcal{T}^{\text{train}}$  , training iterations  $N$  , multi-task dataset  $\mathcal{D}$  , per-task batch size  $M$    
multi-task trajectory prompts  $Z$    
1: for  $n = 1$  to  $N$  do   
2: for Each task  $\mathcal{T}_i\in \mathcal{T}^{\text{train}}$  do   
3: Sample action sequences  $\pmb{x}_0^p (\tau_i)$  of length  $H$  and corresponding state history  $S_{i}^{\mathrm{prev}}$  of length  $L$  from  $\mathcal{D}_i$  with batch size  $M$    
4: Compute normalized return  $R(\tau_{i})$  under  $\tau_{i}$    
5: Sample trajectory prompts  $\tau_{i}^{*}$  of length  $J$  from  $Z_{i}$  with batch size  $M$    
6: end for   
7: Get a batch  $\mathcal{B} = \{\pmb{x}_0^p (\tau_i),\tau_i^*,S_i^{\mathrm{prev}},R(\tau_i)\}_{i = 1}^{|\mathcal{T}^{\text{train}}|}$    
8: Randomly sample a diffusion timestep  $k\sim \mathcal{U}(1,K)$  and obtain noisy sequences  $\pmb{x}_k^p (\tau_i)$    
9: Omit the  $R(\tau)$  condition with probability  $\beta \sim \mathrm{Bern}(p)$    
10: Compute  $\mathcal{L}^p (\theta)$  and update MTDIFF-P model   
11: end for   
# Evaluation Process   
1: Given a task, reset the environment and set desired return  $R_{\max}(\tau)$    
2: Obtain the initial state history  $h_0$  , few-shot prompts  $Z$    
3: Set low-temperature sampling scale  $\beta$  , classifier-free guidance scale  $\alpha$    
4: for  $t = 0$  to  $t_\mathrm{max}$  do   
5: Initialize  $\pmb{x}_K^p (\tau)\sim \mathcal{N}(\pmb {0},\beta \pmb {I})$    
6: Sample  $\tau^{*}\sim Z$  , and formulate  $\pmb {y}'(\tau) = [h_t,\tau^* ]$    
7: for  $k = K$  to 1 do   
8:  $\hat{\epsilon} = \epsilon_{\theta}\bigl (\pmb{x}_{k}^{p}(\tau),\pmb{y}^{\prime}(\tau),\emptyset ,k\bigr) + \alpha \bigl (\epsilon_{\theta}(\pmb{x}_{k}^{p}(\tau),\pmb{y}^{\prime}(\tau),R_{\max}(\tau),k) - \epsilon_{\theta}(\pmb{x}_{k}^{p}(\tau),\pmb{y}^{\prime}(\tau),\emptyset ,k)\bigr)$    
9:  $(\mu_{k - 1},\Sigma_{k - 1})\gets \mathrm{Denoise}(\pmb{x}_k^p (\tau),\hat{\epsilon})$    
10:  $\pmb{x}_{k - 1}^{p}(\tau)\sim \mathcal{N}(\mu_{k - 1},\beta \Sigma_{k - 1})$    
11: end for   
12: Execute the first action from  $\pmb{x}_0^p (\tau)$  as the current action to interact with the environment   
13: Obtain the next state, and update  $h_t$    
14: end for