MASTERING ATARI WITH DISCRETE WORLD MODELS (DreamerV2)

Copyright © 2020 調和系工学研究室 - 北海道大学大学院情報科学研究院情報理工学部門複合情報工学分野 – All rights reserved.
DL ゼミ
MASTERING ATARI WITH DISCRETE
WORLD MODELS (DreamerV2)
北海道大学大学院情報科学研究院
情報理工学部門複合情報工学分野調和系工学研究室
修士課程2年織田智矢

2
メタ情報
著者
• Danijar Hafner (Google Research)
• Timothy Lillicrap(DeepMind)
• Mohammad Norouzi (Google Research)
• Jimmy Ba(University of Toronto)
発表
• ICLR 2021
概要
• World Modelの派生系のDreamerの2代目
• 画像入力から学習した世界モデルの潜在空間内のみで学習
• 同じ計算資源・サンプル数でIQN, Rainbow(モデルフリー)を凌駕
リンク
• https://arxiv.org/pdf/2010.02193.pdf
• https://ai.googleblog.com/2021/02/mastering-atari-with-discrete-world.html

3
イントロ
• 世界モデルとは、環境に関するエージェントの知識
を明示的に表現する方法
• 入力よりコンパクトな内部表現に変換するため、環
境を一般的に捉えることができていると考えられる
• この学習された世界モデル上で強化学習を行う
• メモリ使用量が抑えられ、1GPUで数千の軌道が並列
予測ができる
• アルゴリズムを厳密に比較するためのAtariスコアリ
ング手法も提案

4
MuZeroの成功
• 潜在変数を使ったプランニング（MuZeroはMCTS）
が、囲碁、将棋、Atari等のゲームに有効であること
をMuZeroが示した
• しかし重たすぎて実用的じゃない（2ヶ月学習、大量
GPU）
• 学習した世界モデル内で方策を学習したDreamerの
改良手法であるDreamerV2を提案
– 潜在変数をカテゴリカルに
– KL Balancingを提案

5
生成モデルとしての世界モデル
• 観測される刺激が環境から生成されると仮定
• その観測を生成している環境を近似するように、世
界モデルを学習
環境
世界モデル
近似
観測
生成
推論
潜在
変数

6
生成モデル
• データセットは、とある生成器によって生成されてい
ると考える
• 真の分布𝑝𝑑𝑎𝑡𝑎 𝑥 はわからないので、生成モデル 𝑝𝜃 𝑥
を設計して（𝜃はパラメータ）、真の分布と同じモデル
を構築することを目標とする
https://qiita.com/xkumiyu/items/1cc0223486c560062e00
生成器観測されたデータセット
生成
{𝑥1, … , 𝑥𝑁}
𝑝𝑑𝑎𝑡𝑎 𝑥

7
経験分布
• 真の分布となんらかの距離を測って、生成モデルを
近づけたいが真の分布は手に入らない
• 真の分布の代わりにデータセットで決まる経験分布
を定義
• 後はKLダイバージェンス最小化でパラメータ𝜃に関し
て最尤推定などいろいろ（だいぶはしょりました）

8
識別モデルとの違い
• データ集合 𝐱, 𝑦 𝑖=1
𝑁
（𝐱は入力，𝑦はラベル）から
𝑝 𝑦 𝐱 を求める
• イメージ
識別モデル：
 𝑝(𝑦|𝐱)を直接データ集合から学習
 学習した𝑝(𝑦|𝐱) を用いて𝐱から𝑦 を予測
生成モデル：
 同時分布𝑝(𝐱, 𝑦)をデータ集合から学習
 𝑝 𝑦 𝐱 =
𝑝 𝐱,𝑦
𝑝 𝐱,𝑦 𝑑𝑦
を求めて，𝐱から𝑦 を予測
識別モデル生成モデル
𝑝 𝑦 𝐱 𝑝 𝑦 𝐱 =
𝑝 𝐱, 𝑦
𝑝 𝐱, 𝑦 𝑑𝑦

