def test_reward_overwrite():
"""Test that reward wrapper actually overwrites base rewards."""
env_name = "Pendulum-v0"
num_envs = 3
env = util.make_vec_env(env_name, num_envs)
reward_fn = FunkyReward()
wrapped_env = reward_wrapper.RewardVecEnvWrapper(env, reward_fn)
policy = RandomPolicy(env.observation_space, env.action_space)
sample_until = rollout.min_episodes(10)
default_stats = rollout.rollout_stats(
rollout.generate_trajectories(policy, env, sample_until))
wrapped_stats = rollout.rollout_stats(
rollout.generate_trajectories(policy, wrapped_env, sample_until))
# Pendulum-v0 always has negative rewards
assert default_stats["return_max"] < 0
# ours gives between 1 * traj_len and num_envs * traj_len reward
# (trajectories are all constant length of 200 in Pendulum)
steps = wrapped_stats["len_mean"]
assert wrapped_stats["return_min"] == 1 * steps
assert wrapped_stats["return_max"] == num_envs * steps
# check that wrapped reward is negative (all pendulum rewards is negative)
# and other rewards are non-negative
rand_act, _, _, _ = policy.step(wrapped_env.reset())
_, rew, _, infos = wrapped_env.step(rand_act)
assert np.all(rew >= 0)
assert np.all([info_dict["wrapped_env_rew"] < 0 for info_dict in infos])
def eval_policy(
rl_algo: Union[base_class.BaseAlgorithm, policies.BasePolicy],
venv: vec_env.VecEnv,
n_episodes_eval: int,
) -> Mapping[str, float]:
"""Evaluation of imitation learned policy.
Args:
rl_algo: Algorithm to evaluate.
venv: Environment to evaluate on.
n_episodes_eval: The number of episodes to average over when calculating
the average episode reward of the imitation policy for return.
Returns:
A dictionary with two keys. "imit_stats" gives the return value of
`rollout_stats()` on rollouts test-reward-wrapped environment, using the final
policy (remember that the ground-truth reward can be recovered from the
"monitor_return" key). "expert_stats" gives the return value of
`rollout_stats()` on the expert demonstrations loaded from `rollout_path`.
"""
sample_until_eval = rollout.make_min_episodes(n_episodes_eval)
trajs = rollout.generate_trajectories(
rl_algo,
venv,
sample_until=sample_until_eval,
)
return rollout.rollout_stats(trajs)
def test_density_reward(density_type, is_stationary):
# test on Pendulum rather than Cartpole because I don't handle episodes that
# terminate early yet (see issue #40)
env_name = "Pendulum-v0"
env = util.make_vec_env(env_name, 2)
# construct density-based reward from expert rollouts
rollout_path = "tests/data/expert_models/pendulum_0/rollouts/final.pkl"
expert_trajectories_all = types.load(rollout_path)
n_experts = len(expert_trajectories_all)
expert_trajectories_train = expert_trajectories_all[: n_experts // 2]
reward_fn = DensityReward(
trajectories=expert_trajectories_train,
density_type=density_type,
kernel="gaussian",
obs_space=env.observation_space,
act_space=env.action_space,
is_stationary=is_stationary,
kernel_bandwidth=0.2,
standardise_inputs=True,
)
# check that expert policy does better than a random policy under our reward
# function
random_policy = RandomPolicy(env.observation_space, env.action_space)
sample_until = rollout.min_episodes(n_experts // 2)
random_trajectories = rollout.generate_trajectories(
random_policy, env, sample_until=sample_until
)
expert_trajectories_test = expert_trajectories_all[n_experts // 2 :]
random_score = score_trajectories(random_trajectories, reward_fn)
expert_score = score_trajectories(expert_trajectories_test, reward_fn)
assert expert_score > random_score
def _sample_fixed_length_trajectories(
episode_lengths: Sequence[int], min_episodes: int, **kwargs,
) -> Sequence[types.Trajectory]:
venv = vec_env.DummyVecEnv(
[functools.partial(TerminalSentinelEnv, length) for length in episode_lengths]
)
policy = RandomPolicy(venv.observation_space, venv.action_space)
sample_until = rollout.min_episodes(min_episodes)
trajectories = rollout.generate_trajectories(
policy, venv, sample_until=sample_until, **kwargs,
)
return trajectories
def test_policy(self, *, min_episodes: int = 10) -> dict:
"""Test current imitation policy on environment & give some rollout stats.
Args:
min_episodes: Minimum number of rolled-out episodes.
Returns:
rollout statistics collected by `imitation.utils.rollout.rollout_stats()`.
"""
trajs = rollout.generate_trajectories(
self.policy,
self.env,
sample_until=rollout.min_episodes(min_episodes))
reward_stats = rollout.rollout_stats(trajs)
return reward_stats
def test_rollout_stats():
"""Applying `ObsRewIncrementWrapper` halves the reward mean.
`rollout_stats` should reflect this.
"""
env = gym.make("CartPole-v1")
env = bench.Monitor(env, None)
env = ObsRewHalveWrapper(env)
venv = vec_env.DummyVecEnv([lambda: env])
with serialize.load_policy("zero", "UNUSED", venv) as policy:
trajs = rollout.generate_trajectories(policy, venv, rollout.min_episodes(10))
s = rollout.rollout_stats(trajs)
np.testing.assert_allclose(s["return_mean"], s["monitor_return_mean"] / 2)
np.testing.assert_allclose(s["return_std"], s["monitor_return_std"] / 2)
np.testing.assert_allclose(s["return_min"], s["monitor_return_min"] / 2)
np.testing.assert_allclose(s["return_max"], s["monitor_return_max"] / 2)
def generate_trajectories(venv: vec_env.VecEnv, policy: policies.BasePolicy,
trajectory_length: int,
num_trajectories: int) -> Sequence[types.Trajectory]:
"""Rollouts policy in venv collecting num_trajectories segments.
Complete episodes are collected. An episode of length N is split into
(N // trajectory_length) trajectories of trajectory_length. If N is not
exactly divisible by trajectory_length, we start from a random offset between
0 and (N % trajectory_length).
Arguments:
venv: The environment to generate trajectories in.
policy: The policy to generate trajectories with.
trajectory_length: The length of each trajectory.
num_trajectories: The number of trajectories to sample.
Returns:
A Sequence of num_trajectories trajectories, each of length
trajectory_length.
"""
def sample_until(episodes: Sequence[types.Trajectory]):
"""Computes whether a full batch of data has been collected."""
episode_lengths = np.array([len(t.acts) for t in episodes])
num_trajs = episode_lengths // trajectory_length
return np.sum(num_trajs) >= num_trajectories
episodes = rollout.generate_trajectories(policy, venv, sample_until)
trajectories = []
for episode in episodes:
ep_len = len(episode.acts)
remainder = ep_len % trajectory_length
offset = 0 if remainder == 0 else np.random.randint(remainder)
n_trajs = ep_len // trajectory_length
for i in range(n_trajs):
start = offset + i * trajectory_length
end = start + trajectory_length
trajectory = _slice_trajectory(episode, start, end)
trajectories.append(trajectory)
# We may collect too much data due to episode boundaries, truncate.
trajectories = trajectories[:num_trajectories]
assert len(trajectories) == num_trajectories
return trajectories
def test_potential_shaping_cycle(graph,
session,
venv,
potential_cls,
discount: float,
num_episodes: int = 10) -> None:
"""Test that potential shaping is constant on any fixed-length cycle.
Specifically, performs rollouts of a random policy in the environment.
Fixes the starting state for each trajectory at the all-zero state.
Then computes episode return, and checks they're all equal.
Requires environment be fixed length, otherwise the episode return will vary
(except in the undiscounted case).
"""
policy = base.RandomPolicy(venv.observation_space, venv.action_space)
trajectories = rollout.generate_trajectories(
policy, venv, sample_until=rollout.min_episodes(num_episodes))
transitions = rollout.flatten_trajectories(trajectories)
# Make initial state fixed as all-zero.
# Note don't need to change final state, since `dones` being `True` should
# force potential to be zero at those states.
obs = np.array(transitions.obs)
idxs = np.where(transitions.dones)[0] + 1
idxs = np.pad(idxs[:-1], (1, 0), "constant")
obs[idxs, :] = 0
transitions = dataclasses.replace(transitions, obs=obs)
with graph.as_default(), session.as_default():
reward_model = potential_cls(venv.observation_space,
venv.action_space,
discount=discount)
session.run(tf.global_variables_initializer())
rews = rewards.evaluate_models({"m": reward_model}, transitions)
rets = rewards.compute_return_from_rews(rews,
transitions.dones,
discount=discount)["m"]
if discount == 1.0:
assert np.allclose(rets, 0.0, atol=1e-5)
assert np.allclose(rets, np.mean(rets), atol=1e-5)
def test_policy(self, *, n_trajectories=10, true_reward=True):
"""Test current imitation policy on environment & give some rollout stats.
Args:
n_trajectories (int): number of rolled-out trajectories.
true_reward (bool): should this use ground truth reward from underlying
environment (True), or imitation reward (False)?
Returns:
dict: rollout statistics collected by
`imitation.utils.rollout.rollout_stats()`.
"""
self.imitation_trainer.set_env(self.venv)
trajs = rollout.generate_trajectories(
self.imitation_trainer,
self.venv if true_reward else self.wrapped_env,
sample_until=rollout.min_episodes(n_trajectories),
)
reward_stats = rollout.rollout_stats(trajs)
return reward_stats
def test_unwrap_traj():
"""Check that unwrap_traj reverses `ObsRewIncrementWrapper`.
Also check that unwrapping twice is a no-op.
"""
env = gym.make("CartPole-v1")
env = wrappers.RolloutInfoWrapper(env)
env = ObsRewHalveWrapper(env)
venv = vec_env.DummyVecEnv([lambda: env])
with serialize.load_policy("zero", "UNUSED", venv) as policy:
trajs = rollout.generate_trajectories(policy, venv, rollout.min_episodes(10))
trajs_unwrapped = [rollout.unwrap_traj(t) for t in trajs]
trajs_unwrapped_twice = [rollout.unwrap_traj(t) for t in trajs_unwrapped]
for t, t_unwrapped in zip(trajs, trajs_unwrapped):
np.testing.assert_allclose(t.acts, t_unwrapped.acts)
np.testing.assert_allclose(t.obs, t_unwrapped.obs / 2)
np.testing.assert_allclose(t.rews, t_unwrapped.rews / 2)
for t1, t2 in zip(trajs_unwrapped, trajs_unwrapped_twice):
np.testing.assert_equal(t1.acts, t2.acts)
np.testing.assert_equal(t1.obs, t2.obs)
np.testing.assert_equal(t1.rews, t2.rews)
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 f(
sample_until: rollout.GenTrajTerminationFn
) -> Sequence[types.Trajectory]:
return rollout.generate_trajectories(policy,
venv,
sample_until=sample_until)
def f(total_episodes: int) -> types.Transitions:
return rollout.generate_trajectories(
policy, venv, sample_until=rollout.min_episodes(total_episodes))
def train(
_run,
_seed: int,
algorithm: str,
env_name: str,
num_vec: int,
parallel: bool,
max_episode_steps: Optional[int],
rollout_path: str,
n_expert_demos: Optional[int],
log_dir: str,
total_timesteps: int,
n_episodes_eval: int,
init_tensorboard: bool,
checkpoint_interval: int,
gen_batch_size: int,
init_rl_kwargs: Mapping,
algorithm_kwargs: Mapping[str, Mapping],
discrim_net_kwargs: Mapping[str, Mapping],
) -> dict:
"""Train an adversarial-network-based imitation learning algorithm.
Checkpoints:
- DiscrimNets are saved to `f"{log_dir}/checkpoints/{step}/discrim/"`,
where step is either the training round or "final".
- Generator policies are saved to `f"{log_dir}/checkpoints/{step}/gen_policy/"`.
Args:
_seed: Random seed.
algorithm: A case-insensitive string determining which adversarial imitation
learning algorithm is executed. Either "airl" or "gail".
env_name: The environment to train in.
num_vec: Number of `gym.Env` to vectorize.
parallel: Whether to use "true" parallelism. If True, then use `SubProcVecEnv`.
Otherwise, use `DummyVecEnv` which steps through environments serially.
max_episode_steps: If not None, then a TimeLimit wrapper is applied to each
environment to artificially limit the maximum number of timesteps in an
episode.
rollout_path: Path to pickle containing list of Trajectories. Used as
expert demonstrations.
n_expert_demos: The number of expert trajectories to actually use
after loading them from `rollout_path`.
If None, then use all available trajectories.
If `n_expert_demos` is an `int`, then use exactly `n_expert_demos`
trajectories, erroring if there aren't enough trajectories. If there are
surplus trajectories, then use the first `n_expert_demos` trajectories and
drop the rest.
log_dir: Directory to save models and other logging to.
total_timesteps: The number of transitions to sample from the environment
during training.
n_episodes_eval: The number of episodes to average over when calculating
the average episode reward of the imitation policy for return.
init_tensorboard: If True, then write tensorboard logs to `{log_dir}/sb_tb`.
checkpoint_interval: Save the discriminator and generator models every
`checkpoint_interval` rounds and after training is complete. If 0,
then only save weights after training is complete. If <0, then don't
save weights at all.
gen_batch_size: Batch size for generator updates. Sacred automatically uses
this to calculate `n_steps` in `init_rl_kwargs`. In the script body, this
is only used in sanity checks.
init_rl_kwargs: Keyword arguments for `init_rl`, the RL algorithm initialization
utility function.
algorithm_kwargs: Keyword arguments for the `GAIL` or `AIRL` constructor
that can apply to either constructor. Unlike a regular kwargs argument, this
argument can only have the following keys: "shared", "airl", and "gail".
`algorithm_kwargs["airl"]`, if it is provided, is a kwargs `Mapping` passed
to the `AIRL` constructor when `algorithm == "airl"`. Likewise
`algorithm_kwargs["gail"]` is passed to the `GAIL` constructor when
`algorithm == "gail"`. `algorithm_kwargs["shared"]`, if provided, is passed
to both the `AIRL` and `GAIL` constructors. Duplicate keyword argument keys
between `algorithm_kwargs["shared"]` and `algorithm_kwargs["airl"]` (or
"gail") leads to an error.
discrim_net_kwargs: Keyword arguments for the `DiscrimNet` constructor. Unlike a
regular kwargs argument, this argument can only have the following keys:
"shared", "airl", "gail". These keys have the same meaning as they do in
`algorithm_kwargs`.
Returns:
A dictionary with two keys. "imit_stats" gives the return value of
`rollout_stats()` on rollouts test-reward-wrapped environment, using the final
policy (remember that the ground-truth reward can be recovered from the
"monitor_return" key). "expert_stats" gives the return value of
`rollout_stats()` on the expert demonstrations loaded from `rollout_path`.
"""
if gen_batch_size % num_vec != 0:
raise ValueError(
f"num_vec={num_vec} must evenly divide gen_batch_size={gen_batch_size}."
)
allowed_keys = {"shared", "gail", "airl"}
if not discrim_net_kwargs.keys() <= allowed_keys:
raise ValueError(
f"Invalid discrim_net_kwargs.keys()={discrim_net_kwargs.keys()}. "
f"Allowed keys: {allowed_keys}"
)
if not algorithm_kwargs.keys() <= allowed_keys:
raise ValueError(
f"Invalid discrim_net_kwargs.keys()={algorithm_kwargs.keys()}. "
f"Allowed keys: {allowed_keys}"
)
if not os.path.exists(rollout_path):
raise ValueError(f"File at rollout_path={rollout_path} does not exist.")
expert_trajs = types.load(rollout_path)
if n_expert_demos is not None:
if not len(expert_trajs) >= n_expert_demos:
raise ValueError(
f"Want to use n_expert_demos={n_expert_demos} trajectories, but only "
f"{len(expert_trajs)} are available via {rollout_path}."
)
expert_trajs = expert_trajs[:n_expert_demos]
expert_transitions = rollout.flatten_trajectories(expert_trajs)
total_timesteps = int(total_timesteps)
logging.info("Logging to %s", log_dir)
logger.configure(log_dir, ["tensorboard", "stdout"])
os.makedirs(log_dir, exist_ok=True)
sacred_util.build_sacred_symlink(log_dir, _run)
venv = util.make_vec_env(
env_name,
num_vec,
seed=_seed,
parallel=parallel,
log_dir=log_dir,
max_episode_steps=max_episode_steps,
)
# if init_tensorboard:
# tensorboard_log = osp.join(log_dir, "sb_tb")
# else:
# tensorboard_log = None
gen_algo = util.init_rl(
# FIXME(sam): ignoring tensorboard_log is a hack to prevent SB3 from
# re-configuring the logger (SB3 issue #109). See init_rl() for details.
# TODO(shwang): Let's get rid of init_rl after SB3 issue #109 is fixed?
# Besides sidestepping #109, init_rl is just a stub function.
venv,
**init_rl_kwargs,
)
discrim_kwargs_shared = discrim_net_kwargs.get("shared", {})
discrim_kwargs_algo = discrim_net_kwargs.get(algorithm, {})
final_discrim_kwargs = dict(**discrim_kwargs_shared, **discrim_kwargs_algo)
algorithm_kwargs_shared = algorithm_kwargs.get("shared", {})
algorithm_kwargs_algo = algorithm_kwargs.get(algorithm, {})
final_algorithm_kwargs = dict(
**algorithm_kwargs_shared,
**algorithm_kwargs_algo,
)
if algorithm.lower() == "gail":
algo_cls = adversarial.GAIL
elif algorithm.lower() == "airl":
algo_cls = adversarial.AIRL
else:
raise ValueError(f"Invalid value algorithm={algorithm}.")
trainer = algo_cls(
venv=venv,
expert_data=expert_transitions,
gen_algo=gen_algo,
log_dir=log_dir,
discrim_kwargs=final_discrim_kwargs,
**final_algorithm_kwargs,
)
def callback(round_num):
if checkpoint_interval > 0 and round_num % checkpoint_interval == 0:
save(trainer, os.path.join(log_dir, "checkpoints", f"{round_num:05d}"))
trainer.train(total_timesteps, callback)
# Save final artifacts.
if checkpoint_interval >= 0:
save(trainer, os.path.join(log_dir, "checkpoints", "final"))
# Final evaluation of imitation policy.
results = {}
sample_until_eval = rollout.min_episodes(n_episodes_eval)
trajs = rollout.generate_trajectories(
trainer.gen_algo, trainer.venv_train_norm, sample_until=sample_until_eval
)
results["expert_stats"] = rollout.rollout_stats(expert_trajs)
results["imit_stats"] = rollout.rollout_stats(trajs)
return results
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 = rollout.make_sample_until(rollout_save_n_timesteps,
rollout_save_n_episodes)
eval_sample_until = rollout.min_episodes(n_episodes_eval)
with networks.make_session():
tf.logging.set_verbosity(tf.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
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")
rollout.rollout_and_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")
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
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,
)
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 trajs(venv, rand_policy):
return rollout.generate_trajectories(rand_policy,
venv,
sample_until=rollout.min_episodes(5))
def train(
_run,
_seed: int,
env_name: str,
rollout_path: str,
n_expert_demos: Optional[int],
log_dir: str,
init_trainer_kwargs: dict,
total_timesteps: int,
n_episodes_eval: int,
init_tensorboard: bool,
checkpoint_interval: int,
) -> dict:
"""Train an adversarial-network-based imitation learning algorithm.
Plots (turn on using `plot_interval > 0`):
- Plot discriminator loss during discriminator training steps in blue and
discriminator loss during generator training steps in red.
- Plot the performance of the generator policy versus the performance of
a random policy. Also plot the performance of an expert policy if that is
provided in the arguments.
Checkpoints:
- DiscrimNets are saved to f"{log_dir}/checkpoints/{step}/discrim/",
where step is either the training epoch or "final".
- Generator policies are saved to
f"{log_dir}/checkpoints/{step}/gen_policy/".
Args:
_seed: Random seed.
env_name: The environment to train in.
rollout_path: Path to pickle containing list of Trajectories. Used as
expert demonstrations.
n_expert_demos: The number of expert trajectories to actually use
after loading them from `rollout_path`.
If None, then use all available trajectories.
If `n_expert_demos` is an `int`, then use exactly `n_expert_demos`
trajectories, erroring if there aren't enough trajectories. If there are
surplus trajectories, then use the
first `n_expert_demos` trajectories and drop the rest.
log_dir: Directory to save models and other logging to.
init_trainer_kwargs: Keyword arguments passed to `init_trainer`,
used to initialize the trainer.
total_timesteps: The number of transitions to sample from the environment
during training.
n_episodes_eval: The number of episodes to average over when calculating
the average episode reward of the imitation policy for return.
plot_interval: The number of epochs between each plot. If negative,
then plots are disabled. If zero, then only plot at the end of training.
n_plot_episodes: The number of episodes averaged over when
calculating the average episode reward of a policy for the performance
plots.
extra_episode_data_interval: Usually mean episode rewards are calculated
immediately before every plot. Set this parameter to a nonnegative number
to also add episode reward data points every
`extra_episodes_data_interval` epochs.
show_plots: Figures are always saved to `f"{log_dir}/plots/*.png"`. If
`show_plots` is True, then also show plots as they are created.
init_tensorboard: If True, then write tensorboard logs to `{log_dir}/sb_tb`.
checkpoint_interval: Save the discriminator and generator models every
`checkpoint_interval` epochs and after training is complete. If 0,
then only save weights after training is complete. If <0, then don't
save weights at all.
Returns:
A dictionary with two keys. "imit_stats" gives the return value of
`rollout_stats()` on rollouts test-reward-wrapped
environment, using the final policy (remember that the ground-truth reward
can be recovered from the "monitor_return" key). "expert_stats" gives the
return value of `rollout_stats()` on the expert demonstrations loaded from
`rollout_path`.
"""
total_timesteps = int(total_timesteps)
tf.logging.info("Logging to %s", log_dir)
os.makedirs(log_dir, exist_ok=True)
sacred_util.build_sacred_symlink(log_dir, _run)
# Calculate stats for expert rollouts. Used for plot and return value.
expert_trajs = types.load(rollout_path)
if n_expert_demos is not None:
assert len(expert_trajs) >= n_expert_demos
expert_trajs = expert_trajs[:n_expert_demos]
expert_stats = rollout.rollout_stats(expert_trajs)
with networks.make_session():
if init_tensorboard:
sb_tensorboard_dir = osp.join(log_dir, "sb_tb")
kwargs = init_trainer_kwargs
kwargs["init_rl_kwargs"] = kwargs.get("init_rl_kwargs", {})
kwargs["init_rl_kwargs"]["tensorboard_log"] = sb_tensorboard_dir
trainer = init_trainer(env_name,
expert_trajs,
seed=_seed,
log_dir=log_dir,
**init_trainer_kwargs)
def callback(epoch):
if checkpoint_interval > 0 and epoch % checkpoint_interval == 0:
save(trainer,
os.path.join(log_dir, "checkpoints", f"{epoch:05d}"))
trainer.train(total_timesteps, callback)
# Save final artifacts.
if checkpoint_interval >= 0:
save(trainer, os.path.join(log_dir, "checkpoints", "final"))
# Final evaluation of imitation policy.
results = {}
sample_until_eval = rollout.min_episodes(n_episodes_eval)
trajs = rollout.generate_trajectories(trainer.gen_policy,
trainer.venv_test,
sample_until=sample_until_eval)
results["imit_stats"] = rollout.rollout_stats(trajs)
results["expert_stats"] = expert_stats
return results
