Nel 2018, l’idea dell’apprendimento per rinforzo in un modello di mondo neurale è stata introdotta per la prima volta, e presto questo principio fondamentale è stato applicato ai modelli di mondo. Alcuni dei modelli più importanti che implementano l’apprendimento per rinforzo sono stati il framework Dreamer, che ha introdotto l’apprendimento per rinforzo dallo spazio latente di un modello di stato ricorrente. Il DreamerV2 ha dimostrato che l’uso di latenti discreti potrebbe risultare in errori di composizione ridotti, e il framework DreamerV3 è stato in grado di raggiungere prestazioni umane su una serie di attività in diversi domini con iperparametri fissi.

Inoltre, si possono tracciare paralleli tra i modelli di generazione di immagini e i modelli di mondo, indicando che i progressi fatti nei modelli di visione generativa potrebbero essere replicati per beneficare i modelli di mondo. Dalla popolarità dell’uso dei transformer nei framework di elaborazione del linguaggio naturale, sono emersi i framework DALL-E e VQGAN. I framework hanno implementato autoencoder discreti per convertire le immagini in token discreti e sono stati in grado di costruire modelli generativi di testo-immagine molto potenti ed efficienti sfruttando le capacità di modellazione sequenziale dei transformer autoregressivi. Allo stesso tempo, i modelli di diffusione hanno guadagnato popolarità e oggi i modelli di diffusione si sono affermati come un paradigma dominante per la generazione di immagini ad alta risoluzione. Grazie alle capacità offerte dai modelli di diffusione e dall’apprendimento per rinforzo, si stanno facendo tentativi di combinare i due approcci, con l’obiettivo di sfruttare la flessibilità dei modelli di diffusione come modelli di traiettoria, modelli di ricompensa, pianificatori e politiche per l’augmentazione dei dati nell’apprendimento per rinforzo offline.

I modelli di mondo offrono un metodo promettente per addestrare gli agenti di apprendimento per rinforzo in modo sicuro ed efficiente. Tradizionalmente, questi modelli utilizzano sequenze di variabili latenti discrete per simulare la dinamica dell’ambiente. Tuttavia, questa compressione può trascurare i dettagli visivi cruciali per l’apprendimento per rinforzo. Allo stesso tempo, i modelli di diffusione sono diventati popolari per la generazione di immagini, sfidando i metodi tradizionali che utilizzano latenti discreti. Ispirati da questo cambiamento, in questo articolo, parleremo di DIAMOND (DIffusion As a Model Of eNvironment Dreams), un agente di apprendimento per rinforzo addestrato all’interno di un modello di mondo di diffusione. Esploreremo le scelte di progettazione necessarie per rendere la diffusione adatta alla modellazione del mondo e mostreremo che i dettagli visivi migliorati portano a una migliore prestazione dell’agente. DIAMOND stabilisce un nuovo benchmark sul test competitivo Atari 100k, raggiungendo un punteggio medio normalizzato umano di 1,46, il più alto per gli agenti addestrati interamente all’interno di un modello di mondo.

DIAMOND: I Dettagli Visivi Contano in Atari e Diffusione per la Modellazione del Mondo

I modelli di mondo o modelli generativi di ambienti stanno emergendo come uno dei componenti più importanti per gli agenti generativi per pianificare e ragionare sul loro ambiente. Sebbene l’uso dell’apprendimento per rinforzo abbia ottenuto notevoli successi negli ultimi anni, i modelli che implementano l’apprendimento per rinforzo sono noti per essere inefficienti in termini di campioni, il che limita notevolmente le loro applicazioni nel mondo reale. D’altra parte, i modelli di mondo hanno dimostrato la loro capacità di addestrare gli agenti di apprendimento per rinforzo in modo efficiente in diversi ambienti con una notevole efficienza dei campioni, consentendo al modello di imparare dalle esperienze del mondo reale. I framework di modellazione del mondo recenti modellano generalmente la dinamica dell’ambiente come una sequenza di variabili latenti discrete, con il modello che discretizza lo spazio latente per evitare errori di composizione su orizzonti temporali multipli. Sebbene questo approccio possa dare risultati sostanziali, è anche associato a una perdita di informazioni, portando a una perdita di qualità di ricostruzione e di generalità. La perdita di informazioni potrebbe diventare un ostacolo significativo per scenari del mondo reale che richiedono informazioni ben definite, come l’addestramento di veicoli autonomi. In tali attività, piccoli cambiamenti o dettagli nell’input visivo, come il colore del semaforo o l’indicatore di direzione del veicolo davanti, possono cambiare la politica di un agente. Sebbene aumentare il numero di latenti discreti possa aiutare a evitare la perdita di informazioni, aumenta notevolmente i costi di calcolo.

