【強化学習】R2D2（+Retrace）を解説・実装

この記事は自作している強化学習フレームワーク SimpleDistributedRL の解説記事です。

前：Rainbow、分散強化学習
次：Agent57

R2D2については昔記事を書いていますが、知識も更新されているので改めて書いています。

部分観測マルコフ決定過程

詳細に知りたい方はこちらを見てください。

R2D2で必要な知識としては、エピソード最初からすべての履歴を状態に含めた場合、部分観測マルコフ決定過程(POMDP)はマルコフ決定過程(MDP)として扱えるという事です。
（MDPなら強化学習の手法が問題なく使えます）

なので、過去の実行履歴（いわゆる時系列データ）を含めて学習させることを考えます。

RNN(LSTM)

ニューラルネットワークで時系列データを学習させるモデルとしてはRNN(Recurrent Neural Network)があります。
R2D2ではRNNでデファクトスタンダードになっているLSTMを使っています。
RNN/LSTMの詳細については割愛し、R2D2で必要な要素のみ解説します。

RNNの基本的な構造は以下です。

一般的なNNの入力と出力に加え、隠れ状態 h があるのがRNNの特徴です。
(LSTMでは隠れ状態は出力セルと記憶セル2つから成っていますが、本記事ではこれらをまとめて隠れ状態(hidden state)と言っています)

隠れ状態は時系列情報を記憶する役割を持っています。

参考
・LSTMネットワークの概要
・【超初心者向け】これなら分かる！はじめてのLSTM

R2D2

論文
・Recurrent Experience Replay in Distributed Reinforcement Learning
・Deep Recurrent Q-Learning for Partially Observable MDPs

LSTMの導入

深層強化学習にRNNを導入するには以下の課題があります。

1. Experience Replayと相性が悪い
   学習に使う経験ですが、NNの汎化性能を保つためにランダムにサンプリングされます。
   なので時系列情報がありません（エピソードの途中の経験も入ってくる）
2. エピソード全体を学習しようとすると計算が膨大になる
   1を解決するには1エピソードまるまる学習すれば解決します。（エピソード単位でランダムサンプリングする）
   しかし、1エピソードが短ければ計算可能ですが、長いと現実的な計算量ではなくなります。
   （OpenAI GymのAtariは最大27000step）

上記を解決するためにエピソード途中の隠れ状態も覚えておき、それを元に経験を復元する方法を考えます。(Stored state)

隠れ状態の陳腐化とBurn-in

隠れ状態を保存した時とメモリから取り出した時では、学習が進んでいるのでモデルが違っています。
この場合、当然ですが今の隠れ状態と昔の隠れ状態では差異があります。（陳腐化している）

陳腐化を低減するためにR2D2で提案された手法が Burn-in という手法になります。
（Burn-inはメンダコさんのブログであった"ならし運転"という訳がしっくりきました）

少し前のstepから隠れ状態を流して、今の隠れ状態に近くしようというものが Burn-in です。

0stepから Burn-in すれば1エピソード全体を学習する場合と同じになります。

Burn-inの効果

図は論文より引用しています。
上がパックマン、下がDMLabの"EMSTM WATERMAZE"というゲームらしいです（検索しても出てきませんでした）
パックマンはR2D2がSOTAを達成したゲームで、"EMSTM WATERMAZE"はたくさんの履歴が必要なゲーム（多分POMDPなゲーム？）を比較してみたとの事です。

縦軸がスコアで横軸が保持する履歴数です。
0が履歴なし（時系列情報がない）、∞はエピソード開始からすべての履歴を使う場合です。
この状態で隠れ状態を0から始めるか（赤線：Zero State）、復元するか（青線：Stored State）で比較しています。

横軸が0に近づくとスコアは下がっています。（学習出来ていない）
これは全ての組み合わせで起こっていることから Burn-in は効果があるとの事です。

