今回はTRPOとA3C/A2Cの後継にあたるPPOを実装してみました。

※ネット上の情報をかき集めて自分なりに実装しているので正確ではない可能性がある点はご注意ください
※ライブラリはTensowflow2.0(+Keras)を使っています。

追記：自作フレームワークを作成しており、そちらで改めて記事を書いています。そちらの方が正確なコードとなります。
・【強化学習】PPOを解説・実装

コード全体

本記事で作成したコードは以下です。
（実装は分かりやすさを重視して速度を考慮していません）

GoogleColaboratory

PPO(Proximal Policy Optimization)

TRPOは実装がとても複雑でした。
また、ActorとCriticでパラメータ共有をするA3C型のアーキテクチャでの利用が難しい点やCNNやRNNで性能が悪いという課題がありました。

そこでTRPOをシンプルにした手法がPPOとなります。
PPOは実装の簡易さと高パフォーマンスが売りなようです。
また、PPOはA3C/A2Cの手法も参考にしているため並列学習が入っていますが、本実装ではアルゴリズムの理解をメインにするため並列学習は省いて実装していきます。

参考

Clipped Surrogate Objective

TRPOでは以下の式(代理目的関数:Surrogate Objective)の最大化が目的でした。(TRPOに関しては第5回を参照)

$$
\underset{\theta}{\text{maximize}} L(\theta) = \hat{E}[ \frac{\pi_{\theta} (a|s)}{ \pi_{\theta_{old}}(a|s)} \hat{A} ]
$$

TRPOでは制約条件を加えることで上記の更新を大きくしないように＝更新前後で方策が大きく変わらないようにしました。
PPOでは制約条件の代わりにclipすることで大きな更新幅を抑え、計算の複雑さを減らしています。
また、clipされた目的関数はclipされていない目的関数と比べ小さい値を採用します。

$$
\frac{\pi_{\theta} (a|s)}{ \pi_{\theta_{old}}(a|s)} := r(\theta)
$$

と置いた場合のclipの条件は以下です。

$$
L^{CLIP}(\theta) = \hat{E}[ min(r(\theta) \hat{A}, clip(r(\theta), 1 - e, 1+e) \hat{A})]
$$

$e$ はclipの範囲で 0.2 が使われるようです。
コードだと以下です。

old_pi = 更新前の方策確率
new_pi = 更新後の方策確率
CLIPRANGE = 0.2  # clip範囲

# IS(重点サンプリング)の計算
# ISに関しては第5回を参照
ratio = new_pi / old_pi

# clipする
ratio_clipped = tf.clip_by_value(ratio, 1 - CLIPRANGE, 1 + CLIPRANGE)

# loss を計算
loss_unclipped = ratio * advantage
loss_clipped = ratio_clipped * advantage

# 小さいほうのlossを採用
policy_loss = tf.minimum(loss_unclipped, loss_clipped)

Adaptive KLペナルティ

TRPOの制約条件に関する別のアプローチとして、目的関数に以下のペナルティを組み込む手法も提案されています。
ただ、あまり性能向上に貢献できていないようです。

L^{KLPEN}(\theta) = \hat{E}
\begin{bmatrix}
\frac{\pi_{\theta} (a|s)}{ \pi_{\theta_{old}}(a|s)} \hat{A} 
- \beta KL[\pi_{\theta_{old}} ( . |s) || \pi_{\theta} (.|s)]
\end{bmatrix}

$KL$はKL距離で詳細は第5回を見てください。
$\beta$ は以下の式で調整されるパラメータです。

