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

Diffusion for World Modeling: Visual Details Matter in Atari

May 20, 2024

Autori: Eloi Alonso, Adam Jelley, Vincent Micheli, Anssi Kanervisto, Amos Storkey, Tim Pearce, François Fleuret

cs.AI

Abstract

I modelli del mondo rappresentano un approccio promettente per addestrare agenti di apprendimento per rinforzo in modo sicuro ed efficiente in termini di campioni. I recenti modelli del mondo operano prevalentemente su sequenze di variabili latenti discrete per modellare le dinamiche dell'ambiente. Tuttavia, questa compressione in una rappresentazione discreta compatta potrebbe ignorare dettagli visivi importanti per l'apprendimento per rinforzo. Parallelamente, i modelli di diffusione sono diventati un approccio dominante per la generazione di immagini, sfidando i metodi consolidati che modellano latenti discreti. Motivati da questo cambio di paradigma, introduciamo DIAMOND (DIffusion As a Model Of eNvironment Dreams), un agente di apprendimento per rinforzo addestrato in un modello del mondo basato sulla diffusione. Analizziamo le scelte progettuali chiave necessarie per rendere la diffusione adatta alla modellazione del mondo e dimostriamo come dettagli visivi migliorati possano portare a prestazioni superiori dell'agente. DIAMOND raggiunge un punteggio medio normalizzato umano di 1,46 sul benchmark competitivo Atari 100k; un nuovo record per agenti addestrati interamente all'interno di un modello del mondo. Per promuovere future ricerche sulla diffusione per la modellazione del mondo, rilasciamo il nostro codice, gli agenti e i modelli del mondo giocabili all'indirizzo https://github.com/eloialonso/diamond.

English

World models constitute a promising approach for training reinforcement learning agents in a safe and sample-efficient manner. Recent world models predominantly operate on sequences of discrete latent variables to model environment dynamics. However, this compression into a compact discrete representation may ignore visual details that are important for reinforcement learning. Concurrently, diffusion models have become a dominant approach for image generation, challenging well-established methods modeling discrete latents. Motivated by this paradigm shift, we introduce DIAMOND (DIffusion As a Model Of eNvironment Dreams), a reinforcement learning agent trained in a diffusion world model. We analyze the key design choices that are required to make diffusion suitable for world modeling, and demonstrate how improved visual details can lead to improved agent performance. DIAMOND achieves a mean human normalized score of 1.46 on the competitive Atari 100k benchmark; a new best for agents trained entirely within a world model. To foster future research on diffusion for world modeling, we release our code, agents and playable world models at https://github.com/eloialonso/diamond.