def make_agent(is_leaf=True):
## AGENT CONSTRUCTION:
## Modularize this and allow for cascading (obs dim for child policy should be cat of obs and parents output)
actor_critic = OpsPolicy(
envs.observation_space.shape,
envs.action_space if is_leaf else gym.spaces.Discrete(2),
is_leaf=is_leaf,
base_kwargs=dict(recurrent=True,
partial_obs=args.partial_obs,
gate_input=args.gate_input))
actor_critic.to(device)
# wandb.watch(actor_critic.base)
if args.algo == 'a2c':
agent = algo.A2C_ACKTR(actor_critic,
args.value_loss_coef,
args.entropy_coef,
args.pred_loss_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,
args.pred_loss_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)
rollouts = RolloutStorage(args.num_steps,
args.num_processes,
envs.observation_space.shape,
envs.action_space,
actor_critic.recurrent_hidden_state_size,
info_size=2 if is_leaf else 0)
return actor_critic, agent, rollouts
def main(env, scene_path):
try:
os.makedirs(args.log_dir)
except OSError:
files = glob.glob(os.path.join(args.log_dir, '*.monitor.csv'))
for f in files:
os.remove(f)
save_path = os.path.join(args.save_dir, args.algo)
eval_x = []
eval_y = []
torch.set_num_threads(1)
device = torch.device("cuda:0" if args.cuda else "cpu")
initial_policies = torch.load(os.path.join(args.load_dir, args.algo,
args.initial_policy + ".pt")) \
if args.initial_policy else None
if args.reuse_residual:
residual, ob_rms, initial_policies = initial_policies
else:
residual = None
ob_rms = None
pose_estimator = torch.load(os.path.join(args.load_dir, "pe",
args.pose_estimator + ".pt")) \
if args.pose_estimator else None
envs = make_vec_envs(env,
scene_path,
args.seed,
args.num_processes,
args.gamma,
args.log_dir,
device,
False,
initial_policies,
pose_estimator=pose_estimator,
init_control=not args.dense_ip)
if args.reuse_residual:
vec_norm = get_vec_normalize(envs)
if vec_norm is not None:
vec_norm.eval()
vec_norm.ob_rms = ob_rms
base_kwargs = {'recurrent': args.recurrent_policy}
base = residual.base if args.reuse_residual else None
dist = residual.dist if args.reuse_residual else None
actor_critic = Policy(envs.observation_space.shape,
envs.action_space,
base_kwargs=base_kwargs,
zero_last_layer=True,
base=base,
dist=dist)
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,
burn_in=initial_policies is not None
and not args.reuse_residual)
elif args.algo == 'acktr':
agent = algo.A2C_ACKTR(actor_critic,
args.value_loss_coef,
args.entropy_coef,
acktr=True)
rollouts = RolloutStorage(args.num_steps, args.num_processes,
envs.observation_space.shape, envs.action_space,
actor_critic.recurrent_hidden_state_size)
obs = envs.reset()
rollouts.obs[0].copy_(obs)
rollouts.to(device)
episode_rewards = deque(maxlen=64)
num_updates = int(
args.num_env_steps) // args.num_steps // args.num_processes
total_num_steps = 0
j = 0
max_succ = -1
max_mean_rew = -math.inf
mean_ep_rew = -math.inf
evals_without_improv = 0
start = time.time()
start_update = start
while (not use_metric
and j < num_updates) or (use_metric
and max_succ < args.trg_succ_rate):
if args.eval_interval is not None and j % args.eval_interval == 0:
print("Evaluating current policy...")
i = 0
total_successes = 0
max_trials = 50
eval_recurrent_hidden_states = torch.zeros(
args.num_processes,
actor_critic.recurrent_hidden_state_size,
device=device)
eval_masks = torch.zeros(args.num_processes, 1, device=device)
while i + args.num_processes <= max_trials:
with torch.no_grad():
_, action, _, eval_recurrent_hidden_states = actor_critic.act(
obs,
eval_recurrent_hidden_states,
eval_masks,
deterministic=True)
obs, _, dones, infos = envs.step(action)
if np.all(dones): # Rigid - assumes episodes are fixed length
rews = []
for info in infos:
rews.append(info['rew_success'])
i += args.num_processes
rew = sum([int(rew > 0) for rew in rews])
total_successes += rew
p_succ = (100 * total_successes / i)
eval_x += [total_num_steps]
eval_y += [p_succ]
end = time.time()
print(
f"Evaluation: {total_successes} successful out of {i} episodes - "
f"{p_succ:.2f}% successful. Eval length: {end - start_update}")
torch.save([eval_x, eval_y],
os.path.join(args.save_as + "_eval.pt"))
start_update = end
if p_succ > max_succ:
max_succ = p_succ
max_mean_rew = mean_ep_rew
evals_without_improv = 0
elif mean_ep_rew > max_mean_rew:
print("Unimproved success rate, higher reward")
max_mean_rew = mean_ep_rew
evals_without_improv = 0
else:
evals_without_improv += 1
if evals_without_improv == 10 or max_succ >= args.trg_succ_rate:
save_model = actor_critic
if args.cuda:
save_model = copy.deepcopy(actor_critic).cpu()
save_model = [
save_model,
getattr(get_vec_normalize(envs), 'ob_rms', None),
initial_policies
]
extra = "_final" if evals_without_improv == 5 else ""
torch.save(
save_model,
os.path.join(save_path, args.save_as + f"{extra}.pt"))
break
# save for every interval-th episode or for the last epoch
if ((not use_metric and
(j % args.save_interval == 0 or j == num_updates - 1)) or
(use_metric
and evals_without_improv == 0)) and args.save_dir != "":
os.makedirs(save_path, exist_ok=True)
save_model = actor_critic
if args.cuda:
save_model = copy.deepcopy(actor_critic).cpu()
if pose_estimator is not None:
save_model = [save_model, pose_estimator, initial_policies]
else:
save_model = [
save_model,
getattr(get_vec_normalize(envs), 'ob_rms', None),
initial_policies
]
torch.save(save_model, os.path.join(save_path,
args.save_as + ".pt"))
# torch.save(save_model, os.path.join(save_path, args.save_as + f"{j * args.num_processes * args.num_steps}.pt"))
if args.use_linear_lr_decay:
# decrease learning rate linearly
if args.algo == "acktr":
# use optimizer's learning rate since it's hard-coded in kfac.py
update_linear_schedule(agent.optimizer, j, num_updates,
agent.optimizer.lr)
else:
update_linear_schedule(agent.optimizer, j, num_updates,
args.lr)
if args.algo == 'ppo' and args.use_linear_clip_decay:
agent.clip_param = args.clip_param * (1 - j / float(num_updates))
for step in range(args.num_steps):
# Sample actions
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])
# 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'])
# If done then clean the history of observations.
masks = torch.FloatTensor([[0.0] if done_ else [1.0]
for done_ in done])
rollouts.insert(obs, recurrent_hidden_states, action,
action_log_prob, value, reward, masks)
with torch.no_grad():
next_value = actor_critic.get_value(
rollouts.obs[-1], rollouts.recurrent_hidden_states[-1],
rollouts.masks[-1]).detach()
rollouts.compute_returns(next_value, args.use_gae, args.gamma,
args.tau)
value_loss, action_loss, dist_entropy = agent.update(rollouts)
rollouts.after_update()
total_num_steps = (j + 1) * args.num_processes * args.num_steps
if j % args.log_interval == 0 and len(episode_rewards) > 1:
mean_ep_rew = np.mean(episode_rewards)
if mean_ep_rew > max_mean_rew:
print("Improved max mean reward")
max_mean_rew = mean_ep_rew
evals_without_improv = 0
end = time.time()
print(
"Updates {}, num timesteps {}, FPS {} \n Last {} training episodes: mean/median reward {:.1f}/{:.1f}, min/max reward {:.1f}/{:.1f}"
.format(j, total_num_steps,
int(total_num_steps / (end - start)),
len(episode_rewards), mean_ep_rew,
np.median(episode_rewards), np.min(episode_rewards),
np.max(episode_rewards), dist_entropy, value_loss,
action_loss))
print("Update length: ", end - start_update)
start_update = end
if args.vis and (j % args.vis_interval == 0 or
(not use_metric and j == num_updates - 1)):
try:
# Sometimes monitor doesn't properly flush the outputs
visdom_plot(args.log_dir, args.save_as, args.algo,
total_num_steps)
except IOError:
pass
j += 1
if use_metric:
if max_succ >= args.trg_succ_rate:
print(
f"Achieved greater than {args.trg_succ_rate}% success, advancing curriculum."
)
else:
print(
f"Policy converged with max success rate < {args.trg_succ_rate}%"
)
# Copy logs to permanent location so new graphs can be drawn.
copy_tree(args.log_dir, os.path.join('logs', args.save_as))
envs.close()
return total_num_steps
def main():
torch.set_num_threads(1)
device = torch.device("cuda:0" if args.cuda else "cpu")
if args.vis:
from visdom import Visdom
viz = Visdom(port=args.port)
win = None
envs = make_vec_envs(args.env_name, args.seed, args.num_processes,
args.gamma, args.log_dir, args.add_timestep, device,
False)
if args.load_policy is not None:
actor_critic, ob_rms = torch.load(args.load_policy)
vec_norm = get_vec_normalize(envs)
if vec_norm is not None:
vec_norm.eval()
vec_norm.ob_rms = ob_rms
else:
actor_critic = Policy(envs.observation_space.shape,
envs.action_space,
base_kwargs={'recurrent': args.recurrent_policy})
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)
rollouts = RolloutStorage(args.num_steps, args.num_processes,
envs.observation_space.shape, envs.action_space,
actor_critic.recurrent_hidden_state_size)
obs = envs.reset()
rollouts.obs[0].copy_(obs)
rollouts.to(device)
episode_rewards = deque(
maxlen=(args.num_processes if args.num_processes > 10 else 10))
start = time.time()
snapshot_counter = 0
last_delete = -1
try:
os.makedirs(os.path.join(args.save_dir, args.algo))
except OSError:
pass
log_out_file = open(os.path.join(args.save_dir, args.algo, 'log_info.txt'),
'w')
for j in range(num_updates):
if args.use_linear_lr_decay:
# decrease learning rate linearly
if args.algo == "acktr":
# use optimizer's learning rate since it's hard-coded in kfac.py
update_linear_schedule(agent.optimizer, j, num_updates,
agent.optimizer.lr)
else:
update_linear_schedule(agent.optimizer, j, num_updates,
args.lr)
if args.algo == 'ppo' and args.use_linear_clip_decay:
agent.clip_param = args.clip_param * (1 - j / float(num_updates))
for step in range(args.num_steps):
# Sample actions
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])
# 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'])
# If done then clean the history of observations.
masks = torch.FloatTensor([[0.0] if done_ else [1.0]
for done_ in done])
rollouts.insert(obs, recurrent_hidden_states, action,
action_log_prob, value, reward, masks)
with torch.no_grad():
next_value = actor_critic.get_value(
rollouts.obs[-1], rollouts.recurrent_hidden_states[-1],
rollouts.masks[-1]).detach()
rollouts.compute_returns(next_value, args.use_gae, args.gamma,
args.tau)
value_loss, action_loss, dist_entropy = agent.update(rollouts)
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
# A really ugly way to save a model to CPU
save_model = actor_critic
if args.cuda:
save_model = copy.deepcopy(actor_critic).cpu()
save_model = [
save_model,
getattr(get_vec_normalize(envs), 'ob_rms', None)
]
torch.save(
save_model,
os.path.join(save_path,
args.env_name + "epoch_{:07d}.pt".format(j)))
snapshot_counter += 1
last_delete += 1
if snapshot_counter > 100:
os.system('rm ' + os.path.join(
save_path, args.env_name +
'epoch_{:07d}.py'.format(last_delete)))
snapshot_counter -= 1
total_num_steps = (j + 1) * args.num_processes * args.num_steps
if j % args.log_interval == 0 and len(episode_rewards) > 1:
end = time.time()
log_info = "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)
print(log_info)
sys.stdout.flush()
log_out_file.write(log_info)
log_out_file.flush()
if (args.eval_interval is not None and len(episode_rewards) > 1
and j % args.eval_interval == 0):
eval_envs = make_vec_envs(args.env_name,
args.seed + args.num_processes,
args.num_processes, args.gamma,
eval_log_dir, args.add_timestep, device,
True)
vec_norm = get_vec_normalize(eval_envs)
if vec_norm is not None:
vec_norm.eval()
vec_norm.ob_rms = get_vec_normalize(envs).ob_rms
eval_episode_rewards = []
obs = eval_envs.reset()
eval_recurrent_hidden_states = torch.zeros(
args.num_processes,
actor_critic.recurrent_hidden_state_size,
device=device)
eval_masks = torch.zeros(args.num_processes, 1, device=device)
while len(eval_episode_rewards) < 10:
with torch.no_grad():
_, action, _, eval_recurrent_hidden_states = actor_critic.act(
obs,
eval_recurrent_hidden_states,
eval_masks,
deterministic=True)
# Obser reward and next obs
obs, reward, done, infos = eval_envs.step(action)
eval_masks = torch.FloatTensor([[0.0] if done_ else [1.0]
for done_ in done])
for info in infos:
if 'episode' in info.keys():
eval_episode_rewards.append(info['episode']['r'])
eval_envs.close()
print(" Evaluation using {} episodes: mean reward {:.5f}\n".format(
len(eval_episode_rewards), np.mean(eval_episode_rewards)))
log_out_file.write(
" Evaluation using {} episodes: mean reward {:.5f}\n".format(
len(eval_episode_rewards), np.mean(eval_episode_rewards)))
log_out_file.flush()
sys.stdout.flush()
if args.vis and j % args.vis_interval == 0:
try:
# Sometimes monitor doesn't properly flush the outputs
win = visdom_plot(viz, win, args.log_dir, args.env_name,
args.algo, args.num_env_steps)
except IOError:
pass
def main():
args = get_args()
torch.manual_seed(config.seed)
torch.cuda.manual_seed_all(config.seed)
if config.cuda and torch.cuda.is_available() and config.cuda_deterministic:
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
logger, final_output_dir, tb_log_dir = create_logger(config,
args.cfg,
'train',
seed=config.seed)
eval_log_dir = final_output_dir + "_eval"
utils.cleanup_log_dir(final_output_dir)
utils.cleanup_log_dir(eval_log_dir)
logger.info(pprint.pformat(args))
logger.info(pprint.pformat(config))
writer = SummaryWriter(tb_log_dir)
torch.set_num_threads(1)
device = torch.device("cuda:" + config.GPUS if config.cuda else "cpu")
width = height = 84
envs = make_vec_envs(config.env_name,
config.seed,
config.num_processes,
config.gamma,
final_output_dir,
device,
False,
width=width,
height=height,
ram_wrapper=False)
# create agent
actor_critic = Policy(envs.observation_space.shape,
envs.action_space,
base_kwargs={
'recurrent':
config.recurrent_policy,
'hidden_size':
config.hidden_size,
'feat_from_selfsup_attention':
config.feat_from_selfsup_attention,
'feat_add_selfsup_attention':
config.feat_add_selfsup_attention,
'feat_mul_selfsup_attention_mask':
config.feat_mul_selfsup_attention_mask,
'selfsup_attention_num_keypoints':
config.SELFSUP_ATTENTION.NUM_KEYPOINTS,
'selfsup_attention_gauss_std':
config.SELFSUP_ATTENTION.GAUSS_STD,
'selfsup_attention_fix':
config.selfsup_attention_fix,
'selfsup_attention_fix_keypointer':
config.selfsup_attention_fix_keypointer,
'selfsup_attention_pretrain':
config.selfsup_attention_pretrain,
'selfsup_attention_keyp_maps_pool':
config.selfsup_attention_keyp_maps_pool,
'selfsup_attention_image_feat_only':
config.selfsup_attention_image_feat_only,
'selfsup_attention_feat_masked':
config.selfsup_attention_feat_masked,
'selfsup_attention_feat_masked_residual':
config.selfsup_attention_feat_masked_residual,
'selfsup_attention_feat_load_pretrained':
config.selfsup_attention_feat_load_pretrained,
'use_layer_norm':
config.use_layer_norm,
'selfsup_attention_keyp_cls_agnostic':
config.SELFSUP_ATTENTION.KEYPOINTER_CLS_AGNOSTIC,
'selfsup_attention_feat_use_ln':
config.SELFSUP_ATTENTION.USE_LAYER_NORM,
'selfsup_attention_use_instance_norm':
config.SELFSUP_ATTENTION.USE_INSTANCE_NORM,
'feat_mul_selfsup_attention_mask_residual':
config.feat_mul_selfsup_attention_mask_residual,
'bottom_up_form_objects':
config.bottom_up_form_objects,
'bottom_up_form_num_of_objects':
config.bottom_up_form_num_of_objects,
'gaussian_std':
config.gaussian_std,
'train_selfsup_attention':
config.train_selfsup_attention,
'block_selfsup_attention_grad':
config.block_selfsup_attention_grad,
'sep_bg_fg_feat':
config.sep_bg_fg_feat,
'mask_threshold':
config.mask_threshold,
'fix_feature':
config.fix_feature
})
# init / load parameter
if config.MODEL_FILE:
logger.info('=> loading model from {}'.format(config.MODEL_FILE))
state_dict = torch.load(config.MODEL_FILE)
state_dict = OrderedDict(
(_k, _v) for _k, _v in state_dict.items() if 'dist' not in _k)
actor_critic.load_state_dict(state_dict, strict=False)
elif config.RESUME:
checkpoint_file = os.path.join(final_output_dir, 'checkpoint.pth')
if os.path.exists(checkpoint_file):
logger.info("=> loading checkpoint '{}'".format(checkpoint_file))
checkpoint = torch.load(checkpoint_file)
actor_critic.load_state_dict(checkpoint['state_dict'])
logger.info("=> loaded checkpoint '{}' (epoch {})".format(
checkpoint_file, checkpoint['epoch']))
actor_critic.to(device)
if config.algo == 'a2c':
agent = algo.A2C_ACKTR(
actor_critic,
config.value_loss_coef,
config.entropy_coef,
lr=config.lr,
eps=config.eps,
alpha=config.alpha,
max_grad_norm=config.max_grad_norm,
train_selfsup_attention=config.train_selfsup_attention)
elif config.algo == 'ppo':
agent = algo.PPO(actor_critic,
config.clip_param,
config.ppo_epoch,
config.num_mini_batch,
config.value_loss_coef,
config.entropy_coef,
lr=config.lr,
eps=config.eps,
max_grad_norm=config.max_grad_norm)
elif config.algo == 'acktr':
agent = algo.A2C_ACKTR(
actor_critic,
config.value_loss_coef,
config.entropy_coef,
acktr=True,
train_selfsup_attention=config.train_selfsup_attention,
max_grad_norm=config.max_grad_norm)
# rollouts: environment
rollouts = RolloutStorage(
config.num_steps,
config.num_processes,
envs.observation_space.shape,
envs.action_space,
actor_critic.recurrent_hidden_state_size,
keep_buffer=config.train_selfsup_attention,
buffer_size=config.train_selfsup_attention_buffer_size)
if config.RESUME:
if os.path.exists(checkpoint_file):
agent.optimizer.load_state_dict(checkpoint['optimizer'])
obs = envs.reset()
rollouts.obs[0].copy_(obs)
rollouts.to(device)
episode_rewards = deque(maxlen=10)
start = time.time()
num_updates = int(
config.num_env_steps) // config.num_steps // config.num_processes
best_perf = 0.0
best_model = False
print('num updates', num_updates, 'num steps', config.num_steps)
for j in range(num_updates):
if config.use_linear_lr_decay:
# decrease learning rate linearly
utils.update_linear_schedule(
agent.optimizer, j, num_updates,
agent.optimizer.lr if config.algo == "acktr" else config.lr)
for step in range(config.num_steps):
# Sample actions
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])
recurrent_hidden_states, meta = recurrent_hidden_states
# Obser reward and next obs
obs, reward, done, infos = envs.step(action)
objects_locs = []
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])
if objects_locs:
objects_locs = torch.FloatTensor(objects_locs)
objects_locs = objects_locs * 2 - 1 # -1, 1
else:
objects_locs = None
rollouts.insert(obs,
recurrent_hidden_states,
action,
action_log_prob,
value,
reward,
masks,
bad_masks,
objects_loc=objects_locs)
with torch.no_grad():
next_value = actor_critic.get_value(
rollouts.obs[-1],
rollouts.recurrent_hidden_states[-1],
rollouts.masks[-1],
).detach()
rollouts.compute_returns(next_value, config.use_gae, config.gamma,
config.gae_lambda,
config.use_proper_time_limits)
value_loss, action_loss, dist_entropy = agent.update(rollouts)
rollouts.after_update()
if config.train_selfsup_attention and j > 15:
for _iter in range(config.num_steps // 5):
frame_x, frame_y = rollouts.generate_pair_image()
selfsup_attention_loss, selfsup_attention_output, image_b_keypoints_maps = \
agent.update_selfsup_attention(frame_x, frame_y, config.SELFSUP_ATTENTION)
if j % config.log_interval == 0 and len(episode_rewards) > 1:
total_num_steps = (j + 1) * config.num_processes * config.num_steps
end = time.time()
msg = 'Updates {}, num timesteps {}, FPS {} \n' \
'Last {} training episodes: mean/median reward {:.1f}/{:.1f} ' \
'min/max reward {:.1f}/{:.1f} ' \
'dist entropy {:.1f}, value loss {:.1f}, action loss {:.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)
if config.train_selfsup_attention and j > 15:
msg = msg + 'selfsup attention loss {:.5f}\n'.format(
selfsup_attention_loss)
logger.info(msg)
if (config.eval_interval is not None and len(episode_rewards) > 1
and j % config.eval_interval == 0):
total_num_steps = (j + 1) * config.num_processes * config.num_steps
ob_rms = getattr(utils.get_vec_normalize(envs), 'ob_rms', None)
eval_mean_score, eval_max_score, eval_scores = evaluate(
actor_critic,
ob_rms,
config.env_name,
config.seed,
config.num_processes,
eval_log_dir,
device,
width=width,
height=height)
perf_indicator = eval_mean_score
if perf_indicator > best_perf:
best_perf = perf_indicator
best_model = True
else:
best_model = False
# record test scores
with open(os.path.join(final_output_dir, 'test_scores'),
'a+') as f:
out_s = "TEST: {}, {}, {}, {}\n".format(
str(total_num_steps), str(eval_mean_score),
str(eval_max_score),
[str(_eval_scores) for _eval_scores in eval_scores])
print(out_s, end="", file=f)
logger.info(out_s)
writer.add_scalar('data/mean_score', eval_mean_score,
total_num_steps)
writer.add_scalar('data/max_score', eval_max_score,
total_num_steps)
writer.add_scalars('test', {'mean_score': eval_mean_score},
total_num_steps)
# save for every interval-th episode or for the last epoch
if (j % config.save_interval == 0
or j == num_updates - 1) and config.save_dir != "":
logger.info(
"=> saving checkpoint to {}".format(final_output_dir))
epoch = j / config.save_interval
save_checkpoint(
{
'epoch':
epoch + 1,
'model':
get_model_name(config),
'state_dict':
actor_critic.state_dict(),
'perf':
perf_indicator,
'optimizer':
agent.optimizer.state_dict(),
'ob_rms':
getattr(utils.get_vec_normalize(envs), 'ob_rms', None)
}, best_model, final_output_dir)
final_model_state_file = os.path.join(final_output_dir, 'final_state.pth')
logger.info(
'=> saving final model state to {}'.format(final_model_state_file))
torch.save(actor_critic.state_dict(), final_model_state_file)
# export_scalars_to_json needs results from add scalars
writer.export_scalars_to_json(os.path.join(tb_log_dir, 'all_scalars.json'))
writer.close()
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")
#envs = make_vec_envs(args.env_name, args.seed, args.num_processes,
# args.gamma, args.log_dir, device, False)
envs = make_parallel_env(args.env_name, args.num_processes, args.seed, True)
'''
actor_critic = Policy(
envs.observation_space[0].shape,
envs.action_space[0],
agent_num=args.agent_num,
base_kwargs={'recurrent': args.recurrent_policy})
actor_critic.to(device)
'''
actor_critic = []
for i in range(args.agent_num):
ac = Policy(
envs.observation_space[0].shape,
envs.action_space[0],
agent_num=args.agent_num,
agent_i = i,
base_kwargs={'recurrent': args.recurrent_policy})
ac.to(device)
actor_critic.append(ac)
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)
'''
agent = []
for i in range(args.agent_num):
agent.append(algo.PPO(
actor_critic[i],
i,
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,
model_dir = args.model_dir))
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] + envs.action_space.shape[0], 100,
device)
file_name = os.path.join(
args.gail_experts_dir, "trajs_{}.pt".format(
args.env_name.split('-')[0].lower()))
expert_dataset = gail.ExpertDataset(
file_name, num_trajectories=4, subsample_frequency=20)
drop_last = len(expert_dataset) > args.gail_batch_size
gail_train_loader = torch.utils.data.DataLoader(
dataset=expert_dataset,
batch_size=args.gail_batch_size,
shuffle=True,
drop_last=drop_last)
'''
rollouts = RolloutStorage(args.num_steps, args.num_processes,
envs.observation_space[0].shape, envs.action_space[0],
actor_critic.recurrent_hidden_state_size)
obs = envs.reset()
rollouts.obs[0].copy_(torch.tensor(obs[:,0,:]))
rollouts.to(device)
'''
rollouts = []
for i in range(args.agent_num):
rollout = RolloutStorage(args.num_steps, args.num_processes,
envs.observation_space[0].shape, envs.action_space[0],
actor_critic[i].recurrent_hidden_state_size,
args.agent_num, i)
rollouts.append(rollout)
obs = envs.reset()
# pdb.set_trace()
for i in range(args.agent_num):
rollouts[i].share_obs[0].copy_(torch.tensor(obs.reshape(args.num_processes, -1)))
rollouts[i].obs[0].copy_(torch.tensor(obs[:,i,:]))
rollouts[i].to(device)
episode_rewards = deque(maxlen=10)
num_updates = int(
args.num_env_steps) // args.num_steps // args.num_processes
print(num_updates)
for j in range(num_updates):
#pdb.set_trace()
if args.use_linear_lr_decay:
# decrease learning rate linearly
for i in range(args.agent_num):
utils.update_linear_schedule(agent[i].optimizer, j, num_updates, agent[i].optimizer.lr if args.algo == "acktr" else args.lr)
for step in range(args.num_steps):
# Sample actions
value_list, action_list, action_log_prob_list, recurrent_hidden_states_list = [], [], [], []
with torch.no_grad():
for i in range(args.agent_num):
#pdb.set_trace()
value, action, action_log_prob, recurrent_hidden_states = actor_critic[i].act(
rollouts[i].share_obs[step],
rollouts[i].obs[step], rollouts[i].recurrent_hidden_states[step],
rollouts[i].masks[step])
# import pdb; pdb.set_trace()
value_list.append(value)
action_list.append(action)
action_log_prob_list.append(action_log_prob)
recurrent_hidden_states_list.append(recurrent_hidden_states)
# Obser reward and next obs
action = []
for i in range(args.num_processes):
one_env_action = []
for k in range(args.agent_num):
one_hot_action = np.zeros(envs.action_space[0].n)
one_hot_action[action_list[k][i]] = 1
one_env_action.append(one_hot_action)
action.append(one_env_action)
#start = time.time()
#pdb.set_trace()
obs, reward, done, infos = envs.step(action)
# print(obs[0][0])
# pdb.set_trace()
#end = time.time()
#print("step time: ", end-start)
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[0]])
bad_masks = torch.FloatTensor(
[[0.0] if 'bad_transition' in info.keys() else [1.0]
for info in infos[0]])
'''
masks = torch.ones(args.num_processes, 1)
bad_masks = torch.ones(args.num_processes, 1)
'''
rollouts.insert(obs, recurrent_hidden_states, action,
action_log_prob, value, reward, masks, bad_masks)
'''
#import pdb; pdb.set_trace()
for i in range(args.agent_num):
rollouts[i].insert(torch.tensor(obs.reshape(args.num_processes, -1)), torch.tensor(obs[:,i,:]),
recurrent_hidden_states, action_list[i],
action_log_prob_list[i], value_list[i], torch.tensor(reward[:, i].reshape(-1,1)), masks, bad_masks)
#import pdb; pdb.set_trace()
with torch.no_grad():
next_value_list = []
for i in range(args.agent_num):
next_value = actor_critic[i].get_value(
rollouts[i].share_obs[-1],
rollouts[i].obs[-1], rollouts[i].recurrent_hidden_states[-1],
rollouts[i].masks[-1]).detach()
next_value_list.append(next_value)
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):
discr.update(gail_train_loader, rollouts,
utils.get_vec_normalize(envs)._obfilt)
for step in range(args.num_steps):
rollouts.rewards[step] = discr.predict_reward(
rollouts.obs[step], rollouts.actions[step], args.gamma,
rollouts.masks[step])
for i in range(args.agent_num):
rollouts[i].compute_returns(next_value_list[i], args.use_gae, args.gamma,
args.gae_lambda, args.use_proper_time_limits)
#import pdb; pdb.set_trace()
for i in range(args.agent_num):
value_loss, action_loss, dist_entropy = agent[i].update(rollouts[i])
if (i == 0):
print("value loss: " + str(value_loss))
# print(value_loss)
# pdb.set_trace()
#rollouts.after_update()
obs = envs.reset()
# pdb.set_trace()
for i in range(args.agent_num):
rollouts[i].share_obs[0].copy_(torch.tensor(obs.reshape(args.num_processes, -1)))
rollouts[i].obs[0].copy_(torch.tensor(obs[:,i,:]))
rollouts[i].to(device)
# save for every interval-th episode or for the last epoch
#pdb.set_trace()
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)
if not os.path.exists(save_path + args.model_dir):
os.makedirs(save_path + args.model_dir)
for i in range(args.agent_num):
torch.save([
actor_critic[i],
getattr(utils.get_vec_normalize(envs), 'ob_rms', None)
], save_path + args.model_dir + '/agent_%i' % (i+1) + ".pt")
'''
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))
'''
'''
def onpolicy_main():
print("onpolicy main")
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
torch.set_num_threads(1)
device = torch.device("cuda:0" if args.cuda else "cpu")
summary_name = args.log_dir + '{0}_{1}'
writer = SummaryWriter(summary_name.format(args.env_name, args.save_name))
# Make vector env
envs = make_vec_envs(
args.env_name,
args.seed,
args.num_processes,
args.gamma,
args.log_dir,
device,
False,
env_kwargs=env_kwargs,
)
# agly ways to access to the environment attirubutes
if args.env_name.find('doorenv') > -1:
if args.num_processes > 1:
visionnet_input = envs.venv.venv.visionnet_input
nn = envs.venv.venv.nn
env_name = envs.venv.venv.xml_path
else:
visionnet_input = envs.venv.venv.envs[
0].env.env.env.visionnet_input
nn = envs.venv.venv.envs[0].env.env.env.nn
env_name = envs.venv.venv.envs[0].env.env.env.xml_path
dummy_obs = np.zeros(nn * 2 + 3)
else:
dummy_obs = envs.observation_space
visionnet_input = None
nn = None
if pretrained_policy_load:
print("loading", pretrained_policy_load)
actor_critic, ob_rms = torch.load(pretrained_policy_load)
else:
actor_critic = Policy(dummy_obs.shape,
envs.action_space,
base_kwargs={'recurrent': args.recurrent_policy})
if visionnet_input:
visionmodel = load_visionmodel(env_name, args.visionmodel_path,
VisionModelXYZ())
actor_critic.visionmodel = visionmodel.eval()
actor_critic.nn = nn
actor_critic.to(device)
#disable normalizer
vec_norm = get_vec_normalize(envs)
vec_norm.eval()
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)
rollouts = RolloutStorage(args.num_steps, args.num_processes,
dummy_obs.shape, envs.action_space,
actor_critic.recurrent_hidden_state_size)
full_obs = envs.reset()
initial_state = full_obs[:, :envs.action_space.shape[0]]
if args.env_name.find('doorenv') > -1 and visionnet_input:
obs = actor_critic.obs2inputs(full_obs, 0)
else:
if knob_noisy:
obs = add_noise(full_obs, 0)
else:
obs = full_obs
rollouts.obs[0].copy_(obs)
rollouts.to(device)
episode_rewards = deque(maxlen=10)
start = time.time()
num_updates = int(
args.num_env_steps) // args.num_steps // args.num_processes
for j in range(num_updates):
if args.use_linear_lr_decay:
# decrease learning rate linearly
utils.update_linear_schedule(agent.optimizer, j, num_updates,
args.lr)
pos_control = False
total_switches = 0
prev_selection = ""
for step in range(args.num_steps):
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])
next_action = action
if pos_control:
frame_skip = 2
if step % (512 / frame_skip - 1) == 0:
current_state = initial_state
next_action = current_state + next_action
for kk in range(frame_skip):
full_obs, reward, done, infos = envs.step(next_action)
current_state = full_obs[:, :envs.action_space.shape[0]]
else:
full_obs, reward, done, infos = envs.step(next_action)
# convert img to obs if door_env and using visionnet
if args.env_name.find('doorenv') > -1 and visionnet_input:
obs = actor_critic.obs2inputs(full_obs, j)
else:
if knob_noisy:
obs = add_noise(full_obs, j)
else:
obs = full_obs
for info in infos:
if 'episode' in info.keys():
episode_rewards.append(info['episode']['r'])
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)
with torch.no_grad():
next_value = actor_critic.get_value(
rollouts.obs[-1], rollouts.recurrent_hidden_states[-1],
rollouts.masks[-1]).detach()
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()
writer.add_scalar("Value loss", value_loss, j)
writer.add_scalar("action loss", action_loss, j)
writer.add_scalar("dist entropy loss", dist_entropy, j)
writer.add_scalar("Episode rewards", np.mean(episode_rewards), j)
# 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), 'ob_rms', None)
],
os.path.join(
save_path, args.env_name +
"_{}.{}.pt".format(args.save_name, 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))
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)
DR = True #Domain Randomization
################## for multiprocess world change ######################
if DR:
print("changing world")
envs.close_extras()
envs.close()
del envs
envs = make_vec_envs(
args.env_name,
args.seed,
args.num_processes,
args.gamma,
args.log_dir,
device,
False,
env_kwargs=env_kwargs,
)
full_obs = envs.reset()
if args.env_name.find('doorenv') > -1 and visionnet_input:
obs = actor_critic.obs2inputs(full_obs, j)
else:
obs = full_obs
def main():
all_episode_rewards = [] ### 记录 6/29
all_temp_rewards = [] ### 记录 6/29
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")
envs = make_vec_envs(args.env_name, args.seed, args.num_processes,
args.gamma, args.log_dir, device, False)
actor_critic = Policy(envs.observation_space.shape,
envs.action_space,
base_kwargs={'recurrent': args.recurrent_policy})
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] + envs.action_space.shape[0], 100,
device)
file_name = os.path.join(
args.gail_experts_dir,
"trajs_{}.pt".format(args.env_name.split('-')[0].lower()))
gail_train_loader = torch.utils.data.DataLoader(
gail.ExpertDataset(file_name,
num_trajectories=4,
subsample_frequency=20),
batch_size=args.gail_batch_size,
shuffle=True,
drop_last=True)
rollouts = RolloutStorage(args.num_steps, args.num_processes,
envs.observation_space.shape, envs.action_space,
actor_critic.recurrent_hidden_state_size)
obs = envs.reset()
rollouts.obs[0].copy_(obs)
rollouts.to(device)
episode_rewards = deque(maxlen=10)
start = time.time()
num_updates = int(
args.num_env_steps) // args.num_steps // args.num_processes
print('num_updates ', num_updates)
print('num_steps ', args.num_steps)
count = 0
h5_path = './data/' + args.env_name
if not os.path.exists(h5_path):
os.makedirs(h5_path)
h5_filename = h5_path + '/trajs_' + args.env_name + '_%05d.h5' % (count)
data = {}
data['states'] = []
data['actions'] = []
data['rewards'] = []
data['done'] = []
data['lengths'] = []
episode_step = 0
for j in range(num_updates): ### num-steps
temp_states = []
temp_actions = []
temp_rewards = []
temp_done = []
temp_lenthgs = []
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(
rollouts.obs[step], rollouts.recurrent_hidden_states[step],
rollouts.masks[step])
if j == 0 and step == 0:
print('obs ', type(rollouts.obs[step]),
rollouts.obs[step].shape)
print('hidden_states ',
type(rollouts.recurrent_hidden_states[step]),
rollouts.recurrent_hidden_states[step].shape)
print('action ', type(action), action.shape)
print('action prob ', type(action_log_prob),
action_log_prob.shape)
print('-' * 20)
# Obser reward and next obs
obs, reward, done, infos = envs.step(action)
#print(infos)
#print(reward)
temp_states += [np.array(rollouts.obs[step].cpu())]
temp_actions += [np.array(action.cpu())]
#temp_rewards += [np.array(reward.cpu())]
temp_rewards += [np.array([infos[0]['myrewards']])
] ### for halfcheetah不能直接用 reward ！！ 6/29
temp_done += [np.array(done)]
if j == 0 and step == 0:
print('obs ', type(obs), obs.shape)
print('reward ', type(reward), reward.shape)
print('done ', type(done), done.shape)
print('infos ', len(infos))
for k, v in infos[0].items():
print(k, v.shape)
print()
for info in infos:
if 'episode' in info.keys():
episode_rewards.append(info['episode']['r'])
all_episode_rewards += [info['episode']['r']] ### 记录 6/29
# 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)
temp_lengths = len(temp_states)
temp_states = np.concatenate(temp_states)
temp_actions = np.concatenate(temp_actions)
temp_rewards = np.concatenate(temp_rewards)
temp_done = np.concatenate(temp_done)
#print('temp_lengths',temp_lengths)
#print('temp_states', temp_states.shape)
#print('temp_actions', temp_actions.shape)
#print('temp_rewards', temp_rewards.shape)
if j > int(0.4 * num_updates):
data['states'] += [temp_states]
data['actions'] += [temp_actions]
data['rewards'] += [temp_rewards]
data['lengths'] += [temp_lengths]
data['done'] += [temp_done]
#print('temp_lengths',data['lengths'].shape)
#print('temp_states', data['states'].shape)
#print('temp_actions', data['actions'].shape)
#print('temp_rewards', data['rewards'].shape)
if args.save_expert and len(data['states']) >= 100:
with h5py.File(h5_filename, 'w') as f:
f['states'] = np.array(data['states'])
f['actions'] = np.array(data['actions'])
f['rewards'] = np.array(data['rewards'])
f['done'] = np.array(data['done'])
f['lengths'] = np.array(data['lengths'])
#print('f_lengths',f['lengths'].shape)
#print('f_states', f['states'].shape)
#print('f_actions', f['actions'].shape)
#print('f_rewards', f['rewards'].shape)
count += 1
h5_filename = h5_path + '/trajs_' + args.env_name + '_%05d.h5' % (
count)
data['states'] = []
data['actions'] = []
data['rewards'] = []
data['done'] = []
data['lengths'] = []
with torch.no_grad():
next_value = actor_critic.get_value(
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):
discr.update(gail_train_loader, rollouts,
utils.get_vec_normalize(envs)._obfilt)
for step in range(args.num_steps):
rollouts.rewards[step] = discr.predict_reward(
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)
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), 'ob_rms', None)
], os.path.join(save_path, args.env_name + "_%d.pt" % (args.seed)))
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))
#np.save(os.path.join(save_path, args.env_name+"_%d"%(args.seed)), all_episode_rewards) ### 保存记录 6/29
#print(temp_rewards)
print("temp rewards size", temp_rewards.shape, "mean",
np.mean(temp_rewards), "min", np.min(temp_rewards), "max",
np.max(temp_rewards))
all_temp_rewards += [temp_rewards]
np.savez(os.path.join(save_path,
args.env_name + "_%d" % (args.seed)),
episode=all_episode_rewards,
timestep=all_temp_rewards)
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)
'''data['states'] = np.array(data['states'])
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 + args.env_name)
eval_log_dir = log_dir + "_eval"
utils.cleanup_log_dir(log_dir)
utils.cleanup_log_dir(eval_log_dir)
log_dir2 = os.path.expanduser(args.log_dir2 + args.env_name2)
eval_log_dir2 = log_dir + "_eval"
utils.cleanup_log_dir(log_dir2)
utils.cleanup_log_dir(eval_log_dir2)
torch.set_num_threads(1)
device = torch.device("cuda:0" if args.cuda else "cpu")
import json
file_path = "config.json"
setup_json = json.load(open(file_path, 'r'))
env_conf = setup_json["Default"]
for i in setup_json.keys():
if i in args.env_name:
env_conf = setup_json[i]
# 1 game
envs = make_vec_envs(args.env_name, args.seed, args.num_processes,
args.gamma, args.log_dir, device, env_conf, False)
# 2 game
envs2 = make_vec_envs(args.env_name2, args.seed, args.num_processes,
args.gamma, args.log_dir2, device, env_conf, False)
save_model, ob_rms = torch.load('./trained_models/PongNoFrameskip-v4.pt')
from a2c_ppo_acktr.cnn import CNNBase
a = CNNBase(envs.observation_space.shape[0], recurrent=False)
actor_critic = Policy(
envs.observation_space.shape,
envs.action_space,
#(obs_shape[0], ** base_kwargs)
base=a,
#base_kwargs={'recurrent': args.recurrent_policy}
)
#actor_critic.load_state_dict(save_model.state_dict())
actor_critic.to(device)
actor_critic2 = Policy(envs2.observation_space.shape,
envs2.action_space,
base=a)
#base_kwargs={'recurrent': args.recurrent_policy})
#actor_critic2.load_state_dict(save_model.state_dict())
actor_critic2.to(device)
if args.algo == 'a2c':
agent = algo.A2C_ACKTR(actor_critic,
actor_critic2,
args.value_loss_coef,
args.entropy_coef,
lr=args.lr,
eps=args.eps,
alpha=args.alpha,
max_grad_norm=args.max_grad_norm)
rollouts = RolloutStorage(args.num_steps, args.num_processes,
envs.observation_space.shape, envs.action_space,
actor_critic.recurrent_hidden_state_size)
rollouts2 = RolloutStorage(args.num_steps, args.num_processes,
envs2.observation_space.shape,
envs2.action_space,
actor_critic2.recurrent_hidden_state_size)
obs = envs.reset()
rollouts.obs[0].copy_(obs)
rollouts.to(device)
obs2 = envs2.reset()
rollouts2.obs[0].copy_(obs2)
rollouts2.to(device)
episode_rewards = deque(maxlen=10)
episode_rewards2 = deque(maxlen=10)
start = time.time()
num_updates = int(
args.num_env_steps) // args.num_steps // args.num_processes
for j in range(num_updates):
# 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(
rollouts.obs[step], rollouts.recurrent_hidden_states[step],
rollouts.masks[step])
value2, action2, action_log_prob2, recurrent_hidden_states2, _ = actor_critic2.act(
rollouts2.obs[step],
rollouts2.recurrent_hidden_states[step],
rollouts2.masks[step])
# Obser reward and next obs
obs, reward, done, infos = envs.step(action)
obs2, reward2, done2, infos2 = envs2.step(action2)
for info in infos:
if 'episode' in info.keys():
episode_rewards.append(info['episode']['r'])
for info2 in infos2:
if 'episode' in info2.keys():
episode_rewards2.append(info2['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])
masks2 = torch.FloatTensor([[0.0] if done_ else [1.0]
for done_ in done2])
bad_masks2 = torch.FloatTensor(
[[0.0] if 'bad_transition' in info2.keys() else [1.0]
for info2 in infos2])
rollouts.insert(obs, recurrent_hidden_states, action,
action_log_prob, value, reward, masks, bad_masks)
rollouts2.insert(obs2, recurrent_hidden_states2, action2,
action_log_prob2, value2, reward2, masks2,
bad_masks2)
with torch.no_grad():
next_value = actor_critic.get_value(
rollouts.obs[-1], rollouts.recurrent_hidden_states[-1],
rollouts.masks[-1]).detach()
next_value2 = actor_critic2.get_value(
rollouts2.obs[-1], rollouts2.recurrent_hidden_states[-1],
rollouts2.masks[-1]).detach()
rollouts.compute_returns(next_value, args.use_gae, args.gamma,
args.gae_lambda, args.use_proper_time_limits)
rollouts2.compute_returns(next_value2, args.use_gae, args.gamma,
args.gae_lambda, args.use_proper_time_limits)
value_loss, action_loss, dist_entropy, value_loss2, action_loss2, dist_entropy2 = agent.update(
rollouts, rollouts2)
rollouts.after_update()
rollouts2.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), 'ob_rms', None)
], os.path.join(save_path, args.env_name + ".pt"))
torch.save([
actor_critic2,
getattr(utils.get_vec_normalize(envs2), 'ob_rms2', None)
], os.path.join(save_path, args.env_name2 + ".pt"))
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))
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_rewards2), np.mean(episode_rewards2),
np.median(episode_rewards2), np.min(episode_rewards2),
np.max(episode_rewards2), dist_entropy2, value_loss2,
action_loss2))
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)
ob_rms2 = utils.get_vec_normalize(envs2).ob_rms
evaluate(actor_critic2, ob_rms2, args.env_name2, args.seed,
args.num_processes, eval_log_dir2, device)
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")
# coinrun environments need to be treated differently.
coinrun_envs = {
'CoinRun': 'standard',
'CoinRun-Platforms': 'platform',
'Random-Mazes': 'maze'
}
envs = make_vec_envs(args.env_name,
args.seed,
args.num_processes,
args.gamma,
args.log_dir,
device,
False,
coin_run_level=args.num_levels,
difficulty=args.high_difficulty,
coin_run_seed=args.seed)
if args.env_name in coinrun_envs.keys():
observation_space_shape = (3, 64, 64)
args.save_dir = args.save_dir + "/NUM_LEVELS_{}".format(
args.num_levels) # Save the level info in the
else:
observation_space_shape = envs.observation_space.shape
# trained model name
if args.continue_ppo_training:
actor_critic, _ = torch.load(os.path.join(args.check_point,
args.env_name + ".pt"),
map_location=torch.device(device))
elif args.cor_gail:
embed_size = args.embed_size
actor_critic = Policy(observation_space_shape,
envs.action_space,
hidden_size=args.hidden_size,
embed_size=embed_size,
base_kwargs={'recurrent': args.recurrent_policy})
actor_critic.to(device)
correlator = Correlator(observation_space_shape,
envs.action_space,
hidden_dim=args.hidden_size,
embed_dim=embed_size,
lr=args.lr,
device=device)
correlator.to(device)
embeds = torch.zeros(1, embed_size)
else:
embed_size = 0
actor_critic = Policy(observation_space_shape,
envs.action_space,
hidden_size=args.hidden_size,
base_kwargs={'recurrent': args.recurrent_policy})
actor_critic.to(device)
embeds = None
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,
use_clipped_value_loss=True,
ftrl_mode=args.cor_gail or args.no_regret_gail,
correlated_mode=args.cor_gail)
elif args.algo == 'acktr':
agent = algo.A2C_ACKTR(actor_critic,
args.value_loss_coef,
args.entropy_coef,
acktr=True)
if args.gail or args.no_regret_gail or args.cor_gail:
file_name = os.path.join(
args.gail_experts_dir,
"trajs_{}.pt".format(args.env_name.split('-')[0].lower()))
expert_dataset = gail.ExpertDataset(
file_name, num_trajectories=50,
subsample_frequency=1) #if subsample set to a different number,
# grad_pen might need adjustment
drop_last = len(expert_dataset) > args.gail_batch_size
gail_train_loader = torch.utils.data.DataLoader(
dataset=expert_dataset,
batch_size=args.gail_batch_size,
shuffle=True,
drop_last=drop_last)
if args.gail:
discr = gail.Discriminator(observation_space_shape,
envs.action_space,
device=device)
if args.no_regret_gail or args.cor_gail:
queue = deque(
maxlen=args.queue_size
) # Strategy Queues: Each element of a queue is a dicr strategy
agent_queue = deque(
maxlen=args.queue_size
) # Strategy Queues: Each element of a queue is an agent strategy
pruning_frequency = 1
if args.no_regret_gail:
discr = regret_gail.NoRegretDiscriminator(observation_space_shape,
envs.action_space,
device=device)
if args.cor_gail:
discr = cor_gail.CorDiscriminator(observation_space_shape,
envs.action_space,
hidden_size=args.hidden_size,
embed_size=embed_size,
device=device)
discr.to(device)
rollouts = RolloutStorage(args.num_steps, args.num_processes,
observation_space_shape, envs.action_space,
actor_critic.recurrent_hidden_state_size,
embed_size)
obs = envs.reset()
rollouts.obs[0].copy_(obs)
if args.cor_gail:
rollouts.embeds[0].copy_(embeds)
rollouts.to(device)
episode_rewards = deque(maxlen=10)
start = time.time()
num_updates = int(
args.num_env_steps) // args.num_steps // args.num_processes
for j in range(num_updates):
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 # Roll-out
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], rollouts.embeds[step])
obs, reward, done, infos = envs.step(action.to('cpu'))
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(obs, recurrent_hidden_states, action,
action_log_prob, value, reward, masks, bad_masks)
# Sample mediating/correlating actions # Correlated Roll-out
if args.cor_gail:
embeds, embeds_log_prob, mean = correlator.act(
rollouts.obs[step], rollouts.actions[step])
rollouts.insert_embedding(embeds, embeds_log_prob)
with torch.no_grad():
next_value = actor_critic.get_value(
rollouts.obs[-1], rollouts.recurrent_hidden_states[-1],
rollouts.masks[-1], rollouts.embeds[-1]).detach()
if args.gail or args.no_regret_gail or args.cor_gail:
if args.env_name not in {'CoinRun', 'Random-Mazes'}:
if j >= 10:
envs.venv.eval()
gail_epoch = args.gail_epoch
if args.gail:
if j < 10:
gail_epoch = 100 # Warm up
# no need for gail epoch or warm up in the no-regret case and cor_gail.
for _ in range(gail_epoch):
if utils.get_vec_normalize(envs):
obfilt = utils.get_vec_normalize(envs)._obfilt
else:
obfilt = None
if args.gail:
discr.update(gail_train_loader, rollouts, obfilt)
if args.no_regret_gail or args.cor_gail:
last_strategy = discr.update(gail_train_loader, rollouts,
queue, args.max_grad_norm,
obfilt, j)
for step in range(args.num_steps):
if args.gail:
rollouts.rewards[step] = discr.predict_reward(
rollouts.obs[step], rollouts.actions[step], args.gamma,
rollouts.masks[step])
if args.no_regret_gail:
rollouts.rewards[step] = discr.predict_reward(
rollouts.obs[step], rollouts.actions[step], args.gamma,
rollouts.masks[step], queue)
if args.cor_gail:
rollouts.rewards[
step], correlator_reward = discr.predict_reward(
rollouts.obs[step], rollouts.actions[step],
rollouts.embeds[step], args.gamma,
rollouts.masks[step], queue)
rollouts.correlated_reward[step] = correlator_reward
rollouts.compute_returns(next_value, args.use_gae, args.gamma,
args.gae_lambda, args.use_proper_time_limits)
if args.gail:
value_loss, action_loss, dist_entropy = agent.update(rollouts, j)
elif args.no_regret_gail or args.cor_gail:
value_loss, action_loss, dist_entropy, agent_gains, agent_strategy = \
agent.mixed_update(rollouts, agent_queue, j)
if args.cor_gail:
correlator.update(rollouts, agent_gains, args.max_grad_norm)
if args.no_regret_gail or args.cor_gail:
queue, _ = utils.queue_update(queue, pruning_frequency,
args.queue_size, j, last_strategy)
agent_queue, pruning_frequency = utils.queue_update(
agent_queue, pruning_frequency, args.queue_size, j,
agent_strategy)
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
if not args.cor_gail:
torch.save([
actor_critic,
getattr(utils.get_vec_normalize(envs), 'ob_rms', None)
], os.path.join(save_path, args.env_name + ".pt"))
else:
print("saving models in {}".format(
os.path.join(save_path, args.env_name)))
torch.save(
correlator.state_dict(),
os.path.join(save_path, args.env_name + "correlator.pt"))
torch.save([
actor_critic.state_dict(),
getattr(utils.get_vec_normalize(envs), 'ob_rms', None)
], os.path.join(save_path, args.env_name + "actor.pt"))
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},"
" value loss/action loss {:.1f}/{}".format(
j, total_num_steps, int(total_num_steps / (end - start)),
len(episode_rewards), np.mean(episode_rewards),
np.median(episode_rewards), value_loss, action_loss))
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 main():
chrono = exp.chrono()
envs = make_vec_envs(args.env_name, args.seed, args.num_processes,
args.gamma, args.log_dir, device, False)
actor_critic = Policy(envs.observation_space.shape,
envs.action_space,
base_kwargs={'recurrent': args.recurrent_policy})
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)
rollouts = RolloutStorage(args.num_steps, args.num_processes,
envs.observation_space.shape, envs.action_space,
actor_critic.recurrent_hidden_state_size)
obs = envs.reset()
rollouts.obs[0].copy_(obs)
rollouts.to(device)
episode_rewards = deque(maxlen=10)
start = time.time()
num_updates = int(
args.num_env_steps) // args.num_steps // args.num_processes
for j in range(args.repeat):
with chrono.time('train') as t:
for n in range(args.number):
if args.use_linear_lr_decay:
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(
rollouts.obs[step],
rollouts.recurrent_hidden_states[step],
rollouts.masks[step])
# 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'])
# 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)
with torch.no_grad():
next_value = actor_critic.get_value(
rollouts.obs[-1], rollouts.recurrent_hidden_states[-1],
rollouts.masks[-1]).detach()
# ---
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)
exp.log_batch_loss(action_loss)
exp.log_metric('value_loss', value_loss)
rollouts.after_update()
total_num_steps = (j + 1) * args.num_processes * args.num_steps
if j % args.log_interval == 0 and len(episode_rewards) > 1:
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))
# -- number
# -- chrono
exp.show_eta(j, t)
# -- epoch
exp.report()
envs.close()
def main():
torch.set_num_threads(1)
device = torch.device("cuda:0" if args_iko.cuda else "cpu")
if args_iko.vis:
from visdom import Visdom
viz = Visdom(port=args_iko.port)
win = None
envs = make_vec_envs(args_iko.env_name, args_iko.seed,
args_iko.num_processes, args_iko.gamma,
args_iko.log_dir, args_iko.add_timestep, device,
False)
actor_critic = Policy(envs.observation_space.shape,
envs.action_space,
base_kwargs={'recurrent': args_iko.recurrent_policy})
actor_critic.to(device)
action_shape = 3
reward_model = RewardModel(11 * 11 * 6, 1, 64, 64)
reward_model.to(device)
if args_iko.algo == 'a2c':
agent = algo.A2C_ACKTR(actor_critic,
args_iko.value_loss_coef,
args_iko.entropy_coef,
lr=args_iko.lr,
eps=args_iko.eps,
alpha=args_iko.alpha,
max_grad_norm=args_iko.max_grad_norm)
elif args_iko.algo == 'ppo':
agent = algo.PPO(actor_critic,
args_iko.clip_param,
args_iko.ppo_epoch,
args_iko.num_mini_batch,
args_iko.value_loss_coef,
args_iko.entropy_coef,
args_iko.use_singh,
reward_model,
lr=args_iko.lr,
eps=args_iko.eps,
max_grad_norm=args_iko.max_grad_norm)
elif args_iko.algo == 'acktr':
agent = algo.A2C_ACKTR(actor_critic,
args_iko.value_loss_coef,
args_iko.entropy_coef,
acktr=True)
rollouts = RolloutStorage(args_iko.num_steps, args_iko.num_processes,
envs.observation_space.shape, envs.action_space,
actor_critic.recurrent_hidden_state_size)
obs = envs.reset()
rollouts.obs[0].copy_(obs)
rollouts.to(device)
episode_rewards = deque(maxlen=10)
start = time.time()
for j in range(num_updates):
if args_iko.use_linear_lr_decay:
# decrease learning rate linearly
if args_iko.algo == "acktr":
# use optimizer's learning rate since it's hard-coded in kfac.py
update_linear_schedule(agent.optimizer, j, num_updates,
agent.optimizer.lr)
else:
update_linear_schedule(agent.optimizer, j, num_updates,
args_iko.lr)
if args_iko.algo == 'ppo' and args_iko.use_linear_clip_decay:
agent.clip_param = args_iko.clip_param * (1 -
j / float(num_updates))
reward_train = []
reward_block_penalty = []
reward_bel_gt = []
reward_bel_gt_nonlog = []
reward_infogain = []
reward_bel_ent = []
reward_hit = []
reward_dist = []
reward_inv_dist = []
for step in range(args_iko.num_steps):
# Sample actions
# print(step, args_iko.num_steps)
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])
# Obser reward and next obs
obs, reward, done, infos = envs.step(action)
reward_train.append(reward)
# print("infos is ", infos)
# reward_b.append(infos[0]['auxillary_reward'])
# print("infos is ",infos[0]['auxillary_reward'])
reward_block_penalty.append(infos[0]['reward_block_penalty'])
reward_bel_gt.append(infos[0]['reward_bel_gt'])
reward_bel_gt_nonlog.append(infos[0]['reward_bel_gt_nonlog'])
reward_infogain.append(infos[0]['reward_infogain'])
reward_bel_ent.append(infos[0]['reward_bel_ent'])
reward_hit.append(infos[0]['reward_hit'])
reward_dist.append(infos[0]['reward_dist'])
reward_inv_dist.append(infos[0]['reward_inv_dist'])
# print(reward)
reward.to(device)
reward_model.to(device)
if args_iko.use_singh:
# print("using learning IR")
my_reward = reward_model(obs.clone().to(device),
action.clone().float()).detach()
my_reward.to(device)
reward = reward + args_iko.singh_coef * my_reward.type(
torch.FloatTensor)
# for info in infos:
# if 'episode' in info.keys():
# episode_rewards.append(info['episode']['r'])
# print("infos is ",infos[0]['auxillary_reward'])
# print("info is",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])
rollouts.insert(obs, recurrent_hidden_states, action,
action_log_prob, value, reward, masks)
# print("mean reward_a", np.mean(reward_train))
# print("mean reward_block_penalty", np.mean(reward_block_penalty))
# print("mean reward_bel_gt", np.mean(reward_bel_gt))
# print("mean reward_bel_gt_nonlog", np.mean(reward_bel_gt_nonlog))
# print("mean reward_infogain", np.mean(reward_infogain))
# print("mean reward_bel_ent", np.mean(reward_bel_ent))
# print("mean reward_hit", np.mean(reward_hit))
# print("mean reward_dist", np.mean(reward_dist))
# print("mean reward_inv_dist", np.mean(reward_inv_dist))
total_num_steps = (j + 1) * args_iko.num_processes * args_iko.num_steps
writer.add_scalar('mean_reward_train', np.mean(reward_train),
total_num_steps)
writer.add_scalar('mean_reward_block_penalty',
np.mean(reward_block_penalty), total_num_steps)
writer.add_scalar('mean_reward_bel_gt', np.mean(reward_bel_gt),
total_num_steps)
writer.add_scalar('mean_reward_bel_gt_nonlog',
np.mean(reward_bel_gt_nonlog), total_num_steps)
writer.add_scalar('mean_reward_infogain', np.mean(reward_infogain),
total_num_steps)
writer.add_scalar('mean_reward_bel_ent', np.mean(reward_bel_ent),
total_num_steps)
writer.add_scalar('mean_reward_hit', np.mean(reward_hit),
total_num_steps)
writer.add_scalar('mean_reward_dist', np.mean(reward_dist),
total_num_steps)
writer.add_scalar('mean_reward_inv_dist', np.mean(reward_inv_dist),
total_num_steps)
with torch.no_grad():
next_value = actor_critic.get_value(
rollouts.obs[-1], rollouts.recurrent_hidden_states[-1],
rollouts.masks[-1]).detach()
rollouts.compute_returns(next_value, args_iko.use_gae, args_iko.gamma,
args_iko.tau)
value_loss, action_loss, dist_entropy = agent.update(rollouts)
rollouts.after_update()
# save for every interval-th episode or for the last epoch
if (j % args_iko.save_interval == 0
or j == num_updates - 1) and args_iko.save_dir != "":
save_path = os.path.join(args_iko.save_dir, args_iko.algo)
try:
os.makedirs(save_path)
except OSError:
pass
# A really ugly way to save a model to CPU
save_model = actor_critic
if args_iko.cuda:
save_model = copy.deepcopy(actor_critic).cpu()
save_model = [
save_model,
getattr(get_vec_normalize(envs), 'ob_rms', None)
]
torch.save(
save_model,
os.path.join(
save_path, 'ugl' + str(args_iko.use_gt_likelihood) +
'block-pen-' + str(args_iko.penalty_for_block) + '_' +
'explore-' + str(args_iko.rew_explore) + '_' + 'bel-new-' +
str(args_iko.rew_bel_new) + '_' + 'bel-ent-' +
str(args_iko.rew_bel_ent) + '_' + 'infogain-' +
str(args_iko.rew_infogain) + '_' + 'bel-gt-nolog-' +
str(args_iko.rew_bel_gt_nonlog) + '_' + 'bel-gt-' +
str(args_iko.rew_bel_gt) + '_' + 'dist-' +
str(args_iko.rew_dist) + '_' + 'hit-' +
str(args_iko.rew_hit) + '_' + 'inv-dist-' +
str(args_iko.rew_inv_dist) + args_iko.algo + ".pt"))
total_num_steps = (j + 1) * args_iko.num_processes * args_iko.num_steps
if j % args_iko.log_interval == 0 and len(episode_rewards) > 1:
end = time.time()
print("mean reward_a", np.mean(reward_a))
print("mean_reward_b", np.mean(reward_b))
# 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))
# writer.add_scalar('mean_reward', np.mean(episode_rewards), total_num_steps)
# writer.add_scalar('min_reward', np.min(episode_rewards), total_num_steps)
# writer.add_scalar('max_reward', np.max(episode_rewards), total_num_steps)
# writer.add_scalar('success_rate', np.mean(episode_successes), total_num_steps)
if (args_iko.eval_interval is not None and len(episode_rewards) > 1
and j % args_iko.eval_interval == 0):
eval_envs = make_vec_envs(args_iko.env_name,
args_iko.seed + args_iko.num_processes,
args_iko.num_processes, args_iko.gamma,
eval_log_dir, args_iko.add_timestep,
device, True)
vec_norm = get_vec_normalize(eval_envs)
if vec_norm is not None:
vec_norm.eval()
vec_norm.ob_rms = get_vec_normalize(envs).ob_rms
eval_episode_rewards = []
obs = eval_envs.reset()
eval_recurrent_hidden_states = torch.zeros(
args_iko.num_processes,
actor_critic.recurrent_hidden_state_size,
device=device)
eval_masks = torch.zeros(args_iko.num_processes, 1, device=device)
while len(eval_episode_rewards) < 10:
with torch.no_grad():
_, action, _, eval_recurrent_hidden_states = actor_critic.act(
obs,
eval_recurrent_hidden_states,
eval_masks,
deterministic=True)
# Obser reward and next obs
obs, reward, done, infos = eval_envs.step(action)
eval_masks = torch.FloatTensor([[0.0] if done_ else [1.0]
for done_ in done])
for info in infos:
if 'episode' in info.keys():
eval_episode_rewards.append(info['episode']['r'])
eval_envs.close()
print(" Evaluation using {} episodes: mean reward {:.5f}\n".format(
len(eval_episode_rewards), np.mean(eval_episode_rewards)))
if args_iko.vis and j % args_iko.vis_interval == 0:
try:
# Sometimes monitor doesn't properly flush the outputs
win = visdom_plot(viz, win, args_iko.log_dir,
args_iko.env_name, args_iko.algo,
args_iko.num_env_steps)
except IOError:
pass
writer.close()
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
args_dir, logs_dir, models_dir, samples_dir = get_all_save_paths(
args, 'pretrain', combine_action=args.combine_action)
eval_log_dir = logs_dir + "_eval"
utils.cleanup_log_dir(logs_dir)
utils.cleanup_log_dir(eval_log_dir)
_, _, intrinsic_models_dir, _ = get_all_save_paths(args,
'learn_reward',
load_only=True)
if args.load_iter != 'final':
intrinsic_model_file_name = os.path.join(
intrinsic_models_dir,
args.env_name + '_{}.pt'.format(args.load_iter))
else:
intrinsic_model_file_name = os.path.join(
intrinsic_models_dir, args.env_name + '.pt'.format(args.load_iter))
intrinsic_arg_file_name = os.path.join(args_dir, 'command.txt')
# save args to arg_file
with open(intrinsic_arg_file_name, 'w') as f:
json.dump(args.__dict__, f, indent=2)
torch.set_num_threads(1)
device = torch.device("cuda:0" if args.cuda else "cpu")
envs = make_vec_envs(args.env_name, args.seed, args.num_processes,
args.gamma, logs_dir, device, False)
actor_critic = Policy(envs.observation_space.shape,
envs.action_space,
base_kwargs={'recurrent': args.recurrent_policy})
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)
else:
raise NotImplementedError
if args.use_intrinsic:
obs_shape = envs.observation_space.shape
if len(obs_shape) == 3:
action_dim = envs.action_space.n
elif len(obs_shape) == 1:
action_dim = envs.action_space.shape[0]
if 'NoFrameskip' in args.env_name:
file_name = os.path.join(
args.experts_dir, "trajs_ppo_{}.pt".format(
args.env_name.split('-')[0].replace('NoFrameskip',
'').lower()))
else:
file_name = os.path.join(
args.experts_dir,
"trajs_ppo_{}.pt".format(args.env_name.split('-')[0].lower()))
rff = RewardForwardFilter(args.gamma)
intrinsic_rms = RunningMeanStd(shape=())
if args.intrinsic_module == 'icm':
print('Loading pretrained intrinsic module: %s' %
intrinsic_model_file_name)
inverse_model, forward_dynamics_model, encoder = torch.load(
intrinsic_model_file_name)
icm = IntrinsicCuriosityModule(envs, device, inverse_model, forward_dynamics_model, \
inverse_lr=args.intrinsic_lr, forward_lr=args.intrinsic_lr,\
)
if args.intrinsic_module == 'vae':
print('Loading pretrained intrinsic module: %s' %
intrinsic_model_file_name)
vae = torch.load(intrinsic_model_file_name)
icm = GenerativeIntrinsicRewardModule(envs, device, \
vae, lr=args.intrinsic_lr, \
)
rollouts = RolloutStorage(args.num_steps, args.num_processes,
envs.observation_space.shape, envs.action_space,
actor_critic.recurrent_hidden_state_size)
obs = envs.reset()
rollouts.obs[0].copy_(obs)
rollouts.to(device)
episode_rewards = deque(maxlen=10)
start = time.time()
num_updates = int(
args.num_env_steps) // args.num_steps // args.num_processes
for j in range(num_updates):
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):
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])
obs, reward, done, infos = envs.step(action)
next_obs = obs
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(obs, next_obs, recurrent_hidden_states, action,
action_log_prob, value, reward, masks, bad_masks)
with torch.no_grad():
next_value = actor_critic.get_value(
rollouts.obs[-1], rollouts.recurrent_hidden_states[-1],
rollouts.masks[-1]).detach()
if args.use_intrinsic:
for step in range(args.num_steps):
state = rollouts.obs[step]
action = rollouts.actions[step]
next_state = rollouts.next_obs[step]
if args.intrinsic_module == 'icm':
state = encoder(state)
next_state = encoder(next_state)
with torch.no_grad():
rollouts.rewards[
step], pred_next_state = icm.calculate_intrinsic_reward(
state, action, next_state, args.lambda_true_action)
if args.standardize == 'True':
buf_rews = rollouts.rewards.cpu().numpy()
intrinsic_rffs = np.array(
[rff.update(rew) for rew in buf_rews.T])
rffs_mean, rffs_std, rffs_count = mpi_moments(
intrinsic_rffs.ravel())
intrinsic_rms.update_from_moments(rffs_mean, rffs_std**2,
rffs_count)
mean = intrinsic_rms.mean
std = np.asarray(np.sqrt(intrinsic_rms.var))
rollouts.rewards = rollouts.rewards / torch.from_numpy(std).to(
device)
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()
# 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(models_dir, args.algo)
policy_file_name = os.path.join(save_path, args.env_name + '.pt')
try:
os.makedirs(save_path)
except OSError:
pass
torch.save([
actor_critic,
getattr(utils.get_vec_normalize(envs), 'ob_rms', None)
], policy_file_name)
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(args.env_name, 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))
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 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:" + str(args.cuda_id) if args.cuda else "cpu")
envs = make_vec_envs(args.env_name, args.seed, args.num_processes,
args.gamma, args.log_dir, device, False)
actor_critic = Policy(envs.observation_space.shape,
envs.action_space,
base_kwargs={'recurrent': args.recurrent_policy})
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] + envs.action_space.shape[0], 100,
device)
file_name = os.path.join(
args.gail_experts_dir,
"trajs_{}.pt".format(args.env_name.split('-')[0].lower()))
expert_dataset = gail.ExpertDataset(file_name,
num_trajectories=4,
subsample_frequency=20)
drop_last = len(expert_dataset) > args.gail_batch_size
gail_train_loader = torch.utils.data.DataLoader(
dataset=expert_dataset,
batch_size=args.gail_batch_size,
shuffle=True,
drop_last=drop_last)
########## file related
filename = args.env_name + "_" + args.algo + "_n" + str(args.max_episodes)
if args.attack:
filename += "_" + args.type + "_" + args.aim
filename += "_s" + str(args.stepsize) + "_m" + str(
args.maxiter) + "_r" + str(args.radius) + "_f" + str(args.frac)
if args.run >= 0:
filename += "_run" + str(args.run)
logger = get_log(args.logdir + filename + "_" + current_time)
logger.info(args)
rew_file = open(args.resdir + filename + ".txt", "w")
if args.compute:
radius_file = open(
args.resdir + filename + "_radius" + "_s" + str(args.stepsize) +
"_m" + str(args.maxiter) + "_th" + str(args.dist_thres) + ".txt",
"w")
if args.type == "targ" or args.type == "fgsm":
targ_file = open(args.resdir + filename + "_targ.txt", "w")
num_updates = int(
args.num_env_steps) // args.num_steps // args.num_processes
if args.type == "wb":
attack_net = WbAttacker(agent,
envs,
int(args.frac * num_updates),
num_updates,
args,
device=device)
if args.type == "bb":
attack_net = BbAttacker(agent,
envs,
int(args.frac * num_updates),
num_updates,
args,
device=device)
elif args.type == "rand":
attack_net = RandAttacker(envs,
radius=args.radius,
frac=args.frac,
maxat=int(args.frac * num_updates),
device=device)
elif args.type == "semirand":
attack_net = WbAttacker(agent,
envs,
int(args.frac * num_updates),
num_updates,
args,
device,
rand_select=True)
elif args.type == "targ":
if isinstance(envs.action_space, Discrete):
action_dim = envs.action_space.n
target_policy = action_dim - 1
elif isinstance(envs.action_space, Box):
action_dim = envs.action_space.shape[0]
target_policy = torch.zeros(action_dim)
# target_policy[-1] = 1
print("target policy is", target_policy)
attack_net = TargAttacker(agent,
envs,
int(args.frac * num_updates),
num_updates,
target_policy,
args,
device=device)
elif args.type == "fgsm":
if isinstance(envs.action_space, Discrete):
action_dim = envs.action_space.n
target_policy = action_dim - 1
elif isinstance(envs.action_space, Box):
action_dim = envs.action_space.shape[0]
target_policy = torch.zeros(action_dim)
def targ_policy(obs):
return target_policy
attack_net = FGSMAttacker(envs,
agent,
targ_policy,
radius=args.radius,
frac=args.frac,
maxat=int(args.frac * num_updates),
device=device)
# if args.aim == "obs" or aim == "hybrid":
# obs_space = gym.make(args.env_name).observation_space
# attack_net.set_obs_range(obs_space.low, obs_space.high)
rollouts = RolloutStorage(args.num_steps, args.num_processes,
envs.observation_space.shape, envs.action_space,
actor_critic.recurrent_hidden_state_size)
obs = envs.reset()
rollouts.obs[0].copy_(obs)
rollouts.to(device)
episode_rewards = deque(maxlen=10)
episode = 0
start = time.time()
for j in range(num_updates):
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
if args.type == "fgsm":
# print("before", rollouts.obs[step])
rollouts.obs[step] = attack_net.attack(
rollouts.obs[step], rollouts.recurrent_hidden_states[step],
rollouts.masks[step]).clone()
# print("after", rollouts.obs[step])
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])
if args.type == "targ" or args.type == "fgsm":
if isinstance(envs.action_space, Discrete):
num_target = (
action == target_policy).nonzero()[:, 0].size()[0]
targ_file.write(
str(num_target / args.num_processes) + "\n")
print("percentage of target:",
num_target / args.num_processes)
elif isinstance(envs.action_space, Box):
target_action = target_policy.repeat(action.size()[0], 1)
targ_file.write(
str(
torch.norm(action - target_action).item() /
args.num_processes) + "\n")
print("percentage of target:",
torch.sum(action).item() / args.num_processes)
# Obser reward and next obs
obs, reward, done, infos = envs.step(action.cpu())
for info in infos:
if 'episode' in info.keys():
episode_rewards.append(info['episode']['r'])
# rew_file.write("episode: {}, total reward: {}\n".format(episode, info['episode']['r']))
episode += 1
# 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)
with torch.no_grad():
next_value = actor_critic.get_value(
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):
discr.update(gail_train_loader, rollouts,
utils.get_vec_normalize(envs)._obfilt)
for step in range(args.num_steps):
rollouts.rewards[step] = discr.predict_reward(
rollouts.obs[step], rollouts.actions[step], args.gamma,
rollouts.masks[step])
if args.attack and args.type != "fgsm":
if args.aim == "reward":
logger.info(rollouts.rewards.flatten())
rollouts.rewards = attack_net.attack_r_general(
rollouts, next_value).clone().detach()
logger.info("after attack")
logger.info(rollouts.rewards.flatten())
elif args.aim == "obs":
origin = rollouts.obs.clone()
rollouts.obs = attack_net.attack_s_general(
rollouts, next_value).clone().detach()
logger.info(origin)
logger.info("after")
logger.info(rollouts.obs)
elif args.aim == "action":
origin = torch.flatten(rollouts.actions).clone()
rollouts.actions = attack_net.attack_a_general(
rollouts, next_value).clone().detach()
logger.info("attack value")
logger.info(torch.flatten(rollouts.actions) - origin)
elif args.aim == "hybrid":
res_aim, attack = attack_net.attack_hybrid(
rollouts, next_value, args.radius_s, args.radius_a,
args.radius_r)
print("attack ", res_aim)
if res_aim == "obs":
origin = rollouts.obs.clone()
rollouts.obs = attack.clone().detach()
logger.info(origin)
logger.info("attack obs")
logger.info(rollouts.obs)
elif res_aim == "action":
origin = torch.flatten(rollouts.actions).clone()
rollouts.actions = attack.clone().detach()
logger.info("attack action")
logger.info(torch.flatten(rollouts.actions) - origin)
elif res_aim == "reward":
logger.info(rollouts.rewards.flatten())
rollouts.rewards = attack.clone().detach()
logger.info("attack reward")
logger.info(rollouts.rewards.flatten())
if args.compute:
stable_radius = attack_net.compute_radius(rollouts, next_value)
print("stable radius:", stable_radius)
radius_file.write("update: {}, radius: {}\n".format(
j, np.round(stable_radius, decimals=3)))
rollouts.compute_returns(next_value, args.use_gae, args.gamma,
args.gae_lambda, args.use_proper_time_limits)
if args.attack and args.type == "bb":
attack_net.learning(rollouts)
value_loss, action_loss, dist_entropy = agent.update(rollouts)
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), 'ob_rms', None)
], os.path.join(save_path, args.env_name + ".pt"))
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))
rew_file.write("updates: {}, mean reward: {}\n".format(
j, np.mean(episode_rewards)))
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)
# if episode > args.max_episodes:
# print("reach episodes limit")
# break
if args.attack:
logger.info("total attacks: {}\n".format(attack_net.attack_num))
print("total attacks: {}\n".format(attack_net.attack_num))
rew_file.close()
if args.compute:
radius_file.close()
if args.type == "targ" or args.type == "fgsm":
targ_file.close()
def run(self):
args = self.args
torch.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)
print("CUDA is available: ", torch.cuda.is_available())
if args.cuda:
print("CUDA enabled")
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
else:
if args.cuda_deterministic:
print("Warning CUDA is requested but is not available")
else:
print("CUDA disabled")
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)
print("get_num_thread", torch.get_num_threads())
device = torch.device("cuda:0" if args.cuda else "cpu")
envs = make_vec_envs(args.env_name, self.config_parameters, args.seed,
args.num_processes, args.gamma, args.log_dir,
device, False)
actor_critic = create_IAM_model(envs, args, self.config_parameters)
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)
# This algorithm should be used for the reproduction project.
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] + envs.action_space.shape[0],
100, device)
file_name = os.path.join(
args.gail_experts_dir,
"trajs_{}.pt".format(args.env_name.split('-')[0].lower()))
expert_dataset = gail.ExpertDataset(file_name,
num_trajectories=4,
subsample_frequency=20)
drop_last = len(expert_dataset) > args.gail_batch_size
gail_train_loader = torch.utils.data.DataLoader(
dataset=expert_dataset,
batch_size=args.gail_batch_size,
shuffle=True,
drop_last=drop_last)
rollouts = RolloutStorage(args.num_steps, args.num_processes,
envs.observation_space.shape,
envs.action_space,
actor_critic.recurrent_hidden_state_size)
obs = envs.reset()
rollouts.obs[0].copy_(obs)
rollouts.to(device)
# Always return the average of the last 100 steps. This means the average is sampled.
episode_rewards = deque(maxlen=100)
start = time.time()
num_updates = int(
args.num_env_steps) // args.num_steps // args.num_processes
for j in range(num_updates):
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(
rollouts.obs[step],
rollouts.recurrent_hidden_states[step],
rollouts.masks[step])
# 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'])
# 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)
with torch.no_grad():
next_value = actor_critic.get_value(
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):
discr.update(gail_train_loader, rollouts,
utils.get_vec_normalize(envs)._obfilt)
for step in range(args.num_steps):
rollouts.rewards[step] = discr.predict_reward(
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)
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, self.model_file_name))
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()
elapsed_time = end - start
data = [
j, # Updates
total_num_steps, # timesteps
int(total_num_steps / elapsed_time), # FPS
len(episode_rewards), # Only useful for print statement
np.mean(episode_rewards), # mean of rewards
np.median(episode_rewards), # median of rewards
np.min(episode_rewards), # min rewards
np.max(episode_rewards), # max rewards
dist_entropy,
value_loss,
action_loss,
elapsed_time
]
output = ''.join([str(x) + ',' for x in data])
self.data_saver.append(output)
print(
f"Updates {data[0]}, num timesteps {data[1]}, FPS {data[2]}, elapsed time {int(data[11])} sec. Last {data[3]} training episodes: mean/median reward {data[4]:.2f}/{data[5]:.2f}, min/max reward {data[6]:.1f}/{data[7]:.1f}",
end="\r")
if (args.eval_interval is not None and len(episode_rewards) > 1
and j % args.eval_interval == 0):
obs_rms = utils.get_vec_normalize(envs).obs_rms
evaluate(actor_critic, obs_rms, args.env_name, args.seed,
args.num_processes, eval_log_dir, device)
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 main():
parser = argparse.ArgumentParser()
parser.add_argument('--gpu', type=str, default="")
parser.add_argument('--epochs', type=int, default=30)
parser.add_argument('--optimizer', type=str, default="mixture")
parser.add_argument('--batch-size', type=int, default=1000)
parser.add_argument('--worker', type=int, default=8)
parser.add_argument('--dataset', type=str, default="CIFAR10")
parser.add_argument('--log-dir', type=str, default="logs")
parser.add_argument('--num-classes', type=int, help="number of classes")
parser.add_argument('--use-log-loss', action="store_true")
parser.add_argument('--lr-meta',
type=float,
default=7e-4,
help='learning rate (default: 7e-4)')
parser.add_argument('--meta-epochs',
type=int,
default=30,
help='meta epochs')
parser.add_argument('--eps',
type=float,
default=1e-5,
help='RMSprop optimizer epsilon (default: 1e-5)')
parser.add_argument('--alpha',
type=float,
default=0.99,
help='RMSprop optimizer apha (default: 0.99)')
parser.add_argument('--gamma',
type=float,
default=0.99,
help='discount factor for rewards (default: 0.99)')
parser.add_argument('--use-gae',
action='store_true',
default=False,
help='use generalized advantage estimation')
parser.add_argument('--gae-lambda',
type=float,
default=0.95,
help='gae lambda parameter (default: 0.95)')
parser.add_argument('--entropy-coef',
type=float,
default=0.01,
help='entropy term coefficient (default: 0.01)')
parser.add_argument('--value-loss-coef',
type=float,
default=0.5,
help='value loss coefficient (default: 0.5)')
parser.add_argument('--max-grad-norm',
type=float,
default=0.5,
help='max norm of gradients (default: 0.5)')
parser.add_argument('--pretrained', action="store_true")
parser.add_argument('--name', type=str, default="")
parser.add_argument('--data', type=str, default="")
parser.add_argument('--num-steps', type=int, default=3)
parser.add_argument('--val-percent', type=float, default=0.0)
args = parser.parse_args()
task_name = "{}_da{}_ep{}_bs{}_{}".format(args.optimizer, args.dataset,
args.epochs, args.batch_size,
args.name)
writer = tensorboardX.SummaryWriter(os.path.join(args.log_dir, task_name))
data_transforms = {
'train':
transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
]),
'val':
transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
}
if args.dataset == 'CIFAR10':
train_dataset = torchvision.datasets.CIFAR10(
'./cifar', transform=data_transforms['train'])
val_dataset = torchvision.datasets.CIFAR10(
'./cifar', transform=data_transforms['val'])
elif args.dataset == 'CIFAR100':
train_dataset = torchvision.datasets.CIFAR10(
'./cifar-100', transform=data_transforms['train'], download=True)
val_dataset = torchvision.datasets.CIFAR10(
'./cifar-100', transform=data_transforms['val'], download=True)
elif args.dataset == 'tiny':
data_transforms = {
'train':
transforms.Compose([
transforms.RandomSizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225]),
]),
'val':
transforms.Compose([
transforms.Scale(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225]),
])
}
train_dataset = torchvision.datasets.ImageFolder(
'./tiny-imagenet-200/train', transform=data_transforms['train'])
val_dataset = torchvision.datasets.ImageFolder(
'./tiny-imagenet-200/val', transform=data_transforms['val'])
elif args.dataset == 'CUB':
train_transforms, val_transforms, evaluate_transforms = preprocess_strategy(
'CUB')
traindir = os.path.join(args.data, 'train')
valdir = os.path.join(args.data, 'val')
train_dataset = torchvision.datasets.ImageFolder(
traindir, train_transforms)
val_dataset = torchvision.datasets.ImageFolder(valdir, val_transforms)
train_loader = torch.utils.data.DataLoader(train_dataset,
args.batch_size,
num_workers=args.worker,
shuffle=True)
val_loader = torch.utils.data.DataLoader(val_dataset,
args.batch_size,
num_workers=args.worker,
shuffle=False)
#model = SimpleModel()
model = resnet18(pretrained=args.pretrained)
model.fc = nn.Linear(512, args.num_classes)
if args.optimizer == 'mixture':
action_space = np.array([0, 1, 2])
coord_size = len(model.layers())
ob_name_lstm = ["loss", "val_loss", "step"]
ob_name_scalar = []
num_steps = args.num_steps
obs_shape = (len(ob_name_lstm) + len(ob_name_scalar) + coord_size, )
_hidden_size = 20
hidden_size = _hidden_size * len(ob_name_lstm)
actor_critic = Policy(coord_size, input_size=(len(ob_name_lstm), len(ob_name_scalar)), \
action_space=len(action_space), hidden_size=_hidden_size, window_size=1)
agent = algo.A2C_ACKTR(actor_critic,
args.value_loss_coef,
args.entropy_coef,
lr=args.lr_meta,
eps=args.eps,
alpha=args.alpha,
max_grad_norm=args.max_grad_norm)
meta_info = {'coord_size': coord_size, 'action_space': action_space, 'ob_name_lstm': ob_name_lstm, \
'ob_name_scalar': ob_name_lstm, 'obs_shape': obs_shape, 'hidden_size': hidden_size, \
'actor_critic': actor_critic}
if args.optimizer == 'mixture':
rollouts = RolloutStorage(
num_steps,
obs_shape,
action_shape=coord_size,
hidden_size=hidden_size,
num_recurrent_layers=actor_critic.net.num_recurrent_layers)
names = list(map(lambda x: x[0], list(model.named_parameters())))
optimizer = MixtureOptimizer(model.parameters(),
0.001,
writer=writer,
layers=model.layers(),
names=names)
elif args.optimizer == 'adam':
optimizer = torch.optim.Adam(model.parameters(), 0.001)
elif args.optimizer == 'sgd':
optimizer = torch.optim.SGD(model.parameters(), 0.001)
else:
raise NotImplementedError
if len(args.gpu) == 0:
use_cuda = False
else:
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
use_cuda = True
runner_config = {
'USE_CUDA': use_cuda,
'writer': writer,
'epochs': args.meta_epochs,
'val_percent': args.val_percent,
'num_steps': args.num_steps,
'use_gae': True,
'savepath': 'models/' + task_name
}
trainer_config = {
'train_loader': train_loader,
'val_loader': val_loader,
'USE_CUDA': use_cuda,
'writer': writer,
'use_log_loss': args.use_log_loss,
'print_freq': 5,
'epochs': args.epochs
}
if args.optimizer == 'mixture':
trainer = MetaTrainer(model, nn.CrossEntropyLoss(), optimizer,
**trainer_config)
runner = MetaRunner(trainer, rollouts, agent, actor_critic,
**runner_config)
runner.run()
else:
trainer = Trainer(model, nn.CrossEntropyLoss(), optimizer,
**trainer_config)
runner = Runner(trainer, **runner_config)
runner.run()
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 __init__(self,
env_def,
processes=1,
dir='.',
version=0,
lr=2e-4,
architecture='base',
dropout=0,
reconstruct=None,
r_weight=.05):
self.env_def = env_def
self.num_processes = processes #cpu processes
self.lr = lr
self.version = version
self.save_dir = dir + '/trained_models/'
#Setup
pathlib.Path(self.save_dir).mkdir(parents=True, exist_ok=True)
if (self.num_mini_batch > processes):
self.num_mini_batch = processes
self.writer = SummaryWriter()
self.total_steps = 0
#State
torch.manual_seed(self.seed)
torch.cuda.manual_seed_all(self.seed)
if not self.no_cuda and torch.cuda.is_available(
) and self.cuda_deterministic:
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
utils.cleanup_log_dir(self.log_dir)
utils.cleanup_log_dir(self.eval_log_dir)
torch.set_num_threads(1)
self.level_path = None
self.envs = None
self.num_envs = -1
self.set_envs(num_envs=1)
if (version > 0):
self.actor_critic = self.load(path, version)
else:
self.actor_critic = Policy(
self.envs.observation_space.shape,
self.envs.action_space,
base_kwargs={
'recurrent': self.recurrent_policy,
'shapes': list(reversed(self.env_def.model_shape)),
'dropout': dropout
},
model=architecture)
self.actor_critic.to(self.device)
#Reconstruction
self.reconstruct = reconstruct is not None
if (self.reconstruct):
#layers = self.envs.observation_space.shape[0]
#shapes = list(self.env_def.model_shape)
#self.r_model = Decoder(layers, shapes=shapes).to(self.device)
reconstruct.to(self.device)
self.r_model = lambda x: reconstruct.adapter(reconstruct(x))
#self.r_model = lambda x: reconstruct.adapter(reconstruct(x)).clamp(min=1e-6).log()
#self.r_loss = nn.L1Loss() #nn.NLLLoss() #nn.MSELoss()
self.r_loss = lambda pred, true: -r_weight * (true * torch.log(
pred.clamp(min=1e-7, max=1 - 1e-7))).sum(dim=1).mean()
self.r_optimizer = reconstruct.optimizer #optim.Adam(reconstruct.parameters(), lr = .0001)
if self.algo == 'a2c':
self.agent = A2C_ACKTR(self.actor_critic,
self.value_loss_coef,
self.entropy_coef,
lr=self.lr,
eps=self.eps,
alpha=self.alpha,
max_grad_norm=self.max_grad_norm)
elif self.algo == 'ppo':
self.agent = PPO(self.actor_critic,
self.clip_param,
self.ppo_epoch,
self.num_mini_batch,
self.value_loss_coef,
self.entropy_coef,
lr=self.lr,
eps=self.eps,
max_grad_norm=self.max_grad_norm,
use_clipped_value_loss=False)
elif self.algo == 'acktr':
self.agent = algo.A2C_ACKTR(self.actor_critic,
self.value_loss_coef,
self.entropy_coef,
acktr=True)
self.gail = False
self.gail_experts_dir = './gail_experts'
if self.gail:
assert len(self.envs.observation_space.shape) == 1
self.gail_discr = gail.Discriminator(
self.envs.observation_space.shape[0] +
self.envs.action_space.shape[0], 100, self.device)
file_name = os.path.join(
self.gail_experts_dir,
"trajs_{}.pt".format(env_name.split('-')[0].lower()))
self.gail_train_loader = torch.utils.data.DataLoader(
gail.ExpertDataset(file_name,
num_trajectories=4,
subsample_frequency=20),
batch_size=self.gail_batch_size,
shuffle=True,
drop_last=True)
self.rollouts = RolloutStorage(
self.num_steps, self.num_processes,
self.envs.observation_space.shape, self.envs.action_space,
self.actor_critic.recurrent_hidden_state_size)
def main():
torch.set_num_threads(1)
device = torch.device("cuda:0" if args.cuda else "cpu")
if args.vis:
from visdom import Visdom
viz = Visdom(port=args.port)
win = None
envs = make_vec_envs(args.env_name, args.seed, args.num_processes,
args.gamma, args.log_dir, args.add_timestep, device,
True)
frame_skip = 4 # frame skip
if args.tb_dir[-1] != '/':
args.tb_dir = args.tb_dir + '/'
logger = Logger(args.tb_dir)
logger.write_settings(args)
if args.use_tdm:
# beta scheduler
if args.beta_schedule == 'const':
beta_func = lambda x: float(args.beta_int)
elif args.beta_schedule == 'sqrt':
beta_func = lambda x: 1. / np.sqrt(x + 2)
elif args.beta_schedule == 'log':
beta_func = lambda x: 1. / np.log(x + 2)
elif args.beta_schedule == 'linear':
beta_func = lambda x: 1. / (x + 2)
# bonus function variations
if args.bonus_func == 'linear':
bonus_func = lambda x: x + 1
elif args.bonus_func == 'square':
bonus_func = lambda x: (x + 1)**2
elif args.bonus_func == 'sqrt':
bonus_func = lambda x: (x + 1)**(1 / 2)
elif args.bonus_func == 'log':
bonus_func = lambda x: np.log(x + 1)
# temporal difference module
tdm = TemporalDifferenceModule(
inputSize=2 * int(envs.observation_space.shape[0]),
outputSize=args.time_intervals,
num_fc_layers=int(args.num_layers),
depth_fc_layers=int(args.fc_width),
lr=float(args.opt_lr),
buffer_max_length=args.buffer_max_length,
buffer_RL_ratio=args.buffer_RL_ratio,
frame_skip=frame_skip,
tdm_epoch=args.tdm_epoch,
tdm_batchsize=args.tdm_batchsize,
logger=logger,
bonus_func=bonus_func).to(device)
#collect random trajectories
sample_collector = CollectSamples(envs,
args.num_processes,
initial=True)
tdm.buffer_rand = sample_collector.collect_trajectories(
args.num_rollouts, args.steps_per_rollout)
# initial training
tdm.update()
actor_critic = Policy(envs.observation_space.shape,
envs.action_space,
base_kwargs={'recurrent': args.recurrent_policy})
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)
rollouts = RolloutStorage(args.num_steps, args.num_processes,
envs.observation_space.shape, envs.action_space,
actor_critic.recurrent_hidden_state_size)
obs = envs.reset()
rollouts.obs[0].copy_(obs)
rollouts.to(device)
episode_rewards = deque(maxlen=10)
start = time.time()
for j in range(num_updates):
if args.use_linear_lr_decay:
# decrease learning rate linearly
if args.algo == "acktr":
# use optimizer's learning rate since it's hard-coded in kfac.py
update_linear_schedule(agent.optimizer, j, num_updates,
agent.optimizer.lr)
else:
update_linear_schedule(agent.optimizer, j, num_updates,
args.lr)
if args.algo == 'ppo' and args.use_linear_clip_decay:
agent.clip_param = args.clip_param * (1 - j / float(num_updates))
# acting
for step in range(args.num_steps):
# Sample actions
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])
# Obser reward and next obs
# envs.render()
obs_old = obs.clone()
obs, reward, done, infos = envs.step(action)
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])
#compute intrinsic bonus
if args.use_tdm:
tdm.symm_eval = True if step == args.num_steps - 1 else False
reward_int = tdm.compute_bonus(obs_old, obs).float()
reward += beta_func(
step + j * args.num_steps) * reward_int.cpu().unsqueeze(1)
if (j % args.log_interval == 0) and (step
== args.num_steps - 1):
logger.add_reward_intrinsic(reward_int,
(j + 1) * args.num_steps *
args.num_processes)
#saving to buffer.
rollouts.insert(obs, recurrent_hidden_states, action,
action_log_prob, value, reward, masks)
# saving to buffer and periodic updating parameters
if (args.use_tdm):
tdm.buffer_RL_temp.append((rollouts.obs, rollouts.masks))
if (j % args.num_steps == 0 and j > 0):
tdm.update()
with torch.no_grad():
next_value = actor_critic.get_value(
rollouts.obs[-1], rollouts.recurrent_hidden_states[-1],
rollouts.masks[-1]).detach()
rollouts.compute_returns(next_value, args.use_gae, args.gamma,
args.tau)
value_loss, action_loss, dist_entropy = agent.update(rollouts)
rollouts.after_update()
# save for every interval-th episode or for the last epoch
# no
# save every 1-million steps
if (((j + 1) * args.num_steps * args.num_processes) % 1e6 == 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
# A really ugly way to save a model to CPU
save_model = actor_critic
if args.cuda:
save_model = copy.deepcopy(actor_critic).cpu()
save_model = [
save_model,
getattr(get_vec_normalize(envs), 'ob_rms', None)
]
if j == num_updates - 1:
save_here = os.path.join(
save_path, args.env_name + "_step_{}M.pt".format(
(j + 1) * args.num_steps * args.num_processes // 1e6))
else:
save_here = os.path.join(save_path,
args.env_name + "_final.pt")
torch.save(save_model, save_here) # saved policy.
total_num_steps = (j + 1) * args.num_processes * args.num_steps
# printing outputs
if j % args.log_interval == 0 and len(episode_rewards) > 1:
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))
logger.add_reward(episode_rewards,
(j + 1) * args.num_steps * args.num_processes)
#
# if j % args.tb_interval == 0:
# # mean/std or median/1stqt?
# logger.add_tdm_loss(loss, self.epoch_count*i)
# evaluation process
# if (args.eval_interval is not None
# and len(episode_rewards) > 1
# and j % args.eval_interval == 0):
# eval_envs = make_vec_envs(
# args.env_name, args.seed + args.num_processes, args.num_processes,
# args.gamma, eval_log_dir, args.add_timestep, device, True)
#
# vec_norm = get_vec_normalize(eval_envs)
# if vec_norm is not None:
# vec_norm.eval()
# vec_norm.ob_rms = get_vec_normalize(envs).ob_rms
#
# eval_episode_rewards = []
#
# obs = eval_envs.reset()
# eval_recurrent_hidden_states = torch.zeros(args.num_processes,
# actor_critic.recurrent_hidden_state_size, device=device)
# eval_masks = torch.zeros(args.num_processes, 1, device=device)
#
# while len(eval_episode_rewards) < 10:
# with torch.no_grad():
# _, action, _, eval_recurrent_hidden_states = actor_critic.act(
# obs, eval_recurrent_hidden_states, eval_masks, deterministic=True)
#
# # Obser reward and next obs
# # envs.render()
# obs, reward, done, infos = eval_envs.step(action)
#
# eval_masks = torch.FloatTensor([[0.0] if done_ else [1.0]
# for done_ in done])
# for info in infos:
# if 'episode' in info.keys():
# eval_episode_rewards.append(info['episode']['r'])
#
# eval_envs.close()
#
# print(" Evaluation using {} episodes: mean reward {:.5f}\n".
# format(len(eval_episode_rewards),
# np.mean(eval_episode_rewards)))
# # plotting
# if args.vis and j % args.vis_interval == 0:
# try:
# # Sometimes monitor doesn't properly flush the outputs
# win = visdom_plot(viz, win, args.log_dir, args.env_name,
# args.algo, args.num_env_steps)
# except IOError:
# pass
#if done save:::::::::::
logger.save()
def main():
args = get_args()
trace_size = args.trace_size
toke = tokenizer()
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")
envs = make_vec_envs(args.env_name, args.seed, args.num_processes,
args.gamma, args.log_dir, device, False)
actor_critic = Policy(envs.observation_space.shape,
envs.action_space,
base_kwargs={'recurrent': args.recurrent_policy})
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)
rollouts = RolloutStorage(args.num_steps, args.num_processes,
envs.observation_space.shape, envs.action_space,
actor_critic.recurrent_hidden_state_size)
obs = envs.reset()
tobs = torch.zeros((args.num_processes, trace_size), dtype=torch.long)
#print (tobs.dtype)
rollouts.obs[0].copy_(obs)
rollouts.tobs[0].copy_(tobs)
rollouts.to(device)
episode_rewards = deque(maxlen=10)
start = time.time()
num_updates = int(
args.num_env_steps) // args.num_steps // args.num_processes
for j in range(num_updates):
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(
rollouts.obs[step], rollouts.tobs[step],
rollouts.recurrent_hidden_states[step],
rollouts.masks[step])
# Obser reward and next obs
obs, reward, done, infos = envs.step(action)
tobs = []
envs.render()
for info in infos:
if 'episode' in info.keys():
#print ("episode ", info['episode'])
episode_rewards.append(info['episode']['r'])
trace = info['trace'][0:trace_size]
trace = [x[2] for x in trace]
word_to_ix = toke.tokenize(trace)
seq = prepare_sequence(trace, word_to_ix)
if len(seq) < trace_size:
seq = torch.zeros((trace_size), dtype=torch.long)
seq = seq[:trace_size]
#print (seq.dtype)
tobs.append(seq)
tobs = torch.stack(tobs)
#print (tobs)
#print (tobs.size())
# 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, tobs, recurrent_hidden_states, action,
action_log_prob, value, reward, masks, bad_masks)
with torch.no_grad():
next_value = actor_critic.get_value(
rollouts.obs[-1], rollouts.tobs[-1],
rollouts.recurrent_hidden_states[-1],
rollouts.masks[-1]).detach()
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()
# 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), 'ob_rms', None)
], os.path.join(save_path, args.env_name + ".pt"))
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))
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 get_agent(agent_name,
args,
obs_space,
input_structure,
act_space,
save_dir,
n_ref=0,
is_ref=False):
from a2c_ppo_acktr import algo
from a2c_ppo_acktr.model import Policy, AttentionBase, LinearBase
if args.use_attention:
actor_critic = Policy(obs_space.shape,
act_space,
AttentionBase,
base_kwargs={
'recurrent': args.recurrent_policy,
'input_structure': input_structure
})
elif args.use_linear:
actor_critic = Policy(obs_space.shape, act_space, LinearBase)
else:
# if not is_ref:
# print("A")
actor_critic = Policy(obs_space.shape,
act_space,
action_activation=args.action_activation,
base_kwargs={
'recurrent': args.recurrent_policy,
'critic_dim': n_ref * 2 + 1,
'is_ref': is_ref,
'predict_reward': args.use_reward_predictor,
'timestep_mask': args.use_timestep_mask,
"rnd": args.use_rnd,
'hidden_size': args.hidden_size,
'activation': args.activation
})
# if not is_ref:
# print("B")
# print("actor critic got")
# if not is_ref:
# print("!!@!@")
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(agent_name,
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,
clip_grad_norm=not args.no_grad_norm_clip,
task=args.task,
direction=args.direction,
save_dir=save_dir,
args=args,
is_ref=is_ref)
elif args.algo == 'acktr':
agent = algo.A2C_ACKTR(actor_critic,
args.value_loss_coef,
args.entropy_coef,
acktr=True)
elif args.algo == 'loaded-dice':
agent = algo.LoadedDiCE(actor_critic,
args.dice_epoch,
args.num_mini_batch,
args.value_loss_coef,
args.entropy_coef,
args.dice_lambda,
args.episode_steps,
args.dice_task,
lr=args.lr,
eps=args.eps,
save_dir=save_dir)
elif args.algo == 'hessian':
agent = algo.Hessian(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,
clip_grad_norm=not args.no_grad_norm_clip,
task=args.task,
direction=args.direction,
args=args)
else:
raise ValueError("algo {} not supported".format(args.algo))
return actor_critic, agent
obs_shape,
envs.action_space,
base_kwargs={"recurrent": args.recurrent_policy},
navi=args.navi,
base=base,
hidden_size=args.hidden_size,
n_layers=args.n_layers,
)
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,
def main():
tb_path = os.path.join(os.path.expanduser(args.log_dir), "tensorboard_log")
makedir_if_not_exists(tb_path)
writer = SummaryWriter(tb_path)
torch.set_num_threads(1)
device = torch.device("cuda:0" if args.cuda else "cpu")
# p = multiprocessing.Process(target=_tb_task,args=(tb_path,5013) ,daemon=True)
# p.start()
if args.start_tb:
_tb_task(tb_path, port=5013)
if args.vis:
from visdom import Visdom
viz = Visdom(port=args.port)
win = None
envs = make_vec_envs(args.env_name, args.seed, args.num_processes,
args.gamma, args.log_dir, args.add_timestep, device,
False)
actor_critic = Policy(envs.observation_space.shape,
envs.action_space,
base_kwargs={'recurrent': args.recurrent_policy})
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)
rollouts = RolloutStorage(args.num_steps, args.num_processes,
envs.observation_space.shape, envs.action_space,
actor_critic.recurrent_hidden_state_size)
obs = envs.reset()
rollouts.obs[0].copy_(obs)
rollouts.to(device)
episode_rewards = deque(maxlen=10)
start = time.time()
num_eps = 0 # num training eps
num_steps = 0 # num training eps
for j in range(num_updates):
# list of all values all eps in num updates
num_steps_basline_info = defaultdict(list)
if args.use_linear_lr_decay:
# decrease learning rate linearly
if args.algo == "acktr":
# use optimizer's learning rate since it's hard-coded in kfac.py
update_linear_schedule(agent.optimizer, j, num_updates,
agent.optimizer.lr)
else:
update_linear_schedule(agent.optimizer, j, num_updates,
args.lr)
if args.algo == 'ppo' and args.use_linear_clip_decay:
agent.clip_param = args.clip_param * (1 - j / float(num_updates))
env_basline_info = defaultdict(list)
for step in range(args.num_steps):
# Sample actions
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])
# Obser reward and next obs
obs, reward, done, infos = envs.step(action)
for info in infos:
# episode is done
# add addisiotnal baseline rw info in infos:
if 'basline_rw_mse' in info:
env_basline_info['rw_mse'].append(info['basline_rw_mse'])
env_basline_info['rw_rec'].append(info['basline_rw_rec'])
if 'basline_rw_tcn' in info:
env_basline_info['rw_tcn'].append(info['basline_rw_tcn'])
if 'episode' in info.keys():
# end of episode
episode_rewards.append(info['episode']['r'])
num_steps_basline_info['len_episode'].append(
info['episode']['l'])
# distance of the pushed block
num_steps_basline_info['push_distance'].append(
info['basline_rw_push_dist'])
# take mean over eps
for k, step_vals in env_basline_info.items():
num_steps_basline_info[k].append(np.sum(step_vals))
# add baseline infos
num_eps += 1
env_basline_info = defaultdict(list)
# If done then clean the history of observations.
masks = torch.FloatTensor([[0.0] if done_ else [1.0]
for done_ in done])
rollouts.insert(obs, recurrent_hidden_states, action,
action_log_prob, value, reward, masks)
with torch.no_grad():
next_value = actor_critic.get_value(
rollouts.obs[-1], rollouts.recurrent_hidden_states[-1],
rollouts.masks[-1]).detach()
rollouts.compute_returns(next_value, args.use_gae, args.gamma,
args.tau)
value_loss, action_loss, dist_entropy = agent.update(rollouts)
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
# A really ugly way to save a model to CPU
save_model = actor_critic
if args.cuda:
save_model = copy.deepcopy(actor_critic).cpu()
save_model = [
save_model,
getattr(get_vec_normalize(envs), 'ob_rms', None)
]
torch.save(save_model,
os.path.join(save_path, args.env_name + ".pt"))
total_num_steps = (j + 1) * args.num_processes * args.num_steps
# write baseline finfos for tcn
writer_step = total_num_steps
for k, vals_step_eps in num_steps_basline_info.items():
writer.add_scalar('basline/' + k, np.mean(vals_step_eps),
writer_step)
writer.add_scalar('basline/episodes', num_eps, writer_step)
len_eps = np.mean(num_steps_basline_info['len_episode'])
if j % args.log_interval == 0 and len(episode_rewards) > 1:
end = time.time()
log.info(
"Updates {}, num timesteps {}, FPS {} Last {} training episodes: mean/median reward {:.1f}/{:.1f}, min/max reward {:.1f}/{:.1f}, len eps {}"
.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), int(len_eps), dist_entropy,
value_loss, action_loss))
if j == num_updates or (args.eval_interval is not None
and len(episode_rewards) > 1
and j % args.eval_interval == 0):
vid_log_dir = os.getenv('TCN_ENV_VID_LOG_FOLDER',
'/tmp/env_tcn/train_vid')
vid_log_inter = os.getenv('TCN_ENV_VID_LOG_INTERVAL',
train_vid_log_iter)
os.environ[
'TCN_ENV_VID_LOG_FOLDER'] = "eval_vid" # os.path.join(vid_log_dir,"../eval_vid/","interval_"+str(j))
os.environ['TCN_ENV_VID_LOG_INTERVAL'] = '1'
os.environ['TCN_ENV_EVAL_EPISODE'] = '1'
with redirect_stdout(open(os.devnull, "w")): # no stdout
with suppress_logging():
# eval envs
eval_envs = make_vec_envs(args.env_name,
args.seed + args.num_processes,
1, args.gamma, eval_log_dir,
args.add_timestep, device, True)
vec_norm = get_vec_normalize(eval_envs)
if vec_norm is not None:
vec_norm.eval()
vec_norm.ob_rms = get_vec_normalize(envs).ob_rms
eval_episode_rewards = []
obs = eval_envs.reset()
eval_recurrent_hidden_states = torch.zeros(
args.num_processes,
actor_critic.recurrent_hidden_state_size,
device=device)
eval_masks = torch.zeros(args.num_processes,
1,
device=device)
while len(eval_episode_rewards) < 1:
with torch.no_grad():
_, action, _, eval_recurrent_hidden_states = actor_critic.act(
obs,
eval_recurrent_hidden_states,
eval_masks,
deterministic=True)
# Obser reward and next obs
obs, reward, done, infos = eval_envs.step(action)
eval_masks = torch.tensor([[0.0] if done_ else [1.0]
for done_ in done],
dtype=torch.float32,
device=device)
for info in infos:
if 'episode' in info.keys():
eval_episode_rewards.append(
info['episode']['r'])
eval_envs.close()
os.environ['TCN_ENV_VID_LOG_FOLDER'] = vid_log_dir
os.environ['TCN_ENV_EVAL_EPISODE'] = '0'
os.environ['TCN_ENV_VID_LOG_INTERVAL'] = vid_log_inter
writer.add_scalar('eval/rw', np.mean(eval_episode_rewards), j)
log.info(
" Evaluation using {} episodes: mean reward {:.5f}\n".format(
len(eval_episode_rewards), np.mean(eval_episode_rewards)))
if j % args.vis_interval == 0:
try:
td_plot(writer, args.log_dir)
# Sometimes monitor doesn't properly flush the outputs
# win = visdom_plot(viz, win, args.log_dir, args.env_name,
# args.algo, args.num_env_steps)
except IOError:
print("plt error")
pass
def main():
args = get_args()
# Record trajectories
if args.record_trajectories:
record_trajectories()
return
print(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
# Append the model name
log_dir = os.path.expanduser(args.log_dir)
log_dir = os.path.join(log_dir, args.model_name, str(args.seed))
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")
envs = make_vec_envs(args.env_name, args.seed, args.num_processes,
args.gamma, log_dir, device, False)
# Take activation for carracing
print("Loaded env...")
activation = None
if args.env_name == 'CarRacing-v0' and args.use_activation:
activation = torch.tanh
print(activation)
actor_critic = Policy(envs.observation_space.shape,
envs.action_space,
base_kwargs={
'recurrent': args.recurrent_policy,
'env': args.env_name
},
activation=activation)
actor_critic.to(device)
# Load from previous model
if args.load_model_name:
state = torch.load(
os.path.join(args.save_dir, args.load_model_name,
args.load_model_name + '_{}.pt'.format(args.seed)))[0]
try:
actor_critic.load_state_dict(state)
except:
actor_critic = state
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:
if len(envs.observation_space.shape) == 1:
discr = gail.Discriminator(
envs.observation_space.shape[0] + envs.action_space.shape[0],
100, device)
file_name = os.path.join(
args.gail_experts_dir,
"trajs_{}.pt".format(args.env_name.split('-')[0].lower()))
expert_dataset = gail.ExpertDataset(file_name,
num_trajectories=3,
subsample_frequency=1)
expert_dataset_test = gail.ExpertDataset(file_name,
num_trajectories=1,
start=3,
subsample_frequency=1)
drop_last = len(expert_dataset) > args.gail_batch_size
gail_train_loader = torch.utils.data.DataLoader(
dataset=expert_dataset,
batch_size=args.gail_batch_size,
shuffle=True,
drop_last=drop_last)
gail_test_loader = torch.utils.data.DataLoader(
dataset=expert_dataset_test,
batch_size=args.gail_batch_size,
shuffle=False,
drop_last=False)
print(len(expert_dataset), len(expert_dataset_test))
else:
# env observation shape is 3 => its an image
assert len(envs.observation_space.shape) == 3
discr = gail.CNNDiscriminator(envs.observation_space.shape,
envs.action_space, 100, device)
file_name = os.path.join(args.gail_experts_dir, 'expert_data.pkl')
expert_dataset = gail.ExpertImageDataset(file_name, train=True)
test_dataset = gail.ExpertImageDataset(file_name, train=False)
gail_train_loader = torch.utils.data.DataLoader(
dataset=expert_dataset,
batch_size=args.gail_batch_size,
shuffle=True,
drop_last=len(expert_dataset) > args.gail_batch_size,
)
gail_test_loader = torch.utils.data.DataLoader(
dataset=test_dataset,
batch_size=args.gail_batch_size,
shuffle=False,
drop_last=len(test_dataset) > args.gail_batch_size,
)
print('Dataloader size', len(gail_train_loader))
rollouts = RolloutStorage(args.num_steps, args.num_processes,
envs.observation_space.shape, envs.action_space,
actor_critic.recurrent_hidden_state_size)
obs = envs.reset()
rollouts.obs[0].copy_(obs)
rollouts.to(device)
episode_rewards = deque(maxlen=10)
start = time.time()
#num_updates = int(
#args.num_env_steps) // args.num_steps // args.num_processes
num_updates = args.num_steps
print(num_updates)
# count the number of times validation loss increases
val_loss_increase = 0
prev_val_action = np.inf
best_val_loss = np.inf
for j in range(num_updates):
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(
rollouts.obs[step], rollouts.recurrent_hidden_states[step],
rollouts.masks[step])
# Observe 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'])
# 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)
with torch.no_grad():
next_value = actor_critic.get_value(
rollouts.obs[-1], rollouts.recurrent_hidden_states[-1],
rollouts.masks[-1]).detach()
if args.gail:
if j >= 10:
try:
envs.venv.eval()
except:
pass
gail_epoch = args.gail_epoch
#if j < 10:
#gail_epoch = 100 # Warm up
for _ in range(gail_epoch):
#discr.update(gail_train_loader, rollouts,
#None)
pass
for step in range(args.num_steps):
rollouts.rewards[step] = discr.predict_reward(
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)
value_loss = 0
dist_entropy = 0
for data in gail_train_loader:
expert_states, expert_actions = data
expert_states = Variable(expert_states).to(device)
expert_actions = Variable(expert_actions).to(device)
loss = agent.update_bc(expert_states, expert_actions)
action_loss = loss.data.cpu().numpy()
print("Epoch: {}, Loss: {}".format(j, action_loss))
with torch.no_grad():
cnt = 0
val_action_loss = 0
for data in gail_test_loader:
expert_states, expert_actions = data
expert_states = Variable(expert_states).to(device)
expert_actions = Variable(expert_actions).to(device)
loss = agent.get_action_loss(expert_states, expert_actions)
val_action_loss += loss.data.cpu().numpy()
cnt += 1
val_action_loss /= cnt
print("Val Loss: {}".format(val_action_loss))
#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 != "":
if val_action_loss < best_val_loss:
val_loss_increase = 0
best_val_loss = val_action_loss
save_path = os.path.join(args.save_dir, args.model_name)
try:
os.makedirs(save_path)
except OSError:
pass
torch.save([
actor_critic.state_dict(),
getattr(utils.get_vec_normalize(envs), 'ob_rms', None),
getattr(utils.get_vec_normalize(envs), 'ret_rms', None)
],
os.path.join(
save_path,
args.model_name + "_{}.pt".format(args.seed)))
elif val_action_loss > prev_val_action:
val_loss_increase += 1
if val_loss_increase == 10:
print("Val loss increasing too much, breaking here...")
break
elif val_action_loss < prev_val_action:
val_loss_increase = 0
# Update prev val action
prev_val_action = val_action_loss
# log interval
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))
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 __init__(self,
learner,
envs,
maxat,
maxupd,
targ_policy,
args,
device="cpu",
hidden_sizes=(64, 64),
activation=nn.Tanh,
rand_select=False):
super(TargAttacker, self).__init__()
self.args = args
self.targ_policy = targ_policy
self.learner = learner
self.gamma = args.gamma
self.device = device
self.radius = args.radius
self.frac = args.frac
self.stepsize = args.stepsize
self.maxiter = args.maxiter
self.maxat = maxat
self.maxupd = maxupd
self.delta = args.delta
self.dist_thres = args.dist_thres
self.rand_select = rand_select
self.disc_action = isinstance(envs.action_space, Discrete)
if self.disc_action:
self.action_dim = envs.action_space.n
attack_policy = Policy(
envs.observation_space.shape,
envs.action_space,
base_kwargs={'recurrent': args.recurrent_policy})
attack_policy.to(device)
if isinstance(learner, algo.A2C_ACKTR):
self.im_learner = algo.A2C_ACKTR(attack_policy,
args.value_loss_coef,
args.entropy_coef,
lr=args.lr,
eps=args.eps,
alpha=args.alpha,
max_grad_norm=args.max_grad_norm,
acktr=learner.acktr)
elif isinstance(learner, algo.PPO):
self.im_learner = algo.PPO(attack_policy,
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)
self.cp_net()
self.critic = Value(envs.observation_space.shape[0], hidden_sizes,
activation).to(device)
self.critic_optim = optim.Adam(self.critic.parameters(), lr=args.lr)
self.dist_list = np.array([])
self.attack_num = 0
self.update_num = 0
self.state_buffer = None
self.state_buffer_limit = 100
def main():
if not os.path.exists("./plots"):
os.makedirs("./plots")
gbench = read_gbench('./data/gbench.txt')
args = my_get_args()
print(args)
config = dict(sigma=args.sim_sigma,
momentum=args.sim_momentum,
pump_bins=args.sim_bins,
lag=1000 // args.num_steps,
rshift=args.sim_rshift,
pump_scale=args.sim_scale,
reward_kind=args.sim_reward,
continuous=args.sim_continuous,
span=args.sim_span,
percentile=args.sim_percentile,
last_runs=args.sim_perc_len,
add_linear=not args.sim_no_linear,
start_pump=args.sim_start,
static_features=not args.sim_no_static,
extra_features=not args.sim_no_extra,
curiosity_num=args.curiosity)
base_kwargs = {
'hidden_size': args.hidden_size,
'film_size': 800 * (not args.sim_no_static)
}
if args.relu:
base_kwargs['activation'] = 'relu'
base = FILMBase #FILMBase
if args.gset > 0:
test_graphs = [args.gset]
else:
test_graphs = [1, 2, 3, 4, 5]
#---------------------------------------------------------
assert args.algo in ['a2c', 'ppo', 'acktr']
if args.recurrent_policy:
assert args.algo in ['a2c', 'ppo'
], 'Recurrent policy is not implemented for ACKTR'
num_updates = int(
args.num_env_steps) // args.num_steps // args.num_processes
print('Num updates: ', num_updates)
if args.dry_run:
return
random.seed(args.seed)
np.random.seed(args.seed)
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
torch.set_num_threads(1)
device = torch.device("cuda:0" if args.cuda else "cpu")
logdata = defaultdict(list)
if args.gset > 0:
envs = []
for g in test_graphs:
g_ = read_gset('./data/G{}.txt'.format(g), negate=True)
s = SIMCIM(g_,
device=device,
batch_size=args.num_processes,
**config)
s.runpump()
envs.append(s)
envs = SIMCollection(envs, [gbench[g] for g in test_graphs])
logdata['bls_bench'] = [gbench[g] for g in test_graphs]
else:
envs = SIMGeneratorRandom(800,
0.06,
args.num_processes,
config,
keep=args.sim_keep,
n_sims=args.sim_nsim,
device=device)
if args.snapshot is None:
actor_critic = Policy(envs.observation_space.shape,
envs.action_space,
base=base,
base_kwargs=base_kwargs)
else:
actor_critic, _ = torch.load(
os.path.join(args.save_dir, args.algo, args.snapshot + ".pt"))
actor_critic.to(device)
print(actor_critic)
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)
rollouts = RolloutStorage(args.num_steps, args.num_processes,
envs.observation_space.shape, envs.action_space,
actor_critic.recurrent_hidden_state_size)
obs = envs.reset()
print(rollouts.obs.shape, obs.shape)
rollouts.obs[0].copy_(obs)
rollouts.to(device)
eval_envs = []
for g in test_graphs:
g_ = read_gset('./data/G{}.txt'.format(g), negate=True)
s = SIMCIM(g_,
device=device,
batch_size=args.num_val_processes,
**config)
s.runpump()
eval_envs.append(s)
eval_envs = SIMCollection(eval_envs, [gbench[g] for g in test_graphs])
ref_cuts = [s.lastcuts for s in eval_envs.envs]
logdata['ref_cuts'] = [e.tolist() for e in ref_cuts]
stoch_cuts = None
start = time.time()
for j in range(num_updates):
if args.use_linear_lr_decay:
# decrease learning rate linearly
if args.algo == "acktr":
# use optimizer's learning rate since it's hard-coded in kfac.py
update_linear_schedule(agent.optimizer, j, num_updates,
agent.optimizer.lr)
else:
update_linear_schedule(agent.optimizer, j, num_updates,
args.lr)
if args.algo == 'ppo' and args.use_linear_clip_decay:
agent.clip_param = args.clip_param * (1 - j / float(num_updates))
# ROLLOUT DATA
for step in range(args.num_steps):
# Sample actions
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])
# Obser reward and next obs
obs, reward, done, infos = envs.step(action)
if 'episode' in infos[0].keys():
rw = np.mean([e['episode']['r'] for e in infos])
logdata['episode_rewards'].append(rw.item())
if args.gset > 0:
cuts = [e.lastcuts for e in envs.envs]
logdata['train_median'].append(
[np.median(e).item() for e in cuts])
logdata['train_max'].append(
[np.max(e).item() for e in cuts])
# If done then clean the history of observations.
masks = torch.FloatTensor([[0.0] if done_ else [1.0]
for done_ in done])
rollouts.insert(obs, recurrent_hidden_states, action,
action_log_prob, value, reward, masks)
#UPDATE AGENT
with torch.no_grad():
next_value = actor_critic.get_value(
rollouts.obs[-1], rollouts.recurrent_hidden_states[-1],
rollouts.masks[-1]).detach()
rollouts.compute_returns(next_value, args.use_gae, args.gamma,
args.tau)
value_loss, action_loss, _ = agent.update(rollouts)
logdata['alosses'].append(action_loss)
logdata['vlosses'].append(value_loss)
logdata['train_percentiles'].append(envs.perc.tolist())
rollouts.after_update()
#CHECKPOINTS
# 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
# A really ugly way to save a model to CPU
save_model = actor_critic
if args.cuda:
save_model = copy.deepcopy(actor_critic).cpu()
save_model = [
save_model,
getattr(get_vec_normalize(envs), 'ob_rms', None)
]
torch.save(
save_model,
os.path.join(save_path, args.env_name + '-' + str(j) + ".pt"))
total_num_steps = (j + 1) * args.num_processes * args.num_steps
#LOGGING
if j % args.log_interval == 0 and len(logdata['episode_rewards']) > 1:
end = time.time()
print(
"Updates {}, num timesteps {}, FPS {} \n Last {} training episodes: \
mean/median reward {:.3f}/{:.3f}, min/max reward {:.3f}/{:.3f}\n"
.format(j, total_num_steps,
int(total_num_steps / (end - start)),
len(logdata['episode_rewards']),
np.mean(logdata['episode_rewards'][-10:]),
np.median(logdata['episode_rewards'][-10:]),
np.min(logdata['episode_rewards'][-10:]),
np.max(logdata['episode_rewards'][-10:])))
#EVALUATION
if (args.eval_interval is not None and j % args.eval_interval == 0):
logdata['spumps'] = []
vec_norm = get_vec_normalize(eval_envs)
if vec_norm is not None:
vec_norm.eval()
vec_norm.ob_rms = get_vec_normalize(envs).ob_rms
obs = eval_envs.reset()
eval_recurrent_hidden_states = torch.zeros(
args.num_val_processes,
actor_critic.recurrent_hidden_state_size,
device=device)
eval_masks = torch.zeros(args.num_val_processes, 1, device=device)
eval_done = False
while not eval_done:
p = eval_envs.envs[0].old_p
logdata['spumps'].append(p[:10].cpu().numpy().tolist())
with torch.no_grad():
_, action, _, eval_recurrent_hidden_states = actor_critic.act(
obs,
eval_recurrent_hidden_states,
eval_masks,
deterministic=False)
# Obser reward and next obs
obs, reward, done, infos = eval_envs.step(action)
eval_done = np.all(done)
eval_masks = torch.tensor([[0.0] if done_ else [1.0]
for done_ in done],
dtype=torch.float32,
device=device)
stoch_cuts = [e.lastcuts for e in eval_envs.envs]
logdata['stoch_cuts'] = [e.tolist() for e in stoch_cuts]
logdata['eval_median'].append(
[np.median(e).item() for e in stoch_cuts])
logdata['eval_max'].append([np.max(e).item() for e in stoch_cuts])
logdata['test_percentiles'].append(eval_envs.perc.tolist())
rw = np.mean([e['episode']['r'] for e in infos])
logdata['eval_episode_rewards'].append(rw.item())
print(" Evaluation using {} episodes: mean reward {:.5f}\n".format(
len(logdata['eval_episode_rewards']),
np.mean(logdata['eval_episode_rewards'])))
if j % args.log_interval == 0:
fn = os.path.join(save_path, args.env_name + ".res")
with open(fn, 'w') as f:
json.dump(logdata, f, sort_keys=True, indent=2)
#VISUALIZATION
if j % args.vis_interval == 0:
#if False:
plt.figure(figsize=(15, 10))
plt.subplot(231)
plt.title('Rewards')
plt.xlabel('SIM runs')
plt.plot(logdata['episode_rewards'], c='r', label='mean train')
plt.plot(np.linspace(0, len(logdata['episode_rewards']),
len(logdata['eval_episode_rewards'])),
logdata['eval_episode_rewards'],
'b',
label='mean eval')
plt.legend()
plt.subplot(232)
plt.plot(logdata['alosses'])
plt.title('Policy loss')
plt.subplot(233)
plt.plot(logdata['vlosses'])
plt.title('Value loss')
plt.subplot(234)
plt.title('Pumps')
plt.xlabel('SIM iterations / 10')
plt.plot(np.array(logdata['spumps']))
plt.ylim(-0.05, 1.1)
plt.subplot(235)
plt.plot(logdata['train_percentiles'])
plt.title('Train average percentile')
plt.subplot(236)
plt.title('Test percentiles')
plt.plot(logdata['test_percentiles'])
plt.legend([str(e) for e in test_graphs])
plt.tight_layout()
plt.savefig('./plots/agent_' + args.env_name + '.pdf')
plt.clf()
plt.close()
gc.collect()
#plt.show()
if stoch_cuts is not None:
fig, axs = plt.subplots(len(ref_cuts),
1,
sharex=False,
tight_layout=True)
if len(ref_cuts) == 1:
axs = [axs]
for gi in range(len(ref_cuts)):
mn = min(ref_cuts[gi])
axs[gi].hist(ref_cuts[gi], bins=100, alpha=0.7)
dc = stoch_cuts[gi][stoch_cuts[gi] >= mn]
if dc.size > 0:
axs[gi].hist(dc, bins=100, alpha=0.7)
plt.savefig('./plots/cuts_' + args.env_name + '.pdf')
plt.clf()
plt.close()
gc.collect()
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")
if args.env_name.startswith("lab_"):
gym_name, flow_json = make_lab_env(args.env_name)
args.env_name = gym_name
envs = make_vec_envs(args.env_name, args.seed, args.num_processes,
args.gamma, args.log_dir, device, False)
actor_critic = Policy(
envs.observation_space.shape,
envs.action_space,
base_kwargs={'recurrent': args.recurrent_policy})
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] + envs.action_space.shape[0], 100,
device)
file_name = os.path.join(
args.gail_experts_dir, "trajs_{}.pt".format(
args.env_name.split('-')[0].lower()))
expert_dataset = gail.ExpertDataset(
file_name, num_trajectories=4, subsample_frequency=20)
drop_last = len(expert_dataset) > args.gail_batch_size
gail_train_loader = torch.utils.data.DataLoader(
dataset=expert_dataset,
batch_size=args.gail_batch_size,
shuffle=True,
drop_last=drop_last)
rollouts = RolloutStorage(args.num_steps, args.num_processes,
envs.observation_space.shape, envs.action_space,
actor_critic.recurrent_hidden_state_size)
obs = envs.reset()
rollouts.obs[0].copy_(obs)
rollouts.to(device)
episode_rewards = deque(maxlen=10)
start = time.time()
num_updates = int(
args.num_env_steps) // args.num_steps // args.num_processes
for j in range(num_updates):
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(
rollouts.obs[step], rollouts.recurrent_hidden_states[step],
rollouts.masks[step])
# 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'])
# 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)
with torch.no_grad():
next_value = actor_critic.get_value(
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):
discr.update(gail_train_loader, rollouts,
utils.get_vec_normalize(envs)._obfilt)
for step in range(args.num_steps):
rollouts.rewards[step] = discr.predict_reward(
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)
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), 'ob_rms', None)
], os.path.join(save_path, args.env_name + ".pt"))
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))
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 main(args):
try:
os.makedirs(args.log_dir)
except OSError:
files = glob.glob(os.path.join(args.log_dir, '*.monitor.csv'))
for f in files:
os.remove(f)
eval_log_dir = args.log_dir + "_eval"
try:
os.makedirs(eval_log_dir)
except OSError:
files = glob.glob(os.path.join(eval_log_dir, '*.monitor.csv'))
for f in files:
os.remove(f)
assert args.algo in ['a2c', 'ppo', 'acktr']
if args.recurrent_policy:
assert args.algo in ['a2c', 'ppo'], \
'Recurrent policy is not implemented for ACKTR'
if args.eval_render:
render_env = make_vec_envs(args.env_name,
args.seed,
1,
None,
None,
args.add_timestep,
device='cpu',
allow_early_resets=False)
torch.set_num_threads(1)
num_updates = int(
args.num_env_steps) // args.num_steps // args.num_processes
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
# Uses gpu/cuda by default
device = torch.device("cuda:0" if args.cuda else "cpu")
# Only if running visdoom
if args.vis:
from visdom import Visdom
viz = Visdom(port=args.port)
win = None
envs = make_vec_envs(args.env_name, args.seed, args.num_processes,
args.gamma, args.log_dir, args.add_timestep, device,
False)
# Set up actor_critic
actor_critic = Policy(envs.observation_space.shape,
envs.action_space,
base_kwargs={'recurrent': args.recurrent_policy})
actor_critic.to(device)
# Set algorithm with actor critic and use to learn
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)
rollouts = RolloutStorage(args.num_steps, args.num_processes,
envs.observation_space.shape, envs.action_space,
actor_critic.recurrent_hidden_state_size)
obs = envs.reset()
rollouts.obs[0].copy_(obs)
rollouts.to(device)
episode_rewards = deque(maxlen=10)
start = time.time()
for j in range(num_updates):
if args.use_linear_lr_decay:
# decrease learning rate linearly
if args.algo == "acktr":
# use optimizer's learning rate since it's hard-coded in kfac.py
update_linear_schedule(agent.optimizer, j, num_updates,
agent.optimizer.lr)
else:
update_linear_schedule(agent.optimizer, j, num_updates,
args.lr)
if args.algo == 'ppo' and args.use_linear_clip_decay:
agent.clip_param = args.clip_param * (1 - j / float(num_updates))
for step in range(args.num_steps):
# Sample actions
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])
# 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'])
# If done then clean the history of observations.
masks = torch.FloatTensor([[0.0] if done_ else [1.0]
for done_ in done])
rollouts.insert(obs, recurrent_hidden_states, action,
action_log_prob, value, reward, masks)
with torch.no_grad():
next_value = actor_critic.get_value(
rollouts.obs[-1], rollouts.recurrent_hidden_states[-1],
rollouts.masks[-1]).detach()
rollouts.compute_returns(next_value, args.use_gae, args.gamma,
args.tau)
value_loss, action_loss, dist_entropy = agent.update(rollouts)
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
# A really ugly way to save a model to CPU
save_model = actor_critic
if args.cuda:
save_model = copy.deepcopy(actor_critic).cpu()
save_model = [
save_model,
getattr(get_vec_normalize(envs), 'ob_rms', None)
]
torch.save(
save_model,
os.path.join(
save_path, args.env_name + "-AvgRwrd" +
str(int(np.mean(episode_rewards))) + ".pt"))
print("Saving Model")
total_num_steps = (j + 1) * args.num_processes * args.num_steps
# Logs every log_interval steps
if j % args.log_interval == 0 and len(episode_rewards) > 1:
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))
if (args.eval_interval is not None and len(episode_rewards) > 1
and j % args.eval_interval == 0):
eval_envs = make_vec_envs(args.env_name,
args.seed + args.num_processes,
args.num_processes, args.gamma,
eval_log_dir, args.add_timestep, device,
True)
vec_norm = get_vec_normalize(eval_envs)
if vec_norm is not None:
vec_norm.eval()
vec_norm.ob_rms = get_vec_normalize(envs).ob_rms
eval_episode_rewards = []
obs = eval_envs.reset()
eval_recurrent_hidden_states = torch.zeros(
args.num_processes,
actor_critic.recurrent_hidden_state_size,
device=device)
eval_masks = torch.zeros(args.num_processes, 1, device=device)
while len(eval_episode_rewards) < 10:
with torch.no_grad():
_, action, _, eval_recurrent_hidden_states = actor_critic.act(
obs,
eval_recurrent_hidden_states,
eval_masks,
deterministic=True)
# Obser reward and next obs
obs, reward, done, infos = eval_envs.step(action)
eval_masks = torch.FloatTensor([[0.0] if done_ else [1.0]
for done_ in done])
for info in infos:
if 'episode' in info.keys():
eval_episode_rewards.append(info['episode']['r'])
if args.eval_render:
show_model(render_env, actor_critic)
eval_envs.close()
print(" Evaluation using {} episodes: mean reward {:.5f}\n".format(
len(eval_episode_rewards), np.mean(eval_episode_rewards)))
if args.vis and j % args.vis_interval == 0:
try:
# Sometimes monitor doesn't properly flush the outputs
win = visdom_plot(viz, win, args.log_dir, args.env_name,
args.algo, args.num_env_steps)
except IOError:
pass
def main():
args = get_args()
writer = SummaryWriter(os.path.join('logs', args.save_name), )
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")
envs = make_vec_envs(
basic_env.BasicFlatDiscreteEnv,
args.seed,
args.num_processes,
args.gamma,
args.log_dir,
device,
False,
task='lift',
gripper_type='RobotiqThreeFingerDexterousGripper',
robot='Panda',
controller='JOINT_TORQUE' if args.vel else 'JOINT_POSITION',
horizon=1000,
reward_shaping=True)
actor_critic = Policy(
envs.observation_space.shape,
envs.action_space,
base=Surreal,
# base=OpenAI,
# base=MLP_ATTN,
base_kwargs={
'recurrent':
args.recurrent_policy,
# 'dims': basic_env.BasicFlatEnv().modality_dims
'config':
dict(act='relu' if args.relu else 'tanh', rec=args.rec, fc=args.fc)
})
print(actor_critic)
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] + envs.action_space.shape[0], 100,
device)
file_name = os.path.join(
args.gail_experts_dir,
"trajs_{}.pt".format(args.env_name.split('-')[0].lower()))
expert_dataset = gail.ExpertDataset(file_name,
num_trajectories=4,
subsample_frequency=20)
drop_last = len(expert_dataset) > args.gail_batch_size
gail_train_loader = torch.utils.data.DataLoader(
dataset=expert_dataset,
batch_size=args.gail_batch_size,
shuffle=True,
drop_last=drop_last)
rollouts = RolloutStorage(args.num_steps, args.num_processes,
envs.observation_space.shape, envs.action_space,
actor_critic.recurrent_hidden_state_size)
obs = envs.reset()
rollouts.obs[0].copy_(obs)
rollouts.to(device)
episode_rewards = deque(maxlen=100)
start = time.time()
num_updates = int(
args.num_env_steps) // args.num_steps // args.num_processes
best_reward = 0
for j in range(num_updates):
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)
writer.add_scalar('lr', agent.optimizer.param_groups[0]['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(
rollouts.obs[step], rollouts.recurrent_hidden_states[step],
rollouts.masks[step])
# 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'])
# 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)
with torch.no_grad():
next_value = actor_critic.get_value(
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):
discr.update(gail_train_loader, rollouts,
utils.get_vec_normalize(envs)._obfilt)
for step in range(args.num_steps):
rollouts.rewards[step] = discr.predict_reward(
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)
end = time.time()
total_num_steps = (j + 1) * args.num_processes * args.num_steps
if len(episode_rewards) > 1:
writer.add_scalar('loss/value', value_loss, total_num_steps)
writer.add_scalar('loss/policy', action_loss, total_num_steps)
writer.add_scalar('experiment/num_updates', j, total_num_steps)
writer.add_scalar('experiment/FPS',
int(total_num_steps / (end - start)),
total_num_steps)
writer.add_scalar('experiment/EPISODE MEAN',
np.mean(episode_rewards), total_num_steps)
writer.add_scalar('experiment/EPISODE MEDIAN',
np.median(episode_rewards), total_num_steps)
writer.add_scalar('experiment/EPISODE MIN',
np.min(episode_rewards), total_num_steps)
writer.add_scalar('experiment/EPSIDOE MAX',
np.max(episode_rewards), total_num_steps)
rollouts.after_update()
# save for every interval-th episode or for the last epoch
if len(episode_rewards) > 1 and args.save_dir != "":
rew = np.mean(episode_rewards)
if rew > best_reward:
best_reward = rew
print('saved with best reward', rew)
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.save_name + ".pt"))
if j % args.log_interval == 0 and len(episode_rewards) > 1:
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))
if (args.eval_interval is not None and len(episode_rewards) > 1
and j % args.eval_interval == 0):
obs_rms = utils.get_vec_normalize(envs).obs_rms
evaluate(actor_critic, obs_rms, args.env_name, args.seed,
args.num_processes, eval_log_dir, device)
writer.close()
def main():
ARGUMENTS.update(vars(args))
torch.set_num_threads(1)
device = torch.device("cuda:0" if args.cuda else "cpu")
if args.vis:
from visdom import Visdom
viz = Visdom(port=args.port)
win = None
envs = make_vec_envs(args.env_name, args.seed, args.num_processes,
args.gamma, args.log_dir, args.add_timestep, device, False)
actor_critic = Policy(envs.observation_space.shape, envs.action_space,
base_kwargs={'recurrent': args.recurrent_policy})
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)
rollouts = RolloutStorage(args.num_steps, args.num_processes,
envs.observation_space.shape, envs.action_space,
actor_critic.recurrent_hidden_state_size)
obs = envs.reset()
rollouts.obs[0].copy_(obs)
rollouts.to(device)
episode_rewards = deque(maxlen=10)
start = time.time()
for j in range(num_updates):
if args.use_linear_lr_decay:
# decrease learning rate linearly
if args.algo == "acktr":
# use optimizer's learning rate since it's hard-coded in kfac.py
update_linear_schedule(agent.optimizer, j, num_updates, agent.optimizer.lr)
else:
update_linear_schedule(agent.optimizer, j, num_updates, args.lr)
if args.algo == 'ppo' and args.use_linear_lr_decay:
agent.clip_param = args.clip_param * (1 - j / float(num_updates))
for step in range(args.num_steps):
# Sample actions
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])
# 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'])
# If done then clean the history of observations.
masks = torch.FloatTensor([[0.0] if done_ else [1.0]
for done_ in done])
rollouts.insert(obs, recurrent_hidden_states, action, action_log_prob, value, reward, masks)
with torch.no_grad():
next_value = actor_critic.get_value(rollouts.obs[-1],
rollouts.recurrent_hidden_states[-1],
rollouts.masks[-1]).detach()
rollouts.compute_returns(next_value, args.use_gae, args.gamma, args.tau)
value_loss, action_loss, dist_entropy = agent.update(rollouts)
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
# A really ugly way to save a model to CPU
save_model = actor_critic
if args.cuda:
save_model = copy.deepcopy(actor_critic).cpu()
save_model = [save_model,
getattr(get_vec_normalize(envs), 'ob_rms', None)]
torch.save(save_model, os.path.join(save_path, args.env_name + ".pt"))
total_num_steps = (j + 1) * args.num_processes * args.num_steps
if j % args.log_interval == 0 and len(episode_rewards) > 1:
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))
ALL_UPDATES.append(j)
ALL_TIMESTEPS.append(total_num_steps)
ALL_FPS.append(int(total_num_steps / (end - start)))
ALL_MEAN_REWARDS.append(np.mean(episode_rewards))
ALL_MEDIAN_REWARDS.append(np.median(episode_rewards))
ALL_MIN_REWARDS.append(np.min(episode_rewards))
ALL_MAX_REWARDS.append(np.max(episode_rewards))
ALL_DIST_ENTROPY.append(dist_entropy)
ALL_VALUE_LOSS.append(value_loss)
ALL_ACTION_LOSS.append(action_loss)
if (args.eval_interval is not None
and len(episode_rewards) > 1
and j % args.eval_interval == 0):
eval_envs = make_vec_envs(
args.env_name, args.seed + args.num_processes, args.num_processes,
args.gamma, eval_log_dir, args.add_timestep, device, True)
vec_norm = get_vec_normalize(eval_envs)
if vec_norm is not None:
vec_norm.eval()
vec_norm.ob_rms = get_vec_normalize(envs).ob_rms
eval_episode_rewards = []
obs = eval_envs.reset()
eval_recurrent_hidden_states = torch.zeros(args.num_processes,
actor_critic.recurrent_hidden_state_size, device=device)
eval_masks = torch.zeros(args.num_processes, 1, device=device)
while len(eval_episode_rewards) < 10:
with torch.no_grad():
_, action, _, eval_recurrent_hidden_states = actor_critic.act(
obs, eval_recurrent_hidden_states, eval_masks, deterministic=True)
# Obser reward and next obs
obs, reward, done, infos = eval_envs.step(action)
eval_masks = torch.FloatTensor([[0.0] if done_ else [1.0]
for done_ in done])
for info in infos:
if 'episode' in info.keys():
eval_episode_rewards.append(info['episode']['r'])
eval_envs.close()
print(" Evaluation using {} episodes: mean reward {:.5f}\n".
format(len(eval_episode_rewards),
np.mean(eval_episode_rewards)))
if args.vis and j % args.vis_interval == 0:
try:
# Sometimes monitor doesn't properly flush the outputs
win = visdom_plot(viz, win, args.log_dir, args.env_name,
args.algo, args.num_env_steps)
except IOError:
pass
# Save the results
name = ARGUMENTS['env_name'] + '-' + ARGUMENTS['algo'] + '-' + ARGUMENTS['experiment'] + '-grad_noise' + str(ARGUMENTS['gradient_noise'])
experiment = ro.Experiment(name, directory='results')
data = {
'updates': ALL_UPDATES,
'timesteps': ALL_TIMESTEPS,
'fps': ALL_FPS,
'mean_rewards': ALL_MEAN_REWARDS,
'median_rewards': ALL_MEDIAN_REWARDS,
'min_rewards': ALL_MIN_REWARDS,
'max_rewards': ALL_MAX_REWARDS,
'dist_entropy': ALL_DIST_ENTROPY,
'value_loss': ALL_VALUE_LOSS,
'action_loss': ALL_ACTION_LOSS,
}
data.update(ARGUMENTS)
result = data['mean_rewards'][-1]
experiment.add_result(result, data)