Sintesi Tecnica: Modelli di Mondo con Equivarianza di Flusso (Flow Equivariant World Models)

1. Definizione del Problema

Gli agenti incarnati operano in ambienti dinamici dove possiedono solo un campo visivo ristretto, in prima persona. Questa osservabilità parziale crea una sfida significativa per gli attuali modelli di mondo: mentre l'agente osserva una frazione dell'ambiente, le regioni non osservate continuano a evolversi. Gli attuali metodi allo stato dell'arte, inclusi i transformer di diffusione latente su larga scala e i modelli ricorrenti, spesso non riescono a mantenere la coerenza su orizzonti temporali lunghi in questi contesti.

Senza una struttura di memoria che rispetti le simmetrie sottostanti del mondo, questi modelli soffrono di:

• Perdita di Informazioni: Mentre l'agente si muove, le informazioni passate rilevanti escono dalla finestra di contesto (ad esempio, le finestre scorrevoli nei transformer) e vengono scartate.
• Allucinazione: Quando l'ageto torna su una regione precedentemente osservata, il modello spesso non riesce a ricordare lo stato dell'oggetto, generando invece nuovi oggetti allucinati o facendo svanire quelli esistenti.
• Mancanza di Prior Strutturali: I modelli esistenti trattano le osservazioni come sequenze statiche o dati dipendenti dalla vista, ignorando il fatto che sia il movimento proprio dell'agente che il movimento degli oggetti esterni obbediscono a simmetrie fluide e parametrizzate nel tempo (flussi).

2. Metodologia: Modelli di Mondo con Equivarianza di Flusso (FloWM)

Gli autori propongono FloWM, un framework che integra l'equivarianza di flusso nel cuore della memoria del modello di mondo. L'ipotesi centrale è che le rappresentazioni predittive diventino più potenti quando sono organizzate in allineamento con la struttura temporale e dinamica del mondo.

Concetto Centrale: Equivarianza di Flusso

Il framework estende l'equivarianza di gruppo statica alle trasformazioni di sequenza parametrizzate nel tempo, ovvero i "flussi".

• Definizione: Un modello è equivariante al flusso se, quando la sequenza di input subisce un flusso (ad esempio, uno spostamento dovuto al tempo o al moto), la sequenza di output si trasforma secondo lo stesso flusso.
• Equivarianza Globale: In contesti di osservabilità parziale, il modello deve mantenere una memoria latente dell'intero stato del mondo (osservato e non osservato). La rappresentazione latente deve trasformarsi in modo equivariante rispetto allo stato globale del mondo, non solo rispetto all'osservazione corrente.
• Sistema di Riferimento Co-movente: Per ottenere questo, la memoria ricorrente è trasportata dallo stesso flusso dell'input prima di incorporare nuove osservazioni. Ciò assicura che le caratteristiche rimangano spazialmente allineate con gli oggetti in movimento piuttosto che essere lasciate indietro in un sistema di riferimento fisso.

Componenti Architetturali

Il framework è istanziato attraverso una relazione di ricorrenza generalizzata equivariante al flusso:
$$h_{t+1}(\nu) = \psi_1(\nu) \cdot U_\theta [h_t(\nu); E_\theta [f_t; h_t] (\nu)]$$
Dove:

• $h_t(\nu)$ rappresenta lo stato nascosto organizzato in "canali di velocità" ($\nu$), corrispondenti a diversi campi vettoriali di moto.
• $E_\theta$ è un encoder che mappa le osservazioni nello spazio latente.
• $U_\theta$ è l'operazione di aggiornamento.
• $\psi_1(\nu)$ rappresenta l'azione del flusso su un intervallo temporale (timestep).

Equivarianza del Moto Proprio:
Il modello tratta il moto proprio dell'agente come un flusso relativo. Trasformando lo stato nascosto secondo l'azione dell'agente ($T_{a_t}$), la memoria si sposta efficacementmente nel sistema di riferimento co-movente dell'agente. Ciò consente al modello di costruire una "mappa" strutturata dell'ambiente in cui le nuove osservazioni vengono integrate nelle corrette posizioni relative basandosi sulle azioni dell'agente.

Due Istanze

1. Simple Recurrent FloWM (2D): Progettato per ambienti 2D (es. MNIST World). Utilizza operazioni convoluzionali sugli stati nascosti e sugli input spaziali. Impiega una finestra fissa per leggere/scrivere nello stato nascosto, mentre il resto dello stato fluisce attorno al campo visivo dell'agente.
2. Transformer-Based FloWM (3D): Progettato per ambienti 3D complessi (es. Block World). Utilizza un encoder e un decoder Vision Transformer (ViT). La memoria latente è un insieme di embedding di token organizzati spazialmente che agiscono come una mappa astratta top-down.
   • Nota sull'Equivarianza: Sebbene l'equivarianza esatta dell'encoder sia difficile per le proiezioni 3D-2D senza costose deproiezioni di profondità, gli autori dimostrano che la trasformazione del flusso ricorrente ($T^{-1}_{a_t} \psi_1(\nu)$) incoraggia l'encoder a imparare l'equivarianza durante l'addestramento, ottenendo un'alta fedeltà senza vincoli geometrici espliciti.

