GeoWorld: Geometric World Models

Il paper introduce GeoWorld, un modello del mondo geometrico che utilizza un Hyperbolic JEPA e l'apprendimento per rinforzo geometrico per preservare le relazioni strutturali gerarchiche e migliorare la pianificazione visiva a lungo termine rispetto agli approcci esistenti.

L'essenziale

Immagina di dover insegnare a un robot come cambiare la memoria del suo computer. Non basta dirgli "togli la vecchia, metti la nuova"; deve capire l'ordine esatto, cosa succede se sbaglia un passaggio e come tornare indietro se qualcosa va storto.

Fino a poco tempo fa, i robot e le intelligenze artificiali cercavano di imparare questo guardando i video e provando a "disegnare" il futuro fotogramma per fotogramma. Era come cercare di prevedere il meteo di domani disegnando ogni singola nuvola: un lavoro enorme, lento e soggetto a errori che si accumulano. Se sbagliavi il disegno di una nuvola, il cielo intero del giorno dopo diventava sbagliato.

GeoWorld è una nuova intelligenza artificiale che cambia completamente il modo di pensare a questo problema. Ecco come funziona, spiegato in modo semplice:

1. Il Problema: La mappa piatta vs. la mappa a imbuto

Immagina che lo stato mentale del robot (ciò che "capisce" del mondo) sia una mappa.

• I vecchi metodi usavano una mappa piatta (come un foglio di carta). Su un foglio piatto, se devi andare da un punto A a un punto B, il percorso è una linea dritta. Ma il mondo reale è complesso: ci sono compiti semplici e compiti enormi che contengono molti sotto-compiti. Su una mappa piatta, è difficile rappresentare questa gerarchia (come una famiglia: nonni, genitori, figli). Se il robot fa un piccolo errore, sulla mappa piatta può finire per "scivolare" via dal percorso corretto e perdersi completamente dopo pochi passi.
• GeoWorld usa una mappa a imbuto (matematicamente chiamata "spazio iperbolico"). Immagina un imbuto o un albero: in alto c'è il concetto generale (es. "cambiare la memoria"), e man mano che scendi verso il fondo dell'imbuto, i rami si diramano in azioni sempre più specifiche (es. "svitare la vite", "togliere il chip").
  • L'analogia: Su una mappa piatta, due punti lontani sono solo "lontani". Su una mappa a imbuto, due punti che sembrano vicini in superficie potrebbero essere in realtà su rami diversi di un albero, quindi il robot capisce subito che sono concettualmente diversi. Questo aiuta a non confondersi.

2. La Soluzione: Camminare lungo i sentieri naturali

GeoWorld non cerca di "disegnare" il futuro (come farebbe un artista che dipinge il prossimo secondo). Invece, calcola un "sentiero energetico".

