AC using K-FAC TR (ACKTR)
Scalable trust-region method for deep reinforcement learning using Kronecker-factored approximation.
Y. Wu, E. Mansimov, S. Liao, R. Grosse, and J. Ba. 2017
Github （OpenAI提供的baseline）
摘要——ACKTR扩展了自然策略梯度的框架并用Kronecker因子近似曲率（Kronecker-factored approximate curvature, K-FAC），同时在置信域内优化actor和critic，因此该方法称为Actor Critic using Kronecker-Factored Trust Region (ACKTR)。这是第一个针对actor-critic法的可扩展置信域自然梯度方法，也是从原始图像输入直接学习复杂任务的连续控制和离散控制策略的方法。与先前的actor-critic法相比，ACKTR取得了更高的奖励和平均2~3倍的样本采样效率。
1 Introduction
深度强化学习（Deep RL, DRL）用深度神经网络表示控制策略。尽管效果很好，但这些网络仍用随机梯度下降（stochastic gradient descent, SGD）训练。SGD及相关的一阶方法效率不高，所以有了用多个agent同时与环境交互的分布式方法，但随着并行度的增加，采样效率的回报会迅速减少。
采样效率是RL中的主要关注点。有效减少样本量的一种方法是使用更高级的梯度更新优化技术。自然策略梯度利用自然梯度下降进行策略更新。自然梯度法沿着最陡的梯度下降，用Fisher信息矩阵作基础矩阵，该矩阵与坐标轴的选择无关，只与流形相关（i.e. 表面）。
但是，自然梯度的计算涉及Fisher信息矩阵的求逆，所以计算麻烦。TRPO利用Fisheries向量积避免了显式存储和对Fisher矩阵求逆。但是，TRPO通常需要多次共轭梯度才能获取单个参数的更新，并且准确估计曲率需要每个batch中的大量样本，所以TRPO对于大的模型不适用且采样效率低。
K-FAC是自然梯度的可扩展近似。在监督学习中，通过使用更大的mini-batches，K-FAC可以加速很多SOTA的大规模神经网络。与TRPO不同，每次更新的成本和一次SGD更新相当，并且K-FAC保持曲率信息的运行平均值，这使得K-FAC中可以用小的batch。这表明将K-FAC用于策略优化可以提高现在DRL的采样效率。
ACKTR中对自然梯度用Kronecker-factored近似，可以高效地对梯度的协方差矩阵求逆。该文章还提出扩展自然策略梯度算法以通过Gauss-Newton近似来优化价值函数。在实践中，ACKTR每次更新的计算成本仅比基于SGD的方法高10%~25%。
2 Background
2.1 RL & AC methods
无限、折扣MDP为\((\mathcal{X}, \mathcal{A}, \gamma, P, r)\)，策略为\(\pi_{\theta} (a \vert s_t)\)。Agent的目标是最大化期望\(\gamma\)-折扣的累计return\(J(\theta) = \mathbb{E}_{\pi} [R_t] = \mathbb{E}_{\pi} [\sum_{i \ge 0}^{\infty} \gamma^i r(s_{t+i}, a_{t+i})]\)，该return与策略参数\(\theta\)有关。策略梯度方法直接参数化一个策略\(\pi_{\theta} (a \vert s_t)\)并更新\(\theta\)来最大化目标\(J(\theta)\)。在一般形式中，策略梯度定义为
\[\nabla_\theta \mathcal{J}(\theta) = \mathbb{E}_\pi \left[\sum_{t=0}^{\infty} \Psi^t \nabla_\theta \log \pi_\theta (a_t \vert s_t)\right]\]其中，\(\Psi^t\)一般为advantage function \(A^{\pi} (s_t, a_t)\)，ACKTR中用和A3C中一样的定义，即\(k\)-step函数逼近的return
\[A^\pi (s_t, a_t) = \sum_{i=0}^{k-1} \gamma^i r(s_{t+i}, a_{t+i}) + \gamma^k V_\phi^\pi (s_{t+k}) - V_\phi^\pi (s_t)\]其中,\(V_\phi^\pi (s_t)\)为价值网络，提供给定状态下跟随策略\(\pi\)获得的期望奖励和的估计，\(V_\phi^\pi (s_t) = \mathbb{E}_\pi [R_t]\)。为了训练价值网络的参数，我们按照上述advantage函数进行TD更新，以最小化自举\(k\)-step return \(\hat{R}_t\)和预测值之间的平方差\(\frac{1}{2} \| \hat{R}_t - V_\phi^\pi (s_t) \|^2\)。
2.2 Natural gradient using K-FAC approximation
为了最大化非凸函数\(J(\theta)\)，最陡下降法计算能最大化\(J(\theta + \Delta \theta)\)，满足约束\(\| \Delta \theta \|_B \lt 1\)，其中\(\| \cdot \|_B\)是定义为\(\|x \|_B = (x^T B x)^{\frac{1}{2}}\)的范数且\(B\)为半正定阵。该约束优化问题的解具有\(\Delta \theta \propto -B^{-1} \nabla_\theta J\)，其中\(\nabla_\theta J\)为标准梯度。自然梯度法用Fisher信息矩阵构造范数，\(F\)是KL散度的局部二次近似。该范数与概率分布类的模型参数化\(\theta\)无关，提供了更稳定、有效的更新。但是神经网络参数很多，计算和存储Fisher矩阵和它的逆是不现实的，所以我们需要求近似解。
K-FAC通过对Fisher矩阵进行Kronecker-factored近似，实现了高校的近似自然梯度更新。用\(p(y \vert x)\)表示神经网络的输出分布，\(L = \log p(y \vert x)\)为其对数似然。令\(W \in \mathbb{R}^{C_{out} \times C_{in}}\)为第\(\ell\)层的权重矩阵，其中\(C_{out}\)和\(C_{in}\)分别为该层输出/输入层的神经元数量。该层的输入激活向量为\(a \in \mathbb{R}^{C_{in}}\)，下一层的预激活向量为\(s = Wa\)。权重梯度为\(\nabla_W L = (\nabla_s L) a^T\)。K-FAC利用以下事实并进一步将与第\(\ell\)层相关的block \(F_{\ell}\)近似为\(\hat{F}_{\ell}\)
\[\color{green}{\begin{aligned} F_{\ell} &= \mathbb{E} [\operatorname{vec} \{\nabla_W L\} \operatorname{vec} \{\nabla_W L\}^T] = \mathbb{E}\left[ a a^T \otimes \nabla_s L\left(\nabla_s L\right)^T \right] \\ & \approx \mathbb{E}[a a^T] \otimes \mathbb{E}[\nabla_s L (\nabla_s L)^T] := A \otimes S :=\hat{F}_{\ell} \end{aligned}}\]其中，\(A\)为\(\mathbb{E}[a a^T]\)，\(S\)为\(\mathbb{E}[\nabla_s L (\nabla_s L)^T]\)。（向量化计算\(\operatorname{vec}\)、Kronecker积\(\otimes\)以及分层计算Fisher信息矩阵的细节可参考【DRL-20】ACKTR。）该近似的含义是假设激活函数和反向传播的二阶统计量无关。有了这个近似后，利用基本性质\((P \otimes Q)^{-1} = P^{-1} \otimes Q^{-1}\)和\((P \otimes Q) \operatorname{vec}(T) = P T Q^T\)就可以有效计算自然梯度：
\[\color{green}{ \operatorname{vec}(\Delta W) = \hat{F}_{\ell}^{-1} \operatorname{vec}\left\{\nabla_W \mathcal{J}\right\} = \operatorname{vec}\left(A^{-1} \nabla_W \mathcal{J} S^{-1}\right)}\]根据以上方程，我们可以看出用K-FAC近似自然梯度更新只需要计算和\(W\)大小相当的矩阵。
3 Methods
3.1 Natural gradient in AC
自然梯度是Kakade在十多年前提出的一种适用于策略梯度的方法。但仍然没有自然梯度的可扩展、采样高效且通用的实例化。本节将介绍一个基于actor-critic饭的可扩展且采样高效的自然梯度算法：ACKTR。我们用K-FAC近似来计算自然梯度更新，并用自然梯度来更新actor和critic。
为了定义目标的Fisher矩阵，一个自然的选择是用策略函数并对轨迹分布求期望：
\[F = \mathbb{E}_p(\tau) \bigg[\nabla_{\theta} \log \pi(a_t \vert s_t) \big(\nabla_{\theta} \log \pi(a_t \vert s_t)\big)^T \bigg]\]其中，\(p(\tau)\)是轨迹的分布，由\(p(s_0) = \Pi_{t=0}^T \pi(a_t \vert s_t) p(s_{t+1} \vert s_t, a_t)\)。在实际应用中，我们可以近似训练过程中收集到的难以解决的轨迹期望。
下面，描述如何用自然梯度优化critic。学习一个critic可以看作是一个最小二乘函数逼近问题，虽然目标是移动的。在最小二乘函数逼近的设定中，选择的二阶算法一般为Gauss-Newton，将曲线近似为Gauss-Newton矩阵\(G := \mathbb{E}[J^T J]\)，其中\(J\)为参数到输出映射的Jacobian。Gauss-Newton矩阵等价于高斯观测模型的Fisher矩阵，这个等价也使得我们可以将K-FAC用于critic。将critic \(v\)的输出定义为一个高斯分布\(p(v \vert s_t) \sim \mathcal{N}(v; V(s_t), \sigma^2)\)。根据这个高斯输出分布，定义critic的Fisher矩阵。在实际应用中，我们可以直接把\(\sigma\)设为1，即可等价于普通的Gauss-Newton法。
如果actor和critic不联合，则可以用上面定义的指标分别对两个网络用K-FAC更新。但是，为了防止训练的不稳定性，两个网络共享较低层的表示但具有不同输出的结构往往是有益的。在这种情况下，我们可以通过假设两个输出的分布相互独立来定义策略和价值的联合分布，i.e. \(p(a, v \vert s) = \pi(a \vert s) p(v \vert s)\)，并定义关于\(p(a, v \vert s)\)的Fisher矩阵，这与标准的K-FAC无异，只是需要对网络的输出进行独立采样。然后，我们就可以用K-FAC近似Fisher矩阵\(\mathbb{E}_p(\tau) \bigg[\nabla \log p(a, v \vert s) \nabla \log p(a, v \vert s)^T \bigg]\)来同时进行更新。
另外，我们使用正则阻尼进行正则化。我们也按J. Ba等人提出的Distributed second-order optimization using Kroneckerfactored approximations，执行K-FAC近似所需的二阶统计量和求逆的异步计算，以减少计算时间。
3.2 Step-size Selection and TR optimization
传统地，自然梯度的计算都是使用类似SGD更新的方式，\(\theta \leftarrow \theta - \eta F^{-1} \nabla_{\theta}L\)。但在deep RL中，Schulman等人（在TRPO中）发现这样的更新方式会导致策略的大量更新，导致算法过早收敛到近确定性策略。他们主张用置信域方法，即在特定的数量（关于KL散度）内修改策略分布。因此，我们采用K-FAC的置信域构造，用\(\min(\eta_{\max}, \sqrt{\frac{2 \delta}{\Delta \theta \hat{F} \Delta \theta}})\)作为有效步长\(\eta\)，其中学习率\(\eta_{\max}\)和置信域半径\(\delta\)都是超参数。如果actor和critic是不联合的，那么我们需要分别为两个网络调一组不同的\(\eta_{\max}\)和\(\delta\)。
4 Conclusion
在这项工作中，我们提出了一种用于deep RL的高效采样且低计算成本的置信域优化方法。我们将最近提出的K-FAC用于近似Actor-Critic中自然梯度的更新，并使用置信域优化来提高稳定性。据我们所知，我们率先提出使用自然梯度更新来同时优化actor和critic。我们在Atari游戏和MuJoCo环境中测试了我们的方法，与一阶梯度法(A2C)和迭代二阶方法(TRPO)相比，我们观察到采样效率平均提高了2到3倍。由于我们算法具有可扩展性，我们也是第一个直接从原始像素观测空间训练连续控制中的几个非平凡任务。这表明将K-FACK-FAC自然梯度近似扩展到强化学习中的其他算法是一个很有前景的研究方向。