3. Contributi Chiave

• Formalizzazione dell'Equivarianza di Flusso: Il paper formalizza come le simmetrie parametrizzate nel tempo (flussi) possano essere utilizzate per costruire una memoria strutturata e persistente per la modellazione di mondi dinamici con osservabilità parziale.
• Architettura di Memoria Strutturata: Introduzione di una memoria ricorrente con "canali di velocità" che si sposta e si trasforma in modo equivariante sia rispetto al moto proprio che al moto degli oggetti esterni, mantenendo le informazioni sulle regioni fuori dalla vista allineate nel tempo.
• Nuovi Benchmark: Gli autori introducono nuovi benchmark per la modellazione di mondi con osservabilità parziale:
  • 2D MNIST World: Cifre che si muovono a velocità casuali con osservazione parziale.
  • 3D Dynamic Block World: Blocchi colorati che rimbalzano in una stanza 3D con un agente che naviga tramite azioni discrete.
• Validazione Empirica: Dimostrazione che FloWM supera significativamente i modelli di diffusione allo stato dell'arte (DFoT, DFoT-SSM) e i modelli ricorrenti (RSSM) in termini di previsione a lungo termine e generalizzazione della lunghezza.

4. Risultati Sperimentali

Gli autori hanno valutato FloWM rispetto ai baseline, tra cui History-Guided Diffusion Forcing (DFoT), DFoT-SSM e l'RSSM di DreamerV3.

Metriche di Performance

• Coerenza a Lungo Termine: Su 2D MNIST World, FloWM ha mantenuto previsioni coerenti per 150 timestep oltre l'orizzonte di addestramento (20 timestep), mentre i baseline divergevano o allucinavano. Su 3D Block World, FloWM è rimasto coerente per 210 frame, mentre i baseline perdevano traccia degli oggetti o allucinavano nuovi elementi.
• Riduzione dell'Errore: FloWM ha ottenuto un MSE significativamente più basso e punteggi PSIM/SSIM più elevati rispetto a tutti i baseline. Ad esempio, in 3D Block World, l'MSE di FloWM era di circa 0.0006 (70 frame) contro lo 0.0117 di DFoT.
• Efficienza di Apprendimento: I modelli che combinano l'Equivarianza del Moto Proprio (SME) e i Canali di Velocità (VC) hanno richiesto ordini di grandezza in meno di passi di addestramento per convergere rispetto alle versioni con ablazione.
• Pianificazione Downstream: In un compito di pianificazione "Trova il blocco rosso", FloTM è riuscito a navigare verso l'obiettivo. I baseline spesso allucinavano la posizione del target, portando a una scarsa performance di pianificazione.
• Analisi dei Probe: Una rete di sonda addestrata per decodificare le posizioni dei blocchi dallo spazio latente ha mostrato un'accuratezza del 96% per FloVM rispetto a <1% per DFoT, confermando che FloVM apprende una rappresentazione strutturata ed equivariante delle traiettorie degli oggetti.

Studi di Ablazione

La rimozione dei Canali di Velocità (VC) o dell'Equivarianza del Moto Proprio (SME) ha causato un degrado significativo delle prestazioni, spesso portandole al di sotto dei modelli di diffusione baseline, evidenziando che entrambi i componenti sono critici per gestire l'osservabilità parziale e la dinamica.

5. Significato e Rivendicazioni

Il paper sostiene che le rappresentazioni predittive sono fondamentalmente più potenti quando sono organizzate in linea con la struttura temporale e dinamica del mondo.

• Fondamenti Teorici: Il lavoro fornisce un quadro teorico (equivarianza di flusso) che detta la necessità di una memoria strutturata per i modelli di mondo in ambienti dinamici e con osservabilità parziale.
• Risoluzione dell'Osservabilità Parziale: A differenza dei metodi precedenti che si affidano all'estensione delle finestre di contesto o a banche di memoria statiche, FloWM gestisce intrinsecamente l'evoluzione delle regioni non osservate mantenendo una memoria che si muove con lo stato del mondo.
• Generalizzazione: Il framework consente di generalizzare a sequenze significativamente più lunghe di quelle viste durante l'addestramento, una capacità critica per gli agenti incarnati del mondo reale.
• Ambito Modesto: Gli autori riconoscono che le attuali valutazioni avvengono in contesti controllati con azioni geometriche di corpo rigido. Posizionano questo lavoro come una base per future estensioni a dataset di mondo aperto, spazi di azione continui e interazioni fisiche più complesse.

In sintesi, FloWM dimostra che, codificando esplicitamente le simmetrie del moto nell'architettura della memoria, i modelli di mondo possono raggiungere previsioni stabili, accurate e prive di allucinazioni su orizzonti temporali lunghi in ambienti dinamici con osservabilità parziale.