また、青線は∞の状態では高いですが、履歴数を減らすと下がります。
（赤線はあまり変わらない）
これは赤線(0スタート)の方が隠れ状態を有効に学習していない（学習能力に限界がある）証とのことです。
ただ、これはパックマンのように必ずパフォーマンスを向上するわけではありません。
しかし、履歴を効果的に使う必要があるタスク("EMSTM WATERMAZE")ではパフォーマンスが向上するとの事です。

メモリと分散処理の関係

Experience Replayではメモリに蓄積された経験を使いまわしていました。
ここで問題になるのが古い経験です。
そもそも経験が新しければBurn-inを使う必要がありません。

分散学習とそうじゃない場合でこの経験の鮮度に違いがあります。
これは分散学習はActor数に比例して大量の経験が送られてくるためです。
（1つの経験が、DQNでは平均8回ほど学習されるのに対し、Ape-Xでは1回未満になるそうです）

メモリサイズを小さくすることでこの問題に対処はできますが、経験の多様性は失われてしまいます。

その他のテクニック

rescaling function

DQNでは報酬を(-1,0,1)にクリップしていました。
このクリップを止め、以下関数で報酬を丸める手法が rescaling function です。

\begin{align}
h(z) = sign(z)(\sqrt{|z|+1}-1) + \epsilon z \\
h^{-1}(z) = sign(z) \bigg( \Big(\frac{\sqrt{1 + 4\epsilon(|z|+1+\epsilon)}-1)}{2 \epsilon}\Big)^2 -1 \bigg)\\
\end{align}

$sign$ は符号関数、$\epsilon$ が定数です。
適用箇所はQの予測値の計算箇所となります。

\begin{align}
Q'_{\mu}(s, a; \theta_i^-) &= r + \gamma \max_{a'} Q_\pi(s', a'; \theta_i^-)\\
&↓\\
Q'_{\mu}(s, a; \theta_i^-) &= h \Big( r + \gamma h^{-1} \Big( \max_{a'} Q_\pi(s', a'; \theta_i^-) \Big) \Big)
\end{align}

関数の性質については過去に可視化しているのでそちらをご覧ください。
・rescaling関数の可視化

ε-greedy の固定

Rainbowでは、探索率εは最初は高く徐々に低くするアニーリング方式でした。
しかし、Ape-X で提案された分散学習では各Actor毎に探索率を設定できるため、Actor毎に固定し、探索の多様性を表現していました。
R2D2もこれを踏襲しています。

各Actorの探索率は以下で計算されます。

$$
\epsilon_i = \epsilon^{1 + \frac{i}{N-1} \alpha}
$$

$N$ はActor数、$\epsilon$ と $\alpha$ がハイパーパラメータです。

Retrace損失の適用

次の手法であるNGUにて、ベースラインとしてRetrace損失を適用したR2D2を使用しています。
Retrace損失を加えると multisteps と同じように未来の報酬を加味して計算できるようになります。（予測値の正確性が少し上がります）

Retraceの詳細に関しては以前書いた記事を参照してください。

Rainbowからの引継ぎ

Rainbowから引き継いでいる要素は以下です。

• Double DQN
• Prioritized Experience Replay
• Dueling networks

除いた要素の理由は論文にはないので予想です。
Noisy-networkは、探索は複数のActorが固有のεで探索するので不要になったかと思われます。

CategoricalDQNはApe-Xの議論ですが、分散学習の効果をメインに議論したいために除外したとの事です。
（状態価値を確率分布で表現しようというアルゴリズムの性質上、どうしても複雑になります）
（実装すればパフォーマンスが向上するかもしれないと言っています）

Q1: on using all Rainbow components and on using multiple learners.

These are both interesting directions which we agree may help to boost >performance even further. For this paper, we felt that adding extra components would distract from the finding that it is possible to improve results significantly by scaling up, even with a relatively simple algorithm.
（OpenReviewより）

Multi-step learningはApe-Xでは採用しています。
R2D2ではLSTMで学習する時系列情報がMulti-step learningと被っているので除外しているのかと思われます。

実装

フレームワーク上の実装はgithubを見てください。

ステートフルLSTMとステートレスLSTM

