code: def episodes_to_timestep_batched_transitions(
      episode_dataset: tf.data.Dataset,
      return_horizon: int = 10,
      drop_return_horizon: bool = False,
      min_return_filter: Optional[float] = None,
  ) -> tf.data.Dataset:
      """Process an existing dataset converting it to episode to 3-transitions.
  
      A 3-transition is an Transition with each attribute having an extra dimension
      of size 3, representing 3 consecutive timesteps. Each 3-step object will be
      in random order relative to each other.  See `episode_to_timestep_batch` for
      more information.
  
      Args:
        episode_dataset: An RLDS dataset to process.
        return_horizon: The horizon we want calculate Monte-Carlo returns to.
        drop_return_horizon: Whether we should drop the last `return_horizon` steps.
        min_return_filter: Minimum episode return below which we drop an episode.
  
      Returns:
        A tf.data.Dataset of 3-transitions.
      """
      dataset = episode_dataset.interleave(
          functools.partial(
              episode_to_timestep_batch,
              return_horizon=return_horizon,
              drop_return_horizon=drop_return_horizon,
              calculate_episode_return=min_return_filter is not None,
          ),
          num_parallel_calls=tf.data.experimental.AUTOTUNE,
          deterministic=False,
      )
  
      if min_return_filter is not None:
  
          def filter_on_return(step):
              return step[EPISODE_RETURN][0][0] > min_return_filter
  
          dataset = dataset.filter(filter_on_return)
  
      dataset = dataset.map(
          _step_to_transition,
          num_parallel_calls=tf.data.experimental.AUTOTUNE,
      )
  
      return dataset
  
file_path: acme/acme/agents/jax/mbop/dataset.py
imports: import functools
  import itertools
  from typing import Iterator, Optional
  from acme import types
  from acme.jax import running_statistics
  import jax
  import jax.numpy
  import rlds
  import tensorflow
  import tree
label_desc: Use image node_id 6 for calling a global function with example usage: episodes_to_timestep_batched_transitions(episode_dataset, return_horizon, drop_return_horizon, min_return_filter) and returns: dataset, unknown