code: def create_ds_iterator(ds, batch_size=32, repeat=True):
      """Create numpy iterator from a tf dataset `ds`."""
      if repeat:
          ds = ds.repeat()
      batch_ds = ds.batch(batch_size).prefetch(tf.data.AUTOTUNE)
      return batch_ds.as_numpy_iterator()
  
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 5 for calling a global function with example usage: create_ds_iterator(ds, batch_size, repeat) and returns: batch_ds