[論文レビュー] Transformer-based World Models Are Happy With 100k Interactions

Jan Robine, Marc Höftmann+2 more|arXiv (Cornell University)|Mar 13, 2023

Reinforcement Learning in Robotics被引用数 4

ひとこと要約

本論文は、100万回の環境インタラクションのみを用いて、Atari 100kベンチマークで最先端の性能を達成するTransformerベースのワールドモデル（TWM）を提案する。Transformer-XLアーキテクチャを用いて、観測値、行動、予測報酬を自己回帰的にモデリングし、予測報酬を再帰的にフィードバックすることで、サンプル効率の高いポリシー学習を可能にした。これにより、従来のモデルフリーおよびモデルベース手法を上回る性能を発揮しながらも、推論の計算コストを低く抑えられる。

ABSTRACT

Deep neural networks have been successful in many reinforcement learning settings. However, compared to human learners they are overly data hungry. To build a sample-efficient world model, we apply a transformer to real-world episodes in an autoregressive manner: not only the compact latent states and the taken actions but also the experienced or predicted rewards are fed into the transformer, so that it can attend flexibly to all three modalities at different time steps. The transformer allows our world model to access previous states directly, instead of viewing them through a compressed recurrent state. By utilizing the Transformer-XL architecture, it is able to learn long-term dependencies while staying computationally efficient. Our transformer-based world model (TWM) generates meaningful, new experience, which is used to train a policy that outperforms previous model-free and model-based reinforcement learning algorithms on the Atari 100k benchmark.

研究の動機と目的

• 深層強化学習は通常、人間の学習者よりもはるかに多くの（数億回の）環境インタラクションを必要とするという、サンプル非効率性を解消すること。
• 再帰的ワールドモデルの限界を克服し、自己注意機構によって過去の状態への直接的かつ長距離の注目を可能にすること。
• 生成的ワールドモデルを用いて潜在的想像上でのポリシー学習を実施することで、サンプル効率を向上させること。
• 新しい損失関数とサンプリング戦略を用いて訓練を安定化させ、一般化性能とハイパーパramータのロバスト性を向上させること。
• Transformerベースのワールドモデルが、わずか100万回の実環境インタラクションでのみ、強力な性能を発揮できることを実証すること、特にAtariで100万回のインタラクションを対象とする。

提案手法

• 高次元の画像観測値をCNNで潜在状態に符号化し、それをTransformer入力シーケンスに埋め込む。
• 潜在状態、行動、予測報酬を入力トークンとしてTransformerに供給し、3つのモodalの間で柔軟に注目できるようにする。
• Transformer-XLアーキテクチャを用い、相対的位置エンコーディングを導入することで、長期依存関係を効率的にモデル化し、推論時の計算コストを低減する。
• MLPを用いて、Transformerの隠れ状態に条件付けられた形で、次の潜在状態、報酬、割引係数を予測する。
• ワールドモデルの目的関数におけるエントロピー項と交差エントロピー項のトレードオフを微調整するため、バランスの取れたKLダイバージェンス損失を導入する。
• ポリシーのエントロピーを訓練中に安定化させるために、しきい値付きエントロピー損失を適用し、ハイパーパramータ選択への感受性を低減する。
• 最近の経験に重点を置くバランスの取れたサンプリング手順を実装し、訓練データセットにおけるデータ分布のずれを是正する。

実験結果

リサーチクエスチョン

• RQ1 100万回の環境インタラクションでのみ学習されたTransformerベースのワールドモデルが、Atari 100kベンチマークで既存のモデルフリーおよびモデルベースRL手法を上回ることができるか？
• RQ2 予測報酬をTransformer入力にフィードバックすることで、ワールドモデルの品質および下流ポリシーの性能にどのような影響を与えるか？
• RQ3 提案されたバランスの取れたサンプリング戦略は、均一サンプリングと比較して、訓練の安定性と性能をどの程度向上させるか？
• RQ4 しきい値付きエントロピー損失は、多様なAtariゲームにおいてポリシー訓練を安定化させ、ハイパーパramータの感受性を低減できるか？
• RQ5 Transformer-XLアーキテクチャの使用により、RNNベースのワールドモデルよりも優れた長期依存関係のモデル化が可能になり、推論コストを低く保てるか？

主な発見

• 提案されたTWMは、Atari 100kベンチマークで人間正規化スコアにおいて最先端の性能を達成し、従来のモデルフリーおよびモデルベース手法を上回った。
• Abulation studiesにより、Transformerに予測報酬を条件付けた場合、大多数のゲームで顕著な性能向上が確認された。
• バランスの取れたサンプリング手順により、最近の経験に注目を向けることで、データ分布のずれが軽減され、ポリシー性能が向上した。
• しきい値付きエントロピー損失により、ポリシーのエントロピーが訓練中に安定化し、ハイパーパramータチューニングが簡素化され、ゲーム間での一般化性能が向上した。
• バランスの取れたKLダイバージェンス損失により、エントロピーと交差エントロピーのトレードオフを細かく制御可能となり、ワールドモデル学習の向上が達成された。
• 本モデルは、わずか100万回の実環境インタラクションでのみ高い性能を発揮し、ポリシーのロールアウト時にTransformerを使用しないことで、サンプル効率と推論コストの両面で優れた効率性を示した。