agent_traj.add_epis(epis)
agent_traj = ef.add_next_obs(agent_traj)
agent_traj = ef.compute_pseudo_rews(
agent_traj,
rew_giver=rewf if args.rew_type == 'rew' else advf,
state_only=True if args.rew_type == 'rew' else False)
agent_traj = ef.compute_vs(agent_traj, vf)
agent_traj = ef.compute_rets(agent_traj, args.gamma)
agent_traj = ef.compute_advs(agent_traj, args.gamma, args.lam)
agent_traj = ef.centerize_advs(agent_traj)
agent_traj = ef.compute_h_masks(agent_traj)
agent_traj.register_epis()
if args.data_parallel:
pol.dp_run = True
vf.dp_run = True
if args.rew_type == 'rew':
rewf.dp_run = True
shaping_vf.dp_run = True
elif args.rew_type == 'adv':
advf.dp_run = True
if args.rl_type == 'trpo':
result_dict = airl.train(
agent_traj,
expert_traj,
pol,
vf,
optim_vf,
optim_discrim,
rewf=rewf,
with measure('sample'):
epis = sampler.sample(pol, max_steps=args.max_steps_per_iter)
with measure('train'):
traj = Traj()
traj.add_epis(epis)
traj = ef.compute_vs(traj, vf)
traj = ef.compute_rets(traj, args.gamma)
traj = ef.compute_advs(traj, args.gamma, args.lam)
traj = ef.centerize_advs(traj)
traj = ef.compute_h_masks(traj)
traj.register_epis()
if args.data_parallel:
pol.dp_run = True
vf.dp_run = True
if args.ppo_type == 'clip':
result_dict = ppo_clip.train(traj=traj,
pol=pol,
vf=vf,
clip_param=args.clip_param,
optim_pol=optim_pol,
optim_vf=optim_vf,
epoch=args.epoch_per_iter,
batch_size=args.batch_size if
not args.rnn else args.rnn_batch_size,
max_grad_norm=args.max_grad_norm)
else:
result_dict = ppo_kl.train(traj=traj,
pol=pol,
with measure('sample'):
epis = sampler.sample(pol, max_steps=args.max_steps_per_iter)
with measure('train'):
traj = Traj()
traj.add_epis(epis)
traj = ef.compute_vs(traj, vf)
traj = ef.compute_rets(traj, args.gamma)
traj = ef.compute_advs(traj, args.gamma, args.lam)
traj = ef.centerize_advs(traj)
traj = ef.compute_h_masks(traj)
traj.register_epis()
if args.data_parallel:
pol.dp_run = True
vf.dp_run = True
result_dict = ppo_clip.train(traj=traj, pol=pol, vf=vf, clip_param=args.clip_param,
optim_pol=optim_pol, optim_vf=optim_vf, epoch=args.epoch_per_iter, batch_size=args.batch_size if not args.rnn else args.rnn_batch_size, max_grad_norm=args.max_grad_norm)
total_epi += traj.num_epi
step = traj.num_step
total_step += step
rewards = [np.sum(epi['rews']) for epi in epis]
mean_rew = np.mean(rewards)
logger.record_results(args.log, result_dict, score_file,
total_epi, step, total_step,
rewards,
plot_title=args.env_name)
writer.add_scalar('rewards', mean_rew, total_epi)
