Coupled Local and Global World Models for Efficient First Order RL

Sintesi Tecnica: Modelli di Mondo Locali e Globali Accoppiati per un RL di Primo Ordine Efficiente

Definizione del Problema

L'Apprendimento per Rinforzo (RL) ha ottenuto un successo significativo nei compiti di locomozione tramite simulatori paralleli su GPU, tuttavia affronta barriere sostanziali nel controllo percettivo e nella manipolazione. Due problemi primari ostacolano il progresso:

1. Inefficienza dei Campioni: L'apprendimento online pratico direttamente dalle interazioni con robot reali è spesso impraticabile a causa dell'alto costo della raccolta dati.
2. Fedeltà della Simulazione: Simulare compiti di manipolazione è intrinsecamente più complesso rispetto alla locomozione. Lo stato dell'ambiente cambia dinamicamente con le azioni dell'agente (ad es., oggetti deformabili, dinamiche granulari), introducendo una variabilità infinita che gli attuali simulatori fisici faticano a modellare con sufficiente fedeltà per il trasferimento sim-to-real.

Sebbene l'apprendimento di agenti all'interno di modelli di mondo differenziabili tramite metodi di Gradiente di Primo Ordine (FoG) offra una via per risolvere l'inefficienza dei campioni, gli approcci esistenti affrontano un collo di bottiglia di scalabilità. I modelli di mondo su larga scala (ad es., modelli di diffusione) necessari per dinamiche del mondo reale ad alta fedeltà sono computazionalmente costosi da differenziare, rendendo l'errore di backpropagation standard intrattabile per l'addestramento di policy FoG. Al contrario, i modelli leggeri adatti alla differenziazione mancano spesso dell'espressività necessaria per catturare dinamiche del mondo reale complesse.

Metodologia

Il documento propone un framework innovativo che scinde la simulazione forward dalla differenziazione backward per abilitare un RL basato su modello (FoG-MBRL) di primo ordine efficiente utilizzando modelli di mondo su larga scala. L'architettura centrale consiste in due componenti accoppiate:

1. Modello di Mondo Globale (Passaggio Forward):

   • Un modello su larga scala, addestrato offline (sia esso un modello basato sulla diffusione come DIAMOND, sia un modello di diffusione latente basato su transformer come DreamerV4), funge da simulatore "globale".
   • Questo modello opera nello spazio dei pixel (o in uno spazio latente ad alta dimensionalità) per generare traiettorie forward ($s^{\rightarrow}_{t+1}$) ad alta fedeltà che catturano accuratamente complesse dinamiche visive e di contatto.
   • È addestrato su dati di "play" non strutturati e ad alta entropia raccolti da robot reali, eliminando la necessità di simulatori fisici progettati a mano.

2. Modello di Mondo Locale (Passaggio Backward):

   • Un modello leggero Recurrent State Space Model (RSSM) funge da surrogato "locale" in uno spazio latente a bassa dimensionalità.
   • Questo modello approssima le dinamiche locali e fornisce i gradienti necessari per l'ottimizzazione della policy.
   • Fondamentalmente, i gradienti sono calcolati valutando i Jacobiani del modello locale nei punti di stato generati dal modello globale ($s^{\rightarrow}_{t+1}$), piuttosto che nelle predizioni del modello locale stesso. Questo disaccoppiamento permette al sistema di sfruttare l'accuratezza del modello globale per la generazione delle traiettorie, utilizzando il modello locale per una stima dei gradienti computazionalmente trattabile.

3. Protocollo di Addestramento:

   • Raccolta Dati: Dati di "play" non strutturati vengono raccolti su robot reali senza una rigorosa segmentazione del compito.
   • Modellazione della Ricompensa: A seconda del compito, i segnali di ricompensa sono derivati da funzioni di energia contrastiva (addestrate su dati di play/dimostrazione) o da teste di classificazione leggere che identificano comportamenti rilevanti per il compito.
   • Ottimizzazione: La policy è ottimizzata interamente all'interno del modello di mondo utilizzando l'algoritmo Decoupled forward-backward Model-based policy Optimization (DMO). Il modello locale viene continuamente perfezionato durante l'addestramento RL per mantenere l'accuratezza rispetto alla distribuzione delle azioni dell'attuale policy.

