def main(
env_name='AntDirection-v1',
adapt_lr=0.1,
meta_lr=1.0,
adapt_steps=1,
num_iterations=1000,
meta_bsz=40,
adapt_bsz=20,
tau=1.00,
gamma=0.99,
seed=42,
num_workers=2,
cuda=0,
):
cuda = bool(cuda)
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
if cuda:
torch.cuda.manual_seed(seed)
def make_env():
env = gym.make(env_name)
env = ch.envs.ActionSpaceScaler(env)
return env
env = l2l.gym.AsyncVectorEnv([make_env for _ in range(num_workers)])
env.seed(seed)
env.set_task(env.sample_tasks(1)[0])
env = ch.envs.Torch(env)
policy = DiagNormalPolicy(env.state_size, env.action_size)
if cuda:
policy.to('cuda')
baseline = LinearValue(env.state_size, env.action_size)
for iteration in range(num_iterations):
iteration_reward = 0.0
iteration_replays = []
iteration_policies = []
for task_config in tqdm(env.sample_tasks(meta_bsz),
leave=False,
desc='Data'): # Samples a new config
clone = deepcopy(policy)
env.set_task(task_config)
env.reset()
task = ch.envs.Runner(env)
task_replay = []
# Fast Adapt
for step in range(adapt_steps):
train_episodes = task.run(clone, episodes=adapt_bsz)
clone = fast_adapt_a2c(clone,
train_episodes,
adapt_lr,
baseline,
gamma,
tau,
first_order=True)
task_replay.append(train_episodes)
# Compute Validation Loss
valid_episodes = task.run(clone, episodes=adapt_bsz)
task_replay.append(valid_episodes)
iteration_reward += valid_episodes.reward().sum().item(
) / adapt_bsz
iteration_replays.append(task_replay)
iteration_policies.append(clone)
# Print statistics
print('\nIteration', iteration)
adaptation_reward = iteration_reward / meta_bsz
print('adaptation_reward', adaptation_reward)
# TRPO meta-optimization
backtrack_factor = 0.5
ls_max_steps = 15
max_kl = 0.01
if cuda:
policy.to('cuda', non_blocking=True)
baseline.to('cuda', non_blocking=True)
iteration_replays = [[
r.to('cuda', non_blocking=True) for r in task_replays
] for task_replays in iteration_replays]
# Compute CG step direction
old_loss, old_kl = meta_surrogate_loss(iteration_replays,
iteration_policies, policy,
baseline, tau, gamma, adapt_lr)
grad = autograd.grad(old_loss, policy.parameters(), retain_graph=True)
grad = parameters_to_vector([g.detach() for g in grad])
Fvp = trpo.hessian_vector_product(old_kl, policy.parameters())
step = trpo.conjugate_gradient(Fvp, grad)
shs = 0.5 * torch.dot(step, Fvp(step))
lagrange_multiplier = torch.sqrt(shs / max_kl)
step = step / lagrange_multiplier
step_ = [torch.zeros_like(p.data) for p in policy.parameters()]
vector_to_parameters(step, step_)
step = step_
del old_kl, Fvp, grad
old_loss.detach_()
# Line-search
for ls_step in range(ls_max_steps):
stepsize = backtrack_factor**ls_step * meta_lr
clone = deepcopy(policy)
for p, u in zip(clone.parameters(), step):
p.data.add_(-stepsize, u.data)
new_loss, kl = meta_surrogate_loss(iteration_replays,
iteration_policies, clone,
baseline, tau, gamma, adapt_lr)
if new_loss < old_loss and kl < max_kl:
for p, u in zip(policy.parameters(), step):
p.data.add_(-stepsize, u.data)
break
def main(
experiment='dev',
env_name='Particles2D-v1',
adapt_lr=0.1,
meta_lr=0.01,
adapt_steps=1,
num_iterations=200,
meta_bsz=20,
adapt_bsz=20,
tau=1.00,
gamma=0.99,
num_workers=2,
seed=42,
):
random.seed(seed)
np.random.seed(seed)
th.manual_seed(seed)
def make_env():
return gym.make(env_name)
env = l2l.gym.AsyncVectorEnv([make_env for _ in range(num_workers)])
env.seed(seed)
env = ch.envs.Torch(env)
policy = DiagNormalPolicy(env.state_size, env.action_size)
meta_learner = l2l.algorithms.MetaSGD(policy, lr=meta_lr)
baseline = LinearValue(env.state_size, env.action_size)
opt = optim.Adam(policy.parameters(), lr=meta_lr)
all_rewards = []
for iteration in range(num_iterations):
iteration_loss = 0.0
iteration_reward = 0.0
for task_config in tqdm(
env.sample_tasks(meta_bsz)): # Samples a new config
learner = meta_learner.clone()
env.set_task(task_config)
env.reset()
task = ch.envs.Runner(env)
# Fast Adapt
for step in range(adapt_steps):
train_episodes = task.run(learner, episodes=adapt_bsz)
loss = maml_a2c_loss(train_episodes, learner, baseline, gamma,
tau)
learner.adapt(loss)
# Compute Validation Loss
valid_episodes = task.run(learner, episodes=adapt_bsz)
loss = maml_a2c_loss(valid_episodes, learner, baseline, gamma, tau)
iteration_loss += loss
iteration_reward += valid_episodes.reward().sum().item(
) / adapt_bsz
# Print statistics
print('\nIteration', iteration)
adaptation_reward = iteration_reward / meta_bsz
print('adaptation_reward', adaptation_reward)
all_rewards.append(adaptation_reward)
adaptation_loss = iteration_loss / meta_bsz
print('adaptation_loss', adaptation_loss.item())
opt.zero_grad()
adaptation_loss.backward()
opt.step()
def main(
benchmark=ML10, # Choose between ML1, ML10, ML45
adapt_lr=0.1,
meta_lr=0.1,
adapt_steps=1,
num_iterations=1000,
meta_bsz=20,
adapt_bsz=10, # Number of episodes to sample per task
tau=1.00,
gamma=0.99,
seed=42,
num_workers=10, # Currently tasks are distributed evenly so adapt_bsz should be divisible by num_workers
cuda=0):
env = make_env(benchmark, seed, num_workers)
cuda = bool(cuda)
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
if cuda:
torch.cuda.manual_seed(seed)
policy = DiagNormalPolicy(env.state_size,
env.action_size,
activation='tanh')
if cuda:
policy.to('cuda')
baseline = LinearValue(env.state_size, env.action_size)
for iteration in range(num_iterations):
iteration_reward = 0.0
iteration_replays = []
iteration_policies = []
for task_config in tqdm(env.sample_tasks(meta_bsz),
leave=False,
desc='Data'): # Samples a new config
clone = deepcopy(policy)
env.set_task(task_config)
env.reset()
task = ch.envs.Runner(env)
task_replay = []
# Fast Adapt
for step in range(adapt_steps):
train_episodes = task.run(clone, episodes=adapt_bsz)
clone = fast_adapt_a2c(clone,
train_episodes,
adapt_lr,
baseline,
gamma,
tau,
first_order=True)
task_replay.append(train_episodes)
# Compute Validation Loss
valid_episodes = task.run(clone, episodes=adapt_bsz)
task_replay.append(valid_episodes)
iteration_reward += valid_episodes.reward().sum().item(
) / adapt_bsz
iteration_replays.append(task_replay)
iteration_policies.append(clone)
# Print statistics
print('\nIteration', iteration)
validation_reward = iteration_reward / meta_bsz
print('validation_reward', validation_reward)
# TRPO meta-optimization
backtrack_factor = 0.5
ls_max_steps = 15
max_kl = 0.01
if cuda:
policy.to('cuda', non_blocking=True)
baseline.to('cuda', non_blocking=True)
iteration_replays = [[
r.to('cuda', non_blocking=True) for r in task_replays
] for task_replays in iteration_replays]
# Compute CG step direction
old_loss, old_kl = meta_surrogate_loss(iteration_replays,
iteration_policies, policy,
baseline, tau, gamma, adapt_lr)
grad = autograd.grad(old_loss, policy.parameters(), retain_graph=True)
grad = parameters_to_vector([g.detach() for g in grad])
Fvp = trpo.hessian_vector_product(old_kl, policy.parameters())
step = trpo.conjugate_gradient(Fvp, grad)
shs = 0.5 * torch.dot(step, Fvp(step))
lagrange_multiplier = torch.sqrt(shs / max_kl)
step = step / lagrange_multiplier
step_ = [torch.zeros_like(p.data) for p in policy.parameters()]
vector_to_parameters(step, step_)
step = step_
del old_kl, Fvp, grad
old_loss.detach_()
# Line-search
for ls_step in range(ls_max_steps):
stepsize = backtrack_factor**ls_step * meta_lr
clone = deepcopy(policy)
for p, u in zip(clone.parameters(), step):
p.data.add_(-stepsize, u.data)
new_loss, kl = meta_surrogate_loss(iteration_replays,
iteration_policies, clone,
baseline, tau, gamma, adapt_lr)
if new_loss < old_loss and kl < max_kl:
for p, u in zip(policy.parameters(), step):
p.data.add_(-stepsize, u.data)
break
# Evaluate on a set of unseen tasks
evaluate(benchmark, policy, baseline, adapt_lr, gamma, tau, num_workers,
seed)