強化学習フレームワークRLlibを使ってDQNの学習を行う

[RLlib](https://ray.readthedocs.io/en/latest/rllib.html)は強化学習の実験を行うためのライブラリ。Q-Learning, Policy Gradientのいくつかのアルゴリズムが実装されている。今回はDeep Q-Network(DQN)、特にApe-Xを使った学習を行う方法を紹介する。

RLlibは論文で主張している通り様々なアルゴリズムをサポートしているが、その分設定できる項目が多い。ソースコードを読めばコメントや実装から各々の意味を読み取れるが、把握するのは大変なのでメモとして残しておく。

# Ape-Xについて

設定の前提となるApe-Xの概要を説明する。詳しくは論文: [Distributed Prioritized Experience Replay](https://arxiv.org/abs/1803.00933)を参照のこと。

Ape-XはDQNの1種として位置付けられる。過去に提案されたDQNの性能を向上させる方法を組み合わせ、分散学習を導入したものがApe-Xに相当する。

- [Prioritized Experience Replay](https://arxiv.org/abs/1511.05952)
- [Double Q-Learning](https://arxiv.org/abs/1509.06461) with multi-step bootstrap targets
- [Dueling Network](https://arxiv.org/abs/1511.06581)

# RLLibについて

RLlibはPythonの分散実行ライブラリ`Ray`の1つのサブパッケージであり、学習実験を行う際には`ray`に加えて他のサブパッケージ`Tune`と組み合わせて使う。

## Ray

プロセス・ノード間で共有可能なimmutableオブジェクトを介して実行される関数を非同期で実行するためのライブラリ。ray特有のデコレートされたpythonの関数 `foo` を定義すると `foo.remote()` が呼べるようになる。そして `foo.remote()` を実行すると関数が実行されるが、オブジェクトの値そのものは評価されずに計算の実行がスケジューラーに予約される。関数が返すのは共有用のオブジェクトの固有ID。このIDを `ray.get` に渡すと計算が実行される。

複数の計算過程を実際に実行するためのスケジューラーは`ray.init()`で開始する。RLlibで学習を行うということはスケジューラーに計算を予約することに相当する。サーバーは毎回新たに立ち上げるか、すでに走っているサーバーにネットワークを通して計算の実行を登録することもできる。

## Tune

ハイパーパラメータの推定や学習実験を遂行・可視化するためのライブラリ。パラメーター探索のために学習実験を複数回、メタな推論アルゴリズムなどを組み合わせて行うための機能が実装されている。Tuneは学習実験を`Experiment`という単位で管理しており、キューとして複数の実験を溜め込み、順番に実行していく。

# 学習の始め方

今回はRLLibの[学習サンプルスクリプト](https://github.com/ray-project/ray/blob/master/python/ray/rllib/train.py)のように1回の強化学習実験のみを行う。

学習実験を行うためには実験の内容を記述した辞書を`ray.tune.tune.run_experiments()`に渡せばよい。辞書の内容は以下の通り：

- `name`: 実験の名前。ログが残るディレクトリの名前の一部になる。
- `run`: 実行するアルゴリズムの名前。今回は`APEX`。
- `local_dir`: ログを保存する場所。デフォルトは`~/ray_results`
- `checkpoint_freq`: モデルを保存するIterationの間隔
- `checkpoint_at_end`: 学習終了時に保存するか否か
- `resotre`: 過去のチェックポイントから再開する場合にファイル名を指定する
- `config`: 詳細な設定

`run`は`A3C`, `IMPALA`などのアルゴリズムを指定する。詳細は[Rllibのページ](https://ray.readthedocs.io/en/latest/rllib.html#algorithms)を参照。強化学習の場合には名前に対応した`Agent`クラスが初期化され、メイン関数である`_train（）`が複数回数呼び出される。

Ape-XはRLlibではDQNの特別な場合として実装されている。`APEXAgent`は`DQNAgent`をオーバーライドしたクラスであり、主なパラメーターの説明は[rayのソース](https://github.com/ray-project/ray/blob/master/python/ray/rllib/agents/dqn/dqn.py)にデフォルト設定と共に書いてある。デフォルト以外の設定を行いたい場合には上記の`Experiment`の設定`config`に辞書を追加すれば上書きされる。また、`config`のおすすめ設定が結果付きで公開されている: https://github.com/ray-project/rl-experiments

学習をカスタマイズしたい場合には`config`を書き換える必要がある。以下はその項目の説明。

#　学習の設定項目

項目が多いのでカテゴリに分けて説明する。カテゴリは実装と対応していないので注意。

## Workerに関するもの

Workerは分散で探索を行うプロセスのこと。学習を実行するマシンのリソース（CPU, GPU）に応じて設定する項目。
並列して実行する環境の数は`num_workers`x`num_envs_per_worker`になる。

| 名前                | 型  | 説明                                  |
|---------------------|-----|---------------------------------------|
| num_workers         | int | workerの数                            |
| num_cpus_per_worker | int | 1つのworkerあたりに割り当てるCPUの数  |
| num_gpus_per_worker | int | 1つのworkerあたりに割り当てるGPUの数  |
| num_envs_per_worker | int | 1つのworkerあたりに割り当てる環境の数 |

## Workerの探索に関するもの

APE-XではWorkerがε-greedyアルゴリズムを使って環境とインタラクション(Rollout）して経験データを生成する。εの値は1.0から`exploration_final_eps`へと学習が進むにつれて変更される。変更は`schedule_max_timesteps`ステップごとに`exploration_fraction`の割合だけ変更される。つまり、新しい$\epsilon_{i+1}$は以下のように決まる：

$$\epsilon_{i+1} = \epsilon_{i} + (\epsilon_{final} - \epsilon_{0})$$


$\epsilon_{final}$, $\epsilon_{0}$はそれぞれ`exploration_final_eps`、$\epsilon$の初期値(=1.0)。


| 名前                   | 型       | 説明                                                  |
|------------------------|----------|-------------------------------------------------------|
| per_worker_exploration | bool     | workerごとにε-greedyのεの値を変えるか否か             |
| exploration_fraction   | float    | εを減らす大きさ                                       |
| horizon                | int/null | Rolloutが特定のステップ数までが達したら強制的に止める |
| learning_starts        | int      | 何ステップRolloutしたら学習を開始するか               |
| sample_async           | bool     | Rolloutを非同期に実行するか否か                       |
| sample_batch_size      | int      | 何ステップ分を1つのバッチとしてReplay bufferに送るか  |

## Replay　Bufferに関するもの

[Prioritized Experience Replay](https://arxiv.org/abs/1511.05952)はReplay BufferにQ学習にとってどのくらい役立つかの優先度をつけ、優先度が高いものを多く学習する方法。優先度の付け方は以下のアルゴリズムに従う：

<img width="655" alt="Mendeley_Desktop.png" src="https://qiita-image-store.s3.amazonaws.com/0/137648/2c945591-cab8-8ec7-d8ab-11a6159bedb7.png">
[論文](https://arxiv.org/abs/1511.05952)より

9行目の$p_i$は優先度を決めるための確率で、TD誤差の絶対値に定数を加えたもの。

$p_i=|\delta_i|+\epsilon$

単純にTD誤差だけで確率分布をつくり選ぶと特定の経験のみが集中して使われ続けてしまうので分布の調整をハイパーパラメータで行う。下記の$\alpha, \beta$が相当する。特に$\beta$は`prioritized_replay_beta`から`final_prioritized_replay_beta`へと学習が進むにつれて$\epsilon$と同様に大きくなる。

| 名前                          | 型    | 説明                                                     |
|-------------------------------|-------|----------------------------------------------------------|
| buffer_size                   | int   | バッファ大きさ                                           |
| prioritized_replay            | bool  | 優先度付き経験再生をするか否か                           |
| worker_side_prioritization    | bool  | WorkerがRolloutしたときに優先度づけを行うか              |
| prioritized_replay_eps        | float | TD誤差から優先度の確率分布を計算する際に加える定数　ε    |
| prioritized_replay_alpha      | float | TD誤差から優先度の確率分布を計算する際に加える定数　α    |
| beta_annealing_fraction       | float | βの変更幅                                                |
| prioritized_replay_beta       | float | βの初期値                                                |
| final_prioritized_replay_beta | float | βの上限値                                                |
| compress_observations         | bool  | Replay Bufferに貯めるときに観測情報をLZ4で圧縮するか否か |

## 経験データの前処理に関するもの

環境からのデータ処理については[Rayのページ](https://ray.readthedocs.io/en/latest/rllib-models.html)に概要が載っている。

`Preprocessor`は環境（Gym.Env）のラッパーとして働き、主に観測情報の事前処理を担当する。環境がatariの場合には以下の4つが適用される。ただし`custom_preprocessor`を指定した場合は除く。

- clip reward：報酬の値を$[-1, 1]$に抑える
- greys caling：画像のグレースケール化ける
- frame stack: 複数のフレームをチャンネル方向に積んで1つの観測画像とする
- episodic life: resetが呼ばれたタイミングとは異なり、ゲームオーバーになったときだけゲームプレイをリセットする）がかかる


| 名前                  | 型     | 説明                                                                              |
|-----------------------|--------|-----------------------------------------------------------------------------------|
| preprocessor_pref     | string | 観測情報の前処理方法                                                              |
| observation_filter    | string | `Preprocessor`の後に行われる処理。デフォルトは何もしないNoFilterになっている。    |
| synchronize_filters   | bool   | FilterのパラメーターをIterationごとに同期するか否か                               |
| clip_rewards          | bool   | 各ステップの報酬を[-1, 1]に抑えるかどうか。trueの場合にはnumpy.sign()が使われる。 |
| compress_observations | bool   | Replay Bufferに貯めるときに観測情報をLZ4で圧縮するか否か                          |

## Q-Learningに関するもの

1 Iterationは以下のステップに相当する。`checkpoint_freq`はこのIterationが単位になっていることに注意。

1. 学習を行う
    - `timesteps_iteration`分stepが経過するか、`min_iter_time_s`秒たつまでの間以下を繰り返す：
        - Replay Bufferからサンプリング
        - 勾配計算
        - 重み更新
        - workerの重みを更新する
        - 優先度を更新する
2. 優先度の計算・割り当て
3. ログの更新（Tensorboardの更新）

Replay Bufferへのデータ追加・リプレイは非同期で行われる。


| 名前                       | 型    | 説明                                                                  |
|----------------------------|-------|-----------------------------------------------------------------------|
| double_q                   | bool  | [Double Q-Learning](https://arxiv.org/abs/1509.06461)を適用するか否か |
| dueling                    | bool  | [Dueling Network](https://arxiv.org/abs/1511.06581)を適用するか否か   |
| gamma                      | float | 報酬の割引率                                                          |
| n_step                     | int   | n-step Q-Learningのnの値                                              |
| target_network_update_freq | int   | 何ステップごとにQ-networkを更新するか                                 |
| timesteps_per_iteration    | int   | Iterationのステップ間隔                                               |
| min_iter_time_s            | int   | 1 Iterationにかける最小の秒数                                         |
| train_batch_size           | int   | Q-Networkのバッチの大きさ                                             |

## Q-Networkに関するもの

Q-Networkの前向き計算・後ろ向き計算はTensorflowかPyTorchを使って行う。Q-Networkの構成は自分で書くこともできるが、環境のもつ`observation_space`に合わせて自動的に構成を決めてくれる。特に画像の場合にはCNNが適用され、Convolutionのカーネル大きさのみを設定で上書きすることもできる。

| 名前      | 型    | 説明                                                                                                          |
|-----------|-------|---------------------------------------------------------------------------------------------------------------|
| hiddens   | int   | Convolution層のあとのFully-Connectedな層の大きさ                                                              |
| noisy     | bool  | [Noisy Network](https://arxiv.org/abs/1706.10295)を適用するか否か。これがtrueの場合にはε-greedyは使われない。 |
| sigma0    | float | Noisy Networkの初期パラメーター                                                                               |
| num_atoms | int   | Q-networkの出力分布の数。1より大きくすると[distributional Q-learning](https://arxiv.org/abs/1707.06887)になる |
| v_min     | float | distributional Q-learningのパラメーター                                                                       |
| v_max     | float | distributional Q-learningのパラメーター                                                                       |

## Optimizerに関するもの

| 名前                     | 型     | 説明                                                           |
|--------------------------|--------|----------------------------------------------------------------|
| optimizer_class          | string | Optimizerの種類。Ape-Xの場合は`AsyncReplayOptimizer`を使う。   |
| debug                    | bool   | デバックモードの切り替え                                       |
| max_weight_sync_delay    | int    | モデルパラメーターの更新を前回行ってから何ステップ最低で待つか |
| num_replay_buffer_shards | int    | Replayの並列数                                                 |
| lr                       | float  | 学習率                                                         |
| adam_epsilon             | float  | Adam Optimizerのεの値                                          |
| grad_norm_clipping       | float  | 勾配のノルムの最大値                                           |

## その他

`batch_mode`は`truncate_episodes`、`complete_episodes`のどちらか。`complete_episodes`の場合はエピソードが終わるまでBufferに送らない。

| 名前                   | 型     | 説明                                                                                                                    |
|------------------------|--------|-------------------------------------------------------------------------------------------------------------------------|
| schedule_max_timesteps | int    | ハイパーパラメータを学習の進み具合に応じて変える場合のステップ間隔幅                                                    |
| gpu                    | bool   | GPUを使うか                                                                                                               |
| gpu_fraction           | float  | GPUの使用率(0~1, 1が100%)                                                                                               |
| monitor                | bool   | 定期的にRollout結果を動画で保存するか否か                                                                               |
| batch_mode             | string | Batchのサンプリング方法。 |
| tf_session_args        | dict   | TensorflowのSessionを初期化するときのパラメーター                                                                       |

`tf_session_args`は以下が設定できる。詳しいことはTensorflowのドキュメントに書いてある。

- allow_soft_placement
- device_count
    - CPU
- gpu_options
    - allow_growth
- inter_op_parallelism_threads
- intra_op_parallelism_thread
- log_device_placement