概要

　土木の分野では、機械の制御（ダムの制御）や都市開発（交通×AI）などで、強化学習が使われ始めています。初めてニューラルネットワークを用いて強化学習を構築する場合、最初に構築するアルゴリズムは、やはり深層Q学習(DQN)だろうと思われます。

　DQNのようにQ関数だけで行動と評価を行うのではなく、Actor Critic のように方策とQ関数を分けて行動と評価を行うほうが、サンプル効率が良いのではないかと（なんとなく）考えて、調べてみると、Soft Actor Critic の離散バージョンがありました。サンプル効率が良いかどうかはわかりませんが（連続バージョンのSoft Actor Criticよりはサンプル効率が良いらしい）、ゲーム環境によっては、Rainbowより精度が良いとのこと。

　Soft Actor Criticと離散バージョンのSoft Actor Criticとの違いをかなり大雑把に述べると、ガウス分布かカテゴリカル分布かの違いです。今回は、Soft Actor Criticについて簡単に説明し、離散的なSoft Actor Criticについて紹介したいと思います。
　
　
　論文は、以下を引用しています。

　Soft Actor Criticについては、以下の論文を引用しています。

　私より詳しい説明とPyTorchでの実装は@ku2482さんが説明されています。

Maximum entropy の方法

　最初にMaximum entropy の方法について簡単に説明する。
　通常の報酬の期待値は、ポリシー（方策）を$\pi$とし、経験分布$\rho_{\pi}$（と言っていいのだろうか）を使い

\sum_t \mathbb{E}_{(s_t,a_t) \sim \rho_{\pi}  }\left[r(s_t,a_t) \right]

と書ける。$s_t,a_t,r(s_t,a_t) $は状態、行動、報酬である。この式にエントロピー$\mathcal{H}$を加え、目的関数を

J(\pi)=\sum_{t=0}^{T} \mathbb{E}_{(s_t,a_t) \sim \rho_{\pi}  }\left[r(s_t,a_t) + \alpha \mathcal{H}(\pi(\cdot | s_t) )\right]

とする。ハイパーパラメータ$\alpha$は温度である。
　最適な方策$\pi^*$は、目的関数が最大となるような$\pi$であり

\DeclareMathOperator*{\argmax}{arg\,max} %argmax
\pi^* =\argmax_{\pi} \sum_{t=0}^{T} \mathbb{E}_{(s_t,a_t) \sim \rho_{\pi}  }\left[r(s_t,a_t) + \alpha \mathcal{H}(\pi(\cdot | s_t) )\right]

となるように求める。

Energy-Based Model

　Energy-Basedの方策モデルについて説明する。
　Energy-Basedの方策モデルは、物理でおなじみのカノニカル分布の形をしており、エネルギー関数 $\mathcal{E}$ を使い

\pi(a_t| s_t)\propto \exp\left\{-\mathcal{E}(a_t,s_t) \right\}

と書ける。Q学習の場合は、エネルギー関数 $\mathcal{E}$ を

\mathcal{E}(a_t,s_t) = - \frac{1}{\alpha} Q_{\rm{soft}}(a_t,s_t)

とする。

Soft Actor Critic のニューラルネットワークのパラメータの更新

　Soft Actor Critic のニューラルネットワークのパラメータの更新について説明する。つまり、Actor Critic は２つのニューラルネットワーク、方策 $\pi_{\phi}$ とQ関数 $Q_{\theta}$ のパラメータを更新する必要がある。

　最初に状態価値関数 $V(s_t)$ について説明する。Soft ベルマン演算子 $\mathcal{T}^{\pi}$を使い、$k+1$ステップにおけるQ関数を

Q^{k+1} = \mathcal{T}^{\pi} Q^k

とする。そして、状態価値関数 $V(s_t)$ は以下のように定義し、

V(s_t) = \mathbb{E}_{a_{t} \sim \pi} \left[Q(s_t,a_t) -\alpha \log\pi(a_t|s_t) \right]

Soft ベルマン演算子 $\mathcal{T}^{\pi}$は、状態価値関数を用いてQ関数を以下のように変換する。

