Flexible Multitask Learning with Factorized Diffusion Policy

1. Énoncé du problème

Le papier aborde les défis de l'apprentissage par imitation multitâche en robotique, spécifiquement la difficulté de modéliser des distributions d'actions hautement multimodales et diversifiées.

• Limites des modèles monolithiques : Les politiques traditionnelles basées sur un seul réseau (par exemple, la Diffusion Policy standard) peinent à généraliser à travers des tâches diverses, sous-ajustant souvent des distributions complexes ou échouant à représenter des modes comportementaux distincts (par exemple, différentes stratégies de préhension pour différents objets).
• Limites des modèles modulaires existants : Bien que les architectures de type Mélange d'Experts (MoE) tentent de décomposer les tâches, elles souffrent de :
  • Instabilité de l'entraînement : Difficulté à équilibrer l'utilisation des experts.
  • Problèmes de routage : La sélection discrète d'experts peut entraîner un "déséquilibre de routage" (certains experts ne sont jamais utilisés, d'autres sont surutilisés).
  • Mauvaise interprétabilité : Les experts ont souvent des rôles qui se chevauchent ou sont flous, et les compétences peuvent être distribuées à travers les couches plutôt que d'être désintriquées.
  • Oubli catastrophique : L'adaptation à de nouvelles tâches nécessite souvent de réentraîner l'ensemble du modèle, risquant la perte des compétences précédemment acquises.

2. Méthodologie : Politique de Diffusion Factorisée (FDP)

Les auteurs proposent FDP, une architecture de politique modulaire qui factorise la distribution d'actions en une composition de modèles de diffusion spécialisés.

A. Architecture de base

• Factorisation : Au lieu d'une politique unique $p(a_t|o_t)$, FDP modélise la distribution d'actions comme un produit pondéré de $N$ distributions composantes :
  $$p(a_t|o_t) \propto \prod_{i} p_i(a_t|o_t)^{w_{t,i}}$$
  Ici, chaque $p_i$ est un modèle de diffusion spécialisé (expert) capturant un mode comportemental distinct (sous-compétence), et $w_{t,i}$ sont des poids dépendants de l'observation.
• Échantillonnage compositionnel :
  • Routeur : Un MLP léger prend l'observation $o_t$ et produit des poids continus $\{w_{t,i}\}$ pour toutes les composantes.
  • Agrégation de scores : Contrairement au MoE discret (qui sélectionne un expert), FDP agrège les estimations de score (gradients de la log-probabilité) de toutes les composantes :
    $$\nabla \log p(a_t|o_t) \approx \sum_{i} w_{t,i} \nabla \log p_i(a_t|o_t)$$
  • Processus de débruitage : L'étape itérative de débruitage utilise cette somme pondérée de scores pour guider la génération de la trajectoire d'action. Cela correspond à un échantillonnage à partir du produit de distributions, encadrant efficacement la génération d'actions comme un problème de satisfaction de contraintes.

B. Entraînement et adaptation

• Entraînement conjoint : Toutes les composantes de diffusion et le routeur sont entraînés de bout en bout en utilisant une perte d'erreur quadratique moyenne (MSE) sur la prédiction du score agrégé.
• Adaptation efficace (Upcycling) : Pour s'adapter à une nouvelle tâche sans réentraîner tout le modèle :
  1. Une nouvelle composante de diffusion est introduite.
  2. Elle est initialisée par upcycling (copie des poids) à partir d'une composante existante.
  3. Seule la nouvelle composante et le routeur sont affinés ; toutes les composantes précédentes sont figées.
  4. Cela empêche l'oubli catastrophique et réduit considérablement le nombre de paramètres entraînables.

3. Contributions clés

