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

以前の記事で少し触れたDeepMind製の強化学習向けフレームワーク Acme (a research framework for reinforcement learning) についての紹介記事

(いつものように、自ブログに投稿した英語記事の日本語焼き直し。)

1. はじめに

DeepMindの研究者が実際に毎日のように使っている強化学習向けフレームワークのコード（の一部）をOSSとして公開したもの。

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

学習ループ

loop = acme.EnvironmentLoop(environment, agent)
loop.run()

Acmeは、ネットワークとしてTensorFlow (※ 特定のnightlyビルドのみ)とJAXをサポートしている。

また、経験再生用に Reverb を利用してる。(拙著の解説記事1と解説記事2)

テクニカルレポートには、"Acme: A Research Framework for Distributed Reinforcement Learning" と分散強化学習だと表現されているが、FAQを読む限り、残念ながら分散学習のエージェントは公開されておらず、公開する時期も決まっていないらしい。

2. インストール

PyPIに "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のライセンスが必要。)

インストールTensorFlow版

pip install dm-acme[reverb,tf]

インストールJAX版

pip install dm-acme[reverb,jax]

3. 環境

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

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

gym.Envの利用

import acme
import gym

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

4. エージェント

実装済みのエージェントはこちらにある。
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 (TensorFlow)またはHaiku (JAX)を利用している。そのため、（ざっと確認した限りでは）実装済みのエージェントの内部ネットワークを定義するためには、SonnetやHaikuで深層学習モデルを構築して渡すことが想定されている。（もちろんインターフェースには無関係なので、独自実装する際には関係ない。）

5. チュートリアル

Acmeはクイックスタート ノートブックを提供している。また、MuJoCoのライセンスを所持していれば、チュートリアル ノートブックも試すことができる。

6. おわりに

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

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

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

余談

クイックスタートにスクラッチでコードが書いてあるような、Gym環境をノートブック上で表示するためのコードをパッケージ化してGym-Notebook-Wrapperとして提供しています。(解説記事)
こちらも興味を持っていただけたらぜひ使って見てください。