def load_from_directory(trajectory_dir,
                          epoch=None,
                          temperature=None,
                          n_trajectories=None,
                          up_sample=False,
                          sleep_time_secs=0.1,
                          max_tries=100,
                          wait_forever=False):
    """Load trajectories from specified dir and epoch.

    Args:
      trajectory_dir: (string) directory to find trajectories.
      epoch: (int) epoch for which to load trajectories, if None we don't filter
        on an epoch.
      temperature: (float) this is used to filter the trajectory files, if None
        we don't filter on temperature.
      n_trajectories: (int) This is the batch size of the returned
        BatchTrajectory object if one is returned. If set to None, then the
        number of trajectories becomes the batch size. If set to some number,
        then we wait for those many trajectory files to be available.
      up_sample: (bool) If there are fewer than required (n_trajectories) number
        of incomplete trajectories, then we upsample to make up the numbers.
      sleep_time_secs: (float) Sleep time, to wait for min_trajectories. We
        exponentially back-off this up till a maximum of 10 seconds.
      max_tries: (int) The number of tries to get min_trajectories trajectories.
      wait_forever: (bool) If true, overrides max_tries and waits forever.

    Returns:
      A BatchTrajectory object with all the constraints satisfied or None.
    """

    # Modify the format to get a glob with desired epoch and temperature.
    trajectory_file_glob = TRAJECTORY_FILE_FORMAT.format(
        epoch=epoch if epoch is not None else "*",
        env_id="*",
        temperature=temperature if temperature is not None else "*",
        r="*",
    )

    trajectory_files = tf.io.gfile.glob(
        os.path.join(trajectory_dir, trajectory_file_glob))

    if n_trajectories:
      # We need to get `n_trajectories` number of `trajectory_files`.
      # This works out to a maximum ~3hr waiting period.
      while ((max_tries > 0 or wait_forever) and
             len(trajectory_files) < n_trajectories):
        logging.info(
            "Sleeping for %s seconds while waiting for %s trajectories, found "
            "%s right now.", sleep_time_secs, n_trajectories,
            len(trajectory_files))
        time.sleep(sleep_time_secs)
        max_tries -= 1
        sleep_time_secs = min(10.0, sleep_time_secs * 2)
        trajectory_files = tf.io.gfile.glob(
            os.path.join(trajectory_dir, trajectory_file_glob))

      # We can't get the required number of files and we can't up-sample either.
      if (len(trajectory_files) < n_trajectories) and not up_sample:
        return None

      # Sample up or down as the case maybe.
      trajectory_files = list(
          np.random.choice(trajectory_files, n_trajectories))

    # We read and load all the files, revisit if this becomes a problem.
    trajectories_buffer = []
    for trajectory_file in trajectory_files:
      with tf.io.gfile.GFile(trajectory_file, "rb") as f:
        trajectory = get_pickle_module().load(f)
        assert isinstance(trajectory, Trajectory)
        trajectories_buffer.append(trajectory)

    if not trajectories_buffer:
      return None

    # If n_trajectories wasn't set, then set to the number of trajectories we're
    # returning.
    n_trajectories = n_trajectories or len(trajectories_buffer)

    # Construct and return a new BatchTrajectory object.
    return BatchTrajectory(
        batch_size=n_trajectories,
        trajectories=[Trajectory() for _ in range(n_trajectories)],
        completed_trajectories=trajectories_buffer)
Load trajectories from specified dir and epoch.

    Args:
      trajectory_dir: (string) directory to find trajectories.
      epoch: (int) epoch for which to load trajectories, if None we don't filter
        on an epoch.
      temperature: (float) this is used to filter the trajectory files, if None
        we don't filter on temperature.
      n_trajectories: (int) This is the batch size of the returned
        BatchTrajectory object if one is returned. If set to None, then the
        number of trajectories becomes the batch size. If set to some number,
        then we wait for those many trajectory files to be available.
      up_sample: (bool) If there are fewer than required (n_trajectories) number
        of incomplete trajectories, then we upsample to make up the numbers.
      sleep_time_secs: (float) Sleep time, to wait for min_trajectories. We
        exponentially back-off this up till a maximum of 10 seconds.
      max_tries: (int) The number of tries to get min_trajectories trajectories.
      wait_forever: (bool) If true, overrides max_tries and waits forever.

    Returns:
      A BatchTrajectory object with all the constraints satisfied or None.
    
tensor2tensor/envs/trajectory.py