Flexible Multitask Learning with Factorized Diffusion Policy

1. Probleemstelling

Het artikel adresseert de uitdagingen van multitask imitatielearning in de robotica, specifiek de moeilijkheid om zeer multimodale en diverse actie-distributies te modelleren.

• Beperkingen van Monolithische Modellen: Traditionele single-network policies (zoals de standaard Diffusion Policy) worstelen met generalisatie over diverse taken, vaak door onderfitting van complexe distributies of het falen om distincte gedragsmodi te representeren (bijvoorbeeld verschillende grijpstrategieën voor verschillende objecten).
• Beperkingen van Bestaande Modulaire Modellen: Hoewel Mixture-of-Experts (MoE) architecturen proberen taken te decomponeren, lijden ze aan:
  • Trainingsinstabiliteit: Moeite met het balanceren van expertgebruik.
  • Routing-problemen: Discrete expertselectie kan leiden tot "routing-imbalance" (sommige experts worden nooit gebruikt, anderen worden overmatig gebruikt).
  • Slechte Interpretatiebaarheid: Experts hebben vaak overlappende of onduidelijke rollen, en vaardigheden kunnen verspreid zijn over lagen in plaats van ontkoppeld te zijn.
  • Catastrofaal Vergeten: Aanpassing aan nieuwe taken vereist vaak het opnieuw trainen van het hele model, wat het risico op het verlies van eerder geleerde vaardigheden met zich meebrengt.

2. Methodologie: Factorized Diffusion Policy (FDP)

De auteurs stellen FDP voor, een modulaire policy-architectuur die de actie-distributie factoriseert in een compositie van gespecialiseerde diffusiemodellen.

A. Core Architectuur

• Factorisatie: In plaats van een enkele policy $p(a_t|o_t)$, modelleert FDP de actie-distributie als een gewogen product van $N$ component-distributies:
  $$p(a_t|o_t) \propto \prod_{i} p_i(a_t|o_t)^{w_{t,i}}$$
  Hierbij is elke $p_i$ een gespecialiseerd diffusiemodel (expert) dat een distinct gedragsmodus (sub-vaardigheid) vastlegt, en zijn $w_{t,i}$ observatie-afhankelijke gewichten.
• Compositional Sampling:
  • Router: Een lichtgewicht MLP neemt de observatie $o_t$ en geeft continue gewichten $\{w_{t,i}\}$ voor alle componenten.
  • Score Aggregatie: In tegenstelling tot discrete MoE (die één expert selecteert), aggregeert FDP de schattingen van scores (gradienten van de log-kans) van alle componenten:
    $$\nabla \log p(a_t|o_t) \approx \sum_{i} w_{t,i} \nabla \log p_i(a_t|o_t)$$
  • Denoising-proces: De iteratieve denoising-stap gebruikt deze gewogen som van scores om de generatie van de actie-trajectorie te sturen. Dit komt overeen met het samplen uit het product van distributies, wat actiegeneratie effectief vormgeeft als een constraint satisfaction problem.

B. Training en Adaptatie

• Gecombineerde Training: Alle diffusiecomponenten en de router worden end-to-end getraind met een Mean Squared Error (MSE) verlies op de geaggregeerde scorevoorspelling.
• Efficiënte Adaptatie (Upcycling): Om aan te passen aan een nieuwe taak zonder het hele model opnieuw te trainen:
  1. Een nieuwe diffusiecomponent wordt geïntroduceerd.
  2. Deze wordt geïnitieerd door upcycling (kopiëren van gewichten) van een bestaande component.
  3. Alleen de nieuwe component en de router worden fijn afgestemd; alle vorige componenten worden bevroren.
  4. Dit voorkomt catastrofaal vergeten en reduceert het aantal trainbare parameters aanzienlijk.

3. Belangrijkste Bijdragen

