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

2023-10-10

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.

Machine Learning,Artificial Intelligence

What problem does this paper attempt to address?

The paper aims to address two main issues in multi-task reinforcement learning: planning and data synthesis. Specifically: 1. **Multi-Task Planning**: A diffusion model method named Multi-Task Diffusion Model (MTD IFF) is proposed for generative planning in multi-task environments. Through a conditional generation process, this model can extract task-relevant information from a small number of examples and generate high-quality action sequences to tackle new or unseen tasks. 2. **Data Synthesis**: In multi-task environments, MTD IFF can also synthesize high-quality transition data containing states, actions, and rewards, thereby enhancing low-quality datasets and improving policy performance. This helps to extend the original offline datasets, making them applicable to more task scenarios. The paper experimentally validates MTD IFF-P (for planning) and MTD IFF-S (for data synthesis) on Meta-World and Maze2D benchmarks, demonstrating its effectiveness and flexibility in solving multi-task problems. Compared to existing methods, MTD IFF shows better performance across multiple tasks.

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

MetaDiffuser: Diffusion Model as Conditional Planner for Offline Meta-RL

Diffusion Models as Optimizers for Efficient Planning in Offline RL

MADiff: Offline Multi-agent Learning with Diffusion Models.

Task-agnostic Pre-training and Task-guided Fine-tuning for Versatile Diffusion Planner

Improving Efficiency of Diffusion Models via Multi-Stage Framework and Tailored Multi-Decoder Architectures

Instructed Diffuser with Temporal Condition Guidance for Offline Reinforcement Learning

Efficient Diffusion Transformer Policies with Mixture of Expert Denoisers for Multitask Learning

Off-dynamics Conditional Diffusion Planners

Adaptive Online Replanning with Diffusion Models

Diffused Task-Agnostic Milestone Planner

Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning

Simple Hierarchical Planning with Diffusion

One Transformer Fits All Distributions in Multi-Modal Diffusion at Scale

Hierarchical Diffusion Scattering Graph Neural Network

Diffusion World Model: Future Modeling Beyond Step-by-Step Rollout for Offline Reinforcement Learning

Policy Representation via Diffusion Probability Model for Reinforcement Learning

Diffusion Policies creating a Trust Region for Offline Reinforcement Learning

Diffusion Transformer Policy

Multimodal Diffusion Transformer: Learning Versatile Behavior from Multimodal Goals

Learning Multimodal Behaviors from Scratch with Diffusion Policy Gradient