Résumé Technique : DataMIL – Sélection de Données pour l'Apprentissage par Imitation Robotique avec des Datamodels

1. Énoncé du Problème

La communauté de la robotique s'est récemment tournée vers l'entraînement de politiques de fondation généralistes à l'aide de jeux de données massifs et diversifiés (ex. : Open-X Embodiment). Bien que ces politiques atteignent de fortes performances moyennes sur des tâches variées, elles sont fréquemment moins performantes sur des tâches individuelles spécialisées, ce qui nécessite un ajustement fin (fine-tuning) sur de nouvelles données spécifiques à la tâche.

Le défi central abordé par ce travail est la sélection de données : identifier quels sous-ensembles de grands jeux de données antérieurs, lorsqu'ils sont combinés à des démonstrations limitées spécifiques à une tâche, produisent les meilleures politiques spécialisées.

• Limites des méthodes actuelles : Une sélection naïve basée sur des notions humaines de qualité (ex. : similitude sémantique, similitude visuelle ou proximité état-action) échoue souvent. Ces heuristiques supposent que les données les plus « similaires » sont les plus « utiles », ignorant l'impact réel d'un point de donnée sur la performance de la politique en aval.
• La barrière de l'optimisation : Idéalement, on sélectionnerait les données en testant de manière exhaustive des sous-ensembles via un réentraînement et en évaluant la politique résultante. Cependant, en robotique, l'évaluation d'une politique nécessite des déploiements (rollouts) dans le monde réel, ce qui est chronophage, dangereux et informatiquement infaisable pour une recherche à grande échelle.
• La lacune : Les cadres d'attribution de données existants comme les datamodels (qui prédisent la performance d'un modèle en fonction des données d'entraînement) ont réussi dans le NLP et la vision par ordinateur, mais n'ont pas été directement applicables à la robotique en raison de la nature non différentiable et coûteuse de l'évaluation des politiques en conditions réelles.

2. Méthodologie : DataMIL

Les auteurs introduisent DataMIL (Datamodels for Imitation Learning), un cadre qui étend le paradigme des datamodels à la robotique. DataMIL raisonne sur la sélection de données de manière de bout en bout, en utilisant la politique elle-même pour identifier les points de données qui améliorent la performance, plutôt que de s'appuyer sur des heuristiques statiques.

2.1 Cadre Fondamental

DataMIL traite l'algorithme d'apprentissage de la politique comme une « boîte noire » et entraîne un estimateur (un datamodel) pour prédire la performance d'une politique entraînée sur un sous-ensemble spécifique de données.

