Open16

MASTERING ATARI WITH DISCRETE WORLD MODELS (DreamerV2)

tmyodatmyoda

アブスト
モデルフリー系は性能が高いが大量のインタラクションが必要
しかし近年は表現学習等が成功しており、画像入力から世界モデルを学習し、それをプランニングに使えるように

より少ないインタラクションから学習
オフラインデータからの一般化を促進?
forward-lookingな探索
異なるタスク間での知識の再利用可能に

SimPle等の世界モデルは、モデルフリーには性能かなわない
VPN、MuZeroのようなタスク特化のプランニング手法がやっぱり強い
世界モデルは、タスクをより効率的に達成するため、特定のタスクに絞った環境じゃなくて、一般的な環境のパターンを学習することを目的としてる

DreamerV2は画像から得られる一般的な情報のみを使い、それが報酬に影響されていない場合でも、将来の報酬を正確に予測する

単一GPUで同じ計算資源・サンプル数のIQN,Rainbowを凌駕

tmyodatmyoda

イントロ
世界モデルは、環境に関するエージェントの知識を明示的に表現する方法
モデルフリーより一般化を促進し、潜在的な行動の結果を予測してプランニングできる

世界モデルは、環境を一般的に捉えることができ、別タスクや、オフラインデータからの環境の一般化に有効であることが示されている(Sekar et al., 2020) (Yu et al., 2020).

画像の代わりにコンパクトな表現を予測することで、一般化できているという仮設がある
またこれによって、メモリ使用量も小さくできるため、1GPUで数千の並列予測ができる
このアプローチを使ってDreamerは画像入力で連続制御タスクを達成

tmyodatmyoda

MuZeroはプランニングがボードゲームやAtari等の決定的ゲームに有効であることを示している
ただ、MuZeroは重たすぎて実用的じゃない(2ヶ月学習、大量GPU)

学習した世界モデル内で学習しAtariで人間超えしたDreamerV2を提案
成功したってことは、環境を正確に表現できていることを示している
そのためにDreamerからいくつか改善

  • 離散的な潜在空間
  • KL損失の項のバランシング

また、アルゴリズムを厳密に比較するために、人間のゲーマーと人間の世界記録の両方で正規化したスコアを使う方法を提案

tmyodatmyoda

DreamerV2

  • 過去の経験のデータセットから世界モデルを学習
  • コンパクトなモデルの状態から、Actor ctiricで学習
  • 経験データセットを増やすために、学習したActorを環境で動かす
tmyodatmyoda

WORLD MODEL LEARNING
エージェントの経験を予測モデルにまとめたもの
学習した環境モデルで行動の学習ができる

入力が高次元画像だと、コンパクトな状態表現を学習することが有益(Watter et al., 2015; Karl et al., 2016; Ha and Schmidhuber, 2018)
これらのモデルはlatent dynamics modelsと呼ばれる

潜在空間で先を予測することは、長期的な予測を容易にし、画像生成することなく数千のコンパクトな状態のシーケンスを一括して予測でき効率的
V2ではV1のガウス潜在変数をカテゴリ変数に置き換えた

tmyodatmyoda

Experience dataset
世界モデルはエージェントの過去の経験(images, actions, rewards, discount factors)から学習される
\gamma = 0.999固定、終端は0
学習には、エピソードからランダムにサンプルされた、固定長L=50をバッチ
学習でエピソード終わりを十分に観測するために、ランダムサンプルされた開始インデックスからシーケンス(50)がはみ出すときにはクリッピングする

tmyodatmyoda

Model components

世界モデルの構成(詳細はFig.2)

  • Image encoder
  • Recurrent state-space model (RSSM)
  • Predictor (for image, reward, discount factor)

RSSMは決定的recurrent state のシーケンス h_t を使用し、そこから、各ステップで確率的な状態に関する分布を2つ計算

RSSM: https://www.slideshare.net/DeepLearningJP2016/dllearning-latent-dynamics-for-planning-from-pixels

事後状態z_tは現在の画像x_tの状態も含む
事前状態\hat{z}_tは現在の画像にアクセスすることなく事後状態を予測

コンパクトなモデル状態は決定論的状態と確率的状態の結合