かなり昔(TF1.0/Keras)に記事を書いていますが、今でも大きな違いはないようです。

LSTMには時系列情報を保持するために、入力以外に隠れ状態という情報もinputとして使います。

ステートレスLSTM(デフォルト)では隠れ状態は毎回0にリセットされます。
しかし、今回の場合みたいに可変な時系列で隠れ状態を取り扱いたい場合はステートフルLSTMを使って制御します。
ただ、ステートフルLSTMを使うとミニバッチのサイズを固定しないといけない制約がつくので少し扱いが複雑になります。

Model

画像を入力する場合を想定したモデルは以下となります。

timesteps は burn-in の数とその後の学習stepが可変なので 1 に固定してどの数にも対応できるようにしています。
また論文には、状態だけではなく前のアクションと報酬も入力とするとありますが、実装が複雑になるのでフレームワーク上は省略しています。
（多分NGU/Agent57のUVFAの事？）

class _QNetwork(keras.Model):
    def __init__(self):
        super().__init__()

        # 入力は (batch_size, timesteps, width, height) の形
        # stateless LSTMを使うため、batch_size まで指定
        input_shape = (64, 1, 84, 84)
        in_state = c = kl.Input(batch_input_shape=input_shape)

        # 画像形式に変更
        # (batch_size, timesteps, w, h) -> (batch_size, timesteps, w, h, 1)
        c = kl.Reshape(input_shape + (-1,))(c)

        # DQNの画像処理レイヤー、Conv2Dは時系列を意識する必要なし
        c = kl.Conv2D(32, (8, 8), strides=(4, 4), padding="same", activation="relu")(c)
        c = kl.Conv2D(64, (4, 4), strides=(2, 2), padding="same", activation="relu")(c)
        c = kl.Conv2D(64, (3, 3), strides=(1, 1), padding="same", activation="relu")(c)
        
        # 時系列情報は保持したいので TimeDistributed を入れる(いれないと timesteps までflatになる)
        # (batch_size, timesteps, w, h, ch) -> (batch_size, timesteps, flat_num)
        c = kl.TimeDistributed(kl.Flatten())(c)

        # lstm
        c = kl.LSTM(512, stateful=True, name="lstm")(c)

        # 出力層、config.nb_actionsはアクション数
        c = kl.Dense(config.nb_actions)(c)
        self.model = keras.Model(in_state, c)
        self.lstm_layer = self.model.get_layer("lstm")

    # forward は隠れ状態も入力とする
    def call(self, state, hidden_states):
        self.lstm_layer.reset_states(hidden_states)
        return self.model(state), self.get_hidden_state()

    # 現状の隠れ状態を取得する関数
    def get_hidden_state(self):
        return [self.lstm_layer.states[0].numpy(), self.lstm_layer.states[1].numpy()]

    # 隠れ状態を初期化する関数
    def init_hidden_state(self):
        self.lstm_layer.reset_states()
        return self.get_hidden_state()

Worker

隠れ状態を用いて1エピソードを進めます。
LSTMが関係しているとこのみ抜粋しています。

class Worker(DiscreteActionWorker):

    # 1エピソードの最初
    def call_on_reset(self, state: np.ndarray, invalid_actions: List[int]) -> None:
        
        # 隠れ状態を初期化し保存する
        self.hidden_state = self.parameter.q_online.init_hidden_state()

        # 履歴用
        self.recent_hidden_states = [
            [self.hidden_state[0][0], self.hidden_state[1][0]]
            for _ in range(self.config.burnin + self.config.multisteps + 1)
        ]

    # アクションを返す
    def call_policy(self, state: np.ndarray, invalid_actions: List[int]) -> int:
        # NN入力用の状態を作成、stateless LSTMはbatch_sizeが固定のため増やす
        state = np.asarray([[state]] * self.config.batch_size)

        # hidden_stateを使いQ値を取得(新しいhidden_stateを保存)
        q, self.hidden_state = self.parameter.q_online(state, self.hidden_state)
        q = q[0].numpy()

        （Q値を使ってアクションを出す箇所は同じです）

    # 毎step呼ばれる
    def call_on_step(
        self,
        next_state: np.ndarray,
        reward: float,
        done: bool,
        next_invalid_actions: List[int],
    ) -> Dict:

        # 履歴用に、hidden_stateを保存する
        self.recent_hidden_states.pop(0)
        self.recent_hidden_states.append(
            [
                self.hidden_state[0][0],
                self.hidden_state[1][0],
            ]
        )