• Objectif : Trouver un sous-ensemble $D' \subset D$ qui maximise la métrique cible $M(A(D'))$, où $A$ est l'algorithme d'apprentissage et $M$ est la métrique de performance.
• Formulation du Datamodel : Au lieu d'entraîner la politique, DataMIL apprend une fonction $\hat{f}(D') \approx M(A(D'))$. Pour les datamodels linéaires, cela revient à assigner un score scalaire $\tau(z_i)$ à chaque point de donnée $z_i$, où la performance prédite est la somme des scores du sous-ensemble sélectionné.

2.2 Adaptations Clés pour la Robotique

Pour rendre les datamodels tractables pour la robotique, les auteurs introduisent trois modifications critiques :

1. Fonction de Perte de Substitution (Éviter les Rollouts) :
   Les déploiements dans le monde réel sont non-différentiables et coûteux. DataMIL remplace la véritable métrique de taux de succès par une métrique proxy $\tilde{M}$, définie comme le négatif de la perte de clonage de comportement (Behavior Cloning - BC) sur un ensemble de démonstrations cibles tenues à l'écart :
   $$\tilde{M}(\pi, D_{target}) = \frac{1}{|D_{target}|} \sum_{(s,a) \in D_{target}} -L_{BC}(\pi(s), a)$$
   Cette métrique est entièrement différentiable et ne nécessite aucune interaction avec l'environnement, permettant l'utilisation de méthodes d'estimation basées sur le gradient.

2. Estimateurs de Datamodels :
   Le document emploie deux méthodes pour estimer les scores d'influence $\tau(z_i)$ :

• Estimateur de Régression : Échantillonne des sous-ensembles aléatoires de données, entraîne des politiques, évalue la métrique proxy, et ajuste un modèle linéaire pour prédire la performance en fonction de la présence des données. C'est précis mais très coûteux en calcul.
• Estimateur basé sur le Métagradient : Utilise des fonctions d'influence et des métagradients pour calculer la dérivée de la métrique cible par rapport aux poids des données. Cette méthode est plus efficace pour les modèles volumineux (comme Octo) et évite de devoir réentraîner les politiques pour chaque sous-ensemble.

3. Clustering et Atténuation du Décalage de Distribution :

• Clustering : Pour réduire le bruit dans l'estimation de l'influence, les paires état-action individuelles sont regroupées en clusters temporels (sous-trajectoires ou trajectoires complètes). La granularité optimale du cluster dépend de la taille du jeu de données ; les grands jeux de données (comme OXE) bénéficient de l'agrégation au niveau de la trajectoire.
• Décalage de Distribution : Pour éviter des estimations médiocres causées par des décalages entre les domaines antérieurs et cibles, une petite fraction de données de tâches cibles est incluse lors de la phase d'estimation du datamodel afin d'aligner le processus d'apprentissage avec le domaine cible.

2.3 Pipeline d'Entraînement

Une fois les scores d'influence estimés :

1. Les $x\%$ supérieurs des données antérieures (influence positive la plus élevée) sont sélectionnés pour former $D_{sel}$.
2. La politique finale est entraînée via un co-entraînement sur les données cibles $D_{target}$ et les données sélectionnées $D_{sel}$. À chaque étape, un batch est échantillonné de $D_{target}$ avec une probabilité $\alpha$ et de $D_{sel}$ avec une probabilité $1-\alpha$.

3. Contributions Clés

• Extension des Datamodels à la Robotique : DataMIL est le premier cadre qui applique avec succès les datamodels à l'apprentissage par imitation en robotique, surmontant la barrière des rollouts coûteux en conditions réelles grâce à une perte de substitution différentiable.
• Conscience de la Performance de Bout en Bout : Contrairement aux méthodes heuristiques (similitude visuelle, de mouvement ou d'état), DataMIL optimise directement l'impact sur le succès de la politique, lui permettant de sélectionner les données qui améliorent la performance et de rejeter celles qui la dégradent, même si les données paraissent visuellement similaires.
• Estimateurs Scalables : L'introduction de l'estimateur basé sur le métagradient permet au cadre de passer à l'échelle avec des politiques complexes et volumineuses (ex. : Octo) et des jeux de données massifs (ex. : OXE) là où les approches basées sur la régression seraient informatiquement prohibitives.
• Sélection Cross-Embodiment et Multi-Tâche : La méthode démontre sa capacité à sélectionner des données utiles pour des tâches et des embodiments de robots non présents dans le jeu de données antérieur, et à constituer des jeux de données qui soutiennent simultanément plusieurs objectifs en aval.

4. Résultats Expérimentaux

Les auteurs valident DataMIL à travers plus de 60 tâches en simulation et en conditions réelles :

• MetaWorld (50 Tâches) : DataMIL a obtenu un gain de performance de 10 % par rapport aux meilleures méthodes de référence (Behavior Retrieval, Flow Retrieval, STRAP). Il a réussi à filtrer les démonstrations bruitées et sous-optimales issues de l'exploration autonome, là où les bases de référence basées sur la similitude ont souvent récupéré des données non pertinentes ou de faible qualité.
• LIBERO (10 Tâches) : En utilisant la politique transformer Octo, DataMIL a systématiquement surpassé les bases de référence sur diverses tâches à long horizon. Bien que les bases de référence de similitude visuelle aient bien performé sur certaines tâches spécifiques, DataMIL a offert la performance moyenne la plus robuste.
• Open-X Embodiment (OXE) (Monde Réel) :
  • Franka-Ball & Franka-Pouch : DataMIL a sélectionné des données qui ont considérablement amélioré les taux de succès par rapport à la sélection aléatoire et aux bases de référence heuristiques.
  • Tiago-Sink : Dans un contexte complexe de cross-embodiment (sélectionner des données pour un robot Tiago à partir d'un jeu de données ne contenant aucune donnée Tiago), DataMIL a réussi à identifier des démonstrations pertinentes provenant d'autres robots (ex. : RT-1, BC-Z) partageant « l'essence » de la tâche (manipulation sur table), là où les bases de référence ont échoué à généraliser.
  • Droid-Multitask : DataMIL a constitué un jeu de données unique qui a amélioré les performances sur trois tâches distinctes simultanément, surpassant les bases de référence en taux de succès moyen.

Résultats Qualitatifs :

• DataMIL sélectionne un mélange diversifié de jeux de données plutôt que de trop compter sur une seule source (contrairement aux bases de référence qui récupèrent souvent plus de 80 % d'un seul jeu de données).
• Il distingue efficacement les données utiles des données nuisibles, même lorsqu'elles sont visuellement identiques (ex. : même état mais distributions d'actions différentes), un mode de défaillance courant dans la récupération par similitude.

5. Signification et Revendications

L'article affirme que la sélection de données de bout en bout et sensible à la performance est cruciale pour libérer le potentiel des grands jeux de données antérieurs en robotique. Les auteurs soutiennent que si les grands modèles de fondation fournissent une base solide, leur spécialisation dépend fortement de la qualité des données d'ajustement fin.

• Revendications Modestes sur la Scalabilité : Les auteurs reconnaissent que l'estimation des datamodels reste coûteuse en calcul (plusieurs fois le coût de l'entraînement d'une politique sur toutes les données), bien que l'approche par métagradient atténue ce problème. Ils notent que le réglage des hyperparamètres (ex. : taille du cluster, ratio de sélection) manque actuellement d'intuition théorique forte.
• Portée : Le travail se concentre principalement sur la spécialisation par tâche unique, le réglage Droid-Multitask servant de étape préliminaire vers la curation multi-tâche.
• Conclusion : DataMIL démontre que passer des heuristiques définies par l'humain à une sélection pilotée par les données et consciente du modèle permet d'obtenir des gains constants de taux de succès, particulièrement dans des scénarios complexes, hétérogènes et réels où les méthodes de récupération traditionnelles échouent.