updating experiment B vecenv

ToDo: parameterize network architectures

.gitignore

@@ -1,3 +1,10 @@
+*.png
+*.pkl
+*.pt
+*.zip
+**/ray/*
+**/runs/*
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]
@@ -119,6 +126,7 @@ venv.bak/
 
 # VS Code project settings
 .project
+.vscode
 
 # mkdocs documentation
 /site
@@ -140,6 +148,8 @@ expert_data/
 
 # Results
 experiments/results/
+output/
 
 # Dependencies
 InteractionSimulator/
+imitation/

README.md

@@ -24,8 +24,26 @@ Once the repository has been set up, you can process and save expert track demon
 ```
 python src/expert_data.py --loc [LOCNUM] --track [TRACKNUM]
 ```
+You can (and should) process all tracks at once at location 0 with:
+```
+python src/expert_data.py --all-tracks
+```
 
-You can then load the experts actions and observations using
+You can then train a default behavior cloning policy with the following. Be sure to check help for main.py for running options.
+```
+python src/main.py --train
+```
+You can run tensorboard by running the following and opening `localhost:6006` (or alternatively port-forwarding 6006 from the remote server)
+```
+tensorboard --logdir output/
+```
+You can then test the learned policy with the following, and see the animation file in `output/`:
+```
+python src/main.py --test
+```
+
+
+You can load the experts actions manually
 ```
 from src import expert_data
 observations, actions = expert_data.load_expert_data(loc = [LOCNUM], track = [TRACKNUM])

bc-experiment.py

@@ -0,0 +1,212 @@
+# %%
+import os
+
+from tqdm import tqdm
+from src.core.sampling import rollout
+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'])
+
+    # choose validation environment
+    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'],
+            ), 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]
+
+    # load expert data
+
+    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)
+
+    # configure and train policy
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+    policy = SetPolicy(expert_data.actions.shape[-1], 
+        n_hidden_layers=config['policy']['n_hidden_layers'],
+        hidden_layer_size=config['policy']['hidden_layer_size'],
+        activation=activations[config['policy']['activation']] ) # config net architecture
+    policy = policy.to(device)
+    
+    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'])
+
+    expert_states = expert_data.states[~expert_data.dones].to(device)
+    expert_actions = expert_data.actions[~expert_data.dones].to(device)
+
+    for epoch in range(config['train_epochs']):
+        pi_opt.zero_grad()
+        loss = -policy.log_prob(policy(expert_states), expert_actions).mean()
+        loss.backward()
+        pi_opt.step()
+        pi_lr_scheduler.step()
+
+        if epoch % 25 == 0:
+            gen_states, gen_actions, gen_rewards, gen_dones, gen_collisions = rollout(env_fn, policy.cpu(), n_episodes=60, max_steps_per_episode=200)
+            gen_mean_episode_length = (~gen_dones).sum() / gen_states.shape[0]
+            gen_mean_reward_per_episode = gen_rewards[~gen_dones].sum() / gen_states.shape[0]
+            gen_collision_rate = (1. * gen_collisions.any(-1)).mean()
+
+            tune.report(
+                gen_mean_reward_per_episode=gen_mean_reward_per_episode.item(),
+                mean_episode_length=gen_mean_episode_length.item(),
+                gen_collision_rate=gen_collision_rate.item(),
+                loss=loss.item(),
+            )
+        
+        # save model checkpoints
+        ep = epoch + 1
+        if (ep % 50 == 0):
+            torch.save(policy.state_dict(), f'policy_epoch{ep}.pt')
+
+    # 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=1000)
+    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, 
+                },
+                'policy': {
+                    'learning_rate': 3e-4, 
+                    'learning_rate_decay': tune.grid_search([0.001, 1.0]), 
+                    'hidden_layer_size': tune.grid_search([10, 20, 40, 80]),
+                    'n_hidden_layers': tune.grid_search([2, 3, 4]), 
+                    'activation':0,
+                },
+                'train_epochs': args.epochs,
+                'seed': 0,
+            }
+            # TODO resources_per_trial={'gpu': 1}
+        )
+        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'bc_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'))

best_configs/bc_expA.json

@@ -0,0 +1,15 @@
+{
+    "experiment": "A",
+    "trainenv": {
+        "stop_on_collision": false
+    },
+    "policy": {
+        "learning_rate": 0.0003,
+        "learning_rate_decay": 1.0,
+        "hidden_layer_size": 40,
+        "n_hidden_layers": 2,
+        "activation": 0
+    },
+    "train_epochs": 300,
+    "seed": 0
+}

best_configs/bc_expB.json

@@ -0,0 +1,15 @@
+{
+    "experiment": "B",
+    "trainenv": {
+        "stop_on_collision": false
+    },
+    "policy": {
+        "learning_rate": 0.0003,
+        "learning_rate_decay": 1.0,
+        "hidden_layer_size": 40,
+        "n_hidden_layers": 2,
+        "activation": 0
+    },
+    "train_epochs": 300,
+    "seed": 0
+}

best_configs/gail-trpo_expA.json

@@ -0,0 +1,31 @@
+{
+    "experiment": "A",
+    "trainenv": {
+        "stop_on_collision": false
+    },
+    "policy": {
+        "learning_rate": 0.0003,
+        "learning_rate_decay": 1.0,
+        "delta": 0.01,
+        "iterations_per_epoch": 100,
+        "hidden_layer_size": 40,
+        "n_hidden_layers": 2,
+        "activation": 0
+    },
+    "value": {
+        "learning_rate": 0.0001,
+        "weight_decay": 0.001,
+        "iterations_per_epoch": 1000
+    },
+    "discriminator": {
+        "learning_rate": 0.001,
+        "weight_decay": 0.0001,
+        "iterations_per_epoch": 100,
+        "n_hidden_layers_element": 4,
+        "n_hidden_layers_global": 1,
+        "hidden_layer_size": 10,
+        "activation": 0
+    },
+    "train_epochs": 100,
+    "seed": 0
+}

best_configs/gail-trpo_expB.json

@@ -0,0 +1,31 @@
+{
+    "experiment": "B",
+    "trainenv": {
+        "stop_on_collision": false
+    },
+    "policy": {
+        "learning_rate": 0.0003,
+        "learning_rate_decay": 1.0,
+        "delta": 0.01,
+        "iterations_per_epoch": 100,
+        "hidden_layer_size": 40,
+        "n_hidden_layers": 2,
+        "activation": 0
+    },
+    "value": {
+        "learning_rate": 0.0001,
+        "weight_decay": 0.001,
+        "iterations_per_epoch": 1000
+    },
+    "discriminator": {
+        "learning_rate": 0.001,
+        "weight_decay": 0.0001,
+        "iterations_per_epoch": 100,
+        "n_hidden_layers_element": 4,
+        "n_hidden_layers_global": 1,
+        "hidden_layer_size": 10,
+        "activation": 0
+    },
+    "train_epochs": 100,
+    "seed": 0
+}

best_configs/gail_expA.json

@@ -0,0 +1,31 @@
+{
+    "experiment": "A",
+    "trainenv": {
+        "stop_on_collision": false
+    },
+    "policy": {
+        "learning_rate": 0.0003,
+        "learning_rate_decay": 1.0,
+        "clip_ratio": 0.2,
+        "iterations_per_epoch": 100,
+        "hidden_layer_size": 40,
+        "n_hidden_layers": 2,
+        "activation": 0
+    },
+    "value": {
+        "learning_rate": 0.0001,
+        "weight_decay": 0.001,
+        "iterations_per_epoch": 1000
+    },
+    "discriminator": {
+        "learning_rate": 0.001,
+        "weight_decay": 0.0001,
+        "iterations_per_epoch": 100,
+        "n_hidden_layers_element": 4,
+        "n_hidden_layers_global": 1,
+        "hidden_layer_size": 10,
+        "activation": 0
+    },
+    "train_epochs": 100,
+    "seed": 0
+}

best_configs/gail_expB.json

@@ -0,0 +1,31 @@
+{
+    "experiment": "B",
+    "trainenv": {
+        "stop_on_collision": false
+    },
+    "policy": {
+        "learning_rate": 0.0003,
+        "learning_rate_decay": 1.0,
+        "clip_ratio": 0.2,
+        "iterations_per_epoch": 100,
+        "hidden_layer_size": 40,
+        "n_hidden_layers": 2,
+        "activation": 0
+    },
+    "value": {
+        "learning_rate": 0.0001,
+        "weight_decay": 0.001,
+        "iterations_per_epoch": 1000
+    },
+    "discriminator": {
+        "learning_rate": 0.001,
+        "weight_decay": 0.0001,
+        "iterations_per_epoch": 100,
+        "n_hidden_layers_element": 4,
+        "n_hidden_layers_global": 1,
+        "hidden_layer_size": 10,
+        "activation": 0
+    },
+    "train_epochs": 100,
+    "seed": 0
+}

best_configs/hail-etienne2_expA.json

@@ -0,0 +1,33 @@
+{
+    "experiment": "A",
+    "trainenv": {
+        "stop_on_collision": false,
+        "safe_actions_collision_method": null,
+        "abort_unsafe_collision_method": null
+    },
+    "policy": {
+        "learning_rate": 0.0003,
+        "learning_rate_decay": 1.0,
+        "clip_ratio": 0.2,
+        "iterations_per_epoch": 100,
+        "hidden_layer_size": 10,
+        "n_hidden_layers": 3,
+        "activation": 0,
+        "option": 0
+    },
+    "value": {
+        "learning_rate": 0.001,
+        "iterations_per_epoch": 1000
+    },
+    "discriminator": {
+        "learning_rate": 0.001,
+        "weight_decay": 0.0001,
+        "iterations_per_epoch": 100,
+        "n_hidden_layers_element": 3,
+        "n_hidden_layers_global": 2,
+        "hidden_layer_size": 10,
+        "activation": 0
+    },
+    "train_epochs": 100,
+    "seed": 0
+}

best_configs/hail-etienne2_expB.json

@@ -0,0 +1,33 @@
+{
+    "experiment": "B",
+    "trainenv": {
+        "stop_on_collision": false,
+        "safe_actions_collision_method": null,
+        "abort_unsafe_collision_method": null
+    },
+    "policy": {
+        "learning_rate": 0.0003,
+        "learning_rate_decay": 1.0,
+        "clip_ratio": 0.2,
+        "iterations_per_epoch": 100,
+        "hidden_layer_size": 10,
+        "n_hidden_layers": 3,
+        "activation": 0,
+        "option": 0
+    },
+    "value": {
+        "learning_rate": 0.001,
+        "iterations_per_epoch": 1000
+    },
+    "discriminator": {
+        "learning_rate": 0.001,
+        "weight_decay": 0.0001,
+        "iterations_per_epoch": 100,
+        "n_hidden_layers_element": 3,
+        "n_hidden_layers_global": 2,
+        "hidden_layer_size": 10,
+        "activation": 0
+    },
+    "train_epochs": 100,
+    "seed": 0
+}

best_configs/hail-etienne3_expA.json

@@ -0,0 +1,33 @@
+{
+    "experiment": "A",
+    "trainenv": {
+        "stop_on_collision": false,
+        "safe_actions_collision_method": null,
+        "abort_unsafe_collision_method": null
+    },
+    "policy": {
+        "learning_rate": 0.0003,
+        "learning_rate_decay": 1.0,
+        "clip_ratio": 0.2,
+        "iterations_per_epoch": 100,
+        "hidden_layer_size": 40,
+        "n_hidden_layers": 3,
+        "activation": 0,
+        "option": 0
+    },
+    "value": {
+        "learning_rate": 0.001,
+        "iterations_per_epoch": 1000
+    },
+    "discriminator": {
+        "learning_rate": 0.001,
+        "weight_decay": 0.0001,
+        "iterations_per_epoch": 100,
+        "n_hidden_layers_element": 3,
+        "n_hidden_layers_global": 2,
+        "hidden_layer_size": 10,
+        "activation": 0
+    },
+    "train_epochs": 100,
+    "seed": 0
+}

best_configs/hail-etienne3_expB.json

@@ -0,0 +1,33 @@
+{
+    "experiment": "B",
+    "trainenv": {
+        "stop_on_collision": false,
+        "safe_actions_collision_method": null,
+        "abort_unsafe_collision_method": null
+    },
+    "policy": {
+        "learning_rate": 0.0003,
+        "learning_rate_decay": 1.0,
+        "clip_ratio": 0.2,
+        "iterations_per_epoch": 100,
+        "hidden_layer_size": 40,
+        "n_hidden_layers": 3,
+        "activation": 0,
+        "option": 0
+    },
+    "value": {
+        "learning_rate": 0.001,
+        "iterations_per_epoch": 1000
+    },
+    "discriminator": {
+        "learning_rate": 0.001,
+        "weight_decay": 0.0001,
+        "iterations_per_epoch": 100,
+        "n_hidden_layers_element": 3,
+        "n_hidden_layers_global": 2,
+        "hidden_layer_size": 10,
+        "activation": 0
+    },
+    "train_epochs": 100,
+    "seed": 0
+}

best_configs/hail-etienne_expA.json

@@ -0,0 +1,33 @@
+{
+    "experiment": "A",
+    "trainenv": {
+        "stop_on_collision": false,
+        "safe_actions_collision_method": null,
+        "abort_unsafe_collision_method": null
+    },
+    "policy": {
+        "learning_rate": 0.0003,
+        "learning_rate_decay": 1.0,
+        "clip_ratio": 0.2,
+        "iterations_per_epoch": 100,
+        "hidden_layer_size": 20,
+        "n_hidden_layers": 4,
+        "activation": 0,
+        "option": 0
+    },
+    "value": {
+        "learning_rate": 0.001,
+        "iterations_per_epoch": 1000
+    },
+    "discriminator": {
+        "learning_rate": 0.001,
+        "weight_decay": 0.0001,
+        "iterations_per_epoch": 100,
+        "n_hidden_layers_element": 3,
+        "n_hidden_layers_global": 2,
+        "hidden_layer_size": 10,
+        "activation": 0
+    },
+    "train_epochs": 100,
+    "seed": 0
+}

best_configs/hail-etienne_expB.json

@@ -0,0 +1,33 @@
+{
+    "experiment": "B",
+    "trainenv": {
+        "stop_on_collision": false,
+        "safe_actions_collision_method": null,
+        "abort_unsafe_collision_method": null
+    },
+    "policy": {
+        "learning_rate": 0.0003,
+        "learning_rate_decay": 1.0,
+        "clip_ratio": 0.2,
+        "iterations_per_epoch": 100,
+        "hidden_layer_size": 20,
+        "n_hidden_layers": 4,
+        "activation": 0,
+        "option": 0
+    },
+    "value": {
+        "learning_rate": 0.001,
+        "iterations_per_epoch": 1000
+    },
+    "discriminator": {
+        "learning_rate": 0.001,
+        "weight_decay": 0.0001,
+        "iterations_per_epoch": 100,
+        "n_hidden_layers_element": 3,
+        "n_hidden_layers_global": 2,
+        "hidden_layer_size": 10,
+        "activation": 0
+    },
+    "train_epochs": 100,
+    "seed": 0
+}

best_configs/hail_expA.json

@@ -0,0 +1,33 @@
+{
+    "experiment": "A",
+    "trainenv": {
+        "stop_on_collision": false,
+        "safe_actions_collision_method": null,
+        "abort_unsafe_collision_method": null
+    },
+    "policy": {
+        "learning_rate": 0.0003,
+        "learning_rate_decay": 1.0,
+        "clip_ratio": 0.2,
+        "iterations_per_epoch": 100,
+        "hidden_layer_size": 40,
+        "n_hidden_layers": 2,
+        "activation": 0,
+        "option": 0
+    },
+    "value": {
+        "learning_rate": 0.001,
+        "iterations_per_epoch": 1000
+    },
+    "discriminator": {
+        "learning_rate": 0.001,
+        "weight_decay": 0.0001,
+        "iterations_per_epoch": 100,
+        "n_hidden_layers_element": 4,
+        "n_hidden_layers_global": 1,
+        "hidden_layer_size": 10,
+        "activation": 0
+    },
+    "train_epochs": 90,
+    "seed": 0
+}

best_configs/hail_expB.json

@@ -0,0 +1,33 @@
+{
+    "experiment": "B",
+    "trainenv": {
+        "stop_on_collision": false,
+        "safe_actions_collision_method": null,
+        "abort_unsafe_collision_method": null
+    },
+    "policy": {
+        "learning_rate": 0.0003,
+        "learning_rate_decay": 1.0,
+        "clip_ratio": 0.2,
+        "iterations_per_epoch": 100,
+        "hidden_layer_size": 20,
+        "n_hidden_layers": 2,
+        "activation": 0,
+        "option": 0
+    },
+    "value": {
+        "learning_rate": 0.001,
+        "iterations_per_epoch": 1000
+    },
+    "discriminator": {
+        "learning_rate": 0.001,
+        "weight_decay": 0.0001,
+        "iterations_per_epoch": 100,
+        "n_hidden_layers_element": 4,
+        "n_hidden_layers_global": 2,
+        "hidden_layer_size": 10,
+        "activation": 0
+    },
+    "train_epochs": 85,
+    "seed": 0
+}

best_configs/shail-etienne2_expA.json

@@ -0,0 +1,33 @@
+{
+    "experiment": "A",
+    "trainenv": {
+        "stop_on_collision": false,
+        "safe_actions_collision_method": "circle",
+        "abort_unsafe_collision_method": "circle"
+    },
+    "policy": {
+        "learning_rate": 0.0003,
+        "learning_rate_decay": 1.0,
+        "clip_ratio": 0.2,
+        "iterations_per_epoch": 100,
+        "hidden_layer_size": 10,
+        "n_hidden_layers": 3,
+        "activation": 0,
+        "option": 0
+    },
+    "value": {
+        "learning_rate": 0.001,
+        "iterations_per_epoch": 1000
+    },
+    "discriminator": {
+        "learning_rate": 0.001,
+        "weight_decay": 0.0001,
+        "iterations_per_epoch": 100,
+        "n_hidden_layers_element": 3,
+        "n_hidden_layers_global": 2,
+        "hidden_layer_size": 10,
+        "activation": 0
+    },
+    "train_epochs": 100,
+    "seed": 0
+}

best_configs/shail-etienne2_expB.json

@@ -0,0 +1,33 @@
+{
+    "experiment": "B",
+    "trainenv": {
+        "stop_on_collision": false,
+        "safe_actions_collision_method": "circle",
+        "abort_unsafe_collision_method": "circle"
+    },
+    "policy": {
+        "learning_rate": 0.0003,
+        "learning_rate_decay": 1.0,
+        "clip_ratio": 0.2,
+        "iterations_per_epoch": 100,
+        "hidden_layer_size": 10,
+        "n_hidden_layers": 3,
+        "activation": 0,
+        "option": 0
+    },
+    "value": {
+        "learning_rate": 0.001,
+        "iterations_per_epoch": 1000
+    },
+    "discriminator": {
+        "learning_rate": 0.001,
+        "weight_decay": 0.0001,
+        "iterations_per_epoch": 100,
+        "n_hidden_layers_element": 3,
+        "n_hidden_layers_global": 2,
+        "hidden_layer_size": 10,
+        "activation": 0
+    },
+    "train_epochs": 100,
+    "seed": 0
+}

best_configs/shail-etienne3_expA.json

@@ -0,0 +1,33 @@
+{
+    "experiment": "A",
+    "trainenv": {
+        "stop_on_collision": false,
+        "safe_actions_collision_method": "circle",
+        "abort_unsafe_collision_method": "circle"
+    },
+    "policy": {
+        "learning_rate": 0.0003,
+        "learning_rate_decay": 1.0,
+        "clip_ratio": 0.2,
+        "iterations_per_epoch": 100,
+        "hidden_layer_size": 40,
+        "n_hidden_layers": 3,
+        "activation": 0,
+        "option": 0
+    },
+    "value": {
+        "learning_rate": 0.001,
+        "iterations_per_epoch": 1000
+    },
+    "discriminator": {
+        "learning_rate": 0.001,
+        "weight_decay": 0.0001,
+        "iterations_per_epoch": 100,
+        "n_hidden_layers_element": 3,
+        "n_hidden_layers_global": 2,
+        "hidden_layer_size": 10,
+        "activation": 0
+    },
+    "train_epochs": 100,
+    "seed": 0
+}

best_configs/shail-etienne3_expB.json

@@ -0,0 +1,33 @@
+{
+    "experiment": "B",
+    "trainenv": {
+        "stop_on_collision": false,
+        "safe_actions_collision_method": "circle",
+        "abort_unsafe_collision_method": "circle"
+    },
+    "policy": {
+        "learning_rate": 0.0003,
+        "learning_rate_decay": 1.0,
+        "clip_ratio": 0.2,
+        "iterations_per_epoch": 100,
+        "hidden_layer_size": 40,
+        "n_hidden_layers": 3,
+        "activation": 0,
+        "option": 0
+    },
+    "value": {
+        "learning_rate": 0.001,
+        "iterations_per_epoch": 1000
+    },
+    "discriminator": {
+        "learning_rate": 0.001,
+        "weight_decay": 0.0001,
+        "iterations_per_epoch": 100,
+        "n_hidden_layers_element": 3,
+        "n_hidden_layers_global": 2,
+        "hidden_layer_size": 10,
+        "activation": 0
+    },
+    "train_epochs": 100,
+    "seed": 0
+}

best_configs/shail-etienne4_expA.json

@@ -0,0 +1,33 @@
+{
+    "experiment": "A",
+    "trainenv": {
+        "stop_on_collision": false,
+        "safe_actions_collision_method": "circle",
+        "abort_unsafe_collision_method": "circle"
+    },
+    "policy": {
+        "learning_rate": 0.0003,
+        "learning_rate_decay": 1.0,
+        "clip_ratio": 0.2,
+        "iterations_per_epoch": 500,
+        "hidden_layer_size": 40,
+        "n_hidden_layers": 3,
+        "activation": 0,
+        "option": 0
+    },
+    "value": {
+        "learning_rate": 0.001,
+        "iterations_per_epoch": 1000
+    },
+    "discriminator": {
+        "learning_rate": 0.001,
+        "weight_decay": 0.0001,
+        "iterations_per_epoch": 100,
+        "n_hidden_layers_element": 3,
+        "n_hidden_layers_global": 2,
+        "hidden_layer_size": 10,
+        "activation": 0
+    },
+    "train_epochs": 150,
+    "seed": 0
+}

best_configs/shail-etienne_expA.json

@@ -0,0 +1,33 @@
+{
+    "experiment": "A",
+    "trainenv": {
+        "stop_on_collision": false,
+        "safe_actions_collision_method": "circle",
+        "abort_unsafe_collision_method": "circle"
+    },
+    "policy": {
+        "learning_rate": 0.0003,
+        "learning_rate_decay": 1.0,
+        "clip_ratio": 0.2,
+        "iterations_per_epoch": 100,
+        "hidden_layer_size": 20,
+        "n_hidden_layers": 4,
+        "activation": 0,
+        "option": 0
+    },
+    "value": {
+        "learning_rate": 0.001,
+        "iterations_per_epoch": 1000
+    },
+    "discriminator": {
+        "learning_rate": 0.001,
+        "weight_decay": 0.0001,
+        "iterations_per_epoch": 100,
+        "n_hidden_layers_element": 3,
+        "n_hidden_layers_global": 2,
+        "hidden_layer_size": 10,
+        "activation": 0
+    },
+    "train_epochs": 100,
+    "seed": 0
+}

best_configs/shail-etienne_expB.json

@@ -0,0 +1,33 @@
+{
+    "experiment": "B",
+    "trainenv": {
+        "stop_on_collision": false,
+        "safe_actions_collision_method": "circle",
+        "abort_unsafe_collision_method": "circle"
+    },
+    "policy": {
+        "learning_rate": 0.0003,
+        "learning_rate_decay": 1.0,
+        "clip_ratio": 0.2,
+        "iterations_per_epoch": 100,
+        "hidden_layer_size": 20,
+        "n_hidden_layers": 4,
+        "activation": 0,
+        "option": 0
+    },
+    "value": {
+        "learning_rate": 0.001,
+        "iterations_per_epoch": 1000
+    },
+    "discriminator": {
+        "learning_rate": 0.001,
+        "weight_decay": 0.0001,
+        "iterations_per_epoch": 100,
+        "n_hidden_layers_element": 3,
+        "n_hidden_layers_global": 2,
+        "hidden_layer_size": 10,
+        "activation": 0
+    },
+    "train_epochs": 100,
+    "seed": 0
+}

best_configs/shail_expA.json

@@ -0,0 +1,33 @@
+{
+    "experiment": "A",
+    "trainenv": {
+        "stop_on_collision": false,
+        "safe_actions_collision_method": "circle",
+        "abort_unsafe_collision_method": "circle"
+    },
+    "policy": {
+        "learning_rate": 0.0003,
+        "learning_rate_decay": 1.0,
+        "clip_ratio": 0.2,
+        "iterations_per_epoch": 100,
+        "hidden_layer_size": 40,
+        "n_hidden_layers": 2,
+        "activation": 0,
+        "option": 0
+    },
+    "value": {
+        "learning_rate": 0.001,
+        "iterations_per_epoch": 1000
+    },
+    "discriminator": {
+        "learning_rate": 0.001,
+        "weight_decay": 0.0001,
+        "iterations_per_epoch": 100,
+        "n_hidden_layers_element": 4,
+        "n_hidden_layers_global": 1,
+        "hidden_layer_size": 10,
+        "activation": 0
+    },
+    "train_epochs": 90,
+    "seed": 0
+}

best_configs/shail_expB.json

@@ -0,0 +1,33 @@
+{
+    "experiment": "B",
+    "trainenv": {
+        "stop_on_collision": false,
+        "safe_actions_collision_method": "circle",
+        "abort_unsafe_collision_method": "circle"
+    },
+    "policy": {
+        "learning_rate": 0.0003,
+        "learning_rate_decay": 1.0,
+        "clip_ratio": 0.2,
+        "iterations_per_epoch": 100,
+        "hidden_layer_size": 20,
+        "n_hidden_layers": 2,
+        "activation": 0,
+        "option": 0
+    },
+    "value": {
+        "learning_rate": 0.001,
+        "iterations_per_epoch": 1000
+    },
+    "discriminator": {
+        "learning_rate": 0.001,
+        "weight_decay": 0.0001,
+        "iterations_per_epoch": 100,
+        "n_hidden_layers_element": 4,
+        "n_hidden_layers_global": 2,
+        "hidden_layer_size": 10,
+        "activation": 0
+    },
+    "train_epochs": 85,
+    "seed": 0
+}

config/networks.json5

@@ -1,34 +1,42 @@
 {
-    ego_state: {
+    ego_encoder: {
         input_dim: 5, // number of state vars
-        hidden_n: 1,
+        hidden_n: 0,
         hidden_dim: 5,
         output_dim: 5
     },
     deepsets: {
-        input_dim: 5, // number of relative state vars for others
+        input_dim: 6, // number of relative state vars for others
         phi: {
-            hidden_n: 1,
+            hidden_n: 2,
             hidden_dim: 20,
         },
         latent_dim: 20,
         rho: {
-            hidden_n: 1,
+            hidden_n: 2,
             hidden_dim: 10,
         },
         output_dim: 10
     },
     path_encoder: {
         input_dim: 40, // 2 * path length for (x,y) coordinates
-        hidden_n: 2,
+        hidden_n: 0,
         hidden_dim: 20,
         output_dim: 10,
     },
     head: {
         input_dim: 0, // computed in policy constructor
-        hidden_n: 1,
+        hidden_n: 3,
         hidden_dim: 50,
         output_dim: 1, // number of outputs e.g. number of actions, or just one
         final_activation: 'sigmoid',
-    }
+    },
+    optim: {
+            optimizer: 'adam',
+            lr: 1e-3,
+            weight_decay: 0.1,
+    },
+    train_epochs: 200,
+    train_batch_size: 32,
+    loss: 'huber',
 }

config/value_dice.json5

@@ -0,0 +1,85 @@
+{
+    policy_net: {
+        ego_encoder: {
+            input_dim: 5, // number of state vars
+            hidden_n: 0,
+            hidden_dim: 5,
+            output_dim: 5
+        },
+        deepsets: {
+            input_dim: 6, // number of relative state vars for others
+            phi: {
+                hidden_n: 2,
+                hidden_dim: 20,
+            },
+            latent_dim: 20,
+            rho: {
+                hidden_n: 2,
+                hidden_dim: 10,
+            },
+            output_dim: 10
+        },
+        path_encoder: {
+            input_dim: 40, // 2 * path length for (x,y) coordinates
+            hidden_n: 0,
+            hidden_dim: 20,
+            output_dim: 10,
+        },
+        head: {
+            input_dim: 0, // computed in policy constructor
+            hidden_n: 3,
+            hidden_dim: 50,
+            output_dim: 1, // number of outputs e.g. number of actions, or just one
+            final_activation: 'sigmoid',
+        },
+    },
+    value_net: {
+        ego_encoder: {
+            input_dim: 5, // number of state vars
+            hidden_n: 0,
+            hidden_dim: 5,
+            output_dim: 5
+        },
+        deepsets: {
+            input_dim: 6, // number of relative state vars for others
+            phi: {
+                hidden_n: 2,
+                hidden_dim: 20,
+            },
+            latent_dim: 20,
+            rho: {
+                hidden_n: 2,
+                hidden_dim: 10,
+            },
+            output_dim: 10
+        },
+        path_encoder: {
+            input_dim: 40, // 2 * path length for (x,y) coordinates
+            hidden_n: 0,
+            hidden_dim: 20,
+            output_dim: 10,
+        },
+        action_dim: 1, // number of actions
+        head: {
+            input_dim: 0, // computed in policy constructor
+            hidden_n: 3,
+            hidden_dim: 50,
+            output_dim: 1, // number of outputs e.g. number of actions, or just one
+            final_activation: 'id',
+        },
+    },
+    policy_optim: {
+            optimizer: 'adam',
+            lr: 1e-3,
+            weight_decay: 0.1,
+    },
+    value_optim: {
+            optimizer: 'adam',
+            lr: 1e-3,
+            weight_decay: 0.1,
+    },
+    train_epochs: 200,
+    train_batch_size: 32,
+    discount: 0.95,
+    clip_grad_norm: 1.,
+}

eval_experiments.py

@@ -0,0 +1,100 @@
+import os
+from src.eval_main import eval_main
+from src.evaluation.utils import load_and_average
+import torch
+import json
+
+activations = [torch.nn.Tanh, torch.nn.LeakyReLU]
+
+def main(method:str='expert', folder:str=None, locations=[(0,0)], skip_running=False):
+    
+    exclude_keys_from_policy_kwargs = {'learning_rate', 'learning_rate_decay', 'clip_ratio', 'iterations_per_epoch', 'option'}
+    policy_kwargs = {}
+
+    if method in ['expert', 'idm']:
+        env, env_kwargs ='NRasterizedRouteIncrementingAgent', {}
+    elif method in ['bc','gail']:
+        env='NormalizedContinuousEvalEnv' 
+        env_kwargs={'stop_on_collision':True, 'max_episode_steps':1000}
+    elif method in ['hail']:
+        env = 'NormalizedSafeOptionsEvalEnv'
+        env_kwargs={'stop_on_collision':True, 'max_episode_steps':1000, 'safe_actions_collision_method': None, 'abort_unsafe_collision_method': None}
+    elif method in ['shail']:
+        env = 'NormalizedSafeOptionsEvalEnv'
+        env_kwargs={'stop_on_collision':True, 'max_episode_steps':1000}
+    else:
+        raise NotImplementedError
+
+    files = ['']
+
+    if folder is not None:
+        files = [os.path.join(folder, f) for f in os.listdir(folder) if os.path.isfile(os.path.join(folder, f))]
+        files = [f for f in files if f.endswith('.pt')]
+        with open(os.path.join(folder, 'config.json'), 'rb') as f:
+            config = json.load(f)
+        print('%i policy files found in %s folder' %(len(files), folder))
+        print('found policy config', config['policy'])
+
+    policy_config = {k: v for k, v in config['policy'].items() if k not in exclude_keys_from_policy_kwargs}
+    policy_config['activation'] = activations[policy_config['activation']]
+    print('final policy config', policy_config)
+
+    policy_kwargs.update(policy_config)
+    print('final policy kwargs', policy_kwargs)
+
+    if not skip_running:
+        for policy_file in files:
+            # run metrics on that file
+            outbase = eval_main(locations=locations,
+                                method=method, 
+                                policy_file=policy_file, 
+                                policy_kwargs=policy_kwargs,
+                                env=env, 
+                                env_kwargs=env_kwargs)
+        outfolder = os.path.dirname(outbase)
+    else:
+        locstr = 'loc_'+'_'.join([f'r{ro}t{tr}' for (ro,tr) in locations])
+        if folder is None:
+            outfolder = os.path.join('out',method,locstr)
+        else:
+            path_items = folder.split('/')
+            outfolder = os.path.join('out', '/'.join(path_items[1:]), locstr)
+    
+    # load metrics from save_path
+    average_metrics = load_and_average(outfolder)
+    if method in ['expert', 'idm']:
+        latex_print(average_metrics, light=True)
+    else:
+        latex_print(average_metrics)
+
+def latex_print(am, light=False):
+    """
+    print latex line
+
+    am (Dict[str,tuple]): dict mapping metric_name to (mean, std)
+    """
+
+    print('success rate, distance travelled, RWSE_10, |DeltaV|, AccelJSD')
+    if light:
+        if 'rwse_10s' in am.keys():
+            print("%2.1f& %2.1f & %1.2f & %2.1f& "
+            "%0.3f \\\\" %( 100*am['success rate'][0], am['mean travel distance'][0], am['rwse_10s'][0],
+                            am['average absolute average velocity'][0],am['acceleration distribution divergence'][0] ))
+            return
+
+
+        print("%2.1f& %2.1f & $---$ & $---$ & "
+            "$---$ \\\\" %( 100*am['success rate'][0], am['mean travel distance'][0]))
+        return
+    
+    print("%2.1f \\scriptstyle\\pm %2.1f & %2.1f \\scriptstyle\\pm %2.1f & "
+        "%1.2f \\scriptstyle\\pm %1.2f & %2.1f \\scriptstyle\\pm %1.1f & "
+        "%0.3f \\scriptstyle\\pm %0.3f \\\\" %( 100*am['success rate'][0], 100*am['success rate'][1],
+               am['mean travel distance'][0] , am['mean travel distance'][1]  ,
+               am['rwse_10s'][0] , am['rwse_10s'][1]  ,
+               am['average absolute average velocity'][0] , am['average absolute average velocity'][1]  ,
+               am['acceleration distribution divergence'][0] , am['acceleration distribution divergence'][1]  ))
+
+if __name__=='__main__':
+    import fire 
+    fire.Fire(main)

evaluate_models.sh

@@ -0,0 +1,19 @@
+# can add --skip_running if you've run the runs before on the saved policies
+
+python -m eval_experiments 
+python -m eval_experiments --locations='[(0,4)]' 
+python -m eval_experiments --method idm 
+python -m eval_experiments --method idm --locations='[(0,4)]' 
+python -m eval_experiments --method bc --folder='test_policies/bc/expA' 
+python -m eval_experiments --method bc --folder='test_policies/bc/expB' --locations='[(0,4)]' 
+python -m eval_experiments --method gail --folder='test_policies/gail/expA' 
+python -m eval_experiments --method gail --folder='test_policies/gail/expB' --locations='[(0,4)]' 
+python -m eval_experiments --method hail --folder='test_policies/hail/expA'
+python -m eval_experiments --method hail --folder='test_policies/hail/expB' --locations='[(0,4)]'
+python -m eval_experiments --method shail --folder='test_policies/shail/expA'
+python -m eval_experiments --method shail --folder='test_policies/shail/expB' --locations='[(0,4)]'
+
+python -m eval_experiments --method hail --folder='test_policies/hail-etienne/expA'
+python -m eval_experiments --method hail --folder='test_policies/hail-etienne/expB' --locations='[(0,4)]'
+python -m eval_experiments --method shail --folder='test_policies/shail-etienne/expA'
+python -m eval_experiments --method shail --folder='test_policies/shail-etienne/expB' --locations='[(0,4)]'

experiments/experiment.py

@@ -0,0 +1,203 @@
+import json5
+from functools import partial
+import os
+opj = os.path.join
+
+# set up ray tune
+import ray
+from ray import tune
+from ray.tune import Analysis, ExperimentAnalysis
+from ray.tune.schedulers import ASHAScheduler
+from hyperopt import hp
+from ray.tune.suggest.hyperopt import HyperOptSearch
+
+# get graphs
+import intersim
+from intersim.graphs import ConeVisibilityGraph
+
+
+from src.main import basestr, main
+
+def parse_args():
+    """
+    Parse arguments to main
+    Returns:
+        kwargs: dictionary of arguments:
+            train (bool): whether to run train loop
+            test (bool): whether to run test loop
+            method (str): the method to try for imitation
+            loc (int): the location index of the roundabout
+            config (str): config path
+            seed (int): RNG seed
+    """
+    import argparse
+    parser = argparse.ArgumentParser(description='Save Expert Trajectories')
+    parser.add_argument('--loc', default=0, type=int,
+        help='location (default 0)')
+    parser.add_argument("--train", help="train model",
+        action="store_true")
+    parser.add_argument("--ray", help="use ray tune to run multiple experiments",
+        action="store_true")
+    parser.add_argument("--test", help="test model",
+        action="store_true") 
+    parser.add_argument("--method", help="modeling method",
+        choices=['bc', 'gail', 'advil', 'vd'], default='bc')
+    parser.add_argument("--config", help="config file path",
+        default=None, type=str)
+    parser.add_argument('--seed', default=0, type=int,
+        help='seed')    
+    parser.add_argument('--nframes', default=500, type=int,
+        help='frames for test animation') 
+    parser.add_argument('--nsamples', default=200, type=int,
+        help='number of ray samples') 
+    parser.add_argument('--graph', action='store_true',
+        help='whether to mask the relative states based on a ConeVisibilityGraph')
+    parser.add_argument('-d', default='./expert_data', type=str,
+                       help='data directory')
+    parser.add_argument('-o', default=None, type=str,
+                       help='output directory')                   
+    args = parser.parse_args()
+    kwargs = {
+        'train':args.train, 
+        'test':args.test, 
+        'method':args.method,
+        'loc':args.loc,
+        'config_path':args.config,
+        'seed':args.seed,
+        'ray':args.ray,
+        'nframes':args.nframes,
+        'nsamples':args.nsamples,
+        'datadir':os.path.abspath(args.d),
+        'graph':None,
+        'outdir': opj('output',args.method,'loc%02i'%(args.loc)),
+        'train_tracks':[0,1,2],
+        'cv_tracks':[3],
+        'test_tracks':[4],
+    }
+    if args.o:
+        kwargs['outdir'] = args.o
+    if args.graph:
+        kwargs['graph'] = ConeVisibilityGraph(r=20, half_angle=120)
+    return kwargs
+
+def get_full_config(ray_config:dict, method:str)->dict:
+    """
+    Get full model configuration from ray config and method string
+    Args:
+        ray_config (dict): ray config
+        method (str): method to get full configuration for
+    """
+    if method == 'bc':
+        from src.bc import bc_config
+        config = bc_config(ray_config)
+    elif method == 'vd':
+        from src.value_dice import vd_config
+        config = vd_config(ray_config)
+    else:
+        raise NotImplementedError
+    return config
+
+def get_ray_config(method:str)->dict:
+    """
+    Get configuration for ray based on method.
+    Args: 
+        method (str): method to get configuration for
+    Returns:
+        ray_config (dict): configuration for ray
+    """
+    if method == 'bc':
+        ray_config = {
+            "lr": tune.loguniform(1e-5, 1e-3),
+            "weight_decay": tune.choice([0, 0.1]),
+            "loss": tune.choice(['huber', 'mse']),
+            "train_batch_size": tune.choice([16,32,64]),
+            "deepsets_phi_hidden_n": tune.randint(1,5),
+            "deepsets_phi_hidden_dim": tune.lograndint(8,65),
+            "deepsets_latent_dim": tune.lograndint(8,129),
+            "deepsets_rho_hidden_n": tune.randint(0,3),
+            "deepsets_rho_hidden_dim": tune.lograndint(8,129),
+            "deepsets_output_dim": tune.lograndint(4,129),
+            "head_hidden_n": tune.randint(1,6),
+            "head_hidden_dim": tune.lograndint(16,257),
+            "head_final_activation": tune.choice(['sigmoid', None]),
+        }
+    elif method == 'vd':
+        ray_config = {
+            "policy_lr": tune.loguniform(1e-5, 1e-3),
+            "value_lr": tune.loguniform(1e-5, 1e-3),
+            "policy_weight_decay": tune.choice([0, 0.1]),
+            "value_weight_decay": tune.choice([0, 0.1]),
+            "train_batch_size": tune.choice([16,32,64]),
+            "deepsets_phi_hidden_n": tune.randint(1,5),
+            "deepsets_phi_hidden_dim": tune.lograndint(8,65),
+            "deepsets_latent_dim": tune.lograndint(8,129),
+            "deepsets_rho_hidden_n": tune.randint(0,3),
+            "deepsets_rho_hidden_dim": tune.lograndint(8,129),
+            "deepsets_output_dim": tune.lograndint(4,129),
+            "head_hidden_n": tune.randint(1,6),
+            "head_hidden_dim": tune.lograndint(16,257),
+            "head_final_activation": tune.choice(['sigmoid', None]),
+            "clip_grad_norm": tune.choice([.5, 1., 5., 10.]),
+            "discount": tune.choice([.95, .99])
+        }
+    else:
+        raise NotImplementedError
+    return ray_config
+
+if __name__ == '__main__':
+    kwargs = parse_args()
+    
+    # make prefix of output files
+
+    if kwargs['config_path']:
+        # load config
+        with open(kwargs['config_path'], 'r') as cfg:
+            config = json5.load(cfg)
+        if not os.path.isdir(kwargs['outdir']):
+            os.makedirs(kwargs['outdir'])   
+        filestr = opj(kwargs['outdir'], basestr(**kwargs)) 
+        if kwargs['ray']:
+            filestr = kwargs['config_path'].replace('_config.json','')
+        main(config, filestr=filestr, **kwargs)
+
+    elif kwargs['ray'] and kwargs['train']:
+
+        ray.shutdown() 
+        ray.init(log_to_driver=False)
+
+        def ray_train(config, datadir=None):
+            full_config = get_full_config(config, kwargs['method'])
+            main(full_config, filestr='exp', **kwargs)
+        
+        ray_config = get_ray_config(kwargs['method'])
+        search = HyperOptSearch(ray_config, max_concurrent=8, metric='cv_loss',mode="min",)
+        custom_scheduler = ASHAScheduler(metric='cv_loss', mode="min", grace_period=15)
+
+        analysis = tune.run(
+            ray_train, 
+            #config=ray_config,
+            search_alg=search,
+            scheduler=custom_scheduler,
+            local_dir=kwargs['outdir'],
+            #resources_per_trial={"cpu": 2},
+            time_budget_s=120*60,
+            num_samples=kwargs['nsamples'],
+        )
+    elif kwargs['ray'] and kwargs['test']:
+        analysis = Analysis(kwargs['outdir'], default_metric="cv_loss", default_mode="min")
+        config = analysis.get_best_config()
+        filepath = analysis.get_best_logdir()
+        filestr = opj(filepath, 'exp')
+        config_path = filestr+'_config.json'
+        with open(config_path, 'r') as cfg:
+            config = json5.load(cfg)
+        print("Best ray experiment:", filepath)
+        main(config, filestr=filestr, **kwargs)
+    else:
+        raise Exception('No valid config found')
+
+    
+    
+    
+
+    

experiments/experiments.sh

@@ -0,0 +1,9 @@
+#!/bin/sh
+
+python experiments/experiment.py --ray --train -d ./expert_data/base
+python experiments/experiment.py --ray --test -d ./expert_data/base --nframes 1000
+python experiments/experiment.py --ray --train -d ./expert_data/reg
+python experiments/experiment.py --ray --test -d ./expert_data/reg --nframes 1000
+python experiments/experiment.py --ray --train -d ./expert_data/reg_graph --graph
+python experiments/experiment.py --ray --test -d ./expert_data/reg_graph --graph --nframes 1000
+

experiments/train_vd.sh

@@ -0,0 +1,5 @@
+#!/bin/sh
+
+# python experiments/experiment.py --method vd --train --ray -d expert_data/reg -o output/vd/loc00/reg --nsamples 400
+# python experiments/experiment.py --test --ray --method vd -d expert_data/normal -o output/vd/loc00/normal --nframes 1000
+python experiments/experiment.py --train --method vd --config config/value_dice.json5

gail-experiment.py

@@ -0,0 +1,238 @@
+# %%
+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, 
+                },
+                '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-4, 
+                    'iterations_per_epoch': 100, 
+                    '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'))

generate_demos.sh

@@ -0,0 +1,10 @@
+#DEFAULT PARAMETERS:
+# locs:list=None, (default to all locations)
+# tracks:list=None, (default to all tracks)
+# env_class:str='NRasterizedIncrementingAgent',
+# env_args:dict={width:36,height:36,m_per_px:2},
+# expert_class:str='NRasterizedRouteIncrementingAgent',  
+# expert_args:dict={mu:0.001}):
+
+# python -m src.data.expert --locs='[DR_USA_Roundabout_FT]' --tracks='[0]'
+python -m src.data.expert --locs='[DR_USA_Roundabout_FT]' --tracks='[0]'

interimit/__init__.py

@@ -1 +0,0 @@
-from interimit.expert_data import generate_expert_data, load_expert_data

interimit/data_utils.py

@@ -1,103 +0,0 @@
-import torch
-from torch.utils.data import Dataset, DataLoader
-import numpy as np
-#from torchvision import transforms, utils
-from interimit.expert_data import load_expert_data
-import os
-opj = os.path.join
-
-class InteractionDatasetMultiAgent(Dataset):
-    """
-    Class to handle getting full multi-agent observations and actions
-    """
-    pass
-
-class InteractionDatasetSingleAgent(Dataset):
-    """Class to load states and actions for individual agents."""
-
-    def __init__(self, output_dir='expert_data', loc:int = 0, tracks:list = [0], transforms={}):
-        """
-        Args:
-            output_dir (string): Directory with all the images.
-            loc (int): location index
-            tracks (list[int]): track indices
-            transforms (dict): dictionary of transforms to apply to different variables
-        """
-        self.output_dir = output_dir
-        self.loc = loc
-        self.tracks = tracks
-        self.transforms = transforms
-        #self.action_transform = transforms.get('action', None)
-        #self.state_transform = transforms.get('state', None)
-        #self.relative_state_transform = transforms.get('relative_state', None)
-        #self.paths_x_transform = transforms.get('paths_x', None)
-        #self.paths_y_transform = transform.get('paths_y',None)
-
-        self._load_dataset()
-
-    def _load_dataset(self):
-        """
-        Load the full datasets ahead of time
-        """
-        self.raw_data = {'state':[], 'relative_state':[], 'action':[], 'path_x':[], 'path_y':[]}
-        max_nv = 0
-        for track in self.tracks:
-            try:
-                observations, actions = load_expert_data(path=self.output_dir, loc=self.loc, track=track)
-                print('Loaded location {} track {}'.format(self.loc,track))
-            except:
-                print('Failed to load location {} track {}'.format(self.loc,track))
-                continue
-            T = len(actions)
-            for t in range(T):
-                nni = ~torch.isnan(observations[t]['state'][:,0])
-                max_nv = max(max_nv,nni.count_nonzero())
-                self.raw_data['state'].append(observations[t]['state'][nni])
-                self.raw_data['relative_state'].append(observations[t]['relative_state'][nni.nonzero(),nni.nonzero()])
-                self.raw_data['action'].append(actions[t][nni])
-                self.raw_data['path_x'].append(observations[t]['paths'][0][nni])
-                self.raw_data['path_y'].append(observations[t]['paths'][1][nni])
-
-        # cat lists
-        self.raw_data['state'] = torch.cat(self.raw_data['state'])
-        self.raw_data['action'] = torch.cat(self.raw_data['action'])
-        self.raw_data['path_x'] = torch.cat(self.raw_data['path_x'])
-        self.raw_data['path_y'] = torch.cat(self.raw_data['path_y'])
-
-        # pad second dimension of relative state
-        for i in range(len(self.raw_data['relative_state'])):
-            nv1, nv2, d = self.raw_data['relative_state'][i].shape
-            pad = torch.zeros(nv1, max_nv-nv2, d) * np.nan
-            self.raw_data['relative_state'][i] = torch.cat((self.raw_data['relative_state'][i], pad), dim=1)
-        self.raw_data['relative_state'] = torch.cat(self.raw_data['relative_state'])
-
-        # mandate equal length
-        assert len(self.raw_data['state']) == len(self.raw_data['relative_state']) \
-            == len(self.raw_data['action']) \
-            == len(self.raw_data['path_x']) \
-            == len(self.raw_data['path_y']), 'dataset lengths unequal'
-        
-    def __len__(self):
-        return len(self.raw_data['state'])
-
-    def __getitem__(self, idx):
-        """ 
-        Sample from the dataset
-        Args:
-            idx: index or indices of B samples
-        Returns:
-            sample (dict): sample dictionary with the following entries:
-                state (torch.tensor): (B, 5) raw state
-                relative_state (torch.tensor): (B, max_nv, d) relative state (padded with nans)
-                path_x (torch.tensor): (B, P) tensor of P future path x positions
-                path_y (torch.tensor): (B, P) tensor of P future path y positions
-                action (torch.tensor): (B, 1) actions taken from each state
-        """
-        keys = ['state', 'relative_state', 'path_x', 'path_y', 'action']
-        sample = {key:self.raw_data[key][idx] for key in keys}
-
-        for key in keys: 
-            if key in self.transforms.keys():
-                sample[key] = self.transforms[key](sample[key])
-
-        return sample

interimit/expert_data.py

@@ -1,95 +0,0 @@
-import torch
-
-import pickle
-import gym
-import numpy as np
-
-import intersim
-from intersim.utils import get_map_path, get_svt, SVT_to_stateactions
-from intersim import collisions
-import os
-opj = os.path.join
-
-def generate_expert_data(path: str='expert_data', loc: int = 0, track:int = 0, **kwargs):
-    """
-    Function to save (joint) states and observations from simulated frame
-    Args:
-        path (str): directory to save data
-        loc (int): location index
-        track (int): track index
-        kwargs: arguments for environment instantiation
-    """
-    
-    if not os.path.isdir(path):
-        os.mkdir(path)
-    filestr = opj(path,intersim.LOCATIONS[loc]+'_track%03i'%(track))
-    
-    svt, svt_path = get_svt(base='InteractionSimulator', loc=loc, track=track)
-    osm = get_map_path(base='InteractionSimulator', loc=loc)
-    print('SVT path: {}'.format(svt_path))
-    print('Map path: {}'.format(osm))
-    states, actions = SVT_to_stateactions(svt)
-
-    # animate from environment
-    env = gym.make('intersim:intersim-v0', svt=svt, map_path=osm, **kwargs, 
-        min_acc=-np.inf, max_acc=np.inf)
-    
-    env.reset()
-    done = False
-    obs, actions_taken, max_devs = [], [], []
-    i = 0
-    while not done and i < len(actions):
-        # check state deviation
-        env_state = env.projected_state
-        nni = ~torch.isnan(env_state[:,0])
-        norms = torch.norm(env_state[nni,:2]-states[i,nni,:2], dim=1)
-        max_devs.append(norms.max())
-
-        # propagate environment
-        ob, r, done, info = env.step(env.target_state(svt.simstate[i+1]))
-        obs.append(ob)
-        actions_taken.append(info['action_taken'])
-        i += 1
-    
-    print("Maximum environment deviation from track: %f m" %(max(max_devs)))
-
-    # check for collisions
-    x = torch.stack([ob['state'] for ob in obs])
-    cols = collisions.check_collisions_trajectory(x, svt.lengths, svt.widths)
-    assert ~torch.any(cols), 'Error: Collisions found at indices {}'.format(cols.nonzero(as_tuple=True))
-
-    # shift actions
-    actions_taken.pop(0)
-    obs.pop(-1)
-
-    # save observations and actions
-    pickle.dump(obs,open(filestr+'_observations.pkl', 'wb'))
-    torch.save(torch.stack(actions_taken), filestr+'_actions.pt')
-
-def load_expert_data(path='expert_data', loc: int = 0, track:int = 0):
-    """
-    Load expert data from file.
-    Args:
-        path (str): directory to save data
-        loc (int): location index
-        track (int): track index
-    Returns:
-        obs (list[Observations]): list of observations
-        actions (list[torch.tensor]): list of corresponding actions taken in observations
-    """
-    # load observations and actions
-    filestr = opj(path, intersim.LOCATIONS[loc]+'_track%03i'%(track))
-    obs = pickle.load(open(filestr+'_observations.pkl', 'rb'))
-    actions = torch.load(filestr+'_actions.pt')
-    actions = list(torch.unbind(actions))
-    return obs, actions
-
-if __name__ == '__main__':
-    import argparse
-    parser = argparse.ArgumentParser(description='Save Expert Trajectories')
-    parser.add_argument('--loc', default=0, type=int,
-                       help='location (default 0)')
-    parser.add_argument('--track', default=0, type=int,
-                       help='track number (default 0)')
-    args = parser.parse_args()
-    generate_expert_data(loc=args.loc,track=args.track)

interimit/policies/policy.py

@@ -1,42 +0,0 @@
-
-import torch
-from torch import nn
-
-from interimit.nets.deepsets import DeepSetsModule, Phi
-
-class Policy:
-    pass
-
-class DeepSetsPolicy(Policy, nn.Module):
-    def __init__(self, ego_config, dynamic_config, path_config, head_config):
-        """
-        Args:
-            ego_config (dict): dictionary for configuring the ego network
-            dynamic_config (dict): dictionary for configuring the dynamic input (deepsets) network
-            path_config (dict): dictionary for configuring the path network
-            head_config (dict): dictionary for configuring the common head network
-        """
-        super(DeepSetsPolicy, self).__init__()
-        self.ego_net = Phi.from_config(ego_config)
-        self.deepsets = DeepSetsModule.from_config(dynamic_config)
-        self.path_net = Phi.from_config(path_config)
-        cat_dim = self.ego_net.output_dim + self.deepsets.output_dim + self.path_net.output_dim
-        # head has number of concatenated features as input
-        head_config["input_dim"] = cat_dim
-        self.head = Phi.from_config(head_config)
-
-    def forward(self, ego_state, relative_states, path):
-        """
-        Args:
-            ego_state (torch.tensor): (ns,) state of ego vehicle
-            relative_states (torch.tensor): (nv, ns) relative states of other vehicles (dynamic size)
-            path (torch.tensor): (path_length, 2) coordinates (x,y) of path
-        Returns:
-            x (torch.tensor): (head_output_dim,) output of common head network
-        """
-        x_ego = self.ego_net(ego_state)
-        x_relative = self.deepsets(relative_states)
-        x_path = self.path_net(path.flatten())
-        x = torch.cat([x_ego, x_relative, x_path])
-        x = self.head(x)
-        return x