TF-Agents 0.4 Tutorials : 再生バッファ (翻訳/解説)

翻訳 : (株)クラスキャットセールスインフォメーション
作成日時 : 04/21/2020 (0.4)

* 本ページは、TF Agents の以下のドキュメントを翻訳した上で適宜、補足説明したものです：

Replay Buffers

* サンプルコードの動作確認はしておりますが、必要な場合には適宜、追加改変しています。
* ご自由にリンクを張って頂いてかまいませんが、sales-info@classcat.com までご一報いただけると嬉しいです。

再生バッファ

イントロダクション

強化学習アルゴリズムは環境でポリシーを実行するとき経験の軌道 (= trajectories) をストアするために再生バッファを利用します。訓練の間、エージェントの経験を「再生する」ために軌道のサブセット (シークエンシャルなサブセットかサンプル) のために再生バッファが問い合わせがされます。

この colab では、2 つのタイプの再生バッファを調べます : python-支援と tensorflow-支援で、共通 API を共有します。次のセクションで、API、各々のバッファ実装そしてそれらをデータ収集訓練の間にどのように利用するかを説明します。

セットアップ

tf-agents をまだであればインストールします。

!pip install --upgrade tensorflow-probability
!pip install tf-agents
!pip install gym

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow as tf
import numpy as np

from tf_agents import specs
from tf_agents.agents.dqn import dqn_agent
from tf_agents.drivers import dynamic_step_driver
from tf_agents.environments import suite_gym
from tf_agents.environments import tf_py_environment
from tf_agents.networks import q_network
from tf_agents.replay_buffers import py_uniform_replay_buffer
from tf_agents.replay_buffers import tf_uniform_replay_buffer
from tf_agents.specs import tensor_spec
from tf_agents.trajectories import time_step

tf.compat.v1.enable_v2_behavior()

再生バッファ API

再生バッファ・クラスは次の定義とメソッドを持ちます :

class ReplayBuffer(tf.Module):
  """Abstract base class for TF-Agents replay buffer."""

  def __init__(self, data_spec, capacity):
    """Initializes the replay buffer.

    Args:
      data_spec: A spec or a list/tuple/nest of specs describing
        a single item that can be stored in this buffer
      capacity: number of elements that the replay buffer can hold.
    """

  @property
  def data_spec(self):
    """Returns the spec for items in the replay buffer."""

  @property
  def capacity(self):
    """Returns the capacity of the replay buffer."""

  def add_batch(self, items):
    """Adds a batch of items to the replay buffer."""

  def get_next(self,
               sample_batch_size=None,
               num_steps=None,
               time_stacked=True):
    """Returns an item or batch of items from the buffer."""

  def as_dataset(self,
                 sample_batch_size=None,
                 num_steps=None,
                 num_parallel_calls=None):
    """Creates and returns a dataset that returns entries from the buffer."""


  def gather_all(self):
    """Returns all the items in buffer."""
    return self._gather_all()

  def clear(self):
    """Resets the contents of replay buffer"""

再生バッファが初期化されるとき、それがストアする要素の data_spec を要求することに注意してください。この spec はバッファに追加される trajectory 要素の TensorSpec に相当します。この spec は通常はエージェントの agent.collect_data_spec を見ることにより獲得されます、これは訓練の時 (more on that later) エージェントにより想定される shape, 型と構造を定義します。

TFUniformReplayBuffer

TFUniformReplayBuffer は TF-Agents で最も一般的に使用される再生バッファで、そのためここでチュートリアル内でそれを使用します。TFUniformReplayBuffer では支援 (= backing) バッファ・ストレージは tensorflow 変数により成されるため計算グラフの一部です。

バッファは要素のバッチをストアしてバッチセグメント毎に最大容量 max_length 要素を持ちます。このため、総バッファ容量は batch_size x max_length 要素です。バッファにストアされる要素総ては適合するデータ spec を持たなければなりません。再生バッファがデータ収集のために使用されるとき、spec はエージェントの collect data spec です。

バッファを作成する

TFUniformReplayBuffer を作成するために以下を渡します :

バッファがストアするデータ要素の spec
バッファのバッチサイズに対応するバッチサイズ
バッチセグメント毎の要素の max_length 数

ここにサンプル data specs, batch_size 32 と max_length 1000 を持つ TFUniformReplayBuffer を作成するサンプルがあります。

data_spec =  (
        tf.TensorSpec([3], tf.float32, 'action'),
        (
            tf.TensorSpec([5], tf.float32, 'lidar'),
            tf.TensorSpec([3, 2], tf.float32, 'camera')
        )
)

batch_size = 32
max_length = 1000

replay_buffer = tf_uniform_replay_buffer.TFUniformReplayBuffer(
    data_spec,
    batch_size=batch_size,
    max_length=max_length)

バッファに書く

要素を再生バッファに追加するため、add_batch(items) メソッドを使用します、そこでは項目はバッファに追加されるべき項目のバッチを表す tensor のリスト/タプル/ネストです。項目の各要素は batch_size に等しい外側の次元を持たなければならずそして残りの次元は (再生バッファに渡された data specs と同じ) 項目の data spec を忠実に守らなければなりません。

ここに項目のバッチを追加するサンプルがあります :

action = tf.constant(1 * np.ones(
    data_spec[0].shape.as_list(), dtype=np.float32))
lidar = tf.constant(
    2 * np.ones(data_spec[1][0].shape.as_list(), dtype=np.float32))
camera = tf.constant(
    3 * np.ones(data_spec[1][1].shape.as_list(), dtype=np.float32))
  
