TensorFlow が 5 月 14 日の Google I/O に戻ってきます。今すぐ登録

再生バッファ

Copyright 2018 The TF-Agents Authors.

TensorFlow.orgで表示

Google Colabで実行

GitHub でソースを表示{

ノートブックをダウンロード/a0}

はじめに

強化学習アルゴリズムは、環境でポリシーを実行するときに再生バッファを使用して経験のトラジェクトリを格納します。トレーニング中、再生バッファはトラジェクトリのサブセット（シーケンシャルサブセットまたはサンプルのいずれか）がエージェントの経験を「再生」するために照会されます。

このコラボでは、共通の API を共有する 2 種類の再生バッファ（python 方式と tensorflow 方式）について考察します。次のセクションでは API と各バッファの実装、およびデータ収集トレーニングでそれらを使用する方法を説明します。

セットアップ

tf-agents をまだインストールしていない場合はインストールしてください。

pip install -q tf-agents
pip install -q 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 が必要です。この仕様は、バッファに追加されるトラジェクトリ要素の TensorSpec に対応しています。この仕様は通常、トレーニング時にエージェントが期待する形状、タイプ、構造を定義するエージェントの agent.collect_data_spec を調査することで取得できます（詳細は後述します）。

TFUniformReplayBuffer

TFUniformReplayBuffer は TF-Agents で最も一般的に使用される再生バッファであるため、このチュートリアルではこの再生バッファを使用します。TFUniformReplayBuffer では、バッキングバッファの記憶は tensorflow 変数によって行われるため、計算グラフの一部になっています。

バッファには複数の要素がバッチ単位で格納され、バッチセグメントごとに最大容量の max_length 要素があります。したがって、合計バッファ容量は、batch_size x max_length 要素となります。バッファに格納される要素はすべて、対応するデータ仕様を持っている必要があります。再生バッファがデータ収集に使用される場合、仕様はエージェントの収集データ仕様になります。

バッファの作成:

TFUniformReplayBuffer を作成するには、次の値を渡します。

バッファが格納するデータ要素の仕様
バッファのバッチサイズに対応する batch size
バッチセグメントごとの max_length 要素数

サンプルデータ仕様（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)

バッファへの書き込み:

再生バッファに要素を追加するため、items がバッファに追加されるアイテムのバッチを表すテンソルのリスト/タプル/ネストになっている add_batch(items) メソッドを使用します。items の各要素には外側の次元に等しい batch_size が必要であり、残りの次元は（再生バッファのコンストラクタに渡されたデータ仕様と同じ）アイテムのデータ仕様に準拠していなければなりません。

アイテムのバッチを追加する例を以下に示します。

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 として返します。その後、データセットイテレータを作成し、バッファ内のアイテムのサンプルをイテレートできます。
gather_all() - バッファ内のすべてのアイテムを形状 [batch, time, data_spec] を含むテンソルとして返します。

上記の各メソッドを使用して再生バッファから読み込む例を以下に示します。

# 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))

WARNING:tensorflow:From <ipython-input-1-1f9907631cb9>:7: ReplayBuffer.get_next (from tf_agents.replay_buffers.replay_buffer) is deprecated and will be removed in a future version.
Instructions for updating:
Use `as_dataset(..., single_deterministic_pass=False) instead.
Iterator trajectories:
[(TensorShape([4, 2, 3]), (TensorShape([4, 2, 5]), TensorShape([4, 2, 3, 2]))), (TensorShape([4, 2, 3]), (TensorShape([4, 2, 5]), TensorShape([4, 2, 3, 2]))), (TensorShape([4, 2, 3]), (TensorShape([4, 2, 5]), TensorShape([4, 2, 3, 2])))]
WARNING:tensorflow:From <ipython-input-1-1f9907631cb9>:24: ReplayBuffer.gather_all (from tf_agents.replay_buffers.replay_buffer) is deprecated and will be removed in a future version.
Instructions for updating:
Use `as_dataset(..., single_deterministic_pass=True)` instead.
Trajectories from gather all:
(TensorShape([32, 6, 3]), (TensorShape([32, 6, 5]), TensorShape([32, 6, 3, 2])))

PyUniformReplayBuffer

PyUniformReplayBuffer は TFUniformReplayBuffer と同じ機能を持ちますが、データは tf 変数ではなく numpy 配列に格納されます。このバッファはグラフ外のデータ収集に使用されます。numpy にバッキングストレージを用意すると、Tensorflow 変数を使用せずに一部のアプリケーションがデータ操作（優先度を更新するためのインデックス作成など）を実行しやすくなります。ただし、この実装では Tensorflow を使用したグラフ最適化のメリットを享受できません。

エージェントのポリシートラジェクトリの仕様から 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))

トレーニング中に再生バッファを使用する

再生バッファを作成してアイテムを読み書きする方法がわかったので、エージェントのトレーニング中に再生バッファを使用してトラジェクトリを格納できるようになりました。

データ収集

まず、データ収集中に再生バッファを使用する方法を見てみましょう。

TF-Agents では、Driver（詳細は Driver のチュートリアルをご覧ください）を使用して環境内の経験を収集します。Driver を使用するには、Driver がトラジェクトリを受け取ったときに実行する関数である Observer を指定します。

そのため、トラジェクトリ要素を再生バッファに追加するために add_batch(items) を呼び出すオブザーバーを追加してアイテム（のバッチ）を再生バッファに追加します。

TFUniformReplayBuffer を使用した例を以下に示します。まず、環境、ネットワーク、エージェントを作成します。その後、TFUniformReplayBuffer を作成します。再生バッファ内のトラジェクトリ要素の仕様は、エージェントの収集データ仕様に等しくなっていることに注意してください。その後、その 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()

WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tf_agents/drivers/dynamic_step_driver.py:203: calling while_loop_v2 (from tensorflow.python.ops.control_flow_ops) with back_prop=False is deprecated and will be removed in a future version.
Instructions for updating:
back_prop=False is deprecated. Consider using tf.stop_gradient instead.
Instead of:
results = tf.while_loop(c, b, vars, back_prop=False)
Use:
results = tf.nest.map_structure(tf.stop_gradient, tf.while_loop(c, b, vars))

トレーニングステップ用のデータ読み込み

トラジェクトリ要素を再生バッファに追加した後は、再生バッファからトラジェクトリのバッチを読み取り、トレーニングステップの入力データとして使用できます。

トレーニングループ内で再生バッファのトラジェクトリをトレーニングする方法の例を以下に示します。

# 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)

WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow/python/util/dispatch.py:201: calling foldr_v2 (from tensorflow.python.ops.functional_ops) with back_prop=False is deprecated and will be removed in a future version.
Instructions for updating:
back_prop=False is deprecated. Consider using tf.stop_gradient instead.
Instead of:
results = tf.foldr(fn, elems, back_prop=False)
Use:
results = tf.nest.map_structure(tf.stop_gradient, tf.foldr(fn, elems))