1. Architecture modulaire novatrice : Introduction de FDP, qui utilise une agrégation continue de scores plutôt qu'une sélection discrète d'experts. Cela assure un entraînement stable, prévient le déséquilibre de routage et favorise une spécialisation claire parmi les composantes.
2. Interprétation probabiliste : Fondée sur les modèles basés sur l'énergie (EBM) et le matching de scores, offrant une formulation principielle où la politique composée représente l'intersection des contraintes comportementales.
3. Stratégie d'adaptation évolutive : Proposition d'un mécanisme de "réglage sélectif" utilisant l'upcycling et le figement, permettant une expansion efficace des compétences avec un minimum de données et sans oubli catastrophique.
4. Validation empirique : Démonstration d'une performance supérieure par rapport à des bases de référence solides (Diffusion Policy monolithique, Sparse Diffusion Policy et MoDE) aussi bien en simulation que dans des environnements réels.

4. Résultats expérimentaux

Les auteurs ont évalué FDP sur MetaWorld, RLBench, LIBERO, et la manipulation robotique réelle (bras UR5e).

• Performance multitâche :

  • FDP a constamment surpassé la Diffusion Policy (DP) monolithique et les bases de référence modulaires (SDP, MoDE) sur plus de 30 tâches.
  • MetaWorld : FDP a atteint un taux de succès moyen de 74,8 % contre 70,8 % (DP) et 69,8 % (SDP).
  • RLBench : FDP a atteint un taux de succès moyen de 63,9 % contre 45,6 % (DP) et 42,1 % (SDP).
  • Monde réel : FDP a montré des taux de succès plus élevés dans des tâches complexes (par exemple, accrocher des tasses, saisir des cubes colorés) où les bases de référence échouaient en raison d'un surajustement ou de l'incapacité à capturer des distributions multimodales.

• Adaptation et transfert de tâches :

  • Efficacité des données : Lors de l'adaptation à de nouvelles tâches avec des données limitées (10–25 démonstrations), FDP utilisant la stratégie "Ajouter un nouveau module" a surpassé l'affinage de paramètres complets et l'affinage partiel.
  • Apprentissage continu : Dans un benchmark d'apprentissage continu de 12 tâches (LIBERO), FDP a maintenu une performance élevée à mesure que de nouveaux experts étaient ajoutés, tandis que d'autres méthodes se dégradaient.
  • Rétention des connaissances : Le figement des anciennes composantes a efficacement préservé les compétences antérieures. L'utilisation d'une mémoire tampon de replay minimale (5 démonstrations/tâche) a permis au modèle de conserver près de 100 % de la performance originale après adaptation.

• Analyse et mise à l'échelle :

  • Mise à l'échelle des composantes : La performance s'est améliorée à mesure que le nombre de composantes augmentait de 2 à 4, puis s'est stabilisée.
  • Mise à l'échelle des tâches : L'écart de performance entre FDP et les modèles monolithiques s'est élargi à mesure que le nombre de tâches augmentait, soulignant l'évolutivité de FDP.
  • Interprétabilité : La visualisation a montré que les composantes individuelles se spécialisaient dans des sous-compétences distinctes (par exemple, alignement, approche, préhension) plutôt que dans des niveaux de bruit ou des couches arbitraires.
  • Vitesse d'inférence : FDP prend en charge la "reconstruction partielle" (en utilisant uniquement les $k$ meilleures composantes), permettant une accélération de 2x de l'inférence avec une baisse de performance d'environ 19 %.

5. Importance

• Apprentissage robotique : FDP offre une solution pratique pour déployer des robots dans des environnements dynamiques où ils doivent apprendre de nouvelles compétences rapidement sans oublier les anciennes.
• Stabilité : En remplaçant le routage discret par une composition continue de scores, FDP résout les problèmes d'instabilité de l'entraînement qui ont affecté les politiques de diffusion basées sur MoE.
• Interprétabilité : La factorisation permet aux chercheurs d'inspecter et de comprendre des sous-compétences spécifiques (par exemple, "préhension" par rapport à "approche") au sein de la politique, aidant au débogage et à la vérification de la sécurité.
• Perspectives futures : Le cadre ouvre la voie à des conceptions de modules hétérogènes (mélange d'architectures) et à des systèmes d'apprentissage tout au long de la vie où les robots accumulent continuellement une bibliothèque de compétences spécialisées.

En résumé, FDP représente une avancée significative dans l'apprentissage par imitation robotique en combinant la puissance générative des modèles de diffusion avec une architecture modulaire robuste, interprétable et adaptable.