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として提供しています。(解説記事)
こちらも興味を持っていただけたらぜひ使って見てください。