Trainer

Retrace損失を考えない場合のイメージは以下です。

Retrace損失を加える場合は以下です。

Tensorflow の tf.GradientTape はネストしても問題なく動きました。
再帰関数で実現しています。

class Trainer(RLTrainer):
    def train(self):
        
        # メモリからbatchsを取得
        indices, batchs, weights = self.remote_memory.sample(self.train_count, self.config.batch_size)
        td_errors, loss = self._train_on_batchs(batchs, weights)

        # (batch, dict[x], multisteps) -> (multisteps, batch, x)
        states_list = []
        for i in range(self.config.multisteps + 1):
            states_list.append(np.asarray([[b["states"][i]] for b in batchs]))
        actions_list = []
        mu_probs_list = []
        rewards_list = []
        dones_list = []
        for i in range(self.config.multisteps):
            actions_list.append([b["actions"][i] for b in batchs])
            rewards_list.append([b["rewards"][i] for b in batchs])
            mu_probs_list.append([b["probs"][i] for b in batchs])
            dones_list.append([b["dones"][i] for b in batchs])

        # hidden_states
        states_h = []
        states_c = []
        for b in batchs:
            states_h.append(b["hidden_states"][0])
            states_c.append(b["hidden_states"][1])
        hidden_states = [np.asarray(states_h), np.asarray(states_c)]
        hidden_states_t = [np.asarray(states_h), np.asarray(states_c)]

        # burn-in(onlineとtarget両方実施)
        for i in range(self.config.burnin):
            burnin_state = np.asarray([[b["burnin_states"][i]] for b in batchs])
            _, hidden_states = self.parameter.q_online(burnin_state, hidden_states)
            _, hidden_states_t = self.parameter.q_target(burnin_state, hidden_states_t)

        # 再帰構造で学習
        _, _, td_error, _, loss = self._train_steps(
            states_list,
            actions_list,
            mu_probs_list,
            rewards_list,
            dones_list,
            weights,
            hidden_states,
            hidden_states_t,
            0,
        )

        （メモリの更新）
        （targetの同期）
        return {}

    # Q値(LSTM hidden states)の予測はforward、td_error,retraceはbackwardで予測する必要あり
    def _train_steps(
        self,
        states_list,
        actions_list,
        mu_probs_list,
        rewards_list,
        dones_list,
        weights,
        hidden_states,
        hidden_states_t,
        idx,
    ):

        # 最後はQ値のみを返す
        if idx == self.config.multisteps:
            n_states = states_list[idx]
            n_q, _ = self.parameter.q_online(n_states, hidden_states)
            n_q_target, _ = self.parameter.q_target(n_states, hidden_states)
            n_q = tf.stop_gradient(n_q).numpy()
            n_q_target = tf.stop_gradient(n_q_target).numpy()
            # q, target_q, td_error, retrace, loss
            return n_q, n_q_target, 0.0, 1.0, 0.0

        states = states_list[idx]
        actions = actions_list[idx]

        # return用にq_targetを計算
        q_target, n_hidden_states_t = self.parameter.q_target(states, hidden_states_t)
        q_target = tf.stop_gradient(q_target).numpy()

        # --- 勾配 + targetQを計算
        with tf.GradientTape() as tape:
            q, n_hidden_states = self.parameter.q_online(states, hidden_states)

            # 次のQ値を取得
            n_q, n_q_target, n_td_error, retrace, _ = self._train_steps(
                states_list,
                actions_list,
                mu_probs_list,
                rewards_list,
                dones_list,
                weights,
                n_hidden_states,
                idx + 1,
            )

            # targetQを計算
            target_q = np.zeros(self.config.batch_size)
            for i in range(self.config.batch_size):
                reward = rewards_list[idx][i]
                done = dones_list[idx][i]

                if done:
                    gain = reward
                else:
                    # DoubleDQN: indexはonlineQから選び、値はtargetQを選ぶ
                    n_act_idx = np.argmax(n_q[i])
                    maxq = n_q_target[i][n_act_idx]
                    maxq = inverse_rescaling(maxq)  # rescaling
                    gain = reward + self.config.gamma * maxq
                gain = rescaling(gain)  # rescaling
                target_q[i] = gain

            # retrace
            if self.config.enable_retrace:
                _retrace = np.zeros(self.config.batch_size)
                for i in range(self.config.batch_size):
                    action = actions_list[idx][i]
                    mu_prob = mu_probs_list[idx][i]
                    pi_probs = calc_epsilon_greedy_probs(
                        n_q[i],
                        next_invalid_actions_batchs[i],
                        0.0,
                        self.config.nb_actions,
                    )
                    pi_prob = pi_probs[action]
                    _retrace[i] = self.config.retrace_h * np.minimum(1, pi_prob / mu_prob)

                retrace *= _retrace
                target_q += self.config.gamma * retrace * n_td_error

            # 現在選んだアクションのQ値
            action_onehot = tf.one_hot(actions, self.config.nb_actions)
            q_onehot = tf.reduce_sum(q * action_onehot, axis=1)

            loss = self.loss(target_q * weights, q_onehot * weights)

        grads = tape.gradient(loss, self.parameter.q_online.trainable_variables)
        self.optimizer.apply_gradients(zip(grads, self.parameter.q_online.trainable_variables))
        # --- 勾配計算ここまで
        q = tf.stop_gradient(q).numpy()

        n_td_error = (target_q - q_onehot).numpy() + self.config.gamma * retrace * n_td_error
        return q, q_target, n_td_error, retrace, loss.numpy()

