英語で書いたブログ記事の日本語焼き直し。

1. はじめに

5月26日にDeepMindが強化学習における Experience Replay (経験再生) のためのフレームワークとして、Reverbをリリースした。(参照)

Reverb is an efficient and easy-to-use data storage and transport system designed for machine learning research. Reverb is primarily used as an experience replay system for distributed reinforcement learning algorithms but the system also supports multiple data structure representations such as FIFO, LIFO, and priority queues.

(筆者訳) Reverb は機械学習の研究のためにデザインされた効率的で簡単に利用できるデータ・ストレージ兼データ転送システムです。Reverbは主に分散強化学習アルゴリズムのための experience replay (経験再生) に利用されるが、このシステムはFIFO、LIFO、重み付きキューといったデータ構造をサポートしています。 (筆者訳ここまで)

同じくDeepMind製の強化学習フレームワーク Acme (A research framework for reinforcement learning) は、このReverbを利用している。(Acmeについてはまたの機会に)

2. インストール

記事を書いている6月26日現在、Reverbは、LinuxベースのOSのみをサポートし、 production use レベルではないと公式に記載されている。

TensorFlow の開発版が必要で、以下のコマンドで、PyPIからインストールすることができる

pip install tf-nightly==2.3.0.dev20200604 dm-reverb-nightly

3. アーキテクチャ

Reverbはサーバー・クライアント方式を採用しており、その名前付けルールは他のReplay Bufferの実装よりもデータベースの用語に近いと感じる。

3.1 サーバー

サーバー側のサンプルコードは以下のようになる。

import reverb

server = reverb.Server(tables=[
    reverb.Table(
        name='my_table',
        sampler=reverb.selectors.Uniform(),
        remover=reverb.selectors.Fifo(),
        max_size=100,
        rate_limiter=reverb.rate_limiters.MinSize(1)),
    ],
    port=8000
)

この例では、キャパシティ100の普通のReplay Buffer (一様サンプリング、古いものから順次上書き) が 8000 番のポートを待ち受ける。 reverb.rate_limiters.MinSize(1) は最低 1個のアイテムが入ってくるまでは、いかなるサンプリング要求もブロックすることを意味している。

3.1.1 要素選択 (`sampler`/`remover`)

上の例でわかるように、Reverbでは要素のサンプリングと削除(上書き)のロジックをそれぞれ独立して指定することができる様になっている。

Reverbでサポートしているロジックは、 reverb.selectors に実装してあり、次のようなものがある

Uniform: すべてのアイテムから一様分布でランダムに選択
Prioritized: プライオリティに応じてランダムに選択
Fifo: 最も古いデータを選択
Lifo: 最も新しいデータを選択
MinHeap: プライオリティが最低のデータを選択
MaxHeap: プライオリティが最高のデータを選択

3.1.2 制約条件指定 (`rate_limiter`)

rate_limiter 引数は、Replay Bufferの使用条件をを設定することができる。
Reverbでサポートされている条件は reverb.rate_limiters に実装されており、
次のようなものがある

MinSize: サンプル可能になる最低限のアイテム数を設定
SampleToInsertRatio: データ挿入(更新)とデータサンプルの平均的な比率を設定
- 上書きされるまでにおよそ何回サンプルされるという利用法に役立つ(らしい)
Queue: 上書きされる前に、ちょうど1回取り出される (FIFO用)
Stack: 上書きされる前に、ちょうど1回取り出される（LIFO用）

ソースコードのコメントを見ると、 reverb.rate_limiters.Queue と reverb.rate_limiters.Stack は直接利用することが推奨されておらず、代わりに静的メソッドの reverb.Table.queue と reverb.Table.stack がそれぞれ、FIFOとLIFOのロジックを持つReplay Bufferになるように適切に sampler、remover、 rate_limiter を設定してくれる。

3.2 クライアント

クライアントプログラムのサンプルコードは以下

import reverb

client = reverb.Client('localhost:8000') # サーバーとクライアントが同じマシーンの場合

# [0,1] という状態 (observation) を priority 1.0 でReplay Bufferに入れる例
client.insert([0, 1], priorities={'my_table': 1.0})

# サンプリングすると、 generator が返ってくる
client.sample('my_table', num_samples=2))

3.3 保存/読み込み

Reverbは、データの保存/読み込み機能に対応している。
クライアントから、以下のコードを実行することで、現在のサーバー内のデータがファイル上に保存され、保存したファイルパスが得られる

checkpoint_path = client.checkpoint()

保存したデータを利用して、サーバーを作成することでもとのデータの状態を復元することができる。
注意しないといけないことは、コンストラクタのtables引数は、データを保存した元のサーバーと全く同じものを、 ユーザーの責任で 指定しないことである。

checkpointer = reverb.checkpointers.DefaultCheckpointer(path=checkpoint_path)
server = reverb.Server(tables=[...], checkpointer=checkpointer)

最後に

DeepMindが発表したExperience Replay用の新しいフレームワークのReverbは、まだ安定版には到達していないが、柔軟で大規模な強化学習をターゲットとした将来有望なものであると感じた。

拙作のExperience Replay用ライブラリ cpprb にとっては、巨大なライバルが突然出現したわけではあるが、より小規模な強化学習の実験においては、cpprbの方が便利で使いやすい部分もあると思う。 (過去Qiita記事参照)

(更新: 2020.6.29) クライアントの使用方法を調査して書きました！

【強化学習】爆誕: DeepMind製 Experience Replay (経験再生)フレームワーク Reverb