Riepilogo Tecnico: Modelli di Mondo Sensomotorio (SMWM)

Definizione del Problema

Il documento affronta la sfida di apprendere modelli di mondo latenti da osservazioni ad alta dimensionalità (ad esempio, pixel) senza segnali di ricompensa o etichette di task. Le Architetture Predittive di Embedding Congiunto (JEPA) offrono un framework promettente prevedendo direttamente gli stati futuri in uno spazio di embedding (evitando la ricostruzione dei pixel), ma soffrono di collasso della rappresentazione. Quando un encoder e un modello di dinamica forward vengono addestrati congiuntamente per minimizzare l'errore di predizione, l'encoder può apprendere una mappatura di tutte le osservazioni verso un embedding costante. Questa soluzione banale minimizza la perdita forward ma rende il modello inutile per il controllo o la pianificazione a valle.

Le soluzioni esistenti per prevenire il collasso spesso si affidano a meccanismi complessi come:

• L'uso di encoder pre-addestrati congelati (es. DINO-WM).
• L'aggiunta di regolarizzatori distribuzionali (es. far corrispondere gli embedding a una Gaussiana isotropa tramite SIGReg).
• L'impiego di reti target con media mobile esponenziale (EMA) e stop-gradient.
• L'uso di funzioni di perdita multi-termine che coinvolgono vincoli di varianza-covarianza.

Gli autori sostengono che questi approcci potrebbero non affrontare pienamente la questione fondamentale di cosa una rappresentazione utile di un modello di mondo debba preservare. Basandosi sul concetto di "percezione per l'azione" e sulla percezione enattiva, il documento argomenta che le rappresentazioni dovrebbero essere modellate dalla loro rilevanza rispetto alle azioni, preservando le distinzioni che contano per le dinamiche controllabili e scartando i distrattori incontrollabili.

Metodologia: Modello di Mondo Sensomotorio (SMWM)

Gli autori propongono SMWM, un modello di mondo latente addestrato end-to-end da traiettorie offline prive di ricompensa $(o_t, a_t, o_{t+1})$. L'innovazione principale è l'uso della regolarizzazione della dinamica inversa come unico meccanismo per prevenire il collasso della rappresentazione.

Architettura

Il modello consiste in tre componenti addestrate congiuntamente:

1. Encoder ($f_\theta$): Mappa le osservazioni $o_t$ in un embedding latente compatto $z_t$.
2. Modello di Dinamica Forward ($g_\phi$): Predice lo stato latente successivo $\hat{z}_{t+1}$ dato lo stato attuale $z_t$ e l'azione $a_t$.
3. Modello di Dinamica Inversa ($h_\psi$): Predice l'azione $\hat{a}_t$ compiuta tra due embedding consecutivi $z_t$ e $z_{t+1}$.

Obiettivo di Addestramento

Il modello è ottimizzato minimizzando una funzione di perdita congiunta:
$$\mathcal{L} = \mathcal{L}_{fwd} + \lambda \mathcal{L}_{inv}$$

• Perdita Forward ($\mathcal{L}_{fwd}$): Errore Quadratico Medio (MSE) tra l'embedding successivo predetto $\hat{z}_{t+1} = g_\phi(z_t, a_t)$ e l'effettivo stato codificato $z_{t+1} = f_\theta(o_{t+1})$.
• Perdita Inversa ($\mathcal{L}_{inv}$): MSE tra l'azione predetta $\hat{a}_t = h_\psi(z_t, z_{t+1})$ e l'azione reale $a_t$.
• Peso di Regolarizzazione ($\lambda$): Un singolo iperparametro che controlla la forza del segnale di dinamica inversa.

Meccanismo d'Azione

La perdita inversa funge da meccanismo anti-collasso. Se l'encoder collassa verso un embedding costante, il modello inverso riceve input identici per tutte le transizioni e può solo predire un'azione costante (la media della distribuzione delle azioni). Per ridurre la perdita inversa al di sotto di questo rischio di predittore costante, l'encoder deve preservare informazioni sull'azione $a_t$ contenute nella transizione $(o_t, o_{t+1})$. Ciò costringe la rappresentazione a mantenere informazioni rilevanti per l'azione (gradi di libertà controllabili) filtrando naturalmente i distrattori incontrollabili (rumore esogeno) che non influenzano il recupero dell'azione.

