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: 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.
  
  Returns:
    A `Tensor` containing the updated episode ids.  Accessing or executing
    this varaible also adds `items` to the replay buffer.
code: def add_batch(self, items):
      """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.
  
      Returns:
        A `Tensor` containing the updated episode ids.  Accessing or executing
        this varaible also adds `items` to the replay buffer.
      """
      new_episode_ids = self._replay_buffer.add_batch(
          items=items, episode_ids=self._episode_ids_var
      )
      self._episode_ids_var.assign(new_episode_ids)
      return new_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 3 for calling the StatefulEpisodicReplayBuffer obj's underlying member method code with example usage: obj.add_batch(items) and returns: new_episode_ids