重参数化与强化学习

1. 为什么需要重参数化？随机函数的梯度计算问题

在机器学习领域，目标函数常常需要计算如下形式的随机函数：

$$L(\theta )={\mathbb{E}}_{q_{\theta }(\mathbf{z})}[l(\theta ,\mathbf{z})]=\int l(\theta ,\mathbf{z})q_{\theta }(\mathbf{z})\mathrm{d}\mathbf{z}.$$

其梯度为：

$$\begin{array}{rl}{\mathrm{\nabla }}_{\theta }L(\theta )& ={\mathrm{\nabla }}_{\theta }{\mathbb{E}}_{q_{\theta }(\mathbf{z})}[l(\theta ,\mathbf{z})]={\mathrm{\nabla }}_{\theta }\int l(\theta ,\mathbf{z})q_{\theta }(\mathbf{z})\mathrm{d}\mathbf{z}\\ & =\int {\mathrm{\nabla }}_{\theta }l(\theta ,\mathbf{z})q_{\theta }(\mathbf{z})\mathrm{d}\mathbf{z}+\int l(\theta ,\mathbf{z}){\mathrm{\nabla }}_{\theta }{q}_{\theta }(\mathbf{z})\mathrm{d}\mathbf{z}\\ & =I_1+I_2.\end{array}$$

上式中的第一项可以使用蒙特卡洛法近似求解：

$$I_1=\int {\mathrm{\nabla }}_{\theta }l(\theta ,\mathbf{z})q_{\theta }(\mathbf{z})\mathrm{d}\mathbf{z}\approx \frac{1}{S}\sum _{s=1}^S{\mathrm{\nabla }}_{\theta }l(\theta ,{\mathbf{z}}_s),{\mathbf{z}}_s\sim q_{\theta }.$$

但第二项无法使用朴素蒙特卡洛法来近似，因为它是关于一个概率分布的梯度，而蒙特卡洛法只能计算某个函数关于采样的概率分布的期望。这里一般有两种主流办法来处理它。

2. 得分函数估计 Score Function Estimator

得分函数（score function）是对数似然函数关于概率分布参数的梯度，即

$$s_{\theta }(\mathbf{x})={\mathrm{\nabla }}_{\theta }\log q_{\theta }(\mathbf{x}).$$

通过使用对数导数技巧（log-derivative trick），可以将概率密度函数的梯度转化为概率密度函数与得分函数的乘积：

$${\mathrm{\nabla }}_{\theta }{q}_{\theta }(\mathbf{z})=q_{\theta }(\mathbf{z})\frac{{\mathrm{\nabla }}_{\theta }{q}_{\theta }(\mathbf{z})}{q_{\theta }(\mathbf{z})}=q_{\theta }(\mathbf{z}){\mathrm{\nabla }}_{\theta }\log q_{\theta }(\mathbf{z}).$$

上文中第二项 $I_2$ 可以写成：

$$\begin{array}{rl}{I}_2& =\int l(\theta ,\mathbf{z}){\mathrm{\nabla }}_{\theta }{q}_{\theta }(\mathbf{z})\mathrm{d}\mathbf{z}=\int l(\theta ,\mathbf{z})q_{\theta }(\mathbf{z}){\mathrm{\nabla }}_{\theta }\log q_{\theta }(\mathbf{z})\mathrm{d}\mathbf{z}\\ & ={\mathbb{E}}_{q_{\theta }(\mathbf{z})}[l(\theta ,\mathbf{z}){\mathrm{\nabla }}_{\theta }\log q_{\theta }(\mathbf{z})]\\ & \approx \frac{1}{S}\sum _{s=1}^{S}l(\theta ,{\mathbf{z}}_s){\mathrm{\nabla }}_{\theta }\log q_{\theta }\left({\mathbf{z}}_s\right),{\mathbf{z}}_s\sim q_{\theta }.\end{array}$$

