ReMix: Optimizing Data Mixtures for Large Scale Imitation Learning

Auteurs: Joey Hejna, Chethan Bhateja, Yichen Jiang, Karl Pertsch, Dorsa Sadigh · CoRL 2024

Licences

• Accès libre
• Article: CC BY 4.0

Tout ce qui suit s'appuie sur l'article complet sous sa licence CC BY 4.0.

En bref

ReMix applique une optimisation distributionnellement robuste par groupes à la perte excédentaire d'une politique pour apprendre dans quelle proportion échantillonner chaque jeu de données robotique, et bat sur robots réels les poids uniformes et les poids experts de RT-X.

Pourquoi c'est important

Les équipes qui pré-entraînent sur des données robotiques mutualisées peuvent remplacer des poids d'échantillonnage réglés à la main par des poids calculés à partir des données et, dans les tests rapportés, écarter 75 % des données avec une perte de performance limitée. Les poids appris montrent aussi quels jeux de données sous-pondérer et pourquoi, par exemple les données Kuka collectées de façon autonome, ce qui aide à décider quelles données collecter ou conserver.

Problème

Les grands jeux de données robotiques comme Open X-Embodiment regroupent des données issues de robots, caméras, espaces d'actions et contrôleurs différents, et pourtant les mélanges utilisés pour entraîner des politiques généralistes comme RT-X, Octo et OpenVLA ont été choisis à la main. Les méthodes de filtrage de données issues de la vision et du langage se transposent mal, car les n-grammes ne s'appliquent pas et les plongements visuels ignorent les actions et la nature séquentielle des épisodes.

Solution proposée

ReMix formule la curation comme un problème min-max : une politique discrète minimise sa perte excédentaire de clonage de comportement par rapport à un modèle de référence, tandis que les poids de domaine sont mis à jour par montée de gradient exponentiée pour la maximiser, de sorte que le poids va aux domaines sur lesquels la politique peut encore progresser. Pour rendre les pertes comparables d'un robot à l'autre, les actions sont normalisées de façon gaussienne par domaine et discrétisées en intervalles, et le point de contrôle de référence est choisi par arrêt précoce avant qu'il ne surapprenne un domaine. Les poids moyennés sur l'entraînement sont ensuite réutilisés pour entraîner des politiques de diffusion distinctes ou pour sous-échantillonner les données.

Idées clés

• Utiliser la perte excédentaire par rapport à une politique de référence, plutôt que la perte brute, sous-pondère à la fois les domaines déjà bien ajustés et ceux trop difficiles à ajuster.
• Les différences d'échelle des actions faussent l'optimisation robuste : dans un test jouet avec une normalisation par bornes, 94,3 % du poids allait à un domaine d'actions purement aléatoires ; avec une normalisation gaussienne par domaine, les données réelles recevaient 84,2 %.
• Les pertes d'action continues L1/L2 sont influencées par les valeurs aberrantes et les actions multimodales, si bien que discrétiser les actions donne des poids plus fiables même lorsque la politique finale est un modèle de diffusion continu.
• Les jeux de données robotiques (environ 10k-100k démonstrations) sont assez petits pour que le modèle de référence surapprenne ; utiliser un point de contrôle 50K pas après l'arrêt précoce a réduit la performance de plus de 15 % en moyenne.

Principaux résultats

• Sur le mélange de données RT-X, les politiques entraînées avec les poids ReMix ont surpassé les poids uniformes de 38 % et les poids d'experts humains de RT-X de 32 % en moyenne sur quatre tâches réelles WidowX et deux tâches réelles Franka ; les poids d'experts battaient les poids uniformes de 6 %.
• Sur des sous-ensembles de 25 % du mélange RT-X, ReMix a conservé sa performance tandis que les poids humains chutaient de plus de 10 % ; ReMix a conservé 72 % des données UR5 et 12 % de Kuka, contre 30 % et 24 % pour les poids humains.
• ReMix a fait passer le petit jeu de données Toto de 3,42 % (uniforme) à 16,3 % et réduit Kuka de 24,9 % à 12,1 %.
• Sur le jeu de données Bridge V2 complet, il n'y avait pas de différence notable entre ReMix et les poids uniformes ; sur des sous-ensembles de 25 % de Bridge, ReMix a surpassé les poids uniformes et la référence SSP fondée sur des plongements d'images.
• Sur une tâche de cube WidowX avec un objet distracteur, les poids ReMix ont réussi 7/10 contre 0/10 pour les poids humains comme pour les poids uniformes.
• En simulation (Robomimic NutAssemblySquare réparti par opérateur), ReMix a obtenu 77/100 contre 53/100 pour les poids uniformes avec un sous-échantillonnage de 50 %, et 59/100 contre 39/100 à 25 %.