9
潜在変数モデル
• 観測データとして現れないデータの背後に隠れている要
因として考える
• 実は投げるコインはk個あると仮定
• 毎回の試行でk個のうちいずれかのコインが選ばれて投
げられる
• K個のコインのうちどのコインかを表す潜在変数𝑧を導入
問題
表裏のある歪んだコインを複数回投げた結果が以下のと
き
結果:（表、裏、裏、…）
このコインはどのくらいの確率で表がでる？
Q
A

10
グラフィカルモデル
• コイン選択はカテゴリカル分布
• コインの歪みはベルヌーイ分布
• 後は同時分布を周辺化してパラメータ推定（EMアルゴリズム）
𝐱𝑖
𝐱𝑖~𝑝𝝁 (𝐱|𝐳𝑖)
𝐳𝑖
𝐳𝑖 ~𝑝𝛑(𝐳)
𝝅
𝑁
パラメータ
（コインを選択）
潜在変数
𝝁
観測変数
パラメータ
（コインの歪み）
• k番目が選ばれる確率𝜋𝑘 𝝅 = [𝜋1, … , 𝜋𝑘]𝑇
• 𝜇𝑘はk番目のコインが表(𝑥 = 1)となる確率 𝝁 = [𝜇1, … , 𝜇𝑘]𝑇

11
対数尤度の式変形
• 潜在変数に関して周辺化（𝜋, 𝜇をまとめて𝜃とすると）
• そのままだと周辺尤度を最大化できない（2つパラメー
タがある）
• 対数尤度は潜在変数𝐳における任意の分布𝑞(𝐳)を用いて
log 𝑝𝜃 𝐱 = log 𝑝𝜃 𝐱, 𝐳 𝑑𝐳 = log 𝑝𝜃 𝐱|𝐳 𝑝(𝐳)𝑑𝐳
= log𝑝𝜃 𝐱 𝑞(𝐳) 𝑑𝐳
= 𝑞(𝐳)log 𝑝𝜃 𝐱 𝑑𝐳
= 𝑞(𝐳)log
𝑞(𝐳)𝑝𝜃 𝐳|𝐱 𝑝𝜃 𝐱
𝑞(𝐳)𝑝𝜃 𝐳|𝐱
𝑑𝐳
= 𝑞(𝐳)log
𝑝𝜃 𝐳|𝐱 𝑝𝜃 𝐱
𝑞(𝐳)
𝑑𝐳 + 𝑞 𝐳 log
𝑞 𝐳
𝑝𝜃 𝐳|𝐱
𝑑𝐳
= ℒ(𝑞, 𝜃; 𝐱) + 𝐷𝐾𝐿[𝑞(𝐳)||𝑝𝜃(𝐳|𝐱)]
log𝑝𝜃 𝐱
ただし 𝑞 𝐳 log
𝑝𝜃 𝐱,𝐳
𝑞 𝐳
𝑑𝐳 ≡ ℒ(𝑞, 𝜃; 𝐱) とした

12
エビデンス下界（ELBO）
• 前ページより
• 第1項はエビデンス下界（ELBO）と呼ばれる
– 第2項は0以上なので、ELBOは必ず周辺尤度以下になるため
• このELBO代わりに最大化すれば良い
EMアルゴリズム
1. 𝑞について最大化（Eステップ）
• θを固定すると第2項が0となるとき最大
• 𝑞(𝐳) = 𝑝𝜃(𝐳|𝐱)とすれば良い
2. 𝜃について最大化（Mステップ）
log 𝑝𝜃 𝐱 = ℒ(𝑞, 𝜃; 𝐱) + 𝐷𝐾𝐿[𝑞(𝐳)||𝑝𝜃(𝐳|𝐱)]

13
深層生成モデル
• 従来の生成モデルはデータを生成することは意図し
てなかった
• 事後分布𝑝𝜽 (𝐳|𝐱)を計算し、𝑥から事後分布を計算し
てクラス分類等に使うのが主な用途（例: LDA）
• DNNの研究がすすんで、DNNで確率分布を表現でき
るようになった
• 従来の生成モデルとは違い、誤差逆伝播法によって
勾配を求めて学習する

