def main():
torch.set_num_threads(1)
device = torch.device("cuda:0" if args.cuda else "cpu")
## Make environments
envs = make_vec_envs(args, device)
## Setup Policy / network architecture
if args.load_path != '':
if os.path.isfile(os.path.join(args.load_path, "best_model.pt")):
import_name = "best_model.pt"
else:
import_name = "model.pt"
online_actor_critic = torch.load(
os.path.join(args.load_path, import_name))
target_actor_critic = torch.load(
os.path.join(args.load_path, import_name))
if args.cuda:
target_actor_critic = target_actor_critic.cuda()
online_actor_critic = online_actor_critic.cuda()
else:
online_actor_critic = Policy(occ_obs_shape, sign_obs_shape,
args.state_rep, envs.action_space,
args.recurrent_policy)
online_actor_critic.to(device)
target_actor_critic = Policy(occ_obs_shape, sign_obs_shape,
args.state_rep, envs.action_space,
args.recurrent_policy)
target_actor_critic.to(device)
target_actor_critic.load_state_dict(online_actor_critic.state_dict())
if args.penetration_type == "constant":
target_actor_critic = online_actor_critic
## Choose algorithm to use
if args.algo == 'a2c':
agent = algo.A2C_ACKTR(online_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(online_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(online_actor_critic,
args.value_loss_coef,
args.entropy_coef,
acktr=True)
## Initiate memory buffer
rollouts = RolloutStorage(args.num_steps, args.num_processes,
occ_obs_shape, sign_obs_shape, envs.action_space,
target_actor_critic.recurrent_hidden_state_size)
## Start env with first observation
occ_obs, sign_obs = envs.reset()
if args.state_rep == 'full':
rollouts.occ_obs[0].copy_(occ_obs)
rollouts.sign_obs[0].copy_(sign_obs)
rollouts.to(device)
# Last 20 rewards - can set different queue length for different averaging
episode_rewards = deque(maxlen=args.num_steps)
reward_track = []
best_eval_rewards = 0
start = time.time()
## Loop over every policy updatetarget network
for j in range(num_updates):
## Setup parameter decays
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))
## Loop over num_steps environment updates to form trajectory
for step in range(args.num_steps):
# Sample actionspython3 main.py --algo ppo --num-steps 700000 --penetration-rate $i --env-name TrafficLight-simple-dense-v0 --lr 2.5e-4 --num-processes 8 --num-steps 128 --num-mini-batch 4 --use-linear-lr-decay --use-linear-clip-decay
with torch.no_grad():
# Pass observation through network and get outputs
value, action, action_log_prob, recurrent_hidden_states = target_actor_critic.act(
rollouts.occ_obs[step], rollouts.sign_obs[step],
rollouts.recurrent_hidden_states[step],
rollouts.masks[step])
# Do action in environment and save reward
occ_obs, sign_obs, reward, done, _ = envs.step(action)
episode_rewards.append(reward.numpy())
# Masks the processes which are done
masks = torch.FloatTensor([[0.0] if done_ else [1.0]
for done_ in done])
# Insert step information in buffer
rollouts.insert(occ_obs, sign_obs, recurrent_hidden_states, action,
action_log_prob, value, reward, masks)
## Get state value of current env state
with torch.no_grad():
next_value = target_actor_critic.get_value(
rollouts.occ_obs[-1], rollouts.sign_obs[-1],
rollouts.recurrent_hidden_states[-1],
rollouts.masks[-1]).detach()
## Computes the num_step return (next_value approximates reward after num_step) see Supp Material of https://arxiv.org/pdf/1804.02717.pdf
## Can use Generalized Advantage Estimation
rollouts.compute_returns(next_value, args.use_gae, args.gamma,
args.tau)
# Update the policy with the rollouts
value_loss, action_loss, dist_entropy = agent.update(rollouts)
# Clean the rollout by cylcing last elements to first ones
rollouts.after_update()
if (args.penetration_type == "linear") and (j % update_period == 0):
target_actor_critic.load_state_dict(
online_actor_critic.state_dict())
## Save model}, num timesteps {}, FPS {} \n Last {} training episodes: mean/median reward {:.3f}/{:.3f}, min/max reward {:.3f}/{:.3f}\n".
if (j % args.save_interval == 0
or j == num_updates - 1) and args.save_dir != "":
# A really ugly way to save a model to CPU
save_model = target_actor_critic
if args.cuda:
save_model = copy.deepcopy(target_actor_critic).cpu()
torch.save(save_model, os.path.join(save_path, "model.pt"))
total_num_steps = (j + 1) * args.num_processes * args.num_steps
if args.vis:
# Add the average reward of update to reward tracker
reward_track.append(np.mean(episode_rewards))
## Log progress
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 {:.3f}/{:.3f}, min/max reward {:.3f}/{:.3f}\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))
## Evaluate model on new environments for 10 rewards
percentage = 100 * total_num_steps // args.num_env_steps
if (args.eval_interval is not None and percentage > 1
and (j % args.eval_interval == 0 or j == num_updates - 1)):
print("###### EVALUATING #######")
args_eval = copy.deepcopy(args)
args_eval.num_processes = 1
eval_envs = make_vec_envs(args_eval, device, no_logging=True)
eval_episode_rewards = []
occ_obs, sign_obs = eval_envs.reset()
eval_recurrent_hidden_states = torch.zeros(
args_eval.num_processes,
target_actor_critic.recurrent_hidden_state_size,
device=device)
eval_masks = torch.zeros(args_eval.num_processes, 1, device=device)
while len(eval_episode_rewards) < 3000:
with torch.no_grad():
_, action, _, eval_recurrent_hidden_states = target_actor_critic.act(
occ_obs,
sign_obs,
eval_recurrent_hidden_states,
eval_masks,
deterministic=True)
# Obser reward and next obs
occ_obs, sign_obs, reward, done, infos = eval_envs.step(action)
eval_masks = torch.FloatTensor([[0.0] if done_ else [1.0]
for done_ in done])
eval_episode_rewards.append(reward)
eval_envs.close()
if np.mean(eval_episode_rewards) > best_eval_rewards:
best_eval_rewards = np.mean(eval_episode_rewards)
save_model = target_actor_critic
if args.cuda:
save_model = copy.deepcopy(target_actor_critic).cpu()
torch.save(save_model, os.path.join(save_path,
'best_model.pt'))
## Visualize tracked rewards(over num_steps) over time
if args.vis:
visualize(reward_track, args.algo, save_path)
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)
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))
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))
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
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, True) # True to allow early resets
actor_critic = Policy(
envs.observation_space.shape,
envs.action_space,
base_kwargs={'recurrent': args.recurrent_policy, 'is_genesis':args.env_name.endswith("Genesis")})
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
f = open("log.txt","w")
f.write("#Steps\tReturn\n")
f.flush()
cumulative_steps = 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)
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])
#envs.render()
cumulative_steps += 1
# Obser reward and next obs
if device.type == 'cuda': # For some reason, CUDA actions are nested in an extra layer
action = action[0]
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
average_return = evaluate(actor_critic, envs, args.num_processes, device)
f.write("{}\t{}\n".format(cumulative_steps, average_return))
f.flush()
# Close file if training ever finishes
f.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 + 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():
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()
torch.set_num_threads(1)
# device = torch.device("cuda:0" if args.cuda else "cpu")
device = torch.device("cpu")
# args.env_name = 'Pong-ramNoFrameskip-v4'
args.env_name = 'Pong-ram-v0'
args.num_processes = 2
envs = make_vec_envs(args.env_name, args.seed, args.num_processes,
args.gamma, args.log_dir, device, False)
# ss('here')
actor_critic = Policy(envs.observation_space.shape,
envs.action_space,
base_kwargs={'recurrent': args.recurrent_policy})
actor_critic.to(device)
print(args.recurrent_policy)
print(args.clip_param)
print(args.ppo_epoch)
print('ccccccccc')
print(args.num_mini_batch)
print(args.value_loss_coef)
print(args.entropy_coef)
print('dddddddddddd')
print(args.lr)
print(args.eps)
print(args.max_grad_norm)
ss('in main, after actor_critic')
args.num_mini_batch = 2
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)
# ss('out of define ppo')
args.num_steps = 4
rollouts = RolloutStorage(args.num_steps, args.num_processes,
envs.observation_space.shape, envs.action_space,
actor_critic.recurrent_hidden_state_size)
# ss('rollouts')
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(args.num_env_steps)
# print()
# ss('pp')
sum_re = torch.zeros(args.num_processes, 1)
# print(sum_re.shape)
for j in range(num_updates):
# ss('pp')
is_any_done = False
for step in range(args.num_steps):
# for step in range(50000):
# print(step)
# ss('pp')
# 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])
# print(value)
# print(action_log_prob)
# print(action)
# ss('runner')
# Obser reward and next obs
obs, reward, done, infos = envs.step(action)
sum_re += reward
# print('- --'*20)
# print(reward)
# print(sum_re)
# print()
# print(reward.shape)
if any(done):
# print(sum_re)
# print(done)
# input('hi')
# is_any_done = True
for i in range(len(done)):
if done[i]:
# print(i)
# print(*sum_re[i])
# print(sum_re[i].item())
episode_rewards.append(sum_re[i].item())
# print(sum_re[i])
sum_re[i] *= 0
# pass
# episode_rewards.append(reward.item())
# ss('make reward')
# print(infos)
# ss('runner')
for info in infos:
# print(info)
if 'episode' in info.keys():
episode_rewards.append(info['episode']['r'])
print('what env info with episode do?', info.keys())
# ss('break')
# 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)
# ss('runner')
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, is_any_done,
args.use_proper_time_limits)
# ss('runner1')
value_loss, action_loss, dist_entropy = agent.update(rollouts)
# ss('runner1')
rollouts.after_update()
# ss('runner2')
# 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"))
# print(args.log_interval)
args.log_interval = 100
if j % args.log_interval == 0 and len(episode_rewards) > 1:
# 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 Ent {},V {},A {}"
.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 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)
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(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)
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])
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()
# 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(base=IAMBase, num_frame_stack=None):
seed = 1
env_name = "Warehouse-v0"
num_processes = 32
log_dir = './logs/'
eval_interval = None
log_interval = 10
use_linear_lr_decay = False
use_proper_time_limits = False
save_dir = './trained_models/'
use_cuda = True
# PPO
gamma = 0.99 # reward discount factor
clip_param = 0.1 #0.2
ppo_epoch = 3 #4
num_mini_batch = 32
value_loss_coef = 1 #0.5
entropy_coef = 0.01
lr = 2.5e-4 #7e-4
eps = 1e-5
max_grad_norm = float('inf')
use_gae = True
gae_lambda = 0.95
num_steps = 8 #5
# Store
num_env_steps = 4e6
save_interval = 100
# IAM
dset = [
0, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65,
66, 67, 68, 69, 70, 71, 72
]
#gym.envs.register(env_name, entry_point="environments.warehouse.warehouse:Warehouse",
# kwargs={'seed': seed, 'parameters': {"num_frames": 1}})
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
log_dir = os.path.expanduser(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 use_cuda else "cpu")
envs = make_vec_envs(env_name,
seed,
num_processes,
gamma,
log_dir,
device,
False,
num_frame_stack=num_frame_stack)
actor_critic = Policy(envs.observation_space.shape,
envs.action_space,
base=base,
base_kwargs=({
'dset': dset
} if base == IAMBase else {}))
actor_critic.to(device)
agent = algo.PPO(actor_critic,
clip_param,
ppo_epoch,
num_mini_batch,
value_loss_coef,
entropy_coef,
lr=lr,
eps=eps,
max_grad_norm=max_grad_norm)
rollouts = RolloutStorage(num_steps, 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(num_env_steps) // num_steps // num_processes
for j in range(num_updates):
if use_linear_lr_decay:
# decrease learning rate linearly
utils.update_linear_schedule(agent.optimizer, j, num_updates, lr)
for step in range(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, use_gae, gamma, gae_lambda,
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 % save_interval == 0 or j == num_updates - 1) and save_dir != "":
save_path = os.path.join(save_dir, 'PPO')
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, env_name + ".pt"))
if j % log_interval == 0 and len(episode_rewards) > 1:
total_num_steps = (j + 1) * num_processes * 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 (eval_interval is not None and len(episode_rewards) > 1
and j % eval_interval == 0):
obs_rms = utils.get_vec_normalize(envs).obs_rms
evaluate(actor_critic, obs_rms, env_name, seed, 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 train_loop(agent,
envs,
env_name,
num_updates,
num_steps,
curiosity_module=None,
save_interval=None,
eval_interval=None,
log_interval=None,
time_limit=1000,
curiosity_rew_after=0,
curiosity_rew_before=None,
use_linear_lr_decay=True,
lr_decay_horizon=None,
callbacks=None):
# Create rollout storage
num_processes = envs.num_envs
rollouts = RolloutStorage(num_steps, num_processes,
envs.observation_space.shape, envs.action_space,
agent.actor_critic.recurrent_hidden_state_size)
obs = envs.reset()
rollouts.obs[0].copy_(obs)
device = next(agent.actor_critic.parameters()).device
rollouts.to(device)
# Create curiosity statistics saver
if curiosity_module is not None:
curiosity_stats = CuriosityStatistics(num_processes=num_processes,
time_limit=time_limit)
else:
curiosity_stats = None
# Create train statistics saver
stats = TrainStatistics(log_interval,
with_curiosity=(curiosity_module is not None))
# Possibility to use curiosity only for a few epochs
if curiosity_rew_before is None:
curiosity_rew_before = num_updates
# Train loop
start = time.time()
for j in range(num_updates):
if use_linear_lr_decay:
if lr_decay_horizon is None:
lr_decay_horizon = num_updates
utils.update_linear_schedule(agent.optimizer, j, lr_decay_horizon,
agent.optimizer.defaults['lr'])
curiosity_loss = 0
for step in range(num_steps):
# Sample actions
with torch.no_grad():
value, action, action_log_prob, recurrent_hidden_states = agent.actor_critic.act(
rollouts.obs[step], rollouts.recurrent_hidden_states[step],
rollouts.masks[step])
# Observe extrinsic reward and next obs
obs, reward, done, infos = envs.step(action)
# Compute intrinsic rewards
if curiosity_module is not None:
time_limit_mask = np.array([
0 if 'bad_transition' in info.keys() else 1
for info in infos
])
if use_proper_time_limits:
curiosity_done = done * time_limit_mask
else:
curiosity_done = done
curiosity_reward = curiosity_module.get_reward(
rollouts.obs[step], action, obs, curiosity_done)
curiosity_loss += curiosity_module.update(
rollouts.obs[step], action, obs, curiosity_done)
# Update current reward statistics
stats.update_extrinsic_reward(infos)
if curiosity_module is not None:
curiosity_stats.update(curiosity_reward.cpu().numpy().ravel(),
done)
# 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 (curiosity_module is not None) and (j >= curiosity_rew_after and
j <= curiosity_rew_before):
reward = reward + curiosity_reward
rollouts.insert(obs, recurrent_hidden_states, action,
action_log_prob, value, reward, masks, bad_masks)
with torch.no_grad():
next_value = agent.actor_critic.get_value(
rollouts.obs[-1], rollouts.recurrent_hidden_states[-1],
rollouts.masks[-1]).detach()
rollouts.compute_returns(next_value, use_gae, gamma, gae_lambda,
use_proper_time_limits)
value_loss, action_loss, dist_entropy = agent.update(rollouts)
rollouts.after_update()
if curiosity_module is not None:
curiosity_loss /= num_steps
else:
curiosity_loss = None
stats.update_losses(value_loss, action_loss, dist_entropy,
curiosity_loss)
# Save for every interval-th episode or for the last epoch
if (j % save_interval == 0 or j == num_updates - 1) and save_dir != "":
save_path = os.path.join(save_dir, "ppo")
try:
os.makedirs(save_path)
except OSError:
pass
torch.save([
agent.actor_critic,
getattr(utils.get_vec_normalize(envs), 'ob_rms', None)
], os.path.join(save_path, env_name + ".pt"))
if j % eval_interval == 0:
try:
ob_rms = utils.get_vec_normalize(envs).ob_rms
except:
ob_rms = None
stats.eval_episode_rewards.append(
evaluate(agent.actor_critic, env_name, device, ob_rms))
if j % log_interval == 0 and len(stats.episode_rewards) > 1:
total_num_steps = (j + 1) * num_processes * num_steps
end = time.time()
stats.update_log(total_num_steps, curiosity_stats)
if callbacks is not None:
for callback in callbacks:
callback(stats,
agent,
n_updates=j,
total_n_steps=total_num_steps,
fps=int(total_num_steps / (end - start)))
return stats
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():
chrono = exp.chrono()
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 = 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,
chrono=chrono)
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(args.repeat):
with chrono.time('train') as t:
for n in range(args.number):
with chrono.time('one_batch'):
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))
with chrono.time('generate_rollouts'):
generate_rollouts(**locals())
with torch.no_grad():
next_value = actor_critic.get_value(
rollouts.obs[-1],
rollouts.recurrent_hidden_states[-1],
rollouts.masks[-1]).detach()
# ---
with chrono.time('compute_returns'):
rollouts.compute_returns(next_value, args.use_gae,
args.gamma, args.tau)
with chrono.time('agent.update'): # 11.147009023304644
value_loss, action_loss, dist_entropy = agent.update(
rollouts)
exp.log_batch_loss(action_loss)
exp.log_metric('value_loss', value_loss)
with chrono.time('after_update'):
rollouts.after_update()
total_num_steps = (j +
1) * args.num_processes * args.num_steps
if args.eval_interval is not None and len(
episode_rewards
) > 1 and j % args.eval_interval == 0:
eval_model(**locals())
# -- number
# -- chrono
exp.show_eta(j, t)
# -- epoch
exp.report()
def main():
args = get_args()
import random
random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)
np.random.seed(args.seed)
if args.cuda and torch.cuda.is_available() and args.cuda_deterministic:
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
logdir = args.env_name + '_' + args.algo + '_num_arms_' + str(
args.num_processes) + '_' + time.strftime("%d-%m-%Y_%H-%M-%S")
if args.use_privacy:
logdir = logdir + '_privacy'
elif args.use_noisygrad:
logdir = logdir + '_noisygrad'
elif args.use_pcgrad:
logdir = logdir + '_pcgrad'
elif args.use_testgrad:
logdir = logdir + '_testgrad'
elif args.use_median_grad:
logdir = logdir + '_mediangrad'
logdir = os.path.join('runs', logdir)
logdir = os.path.join(os.path.expanduser(args.log_dir), logdir)
utils.cleanup_log_dir(logdir)
# Ugly but simple logging
log_dict = {
'task_steps': args.task_steps,
'grad_noise_ratio': args.grad_noise_ratio,
'max_task_grad_norm': args.max_task_grad_norm,
'use_noisygrad': args.use_noisygrad,
'use_pcgrad': args.use_pcgrad,
'use_testgrad': args.use_testgrad,
'use_testgrad_median': args.use_testgrad_median,
'testgrad_quantile': args.testgrad_quantile,
'median_grad': args.use_median_grad,
'use_meanvargrad': args.use_meanvargrad,
'meanvar_beta': args.meanvar_beta,
'no_special_grad_for_critic': args.no_special_grad_for_critic,
'use_privacy': args.use_privacy,
'seed': args.seed,
'recurrent': args.recurrent_policy,
'obs_recurrent': args.obs_recurrent,
'cmd': ' '.join(sys.argv[1:])
}
for eval_disp_name, eval_env_name in EVAL_ENVS.items():
log_dict[eval_disp_name] = []
summary_writer = SummaryWriter()
summary_writer.add_hparams(
{
'task_steps': args.task_steps,
'grad_noise_ratio': args.grad_noise_ratio,
'max_task_grad_norm': args.max_task_grad_norm,
'use_noisygrad': args.use_noisygrad,
'use_pcgrad': args.use_pcgrad,
'use_testgrad': args.use_testgrad,
'use_testgrad_median': args.use_testgrad_median,
'testgrad_quantile': args.testgrad_quantile,
'median_grad': args.use_median_grad,
'use_meanvargrad': args.use_meanvargrad,
'meanvar_beta': args.meanvar_beta,
'no_special_grad_for_critic': args.no_special_grad_for_critic,
'use_privacy': args.use_privacy,
'seed': args.seed,
'recurrent': args.recurrent_policy,
'obs_recurrent': args.obs_recurrent,
'cmd': ' '.join(sys.argv[1:])
}, {})
torch.set_num_threads(1)
device = torch.device("cuda:0" if args.cuda else "cpu")
print('making envs...')
envs = make_vec_envs(args.env_name,
args.seed,
args.num_processes,
args.gamma,
args.log_dir,
device,
False,
steps=args.task_steps,
free_exploration=args.free_exploration,
recurrent=args.recurrent_policy,
obs_recurrent=args.obs_recurrent,
multi_task=True)
val_envs = make_vec_envs(args.val_env_name,
args.seed,
args.num_processes,
args.gamma,
args.log_dir,
device,
False,
steps=args.task_steps,
free_exploration=args.free_exploration,
recurrent=args.recurrent_policy,
obs_recurrent=args.obs_recurrent,
multi_task=True)
eval_envs_dic = {}
for eval_disp_name, eval_env_name in EVAL_ENVS.items():
eval_envs_dic[eval_disp_name] = make_vec_envs(
eval_env_name[0],
args.seed,
args.num_processes,
None,
logdir,
device,
True,
steps=args.task_steps,
recurrent=args.recurrent_policy,
obs_recurrent=args.obs_recurrent,
multi_task=True,
free_exploration=args.free_exploration)
prev_eval_r = {}
print('done')
if args.hard_attn:
actor_critic = Policy(envs.observation_space.shape,
envs.action_space,
base=MLPHardAttnBase,
base_kwargs={
'recurrent':
args.recurrent_policy or args.obs_recurrent
})
else:
actor_critic = Policy(envs.observation_space.shape,
envs.action_space,
base=MLPAttnBase,
base_kwargs={
'recurrent':
args.recurrent_policy or args.obs_recurrent
})
actor_critic.to(device)
if (args.continue_from_epoch > 0) and args.save_dir != "":
save_path = os.path.join(args.save_dir, args.algo)
actor_critic_, loaded_obs_rms_ = torch.load(
os.path.join(
save_path, args.env_name +
"-epoch-{}.pt".format(args.continue_from_epoch)))
actor_critic.load_state_dict(actor_critic_.state_dict())
if args.algo != 'ppo':
raise "only PPO is supported"
agent = algo.PPO(actor_critic,
args.clip_param,
args.ppo_epoch,
args.num_mini_batch,
args.value_loss_coef,
args.entropy_coef,
lr=args.lr,
eps=args.eps,
num_tasks=args.num_processes,
attention_policy=False,
max_grad_norm=args.max_grad_norm,
weight_decay=args.weight_decay)
val_agent = algo.PPO(actor_critic,
args.clip_param,
args.ppo_epoch,
args.num_mini_batch,
args.value_loss_coef,
args.entropy_coef,
lr=args.val_lr,
eps=args.eps,
num_tasks=args.num_processes,
attention_policy=True,
max_grad_norm=args.max_grad_norm,
weight_decay=args.weight_decay)
rollouts = RolloutStorage(args.num_steps, args.num_processes,
envs.observation_space.shape, envs.action_space,
actor_critic.recurrent_hidden_state_size)
val_rollouts = RolloutStorage(args.num_steps, args.num_processes,
val_envs.observation_space.shape,
val_envs.action_space,
actor_critic.recurrent_hidden_state_size)
obs = envs.reset()
rollouts.obs[0].copy_(obs)
rollouts.to(device)
val_obs = val_envs.reset()
val_rollouts.obs[0].copy_(val_obs)
val_rollouts.to(device)
episode_rewards = deque(maxlen=10)
start = time.time()
num_updates = int(
args.num_env_steps) // args.num_steps // args.num_processes
save_copy = True
for j in range(args.continue_from_epoch,
args.continue_from_epoch + num_updates):
# policy rollouts
for step in range(args.num_steps):
# Sample actions
actor_critic.eval()
with torch.no_grad():
value, action, action_log_prob, recurrent_hidden_states = actor_critic.act(
rollouts.obs[step], rollouts.recurrent_hidden_states[step],
rollouts.masks[step])
actor_critic.train()
# Obser reward and next obs
obs, reward, done, infos = envs.step(action)
for info in infos:
if 'episode' in info.keys():
episode_rewards.append(info['episode']['r'])
for k, v in info['episode'].items():
summary_writer.add_scalar(
f'training/{k}', v,
j * args.num_processes * args.num_steps +
args.num_processes * step)
# If done then clean the history of observations.
masks = torch.FloatTensor([[0.0] if done_ else [1.0]
for done_ in done])
bad_masks = torch.FloatTensor(
[[0.0] if 'bad_transition' in info.keys() else [1.0]
for info in infos])
rollouts.insert(obs, recurrent_hidden_states, action,
action_log_prob, value, reward, masks, bad_masks)
actor_critic.eval()
with torch.no_grad():
next_value = actor_critic.get_value(
rollouts.obs[-1], rollouts.recurrent_hidden_states[-1],
rollouts.masks[-1]).detach()
actor_critic.train()
rollouts.compute_returns(next_value, args.use_gae, args.gamma,
args.gae_lambda, args.use_proper_time_limits)
if save_copy:
prev_weights = copy.deepcopy(actor_critic.state_dict())
prev_opt_state = copy.deepcopy(agent.optimizer.state_dict())
prev_val_opt_state = copy.deepcopy(
val_agent.optimizer.state_dict())
save_copy = False
value_loss, action_loss, dist_entropy = agent.update(rollouts)
rollouts.after_update()
# validation rollouts
for val_iter in range(args.val_agent_steps):
for step in range(args.num_steps):
# Sample actions
actor_critic.eval()
with torch.no_grad():
value, action, action_log_prob, recurrent_hidden_states = actor_critic.act(
val_rollouts.obs[step],
val_rollouts.recurrent_hidden_states[step],
val_rollouts.masks[step])
actor_critic.train()
# Obser reward and next obs
obs, reward, done, infos = val_envs.step(action)
# If done then clean the history of observations.
masks = torch.FloatTensor([[0.0] if done_ else [1.0]
for done_ in done])
bad_masks = torch.FloatTensor(
[[0.0] if 'bad_transition' in info.keys() else [1.0]
for info in infos])
val_rollouts.insert(obs, recurrent_hidden_states, action,
action_log_prob, value, reward, masks,
bad_masks)
actor_critic.eval()
with torch.no_grad():
next_value = actor_critic.get_value(
val_rollouts.obs[-1],
val_rollouts.recurrent_hidden_states[-1],
val_rollouts.masks[-1]).detach()
actor_critic.train()
val_rollouts.compute_returns(next_value, args.use_gae, args.gamma,
args.gae_lambda,
args.use_proper_time_limits)
val_value_loss, val_action_loss, val_dist_entropy = val_agent.update(
val_rollouts)
val_rollouts.after_update()
# save for every interval-th episode or for the last epoch
if (j % args.save_interval == 0
or j == num_updates - 1) and args.save_dir != "":
save_path = os.path.join(args.save_dir, args.algo)
try:
os.makedirs(save_path)
except OSError:
pass
torch.save([
actor_critic,
getattr(utils.get_vec_normalize(envs), 'obs_rms', None)
], os.path.join(save_path,
args.env_name + "-epoch-{}.pt".format(j)))
if j % args.log_interval == 0 and len(episode_rewards) > 1:
total_num_steps = (j + 1) * args.num_processes * args.num_steps
end = time.time()
print(
"Updates {}, num timesteps {}, FPS {} \n Last {} training episodes: mean/median reward {:.1f}/{:.1f}, min/max reward {:.1f}/{:.1f}\n"
.format(j, total_num_steps,
int(total_num_steps / (end - start)),
len(episode_rewards), np.mean(episode_rewards),
np.median(episode_rewards), np.min(episode_rewards),
np.max(episode_rewards), dist_entropy, value_loss,
action_loss))
revert = False
if (args.eval_interval is not None and len(episode_rewards) > 1
and j % args.eval_interval == 0):
actor_critic.eval()
obs_rms = utils.get_vec_normalize(envs).obs_rms
eval_r = {}
printout = f'Seed {args.seed} Iter {j} '
for eval_disp_name, eval_env_name in EVAL_ENVS.items():
eval_r[eval_disp_name] = evaluate(
actor_critic,
obs_rms,
eval_envs_dic,
eval_disp_name,
args.seed,
args.num_processes,
eval_env_name[1],
logdir,
device,
steps=args.task_steps,
recurrent=args.recurrent_policy,
obs_recurrent=args.obs_recurrent,
multi_task=True,
free_exploration=args.free_exploration)
if eval_disp_name in prev_eval_r:
diff = np.array(eval_r[eval_disp_name]) - np.array(
prev_eval_r[eval_disp_name])
if eval_disp_name == 'many_arms':
if np.sum(diff > 0) - np.sum(
diff < 0) < args.val_improvement_threshold:
print('no update')
revert = True
summary_writer.add_scalar(f'eval/{eval_disp_name}',
np.mean(eval_r[eval_disp_name]),
(j + 1) * args.num_processes *
args.num_steps)
log_dict[eval_disp_name].append([
(j + 1) * args.num_processes * args.num_steps,
eval_r[eval_disp_name]
])
printout += eval_disp_name + ' ' + str(
np.mean(eval_r[eval_disp_name])) + ' '
# summary_writer.add_scalars('eval_combined', eval_r, (j+1) * args.num_processes * args.num_steps)
if revert:
actor_critic.load_state_dict(prev_weights)
agent.optimizer.load_state_dict(prev_opt_state)
val_agent.optimizer.load_state_dict(prev_val_opt_state)
else:
print(printout)
prev_eval_r = eval_r.copy()
save_copy = True
actor_critic.train()
save_obj(log_dict, os.path.join(logdir, 'log_dict.pkl'))
envs.close()
val_envs.close()
for eval_disp_name, eval_env_name in EVAL_ENVS.items():
eval_envs_dic[eval_disp_name].close()
def train_maml_like_ppo_(
init_model,
args,
learning_rate,
num_episodes=20,
num_updates=1,
vis=False,
run_idx=0,
use_linear_lr_decay=False,
):
num_steps = num_episodes * 100
torch.set_num_threads(1)
device = torch.device("cpu")
envs = make_vec_envs(ENV_NAME, seeding.create_seed(None), NUM_PROC,
args.gamma, None, device, allow_early_resets=True, normalize=args.norm_vectors)
raw_env = navigation_2d.unpeele_navigation_env(envs, 0)
# raw_env.set_arguments(args.rm_nogo, args.reduce_goals, True, args.large_nogos)
new_task = raw_env.sample_tasks(run_idx)
raw_env.reset_task(new_task[0])
# actor_critic = Policy(
# envs.observation_space.shape,
# envs.action_space,
# base_kwargs={'recurrent': args.recurrent_policy})
actor_critic = copy.deepcopy(init_model)
actor_critic.to(device)
agent = algo.PPO(
actor_critic,
args.clip_param,
args.ppo_epoch,
args.num_mini_batch,
args.value_loss_coef,
args.entropy_coef,
lr=learning_rate,
eps=args.eps,
max_grad_norm=args.max_grad_norm)
rollouts = RolloutStorage(num_steps, NUM_PROC,
envs.observation_space.shape, envs.action_space,
actor_critic.recurrent_hidden_state_size)
obs = envs.reset()
rollouts.obs[0].copy_(obs)
rollouts.to(device)
fitnesses = []
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)
min_c_rew = float("inf")
vis = []
offending = []
for step in range(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 done[0]:
c_rew = infos[0]["cummulative_reward"]
vis.append((infos[0]['path'], infos[0]['goal']))
offending.extend(infos[0]['offending'])
if c_rew < min_c_rew:
min_c_rew = c_rew
# 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)
rollouts.after_update()
ob_rms = utils.get_vec_normalize(envs)
if ob_rms is not None:
ob_rms = ob_rms.ob_rms
fits, info = evaluate(actor_critic, ob_rms, envs, NUM_PROC, device)
print(f"fitness {fits} update {j+1}")
if (j+1) % 1 == 0:
vis_path(vis, eval_path_rec=info['path'], offending=offending)
fitnesses.append(fits)
return fitnesses[-1], info[0]['reached'], None
def inner_loop_ppo(args, learning_rate, num_steps, num_updates, inst_on,
visualize, save_dir):
torch.set_num_threads(1)
log_writer = SummaryWriter(save_dir, max_queue=1, filename_suffix="log")
device = torch.device("cpu")
env_name = ENV_NAME # "Safexp-PointGoal1-v0"
envs = make_vec_envs(env_name,
np.random.randint(2**32),
NUM_PROC,
args.gamma,
None,
device,
allow_early_resets=True,
normalize=args.norm_vectors)
eval_envs = make_vec_envs(env_name,
np.random.randint(2**32),
1,
args.gamma,
None,
device,
allow_early_resets=True,
normalize=args.norm_vectors)
actor_critic_policy = init_default_ppo(envs, log(args.init_sigma))
# Prepare modified observation shape for instinct
obs_shape = envs.observation_space.shape
inst_action_space = deepcopy(envs.action_space)
inst_obs_shape = list(obs_shape)
inst_obs_shape[0] = inst_obs_shape[0] + envs.action_space.shape[0]
# Prepare modified action space for instinct
inst_action_space.shape = list(inst_action_space.shape)
inst_action_space.shape[0] = inst_action_space.shape[0] + 1
inst_action_space.shape = tuple(inst_action_space.shape)
actor_critic_instinct = Policy(tuple(inst_obs_shape),
inst_action_space,
init_log_std=log(args.init_sigma),
base_kwargs={'recurrent': False})
actor_critic_policy.to(device)
actor_critic_instinct.to(device)
agent_policy = algo.PPO(actor_critic_policy,
args.clip_param,
args.ppo_epoch,
args.num_mini_batch,
args.value_loss_coef,
args.entropy_coef,
lr=learning_rate,
eps=args.eps,
max_grad_norm=args.max_grad_norm)
agent_instinct = algo.PPO(actor_critic_instinct,
args.clip_param,
args.ppo_epoch,
args.num_mini_batch,
args.value_loss_coef,
args.entropy_coef,
lr=learning_rate,
eps=args.eps,
max_grad_norm=args.max_grad_norm)
rollouts_rewards = RolloutStorage(
num_steps, NUM_PROC, envs.observation_space.shape, envs.action_space,
actor_critic_policy.recurrent_hidden_state_size)
rollouts_cost = RolloutStorage(
num_steps, NUM_PROC, inst_obs_shape, inst_action_space,
actor_critic_instinct.recurrent_hidden_state_size)
obs = envs.reset()
i_obs = torch.cat(
[obs, torch.zeros((NUM_PROC, envs.action_space.shape[0]))],
dim=1) # Add zero action to the observation
rollouts_rewards.obs[0].copy_(obs)
rollouts_rewards.to(device)
rollouts_cost.obs[0].copy_(i_obs)
rollouts_cost.to(device)
fitnesses = []
best_fitness_so_far = float("-Inf")
is_instinct_training = False
for j in range(num_updates):
is_instinct_training_old = is_instinct_training
is_instinct_training = phase_shifter(
j, PHASE_LENGTH,
len(TrainPhases)) == TrainPhases.INSTINCT_TRAIN_PHASE.value
is_instinct_deterministic = not is_instinct_training
is_policy_deterministic = not is_instinct_deterministic
for step in range(num_steps):
# Sample actions
with torch.no_grad():
# (value, action, action_log_probs, rnn_hxs), (instinct_value, instinct_action, instinct_outputs_log_prob, i_rnn_hxs), final_action
value, action, action_log_probs, recurrent_hidden_states = actor_critic_policy.act(
rollouts_rewards.obs[step],
rollouts_rewards.recurrent_hidden_states[step],
rollouts_rewards.masks[step],
deterministic=is_policy_deterministic)
instinct_value, instinct_action, instinct_outputs_log_prob, instinct_recurrent_hidden_states = actor_critic_instinct.act(
rollouts_cost.obs[step],
rollouts_cost.recurrent_hidden_states[step],
rollouts_cost.masks[step],
deterministic=is_instinct_deterministic,
)
# Combine two networks
final_action, i_control = policy_instinct_combinator(
action, instinct_action)
obs, reward, done, infos = envs.step(final_action)
# envs.render()
reward, violation_cost = reward_cost_combinator(
reward, infos, NUM_PROC, i_control)
# 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_rewards.insert(obs, recurrent_hidden_states, action,
action_log_probs, value, reward, masks,
bad_masks)
i_obs = torch.cat([obs, action], dim=1)
rollouts_cost.insert(i_obs, instinct_recurrent_hidden_states,
instinct_action, instinct_outputs_log_prob,
instinct_value, violation_cost, masks,
bad_masks)
with torch.no_grad():
next_value_policy = actor_critic_policy.get_value(
rollouts_rewards.obs[-1],
rollouts_rewards.recurrent_hidden_states[-1],
rollouts_rewards.masks[-1]).detach()
next_value_instinct = actor_critic_instinct.get_value(
rollouts_cost.obs[-1],
rollouts_cost.recurrent_hidden_states[-1],
rollouts_cost.masks[-1].detach())
rollouts_rewards.compute_returns(next_value_policy, args.use_gae,
args.gamma, args.gae_lambda,
args.use_proper_time_limits)
rollouts_cost.compute_returns(next_value_instinct, args.use_gae,
args.gamma, args.gae_lambda,
args.use_proper_time_limits)
if not is_instinct_training:
print("training policy")
# Policy training phase
p_before = deepcopy(agent_instinct.actor_critic)
value_loss, action_loss, dist_entropy = agent_policy.update(
rollouts_rewards)
val_loss_i, action_loss_i, dist_entropy_i = 0, 0, 0
p_after = deepcopy(agent_instinct.actor_critic)
assert compare_two_models(
p_before, p_after), "policy changed when it shouldn't"
else:
print("training instinct")
# Instinct training phase
value_loss, action_loss, dist_entropy = 0, 0, 0
p_before = deepcopy(agent_policy.actor_critic)
val_loss_i, action_loss_i, dist_entropy_i = agent_instinct.update(
rollouts_cost)
p_after = deepcopy(agent_policy.actor_critic)
assert compare_two_models(
p_before, p_after), "policy changed when it shouldn't"
rollouts_rewards.after_update()
rollouts_cost.after_update()
ob_rms = utils.get_vec_normalize(envs)
if ob_rms is not None:
ob_rms = ob_rms.ob_rms
fits, info = evaluate(EvalActorCritic(actor_critic_policy,
actor_critic_instinct),
ob_rms,
eval_envs,
NUM_PROC,
reward_cost_combinator,
device,
instinct_on=inst_on,
visualise=visualize)
instinct_reward = info['instinct_reward']
eval_hazard_collisions = info['hazard_collisions']
print(
f"Step {j}, Fitness {fits.item()}, value_loss = {value_loss}, action_loss = {action_loss}, "
f"dist_entropy = {dist_entropy}")
print(
f"Step {j}, Instinct reward {instinct_reward}, value_loss instinct = {val_loss_i}, action_loss instinct= {action_loss_i}, "
f"dist_entropy instinct = {dist_entropy_i} hazard_collisions = {eval_hazard_collisions}"
)
print(
"-----------------------------------------------------------------"
)
# Tensorboard logging
log_writer.add_scalar("fitness", fits.item(), j)
log_writer.add_scalar("value loss", value_loss, j)
log_writer.add_scalar("action loss", action_loss, j)
log_writer.add_scalar("dist entropy", dist_entropy, j)
log_writer.add_scalar("cost/instinct_reward", instinct_reward, j)
log_writer.add_scalar("cost/hazard_collisions", eval_hazard_collisions,
j)
log_writer.add_scalar("value loss instinct", val_loss_i, j)
log_writer.add_scalar("action loss instinct", action_loss_i, j)
log_writer.add_scalar("dist entropy instinct", dist_entropy_i, j)
fitnesses.append(fits)
if fits.item() > best_fitness_so_far:
best_fitness_so_far = fits.item()
torch.save(actor_critic_policy, join(save_dir,
"model_rl_policy.pt"))
torch.save(actor_critic_instinct,
join(save_dir, "model_rl_instinct.pt"))
if is_instinct_training != is_instinct_training_old:
torch.save(actor_critic_policy,
join(save_dir, f"model_rl_policy_update_{j}.pt"))
torch.save(actor_critic_instinct,
join(save_dir, f"model_rl_instinct_update_{j}.pt"))
torch.save(actor_critic_policy,
join(save_dir, "model_rl_policy_latest.pt"))
torch.save(actor_critic_instinct,
join(save_dir, "model_rl_instinct_latest.pt"))
return (fitnesses[-1]), 0, 0
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
source_path = os.path.join(save_path, f"{args.env_name}-s{args.seed}.pt")
torch.save([actor_critic, getattr(utils.get_vec_normalize(envs), "ob_rms", None)], source_path)
if "gibson" in args.custom_gym and "TwoPlayer" in args.env_name:
# copy over policy
# nasty, nasty, first unwrapped is to get to dummyVecEnv, then to source
def learn(env, max_timesteps, timesteps_per_batch, clip_param):
ppo_epoch = 5
num_step = timesteps_per_batch
save_interval = 100
seed = 1000
batch_size = 64
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
log_dir = os.path.expanduser('/tmp/gym/')
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")
envs = make_vec_envs(env, seed, 8, 0.95, log_dir, device, False)
actor_critic = Policy(envs.observation_space.shape,
envs.action_space,
base_kwargs={'recurrent': False})
actor_critic.to(device)
agent = algo.PPO(actor_critic,
clip_param,
ppo_epoch,
batch_size,
0.5,
0.01,
lr=0.00025,
eps=1e-05,
max_grad_norm=0.5)
rollouts = RolloutStorage(num_step, 8, envs.observation_space.shape,
envs.action_space,
actor_critic.recurrent_hidden_state_size)
obs = envs.reset()
rollouts.obs[0].copy_(torch.tensor(obs))
rollouts.to(device)
episode_rewards = deque(maxlen=10)
start = time.time()
num_updates = int(max_timesteps) // num_step // 8
for j in range(num_updates):
# decrease learning rate linearly
utils.update_linear_schedule(agent.optimizer, j, num_updates, 0.00025)
for step in range(num_step):
# 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, True, 0.99, 0.95, False)
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 % save_interval == 0
or j == num_updates - 1) and "./trained_models/" != "":
save_path = os.path.join("./trained_models/", 'ppo')
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, 'UniversalPolicy' + ".pt"))
if j % 1 == 0 and len(episode_rewards) > 1:
total_num_steps = (j + 1) * 8 * num_step
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():
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():
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 train(train_states,
run_dir,
num_env_steps,
eval_env_steps,
writer,
writer_name,
args,
init_model=None):
envs = make_vec_envs(train_states, args.seed, args.num_processes,
args.gamma, 'cpu', 'train', args)
if init_model:
actor_critic, env_step, model_name = init_model
obs_space = actor_critic.obs_space
obs_process = actor_critic.obs_process
obs_module = actor_critic.obs_module
print(f" [load] Loaded model {model_name} at step {env_step}")
else:
obs_space = envs.observation_space
actor_critic = Policy(obs_space,
args.obs_process,
args.obs_module,
envs.action_space,
base_kwargs={'recurrent': args.recurrent_policy})
env_step = 0
actor_critic.to(args.device)
#print(actor_critic)
run_name = run_dir.replace('/', '_')
vid_save_dir = f"{run_dir}/videos/"
try:
os.makedirs(vid_save_dir)
except OSError:
pass
ckpt_save_dir = f"{run_dir}/ckpts/"
try:
os.makedirs(ckpt_save_dir)
except OSError:
pass
if args.algo == 'ppo':
agent = algo.PPO(actor_critic,
args.clip_param,
args.ppo_epoch,
args.num_mini_batch,
args.value_loss_coef,
args.entropy_coef,
args.device,
lr=args.lr,
eps=args.eps,
max_grad_norm=args.max_grad_norm)
elif args.algo == 'a2c':
agent = algo.A2C_ACKTR(actor_critic,
args.value_loss_coef,
args.entropy_coef,
lr=args.lr,
eps=args.eps,
alpha=args.alpha,
max_grad_norm=args.max_grad_norm,
acktr=False)
elif args.algo == 'acktr':
agent = algo.A2C_ACKTR(actor_critic,
args.value_loss_coef,
args.entropy_coef,
lr=args.lr,
eps=args.eps,
alpha=args.alpha,
max_grad_norm=args.max_grad_norm,
acktr=True)
else:
raise NotImplementedError
rollouts = RolloutStorage(args.num_steps, args.num_processes,
envs.observation_space, envs.action_space,
actor_critic.recurrent_hidden_state_size)
obs = envs.reset()
actor_critic.eval()
"""
try:
writer.add_graph(actor_critic, obs)
except ValueError:
print("Unable to write model graph to tensorboard.")
"""
actor_critic.train()
for k in rollouts.obs.keys():
rollouts.obs[k][0].copy_(obs[k][0])
episode_rewards = deque(maxlen=10)
num_updates = num_env_steps // args.num_steps // args.num_processes
batch_size = args.num_steps * args.num_processes
start = time.time()
while env_step < num_env_steps:
s = time.time()
if args.use_linear_lr_decay:
# decrease learning rate linearly
utils.update_linear_schedule(
agent.optimizer, j, num_updates,
agent.optimizer.lr if args.algo == "acktr" else args.lr)
for step in range(args.num_steps):
# Sample actions
with torch.no_grad():
value, action, action_log_prob, recurrent_hidden_states, _ = actor_critic.act(
{
k: rollouts.obs[k][step].float().to(args.device)
for k in rollouts.obs.keys()
}, rollouts.recurrent_hidden_states[step].to(args.device),
rollouts.masks[step].to(args.device))
value = value.cpu()
action = action.cpu()
action_log_prob = action_log_prob.cpu()
recurrent_hidden_states = recurrent_hidden_states.cpu()
# Observe reward and next obs
obs, reward, dones, infos = envs.step(action)
for done, info in zip(dones, infos):
env_state = info['env_state'][1]
if done:
writer.add_scalar(f'train_episode_x/{env_state}',
info['max_x'], env_step)
writer.add_scalar(f'train_episode_%/{env_state}',
info['max_x'] / info['lvl_max_x'] * 100,
env_step)
writer.add_scalar(f'train_episode_r/{env_state}',
info['sum_r'], env_step)
# If done then clean the history of observations.
masks = torch.FloatTensor([[0.0] if done else [1.0]
for done in dones])
bad_masks = torch.FloatTensor(
[[0.0] if 'bad_transition' in info.keys() else [1.0]
for info in infos])
rollouts.insert(obs, recurrent_hidden_states, action,
action_log_prob, value, reward, masks, bad_masks)
with torch.no_grad():
next_value = actor_critic.get_value(
{
k: rollouts.obs[k][-1].float().to(args.device)
for k in rollouts.obs.keys()
}, rollouts.recurrent_hidden_states[-1].to(args.device),
rollouts.masks[-1].to(args.device)).detach().cpu()
rollouts.compute_returns(next_value, args.use_gae, args.gamma,
args.gae_lambda, args.use_proper_time_limits)
value_loss, action_loss, dist_entropy = agent.update(rollouts)
rollouts.after_update()
env_step += batch_size
fps = batch_size / (time.time() - s)
#res = nvidia_smi.nvmlDeviceGetUtilizationRates(handle)
#writer.add_scalar(f'gpu_usage/{writer_name}', res.gpu, env_step)
#writer.add_scalar(f'gpu_mem/{writer_name}', res.memory, env_step)
total_norm = 0
for p in list(
filter(lambda p: p.grad is not None,
actor_critic.parameters())):
param_norm = p.grad.data.norm(2)
total_norm += param_norm.item()**2
total_norm = total_norm**(1. / 2)
obs_norm = {}
for obs_name in args.obs_keys:
t_norm = 0
if obs_name == 'video':
md = actor_critic.base.video_module
elif obs_name == 'audio':
md = actor_critic.base.audio_module
else:
raise NotImplementedError
for p in list(filter(lambda p: p.grad is not None,
md.parameters())):
param_norm = p.grad.data.norm(2)
t_norm += param_norm.item()**2
obs_norm[obs_name] = t_norm**(1. / 2)
prev_env_step = max(0, env_step + 1 - batch_size)
# write training metrics for this batch, usually takes 0.003s
if (env_step + 1
) // args.write_interval > prev_env_step // args.write_interval:
writer.add_scalar(f'grad_norm/{writer_name}', total_norm, env_step)
writer.add_scalar(f'fps/{writer_name}', fps, env_step)
writer.add_scalar(f'value_loss/{writer_name}',
value_loss / batch_size, env_step)
writer.add_scalar(f'action_loss/{writer_name}',
action_loss / batch_size, env_step)
writer.add_scalar(f'dist_entropy/{writer_name}',
dist_entropy / batch_size, env_step)
writer.add_scalar(f'cpu_usage/{writer_name}', psutil.cpu_percent(),
env_step)
writer.add_scalar(f'cpu_mem/{writer_name}',
psutil.virtual_memory()._asdict()['percent'],
env_step)
for obs_name in args.obs_keys:
writer.add_scalar(f'grad_norm_{obs_name}/{writer_name}',
obs_norm[obs_name], env_step)
# print log to console
if (env_step +
1) // args.log_interval > prev_env_step // args.log_interval:
end = time.time()
print(" [log] Env step {} of {}: {:.1f}s, {:.1f}fps".format(
env_step + 1, num_env_steps, end - start, fps))
if len(episode_rewards) > 0:
print(
" Last {} episodes: mean/med reward {:.1f}/{:.1f}, min/max reward {:.1f}/{:.1f}"
.format(len(episode_rewards), np.mean(episode_rewards),
np.median(episode_rewards),
np.min(episode_rewards), np.max(episode_rewards)))
print(
" dist_entropy {:.5f}, value_loss {:.6f}, action_loss {:.6f}, grad_norm {:.6f}"
.format(dist_entropy, value_loss, action_loss, total_norm))
start = time.time()
# save model to ckpt
if ((env_step + 1) // args.save_interval >
prev_env_step // args.save_interval):
torch.save([
actor_critic,
env_step,
run_name,
], os.path.join(ckpt_save_dir, f"{run_name}-{env_step}.pt"))
print(f" [save] Saved model at step {env_step+1}.")
# save model to ckpt and run evaluation if eval_interval and not final iteration in training loop
if ((env_step + 1) // args.eval_interval >
prev_env_step // args.eval_interval
) and env_step < num_env_steps and eval_env_steps > 0:
torch.save([
actor_critic,
env_step,
run_name,
], os.path.join(ckpt_save_dir, f"{run_name}-{env_step}.pt"))
print(f" [save] Saved model at step {env_step+1}.")
envs.close()
del envs # close does not actually get rid of envs, need to del
actor_critic.eval()
eval_score, e_dict = evaluate(train_states, actor_critic,
eval_env_steps, env_step, writer,
vid_save_dir, args.vid_tb_steps,
args.vid_file_steps,
args.obs_viz_layer, args)
print(f" [eval] Evaluation score: {eval_score}")
writer.add_scalar('eval_score', eval_score, env_step)
actor_critic.train()
envs = make_vec_envs(train_states, args.seed, args.num_processes,
args.gamma, 'cpu', 'train', args)
obs = envs.reset()
# TODO: does this work? do we need to increment env step or something? whydden_states insert at 0
for k in rollouts.obs.keys():
rollouts.obs[k][0].copy_(obs[k][0])
# final model save
final_model_path = os.path.join(ckpt_save_dir, f"{run_name}-{env_step}.pt")
torch.save([
actor_critic,
env_step,
run_name,
], final_model_path)
print(
f" [save] Final model saved at step {env_step+1} to {final_model_path}"
)
# final model eval
envs.close()
del envs
eval_score = None
eval_dict = None
if eval_env_steps > 0:
eval_score, eval_dict = evaluate(train_states, actor_critic,
eval_env_steps, env_step, writer,
vid_save_dir, args.vid_tb_steps,
args.vid_file_steps,
args.obs_viz_layer, args)
print(f" [eval] Final model evaluation score: {eval_score:.3f}")
return (actor_critic, env_step, run_name), eval_score, eval_dict
def main():
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():
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,
num_frame_stack=1) # default is stacking 4 frames
save_path = os.path.join(args.save_dir, args.algo)
if os.path.isfile(os.path.join(save_path, args.env_name + ".pt")):
[start_upd, actor_critic,
agent] = torch.load(os.path.join(save_path, args.env_name + ".pt"))
print(start_upd)
print(actor_critic)
print(agent)
else:
start_upd = 0
if args.attarch:
base_kwargs = {
'recurrent': False,
"w": 7,
"h": 7,
"pad": True,
"n_f_conv1": 12,
"n_f_conv2": 24,
"att_emb_size": 64,
"n_heads": 2,
"n_att_stack": 2,
"n_fc_layers": 4,
"baseline_mode": False
}
actor_critic = Policy(envs.observation_space.shape,
envs.action_space,
base=DRRLBase,
base_kwargs=base_kwargs)
elif args.attarchbaseline:
base_kwargs = {
'recurrent': False,
"w": 7,
"h": 7,
"pad": True,
"n_f_conv1": 12,
"n_f_conv2": 24,
"baseline_mode": True,
"n_baseMods": 2
}
actor_critic = Policy(envs.observation_space.shape,
envs.action_space,
base=DRRLBase,
base_kwargs=base_kwargs)
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)
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(start_upd, num_updates): #global iteration
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
): #a batch update of num_steps for each num_process will be created in this
# loop
# 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 != "":
try:
os.makedirs(save_path)
except OSError:
pass
torch.save(
[
j,
actor_critic,
# getattr(utils.get_vec_normalize(envs), 'ob_rms', None), #input normalization not implemented with 2d
# input anyway
agent,
],
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 = 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():
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():
realEval = True #False
gettrace = getattr(sys, 'gettrace', None)
parser = argparse.ArgumentParser(description='RL')
parser.add_argument('--action-type',
type=int,
default=-1,
help='action type to play (default: -1)')
parser.add_argument('--tasks-difficulty-from',
type=int,
default=0,
help='tasks_difficulty_from')
parser.add_argument('--tasks-difficulty-to',
type=int,
default=100000,
help='tasks-difficulty-to')
parser.add_argument('--verboseLevel',
type=int,
default=5,
help='verboseLevel')
parser.add_argument('--filesNamesSuffix',
default="",
help='filesNamesSuffix')
parser.add_argument('--nobest-exit',
type=int,
default=10000,
help='nobest_exit')
args = get_args(parser)
args.algo = 'ppo'
args.env_name = 'QuadruppedWalk-v1' #'RoboschoolAnt-v1' #'QuadruppedWalk-v1' #'RoboschoolAnt-v1' #'QuadruppedWalk-v1'
args.use_gae = True
args.num_steps = 2048
#args.num_processes = 4
args.num_processes = 4
if gettrace():
args.num_processes = 1
args.lr = 0.0001
args.entropy_coef = 0.0
args.value_loss_coef = 0.5
args.ppo_epoch = 4
args.num_mini_batch = 256
args.gamma = 0.99
args.gae_lambda = 0.95
args.clip_param = 0.2
args.use_linear_lr_decay = True #True #True #True #True
args.use_proper_time_limits = True
args.save_dir = "./trained_models/" + args.env_name + "/"
args.load_dir = "./trained_models/" + args.env_name + "/"
args.log_dir = "./logs/robot"
if gettrace():
args.save_dir = "./trained_models/" + args.env_name + "debug/"
args.load_dir = "./trained_models/" + args.env_name + "debug/"
args.log_dir = "./logs/robot_d"
args.log_interval = 30
args.hidden_size = 64
args.last_hidden_size = args.hidden_size
args.recurrent_policy = False #True
args.save_interval = 20
#args.seed = 1
reward_shaping = 0.01
allowMutate = False
if args.seed == -1:
args.seed = time.clock_gettime_ns(time.CLOCK_REALTIME)
quadruppedEnv.settings.tasks_difficulty_from = args.tasks_difficulty_from
quadruppedEnv.settings.tasks_difficulty_to = args.tasks_difficulty_to
# 0 is a walk
# 1 is a balance
# 2 multitasks
# 3 multitask experiments
trainType = 14
filesNamesSuffix = ""
if args.action_type >= 0:
trainType = args.action_type
makeEnvFunction = makeEnv.make_env_with_best_settings
if trainType == 1:
filesNamesSuffix = "balance_"
makeEnvFunction = makeEnv.make_env_for_balance
if trainType == 2:
filesNamesSuffix = "analytical_"
makeEnvFunction = makeEnv.make_env_with_best_settings_for_analytical
if trainType == 3:
filesNamesSuffix = "analytical2_"
makeEnvFunction = makeEnv.make_env_with_best_settings_for_analytical2
if trainType == 4:
filesNamesSuffix = "frontback_"
makeEnvFunction = makeEnv.make_env_with_best_settings_for_front_back
if trainType == 5:
filesNamesSuffix = "leftright_"
makeEnvFunction = makeEnv.make_env_with_best_settings_for_left_right
if trainType == 6:
filesNamesSuffix = "all_"
makeEnvFunction = makeEnv.make_env_with_best_settings_for_all
if trainType == 7:
filesNamesSuffix = "rotate_"
makeEnvFunction = makeEnv.make_env_with_best_settings_for_rotate
if trainType == 8:
filesNamesSuffix = "compound_"
makeEnvFunction = make_env_multinetwork
if trainType == 9:
import pickle
realEval = False
allowMutate = False
args.use_linear_lr_decay = True #False
args.num_env_steps = 5000000
filesNamesSuffix = "test_"
makeEnvFunction = makeEnv.make_env_with_best_settings_for_test
if trainType == 10:
import pickle
realEval = False
allowMutate = False
args.use_linear_lr_decay = True #False
args.num_env_steps = 5000000
filesNamesSuffix = "zoo_"
makeEnvFunction = makeEnv.make_env_with_best_settings_for_test_zoo
if trainType == 11:
args.hidden_size = 128 #64 #128
args.last_hidden_size = args.hidden_size
import pickle
if gettrace():
args.num_processes = 1
else:
args.num_processes = 8
realEval = False
allowMutate = False
args.lr = 0.00001
args.use_linear_lr_decay = True #False
args.num_env_steps = 10000000
filesNamesSuffix = "zigote2_updown_"
print("Samples preload")
global samplesEnvData
samplesEnvData = pickle.load(
open("./QuadruppedWalk-v1_MoveNoPhys.samples", "rb"))
# samplesEnvData = pickle.load( open( "./QuadruppedWalk-v1.samples", "rb" ) )
makeEnvFunction = makeSamplesEnv
if trainType == 12:
import pickle
args.lr = 0.00001
args.hidden_size = 64
args.last_hidden_size = args.hidden_size
filesNamesSuffix = "zigote2_front_back_"
args.clip_param = 0.9
args.value_loss_coef = 0.9
makeEnvFunction = makeEnv.make_env_with_best_settings_for_train
#makeEnvFunction = makeEnv.make_env_with_best_settings_for_record
#makeEnv.samplesEnvData = pickle.load( open( "./QuadruppedWalk-v1_MoveNoPhys.samples", "rb" ) )
if trainType == 13:
filesNamesSuffix = "all_bytasks_"
makeEnvFunction = makeEnv.make_env_with_best_settings_for_all
if trainType == 14:
#args.lr = 0.00001
#args.num_env_steps = 000000
#args.clip_param = 0.5
#args.value_loss_coef =0.8
#random.seed(time.clock_gettime_ns(time.CLOCK_REALTIME))
#args.num_steps = random.choice([256,512,1024,2048,4096])
#args.num_mini_batch = random.choice([32,64,256,512])
#args.ppo_epoch = random.choice([2,4,8,10])
#args.clip_param = random.choice([0.2,0.4,0.6,0.8])
#args.value_loss_coef =random.choice([0.4,0.5,0.6,0.8])
#args.lr = random.choice([0.00001,0.0001,0.00005,0.0005])
args.num_steps = 2048
args.num_mini_batch = 64
args.ppo_epoch = 8
args.lr = 0.0001
args.hidden_size = 64
args.last_hidden_size = args.hidden_size
#
filesNamesSuffix = args.filesNamesSuffix
makeEnvFunction = makeEnv.make_env_with_best_settings_for_all
'''
num_steps: 1024 num_mini_batch 64 ppo_epoch 2
clip_param: 0.2 value_loss_coef 0.6 lr 0.0001
'''
if trainType == 15:
args.num_env_steps = 5000000
filesNamesSuffix = "zigote_updown_"
makeEnvFunction = makeEnv.make_env_with_best_settings_for_train_analytic
if trainType == 16:
args.lr = 0.00001
filesNamesSuffix = "compound_tasks_"
makeEnvFunction = make_env_multinetwork
reward_shaper = DefaultRewardsShaper(scale_value=reward_shaping)
print("ActionType ", trainType, " ", filesNamesSuffix, "seed", args.seed,
"num env steps:", args.num_env_steps, " tasks_dif",
args.tasks_difficulty_from, args.tasks_difficulty_to)
print("Num processes:", args.num_processes)
print("num_steps:", args.num_steps, "num_mini_batch", args.num_mini_batch,
"ppo_epoch", args.ppo_epoch)
print("clip_param:", args.clip_param, "value_loss_coef",
args.value_loss_coef, "lr", args.lr)
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
args.log_dir = "/tmp/tensorboard/"
#TesnorboardX
writer = SummaryWriter(log_dir=args.log_dir + 'runs/{}_PPO_{}_{}'.format(
datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S"), args.env_name,
"ppo"))
writer.add_scalar('options/num_steps', args.num_steps, 0)
writer.add_scalar('options/num_mini_batch', args.num_mini_batch, 0)
writer.add_scalar('options/ppo_epoch', args.ppo_epoch, 0)
writer.add_scalar('options/clip_param', args.clip_param, 0)
writer.add_scalar('options/value_loss_coef', args.value_loss_coef, 0)
writer.add_scalar('options/lr', args.lr, 0)
device = torch.device("cuda:0" if args.cuda else "cpu")
torch.set_num_threads(1)
load_dir = os.path.join(args.load_dir, args.algo)
multiNetworkName = ["frontback_", "all_", "leftright_", "rotate_"]
if trainType == 8:
for net in multiNetworkName:
bestFilename = os.path.join(
load_dir, "{}_{}{}_best.pt".format(args.env_name, net,
args.hidden_size))
ac, _ = torch.load(bestFilename)
policies.append(PPOPlayer(ac, device))
print("Policy multi loaded: ", bestFilename)
multiNetworkName2 = [
"all_bytasks_0_",
"all_bytasks_1_",
"all_bytasks_2_",
"all_bytasks_3_",
"all_bytasks_4_",
"all_bytasks_5_",
"all_bytasks_6_",
"all_bytasks_7_",
"all_bytasks_8_",
"all_bytasks_9_",
"all_bytasks_10_",
"all_bytasks_11_",
"all_bytasks_12_",
]
if trainType == 16:
for net in multiNetworkName2:
bestFilename = os.path.join(
load_dir, "{}_{}{}_best.pt".format(args.env_name, net,
args.hidden_size))
ac, _ = torch.load(bestFilename)
policies.append(PPOPlayer(ac, device))
print("Policy multi loaded: ", bestFilename)
envs = make_vec_envs(args.env_name,
args.seed,
args.num_processes,
args.gamma,
None,
device,
False,
normalizeOb=False,
normalizeReturns=False,
max_episode_steps=args.num_steps,
makeEnvFunc=makeEnvFunction,
num_frame_stack=1,
info_keywords=(
'episode_steps',
'episode_reward',
'progress',
'servo',
'distToTarget',
))
#print(envs.observation_space.shape,envs.action_space)
actor_critic = Policy(envs.observation_space.shape,
envs.action_space,
base_kwargs={
'recurrent': args.recurrent_policy,
'hidden_size': args.hidden_size,
'last_hidden_size': args.last_hidden_size,
'activation_layers_type': "Tanh"
})
'''
# if args.load_dir not None:
load_path = os.path.join(args.load_dir, args.algo)
actor_critic, ob_rms = torch.load(os.path.join(load_path, args.env_name + ".pt"))
'''
load_path = os.path.join(
load_dir, "{}_{}{}_best.pt".format(args.env_name, filesNamesSuffix,
args.hidden_size))
#load_path = os.path.join(load_path, "{}_{}{}.pt".format(args.env_name,filesNamesSuffix,args.hidden_size))
preptrained_path = "../Train/trained_models/QuadruppedWalk-v1/Train_QuadruppedWalk-v1_256.pth"
loadPretrained = False
if loadPretrained and os.path.isfile(preptrained_path):
print("Load preptrained")
abj = torch.load(preptrained_path)
print(abj)
print(actor_critic.base)
actor_critic.base.load_state_dict()
actor_critic.base.eval()
if os.path.isfile(load_path) and not loadPretrained:
actor_critic, ob_rms = torch.load(load_path)
actor_critic.eval()
print("----NN loaded: ", load_path, " -----")
else:
bestFilename = os.path.join(
load_dir,
"{}_{}{}_best_pretrain.pt".format(args.env_name, filesNamesSuffix,
args.hidden_size))
if os.path.isfile(bestFilename):
actor_critic, ob_rms = torch.load(bestFilename)
actor_critic.eval()
print("----NN loaded: ", bestFilename, " -----")
maxReward = -10000.0
maxSteps = 0
minDistance = 50000.0
actor_critic.to(device)
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)
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)
deque_maxLen = 10
episode_rewards = deque(maxlen=deque_maxLen)
episode_steps = deque(maxlen=deque_maxLen)
episode_rewards_alive = deque(maxlen=deque_maxLen)
episode_rewards_progress = deque(maxlen=deque_maxLen)
episode_rewards_servo = deque(maxlen=deque_maxLen)
episode_dist_to_target = deque(maxlen=deque_maxLen)
'''
load_path = os.path.join(args.load_dir, args.algo)
load_path = os.path.join(load_path, args.env_name + ".pt")
actor_critic, ob_rms = torch.load(load_path)
actor_critic.to(device)
actor_critic.eval()
#ob_rms.eval()
'''
'''
args.use_gym_monitor = 1
args.monitor_dir = "./results/"
monitor_path = os.path.join(args.monitor_dir, args.algo)
monitor_path = os.path.join(monitor_path, args.env_name)
args.
if args.use_gym_monitor:
env = wrappers.Monitor(
env, monitor_path, video_callable=False, force=True)
'''
i_episode = 0
save_path = os.path.join(args.save_dir, args.algo)
try:
os.makedirs(save_path)
except OSError:
pass
trainOnSamplesAndExit = False #False
if trainOnSamplesAndExit:
import pickle
print("---------------------------------------")
print("Samples preload")
data = pickle.load(open("./QuadruppedWalk-v1_UpDown.samples", "rb"))
#data = pickle.load( open( "../QuadruppedWalk-v1_NN.samples", "rb" ) )
learning_rate = 0.0001
max_episodes = 100
max_timesteps = 4000
betas = (0.9, 0.999)
log_interval = 1
envSamples = SamplesEnv(data)
envSamples.numSteps = max_timesteps
# create a stochastic gradient descent optimizer
optimizer = torch.optim.Adam(actor_critic.base.actor.parameters(),
lr=learning_rate,
betas=betas)
#optimizer = optim.SGD(net.parameters(), lr=learning_rate, momentum=0.9)
# create a loss function
criterion = nn.MSELoss(reduction="sum")
# run the main training loop
for epoch in range(max_episodes):
state = envSamples.reset()
time_step = 0
testReward = 0
testSteps = 0
loss_sum = 0
loss_max = 0
for t in range(max_timesteps):
time_step += 1
nn_state = torch.FloatTensor((state).reshape(1, -1)).to(device)
optimizer.zero_grad()
net_out = actor_critic.base.forwardActor(nn_state)
net_out = actor_critic.dist.fc_mean(net_out)
state, reward, done, info = envSamples.step(
net_out.detach().numpy())
sim_action = envSamples.recordedActions
sim_action_t = torch.FloatTensor([sim_action]).to(device)
loss = criterion(net_out, sim_action_t)
loss.backward()
optimizer.step()
loss_sum += loss.mean()
loss_max = max(loss_max, loss.max())
testReward += reward
testSteps += 1
if done:
if epoch % log_interval == 0:
#print(best_action_t*scaleActions-net_out*scaleActions)
if args.verboseLevel > 0:
print(
'Train Episode: {} t:{} Reward:{} Loss: mean:{:.6f} max: {:.6f}'
.format(epoch, t, testReward, loss_sum / t,
loss_max))
print(info)
reward = 0
break
bestFilename = os.path.join(
save_path,
"{}_{}{}_best_pretrain.pt".format(args.env_name, filesNamesSuffix,
args.hidden_size))
torch.save([
actor_critic,
getattr(utils.get_vec_normalize(envs), 'ob_rms', None)
], bestFilename)
exit(0)
skipWriteBest = True
if args.verboseLevel > 0:
printNetwork(actor_critic.base.actor)
lock(actor_critic, first=False, last=False)
#if trainType==9:
#allowMutate = False
#lock(actor_critic,first=True,last=False)
#mutate(actor_critic,power=0.00,powerLast=0.3)
if args.verboseLevel > 0:
printNetwork(actor_critic.base.actor)
#from torchsummary import summary
#summary(actor_critic.base.actor, (1, 48, 64))
start = time.time()
num_updates = int(
args.num_env_steps) // args.num_steps // args.num_processes
episodeBucketIndex = 0
maxReward = -10000000000
numEval = 10
if realEval:
envEval = makeEnvFunction(args.env_name)
if hasattr(envEval.env, "tasks") and len(envEval.env.tasks):
numEval = max(numEval, len(envEval.env.tasks))
maxReward = evaluate_policy(envEval,
actor_critic,
numEval * 2,
render=False,
device=device,
verbose=args.verboseLevel)
print("MaxReward on start", maxReward)
noMaxRewardCount = 0
updateIndex = 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)
episode_r = 0.0
stepsDone = 0
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)
#envs.venv.venv.venv.envs[0].render()
if args.verboseLevel > 0:
index = 0
for d in done:
if d:
print(infos[index], flush=True)
index += 1
episodeDone = False
'''
index = 0
for d in done:
if d:
print("")
print(infos[index])
index+=1
'''
for info in infos:
if 'reward' in info.keys():
episodeDone = True
i_episode += 1
episode_rewards.append(info['reward'])
writer.add_scalar('reward/episode', info['reward'],
i_episode)
#print("E:",i_episode," T:",info['episode_steps'], " R:", info['episode_reward'], " D:",info['distToTarget'])
if 'steps' in info.keys():
episode_steps.append(info['steps'])
writer.add_scalar('reward/steps', info['steps'], i_episode)
if 'alive' in info.keys():
episode_rewards_alive.append(info['alive'])
writer.add_scalar('reward/alive', info['alive'], i_episode)
if 'prog' in info.keys():
episode_rewards_progress.append(info['prog'])
writer.add_scalar('reward/progress', info['prog'],
i_episode)
if 'servo' in info.keys():
episode_rewards_servo.append(info['servo'])
writer.add_scalar('reward/servo', info['servo'], i_episode)
if 'd2T' in info.keys():
episode_dist_to_target.append(info['d2T'])
writer.add_scalar('reward/distToTarget', info['d2T'],
i_episode)
for val in info.keys():
if val not in [
"reward", "steps", "alive", "prog", "servo", "d2T",
'epos', 't'
]:
writer.add_scalar('reward/' + val, info[val],
i_episode)
#if episodeDone and i_episode%10==0:
# print(i_episode,"({:.1f}/{}/{:.2f}) ".format(episode_rewards[-1],episode_steps[-1],episode_dist_to_target[-1]),end='',flush=True)
if episodeDone:
episodeBucketIndex += 1
if args.verboseLevel > 0:
print("Mean:", Fore.WHITE, np.mean(episode_rewards),
Style.RESET_ALL, " Median:", Fore.WHITE,
np.median(episode_rewards), Style.RESET_ALL,
" max reward:", maxReward)
#'''len(episode_rewards) and np.mean(episode_rewards)>maxReward and'''
if realEval:
if episodeBucketIndex % args.log_interval == 0 and episodeBucketIndex > args.log_interval:
print("Step:",
(j + 1) * args.num_processes * args.num_steps)
if skipWriteBest == False:
evalReward = evaluate_policy(
envEval,
actor_critic,
numEval,
device=device,
verbose=args.verboseLevel)
writer.add_scalar('reward/eval', evalReward,
i_episode)
if evalReward > maxReward:
maxReward = evalReward
#maxReward = np.mean(episode_rewards)
bestFilename = os.path.join(
save_path, "{}_{}{}_best.pt".format(
args.env_name, filesNamesSuffix,
args.hidden_size))
print(
"Writing best reward:", Fore.GREEN,
"({:.1f}/{:.1f}/{:.1f}/{}/{:.2f}) ".format(
maxReward, np.mean(episode_rewards),
np.median(episode_rewards),
np.mean(episode_steps),
episode_dist_to_target[-1]),
Style.RESET_ALL, bestFilename)
torch.save([
actor_critic,
getattr(utils.get_vec_normalize(envs),
'ob_rms', None)
], bestFilename)
noMaxRewardCount = 0
else:
noMaxRewardCount += 1
if allowMutate:
if noMaxRewardCount == 5:
print("Mutation low last layer")
lock(actor_critic,
first=False,
last=False)
mutate(actor_critic,
power=0.00,
powerLast=0.01)
if noMaxRewardCount == 8:
print("Mutation low non last")
lock(actor_critic,
first=False,
last=False)
mutate(actor_critic,
power=0.01,
powerLast=0.0)
if noMaxRewardCount == 11:
print("Mutation low all")
lock(actor_critic,
first=False,
last=False)
mutate(actor_critic,
power=0.02,
powerLast=0.2)
if noMaxRewardCount == 14:
print("Mutation hi all")
lock(actor_critic,
first=False,
last=False)
mutate(actor_critic,
power=0.03,
powerLast=0.03)
noMaxRewardCount = 0
if noMaxRewardCount == args.nobest_exit:
exit(0)
else:
skipWriteBest = False
else:
if len(episode_rewards) and np.mean(
episode_rewards
) > maxReward and j > args.log_interval:
if skipWriteBest == False:
maxReward = np.mean(episode_rewards)
writer.add_scalar('reward/maxReward', maxReward,
i_episode)
bestFilename = os.path.join(
save_path, "{}_{}{}_best.pt".format(
args.env_name, filesNamesSuffix,
args.hidden_size))
if len(episode_dist_to_target):
print(
"Writing best reward:", Fore.GREEN,
"({:.1f}/{:.1f}/{}/{:.2f}) ".format(
np.mean(episode_rewards),
np.median(episode_rewards),
np.mean(episode_steps),
episode_dist_to_target[-1]),
Style.RESET_ALL, bestFilename)
else:
print(
"Writing best reward:", Fore.GREEN,
"({:.1f}/{:.1f}/{}) ".format(
np.mean(episode_rewards),
np.median(episode_rewards),
np.mean(episode_steps)),
Style.RESET_ALL, bestFilename)
torch.save([
actor_critic,
getattr(utils.get_vec_normalize(envs),
'ob_rms', None)
], bestFilename)
else:
skipWriteBest = False
# 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])
shaped_reward = reward_shaper(reward)
rollouts.insert(obs, recurrent_hidden_states, action,
action_log_prob, value, shaped_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)
writer.add_scalar('reward/value_loss', value_loss, updateIndex)
writer.add_scalar('reward/action_loss', action_loss, updateIndex)
writer.add_scalar('reward/dist_entropy', dist_entropy, updateIndex)
updateIndex += 1
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 != "":
'''
fileName = os.path.join(save_path, "{}_{}{}.pt".format(args.env_name,filesNamesSuffix,args.hidden_size))
torch.save([
actor_critic,
getattr(utils.get_vec_normalize(envs), 'ob_rms', None)
], fileName)
print("Saved:",fileName, " cur avg rewards:",np.mean(episode_rewards))
fileName = os.path.join(save_path, "{}_{}{}_actor.pt".format(args.env_name,filesNamesSuffix,args.hidden_size))
torch.save(actor_critic.state_dict, fileName)
print("Saved:",fileName)
'''
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()
if args.verboseLevel > 0:
print("")
print("Updates {}, num timesteps {}, FPS {}".format(
j, total_num_steps, int(total_num_steps / (end - start))))
print(" Last {} training episodes:".format(
len(episode_rewards)))
print(
" reward mean/median {:.1f}/{:.1f} min/max {:.1f}/{:.1f}".
format(np.mean(episode_rewards),
np.median(episode_rewards), np.min(episode_rewards),
np.max(episode_rewards)))
print(" steps mean/median {:.1f}/{:.1f} min/max {:.1f}/{:.1f}".
format(np.mean(episode_steps), np.median(episode_steps),
np.min(episode_steps), np.max(episode_steps)))
if len(episode_rewards_alive):
print(
" alive mean/median {:.1f}/{:.1f} min/max {:.1f}/{:.1f}"
.format(np.mean(episode_rewards_alive),
np.median(episode_rewards_alive),
np.min(episode_rewards_alive),
np.max(episode_rewards_alive)))
if len(episode_rewards_progress):
print(
" progress mean/median {:.1f}/{:.1f} min/max {:.1f}/{:.1f}"
.format(np.mean(episode_rewards_progress),
np.median(episode_rewards_progress),
np.min(episode_rewards_progress),
np.max(episode_rewards_progress)))
if len(episode_rewards_servo):
print(
" servo mean/median {:.1f}/{:.1f} min/max {:.1f}/{:.1f}"
.format(np.mean(episode_rewards_servo),
np.median(episode_rewards_servo),
np.min(episode_rewards_servo),
np.max(episode_rewards_servo)))
if len(episode_dist_to_target):
print(
" dist to target mean/median {:.3f}/{:.3f} min/max {:.3f}/{:.3f}"
.format(np.mean(episode_dist_to_target),
np.median(episode_dist_to_target),
np.min(episode_dist_to_target),
np.max(episode_dist_to_target)))
print(
" Reward/Steps {:.3f} Progress/Steps: {:.3f} entropy {:.1f} value_loss {:.5f} action_loss {:.5f}\n"
.format(
np.mean(episode_rewards) / np.mean(episode_steps),
(0 if len(episode_rewards_progress) == 0 else
np.mean(episode_rewards_progress) /
np.mean(episode_steps)), dist_entropy, value_loss,
action_loss))
def train_ppo_fine_tune_joint(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(2)
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, True)
actor_critic = Policy( # 2-layer fully connected network
envs.observation_space.shape,
envs.action_space,
base_kwargs={
'recurrent': False,
'hidden_size': 32
})
# behavioral cloning
model = PusherPolicyModel()
num_epochs = 20
model.train(num_epochs=num_epochs)
actor_critic.base.actor[0].weight.data.copy_(model.net.fc1.weight.data)
actor_critic.base.actor[0].bias.data.copy_(model.net.fc1.bias.data)
actor_critic.base.actor[2].weight.data.copy_(model.net.fc2.weight.data)
actor_critic.base.actor[2].bias.data.copy_(model.net.fc2.bias.data)
actor_critic.base.critic[0].weight.data.copy_(model.net.fc1.weight.data)
actor_critic.base.critic[0].bias.data.copy_(model.net.fc1.bias.data)
actor_critic.base.critic[2].weight.data.copy_(model.net.fc2.weight.data)
actor_critic.base.critic[2].bias.data.copy_(model.net.fc2.bias.data)
actor_critic.dist.fc_mean.weight.data.copy_(model.net.fc3.weight.data)
actor_critic.dist.fc_mean.bias.data.copy_(model.net.fc3.bias.data)
actor_critic.to(device)
dataset = np.load('./expert.npz')
obs_expert = torch.Tensor(dataset['obs'])
actions_expert = torch.Tensor(dataset['action'])
obs_expert.to(device)
actions_expert.to(device)
joint_loss_coef = 0.03
agent = PPOJointLoss(actor_critic,
args.clip_param,
args.ppo_epoch,
args.num_mini_batch,
args.value_loss_coef,
args.entropy_coef,
joint_loss_coef=joint_loss_coef,
obs_expert=obs_expert,
actions_expert=actions_expert,
lr=args.lr,
eps=args.eps,
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)
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
episode_reward_means = []
episode_reward_times = []
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)
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)
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))
episode_reward_means.append(np.mean(episode_rewards))
episode_reward_times.append(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)
print(episode_reward_means, episode_reward_times)
return episode_reward_means, episode_reward_times
def main():
import copy
import glob
import os
import time
from collections import deque
import gym
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from a2c_ppo_acktr import algo
from a2c_ppo_acktr.envs import make_vec_envs
from a2c_ppo_acktr.storage import RolloutStorage
from a2c_ppo_acktr.utils import get_vec_normalize, update_linear_schedule
from a2c_ppo_acktr.visualize import visdom_plot
device = torch.device('cuda')
envs = make_vec_envs(args.env_name, args.seed, args.num_processes,
args.gamma, args.log_dir, args.add_timestep, device,
False)
print(envs.observation_space.shape)
actor_critic = Policy(envs.observation_space.shape, envs.action_space)
actor_critic.to(device)
agent = ProximalPolicyOptimization(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,
chrono=chrono)
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)
num_updates = int(
args.num_env_steps) // args.num_steps // args.num_processes
start = time.time()
for j in range(args.repeat):
with chrono.time('train', verbose=True) as t:
for n in range(args.number):
with chrono.time('one_batch', verbose=True):
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))
with chrono.time('generate_rollouts', verbose=True):
generate_rollouts(**locals())
with torch.no_grad():
next_value = actor_critic.get_value(
rollouts.obs[-1],
rollouts.recurrent_hidden_states[-1],
rollouts.masks[-1]).detach()
# ---
with chrono.time('compute_returns', verbose=True):
rollouts.compute_returns(next_value, args.use_gae,
args.gamma, args.tau)
with chrono.time('agent.update',
verbose=True): # 11.147009023304644
value_loss, action_loss, dist_entropy = agent.update(
rollouts)
#exp.log_batch_loss(action_loss)
#exp.log_metric('value_loss', value_loss)
with chrono.time('after_update', verbose=True):
rollouts.after_update()
total_num_steps = (j +
1) * args.num_processes * args.num_steps
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(args.save_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_vision_noise(full_obs, 0)
elif obs_noisy:
obs = add_joint_noise(full_obs)
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)
# 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 args.pos_control:
# print("main step_skip",args.step_skip)
if step % (512 / args.step_skip - 1) == 0:
current_state = initial_state
next_action = current_state + next_action
for kk in range(args.step_skip):
full_obs, reward, done, infos = envs.step(next_action)
current_state = full_obs[:, :envs.action_space.shape[0]]
else:
for kk in range(args.step_skip):
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_vision_noise(full_obs, j)
elif obs_noisy:
obs = add_joint_noise(full_obs)
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()
# Get total number of timesteps
total_num_steps = (j + 1) * args.num_processes * args.num_steps
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:
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):
opening_rate, opening_timeavg = onpolicy_inference(
seed=args.seed,
env_name=args.env_name,
det=True,
load_name=args.save_name,
evaluation=True,
render=False,
knob_noisy=args.knob_noisy,
visionnet_input=args.visionnet_input,
env_kwargs=env_kwargs_val,
actor_critic=actor_critic,
verbose=False,
pos_control=args.pos_control,
step_skip=args.step_skip)
print(
"{}th update. {}th timestep. opening rate {}%. Average time to open is {}."
.format(j, total_num_steps, opening_rate, opening_timeavg))
writer.add_scalar("Opening rate per envstep", opening_rate,
total_num_steps)
writer.add_scalar("Opening rate per update", opening_rate, j)
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():
torch.set_num_threads(1)
device = torch.device("cuda:0" if args.cuda else "cpu")
print(device)
print(save_folder)
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, args.reward_type)
actor_critic = Policy(envs.observation_space.shape, envs.action_space)
actor_critic.to(device)
curiosity = None
if use_curiosity:
curiosity = ICM(envs.observation_space.shape[0], envs.action_space.n)
curiosity.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,
curiosity,
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_curiosity=use_curiosity)
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()
cum_rew = [0] * args.num_processes
rollouts.obs[0].copy_(obs)
rollouts.to(device)
episode_rewards = deque(maxlen=args.num_processes * 2)
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 = agent.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)
envs.render()
cur_reward = reward
to_write = reward.cpu().numpy()
for i in range(args.num_processes):
cum_rew[i] += to_write[i][0]
if use_curiosity:
action_one_hot = (torch.eye(14)[action]).view(-1, 14).cuda()
_, pred_phi, actual_phi = curiosity(
(rollouts.obs[step], obs, action_one_hot))
cur_reward += 0.2 * ((pred_phi - actual_phi).pow(2)).sum(
-1, keepdim=True).cpu() / 2
for i, finished in enumerate(done):
if finished:
percentile = infos[i]['x_pos'] / norm_pos
episode_rewards.append(percentile)
print(cum_rew[i])
with open(train_file[:-4] + str(i) + train_file[-4:],
'a',
newline='') as sfile:
writer = csv.writer(sfile)
writer.writerows([[cum_rew[i], percentile]])
cum_rew[i] = 0
# 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, cur_reward.detach(), masks)
with torch.no_grad():
next_value = agent.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 = agent.actor_critic
if args.cuda:
save_model = copy.deepcopy(agent.actor_critic).cpu()
save_model = [
save_model,
getattr(get_vec_normalize(envs), 'ob_rms', None)
]
torch.save(save_model,
os.path.join(save_folder, '/' + 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) > args.num_processes:
end = time.time()
print(
"Updates {}, num timesteps {}, FPS {} \n Last {} training episodes: mean/median reward {:.1f}/{:.1f}, min/max reward {:.1f}/{:.1f}, cumulative reward {:.3f}\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), np.mean(cum_rew)))
#Evaluation time :
if (args.eval_interval is not None and len(episode_rewards) > 1
and j % args.eval_interval == 0):
num_proc = 1
eval_envs = make_vec_envs(args.env_name, args.seed + num_proc,
num_proc, args.gamma, args.log_dir,
args.add_timestep, device, True,
args.reward_type)
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 = []
test_rew = 0
finish_this = False
obs = eval_envs.reset()
eval_recurrent_hidden_states = torch.zeros(
num_proc,
actor_critic.recurrent_hidden_state_size,
device=device)
eval_masks = torch.zeros(num_proc, 1, device=device)
positions = deque(maxlen=400)
while not finish_this:
with torch.no_grad():
_, action, _, eval_recurrent_hidden_states = agent.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_envs.render()
eval_masks = torch.FloatTensor([[0.0] if done_ else [1.0]
for done_ in done]).cuda()
# for i, finished in enumerate(done):
# if finished:
# percentile = infos[i]['x_pos']/norm_pos
# eval_episode_rewards.append(percentile)
# with open(eval_file, 'a', newline='') as sfile:
# writer = csv.writer(sfile)
# writer.writerows([[percentile]])
test_rew += reward.cpu().numpy()[0, 0]
for i, finished in enumerate(done):
if finished:
print('he died')
percentile = infos[i]['x_pos'] / norm_pos
eval_episode_rewards.append(percentile)
with open(eval_file, 'a', newline='') as sfile:
writer = csv.writer(sfile)
writer.writerows([[test_rew, percentile]])
finish_this = True
#to prevent the agent from getting stuck
positions.append(infos[0]['x_pos'])
pos_ar = np.array(positions)
if (len(positions) >= 200) and (pos_ar < pos_ar[-1] + 20).all(
) and (pos_ar > pos_ar[-1] - 20).all():
print("he's stuck")
percentile = infos[0]['x_pos'] / norm_pos
eval_episode_rewards.append(percentile)
with open(eval_file, 'a', newline='') as sfile:
writer = csv.writer(sfile)
writer.writerows([[test_rew, percentile]])
finish_this = True
eval_envs.close()
positions.clear()
print(
" Evaluation using {} episodes: reward {:.3f}, distance {:.3f}\n"
.format(len(eval_episode_rewards), test_rew,
np.mean(eval_episode_rewards)))
test_rew = 0
finish_this = False
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
