🤖 machine learning

WorldDiT: A Unified Diffusion Architecture for World and Action Modeling

WorldDiT introduit une architecture de transformeur de diffusion unifiée qui génère simultanément des actions continues et prédit les futures images visuelles, atteignant des performances de pointe sur de multiples bancs d'essai de manipulation robotique sans s'appuyer sur de grands modèles de vision-langage pré-entraînés.

Auteurs originaux : Sen Wang, R. Gnana Praveen, Bidhan Roy, Marcos Villagra

Publié 2026-07-28

📖 7 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Sen Wang, R. Gnana Praveen, Bidhan Roy, Marcos Villagra

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez un monde où les robots ne se contentent pas de suivre une liste rigide d'instructions, mais « rêvent » réellement de ce qui va se passer ensuite. C'est la frontière de l'apprentissage robotique, un domaine qui tente d'apprendre aux machines comment se déplacer dans notre monde réel et désordonné. Pendant longtemps, les plus grandes percées sont venues d'une recette spécifique : prendre un « cerveau » massif et pré-entraîné qui connaît déjà tout sur le langage et les images (comme une encyclopédie super intelligente capable aussi de voir), et y attacher un bras de robot. Cela fonctionne, mais c'est comme essayer d'apprendre à un chiot à rapporter la balle en commençant par le former pour qu'il devienne professeur ; c'est lourd, coûteux et difficile de savoir si le robot est en train d'apprendre la tâche ou s'il se contente de copier les notes du professeur. La grande question que se posent les scientifiques est la suivante : pouvons-nous construire un cerveau de robot plus petit, plus léger, qui apprend à bouger et à prédire l'avenir tout en même temps, sans avoir besoin de cette immense encyclopédie pré-chargée ?

Voici WorldDiT, un nouveau cerveau de robot conçu par des chercheurs de Bagel Labs. Considérez WorldDiT non pas comme un étudiant copiant un manuel, mais comme un jeune artiste apprenant à peindre en regardant un film de son propre futur. Au lieu de s'appuyer sur un gigantesque « Modèle Vision-Langage » pré-entraîné (une IA massive qui sait décrire des images et du texte) pour lui dire quoi faire, WorldDiT utilise un moteur unique et unifié qui fait deux choses simultanément. Premièrement, il détermine les prochains mouvements du bras du robot. Deuxièmement, il essaie de deviner ce que la caméra verra dans quelques secondes. C'est comme un joueur d'échecs qui ne se contente pas de planifier son prochain coup, mais qui visualise aussi le plateau trois coups à l'avance pour s'assurer que son plan est cohérent.

Les chercheurs ont entraîné ce système dans un monde simulé appelé LIBRO, où les robots doivent accomplir des tâches comme empiler des blocs ou déplacer des objets. Ils ont découvert que WorldDiT, qui est étonnamment petit (moins de 400 millions de paramètres, contre les milliards utilisés par les autres méthodes de pointe), peut prédire l'image future de la caméra et les actions du robot avec une précision incroyable. En fait, lorsqu'ils l'ont testé, le robot a réussi environ 95 % de ses tâches à travers quatre types de défis différents. La partie la plus excitante ? Il l'a fait sans le lourd « squelette » (backbone) pré-entraîné que la plupart des autres robots performants utilisent. Il a prouvé que vous n'avez pas besoin d'un cerveau géant et pré-chargé pour être un bon robot ; vous avez juste besoin d'un cerveau qui apprend à voir le futur tout en bougeant.

Le robot qui « rêve »

Alors, comment fonctionne réellement WorldDiT ? Imaginez que vous essayiez d'apprendre à jongler. La plupart des robots d'aujourd'hui sont comme des étudiants qui ont mémorisé mille vidéos de maîtres jongleurs et qui essaient de les imiter exactement. WorldDiT est différent. C'est comme un étudiant en jonglage qui est les yeux bandés, mais qui possède une boule de cristal.

