Диффузия для моделирования мира: визуальные детали имеют значение в играх Atari.

Diffusion for World Modeling: Visual Details Matter in Atari

May 20, 2024

Авторы: Eloi Alonso, Adam Jelley, Vincent Micheli, Anssi Kanervisto, Amos Storkey, Tim Pearce, François Fleuret

cs.AI

Аннотация

Мировые модели представляют собой многообещающий подход для обучения агентов обучения с подкреплением способом, который безопасен и эффективен по использованию образцов. Недавние мировые модели в основном работают с последовательностями дискретных скрытых переменных для моделирования динамики окружения. Однако сжатие в компактное дискретное представление может игнорировать визуальные детали, которые важны для обучения с подкреплением. В то же время модели диффузии стали доминирующим подходом для генерации изображений, вызывая вызов у хорошо установленных методов моделирования дискретных скрытых переменных. Под влиянием этого парадигмального сдвига мы представляем DIAMOND (DIffusion As a Model Of eNvironment Dreams) - агента обучения с подкреплением, обученного в модели мира диффузии. Мы анализируем ключевые проектные решения, необходимые для того, чтобы сделать диффузию подходящей для моделирования мира, и демонстрируем, как улучшенные визуальные детали могут привести к улучшению производительности агента. DIAMOND достигает среднего нормализованного человеческого показателя 1,46 на конкурентном бенчмарке Atari 100k; новый лучший результат для агентов, обученных исключительно в модели мира. Для поощрения будущих исследований по диффузии для моделирования мира мы публикуем наш код, агентов и играбельные мировые модели на 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.