Maximum Entropy RL

最大熵强化学习

基础定义

为了防止策略网络输出的概率集中在一个动作上，平衡探索性能，可以通过熵来衡量策略分布的不确定性：

$\mathcal{H}(\pi \mid s) = -\mathcal{E}_{a \sim \pi(\cdot \mid s)} \ln \pi(a \mid s) = -\sum_{a} \pi(a \mid s) \ln \pi(a \mid s)$

原始的 RL 问题的优化目标为：

$J(\pi) = \mathcal{E}_{s_{0} \sim b_{0}(\cdot)} \mathcal{E}_{a_{0} \sim \pi(\cdot \mid s_{0})} \mathcal{E}_{s_{1} \sim p(\cdot \mid s_{0},\ a_{0})} \mathcal{E}_{a_{1} \sim \pi(\cdot \mid s_{1})} \cdots \mathcal{E}_{s_{\mathrm{T}} \sim p(\cdot \mid s_{\mathrm{T} - 1},\ a_{\mathrm{T} - 1})} \mathcal{E}_{a_{\mathrm{T}} \sim \pi(\cdot \mid s_{\mathrm{T}})} \left[ \sum_{t = 0}^{\mathrm{T}} \gamma^{t} \mathcal{R}(s_{t},\ a_{t}) \right]$

在最大熵 RL 框架下，优化目标加入了熵正则项来增强探索程度，同时减少陷入局部最优的可能性：

$\begin{aligned} \Omega(\pi) &= \mathcal{E}_{s_{0} \sim b_{0}(\cdot)} \mathcal{E}_{a_{0} \sim \pi(\cdot \mid s_{0})} \mathcal{E}_{s_{1} \sim p(\cdot \mid s_{0},\ a_{0})} \mathcal{E}_{a_{1} \sim \pi(\cdot \mid s_{1})} \cdots \mathcal{E}_{s_{\mathrm{T}} \sim p(\cdot \mid s_{\mathrm{T} - 1},\ a_{\mathrm{T} - 1})} \mathcal{E}_{a_{\mathrm{T}} \sim \pi(\cdot \mid s_{\mathrm{T}})} \left[ \sum_{t = 0}^{\mathrm{T}} \gamma^{t} \mathcal{H}(\pi \mid s_{t}) \right] \\[7mm] &= \sum_{t = 0}^{\mathrm{T}} \gamma^{t} \mathcal{E}_{s_{0}} \mathcal{E}_{a_{0}} \mathcal{E}_{s_{1}} \mathcal{E}_{a_{1}} \cdots \mathcal{E}_{s_{t}} \Big[ \mathcal{H}(\pi \mid s_{t}) \Big] = \sum_{t = 0}^{\mathrm{T}} \gamma^{t} \mathcal{E}_{s_{0}} \mathcal{E}_{a_{0}} \mathcal{E}_{s_{1}} \mathcal{E}_{a_{1}} \cdots \mathcal{E}_{s_{t}} \Big[ -\mathcal{E}_{a_{t} \sim \pi(\cdot \mid s_{t})} \ln \pi(a_{t} \mid s_{t}) \Big] \\[7mm] &= -\sum_{t = 0}^{\mathrm{T}} \gamma^{t} \mathcal{E}_{s_{0}} \mathcal{E}_{a_{0}} \mathcal{E}_{s_{1}} \mathcal{E}_{a_{1}} \cdots \mathcal{E}_{s_{t}} \mathcal{E}_{a_{t}} \Big[ \ln \pi(a_{t} \mid s_{t}) \Big] = -\mathcal{E}_{s_{0}} \mathcal{E}_{a_{0}} \mathcal{E}_{s_{1}} \mathcal{E}_{a_{1}} \cdots \mathcal{E}_{s_{\mathrm{T}}} \mathcal{E}_{a_{\mathrm{T}}} \left[ \sum_{t = 0}^{\mathrm{T}} \gamma^{t} \ln \pi(a_{t} \mid s_{t}) \right] \end{aligned}$

总体优化目标为：

