🤖 AI

Flexible Multitask Learning with Factorized Diffusion Policy

This paper introduces a novel modular diffusion policy framework that factorizes complex, multimodal robot action distributions into specialized components, thereby improving policy fitting, enabling flexible adaptation to new tasks, and mitigating catastrophic forgetting while outperforming existing monolithic and modular baselines in both simulation and real-world settings.

Original authors: Chaoqi Liu, Haonan Chen, Sigmund H. Høeg, Shaoxiong Yao, Yunzhu Li, Kris Hauser, Yilun Du

Published 2026-04-27

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Chaoqi Liu, Haonan Chen, Sigmund H. Høeg, Shaoxiong Yao, Yunzhu Li, Kris Hauser, Yilun Du

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot to do many different jobs: opening a drawer, picking up a red cube, hanging a mug on a specific hook, and hammering a nail.

The Problem: The "Swiss Army Knife" Struggle
Traditional robot brains (called "monolithic policies") try to be a single, giant Swiss Army knife. They attempt to learn all these different skills in one big, messy brain. The paper argues that this is hard because the robot's actions are too diverse. It's like trying to teach a single student to be a master chef, a professional pianist, and a race car driver all at once. The student gets confused, tries to do everything at once, and ends up doing none of them very well. They might try to "grasp" a hammer like a spoon, or "open" a drawer like a door.

The Solution: The "Specialized Team" (FDP)
The authors propose a new method called Factorized Diffusion Policy (FDP). Instead of one giant brain, imagine a specialized team of experts.

The Experts (Diffusion Components): The robot has several small, specialized "experts."
- Expert A is great at "approaching" objects.
- Expert B is great at "grasping" things gently.
- Expert C is great at "hammering" or "hanging" things.
- Expert D is great at "aligning" objects.
- Each expert only focuses on one specific type of movement or "sub-skill."
The Manager (The Router): When the robot sees a task (e.g., "Hang the mug"), a smart manager (called a "router") looks at the situation. Instead of picking just one expert to do the whole job, the manager asks all the experts for their advice.
- The manager says, "Expert A, you're 80% right about how to approach. Expert B, you're 90% right about how to grab. Expert C, you're 10% right about the angle."
- The manager then mixes all these pieces of advice together to create the perfect final action. It's like a conductor blending different instruments to create a symphony, rather than asking one instrument to play the whole song.

Why This is Better

Stability: Because the manager blends everyone's advice smoothly (instead of picking one and ignoring the rest), the robot learns much faster and doesn't get confused. It's like a team where everyone contributes, rather than a team where one person shouts over everyone else.
No "Forgetting": This is the paper's big win. If you want to teach the robot a new skill (like "screwing in a lightbulb"), you don't have to retrain the whole team. You just hire a new expert for that specific job and let them join the team. The old experts (who know how to open drawers or hang mugs) stay exactly the same. This means the robot learns the new skill without forgetting the old ones—a problem called "catastrophic forgetting" that plagues other methods.
Flexibility: If the robot needs to do a very complex job, the manager can ask for more advice from more experts. If the job is simple, it can rely on just a few.

Real-World Proof
The authors tested this on robots in both computer simulations and the real world.

In Simulations: The robot team (FDP) beat the "single brain" robots and other "team" robots at tasks like opening drawers, assembling pegs, and picking up umbrellas.
In the Real World: They tested it with a real robot arm. When asked to pick up a red cube or hang a mug, the FDP robot was more successful and precise than the others. The other robots often fumbled or dropped things because they couldn't handle the complexity of the task.

The Bottom Line
The paper claims that by breaking a complex robot brain into a team of specialized, mix-and-match experts, robots can learn many different tasks faster, remember their old skills better, and adapt to new jobs without needing a total system overhaul. It turns the robot from a confused generalist into a highly efficient, adaptable specialist team.

1. Problem Statement

The paper addresses the challenges of multitask imitation learning in robotics, specifically the difficulty of modeling highly multimodal and diverse action distributions.

Limitations of Monolithic Models: Traditional single-network policies (e.g., standard Diffusion Policy) struggle to generalize across diverse tasks, often underfitting complex distributions or failing to represent distinct behavioral modes (e.g., different grasping strategies for different objects).
Limitations of Existing Modular Models: While Mixture-of-Experts (MoE) architectures attempt to decompose tasks, they suffer from:
- Training Instability: Difficulty in balancing expert utilization.
- Routing Issues: Discrete expert selection can lead to "routing imbalance" (some experts are never used, others are overused).
- Poor Interpretability: Experts often have overlapping or unclear roles, and skills may be distributed across layers rather than being disentangled.
- Catastrophic Forgetting: Adapting to new tasks often requires retraining the entire model, risking the loss of previously learned skills.

2. Methodology: Factorized Diffusion Policy (FDP)

