Riepilogo Tecnico: Modello Mondiale Prismatico (PRISM-WM)

1. Enunciato del Problema

L'apprendimento per rinforzo basato su modelli (MBRL) e la pianificazione nei domini robotici sono fondamentalmente limitati dalla natura ibrida delle dinamiche fisiche. I sistemi del mondo reale, in particolare nella locomozione su gambe e nella manipolazione, coinvolgono movimenti continui interrotti da eventi discreti (ad esempio, contatti, impatti, adesione vs. scorrimento). Queste interazioni creano dinamiche non lisce e altamente non lineari con discontinuità nette nel campo vettoriale sottostante.

I modelli mondiali latenti convenzionali impiegano tipicamente reti neurali monolitiche che impongono una continuità globale. Sebbene teoricamente capaci di approssimazione universale, questi modelli apprendono praticamente una rappresentazione media che appiattisce eccessivamente le diverse modalità dinamiche. Per i pianificatori di ottimizzazione delle traiettorie, questo eccessivo livellamento induce errori cumulativi durante le previsioni a lungo termine, in particolare ai confini fisici. Ciò comporta:

• Sfocatura spettrale attraverso i confini di contatto, causando un significativo disallineamento dei gradienti.
• Deriva cumulativa negli spazi di configurazione ad alta dimensionalità, rendendo inaffidabile la pianificazione a lungo termine.
• Collasso della politica in ambienti ricchi di contatti, dove i modelli standard non riescono a distinguere tra regimi stabili e instabili.

2. Metodologia: Il Modello Mondiale Prismatico (PRISM-WM)

Per affrontare queste limitazioni, gli autori introducono PRISM-WM, un'architettura strutturata progettata per decomporre le complesse dinamiche ibride in primitive componibili. Il modello si basa su un framework Mixture-of-Experts (MoE) consapevole del contesto con un obiettivo innovativo di ortogonalizzazione latente.

Architettura Principale

• Gating Consapevole del Contesto: Una rete di gating ($G$) riceve un vettore di contesto $c_t$ (composto dallo stato $z_t$, dall'azione $a_t$ e, opzionalmente, dall'incorporamento del compito $e_{task}$) e produce pesi normalizzati tramite softmax $w_t$. Questo meccanismo identifica implicitamente il regime fisico attivo (ad esempio, volo vs. appoggio) senza supervisione esplicita.
• Esperti Specializzati: Lo stato viene instradato a un insieme di $K$ reti esperte $\{E_k\}$. Ogni esperto si specializza nell'apprendere il flusso continuo all'interno di un specifico regime dinamico locale, evitando interferenze da modalità conflittuali.
• Predizione Residuale: Il prossimo stato latente è previsto come una combinazione pesata delle uscite degli esperti aggiunta allo stato corrente:
  $$\hat{z}_{t+1} = z_t + \sum_{k=1}^{K} w_{t,k} \cdot v_k(x_t)$$
  dove $v_k$ rappresenta le dinamiche residue previste dall'esperto $k$-esimo.

Innovazione Chiave: Ortogonalizzazione Latente

Per prevenire il collasso delle modalità (dove gli esperti apprendono caratteristiche ridondanti) e garantire la diversità, PRISM-WM introduce un strato di ortogonalizzazione di Gram-Schmidt differenziabile.

• Dopo che ogni esperto ha prodotto un vettore del penultimo livello $u_k$, il modello applica un processo di Gram-Schmidt stabilizzato per generare vettori di base ortogonali $v_k$.
• Questo impone alle rappresentazioni degli esperti di giacere su una varietà di Stiefel, assicurando che esse coprano uno spazio sottomesso di rango massimo e non ridondante per le dinamiche residue.
• Sebbene le modalità di contatto fisico siano non ordinate, questa ortogonalizzazione garantisce che gli esperti catturino componenti distinti della varietà residua, permettendo alla rete di gating di ricostruire modalità fisiche non ordinate tramite una somma pesata.

Integrazione con la Pianificazione

PRISM-WM è integrato nel framework di pianificazione TD-MPC. Fattorizzando le dinamiche complesse in esperti non ridondanti, il modello riduce il residuo del modello a un passo da $\epsilon$ a circa $\epsilon/K\alpha$. Questa decomposizione strutturale mitiga la deriva cumulativa, consentendo roll-out a lungo termine ad alta fedeltà e una propagazione stabile dei gradienti per l'apprendimento della politica.

3. Contributi Chiave