\mathcal{T}^{\pi} Q(s_t,a_t) = r(s_t,a_t) + \gamma \mathbb{E}_{s_{t+1} \sim p}\left[V(s_{t+1}) \right]

　Q関数の収束については、報酬をエントロピー付きで定義すれば、

r_{\pi}(s_t,a_t) \equiv r(s_t,a_t)+\gamma \mathbb{E}_{s_{t+1} \sim p}\left[\mathcal{H}(\pi(\cdot | s_t) ) \right]

Soft ベルマン演算子 $\mathcal{T}^{\pi}$によるQ関数の更新は、

Q(s_t,a_t) \longleftarrow  r_{\pi}(s_t,a_t)  + \gamma \mathbb{E}_{s_{t+1} \sim p ,a_{t+1} \sim \pi}\left[Q(s_{t+1},a_{t+1}) \right]

と書け、（Soft ではない）通常のベルマン演算子を用いたQ関数の更新となることから、Q関数の収束について証明される（らしい）。
　
　次に Q 関数$Q_{\theta}(s_t,a_t)$のパラメータ$\theta$の更新について説明する。
　 Q 関数$Q_{\theta}(s_t,a_t)$のパラメータ$\theta$を決めるための目的関数は、推定値 $Q$ とベルマン演算子から更新された値 $\hat{Q}$ との誤差から計算される。

J_Q(\theta) = \mathbb{E}_{(s_t,a_t)\sim D}\left[\frac{1}{2}\left(Q_{\theta}(s_{t},a_{t})-\hat{Q}(s_{t},a_{t}) \right)^2 \right]

ベルマン演算子から更新された値 $\hat{Q}$ とは、

\hat{Q}(s_t,a_t) = r(s_t,a_t) + \gamma \mathbb{E}_{s_{t+1} \sim p}\left[V_{\bar{\theta}}(s_{t+1}) \right]

を指す。$(s_t,a_t)\sim D$は Replay Bufferからサンプルされた状態・行動の組である。

　状態価値関数 $V(s_t)$ には $Q_{\theta}(s,a)$ を使うが、学習で勾配を更新する際は、パラメータ$\theta$は固定する。つまり、強化学習を安定させるために用いるtarget network の値を使う。target network のパラメータには、上付きbar（例えば$\bar{\theta}$）をつける。

　$J_Q(\theta)$ の勾配を計算すると、期待値を計算する前の微分を$\hat{\nabla}$として

\hat{\nabla}_{\theta} J_Q(\theta) = \nabla_{\theta} Q_{\theta}(s_t,a_t) \left( Q_{\theta}(s_{t},a_{t}) - r(s_t,a_t) - \gamma V_{\bar{\theta}}(s_{t+1}) \right)

と計算され、勾配法で$\theta$を更新する。target network のパラメータ$\bar{\theta}$は、

\bar{\theta} \longleftarrow \tau \theta + (1-\tau) \bar{\theta}

で更新する。

　最後に方策$\pi_{\phi}(a_t|s_t)$のパラメータ$\phi$の更新について説明する。
　方策$\pi_{\phi}(a_t|s_t)$のパラメータ$\phi$を決めるための目的関数は、カルバック・ライブラー情報量（KL divergence）を用いて定義される。

J_{\pi}(\phi) = \mathbb{E}_{s_t\sim D} \left[D_{\rm{KL}}\left(\pi_{\phi}(\cdot |s_t) 
 \Biggr|\Biggr| \frac{\exp\{\frac{1}{\alpha}Q_{\theta}(s_t,\cdot) \} }{Z_{\theta}(s_t)} \right) \right]

KL divergence は、２つの確率分布の距離（厳密には違うが）を表し、方策 $\pi$ は、Energy-Basedの方策モデルの確率分布に近づくように更新される。
　ただし、確率分布は誤差逆伝播法が使えないので、reparameterization trickを使う。つまり、行動 $a_t$ はガウス分布からのノイズ $\epsilon_t$ を使い、

a_t = f_{\phi}(\epsilon_t,s_t)

