def main():
args = get_args()
torch.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)
if args.cuda and torch.cuda.is_available() and args.cuda_deterministic:
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
log_dir = os.path.expanduser(args.log_dir)
eval_log_dir = log_dir + "_eval"
utils.cleanup_log_dir(log_dir)
utils.cleanup_log_dir(eval_log_dir)
torch.set_num_threads(1)
device = torch.device("cuda:0" if args.cuda else "cpu")
receipts = StorageReceipt()
make_env = lambda tasks: MiniWoBGraphEnvironment(
base_url=os.environ.get("BASE_URL", f"file://{MINIWOB_HTML}/"),
levels=tasks,
level_tracker=LevelTracker(tasks),
wait_ms=500,
)
task = args.env_name
if args.env_name == "PongNoFrameskip-v4":
args.env_name = "clickbutton"
task = "miniwob/click-button.html"
if task == "levels":
tasks = MINIWOB_CHALLENGES
else:
tasks = [[task]]
print("Selected tasks:", tasks)
NUM_ACTIONS = 1
envs = make_vec_envs(
[make_env(tasks[i % len(tasks)]) for i in range(args.num_processes)],
receipts)
if os.path.exists("./datadir/autoencoder.pt"):
dom_autoencoder = torch.load("./datadir/autoencoder.pt")
dom_encoder = dom_autoencoder.encoder
for param in dom_encoder.parameters():
param.requires_grad = False
else:
print("No dom encoder")
dom_encoder = None
actor_critic = Policy(
envs.observation_space.shape,
gym.spaces.Discrete(NUM_ACTIONS), # envs.action_space,
base=GNNBase,
base_kwargs={
"dom_encoder": dom_encoder,
"recurrent": args.recurrent_policy
},
)
actor_critic.dist = NodeObjective()
actor_critic.to(device)
if args.algo == "a2c":
agent = algo.A2C_ACKTR(
actor_critic,
args.value_loss_coef,
args.entropy_coef,
lr=args.lr,
eps=args.eps,
alpha=args.alpha,
max_grad_norm=args.max_grad_norm,
)
elif args.algo == "ppo":
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,
)
elif args.algo == "acktr":
agent = algo.A2C_ACKTR(actor_critic,
args.value_loss_coef,
args.entropy_coef,
acktr=True)
if args.gail:
assert len(envs.observation_space.shape) == 1
discr = gail.Discriminator(envs.observation_space.shape[0], 100,
device)
rr = ReplayRepository("/code/miniwob-plusplus-demos/*turk/*")
ds = rr.get_dataset()
print("GAIL Replay Dataset", ds)
gail_train_loader = torch_geometric.data.DataLoader(
ds, batch_size=args.gail_batch_size, shuffle=True, drop_last=True)
from tensorboardX import SummaryWriter
import datetime
ts_str = datetime.datetime.fromtimestamp(
time.time()).strftime("%Y-%m-%d_%H-%M-%S")
tensorboard_writer = SummaryWriter(
log_dir=os.path.join("/tmp/log", ts_str))
rollouts = ReceiptRolloutStorage(
args.num_steps,
args.num_processes,
(1, ), # envs.observation_space.shape,
envs.action_space,
actor_critic.recurrent_hidden_state_size,
receipts,
)
# resume from last save
if args.save_dir != "":
save_path = os.path.join(args.save_dir, args.algo)
try:
os.makedirs(save_path)
except OSError:
pass
model_path = os.path.join(save_path, args.env_name + ".pt")
if False and os.path.exists(model_path):
print("Loadng previous model:", model_path)
actor_critic = torch.load(model_path)
actor_critic.train()
obs = envs.reset()
rollouts.obs[0].copy_(torch.tensor(obs))
rollouts.to(device)
start = time.time()
num_updates = int(
args.num_env_steps) // args.num_steps // args.num_processes
print("Iterations:", num_updates, args.num_steps)
for j in range(num_updates):
episode_rewards = deque(maxlen=args.num_steps * args.num_processes)
if j and last_action_time + 5 < time.time():
# task likely timed out
print("Reseting tasks")
obs = envs.reset()
rollouts.obs[0].copy_(torch.tensor(obs))
rollouts.recurrent_hidden_states[0].copy_(
torch.zeros_like(rollouts.recurrent_hidden_states[0]))
rollouts.masks[0].copy_(torch.zeros_like(rollouts.masks[0]))
if args.use_linear_lr_decay:
# decrease learning rate linearly
utils.update_linear_schedule(
agent.optimizer,
j,
num_updates,
agent.optimizer.lr if args.algo == "acktr" else args.lr,
)
for step in range(args.num_steps):
# Sample actions
with torch.no_grad():
value, action, action_log_prob, recurrent_hidden_states = actor_critic.act(
receipts.redeem(rollouts.obs[step]),
rollouts.recurrent_hidden_states[step],
rollouts.masks[step],
)
# Obser reward and next obs
last_action_time = time.time()
obs, reward, done, infos = envs.step(action)
for e, i in enumerate(infos):
if i.get("real_action") is not None:
action[e] = i["real_action"]
if i.get("bad_transition"):
action[e] = torch.zeros_like(action[e])
for info in infos:
if "episode" in info.keys():
episode_rewards.append(info["episode"]["r"])
# 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(
torch.tensor(obs),
recurrent_hidden_states,
action,
action_log_prob,
value,
torch.tensor(reward).unsqueeze(1),
masks,
bad_masks,
)
with torch.no_grad():
next_value = actor_critic.get_value(
receipts.redeem(rollouts.obs[-1]),
rollouts.recurrent_hidden_states[-1],
rollouts.masks[-1],
).detach()
if args.gail:
# if j >= 10:
# envs.venv.eval()
gail_epoch = args.gail_epoch
if j < 10:
gail_epoch = 100 # Warm up
for _ in range(gail_epoch):
obsfilt = lambda x, update: x # utils.get_vec_normalize(envs)._obfilt
gl = discr.update(gail_train_loader, rollouts, obsfilt)
print("Gail loss:", gl)
for step in range(args.num_steps):
rollouts.rewards[step] = discr.predict_reward(
receipts.redeem(rollouts.obs[step]),
rollouts.actions[step],
args.gamma,
rollouts.masks[step],
)
rollouts.compute_returns(
next_value,
args.use_gae,
args.gamma,
args.gae_lambda,
args.use_proper_time_limits,
)
value_loss, action_loss, dist_entropy = agent.update(rollouts)
obs_shape = rollouts.obs.size()[2:]
obs = rollouts.obs[:-1].view(-1, *obs_shape)
obs = obs[torch.randint(0, obs.size(0), (1, 32))]
rollouts.after_update()
receipts.prune(rollouts.obs)
# save for every interval-th episode or for the last epoch
if (j % args.save_interval == 0
or j == num_updates - 1) and args.save_dir != "":
save_path = os.path.join(args.save_dir, args.algo)
try:
os.makedirs(save_path)
except OSError:
pass
model_path = os.path.join(save_path, args.env_name + ".pt")
torch.save(actor_critic, model_path)
print("Saved model:", model_path)
if j % args.log_interval == 0 and len(episode_rewards) > 1:
total_num_steps = (j + 1) * args.num_processes * args.num_steps
end = time.time()
print(
"Updates {}, num timesteps {}, FPS {} \n Last {} training episodes: mean/median reward {:.1f}/{:.1f}, min/max reward {:.1f}/{:.1f}\n"
.format(
j,
total_num_steps,
int(total_num_steps / (end - start)),
len(episode_rewards),
np.mean(episode_rewards),
np.median(episode_rewards),
np.min(episode_rewards),
np.max(episode_rewards),
dist_entropy,
value_loss,
action_loss,
))
from pprint import pprint
pprint(LevelTracker.global_scoreboard)
# tensorboard_writer.add_histogram(
# "task_ranks", torch.tensor(predictor._difficulty_rank), total_num_steps
# )
tensorboard_writer.add_histogram("value", value, total_num_steps)
tensorboard_writer.add_histogram("x", actor_critic.base.last_x,
total_num_steps)
tensorboard_writer.add_histogram("query",
actor_critic.base.last_query,
total_num_steps)
tensorboard_writer.add_histogram("inputs_at",
actor_critic.base.last_inputs_at,
total_num_steps)
tensorboard_writer.add_scalar("mean_reward",
np.mean(episode_rewards),
total_num_steps)
tensorboard_writer.add_scalar("median_reward",
np.median(episode_rewards),
total_num_steps)
tensorboard_writer.add_scalar("min_reward",
np.min(episode_rewards),
total_num_steps)
tensorboard_writer.add_scalar("max_reward",
np.max(episode_rewards),
total_num_steps)
tensorboard_writer.add_scalar("dist_entropy", dist_entropy,
total_num_steps)
tensorboard_writer.add_scalar("value_loss", value_loss,
total_num_steps)
tensorboard_writer.add_scalar("action_loss", action_loss,
total_num_steps)
if (args.eval_interval is not None and len(episode_rewards) > 1
and j % args.eval_interval == 0):
ob_rms = utils.get_vec_normalize(envs).ob_rms
evaluate(
actor_critic,
ob_rms,
args.env_name,
args.seed,
args.num_processes,
eval_log_dir,
device,
)
def train(train_states,
run_dir,
num_env_steps,
eval_env_steps,
writer,
writer_name,
args,
init_model=None):
envs = make_vec_envs(train_states, args.seed, args.num_processes,
args.gamma, 'cpu', 'train', args)
if init_model:
actor_critic, env_step, model_name = init_model
obs_space = actor_critic.obs_space
obs_process = actor_critic.obs_process
obs_module = actor_critic.obs_module
print(f" [load] Loaded model {model_name} at step {env_step}")
else:
obs_space = envs.observation_space
actor_critic = Policy(obs_space,
args.obs_process,
args.obs_module,
envs.action_space,
base_kwargs={'recurrent': args.recurrent_policy})
env_step = 0
actor_critic.to(args.device)
#print(actor_critic)
run_name = run_dir.replace('/', '_')
vid_save_dir = f"{run_dir}/videos/"
try:
os.makedirs(vid_save_dir)
except OSError:
pass
ckpt_save_dir = f"{run_dir}/ckpts/"
try:
os.makedirs(ckpt_save_dir)
except OSError:
pass
if args.algo == 'ppo':
agent = algo.PPO(actor_critic,
args.clip_param,
args.ppo_epoch,
args.num_mini_batch,
args.value_loss_coef,
args.entropy_coef,
args.device,
lr=args.lr,
eps=args.eps,
max_grad_norm=args.max_grad_norm)
elif args.algo == 'a2c':
agent = algo.A2C_ACKTR(actor_critic,
args.value_loss_coef,
args.entropy_coef,
lr=args.lr,
eps=args.eps,
alpha=args.alpha,
max_grad_norm=args.max_grad_norm,
acktr=False)
elif args.algo == 'acktr':
agent = algo.A2C_ACKTR(actor_critic,
args.value_loss_coef,
args.entropy_coef,
lr=args.lr,
eps=args.eps,
alpha=args.alpha,
max_grad_norm=args.max_grad_norm,
acktr=True)
else:
raise NotImplementedError
rollouts = RolloutStorage(args.num_steps, args.num_processes,
envs.observation_space, envs.action_space,
actor_critic.recurrent_hidden_state_size)
obs = envs.reset()
actor_critic.eval()
"""
try:
writer.add_graph(actor_critic, obs)
except ValueError:
print("Unable to write model graph to tensorboard.")
"""
actor_critic.train()
for k in rollouts.obs.keys():
rollouts.obs[k][0].copy_(obs[k][0])
episode_rewards = deque(maxlen=10)
num_updates = num_env_steps // args.num_steps // args.num_processes
batch_size = args.num_steps * args.num_processes
start = time.time()
while env_step < num_env_steps:
s = time.time()
if args.use_linear_lr_decay:
# decrease learning rate linearly
utils.update_linear_schedule(
agent.optimizer, j, num_updates,
agent.optimizer.lr if args.algo == "acktr" else args.lr)
for step in range(args.num_steps):
# Sample actions
with torch.no_grad():
value, action, action_log_prob, recurrent_hidden_states, _ = actor_critic.act(
{
k: rollouts.obs[k][step].float().to(args.device)
for k in rollouts.obs.keys()
}, rollouts.recurrent_hidden_states[step].to(args.device),
rollouts.masks[step].to(args.device))
value = value.cpu()
action = action.cpu()
action_log_prob = action_log_prob.cpu()
recurrent_hidden_states = recurrent_hidden_states.cpu()
# Observe reward and next obs
obs, reward, dones, infos = envs.step(action)
for done, info in zip(dones, infos):
env_state = info['env_state'][1]
if done:
writer.add_scalar(f'train_episode_x/{env_state}',
info['max_x'], env_step)
writer.add_scalar(f'train_episode_%/{env_state}',
info['max_x'] / info['lvl_max_x'] * 100,
env_step)
writer.add_scalar(f'train_episode_r/{env_state}',
info['sum_r'], env_step)
# If done then clean the history of observations.
masks = torch.FloatTensor([[0.0] if done else [1.0]
for done in dones])
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, value, reward, masks, bad_masks)
with torch.no_grad():
next_value = actor_critic.get_value(
{
k: rollouts.obs[k][-1].float().to(args.device)
for k in rollouts.obs.keys()
}, rollouts.recurrent_hidden_states[-1].to(args.device),
rollouts.masks[-1].to(args.device)).detach().cpu()
rollouts.compute_returns(next_value, args.use_gae, args.gamma,
args.gae_lambda, args.use_proper_time_limits)
value_loss, action_loss, dist_entropy = agent.update(rollouts)
rollouts.after_update()
env_step += batch_size
fps = batch_size / (time.time() - s)
#res = nvidia_smi.nvmlDeviceGetUtilizationRates(handle)
#writer.add_scalar(f'gpu_usage/{writer_name}', res.gpu, env_step)
#writer.add_scalar(f'gpu_mem/{writer_name}', res.memory, env_step)
total_norm = 0
for p in list(
filter(lambda p: p.grad is not None,
actor_critic.parameters())):
param_norm = p.grad.data.norm(2)
total_norm += param_norm.item()**2
total_norm = total_norm**(1. / 2)
obs_norm = {}
for obs_name in args.obs_keys:
t_norm = 0
if obs_name == 'video':
md = actor_critic.base.video_module
elif obs_name == 'audio':
md = actor_critic.base.audio_module
else:
raise NotImplementedError
for p in list(filter(lambda p: p.grad is not None,
md.parameters())):
param_norm = p.grad.data.norm(2)
t_norm += param_norm.item()**2
obs_norm[obs_name] = t_norm**(1. / 2)
prev_env_step = max(0, env_step + 1 - batch_size)
# write training metrics for this batch, usually takes 0.003s
if (env_step + 1
) // args.write_interval > prev_env_step // args.write_interval:
writer.add_scalar(f'grad_norm/{writer_name}', total_norm, env_step)
writer.add_scalar(f'fps/{writer_name}', fps, env_step)
writer.add_scalar(f'value_loss/{writer_name}',
value_loss / batch_size, env_step)
writer.add_scalar(f'action_loss/{writer_name}',
action_loss / batch_size, env_step)
writer.add_scalar(f'dist_entropy/{writer_name}',
dist_entropy / batch_size, env_step)
writer.add_scalar(f'cpu_usage/{writer_name}', psutil.cpu_percent(),
env_step)
writer.add_scalar(f'cpu_mem/{writer_name}',
psutil.virtual_memory()._asdict()['percent'],
env_step)
for obs_name in args.obs_keys:
writer.add_scalar(f'grad_norm_{obs_name}/{writer_name}',
obs_norm[obs_name], env_step)
# print log to console
if (env_step +
1) // args.log_interval > prev_env_step // args.log_interval:
end = time.time()
print(" [log] Env step {} of {}: {:.1f}s, {:.1f}fps".format(
env_step + 1, num_env_steps, end - start, fps))
if len(episode_rewards) > 0:
print(
" Last {} episodes: mean/med reward {:.1f}/{:.1f}, min/max reward {:.1f}/{:.1f}"
.format(len(episode_rewards), np.mean(episode_rewards),
np.median(episode_rewards),
np.min(episode_rewards), np.max(episode_rewards)))
print(
" dist_entropy {:.5f}, value_loss {:.6f}, action_loss {:.6f}, grad_norm {:.6f}"
.format(dist_entropy, value_loss, action_loss, total_norm))
start = time.time()
# save model to ckpt
if ((env_step + 1) // args.save_interval >
prev_env_step // args.save_interval):
torch.save([
actor_critic,
env_step,
run_name,
], os.path.join(ckpt_save_dir, f"{run_name}-{env_step}.pt"))
print(f" [save] Saved model at step {env_step+1}.")
# save model to ckpt and run evaluation if eval_interval and not final iteration in training loop
if ((env_step + 1) // args.eval_interval >
prev_env_step // args.eval_interval
) and env_step < num_env_steps and eval_env_steps > 0:
torch.save([
actor_critic,
env_step,
run_name,
], os.path.join(ckpt_save_dir, f"{run_name}-{env_step}.pt"))
print(f" [save] Saved model at step {env_step+1}.")
envs.close()
del envs # close does not actually get rid of envs, need to del
actor_critic.eval()
eval_score, e_dict = evaluate(train_states, actor_critic,
eval_env_steps, env_step, writer,
vid_save_dir, args.vid_tb_steps,
args.vid_file_steps,
args.obs_viz_layer, args)
print(f" [eval] Evaluation score: {eval_score}")
writer.add_scalar('eval_score', eval_score, env_step)
actor_critic.train()
envs = make_vec_envs(train_states, args.seed, args.num_processes,
args.gamma, 'cpu', 'train', args)
obs = envs.reset()
# TODO: does this work? do we need to increment env step or something? whydden_states insert at 0
for k in rollouts.obs.keys():
rollouts.obs[k][0].copy_(obs[k][0])
# final model save
final_model_path = os.path.join(ckpt_save_dir, f"{run_name}-{env_step}.pt")
torch.save([
actor_critic,
env_step,
run_name,
], final_model_path)
print(
f" [save] Final model saved at step {env_step+1} to {final_model_path}"
)
# final model eval
envs.close()
del envs
eval_score = None
eval_dict = None
if eval_env_steps > 0:
eval_score, eval_dict = evaluate(train_states, actor_critic,
eval_env_steps, env_step, writer,
vid_save_dir, args.vid_tb_steps,
args.vid_file_steps,
args.obs_viz_layer, args)
print(f" [eval] Final model evaluation score: {eval_score:.3f}")
return (actor_critic, env_step, run_name), eval_score, eval_dict
def main():
args = get_args()
import random
random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)
np.random.seed(args.seed)
if args.cuda and torch.cuda.is_available() and args.cuda_deterministic:
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
logdir = args.env_name + '_' + args.algo + '_num_arms_' + str(
args.num_processes) + '_' + time.strftime("%d-%m-%Y_%H-%M-%S")
if args.use_privacy:
logdir = logdir + '_privacy'
elif args.use_noisygrad:
logdir = logdir + '_noisygrad'
elif args.use_pcgrad:
logdir = logdir + '_pcgrad'
elif args.use_testgrad:
logdir = logdir + '_testgrad'
elif args.use_median_grad:
logdir = logdir + '_mediangrad'
logdir = os.path.join('runs', logdir)
logdir = os.path.join(os.path.expanduser(args.log_dir), logdir)
utils.cleanup_log_dir(logdir)
# Ugly but simple logging
log_dict = {
'task_steps': args.task_steps,
'grad_noise_ratio': args.grad_noise_ratio,
'max_task_grad_norm': args.max_task_grad_norm,
'use_noisygrad': args.use_noisygrad,
'use_pcgrad': args.use_pcgrad,
'use_testgrad': args.use_testgrad,
'use_testgrad_median': args.use_testgrad_median,
'testgrad_quantile': args.testgrad_quantile,
'median_grad': args.use_median_grad,
'use_meanvargrad': args.use_meanvargrad,
'meanvar_beta': args.meanvar_beta,
'no_special_grad_for_critic': args.no_special_grad_for_critic,
'use_privacy': args.use_privacy,
'seed': args.seed,
'recurrent': args.recurrent_policy,
'obs_recurrent': args.obs_recurrent,
'cmd': ' '.join(sys.argv[1:])
}
for eval_disp_name, eval_env_name in EVAL_ENVS.items():
log_dict[eval_disp_name] = []
summary_writer = SummaryWriter()
summary_writer.add_hparams(
{
'task_steps': args.task_steps,
'grad_noise_ratio': args.grad_noise_ratio,
'max_task_grad_norm': args.max_task_grad_norm,
'use_noisygrad': args.use_noisygrad,
'use_pcgrad': args.use_pcgrad,
'use_testgrad': args.use_testgrad,
'use_testgrad_median': args.use_testgrad_median,
'testgrad_quantile': args.testgrad_quantile,
'median_grad': args.use_median_grad,
'use_meanvargrad': args.use_meanvargrad,
'meanvar_beta': args.meanvar_beta,
'no_special_grad_for_critic': args.no_special_grad_for_critic,
'use_privacy': args.use_privacy,
'seed': args.seed,
'recurrent': args.recurrent_policy,
'obs_recurrent': args.obs_recurrent,
'cmd': ' '.join(sys.argv[1:])
}, {})
torch.set_num_threads(1)
device = torch.device("cuda:0" if args.cuda else "cpu")
print('making envs...')
envs = make_vec_envs(args.env_name,
args.seed,
args.num_processes,
args.gamma,
args.log_dir,
device,
False,
steps=args.task_steps,
free_exploration=args.free_exploration,
recurrent=args.recurrent_policy,
obs_recurrent=args.obs_recurrent,
multi_task=True)
val_envs = make_vec_envs(args.val_env_name,
args.seed,
args.num_processes,
args.gamma,
args.log_dir,
device,
False,
steps=args.task_steps,
free_exploration=args.free_exploration,
recurrent=args.recurrent_policy,
obs_recurrent=args.obs_recurrent,
multi_task=True)
eval_envs_dic = {}
for eval_disp_name, eval_env_name in EVAL_ENVS.items():
eval_envs_dic[eval_disp_name] = make_vec_envs(
eval_env_name[0],
args.seed,
args.num_processes,
None,
logdir,
device,
True,
steps=args.task_steps,
recurrent=args.recurrent_policy,
obs_recurrent=args.obs_recurrent,
multi_task=True,
free_exploration=args.free_exploration)
prev_eval_r = {}
print('done')
if args.hard_attn:
actor_critic = Policy(envs.observation_space.shape,
envs.action_space,
base=MLPHardAttnBase,
base_kwargs={
'recurrent':
args.recurrent_policy or args.obs_recurrent
})
else:
actor_critic = Policy(envs.observation_space.shape,
envs.action_space,
base=MLPAttnBase,
base_kwargs={
'recurrent':
args.recurrent_policy or args.obs_recurrent
})
actor_critic.to(device)
if (args.continue_from_epoch > 0) and args.save_dir != "":
save_path = os.path.join(args.save_dir, args.algo)
actor_critic_, loaded_obs_rms_ = torch.load(
os.path.join(
save_path, args.env_name +
"-epoch-{}.pt".format(args.continue_from_epoch)))
actor_critic.load_state_dict(actor_critic_.state_dict())
if args.algo != 'ppo':
raise "only PPO is supported"
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,
num_tasks=args.num_processes,
attention_policy=False,
max_grad_norm=args.max_grad_norm,
weight_decay=args.weight_decay)
val_agent = algo.PPO(actor_critic,
args.clip_param,
args.ppo_epoch,
args.num_mini_batch,
args.value_loss_coef,
args.entropy_coef,
lr=args.val_lr,
eps=args.eps,
num_tasks=args.num_processes,
attention_policy=True,
max_grad_norm=args.max_grad_norm,
weight_decay=args.weight_decay)
rollouts = RolloutStorage(args.num_steps, args.num_processes,
envs.observation_space.shape, envs.action_space,
actor_critic.recurrent_hidden_state_size)
val_rollouts = RolloutStorage(args.num_steps, args.num_processes,
val_envs.observation_space.shape,
val_envs.action_space,
actor_critic.recurrent_hidden_state_size)
obs = envs.reset()
rollouts.obs[0].copy_(obs)
rollouts.to(device)
val_obs = val_envs.reset()
val_rollouts.obs[0].copy_(val_obs)
val_rollouts.to(device)
episode_rewards = deque(maxlen=10)
start = time.time()
num_updates = int(
args.num_env_steps) // args.num_steps // args.num_processes
save_copy = True
for j in range(args.continue_from_epoch,
args.continue_from_epoch + num_updates):
# policy rollouts
for step in range(args.num_steps):
# Sample actions
actor_critic.eval()
with torch.no_grad():
value, action, action_log_prob, recurrent_hidden_states = actor_critic.act(
rollouts.obs[step], rollouts.recurrent_hidden_states[step],
rollouts.masks[step])
actor_critic.train()
# Obser reward and next obs
obs, reward, done, infos = envs.step(action)
for info in infos:
if 'episode' in info.keys():
episode_rewards.append(info['episode']['r'])
for k, v in info['episode'].items():
summary_writer.add_scalar(
f'training/{k}', v,
j * args.num_processes * args.num_steps +
args.num_processes * step)
# 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, value, reward, masks, bad_masks)
actor_critic.eval()
with torch.no_grad():
next_value = actor_critic.get_value(
rollouts.obs[-1], rollouts.recurrent_hidden_states[-1],
rollouts.masks[-1]).detach()
actor_critic.train()
rollouts.compute_returns(next_value, args.use_gae, args.gamma,
args.gae_lambda, args.use_proper_time_limits)
if save_copy:
prev_weights = copy.deepcopy(actor_critic.state_dict())
prev_opt_state = copy.deepcopy(agent.optimizer.state_dict())
prev_val_opt_state = copy.deepcopy(
val_agent.optimizer.state_dict())
save_copy = False
value_loss, action_loss, dist_entropy = agent.update(rollouts)
rollouts.after_update()
# validation rollouts
for val_iter in range(args.val_agent_steps):
for step in range(args.num_steps):
# Sample actions
actor_critic.eval()
with torch.no_grad():
value, action, action_log_prob, recurrent_hidden_states = actor_critic.act(
val_rollouts.obs[step],
val_rollouts.recurrent_hidden_states[step],
val_rollouts.masks[step])
actor_critic.train()
# Obser reward and next obs
obs, reward, done, infos = val_envs.step(action)
# 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])
val_rollouts.insert(obs, recurrent_hidden_states, action,
action_log_prob, value, reward, masks,
bad_masks)
actor_critic.eval()
with torch.no_grad():
next_value = actor_critic.get_value(
val_rollouts.obs[-1],
val_rollouts.recurrent_hidden_states[-1],
val_rollouts.masks[-1]).detach()
actor_critic.train()
val_rollouts.compute_returns(next_value, args.use_gae, args.gamma,
args.gae_lambda,
args.use_proper_time_limits)
val_value_loss, val_action_loss, val_dist_entropy = val_agent.update(
val_rollouts)
val_rollouts.after_update()
# save for every interval-th episode or for the last epoch
if (j % args.save_interval == 0
or j == num_updates - 1) and args.save_dir != "":
save_path = os.path.join(args.save_dir, args.algo)
try:
os.makedirs(save_path)
except OSError:
pass
torch.save([
actor_critic,
getattr(utils.get_vec_normalize(envs), 'obs_rms', None)
], os.path.join(save_path,
args.env_name + "-epoch-{}.pt".format(j)))
if j % args.log_interval == 0 and len(episode_rewards) > 1:
total_num_steps = (j + 1) * args.num_processes * args.num_steps
end = time.time()
print(
"Updates {}, num timesteps {}, FPS {} \n Last {} training episodes: mean/median reward {:.1f}/{:.1f}, min/max reward {:.1f}/{:.1f}\n"
.format(j, total_num_steps,
int(total_num_steps / (end - start)),
len(episode_rewards), np.mean(episode_rewards),
np.median(episode_rewards), np.min(episode_rewards),
np.max(episode_rewards), dist_entropy, value_loss,
action_loss))
revert = False
if (args.eval_interval is not None and len(episode_rewards) > 1
and j % args.eval_interval == 0):
actor_critic.eval()
obs_rms = utils.get_vec_normalize(envs).obs_rms
eval_r = {}
printout = f'Seed {args.seed} Iter {j} '
for eval_disp_name, eval_env_name in EVAL_ENVS.items():
eval_r[eval_disp_name] = evaluate(
actor_critic,
obs_rms,
eval_envs_dic,
eval_disp_name,
args.seed,
args.num_processes,
eval_env_name[1],
logdir,
device,
steps=args.task_steps,
recurrent=args.recurrent_policy,
obs_recurrent=args.obs_recurrent,
multi_task=True,
free_exploration=args.free_exploration)
if eval_disp_name in prev_eval_r:
diff = np.array(eval_r[eval_disp_name]) - np.array(
prev_eval_r[eval_disp_name])
if eval_disp_name == 'many_arms':
if np.sum(diff > 0) - np.sum(
diff < 0) < args.val_improvement_threshold:
print('no update')
revert = True
summary_writer.add_scalar(f'eval/{eval_disp_name}',
np.mean(eval_r[eval_disp_name]),
(j + 1) * args.num_processes *
args.num_steps)
log_dict[eval_disp_name].append([
(j + 1) * args.num_processes * args.num_steps,
eval_r[eval_disp_name]
])
printout += eval_disp_name + ' ' + str(
np.mean(eval_r[eval_disp_name])) + ' '
# summary_writer.add_scalars('eval_combined', eval_r, (j+1) * args.num_processes * args.num_steps)
if revert:
actor_critic.load_state_dict(prev_weights)
agent.optimizer.load_state_dict(prev_opt_state)
val_agent.optimizer.load_state_dict(prev_val_opt_state)
else:
print(printout)
prev_eval_r = eval_r.copy()
save_copy = True
actor_critic.train()
save_obj(log_dict, os.path.join(logdir, 'log_dict.pkl'))
envs.close()
val_envs.close()
for eval_disp_name, eval_env_name in EVAL_ENVS.items():
eval_envs_dic[eval_disp_name].close()