Rescale function

コード化すると以下です。

def rescaling(x, eps=0.001):
    return np.sign(x) * (np.sqrt(np.abs(x) + 1.0) - 1.0) + eps * x

def inverse_rescaling(x, eps=0.001):
    n = np.sqrt(1.0 + 4.0 * eps * (np.abs(x) + 1.0 + eps)) - 1.0
    n = n / (2.0 * eps)
    return np.sign(x) * ((n**2) - 1.0)

探索率

コード化すると以下です。
1は定義されていないので適当に置いています。

def create_epsilon_list(policy_num: int, epsilon=0.4, alpha=8.0):
    assert policy_num > 0
    if policy_num == 1:
        return [epsilon / 4]

    epsilon_list = []
    for i in range(policy_num):
        e = epsilon ** (1 + (i / (policy_num - 1)) * alpha)
        epsilon_list.append(e)
    return epsilon_list

人数による値の違いは以下です。

policy_num	list
2	[0.4, 0.0006553600000000003]
3	[0.4, 0.016190861620062107, 0.0006553600000000003]
4	[0.4, 0.04715560318259697, 0.005559127278786374, 0.0006553600000000003]

実験

Open AI Gymで提供されているPendulum-v1を元に実験しています。
（1つの環境しか見ていないので参考程度に見てください）

Retrace損失

Retrace損失がある場合とない場合の違いです。

横軸がエピソード数で縦軸が報酬です。
少し安定しているような・・・？

ここで使ったコードは github を見てください。

学習の速さ比較(DQN,Rainbow,R2D2)

Google Colaboratory 上で比較しました（CPU）

横軸が時間(秒)で縦軸が報酬です。
LSTMは計算に時間がかかるのである程度スペックがないと学習時間で負けてしまいますね。

学習終了後のエピソード数と学習回数の比較です。
R2D2は学習回数が少なすぎますね…。

使用コード

※古いsrlのバージョンで書いたコードなので動かない可能性があります