　KL divergenceの定義から目的関数は、以下のように書く。

J_{\pi}(\phi) = \mathbb{E}_{s_t\sim D ,\epsilon_t \sim \mathcal{N}} \left[\alpha\log\pi_{\phi}(f_{\phi}(\epsilon_t,s_t) |s_t) -
  Q_{\theta}(s_t,f_{\phi}(\epsilon_t,s_t)) +\alpha\log Z_{\theta}(s_t)  \right]

$\log Z_{\theta}(s_t)$は、パラメータの更新に効いてこないのでプログラムでは省く。
　$J_{\pi}(\phi)$ の勾配を計算すると、以下となる。

\hat{\nabla}_{\phi} J_{\pi}(\phi)= \alpha\nabla_{\phi}\log\pi_{\phi}(a_t|s_t) + \left(\alpha \nabla_{a_t}\log\pi_{\phi}(a_t|s_t)- \nabla_{a_t}Q_{\theta}(s_t,a_t)  \right)\nabla_{\phi}f_{\phi}(\epsilon_t,s_t)

また、$\nabla_{\phi}$は、$\pi_{\phi}$とreparameterization trickで使った$f_{\phi}$に作用する。そして、勾配法で$\phi$を更新する。

正規化流

　複雑なモデルにおいて、潜在変数の事後分布は、単純なガウス分布で近似することは難しい。したがって、複数回、微分可能で可逆な関数 $f$ を使い変換する事で、複雑な分布からのサンプルを取得する方法を正規化流という。正規化流について説明する。
　関数 $f$ で変換前の値を $z$ 返還後を $\hat{z}$とする。つまり、

\hat{z} = f(z)

さらに $f$ は可逆なので、

z = f^{-1}(\hat{z})

となる。$\hat{z}$で微分すると

\mathrm{d}z = \frac{\mathrm{d}f^{-1}}{\mathrm{d}\hat{z}} \mathrm{d}\hat{z}

となる。次に確率分布$q(z)$として座標変換を考えると

\int \mathrm{d}\hat{z} q(\hat{z})= \int \mathrm{d}z q(z)= \int  \mathrm{d}\hat{z}  q(z)\left| \mathrm{det}\left\{\frac{\mathrm{d}f^{-1}}{\mathrm{d}\hat{z}} \right\} \right|

つまり、（逆関数の定理を使うと？）

q(\hat{z})= q(z)\left| \mathrm{det}\left\{\frac{\mathrm{d}f^{-1}}{\mathrm{d}\hat{z}} \right\} \right| =  q(z)\left| \mathrm{det}\left\{\frac{\mathrm{d}f}{\mathrm{d}z} \right\}^{-1} \right|

さらに、対数をとり

\log q(\hat{z})= \log q(z) - \log  \mathrm{det}\left\{\frac{\mathrm{d}f}{\mathrm{d}z} \right\}

が得られる。
　可逆で微分可能な関数として$f(u)=\tanh(u)$を採用し

\frac{\mathrm{d}}{\mathrm{d}u} \tanh(u) = 1- \tanh^2(u)

$\mu(u|s)$ をガウス分布として、方策 $\pi(a|s)$ を

\log \pi(a|s)= \log \mu(u|s) - \log  \mathrm{det}\left\{\frac{\mathrm{d}f(u)}{\mathrm{d}u} \right\}

と表せば、$\log \mathrm{det} A = \mathrm{Tr}\log A $を用いて、最終的に方策$\pi(a|s)$は以下のように表せる。

\log \pi(a|s)= \log \mu(u|s) - \sum_{i=1}^D \log\left(1- \tanh^2(u_i) \right)

温度の調整

　最後にハイパーパラメータであった温度$\alpha (\geq0)$の調整について説明する。強化学習は、ハイパーパラメータの調整がとても難しいので、学習で調整してくれるのはうれしい。

　任意の$t$に対して条件

 \mathbb{E}_{(s_t,a_t)\sim\rho_{\pi}} \left[-\log\pi_t(a_t|s_t)   \right] \geq \mathcal{H}

のもと

