Add tuning script for GAIL (PPO)

gail-experiment.py

@@ -0,0 +1,240 @@
+# %%
+import os
+
+import gym
+from src.core.gail import gail_ppo, Buffer
+from src.core.value import SetValue
+from src.core.policy import SetPolicy
+from src.core.discriminator import DeepsetDiscriminator
+import torch
+
+from intersim.envs import IntersimpleLidarFlatRandom
+from intersim.envs.intersimple import speed_reward
+import functools
+from src.util.wrappers import CollisionPenaltyWrapper, TransformObservation, Setobs
+import numpy as np
+from torch.utils.tensorboard import SummaryWriter
+from ray import tune
+from datetime import datetime
+import json
+
+DIR = os.path.dirname(os.path.abspath(__file__))
+activations = [torch.nn.Tanh, torch.nn.LeakyReLU]
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+def training_function(config):
+    np.random.seed(config['seed'])
+    torch.manual_seed(config['seed'])
+
+    if config['experiment'] == 'A':
+        envs = [Setobs(TransformObservation(CollisionPenaltyWrapper(
+            IntersimpleLidarFlatRandom(
+                n_rays=5,
+                reward=functools.partial(
+                    speed_reward,
+                    collision_penalty=0
+                ),
+                check_collisions=True,
+                stop_on_collision=config['trainenv']['stop_on_collision'],
+            ), collision_distance=6, collision_penalty=100),
+            lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10)
+        )) for _ in range(60)]
+
+    elif config['experiment'] == 'B': 
+        envs = sum([[Setobs(TransformObservation(CollisionPenaltyWrapper(
+            IntersimpleLidarFlatRandom(
+                n_rays=5,
+                reward=functools.partial(
+                    speed_reward,
+                    collision_penalty=0
+                ),
+                check_collisions=True,
+                stop_on_collision=config['trainenv']['stop_on_collision'],
+                track=track,
+            ), collision_distance=6, collision_penalty=100),
+            lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10)
+        )) for _ in range(15)] for track in range(4)],[])
+    
+    else:
+        raise NotImplementedError
+
+    env_fn = lambda i: envs[i]
+
+    policy = SetPolicy(env_fn(0).action_space.shape[0], 
+        n_hidden_layers=config['policy']['n_hidden_layers'],
+        hidden_layer_size=config['policy']['hidden_layer_size'],
+        activation=activations[config['policy']['activation']] ) # config net architecture
+    pi_opt = torch.optim.Adam(policy.parameters(), lr=config['policy']['learning_rate'])
+    pi_lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(pi_opt, gamma=config['policy']['learning_rate_decay'])
+
+    value = SetValue() # config net architecture
+    v_opt = torch.optim.Adam(value.parameters(), lr=config['value']['learning_rate'], weight_decay=config['value']['weight_decay'])
+
+    discriminator = DeepsetDiscriminator(
+        n_hidden_layers_element=config['discriminator']['n_hidden_layers_element'],
+        n_hidden_layers_global=config['discriminator']['n_hidden_layers_global'],
+        hidden_layer_size=config['discriminator']['hidden_layer_size'],
+        activation=activations[config['discriminator']['activation']],
+    )
+    disc_opt = torch.optim.Adam(discriminator.parameters(), lr=config['discriminator']['learning_rate'], weight_decay=config['discriminator']['weight_decay'])
+
+    if config['experiment'] == 'A': 
+        expert_data = torch.load(os.path.join(DIR, 'intersimple-expert-data-setobs2-loc0-track0.pt'))
+    elif config['experiment'] == 'B': 
+        expert_data = [
+            torch.load(os.path.join(DIR, 'intersimple-expert-data-setobs2-loc0-track0.pt')),
+            torch.load(os.path.join(DIR, 'intersimple-expert-data-setobs2-loc0-track1.pt')),
+            torch.load(os.path.join(DIR, 'intersimple-expert-data-setobs2-loc0-track2.pt')),
+            torch.load(os.path.join(DIR, 'intersimple-expert-data-setobs2-loc0-track3.pt')),
+        ]
+        d0 = [d[0] for d in expert_data]
+        d1 = [d[1] for d in expert_data]
+        d2 = [d[2] for d in expert_data]
+        d3 = [d[3] for d in expert_data]
+        expert_data = (torch.cat(d0), torch.cat(d1), torch.cat(d2), torch.cat(d3))
+    
+    expert_data = Buffer(*expert_data)
+
+    def callback(info):
+        tune.report(gen_mean_reward_per_episode=info['gen/mean_reward_per_episode'],
+                    disc_mean_reward_per_episode=info['disc/mean_reward_per_episode'], 
+                    mean_episode_length=info['gen/mean_episode_length'],
+                    gen_collision_rate=info['gen/collision_rate'])
+        
+        # save model checkpoints
+        ep = info['epoch'] + 1
+        if (ep % 25 == 0):
+            torch.save(info['policy'].state_dict(), f'policy_epoch{ep}.pt')
+
+    value, policy = gail_ppo(
+        env_fn=env_fn,
+        expert_data=expert_data,
+        discriminator=discriminator,
+        disc_opt=disc_opt,
+        disc_iters=config['discriminator']['iterations_per_epoch'],
+        policy=policy,
+        value=value,
+        v_opt=v_opt,
+        v_iters=config['value']['iterations_per_epoch'],
+        epochs=config['train_epochs'],
+        rollout_episodes=60, 
+        rollout_steps=200,
+        gamma=0.99,
+        gae_lambda=0.9,
+        clip_ratio=config['policy']['clip_ratio'],
+        pi_opt=pi_opt,
+        pi_iters=config['policy']['iterations_per_epoch'],
+        logger=SummaryWriter(comment='gail-ppo-options-setobs2'),
+        callback=callback,
+        lr_schedulers=[pi_lr_scheduler],
+    )
+
+    # save model
+    torch.save(policy.state_dict(), 'policy_final.pt')
+
+if __name__ == '__main__':
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--train', choices=['A', 'B'])
+    parser.add_argument('--epochs', type=int, default=200)
+    parser.add_argument('--test', type=str, help='path to config file to run final training on')
+    parser.add_argument('--test_seeds', type=int, default=5)
+    parser.add_argument('--test_cpus', type=int, help='number of cpus available to split test seed training over')
+    args = parser.parse_args()
+
+    assert (args.train is None) ^ (args.test is None), 'Must either train on an experiment or test with a config file'
+
+    # if no test config specified, train
+    if args.test is None:
+        print('Running Tuning for Experiment %s'%(args.train))
+        analysis = tune.run(
+            training_function,
+            config={
+                'experiment': args.train,
+                'trainenv': {
+                    'stop_on_collision': False, 
+                    'safe_actions_collision_method': 'circle', 
+                    'abort_unsafe_collision_method': 'circle',
+                },
+                'policy': {
+                    'learning_rate': 3e-4, 
+                    'learning_rate_decay': 1.0, 
+                    'clip_ratio': 0.2, 
+                    'iterations_per_epoch': 100,
+                    'hidden_layer_size': tune.grid_search([20, 40]),
+                    'n_hidden_layers': tune.grid_search([2, 3]), 
+                    'activation':0,
+                },
+                'value': {
+                    'learning_rate': 1e-4,
+                    'weight_decay': 1e-3, 
+                    'iterations_per_epoch': 1000, 
+                },
+                'discriminator': {
+                    'learning_rate': 1e-3, 
+                    'weight_decay': 1e-5, 
+                    'iterations_per_epoch': 500, 
+                    'n_hidden_layers_element': tune.grid_search([3,4]),
+                    'n_hidden_layers_global': tune.grid_search([1,2]),
+                    'hidden_layer_size': 10, 
+                    'activation': 0,
+                },
+                'train_epochs': args.epochs,
+                'seed': 0,
+            }
+        )
+        best_config = analysis.get_best_config(metric='gen_collision_rate', mode='min')
+        print('Best config: ', best_config)
+
+        # safe best_config
+        if not os.path.isdir(os.path.join(DIR, 'best_configs')):
+            os.mkdir(os.path.join(DIR, 'best_configs'))
+        
+        # save gail
+        with open(os.path.join(DIR, 'best_configs',f'gail_exp{args.train}.json'), 'w', encoding='utf-8') as f:
+            json.dump(best_config, f, ensure_ascii=False, indent=4)
+
+    # if config file specified, rerun it with appropriate number of seeds
+    else:
+        with open(args.test, 'rb') as f:
+            config = json.load(f)
+
+        print(f'Retraining {args.test} with {args.test_seeds} seeds on experiment {config["experiment"]}')
+
+        # rerun with appropriate number of seeds
+        rpt = {'cpu': int(args.test_cpus/args.test_seeds)} if (args.test_cpus is not None) else None
+        config['seed'] = tune.grid_search(list(range(1,args.test_seeds+1)))
+        analysis = tune.run(training_function, config=config, resources_per_trial=rpt)
+
+        # move final policies to appropriate directory
+        split_ = os.path.basename(args.test).split('_')
+        model = split_[0]
+        exper = split_[-1].split('.')[0]
+        savepath = os.path.join('test_policies',model,exper)
+        
+        if not os.path.isdir(savepath):
+            os.makedirs(savepath)
+        
+        import shutil
+        for i in range(args.test_seeds):
+            s = analysis._checkpoints[i]['config']['seed']
+            check_dir = analysis._checkpoints[i]['logdir']
+            shutil.copyfile(os.path.join(check_dir,'policy_final.pt'), 
+                os.path.join(savepath, f'policy_seed{s}.pt'))