1. Decomposizione Architettonica: Introduzione di PRISM-WM, che scompone le dinamiche ibride utilizzando un'architettura MoE consapevole del contesto accoppiata all'ortogonalizzazione latente, allontanandosi dalle approssimazioni globali monolitiche.
2. Obiettivo di Ortogonalizzazione: Un obiettivo di addestramento innovativo che utilizza l'ortogonalizzazione di Gram-Schmidt differenziabile per imporre la diversità degli esperti e prevenire il collasso delle modalità, allineando strutturalmente lo spazio latente alle transizioni fisiche discrete.
3. Validazione Teorica ed Empirica: Dimostrazione che questa decomposizione riduce gli errori cumulativi negli spazi ad alta dimensionalità e migliora la geometria del paesaggio di perdita sia per la pianificazione (TD-MPC) che per l'apprendimento diretto della politica (PWM).
4. Capacità di Transfer Zero-Shot: Evidenza che la base ortogonale permette al modello di generalizzare tra simulatori con fisica di contatto diversa (ad esempio, da IsaacGym a contatto morbido a MuJoCo a contatto rigido) isolando specifiche modalità dinamiche (ad esempio, attrito vs. supporto verticale).

4. Risultati Sperimentali

Gli autori hanno valutato PRISM-WM su benchmark di controllo continuo, inclusi umanoidi ad alta dimensionalità, impostazioni multi-compito e transfer sim-to-sim zero-shot.

• Locomozione ad Alta Dimensionalità (DiffRL & Humanoid-Bench): PRISM-WM ha raggiunto una maggiore efficienza nel campionamento e prestazioni asintotiche superiori rispetto alle basi monolitiche (TD-MPC2, DreamerV3, PWM, SAC). In particolare, mentre i modelli monolitici hanno sperimentato un collasso della politica nei compiti di locomozione bipede (Humanoid, SNU-Humanoid) a causa della media di eventi di contatto non lisci, PRISM-WM ha mantenuto un'ottimizzazione stabile.
• Generalizzazione Multi-Compito: Su un benchmark della DeepMind Control Suite con 30 compiti, PRISM-WM ha ottenuto un punteggio normalizzato medio di 0,531, un miglioramento del 23,5% rispetto alla base monolitica TD-MPC2 (0,430). Ha superato la base in 26 compiti su 30, mitigando efficacemente il transfer negativo tra dinamiche eterogenee.
• Transfer Sim-to-Sim Zero-Shot: Nel trasferimento da IsaacGym a MuJoCo, la base MLP monolitica ha subito un calo delle prestazioni del 47,8% con un tasso di successo del 12%. PRISM-WM ha ridotto questo calo all'11,8% con un tasso di successo dell'88%. L'ortogonalizzazione ha permesso al modello di regolare i pesi per esperti specifici (ad esempio, attrito orizzontale) senza compromettere la varietà delle dinamiche globali.
• Stabilità ed Efficienza: Studi di ablazione hanno confermato che l'ortogonalizzazione previene il degrado della stima della ricompensa osservato nei modelli MoE non vincolati. I test di efficienza di inferenza su una NVIDIA RTX 4090 hanno mostrato che PRISM-WM (K=4) ha raggiunto 7.712 FPS, superando leggermente la base monolitica (7.499 FPS) con un sovraccarico di memoria minimo.

5. Significato e Affermazioni

Il documento afferma che PRISM-WM fornisce un substrato ad alta fedeltà per gli algoritmi di ottimizzazione delle traiettorie, agendo come un modello fondazionale per agenti basati su modelli. Scomponendo la fisica complessa in primitive ortogonali, il modello:

• Mitiga la deriva a lungo termine negli orizzonti di pianificazione, consentendo un ragionamento affidabile a lungo termine in ambienti ricchi di contatti.
• Stabilizza la propagazione dei gradienti per l'apprendimento della politica, trasformando paesaggi di perdita fratturati in miscele convesse funzionalmente lisce.
• Abilita un processo decisionale robusto in sistemi ibridi dove i modelli mondiali standard falliscono a causa dell'eccessivo livellamento.

Gli autori concludono che la scomposizione della fisica complessa in primitive ortogonali è "fondamentalmente essenziale per padroneggiare il controllo ricco di contatti". Sebbene l'attuale framework richieda la regolazione manuale del numero di esperti $K$, il lavoro stabilisce una base strutturale per la costruzione di modelli mondiali di scopo generale e fondati sulla fisica per futuri agenti incarnati.