• Immagina di essere su una collina. Il robot vuole arrivare in fondo alla valle (l'obiettivo finale, es. il computer riparato).
• Invece di guardare ogni singolo sasso (ogni pixel dell'immagine), GeoWorld guarda la forma della collina. Sa che la strada migliore è seguire la pendenza naturale (la "geodetica").
• Grazie alla sua mappa a imbuto, questo sentiero è molto più stabile. Anche se il robot inciampa un po' (fa un piccolo errore di calcolo), la forma dell'imbuto lo "riporta" sulla strada giusta, invece di farlo cadere in un burrone.

3. L'Allenamento: Il "Reinforcement Learning Geometrico"

Come si allena GeoWorld?

• Fase 1 (Studio): Il robot guarda migliaia di video di persone che fanno cose (come montare un PC o cucinare). Impara a riconoscere i "sentieri energetici" giusti su quella mappa a imbuto.
• Fase 2 (Allenamento con premi): Qui entra in gioco la parte intelligente. Il robot prova a pianificare un percorso. Se il percorso che ha immaginato è coerente e segue le regole della sua mappa a imbuto (senza saltare rami dell'albero), riceve un "premio". Se si perde o fa un percorso che non ha senso geometrico, riceve una "penalità".
• Questo processo, chiamato Reinforcement Learning Geometrico, insegna al robot a non solo "sapere" cosa fare, ma a "sentire" quando sta per sbagliare strada molto prima che accada realmente.

Perché è importante?

I robot attuali sono bravi a fare un passo alla volta, ma quando devono pianificare una sequenza lunga (es. "prepara una torta" che richiede 10 passaggi), tendono a dimenticare i primi passi o a confondersi.
GeoWorld è come un architetto che non guarda solo i mattoni, ma vede l'intera struttura dell'edificio.

• Risultato: Nei test, GeoWorld è riuscito a pianificare compiti complessi (come cambiare la memoria di un PC) con una precisione molto più alta rispetto ai metodi precedenti, specialmente quando i compiti richiedevano molti passaggi.

In sintesi:
GeoWorld ha smesso di cercare di "disegnare" il futuro fotogramma per fotogramma. Invece, ha imparato a "sentire" la forma del mondo come se fosse un imbuto o un albero, permettendogli di trovare la strada migliore per raggiungere un obiettivo senza perdersi, proprio come un esploratore esperto che conosce la mappa del territorio meglio di chiunque altro.

Sommario tecnico

1. Il Problema

I modelli predittivi del mondo basati sull'energia (Energy-Based Predictive World Models) offrono un approccio potente per la pianificazione visiva multi-step, ragionando su paesaggi energetici latenti invece di generare pixel. Tuttavia, le approcci esistenti affrontano due sfide fondamentali:

1. Negligenza Geometrica: Le rappresentazioni latenti sono tipicamente apprese in uno spazio euclideo. Questo ignora la struttura geometrica sottostante e le relazioni gerarchiche tra gli stati. Di conseguenza, il paesaggio energetico appreso non riesce a catturare distanze geodetiche significative o embedding gerarchici, indebolendo la capacità di pianificazione coerente su orizzonti temporali lunghi.
2. Limiti nella Pianificazione Multi-Step: I modelli predittivi sono spesso addestrati su transizioni video a singolo passo. Sebbene concettualmente possano pianificare su traiettorie complete, le loro prestazioni degradano rapidamente all'aumentare dell'orizzonte di pianificazione a causa dell'accumulo di errori e della difficoltà nel modellare dipendenze temporali a lungo termine.

2. Metodologia: GeoWorld

Per risolvere questi problemi, gli autori introducono GeoWorld, un modello del mondo geometrico che preserva la struttura geometrica e le relazioni gerarchiche nello spazio latente. L'architettura si basa su due componenti principali:

A. Hyperbolic JEPA (H-JEPA)

Invece di operare nello spazio euclideo, GeoWorld mappa le rappresentazioni latenti su una varietà iperbolica (specificamente il modello della palla di Poincaré).

• Mappatura: L'output del codificatore (da $R^n$) viene proiettato sulla varietà iperbolica $\mathcal{H}^n$ tramite una mappa esponenziale.
• Dinamica: Il predittore apprende le dinamiche lungo le geodetiche iperboliche. In uno spazio iperbolico, le distanze crescono esponenzialmente con il raggio, il che permette di rappresentare naturalmente strutture gerarchiche ad albero (tipiche della pianificazione multi-step dove il numero di futuri possibili cresce esponenzialmente con la profondità).
• Obiettivo: Minimizzare la distanza geodetica iperbolica tra lo stato previsto e lo stato reale, garantendo che le transizioni seguano il percorso a minima energia sulla varietà.

B. Geometric Reinforcement Learning (GRL)

Per migliorare la stabilità nella pianificazione multi-step, viene proposto un framework di ottimizzazione basato sull'energia.

• Funzione di Valore: La pianificazione è riformulata come la minimizzazione di una funzione di valore basata sull'energia iperbolica, dove un'energia più bassa corrisponde a una ricompensa cumulativa più alta.
• Regolarizzazione: Viene introdotta una regolarizzazione basata sulla disuguaglianza triangolare iperbolica. Questo vincolo forza le traiettorie previste a rispettare le proprietà delle geodetiche, prevenendo deviazioni degenerate e riducendo l'accumulo di errori durante i roll-out a lungo termine.
• Ottimizzazione: Il predittore viene affinato direttamente attraverso la minimizzazione dell'energia iperbolica e la regolarizzazione, senza bisogno di addestrare modelli di politica o ricompensa aggiuntivi.

C. Pianificazione Basata sull'Energia

Durante l'inferenza, la pianificazione viene eseguita utilizzando il Metodo Cross-Entropy (CEM). L'algoritmo cerca la sequenza di azioni che minimizza il costo energetico (distanza geodetica iperbolica) tra lo stato latente corrente e lo stato obiettivo, seguendo i percorsi geodetici sulla varietà.

3. Contributi Chiave

1. GeoWorld e H-JEPA: Introduzione di un modello del mondo che mappa le rappresentazioni latenti su una varietà iperbolica, preservando la struttura geometrica e le relazioni gerarchiche per una pianificazione multi-step coerente.
2. Geometric Reinforcement Learning (GRL): Un framework di ottimizzazione che affina il predittore tramite minimizzazione dell'energia iperbolica e regolarizzazione della disuguaglianza triangolare, migliorando la stabilità a lungo termine.
3. Risultati Sperimentali: Dimostrazione di prestazioni superiori su benchmark di pianificazione visiva a lungo termine (CrossTask e COIN), con miglioramenti significativi rispetto allo stato dell'arte (V-JEPA 2).

4. Risultati Sperimentali

Il modello è stato valutato sui dataset CrossTask e COIN per la pianificazione visiva condizionata all'obiettivo, sia in modalità procedurale (immagini) che con video.

• Prestazioni: GeoWorld supera costantemente V-JEPA 2 (lo stato dell'arte attuale).
  • Miglioramento di circa 3% nel Success Rate (SR) per la pianificazione a 3 passi.
  • Miglioramento di circa 2% nel SR per la pianificazione a 4 passi.
• Orizzonti Lunghi: In scenari con orizzonti estesi (fino a $T=6$), i modelli esistenti mostrano un rapido degrado delle prestazioni. GeoWorld mantiene una stabilità superiore, ottenendo il miglior SR in tutti gli orizzonti testati.
• Ablazione: Gli studi dimostrano che:
  • L'uso della geometria iperbolica da sola migliora la stabilità rispetto allo spazio euclideo.
  • L'aggiunta di GRL porta a ulteriori guadagni significativi.
  • La combinazione di SFT (Supervised Fine-Tuning) e GRL offre le prestazioni migliori.
  • La curvatura appresa ($c$) tende a stabilizzarsi su valori che bilanciano la struttura gerarchica e la stabilità della propagazione del valore.

5. Significato e Impatto

Il lavoro di GeoWorld segna un passo avanti fondamentale nell'integrazione dei principi geometrici nei modelli predittivi del mondo.

• Superamento dei Limiti Euclidei: Dimostra che lo spazio euclideo è inadeguato per rappresentare la complessità gerarchica e l'espansione esponenziale dei futuri possibili nella pianificazione robotica e visiva.
• Efficienza Computazionale: Rispetto ai modelli generativi che devono decodificare pixel (rumorosi e costosi), GeoWorld opera nello spazio latente, rendendo la pianificazione più efficiente e robusta.
• Pianificazione a Lungo Termine: La capacità di mantenere la coerenza temporale su orizzonti estesi apre nuove possibilità per l'addestramento di agenti autonomi in ambienti complessi e dinamici, riducendo l'accumulo di errori tipico dei metodi autoregressivi.

In sintesi, GeoWorld propone che la comprensione della geometria sottostante (in questo caso iperbolica) sia essenziale per costruire modelli del mondo capaci di ragionamento e pianificazione robusti e scalabili.