A differenza dei regolarizzatori distribuzionali che prescrivono la geometria dello spazio di embedding (ad esempio, forzandolo a essere Gaussiano), la dinamica inversa ancora la rappresentazione a una quantità legata al compito (l'azione) senza imporre specifici priori geometrici.

Contributi Chiave

1. Architettura SMWM: Un modello di mondo sensomotorio semplice, addestrabile end-to-end, che utilizza la dinamica inversa come unico meccanismo anti-collasso, richiedendo un solo iperparametro aggiuntivo ($\lambda$).
2. Proprietà della Rappresentazione: Evidenza empirica che le rappresentazioni apprese non sono solo non-collassate, ma tracciano specificamente i gradi di libertà controllabili, preservano la geometria spaziale (sia cartesiana che periodica) e filtrano i distrattori incontrollabili.
3. Performance a Valle: Dimostrazione di una performance di pianificazione competitiva in compiti di controllo 2D e 3D, mostrando che i vantaggi strutturali della rappresentazione si traducono in una pianificazione guidata dagli obiettivi efficace.

Risultati Sperimentali

Testbed Controllati (Dot World & Sprite World)

• Recupero della Dimensione Intrinseca: In un ambiente "dot world", il modello ha recuperato con successo la vera dimensione intrinseca 2D dello spazio di stato partendo dai pixel, anche quando immersa in uno spazio latente a 64 dimensioni.
• Filtraggio dei Distrattori: Quando sono stati introdotti oggetti "distrattori" incontrollabili (che si muovevano casualmente), l'encoder li ha ignorati. La dimensione latente effettiva corrispondeva ai gradi di libertà controllabili, non alla dimensione totale dell'osservazione.
• Percezione Dipendente dal Controllo: In un esperimento su un mondo di sprite, il modello ha imparato a preservare le variabili di posa (posizione, orientamento) solo quando erano controllabili. Quando l'orientamento era incontrollato, la rappresentazione lo mediava, risultando in una ricostruzione sfocata dell'orientamento.

Benchmark di Pianificazione

Gli autori hanno valutato SMWM su quattro compiti di controllo offline: TwoRoom (navigazione 2D), Reacher (controllo continuo), Push-T (manipolazione 2D) e OGBench-Cube (manipolazione 3D).

• Baseline: Confrontato con SIGReg (una forte baseline JEPA che utilizza regolarizzatori Gaussiani), Forward-only (senza perdita inversa) e campionamento di azioni Random.
• Performance:
  • SMWM ha eguagliato le performance di SIGReg nei tre compiti 2D.
  • SMWM ha superato SIGReg nel compito 3D OGBench-Cube (84% di successo contro il 59%).
  • Entrambi i metodi regolarizzati hanno superato significativamente la ablazione Forward-only, confermando la necessità di un meccanismo anti-collasso.
• Geometria Latente: Esperimenti di probing hanno rivelato che gli embedding di SMWM formano sottospazi a bassa dimensione che rispecchiano la geometria fisica dello stato (ad esempio, strutture cilindriche per angoli giunti periodici in Reacher, mappe cartesiane per la navigazione).

Significato e Rivendicazioni

Il documento sostiene che SMWM fornisca un'istanzai principata e semplice di "percezione per l'azione" nell'apprendimento di modelli di mondo latenti. Obbligando gli stati latenti a preservare informazioni sufficienti per recuperare le azioni dell'agente, il modello orienta naturalmente la rappresentazione verso i gradi di libertà controllabili dell'ambiente.

I punti chiave includono:

• Semplicità: La sola dinamica inversa è sufficiente a stabilizzare l'addestramento end-to-end di una JEPA senza ricostruzione, segnali di ricompensa, encoder congelati o complessi priori distribuzionali.
• Bias Strutturale: Il metodo induce rappresentazioni che sono interpretabili e allineate con la struttura causale dell'ambiente (azioni e i loro effetti), piuttosto che semplici correlazioni predittive.
• Utilità Pratica: Le rappresentazioni apprese supportano una pianificazione latente efficace, particolarmente in ambienti 3D complessi dove altre baseline regolarizzate faticano.

Gli autori riconoscono limitazioni, come l'assunzione che le azioni siano recuperabili da osservazioni consecutive e l'attuale incapacità di catturare quantità non identificabili da singoli frame (come la velocità) senza storie di osservazione. Tuttavia, pongono SMWM come un approccio robusto e privo di ricompensa per apprendere modelli di mondo azionabili da dati offline.