def _wrap_env(self, env: GymEnv) -> VecEnv: if not isinstance(env, VecEnv): if self.verbose >= 1: print("Wrapping the env in a DummyVecEnv.") env = DummyVecEnv([lambda: env]) if is_image_space(env.observation_space) and not isinstance( env, VecTransposeImage): if self.verbose >= 1: print("Wrapping the env in a VecTransposeImage.") env = VecTransposeImage(env) return env
def create_env(n_envs, eval_env=False): """ Create the environment and wrap it if necessary :param n_envs: (int) :param eval_env: (bool) Whether is it an environment used for evaluation or not :return: (Union[gym.Env, VecEnv]) """ global hyperparams global env_kwargs # Do not log eval env (issue with writing the same file) log_dir = None if eval_env else save_path if n_envs == 1: env = DummyVecEnv([ make_env(env_id, 0, args.seed, wrapper_class=env_wrapper, log_dir=log_dir, env_kwargs=env_kwargs) ]) else: # env = SubprocVecEnv([make_env(env_id, i, args.seed) for i in range(n_envs)]) # On most env, SubprocVecEnv does not help and is quite memory hungry env = DummyVecEnv([ make_env(env_id, i, args.seed, log_dir=log_dir, env_kwargs=env_kwargs, wrapper_class=env_wrapper) for i in range(n_envs) ]) if normalize: if args.verbose > 0: if len(normalize_kwargs) > 0: print(f"Normalization activated: {normalize_kwargs}") else: print("Normalizing input and reward") env = VecNormalize(env, **normalize_kwargs) # Optional Frame-stacking if hyperparams.get('frame_stack', False): n_stack = hyperparams['frame_stack'] env = VecFrameStack(env, n_stack) print(f"Stacking {n_stack} frames") if is_image_space(env.observation_space): if args.verbose > 0: print("Wrapping into a VecTransposeImage") env = VecTransposeImage(env) return env
def create_env(n_envs, eval_env=False, no_log=False): """ Create the environment and wrap it if necessary :param n_envs: (int) :param eval_env: (bool) Whether is it an environment used for evaluation or not :param no_log: (bool) Do not log training when doing hyperparameter optim (issue with writing the same file) :return: (Union[gym.Env, VecEnv]) """ global hyperparams global env_kwargs # Do not log eval env (issue with writing the same file) log_dir = None if eval_env or no_log else save_path if n_envs == 1: env = DummyVecEnv([make_env(env_id, 0, args.seed, wrapper_class=env_wrapper, log_dir=log_dir, env_kwargs=env_kwargs)]) else: # env = SubprocVecEnv([make_env(env_id, i, args.seed) for i in range(n_envs)]) # On most env, SubprocVecEnv does not help and is quite memory hungry env = DummyVecEnv([make_env(env_id, i, args.seed, log_dir=log_dir, env_kwargs=env_kwargs, wrapper_class=env_wrapper) for i in range(n_envs)]) # if normalize: # # Copy to avoid changing default values by reference # local_normalize_kwargs = normalize_kwargs.copy() # # Do not normalize reward for env used for evaluation # if eval_env: # if len(local_normalize_kwargs) > 0: # local_normalize_kwargs['norm_reward'] = False # else: # local_normalize_kwargs = {'norm_reward': False} # if args.verbose > 0: # if len(local_normalize_kwargs) > 0: # print(f"Normalization activated: {local_normalize_kwargs}") # else: # print("Normalizing input and reward") # env = VecNormalize(env, **local_normalize_kwargs) # Optional Frame-stacking if hyperparams.get('frame_stack', False): n_stack = hyperparams['frame_stack'] env = VecFrameStack(env, n_stack) print(f"Stacking {n_stack} frames") if is_image_space(env.observation_space): if args.verbose > 0: print("Wrapping into a VecTransposeImage") env = VecTransposeImage(env) return env
def create_vectorized_environment( n_envs: int, frame_stack: int, env_creation_func: t.Callable) -> VecTransposeImage: """Creates a vectorized environment for image-based models. :param n_envs: The number of parallel environment to run. :param frame_stack: The number of frame to stack in each environment. :param env_creation_func: A callable returning a Gym environment. :return: A vectorized environment with frame stacking and image transposition. """ return VecTransposeImage( VecFrameStack(SubprocVecEnv([env_creation_func] * n_envs), frame_stack))
def make_env(env_id, n_envs, frame_stack = True, clip_reward = False, terminal_on_life_loss = False, monitor_dir = './log/monitors', vec_env_cls = SubprocVecEnv): wrapper_kwargs = {'terminal_on_life_loss':terminal_on_life_loss, 'clip_reward':clip_reward} env = make_atari_env(env_id, n_envs, monitor_dir = monitor_dir, vec_env_cls = vec_env_cls, wrapper_kwargs = wrapper_kwargs) env = VecFrameStack(env, 4) env = VecTransposeImage(env)
def create_envs(self, n_envs: int, eval_env: bool = False, no_log: bool = False) -> VecEnv: """ Create the environment and wrap it if necessary. :param n_envs: :param eval_env: Whether is it an environment used for evaluation or not :param no_log: Do not log training when doing hyperparameter optim (issue with writing the same file) :return: the vectorized environment, with appropriate wrappers """ # Do not log eval env (issue with writing the same file) log_dir = None if eval_env or no_log else self.save_path monitor_kwargs = {} # Special case for GoalEnvs: log success rate too if "Neck" in self.env_id or self.is_robotics_env(self.env_id) or "parking-v0" in self.env_id: monitor_kwargs = dict(info_keywords=("is_success",)) # On most env, SubprocVecEnv does not help and is quite memory hungry # therefore we use DummyVecEnv by default env = make_vec_env( env_id=self.env_id, n_envs=n_envs, seed=self.seed, env_kwargs=self.env_kwargs, monitor_dir=log_dir, wrapper_class=self.env_wrapper, vec_env_cls=self.vec_env_class, vec_env_kwargs=self.vec_env_kwargs, monitor_kwargs=monitor_kwargs, ) # Wrap the env into a VecNormalize wrapper if needed # and load saved statistics when present env = self._maybe_normalize(env, eval_env) # Optional Frame-stacking if self.frame_stack is not None: n_stack = self.frame_stack env = VecFrameStack(env, n_stack) if self.verbose > 0: print(f"Stacking {n_stack} frames") # Wrap if needed to re-order channels # (switch from channel last to channel first convention) if is_image_space(env.observation_space) and not is_image_space_channels_first(env.observation_space): if self.verbose > 0: print("Wrapping into a VecTransposeImage") env = VecTransposeImage(env) return env
def _wrap_env(env: GymEnv, verbose: int = 0, monitor_wrapper: bool = True) -> VecEnv: """ " Wrap environment with the appropriate wrappers if needed. For instance, to have a vectorized environment or to re-order the image channels. :param env: :param verbose: :param monitor_wrapper: Whether to wrap the env in a ``Monitor`` when possible. :return: The wrapped environment. """ if not isinstance(env, VecEnv): if not is_wrapped(env, Monitor) and monitor_wrapper: if verbose >= 1: print("Wrapping the env with a `Monitor` wrapper") env = Monitor(env) if verbose >= 1: print("Wrapping the env in a DummyVecEnv.") env = DummyVecEnv([lambda: env]) # Make sure that dict-spaces are not nested (not supported) check_for_nested_spaces(env.observation_space) if isinstance(env.observation_space, gym.spaces.Dict): for space in env.observation_space.spaces.values(): if isinstance(space, gym.spaces.Dict): raise ValueError("Nested observation spaces are not supported (Dict spaces inside Dict space).") if not is_vecenv_wrapped(env, VecTransposeImage): wrap_with_vectranspose = False if isinstance(env.observation_space, gym.spaces.Dict): # If even one of the keys is a image-space in need of transpose, apply transpose # If the image spaces are not consistent (for instance one is channel first, # the other channel last), VecTransposeImage will throw an error for space in env.observation_space.spaces.values(): wrap_with_vectranspose = wrap_with_vectranspose or ( is_image_space(space) and not is_image_space_channels_first(space) ) else: wrap_with_vectranspose = is_image_space(env.observation_space) and not is_image_space_channels_first( env.observation_space ) if wrap_with_vectranspose: if verbose >= 1: print("Wrapping the env in a VecTransposeImage.") env = VecTransposeImage(env) return env
def _wrap_env(env: GymEnv, verbose: int = 0) -> VecEnv: if not isinstance(env, VecEnv): if verbose >= 1: print("Wrapping the env in a DummyVecEnv.") env = DummyVecEnv([lambda: env]) if is_image_space(env.observation_space) and not is_wrapped(env, VecTransposeImage): if verbose >= 1: print("Wrapping the env in a VecTransposeImage.") env = VecTransposeImage(env) # check if wrapper for dict support is needed when using HER if isinstance(env.observation_space, gym.spaces.dict.Dict): env = ObsDictWrapper(env) return env
def maybe_transpose(observation: np.ndarray, observation_space: spaces.Space) -> np.ndarray: """ Handle the different cases for images as PyTorch use channel first format. :param observation: :param observation_space: :return: channel first observation if observation is an image """ # Avoid circular import from stable_baselines3.common.vec_env import VecTransposeImage if is_image_space(observation_space): if not (observation.shape == observation_space.shape or observation.shape[1:] == observation_space.shape): # Try to re-order the channels transpose_obs = VecTransposeImage.transpose_image(observation) if transpose_obs.shape == observation_space.shape or transpose_obs.shape[1:] == observation_space.shape: observation = transpose_obs return observation
def check_env(env: GymEnv, observation_space: gym.spaces.Space, action_space: gym.spaces.Space): """ Checks the validity of the environment to load vs the one used for training. Checked parameters: - observation_space - action_space :param env: (GymEnv) :param observation_space: (gym.spaces.Space) :param action_space: (gym.spaces.Space) """ if (observation_space != env.observation_space # Special cases for images that need to be transposed and not (is_image_space(env.observation_space) and observation_space == VecTransposeImage.transpose_space(env.observation_space))): raise ValueError(f'Observation spaces do not match: {observation_space} != {env.observation_space}') if action_space != env.action_space: raise ValueError(f'Action spaces do not match: {action_space} != {env.action_space}')
def _wrap_env(self, env: GymEnv) -> VecEnv: """ This overrides _wrap_env from stable_baselines3.common.base_class.BaseAlgorithm Now the DummyVecEnv gets wrapped inside a VecNormalize environment to normalize rewards and observations """ if not isinstance(env, VecEnv): if self.verbose >= 1: print("Wrapping the env in a DummyVecEnv.") env = DummyVecEnv([lambda: env]) # Automatically normalize the input features and reward # norm_obs/norm_reward: True if obs/rew should be normalized # clip_obs: Max absolute value for observation # clip_reward: Max value absolute for discounted reward # gamma: discount factor env = VecNormalize(env, training=True, norm_obs=True, norm_reward=True, gamma=0.99) # clip_obs=10., clip_reward=10.0, if is_image_space(env.observation_space) and not isinstance(env, VecTransposeImage): if self.verbose >= 1: print("Wrapping the env in a VecTransposeImage.") env = VecTransposeImage(env) return env
def check_for_correct_spaces(env: GymEnv, observation_space: gym.spaces.Space, action_space: gym.spaces.Space) -> None: """ Checks that the environment has same spaces as provided ones. Used by BaseAlgorithm to check if spaces match after loading the model with given env. Checked parameters: - observation_space - action_space :param env: Environment to check for valid spaces :param observation_space: Observation space to check against :param action_space: Action space to check against """ if ( observation_space != env.observation_space # Special cases for images that need to be transposed and not ( is_image_space(env.observation_space) and observation_space == VecTransposeImage.transpose_space(env.observation_space) ) ): raise ValueError(f"Observation spaces do not match: {observation_space} != {env.observation_space}") if action_space != env.action_space: raise ValueError(f"Action spaces do not match: {action_space} != {env.action_space}")
def image_transpose(env): if is_image_space(env.observation_space) and not is_image_space_channels_first( env.observation_space ): env = VecTransposeImage(env) return env
if custom_params['algo'] == 'dqn': env = DiscreteWrapper(env) if custom_params['USING_VAE']: env = NormalizeWrapper(env) # No need to use normalization if image env = FinalLayerObservationWrapper(env, latent_dim=1028, map="map3") # Step 3.b. To make Vectorized Environment to be able to use Normalize or FramStack (Optional) env = make_vec_env(lambda: env, n_envs=1) # Step 3.b Passing through Normalization and stack frame (Optional) env = VecFrameStack( env, n_stack=custom_params['FRAME_STACK']) # Use 1 for now because we use image if not custom_params['USING_VAE']: env = VecTransposeImage(env) # Uncomment if using 3d obs if custom_params['USING_NORMALIZATION']: env = VecNormalize.load(osp.join(results_dir, "vec_normalization.pkl"), env) # Load the agent if custom_params['algo'] == 'sac': model = SAC.load(osp.join(results_dir, "best_model", "best_model.zip")) elif custom_params['algo'] == 'a2c': model = A2C.load(osp.join(results_dir, "best_model", "best_model.zip")) elif custom_params['algo'] == 'dqn': model = DQN.load(osp.join(results_dir, "best_model", "best_model.zip")) elif custom_params['algo'] == 'ppo': model = PPO.load(osp.join(results_dir, "best_model", "best_model.zip")) else:
if env_name == "takeoff-aviary-v0": eval_env = make_vec_env(TakeoffAviary, env_kwargs=sa_env_kwargs, n_envs=1, seed=0) if env_nam == "hover-aviary-v0": eval_env = make_vec_env(HoverAviary, env_kwargs=sa_env_kwargs, n_envs=1, seed=0) if env_name == "flythrugate-aviary-v0": eval_env = make_vec_env(FlyThruGateAviary, env_kwargs=sa_env_kwargs, n_envs=1, seed=0) eval_env = VecTransposeImage(eval_env) #### Train the model ####################################### # checkpoint_callback = CheckpointCallback(save_freq=1000, save_path=filename+'-logs/', name_prefix='rl_model') callback_on_best = StopTrainingOnRewardThreshold( reward_threshold=EPISODE_REWARD_THRESHOLD, verbose=1) eval_callback = EvalCallback(eval_env, callback_on_new_best=callback_on_best, verbose=1, best_model_save_path=filename + '/', log_path=filename + '/', eval_freq=int(5000 / ARGS.cpu), deterministic=True, render=False) model.learn(total_timesteps=int(1e12), callback=eval_callback,
def create_envs(self, n_envs: int, eval_env: bool = False, no_log: bool = False) -> VecEnv: """ Create the environment and wrap it if necessary. :param n_envs: :param eval_env: Whether is it an environment used for evaluation or not :param no_log: Do not log training when doing hyperparameter optim (issue with writing the same file) :return: the vectorized environment, with appropriate wrappers """ # Do not log eval env (issue with writing the same file) log_dir = None if eval_env or no_log else self.save_path monitor_kwargs = {} # Special case for GoalEnvs: log success rate too if "Neck" in self.env_id or self.is_robotics_env( self.env_id) or "parking-v0" in self.env_id: monitor_kwargs = dict(info_keywords=("is_success", )) # Note: made custom to support Gazebo Runtime wrapping def make_env(): def _init(): env = self.env_wrapper(env=self.env_id, **self.env_kwargs) env.seed(self.seed) env.action_space.seed(self.seed) monitor_path = log_dir if log_dir is not None else None if monitor_path is not None: os.makedirs(log_dir, exist_ok=True) env = Monitor(env, filename=monitor_path, **monitor_kwargs) return env return _init if self.vec_env_class is None: self.vec_env_class = DummyVecEnv env = self.vec_env_class([make_env()], **self.vec_env_kwargs) # Wrap the env into a VecNormalize wrapper if needed # and load saved statistics when present env = self._maybe_normalize(env, eval_env) # Optional Frame-stacking if self.frame_stack is not None: n_stack = self.frame_stack env = VecFrameStack(env, n_stack) if self.verbose > 0: print(f"Stacking {n_stack} frames") # Wrap if needed to re-order channels # (switch from channel last to channel first convention) if is_image_space( env.observation_space) and not is_image_space_channels_first( env.observation_space): if self.verbose > 0: print("Wrapping into a VecTransposeImage") env = VecTransposeImage(env) # check if wrapper for dict support is needed if self.algo == "her": if self.verbose > 0: print("Wrapping into a ObsDictWrapper") env = ObsDictWrapper(env) return env
def forward(self, obs): obs_transposed = VecTransposeImage.transpose_image(obs) latent, _, _ = self.ac_model._get_latent( th.tensor(obs_transposed).to(self.device)) return self.reward_net(latent)
def evaluate(individual: Individual, device: Union[torch.device, str] = "auto") -> Tuple[int]: """ Evaluate a single individual model and return it's mean score after the training time is elapsed. Models are trained and evaluated for a number of timestamps as parameterized in the constants at the top of the file. :param individual: The individual to evaluate. :return: """ t_start = time() layers = individual.weights name = individual.encode() checkpoint_path = os.path.join(BASE_CHECKPOINT_PATH, "PPO", ENV_NAME, name) if os.path.exists(checkpoint_path): return (random.randint(MIN_SCORE, MAX_SCORE), ) os.makedirs(checkpoint_path, exist_ok=True) log_path = os.path.join(BASE_LOG_PATH, "PPO", ENV_NAME, name) os.makedirs(log_path, exist_ok=True) results_path = os.path.join(checkpoint_path, "results.json") if not os.path.exists(results_path): env_args = dict( frame_skip=4, screen_size=84, terminal_on_life_loss=True, clip_reward=True, ) # Creates a gym environment for an atari game using the specified seed and number of environments # This is a "vectorized environment", which means Stable Baselines batches the updates into vectors # for improved performance.. def atari_wrapper(env: gym.Env) -> gym.Env: env = AtariWrapper(env, **env_args) return env def make_env(rank: int, count: int) -> VecEnv: return make_vec_env( ENV_NAME, n_envs=count, seed=RANDOM_SEED + rank, start_index=0, monitor_dir=None, wrapper_class=atari_wrapper, env_kwargs=None, vec_env_cls=SubprocVecEnv, vec_env_kwargs=None, monitor_kwargs=None, ) train_env = make_env(0, N_ENVS) eval_env = make_env(1, 1) # required by models in baselines train_env = VecTransposeImage(train_env) eval_env = VecTransposeImage(eval_env) # setup callback to save model at fixed intervals save_callback = CheckpointCallback(save_freq=CHECKPOINT_FREQ, save_path=checkpoint_path, name_prefix=name) stop_callback = StopTrainingOnRewardThreshold( reward_threshold=EVAL_THRESHOLD) time_callback = TimeLimitCallback(max_time=TIME_LIMIT) best_callback = EvalCallback( eval_env, eval_freq=EVAL_FREQ, best_model_save_path=checkpoint_path, callback_on_new_best=stop_callback, ) list_callback = CallbackList( [save_callback, best_callback, time_callback]) model = PPO( CnnPolicy, train_env, verbose=VERBOSE, batch_size=BATCH_SIZE, seed=RANDOM_SEED * 7, tensorboard_log=log_path, learning_rate=LEARNING_RATE, n_steps=UPDATE_STEPS, n_epochs=N_EPOCHS, ent_coef=ENT_COEF, vf_coef=VF_COEF, clip_range=CLIP_RANGE, device=device, policy_kwargs=dict(features_extractor_class=VariableBenchmark, features_extractor_kwargs=dict(layers=layers)), ) config_path = os.path.join(checkpoint_path, "cnn_config") zip_path = os.path.join(checkpoint_path, "model.zip") # output the model config to a file for easier viewing with open(config_path, "w") as file: file.write(f"{name}\n") file.write(str(model.policy.features_extractor.cnn)) print("Beginning training...") model.learn(TRAIN_STEPS, callback=list_callback, tb_log_name="run") model.save(zip_path) del train_env del eval_env time_taken = time() - t_start print("Beginning evaluation...") # score of the game, standard deviation of multiple runs reward_mean, reward_std = evaluate_policy(model, make_env(2, 1)) with open(results_path, "w") as handle: handle.write(json.dumps((reward_mean, reward_std, time_taken))) else: reward_mean, reward_std, time_taken = json.load(open( results_path, "r")) reward_mean = abs(MIN_SCORE) + reward_mean value = (reward_mean * weighted_time(time_taken), ) print(f"Evaluated {name} with a score of {value} in {(time_taken):.2f}s") return value
def predict( self, observation: np.ndarray, partner_idx: int = 0, state: Optional[np.ndarray] = None, mask: Optional[np.ndarray] = None, deterministic: bool = False, ) -> Tuple[np.ndarray, Optional[np.ndarray]]: """ Get the policy action and state from an observation (and optional state). Includes sugar-coating to handle different observations (e.g. normalizing images). :param observation: (np.ndarray) the input observation :param state: (Optional[np.ndarray]) The last states (can be None, used in recurrent policies) :param mask: (Optional[np.ndarray]) The last masks (can be None, used in recurrent policies) :param deterministic: (bool) Whether or not to return deterministic actions. :return: (Tuple[np.ndarray, Optional[np.ndarray]]) the model's action and the next state (used in recurrent policies) """ # TODO (GH/1): add support for RNN policies # if state is None: # state = self.initial_state # if mask is None: # mask = [False for _ in range(self.n_envs)] observation = np.array(observation) # Handle the different cases for images # as PyTorch use channel first format if is_image_space(self.observation_space) and not ( observation.shape == self.observation_space.shape or observation.shape[1:] == self.observation_space.shape): # Try to re-order the channels transpose_obs = VecTransposeImage.transpose_image(observation) if (transpose_obs.shape == self.observation_space.shape or transpose_obs.shape[1:] == self.observation_space.shape): observation = transpose_obs vectorized_env = is_vectorized_observation(observation, self.observation_space) observation = observation.reshape((-1, ) + self.observation_space.shape) observation = th.as_tensor(observation).to(self.device) with th.no_grad(): actions = self._predict(observation, partner_idx=partner_idx, deterministic=deterministic) # Convert to numpy actions = actions.cpu().numpy() if isinstance(self.action_space, gym.spaces.Box): if self.squash_output: # Rescale to proper domain when using squashing actions = self.unscale_action(actions) else: # Actions could be on arbitrary scale, so clip the actions to avoid # out of bound error (e.g. if sampling from a Gaussian distribution) actions = np.clip(actions, self.action_space.low, self.action_space.high) if not vectorized_env: if state is not None: raise ValueError( "Error: The environment must be vectorized when using recurrent policies." ) actions = actions[0] return actions, state
instance = 'P12' timestamp = datetime.datetime.now().strftime("%y%m%d_%H%M") environment = 'ofp' algo = 'ppo' mode = 'rgb_array' train_steps = np.append( np.outer(10.0**(np.arange(4, 6)), np.arange(1, 10, 1)).flatten(), 10**6) train_steps = [5e6] vec_env = make_vec_env('ofp-v0', env_kwargs={ 'mode': mode, "instance": instance }, n_envs=1) wrap_env = VecTransposeImage(vec_env) vec_eval_env = make_vec_env('ofp-v0', env_kwargs={ 'mode': mode, "instance": instance }, n_envs=1) wrap_eval_env = VecTransposeImage(vec_eval_env) experiment_results = {} for ts in train_steps: ts = int(ts) print(ts) save_path = f"{timestamp}_{instance}_{algo}_{mode}_{environment}_movingavg_nocollisions_{ts}"
from stable_baselines3.common.vec_env import DummyVecEnv, VecTransposeImage from stable_baselines3.common.evaluation import evaluate_policy from stable_baselines3.common.callbacks import EvalCallback # Create a DummyVecEnv for main airsim gym env env = DummyVecEnv([ lambda: Monitor( gym.make( "airgym:airsim-car-sample-v0", ip_address="127.0.0.1", image_shape=(84, 84, 1), )) ]) # Wrap env as VecTransposeImage to allow SB to handle frame observations env = VecTransposeImage(env) # Initialize RL algorithm type and parameters model = DQN( "CnnPolicy", env, learning_rate=0.00025, verbose=1, batch_size=32, train_freq=4, target_update_interval=10000, learning_starts=200000, buffer_size=500000, max_grad_norm=10, exploration_fraction=0.1, exploration_final_eps=0.01,
def main(): set_random_seed(RANDOM_SEED) t_start = time() name = "LargeFinalLayer" checkpoint_path = os.path.join(BASE_CHECKPOINT_PATH, "PPO", ENV_NAME, name) os.makedirs(checkpoint_path, exist_ok=True) log_path = os.path.join(BASE_LOG_PATH, "PPO", ENV_NAME, name) os.makedirs(log_path, exist_ok=True) results_path = os.path.join(checkpoint_path, "results.json") env_args = dict( frame_skip=4, screen_size=84, terminal_on_life_loss=True, clip_reward=True, ) # Creates a gym environment for an atari game using the specified seed and number of environments # This is a "vectorized environment", which means Stable Baselines batches the updates into vectors # for improved performance.. # train_env = make_atari_env(ENV_NAME, n_envs=N_ENVS, seed=RANDOM_SEED, wrapper_kwargs=env_args) def atari_wrapper(env: gym.Env) -> gym.Env: env = AtariWrapper(env, **env_args) return env def make_env(rank: int, count: int) -> VecEnv: return make_vec_env( ENV_NAME, n_envs=count, seed=RANDOM_SEED + rank, start_index=0, monitor_dir=None, wrapper_class=atari_wrapper, env_kwargs=None, vec_env_cls=None, vec_env_kwargs=None, monitor_kwargs=None, ) train_env = make_env(0, N_ENVS) eval_env = make_env(1, 1) # required by models in baselines train_env = VecTransposeImage(train_env) eval_env = VecTransposeImage(eval_env) # setup callback to save model at fixed intervals save_callback = CheckpointCallback(save_freq=CHECKPOINT_FREQ, save_path=checkpoint_path, name_prefix=name) stop_callback = StopTrainingOnRewardThreshold( reward_threshold=EVAL_THRESHOLD) time_callback = TimeLimitCallback(max_time=TIME_LIMIT) best_callback = EvalCallback( eval_env, eval_freq=EVAL_FREQ, best_model_save_path=checkpoint_path, callback_on_new_best=stop_callback, ) list_callback = CallbackList([save_callback, best_callback, time_callback]) model = PPO( CnnPolicy, train_env, verbose=VERBOSE, batch_size=BATCH_SIZE, seed=RANDOM_SEED, tensorboard_log=log_path, learning_rate=LEARNING_RATE, n_steps=UPDATE_STEPS, n_epochs=N_EPOCHS, ent_coef=ENT_COEF, vf_coef=VF_COEF, clip_range=CLIP_RANGE, device=DEVICE_TYPE, policy_kwargs=dict(features_extractor_class=FeatureExtractor), ) config_path = os.path.join(checkpoint_path, "cnn_config") zip_path = os.path.join(checkpoint_path, "model.zip") # output the model config to a file for easier viewing with open(config_path, "w") as file: file.write(f"{name}\n") file.write(str(model.policy.features_extractor.cnn)) print("Beginning training...") model.learn(TRAIN_STEPS, callback=list_callback, tb_log_name="run") # model.learn(TRAIN_STEPS, tb_log_name="run") model.save(zip_path) del train_env # del eval_env time_taken = time() - t_start print("Beginning evaluation...") # score of the game, standard deviation of multiple runs reward_mean, reward_std = evaluate_policy(model, make_env(2, 1)) with open(results_path, "w") as handle: handle.write(json.dumps((reward_mean, reward_std, time_taken)))