huangfu/InteractionImitation
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Set up for ray tune

Browse Source
This commit is contained in:
ebuehrle
2022-02-23 16:55:13 +01:00
parent 91f88983b0
commit f037c119cc
2 changed files with 58 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files

40
scratch/etienne/trpo/experiments/sgail-ppo-options-setobs2.py
View File

@@ -15,6 +15,7 @@ from src.util.wrappers import CollisionPenaltyWrapper, TransformObservation, Set
import numpy as np
from src.safe_options.options import SafeOptionsEnv
from torch.utils.tensorboard import SummaryWriter
from ray import tune
obs_min = np.array([
[-1000, -1000, 0, -np.pi, -1e-1, 0.],
@@ -47,49 +48,54 @@ envs = [SafeOptionsEnv(Setobs(
env_fn = lambda i: envs[i]
policy = SetMaskedDiscretePolicy(env_fn(0).action_space.n)
pi_opt = torch.optim.Adam(policy.parameters(), lr=3e-4)
pi_lr_scheduler = torch.optim.lr_scheduler.StepLR(pi_opt, step_size=50, gamma=0.2)
def training_function(config):
policy = SetMaskedDiscretePolicy(env_fn(0).action_space.n) # config net architecture
pi_opt = torch.optim.Adam(policy.parameters(), lr=3e-5) # config learning rate
pi_lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(pi_opt, gamma=0.98) # config lr decay
value = SetValue()
v_opt = torch.optim.Adam(value.parameters(), lr=1e-3)
value = SetValue() # config net architecture
v_opt = torch.optim.Adam(value.parameters(), lr=1e-3) # config lr
discriminator = DeepsetDiscriminator()
disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3, weight_decay=1e-4)
discriminator = DeepsetDiscriminator() # config net architecture
disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3, weight_decay=1e-3) # config lr, weight decay
expert_data = torch.load('intersimple-expert-data-setobs2.pt')
expert_data = Buffer(*expert_data)
# %%
def callback(epoch, value, policy):
if not epoch % 10:
torch.save(policy.state_dict(), f'sgail-ppo-options-setobs2-{epoch}.pt')
torch.save(value.state_dict(), f'sgail-ppo-options-setobs2-value-{epoch}.pt')
def callback(info):
tune.report(gen_mean_reward_per_episode=info['gen/mean_reward_per_episode'])
value, policy = gail_ppo(
env_fn=env_fn,
expert_data=expert_data,
discriminator=discriminator,
disc_opt=disc_opt,
disc_iters=100,
disc_iters=100, # config
policy=policy,
value=value,
v_opt=v_opt,
v_iters=1000,
v_iters=1000, # config
epochs=200,
rollout_episodes=60,
rollout_steps=60,
gamma=0.99,
gae_lambda=0.9,
clip_ratio=0.2,
clip_ratio=0.2, # config
pi_opt=pi_opt,
pi_iters=100,
pi_iters=100, # config
logger=SummaryWriter(comment='sgail-ppo-options-setobs2'),
callback=callback,
lr_schedulers=[pi_lr_scheduler],
)
torch.save(policy.state_dict(), 'sgail-ppo-options-setobs2.pt')
analysis = tune.run(
training_function,
config={
'dummy': tune.grid_search([0.001, 0.01, 0.1]),
}
)
print('Best config: ', analysis.get_best_config(metric='gen_mean_reward_per_episode', mode='min'))
# %%
policy = SetMaskedDiscretePolicy(env_fn(0).action_space.n)

20
src/safe_options/options.py
View File

@@ -53,7 +53,8 @@ def gail(env_fn, expert_data, discriminator, disc_opt, disc_iters, policy, value
generator_data.ll.actions += 0.1 * torch.randn_like(generator_data.ll.actions)
logger.add_scalar('gen/mean_episode_length', (~generator_data.ll.dones).sum() / generator_data.ll.states.shape[0], epoch)
logger.add_scalar('gen/mean_reward_per_episode', generator_data.hl.rewards[~generator_data.hl.dones].sum() / generator_data.hl.states.shape[0], epoch)
gen_mean_reward_per_episode = generator_data.hl.rewards[~generator_data.hl.dones].sum() / generator_data.hl.states.shape[0]
logger.add_scalar('gen/mean_reward_per_episode', gen_mean_reward_per_episode, epoch)
logger.add_scalar('gen/unsafe_probability_mass', policy.unsafe_probability_mass(policy(generator_data.hl.states[~generator_data.hl.dones], generator_data.hl.safe_actions[~generator_data.hl.dones])).mean(), epoch)
discriminator, loss = train_discriminator(expert_data, generator_data.ll, discriminator, disc_opt, disc_iters, wasserstein, wasserstein_c)
@@ -71,7 +72,12 @@ def gail(env_fn, expert_data, discriminator, disc_opt, disc_iters, policy, value
expert_data = roll_buffer(expert_data, shifts=-3, dims=0)
if callback is not None:
callback(epoch, value, policy)
callback({
'epoch': epoch,
'value': value,
'policy': policy,
'gen/mean_reward_per_episode': gen_mean_reward_per_episode,
})
return value, policy
@@ -89,7 +95,8 @@ def gail_ppo(env_fn, expert_data, discriminator, disc_opt, disc_iters, policy, v
generator_data.ll.actions += 0.1 * torch.randn_like(generator_data.ll.actions)
logger.add_scalar('gen/mean_episode_length', (~generator_data.ll.dones).sum() / generator_data.ll.states.shape[0], epoch)
logger.add_scalar('gen/mean_reward_per_episode', generator_data.hl.rewards[~generator_data.hl.dones].sum() / generator_data.hl.states.shape[0], epoch)
gen_mean_reward_per_episode = generator_data.hl.rewards[~generator_data.hl.dones].sum() / generator_data.hl.states.shape[0]
logger.add_scalar('gen/mean_reward_per_episode', gen_mean_reward_per_episode, epoch)
logger.add_scalar('gen/unsafe_probability_mass', policy.unsafe_probability_mass(policy(generator_data.hl.states[~generator_data.hl.dones], generator_data.hl.safe_actions[~generator_data.hl.dones])).mean(), epoch)
discriminator, loss = train_discriminator(expert_data, generator_data.ll, discriminator, disc_opt, disc_iters, wasserstein, wasserstein_c)
@@ -107,7 +114,12 @@ def gail_ppo(env_fn, expert_data, discriminator, disc_opt, disc_iters, policy, v
expert_data = roll_buffer(expert_data, shifts=-3, dims=0)
if callback is not None:
callback(epoch, value, policy)
callback({
'epoch': epoch,
'value': value,
'policy': policy,
'gen/mean_reward_per_episode': gen_mean_reward_per_episode,
})
for lr_scheduler in lr_schedulers:
lr_scheduler.step()