Make frame chooser.

  Args:
    real_env: T2TEnv to take initial frames from.
    frame_stack_size (int): Number of consecutive frames to extract.
    simulation_random_starts (bool): Whether to choose frames at random.
    simulation_flip_first_random_for_beginning (bool): Whether to flip the first
      frame stack in every batch for the frames at the beginning.
    split (tf.estimator.ModeKeys or None): Data split to take the frames from,
      None means use all frames.

  Returns:
    Function batch_size -> initial_frames.

def make_initial_frame_chooser(
    real_env, frame_stack_size, simulation_random_starts,
    simulation_flip_first_random_for_beginning,
    split=tf.estimator.ModeKeys.TRAIN,
):
  """Make frame chooser.

  Args:
    real_env: T2TEnv to take initial frames from.
    frame_stack_size (int): Number of consecutive frames to extract.
    simulation_random_starts (bool): Whether to choose frames at random.
    simulation_flip_first_random_for_beginning (bool): Whether to flip the first
      frame stack in every batch for the frames at the beginning.
    split (tf.estimator.ModeKeys or None): Data split to take the frames from,
      None means use all frames.

  Returns:
    Function batch_size -> initial_frames.
  """
  initial_frame_rollouts = real_env.current_epoch_rollouts(
      split=split, minimal_rollout_frames=frame_stack_size,
  )
  def initial_frame_chooser(batch_size):
    """Frame chooser."""

    deterministic_initial_frames =\
        initial_frame_rollouts[0][:frame_stack_size]
    if not simulation_random_starts:
      # Deterministic starts: repeat first frames from the first rollout.
      initial_frames = [deterministic_initial_frames] * batch_size
    else:
      # Random starts: choose random initial frames from random rollouts.
      initial_frames = random_rollout_subsequences(
          initial_frame_rollouts, batch_size, frame_stack_size
      )
      if simulation_flip_first_random_for_beginning:
        # Flip first entry in the batch for deterministic initial frames.
        initial_frames[0] = deterministic_initial_frames

    return np.stack([
        [frame.observation.decode() for frame in initial_frame_stack]  # pylint: disable=g-complex-comprehension
        for initial_frame_stack in initial_frames
    ])
  return initial_frame_chooser