Limites

• Les auteurs notent que l'évaluation réelle n'a porté que sur deux bras, WidowX et Franka, et laissent d'autres incarnations robotiques à des travaux futurs.
• ReMix surpondère des jeux de données aux distributions d'actions anormales et multimodales comme Toto, ce qui, selon les auteurs, n'est pas nécessairement souhaitable.
• Le calcul des poids exige deux entraînements sur l'ensemble des données, l'un pour le modèle de référence et l'autre pour l'optimisation robuste.
• Testé sur Bridge V2 et sur le sous-ensemble RT-X d'OpenX, et non sur le jeu de données OpenX complet de plus de 2M épisodes. La section 4.1 décrit Bridge V2 comme comptant 50k démonstrations, tandis que l'annexe B.2 indique qu'environ 45 000 trajectoires ont été utilisées sur les 60K du jeu de données complet.
• Les politiques entraînées uniquement sur les données de pré-entraînement n'ont pas généralisé en zero-shot aux dispositifs des auteurs, de sorte que chaque évaluation repose sur un co-entraînement avec 25 démonstrations dans le domaine par tâche.

Impact sur le domaine

Finaliste du CoRL 2024 Outstanding Paper Award et cité 74 fois. Les auteurs publient le code sur github.com/jhejna/remix. L'article publie aussi ses poids de domaine appris pour le mélange RT-X (tableau 2) et pour les 32 domaines Bridge V2 (tableau 5).

Résumé

Des jeux de données d'apprentissage par imitation de plus en plus volumineux sont collectés dans le but d'entraîner des modèles de fondation pour la robotique. Pourtant, alors que la sélection des données a été d'une importance capitale en vision et en traitement du langage naturel, peu de travaux en robotique se sont demandé sur quelles données ces modèles devraient réellement être entraînés. Dans ce travail, nous étudions comment pondérer différents sous-ensembles, ou « domaines », de jeux de données robotiques pour le pré-entraînement de modèles de fondation robotiques. Concrètement, nous utilisons l'optimisation distributionnellement robuste (DRO) pour maximiser la performance dans le pire cas sur l'ensemble des domaines en aval possibles. Notre méthode, Re-Mix, traite le large éventail de difficultés qui surgissent lorsqu'on applique la DRO aux jeux de données robotiques, notamment la variabilité des espaces d'actions et de la dynamique d'un jeu de données à l'autre. Re-Mix recourt à l'arrêt précoce, à la normalisation des actions et à la discrétisation pour contrer ces problèmes. Grâce à des expériences approfondies sur le plus grand jeu de données open source de manipulation robotique, le jeu de données Open X-Embodiment, nous démontrons que la curation des données peut avoir un impact démesuré sur la performance en aval. Plus précisément, les poids de domaine appris par Re-Mix surpassent les poids uniformes de 38 % en moyenne et les poids choisis par des humains de 32 % sur les jeux de données utilisés pour entraîner des politiques robotiques généralistes existantes, à savoir les modèles RT-X.

Liens

• arXiv 2408.14037: https://arxiv.org/abs/2408.14037
• DOI 10.48550/arXiv.2408.14037: https://doi.org/10.48550/arXiv.2408.14037
• PDF: https://raw.githubusercontent.com/mlresearch/v270/main/assets/hejna25a/hejna25a.pdf
• Code: https://github.com/jhejna/remix