Simulus: Combining Improvements in Sample-Efficient World Model Agents

Riepilogo Tecnico: Simulus – Combinazione di Miglioramenti in Agenti con Modello del Mondo Efficiente dal Punto di Vista del Campionamento

Enunciato del Problema
I modelli del mondo (WM) rappresentano una frontiera nell'apprendimento per rinforzo (RL) efficiente dal punto di vista del campionamento, ma il loro potenziale è spesso irrealizzato a causa della complessità nell'integrare miglioramenti promettenti. Sebbene esistano singoli avanzamenti, combinarli richiede competenze significative e sforzi ingegneristici. Inoltre, i modelli del mondo basati su token esistenti (TBWM) sono in gran parte limitati a osservazioni di immagini e azioni discrete, riducendone l'applicabilità al controllo continuo e agli ambienti multi-modali. Inoltre, componenti come la motivazione intrinseca (per la riduzione dell'incertezza epistemica) e il replay prioritizzato sono frequentemente trascurati negli agenti efficienti dal punto di vista del campionamento, in parte a causa dell'alto costo di integrazione e del rischio che l'esplorazione intrinca possa sprecare budget di interazione scarsi su esperienze irrilevanti per il compito. Questo scenario rispecchia l'era pre-Rainbow nel deep RL, dove miglioramenti validati singolarmente non erano stati combinati sistematicamente.

Metodologia
Gli autori introducono Simulus, un agente modulare basato su token con modello del mondo progettato per integrare e validare una serie di miglioramenti complementari. L'agente è costruito sull'architettura REM e comprende quattro moduli distinti: un modulo di rappresentazione ($V$), un modello del mondo ($M$), un controllore ($C$) e un buffer di replay.

1. Framework di Tokenizzazione Flessibile ($V$): Simulus impiega un sistema di tokenizzazione modulare che supporta combinazioni arbitrarie di modalità di osservazione e azione. Gestisce nativamente immagini, vettori continui, variabili categoriche e "categorici 2D" (griglie simili a immagini di variabili categoriche). Gli encoder producono sequenze di token a lunghezza fissa, che vengono concatenate in una rappresentazione unificata. Ciò consente all'agente di operare in ambienti diversificati, inclusi il controllo continuo e i compiti multi-modali simbolici.
2. Modello del Mondo ($M$): Il modello di dinamiche utilizza una Rete Retentiva (RetNet) potenziata da un meccanismo di Predizione Parallela delle Osservazioni (POP) per modellare le dinamiche dell'ambiente come modellazione di sequenza su token discreti.
   • Motivazione Intrinseca: Per ridurre l'incertezza epistemica, Simulus impiega un ensemble di testate di previsione della prossima osservazione. Il disaccordo tra queste testate, misurato tramite la Divergenza di Jensen-Shannon (JSD), genera ricompense intrinseche ($r_{int}$). Queste vengono combinate con ricompense estrinseche per guidare il controllore verso regioni ad alta incertezza durante l'"immaginazione" (pianificazione all'interno del modello).
   • Replay Prioritario: Il buffer di replay mantiene i valori di perdita del modello del mondo per gli esempi archiviati. L'addestramento utilizza una miscela di campionamento uniforme e prioritario (controllato da un singolo parametro $\alpha$), dove i campioni vengono selezionati in proporzione alle loro perdite trasformate tramite softmax.
3. Controllore ($C$): Il controllore è un'architettura estesa attore-critico che opera interamente su esperienze simulate generate dal modello del mondo.
   • Regressione come Classificazione (RaC): Per gestire in modo robusto segnali di ricompensa e ritorno illimitati, Simulus adotta una rete di classificazione con bin esponenziali per la previsione del valore. Utilizza il metodo di perdita HL-Gauss (basato sulla cross-entropia) invece del classico errore quadratico medio, migliorando la stabilità nella stima del valore.
4. Ciclo di Addestramento: L'agente segue un ciclo di raccolta dati, apprendimento della rappresentazione, apprendimento del modello del mondo e apprendimento del controllo nell'immaginazione. Il controllore viene addestrato utilizzando $\lambda$-ritorni e un obiettivo REINFORCE con una baseline di valore.

Contributi Chiave

• Integrazione di Componenti Trascurati: Il documento identifica e integra la motivazione intrinseca, il replay prioritizzato del modello del mondo e un approccio ibrido RaC. Dimostra che questi componenti non sono solo efficaci singolarmente, ma producono guadagni sinergici quando combinati.
• Agenti Versatili Basati su Token: Adattando la tokenizzazione multi-modale al regime efficiente dal punto di vista del campionamento, Simulus stabilisce i TBWM come agenti versatili in grado di gestire input visivi, continui e multi-modali, superando i limiti dei modelli con azioni discrete basati solo su immagini.
• Validazione dell'Esplorazione Indotta dal Modello: Il lavoro dimostra che l'esplorazione indotta dal modello del mondo migliora significativamente le prestazioni anche sotto budget di interazione ristretti, sfidando l'idea che la motivazione intrinseca sia troppo costosa per gli ambienti efficienti dal punto di vista del campionamento.
• Fondamento Open-Source: Gli autori forniscono un solido punto di partenza per la ricerca sui TBWM rendendo open-source codice e pesi del modello, abbassando la barriera all'ingresso per lo sviluppo di agenti con modello del mondo competitivi.

Risultati
Simulus è stato valutato su tre benchmark diversificati con budget di interazione rigorosi:

• Visual Atari 100K: Simulus ha raggiunto un'efficienza di campionamento allo stato dell'arte (SOTA) tra i metodi privi di pianificazione, superando baseline come DreamerV3, TWM, STORM e DIAMOND. In particolare, è il primo modello del mondo privo di pianificazione a raggiungere punteggi medi interquantilici (IQM) e mediani a livello umano, ottenendo prestazioni sovrumane su 13 dei 26 giochi.
• Controllo Continuo (DMC Proprioception 500K): Nei compiti di controllo continuo, Simulus ha costantemente eguagliato o leggermente superato DreamerV3, dimostrando l'efficacia delle architetture basate su token nei domini continui con vocabolari compatti.
• Simbolico Multi-Modale (Craftax-1M): In un ambiente di sopravvivenza 2D parzialmente osservabile che combina griglie simboliche e caratteristiche continue, Simulus ha superato i metodi concorrenti basati su modello del mondo e tutte le baseline prive di modello (inclusi RND, PPO-RNN ed E3B), gestendo in modo efficiente sequenze di token lunghe.

Studi di ablazione hanno confermato che ogni componente contribuisce individualmente, con le ricompense intrinseche che si rivelano particolarmente cruciali. La combinazione di tutti i componenti ha prodotto le prestazioni più forti, indicando una significativa sinergia.

Significato
Il documento sostiene che il potenziale irrealizzato degli agenti con modello del mondo deriva dalla difficoltà di combinare miglioramenti modulari. Traendo ispirazione dal framework Rainbow in DQN, Simulus dimostra che una combinazione sistematica di flessibilità di tokenizzazione, motivazione intrinseca, replay prioritario e tecniche avanzate di previsione del valore porta a un'efficienza di campionamento allo stato dell'arte. Il lavoro stabilisce che i modelli del mondo basati su token sono un approccio versatile e scalabile per diversi domini RL, a condizione che i componenti siano integrati con cura per evitare le interferenze spesso osservate nelle architetture monolitiche. Gli autori sostengono che, sebbene gli agenti basati su token affrontino attualmente limitazioni nella velocità di addestramento rispetto ad alcune baseline, il loro design modulare offre vantaggi distinti per l'inferenza della politica e la scalabilità della ricerca futura.