code: def choose_indices(
      game,
      dataset_expertise,
      dataset_policies,
      num_datasets=5,
      dataset_episodes=10000,
      returns_filter=None,
  ):
      """Create indices to be used for dataset creation."""
      episodes_per_policy = dataset_episodes // dataset_policies
      (policy_value_lookup, policy_ep_idx_lookup) = _parallel_lookup(
          game,
          episodes_per_policy,
          num_datasets=num_datasets,
          max_checkpoints=dataset_policies,
      )
      policy_ids = list(policy_value_lookup.keys())
      num_ids = len(policy_ids)
      sorted_policy_ids = list(
          sorted(policy_ids, key=lambda pid: policy_value_lookup[pid])
      )
      inexpert_policy_ids = sorted_policy_ids[
          : int(dataset_expertise * num_ids)
      ]
      if returns_filter is not None:
          inexpert_policy_ids = inexpert_policy_ids[
              int(returns_filter * len(inexpert_policy_ids)) :
          ]
  
      if len(inexpert_policy_ids) < dataset_policies:
          selected_policy_ids = inexpert_policy_ids
      else:
          selected_policy_ids = [
              inexpert_policy_ids[int(i)]
              for i in np.linspace(
                  0, len(inexpert_policy_ids) - 1, dataset_policies
              )
          ]
      min_score, max_score = float("inf"), float("-inf")
      total_steps = 0
  
      all_episode_idxs = []
      for policy_id in selected_policy_ids:
          l, ep_lengths, ep_returns = policy_ep_idx_lookup[policy_id]
          x1 = episodes_per_policy // 2
          start_idx, end_idx = max(l // 2 - x1, 0), min(l // 2 + x1, l)
          all_episode_idxs.append(
              (policy_id[0], policy_id[1], f"{start_idx}:{end_idx}")
          )
          # Book keeping
          total_steps += sum(ep_lengths[start_idx:end_idx])
          for ep_return in ep_returns[start_idx:end_idx]:
              min_score = min(min_score, ep_return)
              max_score = max(max_score, ep_return)
      return all_episode_idxs, (total_steps, min_score, max_score)
  
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 2 for calling a global function with example usage: choose_indices(game, dataset_expertise, dataset_policies, num_datasets, dataset_episodes, returns_filter) and returns: all_episode_idxs,  total_steps,  min_score,  max_score