どこから見てもメンダコ

軟体動物門頭足綱八腕類メンダコ科

DreamerV2の実装: 世界モデルベース強化学習でブロック崩し

世界モデル系強化学習の先端手法であるDreamerV2ブロック崩し(BreakoutDeterministic-v4)向けに実装しました。


論文
DreamerV2: [2010.02193] Mastering Atari with Discrete World Models

Google AI Blog: Mastering Atari with Discrete World Models

DreamerV1: [1912.01603] Dream to Control: Learning Behaviors by Latent Imagination

Google AI Blog: Introducing Dreamer: Scalable Reinforcement Learning Using World Models


はじめに

2012年にDeepMindDQNが世界を驚かせてから、たったの10年間で深層強化学習研究は目覚ましい進歩を遂げました。2012年のDQNの頃は深層強化学習エージェントが”人間レベル”でブロック崩しをプレイ可能というだけで驚くべき成果でしたが、2022年現在のモデルフリー強化学習の先端手法である R2D2MuZero ではブロック崩し環境において超人的パフォーマンスに到達しており余裕のパーフェクトプレイを見せつけてきます。


Atari2600 Breakout (ブロック崩し

しかし、このような手法の飛躍的発展にも関わらず深層強化学習の産業応用はそれほど進んでいません。原因の一つは、深層強化学習"劣悪なサンプル効率" 汎化性能の低さという課題を抱えていることでしょう。 たとえばブロック崩し環境について、人間であれば4, 5回の試行錯誤で到達できるようなパフォーマンスに深層強化学習エージェントが到達するためには最先端の手法でさえ何千回の試行錯誤が必要となるのです。また、何千何万回の試行錯誤の末に到達した超人的な強化学習エージェントでさえ、背景画像を変えられてしまうだけでパフォーマンスが崩壊してしまいます。

もし人間のように少数の試行でもそこそこのパフォーマンスを発揮できるような高度な学習能力を持つ強化学習手法を開発できたのならば強化学習の産業実用性は大きく高まります。この実現を目指すアプローチのひとつが世界モデルベース強化学習です。


世界モデルベース強化学習とは

「人間の高度な学習能力を支えているのは脳内シミュレータ―に基づく未来予測と行動計画である」という仮説のもと、脳内シミュレータを備えた強化学習エージェントを訓練しようというのが世界モデルベース強化学習のコンセプトです。

野球を例にとってみましょう。野球の打者は、バットをどのように振るかを決定するのに数ミリ秒かかります。これは、目からの視覚信号が脳に到達するのにかかる時間よりも短い時間です。100マイルの速球を打つことができる理由は、ボールがいつどこに行くかを本能的に予測できるためです。プロのプレーヤーにとって、これはすべて無意識のうちに起こります。彼らの筋肉は、内部モデルの予測に沿って、適切な時間と場所でバットを反射的に振ります。計画を立てるために考えられる将来のシナリオを意識的に展開する必要なしに、将来の予測にすばやく対応できます。(Google翻訳)  ― World Models より

つまりは、『World Models (世界モデル)』ベースの強化学習手法が目指すのは環境のダイナミクス(状態遷移)のモデル化です。もし系の状態遷移を完全にモデル化できたのならばその状態遷移モデルは系の(微分可能な)シミュレータ―として機能します。もし完全なシミュレータ―が得られたらばあとはそのシミュレータ上で満足いく結果が得られるまで行動計画を続ければよいのです。

https://worldmodels.github.io/


DreamerV2:Atari環境で初めてモデルフリー手法に並んだ世界モデルベース強化学習

世界モデルベース強化学習はコンセプトは面白いのですが、モデルフリー手法と比較するとパフォーマンスはいまいちというのがこれまでの傾向であり、とくにAtari2600ベンチマークレトロゲーム)のように視覚的ノイズの大きい環境が苦手でした。

