1. ACKTR 概念
本文 [1] 涉及到一个新术语:Kronecker 因子的近似曲率,以下简称:K-FAC。本文结合三个东西,一个是 K-FAC,一个是置信域优化,另一个是 AC。作者称其为第一个可扩展的 AC 框架的置信域优化方法,可以理解为精度和速度可以调节。
置信域优化方法是为了应对 SGD 及相关一阶优化方法的探索效率不足的问题,这往往导致训练时间过长。另外 A3C 也采用多线程的异步训练方法来应对探索效率问题,但是异步训练对样本的利用效率太低。样本的利用效率往往也十分关键,因为在现实世界,机器人与环境交互的数据量远远不满足训练数据需要的规模和增长速率,即使在仿真环境中,仿真器的运行时间也可以超过计算时间。
在这篇文章之前,采用过高级优化方法的技术有自然策略梯度方法以及 TRPO 方法。这两种方法的计算复杂度都很高。有学者提出 K-FAC 方法 [2, 3],这是对自然梯度方法的一种可扩展的近似,已经用来加速各种大型网络的监督训练。
2. 相关知识
2.1 策略梯度和AC 方法
策略梯度的一般形式可以表示为:
其中 $\Psi^t$ 通常采用优势函数 $A^\pi(s_t,a_t)$。本文对优势函数的估计采用 A3C 中的 k-step 估计方法:
其中 $V_\Phi^\pi$ 表示价值网络,用来逼近策略的回报总和:$V_\Phi^\pi = \mathbb{E}_\pi[R_t]$ 。对价值网络的训练方法采用最小化均方根误差:$\frac{1}{2}| \hat{R}_t - V_\Phi^\pi(s_t) |$ 。
2.2 采用 K-FAC 逼近自然梯度
为了更新非凸目标函数 $\mathcal{J}(\theta)$,最速下降方法通过更新步长 $\Delta \theta$ 来最小化 $\mathcal{J}(\theta+\Delta\theta)$,同时满足约束条件 $| \Delta\theta |_B < 1$。$| \cdot |_B$ 定义为 $| x |_B = (x^T B x)^{\frac{1}{2}}$,其中 B 是半正定矩阵。解以上的待约束优化问题,可以得到 $\Delta \theta \propto -B^{-1} \nabla_\theta \mathcal{J}$,并且 $\nabla_\theta \mathcal{J}$ 就是标准的梯度。假设 $B=I$,那么就是标准的梯度下降方法,$| x |_I$ 就是欧几里德范式(二范式)。论文中说,欧几里德范数的变化取决于参数 $\theta$,但这对于随机参数化的模型来说不是一个好的选择,这会影响到优化轨迹。
自然梯度与随机梯度方法的差别在于,它利用费舍尔信息矩阵 $F$ (KL 散度的本地二次型近似)来构造范式,而不是标准矩阵 $I$。这种范式独立于模型参数 $\theta$,属于一种概率分布。自然梯度可以提供更加稳定和有效的梯度更新。但是对于现在的深度神经网络,参数高达百万级别,计算费舍尔矩阵和它的逆基本上是不现实的。
K-FAC 方法利用 Kronecker-factored 来近似费舍尔信息矩阵,从而能够实现自然梯度更新。
令 $p(y|x)$ 表示神经网络输出的概率分布,$L = \log p(y|x)$ 表示对数似然函数。令 $W \in \mathbb{R}^{C_{out} \times C_{in}}$ 表示神经网络第 $\ell$ 层的参数矩阵,其中 $C_{out}$ 和 $C_{in}$ 表示该层神经网络的输出和输入神经元。令输入的激活向量表示为 $a \in \mathbb{R}^{C_{in}}$ ,输出到下一层的未激活的向量表示为 $s = W a$ 。
根据求导的链式法则,可以知道 $\nabla_W L = (\nabla_s L) a^T$,因此对于 $\ell$ 层参数的费舍尔矩阵 $F_\ell$ 可以近似比表示为 $\hat{F}_\ell$:
上式中 $\otimes$ 表示矩阵 Kronecker(克罗内克)积,也是本方法名字的由来。令 $A$ 表示 $\mathbb{E}[aa^T]$,$S$ 表示 $\mathbb{E}[\nabla_s L(\nabla_s L)^T]$。这个近似可以解释为:假设输入的激活向量的二阶统计与后向传播的导数无关。
根据矩阵克罗内克积的性质,有:$(P \otimes Q) ^{-1} = P^{-1} \otimes Q^{-1}$,$(P \otimes Q) \mathscr{vec}(T) = PTQ^{T}$ 。因此对于自然梯度更新:
公式 $\eqref{4}$ 的计算复杂度与 $|W|$ 的规模相当,而直接计算费舍尔矩阵向量乘积和求逆复杂度是 $|W|^2$ 的规模。
3. 算法框架
3.1 K-FAC + AC
本文利用 K-FAC 方法逼近自然梯度,然后将自然梯度更新用于 AC 框架(既用于 actor,也用于 critic 的更新)。
强化学习目标函数(即 actor 的目标函数)的费舍尔矩阵可以定义为:
其中 $p(\tau)$ 表示轨迹的分布,完整表达式为 $p(\tau) = p(s_0) \prod_{t=0}^T \pi(a_t|s_t)p(s_{t+1}|s_t,a_t)$ 。实际上都是通过在训练过程采样轨迹求平均的方法来求以上的期望,通过表达式求期望太难了。
critic 的目标函数是均方根误差函数,可视为最小二乘逼近问题。对于最小二乘逼近问题,常见的二阶优化方法是 Gauss-Newton 方法,Gauss-Newton 对函数曲率的逼近定义为 $G := \mathbb{E}[J^TJ]$,其中 $J$ 是参数的雅克比矩阵。
假设观测模型为高斯观测模型,则 Gauss-Newton 矩阵等价于费舍尔矩阵。因此也可以将 K-FAC 方法计算 Gauss-Newton 矩阵。为了建立高斯观测模型,本文将 critic 的输出 $v$ 定义为高斯分布:$p(v|s_t) \sim \mathcal{N}(v;V(s_t), \sigma^2)$ 。另外,为了简便,将 $\sigma$ 设置为 1,等价于 vanilla Gauss-Newton 方法。此时 critic 的费舍尔矩阵就是 Gauss-Newton 矩阵。
为了可以将 K-FAC 方法同时运用于 actor 和 critic 的更新中,可以让 actor 和 critic 共享同一个网络,除了输出层不一样。可以定义一个策略分布(actor 的输出)和状态分布(critic 的输出)的联合分布:$p(a,v|s)=\pi(a|s)p(v|s)$,然后构造关于 $p(a,v|s)$ 费舍尔矩阵:$\mathbb{E}_{p(\tau)} [\nabla \log p(a,v|s)\nabla \log p(a,v|s)^T]$ 。
本文还对目标函数运用 Tikhonov 正则项 [4]。
3.2 置信域优化
自然梯度在更新时和随机梯度类似:$\theta \leftarrow \theta - \eta F^{-1} \nabla_\theta L$,这种更新方式有时会导致较大的更新步长。论文 [5] 中将置信域优化和 K-FAC 结合起来,以置信域的方式进行更新。置信域步长:$\min(\eta_{\max}, \sqrt{\frac{2\delta}{\Delta\theta \hat{F} \Delta\theta}})$,其中学习率 $\eta_{\max}$ 和置信域范围参数 $\delta$ 都是超参数。
4. 工程设置
对比算法:A2C、TRPO,采用相同的模型结构。
训练时长:1e7 以上的时间步长,1 个时间步长等于 4 帧。
Adaptive Gauss-Newton:对 critic 输出分布的方差用贝尔曼误差的方差进行估计。贝尔曼误差反映的是当前的状态价值与更新后的状态价值差的绝对值。
4.1 离散控制
网络结构:第一层卷积层采用 32 个 $8 \times 8$ ,步长为 4 的滤波器。第二层卷积层采用 64 个 $4 \times 4$ ,步长为 2 的滤波器。第三层卷积层采用 32 个 $3 \times 3$ ,步长为 1 的滤波器。第四层全连接层采用 512 个神经元。softmax 输出层输出动作,线性输出层输出状态值。
actor 和 critic 共享网络结构(除了输出层)。
最大学习率 $\eta_{\max}$ 在 {0.7,0.2,0.07,0.02} 中选取,置信域半径 $\delta$ 设置为 0.001。
4.2 连续控制
对于低维度状态空间作为输入时,采用两个独立的网络来估计 actor 和 critic 。两层全连接网络,每层 64 个隐藏单元。actor 网络的激活函数采用 Tanh,critic 网络的激活函数采用 ELU。输出层不采用激活函数。actor 网络输出高斯策略分布,标准差作为最后一层的偏置加入,与状态输入无关。
对于高维度 $42 \times 42$ RGB 图像状态空间作为输入时,前面两层卷积层采用 32 个 $3 \times 3$,步长为 2 的滤波器,最后一层全连接层采用 256 个神经元。actor 和 critic 采用独立网络。actor 网络的激活函数采用 ReLU,critic 网络的激活函数采用 ELU。
最大学习率 $\eta_{\max}$ 在 {0.3,0.03,0.003} 中选取,置信域半径 $\delta$ 设置为 0.001。
actor 网络和 critic 网络都采用正交初始化。
参考文献
[1] Wu, Y., Grosse, R., & Ba, J. (n.d.). Scalable trust-region method for deep reinforcement learning using Kronecker-factored approximation. 1–14.
[2] Martens J , Grosse R . Optimizing Neural Networks with Kronecker-factored Approximate Curvature[J]. 2015.
[3] R. Grosse and J. Martens. A Kronecker-factored approximate Fisher matrix for convolutional layers. In ICML, 2016
[4] Martens J , Grosse R . Optimizing Neural Networks with Kronecker-factored Approximate Curvature[J]. 2015.
[5] J. Ba, R. Grosse, and J. Martens. Distributed second-order optimization using Kronecker- factored approximations. In ICLR, 2017.