🤖 machine learning

WorldDiT: A Unified Diffusion Architecture for World and Action Modeling

WorldDiT introduce una arquitectura de transformador de difusión unificada que genera simultáneamente acciones continuas y predice fotogramas visuales futuros, logrando un rendimiento de vanguardia en múltiples evaluaciones de manipulación robótica sin depender de grandes modelos de visión y lenguaje preentrenados.

Autores originales: Sen Wang, R. Gnana Praveen, Bidhan Roy, Marcos Villagra

Publicado 2026-07-28

📖 6 min de lectura🧠 Análisis profundo

CC BY 4.0

Autores originales: Sen Wang, R. Gnana Praveen, Bidhan Roy, Marcos Villagra

Artículo original bajo licencia CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta es una explicación generada por IA del artículo a continuación. No ha sido escrita ni avalada por los autores. Para mayor precisión técnica, consulte el artículo original. Leer descargo de responsabilidad completo

Imagina un mundo donde los robots no solo siguen una lista rígida de instrucciones, sino que realmente "sueñan" con lo que sucederá después. Esta es la frontera del aprendizaje robótico, un campo que intenta enseñar a las máquinas cómo moverse a través de nuestro mundo real y desordenado. Durante mucho tiempo, los mayores avances provinieron de una receta específica: toma un "cerebro" masivo y preentrenado que ya lo sabe todo sobre lenguaje e imágenes (como una enciclopedia superinteligente que también puede ver), y conéctale un brazo robótico. Funciona, pero es como intentar enseñarle a un cachorro a buscar la pelota entrenándolo primero para ser un profesor; es pesado, costoso y difícil saber si el robot está aprendiendo la tarea o simplemente copiando las notas del profesor. La gran pregunta que se hacen los científicos es: ¿Podemos construir un cerebro robótico que sea más pequeño, ligero y que aprenda a moverse y a predecir el futuro al mismo tiempo, sin necesidad de esa enciclopedia gigante y precargada?

Presentamos WorldDiT, un nuevo cerebro robótico diseñado por investigadores de Bagel Labs. Piensa en WorldDiT no como un estudiante copiando un libro de texto, sino como un joven artista que aprende a pintar observando una película de su propio futuro. En lugar de depender de un enorme "Modelo de Visión-Lenguaje" preentrenado (una IA masiva que sabe describir imágenes y texto) para decirle qué hacer, WorldDiT utiliza un único motor unificado que hace dos cosas simultáneamente. Primero, descifra los siguientes movimientos del brazo del robot. Segundo, intenta adivinar qué verá la cámara dentro de unos segundos. Es como un jugador de ajedrez que no solo planea su siguiente movimiento, sino que también visualiza el tablero tres turnos por delante para asegurarse de que su plan tenga sentido.

Los investigadores entrenaron este sistema en un mundo simulado llamado LIBERO, donde los robots tienen que realizar tareas como apilar bloques o mover objetos. Descubrieron que WorldDiT, que es sorprendentemente pequeño (menos de 400 millones de parámetros, comparado con los miles de millones usados por otros métodos destacados), puede predecir el futuro de la vista de la cámara y las acciones del robot con una precisión increíble. De hecho, cuando lo probaron, el robot tuvo éxito en aproximadamente el 95% de sus tareas en cuatro tipos diferentes de desafíos. ¿La parte más emocionante? Lo hizo sin el pesado "backbone" preentrenado que utilizan la mayoría de los otros robots exitosos. Demostró que no necesitas un cerebro gigante y precargado para ser un buen robot; solo necesitas un cerebro que aprenda a ver el futuro mientras se mueve.

El Robot que "Sueña"

Entonces, ¿cómo funciona realmente WorldDiT? Imagina que estás tratando de aprender a hacer malabares. La mayoría de los robots actuales son como estudiantes que han memorizado mil videos de maestros de malabarismo e intentan imitarlos exactamente. WorldDiT es diferente. Es como un estudiante de malabares que tiene los ojos vendados pero posee una bola de cristal.

