PlaNET：ICML2019论文分析

Learning Latent Dynamics for Planning from Pixels是ICML2019的一篇论文，提出了一种基于像素输入的强化学习方法，称为PlaNET（Planning Networks）。该方法通过学习潜在动态模型来进行规划和决策，解决了部分观测和连续控制问题。

算法主要贡献：

1. 相较于通过observation的图像来进行规划，根据编码后的图像进行在潜在空间中的规划在时间复杂度上更具有优势，并且更加高效。
2. Recurrent state space model：本文提出的一种新型的网络单元结构，既包含了确定性因素也考虑了随机因素
3. Latent overshooting：将standard variational bound推广到多步预测上来改善模型的长期预测

核心流程：

总体框架

训练过程的框架如图所示：

建立一个POMDP问题，定义观测模型、转移模型、奖励模型和encoder，利用可观测量来估计状态隐变量及其分布。

decode的图像用于训练的过程，在部署规划的时候只使用绿色部分的隐藏空间。在规划部分使用了交叉熵方法（Cross-Entropy Method, CEM），这是一类应用于重要性采样与优化的蒙特卡洛方法，计算过程为：

1. 从一个指定的概率分布中抽样出N的样本。
2. 通过最小化当前分布与目标分布(target distribution)的KL散度(KL-divergence),对当前的概率分布进行更新，以便在下一次迭代中产生效果最好的样本。

规划的流程为：使用学习到的模型来针对不同的行动预测收益，并选取最好的下一步的行动。可表示为：

$$o_{<=t}, a_{<t} \rightarrow s_t \rightarrow transition model \rightarrow s_{t+1}, s_{t+2}, ...\\ \rightarrow reward\ model \rightarrow R \rightarrow CEM \rightarrow a_t$$

基于先前的观察使用encoder得到当前状态的估计，然后对于未来的行动序列进行采样，并结合transition model来预测未来的状态序列，同时使用reward model来估计获得奖励的期望值。这样对于一个给定的未来行动序列就能够得到一个期望收益了，把这一套当做一个黑盒子我们就能利用CEM来规划出一个好的未来行动序列了。取这个行动序列的第一个行动作为下一步我们选择的行动即可。

模型学习

通过采样得到的数据是观测$o_{1:T}$和行动$a_{1:T}$，我们需要学习一个模型来估计状态隐变量$s_{1:T}$的分布：

$$p(o_{1:T},s_{1:T}\mid a_{1:T}) = \prod_{t=1}^T p(o_t\mid s_t) p(s_t\mid s_{t-1}, a_{t-1})$$

使用一个近似后验$q(s_{1:T}\mid o_{1:T}, a_{1:T})=\prod_{t=1}^T q(s_t\mid o_{1:T}, a_{1:T})$来逼近真实后验$p(s_{1:T}\mid o_{1:T}, a_{1:T})$。

在原文中有一个很复杂的公式：

$$\begin{aligned} \ln p(o_{1:T} \mid a_{1:T}) &\triangleq \ln \mathbb{E}_{p(s_{1:T} \mid a_{1:T})}\left[ \prod_{t=1}^{T} p(o_t \mid s_t) \right] \\ &= \ln \mathbb{E}_{q(s_{1:T} \mid o_{1:T}, a_{1:T})}\left[ \prod_{t=1}^{T} \frac{p(o_t \mid s_t) p(s_t \mid s_{t-1}, a_{t-1})}{q(s_t \mid o_{\le t}, a_{<t})} \right] \\ &\ge \mathbb{E}_{q(s_{1:T} \mid o_{1:T}, a_{1:T})}\left[ \sum_{t=1}^{T} \ln p(o_t \mid s_t) + \ln p(s_t \mid s_{t-1}, a_{t-1}) - \ln q(s_t \mid o_{\le t}, a_{<t}) \right] \tag{8} \\ &= \sum_{t=1}^{T} \left( \underbrace{\mathbb{E}_{q(s_t \mid o_{\le t}, a_{<t})}\left[ \ln p(o_t \mid s_t) \right]}_{\text{reconstruction}} - \underbrace{\mathbb{E}_{q(s_{t-1} \mid o_{\le t-1}, a_{<t-1})}\left[ \text{KL}\left[ q(s_t \mid o_{\le t}, a_{<t}) \parallel p(s_t \mid s_{t-1}, a_{t-1}) \right] \right]}_{\text{complexity}} \right). \end{aligned}$$

