Unified World Models: Coupling Video and Action Diffusion for Pretraining on Large Robotic Datasets

Auteurs: Chuning Zhu, Raymond Yu, Siyuan Feng, Benjamin Burchfiel, Paarth Shah, Abhishek Gupta · RSS 2025

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

Un seul transformer de diffusion, avec des pas de temps de bruit distincts pour les actions et les images futures, joue le rôle de politique, de modèle de dynamique ou de prédicteur vidéo, et peut s'entraîner sur des vidéos sans actions.

Pourquoi c'est important

Un même modèle peut s'entraîner à la fois sur des données robotiques étiquetées en actions et sur des vidéos sans actions, bien plus faciles à obtenir. Lors de tests sur robot réel portant sur cinq tâches avec un dispositif Franka de type DROID, UWM pré-entraîné a obtenu la meilleure réussite sur chaque tâche, et le co-entraînement sur des vidéos supplémentaires l'a encore augmentée, y compris en présence de distractions visuelles.

Problème

L'apprentissage par imitation pour de grands modèles robotiques dépend de démonstrations d'experts de haute qualité, coûteuses et lentes à collecter par téléopération. Les vidéos montrant la dynamique du monde réel abondent, mais la plupart des méthodes actuelles d'apprentissage par imitation ne peuvent pas les utiliser faute d'étiquettes d'action, et le behavior cloning ne capture pas explicitement la dynamique temporelle présente dans les trajectoires.

Solution proposée

UWM couple un processus de diffusion d'actions et un processus de diffusion d'images futures dans un seul transformer, conditionné par l'observation courante, avec un pas de temps de diffusion indépendant pour chaque modalité. Fixer le pas de temps d'une modalité à T la marginalise et le fixer à 0 conditionne sur elle, de sorte que le même modèle échantillonne une politique p(a|o), une dynamique directe p(o′|o,a), une dynamique inverse p(a|o,o′) ou une prédiction vidéo p(o′|o). Les vidéos sans actions sont utilisées à l'entraînement en fixant le pas de temps d'action à T et en remplissant les actions manquantes avec du bruit gaussien.

Idées clés

• Des pas de temps de diffusion indépendants par modalité agissent comme un masquage partiel, de sorte qu'un seul réseau couvre les distributions marginales et conditionnelles sur les actions et les images futures.
• C'est la prédiction des images futures, et non la simple reconstruction de l'image courante, qui produit le gain : 0,86/0,76 de réussite sur Stack-Bowls/Block-Cabinet contre 0,70/0,66 pour la reconstruction de l'image courante et 0,48/0,60 pour Diffusion Policy.
• Les tokens de registre (8 par défaut) facilitent l'échange d'informations entre tokens d'action et tokens d'image, et le conditionnement AdaLN a surpassé l'attention croisée dans l'ablation sur LIBERO.
• UWM et Diffusion Policy obtiennent des résultats à peu près équivalents lorsqu'ils sont entraînés à partir de zéro ; UWM tire davantage profit du pré-entraînement.
• Le co-entraînement sur des vidéos robotiques du domaine a davantage aidé que le co-entraînement sur des vidéos d'activités humaines d'Internet (Kinetics-400 et Something-Something v2), bien que les deux battent les seules données robotiques.

Principaux résultats

• Robot réel, pré-entraîné sur 2 000 trajectoires DROID, réussite en distribution : UWM 0,86, 0,76, 0,78, 0,82 et 0,60 sur Stack-Bowls, Block-Cabinet, Paper-Towel, Hang-Towel et Rice-Cooker, contre 0,48, 0,60, 0,52, 0,64, 0,35 pour Diffusion Policy et 0,66, 0,66, 0,60, 0,66, 0,40 pour GR1.
• Le co-entraînement sur 2 000 trajectoires DROID supplémentaires utilisées comme vidéos sans actions a porté la réussite en distribution d'UWM à 0,92, 0,84, 0,86, 0,86 et 0,65, et la réussite hors distribution (quatre premières tâches) de 0,76/0,60/0,78/0,64 à 0,84/0,72/0,84/0,76.
• Scènes OOD catégorisées (éclairage, arrière-plan, encombrement) : UWM co-entraîné 21/30 et 15/30 sur Stack-Bowls et Block-Cabinet, contre 15/30 et 8/30 pour UWM pré-entraîné et 12/30 et 6/30 pour Diffusion Policy.
• LIBERO (pré-entraînement sur LIBERO-90, affinage sur cinq tâches LIBERO-10 avec décalages de distribution) : réussite moyenne de 0,79 ± 0,11 pour UWM contre 0,71 ± 0,12 pour DP, 0,57 ± 0,19 pour PAD et 0,58 ± 0,14 pour GR1.
• Mode de dynamique inverse suivant 50 trajectoires LIBERO dans la limite de la longueur de trajectoire : 0,73 et 0,45 de réussite sur Book-Caddy et Soup-Cheese, contre 0,37 et 0,33 pour le mode politique.
• Entraîner un UWM sur DROID pendant 100K pas prend 24 heures sur 4 GPU NVIDIA A100.

