Difusión para Modelado del Mundo: Los Detalles Visuales Importan en Atari

Diffusion for World Modeling: Visual Details Matter in Atari

May 20, 2024

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

cs.AI

Resumen

Los modelos del mundo representan un enfoque prometedor para entrenar agentes de aprendizaje por refuerzo de manera segura y eficiente en términos de muestras. Los modelos del mundo recientes operan principalmente sobre secuencias de variables latentes discretas para modelar la dinámica del entorno. Sin embargo, esta compresión en una representación discreta compacta puede ignorar detalles visuales que son importantes para el aprendizaje por refuerzo. Paralelamente, los modelos de difusión se han convertido en un enfoque dominante para la generación de imágenes, desafiando métodos bien establecidos que modelan latentes discretos. Motivados por este cambio de paradigma, presentamos DIAMOND (DIffusion As a Model Of eNvironment Dreams), un agente de aprendizaje por refuerzo entrenado en un modelo del mundo basado en difusión. Analizamos las decisiones clave de diseño necesarias para hacer que la difusión sea adecuada para el modelado del mundo, y demostramos cómo los detalles visuales mejorados pueden llevar a un mejor rendimiento del agente. DIAMOND logra una puntuación media normalizada humana de 1.46 en el competitivo benchmark Atari 100k; un nuevo récord para agentes entrenados completamente dentro de un modelo del mundo. Para fomentar futuras investigaciones sobre la difusión en el modelado del mundo, publicamos nuestro código, agentes y modelos del mundo jugables en 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.