Flexible Multitask Learning with Factorized Diffusion Policy

TL;DR

This work tackles multitask imitation learning in robotics by addressing the highly multimodal action distributions that arise across tasks. It introduces Factorized Diffusion Policy (FDP), a modular framework that factorizes the policy into multiple diffusion components whose outputs are continuously composed via an observation-conditioned router, enabling stable training and scalable adaptation. FDP enables efficient task transfer by upcycling new diffusion components and freezing existing ones to mitigate forgetting, while maintaining strong multitask performance across simulated and real-world robotic manipulation tasks. Empirical results on MetaWorld, RLBench, and real-world experiments demonstrate superior performance and data efficiency compared to strong baselines, with analyses highlighting component specialization, scalable adaptation, and convergence benefits. These findings suggest FDP as a practical, extensible approach for modular, diffusion-based multitask policies in robotics with strong potential for lifelong learning and deployment in evolving domains.

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.

Figures (9)

• Figure 1: Overview of FDP. (a) Given an observation $\mathbf{o}_t$, multiple diffusion experts predict score estimates $\boldsymbol{\varepsilon}_i(\mathbf{a}_t^K, \mathbf{o}_t)$ at each denoising step. A lightweight router network computes observation-dependent weights $\{w_i\}$, which are used to compose the final score as a weighted sum (see (c)). The composed score guides the iterative denoising process over $K$ steps to generate an action $\mathbf{a}_t$. (b) This compositional structure enables FDP to model complex multimodal distributions and supports modular adaptation via selective tuning or addition of diffusion components.
• Figure 2: Real-world setup and task illustrations. (a) Workspace setup with a UR5e arm, Robotiq gripper, and RealSense D415 camera. (b) High-level task illustrations.
• Figure 3: Real-world rollouts. Top: cube-X. Bottom: hang-X. Top and bottom rows show success cases and baseline failure modes.
• Figure 4: Relative success rate improvement of FDP over DP.FDP 's advantage increases as the number of tasks grows. Selected from RLBench. We report mean and standard error over 5 seeds.
• Figure 5: Performance scaling with number of demonstrations.(a) Metaworld tasks door open, drawer open, assembly, window close, peg insert, hammer; RLBench tasks door open, drawer open, assembly, window close, peg insert, hammer. (b) Metaworld tasks door close, drawer close, disassemble, window open; RLBench tasks toilet seat down, close box