def test_categorial_mlp_policy(self):
categorical_mlp_policy = CategoricalMLPPolicy(env_spec=self.env,
hidden_sizes=(1, ))
self.sess.run(tf.compat.v1.global_variables_initializer())
obs = self.env.observation_space.high
assert categorical_mlp_policy.get_action(obs)
def test_get_regularizable_vars(self, obs_dim, action_dim):
env = TfEnv(DummyDiscreteEnv(obs_dim=obs_dim, action_dim=action_dim))
policy = CategoricalMLPPolicy(env_spec=env.spec)
reg_vars = policy.get_regularizable_vars()
assert len(reg_vars) == 2
for var in reg_vars:
assert ('bias' not in var.name) and ('output' not in var.name)
def fixture_exp(snapshot_config, sess):
"""Dummy fixture experiment function.
Args:
snapshot_config (metarl.experiment.SnapshotConfig): The snapshot
configuration used by LocalRunner to create the snapshotter.
If None, it will create one with default settings.
sess (tf.Session): An optional TensorFlow session.
A new session will be created immediately if not provided.
Returns:
np.ndarray: Values of the parameters evaluated in
the current session
"""
with LocalTFRunner(snapshot_config=snapshot_config, sess=sess) as runner:
env = MetaRLEnv(env_name='CartPole-v1')
policy = CategoricalMLPPolicy(name='policy',
env_spec=env.spec,
hidden_sizes=(8, 8))
baseline = LinearFeatureBaseline(env_spec=env.spec)
algo = VPG(env_spec=env.spec,
policy=policy,
baseline=baseline,
max_path_length=100,
discount=0.99,
optimizer_args=dict(learning_rate=0.01, ))
runner.setup(algo, env)
runner.train(n_epochs=5, batch_size=100)
return policy.get_param_values()
def test_is_pickleable(self, obs_dim, action_dim):
env = MetaRLEnv(
DummyDiscreteEnv(obs_dim=obs_dim, action_dim=action_dim))
obs_var = tf.compat.v1.placeholder(
tf.float32,
shape=[None, None, env.observation_space.flat_dim],
name='obs')
policy = CategoricalMLPPolicy(env_spec=env.spec)
policy.build(obs_var)
obs = env.reset()
with tf.compat.v1.variable_scope(
'CategoricalMLPPolicy/CategoricalMLPModel', reuse=True):
bias = tf.compat.v1.get_variable('mlp/hidden_0/bias')
# assign it to all one
bias.load(tf.ones_like(bias).eval())
output1 = self.sess.run(
[policy.distribution.probs],
feed_dict={policy.model.input: [[obs.flatten()]]})
p = pickle.dumps(policy)
with tf.compat.v1.Session(graph=tf.Graph()) as sess:
policy_pickled = pickle.loads(p)
obs_var = tf.compat.v1.placeholder(
tf.float32,
shape=[None, None, env.observation_space.flat_dim],
name='obs')
policy_pickled.build(obs_var)
output2 = sess.run(
[policy_pickled.distribution.probs],
feed_dict={policy_pickled.model.input: [[obs.flatten()]]})
assert np.array_equal(output1, output2)
def test_get_regularizable_vars(self, obs_dim, action_dim):
env = MetaRLEnv(
DummyDiscreteEnv(obs_dim=obs_dim, action_dim=action_dim))
obs_var = tf.compat.v1.placeholder(
tf.float32,
shape=[None, None, env.observation_space.flat_dim],
name='obs')
policy = CategoricalMLPPolicy(env_spec=env.spec)
policy.build(obs_var)
reg_vars = policy.get_regularizable_vars()
assert len(reg_vars) == 2
for var in reg_vars:
assert ('bias' not in var.name) and ('output' not in var.name)
def test_dist_info(self, obs_dim, action_dim):
env = TfEnv(DummyDiscreteEnv(obs_dim=obs_dim, action_dim=action_dim))
with mock.patch(('metarl.tf.policies.'
'categorical_mlp_policy.MLPModel'),
new=SimpleMLPModel):
policy = CategoricalMLPPolicy(env_spec=env.spec)
env.reset()
obs, _, _, _ = env.step(1)
expected_prob = np.full(action_dim, 0.5)
policy_probs = policy.dist_info([obs.flatten()])
assert np.array_equal(policy_probs['prob'][0], expected_prob)
def run_task(snapshot_config, *_):
"""Run task."""
with LocalTFRunner(snapshot_config=snapshot_config,
max_cpus=n_envs) as runner:
env = TfEnv(env_name='CartPole-v1')
policy = CategoricalMLPPolicy(name='policy',
env_spec=env.spec,
hidden_sizes=(32, 32))
baseline = LinearFeatureBaseline(env_spec=env.spec)
algo = TRPO(env_spec=env.spec,
policy=policy,
baseline=baseline,
max_path_length=max_path_length,
discount=0.99,
max_kl_step=0.01)
runner.setup(algo=algo,
env=env,
sampler_cls=BatchSampler,
sampler_args={'n_envs': n_envs})
runner.train(n_epochs=100, batch_size=4000, plot=False)
def test_process_samples_discrete_non_recurrent(self):
env = TfEnv(DummyDiscreteEnv())
policy = CategoricalMLPPolicy(env_spec=env.spec)
baseline = LinearFeatureBaseline(env_spec=env.spec)
max_path_length = 100
with LocalTFRunner(snapshot_config, sess=self.sess) as runner:
algo = BatchPolopt2(env_spec=env.spec,
policy=policy,
baseline=baseline,
max_path_length=max_path_length,
flatten_input=True)
runner.setup(algo, env, sampler_args=dict(n_envs=1))
runner.train(n_epochs=1, batch_size=max_path_length)
paths = runner.obtain_samples(0)
samples = algo.process_samples(0, paths)
# Since there is only 1 vec_env in the sampler and DummyDiscreteEnv
# always terminate, number of paths must be max_path_length, and
# batch size must be max_path_length as well, i.e. 100
assert samples['observations'].shape == (
max_path_length, env.observation_space.flat_dim)
assert samples['actions'].shape == (max_path_length,
env.action_space.n)
assert samples['rewards'].shape == (max_path_length, )
assert samples['baselines'].shape == (max_path_length, )
assert samples['returns'].shape == (max_path_length, )
# there is 100 path
assert samples['lengths'].shape == (max_path_length, )
# non-recurrent policy has empty agent info
assert samples['agent_infos'] == {}
# non-recurrent policy has empty env info
assert samples['env_infos'] == {}
assert isinstance(samples['average_return'], float)
def cem_cartpole(ctxt=None, seed=1):
"""Train CEM with Cartpole-v1 environment.
Args:
ctxt (metarl.experiment.ExperimentContext): The experiment
configuration used by LocalRunner to create the snapshotter.
seed (int): Used to seed the random number generator to produce
determinism.
"""
set_seed(seed)
with LocalTFRunner(snapshot_config=ctxt) as runner:
env = TfEnv(env_name='CartPole-v1')
policy = CategoricalMLPPolicy(name='policy',
env_spec=env.spec,
hidden_sizes=(32, 32))
baseline = LinearFeatureBaseline(env_spec=env.spec)
n_samples = 20
algo = CEM(env_spec=env.spec,
policy=policy,
baseline=baseline,
best_frac=0.05,
max_path_length=100,
n_samples=n_samples)
runner.setup(algo, env, sampler_cls=OnPolicyVectorizedSampler)
runner.train(n_epochs=100, batch_size=1000)
def erwr_cartpole(ctxt=None, seed=1):
"""Train with ERWR on CartPole-v1 environment.
Args:
ctxt (metarl.experiment.ExperimentContext): The experiment
configuration used by LocalRunner to create the snapshotter.
seed (int): Used to seed the random number generator to produce
determinism.
"""
set_seed(seed)
with LocalTFRunner(snapshot_config=ctxt) as runner:
env = TfEnv(env_name='CartPole-v1')
policy = CategoricalMLPPolicy(name='policy',
env_spec=env.spec,
hidden_sizes=(32, 32))
baseline = LinearFeatureBaseline(env_spec=env.spec)
algo = ERWR(env_spec=env.spec,
policy=policy,
baseline=baseline,
max_path_length=100,
discount=0.99)
runner.setup(algo=algo, env=env)
runner.train(n_epochs=100, batch_size=10000, plot=False)
def test_is_pickleable(self, obs_dim, action_dim):
env = TfEnv(DummyDiscreteEnv(obs_dim=obs_dim, action_dim=action_dim))
with mock.patch(('metarl.tf.policies.'
'categorical_mlp_policy.MLPModel'),
new=SimpleMLPModel):
policy = CategoricalMLPPolicy(env_spec=env.spec)
env.reset()
obs, _, _, _ = env.step(1)
with tf.compat.v1.variable_scope('CategoricalMLPPolicy/MLPModel',
reuse=True):
return_var = tf.compat.v1.get_variable('return_var')
# assign it to all one
return_var.load(tf.ones_like(return_var).eval())
output1 = self.sess.run(
policy.model.outputs,
feed_dict={policy.model.input: [obs.flatten()]})
p = pickle.dumps(policy)
with tf.compat.v1.Session(graph=tf.Graph()) as sess:
policy_pickled = pickle.loads(p)
output2 = sess.run(
policy_pickled.model.outputs,
feed_dict={policy_pickled.model.input: [obs.flatten()]})
assert np.array_equal(output1, output2)
def test_vpg_cartpole(self):
"""Test VPG with CartPole-v1 environment."""
with LocalTFRunner(snapshot_config, sess=self.sess) as runner:
env = TfEnv(env_name='CartPole-v1')
policy = CategoricalMLPPolicy(name='policy',
env_spec=env.spec,
hidden_sizes=(32, 32))
baseline = LinearFeatureBaseline(env_spec=env.spec)
algo = VPG(env_spec=env.spec,
policy=policy,
baseline=baseline,
max_path_length=100,
discount=0.99,
optimizer_args=dict(
tf_optimizer_args=dict(learning_rate=0.01, )))
runner.setup(algo, env)
last_avg_ret = runner.train(n_epochs=10, batch_size=10000)
assert last_avg_ret > 90
env.close()
def trpo_gym_tf_cartpole(ctxt=None, seed=1):
"""Train TRPO with CartPole-v0 environment.
Args:
ctxt (metarl.experiment.ExperimentContext): The experiment
configuration used by LocalRunner to create the snapshotter.
seed (int): Used to seed the random number generator to produce
determinism.
"""
set_seed(seed)
with LocalTFRunner(snapshot_config=ctxt) as runner:
env = MetaRLEnv(gym.make('CartPole-v0'))
policy = CategoricalMLPPolicy(name='policy',
env_spec=env.spec,
hidden_sizes=(32, 32))
baseline = LinearFeatureBaseline(env_spec=env.spec)
algo = TRPO(
env_spec=env.spec,
policy=policy,
baseline=baseline,
max_path_length=200,
discount=0.99,
max_kl_step=0.01,
)
runner.setup(algo, env)
runner.train(n_epochs=120, batch_size=4000)
def run_task(snapshot_config, *_):
"""Train CMA_ES with Cartpole-v1 environment.
Args:
snapshot_config (metarl.experiment.SnapshotConfig): The snapshot
configuration used by LocalRunner to create the snapshotter.
*_ (object): Ignored by this function.
"""
with LocalTFRunner(snapshot_config=snapshot_config) as runner:
env = TfEnv(env_name='CartPole-v1')
policy = CategoricalMLPPolicy(name='policy',
env_spec=env.spec,
hidden_sizes=(32, 32))
baseline = LinearFeatureBaseline(env_spec=env.spec)
n_samples = 20
algo = CMAES(env_spec=env.spec,
policy=policy,
baseline=baseline,
max_path_length=100,
n_samples=n_samples)
runner.setup(algo, env, sampler_cls=OnPolicyVectorizedSampler)
runner.train(n_epochs=100, batch_size=1000)
def test_cem_cartpole(self):
"""Test CEM with Cartpole-v1 environment."""
with LocalTFRunner(snapshot_config) as runner:
env = TfEnv(env_name='CartPole-v1')
policy = CategoricalMLPPolicy(name='policy',
env_spec=env.spec,
hidden_sizes=(32, 32))
baseline = LinearFeatureBaseline(env_spec=env.spec)
n_samples = 10
algo = CEM(env_spec=env.spec,
policy=policy,
baseline=baseline,
best_frac=0.1,
max_path_length=100,
n_samples=n_samples)
runner.setup(algo, env, sampler_cls=OnPolicyVectorizedSampler)
rtn = runner.train(n_epochs=10, batch_size=2048)
assert rtn > 40
env.close()
def test_ppo_pendulum_flatten_input(self):
"""Test PPO with CartPole to test observation flattening."""
with LocalTFRunner(snapshot_config, sess=self.sess) as runner:
env = MetaRLEnv(
normalize(ReshapeObservation(gym.make('CartPole-v1'), (2, 2))))
policy = CategoricalMLPPolicy(
env_spec=env.spec,
hidden_nonlinearity=tf.nn.tanh,
)
baseline = LinearFeatureBaseline(env_spec=env.spec)
algo = PPO(env_spec=env.spec,
policy=policy,
baseline=baseline,
max_path_length=100,
discount=0.99,
gae_lambda=0.95,
lr_clip_range=0.2,
policy_ent_coeff=0.0,
optimizer_args=dict(
batch_size=32,
max_epochs=10,
learning_rate=1e-3,
))
runner.setup(algo, env)
last_avg_ret = runner.train(n_epochs=10, batch_size=2048)
assert last_avg_ret > 80
def test_dist_info_sym(self, obs_dim, action_dim):
env = TfEnv(DummyDiscreteEnv(obs_dim=obs_dim, action_dim=action_dim))
with mock.patch(('metarl.tf.policies.'
'categorical_mlp_policy.MLPModel'),
new=SimpleMLPModel):
policy = CategoricalMLPPolicy(env_spec=env.spec)
env.reset()
obs, _, _, _ = env.step(1)
expected_prob = np.full(action_dim, 0.5)
obs_dim = env.spec.observation_space.flat_dim
state_input = tf.compat.v1.placeholder(tf.float32,
shape=(None, obs_dim))
dist1 = policy.dist_info_sym(state_input, name='policy2')
prob = self.sess.run(dist1['prob'],
feed_dict={state_input: [obs.flatten()]})
assert np.array_equal(prob[0], expected_prob)
def fixture_exp(snapshot_config, sess):
with LocalTFRunner(snapshot_config=snapshot_config, sess=sess) as runner:
env = TfEnv(env_name='CartPole-v1')
policy = CategoricalMLPPolicy(name='policy',
env_spec=env.spec,
hidden_sizes=(8, 8))
baseline = LinearFeatureBaseline(env_spec=env.spec)
algo = VPG(
env_spec=env.spec,
policy=policy,
baseline=baseline,
max_path_length=100,
discount=0.99,
optimizer_args=dict(tf_optimizer_args=dict(learning_rate=0.01, )))
runner.setup(algo, env)
runner.train(n_epochs=5, batch_size=100)
return policy.get_param_values()
def test_get_action(self, mock_rand, obs_dim, action_dim):
mock_rand.return_value = 0
env = TfEnv(DummyDiscreteEnv(obs_dim=obs_dim, action_dim=action_dim))
with mock.patch(('metarl.tf.policies.'
'categorical_mlp_policy.MLPModel'),
new=SimpleMLPModel):
policy = CategoricalMLPPolicy(env_spec=env.spec)
env.reset()
obs, _, _, _ = env.step(1)
action, prob = policy.get_action(obs)
expected_prob = np.full(action_dim, 0.5)
assert env.action_space.contains(action)
assert action == 0
assert np.array_equal(prob['prob'], expected_prob)
actions, probs = policy.get_actions([obs, obs, obs])
for action, prob in zip(actions, probs['prob']):
assert env.action_space.contains(action)
assert action == 0
assert np.array_equal(prob, expected_prob)
def run_task(snapshot_config, *_):
"""Run task."""
with LocalTFRunner(snapshot_config=snapshot_config) as runner:
env = TfEnv(gym.make('CartPole-v0'))
policy = CategoricalMLPPolicy(env_spec=env.spec, hidden_sizes=[32, 32])
baseline = LinearFeatureBaseline(env_spec=env.spec)
algo = REPS(env_spec=env.spec,
policy=policy,
baseline=baseline,
max_path_length=100,
discount=0.99)
runner.setup(algo, env)
runner.train(n_epochs=100, batch_size=4000, plot=False)
def test_train(self):
with LocalTFRunner(snapshot_config) as runner:
env = MetaRLEnv(env_name='CartPole-v1')
policy = CategoricalMLPPolicy(name='policy',
env_spec=env.spec,
hidden_sizes=(8, 8))
baseline = LinearFeatureBaseline(env_spec=env.spec)
algo = VPG(env_spec=env.spec,
policy=policy,
baseline=baseline,
max_path_length=100,
discount=0.99,
optimizer_args=dict(learning_rate=0.01, ))
runner.setup(algo, env)
runner.train(n_epochs=1, batch_size=100)
def run_task(snapshot_config, *_):
"""Run task."""
with LocalTFRunner(snapshot_config=snapshot_config) as runner:
env = TfEnv(env_name='CartPole-v1')
policy = CategoricalMLPPolicy(name='policy',
env_spec=env.spec,
hidden_sizes=(32, 32))
baseline = LinearFeatureBaseline(env_spec=env.spec)
algo = TRPO(env_spec=env.spec,
policy=policy,
baseline=baseline,
max_path_length=100,
discount=0.99,
max_kl_step=0.01)
runner.setup(algo, env)
runner.train(n_epochs=100, batch_size=4000)
def test_make_sampler_ray_sampler(self):
with LocalTFRunner(snapshot_config) as runner:
env = TfEnv(env_name='CartPole-v1')
policy = CategoricalMLPPolicy(name='policy',
env_spec=env.spec,
hidden_sizes=(8, 8))
baseline = LinearFeatureBaseline(env_spec=env.spec)
algo = VPG(env_spec=env.spec,
policy=policy,
baseline=baseline,
max_path_length=100,
discount=0.99,
optimizer_args=dict(
tf_optimizer_args=dict(learning_rate=0.01, )))
runner.setup(algo, env, sampler_cls=RaySampler)
assert isinstance(runner._sampler, RaySampler)
runner.train(n_epochs=1, batch_size=10)
def test_reps_cartpole(self):
"""Test REPS with gym Cartpole environment."""
with LocalTFRunner(snapshot_config, sess=self.sess) as runner:
env = TfEnv(gym.make('CartPole-v0'))
policy = CategoricalMLPPolicy(env_spec=env.spec,
hidden_sizes=[32, 32])
baseline = LinearFeatureBaseline(env_spec=env.spec)
algo = REPS(env_spec=env.spec,
policy=policy,
baseline=baseline,
max_path_length=100,
discount=0.99)
runner.setup(algo, env)
last_avg_ret = runner.train(n_epochs=10, batch_size=4000)
assert last_avg_ret > 5
env.close()
def trpo_cartpole_batch_sampler(ctxt=None,
seed=1,
batch_size=4000,
max_path_length=100):
"""Train TRPO with CartPole-v1 environment.
Args:
ctxt (metarl.experiment.ExperimentContext): The experiment
configuration used by LocalRunner to create the snapshotter.
seed (int): Used to seed the random number generator to produce
determinism.
batch_size (int): Number of timesteps to use in each training step.
max_path_length (int): Number of timesteps to truncate paths to.
"""
set_seed(seed)
n_envs = batch_size // max_path_length
with LocalTFRunner(ctxt, max_cpus=n_envs) as runner:
env = MetaRLEnv(env_name='CartPole-v1')
policy = CategoricalMLPPolicy(name='policy',
env_spec=env.spec,
hidden_sizes=(32, 32))
baseline = LinearFeatureBaseline(env_spec=env.spec)
algo = TRPO(env_spec=env.spec,
policy=policy,
baseline=baseline,
max_path_length=max_path_length,
discount=0.99,
max_kl_step=0.01)
runner.setup(algo=algo,
env=env,
sampler_cls=BatchSampler,
sampler_args={'n_envs': n_envs})
runner.train(n_epochs=100, batch_size=4000, plot=False)
def multi_env_ppo(ctxt=None, seed=1):
"""Train PPO on two Atari environments simultaneously.
Args:
ctxt (metarl.experiment.ExperimentContext): The experiment
configuration used by LocalRunner to create the snapshotter.
seed (int): Used to seed the random number generator to produce
determinism.
"""
set_seed(seed)
with LocalTFRunner(ctxt) as runner:
env1 = MetaRLEnv(normalize(gym.make('Adventure-ram-v4')))
env2 = MetaRLEnv(normalize(gym.make('Alien-ram-v4')))
env = MultiEnvWrapper([env1, env2])
policy = CategoricalMLPPolicy(
env_spec=env.spec,
hidden_nonlinearity=tf.nn.tanh,
)
baseline = LinearFeatureBaseline(env_spec=env.spec)
algo = PPO(env_spec=env.spec,
policy=policy,
baseline=baseline,
max_path_length=100,
discount=0.99,
gae_lambda=0.95,
lr_clip_range=0.2,
policy_ent_coeff=0.0,
optimizer_args=dict(
batch_size=32,
max_epochs=10,
learning_rate=1e-3,
))
runner.setup(algo, env)
runner.train(n_epochs=120, batch_size=2048, plot=False)
def run_task(snapshot_config, *_):
"""Run task.
Args:
snapshot_config (metarl.experiment.SnapshotConfig): The snapshot
configuration used by LocalRunner to create the snapshotter.
_ (object): Ignored by this function.
"""
with LocalTFRunner(snapshot_config=snapshot_config) as runner:
env1 = TfEnv(normalize(gym.make('Adventure-ram-v4')))
env2 = TfEnv(normalize(gym.make('Alien-ram-v4')))
env = MultiEnvWrapper([env1, env2])
policy = CategoricalMLPPolicy(
env_spec=env.spec,
hidden_nonlinearity=tf.nn.tanh,
)
baseline = LinearFeatureBaseline(env_spec=env.spec)
algo = PPO(env_spec=env.spec,
policy=policy,
baseline=baseline,
max_path_length=100,
discount=0.99,
gae_lambda=0.95,
lr_clip_range=0.2,
policy_ent_coeff=0.0,
optimizer_args=dict(
batch_size=32,
max_epochs=10,
tf_optimizer_args=dict(learning_rate=1e-3),
))
runner.setup(algo, env)
runner.train(n_epochs=120, batch_size=2048, plot=False)
def test_cma_es_cartpole(self):
"""Test CMAES with Cartpole-v1 environment."""
with LocalTFRunner(snapshot_config) as runner:
env = MetaRLEnv(env_name='CartPole-v1')
policy = CategoricalMLPPolicy(name='policy',
env_spec=env.spec,
hidden_sizes=(32, 32))
baseline = LinearFeatureBaseline(env_spec=env.spec)
n_samples = 20
algo = CMAES(env_spec=env.spec,
policy=policy,
baseline=baseline,
max_path_length=100,
n_samples=n_samples)
runner.setup(algo, env, sampler_cls=OnPolicyVectorizedSampler)
runner.train(n_epochs=1, batch_size=1000)
# No assertion on return because CMAES is not stable.
env.close()
def test_erwr_cartpole(self):
"""Test ERWR with Cartpole-v1 environment."""
with LocalTFRunner(snapshot_config, sess=self.sess) as runner:
env = TfEnv(env_name='CartPole-v1')
policy = CategoricalMLPPolicy(name='policy',
env_spec=env.spec,
hidden_sizes=(32, 32))
baseline = LinearFeatureBaseline(env_spec=env.spec)
algo = ERWR(env_spec=env.spec,
policy=policy,
baseline=baseline,
max_path_length=100,
discount=0.99)
runner.setup(algo, env)
last_avg_ret = runner.train(n_epochs=10, batch_size=10000)
assert last_avg_ret > 80
env.close()
def test_on_policy_vectorized_sampler_n_envs(self, cpus, n_envs,
expected_n_envs):
with LocalTFRunner(snapshot_config, sess=self.sess,
max_cpus=cpus) as runner:
env = TfEnv(gym.make('CartPole-v0'))
policy = CategoricalMLPPolicy(env_spec=env.spec,
hidden_sizes=[32, 32])
baseline = LinearFeatureBaseline(env_spec=env.spec)
algo = REPS(env_spec=env.spec,
policy=policy,
baseline=baseline,
max_path_length=100,
discount=0.99)
runner.setup(algo, env, sampler_args=dict(n_envs=n_envs))
assert isinstance(runner._sampler, OnPolicyVectorizedSampler)
assert runner._sampler._n_envs == expected_n_envs
env.close()