DINO-WM: Modelos Mundiais em Recursos Visuais Pré-treinados permitem Planeamento sem Necessidade de Treino

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

Resumo

A capacidade de prever resultados futuros dadas ações de controle é fundamental para o raciocínio físico. No entanto, tais modelos preditivos, frequentemente chamados de modelos do mundo, têm se mostrado desafiadores de aprender e geralmente são desenvolvidos para soluções específicas de tarefas com aprendizado de políticas online. Argumentamos que o verdadeiro potencial dos modelos do mundo está em sua capacidade de raciocinar e planejar em diversos problemas usando apenas dados passivos. Concretamente, exigimos que os modelos do mundo tenham as seguintes três propriedades: 1) ser treináveis em trajetórias offline pré-coletadas, 2) suportar a otimização do comportamento no momento do teste e 3) facilitar o raciocínio independente da tarefa. Para realizar isso, apresentamos o Modelo do Mundo DINO (DINO-WM), um novo método para modelar dinâmicas visuais sem reconstruir o mundo visual. O DINO-WM aproveita características de patches espaciais pré-treinadas com o DINOv2, permitindo que aprenda a partir de trajetórias comportamentais offline ao prever características de patches futuras. Esse design permite que o DINO-WM alcance objetivos observacionais por meio da otimização de sequências de ações, facilitando o planejamento de comportamento independente da tarefa ao tratar as características de patches desejadas como alvos de previsão. Avaliamos o DINO-WM em vários domínios, incluindo navegação em labirintos, empurrar objetos em mesas e manipulação de partículas. Nossos experimentos demonstram que o DINO-WM pode gerar soluções comportamentais de zero-shot no momento do teste sem depender de demonstrações de especialistas, modelagem de recompensas ou modelos inversos pré-aprendidos. Notavelmente, o DINO-WM exibe fortes capacidades de generalização em comparação com trabalhos anteriores de ponta, adaptando-se a diversas famílias de tarefas, como labirintos configurados arbitrariamente, manipulação por empurrão com formas de objetos variadas e cenários com múltiplas 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.