code: def _get_dataset_with_idxs(game, idxs, num_datasets):
      """Create a game dataset using splits based on idxs."""
  
      def _generate_data_splits(dataset_name, idxs):
          ds_builder = tfds.builder(dataset_name)
          split_lookup = list(ds_builder.info.splits.keys())
          data_splits = [f"{split_lookup[j]}[{k}]" for _, j, k in idxs]
          if not data_splits:
              return None
          return load_data_splits(dataset_name, data_splits)
  
      full_dataset = None
      for dataset_i in range(1, num_datasets + 1):
          # Create a dataset of episodes
          filtered_idxs = [idx for idx in idxs if idx[0] == dataset_i]
          dataset_name = (
              f"rlu_atari_checkpoints_ordered/{game}_run_{dataset_i}"
          )
          dataset = _generate_data_splits(
              dataset_name, list(filtered_idxs)
          )
          if dataset is None:
              continue
          if full_dataset is None:
              full_dataset = dataset
          else:
              full_dataset = full_dataset.concatenate(dataset)
      return full_dataset
  
file_path: dopamine/dopamine/labs/offline_rl/rlu_tfds/scaling_dataset_utils.py
imports: import concurrent.futures
  from dopamine.labs.offline_rl.rlu_tfds import tfds_atari_utils
  import gin
  import numpy
  import tensorflow_datasets
label_desc: Use image node_id 3 for calling a global function with example usage: _get_dataset_with_idxs(game, idxs, num_datasets) and returns: full_dataset, load_data_splits, None