14
Amortized（償却）変分推論
• 生成モデルの事後分布（推論分布、推論モデル）
• しかし生成モデルが複雑な場合（DNNとか）でモデ
ル化されている場合、分母の周辺化の計算が解析的に
できない
• 𝜙（学習可能なパラメータ）を使った𝑞𝜙(𝐳|𝐱)を導入
• するとELBOは𝜙と𝜃の関数になる
• 𝜙について最大化することで推論が求まる
• ELBO最大化は𝜙と𝜃について最大化
𝑝𝜃 𝐳 𝐱 =
𝑝𝜃 𝐱, 𝐳
𝑝𝜃 𝐱, 𝐳 𝑑𝐳
ℒ 𝜙, 𝜃; 𝐱 ≡ 𝑞𝜙(𝐳|𝐱)log
𝑝𝜃 𝐱, 𝐳
𝑞𝜙(𝐳|𝐱)
𝑑𝐳

15
Variational Autoencoder (VAE)
• 潜在変数モデルの確率分布をDNNで表現
• ELBO最大化によって学習（ELBOを目的関数に）
• 第1項は解析的に求まらないので近似する
ℒ 𝜃, 𝜙; 𝐱 = 𝔼𝑞𝜙 𝐳 𝐱 log 𝑝𝜃 𝐱|𝐳 − 𝐷𝐾𝐿[𝑞𝜙 𝐳 𝐱 ∥ 𝑝 𝐳 ]
解析的にもとまる
𝔼𝑞𝜙 𝐳 𝐱 log 𝑝𝜃 𝐱|𝐳 ≅
1
𝐿
𝑙=1
𝐿
log 𝑝𝜃 𝐱|𝐳(𝑙)
, 𝐳(𝑙)
∼ 𝑞𝜙 𝐳 𝐱
𝐱
𝝁 𝝈
𝝀
勾配が通らない
• しかし入力までの計算グラフに確率的
な仮定が入り、自動微分できない
• 再パラメータトリックを使う
（通称リパラ）

16
再パラメータ化トリック
• 推論分布を次のように再パラメータ化
• よく知られている分布なら大体できる
• ガウス分布の場合
• するとELBOの第1項の勾配は
http://blog.shakirm.com/2015/10/machine-learning-trick-of-the-day-4-reparameterisation-tricks/
𝐳 = 𝝁 + 𝝈 ⊙ 𝜺，ただし𝜺~𝒩(0, 𝐈)
𝛻𝜃,𝜙𝔼𝑞𝜙 𝐳 𝐱 log 𝑝𝜃 𝐱|𝐳 = 𝔼𝜺∼𝒩(0,𝐈) 𝛻𝜃,𝜙log 𝑝𝜃 𝐱|𝐳 = 𝝁 + 𝝈 ⊙ 𝜺
≅
1
𝐿
𝑙=1
𝐿
𝛻𝜃,𝜙log 𝑝𝜃 𝐱|𝐳 = 𝝁 + 𝝈 ⊙ 𝜺(𝑙)
, 𝜺(𝑙)
∼ 𝒩(0, 𝐈)
𝜺
𝑞𝜙 𝐳 𝐱 𝑝𝜃 𝐱 𝐳
𝐱
𝝁 𝝈
𝝀
リパラ