values = (action, (lidar, camera))
values_batched = tf.nest.map_structure(lambda t: tf.stack([t] * batch_size),
                                       values)
  
replay_buffer.add_batch(values_batched)

バッファから読む

TFUniformReplayBuffer からデータを読む 3 つの方法があります :

get_next() – バッファから 1 つのサンプルを返します。返されるサンプル・バッチサイズと時間ステップの数はこのメソッドへの引数を通して指定できます。
as_dataset() – 再生バッファを tf.data.Dataset として返します。それからデータセット iterator を作成できてバッファの項目のサンプルを通して反復できます。
gather_all() – バッファの総ての項目を shape [batch, time, data_spec] を持つ Tensor として返します。

下はこれらのメソッドの各々を使用して再生バッファからどのように読むかのサンプルです :

# add more items to the buffer before reading
for _ in range(5):
  replay_buffer.add_batch(values_batched)

# Get one sample from the replay buffer with batch size 10 and 1 timestep:

sample = replay_buffer.get_next(sample_batch_size=10, num_steps=1)

# Convert the replay buffer to a tf.data.Dataset and iterate through it
dataset = replay_buffer.as_dataset(
    sample_batch_size=4,
    num_steps=2)

iterator = iter(dataset)
print("Iterator trajectories:")
trajectories = []
for _ in range(3):
  t, _ = next(iterator)
  trajectories.append(t)
  
print(tf.nest.map_structure(lambda t: t.shape, trajectories))

# Read all elements in the replay buffer:
trajectories = replay_buffer.gather_all()

print("Trajectories from gather all:")
print(tf.nest.map_structure(lambda t: t.shape, trajectories))

PyUniformReplayBuffer

PyUniformReplayBuffer は TFUniformReplayBuffer と同じ機能を持ちますが、tf 変数の代わりに、そのデータは numpy 配列でストアされます。このバッファは out-of-graph データ収集のために使用できます。TensorFlow 変数を使用することなく numpy で支援ストレージを持つことは幾つかのアプリケーションにとっては (優先度を更新するためのインデキシングのような) データ操作を行なうことを容易にします。けれども、この実装は TensorFlow によるグラフ最適化の恩恵を持ちません。

下はエージェントのポリシー trajectory specs から PyUniformReplayBuffer をインスタンス化するサンプルです :

replay_buffer_capacity = 1000*32 # same capacity as the TFUniformReplayBuffer

py_replay_buffer = py_uniform_replay_buffer.PyUniformReplayBuffer(
    capacity=replay_buffer_capacity,
    data_spec=tensor_spec.to_nest_array_spec(data_spec))

訓練の間に再生バッファを使用する

再生バッファをどのように作成し、項目をそれに書いてそれから読むかを知った今、エージェントの訓練の間に trajectories をストアするためにそれを利用できます。

データ・コレクション

最初に、データコレクションの間に再生バッファをどのように使用するかを見ましょう。

TF-Agents では環境での経験を収集するために Driver (より多くの詳細については Driver チュートリアル参照) を使用します。Driver を使用するために、Observer を指定します、これは trajectory を受け取るとき Driver が実行するための関数です。

こうして、trajectory 要素を再生バッファに追加するため、再生バッファに項目 (のバッチ) を追加するために add_batch(items) を呼び出す観測者を追加します。

下は TFUniformReplayBuffer を伴うこれのサンプルです。最初に環境、ネットワークとエージェントを作成します。それからTFUniformReplayBuffer を作成します。再生バッファの trajectory 要素の specs はエージェントの collect data spec に等しいことに注意してください。そしてその add_batch メソッドをドライバーのための観測者として設定します、これは訓練の間にデータ収集を行ないます :

env = suite_gym.load('CartPole-v0')
tf_env = tf_py_environment.TFPyEnvironment(env)

q_net = q_network.QNetwork(
    tf_env.time_step_spec().observation,
    tf_env.action_spec(),
    fc_layer_params=(100,))

agent = dqn_agent.DqnAgent(
    tf_env.time_step_spec(),
    tf_env.action_spec(),
    q_network=q_net,
    optimizer=tf.compat.v1.train.AdamOptimizer(0.001))

replay_buffer_capacity = 1000

replay_buffer = tf_uniform_replay_buffer.TFUniformReplayBuffer(
    agent.collect_data_spec,
    batch_size=tf_env.batch_size,
    max_length=replay_buffer_capacity)

# Add an observer that adds to the replay buffer:
replay_observer = [replay_buffer.add_batch]

collect_steps_per_iteration = 10
collect_op = dynamic_step_driver.DynamicStepDriver(
  tf_env,
  agent.collect_policy,
  observers=replay_observer,
  num_steps=collect_steps_per_iteration).run()

訓練ステップのためにデータを読む

再生バッファに trajectory 要素を追加した後、訓練ステップのための入力データとして使用するために再生バッファから trajectory のバッチを読むことができます。

ここに訓練ループで再生バッファからの trajectory 上でどのように訓練するかのサンプルがあります :

# Read the replay buffer as a Dataset,
# read batches of 4 elements, each with 2 timesteps:
dataset = replay_buffer.as_dataset(
    sample_batch_size=4,
    num_steps=2)

iterator = iter(dataset)

num_train_steps = 10

for _ in range(num_train_steps):
  trajectories, _ = next(iterator)
  loss = agent.train(experience=trajectories)

以上

月	火	水	木	金	土	日
		1	2	3	4	5
6	7	8	9	10	11	12
13	14	15	16	17	18	19
20	21	22	23	24	25	26
27	28	29	30