Voici l'astuce : chaque fois que WorldDiT planifie un mouvement, il essaie aussi de « rêver » de ce que la caméra verra dans le futur. Il observe l'état actuel du robot, l'instruction qui lui a été donnée (comme « ramasse le bloc rouge ») et les dernières secondes de vidéo. Ensuite, il se pose deux questions en même temps :

« Que doit faire le robot ensuite ? »
« Si je fais cela, à quoi ressemblera l'image un instant plus tard ? »

Les chercheurs appellent cela un transformeur de diffusion (diffusion transformer). Si cela ressemble à du jargon de science-fiction, voyez cela comme un processus de « débruitage ». Imaginez une photo couverte de neige statique. WorldDiT part d'une supposition totalement aléatoire et bruitée de ce que devrait être le prochain mouvement du robot et de ce à quoi devrait ressembler l'image future. Ensuite, étape par étape, il nettoie ce bruit, affinant sa supposition jusqu'à trouver un chemin clair et fluide vers l'objectif. Il fait cela pour les actions du robot et pour les images futures.

La magie opère parce que ces deux tâches s'entraident. En essayant de prédire l'image future, le robot est forcé de comprendre la physique du monde. S'il prévoit de pousser un bloc, mais que son « rêve » du futur montre le bloc flottant dans les airs, il sait que son plan est erroné. Il doit ajuster son action jusqu'à ce que le « rêve » corresponde à la réalité. Cela crée une boucle de rétroaction où le robot apprend les règles du monde simplement en essayant de prédire ce qui vient ensuite.

Les résultats : Petit cerveau, grand succès

