[論文レビュー] Mastering Atari with Discrete World Models

Danijar Hafner, Timothy Lillicrap+2 more|arXiv (Cornell University)|Oct 5, 2020

Reinforcement Learning in Robotics参考文献 74被引用数 23

ひとこと要約

DreamerV2 は、別々に訓練された世界モデル内で完全に離散的ラティス表現を用いて行動を学習するモデルベース強化学習エージェントであり、アタリ55ゲームベンチマークで人間水準のパフォーマンスを達成した。同じ計算リソースとウォールクロック時間内において、Rainbow や IQN といったトップクラスの単一GPUモデルフリー代理を上回り、正確な世界モデルが画素から効率的で高パフォーマンスな意思決定を可能にすることを示した。

ABSTRACT

Intelligent agents need to generalize from past experience to achieve goals in complex environments. World models facilitate such generalization and allow learning behaviors from imagined outcomes to increase sample-efficiency. While learning world models from image inputs has recently become feasible for some tasks, modeling Atari games accurately enough to derive successful behaviors has remained an open challenge for many years. We introduce DreamerV2, a reinforcement learning agent that learns behaviors purely from predictions in the compact latent space of a powerful world model. The world model uses discrete representations and is trained separately from the policy. DreamerV2 constitutes the first agent that achieves human-level performance on the Atari benchmark of 55 tasks by learning behaviors inside a separately trained world model. With the same computational budget and wall-clock time, Dreamer V2 reaches 200M frames and surpasses the final performance of the top single-GPU agents IQN and Rainbow. DreamerV2 is also applicable to tasks with continuous actions, where it learns an accurate world model of a complex humanoid robot and solves stand-up and walking from only pixel inputs.

研究の動機と目的

• 学習済みの世界モデル内での想像的ロールアウトのみを用いて、アタリベンチマークで人間水準のパフォーマンスを達成するモデルベース強化学習エージェントを開発すること。
• ポリシーとは別に訓練された世界モデルが、複雑な視覚環境において高精度な計画と一般化を可能にすることを示すこと。
• 単一GPUと10日間の訓練時間のみを用いて、Rainbow や IQN といった最先端のモデルフリー代理と競合可能なパフォーマンスを達成すること。
• 離散的ラティス空間表現が視覚制御タスクにおける世界モデルの正確性とサンプル効率を向上させることを検証すること。
• 離散的および連続的制御タスクに適用可能な再現可能で効率的かつスケーラブルなモデルベース強化学習フレームワークを確立すること。

提案手法

• 世界モデルは、生のピクセル入力を用いてコンパクトでカテゴリカルなラティス表現を学習するため、離散的変分オートエンコーダ（dVAE）を用いる。
• エージェントは、現在の状態と行動を入力として、将来の状態を予測する離散的ラティス空間における遷移モデルを学習する。
• 世界モデルが生成する想像的トラジェクトリを用いて、別個のアクター・クリティックネットワークを訓練する。価値関数とKLバランス損失を最適化するポリシーネットワークを用いる。
• ポリシーは、世界モデルが計画のための環境を提供する中で、内発的興味（curiosity）と外的報酬予測の組み合わせにより訓練される。
• ポリシー訓練の前に、世界モデルは大量の環境相互作用データセットで事前訓練されるため、サンプル効率の高いポリシー学習が可能になる。
• 訓練の安定化と離散的ラティス空間内での表現品質の向上を図るため、新しいKLバランス機構を用いる。

実験結果

リサーチクエスチョン

• RQ1 ポリシーとは別に訓練された世界モデルは、アタリのような複雑な視覚環境において、サンプル効率的で高パフォーマンスな意思決定を可能にするか？
• RQ2 連続的ラティス空間と比較して、離散的ラティス表現を用いることで、世界モデルの正確性と一般化性能が向上するか？
• RQ3 ポリシー訓練中に直接的な環境相互作用を一切行わず、世界モデル内の想像的ロールアウトのみを用いて、アタリ55ゲームベンチマークで人間水準のパフォーマンスを達成できるか？
• RQ4 同じ計算制約下で、モデルベースエージェントのパフォーマンスは、Rainbow や IQN といった最先端のモデルフリー代理と比較してどうなるか？
• RQ5 世界モデルが画像と報酬を予測する能力が、ポリシーの一般化とサンプル効率にどの程度寄与するか？

主な発見

• DreamerV2 は、世界モデル内の想像的ロールアウトのみを用いて、アタリ55ゲームベンチマークで人間水準のパフォーマンスを達成し、そのようなモデルベースエージェントとして初めての達成となった。
• 同じ計算リソースとウォールクロック時間内において、DreamerV2 は2億フレーム評価でのトップクラスの単一GPUモデルフリー代理（Rainbow および IQN）の最終パフォーマンスを上回った。
• 離散的ラティス表現の使用は、世界モデルの正確性を顕著に向上させ、ラティス空間内での安定的かつ高品質な計画を可能にした。
• 画像再構成と報酬予測信号を用いて訓練された世界モデルは、多様な制御タスクに有用な汎用的表現を学習できることを示した。
• DreamerV2 は連続的制御タスクに対しても効果的に一般化でき、ピクセル入力と単一GPUのみを用いて、ヒューマノイドロボットの複雑な歩行行動（立ち上がりや歩行）を解消した。
• 離散的ラティスとKLバランスのパフォーマンス向上効果は、アブレーションスタディにより実証され、最終的なエージェントパフォーマンスに与える影響が確認された。