DINO-WM: Modelos del Mundo en Características Visuales Pre-entrenadas permiten Planificación sin Necesidad de Datos

DINO-WM: World Models on Pre-trained Visual Features enable Zero-shot Planning

November 7, 2024

Autores: Gaoyue Zhou, Hengkai Pan, Yann LeCun, Lerrel Pinto

cs.AI

Resumen

La capacidad de predecir resultados futuros dados ciertas acciones de control es fundamental para el razonamiento físico. Sin embargo, dichos modelos predictivos, a menudo llamados modelos del mundo, han demostrado ser difíciles de aprender y suelen desarrollarse para soluciones específicas de tareas con aprendizaje de políticas en línea. Sostenemos que el verdadero potencial de los modelos del mundo radica en su capacidad para razonar y planificar a través de problemas diversos utilizando solo datos pasivos. Concretamente, requerimos que los modelos del mundo tengan las siguientes tres propiedades: 1) ser entrenables en trayectorias offline pre-recopiladas, 2) admitir la optimización del comportamiento en tiempo de prueba, y 3) facilitar el razonamiento sin tener en cuenta la tarea. Para lograr esto, presentamos DINO World Model (DINO-WM), un nuevo método para modelar la dinámica visual sin reconstruir el mundo visual. DINO-WM aprovecha las características de parches espaciales pre-entrenadas con DINOv2, lo que le permite aprender de trayectorias de comportamiento offline prediciendo futuras características de parches. Este diseño permite a DINO-WM lograr objetivos observacionales a través de la optimización de secuencias de acciones, facilitando la planificación de comportamientos sin tener en cuenta la tarea al tratar las características deseadas de los parches como objetivos de predicción. Evaluamos DINO-WM en varios dominios, incluida la navegación de laberintos, empuje en mesas y manipulación de partículas. Nuestros experimentos demuestran que DINO-WM puede generar soluciones de comportamiento de cero disparo en tiempo de prueba sin depender de demostraciones de expertos, modelado de recompensas o modelos inversos preaprendidos. Es notable que DINO-WM exhibe fuertes capacidades de generalización en comparación con trabajos previos de vanguardia, adaptándose a diversas familias de tareas como laberintos configurados arbitrariamente, manipulación de empuje con formas de objetos variadas y escenarios de múltiples partículas.

English

The ability to predict future outcomes given control actions is fundamental for physical reasoning. However, such predictive models, often called world models, have proven challenging to learn and are typically developed for task-specific solutions with online policy learning. We argue that the true potential of world models lies in their ability to reason and plan across diverse problems using only passive data. Concretely, we require world models to have the following three properties: 1) be trainable on offline, pre-collected trajectories, 2) support test-time behavior optimization, and 3) facilitate task-agnostic reasoning. To realize this, we present DINO World Model (DINO-WM), a new method to model visual dynamics without reconstructing the visual world. DINO-WM leverages spatial patch features pre-trained with DINOv2, enabling it to learn from offline behavioral trajectories by predicting future patch features. This design allows DINO-WM to achieve observational goals through action sequence optimization, facilitating task-agnostic behavior planning by treating desired goal patch features as prediction targets. We evaluate DINO-WM across various domains, including maze navigation, tabletop pushing, and particle manipulation. Our experiments demonstrate that DINO-WM can generate zero-shot behavioral solutions at test time without relying on expert demonstrations, reward modeling, or pre-learned inverse models. Notably, DINO-WM exhibits strong generalization capabilities compared to prior state-of-the-art work, adapting to diverse task families such as arbitrarily configured mazes, push manipulation with varied object shapes, and multi-particle scenarios.