Trust Region Policy Optimization (TRPO) 背后的数学原理

  本文是自己的TRPO算法学习笔记，在数学原理推导核心部分附有自己的理解与解释。整篇文章逻辑清晰，思路顺畅。有想推导的同学可以一起学习。

  TRPO和PPO都是基于Minorize-Maximization MM的算法。

Surrogate function

  RL中期望maximizing the expected discounted rewards，期望折扣奖励 $\eta$η 可用如下数学公式表示：

$\eta\left(\pi_{\theta}\right)=\underset{\tau \sim \pi_{\theta}}{\mathrm{E}}\left[\sum_{t=0}^{\infty} \gamma^{t} r_{t}\right]$η(πθ​)=τ∼πθ​E​[t=0∑∞​γtrt​]

  我们希望去找到一个surrogate function(替代函数)，替代函数具有如下性质：

• a lower bound function for $\eta$η；(是$\eta$η的一个下界函数)
• approximate $\eta$η at the current policy (在当前策略下能够近似等于 $\eta$η)
• easy to optimize. (找这样一个替代函数的目的就是方便优化)

  其图形表示为下图所示(蓝色表示下界函数，红色表示期望折扣奖励)：

  在每一次迭代中，希望去找到一个对于下界函数$M$M来说的最优点，并把它当作当前的策略，如下图所示：

  之后，我们基于新的policy，重新评估下界(re-evaluate a lower bound)，并重复迭代。重复上述过程，策略将会被持续改进。由于策略集有限，所以所以最终将会收敛到局部最优或者全局最优，整个流程如下图所示：

  整体的目标就是在原有的参数空间 $\theta$θ上很难去计算最优值，我们期望用一个替代函数来作为它的lower bound，下界函数比较好优化，然后通过迭代的方式让其逼近原始最优解。

Objective function

  上述想法就很不错，现在我们需要去寻找目标函数和替代函数了。

  首先我们去定义Q-Value function，Value function和Advantage function，如下：

$Q_{\pi}(s_{t},a_{t}) = \mathbb{E_{s_{t+1},a_{t+1},\cdots}}\left[ \sum_{l=0}^{\infty} \gamma^{l}r(s_{t+l}) \right]$Qπ​(st​,at​)=Est+1​,at+1​,⋯​[l=0∑∞​γlr(st+l​)]

$V_{\pi}(s_{t}) = \mathbb{E}_{a_{t},s_{t+1},\cdots} \left[ \sum_{l=0}^{\infty} \gamma^{l}r(s_{t+1}) \right]$Vπ​(st​)=Eat​,st+1​,⋯​[l=0∑∞​γlr(st+1​)]

$A_{\pi}(s,a) = Q_{\pi}(s,a) -V_{\pi}(s)$Aπ​(s,a)=Qπ​(s,a)−Vπ​(s)

  其中 $a_{t} \sim \pi(a_{t}|s_{t})$at​∼π(at​∣st​)，$s_{t+1} \sim P(s_{t+1}|s_{t},a_{t})$st+1​∼P(st+1​∣st​,at​) ,$t \geq 0$t≥0。

Expected discounted reward

  期望折扣奖励 $\eta$η 可被表示为：

$\eta(\pi) = \mathbb{E}_{s_{0},a_{0},\cdots} \left[ \sum_{t=0}^{\infty} \gamma^{t} r(s_{t}) \right]$η(π)=Es0​,a0​,⋯​[t=0∑∞​γtr(st​)]

  其中，$s_{0} \sim \rho_{0}(s_{0})$s0​∼ρ0​(s0​)，$a_{t} \sim \pi(a_{t}|s_{t})$at​∼π(at​∣st​)，$s_{t+1} \sim P(s_{t+1}|s_{t},a_{t})$st+1​∼P(st+1​∣st​,at​)

  由于我们要比较更新前后的两个策略，从而保证策略一直在进步，所以作者这里是将新老策略写到了一个公式中：

$\eta(\tilde{\pi})=\eta(\pi)+\sum_{s} \rho_{\tilde{\pi}}(s) \sum_{a} \tilde{\pi}(a | s) A_{\pi}(s, a)$η(π~)=η(π)+s∑​ρπ~​(s)a∑​π~(a∣s)Aπ​(s,a)

  其中 $\rho_{\pi}$ρπ​ be the (unnormalized) discounted visitation frequencies (任意时刻状态$s$s的访问概率和)，其展开形式表示为：

