def test_onehots_consistent_with_task_sampler():
# Import, construct environments here to avoid using up too much
# resources if this test isn't run.
# pylint: disable=import-outside-toplevel
import metaworld
mt10 = metaworld.MT10()
env = MetaWorldSetTaskEnv(mt10, 'train', add_env_onehot=True)
policy = RandomPolicy(env.action_space)
workers = WorkerFactory(seed=100, max_episode_length=1, n_workers=10)
sampler1 = LocalSampler.from_worker_factory(workers, policy, env)
env_ups = [
SetTaskUpdate(MetaWorldSetTaskEnv, task, None)
for task in env.sample_tasks(10)
]
samples1 = sampler1.obtain_exact_episodes(1, policy, env_ups)
task_sampler = MetaWorldTaskSampler(mt10, 'train', add_env_onehot=True)
env_ups = task_sampler.sample(10)
sampler2 = LocalSampler.from_worker_factory(workers, policy, env_ups)
samples2 = sampler2.obtain_exact_episodes(1, policy, env_ups)
name_to_obs1 = {}
for obs1, name1 in zip(samples1.observations,
samples1.env_infos['task_name']):
name_to_obs1[name1] = obs1
for obs2, name2 in zip(samples2.observations,
samples2.env_infos['task_name']):
assert (name_to_obs1[name2][-10:] == obs2[-10:]).all()
def test_forbidden_cases():
# Import, construct environments here to avoid using up too much
# resources if this test isn't run.
# pylint: disable=import-outside-toplevel
import metaworld
ml1 = metaworld.ML1('push-v1')
with pytest.raises(ValueError):
MetaWorldSetTaskEnv(ml1, 'train', add_env_onehot=True)
with pytest.raises(ValueError):
MetaWorldSetTaskEnv(ml1, 'Test')
def maml_trpo_metaworld_ml1_push(ctxt, seed, epochs, rollouts_per_task,
meta_batch_size):
"""Set up environment and algorithm and run the task.
Args:
ctxt (garage.experiment.ExperimentContext): The experiment
configuration used by Trainer to create the snapshotter.
seed (int): Used to seed the random number generator to produce
determinism.
epochs (int): Number of training epochs.
rollouts_per_task (int): Number of rollouts per epoch per task
for training.
meta_batch_size (int): Number of tasks sampled per batch.
"""
set_seed(seed)
ml1 = metaworld.ML1('push-v1')
tasks = MetaWorldTaskSampler(ml1, 'train')
env = tasks.sample(1)[0]()
test_sampler = SetTaskSampler(MetaWorldSetTaskEnv,
env=MetaWorldSetTaskEnv(ml1, 'test'))
policy = GaussianMLPPolicy(
env_spec=env.spec,
hidden_sizes=(100, 100),
hidden_nonlinearity=torch.tanh,
output_nonlinearity=None,
)
value_function = GaussianMLPValueFunction(env_spec=env.spec,
hidden_sizes=[32, 32],
hidden_nonlinearity=torch.tanh,
output_nonlinearity=None)
meta_evaluator = MetaEvaluator(test_task_sampler=test_sampler,
n_test_tasks=1,
n_exploration_eps=rollouts_per_task)
sampler = RaySampler(agents=policy,
envs=env,
max_episode_length=env.spec.max_episode_length,
n_workers=meta_batch_size)
trainer = Trainer(ctxt)
algo = MAMLTRPO(env=env,
policy=policy,
sampler=sampler,
task_sampler=tasks,
value_function=value_function,
meta_batch_size=meta_batch_size,
discount=0.99,
gae_lambda=1.,
inner_lr=0.1,
num_grad_updates=1,
meta_evaluator=meta_evaluator)
trainer.setup(algo, env)
trainer.train(n_epochs=epochs,
batch_size=rollouts_per_task * env.spec.max_episode_length)
def maml_trpo_metaworld_ml45(ctxt, seed, epochs, episodes_per_task,
meta_batch_size):
"""Set up environment and algorithm and run the task.
Args:
ctxt (ExperimentContext): The experiment configuration used by
:class:`~Trainer` to create the :class:`~Snapshotter`.
seed (int): Used to seed the random number generator to produce
determinism.
epochs (int): Number of training epochs.
episodes_per_task (int): Number of episodes per epoch per task
for training.
meta_batch_size (int): Number of tasks sampled per batch.
"""
set_seed(seed)
ml45 = metaworld.ML45()
# pylint: disable=missing-return-doc,missing-return-type-doc
def wrap(env, _):
return normalize(env, expected_action_scale=10.0)
train_task_sampler = MetaWorldTaskSampler(ml45, 'train', wrap)
test_env = wrap(MetaWorldSetTaskEnv(ml45, 'test'), None)
test_task_sampler = SetTaskSampler(MetaWorldSetTaskEnv,
env=test_env,
wrapper=wrap)
env = train_task_sampler.sample(45)[0]()
policy = GaussianMLPPolicy(
env_spec=env.spec,
hidden_sizes=(100, 100),
hidden_nonlinearity=torch.tanh,
output_nonlinearity=None,
)
value_function = GaussianMLPValueFunction(env_spec=env.spec,
hidden_sizes=(32, 32),
hidden_nonlinearity=torch.tanh,
output_nonlinearity=None)
meta_evaluator = MetaEvaluator(test_task_sampler=test_task_sampler)
trainer = Trainer(ctxt)
algo = MAMLTRPO(env=env,
task_sampler=train_task_sampler,
policy=policy,
value_function=value_function,
meta_batch_size=meta_batch_size,
discount=0.99,
gae_lambda=1.,
inner_lr=0.1,
num_grad_updates=1,
meta_evaluator=meta_evaluator)
trainer.setup(algo, env, n_workers=meta_batch_size)
trainer.train(n_epochs=epochs,
batch_size=episodes_per_task * env.spec.max_episode_length)
def test_pearl_ml1_push(self):
"""Test PEARL with ML1 Push environment."""
params = dict(seed=1,
num_epochs=1,
num_train_tasks=5,
latent_size=7,
encoder_hidden_sizes=[10, 10, 10],
net_size=30,
meta_batch_size=16,
num_steps_per_epoch=40,
num_initial_steps=40,
num_tasks_sample=15,
num_steps_prior=15,
num_extra_rl_steps_posterior=15,
batch_size=256,
embedding_batch_size=8,
embedding_mini_batch_size=8,
reward_scale=10.,
use_information_bottleneck=True,
use_next_obs_in_context=False,
use_gpu=False)
net_size = params['net_size']
set_seed(params['seed'])
# create multi-task environment and sample tasks
ml1 = metaworld.ML1('push-v1')
train_env = MetaWorldSetTaskEnv(ml1, 'train')
env_sampler = SetTaskSampler(MetaWorldSetTaskEnv,
env=train_env,
wrapper=lambda env, _: normalize(env))
env = env_sampler.sample(params['num_train_tasks'])
test_env = MetaWorldSetTaskEnv(ml1, 'test')
test_env_sampler = SetTaskSampler(
MetaWorldSetTaskEnv,
env=test_env,
wrapper=lambda env, _: normalize(env))
augmented_env = PEARL.augment_env_spec(env[0](), params['latent_size'])
qf = ContinuousMLPQFunction(
env_spec=augmented_env,
hidden_sizes=[net_size, net_size, net_size])
vf_env = PEARL.get_env_spec(env[0](), params['latent_size'], 'vf')
vf = ContinuousMLPQFunction(
env_spec=vf_env, hidden_sizes=[net_size, net_size, net_size])
inner_policy = TanhGaussianMLPPolicy(
env_spec=augmented_env,
hidden_sizes=[net_size, net_size, net_size])
pearl = PEARL(
env=env,
policy_class=ContextConditionedPolicy,
encoder_class=MLPEncoder,
inner_policy=inner_policy,
qf=qf,
vf=vf,
num_train_tasks=params['num_train_tasks'],
latent_dim=params['latent_size'],
encoder_hidden_sizes=params['encoder_hidden_sizes'],
test_env_sampler=test_env_sampler,
meta_batch_size=params['meta_batch_size'],
num_steps_per_epoch=params['num_steps_per_epoch'],
num_initial_steps=params['num_initial_steps'],
num_tasks_sample=params['num_tasks_sample'],
num_steps_prior=params['num_steps_prior'],
num_extra_rl_steps_posterior=params[
'num_extra_rl_steps_posterior'],
batch_size=params['batch_size'],
embedding_batch_size=params['embedding_batch_size'],
embedding_mini_batch_size=params['embedding_mini_batch_size'],
reward_scale=params['reward_scale'],
)
set_gpu_mode(params['use_gpu'], gpu_id=0)
if params['use_gpu']:
pearl.to()
trainer = Trainer(snapshot_config)
trainer.setup(algo=pearl,
env=env[0](),
sampler_cls=LocalSampler,
n_workers=1,
worker_class=PEARLWorker)
trainer.train(n_epochs=params['num_epochs'],
batch_size=params['batch_size'])
def rl2_ppo_metaworld_ml1_push(ctxt, seed, meta_batch_size, n_epochs,
episode_per_task):
"""Train RL2 PPO with ML1 environment.
Args:
ctxt (ExperimentContext): The experiment configuration used by
:class:`~Trainer` to create the :class:`~Snapshotter`.
seed (int): Used to seed the random number generator to produce
determinism.
meta_batch_size (int): Meta batch size.
n_epochs (int): Total number of epochs for training.
episode_per_task (int): Number of training episode per task.
"""
set_seed(seed)
ml1 = metaworld.ML1('push-v1')
task_sampler = MetaWorldTaskSampler(ml1, 'train',
lambda env, _: RL2Env(env))
env = task_sampler.sample(1)[0]()
test_task_sampler = SetTaskSampler(MetaWorldSetTaskEnv,
env=MetaWorldSetTaskEnv(ml1, 'test'),
wrapper=lambda env, _: RL2Env(env))
env_spec = env.spec
with TFTrainer(snapshot_config=ctxt) as trainer:
policy = GaussianGRUPolicy(name='policy',
hidden_dim=64,
env_spec=env_spec,
state_include_action=False)
meta_evaluator = MetaEvaluator(test_task_sampler=test_task_sampler)
baseline = LinearFeatureBaseline(env_spec=env_spec)
algo = RL2PPO(meta_batch_size=meta_batch_size,
task_sampler=task_sampler,
env_spec=env_spec,
policy=policy,
baseline=baseline,
discount=0.99,
gae_lambda=0.95,
lr_clip_range=0.2,
optimizer_args=dict(batch_size=32,
max_optimization_epochs=10),
stop_entropy_gradient=True,
entropy_method='max',
policy_ent_coeff=0.02,
center_adv=False,
meta_evaluator=meta_evaluator,
episodes_per_trial=episode_per_task)
trainer.setup(algo,
task_sampler.sample(meta_batch_size),
sampler_cls=LocalSampler,
n_workers=meta_batch_size,
worker_class=RL2Worker,
worker_args=dict(n_episodes_per_trial=episode_per_task))
trainer.train(n_epochs=n_epochs,
batch_size=episode_per_task *
env_spec.max_episode_length * meta_batch_size)
def pearl_metaworld_ml10(ctxt=None,
seed=1,
num_epochs=1000,
num_train_tasks=10,
latent_size=7,
encoder_hidden_size=200,
net_size=300,
meta_batch_size=16,
num_steps_per_epoch=4000,
num_initial_steps=4000,
num_tasks_sample=15,
num_steps_prior=750,
num_extra_rl_steps_posterior=750,
batch_size=256,
embedding_batch_size=64,
embedding_mini_batch_size=64,
reward_scale=10.,
use_gpu=False):
"""Train PEARL with ML10 environments.
Args:
ctxt (garage.experiment.ExperimentContext): The experiment
configuration used by Trainer to create the snapshotter.
seed (int): Used to seed the random number generator to produce
determinism.
num_epochs (int): Number of training epochs.
num_train_tasks (int): Number of tasks for training.
latent_size (int): Size of latent context vector.
encoder_hidden_size (int): Output dimension of dense layer of the
context encoder.
net_size (int): Output dimension of a dense layer of Q-function and
value function.
meta_batch_size (int): Meta batch size.
num_steps_per_epoch (int): Number of iterations per epoch.
num_initial_steps (int): Number of transitions obtained per task before
training.
num_tasks_sample (int): Number of random tasks to obtain data for each
iteration.
num_steps_prior (int): Number of transitions to obtain per task with
z ~ prior.
num_extra_rl_steps_posterior (int): Number of additional transitions
to obtain per task with z ~ posterior that are only used to train
the policy and NOT the encoder.
batch_size (int): Number of transitions in RL batch.
embedding_batch_size (int): Number of transitions in context batch.
embedding_mini_batch_size (int): Number of transitions in mini context
batch; should be same as embedding_batch_size for non-recurrent
encoder.
reward_scale (int): Reward scale.
use_gpu (bool): Whether or not to use GPU for training.
"""
set_seed(seed)
encoder_hidden_sizes = (encoder_hidden_size, encoder_hidden_size,
encoder_hidden_size)
ml10 = metaworld.ML10()
train_env = MetaWorldSetTaskEnv(ml10, 'train')
env_sampler = SetTaskSampler(MetaWorldSetTaskEnv,
env=train_env,
wrapper=lambda env, _: normalize(env))
env = env_sampler.sample(num_train_tasks)
test_env = MetaWorldSetTaskEnv(ml10, 'test')
test_env_sampler = SetTaskSampler(MetaWorldSetTaskEnv,
env=test_env,
wrapper=lambda env, _: normalize(env))
trainer = Trainer(ctxt)
# instantiate networks
augmented_env = PEARL.augment_env_spec(env[0](), latent_size)
qf = ContinuousMLPQFunction(env_spec=augmented_env,
hidden_sizes=[net_size, net_size, net_size])
vf_env = PEARL.get_env_spec(env[0](), latent_size, 'vf')
vf = ContinuousMLPQFunction(env_spec=vf_env,
hidden_sizes=[net_size, net_size, net_size])
inner_policy = TanhGaussianMLPPolicy(
env_spec=augmented_env, hidden_sizes=[net_size, net_size, net_size])
pearl = PEARL(
env=env,
policy_class=ContextConditionedPolicy,
encoder_class=MLPEncoder,
inner_policy=inner_policy,
qf=qf,
vf=vf,
num_train_tasks=num_train_tasks,
latent_dim=latent_size,
encoder_hidden_sizes=encoder_hidden_sizes,
test_env_sampler=test_env_sampler,
meta_batch_size=meta_batch_size,
num_steps_per_epoch=num_steps_per_epoch,
num_initial_steps=num_initial_steps,
num_tasks_sample=num_tasks_sample,
num_steps_prior=num_steps_prior,
num_extra_rl_steps_posterior=num_extra_rl_steps_posterior,
batch_size=batch_size,
embedding_batch_size=embedding_batch_size,
embedding_mini_batch_size=embedding_mini_batch_size,
reward_scale=reward_scale,
)
set_gpu_mode(use_gpu, gpu_id=0)
if use_gpu:
pearl.to()
trainer.setup(algo=pearl,
env=env[0](),
sampler_cls=LocalSampler,
n_workers=1,
worker_class=PEARLWorker)
trainer.train(n_epochs=num_epochs, batch_size=batch_size)
def tcl_pearl_ml1(ctxt=None,
seed=1,
num_epochs=200,
num_train_tasks=50,
num_test_tasks=10,
latent_size=7,
encoder_hidden_size=200,
net_size=300,
meta_batch_size=16,
num_steps_per_epoch=4000,
num_initial_steps=4000,
num_tasks_sample=15,
num_steps_prior=750,
num_extra_rl_steps_posterior=750,
batch_size=256,
embedding_batch_size=64,
embedding_mini_batch_size=64,
max_path_length=200,
reward_scale=10.,
replay_buffer_size=1000000,
use_next_obs=False,
in_sequence_path_aug=True,
emphasized_network=False,
use_kl_loss=True,
use_q_loss=True,
encoder_common_net=True,
single_alpha=False,
use_task_index_label=False,
use_wasserstein_distance=True,
gpu_id=0,
name='push-v1',
prefix='curl_fine_tune',
use_gpu=True):
"""Train TCL-PEARL with ML1 environments.
Args:
ctxt (garage.experiment.ExperimentContext): The experiment
configuration used by LocalRunner to create the snapshotter.
seed (int): Used to seed the random number generator to produce
determinism.
num_epochs (int): Number of training epochs.
num_train_tasks (int): Number of tasks for training.
num_test_tasks (int): Number of tasks for testing.
latent_size (int): Size of latent context vector.
encoder_hidden_size (int): Output dimension of dense layer of the
context encoder.
net_size (int): Output dimension of a dense layer of Q-function and
value function.
meta_batch_size (int): Meta batch size.
num_steps_per_epoch (int): Number of iterations per epoch.
num_initial_steps (int): Number of transitions obtained per task before
training.
num_tasks_sample (int): Number of random tasks to obtain data for each
iteration.
num_steps_prior (int): Number of transitions to obtain per task with
z ~ prior.
num_extra_rl_steps_posterior (int): Number of additional transitions
to obtain per task with z ~ posterior that are only used to train
the policy and NOT the encoder.
batch_size (int): Number of transitions in RL batch.
embedding_batch_size (int): Number of transitions in context batch.
embedding_mini_batch_size (int): Number of transitions in mini context
batch; should be same as embedding_batch_size for non-recurrent
encoder.
max_path_length (int): Maximum path length.
reward_scale (int): Reward scale.
use_gpu (bool): Whether or not to use GPU for training.
"""
set_seed(seed)
encoder_hidden_sizes = (encoder_hidden_size, encoder_hidden_size,
encoder_hidden_size)
print("Running experiences on {}/{}".format(prefix, name))
# create multi-task environment and sample tasks
ml1 = metaworld.ML1(name)
train_env = MetaWorldSetTaskEnv(ml1, 'train')
env_sampler = SetTaskSampler(MetaWorldSetTaskEnv,
env=train_env,
wrapper=lambda env, _: normalize(env))
env = env_sampler.sample(num_train_tasks)
test_env = MetaWorldSetTaskEnv(ml1, 'test')
test_env_sampler = SetTaskSampler(MetaWorldSetTaskEnv,
env=test_env,
wrapper=lambda env, _: normalize(env))
sampler = LocalSampler(agents=None,
envs=env[0](),
max_episode_length=max_path_length,
n_workers=1,
worker_class=TCLPEARLWorker)
trainer = Trainer(ctxt)
# instantiate networks
augmented_env = TCLPEARL.augment_env_spec(env[0](), latent_size)
qf_1 = ContinuousMLPQFunction(env_spec=augmented_env,
hidden_sizes=[net_size, net_size, net_size])
qf_2 = ContinuousMLPQFunction(env_spec=augmented_env,
hidden_sizes=[net_size, net_size, net_size])
inner_policy = TanhGaussianMLPPolicy(
env_spec=augmented_env, hidden_sizes=[net_size, net_size, net_size])
tcl_pearl = TCLPEARL(
env=env,
policy_class=TCLPolicy,
encoder_class=ContrastiveEncoder,
inner_policy=inner_policy,
qf1=qf_1,
qf2=qf_2,
sampler=sampler,
num_train_tasks=num_train_tasks,
num_test_tasks=num_test_tasks,
latent_dim=latent_size,
encoder_hidden_sizes=encoder_hidden_sizes,
test_env_sampler=test_env_sampler,
meta_batch_size=meta_batch_size,
num_steps_per_epoch=num_steps_per_epoch,
num_initial_steps=num_initial_steps,
num_tasks_sample=num_tasks_sample,
num_steps_prior=num_steps_prior,
num_extra_rl_steps_posterior=num_extra_rl_steps_posterior,
batch_size=batch_size,
embedding_batch_size=embedding_batch_size,
embedding_mini_batch_size=embedding_mini_batch_size,
max_path_length=max_path_length,
reward_scale=reward_scale,
replay_buffer_size=replay_buffer_size,
use_next_obs_in_context=use_next_obs,
embedding_batch_in_sequence=in_sequence_path_aug,
use_kl_loss=use_kl_loss,
use_q_loss=use_q_loss,
encoder_common_net=encoder_common_net,
single_alpha=single_alpha,
use_task_index_label=use_task_index_label,
use_wasserstein_distance=use_wasserstein_distance)
set_gpu_mode(use_gpu, gpu_id=gpu_id)
if use_gpu:
tcl_pearl.to()
trainer.setup(algo=tcl_pearl, env=env[0]())
trainer.train(n_epochs=num_epochs, batch_size=batch_size)
def test_sample_and_step():
# Import, construct environments here to avoid using up too much
# resources if this test isn't run.
# pylint: disable=import-outside-toplevel
import metaworld
ml1 = metaworld.ML1('push-v1')
env = MetaWorldSetTaskEnv(ml1, 'train')
assert env.num_tasks == 50
task = env.sample_tasks(1)[0]
env.set_task(task)
env.step(env.action_space.sample())
env.close()
env2 = MetaWorldSetTaskEnv()
env2.set_task(task)
env2.step(env.action_space.sample())
env2.close()
tasks = env.sample_tasks(100)
assert len(tasks) == 100