1. Nieuwe Modulaire Architectuur: Introductie van FDP, dat continue score-aggregatie gebruikt in plaats van discrete expertselectie. Dit zorgt voor stabiele training, voorkomt routing-imbalance en bevordert duidelijke specialisatie onder de componenten.
2. Probabilistische Interpretatie: Gebaseerd op Energy-Based Models (EBM) en score matching, wat een principiële formulering biedt waarbij de samengestelde policy de doorsnede van gedragsbeperkingen vertegenwoordigt.
3. Schalbare Adaptatiestrategie: Voorstellen van een "selectieve tuning"-mechanisme met upcycling en bevriezing, wat efficiënte vaardigheidsuitbreiding mogelijk maakt met minimale data en zonder catastrofaal vergeten.
4. Empirische Validatie: Aangetoond superieure prestaties ten opzichte van sterke baselines (Monolithic DP, Sparse Diffusion Policy, en MoDE) in zowel simulatie als real-world settings.

4. Experimentele Resultaten

De auteurs evalueerden FDP op MetaWorld, RLBench, LIBERO, en real-world robotmanipulatie (UR5e-arm).

• Multitask Prestaties:

  • FDP presteerde consistent beter dan monolithische Diffusion Policy (DP) en modulaire baselines (SDP, MoDE) over 30+ taken.
  • MetaWorld: FDP behaalde 74,8% gemiddeld succes versus 70,8% (DP) en 69,8% (SDP).
  • RLBench: FDP behaalde 63,9% gemiddeld succes versus 45,6% (DP) en 42,1% (SDP).
  • Real-World: FDP toonde hogere successpercentages in complexe taken (bijvoorbeeld mokken ophangen, gekleurde blokjes oppakken) waar baselines faalden door overfitting of onvermogen om multimodale distributies te vangen.

• Taakadaptatie & Transfer:

  • Data-efficiëntie: Bij aanpassing aan nieuwe taken met beperkte data (10–25 demos) presteerde FDP met de "Add New Module"-strategie beter dan full-parameter fine-tuning en partial fine-tuning.
  • Continual Learning: In een continual learning-benchmark van 12 taken (LIBERO) behield FDP hoge prestaties terwijl nieuwe experts werden toegevoegd, terwijl andere methoden degradeerden.
  • Kennisbehoud: Het bevriezen van oude componenten behield effectief eerdere vaardigheden. Met een minimale replay buffer (5 demos/tak) kon het model bijna 100% van de oorspronkelijke prestaties behouden na adaptatie.

• Analyse & Scaling:

  • Component Scaling: De prestaties verbeterden naarmate het aantal componenten toenam van 2 naar 4, en stabiliseerde daarna.
  • Task Scaling: Het prestatieverschil tussen FDP en monolithische modellen nam toe naarmate het aantal taken toenam, wat de schaalbaarheid van FDP benadrukt.
  • Interpretatiebaarheid: Visualisatie toonde aan dat individuele componenten zich specialiseerden in distincte sub-vaardigheden (bijvoorbeeld uitlijning, nadering, grijpen) in plaats van ruisniveaus of willekeurige lagen.
  • Inferentiesnelheid: FDP ondersteunt "partial reconstruction" (alleen de top-$k$ componenten gebruiken), wat een 2x versnelling in inferentie oplevert met slechts een ~19% prestatieverlies.

5. Betekenis

• Robotlearning: FDP biedt een praktische oplossing voor het inzetten van robots in dynamische omgevingen waar ze snel nieuwe vaardigheden moeten leren zonder oude te vergeten.
• Stabiliteit: Door discrete routing te vervangen door continue score-compositie, lost FDP de trainingsinstabiliteitsproblemen op die MoE-gebaseerde diffusiepolities hebben geplaagd.
• Interpretatiebaarheid: De factorisatie stelt onderzoekers in staat om specifieke sub-vaardigheden (bijvoorbeeld "grijpen" versus "naderen") binnen de policy te inspecteren en te begrijpen, wat helpt bij debugging en veiligheidsverificatie.
• Toekomstige Richtingen: Het kader opent wegen voor heterogene moduleontwerpen (mixen van architecturen) en levenslange leersystemen waarbij robots continu een bibliotheek van gespecialiseerde vaardigheden opbouwen.

Kortom, FDP vertegenwoordigt een aanzienlijke vooruitgang in robot imitatielearning door het generatieve vermogen van diffusiemodellen te combineren met een robuuste, interpreteerbare en aanpasbare modulaire architectuur.