import matplotlib.pyplot as plt
import srl
from srl import runner

# --- env & algorithm load
import gym  # isort: skip # noqa F401
from srl.algorithms import dqn, rainbow, r2d2  # isort: skip


env_config = srl.EnvConfig("Pendulum-v1")
rl_configs = []

# DQN
rl_configs.append(("DQN", dqn.Config(hidden_block_kwargs=dict(hidden_layer_sizes=(64, 64)))))

# Rainbow
rl_configs.append(
    (
        "Rainbow",
        rainbow.Config(
            hidden_layer_sizes=(64, 64),
            memory_name="ReplayMemory",
            multisteps=5,
        ),
    )
)

# R2D2(軽量)
rl_configs.append(
    (
        "R2D2",
        r2d2.Config(
            lstm_units=64,
            hidden_layer_sizes=(64,),
            enable_dueling_network=False,
            memory_name="ReplayMemory",
            enable_rescale=False,
            burnin=5,
            sequence_length=5,
        ),
    )
)

# train
results = []
for name, rl_config in rl_configs:
    print(name)
    config = runner.Config(env_config, rl_config)
    _, _, history = runner.train(config, timeout=60 * 40, enable_evaluation=False)
    results.append((name, history))

# plot
plt.figure(figsize=(8, 4))
plt.xlabel("time")
plt.ylabel("reward")
for name, h in results:
    df = h.get_df()
    plt.plot(df["time"], df["episode_reward0"].rolling(10).mean(), label=name)
plt.grid()
plt.legend()
plt.tight_layout()
plt.show()

fig = plt.figure(figsize=(8, 4))
ax1 = fig.add_subplot(1, 2, 1)
ax1.set_ylabel("episode")
names = [n for n, h in results]
episodes = [h.get_df()["episode_count"][-1] for n, h in results]
ax1.bar(names, episodes)
ax2 = fig.add_subplot(1, 2, 2)
ax2.set_ylabel("train")
names = [n for n, h in results]
trains = [h.get_df()["train_count"][-1] for n, h in results]
ax2.bar(names, trains)
plt.tight_layout()
plt.show()

分散処理の違い

経験の陳腐化の影響を見てみました。

上記はburninをなしにして、分散なし(sequence)と分散あり(mp)の報酬の遷移です。
あまり変わらないですね。

上記はメモリにたまった経験の推移です。
分散あり(mp)ではかなり早い段階でメモリがいっぱいになっています。

最後に学習終了時のエピソード回数と学習回数です。
分散あり(mp)のほうは少しエピソード回数に偏ってますね。

使用コード

※古いsrlのバージョンで書いたコードなので動かない可能性があります

import matplotlib.pyplot as plt
import srl
from srl import runner

# --- env & algorithm load
import gym  # isort: skip # noqa F401
from srl.algorithms import r2d2  # isort: skip

def main():

    env_config = srl.EnvConfig("Pendulum-v1")

    rl_config = r2d2.Config(
        lstm_units=64,
        hidden_layer_sizes=(64,),
        enable_dueling_network=False,
        memory_name="ReplayMemory",
        capacity=10000,
        target_model_update_interval=100,
        enable_rescale=False,
        burnin=0,
        sequence_length=10,
    )
    config = runner.Config(env_config, rl_config)

    # train
    timeout = 60 * 20
    _, _, history1 = runner.train(config, timeout=timeout, enable_evaluation=False)
    _, _, history2 = runner.mp_train(config, timeout=timeout, enable_evaluation=False)

    # plot
    plt.figure(figsize=(8, 4))
    plt.xlabel("time")
    plt.ylabel("reward")
    df = history1.get_df()
    plt.plot(df["time"], df["episode_reward0"], label="sequence")
    df = history2.get_df()
    plt.plot(df["time"], df["episode_reward0"].rolling(30).mean(), label="mp")
    plt.grid()
    plt.legend()
    plt.tight_layout()
    plt.show()

    plt.figure(figsize=(8, 4))
    plt.xlabel("time")
    plt.ylabel("memory")
    df = history1.get_df()
    plt.plot(df["time"], df["remote_memory"], label="sequence")
    df = history2.get_df()
    plt.plot(df["time"], df["remote_memory"], label="mp")
    plt.grid()
    plt.legend()
    plt.tight_layout()
    plt.show()

    fig = plt.figure(figsize=(8, 4))
    ax1 = fig.add_subplot(1, 2, 1)
    ax1.set_ylabel("episode")
    episodes = [
        history1.get_df()["episode_count"][-1],
        history2.get_df()["episode_count"][-1],
    ]
    ax1.bar(["sequence", "mp"], episodes)
    ax2 = fig.add_subplot(1, 2, 2)
    ax2.set_ylabel("train")
    episodes = [
        history1.get_df()["train_count"][-1],
        history2.get_df()["train_count"][-1],
    ]
    ax2.bar(["sequence", "mp"], episodes)
    plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    main()

