def test_critique(self):
network = CriticNetwork(hidden_layer_dims=[2])
rx = np.zeros((10,2))
ra = np.zeros((10,1))
dataset = CriticDataset(dict(observations=rx, actions=ra), batch_size=10)
with tf.Session() as session:
critic = WassersteinCritic(
obs_dim=1,
act_dim=1,
dataset=dataset,
network=network
)
session.run(tf.global_variables_initializer())
paths = [
dict(observations=[[1],[2]], actions=[[1],[2]], rewards=[[1],[2]]),
dict(observations=[[1],[2],[3]], actions=[[1],[2],[3]], rewards=[[1],[2],[3]]),
dict(observations=[[1]], actions=[[1]], rewards=[[1]]),
]
rewards = critic.critique(1, paths)
self.assertTrue(len(rewards[0]) == 2)
self.assertTrue(len(rewards[1]) == 3)
self.assertTrue(len(rewards[2]) == 1)
def build_critic(args, data, env, writer=None):
if args.use_critic_replay_memory:
critic_replay_memory = hgail.misc.utils.KeyValueReplayMemory(maxsize=3 * args.batch_size)
else:
critic_replay_memory = None
critic_dataset = CriticDataset(
data,
replay_memory=critic_replay_memory,
batch_size=args.critic_batch_size,
flat_recurrent=args.policy_recurrent
)
critic_network = ObservationActionMLP(
name='critic',
hidden_layer_dims=args.critic_hidden_layer_dims,
dropout_keep_prob=args.critic_dropout_keep_prob
)
critic = WassersteinCritic(
obs_dim=env.observation_space.flat_dim,
act_dim=env.action_space.flat_dim,
dataset=critic_dataset,
network=critic_network,
gradient_penalty=args.gradient_penalty,
optimizer=tf.train.RMSPropOptimizer(args.critic_learning_rate),
n_train_epochs=args.n_critic_train_epochs,
summary_writer=writer,
grad_norm_rescale=args.critic_grad_rescale,
verbose=2,
debug_nan=True
)
return critic
def train_gail(session,
env,
dataset,
obs_dim=1,
act_dim=2,
n_itr=20,
use_env_rewards=False,
discount=.99,
batch_size=4000,
critic_scale=1.,
gail_step_size=.01,
critic_learning_rate=.001,
policy_hid_layer_dims=[32, 32],
gradient_penalty=.1,
critic_n_train_epochs=1,
sampler_args=dict(),
return_algo=False):
network = CriticNetwork(hidden_layer_dims=[32, 32])
critic = WassersteinCritic(obs_dim=obs_dim,
act_dim=act_dim,
dataset=dataset,
network=network,
verbose=2,
gradient_penalty=gradient_penalty,
optimizer=tf.train.AdamOptimizer(
critic_learning_rate, beta1=.5, beta2=.9),
n_train_epochs=critic_n_train_epochs)
policy = CategoricalMLPPolicy(name="policy",
env_spec=env.spec,
hidden_sizes=policy_hid_layer_dims)
baseline = LinearFeatureBaseline(env_spec=env.spec)
reward_handler = RewardHandler(use_env_rewards=use_env_rewards,
critic_final_scale=critic_scale)
algo = GAIL(critic=critic,
reward_handler=reward_handler,
env=env,
policy=policy,
baseline=baseline,
batch_size=batch_size,
max_path_length=200,
n_itr=n_itr,
discount=discount,
step_size=gail_step_size,
sampler_args=sampler_args)
session.run(tf.global_variables_initializer())
if return_algo:
return algo
algo.train(sess=session)
return policy, critic
def test_train(self):
with tf.Session() as session:
network = CriticNetwork(hidden_layer_dims=[24])
batch_size = 10
obs_dim = 2
act_dim = 1
real_data = dict(
observations=np.ones((batch_size, obs_dim)) * .5,
actions=np.ones((batch_size, act_dim)) * .5)
fake_data = dict(
observations=np.ones((batch_size, obs_dim)) * -.5,
actions=np.ones((batch_size, act_dim)) * -.5)
dataset = CriticDataset(real_data, batch_size=batch_size)
critic = WassersteinCritic(
obs_dim=obs_dim,
act_dim=act_dim,
dataset=dataset,
network=network,
gradient_penalty=.01
)
session.run(tf.global_variables_initializer())
n_epochs = 500
for epoch in range(n_epochs):
critic.train(epoch, fake_data)
real_rewards = critic.network.forward(
real_data['observations'], real_data['actions'])
fake_rewards = critic.network.forward(
fake_data['observations'], fake_data['actions'])
self.assertTrue(real_rewards[0] > 1)
self.assertTrue(fake_rewards[0] < -1)
# summary writer
summary_writer = tf.summary.FileWriter(
os.path.join(exp_dir, 'imitate', 'summaries'))
# build the critic
critic_network = ObservationActionMLP(
name='critic',
hidden_layer_dims=[64, 64],
dropout_keep_prob=critic_dropout_keep_prob)
critic = WassersteinCritic(
obs_dim=env.observation_space.flat_dim,
act_dim=env.action_space.n,
dataset=critic_dataset,
network=critic_network,
gradient_penalty=1.,
optimizer=tf.train.RMSPropOptimizer(critic_learning_rate),
n_train_epochs=n_critic_train_epochs,
summary_writer=summary_writer,
grad_norm_rescale=50.,
verbose=2,
)
if use_infogail:
# recognition model
recognition_dataset = RecognitionDataset(batch_size)
recognition_network = ObservationActionMLP(name='recog',
hidden_layer_dims=[32, 32],
output_dim=latent_dim)
recognition_model = RecognitionModel(
obs_dim=env.observation_space.flat_dim,
act_dim=env.action_space.n,
def build_hgail(env, critic_dataset, batch_size):
# critic
with tf.variable_scope('critic'):
critic_network = ObservationActionMLP(name='critic',
hidden_layer_dims=[32, 32])
critic = WassersteinCritic(obs_dim=3,
act_dim=2,
dataset=critic_dataset,
network=critic_network,
gradient_penalty=.01,
optimizer=tf.train.AdamOptimizer(.001,
beta1=.5,
beta2=.9),
n_train_epochs=50)
# base latent variable sampler
base_latent_sampler = UniformlyRandomLatentSampler(
scheduler=ConstantIntervalScheduler(),
name='base_latent_sampler',
dim=3)
with tf.variable_scope('level_1'):
recog_dataset_1 = RecognitionDataset(batch_size=batch_size)
recog_network_1 = ObservationActionMLP(name='recog',
hidden_layer_dims=[32, 32],
output_dim=3)
recog_1 = RecognitionModel(obs_dim=3,
act_dim=2,
dataset=recog_dataset_1,
network=recog_network_1,
variable_type='categorical',
latent_dim=3,
name='recognition_1')
latent_sampler = CategoricalLatentSampler(
scheduler=ConstantIntervalScheduler(k=1),
name='latent_sampler',
policy_name='latent_sampler_policy',
dim=2,
env_spec=env.spec,
latent_sampler=base_latent_sampler,
max_n_envs=20)
baseline_1 = LinearFeatureBaseline(env_spec=env.spec)
algo_1 = TRPO(
env=env,
policy=latent_sampler,
baseline=baseline_1,
batch_size=4000,
max_path_length=100,
n_itr=15,
discount=0.99,
step_size=0.01,
sampler_cls=HierarchySampler,
)
reward_handler_1 = RewardHandler(use_env_rewards=False,
critic_final_scale=1.)
level_1 = Level(depth=1,
algo=algo_1,
reward_handler=reward_handler_1,
recognition_model=recog_1)
with tf.variable_scope('level_0'):
# recognition model
recog_dataset_0 = RecognitionDataset(batch_size=batch_size)
recog_network_0 = ObservationActionMLP(name='recog',
hidden_layer_dims=[32, 32],
output_dim=2)
recog_0 = RecognitionModel(obs_dim=3,
act_dim=2,
dataset=recog_dataset_0,
network=recog_network_0,
variable_type='categorical',
latent_dim=2,
name='recognition_0')
policy = CategoricalLatentVarMLPPolicy(policy_name="policy",
latent_sampler=latent_sampler,
env_spec=env.spec)
baseline_0 = LinearFeatureBaseline(env_spec=env.spec)
algo_0 = TRPO(env=env,
policy=policy,
baseline=baseline_0,
batch_size=4000,
max_path_length=100,
n_itr=5,
discount=0.99,
step_size=0.1,
sampler_args=dict(n_envs=1))
reward_handler_0 = RewardHandler(use_env_rewards=False,
critic_final_scale=1.)
level_0 = Level(depth=0,
algo=algo_0,
reward_handler=reward_handler_0,
recognition_model=recog_0)
hierarchy = [level_0, level_1]
algo = HGAIL(
critic=critic,
hierarchy=hierarchy,
)
return algo
def test_infogail_two_round_stochastic_env(self):
env = TfEnv(TwoRoundNondeterministicRewardEnv())
# dataset of one-hot obs and acts
# optimal actions: 0, 1
# first state
n_expert_samples = 1000
batch_size = 1000
half = int(n_expert_samples / 2)
rx = np.zeros((n_expert_samples, 3))
rx[:half, 2] = 1
rx[half:, 0] = 1
ra = np.zeros((n_expert_samples, 2))
ra[:half, 0] = 1
ra[half:, 1] = 1
with tf.Session() as session:
# critic
critic_dataset = CriticDataset(dict(observations=rx, actions=ra),
batch_size=batch_size)
critic_network = ObservationActionMLP(name='critic',
hidden_layer_dims=[32, 32])
critic = WassersteinCritic(obs_dim=3,
act_dim=2,
dataset=critic_dataset,
network=critic_network,
gradient_penalty=.01,
optimizer=tf.train.AdamOptimizer(
.001, beta1=.5, beta2=.9),
n_train_epochs=50)
# recognition model
recog_dataset = RecognitionDataset(batch_size=batch_size)
recog_network = ObservationActionMLP(name='recog',
hidden_layer_dims=[32, 32],
output_dim=2)
recog = RecognitionModel(obs_dim=3,
act_dim=2,
dataset=recog_dataset,
network=recog_network,
variable_type='categorical',
latent_dim=2)
# policy
env.spec.num_envs = 10
latent_sampler = UniformlyRandomLatentSampler(
scheduler=ConstantIntervalScheduler(),
name='latent_sampler',
dim=2)
policy = CategoricalLatentVarMLPPolicy(
policy_name="policy",
latent_sampler=latent_sampler,
env_spec=env.spec)
# gail
reward_handler = RewardHandler(use_env_rewards=False,
critic_final_scale=1.)
baseline = LinearFeatureBaseline(env_spec=env.spec)
algo = GAIL(critic=critic,
recognition=recog,
reward_handler=reward_handler,
env=env,
policy=policy,
baseline=baseline,
batch_size=4000,
max_path_length=200,
n_itr=15,
discount=.99,
step_size=.01,
sampler_args=dict(n_envs=env.spec.num_envs))
session.run(tf.global_variables_initializer())
# run it!
algo.train(sess=session)
# evaluate
l0_state_infos = dict(latent=[[1, 0]])
l0_dist_2 = policy.dist_info([[0., 0., 1.]],
l0_state_infos)['prob']
l0_dist_0 = policy.dist_info([[1., 0., 0.]],
l0_state_infos)['prob']
l1_state_infos = dict(latent=[[0, 1]])
l1_dist_2 = policy.dist_info([[0., 0., 1.]],
l1_state_infos)['prob']
l1_dist_0 = policy.dist_info([[1., 0., 0.]],
l1_state_infos)['prob']
np.testing.assert_array_almost_equal(l0_dist_2, [[1, 0]], 1)
np.testing.assert_array_almost_equal(l0_dist_0, [[0, 1]], 1)
np.testing.assert_array_almost_equal(l1_dist_2, [[1, 0]], 1)
np.testing.assert_array_almost_equal(l1_dist_0, [[0, 1]], 1)