Riepilogo Tecnico: Ritorno ai Latenti Parsimoniosi: Apprendimento di Modelli del Mondo Centrati sul Compito da Fondamenti Visivi

Enunciato del Problema

I modelli del mondo permettono agli agenti di prevedere le dinamiche future condizionate dalle azioni, una capacità centrale per la pianificazione e il controllo. Tuttavia, gli approcci esistenti affrontano un compromesso fondamentale nella rappresentazione latente:

1. Modelli basati sui pixel (ad es. Dreamer, MuZero) apprendono le dinamiche direttamente dai pixel o da latenti compatti modellati dai segnali di ricompensa. Sebbene efficaci, spesso comportano elevati costi computazionali o richiedono interazione online/supervisione tramite ricompensa per modellare lo spazio latente.
2. Modelli di embedding fondativi (ad es. DINO-WM, V-JEPA) sfruttano encoder visivi pre-addestrati congelati per prevedere direttamente in uno spazio semantico ad alta dimensionalità. Sebbene offrano ricchi prior semantici e generalizzazione zero-shot, gli embedding contengono dettagli eccessivamente irrilevanti per il compito (texture, illuminazione, sfondo) che non influenzano il controllo.

In ambienti offline privi di ricompensa, dove gli agenti devono apprendere da traiettorie fisse senza supervisione tramite ricompensa, questi problemi si aggravano. Utilizzare embedding fondativi grezzi come spazio degli stati costringe il modello a sprecare capacità su variazioni irrilevanti, portando a una scarsa efficienza di modellazione, ridotta interpretabilità del controllo e potenziale collasso della rappresentazione, specialmente in spazi di azione ad alta dimensionalità (ad es. bracci robotici a 7 gradi di libertà). La sfida principale è estrarre una rappresentazione latente compatta e sufficiente al compito da questi embedding fondativi ricchi ma rumorosi senza interazione online.

Metodologia: TC-WM

Gli autori propongono TC-WM (Task-Centric World Model), un framework che tratta gli embedding fondativi pre-addestrati non come lo spazio degli stati finale, ma come un "impalcatura semantica" da cui vengono distillate le dinamiche rilevanti per il compito.

Architettura Principale

TC-WM opera su una pipeline di embedding visivi ($x_t$) e segnali propriocettivi ($s^p_t$):

1. Embedding Congiunto: Un encoder visivo congelato (ad es. DINOv2) mappa le immagini in caratteristiche visive, mentre un layer lineare addestrabile codifica la propriocettione. Questi vengono concatenati per formare un embedding congiunto $x_t$.
2. Encoder Latente (Proiezione): Un encoder lineare proietta l'embedding congiunto ad alta dimensionalità in uno stato latente compatto $z_t$. Crucialmente, $z_t$ è fattorizzato in due sottospazi:
   • $z^s_t$ (Rilevante per lo stato): Allineato allo stato fisico.
   • $z^c_t$ (Complementare): Mantiene la struttura semantica residua (aspetto, sfondo) per soddisfare i vincoli di ricostruzione.
3. Dinamiche Latenti: Un predittore basato su Transformer ($f_\theta$) modella le transizioni temporali nello spazio latente compatto ($\hat{z}_{t+1} = f_\theta(z_{t-H:t}, a_{t-H:t})$).
4. Allineamento Centrato sul Compito: Il modello impiega un obiettivo contrastivo (InfoNCE) per allineare il sottospazio $z^s_t$ con lo stato propriocettivo $s^p_t$. Questo forza attivamente l'encoder a separare i fattori fisici rilevanti per il compito dalle variazioni di disturbo.
5. Decoder di Embedding: Un decoder lineare ricostruisce l'embedding visivo originale da $z_t$. Questo agisce come regolarizzatore, assicurando che il latente compatto rimanga ancorato alla varietà fondativa e prevenendo soluzioni degeneri.

Obiettivo di Addestramento

La perdita totale combina quattro componenti:

• Perdita delle Dinamiche Latenti ($L^z_{dyn}$): Minimizza l'errore di previsione nello spazio latente.
• Perdita delle Dinamiche TC ($L^s_{dyn}$): Prevede la propriocettione futura direttamente dalla storia per imporre dinamiche centrate sul compito.
• Perdita di Allineamento ($L_{align}$): Perdita contrastiva che allinea $z^s_t$ con $s^p_t$.
• Perdita di Ricostruzione ($L_{rec}$): Ricostruisce l'embedding visivo da $z_t$ per preservare la struttura visiva.

Pianificazione al Tempo di Test

Lo spazio latente centrato sul compito appreso supporta diverse strategie di pianificazione:

