Buffers de repetição

Ver no TensorFlow.org Executar no Google Colab Ver fonte no GitHub Baixar caderno

Introdução

Algoritmos de aprendizagem por reforço usam buffers de reprodução para armazenar trajetórias de experiência ao executar uma política em um ambiente. Durante o treinamento, os buffers de reprodução são consultados para um subconjunto das trajetórias (um subconjunto sequencial ou uma amostra) para "reproduzir" a experiência do agente.

Neste colab, exploramos dois tipos de buffers de reprodução: com base em python e com base em tensorflow, compartilhando uma API comum. Nas seções a seguir, descrevemos a API, cada uma das implementações de buffer e como usá-las durante o treinamento de coleta de dados.

Configurar

Instale tf-agents, se ainda não o fez.

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 Replay Buffer

A classe Replay Buffer tem a seguinte definição e 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"""

Note-se que quando o objecto tampão de repetição é inicializado, ele requer o data_spec dos elementos que vai armazenar. Esta especificação corresponde à TensorSpec de elementos de trajectória que vai ser adicionado ao tampão. Esta especificação é geralmente adquirido por olhando de um agente agent.collect_data_spec que define as formas, tipos e estruturas esperadas pelo agente quando a formação (mais sobre isso mais tarde).

TFUniformReplayBuffer

TFUniformReplayBuffer é o buffer de reprodução mais comumente usado em TF-agentes, portanto, vamos usá-lo em nosso tutorial aqui. Em TFUniformReplayBuffer o armazém intermediário de suporte é feito por variáveis tensorflow e assim faz parte do gráfico de computação.

As lojas de tampão lotes de elementos e tem um máximo de capacidade max_length elementos por segmento lote. Assim, a capacidade de tamponamento total é batch_size x max_length elementos. Todos os elementos armazenados no buffer devem ter uma especificação de dados correspondente. Quando o buffer de reprodução é usado para coleta de dados, a especificação é a especificação de dados de coleta do agente.

Criando o buffer:

Para criar um TFUniformReplayBuffer passamos em:

  1. a especificação dos elementos de dados que o buffer irá armazenar
  2. o batch size correspondente ao tamanho do lote do tampão
  3. o max_length número de elementos por segmento lote

Aqui é um exemplo de criar um TFUniformReplayBuffer com características de amostras de dados, batch_size 32 e 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)

Escrevendo no buffer:

Para adicionar elementos ao tampão de repetição, usamos o add_batch(items) método onde items é uma lista / tupla / ninho de tensores representando o lote de artigos a ser adicionado ao tampão. Cada elemento de items devem ter uma dimensão igual exterior batch_size e as restantes dimensões devem estar de acordo com a especificação de dados do item (o mesmo que as especificações de dados passados para o construtor tampão de repetição).

Aqui está um exemplo de adição de um lote de itens

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)

Lendo do buffer

Há três maneiras de ler os dados do TFUniformReplayBuffer :

  1. get_next() - retorna uma amostra a partir da memória intermédia. O tamanho do lote de amostra e o número de passos de tempo retornados podem ser especificados por meio de argumentos para este método.
  2. as_dataset() - devolve o tampão de repetição como um tf.data.Dataset . Pode-se então criar um iterador do conjunto de dados e iterar por meio das amostras dos itens no buffer.
  3. gather_all() - devolve todos os itens na memória intermédia como um tensor com a forma [batch, time, data_spec]

Abaixo estão alguns exemplos de como ler o buffer de reprodução usando cada um desses 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 tem o mesmo functionaly como o TFUniformReplayBuffer mas em vez de variáveis tf, seus dados são armazenados em matrizes numpy. Este buffer pode ser usado para coleta de dados fora do gráfico. Ter o armazenamento de apoio em numpy pode facilitar para alguns aplicativos a manipulação de dados (como indexação para atualizar prioridades) sem usar variáveis ​​do Tensorflow. No entanto, essa implementação não terá o benefício de otimizações de gráfico com Tensorflow.

Abaixo está um exemplo de instanciar um PyUniformReplayBuffer de especificações política trajetória do 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))

Usando buffers de repetição durante o treinamento

Agora que sabemos como criar um buffer de reprodução, gravar itens nele e ler a partir dele, podemos usá-lo para armazenar trajetórias durante o treinamento de nossos agentes.

Coleção de dados

Primeiro, vamos ver como usar o buffer de reprodução durante a coleta de dados.

Em TF-agentes usamos um Driver (veja o tutorial Driver para mais detalhes) a experiência de coleta em um ambiente. Para usar um Driver , nós especificar um Observer que é uma função para o Driver para executar quando recebe uma trajetória.

Assim, para adicionar elementos de trajetória para o buffer de repetição, nós adicionamos um observador que as chamadas add_batch(items) para adicionar um lote de itens no buffer de replay.

Abaixo está um exemplo desta com TFUniformReplayBuffer . Primeiro criamos um ambiente, uma rede e um agente. Então criamos um TFUniformReplayBuffer . Observe que as especificações dos elementos de trajetória no buffer de reprodução são iguais às especificações de dados coletados do agente. Em seguida, definir o seu add_batch método como o observador para o motorista que vai fazer a coleta de dados durante a nossa formação:

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

Leitura de dados para uma etapa de trem

Depois de adicionar elementos de trajetória ao buffer de reprodução, podemos ler lotes de trajetórias do buffer de reprodução para usar como dados de entrada para uma etapa do trem.

Aqui está um exemplo de como treinar nas trajetórias do buffer de reprodução em um loop de treinamento:

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