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","(",")","``","`"]
code: def _as_dataset(
      self,
      sample_batch_size=None,
      num_steps=None,
      sequence_preprocess_fn=None,
      num_parallel_calls=None,
  ):
      return self._replay_buffer.as_dataset(
          sample_batch_size,
          num_steps,
          sequence_preprocess_fn=sequence_preprocess_fn,
          num_parallel_calls=num_parallel_calls,
      )
  
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 7 for calling the StatefulEpisodicReplayBuffer obj's underlying member method code with example usage: obj._as_dataset(sample_batch_size, num_steps, sequence_preprocess_fn, num_parallel_calls) and returns: self