class_docstr: A simple out-of-graph Replay Buffer.
  
  Stores transitions, state, action, reward, next_state, terminal (and any
  extra contents specified) in a circular buffer and provides a uniform
  transition sampling function.
  
  When the states consist of stacks of observations storing the states is
  inefficient. This class writes observations and constructs the stacked states
  at sample time.
  
  Attributes:
    add_count: int, counter of how many transitions have been added (including
      the blank ones at the beginning of an episode).
    invalid_range: np.array, an array with the indices of cursor-related invalid
      transitions
    episode_end_indices: set[int], a set of indices corresponding to the
      end of an episode.
docstr: Returns a batch of transitions (including any extra contents).
  
  If get_transition_elements has been overridden and defines elements not
  stored in self._store, an empty array will be returned and it will be
  left to the child class to fill it. For example, for the child class
  OutOfGraphPrioritizedReplayBuffer, the contents of the
  sampling_probabilities are stored separately in a sum tree.
  
  When the transition is terminal next_state_batch has undefined contents.
  
  NOTE: This transition contains the indices of the sampled elements. These
  are only valid during the call to sample_transition_batch, i.e. they may
  be used by subclasses of this replay buffer but may point to different data
  as soon as sampling is done.
  
  Args:
    batch_size: int, number of transitions returned. If None, the default
      batch_size will be used.
    indices: None or list of ints, the indices of every transition in the
      batch. If None, sample the indices uniformly.
  
  Returns:
    transition_batch: tuple of np.arrays with the shape and type as in
      get_transition_elements().
  
  Raises:
    ValueError: If an element to be sampled is missing from the replay buffer.
docstr_tok: ["Returns","a","batch","of","transitions","(","including","any","extra","contents",")",".","If","get_transition_elements","has","been","overridden","and","defines","elements","not","stored","in","self._store",",","an","empty","array","will","be","returned","and","it","will","be","left","to","the","child","class","to","fill","it",".","For","example",",","for","the","child","class","OutOfGraphPrioritizedReplayBuffer",",","the","contents","of","the","sampling_probabilities","are","stored","separately","in","a","sum","tree",".","When","the","transition","is","terminal","next_state_batch","has","undefined","contents",".","NOTE",":","This","transition","contains","the","indices","of","the","sampled","elements",".","These","are","only","valid","during","the","call","to","sample_transition_batch",",","i.e",".","they","may","be","used","by","subclasses","of","this","replay","buffer","but","may","point","to","different","data","as","soon","as","sampling","is","done",".","Args",":","batch_size",":","int",",","number","of","transitions","returned",".","If","None",",","the","default","batch_size","will","be","used",".","indices",":","None","or","list","of","ints",",","the","indices","of","every","transition","in","the","batch",".","If","None",",","sample","the","indices","uniformly",".","Returns",":","transition_batch",":","tuple","of","np.arrays","with","the","shape","and","type","as","in","get_transition_elements","(",")",".","Raises",":","ValueError",":","If","an","element","to","be","sampled","is","missing","from","the","replay","buffer","."]
code: def sample_transition_batch(self, batch_size=None, indices=None):
      """Returns a batch of transitions (including any extra contents).
  
      If get_transition_elements has been overridden and defines elements not
      stored in self._store, an empty array will be returned and it will be
      left to the child class to fill it. For example, for the child class
      OutOfGraphPrioritizedReplayBuffer, the contents of the
      sampling_probabilities are stored separately in a sum tree.
  
      When the transition is terminal next_state_batch has undefined contents.
  
      NOTE: This transition contains the indices of the sampled elements. These
      are only valid during the call to sample_transition_batch, i.e. they may
      be used by subclasses of this replay buffer but may point to different data
      as soon as sampling is done.
  
      Args:
        batch_size: int, number of transitions returned. If None, the default
          batch_size will be used.
        indices: None or list of ints, the indices of every transition in the
          batch. If None, sample the indices uniformly.
  
      Returns:
        transition_batch: tuple of np.arrays with the shape and type as in
          get_transition_elements().
  
      Raises:
        ValueError: If an element to be sampled is missing from the replay buffer.
      """
      if batch_size is None:
          batch_size = self._batch_size
      if indices is None:
          indices = self.sample_index_batch(batch_size)
      assert len(indices) == batch_size
  
      transition_elements = self.get_transition_elements(batch_size)
      batch_arrays = self._create_batch_arrays(batch_size)
      for batch_element, state_index in enumerate(indices):
          trajectory_indices = [
              (state_index + j) % self._replay_capacity
              for j in range(self._update_horizon)
          ]
          trajectory_terminals = self._store["terminal"][
              trajectory_indices
          ]
          is_terminal_transition = trajectory_terminals.any()
          if not is_terminal_transition:
              trajectory_length = self._update_horizon
          else:
              # np.argmax of a bool array returns the index of the first True.
              trajectory_length = (
                  np.argmax(trajectory_terminals.astype(bool), 0) + 1
              )
          next_state_index = state_index + trajectory_length
          trajectory_discount_vector = self._cumulative_discount_vector[
              :trajectory_length
          ]
          trajectory_rewards = self.get_range(
              self._store["reward"], state_index, next_state_index
          )
  
          # Fill the contents of each array in the sampled batch.
          assert len(transition_elements) == len(batch_arrays)
          for element_array, element in zip(
              batch_arrays, transition_elements
          ):
              if element.name == "state":
                  element_array[
                      batch_element
                  ] = self.get_observation_stack(state_index)
              elif element.name == "reward":
                  # compute the discounted sum of rewards in the trajectory.
                  element_array[batch_element] = np.sum(
                      trajectory_discount_vector * trajectory_rewards,
                      axis=0,
                  )
              elif element.name == "next_state":
                  element_array[
                      batch_element
                  ] = self.get_observation_stack(
                      (next_state_index) % self._replay_capacity
                  )
              elif element.name in ("next_action", "next_reward"):
                  element_array[batch_element] = self._store[
                      element.name.lstrip("next_")
                  ][(next_state_index) % self._replay_capacity]
              elif element.name == "terminal":
                  element_array[batch_element] = is_terminal_transition
              elif element.name == "indices":
                  element_array[batch_element] = state_index
              elif element.name in self._store.keys():
                  element_array[batch_element] = self._store[
                      element.name
                  ][state_index]
              # We assume the other elements are filled in by the subclass.
  
      return batch_arrays
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  import gzip
  import math
  import os
  import pickle
  from absl import logging
  import gin.tf
  import numpy
  import tensorflow
label_desc: Use image node_id 22 for calling the OutOfGraphReplayBuffer obj's underlying member method code with example usage: obj.sample_transition_batch(batch_size, indices) and returns: batch_arrays