PaperlayerView exact revision on arXiv

Source-linked AI summary

Diffusion Model is an Effective Planner and Data Synthesizer for Multi-Task Reinforcement Learning

Haoran He, Chenjia Bai, Kang Xu, Zhuoran Yang, Weinan Zhang, Dong Wang, Bin Zhao, Xuelong Li

arXiv:2305.18459v2cs.LGcs.AI

TL;DR

Existing diffusion-based offline RL methods largely focus on single tasks, leaving the challenge of generalist multi-task decision-making over diverse, multimodal data. MTDIFF combines GPT-based diffusion modeling with prompt learning for generative planning and data synthesis, achieving strong results across multi-task benchmarks and augmenting datasets for seen and unseen tasks. Its main practical limitation is slow diffusion sampling for real-time control.

  • Problem

    Diffusion-based offline RL has mainly addressed single-task settings, leaving limited evidence for one generalist model handling diverse, multimodal multi-task data.

  • Method

    MTDIFF uses a GPT-based diffusion model with prompt learning to perform generative planning and synthesize multi-task offline data.

  • Results

    MTDIFF provides effective planning across Meta-World and Maze2D tasks and improves offline policies through synthesized data, including approximately 180% improvement on box-close.

  • Takeaways & Limitations

    A single diffusion model can share multi-task knowledge for planning and data augmentation, including adaptation to unseen tasks from task-informed prompts.

  • Takeaways & Limitations

    Diffusion sampling is slow for real-time control, taking an average of 1.9s to generate one MetaWorld action sequence on a 3090 GPU.

Abstract

from arXiv · show

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.

1 Introduction

MTDIFF addresses multi-task offline RL by combining diffusion-based trajectory modeling, GPT sequential modeling, and prompt learning. It supports generative planning and data synthesis across diverse tasks.

  • Multi-task offline RL requires modeling diverse, multimodal data across tasks rather than relying on single-task decision-making methods.
  • MTDIFF formulates multi-task learning as diffusion-based denoising and uses a GPT backbone to model sequential trajectories with lower computation burden than U-Net-based alternatives.
  • MTDIFF-P learns task-relevant prompt embeddings and conditions action generation on normalized return and historical states for generative planning.
  • MTDIFF-S adapts the same diffusion framework for data synthesis, while prompt learning enables adaptation to unseen tasks from task-informed demonstrations.
  • MTDIFF combines diffusion generation, GPT sequence modeling, and prompt learning to provide planning and data synthesis for multi-task RL.

2 Preliminaries

The paper frames multi-task offline RL as conditional generative modeling over static, task-partitioned datasets. MTDIFF uses diffusion to model multimodal trajectories and avoid value-function fitting and associated distribution-shift risks.

  • A Markov Decision Process is defined by state and action spaces, transition and reward functions, an initial-state distribution, and a discount factor.
  • Multi-task RL considers task-specific MDPs that may differ in rewards, states, and transitions while sharing an action space and embodied agent.
  • Offline multi-task datasets are static collections partitioned into per-task subsets generated by an unknown behavior policy.
  • MTDIFF treats multi-task policy learning as conditional diffusion generation, exploiting distribution modeling to avoid fitting a value function and reduce distribution-shift risk.
  • The forward diffusion chain gradually adds noise to trajectory data over K steps, while the reverse model learns to denoise conditioned on task attributes and trajectory optimality.
  • Classifier-free guidance combines conditional and unconditional noise predictions to generate samples according to a chosen guidance scale.

3 Methodology

MTDIFF formulates multi-task trajectory modeling as conditional diffusion, using task prompts to support planning and data synthesis. Its GPT-based architecture represents heterogeneous inputs as tokens and predicts noise for reverse denoising.

  • Diffusion formulation: MTDIFF models multimodal multi-task trajectories as a conditional generative process, producing sequences through reverse denoising.The generated sequence is used for either planning or data synthesis.
  • Generative planning: For planning, the model generates action sequences conditioned on task prompts, historical states, and normalized trajectory return.Return conditioning uses classifier-free guidance to extract high-return behavior.
  • Data synthesis: For data synthesis, MTDIFF generates transition sequences containing states, actions, and rewards without return guidance.This variant targets diverse trajectory synthesis for data augmentation.
  • Prompt learning: Trajectory prompts contain states and actions, allowing MTDIFF to identify tasks by implicitly capturing transition and reward information.The paper states that this prompt-based task specification supports generalization to unseen tasks without additional parameter tuning.
  • Architecture design: MTDIFF uses GPT2 self-attention as a trainable sequential backbone, with separate MLP embeddings and a prediction head for diffusion-noise estimation.Different inputs are converted into same-sized embeddings and treated as tokens in a unified architecture.
  • Architecture design: The architecture combines a Transformer backbone with MLPs and a prediction head to model sequential inputs and support efficient training.The predicted noise shares the original input space and drives reverse denoising during inference.

