動画の状態推移を予測する強化学習の論文を読んでみる

動画から系列型のVAE方式による状態推移を予測する強化学習として、優秀な学習効率を達成したモデルの論文である。

[1811.04551] Learning Latent Dynamics for Planning from Pixels

系列型VAEの強化学習モデルWorld Modelを最初に発表したDavid Haも著者になっているので、このモデルはWorld Modelの発展型もしれない。

この成果については下記のICMLの口答発表がある。

Learning Latent Dynamics for Planning from Pixels

GitHubにも著者らの実装と動作した動画が揚がっている

GitHub - google-research/planet: Learning Latent Dynamics for Planning from Pixels

この強化学習ではエージェントが一部しか観察 $o_t$ できないPOMDP(Partial Observable Markov Dicision Process)となっていて、次の分布のパラメータ推定をELBO(evidence lower bound)の最大化で推定している。従ってこれらの分布は全て多変量ガウシアンである。

推移モデル $p(s_t|s_{t-1},a_{t-1})$

観察モデル $p(o_t | s_t)$

報酬モデル $p(r_t|s_t)$

Encoder $p(s_t| o_{\leq t},a_{\leq t})$ 　状態 $s_t$ は潜在変数となっている

摂動 $p(\epsilon)$

論文にある全体のアルゴリズムとしては潜在変数の最適化とサンプリングの組合せで以下となっている。

アルゴリズム1

以下を収束まで繰り返す

　　次の潜在変数の最適化を繰り返す

　　　保存された行動軌跡図からバッチ分 $D$ をランダムに取り出す

　　　ELBO(Equation3)で $\mathcal{L}(\theta)$ を解き事後分布の状態潜在変数のパラメータ $\theta$ を求める

　　パラメータを更新する $\theta \leftarrow \theta - \alpha\nabla_\theta\mathcal{L}(\theta)$

サンプリングを繰り返す

　　　状態潜在変数をEncoder $q(s_t|o_{\leq t},a_{\leq t})$ で推定する

　　　推定された状態に行動計画 $planner$ より最適行動を求め、これに摂動を与える

$a_t \leftarrow planner(q(s_t|o_{\leq t},a_{\leq t}),p)$ アルゴリズム2参照

　　　行動をR回繰り返し、

　　　　　報酬 $r_t$ と観察図 $o_t$ を得る。

　　　　繰返した観察図より軌跡図として保存する

アルゴリズム１中にある $s_t$ の推定の理論は(Equation3)として以下に手案されている。（式の展開は付録8式参照)

左辺の対数尤度を右辺では潜在変数と $s_t$ を介在させて $q(s_t| o_{\leq t},a_{\leq t})$ が一致する様にELBOを解いている。

f:id:mabonki0725:20200217073926p:plain

しかし長い系列のELBOは容易でないので、RNNを採用したRMMSで実装している。

著者らは下図の状態遷移SMMとRNNを合成したRSSM（Recurent State Space Model)を提案している。即ち非線形の状態遷移では誤差が大きく再帰的なRNNの経路を追加して安定化を図っている。ここでSSM（統計的過程）とRNNを（決定的過程）を合成したと言っている。

f:id:mabonki0725:20200216211809p:plain

よってRNNの隠れ変数 $h_t$ を使って以下で推定している。

　　 f:id:mabonki0725:20200216211911p:plain

上記を使って便宜的に変分 $q$ を以下に簡略している。

$q(s_t| o_{\leq t},a_{\leq t})=\Pi_{t+1}^T q(s_t|h_t,o_t)$

さらにEquation3での１step毎の推論をｄステップ毎の推論（Latent Overshooting)に拡張している。

f:id:mabonki0725:20200217081424p:plain

下図はlatent Overshootingのd=3の実装図である。左端が１step毎で右端が3step毎の予測となっている。

f:id:mabonki0725:20200217081200p:plain

・行動計画 $(planner)$ 行動分布の推定

　探索型の価値関数を使わず今まで良い報酬が得られた行動の分布 $q(a)$ とd-step先の予測から、最適な行動分布を求めている。これについては付録にアルゴリズム２として次の様に記述されているが、要は報酬が良かった行動を集め、このガウス分布を求めているだけである。

　行動軌跡を記録する

良い報酬が得られた行動のガウス分布 $q$ から行動 $a$ をサンプリング

　　　 $a_t \sim q(a_{t:t+H})$

　　d-step先の予測

　　　 $s_{t:t+H+1} \sim q(s_t| o_{\leq t},a_{\leq t}) \Pi _{\tau=t+1}^{t+H+1} p(s_\tau | s_{\tau-1} , a_{\tau-1})$

d-step迄の報酬

$R=\sum_{\tau=t+1}^{t+H+1} \mathbb{E}(p(r_\tau | s_\tau)$

行動軌跡から報酬が高いK個のベストの行動軌跡を取り出す。

ガウス行動推分布 $q(a_{t:t+H})$ の平均と分散を計算する

　　　平均 $\mu_{t:t+H} = \frac{1}{K} \sum_{k \in K} a_{t:t+H}$

分散 $\sigma_{t:t+H} = \frac{1}{K-1} \sum_{k \in K} |a:{t:t+H}-\mu_{t:+H}|$

　　　 $q(a_{t:t+H}) \leftarrow \mathcal{N}(\mu_{t:t+H},\sigma_{t:t+H})$

・実験結果

　実験対象としては行動が実数領域で制御できる下図のDeepMindの６個の課題としている。

f:id:mabonki0725:20200217081841p:plain

実験結果は何れの課題でも200個程度の学習で10万回学習したD4PGのレベルに達しており学習効率が優れていることを示している。

f:id:mabonki0725:20200217082529p:plain

注)DP4Gは実数領域のActor-Critic版のDDPG(Deep Determistic Policy Gradient)を改良したもの

・感想

　このモデルの特質は動画より状況を予測して効率のよい強化学習を達成していることである。

　しかしモデルはVAEのELBOで解く為、潜在変数も行動分布も全てガウシアンになっている。この様な分布が適用できるのは状態推移が一様な課題しか適用できず、ゴールがある様なゲームは容易でないと思われる。しかし実務上の課題は最適な一様状態に達する工夫であり、この様な状態に移行する学習には適用できると考えられる。

　　