DeepMind製強化学習フレームワーク Acme

[以前の記事](https://qiita.com/ymd_h/items/97b9c1ff6809242db211)で少し触れたDeepMind製の強化学習向けフレームワーク [Acme](https://github.com/deepmind/acme) (a research framework for reinforcement learning) についての紹介記事

(いつものように、自ブログに投稿した[英語記事](https://ymd_h.gitlab.io/ymd_blog/posts/acme_by_deepmind/)の日本語焼き直し。)

# 1. はじめに

DeepMindの研究者が実際に毎日のように使っている強化学習向けフレームワークのコード（の一部）をOSSとして[公開](https://deepmind.com/research/publications/Acme)したもの。

Acmeはシンプルな学習ループAPIを提供しており、ざっくりとは以下のようなコードになる。

```python:学習ループ
loop = acme.EnvironmentLoop(environment, agent)
loop.run()
```

Acmeは、ネットワークとして[TensorFlow](https://www.tensorflow.org/) (※ 特定のnightlyビルドのみ)と[JAX](https://github.com/google/jax)をサポートしている。

また、経験再生用に [Reverb](https://github.com/deepmind/reverb) を利用してる。(拙著の[解説記事1](https://qiita.com/ymd_h/items/97b9c1ff6809242db211)と[解説記事2](https://qiita.com/ymd_h/items/7eedf8f673aef5d0fb8e))

[テクニカルレポート](https://arxiv.org/abs/2006.00979)には、"Acme: A Research Framework for Distributed Reinforcement Learning" と分散強化学習だと表現されているが、[FAQ](https://github.com/deepmind/acme/blob/master/docs/faq.md)を読む限り、残念ながら分散学習のエージェントは公開されておらず、公開する時期も決まっていないらしい。

# 2. インストール
[PyPI](https://pypi.org/)に "dm-acme" といパッケージ名で[公開](https://pypi.org/project/dm-acme/)されている。

指定可能なオプションは、"jax", "tf", "envs", "reverb", "testing" の5つ。
個人的には、経験再生用の"reverb"は指定必須、"tf"または"jax"のどちらかひとつも必須だと考える。
特にTensorFlowは `'tf-nightly==2.4.0.dev20200708'` が指定してあったりと、環境の互換性など投げ捨てられているので、おとなしく綺麗な環境にAcmeのインストールオプションでインストールすることをおすすめする。

"envs"は指定することで、"dm-control" や "gym" などの強化学習向け環境がインストールされる。("dm-control" は [MuJoCo](http://www.mujoco.org/)のライセンスが必要。)

```shell:インストールTensorFlow版
pip install dm-acme[reverb,tf]
```

```shell:インストールJAX版
pip install dm-acme[reverb,jax]
```


# 3. 環境
[DeepMind Environment API](https://github.com/deepmind/dm_env)に準拠する環境を利用することができる。`gym.Env` との大きな違いは、`step(action)`等のメンバー関数が、`tuple`ではなく`dm_env.TimeStep`という独自クラスを返すことである。`TimeStep`クラスでは、観測値や報酬などに加えて、ステップ（遷移）が、エピソードの最初(`StepType.FIRST`)、途中（`StepType.MID`）、最後（`StepType.LAST`）のどれであるかの情報も保持している。

`gym.Env`向けに`acme.wrappers.GymWrapper`が提供されており、ラップするだけでAcmeで使うことができる

```python:gym.Envの利用
import acme
import gym

env = acme.wrappers.GymWrapper(gym.make("MountainCarContinuous-v0"))
```


# 4. エージェント
実装済みのエージェントは[こちら](https://github.com/deepmind/acme/tree/master/acme/agents)にある。
2020年8月現在、11個のエージェントが提供されている。

* 連続値制御
    * Deep Deterministic Policy Gradient (DDPG)
    * Distributed Distributional Deep Determinist (D4PG)
    * Maximum a posteriori Policy Optimisation (MPO)
    * Distributional Maximum a posteriori Policy Optimisation (DMPO)	
* 離散値制御
    * Deep Q-Networks (DQN)
    * Importance-Weighted Actor-Learner Architectures (IMPALA)
    * Recurrent Replay Distributed DQN (R2D2)
* Batch RL
    * Behavior Cloning (BC)
* デモンストレーションからの学習
    * Deep Q-Learning from Demonstrations (DQfD)
    * Recurrent Replay Distributed DQN from Demonstratinos (R2D3)
* モデルベースRL
    * Monte-Carlo Tree Search (MCTS)

これらの実装は`acme.agents.agent.Agent`を拡張することで実装されているので、カスタムアルゴリズムを実装する際は、同様に拡張することで可能である（と思われる）。

上の提供されているエージェントは、DeepMind製の深層学習ライブラリの[Sonnet](https://github.com/deepmind/sonnet) (TensorFlow)または[Haiku](https://github.com/deepmind/dm-haiku) (JAX)を利用している。そのため、（ざっと確認した限りでは）実装済みのエージェントの内部ネットワークを定義するためには、SonnetやHaikuで深層学習モデルを構築して渡すことが想定されている。（もちろんインターフェースには無関係なので、独自実装する際には関係ない。）

# 5. チュートリアル
Acmeは[クイックスタート ノートブック](https://github.com/deepmind/acme/blob/master/examples/quickstart.ipynb)を提供している。また、[MuJoCo](http://www.mujoco.org/)のライセンスを所持していれば、[チュートリアル ノートブック](https://github.com/deepmind/acme/blob/master/examples/tutorial.ipynb)も試すことができる。

# 6. おわりに

DeepMindが公開している強化学習用フレームワークのAcmeの概要を調査した。第一線級の企業が普段から利用しているコードと同じものが（一部とは言え）提供されており、今後の発展も含めて非常に興味深い。

ただし、TensorFlowのバージョンがnightlyのとあるビルドに固定されていたりと、一般的なユーザーが利用するには少し難があるようにも感じられる。

引き続き、状況をウォッチしていきたい。

# 余談
[クイックスタート](https://github.com/deepmind/acme/blob/master/examples/quickstart.ipynb)にスクラッチでコードが書いてあるような、Gym環境をノートブック上で表示するためのコードをパッケージ化して[Gym-Notebook-Wrapper](https://gitlab.com/ymd_h/gym-notebook-wrapper)として提供しています。([解説記事](https://qiita.com/ymd_h/items/c393797deb72e1779269))
こちらも興味を持っていただけたらぜひ使って見てください。