4 Related Work

Related work spans diffusion-based offline RL, multi-task and few-shot RL, and RL data augmentation. MTDIFF differs by using conditional generative modeling for multi-task offline learning and synthesizing new experiences rather than perturbing observations.

  • Diffusion models in RL: Prior RL studies use diffusion models to learn multimodal offline policies, human behaviors, or conditional sequential decision processes.These works establish diffusion models as generative tools for offline reinforcement learning.
  • Multi-task and few-shot RL: Multi-task RL methods address conflicting task gradients through gradient surgery, conflict-averse learning, or parameter composition.These approaches primarily concern online multi-task learning.
  • Multi-task and few-shot RL: Earlier multi-task decision-transformer methods often rely on expert trajectories and substantial training expenses, while Scaled-QL uses separate networks.MTDIFF instead addresses multi-task learning offline through conditional generation and a Transformer architecture.
  • Data augmentation for RL: Prior RL augmentation methods commonly perturb observations, whereas MTDIFF focuses on synthesizing new experiences for data augmentation.The distinction is between creating new trajectories and generating multiple views of existing observations.

5 Experiments

Experiments evaluate MTDIFF for multi-task planning and data synthesis across Meta-World and Maze2D, including seen and unseen tasks. Results indicate strong planning performance, policy improvement from synthesized data, benefits from multi-task training, and distributional consistency of generated data.

  • Experimental setup: Experiments compare MTDIFF-P and MTDIFF-S with offline, online, and augmentation baselines across Meta-World and Maze2D settings.The evaluation includes 50 Meta-World manipulation tasks, Maze2D maps, and selected seen and unseen tasks.
  • Planning: About 60% success rate is achieved by MTDIFF-P and MTDIFF-P-ONEHOT on near-optimal Meta-World datasets, outperforming other methods and remaining strongest on sub-optimal datasets.MTDIFF-P also outperforms CARE and PaCo when given near-optimal offline data, while MTBC degrades substantially on sub-optimal data.
  • Planning: MTDIFF-P generalizes from 8 training maps to an unseen Maze2D map, where it is compared with PromptDT on seen and unseen maps.Average scores are reported across the 8 training maps, with additional evaluation on a new map.
  • Data synthesis: 180%, 131%, and 161% performance increases are reported for box-close, hand-insert, and coffee-push when MTDIFF-S augments offline datasets.MTDIFF-S synthesizes 2M transitions to expand each original 1M dataset, and policy improvement is observed across all evaluated tasks.
  • Data synthesis: MTDIFF-S outperforms single-task augmentation methods, and its data-synthesis performance progressively improves as training task diversity increases.The comparison includes S4RL, RAD, MTDIFF-S-SINGLE, and SynthER; experiments retrain MTDIFF-S with 10, 20, and 30 tasks.
  • Data synthesis: On unseen tasks, MTDIFF-S improves policy performance by 131%, 180%, and 32% for hand-insert, box-close, and bin-picking, respectively.The method uses much less information for augmentation than S4RL, relying on a single trajectory as a prompt.
  • Data synthesis: Synthetic data overlaps with and expands the original data distribution while remaining consistent with the underlying MDP.This finding is based on T-SNE visualizations and further analyses for four tasks.

6 Conclusion

MTDIFF is presented as a unified diffusion-based method for multi-task RL that supports planning and data synthesis, using trajectory prompts and a GPT-based architecture. The implementation uses separate MTDIFF-P and MTDIFF-S procedures with diffusion training and prompt-conditioned generation.

  • Conclusion: MTDIFF combines diffusion modeling, GPT-based sequential modeling, and trajectory prompts for multi-task RL planning and data synthesis.The paper describes MTDIFF-P for planning and MTDIFF-S for data synthesis.
  • Implementation: The implementation specifies a 20-step trajectory prompt, 200 diffusion steps, a 0.25 conditioning-drop probability, and distinct state-history lengths for Meta-World and Maze2D.State-history lengths are 2 for Meta-World and 5 for Maze2D.
  • MTDIFF-P: MTDIFF-P trains on noisy action sequences conditioned on normalized return, state history, and sampled trajectory prompts.Training randomly selects diffusion timesteps and updates the model using a denoising objective.
  • MTDIFF-P: MTDIFF-P generates actions by iteratively denoising a sampled sequence and executing its first action before updating the state history.Inference uses a desired return, initial state history, trajectory prompts, sampling temperature, and classifier-free guidance.
  • MTDIFF-S: MTDIFF-S trains on transition sequences and synthesizes new data by denoising sequences initialized from trajectory prompts.The generated sequences are accumulated into a synthetic dataset for augmentation.

