从Maximum Entropy Reinforcement Learning出发,MERL的思路是在reward上加入熵作为额外奖励:
$V^{(MERL)}(s_0)= \sum\limits_{t=0}^{\infty} \mathbb{E}_{s_t, a_t \sim \tau} \left[ \gamma^t \left( r(s_t, a_t) + \alpha H(\pi(\cdot \mid s_t)) \right) \right]$
soft V的定义也是一样的,然后需要定义soft Q:
$Q_{soft}(s, a) = \mathbb{E}_{s’ \sim p(s’ \mid s, a)}[r(s, a) + \gamma V_{soft}(s’)]$
$=\mathbb{E}_{s_t, a_t \sim \tau}\left[ \sum\limits_{t=0}^{\infty} \gamma^t r(s_t, a_t) + \sum\limits_{t=1}^{\infty} \gamma^t \alpha H(\pi(\cdot \mid s_t)) \right]$
其中entropy项在外面的expectation下可直接转化为$\alpha H = \mathbb{E}_{\pi}[-\alpha log \pi(\cdot \mid s_t)]$
再对比$V_{soft}(s)$与$\mathbb{E}_{\pi}[Q_{soft}(s, a)]$,发现只少了初始步s时的H项,所以
$V_{soft}(s) = \mathbb{E}_{a \sim \pi}[Q_{soft}(s, a) – \alpha log \pi(a \mid s)]$
于是一步TD:
$\delta_{soft} = r_t + \gamma V_{soft}(s’) – V_{soft}(s) – \alpha log\pi(\bar{a} \mid s)$
这里$\bar{a} \sim \pi(\cdot \mid s)$需要从策略中采样得到,并不是轨迹中的记录$a$,所以用了一个稍微不同的标记。
值的迭代基本上就是与$\delta_{soft}$相关的最小二乘,或者搞成类似v-trace类似的value estimator应该也是可行的。
接下来就是SAC针对actor部分的修改了:
使用了能量模型(Boltzmann分布)而非正态分布。
$\pi(a_t \mid s_t) \propto exp\{ \frac{Q_{soft}(s_t, a_t)}{\beta} \}$
这样的好处是避免了高斯分布策略里的线性“回弹力”,或唯一最优解附近的抛物线势能井的预先假设,使之后batch输入产生的策略更无偏,因为Boltzmann本身就是在特定条件下满足熵最大化的一种分布。
这些条件包括“能量”Q守恒(暂且认为未观测到的$<s, a>$计算出的$Q$值,是某个已经确定的固定值,例如收敛后的最优策略下的$Q^{(\pi^*)}$打了一个固定折扣之后得到的常量。),且在未来的batch内即将出现的不同$a$操作(即宏正则观测)下的微正则态($<s,a>$组合出现事件模拟)可携带自身的能量切换于不同宏正则操作,最终形成热平衡。
这样保证了对未来batch输入的信息(状态)保持最大不确定性的无偏分布,是能够让模型具有更高鲁棒性的。
$\beta$是于温度相关的超参。
根据上面的一步td公式,最优解时它的预期是0,于是有
$\alpha log \pi^*(a_t \mid s_t) = Q_{soft}(s_t, a_t) – V_{soft}(s_t)$
$\pi^*(a_t \mid s_t) = exp \{Q_{soft}(s_t, a_t) / \alpha – V_{soft}(s_t)/\alpha \}$
对上式的操作空间求sum$\sum\limits_{\bar{a}}$有
$\sum\limits_{\bar{a}} \pi^*(\bar{a} \mid s_t) = \sum\limits_{\bar{a}} [ exp \{Q_{soft}(s_t, \bar{a}) /\alpha – V_{soft}(s_t)/\alpha\}]$
左边是1,同时改写右边得到
$1 = \frac{\sum\limits_{\bar{a}} exp \{Q_{soft}(s_t, \bar{a}) /\alpha \} } { exp\{V_{soft}(s_t)/\alpha\ \}}$
因为归一化条件$\sum\limits_{\bar{a}} \pi^* = 1$
对比Boltzmann分布下,每一种宏正则的概率,归一化形式是一样的,只不过温度$\beta$替换了$\alpha$
可知配分函数$Z=\int\limits_{\bar{a}} exp\{ Q(s_t, \bar{a})/\beta \} da = exp\{ \frac{V_{soft}(s_t)}{\alpha} \}$
且$\alpha = \beta$。这时原本从热力学角度出发对actor进行的修改,与最初针对Value函数的修改就对应上了。
$\pi^*(s_t \mid s_t) = \frac {exp \{Q_{soft}(s_t, a_t) / \alpha \}} {exp \{V_{soft}(s_t) / \alpha \}} $
所以策略的更新迭代需要:
$\pi^{new}(s_t \mid s_t) = \frac {exp \{Q_{soft}^{\pi_{old}}(s_t, \bar{a}_t) / \alpha \}} {exp \{V_{soft}^{\pi_{old}}(s_t) / \alpha \}} \forall s_t$
其中$Q_{soft}$需要使用旧策略$\pi_{old}$采样出操作$\bar{a}$,之后$V_{soft}$还需要对$Q_{soft}$计算积分。
实际操作中积分计算在连续空间里几乎不可能,而如果在离散空间里,计算$\mathbb{E}_{\bar{a} \sim \pi_{old}}[Q_{soft}(s_t, \bar{a}_t)]$是可行的。
Soft Q Learning:
有了评估Q与V的公式:
$Q_{soft}(s_t, a_t) = r(s_t, a_t) + \gamma \mathbb{E}_{p(s_{t+1} \mid s_t, a_t)} [V(s_{t+1})]$
$V_{soft}(s_t) = \alpha log \int\limits_{\bar{a}} exp\{ \frac{Q_{soft}(s_t, \bar{a}_t)}{\alpha} \} d\bar{a}$
其中V里面的积分实际上就是softmax操作,当$\alpha \rightarrow 0$时,趋近于hard max操作。
我们可以从$V_{soft}(s_T)=0$开始,倒推回t=0,计算历史轨迹上的$Q_t$与$V_t$,这样我们只需要一个神经网络拟合$Q_{soft}$即可。
之后计算并记录Q网络在TD上拟合需要更新的梯度方向。
再之后新策略可以根据boltzmann分布算出:
$\pi^{new}(s_t \mid s_t) = \frac {exp \{Q_{soft}^{\pi_{old}}(s_t, \bar{a}_t) / \alpha \}} {exp \{V_{soft}^{\pi_{old}}(s_t) / \alpha \}} \forall s_t$
最后更新Q网络的梯度
不断循环【评估Q-得到新策略-更新TD梯度】的过程,即可对离散操作空间进行soft q learning。
推理时直接拿$Q(s, \bar{a})$做$\bar{a}$上的greedy就可以了。
然而对于连续空间要怎样处理呢?连续时无法评估
$V_{soft}(s_t) = \alpha log \int\limits_{\bar{a}} exp\{ \frac{Q_{soft}(s_t, \bar{a}_t)}{\alpha} \} d\bar{a}$
且我们也需要一个actor,让它产生的策略无限趋近于
$\pi^{new}(s_t \mid s_t) = \frac {exp \{Q_{soft}^{\pi_{old}}(s_t, \bar{a}_t) / \alpha \}} {exp \{V_{soft}^{\pi_{old}}(s_t) / \alpha \}} \forall s_t$
通过采样来估算$V_{soft}(s_t)$里的积分,然后用下面的loss来更新actor
$J_{\pi}= D_{KL}(\pi_{\theta}(\cdot \mid s_t) || \frac {exp \{Q_{soft}^{\pi_{old}}(s_t, \cdot) / \alpha \}} {exp \{V_{soft}^{\pi_{old}}(s_t) / \alpha \}})$
是一种方案,于是后面提出了SAC对该方案进行了优化。
SAC:
1步TD的Q网络损失函数:
$J_Q = \mathbb{E}_{(s_t, a_t, s_{t+1}, r_t) \sim \tau}\left[ \frac{1}{2} \left( Q(s_t, a_t) – r_t – \gamma Q(s_{t+1}, a_{t+1}) – \alpha log \pi_{\theta}(\bar{a}_{t+1} \mid s_{t+1}) \right) \right]$
其中$\bar{a}_{t+1} = \mu_{\theta}(s_{t+1}) + \epsilon_t \sigma_{\theta}(s_{t+1})$是使用重采样方法从$\pi_{\theta}$中采样获得的action。
$\epsilon$仍然是从$Norm(0, 1)$里面采样获得。
实际操作时会对action的范围进行限制,原本从Gaussian采样得到的$\bar{a}$,会在外面套上tanh变成$\bar{a}^{(bound)} = tanh(\bar{a})$。要让受限action与不受限action的概率密度相等,就需要改变分布,让$\pi^{(bound)}(\cdot \mid s) = \pi(\cdot \mid s) \Big| det\left( \frac{d\bar{a}^{(bound)}}{d\bar{a}} \right) \Big|^{-1}$
先取log,之后由于action不同维度之间相互独立,取det操作下的Jacobian $\frac{d\bar{\boldsymbol{a}}^{(bound)}}{d\bar{\boldsymbol{a}}}$是对角矩阵,以及$\frac{dtanh(x)}{dx}=\frac{1}{cosh^2(x)}=1-tanh^2(x)$,可得
$log \pi^{(bound)}(\bar{\boldsymbol{a}} \mid s) = log \pi(\bar{\boldsymbol{a}} \mid s) – \sum\limits_{i=1}^{D} log (1-tanh^2(\bar{a}_i))$
其中D是操作空间的维数。类似的转化操作在后面的优化项上也适用。
之后更新actor时的loss:
$J_{\pi} = D_{KL} \left[ \pi_{\theta}(\cdot \mid s) || exp\{ \frac{1} {\alpha}Q(s, \bar{a}) – logZ \} \right]$
因为$Z(s_t)=exp\{V(s_t)/\alpha\}$不受$\theta$影响,所以对$\theta$求导时可以忽略。
$J_{\pi} = \mathbb{E}_{s_t \sim \tau, \bar{a} \sim \pi} \left[ log \left( \frac{\pi_{\theta}(\bar{a} \mid s_t)} {exp\{ Q(s_t, \bar{a}_t) / \alpha \}} \right)\right]$
$ = \frac{1}{\alpha} \mathbb{E}_{s_t \sim \tau, \bar{a} \sim \pi} \left[ \alpha log \pi_{\theta}(\bar{a} \mid s_t) – Q(s, \bar{a})\right]$
最后有一个比较重要的改进,关于如何自动调整温度$\alpha$。之前是把熵$H$作为了某种“惩罚项”,现在改成约束项,即
Lagrange Multiplier的条件公式变成了$ \alpha \left( H(\pi_{\theta}) – H_0\right) $ 也应该作为无约束优化项的一部分,其中$H_0$是熵的最小值的阈值, $\alpha$是变量。
于是可以定义一个优化目标
$J_{\alpha} = \mathbb{E}_{\bar{a} \sim \pi_{\theta}} [- \alpha log \pi(\bar{a} \mid s_t) – \alpha H_0]$
这里由于$\alpha$是一个大于0的正数,所以一般训练参数实际上代表了$log(\alpha)$,以确保无论参数的正负,$\alpha$始终是正。
当训练不够稳定,计算出的同一个batch内的$log \pi$差异非常大时,最好让自动调节温度的学习速率降低。
每次迭代更新时根据buffer内的轨迹,依次最小化$J_Q$, $J_{\pi}$, $J_{\alpha}$,计算出对应的梯度,进行更新即可。