diff --git a/scratch/etienne/trpo/experiments/sgail-ppo-options-setobs2.py b/scratch/etienne/trpo/experiments/sgail-ppo-options-setobs2.py index 11bb889..52a05a0 100644 --- a/scratch/etienne/trpo/experiments/sgail-ppo-options-setobs2.py +++ b/scratch/etienne/trpo/experiments/sgail-ppo-options-setobs2.py @@ -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) + 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, - policy=policy, - value=value, - v_opt=v_opt, - v_iters=1000, - epochs=200, - rollout_episodes=60, - rollout_steps=60, - gamma=0.99, - gae_lambda=0.9, - clip_ratio=0.2, - pi_opt=pi_opt, - pi_iters=100, - logger=SummaryWriter(comment='sgail-ppo-options-setobs2'), - callback=callback, - lr_schedulers=[pi_lr_scheduler], + value, policy = gail_ppo( + env_fn=env_fn, + expert_data=expert_data, + discriminator=discriminator, + disc_opt=disc_opt, + disc_iters=100, # config + policy=policy, + value=value, + v_opt=v_opt, + v_iters=1000, # config + epochs=200, + rollout_episodes=60, + rollout_steps=60, + gamma=0.99, + gae_lambda=0.9, + clip_ratio=0.2, # config + pi_opt=pi_opt, + pi_iters=100, # config + logger=SummaryWriter(comment='sgail-ppo-options-setobs2'), + callback=callback, + lr_schedulers=[pi_lr_scheduler], + ) + +analysis = tune.run( + training_function, + config={ + 'dummy': tune.grid_search([0.001, 0.01, 0.1]), + } ) -torch.save(policy.state_dict(), 'sgail-ppo-options-setobs2.pt') +print('Best config: ', analysis.get_best_config(metric='gen_mean_reward_per_episode', mode='min')) # %% policy = SetMaskedDiscretePolicy(env_fn(0).action_space.n) diff --git a/src/safe_options/options.py b/src/safe_options/options.py index 1042735..e92128b 100644 --- a/src/safe_options/options.py +++ b/src/safe_options/options.py @@ -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()