これまで、価値ベース法、ポリシー勾配法、モデルベース法など、様々な強化学習手法について議論してきました。 しかし、単純な手法では対処できない幅広い強化学習問題に対応するために、これまで見てきたものとは異なる方法で機械学習の技術を組み合わせることができます。 この記事では、そのような発展的な手法として、R2D2とDreamerアプローチの2つを取り上げます。

R2D2

プロセスがマルコフ決定過程（MDP）であるためには、すべての状態遷移がマルコフ性を満たす必要があります。つまり、すべての状態遷移確率は現在の状態と行動のみに依存する必要があります。 しかし、残念ながら一部の強化学習問題の定義はMDPではなく、部分観測可能なMDP（POMDP）です。例えば、ゲームの状態観測として画面のスクリーンショットを撮ることは一般的ですが、 スクリーンショットだけでは状態遷移に大きな影響を与えるゲーム内のオブジェクトの速度を捉えることができません。

この問題を解決するために、RNNベースのエージェントを使用できます。RNNは自然に連続入力を捉えることができるため、一連のスクリーンショットを状態として扱い、 MDP要件を満たすことができます。連続データから学習する必要がある強化学習にRNNを活用することは非常に直感的ですが、 RNNを直接適用することを妨げるいくつかの困難がありました。その一つは、個々の状態遷移を保存する優先経験再生（prioritized experience replay）との互換性の問題でした。 優先経験再生は、サンプリングを複数のワーカーに分散させ、DQNアプローチで高いパフォーマンスを達成するための重要なテクニックです。 RNNを優先経験再生と互換性を持たせるために、バッファにエピソード全体を保存する方法が考えられますが、一部の環境では実現不可能です。 代わりに、シーケンスから固定長（例えば80遷移）のセグメントを保存する方法を考えられます。

しかし、固定長セグメントを保存する場合、別の困難が生じます。勾配を計算するためには、セグメントが始まる直前の隠れ状態を復元する必要があります。 単純なアプローチとして、ゼロの初期隠れ状態からセグメントが開始するタイムステップまで実行することが考えられますが、これは計算効率が悪いです。 そこで、バッファに隠れ状態を保存し、それを初期隠れ状態として使用することができます。しかし、隠れ状態はワーカーによって計算され、 ターゲットによって計算されるわけではないため、不一致が生じます。これを回避するために、保存された隠れ状態を使用して、 ターゲットを一定数のタイムステップ（例えば20タイムステップ）実行するBurn-Inと呼ばれるプロセスを行い、 バーンイン後の隠れ状態から学習を開始することができます。

Recurrent Replay Distributed DQN（R2D2）は、保存された隠れ状態から初期隠れ状態を生成するためにBurn-Inプロセスを使用し、 複数のワーカーによって収集された固定長セグメントから学習します。 ここでの優先度は、最大$n$ステップTD誤差$\max_i \delta_i$と平均絶対$n$ステップTD誤差$\bar{\delta}$の混合（$p = 0.9 \max_i \delta_i + 0.1 \bar{\delta}$）によって計算されます。 また、LSTMの使用、報酬のリスケーリング（クリッピングせずに報酬を小さくして学習を安定させながら分散を維持する）、 高い割引率（$\gamma = 0.997$）の使用、各ワーカーに固定$\epsilon$を持つ$\epsilon$-greedyポリシー（$\epsilon_i = \epsilon^{1+\frac{i}{N-1} \alpha}$で計算）を使用し様々な探索度を実現するなどの多くの工夫も導入しています。 これらの工夫により、R2D2は2018年にAtariゲームのプレイにおいて比類ないパフォーマンスを達成しました。 拡張として、Never Give Up（NGU）は、Dyna-Q+と同様に、R2D2のより良い探索のために内部報酬$R'$を報酬関数$R_t = R^{env}_t + \beta R'_t$に導入しています。

Dreamers