$\rho_{\pi}(s) = P(s_{0}=s) +\gamma P(s_{1} =s) + \gamma^{2} P(s_{2}=s) + \cdots$ρπ​(s)=P(s0​=s)+γP(s1​=s)+γ2P(s2​=s)+⋯

  如果能保证$\sum_{s} \rho_{\tilde{\pi}}(s) \sum_{a} \tilde{\pi}(a | s) A_{\pi}(s, a)$∑s​ρπ~​(s)∑a​π~(a∣s)Aπ​(s,a)大于0，那新的策略下的期望折扣奖励就一直是在进步的。那这里就有两个问题了，1. 上述新老策略写在一个公式里面的证明呢？2. 如何去保证后面的附加项大于0?

  证明：

$\begin{array}{l} \underset{\tau \sim \tilde{\pi}}{\mathrm{E}}\left[\sum_{t=0}^{\infty} \gamma^{t} A^{\pi}\left(s_{t}, a_{t}\right)\right] \\ =\underset{\tau \sim \tilde{\pi}}{\mathrm{E}}\left[\sum_{t=0}^{\infty} \gamma^{t}\left(R\left(s_{t}, a_{t}, s_{t+1}\right)+\gamma V^{\pi}\left(s_{t+1}\right)-V^{\pi}\left(s_{t}\right)\right)\right] \\ =\eta\left(\tilde{\pi}\right)+\underset{\tau \sim \tilde{\pi}}{\mathrm{E}}\left[\sum_{t=0}^{\infty} \gamma^{t+1} V^{\pi}\left(s_{t+1}\right)-\sum_{t=0}^{\infty} \gamma^{t} V^{\pi}\left(s_{t}\right)\right] \\ =\eta\left(\tilde{\pi}\right)+\underset{\tau \sim \tilde{\pi}}{\mathrm{E}}\left[\sum_{t=1}^{\infty} \gamma^{t} V^{\pi}\left(s_{t}\right)-\sum_{t=0}^{\infty} \gamma^{t} V^{\pi}\left(s_{t}\right)\right] \\ =\eta\left(\tilde{\pi}\right)-\underset{\tau \sim \tilde{\pi}}{\mathrm{E}}\left[V^{\pi}\left(s_{0}\right)\right] \\ =\eta\left(\tilde{\pi}\right)-\eta(\pi) \end{array}$τ∼π~E​[∑t=0∞​γtAπ(st​,at​)]=τ∼π~E​[∑t=0∞​γt(R(st​,at​,st+1​)+γVπ(st+1​)−Vπ(st​))]=η(π~)+τ∼π~E​[∑t=0∞​γt+1Vπ(st+1​)−∑t=0∞​γtVπ(st​)]=η(π~)+τ∼π~E​[∑t=1∞​γtVπ(st​)−∑t=0∞​γtVπ(st​)]=η(π~)−τ∼π~E​[Vπ(s0​)]=η(π~)−η(π)​

  由此我们只剩下了第二点，从某个策略$\pi$π出发，通过计算找到一个策略$\tilde{\pi}$π~，使得：

$\sum_{s} \rho_{\tilde{\pi}}(s) \sum_{a} \tilde{\pi}(a | s) A_{\pi}(s, a) \geq 0$s∑​ρπ~​(s)a∑​π~(a∣s)Aπ​(s,a)≥0

  即可使得 $\eta(\tilde{\pi}) \geq \eta(\pi)$η(π~)≥η(π)，也就是说策略改变之后，整体的收益也会增加，从而实现单调递增。那现在所有的问题都转化到了如何使得$\sum_{s} \rho_{\tilde{\pi}}(s) \sum_{a} \tilde{\pi}(a | s) A_{\pi}(s, a) \geq 0$∑s​ρπ~​(s)∑a​π~(a∣s)Aπ​(s,a)≥0？

Function $\mathcal{L}$L

  $\sum_{s} \rho_{\tilde{\pi}}(s) \sum_{a} \tilde{\pi}(a | s) A_{\pi}(s, a)$∑s​ρπ~​(s)∑a​π~(a∣s)Aπ​(s,a)在实际中几乎是不可行的，因为公式中包含$\rho_{\tilde{\pi}}(s)$ρπ~​(s)，也就是说我们需要按照新的策略$\tilde{\pi}$π~与环境交互得到状态$s$s的访问频率，但是这个新的策略$\tilde{\pi}$π~是我们需要去求解的策略。也就是说如果要做的话，我们需要先确定新的策略$\tilde{\pi}$π~，然后使用这个新的策略得到一定量样本，并最终通过这些样本统计判断这个策略能够满足上述要求，使得策略递增。我们需要不断地去尝试每一个可能的新策略。显然这种做法非常低效。

  于是需要去找与上述公式的近似且可解的形式，定义function $\mathcal{L}$L：