Aquí está el truco: Cada vez que WorldDiT planea un movimiento, también intenta "soñar" lo que la cámara verá en el futuro. Observa el estado actual del robot, la instrucción que recibió (como "recoge el bloque rojo") y los últimos segundos de video. Luego, se hace dos preguntas al mismo tiempo:

"¿Qué debe hacer el robot a continuación?"
"Si hago eso, ¿cómo se verá la imagen un momento después?"

Los investigadores llaman a esto un difusión transformadora (diffusion transformer). Si esto suena a jerga de ciencia ficción, piénsalo como un proceso de "eliminación de ruido" (denoising). Imagina una foto cubierta de estática. WorldDiT comienza con una suposición completamente aleatoria y ruidosa de cuál debería ser el siguiente movimiento del robot y cómo debería verse la imagen futura. Luego, paso a paso, limpia ese ruido, refinando su suposición hasta encontrar un camino claro y fluido hacia el objetivo. Hace esto tanto para las acciones del robot como para las imágenes futuras.

La magia ocurre porque estas dos tareas se ayudan mutuamente. Al intentar predecir la imagen futura, el robot se ve obligado a comprender la física del mundo. Si planea empujar un bloque, pero su "sueño" del futuro muestra el bloque flotando en el aire, sabe que su plan es erróneo. Tiene que ajustar su acción hasta que el "sueño" coincida con la realidad. Esto crea un bucle de retroalimentación donde el robot aprende las reglas del mundo simplemente intentando predecir qué viene después.

Los Resultados: Cerebro Pequeño, Gran Éxito

Los investigadores probaron WorldDiT en cuatro suites de simulación (piensa en ellas como cuatro pistas de obstáculos diferentes: Espacial, Objeto, Objetivo y Largo). Lo compararon con otros 24 métodos famosos de aprendizaje robótico.

Aquí está la clave: WorldDiT es diminuto. Tiene alrededor de 399 millones de parámetros (las "neuronas" en su cerebro). Muchos de los otros robots de alto rendimiento utilizan modelos con 1 mil millones o incluso 10 mil millones de parámetros. Usualmente, los cerebros más grandes significan un mejor rendimiento. Pero WorldDiT rompió esa regla.

En estas simulaciones, WorldDiT logró una tasa de éxito del 94.9% en promedio en los cuatro cursos.

Espacial: 98.0% de éxito
Objeto: 97.0% de éxito
Objetivo: 92.8% de éxito
Largo: 91.8% de éxito

Cuando graficas estos resultados, WorldDiD se sitúa en la "frontera de Pareto". En lenguaje sencillo, esto significa que es el robot más eficiente de la sala. Si quieres un robot que sea tan inteligente como WorldDT pero que use un cerebro más pequeño, no puedes encontrar uno. Si quieres un robot que sea tan pequeño como WorldDT pero más inteligente, tampoco puedes encontrarlo. Cualquier otro robot que haya logrado un rendimiento superior al 94.9% era significativamente más grande y pesado.

Por Qué Esto Importa

El artículo argumenta explícitamente en contra de la idea de que debes usar un Modelo de Visión-Lenguaje (VLA) masivo y preentrenado para obtener un buen control robótico. Durante un tiempo, el campo pensó: "Necesitamos un cerebro gigante para entender el mundo". WorldDT dice: "No necesariamente".

Los investigadores demostraron que, al acoplar la generación de acciones (moverse) con el modelado del mundo (predecir el futuro), una arquitectura única y unificada puede aprender a controlar un robot de manera muy efectiva. No solo lo sugirieron; lo midieron en miles de episodios simulados. El robot no solo "tuvo suerte"; superó consistentemente a modelos mucho más grandes.

