class_docstr: An episodic ReplayBuffer with Uniform sampling.
class_docstr_tok: ["An","episodic","ReplayBuffer","with","Uniform","sampling","."]
docstr: Extends a batch of episodes in this buffer.
  
  Args:
    episode_ids: A int64 vector containing the ids of the episodes the items
      are being added to.  Shaped `(max_num_episodes,)`.
    episode_ids_indices: An int64 vector containing the locations in
      `episode_ids` that are being extended.  Shaped `(num_episodes,)`, where
      `num_episodes <= max_num_episodes`.  Rows in `episodes` correspond to
      locations in `episode_ids_indices`.
    episodes: An `Episodes` tuple containing the extension data. Tensors must
      have outer dimension of size `num_episodes`.
  
  Returns:
    A `Tensor` containing the updated episode ids.  Accessing or executing
    this tensor also extends episodes in the replay buffer.
code: def extend_episodes(self, episode_ids, episode_ids_indices, episodes):
      """Extends a batch of episodes in this buffer.
  
      Args:
        episode_ids: A int64 vector containing the ids of the episodes the items
          are being added to.  Shaped `(max_num_episodes,)`.
        episode_ids_indices: An int64 vector containing the locations in
          `episode_ids` that are being extended.  Shaped `(num_episodes,)`, where
          `num_episodes <= max_num_episodes`.  Rows in `episodes` correspond to
          locations in `episode_ids_indices`.
        episodes: An `Episodes` tuple containing the extension data. Tensors must
          have outer dimension of size `num_episodes`.
  
      Returns:
        A `Tensor` containing the updated episode ids.  Accessing or executing
        this tensor also extends episodes in the replay buffer.
      """
      episode_ids.shape.assert_has_rank(1)
      episode_ids_indices = tf.convert_to_tensor(
          value=episode_ids_indices, name="episode_ids_indices"
      )
      episode_ids_indices.shape.assert_has_rank(1)
  
      def _extend_locked(episode_ids, expanded_episode_ids):
          """Does the above within the buffer's critical section."""
          episode_locations = self._get_episode_id_location(episode_ids)
          episode_valid = tf.equal(
              self._episodes_loc_to_id_map.sparse_read(
                  episode_locations
              ),
              episode_ids,
          )
          episode_valid_idx = tf.reshape(
              tf.compat.v1.where(episode_valid), [-1]
          )
          episode_locations = tf.gather(
              episode_locations, episode_valid_idx
          )
          increment_lengths = self._increment_episode_length_locked(
              episode_locations,
              tf.gather(episodes.length, episode_valid_idx),
          )
          set_completed = tf.compat.v1.scatter_update(
              self._episode_completed,
              episode_locations,
              tf.gather(episodes.completed, episode_valid_idx),
          )
          extend = self._data_table.extend(
              episode_locations,
              tf.nest.map_structure(
                  lambda tl: tf.gather(tl, episode_valid_idx),
                  episodes.tensor_lists,
              ),
          )
          with tf.control_dependencies(
              [increment_lengths, set_completed, extend]
          ):
              return tf.identity(expanded_episode_ids)
  
      with tf.device(self._device):
          with tf.name_scope("extend"):
              episode_ids_indices_shape = tf.shape(episode_ids_indices)
              begin_episode = self._begin_episode_fn(episodes)
              begin_episode = tf.broadcast_to(
                  begin_episode,
                  episode_ids_indices_shape,
                  name="begin_episode",
              )
              column_indices = tf.reshape(episode_ids_indices, [-1, 1])
              episode_ids_shape = tf.shape(input=episode_ids)
              # We expand the tensors below from size `num_episodes` (the size of
              # episode_ids_indices) to tensors of  size `max_num_episodes` (the size
              # of episode_ids).
              expanded_begin_episode = tf.scatter_nd(
                  column_indices, begin_episode, episode_ids_shape
              )
              expanded_mask = tf.scatter_nd(
                  column_indices,
                  tf.fill(episode_ids_indices_shape, True),
                  episode_ids_shape,
              )
              expanded_episode_ids = self._get_batch_episode_ids(
                  episode_ids,
                  begin_episode=expanded_begin_episode,
                  mask=expanded_mask,
              )
              episode_ids = tf.gather(
                  expanded_episode_ids, episode_ids_indices
              )
              return self._add_episode_critical_section.execute(
                  lambda: _extend_locked(
                      episode_ids, expanded_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 28 for calling the EpisodicReplayBuffer obj's underlying member method code with example usage: obj.extend_episodes(episode_ids, episode_ids_indices, episodes) and returns: tf, self