3. 重参数化技巧 Reparameterization Trick

实践中，得分函数估计带来的计算方差往往太大。为了解决这个问题，可以使用重参数化方法。该方法被广泛应用到例如变分自编码器（VAE）这样的经典机器学习模型中。

3.1 重参数化怎么做？具体原理详解

对于高斯分布这类可参数化的标准分布，其分布参数与采样过程是可以独立分离的。它可以被分离为一组独立的参数与一个新的随机变量，分离后各个参数与新的随机变量无关。因而新的目标函数可以表示为新随机变量分布的期望，这种分离方法称为重参数化技巧。

3.2 高斯分布的参数化

这里以常用的高斯分布为例解释重参数化。假设随机变量 $\mathbf{x}\in {\mathbb{R}}^D$ 服从高斯分布：

$$\mathbf{x}\sim \mathcal{N}(\mathit{\mu },\mathbf{\Sigma }).$$

我们可以将这个高斯分布进一步拆解为来自于一个标准高斯分布 $\mathit{ϵ}\sim \mathcal{N}(\mathbf{0},\mathbf{I})$ 的线性变换：

$$\mathbf{x}=\mathit{\mu }+\mathbf{L}\mathit{ϵ}\sim \mathcal{N}(\mathit{\mu },\mathbf{L}{\mathbf{L}}^{\mathsf{T}}),$$

式中 $\mathbf{L}$ 是一个下三角方阵。当进一步简化假设各维度变量独立分布时，可以将 $\mathbf{L}$ 退化为一个对角矩阵：

$$\mathbf{L}=\left[\begin{array}{cccc}{\sigma }_{11}& & & \\ {\sigma }_{21}& {\sigma }_{22}& & \\ ⋮& ⋮& \ddots & \\ {\sigma }_{D1}& {\sigma }_{D2}& \cdots & {\sigma }_{DD}\end{array}\right]\approx \left[\begin{array}{cccc}{\sigma }_{11}& & & \\ & {\sigma }_{22}& & \\ & & \ddots & \\ & & & {\sigma }_{DD}\end{array}\right]=\mathrm{diag}(\mathit{\sigma }).$$

通过这种参数化方法，可以使用有限个自然参数来近似表征复杂连续的概率分布。在上述例子中，对于这个高斯分布，仅需要 $2D$ 个参数来描述即可。

当然，我们可以根据数据分布和需要，对概率分布进行多种形式的分解，比如 Gamma 分布、Beta 分布以及 Student-$t$ 分布等等。这样可以把概率分布簇的参数与概率分布自身的随机性拆分和剥离，从而有益于计算和求导。

具体的实用例子：在 VAE 中，假设算法中的隐变量服从高斯分布 $\mathbf{z}\sim q_{\varphi }(\mathbf{z}\mid \mathbf{x})$，那么这个隐变量可以被重参数化表达为

$$\mathbf{z}={\mathit{\mu }}_{\varphi }(\mathbf{x})+\mathit{ϵ}\odot {\mathit{\sigma }}_{\varphi }(\mathbf{x}),$$

从而使目标函数从原来的 ${\mathbb{E}}_{q_{\varphi }(\mathbf{z}\mid \mathbf{x})}[f(\mathbf{z})]$ 变为 ${\mathbb{E}}_{q(\mathit{ϵ})}[f(\mathbf{z})]$。

3.3 使用条件

重参数化法的使用条件是：

• 目标函数 $l(\theta ,z)$ 对参数 $z$ 和 $\theta$ 都可微分；
• 分布 $z$ 的采样可以从某个独立分布 $ϵ$ 获得，然后使用一组参数 $\theta$ 将其转化为分布 $z$，即 $z=r(\theta ,ϵ)$。

3.4 重参数化梯度

将重参数化后的形式代入，原先的随机目标函数可以变为：

