[paper reading] IMPALA: V-trace

IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures

https://arxiv.org/abs/1802.01561

IMPALA 原作者代码：https://github.com/deepmind/scala ble_agent

V-trace

考虑一个由策略 $\mu$ 产生的经验轨迹 $(x_t,a_t,r_t{)}_{t=s}^{t=s+n}$。定以 n 步 V-trace 目标如下

$$v_s\stackrel{def}{=}V(x_s)+\sum _{t=s}^{s+n-1}{\gamma }^{t-s}\left(\prod _{i=s}^{t-1}{c}_i\right){\delta }_{t}V$$

其中 ${\delta }_{t}V\stackrel{def}{=}{\rho }_t(r_t+\gamma V(x_{t+1})-V(x_t))$ 是 TD 误差，${\rho }_t\stackrel{def}{=}\min \left(\bar{\rho },\frac{\pi (a_t|x_t)}{\mu (a_t|x_t)}\right)$ 和 $c_i\stackrel{def}{=}\min \left(\bar{\rho },\frac{\pi (a_i|x_i)}{\mu (a_i|x_i)}\right)$ 是截断重要性采样，假设 $\bar{\rho }\ge \bar{c}$。$s=t$ 的时候，注意 ${\prod }_{i=s}^{t-1}{c}_i=1$。

在完全 on-policy 情况下，${\rho }_t=c_i=1$，因此上式可以写为

$$\begin{array}{rl}{v}_s& =V(x_s)+\sum _{t=s}^{s+n-1}{\gamma }^{t-s}(r_t+\gamma V(x_{t+1})-V(x_t))\\ & =r_s+\gamma r_{s+1}+{\gamma }^2{r}_{s+1}+\cdots +{\gamma }^{n-1}{r}_{s+n-1}+{\gamma }^{n}V(x_{s+n})\end{array}$$

这就是 on-policy 时的 n 步贝尔曼目标。这是 Retrace 所没有的特性。

注意，截断重要性采样中 $c_i$ 和 ${\rho }_t$ 的作用是不同的。

代码中 advantage 的计算：https://github.com/deepmind/scalable_agent/blob/master/vtrace.py#L275

GAE-Vtrace

首先计算 $gae{V}_s$：

$$\{\begin{array}{l}\begin{array}{rlrlrl}& gae{V}_s& & =& & {\delta }_s+\gamma {\rho }_{s+1}{\delta }_{s+1}+{\gamma }^2{\rho }_{s+1}{\rho }_{s+2}{\delta }_{s+2}+\cdots +{\gamma }^{n-1}\left(\prod _{i=s+1}^{s+n-1}{\rho }_i\right){\delta }_{s+n-1}\\ & gae{V}_{s+1}& & =& & {\delta }_{s+1}+\gamma {\rho }_{s+2}{\delta }_{s+2}+\cdots +{\gamma }^{n-2}\left(\prod _{i=s+2}^{s+n-1}{\rho }_i\right){\delta }_{s+n-1}\\ & \cdots & & & & \\ & gae{V}_{s+n-1}& & =& & {\delta }_{s+n-1}\end{array}\end{array}$$

然后，计算 $v_s$：

$$\{\begin{array}{l}\begin{array}{rlrlrlrlrlrlrl}& v_s& & =& & {\rho }_s& & \cdot & & gae{V}_s& & +& & V_s\\ & v_{s+1}& & =& & {\rho }_{s+1}& & \cdot & & gae{V}_{s+1}& & +& & V_{s+1}\\ & \cdots & & & & & & & & & & & & \\ & v_{s+n-1}& & =& & {\rho }_{s+n-1}& & \cdot & & gae{V}_{s+n-1}& & +& & V_{s+n-1}\end{array}\end{array}$$

IMPALA 中计算 advantage 的方式为 $r_s+\gamma v_{s+1}-V_s$，我们进行一下化简：

$$\begin{array}{rl}{r}_s+\gamma v_{s+1}-V_s& =r_s+\gamma ({\rho }_{s+1}\cdot gae{V}_{s+1}+V_{s+1})-V_s\\ & =\gamma {\rho }_{s+1}\cdot gae{V}_{s+1}+(r_s+\gamma V_{s+1}-V_s)\\ & =\gamma {\rho }_{s+1}\cdot gae{V}_{s+1}+{\delta }_s\\ & =gae{V}_s\end{array}$$