回放緩衝區

Copyright 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

Replay Buffer 類別具有以下定義和方法

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 的外部維度，且其餘維度必須符合項目的資料規格 (與傳遞至回放緩衝區建構函式的資料規格相同)。

以下範例說明如何新增項目批次

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 陣列中，而不是 tf 變數中。此緩衝區可用於圖形外資料收集。在 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))