def run_metarl(env, seed, log_dir):
'''
Create metarl model and training.
Replace the ddpg with the algorithm you want to run.
:param env: Environment of the task.
:param seed: Random seed for the trial.
:param log_dir: Log dir path.
:return:
'''
deterministic.set_seed(seed)
runner = LocalRunner(snapshot_config)
# Set up params for ddpg
policy = TanhGaussianMLPPolicy2(env_spec=env.spec,
hidden_sizes=params['policy_hidden_sizes'],
hidden_nonlinearity=nn.ReLU,
output_nonlinearity=None)
qf1 = ContinuousMLPQFunction(env_spec=env.spec,
hidden_sizes=params['qf_hidden_sizes'],
hidden_nonlinearity=F.relu)
qf2 = ContinuousMLPQFunction(env_spec=env.spec,
hidden_sizes=params['qf_hidden_sizes'],
hidden_nonlinearity=F.relu)
replay_buffer = SACReplayBuffer(env_spec=env.spec,
max_size=params['replay_buffer_size'])
sampler_args = {
'agent': policy,
'max_path_length': 1000,
}
sac = SAC(env_spec=env.spec,
policy=policy,
qf1=qf1,
qf2=qf2,
gradient_steps_per_itr=params['gradient_steps_per_itr'],
replay_buffer=replay_buffer,
buffer_batch_size=params['buffer_batch_size'])
# Set up logger since we are not using run_experiment
tabular_log_file = osp.join(log_dir, 'progress.csv')
tensorboard_log_dir = osp.join(log_dir)
dowel_logger.add_output(dowel.StdOutput())
dowel_logger.add_output(dowel.CsvOutput(tabular_log_file))
dowel_logger.add_output(dowel.TensorBoardOutput(tensorboard_log_dir))
runner.setup(algo=sac,
env=env,
sampler_cls=SimpleSampler,
sampler_args=sampler_args)
runner.train(n_epochs=params['n_epochs'],
batch_size=params['gradient_steps_per_itr'])
dowel_logger.remove_all()
return tabular_log_file
def test_to():
"""Test the torch function that moves modules to GPU.
Test that the policy and qfunctions are moved to gpu if gpu is
available.
"""
env_names = ['CartPole-v0', 'CartPole-v1']
task_envs = [MetaRLEnv(env_name=name) for name in env_names]
env = MultiEnvWrapper(task_envs, sample_strategy=round_robin_strategy)
deterministic.set_seed(0)
policy = TanhGaussianMLPPolicy(
env_spec=env.spec,
hidden_sizes=[1, 1],
hidden_nonlinearity=torch.nn.ReLU,
output_nonlinearity=None,
min_std=np.exp(-20.),
max_std=np.exp(2.),
)
qf1 = ContinuousMLPQFunction(env_spec=env.spec,
hidden_sizes=[1, 1],
hidden_nonlinearity=F.relu)
qf2 = ContinuousMLPQFunction(env_spec=env.spec,
hidden_sizes=[1, 1],
hidden_nonlinearity=F.relu)
replay_buffer = PathBuffer(capacity_in_transitions=int(1e6), )
num_tasks = 2
buffer_batch_size = 2
mtsac = MTSAC(policy=policy,
qf1=qf1,
qf2=qf2,
gradient_steps_per_itr=150,
max_path_length=150,
eval_env=env,
env_spec=env.spec,
num_tasks=num_tasks,
steps_per_epoch=5,
replay_buffer=replay_buffer,
min_buffer_size=1e3,
target_update_tau=5e-3,
discount=0.99,
buffer_batch_size=buffer_batch_size)
set_gpu_mode(torch.cuda.is_available())
mtsac.to()
device = global_device()
for param in mtsac._qf1.parameters():
assert param.device == device
for param in mtsac._qf2.parameters():
assert param.device == device
for param in mtsac._qf2.parameters():
assert param.device == device
for param in mtsac._policy.parameters():
assert param.device == device
assert mtsac._log_alpha.device == device
def ml1_push_v1_sac(ctxt=None, seed=1):
"""Set up environment and algorithm and run the task."""
runner = LocalRunner(ctxt)
Ml1_reach_envs = get_ML1_envs("push-v1")
Ml1_reach_test_envs = get_ML1_envs_test("push-v1")
env = MTMetaWorldWrapper(Ml1_reach_envs)
policy = TanhGaussianMLPPolicy2(
env_spec=env.spec,
hidden_sizes=[400, 400, 400],
hidden_nonlinearity=nn.ReLU,
output_nonlinearity=None,
min_std=np.exp(-20.),
max_std=np.exp(2.),
)
qf1 = ContinuousMLPQFunction(env_spec=env.spec,
hidden_sizes=[400, 400, 400],
hidden_nonlinearity=F.relu)
qf2 = ContinuousMLPQFunction(env_spec=env.spec,
hidden_sizes=[400, 400, 400],
hidden_nonlinearity=F.relu)
replay_buffer = SACReplayBuffer(env_spec=env.spec, max_size=int(1e6))
sampler_args = {'agent': policy, 'max_path_length': 150}
timesteps = 100000000
batch_size = int(150 * env.num_tasks)
num_evaluation_points = 500
epochs = timesteps // batch_size
epoch_cycles = epochs // num_evaluation_points
epochs = epochs // epoch_cycles
sac = MTSAC(env=env,
eval_env_dict=Ml1_reach_test_envs,
env_spec=env.spec,
policy=policy,
qf1=qf1,
qf2=qf2,
gradient_steps_per_itr=250,
epoch_cycles=epoch_cycles,
use_automatic_entropy_tuning=True,
replay_buffer=replay_buffer,
min_buffer_size=7500,
target_update_tau=5e-3,
discount=0.99,
buffer_batch_size=6400)
tu.set_gpu_mode(True)
sac.to('cuda:0')
runner.setup(algo=sac,
env=env,
sampler_cls=SimpleSampler,
sampler_args=sampler_args)
runner.train(n_epochs=epochs, batch_size=batch_size)
def test_fixed_alpha():
"""Test if using fixed_alpha ensures that alpha is non differentiable."""
env_names = ['InvertedDoublePendulum-v2', 'InvertedDoublePendulum-v2']
task_envs = [MetaRLEnv(env_name=name) for name in env_names]
env = MultiEnvWrapper(task_envs, sample_strategy=round_robin_strategy)
test_envs = MultiEnvWrapper(task_envs,
sample_strategy=round_robin_strategy)
deterministic.set_seed(0)
runner = LocalRunner(snapshot_config=snapshot_config)
policy = TanhGaussianMLPPolicy(
env_spec=env.spec,
hidden_sizes=[32, 32],
hidden_nonlinearity=torch.nn.ReLU,
output_nonlinearity=None,
min_std=np.exp(-20.),
max_std=np.exp(2.),
)
qf1 = ContinuousMLPQFunction(env_spec=env.spec,
hidden_sizes=[32, 32],
hidden_nonlinearity=F.relu)
qf2 = ContinuousMLPQFunction(env_spec=env.spec,
hidden_sizes=[32, 32],
hidden_nonlinearity=F.relu)
replay_buffer = PathBuffer(capacity_in_transitions=int(1e6), )
num_tasks = 2
buffer_batch_size = 128
mtsac = MTSAC(policy=policy,
qf1=qf1,
qf2=qf2,
gradient_steps_per_itr=100,
max_path_length=100,
eval_env=test_envs,
env_spec=env.spec,
num_tasks=num_tasks,
steps_per_epoch=1,
replay_buffer=replay_buffer,
min_buffer_size=1e3,
target_update_tau=5e-3,
discount=0.99,
buffer_batch_size=buffer_batch_size,
fixed_alpha=np.exp(0.5))
if torch.cuda.is_available():
set_gpu_mode(True)
else:
set_gpu_mode(False)
mtsac.to()
assert torch.allclose(torch.Tensor([0.5] * num_tasks),
mtsac._log_alpha.to('cpu'))
runner.setup(mtsac, env, sampler_cls=LocalSampler)
runner.train(n_epochs=1, batch_size=128, plot=False)
assert torch.allclose(torch.Tensor([0.5] * num_tasks),
mtsac._log_alpha.to('cpu'))
assert not mtsac._use_automatic_entropy_tuning
def sac_half_cheetah_batch(ctxt=None, seed=1):
"""Set up environment and algorithm and run the task.
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.
"""
deterministic.set_seed(seed)
runner = LocalRunner(snapshot_config=ctxt)
env = MetaRLEnv(normalize(gym.make('HalfCheetah-v2')))
policy = TanhGaussianMLPPolicy(
env_spec=env.spec,
hidden_sizes=[256, 256],
hidden_nonlinearity=nn.ReLU,
output_nonlinearity=None,
min_std=np.exp(-20.),
max_std=np.exp(2.),
)
qf1 = ContinuousMLPQFunction(env_spec=env.spec,
hidden_sizes=[256, 256],
hidden_nonlinearity=F.relu)
qf2 = ContinuousMLPQFunction(env_spec=env.spec,
hidden_sizes=[256, 256],
hidden_nonlinearity=F.relu)
replay_buffer = PathBuffer(capacity_in_transitions=int(1e6))
sac = SAC(env_spec=env.spec,
policy=policy,
qf1=qf1,
qf2=qf2,
gradient_steps_per_itr=1000,
max_path_length=500,
replay_buffer=replay_buffer,
min_buffer_size=1e4,
target_update_tau=5e-3,
discount=0.99,
buffer_batch_size=256,
reward_scale=1.,
steps_per_epoch=1)
if torch.cuda.is_available():
set_gpu_mode(True)
else:
set_gpu_mode(False)
sac.to()
runner.setup(algo=sac, env=env, sampler_cls=LocalSampler)
runner.train(n_epochs=1000, batch_size=1000)
def test_sac_inverted_double_pendulum():
"""Test Sac performance on inverted pendulum."""
# pylint: disable=unexpected-keyword-arg
env = MetaRLEnv(normalize(gym.make('InvertedDoublePendulum-v2')))
deterministic.set_seed(0)
policy = TanhGaussianMLPPolicy(
env_spec=env.spec,
hidden_sizes=[32, 32],
hidden_nonlinearity=torch.nn.ReLU,
output_nonlinearity=None,
min_std=np.exp(-20.),
max_std=np.exp(2.),
)
qf1 = ContinuousMLPQFunction(env_spec=env.spec,
hidden_sizes=[32, 32],
hidden_nonlinearity=F.relu)
qf2 = ContinuousMLPQFunction(env_spec=env.spec,
hidden_sizes=[32, 32],
hidden_nonlinearity=F.relu)
replay_buffer = PathBuffer(capacity_in_transitions=int(1e6), )
runner = LocalRunner(snapshot_config=snapshot_config)
sac = SAC(env_spec=env.spec,
policy=policy,
qf1=qf1,
qf2=qf2,
gradient_steps_per_itr=100,
max_path_length=100,
replay_buffer=replay_buffer,
min_buffer_size=1e3,
target_update_tau=5e-3,
discount=0.99,
buffer_batch_size=64,
reward_scale=1.,
steps_per_epoch=2)
runner.setup(sac, env, sampler_cls=LocalSampler)
if torch.cuda.is_available():
set_gpu_mode(True)
else:
set_gpu_mode(False)
sac.to()
ret = runner.train(n_epochs=12, batch_size=200, plot=False)
# check that automatic entropy tuning is used
assert sac._use_automatic_entropy_tuning
# assert that there was a gradient properly connected to alpha
# this doesn't verify that the path from the temperature objective is
# correct.
assert not torch.allclose(torch.Tensor([1.]), sac._log_alpha.to('cpu'))
# check that policy is learning beyond predecided threshold
assert ret > 85
def run_task(snapshot_config, *_):
"""Set up environment and algorithm and run the task."""
runner = LocalRunner(snapshot_config)
env = MetaRLEnv(normalize(gym.make('HalfCheetah-v2')))
policy = TanhGaussianMLPPolicy2(
env_spec=env.spec,
hidden_sizes=[256, 256],
hidden_nonlinearity=nn.ReLU,
output_nonlinearity=None,
min_std=np.exp(-20.),
max_std=np.exp(2.),
)
qf1 = ContinuousMLPQFunction(env_spec=env.spec,
hidden_sizes=[256, 256],
hidden_nonlinearity=F.relu)
qf2 = ContinuousMLPQFunction(env_spec=env.spec,
hidden_sizes=[256, 256],
hidden_nonlinearity=F.relu)
replay_buffer = SimpleReplayBuffer(env_spec=env.spec,
size_in_transitions=int(1e6),
time_horizon=1)
# replay_buffer = SACReplayBuffer(env_spec=env.spec,
# max_size=int(1e6))
sampler_args = {
'agent': policy,
'max_path_length': 1000,
}
sac = SAC(env_spec=env.spec,
policy=policy,
qf1=qf1,
qf2=qf2,
gradient_steps_per_itr=1000,
use_automatic_entropy_tuning=True,
replay_buffer=replay_buffer,
min_buffer_size=1e4,
target_update_tau=5e-3,
discount=0.99,
buffer_batch_size=256,
reward_scale=1.)
runner.setup(algo=sac,
env=env,
sampler_cls=SimpleSampler,
sampler_args=sampler_args)
runner.train(n_epochs=1000, batch_size=1000)
def test_mtsac_get_log_alpha(monkeypatch):
"""Check that the private function _get_log_alpha functions correctly.
MTSAC uses disentangled alphas, meaning that
"""
env_names = ['CartPole-v0', 'CartPole-v1']
task_envs = [MetaRLEnv(env_name=name) for name in env_names]
env = MultiEnvWrapper(task_envs, sample_strategy=round_robin_strategy)
deterministic.set_seed(0)
policy = TanhGaussianMLPPolicy(
env_spec=env.spec,
hidden_sizes=[1, 1],
hidden_nonlinearity=torch.nn.ReLU,
output_nonlinearity=None,
min_std=np.exp(-20.),
max_std=np.exp(2.),
)
qf1 = ContinuousMLPQFunction(env_spec=env.spec,
hidden_sizes=[1, 1],
hidden_nonlinearity=F.relu)
qf2 = ContinuousMLPQFunction(env_spec=env.spec,
hidden_sizes=[1, 1],
hidden_nonlinearity=F.relu)
replay_buffer = PathBuffer(capacity_in_transitions=int(1e6), )
num_tasks = 2
buffer_batch_size = 2
mtsac = MTSAC(policy=policy,
qf1=qf1,
qf2=qf2,
gradient_steps_per_itr=150,
max_path_length=150,
eval_env=env,
env_spec=env.spec,
num_tasks=num_tasks,
steps_per_epoch=5,
replay_buffer=replay_buffer,
min_buffer_size=1e3,
target_update_tau=5e-3,
discount=0.99,
buffer_batch_size=buffer_batch_size)
monkeypatch.setattr(mtsac, '_log_alpha', torch.Tensor([1., 2.]))
for i, _ in enumerate(env_names):
obs = torch.Tensor([env.reset()] * buffer_batch_size)
log_alpha = mtsac._get_log_alpha(dict(observation=obs))
assert (log_alpha == torch.Tensor([i + 1, i + 1])).all().item()
assert log_alpha.size() == torch.Size([mtsac.buffer_batch_size])
def test_mtsac_inverted_double_pendulum():
"""Performance regression test of MTSAC on 2 InvDoublePendulum envs."""
env_names = ['InvertedDoublePendulum-v2', 'InvertedDoublePendulum-v2']
task_envs = [MetaRLEnv(env_name=name) for name in env_names]
env = MultiEnvWrapper(task_envs, sample_strategy=round_robin_strategy)
test_envs = MultiEnvWrapper(task_envs,
sample_strategy=round_robin_strategy)
deterministic.set_seed(0)
runner = LocalRunner(snapshot_config=snapshot_config)
policy = TanhGaussianMLPPolicy(
env_spec=env.spec,
hidden_sizes=[32, 32],
hidden_nonlinearity=torch.nn.ReLU,
output_nonlinearity=None,
min_std=np.exp(-20.),
max_std=np.exp(2.),
)
qf1 = ContinuousMLPQFunction(env_spec=env.spec,
hidden_sizes=[32, 32],
hidden_nonlinearity=F.relu)
qf2 = ContinuousMLPQFunction(env_spec=env.spec,
hidden_sizes=[32, 32],
hidden_nonlinearity=F.relu)
replay_buffer = PathBuffer(capacity_in_transitions=int(1e6), )
num_tasks = 2
buffer_batch_size = 128
mtsac = MTSAC(policy=policy,
qf1=qf1,
qf2=qf2,
gradient_steps_per_itr=100,
max_path_length=100,
eval_env=test_envs,
env_spec=env.spec,
num_tasks=num_tasks,
steps_per_epoch=5,
replay_buffer=replay_buffer,
min_buffer_size=1e3,
target_update_tau=5e-3,
discount=0.99,
buffer_batch_size=buffer_batch_size)
runner.setup(mtsac, env, sampler_cls=LocalSampler)
ret = runner.train(n_epochs=8, batch_size=128, plot=False)
assert ret > 130
def test_fixed_alpha():
"""Test if using fixed_alpha ensures that alpha is non differentiable."""
# pylint: disable=unexpected-keyword-arg
env = MetaRLEnv(normalize(gym.make('InvertedDoublePendulum-v2')))
deterministic.set_seed(0)
policy = TanhGaussianMLPPolicy(
env_spec=env.spec,
hidden_sizes=[32, 32],
hidden_nonlinearity=torch.nn.ReLU,
output_nonlinearity=None,
min_std=np.exp(-20.),
max_std=np.exp(2.),
)
qf1 = ContinuousMLPQFunction(env_spec=env.spec,
hidden_sizes=[32, 32],
hidden_nonlinearity=F.relu)
qf2 = ContinuousMLPQFunction(env_spec=env.spec,
hidden_sizes=[32, 32],
hidden_nonlinearity=F.relu)
replay_buffer = PathBuffer(capacity_in_transitions=int(1e6), )
runner = LocalRunner(snapshot_config=snapshot_config)
sac = SAC(env_spec=env.spec,
policy=policy,
qf1=qf1,
qf2=qf2,
gradient_steps_per_itr=100,
max_path_length=100,
replay_buffer=replay_buffer,
min_buffer_size=100,
target_update_tau=5e-3,
discount=0.99,
buffer_batch_size=64,
reward_scale=1.,
steps_per_epoch=1,
fixed_alpha=np.exp(0.5))
runner.setup(sac, env, sampler_cls=LocalSampler)
sac.to()
runner.train(n_epochs=1, batch_size=100, plot=False)
assert torch.allclose(torch.Tensor([0.5]), sac._log_alpha)
assert not sac._use_automatic_entropy_tuning
def test_pickling(self):
"""Test pickle and unpickle."""
net_size = 10
env_sampler = SetTaskSampler(PointEnv)
env = env_sampler.sample(5)
test_env_sampler = SetTaskSampler(PointEnv)
augmented_env = PEARL.augment_env_spec(env[0](), 5)
qf = ContinuousMLPQFunction(
env_spec=augmented_env,
hidden_sizes=[net_size, net_size, net_size])
vf_env = PEARL.get_env_spec(env[0](), 5, '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,
inner_policy=inner_policy,
qf=qf,
vf=vf,
num_train_tasks=5,
num_test_tasks=5,
latent_dim=5,
encoder_hidden_sizes=[10, 10],
test_env_sampler=test_env_sampler)
# This line is just to improve coverage
pearl.to()
pickled = pickle.dumps(pearl)
unpickled = pickle.loads(pickled)
assert hasattr(unpickled, '_replay_buffers')
assert hasattr(unpickled, '_context_replay_buffers')
assert unpickled._is_resuming
def test_output_shape(self, batch_size, hidden_sizes):
env_spec = TfEnv(DummyBoxEnv())
obs_dim = env_spec.observation_space.flat_dim
act_dim = env_spec.action_space.flat_dim
obs = torch.ones(batch_size, obs_dim, dtype=torch.float32)
act = torch.ones(batch_size, act_dim, dtype=torch.float32)
qf = ContinuousMLPQFunction(env_spec=env_spec,
hidden_nonlinearity=None,
hidden_sizes=hidden_sizes,
hidden_w_init=nn.init.ones_,
output_w_init=nn.init.ones_)
output = qf(obs, act)
assert output.shape == (batch_size, 1)
def ddpg_pendulum(ctxt=None, seed=1, lr=1e-4):
"""Train DDPG with InvertedDoublePendulum-v2 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.
lr (float): Learning rate for policy optimization.
"""
set_seed(seed)
runner = LocalRunner(ctxt)
env = MetaRLEnv(normalize(gym.make('InvertedDoublePendulum-v2')))
policy = DeterministicMLPPolicy(env_spec=env.spec,
hidden_sizes=[64, 64],
hidden_nonlinearity=F.relu,
output_nonlinearity=torch.tanh)
exploration_policy = AddOrnsteinUhlenbeckNoise(env.spec, policy, sigma=0.2)
qf = ContinuousMLPQFunction(env_spec=env.spec,
hidden_sizes=[64, 64],
hidden_nonlinearity=F.relu)
replay_buffer = PathBuffer(capacity_in_transitions=int(1e6))
policy_optimizer = (torch.optim.Adagrad, {'lr': lr, 'lr_decay': 0.99})
ddpg = DDPG(env_spec=env.spec,
policy=policy,
qf=qf,
replay_buffer=replay_buffer,
steps_per_epoch=20,
n_train_steps=50,
min_buffer_size=int(1e4),
exploration_policy=exploration_policy,
target_update_tau=1e-2,
discount=0.9,
policy_optimizer=policy_optimizer,
qf_optimizer=torch.optim.Adam)
runner.setup(algo=ddpg, env=env)
runner.train(n_epochs=500, batch_size=100)
def run_task(snapshot_config, *_):
"""Set up environment and algorithm and run the task.
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.
_ : Unused parameters
"""
runner = LocalRunner(snapshot_config)
env = MetaRLEnv(normalize(gym.make('InvertedDoublePendulum-v2')))
action_noise = OUStrategy(env.spec, sigma=0.2)
policy = DeterministicMLPPolicy(env_spec=env.spec,
hidden_sizes=[64, 64],
hidden_nonlinearity=F.relu,
output_nonlinearity=torch.tanh)
qf = ContinuousMLPQFunction(env_spec=env.spec,
hidden_sizes=[64, 64],
hidden_nonlinearity=F.relu)
replay_buffer = SimpleReplayBuffer(env_spec=env.spec,
size_in_transitions=int(1e6),
time_horizon=100)
policy_optimizer = (torch.optim.Adagrad, {'lr': 1e-4, 'lr_decay': 0.99})
ddpg = DDPG(env_spec=env.spec,
policy=policy,
qf=qf,
replay_buffer=replay_buffer,
steps_per_epoch=20,
n_train_steps=50,
min_buffer_size=int(1e4),
exploration_strategy=action_noise,
target_update_tau=1e-2,
discount=0.9,
policy_optimizer=policy_optimizer,
qf_optimizer=torch.optim.Adam)
runner.setup(algo=ddpg, env=env)
runner.train(n_epochs=500, batch_size=100)
def test_forward(self, hidden_sizes):
env_spec = TfEnv(DummyBoxEnv())
obs_dim = env_spec.observation_space.flat_dim
act_dim = env_spec.action_space.flat_dim
obs = torch.ones(obs_dim, dtype=torch.float32).unsqueeze(0)
act = torch.ones(act_dim, dtype=torch.float32).unsqueeze(0)
qf = ContinuousMLPQFunction(env_spec=env_spec,
hidden_nonlinearity=None,
hidden_sizes=hidden_sizes,
hidden_w_init=nn.init.ones_,
output_w_init=nn.init.ones_)
output = qf(obs, act)
expected_output = torch.full([1, 1],
fill_value=(obs_dim + act_dim) *
np.prod(hidden_sizes),
dtype=torch.float32)
assert torch.eq(output, expected_output)
def test_is_pickleable(self, hidden_sizes):
env_spec = TfEnv(DummyBoxEnv())
obs_dim = env_spec.observation_space.flat_dim
act_dim = env_spec.action_space.flat_dim
obs = torch.ones(obs_dim, dtype=torch.float32).unsqueeze(0)
act = torch.ones(act_dim, dtype=torch.float32).unsqueeze(0)
qf = ContinuousMLPQFunction(env_spec=env_spec,
hidden_nonlinearity=None,
hidden_sizes=hidden_sizes,
hidden_w_init=nn.init.ones_,
output_w_init=nn.init.ones_)
output1 = qf(obs, act)
p = pickle.dumps(qf)
qf_pickled = pickle.loads(p)
output2 = qf_pickled(obs, act)
assert torch.eq(output1, output2)
def test_ddpg_pendulum(self):
"""Test DDPG with Pendulum environment.
This environment has a [-3, 3] action_space bound.
"""
deterministic.set_seed(0)
runner = LocalRunner(snapshot_config)
env = MetaRLEnv(normalize(gym.make('InvertedPendulum-v2')))
policy = DeterministicMLPPolicy(env_spec=env.spec,
hidden_sizes=[64, 64],
hidden_nonlinearity=F.relu,
output_nonlinearity=torch.tanh)
exploration_policy = AddOrnsteinUhlenbeckNoise(env.spec,
policy,
sigma=0.2)
qf = ContinuousMLPQFunction(env_spec=env.spec,
hidden_sizes=[64, 64],
hidden_nonlinearity=F.relu)
replay_buffer = PathBuffer(capacity_in_transitions=int(1e6))
algo = DDPG(env_spec=env.spec,
policy=policy,
qf=qf,
replay_buffer=replay_buffer,
steps_per_epoch=20,
n_train_steps=50,
min_buffer_size=int(1e4),
exploration_policy=exploration_policy,
target_update_tau=1e-2,
discount=0.9)
runner.setup(algo, env)
last_avg_ret = runner.train(n_epochs=10, batch_size=100)
assert last_avg_ret > 10
env.close()
def test_ddpg_double_pendulum(self):
"""Test DDPG with Pendulum environment."""
deterministic.set_seed(0)
runner = LocalRunner(snapshot_config)
env = MetaRLEnv(gym.make('InvertedDoublePendulum-v2'))
action_noise = OUStrategy(env.spec, sigma=0.2)
policy = DeterministicMLPPolicy(env_spec=env.spec,
hidden_sizes=[64, 64],
hidden_nonlinearity=F.relu,
output_nonlinearity=torch.tanh)
qf = ContinuousMLPQFunction(env_spec=env.spec,
hidden_sizes=[64, 64],
hidden_nonlinearity=F.relu)
replay_buffer = SimpleReplayBuffer(env_spec=env.spec,
size_in_transitions=int(1e6),
time_horizon=100)
algo = DDPG(env_spec=env.spec,
policy=policy,
qf=qf,
replay_buffer=replay_buffer,
steps_per_epoch=20,
n_train_steps=50,
min_buffer_size=int(1e4),
exploration_strategy=action_noise,
target_update_tau=1e-2,
discount=0.9)
runner.setup(algo, env)
last_avg_ret = runner.train(n_epochs=10, batch_size=100)
assert last_avg_ret > 45
env.close()
def mt50_sac_normalize_all(ctxt=None, seed=1):
"""Set up environment and algorithm and run the task."""
runner = LocalRunner(ctxt)
envs = MT50.get_train_tasks(sample_all=True)
test_envs = MT50.get_test_tasks(sample_all=True)
MT50_envs_by_id = {
name: MetaRLEnv(
normalize(env,
normalize_reward=True,
normalize_obs=True,
flatten_obs=False))
for (name, env) in zip(envs._task_names, envs._task_envs)
}
MT50_envs_test = {
name: MetaRLEnv(normalize(env, normalize_obs=True, flatten_obs=False))
for (name, env) in zip(test_envs._task_names, test_envs._task_envs)
}
env = MTMetaWorldWrapper(MT50_envs_by_id)
policy = TanhGaussianMLPPolicy2(
env_spec=env.spec,
hidden_sizes=[400, 400, 400],
hidden_nonlinearity=nn.ReLU,
output_nonlinearity=None,
min_std=np.exp(-20.),
max_std=np.exp(2.),
)
qf1 = ContinuousMLPQFunction(env_spec=env.spec,
hidden_sizes=[400, 400, 400],
hidden_nonlinearity=F.relu)
qf2 = ContinuousMLPQFunction(env_spec=env.spec,
hidden_sizes=[400, 400, 400],
hidden_nonlinearity=F.relu)
replay_buffer = SACReplayBuffer(env_spec=env.spec, max_size=int(1e6))
sampler_args = {'agent': policy, 'max_path_length': 150}
timesteps = 100000000
batch_size = int(150 * env.num_tasks)
num_evaluation_points = 500
epochs = timesteps // batch_size
epoch_cycles = epochs // num_evaluation_points
epochs = epochs // epoch_cycles
sac = MTSAC(env=env,
eval_env_dict=MT50_envs_test,
env_spec=env.spec,
policy=policy,
qf1=qf1,
qf2=qf2,
gradient_steps_per_itr=250,
epoch_cycles=epoch_cycles,
use_automatic_entropy_tuning=True,
replay_buffer=replay_buffer,
min_buffer_size=7500,
target_update_tau=5e-3,
discount=0.99,
buffer_batch_size=6400)
tu.set_gpu_mode(True)
sac.to('cuda:0')
runner.setup(algo=sac,
env=env,
sampler_cls=SimpleSampler,
sampler_args=sampler_args)
runner.train(n_epochs=epochs, batch_size=batch_size)
def run_task(snapshot_config, *_):
"""Set up environment and algorithm and run the task.
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.
_ : Unused parameters
"""
# create multi-task environment and sample tasks
env_sampler = SetTaskSampler(
lambda: MetaRLEnv(normalize(ML1.get_train_tasks('push-v1'))))
env = env_sampler.sample(params['num_train_tasks'])
test_env_sampler = SetTaskSampler(
lambda: MetaRLEnv(normalize(ML1.get_test_tasks('push-v1'))))
test_env = test_env_sampler.sample(params['num_train_tasks'])
runner = LocalRunner(snapshot_config)
obs_dim = int(np.prod(env[0]().observation_space.shape))
action_dim = int(np.prod(env[0]().action_space.shape))
reward_dim = 1
# instantiate networks
encoder_in_dim = obs_dim + action_dim + reward_dim
encoder_out_dim = params['latent_size'] * 2
net_size = params['net_size']
context_encoder = MLPEncoder(input_dim=encoder_in_dim,
output_dim=encoder_out_dim,
hidden_sizes=[200, 200, 200])
space_a = akro.Box(low=-1,
high=1,
shape=(obs_dim + params['latent_size'], ),
dtype=np.float32)
space_b = akro.Box(low=-1, high=1, shape=(action_dim, ), dtype=np.float32)
augmented_env = EnvSpec(space_a, space_b)
qf1 = ContinuousMLPQFunction(env_spec=augmented_env,
hidden_sizes=[net_size, net_size, net_size])
qf2 = ContinuousMLPQFunction(env_spec=augmented_env,
hidden_sizes=[net_size, net_size, net_size])
obs_space = akro.Box(low=-1, high=1, shape=(obs_dim, ), dtype=np.float32)
action_space = akro.Box(low=-1,
high=1,
shape=(params['latent_size'], ),
dtype=np.float32)
vf_env = EnvSpec(obs_space, action_space)
vf = ContinuousMLPQFunction(env_spec=vf_env,
hidden_sizes=[net_size, net_size, net_size])
policy = TanhGaussianMLPPolicy2(
env_spec=augmented_env, hidden_sizes=[net_size, net_size, net_size])
context_conditioned_policy = ContextConditionedPolicy(
latent_dim=params['latent_size'],
context_encoder=context_encoder,
policy=policy,
use_ib=params['use_information_bottleneck'],
use_next_obs=params['use_next_obs_in_context'],
)
pearlsac = PEARLSAC(
env=env,
test_env=test_env,
policy=context_conditioned_policy,
qf1=qf1,
qf2=qf2,
vf=vf,
num_train_tasks=params['num_train_tasks'],
num_test_tasks=params['num_test_tasks'],
latent_dim=params['latent_size'],
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'],
num_evals=params['num_evals'],
num_steps_per_eval=params['num_steps_per_eval'],
batch_size=params['batch_size'],
embedding_batch_size=params['embedding_batch_size'],
embedding_mini_batch_size=params['embedding_mini_batch_size'],
max_path_length=params['max_path_length'],
reward_scale=params['reward_scale'],
)
tu.set_gpu_mode(params['use_gpu'], gpu_id=0)
if params['use_gpu']:
pearlsac.to()
runner.setup(algo=pearlsac,
env=env,
sampler_cls=PEARLSampler,
sampler_args=dict(max_path_length=params['max_path_length']))
runner.train(n_epochs=params['num_epochs'],
batch_size=params['batch_size'])
def pearl_metaworld_ml1_push(ctxt=None,
seed=1,
num_epochs=1000,
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=150,
reward_scale=10.,
use_gpu=False):
"""Train PEARL with ML1 environments.
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.
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)
# create multi-task environment and sample tasks
env_sampler = SetTaskSampler(lambda: MetaRLEnv(
normalize(mwb.ML1.get_train_tasks('push-v1'))))
env = env_sampler.sample(num_train_tasks)
test_env_sampler = SetTaskSampler(lambda: MetaRLEnv(
normalize(mwb.ML1.get_test_tasks('push-v1'))))
runner = LocalRunner(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,
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,
)
set_gpu_mode(use_gpu, gpu_id=0)
if use_gpu:
pearl.to()
runner.setup(algo=pearl,
env=env[0](),
sampler_cls=LocalSampler,
sampler_args=dict(max_path_length=max_path_length),
n_workers=1,
worker_class=PEARLWorker)
runner.train(n_epochs=num_epochs, batch_size=batch_size)
def mtsac_metaworld_mt50(ctxt=None, seed=1, use_gpu=False, _gpu=0):
"""Train MTSAC with MT50 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.
use_gpu (bool): Used to enable ussage of GPU in training.
_gpu (int): The ID of the gpu (used on multi-gpu machines).
"""
deterministic.set_seed(seed)
runner = LocalRunner(ctxt)
task_names = mwb.MT50.get_train_tasks().all_task_names
train_envs = []
test_envs = []
for task_name in task_names:
train_env = normalize(MetaRLEnv(mwb.MT50.from_task(task_name)),
normalize_reward=True)
test_env = normalize(MetaRLEnv(mwb.MT50.from_task(task_name)))
train_envs.append(train_env)
test_envs.append(test_env)
mt50_train_envs = MultiEnvWrapper(train_envs,
sample_strategy=round_robin_strategy,
mode='vanilla')
mt50_test_envs = MultiEnvWrapper(test_envs,
sample_strategy=round_robin_strategy,
mode='vanilla')
policy = TanhGaussianMLPPolicy(
env_spec=mt50_train_envs.spec,
hidden_sizes=[400, 400, 400],
hidden_nonlinearity=nn.ReLU,
output_nonlinearity=None,
min_std=np.exp(-20.),
max_std=np.exp(2.),
)
qf1 = ContinuousMLPQFunction(env_spec=mt50_train_envs.spec,
hidden_sizes=[400, 400, 400],
hidden_nonlinearity=F.relu)
qf2 = ContinuousMLPQFunction(env_spec=mt50_train_envs.spec,
hidden_sizes=[400, 400, 400],
hidden_nonlinearity=F.relu)
replay_buffer = PathBuffer(capacity_in_transitions=int(1e6), )
timesteps = 100000000
batch_size = int(150 * mt50_train_envs.num_tasks)
num_evaluation_points = 500
epochs = timesteps // batch_size
epoch_cycles = epochs // num_evaluation_points
epochs = epochs // epoch_cycles
mtsac = MTSAC(policy=policy,
qf1=qf1,
qf2=qf2,
gradient_steps_per_itr=150,
max_path_length=250,
eval_env=mt50_test_envs,
env_spec=mt50_train_envs.spec,
num_tasks=10,
steps_per_epoch=epoch_cycles,
replay_buffer=replay_buffer,
min_buffer_size=7500,
target_update_tau=5e-3,
discount=0.99,
buffer_batch_size=6400)
set_gpu_mode(use_gpu, _gpu)
mtsac.to()
runner.setup(algo=mtsac, env=mt50_train_envs, sampler_cls=LocalSampler)
runner.train(n_epochs=epochs, batch_size=batch_size)
def test_pearl_ml1_push(self):
"""Test PEARL with ML1 Push environment."""
params = dict(seed=1,
num_epochs=1,
num_train_tasks=5,
num_test_tasks=1,
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,
max_path_length=50,
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'])
env_sampler = SetTaskSampler(
lambda: MetaRLEnv(normalize(ML1.get_train_tasks('push-v1'))))
env = env_sampler.sample(params['num_train_tasks'])
test_env_sampler = SetTaskSampler(
lambda: MetaRLEnv(normalize(ML1.get_test_tasks('push-v1'))))
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'],
num_test_tasks=params['num_test_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'],
max_path_length=params['max_path_length'],
reward_scale=params['reward_scale'],
)
set_gpu_mode(params['use_gpu'], gpu_id=0)
if params['use_gpu']:
pearl.to()
runner = LocalRunner(snapshot_config)
runner.setup(
algo=pearl,
env=env[0](),
sampler_cls=LocalSampler,
sampler_args=dict(max_path_length=params['max_path_length']),
n_workers=1,
worker_class=PEARLWorker)
runner.train(n_epochs=params['num_epochs'],
batch_size=params['batch_size'])
def run_metarl(env, test_env, seed, log_dir):
"""Create metarl model and training."""
deterministic.set_seed(seed)
snapshot_config = SnapshotConfig(snapshot_dir=log_dir,
snapshot_mode='gap',
snapshot_gap=10)
runner = LocalRunner(snapshot_config)
obs_dim = int(np.prod(env[0]().observation_space.shape))
action_dim = int(np.prod(env[0]().action_space.shape))
reward_dim = 1
# instantiate networks
encoder_in_dim = obs_dim + action_dim + reward_dim
encoder_out_dim = params['latent_size'] * 2
net_size = params['net_size']
context_encoder = MLPEncoder(input_dim=encoder_in_dim,
output_dim=encoder_out_dim,
hidden_sizes=[200, 200, 200])
space_a = akro.Box(low=-1,
high=1,
shape=(obs_dim + params['latent_size'], ),
dtype=np.float32)
space_b = akro.Box(low=-1, high=1, shape=(action_dim, ), dtype=np.float32)
augmented_env = EnvSpec(space_a, space_b)
qf1 = ContinuousMLPQFunction(env_spec=augmented_env,
hidden_sizes=[net_size, net_size, net_size])
qf2 = ContinuousMLPQFunction(env_spec=augmented_env,
hidden_sizes=[net_size, net_size, net_size])
obs_space = akro.Box(low=-1, high=1, shape=(obs_dim, ), dtype=np.float32)
action_space = akro.Box(low=-1,
high=1,
shape=(params['latent_size'], ),
dtype=np.float32)
vf_env = EnvSpec(obs_space, action_space)
vf = ContinuousMLPQFunction(env_spec=vf_env,
hidden_sizes=[net_size, net_size, net_size])
policy = TanhGaussianMLPPolicy2(
env_spec=augmented_env, hidden_sizes=[net_size, net_size, net_size])
context_conditioned_policy = ContextConditionedPolicy(
latent_dim=params['latent_size'],
context_encoder=context_encoder,
policy=policy,
use_ib=params['use_information_bottleneck'],
use_next_obs=params['use_next_obs_in_context'],
)
train_task_names = ML10.get_train_tasks()._task_names
test_task_names = ML10.get_test_tasks()._task_names
pearlsac = PEARLSAC(
env=env,
test_env=test_env,
policy=context_conditioned_policy,
qf1=qf1,
qf2=qf2,
vf=vf,
num_train_tasks=params['num_train_tasks'],
num_test_tasks=params['num_test_tasks'],
latent_dim=params['latent_size'],
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'],
num_evals=params['num_evals'],
num_steps_per_eval=params['num_steps_per_eval'],
batch_size=params['batch_size'],
embedding_batch_size=params['embedding_batch_size'],
embedding_mini_batch_size=params['embedding_mini_batch_size'],
max_path_length=params['max_path_length'],
reward_scale=params['reward_scale'],
train_task_names=train_task_names,
test_task_names=test_task_names,
)
tu.set_gpu_mode(params['use_gpu'], gpu_id=0)
if params['use_gpu']:
pearlsac.to()
tabular_log_file = osp.join(log_dir, 'progress.csv')
tensorboard_log_dir = osp.join(log_dir)
dowel_logger.add_output(dowel.StdOutput())
dowel_logger.add_output(dowel.CsvOutput(tabular_log_file))
dowel_logger.add_output(dowel.TensorBoardOutput(tensorboard_log_dir))
runner.setup(algo=pearlsac,
env=env,
sampler_cls=PEARLSampler,
sampler_args=dict(max_path_length=params['max_path_length']))
runner.train(n_epochs=params['num_epochs'],
batch_size=params['batch_size'])
dowel_logger.remove_all()
return tabular_log_file
def mt10_sac(ctxt=None, seed=1):
"""Set up environment and algorithm and run the task."""
runner = LocalRunner(ctxt)
MT10_envs_by_id = {}
MT10_envs_test = {}
for (task, env) in EASY_MODE_CLS_DICT.items():
MT10_envs_by_id[task] = MetaRLEnv(
env(*EASY_MODE_ARGS_KWARGS[task]['args'],
**EASY_MODE_ARGS_KWARGS[task]['kwargs']))
# python 3.6 dicts are ordered
MT10_envs_test[task] = MetaRLEnv(
env(*EASY_MODE_ARGS_KWARGS[task]['args'],
**EASY_MODE_ARGS_KWARGS[task]['kwargs']))
env = IgnoreDoneWrapper(MTMetaWorldWrapper(MT10_envs_by_id))
policy = TanhGaussianMLPPolicy2(
env_spec=env.spec,
hidden_sizes=[400, 400, 400],
hidden_nonlinearity=nn.ReLU,
output_nonlinearity=None,
min_std=np.exp(-20.),
max_std=np.exp(2.),
)
qf1 = ContinuousMLPQFunction(env_spec=env.spec,
hidden_sizes=[400, 400, 400],
hidden_nonlinearity=F.relu)
qf2 = ContinuousMLPQFunction(env_spec=env.spec,
hidden_sizes=[400, 400, 400],
hidden_nonlinearity=F.relu)
replay_buffer = SACReplayBuffer(env_spec=env.spec, max_size=int(1e6))
sampler_args = {'agent': policy, 'max_path_length': 150}
timesteps = 20000000
batch_size = int(150 * env.num_tasks)
num_evaluation_points = 500
epochs = timesteps // batch_size
epoch_cycles = epochs // num_evaluation_points
epochs = epochs // epoch_cycles
sac = MTSAC(env=env,
eval_env_dict=MT10_envs_test,
env_spec=env.spec,
policy=policy,
qf1=qf1,
qf2=qf2,
gradient_steps_per_itr=150,
epoch_cycles=epoch_cycles,
use_automatic_entropy_tuning=True,
replay_buffer=replay_buffer,
min_buffer_size=1500,
target_update_tau=5e-3,
discount=0.99,
buffer_batch_size=1280)
tu.set_gpu_mode(True)
sac.to('cuda:0')
runner.setup(algo=sac,
env=env,
sampler_cls=SimpleSampler,
sampler_args=sampler_args)
runner.train(n_epochs=epochs, batch_size=batch_size)
def mtsac_metaworld_ml1_pick_place(ctxt=None, seed=1, _gpu=None):
"""Train MTSAC with the ML1 pick-place-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.
_gpu (int): The ID of the gpu to be used (used on multi-gpu machines).
"""
deterministic.set_seed(seed)
runner = LocalRunner(ctxt)
train_envs = []
test_envs = []
env_names = []
for i in range(50):
train_env = MetaRLEnv(
normalize(mwb.ML1.get_train_tasks('pick-place-v1'),
normalize_reward=True))
test_env = pickle.loads(pickle.dumps(train_env))
env_names.append('pick_place_{}'.format(i))
train_envs.append(train_env)
test_envs.append(test_env)
ml1_train_envs = MultiEnvWrapper(train_envs,
sample_strategy=round_robin_strategy,
env_names=env_names)
ml1_test_envs = MultiEnvWrapper(test_envs,
sample_strategy=round_robin_strategy,
env_names=env_names)
policy = TanhGaussianMLPPolicy(
env_spec=ml1_train_envs.spec,
hidden_sizes=[400, 400, 400],
hidden_nonlinearity=nn.ReLU,
output_nonlinearity=None,
min_std=np.exp(-20.),
max_std=np.exp(2.),
)
qf1 = ContinuousMLPQFunction(env_spec=ml1_train_envs.spec,
hidden_sizes=[400, 400, 400],
hidden_nonlinearity=F.relu)
qf2 = ContinuousMLPQFunction(env_spec=ml1_train_envs.spec,
hidden_sizes=[400, 400, 400],
hidden_nonlinearity=F.relu)
replay_buffer = PathBuffer(capacity_in_transitions=int(1e6), )
timesteps = 10000000
batch_size = int(150 * ml1_train_envs.num_tasks)
num_evaluation_points = 500
epochs = timesteps // batch_size
epoch_cycles = epochs // num_evaluation_points
epochs = epochs // epoch_cycles
mtsac = MTSAC(policy=policy,
qf1=qf1,
qf2=qf2,
gradient_steps_per_itr=150,
max_path_length=150,
eval_env=ml1_test_envs,
env_spec=ml1_train_envs.spec,
num_tasks=50,
steps_per_epoch=epoch_cycles,
replay_buffer=replay_buffer,
min_buffer_size=1500,
target_update_tau=5e-3,
discount=0.99,
buffer_batch_size=1280)
if _gpu is not None:
set_gpu_mode(True, _gpu)
mtsac.to()
runner.setup(algo=mtsac, env=ml1_train_envs, sampler_cls=LocalSampler)
runner.train(n_epochs=epochs, batch_size=batch_size)