$J_{\mathcal{H}}(\pi) = J(\pi) + \alpha \Omega(\pi) = \mathcal{E}_{s_{0}} \mathcal{E}_{a_{0}} \mathcal{E}_{s_{1}} \mathcal{E}_{a_{1}} \cdots \mathcal{E}_{s_{\mathrm{T}}} \mathcal{E}_{a_{\mathrm{T}}} \left[ \sum_{t = 0}^{\mathrm{T}} \gamma^{t} \Big( \mathcal{R}(s_{t},\ a_{t}) - \alpha \ln \pi(a_{t} \mid s_{t}) \Big) \right]$

Soft Bellman Expectation Equation

为了方便后续推导，定义带熵正则的动作价值函数为：

$q_{\mathcal{H}}^{(t)}(s_{t},\ a_{t}) = \mathcal{E}_{s_{t + 1}} \mathcal{E}_{a_{t + 1}} \cdots \mathcal{E}_{s_{\mathrm{T}}} \mathcal{E}_{a_{\mathrm{T}}} \left[ \sum_{\tau = t}^{\mathrm{T}} \gamma^{\tau - t} \mathcal{R}(s_{\tau},\ a_{\tau}) - \alpha \sum_{\tau = t + 1}^{\mathrm{T}} \gamma^{\tau - t} \ln \pi(a_{\tau} \mid s_{\tau}) \right]$

以上动作价值的定义中没有加入 $\ln \pi(a_{t} \mid s_{t})$ 项，原因是在状态 $s_{t}$ 下给动作 $a_{t}$ 后该项失去了指导策略熵提升的作用。类似地，定义带熵正则的状态价值函数为：

$v_{\mathcal{H}}^{(t)}(s_{t}) = \mathcal{E}_{a_{t}} \mathcal{E}_{s_{t + 1}} \mathcal{E}_{a_{t + 1}} \cdots \mathcal{E}_{s_{\mathrm{T}}} \mathcal{E}_{a_{\mathrm{T}}} \left[ \sum_{\tau = t}^{\mathrm{T}} \gamma^{\tau - t} \mathcal{R}(s_{\tau},\ a_{\tau}) - \alpha \sum_{\tau = t}^{\mathrm{T}} \gamma^{\tau - t} \ln \pi(a_{\tau} \mid s_{\tau}) \right]$

通过定义可以将二者进行展开，得到 MaxEntRL 下的贝尔曼期望方程（soft）：

$q_{\mathcal{H}}^{(t)}(s_{t},\ a_{t}) = \mathcal{R}(s_{t},\ a_{t}) + \gamma \mathcal{E}_{s_{t + 1}} \Big[ v_{\mathcal{H}}^{(t + 1)}(s_{t + 1}) \Big] \qquad v_{\mathcal{H}}^{(t)}(s_{t}) = \alpha \mathcal{H}(\pi \mid s_{t}) + \mathcal{E}_{a_{t}} \Big[ q_{\mathcal{H}}^{(t)}(s_{t},\ a_{t}) \Big]$

在无限期规划下可以重写为：

$q_{\mathcal{H}}(s,\ a) = \mathcal{R}(s,\ a) + \gamma \mathcal{E}_{s'} \Big[ v_{\mathcal{H}}(s') \Big] \qquad v_{\mathcal{H}}(s) = \alpha \mathcal{H}(\pi \mid s) + \mathcal{E}_{a} \Big[ q_{\mathcal{H}}(s,\ a) \Big]$

在该定义下的贝尔曼期望方程（soft）算子 $\mathscr{L}_{\pi} : \mathcal{Q} \mapsto \mathcal{Q}$ 为：

$\mathscr{L}_{\pi} \{ q_{\mathcal{H}} \} = \mathcal{R}(s,\ a) + \gamma \mathcal{E}_{s'} \Big[ v_{\mathcal{H}}(s') \Big] = \mathcal{R}(s,\ a) +\gamma \mathcal{E}_{s'} \mathcal{E}_{a'} \Big[ q_{\mathcal{H}}(s',\ a') - \alpha \ln \pi(a' \mid s') \Big]$

利用类似的方法可以证明该算子为压缩映射，贝尔曼期望方程（soft）存在唯一解。

Soft Bellman Optimal Equation

在有限期规划中 $t$ 时间步下，利用动态规划的思想，得到当前的最优策略和最优动作价值函数之间的关系：