$\mathcal{L}_{\pi}(\tilde{\pi}) =\eta({\pi}) + \sum_{s} \rho_{\pi}(s) \sum_{a} \tilde{\pi}(a | s) A_{\pi}(s, a)$Lπ​(π~)=η(π)+s∑​ρπ​(s)a∑​π~(a∣s)Aπ​(s,a)

  与之前的公式：

$\eta(\tilde{\pi})=\eta(\pi)+\sum_{s} \rho_{\tilde{\pi}}(s) \sum_{a} \tilde{\pi}(a | s) A_{\pi}(s, a)$η(π~)=η(π)+s∑​ρπ~​(s)a∑​π~(a∣s)Aπ​(s,a)

  对比，我们可以发现，两者的不同仅仅在于状态访问概率$\rho_{\tilde{\pi}}(s)$ρπ~​(s)、$\rho_{\pi}(s)$ρπ​(s)的不同。那$\mathcal{L}_{\pi}(\tilde{\pi})$Lπ​(π~)能否满足要求呢？其实两者的数值和导数方向都是相同的，那么用$\mathcal{L}_{\pi}(\tilde{\pi})$Lπ​(π~)代替原始目标函数也是可以的，要求更新的幅度不要太大就好。

  上面说了这么多，其实就是为了说明我们更新的幅度不要太大，因为更新大了之后，上述近似函数就无法成立，$\mathcal{L}_{\pi}(\tilde{\pi})$Lπ​(π~)无法成立的话，你所拿策略$\pi$π采样得到的样本就没用了，因为实际上样本是要去新的策略$\tilde{\pi}$π~里面去采样的，只是因为做了近似才可以用老的策略$\pi$π去采样。所以那我们怎么来保证其更新幅度不要太大呢？

  从之前的分析中可以知道$\mathcal{L}$L是下界函数(bound function)$M$M中的一部分，$M$M中的另外一项是KL散度(KL-divergence)：

$D_{KL}(P|Q|) = \mathbb{E}_{x} \text{log} \frac{P(x)}{Q(x)}$DKL​(P∣Q∣)=Ex​logQ(x)P(x)​

  因此我们把策略模型看成一个概率分布，使用KL散度表示两个分布的距离。那两者之间有什么关系呢？原论文中作者用两页纸证明了 $\boldsymbol{\eta}$η 的下界：

$\eta (\pi_{new}) \geq L_{\pi_{old}}(\pi_{new}) -\frac{4 \varepsilon \gamma}{(1-\gamma)^{2}} \alpha^{2}$η(πnew​)≥Lπold​​(πnew​)−(1−γ)24εγ​α2

  其中 $\varepsilon = \max_{s,a}|A_{\pi}(s,a)|$ε=maxs,a​∣Aπ​(s,a)∣，$\alpha = D_{TV}^{max}(\pi_{old},\pi_{new})$α=DTVmax​(πold​,πnew​)，$D_{TV}^{max}(\pi, \tilde{\pi} ) = max D_{TV}(\pi(\cdot|s)||\tilde{\pi}(\cdot|s))$DTVmax​(π,π~)=maxDTV​(π(⋅∣s)∣∣π~(⋅∣s))，$D_{TV}(p||q) = \frac{1}{2} \sum_{i}|p_{i}-q_{i}|$DTV​(p∣∣q)=21​∑i​∣pi​−qi​∣。

  $D_{TV}$DTV​ 是total variation divergence，由于：

$D_{TV}(p||q)^{2} \leq D_{KL}(p||q)$DTV​(p∣∣q)2≤DKL​(p∣∣q)

  得到新的下界函数(lower bound function)：

$\eta (\tilde{\pi}) \geq L_{\pi}(\tilde{\pi}) -CD_{KL}^{max}(\pi,\tilde{\pi})$η(π~)≥Lπ​(π~)−CDKLmax​(π,π~)

  其中 $C = \frac{4 \varepsilon \gamma}{(1-\gamma)^{2}}$C=(1−γ)24εγ​，$D_{KL}^{max}(\pi,\tilde{\pi}) = max_{s}D_{KL}(\pi(\cdot|s) || \tilde{\pi}(\cdot |s))$DKLmax​(π,π~)=maxs​DKL​(π(⋅∣s)∣∣π~(⋅∣s))

