class_docstr: Gym environment, running with world model.
  
  Allows passing custom initial frames.
  
  Examples:
    Setup simulated env from some point of real rollout.
      >>> sim_env = SimulatedGymEnv(setable_initial_frames=True, **kwargs)
      >>> real_env = FlatBatchEnv(T2TGymEnv(...))
      >>> while ...:
      >>>   ob, _, _, _ = real_env.step(action)
      >>>   sim_env.add_to_initial_stack(ob)
      >>> sim_env.reset()
      >>> # Continue sim_env rollout.
class_docstr_tok: ["Gym","environment",",","running","with","world","model",".","Allows","passing","custom","initial","frames",".","Examples",":","Setup","simulated","env","from","some","point","of","real","rollout",".",">",">",">","sim_env","=","SimulatedGymEnv","(","setable_initial_frames=True",",","*","*","kwargs",")",">",">",">","real_env","=","FlatBatchEnv","(","T2TGymEnv","(","...",")",")",">",">",">","while","...",":",">",">",">","ob",",","_",",","_",",","_","=","real_env.step","(","action",")",">",">",">","sim_env.add_to_initial_stack","(","ob",")",">",">",">","sim_env.reset","(",")",">",">",">","#","Continue","sim_env","rollout","."]
docstr: 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`.
docstr_tok: ["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","`","."]
code: 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
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import copy
  import os
  import re
  import gym
  import numpy
  import six
  from tensor2tensor.models.research.rl import get_policy
  from tensor2tensor.models.research.rl import make_simulated_env_fn_from_hparams
  from tensor2tensor.rl import rl_utils
  from tensor2tensor.rl.envs.simulated_batch_gym_env import FlatBatchEnv
  from tensor2tensor.utils import hparam
  from tensor2tensor.utils import trainer_lib
  from tensor2tensor.utils.misc_utils import camelcase_to_snakecase
  import tensorflow.compat.v1
label_desc: Use image node_id 1 to create a new SimulatedGymEnv object from inherited base classes: gym with example: obj = SimulatedGymEnv(real_env, world_model_dir, hparams, random_starts, setable_initial_frames)