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)