深度强化学习算法：SAC
前置知识：深度强化学习算法一图流
参考资料：
- 论文 - SAC (v1)
- [论文简析]SAC: Soft Actor-Critic Part 1[1801.01290]
- 论文 - SAC (v2)
- [论文简析]SAC: Soft Actor-Critic Part 2[1812.05905]
本文介绍的 SAC 版本是论文 - SAC (v2)中的新版本。
1 Soft Actor-Critic (SAC) 的优点
| 算法 | 策略类型 | 算法类型 | 是否适用 Replay Buffer | Sample Efficiency | Stability |
|---|---|---|---|---|---|
| TRPO & PPO | stochastic policy | On-policy | N | Low | High |
| DDPG & TD3 | deterministic policy | Off-policy | Y | High | Low |
| SAC | stochastic policy | Off-policy | Y | High | High |
2 熵正则 (Entropy Regularization)
在强化学习中，我们希望智能体能够尽可能多地探索环境，以便找到最优策略。这就需要策略函数 \(\pi(\textcolor{red}{a}|\textcolor{green}{s})\) 的决策尽可能随机。
熵 (Entropy) 用来衡量随机变量的不确定性。随机变量 $X \sim P(x)$ 的熵定义为：
\[H(X) = \mathbb{E}_{X \sim P(x)} \left[- \ln P(x)\right] = -\sum_{x \in \mathcal{X}} P(x) \cdot \ln P(x)\]熵越大，说明随机变量的不确定性越高。
所以智能体动作 \(\textcolor{red}{A} \sim \pi(\textcolor{red}{\cdot}|\textcolor{green}{s})\) 的熵定义为：
\[H_\pi (\textcolor{green}{s}) = \mathbb{E}_{\textcolor{red}{A} \sim \pi(\textcolor{red}{\cdot}|\textcolor{green}{s})} \left[- \ln \pi (\textcolor{red}{A}|\textcolor{green}{s})\right] = -\sum_{\textcolor{red}{a} \in \textcolor{red}{\mathcal{A}}} \pi(\textcolor{red}{a}|\textcolor{green}{s}) \cdot \ln \pi(\textcolor{red}{a}|\textcolor{green}{s})\]常规的策略学习目标是最大化状态价值函数的期望 \(J(\pi) = \mathbb{E}_{\textcolor{green}{S}} \left[V_\pi(\textcolor{green}{S})\right]\) 。SAC 将动作 $\textcolor{red}{A}$ 的熵的期望 \(\mathbb{E}_{\textcolor{green}{S}} \left[H_\pi (\textcolor{green}{S})\right]\) 作为正则项加入到目标函数中：
\[J(\pi) \triangleq \mathbb{E}_{\textcolor{green}{S}} \left[V_\pi(\textcolor{green}{S})\right] + \lambda \cdot \mathbb{E}_{\textcolor{green}{S}} \left[H_\pi (\textcolor{green}{S})\right] = \mathbb{E}_{\textcolor{green}{S}} \left[V_\pi(\textcolor{green}{S}) + \lambda \cdot H_\pi (\textcolor{green}{S})\right]\]其中，$\lambda$ 是熵正则项的权重。
SAC 将引入熵正则项后期望中的部分定义为 Soft State Value Function ：
\[\begin{split} V_\pi(\textcolor{green}{s}) &\triangleq V_\pi(\textcolor{green}{s}) + \lambda \cdot H_\pi (\textcolor{green}{s}) \\ &= \mathbb{E}_{\textcolor{red}{A}} \left[Q_\pi(\textcolor{green}{s}, \textcolor{red}{A})\right] + \lambda \cdot \mathbb{E}_{\textcolor{red}{A}} \left[- \ln \pi (\textcolor{red}{A}|\textcolor{green}{s})\right] \\ &= \mathbb{E}_{\textcolor{red}{A}} \left[Q_\pi(\textcolor{green}{s}, \textcolor{red}{A}) - \lambda \cdot \ln \pi (\textcolor{red}{A}|\textcolor{green}{s})\right] \end{split}\]同时，SAC 将上式中的动作价值函数 $Q_\pi(\textcolor{green}{s}, \textcolor{red}{a})$ 定义为 Soft Action Value Function 。
3 Soft Policy Iteration
本节是 SAC 的理论依据，其保证了策略可以通过 Soft Policy Iteration 收敛到最优最大熵策略。Soft Policy Iteration 是 Soft Policy Evaluation 和 Soft Policy Improvement 的交替迭代。
3.1 Soft Policy Evaluation (Critic 相关)
SAC 将贝尔曼算子 (Bellman backup operator) $\mathcal{T}^\pi$ 定义为：
\[\mathcal{T}^\pi Q(\textcolor{green}{s_t}, \textcolor{red}{a_t}) \triangleq \textcolor{blue}{R_t} + \gamma \cdot \mathbb{E}_{\textcolor{green}{S_{t+1}}} \left[ V(\textcolor{green}{s_{t+1}})\right]\]其中， \(V(\textcolor{green}{s_{t+1}}) = \mathbb{E}_{\textcolor{red}{A}} \left[Q(\textcolor{green}{s_{t+1}}, \textcolor{red}{A}) - \lambda \cdot \ln \pi (\textcolor{red}{A}|\textcolor{green}{s_{t+1}})\right]\) 是 Soft State Value Function ，$Q(\textcolor{green}{s_t}, \textcolor{red}{a_t})$ 是 Soft Action Value Function 。
论文中使用引理 1保证：对于任意给定的策略函数 $\pi$ ，使用递推关系 $Q^{k+1} = \mathcal{T}^\pi Q^k$ 定义的数列 ${Q^k}$ 收敛。并且当 $k \to \infty$ 时，${Q^k}$ 收敛到 $\pi$ 的 soft Q-value 。
Lemma 1 (Soft Policy Evaluation) 的证明见论文 - SAC (v2)的 Appendix B.1 。
贝尔曼算子的定义来源于深度强化学习算法一图流中状态价值函数定义处右下箭头所指位置。
3.2 Soft Policy Improvement (Actor 相关)
SAC 更新策略的目标是最大化 Soft Q-value 。为了方便策略更新，SAC 将策略函数限制在高斯分布的范围内。更新策略的方法是让新策略的分布尽可能接近原策略的动作价值函数的指数的形状：
\[\pi_{\text{new}} = \mathop{\text{argmin}}\limits_{\pi' \in \Pi} D_{KL} \left(\pi' (\textcolor{red}{\cdot} | \textcolor{green}{s_t}) \middle|\middle| \frac{\exp (\frac{1}{\lambda} Q_{\pi_{\text{old}}}(\textcolor{green}{s_t}, \textcolor{red}{\cdot}))}{Z_{\pi_{\text{old}}}(\textcolor{green}{s_t})}\right)\]其中，$\Pi$ 是所有高斯分布策略函数的集合；$Z_{\pi_{\text{old}}}(\textcolor{green}{s_t})$ 是配分函数 (Partition function) ，它是平衡态统计物理学中的概念，此处用于归一化动作价值函数的指数。
论文中使用引理 2保证：使用 KL 散度最小化的方法得到的新策略的动作价值函数的值 $Q_{\pi_{\text{new}}}(\textcolor{green}{s_t}, \textcolor{red}{a_t})$ 不会低于原策略的动作价值函数的值 $Q_{\pi_{\text{old}}}(\textcolor{green}{s_t}, \textcolor{red}{a_t})$ ，即： $Q_{\pi_{\text{new}}}(\textcolor{green}{s_t}, \textcolor{red}{a_t}) \geqslant Q_{\pi_{\text{old}}}(\textcolor{green}{s_t}, \textcolor{red}{a_t})$ for all $(\textcolor{green}{s_t}, \textcolor{red}{a_t}) \in \mathcal{S} \times \mathcal{A}$ 。
Lemma 2 (Soft Policy Improvement) 的证明见论文 - SAC (v2)的 Appendix B.2 。
3.3 Soft Policy Iteration 的收敛性
论文中的定理 1说明：对于任意一个高斯分布策略 $\pi \in \Pi$ ，交替迭代 Soft Policy Evaluation 和 Soft Policy Improvement ，会使策略 $\pi$ 收敛到最优最大熵策略 $\pi^\star$ ，满足 $Q_{\pi^\star}(\textcolor{green}{s_t}, \textcolor{red}{a_t}) \geqslant Q_{\pi}(\textcolor{green}{s_t}, \textcolor{red}{a_t})$ for all $\pi \in \Pi$ and $(\textcolor{green}{s_t}, \textcolor{red}{a_t}) \in \mathcal{S} \times \mathcal{A}$ 。
Theorem 1 (Soft Policy Iteration) 的证明见论文 - SAC (v2)的 Appendix B.3 。
4 Soft Actor-Critic (SAC) 算法
论文 - SAC (v1)旧版本的 SAC 算法使用神经网络近似了 Soft State Value Function 。
论文 - SAC (v2)新版本的 SAC 算法不对 Soft State Value Function 进行近似。
4.1 使用神经网络近似策略函数和软动作价值函数
SAC 使用两个 Soft Q-network $q(\textcolor{green}{s}, \textcolor{red}{a}; \omega_1)$ 和 $q(\textcolor{green}{s}, \textcolor{red}{a}; \omega_2)$ 近似 Soft Action Value Function $Q(\textcolor{green}{s}, \textcolor{red}{a})$ 。它们各自对应的 Target Q-network 为 $q(\textcolor{green}{s}, \textcolor{red}{a}; \omega_1^-)$ 和 $q(\textcolor{green}{s}, \textcolor{red}{a}; \omega_2^-)$ 。
SAC 使用神经网络近似策略函数 \(\textcolor{red}{a} \sim \pi (\textcolor{red}{\cdot}|\textcolor{green}{s}; \theta)\) 。
4.2 Soft Q-network 的更新
SAC 根据 3.1 节中的贝尔曼算子定义式，Soft Action Value Function 可以用轨迹 $(\textcolor{green}{s_t}, \textcolor{red}{a_t}, \textcolor{blue}{r_t}, \textcolor{green}{s_{t+1}})$ 进行蒙特卡洛近似：
\[Q(\textcolor{green}{s_t}, \textcolor{red}{a_t}) \approx \textcolor{blue}{r_t} + \gamma \cdot V(\textcolor{green}{s_{t+1}}) = \textcolor{blue}{r_t} + \gamma \cdot \left[Q(\textcolor{green}{s_{t+1}}, \textcolor{red}{a_{t+1}}) - \lambda \cdot \ln \pi(\textcolor{red}{a_{t+1}}|\textcolor{green}{s_{t+1}})\right]\]因此可以使用 TD 算法更新 Soft Q-network 。
更新流程：
- 智能体与环境交互，得到轨迹 $(\textcolor{green}{s_t}, \textcolor{red}{a_t}, \textcolor{blue}{r_t}, \textcolor{green}{s_{t+1}})$
- 计算价值 $q_t = q(\textcolor{green}{s_t}, \textcolor{red}{a_t}; \omega_t)$
- 针对新状态选择下一个动作，但不执行 \(\textcolor{red}{\tilde{a}_{t+1}} \sim \pi(\textcolor{red}{\cdot}|\textcolor{green}{s_{t+1}}; \theta_t)\)
- 计算 Soft State Value \(v_{t+1} = q(\textcolor{green}{s_{t+1}}, \textcolor{red}{\tilde{a}_{t+1}}; \omega_t^-) - \lambda \cdot \ln \pi(\textcolor{red}{\tilde{a}_{t+1}}|\textcolor{green}{s_{t+1}}; \theta_t)\)
- 计算 TD target $y_t = \textcolor{blue}{r_t} + \gamma \cdot v_{t+1}$
- 计算 TD error $\delta_t = q_t - y_t$
- 最小化 TD error ，定义损失函数 $L_t = \frac{1}{2} \delta_t^2$
- 使用梯度下降最小化损失函数 \(\omega_{t+1} \leftarrow \omega_t - \alpha \cdot \frac{\partial L_t}{\partial \omega}|_{\omega = \omega_t} = \omega_t - \alpha \cdot \delta_t \cdot \frac{\partial q(\textcolor{green}{s_t}, \textcolor{red}{a_t}; \omega)}{\partial \omega}|_{\omega = \omega_t}\)
4.3 策略网络的更新
SAC 更新策略网络的目标是最大化 Soft State Value Function ：
\[\theta^\star = \mathop{\text{argmax}}\limits_{\theta} \mathbb{E}_{\textcolor{green}{S}} \left[V(\textcolor{green}{S})\right] = \mathop{\text{argmax}}\limits_{\theta} \mathbb{E}_{\textcolor{green}{S}} \left\{\mathbb{E}_{\textcolor{red}{A}} \left[Q(\textcolor{green}{S}, \textcolor{red}{A}) - \lambda \cdot \ln \pi (\textcolor{red}{A}|\textcolor{green}{S}; \theta)\right]\right\}\]注：最小化 3.2 节中的 KL 散度的目标函数与最大化 Soft State Value Function 的目标函数等价：
\[\begin{split} J(\theta) &= D_{KL} \left(\pi(\textcolor{red}{\cdot} | \textcolor{green}{s}; \theta) \middle|\middle| \frac{\exp (\frac{1}{\lambda} Q(\textcolor{green}{s}, \textcolor{red}{\cdot}; \omega))}{Z(\textcolor{green}{s})}\right) \\ &= D_{KL} \left(\pi (\textcolor{red}{\cdot} | \textcolor{green}{s}; \theta) \middle|\middle| \exp \left[\frac{1}{\lambda} Q(\textcolor{green}{s}, \textcolor{red}{\cdot}; \omega) - \ln Z(\textcolor{green}{s})\right]\right) \\ &= \mathbb{E}_{\textcolor{red}{A}, \textcolor{green}{S}} \left\{\ln \frac{\pi (\textcolor{red}{A} | \textcolor{green}{S}; \theta)}{\exp \left[\frac{1}{\lambda} Q(\textcolor{green}{S}, \textcolor{red}{A}; \omega) - \ln Z(\textcolor{green}{S})\right]}\right\} \\ &= \mathbb{E}_{\textcolor{red}{A}, \textcolor{green}{S}} \left[\ln \pi (\textcolor{red}{A} | \textcolor{green}{S}; \theta) - \frac{1}{\lambda} Q(\textcolor{green}{S}, \textcolor{red}{A}; \omega) + \ln Z(\textcolor{green}{S})\right] \end{split}\]其中，$Z(\textcolor{green}{s})$ 在更新策略时与 $\theta$ 无关，所以可以忽略。
为了实现连续控制，SAC 将策略网络重参数化 (Reparameterization Trick) ，使用 squashed Gaussian policy 作为策略函数：
\[\textcolor{red}{a} = f(\textcolor{green}{s}, \epsilon; \theta) = \tanh (\mu(\textcolor{green}{s}; \theta^\mu) + \sigma(\textcolor{green}{s}; \theta^\sigma) \odot \epsilon)\]对于 $n$ 维连续动作空间，神经网络 $\mu(\textcolor{green}{s}; \theta^\mu)$ 和 $\sigma(\textcolor{green}{s}; \theta^\sigma)$ 的输出都是 $n$ 维向量。$\theta = (\theta^\mu, \theta^\sigma)$ 是神经网络的参数。$n$ 维向量 $\epsilon$ 的各个分量服从标准正态分布，即：$\epsilon \sim \mathcal{N}(0, I)$ 。$\odot$ 表示 Hadamard 乘积，即对应位置元素相乘。
更新流程（最小化 KL 散度）：
- 智能体与环境交互，得到轨迹 $(\textcolor{green}{s_t}, \textcolor{red}{a_t}, \textcolor{blue}{r_t}, \textcolor{green}{s_{t+1}})$
- 定义重参数化的损失函数 \(L(\theta) = \mathbb{E}_{\textcolor{green}{S}, \epsilon} \left[\lambda \cdot \ln \pi (f(\textcolor{green}{S}, \epsilon; \theta) | \textcolor{green}{S}; \theta) - Q(\textcolor{green}{S}, f(\textcolor{green}{S}, \epsilon; \theta); \omega)\right]\)
- 使用梯度下降最小化损失函数 \(\theta_{t+1} \leftarrow \theta_t - \alpha \cdot \frac{\partial L(\theta)}{\partial \theta}|_{\theta = \theta_t}\)
其中，使用 $(\textcolor{green}{s_t}, \textcolor{red}{a_t}, \textcolor{blue}{r_t}, \textcolor{green}{s_{t+1}})$ 进行蒙特卡洛近似的损失函数的梯度为：
\[\frac{\partial L(\theta)}{\partial \theta} = \lambda \cdot \frac{\partial \ln \pi (\textcolor{red}{a_t}|\textcolor{green}{s_t}; \theta)}{\partial \theta} + \left[\lambda \cdot \frac{\partial \ln \pi (\textcolor{red}{a}|\textcolor{green}{s_t}; \theta)}{\partial \textcolor{red}{a}} \middle|_{\textcolor{red}{a}=\textcolor{red}{a_t}} - \frac{\partial Q(\textcolor{green}{s_t}, \textcolor{red}{a}; \omega)}{\partial \textcolor{red}{a}} \middle|_{\textcolor{red}{a}=\textcolor{red}{a_t}}\right] \cdot \frac{\partial f(\textcolor{green}{s_t}, \epsilon_t; \theta)}{\partial \theta}\]求梯度的过程参考这篇文章。
4.5 算法流程
- 初始化两个 Soft Q-network $q(\textcolor{green}{s}, \textcolor{red}{a}; \omega_1)$ 和 $q(\textcolor{green}{s}, \textcolor{red}{a}; \omega_2)$ ，以及策略网络 \(\textcolor{red}{a} \sim \pi (\textcolor{red}{\cdot}|\textcolor{green}{s}; \theta)\)
- 初始化 Soft Target Q-network $q(\textcolor{green}{s}, \textcolor{red}{a}; \omega_1^-)$ 和 $q(\textcolor{green}{s}, \textcolor{red}{a}; \omega_2^-)$ ，并设置参数 $\omega_1^- \leftarrow \omega_1$ 和 $\omega_2^- \leftarrow \omega_2$
- 初始化 Replay Buffer $\mathcal{D}$
- 与环境交互直到结束，将所有轨迹 $(\textcolor{green}{s_t}, \textcolor{red}{a_t}, \textcolor{blue}{r_t}, \textcolor{green}{s_{t+1}})$ 存入 Replay Buffer $\mathcal{D}$
- 从 Replay Buffer $\mathcal{D}$ 中采样一批数据 $B = {(\textcolor{green}{s_i}, \textcolor{red}{a_i}, \textcolor{blue}{r_i}, \textcolor{green}{s_{i+1}})}$
- 计算所有 TD target \(y_i = \textcolor{blue}{r_i} + \gamma \cdot \left[\min\limits_{j=1,2} q(\textcolor{green}{s_{i+1}}, \textcolor{red}{\tilde{a}_{i+1}}; \omega_j^-) - \lambda \cdot \ln \pi(\textcolor{red}{\tilde{a}_{i+1}}|\textcolor{green}{s_{i+1}}; \theta)\right]\) ，其中 \(\textcolor{red}{\tilde{a}_{i+1}} \sim \pi(\textcolor{red}{\cdot}|\textcolor{green}{s_{i+1}}; \theta)\)
- 使用 TD error 的梯度 \(\nabla_{\omega_j} \frac{1}{|B|} \sum\limits_{i=1}^{|B|} \left[q(\textcolor{green}{s_i}, \textcolor{red}{a_i}; \omega_j) - y_i\right]^2\) 进行梯度下降更新两个 Soft Q-network ，其中 $j=1,2$
- 使用 Soft State Value Function 的梯度 \(\nabla_{\theta} \frac{1}{|B|} \sum\limits_{i=1}^{|B|} \left[\min\limits_{j=1,2} q(\textcolor{green}{s_i}, f(\textcolor{green}{s_i}, \epsilon_i; \theta); \omega_j) - \lambda \cdot \ln \pi (f(\textcolor{green}{s_i}, \epsilon_i; \theta)|\textcolor{green}{s_i}; \theta)\right]\) 进行梯度上升更新策略网络
- 调整熵的权重 $\lambda$ （见论文 - SAC (v2)第 5 章）
- 更新 Soft Target Q-network 的参数 $\omega_j^- \leftarrow \tau \cdot \omega_j + (1 - \tau) \cdot \omega_j^-$ ，其中 $j=1,2$
- 重复步骤 5-10 直到策略网络收敛
- 返回步骤 4 进行下一轮训练