追記 2022/8/14：ステートレスLSTMでも可変な時系列情報を取り扱えるようでした

調査結果は以下の記事です。
LSTM(RNN)で可変長な時系列と隠れ状態について調べてみた

githubのコードは更新してあります。
更新後で速度を比較してみた結果は以下です。

使用コード

※古いsrlのバージョンで書いたコードなので動かない可能性があります

import matplotlib.pyplot as plt
import srl
from srl.runner import sequence

# --- env & algorithm load
import gym  # isort: skip # noqa F401
from srl.algorithms import dqn, rainbow, r2d2, r2d2_stateful  # isort: skip


env_config = srl.EnvConfig("Pendulum-v1")
rl_configs = []

# DQN
rl_configs.append(("DQN", dqn.Config(hidden_block_kwargs=dict(hidden_layer_sizes=(64, 64)))))

# Rainbow
rl_configs.append(
    (
        "Rainbow",
        rainbow.Config(
            hidden_layer_sizes=(64, 64),
            memory_name="ReplayMemory",
            multisteps=5,
        ),
    )
)

# R2D2(stateful)
rl_configs.append(
    (
        "R2D2 stateful",
        r2d2_stateful.Config(
            lstm_units=64,
            hidden_layer_sizes=(64,),
            enable_dueling_network=False,
            memory_name="ReplayMemory",
            enable_rescale=False,
            burnin=5,
            sequence_length=5,
        ),
    )
)

# R2D2(修正後)
rl_configs.append(
    (
        "R2D2",
        r2d2.Config(
            lstm_units=64,
            hidden_layer_sizes=(64,),
            enable_dueling_network=False,
            memory_name="ReplayMemory",
            enable_rescale=False,
            burnin=5,
            sequence_length=5,
        ),
    )
)

# train
results = []
for name, rl_config in rl_configs:
    print(name)
    config = sequence.Config(env_config, rl_config)
    _, _, history = sequence.train(config, timeout=60 * 30, enable_evaluation=False)
    results.append((name, history))

# plot
plt.figure(figsize=(8, 4))
plt.xlabel("time")
plt.ylabel("reward")
for name, h in results:
    df = h.get_df()
    plt.plot(df["time"], df["episode_reward0"].rolling(10).mean(), label=name)
plt.grid()
plt.legend()
plt.tight_layout()
plt.show()

fig = plt.figure(figsize=(8, 4))
ax1 = fig.add_subplot(1, 2, 1)
ax1.set_ylabel("episode")
names = [n for n, h in results]
episodes = [h.get_df()["episode_count"][-1] for n, h in results]
ax1.bar(names, episodes)
ax2 = fig.add_subplot(1, 2, 2)
ax2.set_ylabel("train")
names = [n for n, h in results]
trains = [h.get_df()["train_count"][-1] for n, h in results]
ax2.bar(names, trains)
plt.tight_layout()
plt.show()

倍以上早くなりました。

最後に

ここら辺からマシンスペックがものをいうようになる印象です。
LSTMもGPUパワーがないとあまり効果は得られない気がします。