Equivariant Diffusion Policy

Auteurs: Dian Wang, Stephen Hart, David Surovik, Tarik Kelestemur, Haojie Huang, Haibo Zhao, Mark Yeatman, Jiuguang Wang, Robin Walters, Robert Platt · 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 ; sa figure est montrée en citant les auteurs.

En bref

Equivariant Diffusion Policy intègre la symétrie de rotation SO(2) dans le débruiteur de Diffusion Policy, augmentant la réussite moyenne en simulation de 21,9 points avec 100 démonstrations et apprenant des tâches réelles à partir de 20-60 démonstrations.

Pourquoi c'est important

Les démonstrations réelles coûtent cher, et intégrer la symétrie de rotation sur table a permis à la politique de réussir six tâches réelles, dont une tâche à long horizon de cuisson de bagel, avec 20 à 60 démonstrations, là où une référence Diffusion Policy voxelique échouait le plus souvent. Comme l'analyse couvre des actions de pince SE(3) complètes, elle s'applique aux tâches nécessitant une rotation hors plan, qu'une variante SE(2) ne pouvait pas résoudre.

Problème

Diffusion Policy capture un comportement de démonstration multimodal, mais doit apprendre une fonction de débruitage sur les observations, les actions bruitées et les niveaux de bruit, ce qui est plus difficile que d'apprendre une politique directe de l'état vers l'action et demande plus de données. Les politiques de manipulation équivariantes antérieures utilisaient soit des espaces d'actions SE(3) en boucle ouverte, soit SE(2) en boucle fermée, et aucune n'appliquait la symétrie aux politiques de diffusion.

Solution proposée

Les auteurs prouvent que la fonction de prédiction du bruit est SO(2)-équivariante dès lors que la politique experte l'est, et établissent comment la rotation autour de l'axe de gravité agit sur des actions de pince 6-DoF complètes, en commande de pose absolue comme relative. Le débruiteur utilise des encodeurs d'observation et d'action équivariants (construits avec la bibliothèque escnn) qui produisent des plongements pour un groupe de rotations discret, un U-Net temporel 1D partagé appliqué à chaque élément du groupe, et un décodeur de bruit équivariant. Les entrées peuvent être des images RGB ou une grille de voxels de 64×64×64.

Idées clés

• Si la politique experte est SO(2)-équivariante, la fonction de prédiction du bruit de la diffusion l'est aussi, ce qui fait des réseaux équivariants un biais inductif correct pour le débruiteur.
• Les gains sont les plus importants avec peu de démonstrations et sur les tâches où les poses initiales des objets sont largement aléatoires ; sur les tâches Robomimic peu aléatoires, la méthode était équivalente ou légèrement meilleure que Diffusion Policy.
• Dans les ablations, la structure équivariante comptait davantage que l'entrée voxelique : avec 100 démonstrations, supprimer l'équivariance réduisait la réussite moyenne de 17,6 points et supprimer les voxels de 10,3.
• L'augmentation de données par rotation sur un CNN standard aidait sur des tâches simples comme Stack, mais restait bien en deçà du réseau équivariant sur les plus difficiles.
• Un espace d'actions SE(3) complet compte : une variante SE(2) a obtenu 0,0 % sur Coffee Preparation contre 85,3 % pour SE(3), et 12,7 % contre 40,0 % sur Threading.

Principaux résultats

• Sur 12 tâches MimicGen en commande absolue avec 100 démonstrations, la version voxelique a atteint en moyenne 63,9 % de réussite contre 42,0 % pour Diffusion Policy (DiffPo-C), soit un gain de 21,9 points ; la version RGB a atteint en moyenne 53,7 %.
• Avec 200 démonstrations, la version voxelique (72,6 %) a battu toutes les références entraînées avec 1000 démonstrations (meilleure 71,4 %).
• En commande absolue, la version voxelique a surpassé les références dans 11 environnements sur 12, l'exception étant Hammer Cleanup D1.
• Sur un bras Franka réel, la réussite face à une référence Diffusion Policy voxelique était : Oven Opening 95 % contre 60 % (20 démonstrations) ; Banana in Bowl 95 % contre 30 %, Letter Alignment 95 % contre 0 %, Trash Sweeping 90 % contre 5 % (40 démonstrations chacune) ; Hammer to Drawer 85 % contre 5 % (60 démonstrations) ; Bagel Baking 80 % contre 10 % (58 démonstrations).
• Tous les échecs de Bagel Baking provenaient des butées articulaires du robot, et la politique a généralisé en zero-shot à des rotations du four jamais vues, sauf une pose bloquée par les butées articulaires.

Limites

• Les auteurs notent que l'équivariance n'est que partiellement exploitée en raison d'un défaut de symétrie dans le système de vision ; le bras visible dans la grille de voxels et le bruit de la caméra peuvent briser la symétrie.
• Une « équivariance incorrecte » peut dégrader la performance lorsque la symétrie du modèle entre en conflit avec les démonstrations.
• La théorie des actions SO(2) n'est montrée qu'avec des politiques de diffusion, et non avec d'autres chaînes comme BC RNN, qui s'en sortait bien en commande relative.
• La méthode n'est pas encore étendue à la navigation, à la locomotion ni à la manipulation mobile.
• Avec une entrée RGB, la caméra en vue agent n'est pas en vue de dessus, si bien que la rotation de l'image ne correspond pas à la rotation réelle de l'état ; en commande relative, la version RGB n'était que marginalement meilleure que les références.

Impact sur le domaine

Finaliste du CoRL 2024 Outstanding Paper Award et cité 107 fois. L'article renvoie à une page de projet sur equidiff.github.io. Il ne nomme aucun dépôt de code dans son texte.

Résumé

Des travaux récents ont montré que les modèles de diffusion constituent une approche efficace pour apprendre les distributions multimodales issues des données de démonstration en clonage de comportement. Cependant, un inconvénient de cette approche est la nécessité d'apprendre une fonction de débruitage, nettement plus complexe que l'apprentissage d'une politique explicite. Dans ce travail, nous proposons Equivariant Diffusion Policy, une nouvelle méthode d'apprentissage de politiques de diffusion qui exploite les symétries du domaine pour obtenir une meilleure efficacité en échantillons et une meilleure généralisation de la fonction de débruitage. Nous analysons théoriquement la symétrie SO(2) de la commande 6-DoF complète et caractérisons les conditions dans lesquelles un modèle de diffusion est SO(2)-équivariant. Nous évaluons en outre empiriquement la méthode sur un ensemble de 12 tâches de simulation dans MimicGen et montrons qu'elle obtient un taux de réussite supérieur en moyenne de 21,9 % à celui de la référence Diffusion Policy. Nous évaluons aussi la méthode sur un système réel pour montrer que des politiques efficaces peuvent être apprises avec relativement peu d'échantillons d'entraînement, ce dont la référence Diffusion Policy est incapable.

Liens

• arXiv 2407.01812: https://arxiv.org/abs/2407.01812
• DOI 10.48550/arXiv.2407.01812: https://doi.org/10.48550/arXiv.2407.01812
• PDF: https://raw.githubusercontent.com/mlresearch/v270/main/assets/wang25a/wang25a.pdf
• Page du projet: https://equidiff.github.io/