$$\begin{array}{rl}L(\theta )& ={\mathbb{E}}_{q_{\theta }(z)}[l(\theta ,z)]=\int l(\theta ,z)q_{\theta }(z)\mathrm{d}z\\ & =\int l(\theta ,r(\theta ,ϵ))q(ϵ)\mathrm{d}ϵ={\mathbb{E}}_{q(ϵ)}[l(\theta ,r(\theta ,ϵ))].\end{array}$$

那么我们就可以计算其梯度：

$$\begin{array}{rl}{\mathrm{\nabla }}_{\theta }L(\theta )& ={\mathrm{\nabla }}_{\theta }{\mathbb{E}}_{q(ϵ)}[l(\theta ,r(\theta ,ϵ))]={\mathbb{E}}_{q(ϵ)}[{\mathrm{\nabla }}_{\theta }l(\theta ,r(\theta ,ϵ))]\\ & \approx \frac{1}{S}\sum _{s=1}^S{\mathrm{\nabla }}_{\theta }l(\theta ,r(\theta ,{ϵ}_s)),{ϵ}_s\sim q.\end{array}$$

这个梯度称为重参数化梯度（reparameterization gradient / pathwise derivative）。

4. 为什么要用重参数化？两类优势

4.1 优势一：穿透随机性的求导

我们可以把概率分布重参数化后的参数理解为概率分布的“坐标”，因而可以把目标函数对于概率分布的梯度，转移到目标函数对这些坐标分量的梯度计算上去。对这些自然参数的求导有着明确的数学含义，比如 $\frac{\partial }{\partial \mu }$ 是指概率分布的均值变化单位数值所带来的目标函数的变化量。

通过这种方法，可以近似认为：计算某个函数对该分布的导数，等价于计算该函数对该概率分布自然参数的导数。而在很多机器学习算法的实际使用中，假如最终的目标函数与某个概率分布有关，那么对目标函数的梯度就可以一直向后传递至对该概率分布的自然参数中，而不受随机性带来的影响。由此，可以便捷地构建基于梯度下降优化的各种机器学习模型（例如 VAE），从而解决一系列实用性问题。

（图 1：重参数化求导的原理解释图）

4.2 优势二：降低方差

通过重参数化，可以解决使用 Score Function Method 进行求取梯度时可能导致的高方差问题 [1]。具体来说，分别展开两种方法的梯度公式。在 Score Function Method 中：

$$\begin{array}{rl}{\mathrm{\nabla }}_{\theta }{L}_{\mathrm{SF}}(\theta )& =\int {\mathrm{\nabla }}_{\theta }l(\theta ,z)q_{\theta }(z)\mathrm{d}z+\int l(\theta ,z){\mathrm{\nabla }}_{\theta }{q}_{\theta }(z)\mathrm{d}z\\ & ={\mathbb{E}}_{q_{\theta }(z)}[{\mathrm{\nabla }}_{\theta }l(\theta ,z)+l(\theta ,z){\mathrm{\nabla }}_{\theta }\log q_{\theta }(z)]\\ & ={\mathbb{E}}_{q_{\theta }(z)}[\frac{\partial l(\theta ,z)}{\partial \theta }+\frac{\partial l(\theta ,z)}{\partial z}\frac{\partial z}{\partial \theta }+l(\theta ,z)\frac{\partial \log q_{\theta }(z)}{\partial \theta }].\end{array}$$

在 Reparameterization Method 中：

$$\begin{array}{rl}{\mathrm{\nabla }}_{\theta }{L}_{\mathrm{Reparam}}(\theta )& ={\mathrm{\nabla }}_{\theta }\int l(\theta ,z)q_{\theta }(z)\mathrm{d}z={\mathrm{\nabla }}_{\theta }\int l(\theta ,r(\theta ,ϵ))q(ϵ)\mathrm{d}ϵ\\ & ={\mathbb{E}}_{q(ϵ)}[{\mathrm{\nabla }}_{\theta }l(\theta ,r(\theta ,ϵ))]\\ & ={\mathbb{E}}_{q(ϵ)}[\frac{\partial l(\theta ,r)}{\partial \theta }+\frac{\partial l(\theta ,r)}{\partial r}\frac{\partial r(\theta ,ϵ)}{\partial \theta }].\end{array}$$