Les chercheurs ont testé WorldDiT dans quatre suites de simulation différentes (considérez-les comme quatre parcours d'obstacles différents : Spatial, Objet, But, et Long). Ils l'ont comparé à 24 autres méthodes célèbres d'apprentissage robotique.

Voici le plus frappant : WorldDiT est minuscule. Il possède environ 399 millions de paramètres (les « neurones » de son cerveau). Beaucoup des autres robots les plus performants utilisent des modèles de 1 milliard ou même 10 milliards de paramètres. Généralement, des cerveaux plus gros signifient de meilleures performances. Mais WorldDiT a brisé cette règle.

Dans ces simulations, WorldDiT a atteint un taux de réussite de 94,9 % en moyenne sur les quatre parcours.

Spatial : 98,0 % de réussite
Objet : 97,0 % de réussite
But : 92,8 % de réussite
Long : 91,8 % de réussite

Lorsque vous tracez ces résultats sur un graphique, WorldDiT se situe sur la « frontière de Pareto ». En langage clair, cela signifie qu'il est le robot le plus efficace de la pièce. Si vous voulez un robot aussi intelligent que WorldDT mais qui utilise un cerveau plus petit, vous n'en trouverez pas. Si vous voulez un robot aussi petit que WorldDT mais plus intelligent, vous n'en trouverez pas non plus. Tous les autres robots qui ont fait mieux que 94,9 % étaient nettement plus grands et plus lourds.

Pourquoi cela importe

L'article soutient explicitement l'idée qu'il n'est pas nécessaire d'utiliser un modèle Vision-Langage (VLA) massif et pré-entraîné pour obtenir un bon contrôle robotique. Pendant un certain temps, le domaine pensait : « Nous avons besoin d'un cerveau géant pour comprendre le monde ». WorldDT répond : « Pas nécessairement ».

Les chercheurs ont montré qu'en couplant la génération d'actions (le mouvement) avec la modélisation du monde (la prédiction du futur), une architecture unique et unifiée peut apprendre à contrôler un robot très efficacement. Ils ne se sont pas contentés de le suggérer ; ils l'ont mesuré à travers des milliers d'épisodes simulés. Le robot n'a pas seulement eu de la « chance » ; il a systématiquement surpassé des modèles beaucoup plus grands.

Cependant, il y a un bémol. Ces résultats proviennent de simulations (des jeux informatiques où vivent les robots). L'article note que bien que la performance soit forte, elle constitue une base pour de futures études. Nous ne savons pas encore si ce robot qui « rêve » fonctionnera parfaitement dans une cuisine réelle, désordonnée, avec une vraie gravité et des sols réellement glissants. Mais les résultats de la simulation sont un indice fort que nous pourrions être capables de construire des robots plus petits, moins chers et plus efficaces à l'avenir — des robots qui apprennent en imaginant le futur, plutôt qu'en se contentant de mémoriser le passé.

En fin de compte, WorldDT suggère que le secret d'un robot intelligent n'est pas seulement d'avoir une immense bibliothèque de connaissances, mais d'avoir un cerveau capable de jouer le futur dans son esprit pour prendre de meilleures décisions dès maintenant.

Résumé Technique : WorldDiT

Problématique

Les avancées récentes dans l'apprentissage de politiques robotiques reposent souvent sur l'attachement de la génération d'actions à de grands modèles de vision-langage pré-entraînés (VLM). Bien que cette approche tire parti d'une large compréhension de la perception et du langage, elle occulte les contributions spécifiques de l'échelle et de l'architecture du modèle. Les auteurs posent la question suivante : une architecture de transformeur de diffusion (Diffusion Transformer) unifiée peut-elle combiner la génération d'actions continues avec un objectif auxiliaire de prédiction visuelle future pour atteindre une performance de contrôle robuste sans s'appuyer sur un grand VLM pré-entraîné comme squelette d'action ?

Méthodologie : WorldDiT

WorldDiT est une architecture de transformeur de diffusion unifiée conçue pour coupler la génération d'actions avec la modélisation d'un monde visuel. Contrairement aux modèles monde-action autorégressifs qui génèrent des jetons d'action séquentiellement à l'aide d'un grand VLM, WorldDiT utilise un unique squelette de Diffusion Transformer (DiT) partagé pour modéliser à la fois les actions robotiques continues et des patchs RGB normalisés sélectionnés à partir de cadres de caméras futurs.

Architecture et Encodage

Le système utilise une approche hybride « gelée » et « entraînable » :

Encodeurs : Des encodeurs visuels gelés (MAE pour les images, CLIP pour le texte) et un encodeur d'état de robot entraînable (encodant séparément les états du bras et de la pince) traitent les entrées.
Tokenisation : Le système construit des jetons multimodaux comprenant :
- 34 jetons visuels (34 patchs provenant des caméras primaire et du poignet).
- 1 jeton de langage (issu de l'instruction).
- 1 jeton d'état (état du robot).
- Cela produit 36 jetons de contexte nets par créneau temporel.
Squelette (Backbone) : Le cœur est un squelette DiT partagé avec une profondeur de 4, une taille cachée de 1024, 16 têtes d'attention et 4 jetons de registre. Il utilise un AdaLN-Zero par jeton et une attention strictement bloc-causale et sûre pour l'action, garantissant que les requêtes d'action ne puissent pas lire les cibles RGB bruitées.

Objectif d'Entraînement : Flow Matching

WorldDiT est entraîné via le flow matching (appariement de flux). La fenêtre d'entraînement ( $N=10$ ) consiste en :

Contexte ( $C=3$ ) : Trois étapes observées contenant le langage, l'RGB multi-vues et les états du robot.
Cibles ( $H=7$ ) : Un bloc d'actions de 7 étapes et un patch RGB futur normalisé.

La cible visuelle future est construite à partir des cadres des caméras primaire et du poignet à un décalage temporel de $H$ . Chaque cadre de 224x224 est divisé en patchs de 16x16 ; 64 patchs espacés uniformément sont conservés de chaque caméra, normalisés par leur propre moyenne et variance, et concaténés pour former une cible de 128 jetons ( $Y_{rgb}$ ).

Le modèle reçoit des jetons de contexte propres et des jetons de cible corrompus (actions et patchs RGB) avec un bruit gaussien. Il prédit la vitesse de flux pour un chemin rectiligne du bruit vers la cible propre. La perte totale est une somme pondérée de la perte de vitesse d'action et de la perte de vitesse des patchs RGB :
$L_{total} = w_{action}L_{action}^{flow} + w_{rgb}L_{rgb}^{flow}$
Crucialement, la prédiction des patchs RGB sert de signal de supervision auxiliaire pendant l'entraînement mais est absente lors de l'inférence.

Inférence

Au déploiement, le système opère sur un horizon fuyant :

Le squelette unifié reçoit l'historique (langage, observations, états) jusqu'à l'étape actuelle.
Aucune cible RGB future ni étiquette d'action n'est fournie.
Le modèle génère un bloc d'actions de 7 étapes en intégrant numériquement le champ de vitesse d'action appris, d'un bruit gaussien ( $\tau=0$ ) vers l'état propre ( $\tau=1$ ), sur 20 étapes d'Euler.
Les trois premières actions du bloc prédit sont exécutées, l'état est observé, et la fenêtre est décalée pour la replanification.

Principales Contributions

Architecture Unifiée : Un seul transformeur de diffusion qui apprend conjointement la génération d'actions continues et la prédiction de patchs RGB futurs normalisés sans s'appuyer sur un grand squelette VLM pré-entraîné.
Conception Entraînement/Déploiement : Une configuration novatrice où la prédiction de patchs RGB fournit la supervision pendant l'entraînement mais est retirée du graphe d'inférence, permettant un chemin de déploiement compact dédié uniquement à l'action.
Baseline de Moins d'un Milliard de Paramètres : Une configuration avec environ 399 millions de paramètres totaux (135 millions entraînables) qui établit une base solide pour les futures études de mise à l'échelle (scaling).

Résultats

Les auteurs ont évalué WorldDiT sur le benchmark de manipulation LIBERO, spécifiquement à travers quatre suites : Spatial, Object, Goal, et Long (LIBERO-10).

Performance : WorldDiT a atteint les taux de succès moyens suivants :
- Spatial : 98,0 %
- Object : 97,0 %
- Goal : 92,8 %
- Long : 91,8 %
- Moyenne Globale : 94,9 %
Frontière de Pareto : Lorsqu'il est tracé par rapport au nombre total de paramètres du modèle instancié, WorldDiT se situe sur la frontière de Pareto rapportée. Parmi les méthodes rapportant des résultats sur les quatre suites, chaque méthode ayant un taux de succès moyen plus élevé utilise plus de paramètres totaux. Inversement, les méthodes ayant des paramètres similaires ou moindres rapportent des taux de succès plus bas.
Comparaison : WorldDiT surpasse les autres méthodes qui n'utilisent pas de grands squelettes VLM pré-entraînés (ex: DreamVLA, FlowVLA, Octo) et rivalise favorablement avec les méthodes qui en utilisent, malgré un nombre de paramètres nettement inférieur.

Signification et Revendications

L'article affirme que WorldDiT démontre qu'une forte performance sur les benchmarks est atteignable sans placer un grand modèle de vision-langage pré-entraîné dans le squelette d'action. En couplant la génération d'actions à un objectif de prédiction de patchs RGB futurs, l'architecture atteint un compromis d'efficacité de paramètres élevé.

Les auteurs positionnent WorldDiT comme une « base solide de moins d'un milliard de paramètres » pour les futures études de mise à l'échelle. Ils suggèrent que ce résultat soutient la modélisation unifiée du monde et de l'action comme une base compacte pour le passage à l'échelle en termes de capacité de modèle, de diversité de données et de contextes de déploiement. L'article note modestement que si ces résultats caractérisent le compromis en dessous d'un milliard de paramètres, des travaux futurs sont nécessaires pour déterminer si ce compromis persiste à des échelles plus grandes et comment l'objectif de patchs RGB futurs normalisés façonne spécifiquement le contrôle. Les auteurs soulignent également le potentiel de décomposition de la formulation de flow-matching partagée en experts entraînés indépendamment pour du matériel hétérogène.