17
VAEのELBO（目的関数）
ℒ 𝜃, 𝜙; 𝐱
= 𝔼𝑞𝜙 𝐳 𝐱 log 𝑝𝜃 𝐱|𝐳 − 𝐷𝐾𝐿[𝑞𝜙 𝐳 𝐱 ∥ 𝑝 𝐳 ]
=
1
𝐿
𝑙=1
𝐿
log 𝑝𝜃 𝐱|𝐳(𝑙)
− 𝐷𝐾𝐿[𝑞𝜙 𝐳 𝐱 ∥ 𝑝 𝐳 ]
=
1
𝐿
𝑙=1
𝐿
(𝐱 log 𝝀 𝑙
+ 1 − 𝐱 log(1− 𝝀 𝑙
)) +
1
2
𝑗=1
𝐽
(1 + log 𝜎𝑗
2
− 𝜇𝑗
2
− 𝜎𝑗
2
)
負の再構成誤差
（生成モデルがベルヌーイ分布の場合）
正則化項
（事前分布𝑝 𝐳 が標準ガウス分布𝒩(0, 𝐈) の場合）
ただし， 𝝁 = 𝑔𝜙
𝜇
𝐱 , 𝝈 = 𝑔𝜙
𝜎
𝐱 ,
𝝀 𝑙 = 𝑓𝜃 𝐳 𝑙 , 𝐳 𝑙 = 𝝁 + 𝝈 ⊙ 𝜺 𝑙 , 𝜺 𝑙 ~𝒩 0, 𝐈
𝜺
𝑞𝜙 𝐳 𝐱
𝑝𝜃 𝐱 𝐳
𝐱
𝝁 𝝈
𝝀
• VAEでは推論モデルで入力𝐱を𝐳にエンコードし，生成モデルで𝐳から
𝐱をデコード
→オートエンコーダーとみなせる
𝐳
推論モデル 𝑞𝜙(𝐳|𝐱) 生成モデル 𝑝𝜃(𝐱|𝐳)
𝐱 𝐱
再構成
入力

18
VAEと表現学習
• VAEでは再構成だけでなく、表現𝐳~𝑞𝜙(𝐳|𝐱)も学習し
ているとみなせる
• 深層生成モデルにおいては，表現学習は推論と等価
– エンコーダで入力から表現へ写像
表現学習？
• データから「良い表現」を（できれば教師なしで）
獲得する学習
良い表現？
• 元のデータの性質をある程度保持しつつ，他のタス
クにも使い回せるような表現

19
強化学習主要アルゴリズム系譜
• あ
出典: https://twitter.com/ImAI_Eruel/status/1303677795806056451/photo/1

20
強化学習の定式化
• 環境にマルコフ性を仮定し，マルコフ決定過程で定式化
• マルコフ性
- 現在の状態𝑠tと行動𝑎tのみに依存して𝑡 + 1の状態𝑠t+1が決ま
る性質
- 次の状態が現在の状態にのみ依存する
https://ja.wikipedia.org/wiki/マルコフ決定過程

21
目的関数
• エージェントが環境内で得る報酬の総和を最大化する方策
を獲得することが，強化学習の目的
• 無限ステップの行動で発散するため，割引率𝛾を導入
• 収益に関連した目的関数を最大化する方策を学習する問題
に帰着
• 以下価値関数を、方策の良さを示す目的関数とする場合が
多い
𝑅𝑡 = 𝑟𝑡+1 + 𝛾𝑟𝑡+2 + 𝛾2𝑟𝑡+3 … =
𝑘=0
∞
𝛾𝑘𝑟𝑡+𝑘+1
V𝜋 𝑠 = 𝔼 𝑅𝑡+1|𝑆𝑡 = 𝑠

22
価値関数と方策
• 状態価値関数: V
- 状態𝑠から方策𝜋に従って行動したときに得る収益の期待値
• 行動価値関数: 𝑄
- ある状態𝑠で行動𝑎を選択して，その後方策𝜋に従って行動し
たときに得る収益の期待値
• 方策:𝜋(𝑎|s)
– ある状態𝑠で行動𝑎を選択する確率を出力する関数
– 他にGreedy, ε-greedy, ソフトマックス方策, ボルツマン方策が
ある
V𝜋 𝑠 = 𝔼 𝑅𝑡+1|𝑆𝑡 = 𝑠
𝑄𝜋 = 𝔼 𝑅𝑡+1|𝑆𝑡 = 𝑠, 𝐴𝑡 = 𝑎

23
ベルマン方程式
• 価値関数が満たすべき式
• 以下の再帰的な式をベルマン方程式という
– ベルマン方程式は価値関数の定義から再帰的に導かれる
https://www.slideshare.net/ShotaImai3/rlssdeepreinforcementlearning から抜粋

