Stochastic Latent Actor-Critic: Deep Reinforcement Learning with a Latent Variable Model

NIPS 2020
paper

Intro

面对图像任务下RL存在两个挑战:表征学习以及任务学习。本文提出学习潜在变量模型并基于该模型执行RL。

Method

考虑部分可观测MDP(POMDP), 其概率图模型如图所示。该模型下智能体无法影响过去$\tau$步的状态动作，而是预测未来最有动作，直至回合结束。

其中$O_{\tau +1}$为新引入的随机变量分布，$p({\mathcal{O}}_t=1|{\mathbf{z}}_t,{\mathbf{a}}_t)=\exp (r({\mathbf{z}}_t,{\mathbf{a}}_t))$。算法对观测以及奖励联合构建时序模型，并通过最大化似然概率$p({\mathbf{x}}_{1:\tau +1},{\mathcal{O}}_{\tau +1:T}|{\mathbf{a}}_{1:\tau })$分布优化策略。进一步通过变分推断得到该分布的ELBO，

其中$r({\mathbf{z}}_t,{\mathbf{a}}_t)=\log p({\mathcal{O}}_t=1|{\mathbf{z}}_t,{\mathbf{a}}_t)$
$$q({\mathbf{z}}_{1:T},{\mathbf{a}}_{\tau +1:T}|{\mathbf{x}}_{1:\tau +1},{\mathbf{a}}_{1:\tau })=\prod _{t=0}^{\tau }q({\mathbf{z}}_{t+1}|{\mathbf{x}}_{t+1},{\mathbf{z}}_t,{\mathbf{a}}_t)\prod _{t=\tau +1}^{T-1}p({\mathbf{z}}_{t+1}|{\mathbf{z}}_t,{\mathbf{a}}_t)\prod _{t=\tau +1}^T\pi ({\mathbf{a}}_t|{\mathbf{x}}_{1:t},{\mathbf{a}}_{1:t-1})$$
$$p({\mathbf{x}}_{1:\tau +1},{\mathcal{O}}_{\tau +1:T},{\mathbf{z}}_{1:T},{\mathbf{a}}_{\tau +1:T}|{\mathbf{a}}_{1:\tau })=\prod _{t=1}^{\tau +1}p({\mathbf{x}}_t|{\mathbf{z}}_t)\prod _{t=0}^{T-1}p({\mathbf{z}}_{t+1}|{\mathbf{z}}_t,{\mathbf{a}}_t)\prod _{t=\tau +1}^{T}p({\mathcal{O}}_t|{\mathbf{z}}_t,{\mathbf{a}}_t)\prod _{t=\tau +1}^{T}p({\mathbf{a}}_t)$$

因此，问题转化为最大化这个ELBO。ELBO分为两个部分第一个部分学习隐变量模型，通过最小化下列损失函数实现模型参数训练
$$J_M(\psi )={\mathbb{E}}_{{\mathbf{z}}_{1:\tau +1}\sim q_{\psi }}\left[\sum _{t=0}^{\tau }-\log p_{\psi }({\mathbf{x}}_{t+1}|{\mathbf{z}}_{t+1})+{\mathrm{D}}_{\mathbf{K}.}(q_{\psi }({\mathbf{z}}_{t+1}|{\mathbf{x}}_{t+1},{\mathbf{z}}_t,{\mathbf{a}}_t)\Vert p_{\psi }({\mathbf{z}}_{t+1}|{\mathbf{z}}_t,{\mathbf{a}}_t))\right]$$

第二部分，文章假设$a_t$先验为服从均匀分布因此$\log p(a_t)$该项为常数项，可以舍弃去。则该目标变为最大化熵的RL。本文采用了SAC。首先对q函数最小化soft bellman 均方误差优化
$$\begin{array}{c}{J}_Q(\theta )={\mathbb{E}}_{{\mathbf{z}}_{1:\tau +1}\sim q_{\psi }}\left[\frac{1}{2}{\left(Q_{\theta }({\mathbf{z}}_{\tau },{\mathbf{a}}_{\tau })-(r_{\tau }+\gamma V_{\bar{\theta }}({\mathbf{z}}_{\tau +1}))\right)}^2\right],\\ V_{\theta }({\mathbf{z}}_{\tau +1})=\underset{{\mathbf{a}}_{\tau +1}\sim {\pi }_{\varphi }}{E}[Q_{\theta }({\mathbf{z}}_{\tau +1},{\mathbf{a}}_{\tau +1})-\alpha \log {\pi }_{\varphi }({\mathbf{a}}_{\tau +1}|{\mathbf{x}}_{1:\tau +1},{\mathbf{a}}_{1:\tau })]\end{array}$$
其中${\mathbf{z}}_{\tau +1}\sim q_{\psi }({\mathbf{z}}_{\tau +1}|{\mathbf{x}}_{\tau +1},{\mathbf{z}}_{\tau },{\mathbf{a}}_{\tau })$,然后策略优化采用SAC
$$\begin{array}{r}{J}_{\pi }(\varphi )={\mathbb{E}}_{{\mathbf{z}}_{1:\tau +1}\sim q_{\psi }}\left[{\mathbb{E}}_{{\mathbf{a}}_{\tau +1}\sim {\pi }_{\varphi }}\left[\alpha \log {\pi }_{\varphi }({\mathbf{a}}_{\tau +1}|{\mathbf{x}}_{1:\tau +1},{\mathbf{a}}_{1:\tau })-Q_{\theta }({\mathbf{z}}_{\tau +1},{\mathbf{a}}_{\tau +1})\right]\right]\end{array}$$

伪代码