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

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

Licencias

• Acceso abierto
• Artículo: CC BY 4.0

Todo lo que sigue se basa en el artículo completo bajo su licencia CC BY 4.0; su figura se muestra con crédito a los autores.

En resumen

Un único Diffusion Transformer con pasos de tiempo de ruido separados para las acciones y las imágenes futuras actúa como política, modelo de dinámica o predictor de vídeo, y puede entrenarse con vídeo sin acciones.

Por qué importa

Un solo modelo puede entrenarse tanto con datos robóticos etiquetados con acciones como con vídeo sin acciones, que es mucho más fácil de obtener. En pruebas con robot real en cinco tareas con un montaje Franka al estilo de DROID, UWM preentrenado obtuvo el mayor éxito en todas las tareas, y el coentrenamiento con vídeos adicionales lo elevó aún más, también con distracciones visuales.

Problema

El aprendizaje por imitación para grandes modelos robóticos depende de demostraciones expertas de alta calidad, que son caras y lentas de recoger mediante teleoperación. El vídeo que muestra la dinámica del mundo real abunda, pero la mayoría de los métodos actuales de aprendizaje por imitación no pueden usarlo porque no tiene etiquetas de acciones, y la clonación de comportamiento no captura de forma explícita la dinámica temporal presente en las trayectorias.

Solución propuesta

UWM acopla un proceso de difusión de acciones y un proceso de difusión de imágenes futuras en un solo transformador, condicionado por la observación actual, con un paso de tiempo de difusión independiente para cada modalidad. Fijar el paso de tiempo de una modalidad en T la marginaliza y fijarlo en 0 condiciona sobre ella, de modo que el mismo modelo muestrea una política p(a|o), una dinámica directa p(o′|o,a), una dinámica inversa p(a|o,o′) o una predicción de vídeo p(o′|o). El vídeo sin acciones se usa en el entrenamiento fijando el paso de tiempo de las acciones en T y rellenando las acciones que faltan con ruido gaussiano.

Ideas clave

• Los pasos de tiempo de difusión independientes por modalidad actúan como un enmascaramiento parcial, de modo que una sola red cubre las distribuciones marginales y condicionales sobre las acciones y los fotogramas futuros.
• Lo que impulsa la mejora es predecir fotogramas futuros, no simplemente reconstruir el fotograma actual: 0,86/0,76 de éxito en Stack-Bowls/Block-Cabinet frente a 0,70/0,66 con la reconstrucción del fotograma actual y 0,48/0,60 con Diffusion Policy.
• Los register tokens (8 por defecto) ayudan a intercambiar información entre los tokens de acción y de imagen, y el condicionamiento AdaLN superó a la atención cruzada en la ablación en LIBERO.
• UWM y Diffusion Policy rinden más o menos igual cuando se entrenan desde cero; UWM gana más con el preentrenamiento.
• El coentrenamiento con vídeos de robots del propio dominio ayudó más que el coentrenamiento con vídeos de actividad humana de Internet (Kinetics-400 y Something-Something v2), aunque ambos superaron a los datos robóticos solos.

Resultados principales

• Robot real, preentrenado con 2000 trayectorias de DROID, éxito dentro de distribución: UWM 0,86, 0,76, 0,78, 0,82 y 0,60 en Stack-Bowls, Block-Cabinet, Paper-Towel, Hang-Towel y Rice-Cooker, frente a 0,48, 0,60, 0,52, 0,64, 0,35 con Diffusion Policy y 0,66, 0,66, 0,60, 0,66, 0,40 con GR1.
• El coentrenamiento con otras 2000 trayectorias de DROID usadas como vídeo sin acciones elevó el éxito de UWM dentro de distribución a 0,92, 0,84, 0,86, 0,86 y 0,65, y el éxito fuera de distribución (las cuatro primeras tareas) de 0,76/0,60/0,78/0,64 a 0,84/0,72/0,84/0,76.
• Escenas OOD categorizadas (iluminación, fondo, desorden): UWM coentrenado 21/30 y 15/30 en Stack-Bowls y Block-Cabinet, frente a 15/30 y 8/30 con UWM preentrenado y 12/30 y 6/30 con Diffusion Policy.
• LIBERO (preentrenamiento en LIBERO-90, ajuste fino en cinco tareas de LIBERO-10 con cambios de distribución): éxito medio de 0,79 ± 0,11 para UWM frente a 0,71 ± 0,12 para DP, 0,57 ± 0,19 para PAD y 0,58 ± 0,14 para GR1.
• Modo de dinámica inversa siguiendo 50 trayectorias de LIBERO dentro de la longitud de la trayectoria: 0,73 y 0,45 de éxito en Book-Caddy y Soup-Cheese, frente a 0,37 y 0,33 en el modo política.
• Entrenar un UWM con DROID durante 100K pasos lleva 24 horas en 4 GPU NVIDIA A100.

