Riepilogo Tecnico: FF-JEPA: Pianificazione a lungo termine in modelli del mondo con pianificatori latenti

Definizione del Problema

I recenti progressi nei modelli del mondo visivi, in particolare le Joint Embedding Predictive Architectures (JEPA) come LeWorldModel (LeWM), hanno dimostrato la capacità di prevedere osservazioni future e supportare il controllo basato su modelli tramite pianificatori basati sul campionamento come il Cross-Entropy Method (CEM). Tuttavia, gli autori identificano tre criticità fondamentali che impediscono l'implementazione di questi modelli "piatti" (flat) in compiti reali a lungo termine:

1. Proibizioni Computazionali: La pianificazione su orizzonti temporali lunghi richiede rollout ripetuti e sequenze di azioni sempre più ampie, rendendo l'ottimizzazione tramite CEM computazionalmente intrattabile.
2. Accumulo di Errore: Gli errori di previsione autoregressiva si accumulano nel corso di molti passi, causando la divergenza del pianificatore per traiettorie complesse.
3. Dipendenza dall'Obiettivo: I metodi esistenti richiedono tipicamente un'immagine o uno stato di obiettivo esplicito ($o_g$) per guidare la pianificazione. In molti scenari reali, una specifica di obiettivo così concreta non è disponibile o è impraticabile.

Sebbene i metodi gerarchici e quelli generativi abbiano tentato di affrontare questi problemi, essi operano spesso ancora nello spazio dei pixel o si affidano a immagini/traiettorie di obiettivo specificate esternamente, limitandone l'applicabilità in ambienti aperti.

Metodologia: Forward-Forward JEPA (FF-JEPA)

Il paper propone FF-JEPA, un framework gerarchico che estende i modelli del mondo in stile JEPA introducendo un pianificatore latente ($G$) che opera accanto a un modello forward standard condizionato dalle azioni ($P$). Questo approccio decompone la pianificazione a lungo termine in una sequenza di problemi di ottimizzazione a breve termine trattabili, senza richiedere immagini di obiettivo esplicite.

Architettura Core

Il sistema è composto da due componenti primarie addestrate all'interno dello spazio latente definito da un encoder di un modello del mondo pre-addestrato e congelato ($E$):

1. Il Modello del Mondo ($P$): Un modello di dinamica forward standard condizionato dalle azioni (basato su LeWM) che predice il successivo stato latente dato una storia di stati e un'azione: $\hat{z}_{t+1} = P(z_{t-W_P+1:t}, a_t)$.
2. Il Pianificatore Latente ($G$): Un modello di dinamica forward privo di condizionamento ad azioni che predice il prossimo subgoal latente ($\hat{z}_{sg}$) basandosi esclusivamente sulla storia dei precedenti subgoal latenti. Opera a una risoluzione temporale superiore (ogni $H$ passi dell'ambiente).
   • Pianificatore Deterministico ($G_{Det}$): Un'architettura transformer (simile al predittore del modello del mondo ma senza condizionamento all'azione) addestrata per minimizzare l'errore quadratico medio tra i futuri subgoal previsti e quelli reali.
   • Pianificatore di Diffusione ($G_{DM}$): Un backbone DiT (Diffusion Transformer) addestrato con un obiettivo di denoising score-matching per predire una sequenza di futuri subgoal.

Inferenza della Pianificazione

Il processo di inferenza è gerarchico:

1. Predizione del Subgoal: Il pianificatore latente $G$ predice il prossimo stato di subgoal $\hat{z}_{sg, m+1}$ data la storia attuale dei subgoal.
2. Ottimizzazione delle Azioni: Il modello del mondo $P$ utilizza il CEM per trovare una sequenza di azioni ($a_{1:H}$) che guidi lo stato predetto verso il subgoal predetto $\hat{z}_{sg, m+1}$ entro un breve orizzonte $H$.
3. Iterazione: Questo ciclo si ripete, concatenando efficacemente piani a breve termine per raggiungere obiettivi a lungo termine.

