Riepilogo Tecnico: Scalabilità dell'Apprendimento per Rinforzo con Modelli del Mondo attraverso l'Ottimizzazione della Politica Diffusiva

1. Enunciato del Problema

L'apprendimento per rinforzo (RL) basato su modelli che utilizza modelli del mondo offre un paradigma promettente per la pianificazione efficiente in termini di campioni, consentendo agli agenti di "sognare" e pianificare all'interno di ambienti interni appresi. Tuttavia, la scalabilità di questi approcci rimane fondamentalmente limitata. Mentre i lavori precedenti hanno attribuito le lacune nelle prestazioni al bias del modello e agli errori di previsione cumulativi, questo documento identifica un collo di bottiglia più critico e poco esplorato: il disallineamento strutturale tra la ricerca e l'apprendimento del valore.

Negli approcci esistenti basati su modelli del mondo (ad esempio, la serie TD-MPC), il miglioramento della politica si basa tipicamente su un pianificatore basato sul campionamento (come il Model Predictive Path Integral, MPPI) che opera sul modello del mondo appreso. Al contrario, la funzione di valore viene addestrata su traiettorie generate da una politica separata, non basata sulla ricerca (spesso la politica comportamentale o una politica del buffer di replay). Ciò crea una discrepanza distribuzionale:

1. Lacuna nell'Iterazione del Valore: La funzione di valore viene addestrata su dati provenienti dalla politica $\beta$, ma viene interrogata dalla politica di ricerca $\pi$. Quando $\pi$ si discosta significativamente da $\beta$, le stime del valore diventano inaffidabili, portando spesso a una sovrastima spuria.
2. Lacuna nel Miglioramento della Politica: La procedura di ricerca sfrutta queste sovrastime, rafforzando gli errori e causando un ciclo vizioso di spostamento distribuzionale e amplificazione dell'errore. Di conseguenza, la ricerca non costituisce una vera ottimizzazione della politica nel senso convenzionale, non garantendo un miglioramento monotono o la convergenza verso la politica ottimale.

2. Metodologia: Ottimizzazione della Politica Diffusiva Basata su Modelli (MBDPO)

Per affrontare queste limitazioni, gli autori propongono la Ottimizzazione della Politica Diffusiva Basata su Modelli (MBDPO), un framework che unifica la ricerca e l'ottimizzazione della politica attraverso rappresentazioni di politica diffusiva all'interno di modelli del mondo latenti.

Framework Principale

MBDPO riformula l'ottimizzazione della politica come un processo diffusivo su traiettorie immaginate nello spazio latente, invece di affidarsi a pianificatori espliciti come MPPI. Il processo prevede:

• Modello del Mondo Latente: L'agente apprende un encoder $E$, dinamiche latenti $F$ e un predittore di ricompensa $R$ per operare in uno spazio latente compatto $Z$.
• Funzione di Energia Implicita: Per ancorare gli aggiornamenti della politica e prevenire lo spostamento distribuzionale, MBDPO apprende una funzione di energia implicita $E_\theta$ dal dataset del buffer tramite un obiettivo contrastivo. Questa funzione approssima la densità logaritmica della politica comportamentale $\beta$, agendo efficacemente come un vincolo KL.
• Corrispondenza del Punteggio tramite Traiettorie Immaginate: Invece di un processo diffusivo in avanti, la funzione di punteggio $\phi$ viene stimata direttamente dalle traiettorie immaginate all'interno del modello del mondo appreso.
  • L'algoritmo campiona sequenze di azioni pulite candidate da una distribuzione di proposta.
  • Queste sequenze vengono propagate attraverso le dinamiche latenti $F$ per calcolare un ritorno cumulativo regolarizzato dall'energia $\tilde{G}$.
  • I campioni vengono ridimensionati in base a $\tilde{G}$ per stimare la funzione di punteggio, che guida il processo diffusivo inverso.

Unificazione Teorica

Il documento stabilisce un'equivalenza teorica (Teorema 1) tra la condizione di corrispondenza del punteggio nei modelli diffusivi e la soluzione di un problema di ottimizzazione della politica vincolato da KL.

• L'obiettivo è massimizzare il ritorno atteso soggetto a un vincolo di divergenza KL rispetto alla politica comportamentale: $\max_{\pi_\phi} \mathbb{E}[G] \text{ s.t. } D_{KL}(\pi_\phi \| \beta) \leq \epsilon$.
• MBDPO dimostra che la politica ottimale per questo problema vincolato segue una distribuzione di Gibbs.
• Apprendendo la funzione di punteggio per corrispondere a questa distribuzione, il processo diffusivo inverso esegue implicitamente l'ottimizzazione vincolata, unificando così la ricerca (generazione di traiettorie) e l'apprendimento della politica (massimizzazione del valore) in un unico framework coerente.

Algoritmo Pratico

MBDPO adotta un paradigma di addestramento congiunto:

