def train(args, seeds):
args.cuda = not args.no_cuda and torch.cuda.is_available()
device = torch.device("cuda:0" if args.cuda else "cpu")
if 'cuda' in device.type:
print('Using CUDA\n')
torch.set_num_threads(1)
utils.seed(args.seed)
# Configure logging
if args.xpid is None:
args.xpid = "lr-%s" % time.strftime("%Y%m%d-%H%M%S")
log_dir = os.path.expandvars(os.path.expanduser(args.log_dir))
plogger = FileWriter(
xpid=args.xpid, xp_args=args.__dict__, rootdir=log_dir,
seeds=seeds,
)
stdout_logger = HumanOutputFormat(sys.stdout)
checkpointpath = os.path.expandvars(
os.path.expanduser("%s/%s/%s" % (log_dir, args.xpid, "model.tar"))
)
global last_checkpoint_time
current_update_count = 0
initial_update_count = 0
last_logged_update_count_at_restart = -1
start_level = 0
num_levels = 1
level_sampler_args = dict(
num_actors=args.num_processes,
strategy=args.level_replay_strategy,
max_score_coef=args.level_replay_max_score_coef,
replay_schedule=args.level_replay_schedule,
score_transform=args.level_replay_score_transform,
temperature=args.level_replay_temperature,
eps=args.level_replay_eps,
rho=args.level_replay_rho,
replay_prob=args.level_replay_prob,
alpha=args.level_replay_alpha,
staleness_coef=args.staleness_coef,
staleness_transform=args.staleness_transform,
staleness_temperature=args.staleness_temperature,
sample_full_distribution=args.train_full_distribution,
seed_buffer_size=args.level_replay_seed_buffer_size,
seed_buffer_priority=args.level_replay_seed_buffer_priority,
tscl_window_size=args.tscl_window_size)
level_sampler_secondary_args = {}
if args.level_replay_secondary_strategy:
level_sampler_secondary_args = dict(
strategy=args.level_replay_secondary_strategy,
score_transform=args.level_replay_secondary_score_transform,
temperature=args.level_replay_secondary_temperature,
eps=args.level_replay_secondary_eps,
staleness_coef=args.secondary_staleness_coef,
staleness_transform=args.secondary_staleness_transform,
staleness_temperature=args.secondary_staleness_temperature,)
args_tmp = level_sampler_args.copy()
args_tmp.update(level_sampler_secondary_args)
level_sampler_secondary_args = args_tmp
envs, level_sampler, secondary_level_sampler = make_lr_venv(
num_envs=args.num_processes, env_name=args.env_name,
seeds=seeds, device=device,
num_levels=num_levels,
start_level=start_level,
no_ret_normalization=args.no_ret_normalization,
distribution_mode=args.distribution_mode,
paint_vel_info=args.paint_vel_info,
level_sampler_args=level_sampler_args,
level_sampler_secondary_args=level_sampler_secondary_args,
level_replay_strategy_mix_coef=args.level_replay_strategy_mix_coef)
is_minigrid = args.env_name.startswith('MiniGrid')
actor_critic = model_for_env_name(args, envs)
actor_critic.to(device)
rollouts = RolloutStorage(args.num_steps, args.num_processes,
envs.observation_space.shape, envs.action_space,
actor_critic.recurrent_hidden_state_size)
batch_size = int(args.num_processes * args.num_steps / args.num_mini_batch)
aug_id = None
if args.algo == 'ucb' or args.use_ucb:
print('Using UCB')
aug_id = data_augs.Identity
aug_list = [aug_to_func[t](batch_size=batch_size)
for t in list(aug_to_func.keys())]
mix_alpha = None if not args.use_mixreg else args.mixreg_alpha
agent = algo.UCBDrAC(
actor_critic,
args.clip_param,
args.ppo_epoch,
args.num_mini_batch,
args.value_loss_coef,
args.entropy_coef,
lr=args.lr,
eps=args.eps,
max_grad_norm=args.max_grad_norm,
aug_list=aug_list,
aug_id=aug_id,
aug_coef=args.aug_coef,
num_aug_types=len(list(aug_to_func.keys())),
ucb_exploration_coef=args.ucb_exploration_coef,
ucb_window_length=args.ucb_window_length,
mix_alpha=mix_alpha,
log_grad_norm=args.log_grad_norm)
elif args.algo == 'mixreg' or args.use_mixreg:
agent = algo.MixRegPPO(
actor_critic,
args.clip_param,
args.ppo_epoch,
args.num_mini_batch,
args.value_loss_coef,
args.entropy_coef,
lr=args.lr,
eps=args.eps,
max_grad_norm=args.max_grad_norm,
log_grad_norm=args.log_grad_norm,
mix_alpha=args.mixreg_alpha)
else:
agent = algo.PPO(
actor_critic,
args.clip_param,
args.ppo_epoch,
args.num_mini_batch,
args.value_loss_coef,
args.entropy_coef,
lr=args.lr,
eps=args.eps,
max_grad_norm=args.max_grad_norm,
log_grad_norm=args.log_grad_norm)
level_seeds = torch.zeros(args.num_processes)
if level_sampler:
obs, level_seeds = envs.reset()
else:
obs = envs.reset()
level_seeds = level_seeds.unsqueeze(-1)
rollouts.obs[0].copy_(obs)
rollouts.to(device)
episode_rewards = deque(maxlen=10)
num_updates = int(
args.num_env_steps) // args.num_steps // args.num_processes
def checkpoint():
if args.disable_checkpoint:
return
logging.info("Saving checkpoint to %s", checkpointpath)
checkpoint_states = {
"model_state_dict": actor_critic.state_dict(),
"optimizer_state_dict": agent.optimizer.state_dict(),
"rollouts": rollouts,
"episode_rewards": episode_rewards,
"level_sampler": level_sampler,
"current_update_count": current_update_count
}
if hasattr(agent, 'bandit_state_dict'):
checkpoint_states.update({
"bandit_state_dict": agent.bandit_state_dict()
})
torch.save(
checkpoint_states,
checkpointpath
)
# Load checkpoint
if args.checkpoint and os.path.exists(checkpointpath):
checkpoint_states = torch.load(checkpointpath)
actor_critic.load_state_dict(checkpoint_states['model_state_dict'])
agent.optimizer.load_state_dict(checkpoint_states["optimizer_state_dict"])
rollouts = checkpoint_states["rollouts"]
episode_rewards = checkpoint_states["episode_rewards"]
level_sampler = checkpoint_states["level_sampler"]
current_update_count = checkpoint_states["current_update_count"]
initial_update_count = current_update_count
last_logged_update_count_at_restart = plogger.latest_tick() + 1 # ticks are 0-indexed updates
if hasattr(agent, 'load_bandit_state_dict'):
agent.load_bandit_state_dict(checkpoint_states["bandit_state_dict"])
logging.info(f"Resuming preempted job from update: {current_update_count}\n")
timer = timeit.default_timer
update_start_time = timer()
agent_id = 0 # np.random.choice(range(actor_critic.ensemble_size))
for j in range(initial_update_count, num_updates):
actor_critic.train()
for step in range(args.num_steps):
# Sample actions
with torch.no_grad():
obs_id = rollouts.obs[step]
if aug_id:
obs_id = aug_id(obs_id)
value, action, action_log_dist, recurrent_hidden_states = actor_critic.act(
obs_id, rollouts.recurrent_hidden_states[step], rollouts.masks[step], agent_id=agent_id)
action_log_prob = action_log_dist.gather(-1, action)
uncertainties = actor_critic.get_uncertainty(obs_id, rollouts.recurrent_hidden_states[step], rollouts.masks[step])
# Observed reward and next obs
obs, reward, done, infos = envs.step(action)
# Reset all done levels by sampling from level sampler
for i, info in enumerate(infos):
if 'episode' in info.keys():
episode_rewards.append(info['episode']['r'])
if level_sampler:
level_seeds[i][0] = info['level_seed']
# If done then clean the history of observations.
masks = torch.FloatTensor(
[[0.0] if done_ else [1.0] for done_ in done])
bad_masks = torch.FloatTensor(
[[0.0] if 'bad_transition' in info.keys() else [1.0]
for info in infos])
rollouts.insert(
obs, recurrent_hidden_states,
action, action_log_prob, action_log_dist,
value, reward, masks, bad_masks, uncertainties, level_seeds)
with torch.no_grad():
obs_id = rollouts.obs[-1]
if aug_id:
obs_id = aug_id(obs_id)
next_value = actor_critic.get_value(
obs_id, rollouts.recurrent_hidden_states[-1],
rollouts.masks[-1]).detach()
rollouts.compute_returns(next_value, args.use_gae, args.gamma, args.gae_lambda)
# Update level sampler
if level_sampler:
level_sampler.update_with_rollouts(rollouts)
if secondary_level_sampler:
secondary_level_sampler.update_with_rollouts(rollouts)
if args.use_ucb and j > 0:
agent.update_ucb_values(rollouts)
value_loss, action_loss, dist_entropy, info = agent.update(rollouts)
rollouts.after_update()
if level_sampler:
level_sampler.after_update()
if secondary_level_sampler:
secondary_level_sampler.after_update()
current_update_count = j + 1
# ==== Everything below here is for logging + checkpointing ====
if current_update_count <= last_logged_update_count_at_restart:
continue
# Log stats every log_interval updates or if it is the last update
if (j % args.log_interval == 0 and len(episode_rewards) > 1) or j == num_updates - 1:
total_num_steps = (j + 1) * args.num_processes * args.num_steps
update_end_time = timer()
num_interval_updates = 1 if j == 0 else args.log_interval
sps = num_interval_updates*(args.num_processes * args.num_steps) / (update_end_time - update_start_time)
update_start_time = update_end_time
logging.info(f"\nUpdate {j} done, {total_num_steps} steps\n ")
logging.info(f"\nEvaluating on {args.num_test_seeds} test levels...\n ")
eval_episode_rewards = evaluate(
args,
actor_critic,
args.num_test_seeds,
device,
aug_id=aug_id)
logging.info(f"\nEvaluating on {args.num_test_seeds} train levels...\n ")
train_eval_episode_rewards = evaluate(
args,
actor_critic,
args.num_test_seeds, # Use same number of levels for evaluating train and test seeds
device,
start_level=0,
num_levels=args.num_train_seeds,
seeds=seeds,
aug_id=aug_id)
stats = {
"step": total_num_steps,
"pg_loss": action_loss,
"value_loss": value_loss,
"dist_entropy": dist_entropy,
"train:mean_episode_return": np.mean(episode_rewards),
"train:median_episode_return": np.median(episode_rewards),
"test:mean_episode_return": np.mean(eval_episode_rewards),
"test:median_episode_return": np.median(eval_episode_rewards),
"train_eval:mean_episode_return": np.mean(train_eval_episode_rewards),
"train_eval:median_episode_return": np.median(train_eval_episode_rewards),
"sps": sps,
}
if args.log_grad_norm:
stats.update({
"mean_grad_norm": np.mean(info['grad_norms'])
})
if is_minigrid:
stats["train:success_rate"] = np.mean(np.array(episode_rewards) > 0)
stats["train_eval:success_rate"] = np.mean(np.array(train_eval_episode_rewards) > 0)
stats["test:success_rate"] = np.mean(np.array(eval_episode_rewards) > 0)
if j == num_updates - 1:
logging.info(f"\nLast update: Evaluating on {args.num_test_seeds} test levels...\n ")
final_eval_episode_rewards = evaluate(args, actor_critic, args.final_num_test_seeds, device)
mean_final_eval_episode_rewards = np.mean(final_eval_episode_rewards)
median_final_eval_episide_rewards = np.median(final_eval_episode_rewards)
plogger.log_final_test_eval({
'num_test_seeds': args.final_num_test_seeds,
'mean_episode_return': mean_final_eval_episode_rewards,
'median_episode_return': median_final_eval_episide_rewards
})
plogger.log(stats)
if args.verbose:
stdout_logger.writekvs(stats)
# Log level weights
if level_sampler and j % args.weight_log_interval == 0:
plogger.log_level_weights(level_sampler.sample_weights(), level_sampler.seeds)
# Checkpoint
timer = timeit.default_timer
if last_checkpoint_time is None:
last_checkpoint_time = timer()
try:
if j == num_updates - 1 or \
(args.save_interval > 0 and timer() - last_checkpoint_time > args.save_interval * 60): # Save every 60 min.
checkpoint()
last_checkpoint_time = timer()
logging.info(f"\nSaved checkpoint after update {current_update_count}")
except KeyboardInterrupt:
return
def train(args, seeds):
global last_checkpoint_time
args.cuda = not args.no_cuda and torch.cuda.is_available()
device = torch.device("cuda:0" if args.cuda else "cpu")
if 'cuda' in device.type:
print('Using CUDA\n')
torch.set_num_threads(1)
utils.seed(args.seed)
# Configure logging
if args.xpid is None:
args.xpid = "lr-%s" % time.strftime("%Y%m%d-%H%M%S")
log_dir = os.path.expandvars(os.path.expanduser(args.log_dir))
plogger = FileWriter(
xpid=args.xpid,
xp_args=args.__dict__,
rootdir=log_dir,
seeds=seeds,
)
stdout_logger = HumanOutputFormat(sys.stdout)
checkpointpath = os.path.expandvars(
os.path.expanduser("%s/%s/%s" % (log_dir, args.xpid, "model.tar")))
# Configure actor envs
start_level = 0
if args.full_train_distribution:
num_levels = 0
level_sampler_args = None
seeds = None
else:
num_levels = 1
level_sampler_args = dict(
num_actors=args.num_processes,
strategy=args.level_replay_strategy,
replay_schedule=args.level_replay_schedule,
score_transform=args.level_replay_score_transform,
temperature=args.level_replay_temperature,
eps=args.level_replay_eps,
rho=args.level_replay_rho,
nu=args.level_replay_nu,
alpha=args.level_replay_alpha,
staleness_coef=args.staleness_coef,
staleness_transform=args.staleness_transform,
staleness_temperature=args.staleness_temperature)
envs, level_sampler = make_lr_venv(
num_envs=args.num_processes,
env_name=args.env_name,
seeds=seeds,
device=device,
num_levels=num_levels,
start_level=start_level,
no_ret_normalization=args.no_ret_normalization,
distribution_mode=args.distribution_mode,
paint_vel_info=args.paint_vel_info,
level_sampler_args=level_sampler_args)
is_minigrid = args.env_name.startswith('MiniGrid')
actor_critic = model_for_env_name(args, envs)
actor_critic.to(device)
print(actor_critic)
rollouts = RolloutStorage(args.num_steps, args.num_processes,
envs.observation_space.shape, envs.action_space,
actor_critic.recurrent_hidden_state_size)
batch_size = int(args.num_processes * args.num_steps / args.num_mini_batch)
def checkpoint():
if args.disable_checkpoint:
return
logging.info("Saving checkpoint to %s", checkpointpath)
torch.save(
{
"model_state_dict": actor_critic.state_dict(),
"optimizer_state_dict": agent.optimizer.state_dict(),
"args": vars(args),
},
checkpointpath,
)
agent = algo.PPO(actor_critic,
args.clip_param,
args.ppo_epoch,
args.num_mini_batch,
args.value_loss_coef,
args.entropy_coef,
lr=args.lr,
eps=args.eps,
max_grad_norm=args.max_grad_norm,
env_name=args.env_name)
level_seeds = torch.zeros(args.num_processes)
if level_sampler:
obs, level_seeds = envs.reset()
else:
obs = envs.reset()
level_seeds = level_seeds.unsqueeze(-1)
rollouts.obs[0].copy_(obs)
rollouts.to(device)
episode_rewards = deque(maxlen=10)
num_updates = int(
args.num_env_steps) // args.num_steps // args.num_processes
timer = timeit.default_timer
update_start_time = timer()
for j in range(num_updates):
actor_critic.train()
for step in range(args.num_steps):
# Sample actions
with torch.no_grad():
obs_id = rollouts.obs[step]
value, action, action_log_dist, recurrent_hidden_states = actor_critic.act(
obs_id, rollouts.recurrent_hidden_states[step],
rollouts.masks[step])
action_log_prob = action_log_dist.gather(-1, action)
# Obser reward and next obs
obs, reward, done, infos = envs.step(action)
# Reset all done levels by sampling from level sampler
for i, info in enumerate(infos):
if 'episode' in info.keys():
episode_rewards.append(info['episode']['r'])
if level_sampler:
level_seeds[i][0] = info['level_seed']
# If done then clean the history of observations.
masks = torch.FloatTensor([[0.0] if done_ else [1.0]
for done_ in done])
bad_masks = torch.FloatTensor(
[[0.0] if 'bad_transition' in info.keys() else [1.0]
for info in infos])
rollouts.insert(obs, recurrent_hidden_states, action,
action_log_prob, action_log_dist, value, reward,
masks, bad_masks, level_seeds)
with torch.no_grad():
obs_id = rollouts.obs[-1]
next_value = actor_critic.get_value(
obs_id, rollouts.recurrent_hidden_states[-1],
rollouts.masks[-1]).detach()
rollouts.compute_returns(next_value, args.gamma, args.gae_lambda)
# Update level sampler
if level_sampler:
level_sampler.update_with_rollouts(rollouts)
value_loss, action_loss, dist_entropy = agent.update(rollouts)
rollouts.after_update()
if level_sampler:
level_sampler.after_update()
# Log stats every log_interval updates or if it is the last update
if (j % args.log_interval == 0
and len(episode_rewards) > 1) or j == num_updates - 1:
total_num_steps = (j + 1) * args.num_processes * args.num_steps
update_end_time = timer()
num_interval_updates = 1 if j == 0 else args.log_interval
sps = num_interval_updates * (args.num_processes *
args.num_steps) / (update_end_time -
update_start_time)
update_start_time = update_end_time
logging.info(f"\nUpdate {j} done, {total_num_steps} steps\n ")
logging.info(
f"\nEvaluating on {args.num_test_seeds} test levels...\n ")
eval_episode_rewards, transitions = evaluate(
args, actor_critic, args.num_test_seeds, device)
plogger._save_data(transitions, f'test_trajectories_{j}.pkl')
logging.info(
f"\nEvaluating on {args.num_test_seeds} train levels...\n ")
train_eval_episode_rewards, transitions = evaluate(
args,
actor_critic,
args.num_test_seeds,
device,
start_level=0,
num_levels=args.num_train_seeds,
seeds=seeds,
level_sampler=level_sampler)
stats = {
"step":
total_num_steps,
"pg_loss":
action_loss,
"value_loss":
value_loss,
"dist_entropy":
dist_entropy,
"train:mean_episode_return":
np.mean(episode_rewards),
"train:median_episode_return":
np.median(episode_rewards),
"test:mean_episode_return":
np.mean(eval_episode_rewards),
"test:median_episode_return":
np.median(eval_episode_rewards),
"train_eval:mean_episode_return":
np.mean(train_eval_episode_rewards),
"train_eval:median_episode_return":
np.median(train_eval_episode_rewards),
"sps":
sps,
}
if is_minigrid:
stats["train:success_rate"] = np.mean(
np.array(episode_rewards) > 0)
stats["train_eval:success_rate"] = np.mean(
np.array(train_eval_episode_rewards) > 0)
stats["test:success_rate"] = np.mean(
np.array(eval_episode_rewards) > 0)
if j == num_updates - 1:
logging.info(
f"\nLast update: Evaluating on {args.num_test_seeds} test levels...\n "
)
final_eval_episode_rewards, transitions = evaluate(
args, actor_critic, args.final_num_test_seeds, device)
mean_final_eval_episode_rewards = np.mean(
final_eval_episode_rewards)
median_final_eval_episide_rewards = np.median(
final_eval_episode_rewards)
plogger.log_final_test_eval({
'num_test_seeds':
args.final_num_test_seeds,
'mean_episode_return':
mean_final_eval_episode_rewards,
'median_episode_return':
median_final_eval_episide_rewards
})
plogger.log(stats)
if args.verbose:
stdout_logger.writekvs(stats)
# Log level weights
if level_sampler and j % args.weight_log_interval == 0:
plogger.log_level_weights(level_sampler.sample_weights())
# Checkpoint
timer = timeit.default_timer
if last_checkpoint_time is None:
last_checkpoint_time = timer()
try:
if j == num_updates - 1 or \
(args.save_interval > 0 and timer() - last_checkpoint_time > args.save_interval * 60): # Save every 10 min.
checkpoint()
last_checkpoint_time = timer()
except KeyboardInterrupt:
return