しかしICLR2021にて発表された"DreamerV2"は、非分散並列手法としてはatari環境で長くトップにあったモデルフリー手法Rainbowを超えるパフォーマンスを示したことにより、世界モデルベース強化学習への期待感を大きく高めることとなりました。本稿ではこのDreamerV2のブロック崩し環境向け実装例を簡単な解説とともに紹介します。

論文Fig1 DreamerV2は世界モデルベース強化学習で初めてアタリ環境でのRainbow越えを達成

[2010.02193] Mastering Atari with Discrete World Models(DreamerV2) 


世界モデル(World Models)について

世界モデルがやりたいこととは、t時刻における状態xとアクションaからt+1時刻における状態x を予測することです。しかし画像のような高次元かつ密なデータにおいて精度よく遷移予測を行うのはなかなかに困難な課題です。

画像観測から次の画像観測を直接予測するのは難しい

そこで、系の直接観測はより少数の潜在変数に支配されているはずだ(多様体仮説)という考えにもとづき、Variational AutoEncoder(VAE) の訓練によって得られる潜在変数空間における遷移モデルを関数近似しようというのがWorld Modelの基本コンセプトです。

下図について、潜在変数空間における状態遷移モデル  \displaystyle{
 q( z_{t+1} | z_{t}, a_{t})
} は確率モデルであることに注意してください。これは状態遷移を決定論的な関数で近似してしまうと、(極端には)ブラックジャックのようなランダム系にまったく対応できなくなってしまうためです

World Modelは潜在変数空間における状態遷移モデルを近似する

このような「潜在変数空間における状態遷移のモデル化」というコンセプト自体は実際のところ何ら新しいものではありません。Dreamer論文内でも言及されているようにこれは非線形カルマンフィルタや隠れマルコフモデルのバリエーションと捉えることができるためです。

それにも関わらずWorld Modelを魅力的な手法にしている理由は主に2点です

  1. 高い説明可能性
  2. 微分可能シミュレータの獲得

まず前者について、World ModelではVAEアーキテクチャを用いるので潜在変数から実観測を復元することができるため、高い説明可能性を有しています*1。すなわち「今この行動をしたらこの後はこういう風になるよ」ということを視覚的に説明できるわけです。説明可能性の高さは産業活用へのハードルを大きく低減するため大変重要な特性です。

VAEなので潜在変数を実空間に復元できる

さらに後者について、World Model が獲得する状態遷移モデル=系のシミュレータの実体はニューラルネットワークであるため微分可能です。よって方策も関数近似している場合は、シミュレータから獲得できる報酬を最大化するように方策ネットワークを直接最適化できることとなります。イメージとしてはDDPGにおけるQ関数がシミュレータに置き換わった感じでしょうか。

参考: 微分可能なシミュレータ上での方策最適化 - Preferred Networks Research & Development


微分可能であることを置いておいたとしても、軽量なシミュレータを獲得できることは一つの系で様々なタスクを学習させたい場合や試行の金銭/時間的なコストが重い系においては魅力的な特性です。また、世界モデルの構築は強化学習ではなく単に教師あり学習なのでデータの再利用が容易であるためオフラインセッティングと相性がよいことも特筆すべき利点です。


方策の獲得

世界モデル=系のシミュレータがうまく構築できているならば方策の獲得はどのように行ってもOKです。世界モデルを普通にシミュレータとして使って方策勾配法で最適化してもよいですし、微分可能シミュレータであることを生かしてダイレクトに方策ネットワークを最適化してもよいです。あるいは軽量シミュレータであることを生かしてCMA-ESのような進化戦略アルゴリズムで方策パラメータをブラックボックス最適化してしまうという方法もよく用いられているようです。


Dreamerへの系譜

DreamerV2に至るまでの先行研究の流れをまとめます。

Wolrd Models (2018)

[1803.10122] World Models ( David Ha, Jürgen Schmidhuber, 2018)

"World Models" という用語の初出は(たぶん)この論文です。コンセプトの似た研究はこの論文以前からありましたが、World Models というキャッチーなネーミングによってVAEベースの状態遷移モデルによるアプローチが広く認識されるようになりました。

