code: def episode_to_timestep_batch(
      episode: rlds.BatchedStep,
      return_horizon: int = 0,
      drop_return_horizon: bool = False,
      calculate_episode_return: bool = False,
  ) -> tf.data.Dataset:
      """Converts an episode into multi-timestep batches.
  
      Args:
        episode: Batched steps as provided directly by RLDS.
        return_horizon: int describing the horizon to which we should accumulate the
          return.
        drop_return_horizon: bool whether we should drop the last `return_horizon`
          steps to avoid mis-calculated returns near the end of the episode.
        calculate_episode_return: Whether to calculate episode return.  Can be an
          expensive operation on datasets with many episodes.
  
      Returns:
        rl_dataset.DatasetType of 3-batched transitions, with scalar rewards
          expanded to 1D rewards
  
      This means that for every step, the corresponding elements will be a batch of
      size 3, with the first batched element corresponding to *_t-1, the second to
      *_t and the third to *_t+1,  e.g. you can access the previous observation as:
      ```
      o_tm1 = el[types.OBSERVATION][0]
      ```
      Two additional keys can be added: 'R_t' which corresponds to the undiscounted
      return for horizon `return_horizon` from time t (always present), and
      'R_total' which corresponds to the total return of the associated episode (if
      `calculate_episode_return` is True). Rewards are converted to be (at least)
      one-dimensional, prior to batching (to avoid ()-shaped elements).
  
      In this example, 0-valued observations correspond to o_{t-1}, 1-valued
      observations correspond to o_t, and 2-valued observations correspond to
      s_{t+1}.  This same structure is true for all keys, except 'R_t' and 'R_total'
      which are both scalars.
      ```
      ipdb> el[types.OBSERVATION]
      <tf.Tensor: shape=(3, 11), dtype=float32, numpy=
      array([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
           [1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
           [2., 2., 2., 2., 2., 2., 2., 2., 2., 2., 2.]], dtype=float32)>
      ```
      """
      steps = episode[rlds.STEPS]
  
      if drop_return_horizon:
          episode_length = steps.cardinality()
          steps = steps.take(episode_length - return_horizon)
  
      # Calculate n-step return:
      rewards = steps.map(lambda step: step[rlds.REWARD])
      batched_rewards = rlds.transformations.batch(
          rewards,
          size=return_horizon,
          shift=1,
          stride=1,
          drop_remainder=True,
      )
      returns = batched_rewards.map(tf.math.reduce_sum)
      output = tf.data.Dataset.zip((steps, returns)).map(
          _append_n_step_return
      )
  
      # Calculate total episode return for potential filtering, use total # of steps
      # to calculate return.
      if calculate_episode_return:
          dtype = (
              jnp.float64 if jax.config.jax_enable_x64 else jnp.float32
          )
          # Need to redefine this here to avoid a tf.data crash.
          rewards = steps.map(lambda step: step[rlds.REWARD])
          episode_return = rewards.reduce(dtype(0), lambda x, y: x + y)
          output = output.map(
              functools.partial(
                  _append_episode_return, episode_return=episode_return
              )
          )
  
      output = output.map(_expand_scalars)
  
      output = rlds.transformations.batch(
          output, size=3, shift=1, drop_remainder=True
      )
      return output
  
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 4 for calling a global function with example usage: episode_to_timestep_batch(episode, return_horizon, drop_return_horizon, calculate_episode_return) and returns: output