code: def create_atari_ds_loader(
      dataset,
      transition_fn=None,
      shuffle_num_episodes=1000,
      shuffle_num_steps=50000,
      cycle_length=100,
  ):
      """Create uniformly subsampled Atari `game` dataset."""
      if transition_fn is None:
          transition_fn = get_transition_dataset_fn(4)
      # Shuffle the episodes to avoid consecutive episodes
      dataset = dataset.shuffle(shuffle_num_episodes)
      # Interleave the steps across many different episodes
      dataset = dataset.interleave(
          transition_fn,
          cycle_length=cycle_length,
          block_length=1,
          deterministic=False,
          num_parallel_calls=tf.data.AUTOTUNE,
      )
      # Shuffle steps in the dataset
      shuffled_dataset = dataset.shuffle(
          shuffle_num_steps, reshuffle_each_iteration=True
      )
      return shuffled_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 4 for calling a global function with example usage: create_atari_ds_loader(dataset, transition_fn, shuffle_num_episodes, shuffle_num_steps, cycle_length) and returns: shuffled_dataset