事後状態から、現在の画像の状態x_tを再構成し、報酬r_t, 割引係数\gamma_tを予測

すべてのコンポーネントはNNで実装
\phiはそれらを組み合わせたパラメータベクトル
遷移予測器は、次の状態を今の状態と行動から予測(画像の観測や生成しなくてもよい)
割引率予測器は、モデル予測から行動を学習する際に、エピソードが終了する確率を推定できる??

tmyodatmyoda

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と呼ぶ
事後のエントロピー増加より、正確な事前の学習を促す

tmyodatmyoda

BEHAVIOR LEARNING

actor-criticの学習するときはrepresentationは固定
1GPUで2500の潜在的なtrajectoryをシミュレート

Imagination MDP

世界モデルの潜在空間の行動を学習するために、imagination MDPを次のように定義

初期状態分布\hat{z}_0は世界モデル学習中に出会ったコンパクトなモデル状態の分布
transition は\hat{z}_{1:H}のコンパクトなモデル状態のシーケンスを出力(H=15)
rewardの平均値を報酬シーケンス\hat{r}_{1:H}
discountは割引率のシーケンスを出力\hat{\gamma}_{1:H}, down-weight rewardsに使用

エピソード終了の可能性を考慮して、予測された割引率でactor-ctiricのロスを重み付け


Fig2キャプション約
trajectoryはモデル学習時に計算された事後状態から始まり、actorから行動をサンプルして、次を予測
criticは各状態の将来の報酬の期待値の合計を予測
ctiricは予測された報酬からTD学習
actorは強化学習の勾配と、世界モデルのstraight-through勾配を使ってctiricの予測を最大化するように学習

tmyodatmyoda

Model components

imagenation MDPにおける長期的な行動を学習するため、行動を選択するactorと決定論的なctiricを学習する

actorはcriticの出力を最大化する行動を学習することを目指す
criticは想像上の各状態からactorが獲得する将来の報酬の合計を正確に予測することを目指す
actorのパラメータ\psi, ctiricは\xi

実際の環境とは対照的に、潜在的な状態はマルコフ性があるので、追加条件は必要ない
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)を使用

\lambda=0.95 \lambdaターゲットは異なるhorizonに対するn-step returnの加重平均であり、長いhorizonは指数関数的に加重が小さくなる
2乗損失で学習

Adamで最適化
最終ステップではtarget=ctiricなので損失項はなし
sg関数で勾配を止める
100stepのtaget networkで価値学習を安定化(DQNのやつ)(Mnih et al., 2015)

Actor loss
ctiricの学習で計算した\lambda-returnを最大化するように学習
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

tmyodatmyoda

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日かかったって

tmyodatmyoda

ATARI PERFORMANCE

4つの比較方法を検討

  • Gamer Median
    ランダムポリシーとプロゲーマーで正規化されて、複数シードで平均とり、すべてのタスクの中央値
    しかし半分のゲームのスコアが0でも値が変わらずロバストじゃない
  • Gamer Mean
    すべてのタスクを考慮しているが、ゲーマーが苦手(RLが得意)なゲームで高い正規化されたスコア出せるので、平均で人間超えは余裕
  • Record Mean
    プロゲーマーではなく人間の世界記録で正規化するとたしかに改善するが、やはりRLが得意なゲームが支配的
  • Chipped Record Mean
    人間の世界記録でクリップすることで、1を超えないようにする
    これでRLのイカれた記録がででも堅牢

rainbowがmedianでは良いスコアだけど、meansだとIQNに負ける

scheduledはアニーリング

tmyodatmyoda

Individual games

表K.1にスコア一覧
世界記録がないゲーム2つあったけど適当な値いれた

図E.1はそれぞれのゲームについて各アルゴリズムと比較
Video Pinaballが全負けした
理由としては、ボールが1pixelしかなかったから、世界モデルで再構成できなかったからと考察

tmyodatmyoda

ABLATION STUDY

要素の検証

  1. カテゴリカル潜在変数 vs ガウシアン潜在変数(V1)
  2. KL balancing
  3. 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があるためポリシーの最適化には適さないからと考察

tmyodatmyoda

DIscussion

DreamerV2はV1にカテゴリカル潜在変数とkl balancingを導入したもの
報酬によらず画像から表現学習できている