Limites

• Le modèle n'apprend pas encore à partir de vidéos humaines à grande échelle, de sorte que l'écart d'incarnation entre humain et robot n'est pas traité.
• Les prédictions de dynamique directe peuvent contenir des artefacts, ce qui peut réduire leur utilité pour la planification.
• Les auteurs attendent des gains supplémentaires d'une prédiction vidéo plus dense, que cette version n'utilise pas.
• Les gains par rapport aux références dans la configuration hors distribution de LIBERO sont bien plus faibles que sur le robot réel ; les auteurs suggèrent que la dynamique en simulation est plus simple.
• Rice-Cooker n'a été évalué que sur 20 initialisations proches de la distribution et pas dans des configurations OOD, et tous les tests sur robot réel utilisent un seul dispositif Franka.

Impact sur le domaine

L'article compte 244 citations. Il indique que les vidéos et le code sont disponibles sur sa page de projet, qui renvoie au dépôt de code WEIRDLabUW/unified-world-model sur GitHub.

Résumé

L'apprentissage par imitation s'est imposé comme une approche prometteuse pour construire des robots généralistes. Cependant, faire passer l'apprentissage par imitation à l'échelle pour de grands modèles de fondation robotiques reste difficile en raison de sa dépendance à des démonstrations d'experts de haute qualité. Parallèlement, de grandes quantités de données vidéo montrant un large éventail d'environnements et des comportements variés sont facilement disponibles. Ces données constituent une riche source d'informations sur la dynamique du monde réel et les interactions agent-environnement. Exploiter directement ces données pour l'apprentissage par imitation s'est toutefois révélé difficile en raison de l'absence d'annotations d'actions, requises par la plupart des méthodes actuelles. Dans ce travail, nous présentons Unified World Models (UWM), un cadre qui permet d'exploiter à la fois des données vidéo et des données d'action pour l'apprentissage de politiques. Plus précisément, un UWM intègre un processus de diffusion d'actions et un processus de diffusion vidéo dans une architecture transformer unifiée, où des pas de temps de diffusion indépendants régissent chaque modalité. Nous montrons qu'en contrôlant simplement chaque pas de temps de diffusion, UWM peut représenter de manière flexible une politique, une dynamique directe, une dynamique inverse et un générateur vidéo. Grâce à des expériences en simulation et dans le monde réel, nous montrons que : (1) UWM permet un pré-entraînement efficace sur des jeux de données robotiques multitâches à grande échelle avec des prédictions à la fois de dynamique et d'actions, ce qui donne des politiques plus généralisables et plus robustes que l'apprentissage par imitation, (2) UWM facilite naturellement l'apprentissage à partir de données vidéo sans actions grâce au contrôle indépendant des pas de temps de diffusion propres à chaque modalité, améliorant encore les performances des politiques affinées. Nos résultats suggèrent qu'UWM constitue une étape prometteuse vers l'exploitation de grands jeux de données hétérogènes pour un apprentissage robotique passant à l'échelle, et fournit une unification simple entre les paradigmes souvent disjoints de l'apprentissage par imitation et de la modélisation du monde. Les vidéos et le code sont disponibles sur https://weirdlabuw.github.io/uwm/

Liens

• arXiv 2504.02792: https://arxiv.org/abs/2504.02792
• DOI https://doi.org/10.15607/rss.2025.xxi.015: https://doi.org/https://doi.org/10.15607/rss.2025.xxi.015
• PDF: https://arxiv.org/pdf/2504.02792v1
• Code: https://github.com/WEIRDLabUW/unified-world-model
• Page du projet: https://weirdlabuw.github.io/uwm/