$\pi_{\star}^{(t)}(\cdot \mid s) = \argmax_{\pi} \Big[ \alpha \mathcal{H}(\pi \mid s) + \mathcal{E}_{a \sim \pi(\cdot \mid s)} q_{\star}^{(t)}(s,\ a) \Big] = \argmax_{\pi} \mathcal{E}_{a \sim \pi(\cdot \mid s)} \Big[ q_{\star}^{(t)}(s,\ a) - \alpha \ln \pi(a \mid s) \Big]$

考虑优化问题

$\begin{gathered} \max_{p} \sum_{x} p(x) \Big[ \phi(x) - \alpha \ln p(x) \Big] \\[5mm] \mathrm{s.t.} \quad \sum_{x} p(x) = 1 \quad 0 \le p(x) \le 1 \end{gathered} \quad \Rightarrow \quad \begin{gathered} \min_{p} \sum_{x} p(x) \Big[ \alpha \ln p(x) - \phi(x) \Big] \\[5mm] \mathrm{s.t.} \quad \sum_{x} p(x) - 1 = 0 \quad -p(x) \le 0 \quad p(x) - 1 \le 0 \end{gathered}$

可以看出目标函数和不等式约束函数均为凸函数，并且等式约束函数为仿射变换，构造拉格朗日函数：

$\mathcal{L}(p,\ \lambda,\ \mu,\ \nu) = \sum_{x} p(x) \Big[ \alpha \ln p(x) - \phi(x) \Big] + \lambda \left[ \sum_{x} p(x) - 1 \right] + \sum_{x} -\mu(x) p(x) + \sum_{x} \nu(x) \Big[ p(x) - 1 \Big]$

最优分布需要满足：

$\frac{\partial \mathcal{L}}{\partial p(x)} = \alpha \Big[ \ln p(x) + 1 \Big] - \phi(x) + \lambda - \mu(x) + \nu(x) = 0$

解得：

$p_{\star}(x) = \exp \left[ \frac{1}{\alpha} \Big( \phi(x) - \lambda + \mu(x) - \nu(x) \Big) - 1 \right] \propto \exp \left[ \frac{1}{\alpha} \phi(x) \right]$

在最优分布下的，目标函数取到最大值：

$J_{\star} = \sum_{x} p_{\star}(x) \Big[ \phi(x) - \alpha \ln p_{\star}(x) \Big] = \sum_{x} p_{\star}(x) \Big[ \phi(x) - \phi(x) + \alpha \ln Z \Big] = \alpha \ln Z = \underset{x}{\operatorname{softmax}_{\alpha}} \phi(x)$

其中，归一化因子 $Z = \sum_{x} \exp \left[ \dfrac{1}{\alpha} \phi(x) \right]$。因此当前的最优策略为：

$\pi_{\star}^{(t)}(a \mid s) = \frac{1}{Z} \exp \left[ \frac{1}{\alpha} q_{\star}^{(t)}(s,\ a) \right] \quad Z = \sum_{a} \exp \left[ \frac{1}{\alpha} q_{\star}^{(t)}(s,\ a) \right]$

进而得到最优状态价值函数和最优动作价值函数之间的关系，即贝尔曼最优方程（soft）：

$v_{\star}^{(t)}(s) = \alpha \mathcal{H}(\pi_{\star}^{(t)} \mid s) + \mathcal{E}_{a \sim \pi_{\star}^{(t)}(\cdot \mid s)} q_{\star}^{(t)}(s,\ a) = \underset{a}{\operatorname{softmax}_{\alpha}} q_{\star}^{(t)}(s,\ a)$

在该定义下的贝尔曼最优方程（soft）算子 $\mathscr{L} : \mathcal{Q} \mapsto \mathcal{Q}$ 为：

$\mathscr{L} \{ q_{\mathcal{H}} \} = \mathcal{R}(s,\ a) + \gamma \mathcal{E}_{s' \sim p(\cdot \mid s,\ a)} \left[ \underset{b}{\operatorname{softmax}_{\alpha}} q_{\mathcal{H}}(s',\ b) \right] = \mathcal{R}(s,\ a) + \gamma \mathcal{E}_{s' \sim p(\cdot \mid s,\ a)} \left[ \alpha \ln \sum_{b} \exp \left[ \frac{1}{\alpha} q_{\mathcal{H}}(s',\ b) \right] \right]$

