Tampons de relecture

Copyright 2021 Les auteurs TF-Agents.

introduction

Les algorithmes d'apprentissage par renforcement utilisent des tampons de relecture pour stocker les trajectoires d'expérience lors de l'exécution d'une politique dans un environnement. Pendant la formation, les tampons de relecture sont interrogés pour un sous-ensemble des trajectoires (soit un sous-ensemble séquentiel, soit un échantillon) pour "rejouer" l'expérience de l'agent.

Dans cette collaboration, nous explorons deux types de tampons de relecture : basés sur Python et sur Tensorflow, partageant une API commune. Dans les sections suivantes, nous décrivons l'API, chacune des implémentations de tampon et comment les utiliser pendant la formation à la collecte de données.

Installer

Installez tf-agents si vous ne l'avez pas déjà fait.

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 tampon de relecture

La classe Replay Buffer a la définition et les méthodes suivantes :

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

Notez que lorsque l'objet de la mémoire tampon de lecture est initialisé, il faut l' data_spec des éléments qu'il va stocker. Cette spec correspond à la TensorSpec d'éléments de trajectoire qui seront ajoutés au tampon. Cette spécification est généralement acquise en regardant d'un agent agent.collect_data_spec qui définit les formes, les types et les structures attendues par l'agent lors de la formation (plus tard).

TFUniformReplayBuffer

TFUniformReplayBuffer est le plus couramment utilisé dans le tampon de lecture TF-agents, nous allons donc l' utiliser dans notre tutoriel ici. Dans TFUniformReplayBuffer la mémoire tampon de support se fait par des variables tensorflow et fait donc partie du graphe de calcul.

Les magasins tampons lots d'éléments et a une capacité maximale max_length éléments par segment discontinu. Ainsi, la capacité totale de la mémoire tampon est batch_size x max_length éléments. Les éléments stockés dans le tampon doivent tous avoir une spécification de données correspondante. Lorsque le tampon de relecture est utilisé pour la collecte de données, la spécification est la spécification de collecte de données de l'agent.

Création du tampon :

Pour créer un TFUniformReplayBuffer on passe:

1. la spécification des éléments de données que le tampon va stocker
2. la batch size correspondant à la taille du lot de la mémoire tampon
3. le max_length nombre d'éléments par segment discontinu

Voici un exemple de création d' un TFUniformReplayBuffer avec des échantillons des spécifications de données, batch_size 32 et 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)

Écriture dans le tampon :

Pour ajouter des éléments dans la mémoire tampon de lecture, nous utilisons le add_batch(items) méthode où les items est une liste / tuple / nid de tenseurs représentant le lot d'articles à ajouter au tampon. Chaque élément d' items doit avoir une dimension extérieure égale batch_size et les dimensions restantes doit être conforme à la spécification de données de l'élément ( le même que les spécifications de données transmises au constructeur de mémoire de reprise).

Voici un exemple d'ajout d'un lot d'articles

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)

Lecture du tampon

Il y a trois façons de lire les données du TFUniformReplayBuffer :

1. get_next() - Rendements un échantillon à partir du tampon. La taille du lot d'échantillons et le nombre de pas de temps renvoyés peuvent être spécifiés via les arguments de cette méthode.
2. as_dataset() - renvoie la mémoire de reprise en tant que tf.data.Dataset . On peut ensuite créer un itérateur d'ensemble de données et parcourir les échantillons des éléments du tampon.
3. gather_all() - retourne tous les éléments dans la mémoire tampon comme un tenseur de forme [batch, time, data_spec] le [batch, time, data_spec]

Vous trouverez ci-dessous des exemples de lecture à partir du tampon de relecture à l'aide de chacune de ces méthodes :

# 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 a le même functionaly que le TFUniformReplayBuffer mais au lieu de variables Tf, ses données sont stockées dans des tableaux numpy. Ce tampon peut être utilisé pour la collecte de données hors graphique. Le fait d'avoir le stockage de sauvegarde dans numpy peut permettre à certaines applications de manipuler plus facilement les données (telles que l'indexation pour la mise à jour des priorités) sans utiliser de variables Tensorflow. Cependant, cette implémentation ne bénéficiera pas des optimisations de graphes avec Tensorflow.

Voici un exemple d'instancier un PyUniformReplayBuffer de specs de la trajectoire de la politique de l'agent:

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

Utilisation des tampons de relecture pendant l'entraînement

Maintenant que nous savons créer un tampon de relecture, y écrire des éléments et y lire des éléments, nous pouvons l'utiliser pour stocker des trajectoires lors de la formation de nos agents.

Collecte de données

Voyons d'abord comment utiliser le tampon de relecture pendant la collecte de données.

Dans TF-agents , nous utilisons un Driver (voir le tutoriel du pilote pour plus de détails) à l' expérience Collect dans un environnement. Pour utiliser un Driver , nous spécifions un Observer qui est une fonction pour le Driver d'exécuter lorsqu'il reçoit une trajectoire.

Ainsi, pour ajouter des éléments de trajectoire dans la mémoire tampon de lecture, on ajoute un observateur que les appels add_batch(items) pour ajouter un lot d'articles sur la mémoire tampon de lecture.

Ci - dessous un exemple avec TFUniformReplayBuffer . Nous créons d'abord un environnement, un réseau et un agent. Ensuite , nous créons un TFUniformReplayBuffer . Notez que les spécifications des éléments de trajectoire dans le tampon de relecture sont égales aux spécifications de collecte de données de l'agent. Nous avons ensuite mis sa add_batch méthode que l'observateur pour le conducteur qui fera les données recueillent au cours de notre formation:

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

Lecture des données d'un pas de train

Après avoir ajouté des éléments de trajectoire au tampon de relecture, nous pouvons lire des lots de trajectoires à partir du tampon de relecture à utiliser comme données d'entrée pour une étape de train.

Voici un exemple de la façon de s'entraîner sur des trajectoires à partir du tampon de relecture dans une boucle d'entraînement :

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