The authors propose FDP, a modular policy architecture that factorizes the action distribution into a composition of specialized diffusion models.

A. Core Architecture

Factorization: Instead of a single policy $p(a_t|o_t)$ , FDP models the action distribution as a weighted product of $N$ component distributions:
$p(a_t|o_t) \propto \prod_{i} p_i(a_t|o_t)^{w_{t,i}}$
Here, each $p_i$ is a specialized diffusion model (expert) capturing a distinct behavioral mode (sub-skill), and $w_{t,i}$ are observation-dependent weights.
Compositional Sampling:
- Router: A lightweight MLP takes the observation $o_t$ and outputs continuous weights $\{w_{t,i}\}$ for all components.
- Score Aggregation: Unlike discrete MoE (which selects one expert), FDP aggregates the score estimates (gradients of the log-probability) from all components:
  $\nabla \log p(a_t|o_t) \approx \sum_{i} w_{t,i} \nabla \log p_i(a_t|o_t)$
- Denoising Process: The iterative denoising step uses this weighted sum of scores to guide the generation of the action trajectory. This corresponds to sampling from the product of distributions, effectively framing action generation as a constraint satisfaction problem.

B. Training and Adaptation

Joint Training: All diffusion components and the router are trained end-to-end using a Mean Squared Error (MSE) loss on the aggregated score prediction.
Efficient Adaptation (Upcycling): To adapt to a new task without retraining the whole model:
1. A new diffusion component is introduced.
2. It is initialized by upcycling (copying weights) from an existing component.
3. Only the new component and the router are fine-tuned; all previous components are frozen.
4. This prevents catastrophic forgetting and significantly reduces the number of trainable parameters.

3. Key Contributions

Novel Modular Architecture: Introduced FDP, which uses continuous score aggregation rather than discrete expert selection. This ensures stable training, prevents routing imbalance, and promotes clear specialization among components.
Probabilistic Interpretation: Grounded in Energy-Based Models (EBM) and score matching, providing a principled formulation where the composed policy represents the intersection of behavioral constraints.
Scalable Adaptation Strategy: Proposed a "selective tuning" mechanism using upcycling and freezing, enabling efficient skill expansion with minimal data and no catastrophic forgetting.
Empirical Validation: Demonstrated superior performance over strong baselines (Monolithic DP, Sparse Diffusion Policy, and MoDE) in both simulation and real-world settings.

4. Experimental Results

The authors evaluated FDP on MetaWorld, RLBench, LIBERO, and real-world robotic manipulation (UR5e arm).

Multitask Performance:
- FDP consistently outperformed monolithic Diffusion Policy (DP) and modular baselines (SDP, MoDE) across 30+ tasks.
- MetaWorld: FDP achieved 74.8% average success vs. 70.8% (DP) and 69.8% (SDP).
- RLBench: FDP achieved 63.9% average success vs. 45.6% (DP) and 42.1% (SDP).
- Real-World: FDP showed higher success rates in complex tasks (e.g., hanging mugs, picking colored cubes) where baselines failed due to overfitting or inability to capture multimodal distributions.
Task Adaptation & Transfer:
- Data Efficiency: When adapting to new tasks with limited data (10–25 demos), FDP using the "Add New Module" strategy outperformed full-parameter fine-tuning and partial fine-tuning.
- Continual Learning: In a 12-task continual learning benchmark (LIBERO), FDP maintained high performance as new experts were added, whereas other methods degraded.
- Knowledge Retention: Freezing old components effectively preserved prior skills. Using a minimal replay buffer (5 demos/task) allowed the model to retain nearly 100% of original performance after adaptation.
Analysis & Scaling:
- Component Scaling: Performance improved as the number of components increased from 2 to 4, then plateaued.
- Task Scaling: The performance gap between FDP and monolithic models widened as the number of tasks increased, highlighting FDP's scalability.
- Interpretability: Visualization showed that individual components specialized in distinct sub-skills (e.g., alignment, approach, grasping) rather than noise levels or arbitrary layers.
- Inference Speed: FDP supports "partial reconstruction" (using only top- $k$ components), achieving a 2x speedup in inference with only a ~19% performance drop.

5. Significance

Robotic Learning: FDP offers a practical solution for deploying robots in dynamic environments where they must learn new skills rapidly without forgetting old ones.
Stability: By replacing discrete routing with continuous score composition, FDP solves the training instability issues that have plagued MoE-based diffusion policies.
Interpretability: The factorization allows researchers to inspect and understand specific sub-skills (e.g., "grasping" vs. "approaching") within the policy, aiding in debugging and safety verification.
Future Directions: The framework opens avenues for heterogeneous module designs (mixing architectures) and lifelong learning systems where robots continuously accumulate a library of specialized skills.

In summary, FDP represents a significant advancement in robotic imitation learning by combining the generative power of diffusion models with a robust, interpretable, and adaptable modular architecture.