Inoltre, negli ultimi anni, i modelli di diffusione sono emersi come l’approccio dominante per la generazione di immagini di alta qualità, poiché i framework costruiti sui modelli di diffusione imparano a invertire un processo di rumorosità e competono direttamente con alcuni degli approcci più stabiliti che modellano token discreti, offrendo così un’alternativa promettente per eliminare la necessità di discretizzazione nella modellazione del mondo. I modelli di diffusione sono noti per la loro capacità di essere facilmente condizionati e di modellare in modo flessibile distribuzioni complesse e multimodali senza collasso di modalità. Questi attributi sono cruciali per la modellazione del mondo, poiché la condizionamento consente a un modello di mondo di riflettere accuratamente le azioni di un agente, portando a una più affidabile assegnazione di credito. Inoltre, la modellazione di distribuzioni multimodali offre una maggiore diversità di scenari di addestramento per l’agente, migliorando così le sue prestazioni generali.

Basandosi su queste caratteristiche, DIAMOND, (DIffusion As a Model Of eNvironment Dreams), un agente di apprendimento per rinforzo addestrato all’interno di un modello di mondo di diffusione. Il framework DIAMOND fa scelte di progettazione accurate per garantire che il suo modello di mondo di diffusione rimanga efficiente e stabile su orizzonti temporali lunghi. Il framework fornisce un’analisi qualitativa per dimostrare l’importanza di queste scelte di progettazione. DIAMOND stabilisce un nuovo stato dell’arte con un punteggio medio normalizzato umano di 1,46 sul benchmark Atari 100k, il più alto per gli agenti addestrati interamente all’interno di un modello di mondo. Operare nello spazio delle immagini consente al modello di mondo di diffusione di DIAMOND di sostituire in modo trasparente l’ambiente, offrendo una maggiore comprensione del comportamento del modello di mondo e dell’agente. Notabilmente, le prestazioni migliorate in alcuni giochi sono attribuibili a una migliore modellazione dei dettagli visivi critici. Il framework DIAMOND modella l’ambiente come un POMDP standard o un processo di decisione markoviano parzialmente osservabile con un insieme di stati, un insieme di azioni discrete e un insieme di osservazioni di immagini. Le funzioni di transizione descrivono la dinamica dell’ambiente, e la funzione di ricompensa mappa le transizioni in ricompense scalari.

DIAMOND: Metodologia e Architettura

In sostanza, i modelli di diffusione sono una classe di modelli generativi che generano un campione invertendo il processo di rumorosità e traggono ispirazione pesante dalla termodinamica non in equilibrio. Il framework DIAMOND considera un processo di diffusione indicizzato da una variabile temporale continua con marginali e condizioni al contorno con una distribuzione priori non strutturata tracciabile. Inoltre, per ottenere un modello generativo, che mappa dal rumore ai dati, il framework DIAMOND deve invertire il processo, con il processo di inversione che è anch’esso un processo di diffusione, che si esegue all’indietro nel tempo. Inoltre, in qualsiasi punto nel tempo, non è banale stimare la funzione di punteggio poiché il framework DIAMOND non ha accesso alla vera funzione di punteggio, e il modello supera questo ostacolo implementando un obiettivo di corrispondenza dei punteggi, un approccio che consente a un framework di addestrare un modello di punteggio senza conoscere la funzione di punteggio sottostante. Il modello di diffusione basato sui punteggi fornisce un modello generativo uncondizionato. Tuttavia, un modello generativo condizionato della dinamica dell’ambiente è necessario per servire come modello di mondo, e per servire a questo scopo, il framework DIAMOND guarda al caso generale dell’approccio POMDP, in cui il framework può utilizzare le osservazioni e le azioni passate per approssimare lo stato markoviano sconosciuto. Come dimostrato nella Figura 1, il framework DIAMOND utilizza questa storia per condizionare un modello di diffusione, per stimare e generare la prossima osservazione direttamente. Sebbene il framework DIAMOND possa in teoria ricorrere a qualsiasi risolutore SDE o ODE, c’è un compromesso tra NFE o numero di valutazioni della funzione e qualità del campione che impatta notevolmente il costo di inferenza dei modelli di diffusione.

