code: def get_transition_dataset_fn(
      stack_size, update_horizon=1, gamma=0.99, return_to_go=False
  ):
      """Creates a dataset of (s, a, r, s', a') transitions."""
      batch_fn_cls = BatchToTransition(
          stack_size, update_horizon, gamma, return_to_go=return_to_go
      )
  
      def _add_rtg(episode_data):
          episode_data["return_to_go"] = tf.cumsum(
              episode_data[rlds.REWARD], reverse=True
          )
          return episode_data
  
      def make_transition_dataset(episode_data):
          """Converts an episode of steps to a dataset of custom transitions."""
          # Create a dataset of 2-step sequences with overlap of 1.
          episode = episode_data[rlds.STEPS]
          if return_to_go:
              episode = episode.map(_add_rtg)
          batched_steps = rlds.transformations.batch(
              episode,
              size=stack_size + update_horizon,
              shift=1,
              drop_remainder=True,
          )
          # pylint: disable=g-long-lambda
          batch_fn = lambda x: batch_fn_cls.create_transitions(
              x, episode_return=episode_data["episode_return"]
          )
          # pylint: enable=g-long-lambda
          return batched_steps.map(
              batch_fn, num_parallel_calls=tf.data.AUTOTUNE
          )
  
      return make_transition_dataset
  
file_path: dopamine/dopamine/labs/offline_rl/rlu_tfds/tfds_atari_utils.py
imports: import rlds
  import tensorflow
  import tensorflow_datasets
label_desc: Use image node_id 1 for calling a global function with example usage: get_transition_dataset_fn(stack_size, update_horizon, gamma, return_to_go) and returns: make_transition_dataset, episode_data, batched_steps