def test_serialize_identity(env_name, net_cls, tmpdir):
"""Does output of deserialized reward network match that of original?"""
logging.info(f"Testing {net_cls}")
venv = util.make_vec_env(env_name, n_envs=1, parallel=False)
original = net_cls(venv.observation_space, venv.action_space)
random = base.RandomPolicy(venv.observation_space, venv.action_space)
tmppath = os.path.join(tmpdir, "reward.pt")
th.save(original, tmppath)
loaded = th.load(tmppath)
assert original.observation_space == loaded.observation_space
assert original.action_space == loaded.action_space
transitions = rollout.generate_transitions(random, venv, n_timesteps=100)
unshaped_fn = serialize.load_reward("RewardNet_unshaped", tmppath, venv)
shaped_fn = serialize.load_reward("RewardNet_shaped", tmppath, venv)
rewards = {
"train": [],
"test": [],
}
for net in [original, loaded]:
trans_args = (
transitions.obs,
transitions.acts,
transitions.next_obs,
transitions.dones,
)
rewards["train"].append(net.predict_reward_train(*trans_args))
rewards["test"].append(net.predict_reward_test(*trans_args))
args = (
transitions.obs,
transitions.acts,
transitions.next_obs,
transitions.dones,
)
rewards["train"].append(shaped_fn(*args))
rewards["test"].append(unshaped_fn(*args))
for key, predictions in rewards.items():
assert len(predictions) == 3
assert np.allclose(predictions[0], predictions[1])
assert np.allclose(predictions[0], predictions[2])
def test_serialize_identity(session, env_name, reward_net):
"""Does output of deserialized reward network match that of original?"""
net_name, net_cls = reward_net
print(f"Testing {net_name}")
venv = util.make_vec_env(env_name, n_envs=1, parallel=False)
with tf.variable_scope("original"):
original = net_cls(venv.observation_space, venv.action_space)
random = base.RandomPolicy(venv.observation_space, venv.action_space)
session.run(tf.global_variables_initializer())
with tempfile.TemporaryDirectory(
prefix='imitation-serialize-rew') as tmpdir:
original.save(tmpdir)
with tf.variable_scope("loaded"):
loaded = net_cls.load(tmpdir)
assert original.observation_space == loaded.observation_space
assert original.action_space == loaded.action_space
rollouts = rollout.generate_transitions(random, venv, n_timesteps=100)
feed_dict = {}
outputs = {'train': [], 'test': []}
for net in [original, loaded]:
feed_dict.update(_make_feed_dict(net, rollouts))
outputs['train'].append(net.reward_output_train)
outputs['test'].append(net.reward_output_test)
unshaped_name = f"{net_name}_unshaped"
shaped_name = f"{net_name}_shaped"
with serialize.load_reward(unshaped_name, tmpdir, venv) as unshaped_fn:
with serialize.load_reward(shaped_name, tmpdir, venv) as shaped_fn:
rewards = session.run(outputs, feed_dict=feed_dict)
old_obs, actions, new_obs, _ = rollouts
steps = np.zeros((old_obs.shape[0], ))
rewards['train'].append(
shaped_fn(old_obs, actions, new_obs, steps))
rewards['test'].append(
unshaped_fn(old_obs, actions, new_obs, steps))
for key, predictions in rewards.items():
assert len(predictions) == 3
assert np.allclose(predictions[0], predictions[1])
assert np.allclose(predictions[0], predictions[2])
def test_serialize_identity(session, env_name, net_cls, tmpdir):
"""Does output of deserialized reward network match that of original?"""
logging.info(f"Testing {net_cls}")
venv = util.make_vec_env(env_name, n_envs=1, parallel=False)
with tf.variable_scope("original"):
original = net_cls(venv.observation_space, venv.action_space)
random = base.RandomPolicy(venv.observation_space, venv.action_space)
session.run(tf.global_variables_initializer())
original.save(tmpdir)
with tf.variable_scope("loaded"):
loaded = reward_net.RewardNet.load(tmpdir)
assert original.observation_space == loaded.observation_space
assert original.action_space == loaded.action_space
transitions = rollout.generate_transitions(random, venv, n_timesteps=100)
feed_dict = {}
outputs = {"train": [], "test": []}
for net in [original, loaded]:
feed_dict.update(_make_feed_dict(net, transitions))
outputs["train"].append(net.reward_output_train)
outputs["test"].append(net.reward_output_test)
with serialize.load_reward("RewardNet_unshaped", tmpdir, venv) as unshaped_fn:
with serialize.load_reward("RewardNet_shaped", tmpdir, venv) as shaped_fn:
rewards = session.run(outputs, feed_dict=feed_dict)
args = (
transitions.obs,
transitions.acts,
transitions.next_obs,
transitions.dones,
)
rewards["train"].append(shaped_fn(*args))
rewards["test"].append(unshaped_fn(*args))
for key, predictions in rewards.items():
assert len(predictions) == 3
assert np.allclose(predictions[0], predictions[1])
assert np.allclose(predictions[0], predictions[2])
def test_reward_valid(env_name, reward_type):
"""Test output of reward function is appropriate shape and type."""
venv = util.make_vec_env(env_name, n_envs=1, parallel=False)
TRAJECTORY_LEN = 10
obs = _sample(venv.observation_space, TRAJECTORY_LEN)
actions = _sample(venv.action_space, TRAJECTORY_LEN)
next_obs = _sample(venv.observation_space, TRAJECTORY_LEN)
steps = np.arange(0, TRAJECTORY_LEN)
reward_fn = serialize.load_reward(reward_type, "foobar", venv)
pred_reward = reward_fn(obs, actions, next_obs, steps)
assert pred_reward.shape == (TRAJECTORY_LEN, )
assert isinstance(pred_reward[0], numbers.Number)
def rollouts_and_policy(
_run,
_seed: int,
env_name: str,
total_timesteps: int,
*,
log_dir: str,
num_vec: int,
parallel: bool,
max_episode_steps: Optional[int],
normalize: bool,
normalize_kwargs: dict,
init_rl_kwargs: dict,
n_episodes_eval: int,
reward_type: Optional[str],
reward_path: Optional[str],
rollout_save_interval: int,
rollout_save_final: bool,
rollout_save_n_timesteps: Optional[int],
rollout_save_n_episodes: Optional[int],
policy_save_interval: int,
policy_save_final: bool,
init_tensorboard: bool,
) -> dict:
"""Trains an expert policy from scratch and saves the rollouts and policy.
Checkpoints:
At applicable training steps `step` (where step is either an integer or
"final"):
- Policies are saved to `{log_dir}/policies/{step}/`.
- Rollouts are saved to `{log_dir}/rollouts/{step}.pkl`.
Args:
env_name: The gym.Env name. Loaded as VecEnv.
total_timesteps: Number of training timesteps in `model.learn()`.
log_dir: The root directory to save metrics and checkpoints to.
num_vec: Number of environments in VecEnv.
parallel: If True, then use DummyVecEnv. Otherwise use SubprocVecEnv.
max_episode_steps: If not None, then environments are wrapped by
TimeLimit so that they have at most `max_episode_steps` steps per
episode.
normalize: If True, then rescale observations and reward.
normalize_kwargs: kwargs for `VecNormalize`.
init_rl_kwargs: kwargs for `init_rl`.
n_episodes_eval: The number of episodes to average over when calculating
the average ground truth reward return of the final policy.
reward_type: If provided, then load the serialized reward of this type,
wrapping the environment in this reward. This is useful to test
whether a reward model transfers. For more information, see
`imitation.rewards.serialize.load_reward`.
reward_path: A specifier, such as a path to a file on disk, used by
reward_type to load the reward model. For more information, see
`imitation.rewards.serialize.load_reward`.
rollout_save_interval: The number of training updates in between
intermediate rollout saves. If the argument is nonpositive, then
don't save intermediate updates.
rollout_save_final: If True, then save rollouts right after training is
finished.
rollout_save_n_timesteps: The minimum number of timesteps saved in every
file. Could be more than `rollout_save_n_timesteps` because
trajectories are saved by episode rather than by transition.
Must set exactly one of `rollout_save_n_timesteps`
and `rollout_save_n_episodes`.
rollout_save_n_episodes: The number of episodes saved in every
file. Must set exactly one of `rollout_save_n_timesteps` and
`rollout_save_n_episodes`.
policy_save_interval: The number of training updates between saves. Has
the same semantics are `rollout_save_interval`.
policy_save_final: If True, then save the policy right after training is
finished.
init_tensorboard: If True, then write tensorboard logs to {log_dir}/sb_tb
and "output/summary/...".
Returns:
The return value of `rollout_stats()` using the final policy.
"""
os.makedirs(log_dir, exist_ok=True)
sacred_util.build_sacred_symlink(log_dir, _run)
sample_until = util.rollout.make_sample_until(rollout_save_n_timesteps,
rollout_save_n_episodes)
eval_sample_until = util.rollout.min_episodes(n_episodes_eval)
with util.make_session():
tf.logging.set_verbosity(tf.logging.INFO)
sb_logger.configure(folder=osp.join(log_dir, 'rl'),
format_strs=['tensorboard', 'stdout'])
rollout_dir = osp.join(log_dir, "rollouts")
policy_dir = osp.join(log_dir, "policies")
os.makedirs(rollout_dir, exist_ok=True)
os.makedirs(policy_dir, exist_ok=True)
if init_tensorboard:
sb_tensorboard_dir = osp.join(log_dir, "sb_tb")
init_rl_kwargs["tensorboard_log"] = sb_tensorboard_dir
venv = util.make_vec_env(env_name, num_vec, seed=_seed,
parallel=parallel, log_dir=log_dir,
max_episode_steps=max_episode_steps)
log_callbacks = []
with contextlib.ExitStack() as stack:
if reward_type is not None:
reward_fn_ctx = load_reward(reward_type, reward_path, venv)
reward_fn = stack.enter_context(reward_fn_ctx)
venv = RewardVecEnvWrapper(venv, reward_fn)
log_callbacks.append(venv.log_callback)
tf.logging.info(
f"Wrapped env in reward {reward_type} from {reward_path}.")
vec_normalize = None
if normalize:
venv = vec_normalize = VecNormalize(venv, **normalize_kwargs)
policy = util.init_rl(venv, verbose=1, **init_rl_kwargs)
# Make callback to save intermediate artifacts during training.
step = 0
def callback(locals_: dict, _) -> bool:
nonlocal step
step += 1
policy = locals_['self']
# TODO(adam): make logging frequency configurable
for callback in log_callbacks:
callback(sb_logger)
if rollout_save_interval > 0 and step % rollout_save_interval == 0:
save_path = osp.join(rollout_dir, f"{step}.pkl")
util.rollout.save(save_path, policy, venv, sample_until)
if policy_save_interval > 0 and step % policy_save_interval == 0:
output_dir = os.path.join(policy_dir, f'{step:05d}')
serialize.save_stable_model(output_dir, policy, vec_normalize)
policy.learn(total_timesteps, callback=callback)
# Save final artifacts after training is complete.
if rollout_save_final:
save_path = osp.join(rollout_dir, "final.pkl")
util.rollout.save(save_path, policy, venv, sample_until)
if policy_save_final:
output_dir = os.path.join(policy_dir, "final")
serialize.save_stable_model(output_dir, policy, vec_normalize)
# Final evaluation of expert policy.
trajs = util.rollout.generate_trajectories(
policy, venv, eval_sample_until)
stats = util.rollout.rollout_stats(trajs)
return stats
def eval_policy(
_run,
_seed: int,
env_name: str,
eval_n_timesteps: Optional[int],
eval_n_episodes: Optional[int],
num_vec: int,
parallel: bool,
render: bool,
render_fps: int,
log_dir: str,
policy_type: str,
policy_path: str,
reward_type: Optional[str] = None,
reward_path: Optional[str] = None,
max_episode_steps: Optional[int] = None,
):
"""Rolls a policy out in an environment, collecting statistics.
Args:
_seed: generated by Sacred.
env_name: Gym environment identifier.
eval_n_timesteps: Minimum number of timesteps to evaluate for. Set exactly
one of `eval_n_episodes` and `eval_n_timesteps`.
eval_n_episodes: Minimum number of episodes to evaluate for. Set exactly
one of `eval_n_episodes` and `eval_n_timesteps`.
num_vec: Number of environments to run simultaneously.
parallel: If True, use `SubprocVecEnv` for true parallelism; otherwise,
uses `DummyVecEnv`.
max_episode_steps: If not None, then environments are wrapped by
TimeLimit so that they have at most `max_episode_steps` steps per
episode.
render: If True, renders interactively to the screen.
log_dir: The directory to log intermediate output to. (As of 2019-07-19
this is just episode-by-episode reward from bench.Monitor.)
policy_type: A unique identifier for the saved policy,
defined in POLICY_CLASSES.
policy_path: A path to the serialized policy.
reward_type: If specified, overrides the environment reward with
a reward of this.
reward_path: If reward_type is specified, the path to a serialized reward
of `reward_type` to override the environment reward with.
Returns:
Return value of `imitation.util.rollout.rollout_stats()`.
"""
os.makedirs(log_dir, exist_ok=True)
sacred_util.build_sacred_symlink(log_dir, _run)
tf.logging.set_verbosity(tf.logging.INFO)
tf.logging.info('Logging to %s', log_dir)
sample_until = rollout.make_sample_until(eval_n_timesteps, eval_n_episodes)
venv = util.make_vec_env(env_name,
num_vec,
seed=_seed,
parallel=parallel,
log_dir=log_dir,
max_episode_steps=max_episode_steps)
venv = VecNormalize(venv, training=False, norm_reward=False)
venv = venv.load(policy_path + "/vec_normalize.pkl", venv)
if render:
venv = InteractiveRender(venv, render_fps)
# TODO(adam): add support for videos using VideoRecorder?
with contextlib.ExitStack() as stack:
if reward_type is not None:
reward_fn_ctx = load_reward(reward_type, reward_path, venv)
reward_fn = stack.enter_context(reward_fn_ctx)
venv = reward_wrapper.RewardVecEnvWrapper(venv, reward_fn)
tf.logging.info(
f"Wrapped env in reward {reward_type} from {reward_path}.")
with serialize.load_policy(policy_type, policy_path, venv) as policy:
trajs = rollout.generate_trajectories(policy, venv, sample_until)
return rollout.rollout_stats(trajs)
def batch_reward_heatmaps(
checkpoints_dir: Union[str, pathlib.Path],
n_gen_trajs: int = 50,
exp_trajs: Optional[List[types.Trajectory]] = None,
) -> Dict[pathlib.Path, plt.Figure]:
"""Build multiple mountain car reward heatmaps from a checkpoint directory.
One plot is generated for every combination of action and checkpoint timestep.
Args:
checkpoints_dir: Path to `checkpoint/` directory from AIRL or GAIL output
directory.
n_gen_trajs: The number of trajectories to rollout using each generator
checkpoint. The transitions in the trajectory are scatterplotted on top of
the heatmap from the same checkpoint timestamp. Nonpositive indicates that
no trajectories should be plotted.
exp_trajs: Expert trajectories for scatterplotting. Generator trajectories
are dynamically generated from generator checkpoints.
Returns:
A dictionary mapping relative paths to `plt.Figure`. Every key is of the
form "{action_name}/{checkpoint_step}" where action_name is "left",
"neutral", or "right".
"""
result = {}
venv = vec_env.DummyVecEnv([lambda: gym.make("MountainCar-v0")])
checkpoints_dir = pathlib.Path(checkpoints_dir)
for checkpoint_dir in sorted(checkpoints_dir.iterdir()):
vec_normalize_path = checkpoint_dir / "gen_policy" / "vec_normalize.pkl"
discrim_path = checkpoint_dir / "discrim.pt"
policy_path = checkpoint_dir / "gen_policy"
if n_gen_trajs > 0:
# `load_policy` automatically loads VecNormalize for policy evaluation.
gen_policy = policies_serialize.load_policy(
"ppo", str(policy_path), venv)
gen_trajs = rollout.generate_trajectories(
gen_policy,
venv,
sample_until=rollout.min_episodes(n_gen_trajs))
else:
gen_trajs = None
# `gen_trajs` contains unnormalized observations.
# Load VecNormalize for use in RewardFn, which doesn't automatically
# normalize input observations.
with open(vec_normalize_path, "rb") as f:
vec_normalize = pickle.load(f) # type: vec_env.VecNormalize
vec_normalize.training = False
reward_fn = rewards_serialize.load_reward("DiscrimNet", discrim_path,
venv)
norm_rew_fn = common.build_norm_reward_fn(reward_fn=reward_fn,
vec_normalize=vec_normalize)
for act in range(MC_NUM_ACTS):
fig = make_heatmap(
act=act,
reward_fn=norm_rew_fn,
gen_trajs=gen_trajs,
exp_trajs=exp_trajs,
)
path = pathlib.Path(ACT_NAMES[act], checkpoint_dir.name)
result[path] = fig
return result
def eval_policy(
_run,
_seed: int,
env_name: str,
eval_n_timesteps: Optional[int],
eval_n_episodes: Optional[int],
num_vec: int,
parallel: bool,
render: bool,
render_fps: int,
videos: bool,
video_kwargs: Mapping[str, Any],
log_dir: str,
policy_type: str,
policy_path: str,
reward_type: Optional[str] = None,
reward_path: Optional[str] = None,
max_episode_steps: Optional[int] = None,
):
"""Rolls a policy out in an environment, collecting statistics.
Args:
_seed: generated by Sacred.
env_name: Gym environment identifier.
eval_n_timesteps: Minimum number of timesteps to evaluate for. Set exactly
one of `eval_n_episodes` and `eval_n_timesteps`.
eval_n_episodes: Minimum number of episodes to evaluate for. Set exactly
one of `eval_n_episodes` and `eval_n_timesteps`.
num_vec: Number of environments to run simultaneously.
parallel: If True, use `SubprocVecEnv` for true parallelism; otherwise,
uses `DummyVecEnv`.
max_episode_steps: If not None, then environments are wrapped by
TimeLimit so that they have at most `max_episode_steps` steps per
episode.
render: If True, renders interactively to the screen.
render_fps: The target number of frames per second to render on screen.
videos: If True, saves videos to `log_dir`.
video_kwargs: Keyword arguments passed through to `video_wrapper.VideoWrapper`.
log_dir: The directory to log intermediate output to, such as episode reward.
policy_type: A unique identifier for the saved policy,
defined in POLICY_CLASSES.
policy_path: A path to the serialized policy.
reward_type: If specified, overrides the environment reward with
a reward of this.
reward_path: If reward_type is specified, the path to a serialized reward
of `reward_type` to override the environment reward with.
Returns:
Return value of `imitation.util.rollout.rollout_stats()`.
"""
os.makedirs(log_dir, exist_ok=True)
sacred_util.build_sacred_symlink(log_dir, _run)
logging.basicConfig(level=logging.INFO)
logging.info("Logging to %s", log_dir)
sample_until = rollout.make_sample_until(eval_n_timesteps, eval_n_episodes)
post_wrappers = [video_wrapper_factory(log_dir, **video_kwargs)
] if videos else None
venv = util.make_vec_env(
env_name,
num_vec,
seed=_seed,
parallel=parallel,
log_dir=log_dir,
max_episode_steps=max_episode_steps,
post_wrappers=post_wrappers,
)
try:
if render:
# As of July 31, 2020, DummyVecEnv rendering only works with num_vec=1
# due to a bug on Stable Baselines 3.
venv = InteractiveRender(venv, render_fps)
if reward_type is not None:
reward_fn = load_reward(reward_type, reward_path, venv)
venv = reward_wrapper.RewardVecEnvWrapper(venv, reward_fn)
logging.info(
f"Wrapped env in reward {reward_type} from {reward_path}.")
policy = serialize.load_policy(policy_type, policy_path, venv)
trajs = rollout.generate_trajectories(policy, venv, sample_until)
return rollout.rollout_stats(trajs)
finally:
venv.close()
def rollouts_and_policy(
_seed: int,
env_name: str,
total_timesteps: int,
*,
log_dir: str = None,
num_vec: int = 8,
parallel: bool = False,
max_episode_steps: Optional[int] = None,
normalize: bool = True,
make_blank_policy_kwargs: dict = {},
reward_type: Optional[str] = None,
reward_path: Optional[str] = None,
rollout_save_interval: int = 0,
rollout_save_final: bool = False,
rollout_save_n_timesteps: Optional[int] = None,
rollout_save_n_episodes: Optional[int] = None,
policy_save_interval: int = -1,
policy_save_final: bool = True,
) -> None:
"""Trains an expert policy from scratch and saves the rollouts and policy.
At applicable training steps `step` (where step is either an integer or
"final"):
- Policies are saved to `{log_dir}/policies/{step}.pkl`.
- Rollouts are saved to `{log_dir}/rollouts/{step}.pkl`.
Args:
env_name: The gym.Env name. Loaded as VecEnv.
total_timesteps: Number of training timesteps in `model.learn()`.
log_dir: The root directory to save metrics and checkpoints to.
num_vec: Number of environments in VecEnv.
parallel: If True, then use DummyVecEnv. Otherwise use SubprocVecEnv.
max_episode_steps: If not None, then environments are wrapped by
TimeLimit so that they have at most `max_episode_steps` steps per
episode.
normalize: If True, then rescale observations and reward.
make_blank_policy_kwargs: Kwargs for `make_blank_policy`.
reward_type: If provided, then load the serialized reward of this type,
wrapping the environment in this reward. This is useful to test
whether a reward model transfers. For more information, see
`imitation.rewards.serialize.load_reward`.
reward_path: A specifier, such as a path to a file on disk, used by
reward_type to load the reward model. For more information, see
`imitation.rewards.serialize.load_reward`.
rollout_save_interval: The number of training updates in between
intermediate rollout saves. If the argument is nonpositive, then
don't save intermediate updates.
rollout_save_final: If True, then save rollouts right after training is
finished.
rollout_save_n_timesteps: The minimum number of timesteps saved in every
file. Could be more than `rollout_save_n_timesteps` because
trajectories are saved by episode rather than by transition.
Must set exactly one of `rollout_save_n_timesteps`
and `rollout_save_n_episodes`.
rollout_save_n_episodes: The number of episodes saved in every
file. Must set exactly one of `rollout_save_n_timesteps` and
`rollout_save_n_episodes`.
policy_save_interval: The number of training updates between saves. Has
the same semantics are `rollout_save_interval`.
policy_save_final: If True, then save the policy right after training is
finished.
"""
_validate_traj_generate_params(rollout_save_n_timesteps,
rollout_save_n_episodes)
with util.make_session():
tf.logging.set_verbosity(tf.logging.INFO)
sb_logger.configure(folder=osp.join(log_dir, 'rl'),
format_strs=['tensorboard', 'stdout'])
rollout_dir = osp.join(log_dir, "rollouts")
policy_dir = osp.join(log_dir, "policies")
os.makedirs(rollout_dir, exist_ok=True)
os.makedirs(policy_dir, exist_ok=True)
venv = util.make_vec_env(env_name,
num_vec,
seed=_seed,
parallel=parallel,
log_dir=log_dir,
max_episode_steps=max_episode_steps)
log_callbacks = []
with contextlib.ExitStack() as stack:
if reward_type is not None:
reward_fn_ctx = load_reward(reward_type, reward_path, venv)
reward_fn = stack.enter_context(reward_fn_ctx)
venv = RewardVecEnvWrapper(venv, reward_fn)
log_callbacks.append(venv.log_callback)
tf.logging.info(
f"Wrapped env in reward {reward_type} from {reward_path}.")
vec_normalize = None
if normalize:
venv = vec_normalize = VecNormalize(venv)
policy = util.init_rl(venv, verbose=1, **make_blank_policy_kwargs)
# Make callback to save intermediate artifacts during training.
step = 0
def callback(locals_: dict, _) -> bool:
nonlocal step
step += 1
policy = locals_['self']
# TODO(adam): make logging frequency configurable
for callback in log_callbacks:
callback(sb_logger)
if rollout_save_interval > 0 and step % rollout_save_interval == 0:
util.rollout.save(rollout_dir,
policy,
venv,
step,
n_timesteps=rollout_save_n_timesteps,
n_episodes=rollout_save_n_episodes)
if policy_save_interval > 0 and step % policy_save_interval == 0:
output_dir = os.path.join(policy_dir, f'{step:05d}')
serialize.save_stable_model(output_dir, policy,
vec_normalize)
return True # Continue training.
policy.learn(total_timesteps, callback=callback)
# Save final artifacts after training is complete.
if rollout_save_final:
util.rollout.save(rollout_dir,
policy,
venv,
"final",
n_timesteps=rollout_save_n_timesteps,
n_episodes=rollout_save_n_episodes)
if policy_save_final:
output_dir = os.path.join(policy_dir, "final")
serialize.save_stable_model(output_dir, policy, vec_normalize)
def rollouts_and_policy(
_run,
_seed: int,
env_name: str,
total_timesteps: int,
*,
log_dir: str,
num_vec: int,
parallel: bool,
max_episode_steps: Optional[int],
normalize: bool,
normalize_kwargs: dict,
init_rl_kwargs: dict,
n_episodes_eval: int,
reward_type: Optional[str],
reward_path: Optional[str],
rollout_save_final: bool,
rollout_save_n_timesteps: Optional[int],
rollout_save_n_episodes: Optional[int],
policy_save_interval: int,
policy_save_final: bool,
init_tensorboard: bool,
) -> dict:
"""Trains an expert policy from scratch and saves the rollouts and policy.
Checkpoints:
At applicable training steps `step` (where step is either an integer or
"final"):
- Policies are saved to `{log_dir}/policies/{step}/`.
- Rollouts are saved to `{log_dir}/rollouts/{step}.pkl`.
Args:
env_name: The gym.Env name. Loaded as VecEnv.
total_timesteps: Number of training timesteps in `model.learn()`.
log_dir: The root directory to save metrics and checkpoints to.
num_vec: Number of environments in VecEnv.
parallel: If True, then use DummyVecEnv. Otherwise use SubprocVecEnv.
max_episode_steps: If not None, then environments are wrapped by
TimeLimit so that they have at most `max_episode_steps` steps per
episode.
normalize: If True, then rescale observations and reward.
normalize_kwargs: kwargs for `VecNormalize`.
init_rl_kwargs: kwargs for `init_rl`.
n_episodes_eval: The number of episodes to average over when calculating
the average ground truth reward return of the final policy.
reward_type: If provided, then load the serialized reward of this type,
wrapping the environment in this reward. This is useful to test
whether a reward model transfers. For more information, see
`imitation.rewards.serialize.load_reward`.
reward_path: A specifier, such as a path to a file on disk, used by
reward_type to load the reward model. For more information, see
`imitation.rewards.serialize.load_reward`.
rollout_save_interval: The number of training updates in between
intermediate rollout saves. If the argument is nonpositive, then
don't save intermediate updates.
rollout_save_final: If True, then save rollouts right after training is
finished.
rollout_save_n_timesteps: The minimum number of timesteps saved in every
file. Could be more than `rollout_save_n_timesteps` because
trajectories are saved by episode rather than by transition.
Must set exactly one of `rollout_save_n_timesteps`
and `rollout_save_n_episodes`.
rollout_save_n_episodes: The number of episodes saved in every
file. Must set exactly one of `rollout_save_n_timesteps` and
`rollout_save_n_episodes`.
policy_save_interval: The number of training updates between saves. Has
the same semantics are `rollout_save_interval`.
policy_save_final: If True, then save the policy right after training is
finished.
init_tensorboard: If True, then write tensorboard logs to {log_dir}/sb_tb
and "output/summary/...".
Returns:
The return value of `rollout_stats()` using the final policy.
"""
os.makedirs(log_dir, exist_ok=True)
sacred_util.build_sacred_symlink(log_dir, _run)
sample_until = rollout.make_sample_until(rollout_save_n_timesteps,
rollout_save_n_episodes)
eval_sample_until = rollout.min_episodes(n_episodes_eval)
logging.basicConfig(level=logging.INFO)
logger.configure(folder=osp.join(log_dir, "rl"),
format_strs=["tensorboard", "stdout"])
rollout_dir = osp.join(log_dir, "rollouts")
policy_dir = osp.join(log_dir, "policies")
os.makedirs(rollout_dir, exist_ok=True)
os.makedirs(policy_dir, exist_ok=True)
if init_tensorboard:
# sb_tensorboard_dir = osp.join(log_dir, "sb_tb")
# Convert sacred's ReadOnlyDict to dict so we can modify on next line.
init_rl_kwargs = dict(init_rl_kwargs)
# init_rl_kwargs["tensorboard_log"] = sb_tensorboard_dir
# FIXME(sam): this is another hack to prevent SB3 from configuring the
# logger on the first .learn() call. Remove it once SB3 issue #109 is
# fixed.
init_rl_kwargs["tensorboard_log"] = None
venv = util.make_vec_env(
env_name,
num_vec,
seed=_seed,
parallel=parallel,
log_dir=log_dir,
max_episode_steps=max_episode_steps,
)
callback_objs = []
if reward_type is not None:
reward_fn = load_reward(reward_type, reward_path, venv)
venv = RewardVecEnvWrapper(venv, reward_fn)
callback_objs.append(venv.make_log_callback())
logging.info(
f"Wrapped env in reward {reward_type} from {reward_path}.")
vec_normalize = None
if normalize:
venv = vec_normalize = VecNormalize(venv, **normalize_kwargs)
if policy_save_interval > 0:
save_policy_callback = serialize.SavePolicyCallback(
policy_dir, vec_normalize)
save_policy_callback = callbacks.EveryNTimesteps(
policy_save_interval, save_policy_callback)
callback_objs.append(save_policy_callback)
callback = callbacks.CallbackList(callback_objs)
policy = util.init_rl(venv, verbose=1, **init_rl_kwargs)
policy.learn(total_timesteps, callback=callback)
# Save final artifacts after training is complete.
if rollout_save_final:
save_path = osp.join(rollout_dir, "final.pkl")
rollout.rollout_and_save(save_path, policy, venv, sample_until)
if policy_save_final:
output_dir = os.path.join(policy_dir, "final")
serialize.save_stable_model(output_dir, policy, vec_normalize)
# Final evaluation of expert policy.
trajs = rollout.generate_trajectories(policy, venv, eval_sample_until)
stats = rollout.rollout_stats(trajs)
return stats