[論文レビュー] TransDreamer: Reinforcement Learning with Transformer World Models

Chang Chen, Yifu Wu+2 more|arXiv (Cornell University)|Feb 19, 2022

Reinforcement Learning in Robotics被引用数 8

ひとこと要約

TransDreamerは、再帰的ワールドモデルを置き換える新しいトランスフォーマー状態空間モデル（TSSM）を備えた完全なトランスフォーマー基盤のモデルベース強化学習エージェントを提案する。これにより、長期間記憶と計画が向上する。複雑な2Dおよび3Dの視覚的制御タスク（長期間の推論を要するもの）においてDreamerを上回る性能を発揮するが、短期的依存性を持つ単純なタスクではDreamerと同等の性能を示す。

ABSTRACT

The Dreamer agent provides various benefits of Model-Based Reinforcement Learning (MBRL) such as sample efficiency, reusable knowledge, and safe planning. However, its world model and policy networks inherit the limitations of recurrent neural networks and thus an important question is how an MBRL framework can benefit from the recent advances of transformers and what the challenges are in doing so. In this paper, we propose a transformer-based MBRL agent, called TransDreamer. We first introduce the Transformer State-Space Model, a world model that leverages a transformer for dynamics predictions. We then share this world model with a transformer-based policy network and obtain stability in training a transformer-based RL agent. In experiments, we apply the proposed model to 2D visual RL and 3D first-person visual RL tasks both requiring long-range memory access for memory-based reasoning. We show that the proposed model outperforms Dreamer in these complex tasks.

研究の動機と目的

• トランスフォーマー基盤のワールドモデルが、モデルベース強化学習におけるサンプル効率と長期間推論を改善できるかどうかを調査すること。
• モデルベース強化学習において、純トランスフォーマー政策の安定した学習を実現する課題に取り組むこと。先行研究では、疎な報酬から純トランスフォーマー政策を学習することが困難であると示されている。
• トランスフォーマーと互換性があり、並列学習と効果的な想像を可能にする、確率的でアクションに依存するワールドモデルを構築すること。
• 向上したワールドモデルが、複雑な長期間記憶の相互作用を要するタスクでより良い性能を発揮するかどうかを評価すること。
• 提案フレームワークが、RNNベースのモデルが期待される成功を示す単純なタスクでも強力な性能を維持していることを検証すること。

提案手法

• 自己注意機構を用いて潜在空間における長期間依存関係を捉える、確率的でアクションに依存するワールドモデル「トランスフォーマー状態空間モデル（TSSM）」を提案する。
• TSSMを、確率的ダイナミクス予測のための事前分布と事後分布の両方をサポートするように適応させ、想像における正確なロールアウトを可能にする。
• ポリシーとワールドモデルの両方でTSSMワールドモデルを共有することで、一貫した潜在表現を用いたエンドツーエンド学習を実現する。
• TSSMの学習に変分推論の目的関数（ELBO）を用い、確率的潜在状態と決定的隠れ状態を多頭部自己注意によって更新する。
• TSSMが生成する想像された軌道に基づいて学習するトランスフォーマー基盤のアクタクリティック政策を採用し、サンプル効率の高い学習を実現する。
• TSSMをDreamerフレームワークに統合し、RSSMの置き換えを実現しつつ、全体のMBRL学習ループを保持する。

実験結果

リサーチクエスチョン

• RQ1 長期間にわたる視覚的制御タスク（複雑な記憶の相互作用を要するもの）において、トランスフォーマー基盤のワールドモデルが再帰的モデルを上回れるか？
• RQ2 安定したエンドツーエンドのトランスフォーマー基盤MBRLエージェントを構築するにあたり、主なアーキテクチャ的および訓練上の課題は何か？
• RQ3 トランスフォーマー基盤のワールドモデルは、特に画像生成と報酬予測において、想像された軌道の質を向上させられるか？
• RQ4 短期的依存性を持つタスク（例：DMC Cheetah Runやアーケードゲーム）において、トランスフォーマー基盤のMBRLエージェントはRNNベースのベースライン（Dreamer）と比較してどのように性能を発揮するか？
• RQ5 TSSMワールドモデルは、多様な視覚的環境に一般化可能であり、同時に訓練の安定性と計算効率を維持できるか？

主な発見

• 5-Ball Dense環境（長期間記憶と正確なボール順序の推論を要する複雑な2D視覚的制御タスク）において、TransDreamerはDreamerを上回る性能を発揮する。
• TransDreamerは、より鮮明で正確な想像されたフレームを生成し、ボール収集イベントにおける報酬スパイクを正しく予測する。一方、Dreamerはぼやけた画像と誤った報酬予測を生成する。
• TSSMワールドモデルは、特に長期間依存関係が重要な後方の想像ステップにおいて、DreamerのRSSMよりも正確な将来のフレーム予測と報酬推定を実現する。
• 短期的記憶タスク（例：DMC Cheetah Runやアーケードゲーム）では、TransDreamerはDreamerと同等の最終的な性能を達成するが、収束が遅い。
• TransDreamerは、報酬予測の精度と視覚的想像の忠実度が向上しており、これはDreamerのRNNベースの対比モデルよりもより強固で正確なワールドモデルであることを示している。
• TSSMは、適切なアーキテクチャ設計が施された場合、完全にトランスフォーマー基盤のRLエージェントの安定した訓練を可能にし、トランスフォーマー基盤MBRLが実用的で効果的であることを実証している。