Búferes de reproducción

Ver en TensorFlow.org Ejecutar en Google Colab Ver fuente en GitHub Descargar cuaderno

Introducción

Los algoritmos de aprendizaje por refuerzo utilizan búferes de reproducción para almacenar trayectorias de experiencia al ejecutar una política en un entorno. Durante el entrenamiento, se consultan los búferes de reproducción para un subconjunto de las trayectorias (ya sea un subconjunto secuencial o una muestra) para "reproducir" la experiencia del agente.

En esta colab, exploramos dos tipos de búferes de reproducción: respaldados por python y respaldados por tensorflow, que comparten una API común. En las siguientes secciones, describimos la API, cada una de las implementaciones de búfer y cómo usarlas durante el entrenamiento de recolección de datos.

Configuración

Instale tf-agents si aún no lo ha hecho.

pip install tf-agents

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 de búfer de reproducción

La clase Replay Buffer tiene la siguiente definición y métodos:

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"""

Tenga en cuenta que cuando se inicializa el objeto búfer de reproducción, se requiere la data_spec de los elementos que va a almacenar. Esta especificación corresponde a la TensorSpec de elementos de trayecto que se añadirán a la memoria intermedia. Esta especificación se adquiere generalmente examinado de un agente agent.collect_data_spec que define las formas, tipos y estructuras esperadas por el agente cuando la formación (más sobre esto más adelante).

TFUniformReplayBuffer

TFUniformReplayBuffer es el búfer de repetición más comúnmente utilizado en la carretera TF-agentes, por lo tanto vamos a utilizar en nuestro tutorial aquí. En TFUniformReplayBuffer el depósito tampón de respaldo se realiza por variables tensorflow y por lo tanto es parte de la gráfica de cálculo.

Las memorias intermedias lotes de elementos y tiene un máximo de capacidad max_length elementos por segmento lote. Por lo tanto, la capacidad de memoria tampón total es batch_size x max_length elementos. Todos los elementos almacenados en el búfer deben tener una especificación de datos coincidente. Cuando el búfer de reproducción se utiliza para la recopilación de datos, la especificación es la especificación de recopilación de datos del agente.

Creando el búfer:

Para crear un TFUniformReplayBuffer pasamos en:

  1. la especificación de los elementos de datos que almacenará el búfer
  2. el batch size correspondiente al tamaño de lote de la memoria intermedia
  3. la max_length número de elementos por segmento lote

Este es un ejemplo de creación de una TFUniformReplayBuffer con especificaciones de datos de muestra, batch_size 32 y max_length 1000.

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)

Escribiendo en el búfer:

Para añadir elementos a la memoria intermedia de reproducción, utilizamos el add_batch(items) método en el que items es una lista / tupla / nido de tensores que representan el lote de artículos que se añade a la memoria intermedia. Cada elemento de items debe tener una dimensión igual exterior batch_size y las dimensiones restantes debe adherirse a la especificación de datos del elemento (igual que las especificaciones de datos pasados a la memoria intermedia constructor replay).

A continuación, se muestra un ejemplo de cómo agregar un lote de artículos.

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)

Leyendo del búfer

Hay tres maneras de leer los datos de la TFUniformReplayBuffer :

  1. get_next() - devuelve uno de la muestra de la memoria intermedia. El tamaño del lote de muestra y el número de pasos de tiempo devueltos se pueden especificar mediante argumentos para este método.
  2. as_dataset() - devuelve el búfer de reproducción como un tf.data.Dataset . Luego, se puede crear un iterador de conjunto de datos e iterar a través de las muestras de los elementos en el búfer.
  3. gather_all() - devuelve todos los artículos en el búfer como un tensor con forma de [batch, time, data_spec]

A continuación se muestran ejemplos de cómo leer desde el búfer de reproducción utilizando cada uno de estos métodos:

# 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 /tmp/ipykernel_15476/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.
WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/data/experimental/ops/counter.py:66: scan (from tensorflow.python.data.experimental.ops.scan_ops) is deprecated and will be removed in a future version.
Instructions for updating:
Use `tf.data.Dataset.scan(...) 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 /tmp/ipykernel_15476/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 tiene el mismo functionaly como el TFUniformReplayBuffer pero en lugar de las variables del TF, sus datos se almacenan en matrices numpy. Este búfer se puede utilizar para la recopilación de datos fuera del gráfico. Tener el almacenamiento de respaldo en numpy puede facilitar que algunas aplicaciones realicen la manipulación de datos (como la indexación para actualizar las prioridades) sin usar variables de Tensorflow. Sin embargo, esta implementación no tendrá el beneficio de las optimizaciones de gráficos con Tensorflow.

A continuación se muestra un ejemplo de instancias de un PyUniformReplayBuffer de las especificaciones de la trayectoria política del agente:

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

Usar búferes de reproducción durante el entrenamiento

Ahora que sabemos cómo crear un búfer de reproducción, escribir elementos en él y leer de él, podemos usarlo para almacenar trayectorias durante el entrenamiento de nuestros agentes.

Recopilación de datos

Primero, veamos cómo usar el búfer de reproducción durante la recolección de datos.

En la carretera TF-Agentes usamos un Driver (ver el tutorial del conductor para más detalles) para recoger la experiencia en un entorno. Para utilizar un Driver , especificamos un Observer que es una función para el Driver para ejecutar cuando se recibe una trayectoria.

Por lo tanto, para agregar elementos de trayecto en la memoria intermedia de repetición, añadimos un observador que las llamadas add_batch(items) para añadir un lote de artículos en el búfer de reproducción.

A continuación se muestra un ejemplo de esto con TFUniformReplayBuffer . Primero creamos un entorno, una red y un agente. Entonces se crea una TFUniformReplayBuffer . Tenga en cuenta que las especificaciones de los elementos de trayectoria en el búfer de reproducción son iguales a las especificaciones de recopilación de datos del agente. A continuación, establecemos su add_batch método que el observador para el conductor que va a hacer recogen los datos durante nuestro entrenamiento:

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

Leer datos para un paso de tren

Después de agregar elementos de trayectoria al búfer de reproducción, podemos leer lotes de trayectorias del búfer de reproducción para usarlos como datos de entrada para un paso de tren.

A continuación, se muestra un ejemplo de cómo entrenar en trayectorias desde el búfer de reproducción en un ciclo de entrenamiento:

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