replay_buffer.py

"""
Notes
-----

* For priority updates, see https://github.com/deepmind/reverb/issues/28
"""

import logging
import tensorflow as tf

logger = logging.getLogger(__name__)

try:
    import reverb
except ImportError as ex:
    logger.error(
        f'[replay_buffer] Could not import reverb: \n   {ex}   \n'
        + ' >>> If you have issues with libpython3.7m.so.1.0, try running: \n'
        + ' >>> $ export LD_LIBRARY_PATH=$CONDA_PREFIX/lib \n'
        + ' >>> in a conda environment, '
        + ' >>> or see https://github.com/deepmind/acme/issues/47 \n'
        + ' >>> See also https://stackoverflow.com/a/46833531 for how to set \n'
        + ' >>> LD_LIBRARY_PATH automatically when activating a conda environment.'
    )
    exit(1)


class ChunkWriter:
    """Wrapper for reverb's TrajectoryWriter"""

    def __init__(self, reverb_client, chunk_size, entries):
        self.writer = None
        self.chunk_size = chunk_size
        self.client = reverb_client
        self.total_items = 0
        self.entries = set(entries)

    def __enter__(self):
        self.writer = self.client.trajectory_writer(num_keep_alive_refs=self.chunk_size)
        return self

    def __exit__(self, *args, **kwargs):
        self.writer.__exit__(*args, **kwargs)

    def __del__(self):
        self.writer.__del__()

    def end_episode(self):
        return self.writer.end_episode()

    def append(self, *args, **kwargs):
        self.writer.append(*args, **kwargs)
        if self.writer.episode_steps >= self.chunk_size:
            trajectory = dict()
            for key in self.writer.history:
                if key not in self.entries:
                    raise RuntimeError(
                        'Cannot add to replay buffer an item that'
                        f' was not setup with setup_entry() method of ReplayBuffer: {key}')
                trajectory[key] = self.writer.history[key][-self.chunk_size:]
            self.writer.create_item(
                table='replay_buffer',
                priority=1.0,
                trajectory=trajectory)
            self.total_items += 1


class ReplayBuffer:
    """Defines an experience replay using reverb.

    Stores chunks of trajectories.

    TODO:
    * Sampling from different tables (priotirized, uniform etc.)
    * Update priorities.

    """

    def __init__(
            self,
            batch_size: int,
            chunk_size: int,
            max_replay_size: int,
    ):
        if chunk_size < 1:
            raise ValueError('chunk_size needs to be >= 1')

        self._batch_size = batch_size
        self._chunk_size = chunk_size
        self._max_replay_size = max_replay_size

        self._reverb_server = None
        self._reverb_client = None
        self._reverb_dataset = None
        self._batched_dataset = None
        self._chunk_writer = None
        self._signature = dict()

    @property
    def dataset(self):
        return self._batched_dataset

    def get_writer(self):
        self._chunk_writer = ChunkWriter(self._reverb_client, self._chunk_size, list(self._signature.keys()))
        return self._chunk_writer

    def sample(self):
        if self._chunk_writer is None:
            raise RuntimeError('Calling sample() without previous call to get_writer()')
        if self._chunk_writer.total_items < self._batch_size:
            return None
        return next(self.dataset)

    def setup_entry(self, name, shape, dtype):
        """
        Setup new entry in the replay buffer.

        Parameters
        ----------
        name : str
            Entry name.
        shape :
            Shape of the data. Can be nested.
        dtype :
            Type of the data. Can be nested.
        """
        if name in self._signature:
            raise ValueError(f'Entry {name} already added to the replay buffer.')

        self._signature[name] = tf.TensorSpec(
            shape=[self._chunk_size, *shape],
            dtype=dtype,
        )

    def build(self):
        """Creates reverb server, client and dataset."""
        self._reverb_server = reverb.Server(
            tables=[
                reverb.Table(
                    name='replay_buffer',
                    sampler=reverb.selectors.Uniform(),
                    remover=reverb.selectors.Fifo(),
                    max_size=self._max_replay_size,
                    rate_limiter=reverb.rate_limiters.MinSize(1),
                    signature=self._signature,
                ),
            ],
            port=None
        )
        self._reverb_client = reverb.Client(f'localhost:{self._reverb_server.port}')
        self._reverb_dataset = reverb.TrajectoryDataset.from_table_signature(
            server_address=f'localhost:{self._reverb_server.port}',
            table='replay_buffer',
            max_in_flight_samples_per_worker=2 * self._batch_size)
        self._batched_dataset = self._reverb_dataset.batch(self._batch_size, drop_remainder=True).as_numpy_iterator()
        # logger.info(self._reverb_client.server_info())