Monotonically improving guarantee

  What we really prove here is the new policy generated from optimizing $M$M will have a guarantee that it will perform better in $\eta$η (the real expected rewards) than the old policy. Since there are only finite policies, the continuous improvement will only lead us to a local or a global optimal point.

  令$M_{i}(\pi) = L_{\pi_{i}}(\pi) -CD_{KL}^{max}(\pi_{i},\pi)$Mi​(π)=Lπi​​(π)−CDKLmax​(πi​,π)，有$\eta({\pi_{i+1}}) \geq M_{i}(\pi_{i+1})$η(πi+1​)≥Mi​(πi+1​)，$\eta({\pi_{i}}) = M_{i}(\pi_{i})$η(πi​)=Mi​(πi​)，所以可以得到：

$\eta({\pi_{i+1}}) - \eta({\pi_{i}}) \geq M_{i}(\pi_{i+1}) - M_{i}(\pi_{i})$η(πi+1​)−η(πi​)≥Mi​(πi+1​)−Mi​(πi​)

  则$\pi_{i+1} =arg\max_{\pi}M_{i}(\pi)$πi+1​=argmaxπ​Mi​(π)时，期望折扣奖励将在下一次迭代被提升。

  由此我们可以得到保证策略提升的算法。Here is the iteration algorithm that guarantees that the new policy will always perform better than the current one.

  此时算法的目标函数变为：

$maximize_{\tilde{\pi}}[L_{\pi}(\tilde{\pi}) -CD_{KL}^{max}(\pi,\tilde{\pi})]$maximizeπ~​[Lπ​(π~)−CDKLmax​(π,π~)]

  上式过于保守，将其做一些转变，得到有约束条件的优化目标：

$maximize_{\pi}L_{\pi_{old}}(\pi)$maximizeπ​Lπold​​(π)

$s.t. D_{KL}^{max}(\pi_{old},\pi) \leq \delta$s.t.DKLmax​(πold​,π)≤δ

  在公式中需要对最大值进行约束，而最大值表示为KL散度上界，这实际上相当于对所有状态的KL散度进行约束，这样约束条件会变得多而复杂，将最大值变成均值理论上有所放松，但实际效果还好，于是有了：

  优化目标$\mathcal{L}_{\pi}(\tilde{\pi})$Lπ​(π~)为：

$\mathcal{L}_{\pi}(\tilde{\pi}) =\eta({\pi}) + \sum_{s} \rho_{\pi}(s) \sum_{a} \tilde{\pi}(a | s) A_{\pi}(s, a)$Lπ​(π~)=η(π)+s∑​ρπ​(s)a∑​π~(a∣s)Aπ​(s,a)

  而其中$\tilde{\pi}(a | s)$π~(a∣s)与新策略有关，无法对其采样，由此我们引入重要性采样：

$\sum_{a} \pi_{\theta}\left(a | s_{n}\right) A_{\theta_{\text {old }}}\left(s_{n}, a\right)=\mathbb{E}_{a \sim q}\left[\frac{\pi_{\theta}\left(a | s_{n}\right)}{q\left(a | s_{n}\right)} A_{\theta_{\text {old }}}\left(s_{n}, a\right)\right]$a∑​πθ​(a∣sn​)Aθold ​​(sn​,a)=Ea∼q​[q(a∣sn​)πθ​(a∣sn​)​Aθold ​​(sn​,a)]

  the objective can be rewritten as：

$\begin{array}{c} \underset{\theta}{\operatorname{maximize}} \mathbb{E}_{s \sim \rho_{\theta_{\text {old }}}, a \sim q}\left[\frac{\pi_{\theta}(a | s)}{q(a | s)} \hat{A}_{\theta_{\text {old }}}(s, a)\right] \\ \text { subject to } \mathbb{E}_{s \sim \rho_{\theta_{\text {old }}}}\left[D_{\mathrm{KL}}\left(\pi_{\theta_{\text {old }}}(\cdot | s) \| \pi_{\theta}(\cdot | s)\right)\right] \leq \delta \end{array}$θmaximize​Es∼ρθold ​​,a∼q​[q(a∣s)πθ​(a∣s)​A^θold ​​(s,a)] subject to Es∼ρθold ​​​[DKL​(πθold ​​(⋅∣s)∥πθ​(⋅∣s))]≤δ​

  With Lagrangian duality, a constraint for an objective function can be integrated back to the objective function with a multiplier. Both are mathematically the same(利用拉格朗日对偶，把约束项提到目标函数中去):

直觉的看法

  在之前的Gradient ascent方法中都是选择了梯度的方向，如下图中的左图所示，但是更新步长如果选取地不好很容易掉入深渊：

  在TRPO中限制了更新步长，并且在数学上证明了会收敛到局部最优或者全局最优：

• 参考文献：
• https://medium.com/@jonathan_hui/rl-the-math-behind-trpo-ppo-d12f6c745f33
• 强化学习精要核心算法与Tensorflow实现。