Diffusion pour la modélisation du monde : Les détails visuels comptent dans 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

Les modèles du monde constituent une approche prometteuse pour entraîner des agents d'apprentissage par renforcement de manière sûre et efficace en termes d'échantillons. Les modèles du monde récents opèrent principalement sur des séquences de variables latentes discrètes pour modéliser la dynamique de l'environnement. Cependant, cette compression en une représentation discrète compacte peut ignorer des détails visuels importants pour l'apprentissage par renforcement. Parallèlement, les modèles de diffusion sont devenus une approche dominante pour la génération d'images, remettant en question les méthodes bien établies qui modélisent des latents discrets. Motivés par ce changement de paradigme, nous introduisons DIAMOND (DIffusion As a Model Of eNvironment Dreams), un agent d'apprentissage par renforcement entraîné dans un modèle du monde basé sur la diffusion. Nous analysons les choix de conception clés nécessaires pour rendre la diffusion adaptée à la modélisation du monde, et démontrons comment des détails visuels améliorés peuvent conduire à une meilleure performance de l'agent. DIAMOND atteint un score normalisé humain moyen de 1,46 sur le benchmark compétitif Atari 100k ; un nouveau record pour les agents entraînés entièrement dans un modèle du monde. Pour favoriser les recherches futures sur la diffusion pour la modélisation du monde, nous publions notre code, nos agents et nos modèles du monde jouables sur 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.