Diffusion für Weltmodellierung: Visuelle Details sind wichtig bei Atari.

Diffusion for World Modeling: Visual Details Matter in Atari

May 20, 2024

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

cs.AI

papers.abstract

Weltmodelle stellen einen vielversprechenden Ansatz dar, um Trainingsverfahren für Verstärkungslernagenten auf sichere und sample-effiziente Weise zu entwickeln. Aktuelle Weltmodelle arbeiten hauptsächlich mit Sequenzen diskreter latenter Variablen, um Umweltdynamiken zu modellieren. Diese Kompression in eine kompakte diskrete Darstellung kann jedoch visuelle Details vernachlässigen, die für das Verstärkungslernen wichtig sind. Gleichzeitig sind Diffusionsmodelle zu einem dominierenden Ansatz für die Bildgenerierung geworden und stellen etablierte Methoden zur Modellierung diskreter Latenzen in Frage. Angespornt durch diesen Paradigmenwechsel stellen wir DIAMOND (DIffusion As a Model Of eNvironment Dreams) vor, einen Verstärkungslernagenten, der in einem Diffusionsweltmodell trainiert wird. Wir analysieren die wesentlichen Designentscheidungen, die erforderlich sind, um Diffusion für die Weltmodellierung geeignet zu machen, und zeigen auf, wie verbesserte visuelle Details zu einer besseren Leistung des Agenten führen können. DIAMOND erreicht einen durchschnittlichen menschlich normalisierten Score von 1,46 im wettbewerbsfähigen Atari 100k Benchmark; ein neuer Bestwert für Agenten, die ausschließlich innerhalb eines Weltmodells trainiert wurden. Um zukünftige Forschung zur Verwendung von Diffusion für die Weltmodellierung zu fördern, veröffentlichen wir unseren Code, Agenten und spielbaren Weltmodelle unter 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.