Add basic setup.py

.gitignore

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

README.md

@@ -1,22 +1,10 @@
 # InteractionImitation
-Imitation Learning with the [Interaction Dataset](https://interaction-dataset.com/) via the [InteractionSimulator](https://github.com/sisl/InteractionSimulator) gym environments.
+Imitation Learning with the INTERACTION Dataset
-
-Code for "[SHAIL: Safety-Aware Hierarchical Adversarial Imitation Learning for Autonomous Driving in Urban Environments](https://arxiv.org/abs/2204.01922)".
-If you find this repository useful, please cite the paper:
-
-```
-@article{jamgochian2022shail,
-  author = {Arec Jamgochian and Etienne Buehrle and Johannes Fischer and Mykel J. Kochenderfer},
-  title = {{SHAIL}: Safety-Aware Hierarchical Adversarial Imitation Learning for Autonomous Driving in Urban Environments},
-  journal = {arXiv:2204.01922 [cs]},
-  year = {2022}
-}
-```
 
 ## Getting started
-Clone the `InteractionSimulator` with the `shail` tag and pip install the module.
+Clone InteractionSimulator and pip install the module.
 ```
-git clone --branch shail https://github.com/sisl/InteractionSimulator.git
+git clone https://github.com/sisl/InteractionSimulator.git
 cd InteractionSimulator
 pip install -e .
 cd ..
@@ -31,28 +19,35 @@ The INTERACTION dataset contains a two folders which should be copied into a fol
   - the contents of `recorded_trackfiles` should be copied to `./InteractionSimulator/datasets/trackfiles`
   - the contents of `maps` should be copied to `./InteractionSimulator/datasets/maps`
 
-## Processing and saving expert demos
+## Processing, saving, and loading expert demos
-Once the repository has been set up, you need to generate two separate sets of expert demos for tracks 0-4. The first command generates true joint and individual states and actions necessary for evaluating, saving them in `expert_data/`. The second command generates trajectory rollouts according to individual agent observations, which is later used as expert data for the learning models.
+Once the repository has been set up, you can process and save expert track demonstrations with:
 ```
-python -m src.expert --locs='[DR_USA_Roundabout_FT]' --tracks='[0,1,2,3,4]'
+python src/expert_data.py --loc [LOCNUM] --track [TRACKNUM]
-python -m intersimple-expert-rollout-setobs2 --tracks='[0,1,2,3,4]'
 ```
 
-
+You can then load the experts actions and observations using
-## Tuning hyperparameters and training finalized models
+```
-To tune models, we use `ray[tune]` grid searches. You can run see the commands we used to train in the top half of `train_models.sh`, as well as the hyperparameters we search over in `bc-experiment.py`, `gail-experiment.py`, and `shail-experiment.py`. After training the models, configurations get saved in `best_configs/` (the best SHAIL confg gets copied to a HAIL config, with the appropriate environment parameters changed for ablation). However, upon manual inspection of the training runs, we note some better performance than the automatically-set configs at earlier epochs, so we adjust the `best_configs` manually.
+from src import expert_data
-
+observations, actions = expert_data.load_expert_data(loc = [LOCNUM], track = [TRACKNUM])
-After the `best_configs/` are set, we rerun each configuration with multiple seeds. The commands to do so are in the bottom half of `train_models.sh`. This saves different learned policy files to `test_policies/`.
+for (s, a) in zip (observations, actions):
-
+    # do some imitation learning
-
+```
-## Evaluating models
-To evaluate the learned policies, we rerun each model in particular setting, evaluate all our metrics, and average over different trained model seeds. The commands to do so are in `evaluate_models.sh`. 
 
 
 ## Package Structure
 ```
 InteractionImitation
-|- TODO
+|- demos
+|- algorithms
+   |- BC
+   |- AdVIL
+|- nets
+   |- Encoder
+   |- DeepSet
+   |- Decoder
+|- policies
+|- discriminators
+|- demo_generators
 ```
 
 ## Type Definitions

bc-experiment.py

@@ -1,218 +0,0 @@
-# %%
-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'],
-                use_idm=config['trainenv']['use_idm'],
-            ), 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'],
-                use_idm=config['trainenv']['use_idm'],
-                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]
-
-    # 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=500)
-    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, 
-                    'use_idm':True,
-                },
-                'policy': {
-                    'learning_rate': 3e-4, 
-                    'learning_rate_decay': tune.grid_search([0.999, 1.0]), 
-                    'hidden_layer_size': tune.grid_search([10, 20, 40]),
-                    'n_hidden_layers': tune.grid_search([2, 3]), 
-                    'activation':tune.grid_search([0, 1]),
-                },
-                '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'))
-            shutil.copyfile(os.path.join(check_dir,'params.json'),
-                os.path.join(savepath, 'config.json')) # copy config automatically

best_configs/bc_expA.json

@@ -1,16 +0,0 @@
-{
-    "experiment": "A",
-    "trainenv": {
-        "stop_on_collision": false,
-        "use_idm": true
-    },
-    "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

@@ -1,16 +0,0 @@
-{
-    "experiment": "B",
-    "trainenv": {
-        "stop_on_collision": false,
-        "use_idm": true
-    },
-    "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_expA.json

@@ -1,32 +0,0 @@
-{
-    "experiment": "A",
-    "trainenv": {
-        "stop_on_collision": false,
-        "use_idm": true
-    },
-    "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

@@ -1,32 +0,0 @@
-{
-    "experiment": "B",
-    "trainenv": {
-        "stop_on_collision": false,
-        "use_idm": true
-    },
-    "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_expA.json

@@ -1,34 +0,0 @@
-{
-    "experiment": "A",
-    "trainenv": {
-        "stop_on_collision": false,
-        "safe_actions_collision_method": null,
-        "abort_unsafe_collision_method": null,
-        "use_idm": true
-    },
-    "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

@@ -1,34 +0,0 @@
-{
-    "experiment": "B",
-    "trainenv": {
-        "stop_on_collision": false,
-        "safe_actions_collision_method": null,
-        "abort_unsafe_collision_method": null,
-        "use_idm": true
-    },
-    "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_expA.json

@@ -1,34 +0,0 @@
-{
-    "experiment": "A",
-    "trainenv": {
-        "stop_on_collision": false,
-        "safe_actions_collision_method": "circle",
-        "abort_unsafe_collision_method": "circle",
-        "use_idm": true
-    },
-    "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

@@ -1,34 +0,0 @@
-{
-    "experiment": "B",
-    "trainenv": {
-        "stop_on_collision": false,
-        "safe_actions_collision_method": "circle",
-        "abort_unsafe_collision_method": "circle",
-        "use_idm": true
-    },
-    "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

@@ -0,0 +1,34 @@
+{
+    ego_state: {
+        input_dim: 5, // number of state vars
+        hidden_n: 1,
+        hidden_dim: 5,
+        output_dim: 5
+    },
+    deepsets: {
+        input_dim: 5, // number of relative state vars for others
+        phi: {
+            hidden_n: 1,
+            hidden_dim: 20,
+        },
+        latent_dim: 20,
+        rho: {
+            hidden_n: 1,
+            hidden_dim: 10,
+        },
+        output_dim: 10
+    },
+    path_encoder: {
+        input_dim: 40, // 2 * path length for (x,y) coordinates
+        hidden_n: 2,
+        hidden_dim: 20,
+        output_dim: 10,
+    },
+    head: {
+        input_dim: 0, // computed in policy constructor
+        hidden_n: 1,
+        hidden_dim: 50,
+        output_dim: 1, // number of outputs e.g. number of actions, or just one
+        final_activation: 'sigmoid',
+    }
+}

cp-videos.sh

@@ -1,10 +0,0 @@
-# cp-videos videos/ videos/icra23/
-
-agents=( 5 27 39 43 47 53 63 81 83 87 93 96 105 113 124 127 130 134 )
-
-for a in "${agents[@]}"
-do
-    cp "$1/expert_agent/loc0/track0/agent${a}_ani.mp4" "$2/t${a}expert.mp4"
-    cp "$1/idm/loc0/track0/agent${a}_ani.mp4" "$2/t${a}idm.mp4"
-    cp "$1/shail/loc0/track0/agent${a}_ani.mp4" "$2/t${a}shail.mp4"
-done

eval_experiments.py

@@ -1,108 +0,0 @@
-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, save_videos:bool=False, videos_folder:str='videos', first_seed_only:bool=False):
-    
-    exclude_keys_from_policy_kwargs = {'learning_rate', 'learning_rate_decay', 'clip_ratio', 'iterations_per_epoch', 'option'}
-    policy_kwargs = {}
-
-    if method in ['expert', 'expert_agent']:
-        env, env_kwargs ='NRasterizedRouteIncrementingAgent', {}
-    elif method in ['idm']:
-        env, env_kwargs ='NRasterizedRouteIncrementingAgent', {'use_idm':True}
-    elif method in ['bc','gail']:
-        env='NormalizedContinuousEvalEnv' 
-        env_kwargs={'stop_on_collision':True, 'max_episode_steps':1000, 'use_idm':True}
-    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, 'use_idm':True}
-    elif method in ['shail']:
-        env = 'NormalizedSafeOptionsEvalEnv'
-        env_kwargs={'stop_on_collision':True, 'max_episode_steps':1000, 'use_idm':True}
-    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')]
-
-        if first_seed_only:
-            files = files[:1]
-
-        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,
-                                videos_folder=None if not save_videos else videos_folder)
-        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 & %2.1f & %1.2f& "
-            "%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 & "
-        "%2.1f \\scriptstyle\\pm %1.1f & %1.2f \\scriptstyle\\pm %1.2f & "
-        "%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

@@ -1,18 +0,0 @@
-# can add --skip_running if you've already run the saved policies through the test environments and have appropriate
-# metrics in the out folder. Doing so will generate average metrics quickly.
-
-# Experiment A
-python -m eval_experiments 
-python -m eval_experiments --method idm 
-python -m eval_experiments --method bc --folder='test_policies/bc/expA' 
-python -m eval_experiments --method gail --folder='test_policies/gail/expA' 
-python -m eval_experiments --method hail --folder='test_policies/hail/expA'
-python -m eval_experiments --method shail --folder='test_policies/shail/expA' 
-
-# Experiment B
-python -m eval_experiments --locations='[(0,4)]' 
-python -m eval_experiments --method idm --locations='[(0,4)]'
-python -m eval_experiments --method bc --folder='test_policies/bc/expB' --locations='[(0,4)]'
-python -m eval_experiments --method gail --folder='test_policies/gail/expB' --locations='[(0,4)]'
-python -m eval_experiments --method hail --folder='test_policies/hail/expB' --locations='[(0,4)]'
-python -m eval_experiments --method shail --folder='test_policies/shail/expB' --locations='[(0,4)]'

gail-experiment.py

@@ -1,243 +0,0 @@
-# %%
-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'],
-                use_idm=config['trainenv']['use_idm'],
-            ), 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'],
-                use_idm=config['trainenv']['use_idm'],
-                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,
-                    'use_idm': True, 
-                },
-                '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'))
-            shutil.copyfile(os.path.join(check_dir,'params.json'),
-                os.path.join(savepath, 'config.json')) # copy config automatically

generate_videos.sh

@@ -1,20 +0,0 @@
-# can add --skip_running if you've already run the saved policies through the test environments and have appropriate
-# metrics in the out folder. Doing so will generate average metrics quickly.
-
-# Experiment A
-python -m eval_experiments 
-python -m eval_experiments --method expert_agent --save_videos --first_seed_only
-python -m eval_experiments --method idm --save_videos --first_seed_only
-python -m eval_experiments --method bc --folder='test_policies/bc/expA' --save_videos --first_seed_only
-python -m eval_experiments --method gail --folder='test_policies/gail/expA' --save_videos --first_seed_only
-python -m eval_experiments --method hail --folder='test_policies/hail/expA' --save_videos --first_seed_only
-python -m eval_experiments --method shail --folder='test_policies/shail/expA' --save_videos --first_seed_only
-
-# Experiment B
-python -m eval_experiments --locations='[(0,4)]' 
-python -m eval_experiments --method expert_agent --locations='[(0,4)]' --save_videos --first_seed_only
-python -m eval_experiments --method idm --locations='[(0,4)]' --save_videos --first_seed_only
-python -m eval_experiments --method bc --folder='test_policies/bc/expB' --locations='[(0,4)]' --save_videos --first_seed_only
-python -m eval_experiments --method gail --folder='test_policies/gail/expB' --locations='[(0,4)]' --save_videos --first_seed_only
-python -m eval_experiments --method hail --folder='test_policies/hail/expB' --locations='[(0,4)]' --save_videos --first_seed_only
-python -m eval_experiments --method shail --folder='test_policies/shail/expB' --locations='[(0,4)]' --save_videos --first_seed_only

interimit/__init__.py

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

interimit/data_utils.py

@@ -0,0 +1,103 @@
+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

@@ -0,0 +1,95 @@
+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/nets/deepsets.py

@@ -0,0 +1,111 @@
+import torch
+from torch import nn
+
+from interimit.nets.util import parse_functional
+
+class DeepSetsModule(nn.Module):
+    def __init__(self, input_dim, phi_hidden_n, phi_hidden_dim, latent_dim, rho_hidden_n, rho_hidden_dim, output_dim):
+        """
+        Args:
+            input_dim (int): input size of one instance of the set; input size of phi
+            phi_hidden_n (int): number of hidden layers in phi
+            phi_hidden_dim (int): size of hidden layers in phi
+            latent_dim (int): output size of phi network, where sum is taken over instances; input size of rho
+            rho_hidden_n (int): number of hidden layers in rho
+            rho_hidden_dim (int): size of hidden layers in rho
+            output_dim (int): output size of rho
+        """
+        super(DeepSetsModule, self).__init__()
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+        self.phi = Phi(self.input_dim, phi_hidden_n, phi_hidden_dim, latent_dim)
+        self.rho = Phi(latent_dim, rho_hidden_n, rho_hidden_dim, self.output_dim)
+        self.pooling = torch.sum  # torch.max # torch.mean
+
+    @staticmethod
+    def from_config(config):
+        """
+        Args:
+            config (dict): dictionary with network parameters in the form
+                {
+                    "input_dim": 5,
+                    "phi": {
+                        "hidden_n": 1,
+                        "hidden_dim": 10,
+                    },
+                    "latent_dim": 8,
+                    "rho": {
+                        "hidden_n": 1,
+                        "hidden_dim": 10,
+                    },
+                    "output_dim" : 1,
+                }
+        Returns:
+            m (nn.Module): deep sets module
+        """
+        input_dim = config["input_dim"]
+        phi = config["phi"]
+        latent_dim = config["latent_dim"]
+        rho = config["rho"]
+        output_dim = config["output_dim"]
+        m = DeepSetsModule(input_dim, phi["hidden_n"], phi["hidden_dim"], latent_dim, rho["hidden_n"], rho["hidden_dim"], output_dim)
+        return m
+
+    def forward(self, x):
+        """
+        Args:
+            x (torch.tensor): (batch_size, dynamic_size, input_dim)
+        Returns:
+            y (torch.tensor): (batch_size, output_dim)
+        """
+        # use negative dynamic_dim since batch dimensions are inserted at the front
+        dynamic_dim = -2
+        # iterate over dynamic dimension to apply phi to every instance
+        latent = tuple(self.phi(instance) for instance in x.unbind(dynamic_dim))
+        # stack outputs of phi
+        latent = torch.stack(latent, dim=dynamic_dim)
+        # apply pooling function to reduce dynamic dimension
+        latent = self.pooling(latent, dim=dynamic_dim)
+        # apply rho network
+        y = self.rho(latent)
+        return y
+
+
+class Phi(nn.Module):
+    def __init__(self, input_dim, hidden_n, hidden_dim, output_dim, final_activation=None):
+        """
+        Fully connected feedforward network with same size for all hidden layers and ReLU activation
+
+        Args:
+            input_dim (int): input dimension
+            hidden_n (int): number of hidden layers
+            hidden_dim (int): hidden layer dimension
+            output_dim (int): output dimension
+        """
+        super(Phi, self).__init__()
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+        self.layers = [nn.Linear(self.input_dim, hidden_dim)]
+        for _ in range(hidden_n - 1):
+            self.layers.append(nn.Linear(hidden_dim, hidden_dim))
+        self.layers.append(nn.Linear(hidden_dim, self.output_dim))
+        # self.in_layer = nn.Linear(input_dim, hidden_dim)
+        # self.hidden_layers = [nn.Linear(hidden_dim, hidden_dim) for _ in range(hidden_n - 1)]
+        # self.out_layer = nn.Linear(hidden_dim, output_dim)
+        self.activation = nn.functional.relu
+        self.final_activation = final_activation if final_activation else self.activation
+
+    def forward(self, x):
+        for layer in self.layers[:-1]:
+            x = self.activation(layer(x))
+        x = self.final_activation(self.layers[-1](x))
+        return x
+
+    @staticmethod
+    def from_config(config):
+        args = (config["input_dim"], config["hidden_n"], config["hidden_dim"], config["output_dim"])
+        if "final_activation" in config:
+            kwargs = {"final_activation": parse_functional(config["final_activation"])}
+        else:
+            kwargs = {}
+        return Phi(*args, **kwargs)

interimit/nets/util.py

@@ -0,0 +1,14 @@
+import torch
+from torch.nn import functional
+
+def parse_functional(functional_config):
+    if functional_config is None:
+        return None
+    elif isinstance(functional_config, str):
+        if functional_config == 'relu':
+            return functional.relu
+        elif functional_config == 'sigmoid':
+            return functional.sigmoid
+        elif functional_config == 'softmax':
+            return functional.softmax
+    

interimit/policies/policy.py

@@ -0,0 +1,42 @@
+
+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

intersimple-expert-rollout-setobs2.py

@@ -1,65 +0,0 @@
-import torch
-import functools
-from src.core.sampling import rollout_sb3
-from intersim.envs import IntersimpleLidarFlatIncrementingAgent
-from intersim.envs.intersimple import speed_reward
-from intersim.expert import NormalizedIntersimpleExpert
-from src.util.wrappers import CollisionPenaltyWrapper, Setobs
-import numpy as np
-from gym.wrappers import TransformObservation
-
-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 main(track:int, loc:int=0):
-    env = IntersimpleLidarFlatIncrementingAgent(
-        loc=loc,
-        track=track,
-        n_rays=5,
-        reward=functools.partial(
-            speed_reward,
-            collision_penalty=0
-        ),
-    )
-
-    policy = NormalizedIntersimpleExpert(env, mu=0.001)
-
-    env = Setobs(TransformObservation(
-        CollisionPenaltyWrapper(
-            env,
-            collision_distance=6, collision_penalty=100
-        ), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10)
-    ))
-    print(env.nv, 'vehicles')
-    expert_data = rollout_sb3(env, policy, n_episodes=150, max_steps_per_episode=200)
-
-    states, actions, rewards, dones = expert_data
-    print(f'Expert mean episode length {(~dones).sum() / states.shape[0]}')
-    print(f'Expert mean reward per episode {rewards[~dones].sum() / states.shape[0]}')
-    print(f'Observation mean', states[~dones].mean(0))
-    print(f'Observation std', states[~dones].std(0))
-
-    torch.save(expert_data, f'intersimple-expert-data-setobs2-loc{loc}-track{track}.pt')
-
-def loop(tracks:list=[0]):
-    for track in tracks:
-        main(track)
-
-if __name__=='__main__':
-    import fire
-    fire.Fire(loop)