class_docstr: An episodic ReplayBuffer with Uniform sampling.
class_docstr_tok: ["An","episodic","ReplayBuffer","with","Uniform","sampling","."]
docstr: Adds a batch of single steps for the corresponding episodes IDs.
  
  Args:
    items: A batch of items to be added to the buffer. Items will have the
      same structure as the data_spec of this class, but the tensors in items
      will have an extra outer dimension `(num_episodes, ...)` in addition to
      the corresponding spec in data_spec.
    episode_ids: A int64 vector `Tensor` containing the ids of the episodes
      the items are being added to. Shaped `(num_episodes,)`.
  
  Returns:
    A `Tensor` containing the updated episode ids.  Accessing or executing
    this tensor also adds `items` to the replay buffer.
code: def add_batch(self, items, episode_ids):
      """Adds a batch of single steps for the corresponding episodes IDs.
  
      Args:
        items: A batch of items to be added to the buffer. Items will have the
          same structure as the data_spec of this class, but the tensors in items
          will have an extra outer dimension `(num_episodes, ...)` in addition to
          the corresponding spec in data_spec.
        episode_ids: A int64 vector `Tensor` containing the ids of the episodes
          the items are being added to. Shaped `(num_episodes,)`.
  
      Returns:
        A `Tensor` containing the updated episode ids.  Accessing or executing
        this tensor also adds `items` to the replay buffer.
      """
      episode_ids.shape.assert_has_rank(1)
      with tf.device(self._device):
          with tf.name_scope("add_batch"):
              # If begin_episode is True, then the increment of the episode_id happens
              # before trying to add anything to the buffer, regardless of whether the
              # item will actually be added.
  
              begin_episode = self._begin_episode_fn(items)
              end_episode = self._end_episode_fn(items)
              batch_episode_ids = self._get_batch_episode_ids(
                  episode_ids, begin_episode, end_episode
              )
              episodes_locations = tf.math.mod(
                  batch_episode_ids, self._capacity
              )
              # Accessing episode_id may modify self._episodes_loc_to_id_map, so
              # ensure it is executed before
              with tf.control_dependencies([episodes_locations]):
                  episode_valid = tf.equal(
                      self._episodes_loc_to_id_map.sparse_read(
                          episodes_locations
                      ),
                      batch_episode_ids,
                  )
  
              def _add_batch():
                  """Add elements to the appropiate episode_locations."""
                  ids_to_update = tf.reshape(
                      tf.compat.v1.where(episode_valid), [-1]
                  )
                  episodes_locations_ = tf.gather(
                      episodes_locations, ids_to_update
                  )
                  filter_items = lambda item: tf.gather(
                      item, ids_to_update
                  )
                  items_ = tf.nest.map_structure(filter_items, items)
                  write_data_op = self._data_table.add(
                      episodes_locations_, items_
                  )
                  inc_episode_lengths = (
                      self._increment_episode_length_locked(
                          episodes_locations_
                      )
                  )
                  inc_write_counter_op = self._num_writes.assign_add(1)
                  with tf.control_dependencies(
                      [
                          write_data_op,
                          inc_episode_lengths,
                          inc_write_counter_op,
                      ]
                  ):
                      return tf.identity(batch_episode_ids)
  
              num_adds = tf.reduce_sum(
                  input_tensor=tf.cast(episode_valid, tf.int64)
              )
  
              def _maybe_add_batch():
                  return self._add_episode_critical_section.execute(
                      _add_batch
                  )
  
              return tf.cond(
                  pred=num_adds > 0,
                  true_fn=_maybe_add_batch,
                  false_fn=lambda: episode_ids,
              )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  import gin
  import tensorflow
  from tf_agents import specs
  from tf_agents.replay_buffers import episodic_table
  from tf_agents.replay_buffers import replay_buffer
  from tf_agents.replay_buffers import table
  from tf_agents.utils import common
  from tensorflow.python.data.util import nest
  from tensorflow.python.layers import utils
  from tensorflow.python.ops import list_ops
  from tensorflow.python.ops.distributions import util
label_desc: Use image node_id 7 for calling the EpisodicReplayBuffer obj's underlying member method code with example usage: obj.add_batch(items, episode_ids) and returns: tf, self, tf