Task-Agnostic Pre-training and Task-Guided Fine-tuning for Versatile Diffusion Planner
This work addresses the problem of reducing human effort in multi-task planning for robotics and AI systems, though it is incremental as it builds on existing diffusion and RL methods.
The paper tackles the challenge of costly expert data or reward design for multi-task planning by proposing SODP, a two-stage diffusion planner that pre-trains on sub-optimal trajectories and fine-tunes with task-specific rewards, achieving state-of-the-art performance in Meta-World and Adroit domains with minimal fine-tuning data.
Diffusion models have demonstrated their capabilities in modeling trajectories of multi-tasks. However, existing multi-task planners or policies typically rely on task-specific demonstrations via multi-task imitation, or require task-specific reward labels to facilitate policy optimization via Reinforcement Learning (RL). They are costly due to the substantial human efforts required to collect expert data or design reward functions. To address these challenges, we aim to develop a versatile diffusion planner capable of leveraging large-scale inferior data that contains task-agnostic sub-optimal trajectories, with the ability to fast adapt to specific tasks. In this paper, we propose SODP, a two-stage framework that leverages Sub-Optimal data to learn a Diffusion Planner, which is generalizable for various downstream tasks. Specifically, in the pre-training stage, we train a foundation diffusion planner that extracts general planning capabilities by modeling the versatile distribution of multi-task trajectories, which can be sub-optimal and has wide data coverage. Then for downstream tasks, we adopt RL-based fine-tuning with task-specific rewards to quickly refine the diffusion planner, which aims to generate action sequences with higher task-specific returns. Experimental results from multi-task domains including Meta-World and Adroit demonstrate that SODP outperforms state-of-the-art methods with only a small amount of data for reward-guided fine-tuning.