論文では"MVC"の3コンポーネント*2で構成されるWorld Modelアーキテクチャ のコンセプトを提案しました。

  • Vision: 画像観測から潜在状態zを抽出するVariational AutoEncoder
  • Memory RNN: 過去の観測情報を潜在変数空間で記憶し現観測と合わせて次観測を予測
  • Controller:方策関数

World Models論文 Fig4より

Fig6 Memory RNNアーキテクチャ詳細

この手法における状態遷移の予測は次の3ステップで行います。

  1. 時刻tの潜在変数zと行動aをLSTMに入力
  2. LSTM出力から時刻t+1における潜在変数zの分布を予測(混合ガウス分布のパラメータを予測)
  3. 予測した時刻t+1における混合ガウス分布から潜在変数をサンプリング


Controller(=方策)について、World Models における方策関数は超シンプルな線形モデルです。パラメータ数が少ないことを生かして進化戦略アルゴリズム(CMA-ES)を用いて重みをブラックボックス最適化(勾配フリー最適化, derivaritive-free optimization)してしまいます。

World Models の方策関数はシンプルな線形モデル


PlaNet: Deep Planning Network (2019)

Google AI Blog: Introducing PlaNet: A Deep Planning Network for Reinforcement Learning

https://arxiv.org/pdf/1811.04551.pdf

オリジナルのWorld Modelでは系の潜在状態を確率変数(VAEの潜在変数 z)でのみ表現しますが、PlaNetでは確率的変数と決定論的変数の両方で系の潜在状態を表現することによって、World Modelアーキテクチャ時系列video予測を高精度化することに成功しました。

  • 決定論的潜在変数*3: RNN(GRU)の出力h
  • 確率的潜在変数: VAEの潜在変数z

多くのベンチマーク環境に置いて確率的な遷移をするのは画面のごく一部であり背景などはほぼ決定論的な遷移をすることを考えれば、決定論的潜在変数を導入することでvideo predictionが高精度化することは妥当な結果に感じます。このPlaNet版World Modelアーキテクチャはほぼ姿を変えずに以降のDreamerおよびDreamerV2に継承されています。

Google AI blog Dreamer より(※DreamerとPlaNetの世界モデル部はほぼ同じ)

PlaNetでは方策関数を作成せず、軽量高速並列シミュレータとしてのWorld Modelの特性を生かして最適なアクションシークエンスを直接探索します。たとえば離散行動空間4の系であればこの先15stepについて可能なアクションシークエンスは4**15通り存在します。いくら世界モデルが軽量シミュレータとして機能すると言えどもこれを全探索はさすがに厳しいので、PlaNetでは進化戦略アルゴリズムで最適シークエンスを探索します。

探索に使用する進化戦略アルゴリズムは(計算リソースが許すなら)遺伝的アルゴリズムでも何でもいいと思うのですが、PlaNetではシンプルに Cross-entropy method を採用しています。


Dreamer (2020)

Google AI Blog: Introducing Dreamer: Scalable Reinforcement Learning Using World Models

https://arxiv.org/pdf/1912.01603.pdf

Dreamerとは世界モデル上で強化学習エージェントを訓練するという SimPle的なアプローチがPlaNetの高精度なWorldModelと組み合わさることにより大きな成功を収めた手法と要約できます。

DreamerはPlaNetの高精度なWorld Model上で強化学習エージェントを訓練することによって、連続値コントロール環境において当時SotAのモデルフリー手法であるD4PG(Distributinal Distributed Deep Deterministic Policy Gradient)に相当する高いパフォーマンスを達成することに成功しました。しかもサンプル効率は20倍以上です。

モデルフリーSotA手法D4PG 相当/以上のパフォーマンス

Dreamerは世界モデル上で強化学習エージェントを訓練する

Dreamerの特徴はPlaNet版世界モデル上のみでActor-Critic型の強化学習エージェントを訓練する点です。エージェントのネットワーク更新に実環境との相互作用によって得られたサンプルは一切使わないゆえに、夢の中で訓練する"Dreamer"という命名がなされています。ニュアンスとしてはイメージトレーニングといった感じ。

