【強化学習】自作ライブラリでDQN

# TL;DR
自作のReplay Bufferライブラリ[cpprb](https://gitlab.com/ymd_h/cpprb) を使って、DQNを実装してみた。

高い自由度と効率性を兼ね備えている（つもりな）のでおすすめ。

# 1. 背景と経緯

[Open AI/Baselines](https://github.com/openai/baselines) や [Ray/RLlib](https://ray.readthedocs.io/en/latest/rllib.html) のような、強化学習一式の環境を利用すると、ちょっとしたコードで様々なアルゴリズムを試してみることができる。

例えば、Open AI/Baselinesで、AtariのPongをDQNで学習させるには以下のコマンドを実行するだけで良いと公式READMEに記載されている。

``` shell
python -m baselines.run --alg=deepq --env=PongNoFrameskip-v4 --num_timesteps=1e6
```


一方、既存のアルゴリズムをテストするのは簡単だけれども、研究者やライブラリ開発者が新しい独自アルゴリズムを作ろうとした際に、どこから手をつけていいのか大きすぎて大変だと思う。

強化学習の研究をしている友人もTensorFlowなどの深層学習のライブラリは利用するものも、他の部分は独自に実装してた（ようであった）。

そんな2018年の暮れ頃、その友人から「Cythonって興味ある？ Pythonで実装しているReplay Bufferが、（状況によっては）深層学習の学習部分なみに遅くて、Cythonでスピードアップを図りたいんだけど」（記憶）と誘われて実装始めたのが cpprb である。

（その友人は、cpprbとTensorFlow 2.x を利用して、[tf2rl](https://github.com/keiohta/tf2rl)という強化学習ライブラリを公開していて、こちらも超おすすめ！）

# 2. 特徴

そんな背景もあって実装はじめた cpprb なので、高い自由度と効率性に主眼をおいて開発している。

## 2.1 高い自由度
バッファに保存する変数名・サイズ・型を `dict` 形式で指定することで自由に決めることができる。

例えば、極端な例だが、 `next_next_obs`, `previous_act`, `secondary_reward` なんてものも保存することができる。

``` python
import numpy as np
from cpprb import ReplayBuffer

buffer_size = 1024

# shape と dtype を変数ごとに指定できる。デフォルトは、{"shape":1,"dtype": np.float32}
rb = ReplayBuffer(buffer_size,
                  {"obs": {"shape": (3,3)},
                   "act": {"shape": 3, "dtype": np.int},
                   "rew": {},
                   "done": {},
                   "next_obs": {"shape": (3,3)},
                   "next_next_obs": {"shape": (3,3)},
                   "previous_act": {"shape": 3, "dtype": np.int},
                   "secondary_reward": {}})

# Key-Value 形式で指定する （初期化時に指定した変数が不足していると `KeyError`）
rb.add(obs=np.zeros(shape=(3,3)),
       act=np.ones(3,dtype=np.int),
       rew=0.5,
       done=0,
       next_obs=np.zeros(shape=(3,3)),
       next_next_obs=np.ones(shape=(3,3)),
       previous_act=np.ones(3,dtype=np.int),
       secondary_reward=0.3)
```

## 2.2 効率性
Prioritized Experience Replayの遅さの原因であるSegment TreeをCython経由で、C++実装しているため、かなり早い。

![PrioritizedReplayBuffer_sample.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/592611/b51d76c2-8f34-c302-1bc1-95272304585d.png)


ベンチマークを見る限り速度で圧勝している。(2020年4月現在。最新版は[プロジェクトサイト](https://ymd_h.gitlab.io/cpprb/comparison/benchmark/)へ)



注意: 強化学習全体では、Segment Treeの速度だけではなくて、うまく探索を並列化させるなどの対策が重要





# 3. インストール
(情報が古くなってるかもしれないので、最新の[インストール](https://ymd_h.gitlab.io/cpprb/installation/)方法も参照)

## 3.1 バイナリインストール
PyPIに公開しているので、pip(や類似のツール)を用いてインストールすることができる。

Windows/Linux向けには、wheel形式のバイナリを配布しているので、多くの場合は何も考えずに以下のコマンドでインストールすることができる。
(注: `venv` やdocker等仮想環境の利用を推奨。)

```
pip install cpprb
```

注: macOSは標準の開発ツールチェインの一部になっている [clang](https://clang.llvm.org/) が、 C++17 の機能の `std::shared_ptr` の配列型への特殊化を未実装のためコンパイルできずバイナリを配布できていない。

## 3.2 ソースからインストール

ソースコードから各自ビルドする必要がある。ビルドには以下が必要である。

*  GCC >= 7.2(?)

環境変数 `CC` と `CXX` に `g++` を指定してビルドを実行する必要がある。



# 3. DQN実装

(2020年7月26日: cpprbのバージョンアップ及び、[gnwrapper](https://gitlab.com/ymd_h/gym-notebook-wrapper)の利用によりコードを修正)

Google Colab 上で動作するDQNを書いてみた

まずは、必要なライブラリをインストール

``` python
!apt update > /dev/null 2>&1
!apt install -y xvfb > /dev/null
!pip install gym gym-notebook-wrapper cpprb["all"] tensorflow > /dev/null
%load_ext tensorboard
```

``` python
# Standard Library
import os
import datetime

# 3rd Party
import gym
import numpy as np

from scipy.special import softmax

import tensorflow as tf
from tensorflow.keras.models import Sequential,clone_model
from tensorflow.keras.layers import InputLayer,Dense
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.callbacks import EarlyStopping,TensorBoard
from tensorflow.summary import create_file_writer

from tqdm.notebook import tqdm

from cpprb import create_buffer
import gnwrapper
```

``` python
%load_ext tensorboard
%tensorboard --logdir logs
```

``` python
gamma = 0.99
batch_size = 1024
buffer_size = int(1e6)

N_iteration = 12000
per_train = 10

prioritized = True

egreedy = False

loss = "huber_loss"
# loss = "mean_squared_error"


env_name = 'CartPole-v1' # reward = 1
# env_name = 'MountainCar-v0' # reward = -1
# env_name = 'Acrobot-v1' # reward = -1


JST = datetime.timezone(datetime.timedelta(hours=+9), 'JST')
dir_name = datetime.datetime.now(JST).strftime(f"%Y%m%d-%H%M%S-{env_name}")

logdir = os.path.join("logs", dir_name)
writer = create_file_writer(logdir + "/metrics")
writer.set_as_default()

env = gym.make(env_name)
env = gnwrapper.Monitor(env,logdir + "/video/")

observation = env.reset()


model = Sequential([InputLayer(input_shape=(observation.shape)), # 4 for CartPole
                    Dense(64,activation='relu'),
                    Dense(64,activation='relu'),
                    Dense(env.action_space.n)]) # 2 for CartPole

target_model = clone_model(model)


optimizer = Adam()

model.compile(loss =  loss,
              optimizer = optimizer,
              metrics=['accuracy'])


Nstep = {"size":4, "gamma": gamma, "rew":"rew", "next": "next_obs"}

rb = create_buffer(buffer_size,
                   {"obs":{"shape": observation.shape},
                    "act":{"shape": 1,"dtype": np.ubyte},
                    "rew": {},
                    "next_obs": {"shape": observation.shape},
                    "done": {}},
                    prioritized = prioritized,
                   Nstep=Nstep)

action_index = np.arange(env.action_space.n).reshape(1,-1)

n_episode = 0
observation = env.reset()
sum_reward = 0.0

for total_step in tqdm(range(N_iteration)):
    actions = softmax(np.ravel(model.predict(observation.reshape(1,-1),batch_size=1)))
    actions = actions / actions.sum()

    if egreedy:
        if np.random.rand() < 0.1:
            action = env.action_space.sample()
        else:
            action = np.argmax(actions)
    else:    
        action = np.random.choice(actions.shape[0],p=actions)

    next_observation, reward, done, info = env.step(action)
    sum_reward += reward
    rb.add(obs=observation,
           act=action,
           rew=reward,
           next_obs=next_observation,
           done=done)
    observation = next_observation

    if rb.get_stored_size() == 0:
        continue

    sample = rb.sample(batch_size)
    Q_pred = model.predict(sample["obs"])
    Q_true = sample['rew'] + gamma*(sample["discounts"] if Nstep else 1)*(1.0 - sample["done"])*target_model.predict(sample['next_obs']).max(axis=1,keepdims=True)
    target = tf.where(tf.one_hot(tf.cast(tf.reshape(sample["act"],[-1]),dtype=tf.int32),env.action_space.n,True,False),
                      tf.broadcast_to(Q_true,[batch_size,env.action_space.n]),
                      Q_pred)
    
    if prioritized:
        TD = np.square(target - Q_pred).sum(axis=1)
        rb.update_priorities(sample["indexes"],TD)

    weights = sample["weights"] if prioritized else None
    
    model.fit(x=sample['obs'],
              y=target,
              batch_size=batch_size,
              sample_weight=weights,
              verbose = 0)

    if total_step % per_train == 0:
        target_model.set_weights(model.get_weights())

    tf.summary.scalar("step reward",data=reward,step=total_step)
    tf.summary.scalar("episode vs total step",data=n_episode,step=total_step)
    if bool(done):
        tf.summary.scalar("episode reward vs episode",data=sum_reward,step=n_episode)
        tf.summary.scalar("episode reward vs total step",data=sum_reward,step=total_step)
        observation = env.reset()
        rb.on_episode_end()
        n_episode += 1
        sum_reward = 0.0

env.display()
```


# 4. 結果

rewardの結果（x軸:エピソード・y軸:エピソード報酬）。

<img width="1296" alt="SS 2020-07-26 21.13.36.png" src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/592611/0189527f-2fe2-5b91-0eff-4739e215c439.png">



# 6. まとめ

強化学習向けReplay Bufferを提供する自作ライブラリ [cpprb](https://gitlab.com/ymd_h/cpprb) を利用して、DQNを実装した。

cpprb は高い自由度と効率性を重視して開発している。

興味を持ってくれた人は、試してみて issueやマージリクエストをぜひ。（英語が好ましいけど、日本語でもOK）


# 参考リンク

*  cpprb
    *  [GitLabレポジトリ](https://gitlab.com/ymd_h/cpprb)
    *  [GitHubミラー](https://github.com/ymd-h/cpprb)
    *  [PyPIレポジトリ](https://pypi.org/project/cpprb/)
    *  [プロジェクトサイト](https://ymd_h.gitlab.io/cpprb/)
*  [tf2rl](https://github.com/keiohta/tf2rl)