Sin embargo, hay un detalle. Estos resultados provienen de simulaciones (juegos de computadora donde viven los robots). El artículo señala que, aunque el rendimiento es sólido, es una base para estudios futuros. Aún no sabemos si este robot que "sueña" funcionará perfectamente en una cocina real y desordenada, con gravedad real y suelos resbaladizos reales. Pero los resultados de la simulación son un fuerte indicio de que podríamos construir robots más pequeños, más baratos y más eficientes en el futuro: robots que aprenden imaginando el futuro, en lugar de simplemente memorizar el pasado.

Al final, WorldDT sugiere que el secreto de un robot inteligente no es solo tener una enorme biblioteca de conocimientos; es tener un cerebro que pueda representar el futuro en su mente para tomar mejores decisiones ahora mismo.

Resumen Técnico: WorldDiT

Planteamiento del Problema

Los avances recientes en el aprendizaje de políticas robóticas suelen depender de la vinculación de la generación de acciones a grandes modelos de visión y lenguaje preentrenados (VLM). Si bien este enfoque aprovecha una amplia percepción y comprensión del lenguaje, oscurece las contribuciones específicas de la escala y la arquitectura del modelo. Cuando una política contiene varios miles de millones de parámetros, resulta difícil discernir si un fuerte rendimiento de control proviene del diseño de la acción, del backbone preentrenado o de su combinación. Los autores plantean la siguiente pregunta: ¿Puede una arquitectura de transformador de difusión unificada combinar la generación de acciones continuas con un objetivo auxiliar de predicción visual futura para lograr un fuerte rendimiento de control sin depender de un gran VLM preentrenado como backbone de la acción?

Metodología: WorldDiT

WorldDiT es una arquitectura de transformador de difusión unificada diseñada para acoplar la generación de acciones con el modelado de mundo visual. A diferencia de los modelos de mundo-acción autorregresivos que generan tokens de acción secuencialmente utilizando un gran VLM, WorldDiT emplea un único backbone de Transformador de Difusión (DiT) compartido para modelar tanto acciones robóticas continuas como parches RGB normalizados seleccionados de fotogramas de cámara futuros.

Arquitectura y Codificación

El sistema utiliza un enfoque híbrido de "congelado" y "entrenable":

Codificadores: Se utilizan codificadores visuales "congelados" (MAE para imágenes, CLIP para texto) y un codificador de estado robótico entrenable (que codifica por separado los estados del brazo y la pinza) para procesar las entradas.
Tokenización: El sistema construye tokens multimodales que consisten en:
- 34 tokens visuales (34 parches de las cámaras primaria y de la muñera).
- 1 token de lenguaje (de la instrucción).
- 1 token de estado (estado del robot).
- Esto resulta en 36 tokens de contexto limpios por intervalo temporal.
Backbone: El núcleo es un backbone DiT compartido con una profundidad de 4, un tamaño oculto de 1024, 16 cabezales de atención y 4 tokens de registro. Utiliza AdaLN-Zero por token y una atención estrictamente causal por bloques y segura para la acción, lo que garantiza que las consultas de acción no puedan leer los objetivos RGB con ruido.

Objetivo de Entrenamiento: Flow Matching

WorldDiT se entrena utilizando flow matching (ajuste de flujo). La ventana de entrenamiento ( $N=10$ ) consta de:

Contexto ( $C=3$ ): Tres pasos observados que contienen lenguaje, RGB multivista y estados.
Objetivos ( $H=7$ ): Un bloque de acción de 7 pasos y un objetivo de parche RGB futuro normalizado.

El objetivo visual futuro se construye a partir de los fotogramas de las cámaras primaria y de la muñera en un desfase temporal de $H$ . Cada fotograma de 224x224 se divide en parches de 16x16; se retienen 64 parches distribuidos uniformemente de cada cámara, normalizados por su propia media y varianza, y se concatenan para formar un objetivo de 128 tokens ( $Y_{rgb}$ ).

