class_docstr: Wrapper enabling use of `EpisodicReplayBuffer` with a `Driver`.
  
  This wrapper keeps track of episode ids in a `tf.Variable`.
  
  Use:
  ```python
  tf_env = ...
  rb = EpisodicReplayBuffer(...)
  stateful_rb = StatefulEpisodicReplayBuffer(rb, num_episodes=tf_env.batch_size)
  driver = DynamicEpisodeDriver(
      tf_env, policy, ...,
      observers=[
        lambda traj: stateful_rb.add_batch(traj),
      ])
  driver.run()
  ```
class_docstr_tok: ["Wrapper","enabling","use","of","`","EpisodicReplayBuffer","`","with","a","`","Driver","`",".","This","wrapper","keeps","track","of","episode","ids","in","a","`","tf.Variable","`",".","Use",":","``","`","python","tf_env","=","...","rb","=","EpisodicReplayBuffer","(","...",")","stateful_rb","=","StatefulEpisodicReplayBuffer","(","rb",",","num_episodes=tf_env.batch_size",")","driver","=","DynamicEpisodeDriver","(","tf_env",",","policy",",","...",",","observers=","[","lambda","traj",":","stateful_rb.add_batch","(","traj",")",",","]",")","driver.run","(",")","``","`"]
docstr: Create a `StatefulEpisodicReplayBuffer` for `num_episodes` batches.
  
  Args:
    replay_buffer: An instance of `EpisodicReplayBuffer`.
    num_episodes: (Optional) integer, number of episode IDs to create. If
      `None`, a scalar ID variable is returned.
  
  Raises:
    TypeError: If `replay_buffer` is not an `EpisodicReplayBuffer`.
code: def __init__(self, replay_buffer, num_episodes=None):
      """Create a `StatefulEpisodicReplayBuffer` for `num_episodes` batches.
  
      Args:
        replay_buffer: An instance of `EpisodicReplayBuffer`.
        num_episodes: (Optional) integer, number of episode IDs to create. If
          `None`, a scalar ID variable is returned.
  
      Raises:
        TypeError: If `replay_buffer` is not an `EpisodicReplayBuffer`.
      """
      super(StatefulEpisodicReplayBuffer, self).__init__(
          replay_buffer.data_spec, replay_buffer.capacity
      )
  
      if not isinstance(replay_buffer, EpisodicReplayBuffer):
          raise TypeError(
              "Expected an EpisodicReplayBuffer, saw {}".format(
                  replay_buffer
              )
          )
      shape = ()
      if num_episodes and num_episodes > 0:
          if num_episodes > replay_buffer.capacity:
              raise ValueError(
                  "Buffer cannot create episode_ids when "
                  "num_episodes {} > capacity {}.".format(
                      num_episodes, replay_buffer.capacity
                  )
              )
          shape = (num_episodes,)
      self._replay_buffer = replay_buffer
      self._episode_ids_var = common.create_variable(
          "episode_id",
          initial_value=_INVALID_EPISODE_ID,
          shape=shape,
          use_local_variable=True,
      )
  
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 1 to create a new StatefulEpisodicReplayBuffer object from inherited base classes: replay_buffer_base with example: obj = StatefulEpisodicReplayBuffer(replay_buffer, num_episodes)