d = \hat{E}
\begin{bmatrix}
    KL[\pi_{\theta_{old}} ( . |s) || \pi_{\theta} (.|s)
\end{bmatrix}
\\
\beta = \left\{
\begin{array}{ll}
\frac{\beta}{2} & (d < \frac{d_{targ}}{1.5})\\
2 \beta & (d > 1.5 d_{targ})\\
\beta & (otherwise)

\end{array}
\right.

$d_{targ}$ は定数(0.01等)です。
$\beta$ の初期値や2や1.5といった定数はハイパーパラメータですが、更新によって自動で調整されるためあまり気にしなくていいそうです。

コードは以下です。

# KL距離の計算(離散ver)
def compute_kl_discrete(probs1, probs2):
    return tf.reduce_sum(probs1 * tf.math.log(probs1 / probs2), axis=1, keepdims=True)

# KL距離の計算(ガウス分布ver)
def compute_kl_discrete_normal(mean1, stddev1, mean2, stddev2):
    a1 = tf.math.log(stddev2 / stddev1)
    a2 = (tf.square(stddev1) + tf.square(mean1 - mean2)) / (2.0 * tf.square(stddev2))
    a3 = -0.5
    return a1 + a2 + a3

adaptive_kl_targ = 0.01
adaptive_kl_beta = 0.5

# 学習のタイミング
def train():

    states = 状態
    old_probs = 経験取得時の各アクションの確率
    
    # 勾配計算
    with tf.GradientTape() as tape:
        new_probs = 現在のモデルから出したアクション確率

        policy_loss = 計算した後の方策loss

        # KLペナルティ
        kl = compute_kl_discrete(old_probs, action_probs)
        policy_loss -= adaptive_kl_beta * kl

        （略）

        # ミニバッチ
        loss = tf.reduce_mean(loss)

    # KLペナルティβの調整
    kl_mean = tf.reduce_mean(kl)
    if kl_mean < adaptive_kl_targ / 1.5:
        adaptive_kl_beta /= 2
    elif kl_mean > adaptive_kl_targ * 1.5:
        adaptive_kl_beta *= 2

その他のテクニック

TRPO/PPOにおける実装において、細かいテクニックがいくつかあるようです。

Implementation Matters in Deep Policy Gradients: A Case Study on PPO and TRPO(論文)

1. Value Clipping

Clipped Surrogate Objective は方策モデルに対して行っていますが、これを価値関数にも行う方法です。

$$
L^V = \max[
(V_{\theta_{t}} - V_{targ})^2 ,
(clip(V_{\theta_{t}}, V_{\theta_{t-1}} - \epsilon, V_{\theta_{t-1}} + \epsilon) - V_{targ})^2
]
$$


old_v = 経験収集時の状態価値
advantage = 計算した割引報酬率(にbaseline等加味したもの)

pi_clip_range = 0.2  # クリップ幅

# 勾配計算箇所
with tf.GradientTape() as tape:
    # 現在の状態価値
    action_probs, v = model(states, training=True)

    （略）

    #--- Value loss
    v_clipped = tf.clip_by_value(v, old_v - pi_clip_range, old_v + pi_clip_range)
    value_loss = tf.maximum((v - advantage) ** 2, (v_clipped - advantage) ** 2)

    （略）

2. Reward scaling

割引報酬和のスケーリング(baseline)ですが、標準偏差で割る手法です(平均は引きません)
報酬がスパース(報酬がなかなか取得できない場合)は逆に不安定になるようです。

コードはnoneも含め4種類作成しておきました。

advantage = 割引報酬
baseline_type = "std"

# baseline
if baseline_type == "none":
    pass
elif baseline_type == "ave":
    advantage -= np.mean(advantage)
elif baseline_type == "std":
    advantage = advantage / (np.std(advantage) + 1e-8)
elif baseline_type == "normal":
    advantage = (advantage - np.mean(advantage)) / (np.std(advantage) + 1e-8)

3. Orthogonal initialization and layer scaling

方策・価値モデルの重みの初期化ですが、レイヤー毎にバラバラになるように直交行列で初期化したほうがパフォーマンスが向上するようです。

Kerasではレイヤーの kernel_initializer 引数をOrthogonalにする事で実現できます。

init_layer_weight = "orthogonal"

# kernel_initializer 省略時の初期値
if init_layer_weight == "":
    init_layer_weight = "glorot_uniform"

nb_actions = アクション数
dense_units = 16
activation = "tanh"

# 各レイヤーの定義
dense1 = keras.layers.Dense(dense_units, activation=activation, kernel_initializer=init_layer_weight)
dense2 = keras.layers.Dense(dense_units, activation=activation, kernel_initializer=init_layer_weight)
actor_layer = keras.layers.Dense(nb_actions, activation="linear", kernel_initializer=init_layer_weight)
critic_layer = keras.layers.Dense(1, activation="linear", kernel_initializer=init_layer_weight)

4. Adam learning rate annealing

Adamの学習率ですが、徐々に下げていく(焼きなまし法)ほうがパフォーマンスが向上するようです。
Kerasでは optimizer の学習率 lr は直接変更できます。

optimizer_initial_lr = 0.02  # 初期学習率
optimizer_final_lr = 0.01    # 終了学習率
optimizer_lr_step = 200*50   # 終了学習率になるまでの更新回数

model.optimizer = Adam()
update_count = 0

# 1回学習毎
def update():

    if update_count > optimizer_lr_step:
        model.optimizer.lr = optimizer_final_lr
    else:
        model.optimizer.lr = optimizer_initial_lr - \
            (optimizer_initial_lr - optimizer_final_lr) * \
            (update_count / optimizer_lr_step)
    update_count += 1

5. Reward Clipping

報酬は事前設定された範囲([-5,5]または[-10,10]など)でclipします。

reward_clip = [-5, 5]  # clip範囲

# 1stepのタイミング
n_state, reward, done, _ = env.step(action)

# 報酬のclip
if reward_clip is not None:
    reward = np.clip(reward, reward_clip[0], reward_clip[1])

6. Observation Normalization

状態は平均0、分散1に正規化して学習に使います。

# 学習時

states = 状態

# 状態を正規化
states = (states - np.mean(states, axis=0, keepdims=True)) / \
    (np.std(states, axis=0, keepdims=True)  + 1e-8)

7. Observation Clipping

報酬と同様に状態もclipします。(通常は[-10,10])

state_clip = [-10, 10]  # clip範囲

# 1stepのタイミング
n_state, reward, done, _ = env.step(action)

# 状態のclip
if state_clip is not None:
    state = np.clip(state, state_clip[0], state_clip[1])

8. Hyperbolic tan activations

方策/価値モデルの活性化関数は Hyperbolic tan(tanh) が使われるようです。
コードは 3. に記載してあります。

9. Global Gradient Clipping

方策/価値モデルの勾配を計算した後、全てのパラメータはL2(ノルム2)でclipします。

global_gradient_clip_norm = 0.5  # clip範囲

with tf.GradientTape() as tape:
    （省略）
    loss = tf.reduce_mean(loss)

grads = tape.gradient(loss, model.trainable_variables)
# L2でクリップ
grads, _ = tf.clip_by_global_norm(grads, global_gradient_clip_norm)
model.optimizer.apply_gradients(zip(grads, model.trainable_variables))

(おまけ)アクションの(逆?)正規化

論文には書いていないですが実装上のテクニックとして書いておきます。
連続行動空間に限った話ですが、アクションの範囲を制限して学習します。

学習では[-1,1]の範囲で学習し、環境へは適切な範囲に置き換えて渡します。
例えば環境の取りうる値が[0,10]の場合は $5a + 5$ して渡します。

action_centor = (action_space.high + action_space.low)/2
action_scale = action_space.high - action_centor

action = アクション

# 1step
n_state, reward, done, _ = env.step((action * action_scale) + action_centor)

Experience Memory

経験の収集ですが、過去のデータも使えるので Experience Memory を実装します。
PPO の大きな利点として重点サンプリングにより、On-policy学習であるにもかかわらず過去の経験が使える点です。
(重点サンプリングを使わない場合は、過去の経験は違う方策で経験した内容のため学習に使えない)

コードは以下です。

gamma = 0.9          # 割引率
warmup_size = 100    # 最低限貯めるキューサイズ
buffer_size = 1000   # キューの最大サイズ
batch_size = 32      # バッチサイズ

# 収集する経験は上限を決め、古いものから削除する
from collections import deque
experiences = deque(maxlen=buffer_size)

# 学習ループ
for episode in range(300):
    （略）

    # 1エピソードの経験
    episode_exp = []

    # 1episodeループ
    while not done:
        （略）

        # 経験を保存する
        episode_exp.append({
            "state": state,
            "action": action,
            "reward": reward,
            "n_state": n_state,
            "done": done,
            # 学習用に追加
            "action_prob": action_prob,  # アクションの確率
            "v": v,                      # 状態価値
        })

        # wamup_size 分経験が貯まれば学習を開始する
        if len(experiences) >= warmup_size:

            # ランダムに経験を取得してバッチを作成
            batchs = random.sample(experiences, batch_size)
            
            # 学習
            train(batchs)

    #--- 1エピソード終わったので推定価値を計算(GAE or モンテカルロ法)
    # コードはモンテカルロ法の例
    for i,exp in enumerate(episode_exp):
        G = 0
        t = 0
        for j in range(i, len(episode_exp)):
            G += (gamma ** t) * episode_exp[j]["reward"]
            t += 1
        exp["discounted_reward"] = G
        
        # 割引報酬を計算した経験を、experiences に追加する
        experiences.append(exp)

実装1、A2C/A3Cベース

PPOは実装が柔軟で、アクションが離散値・連続値どちらも実装ができ、かつ方策モデルと状態価値モデルが一緒になっているモデル(A2C/A3C)と方策モデルと状態価値モデルが分かれているモデル、どちらの場合も実装できます。
アクションが離散値でモデルが統一されているパターンと、アクションが連続値でモデルが分かれているパターンを実装してみました。

環境	アクション	モデル
CartPole-v0	離散値	方策モデルと状態価値モデルが同じ

実装にあたりいろんなテクニックをハイパーパラメータとして設定できるようにしてみました。
ハイパーパラメータは以下です。

train_interval = 10  # 学習間隔
gamma = 0.9          # 割引率
experience_mode = "GAE"  # 収集した経験の評価方法(MC or GAE)
gae_lambda = 0.9     # GAEの割引率
warmup_size = 512    # 最低限貯めるキューサイズ
buffer_size = 1000   # キューの最大サイズ
batch_size = 32      # バッチサイズ
dense_units = 16     # Denseユニット数
activation = "relu"  # 活性化関数
init_layer_weight = "orthogonal"    # レイヤーの重さの初期化方法(""で省略)
optimizer = Adam()
optimizer_initial_lr = 0.02  # 初期学習率
optimizer_final_lr = 0.01    # 終了学習率
optimizer_lr_step = 200*10   # 終了学習率になるまでの更新回数
baseline_type = "ave"        # baselineの方法
enable_advantage_function = True  # 価値推定で状態価値を引くか
pi_clip_range = 0.2       # PPOにおけるISのクリップ範囲
adaptive_kl_targ = 0.01   # Adapitive KLペナルティ内の定数(0でAdapitive KLペナルティを無効)
enable_value_clip = True  # 価値関数もclipするか
value_loss_weight = 1.0   # 学習時の状態価値の反映率
entropy_weight = 0.1      # 学習時のエントロピーの反映率
global_gradient_clip_norm = 0.5  # 勾配のL2におけるclip値(0で無効)
enable_state_normalized = False  # 状態を正規化
state_clip = None   # 状態のclip(Noneで無効、(-10,10)で指定)
reward_clip = None  # 報酬のclip(Noneで無効、(-10,10)で指定)

モデルは以下です。

実装コードに関して

メインの実装は第2回と変わりません。
それに Clipped Surrogate Objective やその他のテクニックを追加しただけになります。

実行結果

139 step, reward: 139.0
136 step, reward: 136.0
120 step, reward: 120.0
200 step, reward: 200.0
131 step, reward: 131.0

試しに学習が失敗しない範囲で有効にしているパラメータがあるので、ハイパーパラメータを調整すればもっと安定して学習できます。

実装2、連続行動空間

アクションが連続値でモデルが分かれている実装です。

環境	アクション	モデル
Pendulum-v0	連続値	方策モデルと状態価値モデルは別

ハイパーパラメータは以下です。

train_interval = 100  # 学習間隔
stddev = 0.5          # 標準偏差
gamma = 0.9           # 割引率
experience_mode = "GAE"  # 収集した経験の評価方法
gae_lambda = 0.9     # GAEの割引率
warmup_size = 2000   # 最低限貯めるキューサイズ
buffer_size = 4000   # キューの最大サイズ
batch_size = 256     # バッチサイズ
dense_units = 32     # Denseユニット数
activation = "relu"  # 活性化関数
init_layer_weight = "orthogonal"  # レイヤーの重さの初期化方法(""で省略)
optimizer = Adam()
optimizer_initial_lr = 0.003  # 初期学習率
optimizer_final_lr = 0.003    # 終了学習率
optimizer_lr_step = 200       # 終了学習率になるまでの更新回数
baseline_type = "ave"         # baselineの方法
enable_advantage_function = True  # 価値推定で状態価値を引くか
pi_clip_range = 0.2        # PPOにおけるISのクリップ範囲
adaptive_kl_targ = 0       # Adapitive KLペナルティ内の定数(0でAdapitive KLペナルティを無効)
enable_value_clip = False  # 価値関数もclipするか
global_gradient_clip_norm = 0.5  # 勾配のL2におけるclip値(0で無効)
enable_state_normalized = False  # 状態を正規化
state_clip = None        # 状態のclip(Noneで無効、(-10,10)で指定)
reward_clip = [-10, 10]  # 報酬のclip(Noneで無効、(-10,10)で指定)
enable_action_normalization = True  # アクションの正規化
enable_sgp = False                  # Squashed Gaussian Policy

モデルは以下です。

標準偏差ですが、それも含めて学習しようとすると安定してできなかったので固定値にしています。

実行結果

200 step, reward: -1512.8988019581018
200 step, reward: -711.3905029077422
200 step, reward: -390.840581351853
200 step, reward: -703.1192896166007
200 step, reward: -522.5698912223603

やはりPendulum-v0は学習が難しいですね。
悪くない学習はしているんですけど…。

あとがき

方策モデルの学習は何かと不安定だったんですが、PPO は圧倒的な安定さですね。
連続行動空間の学習でもかなり安定して学習できているのはすごいです。