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

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

Lizenzen

Open Access
Paper: CC BY 4.0

Alles Folgende stützt sich auf das vollständige Paper unter seiner Lizenz CC BY 4.0; die Abbildung wird mit Nennung der Autoren gezeigt.

Kurz gesagt

Ein einziger Diffusion-Transformer mit getrennten Rauschzeitschritten für Aktionen und zukünftige Bilder fungiert als Policy, Dynamikmodell oder Videoprädiktor und kann auf aktionsfreien Videos trainieren.

Warum es wichtig ist

Ein Modell kann sowohl auf aktionsgelabelten Roboterdaten als auch auf aktionsfreien Videos trainieren, die weit leichter zu beschaffen sind. In Tests auf dem realen Roboter mit fünf Aufgaben auf einem Franka-Aufbau im Stil von DROID hatte vortrainiertes UWM bei jeder Aufgabe den höchsten Erfolg, und Co-Training auf zusätzlichen Videos steigerte ihn weiter, auch unter visuellen Ablenkungen.

Problem

Imitationslernen für große Robotermodelle hängt von hochwertigen Expertendemonstrationen ab, deren Erhebung per Teleoperation teuer und langsam ist. Videos, die reale Dynamik zeigen, gibt es reichlich, doch die meisten aktuellen Imitationslernmethoden können sie nicht nutzen, weil sie keine Aktionslabels haben, und Behavior Cloning erfasst die in Trajektorien vorhandene zeitliche Dynamik nicht explizit.

Lösungsansatz

UWM koppelt einen Aktions-Diffusionsprozess und einen Diffusionsprozess für zukünftige Bilder in einem Transformer, konditioniert auf die aktuelle Beobachtung, mit einem unabhängigen Diffusionszeitschritt für jede Modalität. Wird der Zeitschritt einer Modalität auf T gesetzt, wird sie marginalisiert, und bei 0 wird auf sie konditioniert, sodass dasselbe Modell eine Policy p(a|o), eine Vorwärtsdynamik p(o′|o,a), eine Inversdynamik p(a|o,o′) oder eine Videovorhersage p(o′|o) sampelt. Aktionsfreie Videos werden im Training verwendet, indem der Aktionszeitschritt auf T fixiert und die fehlenden Aktionen mit gaußschem Rauschen aufgefüllt werden.

Wichtigste Erkenntnisse

Unabhängige Diffusionszeitschritte pro Modalität wirken als partielle Maskierung, sodass ein einziges Netzwerk die Rand- und bedingten Verteilungen über Aktionen und zukünftige Frames abdeckt.
Die Vorhersage zukünftiger Frames, nicht bloß die Rekonstruktion des aktuellen Frames, treibt den Gewinn: 0,86/0,76 Erfolg bei Stack-Bowls/Block-Cabinet gegenüber 0,70/0,66 für die Rekonstruktion des aktuellen Frames und 0,48/0,60 für Diffusion Policy.
Register-Tokens (standardmäßig 8) helfen beim Informationsaustausch zwischen Aktions- und Bild-Tokens, und AdaLN-Konditionierung übertraf in der LIBERO-Ablation Cross-Attention.
UWM und Diffusion Policy schneiden beim Training von Grund auf etwa gleich ab; UWM profitiert stärker vom Vortraining.
Co-Training auf domänenspezifischen Robotervideos half mehr als Co-Training auf Internetvideos menschlicher Aktivitäten (Kinetics-400 und Something-Something v2), doch beide schlugen reine Roboterdaten.

Hauptergebnisse

Realer Roboter, vortrainiert auf 2.000 DROID-Trajektorien, Erfolg innerhalb der Verteilung: UWM 0,86, 0,76, 0,78, 0,82 und 0,60 bei Stack-Bowls, Block-Cabinet, Paper-Towel, Hang-Towel und Rice-Cooker gegenüber 0,48, 0,60, 0,52, 0,64, 0,35 für Diffusion Policy und 0,66, 0,66, 0,60, 0,66, 0,40 für GR1.
Co-Training auf 2.000 weiteren DROID-Trajektorien, die als aktionsfreie Videos genutzt wurden, steigerte den Erfolg von UWM innerhalb der Verteilung auf 0,92, 0,84, 0,86, 0,86 und 0,65 und den Erfolg außerhalb der Verteilung (erste vier Aufgaben) von 0,76/0,60/0,78/0,64 auf 0,84/0,72/0,84/0,76.
Kategorisierte OOD-Szenen (Beleuchtung, Hintergrund, Unordnung): co-trainiertes UWM 21/30 und 15/30 bei Stack-Bowls und Block-Cabinet gegenüber 15/30 und 8/30 für vortrainiertes UWM und 12/30 und 6/30 für Diffusion Policy.
LIBERO (Vortraining auf LIBERO-90, Fine-Tuning auf fünf LIBERO-10-Aufgaben mit Verteilungsverschiebungen): durchschnittlicher Erfolg 0,79 ± 0,11 für UWM gegenüber 0,71 ± 0,12 für DP, 0,57 ± 0,19 für PAD und 0,58 ± 0,14 für GR1.
Inversdynamik-Modus, der 50 LIBERO-Trajektorien innerhalb der Trajektorienlänge folgt: 0,73 und 0,45 Erfolg bei Book-Caddy und Soup-Cheese gegenüber 0,37 und 0,33 für den Policy-Modus.
Das Training eines UWM auf DROID für 100K Schritte dauert 24 Stunden auf 4 NVIDIA-A100-GPUs.

