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)))