src/core/gail.py

@@ -30,7 +30,7 @@ def roll_buffer(buffer, *args, **kwargs):
 
 def gail(env_fn, expert_data, discriminator, disc_opt, disc_iters, policy, value,
          v_opt, v_iters, epochs, rollout_episodes, rollout_steps, gamma,
-         gae_lambda, delta, backtrack_coeff, backtrack_iters, cg_iters=10, cg_damping=0.1, wasserstein=False, wasserstein_c=None, logger=TerminalLogger()):
+         gae_lambda, delta, backtrack_coeff, backtrack_iters, cg_iters=10, cg_damping=0.1, wasserstein=False, wasserstein_c=None, logger=TerminalLogger(), callback=None, lr_schedulers=[]):
 
     policy(torch.zeros(env_fn(0).observation_space.shape))
     policy = ReparamPolicy(policy)
@@ -39,10 +39,15 @@ def gail(env_fn, expert_data, discriminator, disc_opt, disc_iters, policy, value
     logger.add_scalar('expert/mean_reward_per_episode', expert_data.rewards[~expert_data.dones].sum() / expert_data.states.shape[0])
 
     for epoch in tqdm(range(epochs)):
-        generator_data = Buffer(*rollout(env_fn, policy, rollout_episodes, rollout_steps))
+        states, actions, rewards, dones, collisions = rollout(env_fn, policy, rollout_episodes, rollout_steps)
+        generator_data = Buffer(states, actions, rewards, dones)
 
-        logger.add_scalar('gen/mean_episode_length', (~generator_data.dones).sum() / generator_data.states.shape[0], epoch)
+        gen_mean_episode_length = (~generator_data.dones).sum() / generator_data.states.shape[0]
-        logger.add_scalar('gen/mean_reward_per_episode', generator_data.rewards[~generator_data.dones].sum() / generator_data.states.shape[0], epoch)
+        logger.add_scalar('gen/mean_episode_length', gen_mean_episode_length, epoch)
+        gen_mean_reward_per_episode = generator_data.rewards[~generator_data.dones].sum() / generator_data.states.shape[0]
+        logger.add_scalar('gen/mean_reward_per_episode', gen_mean_reward_per_episode, epoch)
+        gen_collision_rate = (1. * collisions.any(-1)).mean()
+        logger.add_scalar('gen/collision_rate', gen_collision_rate, epoch)
 
         discriminator, loss = train_discriminator(expert_data, generator_data, discriminator, disc_opt, disc_iters, wasserstein, wasserstein_c)
         if wasserstein:
@@ -50,25 +55,45 @@ def gail(env_fn, expert_data, discriminator, disc_opt, disc_iters, policy, value
         else:
             generator_data.rewards = -F.logsigmoid(discriminator(generator_data.states, generator_data.actions))
         logger.add_scalar('disc/final_loss', loss, epoch)
-        logger.add_scalar('disc/mean_reward_per_episode', generator_data.rewards[~generator_data.dones].sum() / generator_data.states.shape[0], epoch)
+        disc_mean_reward_per_episode = generator_data.rewards[~generator_data.dones].sum() / generator_data.states.shape[0]
+        logger.add_scalar('disc/mean_reward_per_episode', disc_mean_reward_per_episode, epoch)
 
         value, policy = trpo_step(value, policy, generator_data.states, generator_data.actions, generator_data.rewards, generator_data.dones, gamma, gae_lambda, delta, backtrack_coeff, backtrack_iters, v_opt, v_iters, cg_iters, cg_damping)
         expert_data = roll_buffer(expert_data, shifts=-3, dims=0)
 
+        if callback is not None:
+            callback({
+                'epoch': epoch,
+                'value': value,
+                'policy': policy,
+                'gen/mean_episode_length': gen_mean_episode_length.item(),
+                'gen/mean_reward_per_episode': gen_mean_reward_per_episode.item(),
+                'gen/collision_rate': gen_collision_rate.item(),
+                'disc/mean_reward_per_episode': disc_mean_reward_per_episode.item(),
+            })
+    
+    for lr_scheduler in lr_schedulers:
+        lr_scheduler.step()
+    
     return value, policy
 
 def gail_ppo(env_fn, expert_data, discriminator, disc_opt, disc_iters, policy, value,
          v_opt, v_iters, epochs, rollout_episodes, rollout_steps, gamma,
-         gae_lambda, clip_ratio, pi_opt, pi_iters, target_kl=None, max_grad_norm=None, wasserstein=False, wasserstein_c=None, logger=TerminalLogger()):
+         gae_lambda, clip_ratio, pi_opt, pi_iters, target_kl=None, max_grad_norm=None, wasserstein=False, wasserstein_c=None, logger=TerminalLogger(), callback=None, lr_schedulers=[]):
 
     logger.add_scalar('expert/mean_episode_length', (~expert_data.dones).sum() / expert_data.states.shape[0])
     logger.add_scalar('expert/mean_reward_per_episode', expert_data.rewards[~expert_data.dones].sum() / expert_data.states.shape[0])
 
     for epoch in range(epochs):
-        generator_data = Buffer(*rollout(env_fn, policy, rollout_episodes, rollout_steps))
+        states, actions, rewards, dones, collisions = rollout(env_fn, policy, rollout_episodes, rollout_steps)
+        generator_data = Buffer(states, actions, rewards, dones)
 
-        logger.add_scalar('gen/mean_episode_length', (~generator_data.dones).sum() / generator_data.states.shape[0], epoch)
+        gen_mean_episode_length = (~generator_data.dones).sum() / generator_data.states.shape[0]
-        logger.add_scalar('gen/mean_reward_per_episode', generator_data.rewards[~generator_data.dones].sum() / generator_data.states.shape[0], epoch)
+        logger.add_scalar('gen/mean_episode_length', gen_mean_episode_length, epoch)
+        gen_mean_reward_per_episode = generator_data.rewards[~generator_data.dones].sum() / generator_data.states.shape[0]
+        logger.add_scalar('gen/mean_reward_per_episode', gen_mean_reward_per_episode, epoch)
+        gen_collision_rate = (1. * collisions.any(-1)).mean()
+        logger.add_scalar('gen/collision_rate', gen_collision_rate, epoch)
 
         discriminator, loss = train_discriminator(expert_data, generator_data, discriminator, disc_opt, disc_iters, wasserstein, wasserstein_c)
         if wasserstein:
@@ -76,11 +101,26 @@ def gail_ppo(env_fn, expert_data, discriminator, disc_opt, disc_iters, policy, v
         else:
             generator_data.rewards = -F.logsigmoid(discriminator(generator_data.states, generator_data.actions))
         logger.add_scalar('disc/final_loss', loss, epoch)
-        logger.add_scalar('disc/mean_reward_per_episode', generator_data.rewards[~generator_data.dones].sum() / generator_data.states.shape[0], epoch)
+        disc_mean_reward_per_episode =  generator_data.rewards[~generator_data.dones].sum() / generator_data.states.shape[0]
+        logger.add_scalar('disc/mean_reward_per_episode', disc_mean_reward_per_episode, epoch)
 
         value, policy = ppo_step(value, policy, generator_data.states, generator_data.actions, generator_data.rewards, generator_data.dones, clip_ratio, gamma, gae_lambda, pi_opt, pi_iters, v_opt, v_iters, target_kl, max_grad_norm)
         expert_data = roll_buffer(expert_data, shifts=-3, dims=0)
 
+        if callback is not None:
+            callback({
+                'epoch': epoch,
+                'value': value,
+                'policy': policy,
+                'gen/mean_episode_length': gen_mean_episode_length.item(),
+                'gen/mean_reward_per_episode': gen_mean_reward_per_episode.item(),
+                'gen/collision_rate': gen_collision_rate.item(),
+                'disc/mean_reward_per_episode': disc_mean_reward_per_episode.item(),
+            })
+    
+    for lr_scheduler in lr_schedulers:
+        lr_scheduler.step()
+    
     return value, policy
 
 def train_discriminator(expert_data, generator_data, discriminator, disc_opt, disc_iters, wasserstein, wasserstein_c=None):

src/core/sampling.py

@@ -2,6 +2,7 @@ import torch
 import gym
 from stable_baselines3.common.vec_env import DummyVecEnv as VecEnv
 from tqdm import tqdm
+import numpy as np
 
 def rollout(env_fn, policy, n_episodes, max_steps_per_episode):
     env = env_fn(0)
@@ -9,23 +10,27 @@ def rollout(env_fn, policy, n_episodes, max_steps_per_episode):
     actions = torch.zeros(n_episodes, max_steps_per_episode + 1, *env.action_space.shape)
     rewards = torch.zeros(n_episodes, max_steps_per_episode + 1)
     dones = torch.ones(n_episodes, max_steps_per_episode + 1, dtype=bool)
+    collisions = torch.zeros(n_episodes, max_steps_per_episode, dtype=bool)
 
     env = VecEnv(list(map(lambda i: (lambda: env_fn(i)), range(n_episodes))))
 
     states[:, 0] = torch.tensor(env.reset()).clone().detach()
     dones[:, 0] = False
 
-    for s in range(max_steps_per_episode):
+    for s in tqdm(range(max_steps_per_episode), 'Rollout'):
         actions[:, s] = policy.sample(policy(states[:, s])).clone().detach()
 
         clipped_actions = actions[:, s]
         if isinstance(env.action_space, gym.spaces.Box):
             clipped_actions = torch.clamp(clipped_actions, torch.from_numpy(env.action_space.low), torch.from_numpy(env.action_space.high))
 
-        o, r, d, _ = env.step(clipped_actions)
+        o, r, d, info = env.step(clipped_actions)
         states[:, s + 1] = torch.tensor(o).clone().detach()
         rewards[:, s] = torch.tensor(r).clone().detach()
         dones[:, s + 1] = torch.tensor(d).clone().detach()
+        collisions[:, s] = torch.from_numpy(np.stack([
+            i['collision'] for i in info
+        ])).detach().clone()
 
     dones = dones.cumsum(1) > 0
 
@@ -34,7 +39,7 @@ def rollout(env_fn, policy, n_episodes, max_steps_per_episode):
     rewards = rewards[:, :max_steps_per_episode]
     dones = dones[:, :max_steps_per_episode]
 
-    return states, actions, rewards, dones
+    return states, actions, rewards, dones, collisions
 
 
 def rollout_sb3(env, policy, n_episodes, max_steps_per_episode):