对比上述两式，Score Function Method 会额外引入：

$${\mathbb{E}}_{q_{\theta }(z)}[l(\theta ,z){\mathrm{\nabla }}_{\theta }\log q_{\theta }(z)].$$

当 $l(\theta ,z)$ 和分布无关时，该项为零。但在一般情况下，该项的存在会带来额外的梯度方差。因此一般来说：

$${\mathbb{V}}_{q(ϵ)}(\frac{\partial l(\theta ,r)}{\partial \theta }+\frac{\partial l(\theta ,r)}{\partial r}\frac{\partial r(\theta ,ϵ)}{\partial \theta })\le {\mathbb{V}}_{q_{\theta }(z)}(\frac{\partial l(\theta ,z)}{\partial \theta }+\frac{\partial l(\theta ,z)}{\partial z}\frac{\partial z}{\partial \theta }+l(\theta ,z)\frac{\partial \log q_{\theta }(z)}{\partial \theta }).$$

而重参数化之所以可以避免这一项，是因为使用了独立于分布参数的采样，故此项梯度为 0。

（不过，我们可以构造某些特殊的函数，让上式中的各项产生相关性，从而使得“重参数化法的梯度方差比 Score Function Method 更小”这个命题不成立。但在统计意义上，即实际应用中的绝大多数情况，使用重参数化可以有效降低方差。）

此外，在实践中，较低的采样样本数会带来较大的方差，但较高的采样样本数又会降低计算效率。因此，如果使用重参数化技巧直接对目标函数进行求导，降低梯度的方差，就可以让我们使用更少的采样样本，从而也对提升训练效率有帮助。

5. 强化学习中的重参数化：SAC 实例

在强化学习 SAC 算法中 [2]，策略通过 SAC 中定义的 soft value function（即引入最大熵的价值函数）来指导策略优化。具体来说，SAC 算法希望找到可以最大化 soft value function 的策略，即：

$$\begin{array}{rl}{\pi }^{\ast }& =\arg \underset{\pi }{max}{V}_{\pi }(s),\\ V_{\pi }(s)& ={\mathbb{E}}_{a\sim \pi }[Q_{\pi }(s,a)]+\alpha H(\pi (\cdot \mid s))\\ & ={\mathbb{E}}_{a\sim \pi }[Q_{\pi }(s,a)-\alpha \log \pi (a\mid s)].\end{array}$$

在实现中，首先需要策略部分（Actor）输出当前状态对应的动作，然后交给价值函数（Critic）判断好坏得到目标函数，并用梯度下降进行优化。梯度将会从 Critic 沿着动作一路回传到 Actor。

在连续动作空间上，需要对策略使用高斯分布进行重参数化，并使用 $\tanh$ 函数压缩最终动作的输出范围到 $[-1,1]$，即采样到的动作通过下式获得：

$$a_{\theta }(s,ϵ)=\tanh ({\mu }_{\theta }(s)+ϵ\cdot {\sigma }_{\theta }(s)),ϵ\sim \mathcal{N}(0,\mathbf{I}).$$

需要注意的是，这里需要一个额外的矫正操作，因为使用 $\tanh$ 函数压缩最终动作 $a$ 会改变原本无界动作 $a_u$ 的概率密度函数，即：

$$\pi (a\mid s)=p(a_u\mid s){|det\frac{\mathrm{d}a}{\mathrm{d}{a}_u}|}^{-1},$$

其中

$$\frac{\mathrm{d}a}{\mathrm{d}{a}_u}=\mathrm{diag}(1-{\tanh }^2(a_u)).$$

