Difusão para Modelagem de Mundo: Detalhes Visuais Importam no 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

Resumo

Modelos de mundo representam uma abordagem promissora para o treinamento de agentes de aprendizado por reforço de forma segura e eficiente em termos de amostras. Recentemente, os modelos de mundo operam predominantemente em sequências de variáveis latentes discretas para modelar a dinâmica do ambiente. No entanto, essa compressão em uma representação discreta compacta pode ignorar detalhes visuais que são importantes para o aprendizado por reforço. Paralelamente, os modelos de difusão tornaram-se uma abordagem dominante para a geração de imagens, desafiando métodos bem estabelecidos que modelam variáveis latentes discretas. Motivados por essa mudança de paradigma, introduzimos o DIAMOND (DIffusion As a Model Of eNvironment Dreams), um agente de aprendizado por reforço treinado em um modelo de mundo baseado em difusão. Analisamos as principais escolhas de design necessárias para tornar a difusão adequada para a modelagem de mundo e demonstramos como detalhes visuais aprimorados podem levar a um melhor desempenho do agente. O DIAMOND alcança uma pontuação média normalizada por humanos de 1,46 no benchmark competitivo Atari 100k; um novo recorde para agentes treinados inteiramente dentro de um modelo de mundo. Para fomentar pesquisas futuras sobre difusão para modelagem de mundo, disponibilizamos nosso código, agentes e modelos de mundo jogáveis em 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.