def test_learning(self):
pol_net = PolNet(self.env.observation_space,
self.env.action_space,
h1=32,
h2=32)
pol = GaussianPol(self.env.observation_space, self.env.action_space,
pol_net)
vf_net = VNet(self.env.observation_space)
vf = DeterministicSVfunc(self.env.observation_space, vf_net)
discrim_net = DiscrimNet(self.env.observation_space,
self.env.action_space,
h1=32,
h2=32)
discrim = DeterministicSAVfunc(self.env.observation_space,
self.env.action_space, discrim_net)
sampler = EpiSampler(self.env, pol, num_parallel=1)
optim_vf = torch.optim.Adam(vf_net.parameters(), 3e-4)
optim_discrim = torch.optim.Adam(discrim_net.parameters(), 3e-4)
with open(os.path.join('data/expert_epis', 'Pendulum-v0_2epis.pkl'),
'rb') as f:
expert_epis = pickle.load(f)
expert_traj = Traj()
expert_traj.add_epis(expert_epis)
expert_traj.register_epis()
epis = sampler.sample(pol, max_steps=32)
agent_traj = Traj()
agent_traj.add_epis(epis)
agent_traj = ef.compute_pseudo_rews(agent_traj, discrim)
agent_traj = ef.compute_vs(agent_traj, vf)
agent_traj = ef.compute_rets(agent_traj, 0.99)
agent_traj = ef.compute_advs(agent_traj, 0.99, 0.95)
agent_traj = ef.centerize_advs(agent_traj)
agent_traj = ef.compute_h_masks(agent_traj)
agent_traj.register_epis()
result_dict = gail.train(agent_traj,
expert_traj,
pol,
vf,
discrim,
optim_vf,
optim_discrim,
rl_type='trpo',
epoch=1,
batch_size=32,
discrim_batch_size=32,
discrim_step=1,
pol_ent_beta=1e-3,
discrim_ent_beta=1e-5)
del sampler
parallel_dim=1 if args.rnn else 0)
else:
raise ValueError('Only Box, Discrete, and MultiDiscrete are supported')
if args.rnn:
vf_net = VNetLSTM(ob_space, h_size=256, cell_size=256)
else:
vf_net = VNet(ob_space)
vf = DeterministicSVfunc(ob_space,
vf_net,
args.rnn,
data_parallel=args.data_parallel,
parallel_dim=1 if args.rnn else 0)
discrim_net = DiscrimNet(ob_space,
ac_space,
h1=args.discrim_h1,
h2=args.discrim_h2)
discrim = DeterministicSAVfunc(ob_space,
ac_space,
discrim_net,
data_parallel=args.data_parallel)
sampler = EpiSampler(env, pol, num_parallel=args.num_parallel, seed=args.seed)
optim_pol = torch.optim.Adam(pol_net.parameters(), args.pol_lr)
optim_vf = torch.optim.Adam(vf_net.parameters(), args.vf_lr)
optim_discrim = torch.optim.Adam(discrim_net.parameters(), args.discrim_lr)
with open(os.path.join(args.expert_dir, args.expert_fname), 'rb') as f:
expert_epis = pickle.load(f)
expert_traj = Traj()
vf = DeterministicSVfunc(observation_space, vf_net)
if args.rew_type == 'rew':
rewf_net = VNet(observation_space, h1=args.discrim_h1, h2=args.discrim_h2)
rewf = DeterministicSVfunc(observation_space, rewf_net)
shaping_vf_net = VNet(observation_space,
h1=args.discrim_h1,
h2=args.discrim_h2)
shaping_vf = DeterministicSVfunc(observation_space, shaping_vf_net)
optim_discrim = torch.optim.Adam(
list(rewf_net.parameters()) + list(shaping_vf_net.parameters()),
args.discrim_lr)
advf = None
elif args.rew_type == 'adv':
advf_net = DiscrimNet(observation_space,
action_space,
h1=args.discrim_h1,
h2=args.discrim_h2)
advf = DeterministicSAVfunc(observation_space, action_space, advf_net)
optim_discrim = torch.optim.Adam(advf_net.parameters(), args.discrim_lr)
rewf = None
shaping_vf = None
else:
raise ValueError('Only rew and adv are supported')
sampler = EpiSampler(env, pol, num_parallel=args.num_parallel, seed=args.seed)
optim_pol = torch.optim.Adam(pol_net.parameters(), args.pol_lr)
optim_vf = torch.optim.Adam(vf_net.parameters(), args.vf_lr)
with open(os.path.join(args.expert_dir, args.expert_fname), 'rb') as f:
expert_epis = pickle.load(f)