所以代码实现中，我们需要修改压缩后的概率密度函数的 log prob 为：

$$\log \pi (a\mid s)=\log p(a_u\mid s)-\sum _{d=1}^D\log (1-{\tanh }^2(a_{u,d})).$$

考虑 SAC 算法中使用的 double Q-trick，最终重参数化后的策略学习目标函数为：

$${\theta }^{\ast }=\arg \underset{\theta }{max}{\mathbb{E}}_{\begin{array}{c}s\\ ϵ\sim \mathcal{N}(0,\mathbf{I})\end{array}}[\underset{j=1,2}{min}{Q}_{{\varphi }_j}(s,a_{\theta }(s,ϵ))-\alpha \log {\pi }_{\theta }(a_{\theta }(s,ϵ)\mid s)].$$

对于一个 batch 为 $B$ 的数据，这一批次数据产生的梯度为：

$$\begin{array}{rl}{\mathrm{\nabla }}_{\theta }J(\theta )& ={\mathrm{\nabla }}_{\theta }\frac{1}{|B|}\sum _{s\in B}(\underset{j=1,2}{min}{Q}_{{\varphi }_j}(s,a_{\theta }(s,ϵ))-\alpha \log {\pi }_{\theta }(a_{\theta }(s,ϵ)\mid s))\\ & =\frac{1}{|B|}\sum _{s\in B}({\mathrm{\nabla }}_{\theta }\left[\underset{j=1,2}{min}{Q}_{{\varphi }_j}(s,a_{\theta }(s,ϵ))\right]-\alpha {\mathrm{\nabla }}_{\theta }\log {\pi }_{\theta }(a_{\theta }(s,ϵ)\mid s)).\end{array}$$

6. PPO 与重参数化的关系

回顾 PPO 的目标函数：

$${\mathbb{E}}_t[min(\frac{p_{\theta }(a_t\mid s_t)}{p_{{\theta }_k}(a_t\mid s_t)}{\hat{A}}^{{\theta }_k}(s_t,a_t),\mathrm{clip}(\frac{p_{\theta }(a_t\mid s_t)}{p_{{\theta }_k}(a_t\mid s_t)},1-ϵ,1+ϵ){\hat{A}}^{{\theta }_k}(s_t,a_t))].$$

可以发现，虽然公式中存在着 $p_{\theta }(a\mid s)$ 的形式，但是与 SAC 中不同：训练时用到的动作 $a_t$ 并不由将要更新的策略函数的参数 $\theta$ 生成，而是由旧策略（收集数据的策略）函数的参数 ${\theta }_k$ 生成。因此，在 PPO 中仅是借用重参数化的形式，将策略函数参数化为某个概率函数（比如连续动作空间使用高斯分布），即可让梯度通过概率函数的定义式直接回传。公式中的动作 $a_t$、状态 $s_t$、优势函数估计 ${\hat{A}}^{{\theta }_k}$ 只是为了更新参数 $\theta$ 而存在的数据，它们本身并不回传任何梯度。

$$\text{Data: }⟨s,a,{\hat{A}}_{{\theta }_k},{\mathrm{logit}}_{{\theta }_k},\dots ⟩\text{Loss: }\frac{{\pi }_{\theta }(a\mid s)}{{\pi }_{{\theta }_k}(a\mid s)}{\hat{A}}_{{\theta }_k}.$$

（图 2：PPO 策略部分优化计算图，只有蓝色部分回传梯度，紫色部分只是参与计算。）

参考文献

1. Xu M, Quiroz M, Kohn R, et al. Variance reduction properties of the reparameterization trick[C]//The 22nd International Conference on Artificial Intelligence and Statistics. PMLR, 2019: 2711-2720.
2. Haarnoja T, Zhou A, Abbeel P, et al. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor[C]//International Conference on Machine Learning. PMLR, 2018: 1861-1870.