Contributi Chiave

• Framework FoG-MBRL Disaccoppiato: Gli autori introducono un metodo che separa la simulazione forward ad alta fedeltà (tramite grandi modelli di diffusione/flow) dalla differenziazione backward efficiente (tramite RSSM leggeri). Ciò risolve il compromesso tra complessità di modellazione e costo di calcolo del gradiente.
• Apprendimento Simulator-Free da Dati Reali: L'approccio addestra agenti interamente partendo da zero utilizzando dati di robot reali, bypassando la necessità di simulatori fisici e affrontando direttamente il gap sim-to-real.
• Deployment Zero-Shot: Il framework dimostra la capacità di addestrare policy nel modello di mondo e di distribuirle su robot reali senza ulteriore fine-tuning o guida di behavior cloning.
• Robustezza a Dinamiche Complesse: Il metodo gestisce con successo compiti che coinvolgono osservabilità parziale, dinamiche di contatto complesse e alti gradi di libertà (ad es., manipolazione di umanoidi).

Risultati Sperimentali

Gli autori hanno valutato il loro approccio su tre compiti robotici reali di complessità crescente:

1. Push-T: Un compito di manipolazione su tavolo che consiste nel spingere un oggetto a forma di T.
2. Ego-Centric Push Cube: Un robot quadrupede che spinge un cubo verso un obiettivo, coinvolgendo dinamiche di camminata e occlusioni parziali.
3. Humanoid Grasp and Lift: Un umanoide ad alto DoF che afferra e solleva una scatola usando una mano destra.

Risultati Chiave:

• Efficienza di Campioni e Temporale: Il metodo DMO proposto ha superato significativamente PPO (Proximal Policy Optimization) sia in termini di efficienza dei campioni che di tempo di calcolo (wall-clock time) in tutti i compiti. PPO ha faticato a convergere o ha richiesto campioni eccessivi, cadendo spesso in soluzioni subottimali locali.
• Necessità di Modelli Globali: Gli studi di ablazione ("No Diffusion"), che sostituivano il modello di diffusione globale con un normale RSSM per i passaggi forward, hanno portato al fallimento totale del compito (0% di successo), confermando che i modelli forward ad alta fedeltà sono essenziali per la manipolazione complessa.
• Superiorità rispetto al Behavior Cloning: Le policy apprese hanno superato i baseline di Behavior Cloning (Action Chunking Transformer), dimostrando la capacità di scoprire strategie a lungo termine e gestire regioni dello spazio degli stati non ben coperte dalle dimostrazioni.
• Successo Zero-Shot: Le policy addestrate interamente nel modello di mondo hanno raggiunto alti tassi di successo (ad es., 9/10 su Push-T, 8/10 su Humanoid Grasp) quando distribuite su robot reali senza ulteriore addestramento nel mondo reale.

Significato e Rivendicazioni

Il documento sostiene che l'apprendimento di agenti interamente all'interno di modelli di mondo basati sui dati è una via percorribile e promettente per risolvere compiti difficili da modellare nello spazio delle immagini. Disaccoppiando i requisiti di accuratezza forward e trattabilità backward, gli autori dimostrano che i modelli di mondo su larga scala e computazionalmente pesanti possono essere utilizzati efficacemente per l'ottimizzazione efficiente della policy.

Il lavoro si pone come una soluzione alla dipendenza dai simulatori fisici progettati a mano, offrendo un metodo scalabile verso ambienti complessi e non strutturati attraverso l'acquisizione di dati reali. Gli autori osservano che, sebbene la loro implementazione attuale richieda l'addestramento di modelli globali da zero e si basi su ricompense etichettate manualmente per certi compiti, il framework fornisce un proxy sicuro per l'addestramento di robot reali, consentendo un'ottimizzazione della policy scalabile senza il rischio di danni all'hardware. Il lavoro futuro è identificato nell'uso di modelli video pre-addestrati su scala internet per ridurre ulteriormente i requisiti di dati e automatizzare l'etichettatura delle ricompense.