Presentamos DriftWorld, un modelo del mundo condicionado por acciones basado en drifting que puede generar predicciones de alta calidad de fotogramas futuros en una sola pasada hacia adelante. Este comportamiento de un solo paso hace que el modelo sea significativamente más rápido al generar secuencias de rollouts que los modelos de mundo existentes condicionados por acciones, que se basan principalmente en difusión de múltiples pasos.
Formulación del Problema
El modelo del mundo condicionado por acciones se define como una función que predice observaciones futuras condicionadas a observaciones y acciones actuales, aprendida mediante un objetivo de drifting.
Arquitectura Condicionada por Acciones
Para generar múltiples fotogramas futuros a la vez, la U-Net utiliza una convolución espacio-temporal factorizada: una convolución espacial se aplica de forma independiente e idéntica a cada uno de los T pasos de tiempo, y luego se aplica una convolución temporal de forma independiente e idéntica en cada ubicación espacial.
Evaluación Visual del Modelado del Mundo
En los cuatro entornos, DriftWorld iguala o supera a las líneas base de modelos del mundo condicionados por acciones en términos de calidad de generación visual, mientras opera en una fracción del tiempo de inferencia.
Robomimic. Tanto en configuraciones de una vista como de dos vistas, DriftWorld logra un rendimiento competitivo en comparación con Ctrl-World y el modelo de mundo de difusión GPC, y es significativamente más rápido.
Mejora de Políticas en Tiempo de Inferencia en DriftWorld
DriftWorld permite una mejora de políticas en tiempo de inferencia más rápida y efectiva en comparación con los modelos de mundo de difusión existentes. Evaluamos esta capacidad tomando una política preentrenada, ejecutando sus propuestas de acción en DriftWorld y seleccionando las acciones que resultan en los mejores estados futuros.
En particular, empleamos el enfoque GPC-RANK. En cada paso, la política base muestrea K propuestas de acción candidatas. Cada propuesta se ejecuta en imaginación a través del modelo del mundo. Luego, una función de recompensa puntúa las observaciones futuras predichas, y el fragmento de acción con la recompensa más alta se selecciona para su ejecución en el mundo real.
DriftWorld como Simulador Offline para Evaluación de Políticas
DriftWorld sirve como un simulador offline preciso para la evaluación de políticas, prediciendo con éxito tanto el rendimiento absoluto como la clasificación relativa de las políticas. Validamos esta capacidad en las tareas Robomimic y Push-T.
Robomimic. Validamos adicionalmente DriftWorld en las tareas Robomimic. Para ello, post-entrenamos el modelo en videos adicionales que representan casos de fallo. Ejecutamos nueve políticas de difusión en el modelo del mundo y registramos las tasas de éxito promedio en 50 semillas de inicialización del entorno. Las tasas de éxito simuladas en DriftWorld se asemejan mucho a las del simulador MuJoCo de verdad fundamental con coeficientes de correlación de 0.9916 y 0.9250, superando al modelo base Ctrl-World.
Discusión
Limitaciones. Primero, DriftWorld se basa en un extractor de características preentrenado robusto (DINOv2/v3) para mantener la nitidez visual en escenas complejas. Segundo, el marco de drifting requiere un mayor uso de memoria que la difusión estándar durante el entrenamiento, ya que el modelo necesita generar múltiples muestras negativas (por ejemplo, 64) por pasada hacia adelante para calcular la pérdida de drifting. Como resultado, el número de fotogramas de contexto y generados por muestra negativa está limitado por la VRAM de la GPU. Una dirección para el trabajo futuro es aumentar la longitud de las ventanas de contexto y generación para mejorar la consistencia temporal a largo plazo, potencialmente utilizando un historial disperso o un VAE de video con compresión temporal.
Conclusión. Presentamos DriftWorld, un modelo del mundo condicionado por acciones, rápido y de un solo paso. Para adaptar los modelos generativos de drifting a la generación de video condicionada, propusimos adaptaciones clave, incluyendo una arquitectura U-Net condicionada por acciones, integración del espacio de características y seguimiento acentuado de acciones. DriftWorld logra una aceleración promedio de 17x en comparación con las líneas base de difusión de última generación y genera rollouts precisos y visualmente consistentes, lo que conduce a una mejor búsqueda de acciones en tiempo de inferencia y una evaluación de políticas offline confiable.
Comparación de Videos Generados en Bridge-V2, RT-1 y Language Table
Los videos generados por DriftWorld son muy similares a la verdad fundamental en los conjuntos de datos Language Table, Bridge-V2 y RT-1.
Comparación de los Rollouts de Políticas de DriftWorld
Los rollouts de políticas en DriftWorld son muy similares a los rollouts de verdad fundamental.
Simulación de Fotograma Único Versus a Nivel de Fragmento
DriftWorld admite simulación tanto de fotograma único como a nivel de fragmento en una sola pasada hacia adelante. Generar 4 fotogramas futuros funciona mejor para Push-T.
Ponderación de Movimiento para la Pérdida de Drifting
La técnica de ponderación de movimiento es esencial para la generación autorregresiva de DriftWorld en conjuntos de datos robóticos reales con fondos complejos. Sin ponderación de movimiento, el efector final del robot generado permanece prácticamente estacionario. Con ponderación de movimiento, el efector final generado se mueve según las acciones.
Detalles para la Mejora de Políticas en Tiempo de Inferencia en DriftWorld
Se proporcionan detalles adicionales sobre la mejora de políticas en tiempo de inferencia en DriftWorld mediante GPC-RANK.
Detalles para la Evaluación de Políticas en DriftWorld
Políticas evaluadas. Para Push-T, tomamos seis puntos de control de la política de difusión en las épocas 50, 100 hasta 300, así como un punto de control entrenado más largo en la época 650. Para la tarea Robomimic Lift, tomamos nueve puntos de control de la política de difusión en las épocas 2, 4, 6 hasta 18. Para la tarea Robomimic Can, tomamos nueve puntos de control de la política de difusión en las épocas 50, 100 hasta 400, así como la época 75.
Post-entrenamiento en demostraciones de fallo. Para las tareas Robomimic Lift y Can, DriftWorld y las líneas base se entrenan inicialmente en el conjunto de datos multi-humano, que contiene solo demostraciones exitosas. Dado que estos modelos iniciales proporcionan estimaciones demasiado optimistas de las tasas de éxito de las políticas, post-entrenamos los modelos en un conjunto de datos con un gran porcentaje de demostraciones de fallo, creado ejecutando puntos de control tempranos de una política de difusión. Este post-entrenamiento permite a los modelos simular fallos con precisión.
Preguntas Frecuentes
¿Qué tan rápido es DriftWorld en comparación con los modelos de mundo basados en difusión? DriftWorld logra una aceleración promedio de 17x sobre las líneas base de difusión de última generación al generar fotogramas futuros en una sola pasada hacia adelante en lugar de múltiples pasos iterativos de eliminación de ruido.
¿Qué hace única a la arquitectura de DriftWorld para la generación de video condicionada por acciones? DriftWorld utiliza una U-Net espacio-temporal factorizada combinada con una función de pérdida de drifting, integración del espacio de características y ponderación de movimiento para permitir la generación de video condicionada en un solo paso.
¿Se puede usar DriftWorld para la evaluación de políticas offline? Sí, DriftWorld predice con precisión tanto el rendimiento absoluto como la clasificación relativa de las políticas, logrando coeficientes de correlación de 0.9916 y 0.9250 con simuladores de verdad fundamental en las tareas Robomimic.
¿Cuáles son las principales limitaciones del enfoque DriftWorld? DriftWorld requiere un extractor de características preentrenado robusto para la calidad visual y tiene un alto uso de memoria durante el entrenamiento debido a la generación de múltiples muestras negativas por pasada hacia adelante.