24
基本性質
• 探索と活用のジレンマ
– 現状最も良いと思われる行動を続けるだけでは，他にもっ
と良い行動を見逃す可能性がある
– 活用(Exploitation)，探索(Exploration)のバランスが大事
• 遅延報酬
– 多くの行動を繰り返した後でなければ，大きな報酬を与え
られないような環境もある
– 長期的な利益を考える必要→期待報酬の最大化
• 価値(Value)ベース or 方策(Policy)ベース
– 価値評価を学習するか，戦略を学習するか

25
方策 or 価値関数の学習
価値関数の学習
• ある状態・行動における価値を学習し、それに基づき（価値関数に
従って）行動
• ε-greedy, softmax方策等を使う
• 学習方法: DP, MC, TD学習
• 例: SARSA, Q-Learning, DQN
方策の学習
• 価値関数を参照せず直接方策を学習
• （しかし最近の手法は、状態価値関数を一緒に学習する）
• 行動生成: 確率分布からサンプリング
• 学習方法: 累積報酬和Rの期待値𝐽 𝜃 を勾配上昇法で最大化
• 例: VPG, TRPO, PPO

26
方策勾配定理（方策の学習）
• 方策の良さの指標を𝑓(∙)として目的関数𝐽 𝜃 =
𝐸𝜋𝜃
[𝑓(∙)]を最大化する𝜃を学習
• 累積報酬和の良さの指標にした方策の最適化を考え
る
θで微分して勾配上昇
http://rail.eecs.berkeley.edu/deeprlcourse/static/slides/lec-5.pdf

27
REINFORCE アルゴリズム
欠点
• 収益𝑟(𝜏)の分散によって方策が大きく変わってしまう
CS285 Fa19 9/16/19 https://youtu.be/Ds1trXd6pos?t=2020

28
ベースラインの導入
松尾研強化学習セミナーから抜粋

29
方策勾配計算の改善
エピソード終端
までの報酬和

30
Advantage関数の学習

31
状態価値関数の学習

32
ここまでまとめ

33
モデルフリーとモデルベース
https://www.slideshare.net/ShotaImai3/rlssdeepreinforcementlearning から抜粋

34
世界モデル
• Atariタスクの入力は高次元画像
• コンパクトな状態表現を学習す
る方が有益[1,2,3]
• これらのモデルはlatent
dynamics modelと呼ばれる
https://ai.googleblog.com/2021/02/mastering-atari-with-discrete-world.html
[1] M Watter, J Springenberg, J Boedecker, M Riedmiller. Embed to Control: A Locally Linear Latent Dynamics Model for Control
From Raw Images. Advances in Neural Information Processing Systems, 2015.
[2] M Karl, M Soelch, J Bayer, P van der Smagt. Deep Variational Bayes Filters: Unsupervised Learning of State Space Models
From Raw Data. ArXiv Preprint ArXiv:1605.06432, 2016.
[3] D Ha J Schmidhuber. World Models. ArXiv Preprint ArXiv:1803.10122, 2018.
• 潜在空間で先を予測することは、長期的な予測を容易に
する
• 一括して予測できるので効率的である

35
学習の流れ
データセット
• エージェントの過去の経験から学習
– 方策がある程度学習したら方策に従って環境とインタラク
ションしてデータセットを増やす（Atariでは4回）
• 学習時はエピソードからランダムにサンプルされた
固定長𝐿＝50のシーケンスで取り出し
– エピソード終端を十分に観測するため、エピソードはみ出
す場合はクロップする
DreamerV1より, https://ai.googleblog.com/2020/03/introducing-dreamer-scalable.html

36
世界モデルのコンポーネント
• Image predictor
• Recurrent state-space model (RSSM)
• Predictors (for image, reward, discount factor)
• エンコーダが各画像を確率的な表現に変換
• それを世界モデルの再帰的な状態に組み込む
→表現が確率的なため、必要な情報のみを抽出するようになり、
未知画像に対してロバストに
https://arxiv.org/pdf/1811.04551.pdf

