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: Extends a batch of episodes in this buffer.
  
  Args:
    episode_ids_indices: An int64 vector containing the locations in
      `self.episode_ids` that are being extended.  Shaped `(num_episodes,)`,
      where `num_episodes <= max_num_episodes`.  Rows in `episodes` and
      `begin_episode` 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 the replay buffer.
code: def extend_episodes(self, episode_ids_indices, episodes):
      """Extends a batch of episodes in this buffer.
  
      Args:
        episode_ids_indices: An int64 vector containing the locations in
          `self.episode_ids` that are being extended.  Shaped `(num_episodes,)`,
          where `num_episodes <= max_num_episodes`.  Rows in `episodes` and
          `begin_episode` 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 the replay buffer.
      """
      new_episode_ids = self._replay_buffer.extend_episodes(
          episode_ids=self._episode_ids_var,
          episode_ids_indices=episode_ids_indices,
          episodes=episodes,
      )
      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 5 for calling the StatefulEpisodicReplayBuffer obj's underlying member method code with example usage: obj.extend_episodes(episode_ids_indices, episodes) and returns: new_episode_ids