Google Brain，Deepmind 提出的 PlaNet
論文位址：https://danijar.com/publications/2019-planet.pdf
源碼位址：https://github.com/google-research/planet
摘要
planning（規劃）對已知環境轉移動态的控制任務非常成功，如需在未知環境中使用 planning，則 agent 需要通過環境互動學習轉移動态。然而，學到足夠精确支援 planning 的轉移動态模型是一項長期挑戰，特别是在圖像領域。
本文提出 Deep Planning Network (PlaNet)，是一個model-based agent，能夠：
- 從圖像中學習環境轉移動态；
- 通過在隐空間進行快速線上 planning 來選擇動作。
要實作高性能，轉移動态模型必須能精确預測多個時間步之後的 reward。為達成這一目标，本文使用了一個隐動态模型，包含确定性和随機性元件。同時，本文還提出了一個多步變分推斷目标，記為 latent overshooting.
本文所提出的 agent 僅使用像素觀察，來解決部分可觀測，稀疏 reward 的控制問題，超出了以往 model-based方法解決問題的難度。其表現與 model-free 方法相當，而樣本量更少。
背景介紹
學習轉移動态模型最主要的困難包括：模型的準确性，多步預測累積的誤差，未能捕獲未來豐富的可能性，以及對訓練分布的過拟合造成的泛化失敗。
使用學習得到的模型進行 planning，相比 model-free 方法有以下幾點好處：
- model-based planning 方法更加 data-efficient（資料依賴小），因為它對訓練信号的使用更充分，并且不需要通過貝爾曼方程傳遞 reward；
- 以增加對 action 的搜尋為計算代價，planning 提高了agent 的表現
- 學到的模型獨立于特定任務，有潛力通過轉化服務于相同/似環境下的其他任務。
近期的工作已經展示了在簡單低維情形下學習轉移動态的優勢，但這些工作都假定已知環境的底層狀态和明确的reward函數，這在現實應用中并不可行。在高維情形下，我們希望在一個緊湊的隐空間學習轉移動态以保證快速 planning。這種隐空間模型已經在簡單任務上取得了成功。
一句話總結：PlaNet 發展了model-based planning 方法，解決了更困難的基于像素的連續控制問題。
主要貢獻：潛空間planning，循環狀态空間模型，latent overshooting（給出多步預測變分下界，改善多步預測結果）.
潛空間planning
Framework
假設已有學到的轉移動态模型，算法架構如下：
由于單一圖像無法反映環境狀态的全部資訊，将問題模組化為一個POMDP（部分可觀測 markov 決策過程）， s t s_t st為底層狀态， o t o_t ot為環境觀測， a t a_t at為連續動作向量。
注意到PlaNet所有的模組化都隻針對環境，而後在環境中進行規劃，并沒有政策網絡/值網絡。同時，在互動過程中，每一步都重新做planning，以保證能夠根據新的觀測作出調整（與MPC類似）。
Planning 算法
Planning算法使用了CEM (cross entropy method), CEM是一種基于群體的優化方法，它能夠推斷出動作序列的最優分布。
模型
本文提出一個循環狀态空間模型 (RSSM) 在隐空間進行前向預測，這一模型可被看作非線性Kalman濾波器或者序列VAE.
SSM
潛狀态空間模型如下圖b所示，下面所涉及的模型都是由神經網絡參數化的高斯分布。
由于模型都是非線性的，是以無法直接計算出參數學習所需的狀态後驗。是以，本文使用編碼器 q ( s 1 : T ∣ o 1 : T , a 1 : T ) = ∏ t = 1 T q ( s t ∣ s t − 1 , a t − 1 , o t ) q\left(s_{1 : T} | o_{1 : T}, a_{1 : T} \right)=\prod_{t=1}^{T} q\left(s_{t} | s_{t-1}, a_{t-1}, o_{t}\right) q(s1:T∣o1:T,a1:T)=∏t=1Tq(st∣st−1,at−1,ot)來從過去的觀測和動作中近似推斷狀态後驗，其中 q ( s t ∣ s t − 1 , a t − 1 , o t ) q\left(s_{t} | s_{t-1}, a_{t-1}, o_{t}\right) q(st∣st−1,at−1,ot)是一個由“卷積網絡+前向網絡”參數化的對角高斯分布。
為使用編碼器，本文建立了一個基于資料對數似然的變分下界 (Jensen 不等式)：
ln p ( o 1 : T ∣ a 1 : T ) ≜ ln ∫ ∏ t p ( s t ∣ s t − 1 , a t − 1 ) p ( o t ∣ s t ) d s 1 : T ≥ ∑ t = 1 T ( E q ( s t ∣ o ≤ t , a < t ) [ ln p ( o t ∣ s t ) ] ⎵ reconstruction − E q ( s t − 1 ∣ o ≤ t − 1 , a < t − 1 ) [ K L [ q ( s t ∣ o ≤ t , a < t ) ∥ p ( s t ∣ s t − 1 , a t − 1 ) ] ] ⎵ complexity ) . \ln p\left(o_{1 : T} | a_{1 : T}\right) \triangleq \ln \int \prod_{t} p\left(s_{t} | s_{t-1}, a_{t-1}\right) p\left(o_{t} | s_{t}\right) d s_{1 : T}\geq \sum_{t=1}^{T}\left(\underbrace{\mathrm{E}_{q\left(s_{t} | o _{\leq t}, a_{<t}\right)}\left[\ln p\left(o_{t} | s_{t}\right)\right]}_{\text { reconstruction }} -\underbrace{\mathrm{E}_{q\left(s_{t-1} | o _{\leq t-1}, a_{<t-1}\right)}\left[\mathrm{KL}\left[q\left(s_{t} | o_{ \leq t}, a_{<t}\right) \| p\left(s_{t} | s_{t-1}, a_{t-1}\right)\right]\right]}_{\text { complexity }}\right). lnp(o1:T∣a1:T)≜ln∫∏tp(st∣st−1,at−1)p(ot∣st)ds1:T≥∑t=1T⎝⎜⎛ reconstruction
Eq(st∣o≤t,a<t)[lnp(ot∣st)]− complexity
Eq(st−1∣o≤t−1,a<t−1)[KL[q(st∣o≤t,a<t)∥p(st∣st−1,at−1)]]⎠⎟⎞.
推導比較簡單，包含重要性采樣和Jensen不等式的使用兩個步驟，包含在附錄中，關于變分下界的推導也可參考變分自編碼器VAE(Variational Autoencoders)及示例代碼。使用重參數化技巧保證了随機梯度下降可以順利使用。
RSSM
盡管随機模型具有很強的泛化能力，然而，純粹的随機模型難以記住多個時間步驟所包含的資訊。理論上，某些情況下政策方差可以設定為0（确定性政策），然而上述優化過程難以收斂到這類結果。這促使本文引入了一個确定性激活向量序列 { h t } t = 1 T \left\{h_{t}\right\}_{t=1}^{T} {ht}t=1T 允許模型确定地通路所有先前的狀态。于是就有了本文所使用的模型 RSSM，如上圖c所示。
其中 f ( h t − 1 , s t − 1 , a t − 1 ) f\left(h_{t-1}, s_{t-1}, a_{t-1}\right) f(ht−1,st−1,at−1)是一個RNN。直覺地，我們可以了解為該模型将狀态劃分為一個随即部分 s t s_t st和一個确定狀态 h t h_t ht，它們由之前時間步的随即部分和确定部分決定。本文使用編碼器 q ( s 1 : T ∣ o 1 : T , a 1 : T ) = ∏ t = 1 T q ( s t ∣ h t , o t ) q\left(s_{1 : T} | o_{1 : T}, a_{1 : T}\right)=\prod_{t=1}^{T} q\left(s_{t} | h_{t}, o_{t}\right) q(s1:T∣o1:T,a1:T)=∏t=1Tq(st∣ht,ot)來參數化近似狀态後驗。
Latent Overshooting
Motivation
在上一部分，本文建立了在隐序列空間進行學習和推斷的标準變分下界，包含一個針對觀測的重構項和一個針對近似後驗的KL散度正則項。這個變分推斷目标的一個缺陷是：轉移機率函數 p ( s t ∣ s t − 1 , a t − 1 ) p\left(s_{t} | s_{t-1}, a_{t-1}\right) p(st∣st−1,at−1)的随機部分僅被KL正則中的單步預測所訓練，梯度經過 p ( s t ∣ s t − 1 , a t − 1 ) p\left(s_{t} | s_{t-1}, a_{t-1}\right) p(st∣st−1,at−1)後直接進入 q ( s t − 1 ) q(s_{t-1}) q(st−1)，而不進入 p ( s t − 1 ∣ s t − 2 , a t − 2 ) p\left(s_{t-1} | s_{t-2}, a_{t-2}\right) p(st−1∣st−2,at−2)進而周遊多個 p ( s t i ) p(s_{t_i}) p(sti). 在這一部分，我們将标準變分下界推廣到 latent overshooting，可在潛空間訓練所有的多步預測。
事實上，如果模型能夠做完美的單步預測，那麼它也能做出完美的多步預測。然而，通常使用的模型都容量有限并且限制在确定的分布族上，将模型單步訓練至收斂，與多步預測所需的最佳模型并不一緻。成功規劃需要精确的多步預測，并在所有的多步預測資料上訓練模型。本文把這一想法推廣到潛序列模型，表明多步預測也可以被潛空間上的loss來優化，而不需要添加額外的圖像。
固定距離 d d d的多步預測
我們開始把上面訓練單步預測的标準變分下界推廣到固定距離 d d d的多步預測情形。為省略起見，省略了 s t s_t st關于 a t a_t at的先驗。則：
p ( s t ∣ s t − d ) ≜ ∫ ∏ τ = t − d + 1 t p ( s τ ∣ s τ − 1 ) d s t − d + 1 : t − 1 = E p ( s t − 1 ∣ s t − d ) [ p ( s t ∣ s t − 1 ) ] \begin{aligned} p\left(s_{t} | s_{t-d}\right) & \triangleq \int \prod_{\tau=t-d+1}^{t} p\left(s_{\tau} | s_{\tau-1}\right) d s_{t-d+1 : t-1} \\ &=\mathrm{E}_{p\left(s_{t-1} | s_{t-d}\right)}\left[p\left(s_{t} | s_{t-1}\right)\right] \end{aligned} p(st∣st−d)≜∫τ=t−d+1∏tp(sτ∣sτ−1)dst−d+1:t−1=Ep(st−1∣st−d)[p(st∣st−1)]
于是，關于多步預測分布 p d p_{d} pd 的變分下界為：
ln p d ( o 1 : T ) ≜ ln ∫ ∏ t = 1 T p ( s t ∣ s t − d ) p ( o t ∣ s t ) d s 1 : T ≥ ∑ t = 1 T ( E q ( s t ∣ o ≤ t ) [ ln p ( o t ∣ s t ) ] ⎵ reconstruction − E p ( s t − 1 ∣ s t − d ) q ( s t − d ∣ o ≤ t − d ) [ K L [ q ( s t ∣ o ≤ t ) ∥ p ( s t ∣ s t − 1 ) ] ] ⎵ multi-step prediction ) \ln p_{d}\left(o_{1 : T}\right) \triangleq \ln \int \prod_{t=1}^{T} p\left(s_{t} | s_{t-d}\right) p\left(o_{t} | s_{t}\right) d s_{1 : T}\geq \sum_{t=1}^{T}\left(\underbrace{\mathrm{E}_{q\left(s_{t} | o _{\leq t}\right)}\left[\ln p\left(o_{t} | s_{t}\right)\right]}_{\text { reconstruction }}-\underbrace{\mathrm{E}_{p\left(s_{t-1} | s_{t-d}\right) q\left(s_{t-d} | o_{ \leq t-d}\right)}\left[\mathrm{KL}\left[q\left(s_{t} | o_{ \leq t}\right) \| p\left(s_{t} | s_{t-1}\right)\right]\right]}_{\text { multi-step prediction }}\right) lnpd(o1:T)≜ln∫∏t=1Tp(st∣st−d)p(ot∣st)ds1:T≥∑t=1T⎝⎜⎛ reconstruction
Eq(st∣o≤t)[lnp(ot∣st)]− multi-step prediction
Ep(st−1∣st−d)q(st−d∣o≤t−d)[KL[q(st∣o≤t)∥p(st∣st−1)]]⎠⎟⎞
推導和上一部分類似，最大化這一目标就能使多步預測的準确性得到訓練。這一反應了在多步預測中，模型無需通路先前的所有觀測，僅在潛空間就可進行預測。
Latent overshooting：可變距離多步預測
Latent overshooting就是再把多步預測變分下界從固定 d d d步推廣到 1 ≤ d ≤ D 1\le d \le D 1≤d≤D:
1 D ∑ d = 1 D ln p d ( o 1 : T ) ≥ ∑ t = 1 T ( E q ( s t ∣ o ≤ t ) [ ln p ( o t ∣ s t ) ] ⎵ reconstruction − 1 D ∑ d = 1 D β d E p ( s t − 1 ∣ s t − d ) q ( s t − d ∣ o ≤ t − d ) [ K L [ q ( s t ∣ o ≤ t ) ∥ p ( s t ∣ s t − 1 ) ] ] ⎵ latent overshooting ) \frac{1}{D} \sum_{d=1}^{D} \ln p_{d}\left(o_{1 : T}\right) \geq \sum_{t=1}^{T}\left(\underbrace{\mathrm{E}_{q\left(s_t | o _{\leq t}\right)}\left[\ln p\left(o_{t} | s_{t}\right)\right]}_{\text { reconstruction }}-\underbrace{\frac{1}{D} \sum_{d=1}^{D}\beta_d\mathrm{E}_{p\left(s_{t-1} | s_{t-d}\right) q\left(s_{t-d} | o_{ \leq t-d}\right)}\left[\mathrm{KL}\left[q\left(s_{t} | o_{ \leq t}\right) \| p\left(s_{t} | s_{t-1}\right)\right]\right]}_{\text { latent overshooting }}\right) D1∑d=1Dlnpd(o1:T)≥∑t=1T⎝⎜⎜⎜⎜⎛ reconstruction
Eq(st∣o≤t)[lnp(ot∣st)]− latent overshooting
D1d=1∑DβdEp(st−1∣st−d)q(st−d∣o≤t−d)[KL[q(st∣o≤t)∥p(st∣st−1)]]⎠⎟⎟⎟⎟⎞
Latent overshooting 可以被看作潛空間的正則項，鼓勵單步和多步預測之間的一緻性，因為這兩者在資料集上的期望應該是相等的。同時引入了權重因子 { β d } d = 1 D \{\beta_d\}^D_{d=1} {βd}d=1D，可以被調整來決定是更多關注短期預測還是長期預測。
對比
如上圖所示，a為标準變分下界，包含下方的重構loss，和上方單步預測的KL正則；b為觀測overshooting，将潛空間展開預測多步重構loss，但是在圖像領域計算代價過于昂貴；c可以看作以上兩者的結合和擴充。
從附錄結果看，latent overshooting 對于 RSSM 提升不顯著，對于确定性模型可能意義更大。
實驗