37
コンポーネント詳細
• RSSMは決定論的なシーケンスh𝑡を使用
• Transition predictorで状態𝑧𝑡を更新
• そこから各ステップで確率的な状態に関する分布を
計算
• コンパクトな状態表現とはh𝑡と𝑧𝑡のconcatenation
• すべてのコンポーネントはNNで実装
– 画像のencoder-decoderはCNN
– recurrent modelはGRU
– あとはMLP (4 layer, 400 units)

38
世界モデルの学習
確率分布
image predictor: ガウス分布（ピクセル分)
reward predictor: 1変数ガウス分布
discount predictor: ベルヌーイ分布
transition predictor: カテゴリカル分布
誤差関数
• すべてのコンポーネントは同時に最適化される
カテゴリカル分布の勾配の通し方

39
KL Balancing
• 世界モデルはrepresentation modelが事後状態の近似、
transition predictorが時間的な事前状態である、sequential
VAEであると解釈できる
• ELBOではKL lossが2つの役割を果たす
– representationに対する事前状態の学習
– 事前状態に対するrepresentationの正則化
• transitionの学習は簡単ではないため、学習が不十分なと
きは事前状態に対して正則化してほしくない
• 事前状態に対する学習率を𝛼 = 0.8, 近似事後状態に対する
学習率を1 − 𝛼にする
• これにより正確な事前状態の学習を促す

40
方策の学習（Actor-Critic)
Imagination MDP
• 世界モデルの潜在空間の行動を学習するために、
imagination MDPを次のように定義
– 初期分布は世界モデルの学習中に出会ったもの
– transition, reward, discountはシーケンスで出力
Actor-Criticの学習はrepresentationのパラメータ固定

41
方策の学習（Actor-Critic)
ActorはCriticの出力を最大化するように学習
Criticはactorが獲得する将来の合計を予測するように学習
Critic loss
• n-step learning, general 𝜆 target, 100 step target network
• sg関数は最終ステップだった場合に勾配をストップ
Actor loss
• Reinforceの勾配と、predictorsたちのstraight-through勾配

42
実験
• 4つのモデルフリーアルゴリズムと比較
– IQN, Rainbow, C51, DQN
• Atariでよく比較される55ゲームで評価
• 1GPU, 1環境, 環境とのインタラクションは200Mまで
• 1episodeあたり108,000stepの制限
• モデルフリーアルゴリズムが200Mの中、世界モデル
内で468Bもコンパクトな状態から学習（1万倍）

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

44
実験結果
schedulesはActor lossの係数をアニーリングしたら結果
がよくなった
Video pingballというゲームはDreamerV2が全負け
理由として、ボールが1pixelしかなかったから世界モデ
ルで再構成できなかったためと考察

45
世界モデルの学習結果
• イメージデコーダーで画像を復元したとき、
どれだけ環境を再現できているか比較

カテゴリカル潜在変数(V2) vs ガウシアン潜在変数（V1)
• カテゴリカルが42ゲームが勝ち、8負け、5タイ
• なぜかはハッキリとはわかっていない
KL balancingなし
• 44ゲームが勝ち、6負け、5タイ
• 世界モデルの事前状態を先に正確に学習することは大切
• 世界モデル以外の確率モデルにも有効だと考えられる
Image gradient stop
• 3勝ち, 51負け, 1タイ
reward gradient stop
• 15勝ち, 22負け, 18タイ
reinforce gradient stop
• 18勝ち, 24負け, 13タイ
46
各要素の検証(Ablation Study)

47
まとめ
• モデルベース手法である世界モデルでモデル
フリーの手法を超えたDreamerV2を提案
• 潜在変数をガウス分布ではなく、カテゴリカ
ル分布にすることで性能が向上
• KL Balancingを提案
• Atariのベンチマーク方法を提案
• 1GPUで数千の軌道が並列予測が可能

MASTERING ATARI WITH DISCRETE WORLD MODELS (DreamerV2)

Recommended

More Related Content

What's hot (20)

Similar to MASTERING ATARI WITH DISCRETE WORLD MODELS (DreamerV2) (20)

More from harmonylab (20)

MASTERING ATARI WITH DISCRETE WORLD MODELS (DreamerV2)