Fondamentalmente, questo metodo elimina la necessità di un'immagine di obiettivo target al momento dell'inferenza; il pianificatore genera la traiettoria verso un obiettivo implicitamente definito.

Contributi Chiave

• Pianificazione a Lungo Termine Senza Obiettivo (Goal-Free): Il framework consente la pianificazione in modelli del mondo senza immagini di obiettivo esplicite, utilizzando un pianificatore latente per generare subgoal intermedi.
• Decomposizione Gerarchica: Affronta il "collasso a lungo termine" dei modelli del mondo piatti scomponendo traiettorie complesse in sottoproblemi più brevi e risolvibili, mitigando l'accumulo di errori e i costi computazionali.
• Addestramento Non Etichettato: Il pianificatore latente può essere addestrato su dimostrazioni di successo non etichettate (sottocampionate dallo spazio latente del modello del mondo), riducendo la dipendenza da dati etichettati con le azioni.
• Framework Unificato: Reinterpreta il modello del mondo come un modulo di dinamica inversa che opera su traiettorie latenti immaginate, unificando la modellazione predittiva e il controllo all'interno di un unico spazio latente coerente.

Valutazione Empirica

Gli autori hanno valutato FF-JEPA sul compito PushT utilizzando tre configurazioni sperimentali: breve orizzonte (25 passi), lungo orizzonte (75 passi) e inizializzazione casuale (300 passi con posizioni di partenza randomizzate).

Risultati

• Performance a Lungo Termine: Il baseline piatto LeWM è collassato sui compiti a lungo orizzonte (tasso di successo del 3,52% a $t=75$) ed è fallito completamente sotto inizializzazione casuale (0,00%). Al contrario, FF-JEPA (Diffusion) ha raggiunto il 91,80% di successo a $t=75$ e l'82,42% sotto inizializzazione casuale.
• Vantaggio Goal-Free: A differenza dei baseline (LeWM, DINO-WM) che richiedono un'immagine di obiettivo per l'inferenza, FF-JEPA no. Nonostante ciò, ha superato i baseline gerarchici DINO-WM sia in termini di breve orizzonte (96,09% vs 89,0%) che di lungo orizzonte (91,80% vs 61,0%).
• Efficienza dei Dati: Gli studi di ablazione hanno mostrato che l'addestramento del pianificatore di diffusione su soli 200 esempi di alta qualità (una riduzione di 40 volte dei dati) produce comunque prestazioni competitive (76,17% vs 82,42% con i dati completi), suggerendo un'alta efficienza dei dati.
• Overhead Computazionale: Il pianificatore deterministico ha aggiunto un overhead di tempo di inferenza trascurabile (2,1ms) rispetto al CEM (926ms), mentre il pianificatore di diffusione ha aggiunto ~242ms. Entrambi sono significativamente più veloci dei ripetuti rollout richiesti per la pianificazione piatta a lungo termine.

Analisi dei Fallimenti

I casi di fallimento si sono verificati tipicamente quando l'agente si allontanava troppo dalla posizione del subgoal predetto o quando il pianificatore latente generava uno stato di subgoal di bassa qualità (ad esempio, mancando l'agente).

Significato e Rivendicazioni

Il paper posiziona FF-JEPA come una direzione promettente per l'apprendimento robotico scalabile e goal-free. Dimostrando che i modelli del mondo condizionati dalle azioni, accoppiati con pianificatori latenti addestrati su dati non etichettati, possono superare i limiti della pianificazione piatta, gli autori suggeriscono che questo approccio consente agli agenti di operare in ambienti aperti dove le specifiche di obiettivo esplicite non sono disponibili. Il lavoro evidenzia come la decomposizione della pianificazione in subgoal a breve termine sia una strategia valida per mitigare i problemi di computazione e di accumulo di errore inerenti agli attuali approcci di modelli del mondo.