class_docstr: An episodic ReplayBuffer with Uniform sampling.
class_docstr_tok: ["An","episodic","ReplayBuffer","with","Uniform","sampling","."]
docstr: Increments the episode_id of the elements that have begin_episode True.
  
  Mark as completed those that have end_episode True.
  
  Args:
    batch_episode_ids: A tf.int64 tensor with shape `(num_episodes,)`
      containing a one or more episodes IDs.
    begin_episode: A Boolean Tensor whether should increment each episode ID.
      It can be a scalar or a vector with the same dimensions of
      batch_episode_ids.
    end_episode: A Boolean Tensor whether should end each episode ID. It can
      be a scalar or a vector with the same dimensions of batch_episode_ids.
    mask: An optional Boolean Tensor to select which IDs are updated.
  
  Returns:
    An Tensor shaped `(num_episodes,)` with the updated episode IDs.
  Raises:
    ValueError: If the shape of `begin_episode` is not compatible with
      `batch_episode_ids`.
code: def _get_batch_episode_ids(
      self,
      batch_episode_ids,
      begin_episode=False,
      end_episode=False,
      mask=None,
  ):
      """Increments the episode_id of the elements that have begin_episode True.
  
      Mark as completed those that have end_episode True.
  
      Args:
        batch_episode_ids: A tf.int64 tensor with shape `(num_episodes,)`
          containing a one or more episodes IDs.
        begin_episode: A Boolean Tensor whether should increment each episode ID.
          It can be a scalar or a vector with the same dimensions of
          batch_episode_ids.
        end_episode: A Boolean Tensor whether should end each episode ID. It can
          be a scalar or a vector with the same dimensions of batch_episode_ids.
        mask: An optional Boolean Tensor to select which IDs are updated.
  
      Returns:
        An Tensor shaped `(num_episodes,)` with the updated episode IDs.
      Raises:
        ValueError: If the shape of `begin_episode` is not compatible with
          `batch_episode_ids`.
      """
      if batch_episode_ids.shape.rank != 1:
          raise ValueError(
              "batch_episode_ids must be a vector with 1 dimension"
          )
      # If batch_episode_ids value is still -1 we need to assign a
      # proper value.  Find which IDs need to be updated.
      begin_episode = tf.convert_to_tensor(value=begin_episode)
      end_episode = tf.convert_to_tensor(value=end_episode)
      ids_to_update_mask = (batch_episode_ids < 0) | begin_episode
      if mask is not None:
          ids_to_update_mask &= mask
  
      def _update_batch_episode_ids():
          """Increment the episode_id inside a critical section."""
          num_ids = tf.reduce_sum(
              input_tensor=tf.cast(ids_to_update_mask, tf.int64)
          )
          end_id = self._last_episode.assign_add(num_ids).value() + 1
          start_id = end_id - num_ids
          new_batch_episode_ids = tf.range(start_id, end_id)
          # Update when b/74385543 is fixed.
          ids_to_update = tf.compat.v1.where(ids_to_update_mask)
          scattered_updated_episode_ids = tf.scatter_nd(
              ids_to_update,
              new_batch_episode_ids,
              shape=tf.shape(batch_episode_ids, out_type=tf.int64),
          )
          updated_batch_episode_ids = tf.compat.v1.where(
              ids_to_update_mask,
              scattered_updated_episode_ids,
              batch_episode_ids,
          )
          episode_locations = tf.math.mod(
              new_batch_episode_ids, self._capacity
          )
          update_mapping = tf.compat.v1.scatter_update(
              self._episodes_loc_to_id_map,
              [episode_locations],
              [new_batch_episode_ids],
          )
          reset_completed = tf.compat.v1.scatter_update(
              self._episode_completed, [episode_locations], 0
          )
          reset_data = self._data_table.clear_rows(episode_locations)
          reset_length = tf.compat.v1.scatter_update(
              self._episode_lengths, episode_locations, 0
          )
          with tf.control_dependencies(
              [
                  update_mapping,
                  reset_completed,
                  reset_data,
                  reset_length,
              ]
          ):
              return tf.identity(updated_batch_episode_ids)
  
      episode_ids = self._add_episode_critical_section.execute(
          _update_batch_episode_ids
      )
  
      # _maybe_end_batch_episodes acquires the critical section.
      mark_completed = self._maybe_end_batch_episodes(
          episode_ids, end_episode
      )
  
      with tf.control_dependencies([mark_completed]):
          episode_ids = tf.identity(episode_ids)
      return 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 20 for calling the EpisodicReplayBuffer obj's underlying member method code with example usage: obj._get_batch_episode_ids(batch_episode_ids, begin_episode, end_episode, mask) and returns: episode_ids, tf