def __init__(self, real_env, world_model_dir, hparams, random_starts,
               setable_initial_frames=False):
    """Init.

    Args:
       real_env: gym environment.
       world_model_dir: path to world model checkpoint directory.
       hparams: hparams for rlmb pipeline.
       random_starts: if restart world model from random frames, or only
         from initial ones (from beginning of episodes). Valid only when
         `setable_initial_fames` set to False.
       setable_initial_frames: if True, initial_frames for world model should be
         set by `add_to_initial_stack`.
    """

    self._setable_initial_frames = setable_initial_frames

    if self._setable_initial_frames:
      real_obs_shape = real_env.observation_space.shape
      shape = (1, hparams.frame_stack_size) + real_obs_shape
      self._initial_frames = np.zeros(shape=shape, dtype=np.uint8)
      def initial_frame_chooser(batch_size):
        assert batch_size == 1
        return self._initial_frames

    else:
      initial_frame_chooser = rl_utils.make_initial_frame_chooser(
          real_env, hparams.frame_stack_size,
          simulation_random_starts=random_starts,
          simulation_flip_first_random_for_beginning=False
      )
    env_fn = make_simulated_env_fn_from_hparams(
        real_env, hparams,
        batch_size=1,
        initial_frame_chooser=initial_frame_chooser,
        model_dir=world_model_dir,
    )

    env = env_fn(in_graph=False)
    self.env = FlatBatchEnv(env)

    self.observation_space = self.env.observation_space
    self.action_space = self.env.action_space
Init.

    Args:
       real_env: gym environment.
       world_model_dir: path to world model checkpoint directory.
       hparams: hparams for rlmb pipeline.
       random_starts: if restart world model from random frames, or only
         from initial ones (from beginning of episodes). Valid only when
         `setable_initial_fames` set to False.
       setable_initial_frames: if True, initial_frames for world model should be
         set by `add_to_initial_stack`.
    
tensor2tensor/rl/player_utils.py