以前モデルフリー手法のサンプル効率の低さを克服するために、環境をシミュレートして計画を実行するモデルを訓練するというアイデアを思いつきました。 しかし、現在の状態とアクションのみから次の状態と報酬を予測する環境モデルを訓練することは、POMDPではうまく機能しない可能性があります。 そこで、RNNを活用し、前のステップの隠れ状態（記憶）、現在の状態、およびアクションに基づいて次の状態と報酬の予測を生成するモデルを作ることができます。

しかし現実世界の強化学習問題のほとんどは、大きな状態表現（コンピュータゲームのスクリーンショットなど）を持っており、各RNNセルがこれらの大きな状態表現を処理することは計算コストが高くなります。 代わりに、変分オートエンコーダー（VAE）を訓練して、重要な特徴を保持する小さな潜在状態表現を生成し、セルとアクター・クリティックがこの潜在表現で動作するようにすることができます。 上の図は、この手法を用いて世界モデルがどのように訓練され、計画を通じてアクター・クリティック手法の訓練に使用されるかを示しています。 この方法は、人間が脳内で環境について「夢を見る」ことを学び、その夢からのみ学習することに類似しているため、Dreamerと名付けられています。

DreamerV2は、離散的な潜在状態表現とKLダイバージェンスバランシング（KLダイバージェンス損失を使用）を導入し、 エンコーダーと隠れ状態からの潜在状態生成がより近くなるよう両方向から促しました。これにより、モデルベース手法として、 Rainbow（関数近似に関する記事で紹介した様々な技術を組み込んだDQN）よりも高いパフォーマンスを実現しました。 DreamerV3はモデルをスケールアップし、R2D2のような報酬のリスケーリングを導入し、 クリティック学習のために2ホットエンコードされた$\lambda$リターンに基づく離散回帰アプローチを使用しました。 これにより、大規模なハイパーパラメータ調整なしで、様々なタスクにおいて前身よりも高いパフォーマンスとサンプル効率を実現しました。

AlphaZeroの後継であるMuZeroは、実際の環境ではなくMCTS内で潜在状態表現を扱うことを学習する類似のアプローチを取り、 様々な領域で前身よりも優れたサンプル効率と高いパフォーマンスを達成しました。Atariゲームではさらに、R2D2よりも優れた結果を達成しました。 しかし、Dreamerアプローチとは異なり、MuZeroはVAEを使用せず、特定のシナリオ（環境が確率的である場合や説明可能性が重要な場合）では重要である、 分布的な潜在状態表現を生成したりデコーダーで状態を再構築したりすることを目指していません。 さらに、MuZeroは各フレームを処理するRNNに依存せず、Atariの複数フレームを入力状態として取り込み、よりシンプルなアーキテクチャを実現しています。

結論

この記事では、価値ベースおよびモデルベースの手法でPOMDPをMDPにするためにRNNを活用するR2D2とDreamerモデルの主要なアイデアについて議論しました。 また、MCTSでより高いサンプル効率を実現するために潜在状態表現を使用するMuZeroについても簡単に触れました。実装の詳細については、 以下で引用された論文やリソースを確認することをお勧めします。この記事を持って、強化学習分野の最近の論文のほとんどを理解するために必要な基本的な概念をカバーしている、 強化学習シリーズを（少なくとも一時的に）締めくくります。このシリーズから何かを得ていただければ幸いです。

リソース

• cwkx. 2021. Reinforcement Learning 10: Extended methods. YouTube.
• Hafner, D. et al. 2021. Mastering Atari With Discrete World Models. ICLR 2021.
• Hafner, D. et al. 2023. Mastering Diverse Domains Through World Models. ArXiv.
• Kapturowski, S. et al. 2018. Recurrent Experience Replay In Distributed Reinforcement Learning. Open Review.
• Mayer, E. 2023. Model Based RL Finally Works!. YouTube.
• Schrittwieser, J. et al. 2020. Mastering Atari, Go, chess and shogi by planning with a learned model. Nature.
• メンダコ. 2021. ray で実装する分散強化学習 ④R2D2. Hatena Blog.
• メンダコ. 2022. 世界モデルベース強化学習①： DreamerV2の実装. Hatena Blog.