<commit_before># coding=utf-8
# Copyright 2018 The TF-Agents Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Lint as: python3
"""Reverb as a TF-Agents replay buffer."""

from __future__ import absolute_import
from __future__ import division
# Using Type Annotations.
from __future__ import print_function

import gin
import tensorflow.compat.v2 as tf

from tf_agents.replay_buffers import replay_buffer
from tf_agents.utils import lazy_loader

# Lazy loading since not all users have the reverb package installed.
reverb = lazy_loader.LazyLoader('reverb', globals(), 'reverb')


# The default table name.
DEFAULT_TABLE = 'experience'


@gin.configurable
class ReverbReplayBuffer(replay_buffer.ReplayBuffer):
  """Reverb ReplayBuffer exposed as a TF-Agents replay buffer."""

  def __init__(self,
               data_spec,
               table_name,
               sequence_length,
               server_address=None,
               local_server=None,
               dataset_buffer_size=None,
               max_cycle_length=32,
               num_workers_per_iterator=-1,
               max_samples_per_stream=-1,
               rate_limiter_timeout_ms=-1,
               ):
    """Initializes a reverb replay buffer.

    *NOTE*: If the user calls `as_dataset()` with a value of `num_steps` that
    does not match the `sequence_length` used in the observer (or if
    `sequence_length is None` varies from episode to episode), then individual
    sequences will be truncated to the highest possible multiple of `num_steps`
    before being split. In this case, **data loss may occur**. For example, in
    sparse reward cases where there's a reward only at the final frame and the
    sequence length is not a multiple of `num_steps`, the reward will be
    dropped.

    When using sparse rewards at the end of sequences, the user should prefer
    to either request full episodes (e.g., call `as_dataset` with
    `num_steps=None` and `sample_batch_size=None` (for variable-length
    episodes)); or use `sequence_preprocess_fn` to first propagate the rewards
    from the end of the sequence through to the earlier frames using e.g.
    temporal discounting.

    Args:
      data_spec: Spec for the data held in the replay buffer.
      table_name: Name of the table that will be sampled.
      sequence_length: (can be set to `None`, i.e unknown) The number of
        timesteps that each sample consists of. If not `None`, then the lengths
        of samples received from the server will be validated against this
        number.

        **NOTE** This replay buffer will be at its most performant
        if the `sequence_length` here is equal to `num_steps` passed to
        `as_dataset`, and is also used when writing to the replay buffer
        (for example, see the `sequence_lengths` argument of the
        `Reverb.*Observer` classes).
      server_address: (Optional) Address of the reverb replay server. One of
        `server_address` or `local_server` must be provided.
      local_server: (Optional) An instance of `reverb.Server` that holds
        the replay's data.
      dataset_buffer_size: (Optional) This is the prefetch buffer size
        (in number of items) of the Reverb Dataset object.  A good rule of
        thumb is to set this value to 2-3x times the sample_batch_size you
        will use.
      max_cycle_length: (Optional) The number of sequences used to populate the
        batches of `as_dataset`.  By default, `min(32, sample_batch_size)` is
        used, but the number can be between `1` and `sample_batch_size`.
      num_workers_per_iterator: (Defaults to -1, i.e auto selected) The number
        of worker threads to create per dataset iterator. When the selected
        table uses a FIFO or Heap sampler (i.e a queue) then exactly 1 worker
        must be used to avoid races causing invalid ordering of items. For all
        other samplers, this value should be roughly equal to the number of
        threads available on the CPU.
      max_samples_per_stream: (Defaults to -1, i.e auto selected) The maximum
        number of samples to fetch from a stream before a new call is made.
        Keeping this number low ensures that the data is fetched uniformly from
        all servers.
      rate_limiter_timeout_ms: (Defaults to -1: infinite).  Timeout
        (in milliseconds) to wait on the rate limiter when sampling from the
        table. If `rate_limiter_timeout_ms >= 0`, this is the timeout passed to
        `Table::Sample` describing how long to wait for the rate limiter to
        allow sampling.
    """
    if (server_address is None) == (local_server is None):
      raise ValueError(
          'Exactly one of the server_address or local_server must be provided.')

    self._table_name = table_name
    self._sequence_length = sequence_length
    self._local_server = local_server
    self._server_address = server_address
    self._dataset_buffer_size = dataset_buffer_size
    self._max_cycle_length = max_cycle_length
    # TODO(b/156531956) Remove these
    self._num_workers_per_iterator = num_workers_per_iterator
    self._max_samples_per_stream = max_samples_per_stream
    self._rate_limiter_timeout_ms = rate_limiter_timeout_ms

    if local_server:
      self._server_address = 'localhost:{}'.format(local_server.port)
      self._py_client = self._local_server.in_process_client()
    else:
      self._py_client = reverb.Client(self._server_address)

    self._tf_client = reverb.TFClient(self._server_address, 'rb_tf_client')
    self._table_info = self.get_table_info()
    sampler = self._table_info.sampler_options
    remover = self._table_info.remover_options
    self._deterministic_table = (
        sampler.is_deterministic and remover.is_deterministic)

    capacity = self._table_info.max_size
    super(ReverbReplayBuffer, self).__init__(
        data_spec=data_spec, capacity=capacity, stateful_dataset=True)

  @property
  def py_client(self):
    return self._py_client

  @property
  def local_server(self):
    return self._local_server

  @property
  def tf_client(self):
    return self._tf_client

  def get_table_info(self):
    return self._py_client.server_info()[self._table_name]

  def _num_frames(self):
    """Returns the number of frames in the replay buffer.

    **Note:** This might return inconsistent result if there are pending<commit_msg>[TF-Agents] Never use reverb's in_process_client; it's only there for testing purposes.<commit_after># coding=utf-8
# Copyright 2018 The TF-Agents Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Lint as: python3
"""Reverb as a TF-Agents replay buffer."""

from __future__ import absolute_import
from __future__ import division
# Using Type Annotations.
from __future__ import print_function

import gin
import tensorflow.compat.v2 as tf

from tf_agents.replay_buffers import replay_buffer
from tf_agents.utils import lazy_loader

# Lazy loading since not all users have the reverb package installed.
reverb = lazy_loader.LazyLoader('reverb', globals(), 'reverb')


# The default table name.
DEFAULT_TABLE = 'experience'


@gin.configurable
class ReverbReplayBuffer(replay_buffer.ReplayBuffer):
  """Reverb ReplayBuffer exposed as a TF-Agents replay buffer."""

  def __init__(self,
               data_spec,
               table_name,
               sequence_length,
               server_address=None,
               local_server=None,
               dataset_buffer_size=None,
               max_cycle_length=32,
               num_workers_per_iterator=-1,
               max_samples_per_stream=-1,
               rate_limiter_timeout_ms=-1,
               ):
    """Initializes a reverb replay buffer.

    *NOTE*: If the user calls `as_dataset()` with a value of `num_steps` that
    does not match the `sequence_length` used in the observer (or if
    `sequence_length is None` varies from episode to episode), then individual
    sequences will be truncated to the highest possible multiple of `num_steps`
    before being split. In this case, **data loss may occur**. For example, in
    sparse reward cases where there's a reward only at the final frame and the
    sequence length is not a multiple of `num_steps`, the reward will be
    dropped.

    When using sparse rewards at the end of sequences, the user should prefer
    to either request full episodes (e.g., call `as_dataset` with
    `num_steps=None` and `sample_batch_size=None` (for variable-length
    episodes)); or use `sequence_preprocess_fn` to first propagate the rewards
    from the end of the sequence through to the earlier frames using e.g.
    temporal discounting.

    Args:
      data_spec: Spec for the data held in the replay buffer.
      table_name: Name of the table that will be sampled.
      sequence_length: (can be set to `None`, i.e unknown) The number of
        timesteps that each sample consists of. If not `None`, then the lengths
        of samples received from the server will be validated against this
        number.

        **NOTE** This replay buffer will be at its most performant
        if the `sequence_length` here is equal to `num_steps` passed to
        `as_dataset`, and is also used when writing to the replay buffer
        (for example, see the `sequence_lengths` argument of the
        `Reverb.*Observer` classes).
      server_address: (Optional) Address of the reverb replay server. One of
        `server_address` or `local_server` must be provided.
      local_server: (Optional) An instance of `reverb.Server` that holds
        the replay's data.
      dataset_buffer_size: (Optional) This is the prefetch buffer size
        (in number of items) of the Reverb Dataset object.  A good rule of
        thumb is to set this value to 2-3x times the sample_batch_size you
        will use.
      max_cycle_length: (Optional) The number of sequences used to populate the
        batches of `as_dataset`.  By default, `min(32, sample_batch_size)` is
        used, but the number can be between `1` and `sample_batch_size`.
      num_workers_per_iterator: (Defaults to -1, i.e auto selected) The number
        of worker threads to create per dataset iterator. When the selected
        table uses a FIFO or Heap sampler (i.e a queue) then exactly 1 worker
        must be used to avoid races causing invalid ordering of items. For all
        other samplers, this value should be roughly equal to the number of
        threads available on the CPU.
      max_samples_per_stream: (Defaults to -1, i.e auto selected) The maximum
        number of samples to fetch from a stream before a new call is made.
        Keeping this number low ensures that the data is fetched uniformly from
        all servers.
      rate_limiter_timeout_ms: (Defaults to -1: infinite).  Timeout
        (in milliseconds) to wait on the rate limiter when sampling from the
        table. If `rate_limiter_timeout_ms >= 0`, this is the timeout passed to
        `Table::Sample` describing how long to wait for the rate limiter to
        allow sampling.
    """
    if (server_address is None) == (local_server is None):
      raise ValueError(
          'Exactly one of the server_address or local_server must be provided.')

    self._table_name = table_name
    self._sequence_length = sequence_length
    self._local_server = local_server
    self._server_address = server_address
    self._dataset_buffer_size = dataset_buffer_size
    self._max_cycle_length = max_cycle_length
    # TODO(b/156531956) Remove these
    self._num_workers_per_iterator = num_workers_per_iterator
    self._max_samples_per_stream = max_samples_per_stream
    self._rate_limiter_timeout_ms = rate_limiter_timeout_ms

    if local_server:
      self._server_address = 'localhost:{}'.format(local_server.port)

    self._py_client = reverb.Client(self._server_address)
    self._tf_client = reverb.TFClient(self._server_address, 'rb_tf_client')
    self._table_info = self.get_table_info()
    sampler = self._table_info.sampler_options
    remover = self._table_info.remover_options
    self._deterministic_table = (
        sampler.is_deterministic and remover.is_deterministic)

    capacity = self._table_info.max_size
    super(ReverbReplayBuffer, self).__init__(
        data_spec=data_spec, capacity=capacity, stateful_dataset=True)

  @property
  def py_client(self):
    return self._py_client

  @property
  def local_server(self):
    return self._local_server

  @property
  def tf_client(self):
    return self._tf_client

  def get_table_info(self):
    return self._py_client.server_info()[self._table_name]

  def _num_frames(self):
    """Returns the number of frames in the replay buffer.

    **Note:** This might return inconsistent result if there are pending
{'fraction_non_alphanumeric': 0.06344125136413241, 'fraction_numerical': 0.005092761004001455, 'mean_word_length': 3.266294227188082, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5824252', 'n_tokens_mistral': 3872, 'n_tokens_neox': 3742, 'n_words': 1494}