Flexible Multitask Learning with Factorized Diffusion Policy

TL;DR

This paper presents a modular diffusion policy framework for multitask robotic learning, enabling flexible adaptation and improved performance over traditional monolithic models.

Contribution

A novel factorized diffusion policy approach that captures diverse behaviors and allows for efficient task adaptation while reducing catastrophic forgetting.

Findings

Outperforms strong modular and monolithic baselines in simulation and real-world tasks.

Enables flexible policy adaptation by adding or fine-tuning components.

Effectively models complex, multimodal action distributions in robotics.

Abstract

Multitask learning poses significant challenges due to the highly multimodal and diverse nature of robot action distributions. However, effectively fitting policies to these complex task distributions is often difficult, and existing monolithic models often underfit the action distribution and lack the flexibility required for efficient adaptation. We introduce a novel modular diffusion policy framework that factorizes complex action distributions into a composition of specialized diffusion models, each capturing a distinct sub-mode of the behavior space for a more effective overall policy. In addition, this modular structure enables flexible policy adaptation to new tasks by adding or fine-tuning components, which inherently mitigates catastrophic forgetting. Empirically, across both simulation and real-world robotic manipulation settings, we illustrate how our method consistently outperforms strong modular and monolithic baselines.