code: def generate_data(self, data_dir, tmp_dir=None, task_id=-1):
      
      if not self._rollouts_by_epoch_and_split[self.current_epoch]:
        # Data not loaded from disk.
        self._split_current_epoch()
  
      rollouts_by_split = self._rollouts_by_epoch_and_split[self.current_epoch]
      splits_and_paths = self.splits_and_paths(data_dir)
  
      for (split, paths) in splits_and_paths:
        rollouts = rollouts_by_split[split]
        num_frames = self._calc_num_frames(rollouts)
        shard_size = num_frames // len(paths)
  
        frame_gen = self._generate_frames(rollouts)
        for (path_index, path) in enumerate(paths):
          limit = shard_size
          # Put the remainder in the last shard to preserve the ordering.
          if path_index == len(paths) - 1:
            limit = None
          generator_utils.generate_files(
              itertools.islice(frame_gen, limit), [path],
              cycle_every_n=float("inf")
          )
docs: Saves the current epoch rollouts to disk, split into train/dev sets.