Limitaciones

• El modelo todavía no aprende de vídeos humanos a gran escala, por lo que no se aborda la brecha de encarnación entre humanos y robots.
• Las predicciones de dinámica directa pueden contener artefactos, lo que puede reducir su utilidad para la planificación.
• Los autores esperan mejoras adicionales con una predicción de vídeo más densa, que esta versión no usa.
• Las mejoras respecto a las referencias en el escenario fuera de distribución de LIBERO son mucho menores que en el robot real; los autores sugieren que la dinámica de la simulación es más simple.
• Rice-Cooker solo se evaluó con 20 inicializaciones cercanas a la distribución y no en escenarios OOD, y todas las pruebas con robot real usan un único montaje Franka.

Impacto en el campo

El artículo tiene 244 citas. Indica que los vídeos y el código están disponibles en su página del proyecto, que enlaza al repositorio de código WEIRDLabUW/unified-world-model en GitHub.

Resumen

El aprendizaje por imitación se ha consolidado como un enfoque prometedor para construir robots generalistas. Sin embargo, escalar el aprendizaje por imitación para grandes modelos fundacionales robóticos sigue siendo difícil debido a su dependencia de demostraciones expertas de alta calidad. Mientras tanto, hay disponibles grandes cantidades de datos de vídeo que muestran una amplia variedad de entornos y comportamientos diversos. Estos datos son una fuente rica de información sobre la dinámica del mundo real y las interacciones entre el agente y el entorno. Sin embargo, aprovechar estos datos directamente para el aprendizaje por imitación ha resultado difícil por la falta de las anotaciones de acciones que requieren la mayoría de los métodos actuales. En este trabajo presentamos Unified World Models (UWM), un marco que permite aprovechar tanto datos de vídeo como de acciones para el aprendizaje de políticas. En concreto, un UWM integra un proceso de difusión de acciones y un proceso de difusión de vídeo dentro de una arquitectura de transformador unificada, en la que pasos de tiempo de difusión independientes rigen cada modalidad. Mostramos que, simplemente controlando cada paso de tiempo de difusión, UWM puede representar de forma flexible una política, una dinámica directa, una dinámica inversa y un generador de vídeo. Mediante experimentos en simulación y en el mundo real, mostramos que: (1) UWM permite un preentrenamiento eficaz en conjuntos de datos robóticos multitarea a gran escala con predicciones tanto de dinámica como de acciones, lo que da lugar a políticas más generalizables y robustas que el aprendizaje por imitación, (2) UWM facilita de forma natural el aprendizaje a partir de datos de vídeo sin acciones mediante el control independiente de los pasos de tiempo de difusión específicos de cada modalidad, lo que mejora aún más el rendimiento de las políticas ajustadas. Nuestros resultados sugieren que UWM supone un paso prometedor hacia el aprovechamiento de conjuntos de datos grandes y heterogéneos para un aprendizaje robótico escalable, y ofrece una unificación sencilla entre los paradigmas, a menudo separados, del aprendizaje por imitación y del modelado del mundo. Los vídeos y el código están disponibles en https://weirdlabuw.github.io/uwm/

Enlaces

• 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
• Código: https://github.com/WEIRDLabUW/unified-world-model
• Página del proyecto: https://weirdlabuw.github.io/uwm/