El modelo recibe tokens de contexto limpios y tokens de objetivo corrompidos (parches de acción y RGB) con ruido gaussiano. Predice la velocidad del flujo para un camino recto desde el ruido hacia el objetivo limpio. La pérdida total es una suma ponderada de la pérdida de velocidad de la acción y la pérdida de velocidad del parche RGB:
$L_{total} = w_{action}L_{action}^{flow} + w_{rgb}L_{rgb}^{flow}$
Crucialmente, la predicción de parches RGB sirve como una señal de supervisión auxiliar durante el entrenamiento, pero está ausente en la inferencia.

Inferencia

En el despliegue, el sistema opera con un horizonte de retroceso:

El backbone unificado recibe el historial (lenguaje, observaciones, estados) hasta el paso actual.
No se proporcionan objetivos RGB futuros ni etiquetas de acción.
El modelo genera un bloque de acción de 7 pasos mediante la integración numérica del campo de velocidad de acción aprendido desde el ruido gaussiano ( $\tau=0$ ) hasta el estado limpio ( $\tau=1$ ) a lo largo de 20 pasos de Euler.
Se ejecutan los tres primeros pasos de la acción del bloque predicho, se observa el estado y se desplaza la ventana para el replanteamiento.

Contribuciones Clave

Arquitectura Unificada: Un único transformador de difusión que aprende conjuntamente la generación de acciones continuas y la predicción de parches RGB futuros normalizados sin un gran backbone de VLM preentrenado.
Diseño de Entrenamiento/Despliegue: Una configuración novedosa donde la predicción de parches RGB suministra supervisión durante el entrenamiento pero se elimina del grafo de inferencia, permitiendo una ruta de despliegue compacta y solo de acción.
Línea Base de Menos de Mil Millones de Parámetros: Una configuración con aproximadamente 399 millones de parámetros totales (135 millones entrenables) que establece una línea base sólida para futuros estudios de escalado.

Resultados

Los autores evaluaron WorldDiT en el benchmark de manipulación LIBERO, específicamente a través de cuatro suites: Spatial, Object, Goal y Long (LIBERO-10).

Rendimiento: WorldDiT logró las siguientes tasas de éxito medias:
- Spatial: 98.0%
- Object: 97.0%
- Goal: 92.8%
- Long: 91.8%
- Media General: 94.9%
Frontera de Pareto: Al graficarlo contra los parámetros totales del modelo instanciado, WorldDiT se sitúa en la frontera de Pareto reportada. Entre los métodos que informan resultados en las cuatro suites, cada método con una tasa de éxito media mayor utiliza más parámetros totales. Por el contrario, los métodos con parámetros similares o menores reportan tasas de éxito más bajas.
Comparación: WorldDiT supera a otros métodos que no utilizan backbones de VLM preentrenados (por ejemplo, DreamVLA, FlowVLA, Octo) y compite favorablemente con métodos que sí los utilizan, a pesar de tener significativamente menos parámetros.

Significado y Reivindicaciones

El artículo afirma que WorldDiT demuestra que es posible alcanzar un fuerte rendimiento en los benchmarks sin necesidad de colocar un gran modelo de visión y lenguaje preentrenado en el backbone de la acción. Al acoplar la generación de acciones con un objetivo de predicción de parches RGB futuros, la arquitectura logra un alto equilibrio de eficiencia de parámetros.

Los autores posicionan a WorldDiT como una "línea base sólida de menos de mil millones de parámetros" para futuros estudios de escalado. Sugieren que este resultado respalda el modelado unificado de mundo y acción como una base compacta para el escalado a través de la capacidad del modelo, la diversidad de datos y los entornos de despliegue. El artículo señala modestamente que, si bien los resultados caracterizan el compromiso por debajo de los mil millones de parámetros, es necesario realizar trabajos futuros para determinar si este compromiso persiste a escalas mayores y cómo el objetivo de RGB futuro normalizado moldea específicamente el control. Los autores también destacan el potencial de descomponer la formulación compartida de ajuste de flujo en expertos entrenados independientemente para hardware heterogéneo.