code: def _batch_steps(episode: rlds.Episode) -> tf.data.Dataset:
      return rlds.transformations.batch(
          episode[rlds.STEPS], size=2, shift=1, drop_remainder=True
      )
  
imports: import logging
  from typing import Any, Iterator, Optional, Tuple, Sequence
  from acme import specs
  from acme import types
  from flax import jax_utils
  import jax
  import jax.numpy
  import numpy
  import rlds
  import tensorflow
  import tensorflow_datasets
label_desc: Use image node_id 2 for calling a global function with example usage: _batch_steps(episode) and returns: rlds