Soft Actor-Criticの論文を説明します．
以下の2つの論文を参考にしていて，本文中の図は全て論文からの引用になります．
- Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor
- Soft Actor-Critic Algorithms and Applications
対象読者としましては，
- 実装ではなく，数式を追いたい人
- 証明とまではいかないけど，理論をある程度理解したい人
- 強化学習の初心者〜中級者くらいの人
を想定しています．
また，Soft Actor-CriticのPyTorchでの実装をここに載せています．コメントを詳しめに書いていますので，アルゴリズムの理解の助けになればと思います．特にre-parameterization trickのあたりは，実際にどのような処理をしているのかを参照することで理解が深まると思います．
記事に関して，間違い等がありましたらご指摘ください．
概要
Soft Actor-Critic(SAC)は簡単にいうと，普通のActor-Criticの目的関数(期待収益の最大化)に，方策のエントロピー最大化項を加えたものになります．
先行研究とどこが違うのか？
model-freeのRLアルゴリズムでは主に
- TRPO，PPO，A3Cなどのon-policyなアルゴリズム
- DDPG，TD3などのoff-policyなアルゴリズム
がよく用いられています．
前者のon-policyなアルゴリズムでは，勾配更新のたびに新たなサンプルを必要とする(過去のサンプルの再利用はできない)ため，サンプル効率が悪いことが問題に挙げられます．また後者のoff-policyなアルゴリズムでは，サンプル効率は改善されるものの，既存のアルゴリズムではハイパーパラメータに超鋭敏である(学習が不安定)ということが問題に挙げられます．
そこで本論文で紹介するSACでは，
- off-policyの学習による高いサンプル効率
- 方策のエントロピー最大化項を加えて，よりdiverseな探索による学習の安定化
の2点を実現したアルゴリズムを提案しています．また本論文の良いところとしては，新しく導入したsoft Q関数に関するベルマンオペレータの単調増加性や収束性をきちんと証明しているところが挙げられます．Soft Bellman Equationなどの詳しい解説については，Haarnoja et al.(2017)を参照してください．
技術のキモは？
通常のRLでは，
$$
J(\pi) = \sum_{t=0}^{T} E_{(s_t, a_t) \sim \rho_\pi} [r(s_t, a_t)]
$$
という期待収益を最大化することを目的とします．これは方策$\pi$に従って行動した場合に，時刻Tまでに得られる報酬の総和(=収益)の期待値を最大化することを意味します．
一方SACでは，上記の目的関数にエントロピー最大化項を加えます．
$$
J(\pi) = \sum_{t=0}^{T} E_{(s_t, a_t) \sim \rho_\pi} [r(s_t, a_t) + \alpha H(\cdot | s_t)]
$$
ここで$\alpha$は温度パラメータで，どのくらいエントロピー項$H$を考慮するかを決めるハイパーパラメータになります．SACでは，上記の目的関数を最大化することで，サンプル効率・安定性の2つを兼ね備えたアルゴリズムを実現しています．
有限な状態・行動空間でのSACアルゴリズム
問題設定
今回は通常の方策反復法(Actor-Critic)からSACアルゴリズムを導出するために，まず有限な状態空間・行動空間を持ったMDPを考え，$(S, A, p, r)$で表します．$S$，$A$は状態空間，行動空間を表します．$p: S \times S \times A \to [0, 1]$は状態遷移確率を表し，$r: S \times A \to [r_{min}, r_{max}]$は各時刻での報酬関数を表します．また，時刻tでの状態，および状態・行動の期待確率をそれぞれ$\rho_{\pi}(s_t)$，$\rho_{\pi}(s_t, a_t)$で表します．これはある方策$\pi$に従って行動し続けた場合に，ある時刻tで状態$s = s_t$，および$s=s_t, a=a_t$となる確率です．また割引率$\gamma$を導入します．
soft Q関数の評価
SACアルゴリズムを導出するために，通常のベルマンオペレータを修正したベルマンオペレータ$T^{\pi}$を定義します．
$$
T^{\pi} Q(s_t, a_t) = r(s_t, a_t) + \gamma E_{s_{t+1}\sim p}[V(s_{t+1})]
$$
ここまでは通常のベルマンオペレータと同じなのですが，価値関数$V$にsoft価値関数と呼ばれる，通常の価値関数にエントロピー項を加えたものを用います．
$$
V(s_t) = E_{a_t \sim \pi} [Q(s_t, a_t) - \alpha \log \pi(a_t | s_t)]
$$
ここで負の対数項は，方策のエントロピーを表し，$\alpha$はどのくらいエントロピー項を重視するかを決める温度パラメータになります．ここでは詳しく説明しませんが，上記のベルマンオペレータを任意の関数$Q: S \times A \to R$を適用(e.g. $Q^{k+1}=T^{\pi}Q^k$)していくと，$k \to \infty$で$Q^k$が方策$\pi$のsoft Q関数に収束することが論文の付録で証明されています．
このようなSoft Bellman Equationに基づいたQ-learningをSoft Q-learningと言います．(Haarnoja et al.(2017))
方策の更新
ここでは，前節で推定したsoft Q関数を最大化する方策を求めるための更新式を考えます．(証明は省略します)
付録の証明では，以下の更新式を更新していくことで(有限の状態・行動空間の場合は)soft Q関数を最大化するような方策に収束することが示されています．
$$
\pi_{new} = argmin_{\pi' \sim \Pi} ; D_{KL} (\pi'(\cdot | s_t) ||\frac{exp(\frac{1}{\alpha}Q^{\pi_{old}}(s_t, a_t))}{Z^{\pi_{old}}(s_t)})
$$
つまり，ある方策集合$\Pi$の中から上記の$D_{KL}$を最小化するような方策を選んで更新していきます．$D_{KL}$の中の$Z$は正規化パラメータです．連続空間の場合には，方策の勾配を求める際に定数である$Z$は無視することができる(後述)のでそこまで気にしなくて良いです．
この更新式では，方策$\pi$がsoft Q関数の指数分布に近づくように更新していく，という風に考えることができます．この更新式こそが，soft learningのアイデアのすべてを表現しています．$\alpha\to 0$のとき，方策はsoft Q関数が最大を取るような行動を選択する決定的方策(greedyな方策)に近づいていきます．一方，$\alpha \to \infty$のとき，方策は一様なランダム方策に近づいていきます．このことから，$\alpha$はランダム性の大きさを表すパラメータであると言うことができ，SACでは「(soft)Qの大きい行動を選択しやすくしつつも，その他の行動もQに応じた(非ゼロの)確率で行動を選択する方策」を学習していると言えます．エントロピーは大まかに乱雑さを表す量であり，エントロピー最大化項はできるだけ不確かさ(ランダム性)を増大させている，という直感とも一致します．
(有限空間での)SACアルゴリズム
有限空間でのSACアルゴリズムは，
- soft Q関数の評価を行い，soft Q関数の推定精度を高める
- 方策の更新を行い，推定されたsoft Q関数を最大化するように方策を更新する
の2つを交互に行うことで，いずれは(真の)soft Q関数を最大化する最適方策に収束することが証明されます．(付録参照)
このように有限の状態・行動空間ではsoft Q関数を表形式で表すことができるため，反復を行うことで最適方策に収束していくことが示せます．
連続な状態・行動空間でのSACアルゴリズム
問題設定
ほぼ有限の場合と同じですが，$p: S \times S \times A \to [0, \infty)$が状態遷移確率密度を表し，$\rho$も確率ではなく確率密度となります．
ここでは下記の2つの関数をニューラルネットを用いて関数近似します．
- $Q(\cdot | \theta)$ : soft Q関数
- $\pi (\cdot | \phi)$ : 方策関数
SACの元論文では，soft価値関数$V(\cdot | \psi)$もニューラルネットで関数近似していましたが，著者らの実装や2018年末に出された2本目の論文ではsoft価値関数は明示的に関数近似していません．代わりにSoft Belleman Equationに基づいて，以下のようにsoft価値関数とsoft Q関数の関係を記述出来るので，それに基づいてsoft価値関数を推定します．
$$
V(s_t) = E_{a_t \sim \pi} [Q(S_t, a_t) - \alpha \log \pi(a_t | s_t)]
$$
もう一度，最大化したい目的関数を下記に示しておきます．
$$
J(\pi) = \sum_{t=0}^{T} E_{(s_t, a_t) \sim \rho_\pi} [r(s_t, a_t) + \alpha H(\cdot | s_t)]
$$
更新式
SACでは各エピソードの各時間ステップで，下記の2つの関数近似器(ネットワーク)の勾配を求め，SDGを用いてパラメータを更新していくことで学習を行います．SACでは，下記の更新をReplay Bufferなどを用いてoff-policyで行うという点に注意してください．
soft Q関数
soft Q関数は，下記の二乗誤差を最小化していきます．
$$
J_Q(\theta) = E_{(s_t, a_t) \sim D}[\frac{1}{2}\left( Q_{\theta}(s_t, a_t) - \left(r(s_t, a_t) + \gamma E_{s_{t+1} \sim p}[V_{\bar\theta}(s_{t+1})]\right) \right)^2]
$$
2項目と3項目は，soft Q関数の推定値を表していて，前節で示したsoft価値関数とsoft Q関数の関係から導出します．ただし，$\bar \theta$は関数近似したsoft Q関数のtarget networkのパラメータを表します．target networkの更新則にはmoving averageを用いています．
target networkの詳細はここでは説明しませんが，target networkを用いることによって，更新式のsoft Q関数の推定のバイアスを低減させ，学習を安定化させることができると言われています．より詳細な説明は，DQNの論文などを参照してください．
勾配は下式のように推定できます．($\theta$で偏微分してみてください．)
$$
\hat\nabla_{\theta} J_Q(\theta) = \nabla_{\theta} Q_\theta (s_t, a_t) ( Q_\theta(s_t, a_t) - (r(s_t, a_t) + \gamma (Q_{\bar \theta}(s_{t+1}, a_{t+1}) - \alpha \log \pi_\phi(a_{t+1} | s_{t+1}))))
$$
方策関数
方策関数は，有限空間の場合と同様に下式を最小化します．
$$
J_\pi (\phi) = E_{s_t \sim D} [ D_{KL} \left( \pi_\phi(\cdot | s_t) || \frac{exp(\frac{1}{\alpha} Q_\theta (s_t, \cdot))}{Z_\theta(s_t)} \right) ]
$$
ここで，deterministicな方策$\pi(\cdot|s)$を関数近似する際に，reparameterization trickと呼ばれるテクニックを用いて下記のように関数近似(パラメータ$\phi$)してあげます．
$$
a_t = f_\phi(\epsilon_t ; s_t)
$$
ここで$f_\phi$はパラメータ$\phi$を持ったニューラルネットワークで，$\epsilon$は学習の際にある程度の探索を課すためのノイズを表します．例しては，関数近似器が平均$\mu$と分散$\sigma^2$を出力して，学習時にはランダム性を考慮して$a_t \sim N(\mu, \sigma^2)$，評価時にはdeterministicな方策として$a_t = \mu$などとすることが考えられます．($N(\cdot, \cdot)$はガウス分布)
reparameterization trickでは学習時の確率的な方策を，deterministicな方策から得られた行動を平均とした確率分布からサンプリングすることで(ノイズ項を分離することで)，ノイズ項を含むニューラルネットの勾配を計算できるメリットがあります．
$Z$は$\phi$に依存しないので省略すると，上式は以下のように書き換えることができます．
$$
J_\pi (\phi) = E_{s_t \sim D, \epsilon_t \sim N} [\alpha \log \pi_\phi(f_\phi(\epsilon_t;s_t) | s_t) - Q_\theta (s_t, f_\phi(\epsilon_t;s_t)) ]
$$
reparameterization trickを用いることで，勾配を近似することができます．
$$
\hat\nabla_{\phi}J_\pi (\phi) = \nabla_{\phi} \alpha \log \pi_\phi(a_t | s_t) + \left( \nabla_{a_t} \alpha \log \pi_\phi(a_t | s_t) - \nabla_{a_t} Q(s_t, a_t) \right) \nabla_{\phi}f_\phi(\epsilon; s_t)
$$
エントロピー係数
Soft Actor-CriticをはじめとするMat-Ent RLでは，エントロピー項の係数$\alpha$はハイパーパラメータであり，$\alpha$の設定がヒューリスティックとなってしまう問題が生じます．さらに
- 最適な行動が不確かなときにはよりstochasticな(exploration)，一方で行動の良し悪しが明確なときにはよりdeterministicな(exploitation)方策が好ましい
- 学習が進むにつれ方策も改善されていく
ため，学習全体で固定した$\alpha$を用いるのは適切ではないと考えられます．そこで筆者らはMax-Ent RLを，以下のようなエントロピー制約付きの収益最大化問題に定式化し直しました．
\max_{\pi} \;\; E_{(s_t, a_t) \sim \rho_\pi} [\sum_{t=0}^{T}r(s_t, a_t)] \\
{\rm sub. to.}\;\; E_{(s_t, a_t) \sim \rho_\pi}[- \log(\pi(a_t|s_t))] \ge \tilde{H}
1行目は通常の収益最大化を表していて，2行目はエントロピー制約を表しています．ただし，$\tilde{H}$は期待エントロピーの最小値を表していて，上式はエントロピーの期待値がある定数以上になるような制約付きの収益最大化問題になっています．(簡単のため，論文とは少し定式化を変えています．)また，論文の実装では$\tilde{H} = - |A|$ ( $|A|$ は行動の次元数)としています.
ラグランジュの双対問題を考えると，
$$
\max_\pi \min_{\alpha_t \ge 0} E_{(s_t, a_t) \sim \rho_\pi}[\sum_{t=0}^{T}r(s_t, a_t)] - \alpha_t \left(E_{(s_t, a_t) \sim \rho_\pi}[- \log(\pi(a_t|s_t))] - \tilde{H} \right)
$$
となります．ここで第1項目は期待収益，第2項目はエントロピー制約，$\alpha_t$はラグランジュの双対変数です．ここで上式を最大化する方策はMax-Ent RLの最適方策$\pi^{*}$なので，最適な双対変数は
$$
\alpha_t^* = arg \min_{\alpha_t} E_{(s_t, a_t) \sim \rho_\pi}[- \alpha_t \log(\pi^*(a_t|s_t)) - \alpha_t \tilde{H}]
$$
です．よって，$\alpha$の(最小化する)目的関数を
$$
J(\alpha) = E_{s_t \sim D, a_t \sim \pi_t}[-\alpha \log \pi_t(a_t|s_t) - \alpha \tilde{H}]
$$
と設定します．各更新ステップで，方策を更新した後にその方策の期待エントロピーが設定した最小値に近くなるように$\alpha$を更新します．このことは，その方策の期待エントロピーが設定した最小値よりも大きい(stochasticな)ときには係数$\alpha$を小さく(方策がよりdeterministicになるように)し，その方策の期待エントロピーが最小値よりも小さい(deterministicな)ときには係数$\alpha$を大きく(方策がよりstochasticになるように)更新することを意味しています．
ラグランジュの双対性などに関しては，Boydのconvex optimizationに非常にわかりやすく解説されています．
アルゴリズム全体
SACでは，上で述べたようにsoft Q関数と方策をニューラルネットワークなどで関数近似します．今回は詳しく説明しませんが，TD3の論文(Fujimoto et al.(2018))で紹介されているclipped double Qなどのテクニックも用いられています．簡単に説明すると，Q関数の関数近似器を2つ用意してQ関数の推定の際には2つのQ関数を用いてそれぞれ推定して，小さい方を使うことで過大評価を防ごうぜってやつです．(かなり雑に説明してます，TD3の解説記事もあとで書きます...)
$$
y_\theta = min_{i = 1,2} ;\left( Q_{\bar \theta_i}(s, a) - \alpha \log \pi_\phi(a|s) \right)
$$
アルゴリズムの大まかな流れとしては，
- soft Q関数を2つのパラメータ$\theta_i (i = 1, 2)$と方策のパラメータ$\phi$を初期化
- soft Q関数のtarget networkのパラメータを初期化: $\bar \theta_i = \theta_i (i = 1, 2)$
- 以下を反復
- ある程度Agentを行動させたら，Replay Bufferから$(s_t, a_t, s'_t, r_t)$をサンプリングする．
- 2つのsoft Q関数を$\hat\nabla_{\theta_i}$$ J_Q(\theta_i)(i=1,2)$に基づいて更新．(ただしsoft Q関数の推定値には，2つの推定値のうち小さい方を共通して用いる)
- 方策を$\hat\nabla_{\phi}$$J_\pi (\phi)$に基づいて更新
- $\alpha$を$\hat\nabla_{\alpha} J(\alpha)$に基づいて更新
- target networkのパラメータ$\hat \theta_i (i=1,2)$をweight averageに基づいて更新
という感じになります．
結果
特に行動空間が複雑で，難しいと思えるタスクに関して良い感じに学習しているよ，という結果になりました．(Walker2d-v1での性能が低いのが気になりますが...)
論文ではablation study等も行なっているので，ぜひ参照してみてください．
まとめ
Soft Actor-Criticでは，通常のActor-Criticに方策のエントロピー最大化を加えることで，先行研究の課題であるサンプル効率・安定性(ハイパーパラメータへの鋭敏性)を改善しました．皆さんの勉強の助けになれば嬉しいです．
文章を書くのが苦手なので，内容に全く関係ないご指摘等も頂けたら嬉しいです．