回放缓冲区

版权所有 2023 The TF-Agents Authors。

在 TensorFlow.org 上查看

在 Google Colab 中运行

在 GitHub 上查看源代码

下载笔记本

简介

强化学习算法使用回放缓冲区来存储在环境中执行策略时的经验轨迹。在训练期间，会查询回放缓冲区以获取轨迹的子集（顺序子集或样本），以“回放”代理的经验。

在这个 Colab 中，我们探索了两种类型的回放缓冲区：python 支持的和 tensorflow 支持的，它们共享一个通用的 API。在以下部分，我们将描述 API、每个缓冲区实现以及如何在数据收集训练期间使用它们。

设置

如果您还没有安装 tf-agents，请安装。

pip install tf-agents
pip install tf-keras

import os
# Keep using keras-2 (tf-keras) rather than keras-3 (keras).
os.environ['TF_USE_LEGACY_KERAS'] = '1'

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

回放缓冲区 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，我们传入

1. 缓冲区将存储的数据元素的规范
2. 对应于缓冲区批次大小的 batch size
3. 每个批次段的 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)

写入缓冲区

要将元素添加到回放缓冲区，我们使用 add_batch(items) 方法，其中 items 是表示要添加到缓冲区的项目批次的张量列表/元组/嵌套。items 的每个元素都必须具有等于 batch_size 的外部维度，其余维度必须符合项目的 data spec（与传递给回放缓冲区构造函数的 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 读取数据

1. get_next() - 从缓冲区返回一个样本。可以通过此方法的参数指定返回的样本批次大小和时间步长数量。
2. as_dataset() - 将回放缓冲区作为 tf.data.Dataset 返回。然后，您可以创建一个数据集迭代器并遍历缓冲区中项目的样本。
3. 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 /tmpfs/tmp/ipykernel_27300/1348928897.py: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 /tmpfs/tmp/ipykernel_27300/1348928897.py: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 相同，但它使用 NumPy 数组存储数据，而不是 TensorFlow 变量。此缓冲区可用于图外数据收集。使用 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，我们需要指定一个 Observer，它是一个函数，用于 Driver 在接收到轨迹时执行。

因此，要将轨迹元素添加到回放缓冲区，我们添加一个观察器，该观察器调用 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()

读取用于训练步骤的数据

将轨迹元素添加到回放缓冲区后，我们可以从回放缓冲区读取轨迹批次，用作训练步骤的输入数据。

以下是如何在训练循环中使用回放缓冲区中的轨迹进行训练的示例

# 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.9/site-packages/tensorflow/python/util/dispatch.py:1260: 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))