\underset{\pi[0:T]}{\max} \mathbb{E}_{\rho_{\pi}}\left[\sum_{t=0}^T r(s_t,a_t )\right]

を求めることを考える（条件付最適化）。この式を展開してみると

\underset{\pi_0}{\max}\left\{\mathbb{E}[r(s_0,a_0)]+\underset{\pi_1}{\max}\left\{ \mathbb{E}[r(s_1,a_1)]+\ldots \left\{\underset{\pi_T}{\max}\mathbb{E}[r(s_T,a_T)] \right\}\ldots\right\} \right\}

の形となる。

まず、$t = T$ に注目すると、条件

 \mathbb{E}_{(s_T,a_T)\sim\rho_{\pi}} \left[-\log\pi_T(a_T|s_T)   \right] -\mathcal{H}  \geq 0

から、$\alpha_T$ を双対変数として、最適な方策 $\pi_T^{\ast}$ は以下の式を満たすように求める（双対問題？）。

\underset{\pi_T}{\max}\mathbb{E}[r(s_T,a_T)] = \underset{\alpha_T\geq 0}{\min} \left\{\underset{\pi_T}{\max} \mathbb{E}_{(s_T,a_T)\sim\rho_{\pi}}\left[r(s_T,a_T)- \alpha_T \log\pi_T(a_T|s_T)\right] - \alpha_T \mathcal{H} \right\}

　温度$\alpha_T$におけるMaximum entropy 方策を $\pi^*_T(a_T|s_T:\alpha_T)$ とすれば、最適な温度$\alpha_T^{\ast}$は以下のように求める。

\DeclareMathOperator*{\argmin}{arg\,min} %argmin
\alpha_T^* =\argmin_{\alpha_T} \mathbb{E}_{s_T,a_T\sim \pi_T^*}\left[- \alpha_T \log\pi_T^*(a_T|s_T:\alpha_T) - \alpha_T \mathcal{H} \right]

次に、$t=T-1$に注目すると、$t=T$における最適な Q 関数は以下であることに注意すれば

Q_{T}^{\ast}(s_{T},a_{T}) = \mathbb{E}[r(s_{T},a_{T})]

Q 関数は、$t=T$において最適化された温度$\alpha_T^{\ast}$と方策$\pi^*_T$の結果を使えば、以下のように書ける。

\begin{align}
Q_{T-1}^{\ast}(s_{T-1},a_{T-1}) &= \mathbb{E}[r(s_{T-1},a_{T-1})] + \mathbb{E}\left[ Q_{T-1}^{\ast}(s_{T},a_{T}) - \alpha_T^{\ast} \log\pi_T^{\ast}(a_T|s_T) \right]　\\
&= \mathbb{E}[r(s_{T-1},a_{T-1})] +\underset{\pi_T}{\max}\mathbb{E}[r(s_T,a_T)]＋\alpha_T \mathcal{H}
\end{align}

また、最適な方策 $\pi_{T-1}^{\ast}$ は以下の式を満たすように求める。

\underset{\pi_{T-1}}{\max}Q_{T-1}^{\ast}(s_{T-1},a_{T-1})　= \underset{\alpha_{T-1}\geq 0}{\min} \left\{\underset{\pi_{T-1}}{\max} \mathbb{E}\left[Q_{T-1}^{\ast}(s_{T-1},a_{T-1})- \alpha_{T-1} \log\pi_{T-1}(a_{T-1}|s_{T-1})\right] - \alpha_{T-1} \mathcal{H}\right\}

上２つの式を組み合わせれば

\begin{align}
&\underset{\pi_{T-1}}{\max}Q_{T-1}^{\ast}(s_{T-1},a_{T-1})-\alpha_T^{\ast} \mathcal{H}　= \underset{\pi_{T-1}}{\max}\left\{\mathbb{E}[r(s_{T-1},a_{T-1})] +\underset{\pi_T}{\max}\mathbb{E}[r(s_T,a_T)] \right\} \\
&=\underset{\alpha_{T-1}\geq 0}{\min} \left\{\underset{\pi_{T-1}}{\max} \mathbb{E}\left[Q_{T-1}^{\ast}(s_{T-1},a_{T-1})- \alpha_{T-1} \log\pi_{T-1}(a_{T-1}|s_{T-1})\right] - \alpha_{T-1} \mathcal{H}\right\}-\alpha_T^{\ast} \mathcal{H} 
\end{align}