Basandosi su queste conoscenze, esaminiamo ora la realizzazione pratica del framework DIAMOND di un modello di mondo basato sulla diffusione, compresi i coefficienti di deriva e diffusione corrispondenti a una scelta specifica di approccio di diffusione. Invece di optare per DDPM, un candidato naturalmente adatto al compito, il framework DIAMOND si basa sulla formulazione EDM e considera un kernel di perturbazione con una funzione a valore reale del tempo di diffusione chiamata programma di rumore. Il framework seleziona i precondizionatori per mantenere la varianza di input e output per qualsiasi livello di voce. L’addestramento della rete combina segnale e rumore in modo adattivo a seconda del livello di degrado, e quando il rumore è basso, l’obiettivo diventa la differenza tra il segnale pulito e il segnale perturbato, ovvero il rumore gaussiano aggiunto. Intuitivamente, ciò impedisce all’obiettivo di addestramento di diventare banale nel regime a basso rumore. In pratica, questo obiettivo è ad alta varianza agli estremi del programma di rumore, quindi il modello campiona il livello di rumore da una distribuzione log-normale scelta empiricamente per concatenare l’addestramento intorno alle regioni di rumore medio. Il framework DIAMOND utilizza un componente U-Net 2D standard per il campo vettoriale e mantiene un buffer di osservazioni e azioni passate che il framework utilizza per condizionarsi. Il framework DIAMOND concatena quindi queste osservazioni passate alla prossima osservazione rumorosa e alle azioni di input attraverso layer di normalizzazione di gruppo adattivo nei blocchi residuali dell’U-Net.

DIAMOND: Esperimenti e Risultati

Per una valutazione completa, il framework DIAMOND opta per il benchmark Atari 100k. Il benchmark Atari 100k consiste in 26 giochi progettati per testare una vasta gamma di capacità degli agenti. In ogni gioco, un agente è limitato a 100k azioni nell’ambiente, il che equivale approssimativamente a 2 ore di gioco umano, per imparare il gioco prima della valutazione. Per confronto, gli agenti Atari non vincolati di solito si addestrano per 50 milioni di passi, rappresentando un aumento di 500 volte nell’esperienza. Abbiamo addestrato DIAMOND da zero utilizzando 5 semi casuali per ogni gioco. Ogni esecuzione di addestramento ha richiesto circa 12 GB di VRAM e ha impiegato circa 2,9 giorni su una sola scheda grafica Nvidia RTX 4090, per un totale di 1,03 anni di GPU. La seguente tabella fornisce il punteggio per tutti i giochi, la media e la media interquartile dei punteggi normalizzati umani.

Seguendo i limiti delle stime puntuali, il framework DIAMOND fornisce una fiducia di bootstrap stratificata nella media e nella media interquartile dei punteggi normalizzati umani, insieme a profili di prestazione e metriche aggiuntive, come riassunto nella seguente figura.

I risultati mostrano che DIAMOND si esegue in modo eccezionale in tutto il benchmark, superando i giocatori umani in 11 giochi e raggiungendo un punteggio medio normalizzato umano di 1,46, stabilendo un nuovo record per gli agenti addestrati interamente all’interno di un modello di mondo. Inoltre, la media interquartile di DIAMOND è paragonabile a STORM e supera tutti gli altri punti di riferimento. DIAMOND eccelle in ambienti in cui la cattura di piccoli dettagli è cruciale, come Asterix, Breakout e RoadRunner. Inoltre, come discusso in precedenza, il framework DIAMOND ha la flessibilità di implementare qualsiasi modello di diffusione nella sua pipeline, sebbene opti per l’approccio EDM, sarebbe stato un choix naturale optare per il modello DDPM poiché è già implementato in numerose applicazioni generative di immagini. Per confrontare l’approccio EDM con l’implementazione DDPM, il framework DIAMOND addestra entrambe le varianti con la stessa architettura di rete sullo stesso set di dati statico condiviso con oltre 100k frame raccolti con una politica di esperto. Il numero di passaggi di denoising è direttamente correlato al costo di inferenza del modello di mondo, quindi meno passaggi ridurranno il costo di addestramento di un agente su traiettorie immaginate. Per garantire che il nostro modello di mondo rimanga computazionalmente paragonabile con altri punti di riferimento, come IRIS che richiede 16 NFE per timestep, ci proponiamo di utilizzare non più di decine di passaggi di denoising, preferibilmente meno. Tuttavia, impostare il numero di passaggi di denoising troppo basso può degradare la qualità visiva, portando a errori di composizione. Per valutare la stabilità delle diverse varianti di diffusione, visualizziamo le traiettorie immaginate in modo autoregressivo fino a t = 1000 passi temporali nella seguente figura, utilizzando diversi numeri di passaggi di denoising n ≤ 10.

