Diffusie voor Wereldmodellering: Visuele Details zijn Belangrijk in Atari

Diffusion for World Modeling: Visual Details Matter in Atari

May 20, 2024

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

cs.AI

Samenvatting

Wereldmodellen vormen een veelbelovende aanpak voor het trainen van reinforcement learning-agenten op een veilige en sample-efficiënte manier. Recente wereldmodellen werken voornamelijk met sequenties van discrete latente variabelen om omgevingsdynamiek te modelleren. Deze compressie naar een compacte discrete representatie kan echter visuele details negeren die belangrijk zijn voor reinforcement learning. Tegelijkertijd zijn diffusiemodellen een dominante aanpak geworden voor beeldgeneratie, wat een uitdaging vormt voor gevestigde methoden die discrete latente variabelen modelleren. Gemotiveerd door deze paradigmaverschuiving introduceren we DIAMOND (DIffusion As a Model Of eNvironment Dreams), een reinforcement learning-agent die getraind is in een diffusie-wereldmodel. We analyseren de belangrijkste ontwerpkeuzes die nodig zijn om diffusie geschikt te maken voor wereldmodellering, en demonstreren hoe verbeterde visuele details kunnen leiden tot betere agentprestaties. DIAMOND behaalt een gemiddelde menselijk genormaliseerde score van 1,46 op de competitieve Atari 100k-benchmark; een nieuw record voor agenten die volledig binnen een wereldmodel zijn getraind. Om toekomstig onderzoek naar diffusie voor wereldmodellering te bevorderen, maken we onze code, agenten en speelbare wereldmodellen beschikbaar op 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.