Flexible Multitask Learning with Factorized Diffusion Policy
This addresses multitask learning challenges in robotics, enabling more effective policy fitting and flexible adaptation to new tasks while mitigating catastrophic forgetting.
The paper tackles the challenge of fitting policies to complex multimodal robot action distributions in multitask learning by introducing a modular diffusion policy framework that factorizes action distributions into specialized components. The method consistently outperforms strong baselines in simulation and real-world robotic manipulation settings.
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.