また、Dreamerは世界モデルが即時報酬について微分可能であることを最大限活用し、方策勾配定理を使わず期待報酬和を直接最大化するように方策関数のパラメータを更新します。これはDDPGの更新式とよく似ており、すなわちDDPGではQ関数の出力値を大きくするように方策関数を更新しますが(下式)、Dreamerではシミュレータ(世界モデル)から獲得するN-step報酬和*4が大きくなるように方策関数を更新します。

これはDDPGの更新式


DreamerV2 (2021)

Mastering Atari with Discrete World Models | OpenReview

Google AI Blog: Mastering Atari with Discrete World Models

DreamerV2では、Dreamerに2つのトリックを追加することでAtari環境においてRainbowを超えるパフォーマンスを発揮することを報告しました。

  1. VAEの潜在変数分布に正規分布でなくOneHotCategoricalDistributionを仮定 (DiscreteVAE)
  2. KLバランシング

1についてはVAEの潜在変数分布に単峰正規分布という強い仮定を置くのではなく、より自在な分布を表現できるカテゴリ分布を用いることで系の確率的な遷移をより表現しやすくなるのではないかと考察されています。

DreamerV2ではVAEをDiscreteVAEに差し替える

潜在変数に単峰ガウス分布を仮定するくらいなら多少目が粗くても自由な分布を表現できるカテゴリ分布のほうがいい

世界モデルベースで初めてRainbowのパフォーマンスを超えたのは間違いなく大きな成果なのですが変更があまりにも単純なのでこれ自体は語るところの少ない論文でもあります。また、論文ではサンプル効率がRainbowと比較して良いことが示されているわけでもなく、モデルベースというくくりで見ればパフォーマンスでMuZeroにはまったく敵わないということもあり、OpenReviewでは"動機不明の研究"との辛口コメントも見られます。さらに、DreamerV1では最大限に活用されていた「世界モデルの微分可能性」もAtari環境ではあまり役に立たないようです。

・・・アカデミックにはあまり面白くない手法かもしれませんが、世界モデルでRainbow超えは素直にすごいと思ったので実装してみたというのが本記事のモチベーションです。


非VAEの状態遷移モデル:MuZero(2020)

MuZero: Mastering Go, chess, shogi and Atari without rules

世界モデルとは呼称されませんがコンセプトのよく似た手法であるMuZero についても簡単に説明します。

MuZeroとはAlphaZeroのモンテカルロ木探索アルゴリズム潜在変数空間*5における状態遷移モデルを導入することで、AlphaZeroを古典ボードゲーム以外(=Atari2600環境)にも適用可能にした手法です。”ニューラルネットワークで近似された状態遷移モデルを利用した先読み検索による行動プランニング”という、PlaNetに極めて近いコンセプトによって当時のトップモデルフリー手法 R2D2を超えAtari環境SotAを達成しました。

MuZero: 潜在変数空間の状態遷移モデルを利用したMCTS


コンセプトが似ているにも関わらず、atari環境に置いてMuZeroはDreamerV2よりもはるかに高いパフォーマンスを発揮するのはなぜでしょうか? あくまで個人の考察ですが、これはMuZeroがいろんなものを切り捨ててatari環境でのパフォーマンスに特化したからと思われます。

1. MuZeroは説明可能性を切り捨てている

VAEベース世界モデルの特徴のひとつは、潜在変数を実画像に復元可能であるために未来予測を人間が視覚的に理解可能であることです。これは実用上は本当に重要である一方で、実画像への復元可能性を保証するために潜在変数が冗長なものになってしまっている可能性があります。たとえば、ブロック崩し攻略において重要なのは極論 "ボールの位置" と ”パドルの位置” だけであるに関わらず、VAEベース世界モデルの潜在変数にはブロックの色やスコア表示など余分なノイズ情報が多量に含まれているはずです。

