MASTERING ATARI WITH DISCRETE WORLD MODELS (DreamerV2)
通称DreamerV2
DreamerV2は、同じ計算資源と学習時間でシングルGPUのトップエージェントであるIQNやRainbow超え
アブスト
モデルフリー系は性能が高いが大量のインタラクションが必要
しかし近年は表現学習等が成功しており、画像入力から世界モデルを学習し、それをプランニングに使えるように
より少ないインタラクションから学習
オフラインデータからの一般化を促進?
forward-lookingな探索
異なるタスク間での知識の再利用可能に
SimPle等の世界モデルは、モデルフリーには性能かなわない
VPN、MuZeroのようなタスク特化のプランニング手法がやっぱり強い
世界モデルは、タスクをより効率的に達成するため、特定のタスクに絞った環境じゃなくて、一般的な環境のパターンを学習することを目的としてる
DreamerV2は画像から得られる一般的な情報のみを使い、それが報酬に影響されていない場合でも、将来の報酬を正確に予測する
単一GPUで同じ計算資源・サンプル数のIQN,Rainbowを凌駕
イントロ
世界モデルは、環境に関するエージェントの知識を明示的に表現する方法
モデルフリーより一般化を促進し、潜在的な行動の結果を予測してプランニングできる
世界モデルは、環境を一般的に捉えることができ、別タスクや、オフラインデータからの環境の一般化に有効であることが示されている(Sekar et al., 2020) (Yu et al., 2020).
画像の代わりにコンパクトな表現を予測することで、一般化できているという仮設がある
またこれによって、メモリ使用量も小さくできるため、1GPUで数千の並列予測ができる
このアプローチを使ってDreamerは画像入力で連続制御タスクを達成
MuZeroはプランニングがボードゲームやAtari等の決定的ゲームに有効であることを示している
ただ、MuZeroは重たすぎて実用的じゃない(2ヶ月学習、大量GPU)
学習した世界モデル内で学習しAtariで人間超えしたDreamerV2を提案
成功したってことは、環境を正確に表現できていることを示している
そのためにDreamerからいくつか改善
- 離散的な潜在空間
- KL損失の項のバランシング
また、アルゴリズムを厳密に比較するために、人間のゲーマーと人間の世界記録の両方で正規化したスコアを使う方法を提案
DreamerV2
- 過去の経験のデータセットから世界モデルを学習
- コンパクトなモデルの状態から、Actor ctiricで学習
- 経験データセットを増やすために、学習したActorを環境で動かす
WORLD MODEL LEARNING
エージェントの経験を予測モデルにまとめたもの
学習した環境モデルで行動の学習ができる
入力が高次元画像だと、コンパクトな状態表現を学習することが有益(Watter et al., 2015; Karl et al., 2016; Ha and Schmidhuber, 2018)
これらのモデルはlatent dynamics modelsと呼ばれる
潜在空間で先を予測することは、長期的な予測を容易にし、画像生成することなく数千のコンパクトな状態のシーケンスを一括して予測でき効率的
V2ではV1のガウス潜在変数をカテゴリ変数に置き換えた
Experience dataset
世界モデルはエージェントの過去の経験(images, actions, rewards, discount factors)から学習される
学習には、エピソードからランダムにサンプルされた、固定長L=50をバッチ
学習でエピソード終わりを十分に観測するために、ランダムサンプルされた開始インデックスからシーケンス(50)がはみ出すときにはクリッピングする
Model components
世界モデルの構成(詳細はFig.2)
- Image encoder
- Recurrent state-space model (RSSM)
- Predictor (for image, reward, discount factor)
RSSMは決定的recurrent state のシーケンス
RSSM: https://www.slideshare.net/DeepLearningJP2016/dllearning-latent-dynamics-for-planning-from-pixels
事後状態
事前状態
コンパクトなモデル状態は決定論的状態と確率的状態の結合
事後状態から、現在の画像の状態
すべてのコンポーネントはNNで実装
遷移予測器は、次の状態を今の状態と行動から予測(画像の観測や生成しなくてもよい)
割引率予測器は、モデル予測から行動を学習する際に、エピソードが終了する確率を推定できる??
Neural networks
RSSMはGRUで、決定的なリカレントの状態を計算
モデル状態は、決定的なGRU状態とサンプルされた確率的な状態のconcat
Image predictorは逆conv
84x84 grayscale to 64x64にしてDreamerV1のCNNアーキテクチャに突っ込んでる
transition, reward, discount predictorはMLP
モデルはすべてELU activation
Distributions
Image predictor: diagonal Gaussian likelihood with unit variance(これがわからない)
reward predictor: univariate Gaussian with unit variance
discount predictor: ベルヌーイ分布の尤度
DreamerV1ではモデル状態の潜在変数はガウス分布なので、reparaで勾配通してたが、V2ではカテゴリカル変数のベクトルを使う(Bengio et al., 2013)
以下のアルゴリズムで、自動微分を用いてかんたんに実装できるstraight-through勾配を使用して学習
Loss Function
すべてのコンポーネントは同時に最適化
image, reward, discount, transitionは対数尤度を最大化
representationは以下の期待値により予測タスクを容易にするモデル状態を生成するように学習
更に高いエントロピーを持つモデル状態を生成するように正則化されており、学習中の多くの異なるモデル状態に対してロバスト
OptimizerはAdam, Atariの場合はβ=0.1で、連続制御の場合はβ=1.0でKL lossをスケーリング (Higgins et al., 2016).
KL balancing
世界モデルはrepresentation modelが事後状態、transitionが一時的な事前状態である、sequential VAEであると解釈できる
ELBO(evidence lower bound, 変分下限)
参考: https://www.ccn.yamanashi.ac.jp/~tmiyamoto/img/variational_bayes1.pdf
ELBO objectiveでは、KL lossが2つの役割を果たす
representationに対する事前状態の学習
事前状態に対するrepresentationの正規化
しかし、事前状態の学習が不十分な状態でrepresentationを正規化されたくないため、事前の学習率をα=0.8、近似事後の学習率を1-αとする
これにより、representationより事前状態の方が早くKL lossを最小化できる
このテクニックをKL balancingと呼ぶ
事後のエントロピー増加より、正確な事前の学習を促す
BEHAVIOR LEARNING
actor-criticの学習するときはrepresentationは固定
1GPUで2500の潜在的なtrajectoryをシミュレート
Imagination MDP
世界モデルの潜在空間の行動を学習するために、imagination MDPを次のように定義
初期状態分布
transition は
rewardの平均値を報酬シーケンス
discountは割引率のシーケンスを出力
エピソード終了の可能性を考慮して、予測された割引率でactor-ctiricのロスを重み付け
Fig2キャプション約
trajectoryはモデル学習時に計算された事後状態から始まり、actorから行動をサンプルして、次を予測
criticは各状態の将来の報酬の期待値の合計を予測
ctiricは予測された報酬からTD学習
actorは強化学習の勾配と、世界モデルのstraight-through勾配を使ってctiricの予測を最大化するように学習
Model components
imagenation MDPにおける長期的な行動を学習するため、行動を選択するactorと決定論的なctiricを学習する
actorはcriticの出力を最大化する行動を学習することを目指す
criticは想像上の各状態からactorが獲得する将来の報酬の合計を正確に予測することを目指す
actorのパラメータ
実際の環境とは対照的に、潜在的な状態はマルコフ性があるので、追加条件は必要ない
actor-criticはともにMLP ELU活性化関数で1Mパラメータ
actorは行動に関するカテゴリカル分布
ctiricは決定論的な出力(スカラーかな?)
それぞれロスは別
Citric loss
Imagination MDPでは複数ステップのオンポリシー軌道(actorに従ってるため)を生成できるため、n-step learningの方が効率良い
そのため、 general λ-target (Sutton and Barto, 2018; Schulman et al., 2015)を使用
2乗損失で学習
Adamで最適化
最終ステップではtarget=ctiricなので損失項はなし
sg関数で勾配を止める
100stepのtaget networkで価値学習を安定化(DQNのやつ)(Mnih et al., 2015)
Actor loss
ctiricの学習で計算した
actorのパラメータに関して異なる勾配がある
DreamerV2は次の2つの勾配を組み合わせている
unbiased but high-variance Reinforce gradients
biased but low-variance straight through gradients
更に、actorのエントロピーを正則化することで、探索も促している
DreamerV2では離散的な潜在変数と離散的な行動の用法を使っている
サンプルされた行動と状態のシーケンスをバックプロパゲートするため、ストレートスルー勾配を使う(Bengio et al., 2013)
これによってlow-varianceでbiased勾配の推定値になる
直感的にはlow-varianceでbiasedなダイナミクスをバックプロパゲートする方が、初期の学習速度が早く、unbiased high-varianceな方が良い解に収束する可能性がある
Atariではreinforceの勾配が効果的だったので ρ = 1 and η = 10−3
連続値制御ではダイナミクスが効果的だったのでρ=0とη=10-4
これらのハイパラをアニーリングしたらもっと良い性能出たけど、アニーリングなしのスコアを論文に乗せるわ
Adam
3 EXPERIMENTS
Atariベンチで評価
4つのモデルフリーアルゴリズムと比較(IQN, Rainbow, C51, DQN)
DreamerV2のでは、1台のNVIDIA V100 GPUと1つの環境インスタンスのみを使用し、10日以内に2億回の環境ステップを達成
2億回の環境ステップの間に、DreamerV2はモデルの下で想像される468Bのコンパクトな状態から学習します。
これは、行動を繰り返した後に実環境から受け取る50Mの入力よりも1万倍多い
Experimental setup
(Mnih et al., 2016; Brockman et al., 2016; Hessel et al., 2018; Castro et al., 2018; Badia et al., 2020)で選ばれてる55のゲームで評価(今後はこれ推奨)
200Mの環境に対するステップ
action repeatは4(続けてactionを入力かな)
1episodeあたり108,000stepの時間制限
メタデータへのアクセスなし(画像のみ)
full action space, and sticky actions.
DreamerV2ではframe stackingは行わない
ゲームごとに別々のエージェント
各エージェントは1つの環境インスタンスのみ(並列化禁止←これきつくない?)
Model-free baselines
Rainbowは10日かかったって
ATARI PERFORMANCE
4つの比較方法を検討
- Gamer Median
ランダムポリシーとプロゲーマーで正規化されて、複数シードで平均とり、すべてのタスクの中央値
しかし半分のゲームのスコアが0でも値が変わらずロバストじゃない - Gamer Mean
すべてのタスクを考慮しているが、ゲーマーが苦手(RLが得意)なゲームで高い正規化されたスコア出せるので、平均で人間超えは余裕 - Record Mean
プロゲーマーではなく人間の世界記録で正規化するとたしかに改善するが、やはりRLが得意なゲームが支配的 - Chipped Record Mean
人間の世界記録でクリップすることで、1を超えないようにする
これでRLのイカれた記録がででも堅牢
rainbowがmedianでは良いスコアだけど、meansだとIQNに負ける
scheduledはアニーリング
Individual games
表K.1にスコア一覧
世界記録がないゲーム2つあったけど適当な値いれた
図E.1はそれぞれのゲームについて各アルゴリズムと比較
Video Pinaballが全負けした
理由としては、ボールが1pixelしかなかったから、世界モデルで再構成できなかったからと考察
ABLATION STUDY
要素の検証
- カテゴリカル潜在変数 vs ガウシアン潜在変数(V1)
- KL balancing
- representationを学習する上でimage reconstruction, reward からの勾配の重要度を調査するために、モデル状態に入る前に勾配を止めた
Fig5, Table2に結果がまとまっている
1の結果
カテゴリカルが55game中42ゲームが勝ち、8負け、5タイ
なんでカテゴリカルがいいのかわかってない
2の結果
44ゲームが勝ち、6負け、5タイ
世界モデルのprior dynamicsを学習することは、ポリシーの学習に重要(軌跡を生成するから)
kl balancingは世界モデル以外の学習済みの事前情報を持つ確率モデルにも有効と考えられる
3の結果
imageのgradient stop
3勝ち, 51負け, 1タイ
rewardのgradient stop
15勝ち, 22負け, 18タイ
ちょっと良くなったわけは、経験した報酬に固有じゃない表現が、しらない状況に対してよりよく一般化したから?
reinforce gradintだけ
18勝ち, 24負け, 13タイ
DreamerV2がreinforce gradientを主に使っていることを示している
straight-forward勾配をreinforceに混ぜると大幅に性能up
straight-forward勾配だけを使うと惨敗
これだけだとbiasがあるためポリシーの最適化には適さないからと考察
DIscussion
DreamerV2はV1にカテゴリカル潜在変数とkl balancingを導入したもの
報酬によらず画像から表現学習できている