Einschränkungen

Das Modell lernt noch nicht aus großen Mengen menschlicher Videos, sodass die Embodiment-Lücke zwischen Mensch und Roboter nicht adressiert wird.
Vorhersagen der Vorwärtsdynamik können Artefakte enthalten, was ihren Nutzen für die Planung verringern kann.
Die Autoren erwarten weitere Gewinne durch eine dichtere Videovorhersage, die diese Version nicht verwendet.
Die Gewinne gegenüber Baselines im LIBERO-Setting außerhalb der Verteilung sind viel kleiner als auf dem realen Roboter; die Autoren vermuten, dass die Dynamik in der Simulation einfacher ist.
Rice-Cooker wurde nur an 20 verteilungsnahen Initialisierungen und nicht in OOD-Settings evaluiert, und alle Tests auf dem realen Roboter verwenden einen einzigen Franka-Aufbau.

Wirkung auf das Fachgebiet

Das Paper hat 244 Zitierungen. Es gibt an, dass Videos und Code auf seiner Projektseite verfügbar sind, die auf das Code-Repository WEIRDLabUW/unified-world-model auf GitHub verlinkt.

Zusammenfassung

Imitationslernen hat sich als vielversprechender Ansatz zum Bau generalistischer Roboter herausgebildet. Die Skalierung von Imitationslernen für große Roboter-Foundation-Models bleibt jedoch schwierig, da es auf hochwertige Expertendemonstrationen angewiesen ist. Gleichzeitig sind große Mengen an Videodaten, die ein breites Spektrum von Umgebungen und vielfältige Verhaltensweisen zeigen, leicht verfügbar. Diese Daten sind eine reichhaltige Informationsquelle über reale Dynamik und Interaktionen zwischen Agent und Umgebung. Diese Daten direkt für Imitationslernen zu nutzen, hat sich jedoch als schwierig erwiesen, da die von den meisten heutigen Methoden benötigte Aktionsannotation fehlt. In dieser Arbeit stellen wir Unified World Models (UWM) vor, ein Framework, das es ermöglicht, sowohl Video- als auch Aktionsdaten für das Policy-Lernen zu nutzen. Konkret integriert ein UWM einen Aktions-Diffusionsprozess und einen Video-Diffusionsprozess in einer einheitlichen Transformer-Architektur, in der unabhängige Diffusionszeitschritte jede Modalität steuern. Wir zeigen, dass UWM allein durch die Steuerung der einzelnen Diffusionszeitschritte flexibel eine Policy, eine Vorwärtsdynamik, eine Inversdynamik und einen Videogenerator darstellen kann. Durch Experimente in Simulation und realer Welt zeigen wir, dass: (1) UWM ein effektives Vortraining auf großen Multi-Task-Roboterdatensätzen mit Dynamik- und Aktionsvorhersagen ermöglicht, was zu besser generalisierbaren und robusteren Policies führt als Imitationslernen, (2) UWM durch die unabhängige Steuerung modalitätsspezifischer Diffusionszeitschritte auf natürliche Weise das Lernen aus aktionsfreien Videodaten erleichtert und so die Leistung feinabgestimmter Policies weiter verbessert. Unsere Ergebnisse deuten darauf hin, dass UWM einen vielversprechenden Schritt zur Nutzung großer, heterogener Datensätze für skalierbares Roboterlernen darstellt und eine einfache Vereinheitlichung der oft getrennten Paradigmen Imitationslernen und Weltmodellierung bietet. Videos und Code sind verfügbar unter https://weirdlabuw.github.io/uwm/

Open Source Robotics