一方、MuZeroでは潜在変数を実画像へ復元することはできない代わりに効果的な状態表現獲得ができていると考えられます。

2. MuZeroは確率的遷移を切り捨てている

VAEの潜在変数は確率変数であるため、確率的な状態遷移を自然に取り扱うことができます。一方、MuZeroの潜在変数とはResNetの隠れ層でしかないので決定論的です。同じ状態で同じ行動を取っても確率的に遷移する系(コイントスとか)においては状態表現が確率的であることは必須であるように思いますが、しかしAtari2600のほとんどのゲームの挙動は決定論的であると見なせるため、Dreamerでは確率的な遷移を表現できるVAEベース状態表現のメリットがあまり得られない上に学習を不安定にしていると考えられます。

3. MuZeroはPOMDPを切り捨てている

MuZeroはRNNを採用していませんが、代わりに直近32フレーム(2-3秒相当)をまとめて入力することで短期記憶の代わりとしています。アタリ程度なら直近32フレーム入力すれば十分にMDPと見なせる、と割り切ることでシンプルなアーキテクチャを実現しています。

horomary.hatenablog.com


Tensorflow2による実装例

ここからはtensorflow2によるBreakout向け実装例を紹介します。要点のみの掲載ですのでコード全文はGithubを参照ください。

github.com

世界モデル部

DreamerV2の世界モデルアーキテクチャは論文中では下図のようにシンプルに示されていますが、実装はこの図よりもはるかにややこしいことになっています。

DreamerV2世界モデル

DreamerV2世界モデルのデータの流れをより詳細に図示したものがこちらです。

DreamerV2の世界モデル

ロス計算のための予測ヘッドまで含めると下図のようになります*6

予測ヘッドとDecoderまで含めた場合

ロス関数は通常のVAEと同様な(画像再構成誤差+KL項)に (即時報酬予測誤差+割引率予測誤差)を合わせたものとなっています。なお、割引率予測とはほぼエピソード終了予測同じようなものと考えてOKです。

実装例はこちら。discreteVAEなのでカテゴリ分布向けのReparameterization Trickを使ってるのがポイント


ロールアウト部

実観測から潜在状態s(zとhをconcatしたもの)を計算し、sに基づいて行動決定します。


強化学習エージェント部

上述の通り、Atari向けのDreamerV2では世界モデルの微分可能性を活用せず単に軽量なシミュレータとして使用するため、強化学習エージェントの訓練について技術的に特筆するべきことはありません。適当にActor-Criticを実装するだけです。私はPPOで実装しました。


学習結果

着実にパフォーマンス向上しつづけているのですが、実行速度パフォーマンス(実装レベルのチューニング&GPU性能)が悪すぎるため一週間ほど学習を続けてようやく30-40点台に到達するという何とも微妙な結果になってしまいました。R2D2を実装したときも思いましたがRNNが入ると本当に速度チューニングがつらい。こういうややこしいRNNはどうしたら高速化できるのだろうか?


雑記

※個人の定性的な感想です

  • センシティブなハイパラが多すぎる
  • とくに世界モデルとActorが共進化するように学習率を制御するのが大変
  • KLロス項の変動が大きいとActorの学習が破綻する傾向がある
  • VAEの一般的な弱点として小さくて動きの速い要素が存在するとつらい
  • しかしブロック崩しは小さくて動きの速いボールこそが最重要なオブジェクト
  • 世界モデルはコンテクストを理解している印象はなく、暗記ゲーをしている疑いがある
  • 安直だけどMaskedAutoEncoderのような表現学習と組み合わせるとよいのかも

*1:世界モデルから説明可能性を投げ捨てるとMuZeroになる

*2:WebアプリケーションのMVCアーキテクチャとかけてるのもオシャレポイント

*3:決定論的な隠れ層について潜在変数という呼称が適切かは知らない

*4:実際は単純和ではない

*5:ResNetの中間出力を潜在変数と言うことが適切とは思えないが他によいワードが思いつかない

*6:ただし割引率γの図示は割愛