世界モデリングにおける拡散：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は、競争力のあるAtari 100kベンチマークで1.46の平均人間正規化スコアを達成しました。これは、ワールドモデル内で完全に訓練されたエージェントにとっての新たな最高記録です。拡散を用いたワールドモデリングの将来の研究を促進するため、我々はコード、エージェント、およびプレイ可能なワールドモデルを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.