C Ablation Study on Model Architecture

The ablation compares MTDIFF-P with a similarly sized U-Net architecture to assess the contribution of the proposed transformer-based design. The comparison is conducted on MT50-rand using average success rate across three seeds.

  • Architecture comparison: The architecture ablation replaces the proposed transformer-based model with a similarly sized U-Net and encodes prompts using a Temporal Convolutional Network.Both variants are trained on the near-optimal dataset.
  • Evaluation: Table 3 reports average success rates across three seeds for MTDIFF-P and its U-Net variant on MT50-rand.The caption identifies the evaluation as an MT50-rand comparison.

D Environmental Details of Maze2D

Maze2D uses eight training maps and one unseen map to evaluate generalization. MTDIFF-P succeeds on an illustrated unseen case where PromptDT fails, but also fails on difficult unseen cases.

  • Training environment: Eight different maps are designed for multi-task training in Maze2D, with different tasks using different reward and transition functions.The training-map layouts are shown in Figure 8.
  • Generalization setup: The evaluation adds one new unseen map to test generalization beyond the eight training maps.The setup uses 35k episodes, with episode lengths of 600 for training and 200 for evaluation.
  • Unseen-map result: MTDIFF-P reaches the goal on an illustrated unseen map where PromptDT fails to obtain a reasonable path.The result is presented as evidence of few-shot adaptation from trajectory prompts and the designed architecture.
  • Generalization boundary: MTDIFF may fail on difficult unseen cases when their difficulty or distributional deviation exceeds its generalizability threshold.Figure 9 shows two unseen cases where both PromptDT and MTDIFF fail.
  • Data synthesis: Synthetic data is reported to cover or broaden the original data distribution while preserving consistency with the underlying MDP.The distribution visualization is shown in Figure 11 for four tasks.

E.2 Statistical Analysis

The statistical analysis compares augmented data using L2 distance from the training dataset and dynamics error. It evaluates MTDIFF-S alongside S4RL and RAD using the metrics summarized in Table 4.

  • Metrics: The analysis measures Dynamics Error and L2 Distance from Dataset for augmented data produced by MTDIFF-S, S4RL, and RAD.Dynamics Error is the MSE between augmented and true next states, while L2 Distance is each datapoint’s minimum distance from the dataset.
  • Method comparison: S4RL has the smallest dynamics error because its augmentation adds data within a small epsilon-ball around original data points.RAD instead uses random amplitude scaling, which produces larger deviations according to the comparison discussion.

F Limitations and Discussions

MTDIFF is presented as applicable to multi-task decision-making and data synthesis, but its sampling speed limits real-time control. The authors discuss speed-quality trade-offs and possible acceleration strategies.

  • Diffusion sampling speed limits MTDIFF’s potential for real-time control.
  • Generating one Meta-World action sequence takes 1.9 seconds on average using a 3090 GPU.
  • DPM-Solver, larger batch sizes, and consistency models are identified as routes toward faster inference.
  • MTDIFF can support multi-task decision-making and synthesize additional data for policy improvement.

G Dataset collection

The evaluation uses Meta-World datasets spanning random-to-expert and sub-optimal experience, plus Maze2D trajectories collected across eight training maps. The section also reports strong MTDIFF-P results against PromptDT, though several supplied passages provide only partial table context.

  • Meta-World: Meta-World data are collected from SAC replay buffers, with 1M transitions per task.
  • Meta-World: The near-optimal dataset contains 100M transitions ranging from random experience to converged expert behavior.
  • Meta-World: The sub-optimal dataset contains 50M transitions from the initial half of near-optimal trajectories, with much less expert data.
  • Maze2D: Maze2D data are generated across eight training maps by sampling random goals and using a planner with a PD controller.
  • Maze2D: Maze2D contains 35k episodes in total.
  • Comparison: MTDIFF-P is reported to outperform PromptDT substantially in multi-task performance and data synthesis.

I Single-Task Performance

The supplied passages describe evaluating one MTDIFF-P model across the MT50-rand tasks, reporting average returns over 50 episodes for each task.

  • One MTDIFF-P model is trained on MT50-rand and evaluated separately on each task.
  • Each task’s performance is reported as average evaluated return over 50 episodes.
  • Table 6 reports the mean and standard deviation for every MT50-rand task.
Loading 2305.18459v2…