Diffusion for World Modeling: Visual Details Matter in Atari

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Alonso, Eloi, Jelley, Adam, Micheli, Vincent, Kanervisto, Anssi, Storkey, Amos, Pearce, Tim, Fleuret, François
Format: Preprint
Published: 2024
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866929568111132672
author Alonso, Eloi
Jelley, Adam
Micheli, Vincent
Kanervisto, Anssi
Storkey, Amos
Pearce, Tim
Fleuret, François
author_facet Alonso, Eloi
Jelley, Adam
Micheli, Vincent
Kanervisto, Anssi
Storkey, Amos
Pearce, Tim
Fleuret, François
contents 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. We further demonstrate that DIAMOND's diffusion world model can stand alone as an interactive neural game engine by training on static Counter-Strike: Global Offensive gameplay. To foster future research on diffusion for world modeling, we release our code, agents, videos and playable world models at https://diamond-wm.github.io.
format Preprint
id arxiv_https___arxiv_org_abs_2405_12399
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Diffusion for World Modeling: Visual Details Matter in Atari
Alonso, Eloi
Jelley, Adam
Micheli, Vincent
Kanervisto, Anssi
Storkey, Amos
Pearce, Tim
Fleuret, François
Machine Learning
Artificial Intelligence
Computer Vision and Pattern Recognition
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. We further demonstrate that DIAMOND's diffusion world model can stand alone as an interactive neural game engine by training on static Counter-Strike: Global Offensive gameplay. To foster future research on diffusion for world modeling, we release our code, agents, videos and playable world models at https://diamond-wm.github.io.
title Diffusion for World Modeling: Visual Details Matter in Atari
topic Machine Learning
Artificial Intelligence
Computer Vision and Pattern Recognition
url https://arxiv.org/abs/2405.12399