が得られる。上式の１行目の右辺に注目すると、この操作を $t=0$まで行えば、

\underset{\pi[0:T]}{\max} \mathbb{E}_{\rho_{\pi}}\left[\sum_{t=0}^T r(s_t,a_t )\right]

が得られることが想像できる。したがって、最適な方策 $\pi_{t}^{\ast}$が求まり、条件を満たすように、最適な温度$\alpha_{t}^{\ast}$を決めるには、

\DeclareMathOperator*{\argmin}{arg\,min} %argmin
\alpha_{t}^* =\argmin_{\alpha_{t}} \mathbb{E}_{s_{t},a_{t}\sim \pi_{t}^*}\left[- \alpha_{t} \log\pi_{t}^*(a_{t}|s_{t}:\alpha_{t}) - \alpha_{t} \mathcal{H} \right]

となるように求めればよい。

　したがって、温度$\alpha$に関する目的関数は、以下のように構築し

J(\alpha) = \mathbb{E}_{a_{t} \sim \pi_t} \left[-\alpha \left(\log\pi(a_t|s_t)+  \mathcal{H}\right) \right]

上式を最小となるように$\alpha$を勾配法で求めればよい。

離散的なSoft Actor Critic

　離散的なSoft Actor Criticは、Soft Actor Criticと相違点が５点ありそれを説明する。

(１).　Q 関数の入出力
　Q関数は、$Q: S \times A \to \mathbb R$ から$Q: S \to\mathbb R^{|A|} $に変更。つまり、 $Q(s,a) \to Q(s)$ 。
Soft Actor Critic のQ 関数は状態と行動を入力していた。離散的なSoft Actor CriticのQ関数は状態のみを入力とする。

(２).　方策の出力
　方策は、$\pi: S \to \mathbb R^{2|A|}$ から$ S \to \mathbb [0,1]^{|A|} $に変更。
つまり、離散的なSoft Actor Criticは、Softmax関数を使い直接確率を出力させる。

(３).　状態価値関数
　Soft Actor Criticは、平均・分散を計算し、reparameterization trickを使い行動 $a$ をサンプリングしていた。

V(s_t) = \mathbb{E}_{a_{t} \sim \pi} \left[Q(s_t,a_t) -\alpha \log\pi(a_t|s_t) \right]

離散的なSoft Actor Criticは、方策の出力が直接確率になっているので、

V(s_t) = \pi^T(s_t) \left[Q(s_t) -\alpha \log\pi(s_t) \right]

と変更される。

(４).　温度調整のための目的関数
　行動 $a$ をサンプリングする必要がないので、(３)と同様に

J(\alpha) = \mathbb{E}_{a_{t} \sim \pi_t} \left[-\alpha \left(\log\pi(a_t|s_t)+  \mathcal{H}\right) \right]

から

J(\alpha) =  \pi_t^T(s_t) \left[-\alpha \left(\log\pi(s_t)+  \mathcal{H}\right) \right]

と変更される。

(５).　方策のパラメータを決定するための目的関数
　行動 $a$ をサンプリングする必要がないので、(３)、（４）と同様に

J_{\pi}(\phi) = \mathbb{E}_{s_t\sim D ,\epsilon_t \sim \mathcal{N}} \left[\alpha\log\pi_{\phi}(f_{\phi}(\epsilon_t,s_t) |s_t) -
  Q_{\theta}(s_t,f_{\phi}(\epsilon_t,s_t))  \right]

から

J_{\pi}(\phi) = \mathbb{E}_{s_t\sim D } \left[ \pi_t^T(s_t) \left(\alpha\log\pi_{\phi}(s_t) -
  Q_{\theta}(s_t) \right) \right]

と変更される。

最後に

まだ学習をさせてないので、結果は次回？