Osserviamo che utilizzare DDPM (a) in questo regime comporta errori di composizione gravi, facendo sì che il modello di mondo si allontani rapidamente dalla distribuzione. Al contrario, il modello di mondo di diffusione basato su EDM (b) rimane molto più stabile su orizzonti temporali lunghi, anche con un solo passaggio di denoising. Le traiettorie immaginate con modelli di mondo di diffusione basati su DDPM (a sinistra) e EDM (a destra) sono mostrate. L’osservazione iniziale a t = 0 è la stessa per entrambi, e ogni riga corrisponde a un numero decrescente di passaggi di denoising n. Osserviamo che la generazione basata su DDPM soffre di errori di composizione, con numeri minori di passaggi di denoising che portano a una rapida accumulazione di errori. Al contrario, il modello di mondo di DIAMOND basato su EDM rimane molto più stabile, anche per n = 1. La previsione ottimale a un solo passo è l’aspettativa su possibili ricostruzioni per un dato input rumoroso, che può essere fuori dalla distribuzione se la distribuzione posteriore è multimodale. Sebbene alcuni giochi, come Breakout, abbiano transizioni deterministiche che possono essere modellate con precisione con un solo passaggio di denoising, altri giochi esibiscono parziale osservabilità, risultando in distribuzioni di osservazione multimodali. In questi casi, un solutore iterativo è necessario per guidare la procedura di campionamento verso una modalità specifica, come illustrato nel gioco Boxing nella seguente figura. Di conseguenza, il framework DIAMOND ha impostato n = 3 in tutti i nostri esperimenti.

La figura precedente confronta il campionamento a un solo passo (riga superiore) e il campionamento multiplo (riga inferiore) in Boxing. I movimenti del giocatore nero sono imprevedibili, facendo sì che il denoising a un solo passo interpoli tra risultati possibili, portando a previsioni sfocate. Al contrario, il campionamento multiplo produce un’immagine chiara guidando la generazione verso una modalità specifica. È interessante notare che, poiché la politica controlla il giocatore bianco, le sue azioni sono note al modello di mondo, eliminando l’ambiguità. Pertanto, sia il campionamento a un solo passo che il campionamento multiplo prevedono correttamente la posizione del giocatore bianco.

Nella figura precedente, le traiettorie immaginate da DIAMOND esibiscono generalmente una qualità visiva più alta e sono più fedeli all’ambiente vero rispetto a quelle immaginate da IRIS. Le traiettorie generate da IRIS contengono incoerenze visive tra i frame (evidenziate da caselle bianche), come nemici visualizzati come ricompense e viceversa. Sebbene queste incoerenze possano influenzare solo pochi pixel, possono avere un impatto significativo sull’apprendimento per rinforzo. Ad esempio, un agente di solito si propone di mirare alle ricompense e di evitare i nemici, quindi queste piccole discordanze visive possono rendere più difficile l’apprendimento di una politica ottimale. La figura mostra frame consecutivi immaginati con IRIS (a sinistra) e DIAMOND (a destra). Le caselle bianche evidenziano le incoerenze tra i frame, che appaiono solo nelle traiettorie generate con IRIS. In Asterix (riga superiore), un nemico (arancione) diventa una ricompensa (rossa) nel secondo frame, poi torna a essere un nemico nel terzo e di nuovo una ricompensa nel quarto. In Breakout (riga centrale), i mattoni e il punteggio sono incoerenti tra i frame. In Road Runner (riga inferiore), le ricompense (piccoli punti blu sulla strada) sono rappresentate in modo incoerente tra i frame. Queste incoerenze non si verificano con DIAMOND. In Breakout, il punteggio viene aggiornato in modo affidabile di +7 quando un mattone rosso viene rotto.

Conclusione

In questo articolo, abbiamo parlato di DIAMOND, un agente di apprendimento per rinforzo addestrato all’interno di un modello di mondo di diffusione. Il framework DIAMOND fa scelte di progettazione accurate per garantire che il suo modello di mondo di diffusione rimanga efficiente e stabile su orizzonti temporali lunghi. Il framework fornisce un’analisi qualitativa per dimostrare l’importanza di queste scelte di progettazione. DIAMOND stabilisce un nuovo stato dell’arte con un punteggio medio normalizzato umano di 1,46 sul benchmark Atari 100k, il più alto per gli agenti addestrati interamente all’interno di un modello di mondo. Operare nello spazio delle immagini consente al modello di mondo di diffusione di DIAMOND di sostituire in modo trasparente l’ambiente, offrendo una maggiore comprensione del comportamento del modello di mondo e dell’agente. Notabilmente, le prestazioni migliorate in alcuni giochi sono attribuibili a una migliore modellazione dei dettagli visivi critici. Il framework DIAMOND modella l’ambiente come un POMDP standard o un processo di decisione markoviano parzialmente osservabile con un insieme di stati, un insieme di azioni discrete e un insieme di osservazioni di immagini. Le funzioni di transizione descrivono la dinamica dell’ambiente, e la funzione di ricompensa mappa le transizioni in ricompense scalari. Il framework DIAMOND modella l’ambiente come un modello di mondo standard POMDP o un processo di decisione markoviano parzialmente osservabile con un insieme di stati, un insieme di azioni discrete e un insieme di osservazioni di immagini. Le funzioni di transizione descrivono la dinamica dell’ambiente, e la funzione di ricompensa mappa le transizioni in ricompense scalari.