code: def make_initial_frame_chooser(
      real_env, frame_stack_size, simulation_random_starts,
      simulation_flip_first_random_for_beginning,
      split=tf.estimator.ModeKeys.TRAIN,
  ):
    
    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
docs: 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.