Soft Actor-Critic(SAC)でdonkeycarのシミュを走らせてみる

東大の松尾先生の研究室主催のセミナー(RLSP2021)の復習に書いた記事です.
1回生で参加し, 知識不足ではありましたがとても勉強になりました. ありがとうございました.

f:id:Emile11235:20210810191909p:plain

記号の定義
SACとは?
environmentについて
実装
感想など
出典,参考など

記号の定義

$\mathcal{S, A}$ : 状態, 行動
$p: \mathcal{S \times A} \rightarrow \mathbb{R}$ : 遷移確率関数
$\pi : \mathcal{S \rightarrow A}$ : 方策
$r : \mathcal{S \times A} \rightarrow \mathbb{R}$ : 報酬関数

SACとは?

概要

エントロピー項 $H(X)$ を追加した目的関数 $J(\theta) := \mathbb{E}[\sum_{t} r(s_t, a_t) + \alpha H(\pi_\phi(\cdot | s_t))]$ を用いる手法.
状態空間, 行動空間が連続な場合にも適用可能.

※大抵は目的関数は $J(\theta) := \mathbb{E}[\sum_{t} r(s_t, a_t)]$
※エントロピー項の定義は $H(X) := - \sum_{x \in U} P_X(x) \log P_X(x)$

Criricは $Q$ 関数( $Q_\theta : \mathcal{S \times A} \rightarrow \mathbb{R}$ )を近似し, Actorは $\pi_\phi$ ( $Q$ 関数のsoftmax方策)を近似する.
ActorはGaussian Policyを用い, 分布の平均値と分散を出力する事で行動の分布を生成する.

学習とそれぞれの損失関数

・Critic( $Q_\theta$ )の損失関数

Criticはエントロピー項を考慮した $Q$ 関数を近似する.
具体的にはベルマン作用素 $B^\pi Q_\theta(s_t,a_t) := r(s_t, a_t) + \gamma \mathbb{E}_{s_{t+1}\ \ \sim p}[ V(s_{t+1}) ]$ を用いて更新する.
※状態価値関数は $V(s_t) := \mathbb{E}_{a_t \ \sim \pi}[ Q(s_t, a_t) - \alpha \log \pi(a_t | s_t) ]$ と定義される

このベルマン作用素はエントロピーの最大化を含む様に修正されているが,
$r'(s_t, a_t) := r(s_t, a_t) - \alpha \log \pi (a_t | s_t)$ を報酬関数として用いると解釈すれば,
通常のベルマン期待作用素と同じ形になり, 同等の性質が言える.
(targetの方策に基づいた $Q$ 関数に収束. 方策の評価をすると解釈できる.)

以上から, Criticは損失関数として
$J_Q(\theta) := \mathbb{E}_{(s_t,a_t)\ \sim D} [ \frac{1}{2} (Q_\theta (s_t, a_t) - (r(s_t, a_t) + \gamma \mathbb{E}_{s_{t+1}\ \ \sim p} [ V_{\bar{\theta}} (s_{t+1}) ] )\ )^2 ]$
を用いて学習する. (ベルマン作用素を適用した $Q$ 関数に最も近くなるように最小二乗法をかけてる)

※ $\bar{\theta} := \min_{\theta^{target}_i}\ \ [ V_{\theta^{target}_i}\ \ (s_{t+1}) ] = \min_{\theta^{target}_i}\ \ [ Q_{\theta^{target}_i}\ \ (s_t,a_t) - \alpha \log \pi_\phi (a_t | s_t) ]$ であり,
double networkとTarget Q-networkと呼ばれる手法を用いている.
(不確かな時は楽観的にというヒューリスティックと急激な変化を抑える意図がある)

・Actor( $\pi$ )の損失関数

$\pi$ は $Q$ 関数のsoftmax方策を近似する.
方策は $\pi_{new} := \mathrm{argmin}_{\pi' \in \Pi} D_{KL}(\pi'(\cdot | s_t) || \frac{\exp(Q_\theta \ \ (s_t, \cdot) )}{Z_\theta (s_t)})$ として更新する.
(方策の改善と解釈できる.)

また, $\mathbb{E}_{s_t \ \sim D}[ D_{KL}(\pi'(\cdot | s_t) || \frac{\exp(Q_\theta \ \ (s_t, \cdot) )}{Z_\theta (s_t)}) ] = \mathbb{E}_{a_t\ \ \sim \pi}[ \alpha \log \pi_\phi(a_t | s_t) - Q_\theta (s_t, a_t) ] + const.$
より, Actorは損失関数として $J_\pi(\phi) = \mathbb{E}_{a_t\ \ \sim \pi}[ \alpha \log \pi_\phi(a_t | s_t) - Q_{\theta'} (s_t, a_t) ]$ を用いて学習する.

※ $Z_\theta$ はsoftmax関数の正規化項, $\theta' := \min_{\theta_i}\{ Q_{\theta_i}(s_t, a_t) \}$ .

ただし, $a_t \sim N(\mu_t, \sigma_t)$ として行動を生成すると微分が出来なくなるので,
行動を $\epsilon_t \sim N(0, I)$ を用いて, $a_t := \mu_t + \epsilon_t \circ \sigma_t$ として行動を生成する事で, 微分できる様にする.
(reparameterization trickと呼ばれる手法)

全体の流れ

SACではExperience ReplayとTarget Q-Networkを用いる事に注意.
ActorとCriticの関係は以下の様になる.
f:id:Emile11235:20210810021217p:plain

全体の流れとしては,
1. 現在の方策から行動を生成し, $(s_t, a_t, s_{t+1}, r_t)$ をバッファ $B$ に収集.
2. $\{s_t, a_t, s_{t+1}, a_t\}$ をバッファ $B$ からサンプリング.
3. Q関数を損失関数 $J_Q(\theta)$ を用いて, $\theta \leftarrow \theta - \lambda_Q \nabla_{\theta_i} J_Q(\theta_i) \ \ \ (i=1,\ 2)$ として更新
4. 方策を損失関数 $J_\pi(\phi)$ を用いて, $\phi \leftarrow \phi - \lambda_\pi \nabla_\phi J_\pi(\phi)$ として更新
5. 何回かQ関数と方策を更新したら,
target Q-functionも $\theta^{target}_i \leftarrow \tau \theta^{target}_i + (1-\tau)\theta^{target}_i \ \ \ \ (i = 1,2)$ として更新

environmentについて

gym-donkeycarを利用した.
github.com

ドキュメントに書いてある仕様と実際の仕様が異なっていた気がします.
(どこか忘れてしまいました.)

実装

この実装では, envから得られる状態にVAEで前処理をした物を3frame分集めて1つのstateにした.
また, actionも3回連続で入力した.

(VAEはRLSP2021で同じチームだったikedaさん(Ryunosuke Ikeda (@ImR0305) | Twitter)が実装してくださいました. ありがとうございます.)

github.com

f:id:Emile11235:20210810033059g:plain
↑走行時の映像

※VAEはEncoderとDecoderを持っておいて入力をEncodeした物をDecodeした物と入力の誤差が小さく成る様にEncoder, Decoderを学習させる事で入力をEncodeする事で次元削減できるという手法.

f:id:Emile11235:20210815214554p:plain
↑criticのlossのグラフ

f:id:Emile11235:20210815214558p:plain
↑actorのlossのグラフ

f:id:Emile11235:20210815214602p:plain
↑報酬のグラフ