1. Modello del Mondo e Iterazione del Valore: L'encoder, le dinamiche, la ricompensa, la funzione di valore ($\hat{Q}$) e la funzione di energia implicita vengono aggiornati utilizzando segmenti di traiettoria da un buffer di replay.
2. Ottimizzazione della Politica Diffusiva: La rete di punteggio viene addestrata per prevedere la funzione di punteggio utilizzando il campionamento Monte Carlo di traiettorie immaginate pesate in base ai loro ritorni. Durante l'inferenza, la rete di punteggio funge da soluzione ammortizzata per una generazione efficiente di sequenze di azioni.

3. Contributi Chiave

1. Identificazione del Disallineamento Strutturale: Il documento fornisce un'analisi teorica (Problemi 1 e 2) che dimostra come il disaccoppiamento tra ricerca e apprendimento del valore negli approcci esistenti basati su modelli del mondo porti a subottimalità e impedisca un miglioramento monotono della politica.
2. Framework MBDPO: L'introduzione di un framework unificato che formula la ricerca e l'ottimizzazione della politica come un processo diffusivo. Questo approccio sfrutta una funzione di energia implicita per ancorare la politica alla distribuzione comportamentale, mitigando la lacuna nell'iterazione del valore e la deriva distribuzionale.
3. Prestazioni Scalabili: Lo sviluppo di un algoritmo RL diffusivo scalabile che dimostra guadagni di prestazioni consistenti e monotoni all'aumentare della capacità del modello (da 1,7 milioni a 340 milioni di parametri) in contesti diversi, inclusi il pre-addestramento offline multi-task, l'apprendimento online e il fine-tuning da offline a online.
4. Validazione Empirica: Una valutazione estesa su 121 compiti in quattro benchmark (DMControl, MetaWorld, ManiSkill2, MyoSuite) che mostra prestazioni superiori rispetto alle baseline dello stato dell'arte (SAC, DreamerV3, TD-MPC, TD-MPC2).

4. Risultati Sperimentali

• Pre-addestramento Offline: MBDPO esibisce un chiaro comportamento di scalabilità monotona. All'aumentare dei parametri del modello da 1,7 milioni a 340 milioni, il punteggio normalizzato sui dataset da 80 e 30 compiti migliora in modo coerente, superando significativamente TD-MPC2, specialmente a scale di modello più piccole.
• Apprendimento Online: In contesti online da zero, MBDPO ottiene risultati superiori o competitivi su tutti i 121 compiti. Dimostra un miglioramento accelerato della politica con una varianza inferiore nelle fasi iniziali dell'addestramento rispetto alle baseline.
• Mitigazione dei Colli di Bottiglia: L'analisi empirica mostra che MBDPO riduce significativamente gli errori di differenza temporale (TD) incrociata e la deriva delle azioni rispetto a TD-MPC2. La politica diffusiva rimane stabile e converge in modo fluido, mentre i metodi di ricerca euristica mostrano alta volatilità e una grave deriva della politica.
• Generalizzazione: Negli esperimenti di trasferimento da offline a online (O2O), un agente generalista pre-addestrato su 70 compiti è stato fine-tuned con successo su 10 compiti non visti con soli 40k interazioni per compito, superando l'addestramento da zero.
• Efficienza: MBDPO comporta un sovraccarico trascurabile nell'addestramento online, con tempi di addestramento per passo comparabili o più rapidi rispetto a TD-MPC2 nella maggior parte dei domini.

5. Significato e Affermazioni

Il documento afferma che MBDPO stabilisce un nuovo paradigma per l'apprendimento della politica all'interno di modelli del mondo risolvendo il disallineamento strutturale tra ricerca e apprendimento del valore. Le intuizioni chiave includono:

• Scalabilità della Diffusione: I modelli diffusivi agiscono come un'interfaccia intrinseca che collega i modelli del mondo e l'ottimizzazione della politica, consentendo un'iterazione continua del valore e un miglioramento della politica attraverso traiettorie immaginate.
• La Rappresentazione è Cruciale: La qualità della rappresentazione latente appresa è critica. L'ottimizzazione congiunta di MBDPO aiuta a strutturare lo spazio latente in modo che sia allineato con l'intuizione fisica, risultando in traiettorie in ciclo chiuso strutturate per compiti ciclici e percorsi diretti agli obiettivi per la manipolazione, che sono più significativi dal punto di vista fisico rispetto a quelli prodotti da metodi di ricerca euristica.
• Fondamento Principiato: Il lavoro suggerisce che la scalabilità dell'apprendimento della politica non richiede pipeline di ricerca euristica più complesse, ma piuttosto un fondamento principiato che colleghi intrinsecamente i modelli del mondo ai paradigmi fondamentali dell'RL (in particolare tramite obiettivi vincolati da KL).

Gli autori notano le limitazioni, inclusa l'attuale focalizzazione su ambienti simulati e la mancanza di scalabilità verso dimensioni di modelli fondazionali estremamente grandi a causa di vincoli computazionali. Le direzioni future includono il dispiegamento di MBDPO su robot del mondo reale e la sua integrazione con rappresentazioni visive o multimodali pre-addestrate più ricche.