为了证明 $\mathscr{L}$ 是一个压缩映射，首先证明：

$\exp \left[ \frac{1}{\alpha} q_{1}(s,\ a) - \frac{1}{\alpha} q_{2}(s,\ a) \right] \le \exp \left[ \frac{1}{\alpha} \max_{s,\ a} \Big| q_{1}(s,\ a) - q_{2}(s,\ a) \Big| \right] = \exp \left[ \frac{1}{\alpha} \Big\| q_{1} - q_{2} \Big\|_{\infty} \right]$

因此：

$\exp \left[ \frac{1}{\alpha} q_{1}(s,\ a) \right] \le \exp \left[ \frac{1}{\alpha} q_{2}(s,\ a) + \frac{1}{\alpha} \Big\| q_{1} - q_{2} \Big\|_{\infty} \right]$

进而有：

$\begin{aligned} \underset{a}{\operatorname{softmax}_{\alpha}} q_{1}(s,\ a) &= \alpha \ln \sum_{a} \exp \left[ \frac{1}{\alpha} q_{1}(s,\ a) \right] \le \alpha \ln \sum_{a} \exp \left[ \frac{1}{\alpha} q_{2}(s,\ a) + \frac{1}{\alpha} \Big\| q_{1} - q_{2} \Big\|_{\infty} \right] \\[7mm] &= \alpha \ln \exp \left[ \frac{1}{\alpha} \Big\| q_{1} - q_{2} \Big\|_{\infty} \right] + \alpha \ln \sum_{a} \exp \left[ \frac{1}{\alpha} q_{2}(s,\ a) \right] = \Big\| q_{1} - q_{2} \Big\|_{\infty} + \underset{a}{\operatorname{softmax}_{\alpha}} q_{2}(s,\ a) \end{aligned}$

通过类似的方法可以得到：

$\underset{a}{\operatorname{softmax}_{\alpha}} q_{1}(s,\ a) \ge -\Big\| q_{1} - q_{2} \Big\|_{\infty} + \underset{a}{\operatorname{softmax}_{\alpha}} q_{2}(s,\ a)$

因此：

$\begin{aligned} \Big| \mathscr{L} \{ q_{1} \}(s,\ a) - \mathscr{L} \{ q_{2} \}(s,\ a) \Big| &= \left| \gamma \sum_{s'} p(s' \mid s,\ a) \left[ \underset{a}{\operatorname{softmax}_{\alpha}} q_{1}(s,\ a) - \underset{a}{\operatorname{softmax}_{\alpha}} q_{2}(s,\ a) \right] \right| \\[7mm] &\le \gamma \sum_{s'} p(s' \mid s,\ a) \left| \underset{a}{\operatorname{softmax}_{\alpha}} q_{1}(s,\ a) - \underset{a}{\operatorname{softmax}_{\alpha}} q_{2}(s,\ a) \right| \\[7mm] &\le \gamma \sum_{s'} p(s' \mid s,\ a) \Big\| q_{1} - q_{2} \Big\|_{\infty} = \gamma \Big\| q_{1} - q_{2} \Big\|_{\infty} \end{aligned}$

最终证明 $\mathscr{L}$ 是一个压缩映射 $\| \mathscr{L} \{ q_{1} \} - \mathscr{L} \{ q_{2} \} \|_{\infty} \le \gamma \| q_{1} - q_{2} \|_{\infty}$，贝尔曼最优方程（soft）存在唯一解 $q_{\star}$，即无限期规划下的最优动作价值函数，同时在无限期规划下的最优策略满足：

$\pi_{\star}(a \mid s) = \frac{1}{Z} \exp \left[ \frac{1}{\alpha} q_{\star}(s,\ a) \right] = \exp \left[ \frac{1}{\alpha} q_{\star}(s,\ a) - \frac{1}{\alpha} v_{\star}(s) \right] \propto \exp \left[ \frac{1}{\alpha} q_{\star}(s,\ a) \right]$