通过对不等式右侧做随机梯度上升即可学到这三个模型。

推导步骤

步骤 1：观测对数似然的全概率展开

根据全概率公式，观测的边缘对数似然可以表示为：

$$\ln p(o_{1:T} \mid a_{1:T}) \triangleq \ln \mathbb{E}_{p(s_{1:T} \mid a_{1:T})} \left[ \prod_{t=1}^T p(o_t \mid s_t) \right]$$

这里我们对隐状态序列 $s_{1:T}$ 进行边缘化，将观测似然写成对动力学先验 $p(s_{1:T} \mid a_{1:T})$ 的期望形式。

步骤 2：引入近似后验的重要性加权

利用重要性采样技巧，将对动力学先验的期望转换为对近似后验 $q$ 的期望：

$$= \ln \mathbb{E}_{q(s_{1:T} \mid o_{1:T}, a_{1:T})} \left[ \prod_{t=1}^T \frac{p(o_t \mid s_t) p(s_t \mid s_{t-1}, a_{t-1})}{q(s_t \mid o_{\le t}, a_{<t})} \right]$$

根据重要性加权公式：

$$\mathbb{E}_{p(s)}[f(s)] = \mathbb{E}_{q(s)}\left[ f(s) \cdot \frac{p(s)}{q(s)} \right]$$

其中 $p(s_{1:T} \mid a_{1:T}) = \prod_{t=1}^T p(s_t \mid s_{t-1}, a_{t-1})$，$q(s_{1:T} \mid o_{1:T}, a_{1:T}) = \prod_{t=1}^T q(s_t \mid o_{\le t}, a_{<t})$，因此两者的比值可拆分为每一步的乘积形式。

步骤 3：应用 Jensen 不等式得到下界

由于对数函数是凹函数，根据 Jensen 不等式 $\ln \mathbb{E}[X] \ge \mathbb{E}[\ln X]$，可得：

$$\ge \mathbb{E}_{q(s_{1:T} \mid o_{1:T}, a_{1:T})} \left[ \sum_{t=1}^T \ln p(o_t \mid s_t) + \ln p(s_t \mid s_{t-1}, a_{t-1}) - \ln q(s_t \mid o_{\le t}, a_{<t}) \right]$$

将期望内的乘积取对数后转化为求和，得到对数似然的一个下界（即证据下界 ELBO）。最大化该下界即可间接最大化原始对数似然。

步骤 4：整理为重构项与复杂度项

将上式整理为按时间步分解的形式，可得到最终的目标函数：

$$= \sum_{t=1}^T \left( \underbrace{\mathbb{E}_{q(s_t \mid o_{\le t}, a_{<t})} [\ln p(o_t \mid s_t)]}_{\text{重构项}} - \underbrace{\mathbb{E}_{q(s_{t-1} \mid o_{\le t-1}, a_{<t-1})} \left[ \text{KL}\big(q(s_t \mid o_{\le t}, a_{<t}) \parallel p(s_t \mid s_{t-1}, a_{t-1})\big) \right]}_{\text{复杂度项}} \right)$$

• 重构项（Reconstruction）：
  衡量从隐状态 $s_t$ 重建观测 $o_t$ 的能力，与 VAE 中的重建损失作用相同，目标是让重建观测与真实观测尽可能一致。
• 复杂度项（Complexity）：
  用 KL 散度衡量近似后验 $q(s_t \mid o_{\le t}, a_{<t})$ 与动力学先验 $p(s_t \mid s_{t-1}, a_{t-1})$ 的差异，起到正则化作用，防止后验分布偏离动力学模型的先验规律，保证隐状态的转移符合环境动态。

最终目标函数（ELBO）

整理后的变分下界就是 PlaNET 世界模型的训练目标：

$$\mathcal{L} = \sum_{t=1}^T \left( \text{Recon}_t - \text{KL}_t \right)$$

训练时同时优化观测模型、动力学模型和近似后验模型，即可得到一个能够预测环境动态的世界模型，后续可在此模型中进行规划（如 CEM 算法），大幅提升强化学习的样本效率。