DINO-WM: Modelli del Mondo su Caratteristiche Visuali Pre-allenate abilitano la Pianificazione a Zero-shot

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

November 7, 2024

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

cs.AI

Abstract

La capacità di prevedere risultati futuri date azioni di controllo è fondamentale per il ragionamento fisico. Tuttavia, tali modelli predittivi, spesso chiamati modelli del mondo, si sono dimostrati difficili da apprendere e sono tipicamente sviluppati per soluzioni specifiche del compito con apprendimento della politica online. Sosteniamo che il vero potenziale dei modelli del mondo risieda nella loro capacità di ragionare e pianificare su problemi diversi utilizzando solo dati passivi. Concretamente, richiediamo che i modelli del mondo abbiano le seguenti tre proprietà: 1) essere addestrabili su traiettorie offline pre-raccolte, 2) supportare l'ottimizzazione del comportamento durante il test e 3) facilitare il ragionamento agnostico rispetto al compito. Per realizzare ciò, presentiamo DINO World Model (DINO-WM), un nuovo metodo per modellare la dinamica visiva senza ricostruire il mondo visivo. DINO-WM sfrutta le caratteristiche spaziali a patch pre-addestrate con DINOv2, consentendogli di apprendere da traiettorie comportamentali offline prevedendo le future caratteristiche delle patch. Questo design permette a DINO-WM di raggiungere obiettivi osservazionali attraverso l'ottimizzazione della sequenza di azioni, facilitando la pianificazione del comportamento agnostica rispetto al compito trattando le caratteristiche delle patch desiderate come obiettivi di previsione. Valutiamo DINO-WM in vari domini, tra cui la navigazione in labirinto, il push sul tavolo e la manipolazione di particelle. I nostri esperimenti dimostrano che DINO-WM può generare soluzioni comportamentali zero-shot durante il test senza fare affidamento su dimostrazioni di esperti, modellizzazione del reward o modelli inversi pre-appresi. In particolare, DINO-WM mostra forti capacità di generalizzazione rispetto ai lavori precedenti all'avanguardia, adattandosi a diverse famiglie di compiti come labirinti configurati arbitrariamente, manipolazione push con forme di oggetti varie e scenari multi-particella.

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.