TD-VAE: 时间差分变分自编码器.
1. TD-VAE的动机
时间差分变分自编码器(Temporal Difference VAE)是一种为序列数据设计的变分自编码器。它的设计动机主要有三点：
⚪ 状态空间模型 State-Space Models
状态空间模型是指用一个不能观测的隐状态序列$z=(z_1,…,z_T)$决定观测状态$x=(x_1,…,x_T)$，通常建模为一个Markov链模型。
⚪ 信念状态 Belief State
在强化学习中，一个代理agent会把所有过去的状态编码为对未来的预测，称为信念状态。若在$t$时刻的信念状态为$b_t=belief(x_1,…,x_t)=belief(b_t,x_t)$，则未来状态的分布写作$p(x_{t+1},…,x_T|x_1,…,x_t)≈p(x_{t+1},…,x_T|b_t)$。在TD-VAE中，信念状态通过一个循环策略的隐状态表示$b_t=RNN(b_{t-1},x_t)$。
⚪ 跳跃式预测 Jumpy Prediction
代理agent应该能根据目前收集的信息做出对比较遥远的未来的预测，即预测未来几步的状态。
2. TD-VAE的目标函数
变分自编码器(VAE)的目标函数为对数似然的变分下界:
\[\begin{aligned} \log p(x) &= \log \sum_{z}^{} p(x,z) = \log \sum_{z}^{} \frac{p(x,z)}{q(z|x)}q(z|x) \\ &= \log \Bbb{E}_{z \text{~} q(z|x)}[\frac{p(x,z)}{q(z|x)}] \geq \Bbb{E}_{z \text{~} q(z|x)}[\log \frac{p(x,z)}{q(z|x)}] \\ &= \Bbb{E}_{z \text{~} q(z|x)}[\log p(x,z)-\log q(z|x)] \end{aligned}\]将状态$x_t$的分布建模为过去状态$x_{\lt t}$的条件概率，由于状态空间建模为Markov链，因此状态$x_t$只与当前时刻和前一时刻的隐状态$z_t,z_{t-1}$有关：
\[\log p(x_t | x_{\lt t}) = \Bbb{E}_{z_t,z_{t-1} \text{~} q(z_t,z_{t-1}|x_{\leq t})}[\log p(x_t,z_t,z_{t-1}|x_{\lt t})-\log q(z_t,z_{t-1}|x_{\leq t})]\]上式又可以分解为：
\[\begin{aligned} \log &p(x_t | x_{\lt t}) \\ &=\Bbb{E}_{z_t,z_{t-1} \text{~} q}[\log p(x_t,z_t,z_{t-1}|x_{\lt t})-\log q(z_t,z_{t-1}|x_{\leq t})] \\ &=\Bbb{E}_{z_t,z_{t-1} \text{~} q}[\log p(x_t|z_t,z_{t-1},x_{\lt t}) p(z_t,z_{t-1}|x_{\lt t})-\log q(z_t|x_{\leq t})q(z_{t-1}|z_t,x_{\leq t})] \\ &=\Bbb{E}_{z_t,z_{t-1} \text{~} q}[\log p(x_t|z_t,z_{t-1},x_{\lt t}) p(z_{t-1}|x_{\lt t})p(z_t|z_{t-1},x_{\lt t})-\log q(z_t|x_{\leq t})q(z_{t-1}|z_t,x_{\leq t})] \end{aligned}\]注意到由于Markov假设，$p(x_t|z_t,z_{t-1},x_{\lt t})=p(x_t|z_t)$，$p(z_t|z_{t-1},x_{\lt t})=p(z_t|z_{t-1})$，则上式表示为：
\[\begin{aligned} \log p(x_t | x_{\lt t}) =\Bbb{E}_{z_t,z_{t-1} \text{~} q}[&\log p(x_t|z_t) +\log p(z_{t-1}|x_{\lt t}) +\log p(z_t|z_{t-1}) \\ &-\log q(z_t|x_{\leq t})-\log q(z_{t-1}|z_t,x_{\leq t})] \end{aligned}\]根据跳跃式预测的思想，上述目标不仅在$t-1,t$时刻成立，也在任何时间段$t_1<t_2$成立。该目标函数表示有四种分布需要学习：
(1) 解码分布 decoder distribution
$p(x_t|z_t)$是概率解码器，将其建模为$p_D^{t_2}(x_{t_2})$。
(2) 转移分布 transition distribution
$p(z_t|z_{t-1})$捕捉隐变量之间的顺序依赖关系，将其建模为$p_T^{t_2}(z_{t_2})$。
(3) 信念分布 belief distribution
$p(z_{t-1}|x_{\lt t})=p(z_{t-1}|b_{t-1})$和$q(z_t|x_{\leq t})=q(z_t|b_{t})$都是通过信念状态预测隐状态，将其分别建模为$p_B^{t_1}(z_{t_1})$和$p_B^{t_2}(z_{t_2})$。
(4) 平滑分布 smoothing distribution
$q(z_{t-1}|z_t,x_{\leq t})$是对过去状态的平滑，也可以通过信念状态表示为$q(z_{t-1}|z_t,b_{t-1},b_{t})$，将其建模为$q_S^{t_1|t_2}(z_{t_1})$。
(5) 最终目标函数
TD-VAE最终的目标函数为：
\[\begin{aligned} \log p(x_{t_2} | x_{t_1}) =\Bbb{E}_{z_{t_2},z_{t_1} \text{~} q}[&\log p_D(x_{t_2}|z_{t_2}) +\log p_B(z_{t_1}|b_{t_1}) +\log p_T(z_{t_2}|z_{t_1}) \\ &-\log p_B(z_{t_2}|b_{t_2})-\log q_S(z_{t_1}|z_{t_2},b_{t_1},b_{t_2})] \end{aligned}\]