• Controllo a bassa dimensionalità (Navigazione): Controllo Predittivo del Modello (MPC) che utilizza il Metodo dell'Entropia Incrociata (CEM) per ottimizzare sequenze di azioni nello spazio latente.
• Manipolazione ad alta dimensionalità: Pianificazione con Diffusione Latente (LDP) combinata con un Modello di Dinamica Inversa (IDM) per generare azioni eseguibili da traiettorie latenti.

Contributi Chiave

1. Decomposizione Centrata sul Compito: Il documento introduce un metodo per decomporre embedding fondativi ad alta dimensionalità in uno spazio latente compatto e sufficiente al compito, separando esplicitamente le dinamiche dello stato fisico dalle variabili di disturbo semantiche.
2. Identificabilità Attiva: A differenza di lavori precedenti che assumono passivamente la diversità dei meccanismi per l'identificabilità, TC-WM utilizza l'allineamento contrastivo con la propriocettione per attivare l'induzione di condizionali distinti per diversi stati fisici, garantendo teoricamente il recupero dei fattori centrati sul compito fino a una trasformazione affine.
3. Apprendimento Offline Privo di Ricompensa: Il framework raggiunge una modellazione del mondo e una pianificazione di alta qualità in ambienti offline senza segnali di ricompensa, affidandosi esclusivamente a osservazioni visive e informazioni laterali propriocettive.
4. Garanzia Teorica: Gli autori forniscono un teorema che dimostra come, sotto ipotesi lievi (dinamiche contestuali, variabilità latente-osservata e allineamento sufficiente), il latente compatto appreso sia identificabile e recuperi il vero stato fisico fino a una riparametrizzazione.

Risultati Sperimentali

Gli autori valutano TC-WM su 9 benchmark offline (Maze, Wall, Push-T, Lift, Can, Square, Reacher, Cheetah, Hopper) confrontandolo con baseline tra cui TD-MPC2, DreamerV3, MuZero e DINO-WM.

• Qualità del Modello del Mondo: TC-WM ottiene l'errore di previsione di rollout latente più basso (MSE) in quasi tutti i compiti, in particolare negli ambienti di manipolazione dove sono critiche dinamiche a lungo orizzonte accurate. Sebbene la ricostruzione nello spazio delle immagini sia competitiva, il guadagno principale risiede nella stabilità e accuratezza delle dinamiche latenti.
• Prestazioni di Pianificazione:
  • Nella navigazione (azioni a bassa dimensionalità), TC-WM eguaglia le prestazioni di baseline solide come TD-MPC2.
  • Nella manipolazione (azioni ad alta dimensionalità, ricche di contatti), TC-WM supera significativamente tutte le baseline. Ad esempio, nel compito "Lift" di Robomimic, TC-WM raggiunge un tasso di successo del 54% rispetto al 29% di DINO-WM e al 31% di TD-MPC2.
• Analisi della Rappresentazione:
  • Sondaggio Lineare: Regressori lineari addestrati sullo spazio latente di TC-WM recuperano stati propriocettivi e degli oggetti con punteggi $R^2$ superiori rispetto alle baseline, confermando che il latente è fondato fisicamente.
  • Anti-Collasso: TC-WM dimostra un migliore utilizzo dello spazio latente (rank effettivo più alto) rispetto a DINO-WM, che soffre di collasso della rappresentazione quando si utilizzano direttamente embedding grezzi.
  • Robustezza: TC-WM mantiene le prestazioni sotto perturbazioni visive (rumore gaussiano, jitter di colore) meglio di DINO-WM, con riduzioni minori nell'SSIM.

Significato e Affermazioni

Il documento afferma che gli embedding fondativi sono meglio utilizzati come impalcatura semantica piuttosto che come rappresentazione diretta dello stato. Proiettando questi embedding in un sottospazio centrato sul compito allineato con segnali fisici, gli agenti possono ottenere:

• Latenti Parsimoniosi: Una riduzione della dimensionalità che filtra i dettagli visivi irrilevanti preservando al contempo le dinamiche necessarie per il controllo.
• Controllo Migliorato: Prestazioni di pianificazione potenziate in ambienti complessi e ad alta dimensionalità dove i modelli fondativi standard falliscono a causa del rumore e delle variazioni irrilevanti.
• Solidità Teorica: Una giustificazione formale del motivo per cui proiezioni lineari e allineamento contrastivo sono sufficienti per recuperare fattori rilevanti per il compito da encoder visivi congelati.

Gli autori posizionano TC-WM come un ponte tra il potere di generalizzazione dei modelli fondativi e la controllabilità richiesta per gli agenti fisici, affrontando specificamente il divario negli scenari di apprendimento offline privi di ricompensa.