code: def _parallel_lookup(
      game, episodes_per_policy, max_checkpoints=50, num_datasets=5
  ):
      """Get idxs to use for dataset splits to be loaded."""
  
      def _get_lookups(i, j, ds_chkpt, episodes_per_policy):
          policy_value_lookup = {}
          policy_ep_idx_lookup = {}
          ds = ds_chkpt.map(
              lambda x: (x["episode_return"], len(x["steps"]))
          )
          info = list(ds.as_numpy_iterator())
          policy_id = (i, j)
          episode_returns = [ep_return for ep_return, _ in info]
          episode_lengths = [num_steps for _, num_steps in info]
          if len(info) < episodes_per_policy:
              print(
                  f"Warning! Not enough episodes in this chunk ({len(info)}) to"
                  f" support {episodes_per_policy} episodes per policy."
              )
          policy_value_lookup[policy_id] = np.mean(episode_returns)
          policy_ep_idx_lookup[policy_id] = (
              len(info),
              episode_lengths,
              episode_returns,
          )
          return policy_value_lookup, policy_ep_idx_lookup
  
      policy_value_lookup, policy_ep_idx_lookup = {}, {}
      with concurrent.futures.ThreadPoolExecutor() as executor:
          # Start the load operations and mark each future with its URL
          jobs = []
          for i in range(1, num_datasets + 1):
              dataset_name = (
                  f"rlu_atari_checkpoints_ordered/{game}_run_{i}"
              )
              ds = tfds.load(dataset_name)
              for j, chkpt in enumerate(ds):
                  # Use only the first 25 ckpts.
                  if int(chkpt.split("_")[1]) < max_checkpoints:
                      job = executor.submit(
                          _get_lookups,
                          i,
                          j,
                          ds[chkpt],
                          episodes_per_policy,
                      )
                      jobs.append(job)
          for future in concurrent.futures.as_completed(jobs):
              try:
                  pval, ep_idx = future.result()
                  policy_value_lookup.update(pval)
                  policy_ep_idx_lookup.update(ep_idx)
              except Exception as exc:  # pylint: disable=broad-except
                  print(f"Generated an exception: {exc}")
      return policy_value_lookup, policy_ep_idx_lookup
  
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 1 for calling a global function with example usage: _parallel_lookup(game, episodes_per_policy, max_checkpoints, num_datasets) and returns: policy_value_lookup,  policy_ep_idx_lookup, policy_value_lookup,  policy_ep_idx_lookup