Use ray tune

Merge different methods to sample the policy and collect transitions

.gitignore

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

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])

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.,
+}

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

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

requirements.txt

@@ -1,6 +1,9 @@
 tikzplotlib
 torch
+sklearn
 pytest
 json5
-gym
-git+https://github.com/sisl/InteractionSimulator
+tqdm
+tensorboard
+ray[tune]
+hyperopt

scratch/etienne/intersimple/airl_flat.py

@@ -0,0 +1,64 @@
+# %%
+import pathlib
+import pickle
+import tempfile
+
+import stable_baselines3 as sb3
+from stable_baselines3.common.env_util import make_vec_env
+
+from imitation.algorithms import adversarial, bc
+from imitation.data import rollout
+from imitation.util import logger
+
+from intersim.envs.intersimple import IntersimpleReward
+
+model_name = 'airl_flat'
+
+# Load pickled test demonstrations.
+with open("data/NormalizedIntersimpleExpert_IntersimpleRewardAgent51.pkl", "rb") as f:
+    # This is a list of `imitation.data.types.Trajectory`, where
+    # every instance contains observations and actions for a single expert
+    # demonstration.
+    trajectories = pickle.load(f)
+
+# %%
+# Convert List[types.Trajectory] to an instance of `imitation.data.types.Transitions`.
+# This is a more general dataclass containing unordered
+# (observation, actions, next_observation) transitions.
+transitions = rollout.flatten_trajectories(trajectories)
+
+venv = make_vec_env(IntersimpleReward, n_envs=2, env_kwargs={'agent': 51})
+
+tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
+tempdir_path = pathlib.Path(tempdir.name)
+print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
+
+# Train AIRL on expert data.
+# GAIL, and AIRL also accept as `expert_data` any Pytorch-style DataLoader that
+# iterates over dictionaries containing observations, actions, and next_observations.
+logger.configure(tempdir_path / "AIRL/")
+airl_trainer = adversarial.AIRL(
+    venv,
+    expert_data=transitions,
+    expert_batch_size=64,
+    gen_algo=sb3.PPO("MlpPolicy", venv, verbose=1, n_steps=1024), # n_steps = 2048 ?
+)
+airl_trainer.train(total_timesteps=100000)
+airl_trainer.gen_algo.save(model_name)
+
+del airl_trainer
+
+# %%
+model = sb3.PPO.load(model_name)
+
+env = IntersimpleReward(agent=51)
+
+obs = env.reset()
+while True:
+    action, _states = model.predict(obs)
+    obs, rewards, done, info = env.step(action)
+    env.render(mode='post')
+    if done:
+        break
+
+env.close(filestr='render/'+model_name)

scratch/etienne/intersimple/bc_flat.py

@@ -0,0 +1,59 @@
+# %%
+import pathlib
+import pickle
+import tempfile
+
+import stable_baselines3 as sb3
+from stable_baselines3.common.env_util import make_vec_env
+
+from imitation.algorithms import adversarial, bc
+from imitation.data import rollout
+from imitation.util import logger
+
+from intersim.envs.intersimple import IntersimpleReward
+
+model_name = 'bc_flat'
+
+# Load pickled test demonstrations.
+with open("data/NormalizedIntersimpleExpert_IntersimpleRewardAgent51.pkl", "rb") as f:
+    # This is a list of `imitation.data.types.Trajectory`, where
+    # every instance contains observations and actions for a single expert
+    # demonstration.
+    trajectories = pickle.load(f)
+
+# %%
+# Convert List[types.Trajectory] to an instance of `imitation.data.types.Transitions`.
+# This is a more general dataclass containing unordered
+# (observation, actions, next_observation) transitions.
+transitions = rollout.flatten_trajectories(trajectories)
+
+venv = make_vec_env(IntersimpleReward, n_envs=2, env_kwargs={'agent': 51})
+
+tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
+tempdir_path = pathlib.Path(tempdir.name)
+print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
+
+# Train BC on expert data.
+# BC also accepts as `expert_data` any PyTorch-style DataLoader that iterates over
+# dictionaries containing observations and actions.
+logger.configure(tempdir_path / "BC/")
+bc_trainer = bc.BC(venv.observation_space, venv.action_space, expert_data=transitions)
+bc_trainer.train(n_epochs=1000)
+bc_trainer.save_policy(model_name)
+
+del bc_trainer
+
+# %%
+model = bc.reconstruct_policy(model_name)
+
+env = IntersimpleReward(agent=51)
+
+obs = env.reset()
+while True:
+    action, _states = model.predict(obs)
+    obs, rewards, done, info = env.step(action)
+    env.render(mode='post')
+    if done:
+        break
+
+env.close(filestr='render/'+model_name)

scratch/etienne/intersimple/data/expert.py

@@ -0,0 +1,121 @@
+from intersim.envs.intersimple import Intersimple
+from stable_baselines3.common.policies import BasePolicy
+import gym
+import intersim.envs.intersimple
+import imitation.data.rollout as rollout
+from stable_baselines3.common.vec_env.dummy_vec_env import DummyVecEnv
+from imitation.data.wrappers import RolloutInfoWrapper
+
+class IntersimExpert(BasePolicy):
+    
+    def __init__(self, intersim_env, mu=0, *args, **kwargs):
+        super().__init__(
+            observation_space=gym.spaces.Space(),
+            action_space=gym.spaces.Space(),
+            *args, **kwargs
+        )
+        self._intersim = intersim_env
+        self._mu = mu
+
+    def forward(self, *args, **kwargs):
+        raise NotImplementedError()
+
+    def _predict(self, *args, **kwargs):
+        raise NotImplementedError()
+
+    def _action(self):
+        target_t = min(self._intersim._ind + 1, len(self._intersim._svt.simstate) - 1)
+        target_state = self._intersim._svt.simstate[target_t]
+        return self._intersim.target_state(target_state, mu=self._mu)
+
+    def predict(self, *args, **kwargs):
+        return self._action(), None
+
+class IntersimpleExpert(BasePolicy):
+
+    def __init__(self, intersimple_env, mu=0, *args, **kwargs):
+        super().__init__(
+            observation_space=intersimple_env.observation_space,
+            action_space=intersimple_env.action_space,
+            *args, **kwargs
+        )
+        self._intersimple = intersimple_env
+        self._intersim_expert = IntersimExpert(intersimple_env._env, mu=mu)
+    
+    def forward(self, *args, **kwargs):
+        raise NotImplementedError()
+
+    def _predict(self, *args, **kwargs):
+        raise NotImplementedError()
+
+    def _action(self):
+        return self._intersim_expert._action()[self._intersimple._agent]
+
+    def predict(self, *args, **kwargs):
+        return self._action(), None
+
+class NormalizedIntersimpleExpert(IntersimpleExpert):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def predict(self, *args, **kwargs):
+        action, _ = super().predict(*args, **kwargs)
+        return self._intersimple._normalize(action), None
+
+class DummyVecEnvPolicy():
+
+    def __init__(self, experts):
+        self._experts = [e() for e in experts]
+    
+    def predict(self, *args, **kwargs):
+        predictions = [e.predict() for e in self._experts]
+        actions = [p[0] for p in predictions]
+        states = [p[1] for p in predictions]
+        return actions, states
+
+def save_video(env, expert):
+    env.reset()
+    env.render()
+    done = False
+    while not done:
+        actions, _ = expert.predict()
+        _, _, done, _ = env.step(actions)
+        env.render()
+    env.close()
+
+def demonstrations(expert='NormalizedIntersimpleExpert', env='NRasterizedRandomAgent', path=None, min_timesteps=25000, min_episodes=None, video=False, env_args={}, policy_args={}):
+    """Rollout and save expert demos.
+    
+    Usage:
+        python -m intersimple.expert <flags>
+
+    """
+    Env = intersim.envs.intersimple.__dict__[env]
+    Expert = globals()[expert]
+
+    env = Env(**env_args)
+    info_env = RolloutInfoWrapper(env)
+    venv = DummyVecEnv([lambda: info_env])
+
+    policy = Expert(env, **policy_args)
+    venv_policy = DummyVecEnvPolicy([lambda: policy])
+
+    if video:
+        save_video(env, policy)
+    
+    path = path or (policy.__class__.__name__ + '_' + env.__class__.__name__ + '.pkl')
+
+    rollout.rollout_and_save(
+        path=path,
+        policy=venv_policy,
+        venv=venv,
+        sample_until=rollout.make_sample_until(
+            min_timesteps=min_timesteps,
+            min_episodes=min_episodes,
+        )
+    )
+
+if __name__ == '__main__':
+    import fire
+    fire.Fire(demonstrations)

scratch/etienne/intersimple/data/generate.sh

@@ -0,0 +1,4 @@
+#python -m expert --env=IntersimpleReward --min_timesteps=200 --env_args='{agent:51}' --path='NormalizedIntersimpleExpert_IntersimpleRewardAgent51.pkl'
+#python -m expert --env=IntersimpleReward --min_timesteps=200 --env_args='{agent:51}' --policy_args='{mu:0.005}' --path='NormalizedIntersimpleExpert_IntersimpleRewardAgent51Mu.005.pkl'
+#python -m expert --env=IntersimpleReward --min_timesteps=200 --env_args='{agent:51}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpert_IntersimpleRewardAgent51Mu.001.pkl'
+python -m expert --env=NRasterized --min_timesteps=200 --env_args='{agent:51,width:36,height:36,m_per_px:2}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001_NRasterizedAgent51w36h36mppx2.pkl'

scratch/etienne/intersimple/gail/discriminator.py

@@ -0,0 +1,58 @@
+import torch
+
+# imitation.rewards.discrim_nets.DiscrimNetGAIL is composed of self.discriminator (nn.Module),
+# which gets called with inputs (state, action) when needed.
+
+class CnnDiscriminator(torch.nn.Module):
+    """ConvNet similar to stable_baselines3.common.policies.ActorCriticCnnPolicy."""
+
+    def __init__(self, env):
+        super().__init__()
+
+        obs_channels, _, _ = env.observation_space.shape
+        (action_size,) = env.action_space.shape
+        in_channels = obs_channels + action_size
+
+        self.cnn = torch.nn.Sequential(
+            torch.nn.Conv2d(in_channels, 32, kernel_size=(8, 8), stride=(4, 4)), # 5+1 -> 32
+            torch.nn.ReLU(),
+            torch.nn.Conv2d(32, 64, kernel_size=(4, 4), stride=(2, 2)), # 32 -> 64
+            torch.nn.ReLU(),
+            torch.nn.Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1)), # 64 -> 64
+            torch.nn.ReLU(),
+            torch.nn.Flatten(start_dim=1, end_dim=-1),
+            torch.nn.LazyLinear(512), # 28224 -> 512
+            torch.nn.ReLU(),
+            torch.nn.LazyLinear(1), # 512 -> 1
+        )
+    
+    @staticmethod
+    def _concatenate(state, action):
+        b, _, h, w = state.shape
+        _, a = action.shape
+        act = action.unsqueeze(-1).unsqueeze(-1).expand((b, a, h, w))
+        sa = torch.cat((state, act), -3)
+        return sa
+
+    def forward(self, state, action):
+        sa = self._concatenate(state, action)
+        return self.cnn(sa).squeeze()
+
+class MlpDiscriminator(torch.nn.Module):
+    """MLP similar to stable_baselines3.common.policies.ActorCriticPolicy."""
+
+    def __init__(self, env=None):
+        super().__init__()
+        self.flatten = torch.nn.Flatten(start_dim=1, end_dim=-1)
+        self.mlp = torch.nn.Sequential(
+            torch.nn.LazyLinear(64), # 42 -> 64
+            torch.nn.Tanh(),
+            torch.nn.LazyLinear(64), # 64 -> 64
+            torch.nn.Tanh(),
+            torch.nn.LazyLinear(1),  # 64 -> 1
+        )
+    
+    def forward(self, state, action):
+        flat = self.flatten(state)
+        sa = torch.cat((action, flat), -1)
+        return self.mlp(sa).squeeze()

scratch/etienne/intersimple/gail/test_discriminator.py

@@ -0,0 +1,45 @@
+from intersim.envs.intersimple import NRasterized
+from discriminator import CnnDiscriminator
+import torch
+
+def test_image_concatenation():
+    env = NRasterized()
+    disc = CnnDiscriminator(env)
+    s = torch.tensor(env.reset()).unsqueeze(0)
+    a = torch.tensor([[0.5]])
+    sa = disc._concatenate(s, a)
+
+    assert s.shape == (1, 5, 200, 200)
+    assert a.shape == (1, 1)
+    assert sa.shape == (1, 6, 200, 200)
+    assert torch.allclose(sa[:, :5], 1.0 * s)
+    assert (sa[:, 5] == a.unsqueeze(-1)).all()
+
+def test_image_concatenation3():
+    env = NRasterized()
+    disc = CnnDiscriminator(env)
+
+    s1 = env.reset()
+    a1 = 0.15
+    s2, _, _, _ = env.step(0.9)
+    a2 = 0.25
+    s3, _, _, _ = env.step(-0.9)
+    a3 = 0.35
+    
+    s = torch.stack([
+        torch.tensor(s1),
+        torch.tensor(s2),
+        torch.tensor(s3)
+    ], axis=0)
+    a = torch.tensor([
+        [a1],
+        [a2],
+        [a3],
+    ])
+    sa = disc._concatenate(s, a)
+
+    assert s.shape == (3, 5, 200, 200)
+    assert a.shape == (3, 1)
+    assert sa.shape == (3, 6, 200, 200)
+    assert torch.allclose(sa[:, :5], 1.0 * s)
+    assert (sa[:, 5] == a.unsqueeze(-1)).all()

scratch/etienne/intersimple/gail_flat.py

@@ -0,0 +1,106 @@
+# %%
+import pathlib
+import pickle
+import tempfile
+
+import stable_baselines3 as sb3
+from stable_baselines3.common.env_util import make_vec_env
+
+from imitation.algorithms import adversarial, bc
+from imitation.data import rollout
+from imitation.util import logger
+
+from intersimple.intersimple import IntersimpleReward, speed_reward
+
+from gail.discriminator import MlpDiscriminator
+import numpy as np
+import functools
+from stable_baselines3.common.evaluation import evaluate_policy
+from ray import tune
+import os
+import torch
+
+model_name = 'gail_flat'
+
+# %%
+# Load pickled test demonstrations.
+#with open("data/NormalizedIntersimpleExpert_IntersimpleRewardAgent51.pkl", "rb") as f:
+with open("data/NormalizedIntersimpleExpert_IntersimpleRewardAgent51Mu.001.pkl", "rb") as f:
+    # This is a list of `imitation.data.types.Trajectory`, where
+    # every instance contains observations and actions for a single expert
+    # demonstration.
+    trajectories = pickle.load(f)
+
+# %%
+# Convert List[types.Trajectory] to an instance of `imitation.data.types.Transitions`.
+# This is a more general dataclass containing unordered
+# (observation, actions, next_observation) transitions.
+transitions = rollout.flatten_trajectories(trajectories)
+
+venv = make_vec_env(IntersimpleReward, n_envs=2, env_kwargs={'agent': 51})
+
+tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
+tempdir_path = pathlib.Path(tempdir.name)
+print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
+
+def training_function(config, checkpoint_dir=None):
+    logger.configure(tempdir_path / "GAIL/")
+
+    discriminator = MlpDiscriminator()
+    if checkpoint_dir:
+        discriminator.load_state_dict(torch.load(os.path.join(checkpoint_dir, 'disc_checkpoint')))
+        generator = sb3.PPO.load(os.path.join(checkpoint_dir, 'gen_checkpoint'))
+    else:
+        generator = sb3.PPO("MlpPolicy", venv, verbose=1, n_steps=config['n_steps'])
+
+    gail_trainer = adversarial.GAIL(
+        venv,
+        expert_data=transitions,
+        expert_batch_size=config['expert_batch_size'],
+        n_disc_updates_per_round=config['n_disc_updates_per_round'],
+        discrim_kwargs={'discrim_net': MlpDiscriminator()},
+        gen_algo=generator,
+    )
+
+    def callback(epoch):
+        eval_env = IntersimpleReward(agent=51, reward=functools.partial(speed_reward, collision_penalty=0.))
+        #sync_envs_normalization(self.training_env, self.eval_env)
+        episode_rewards, episode_lengths = evaluate_policy(generator, eval_env)
+        tune.report(progress=np.mean(episode_rewards))
+
+        with tune.checkpoint_dir(step=epoch) as checkpoint_dir:
+            gail_trainer.gen_algo.save(os.path.join(checkpoint_dir, 'gen_checkpoint'))
+            torch.save(discriminator.state_dict(), os.path.join(checkpoint_dir, 'disc_checkpoint'))
+
+    gail_trainer.train(total_timesteps=400000, callback=callback)
+
+analysis = tune.run(
+    training_function,
+    config = {
+        'expert_batch_size': tune.randint(1, 220), #220,
+        'n_disc_updates_per_round': tune.randint(2, 100), #16,
+        'n_steps': tune.randint(1, 10000), #4096,
+    },
+    resources_per_trial={
+        'gpu': 1,
+    }
+)
+
+print('Best config', analysis.get_best_config(metric='progress', mode='max'))
+
+#del gail_trainer
+
+# %%
+model = sb3.PPO.load(model_name)
+
+env = IntersimpleReward(agent=51)
+
+obs = env.reset()
+while True:
+    action, _states = model.predict(obs)
+    obs, rewards, done, info = env.step(action)
+    env.render(mode='post')
+    if done:
+        break
+
+env.close(filestr='render/'+model_name)

scratch/etienne/intersimple/gail_image.py

@@ -0,0 +1,70 @@
+# %%
+import pathlib
+import pickle
+import tempfile
+
+import stable_baselines3 as sb3
+from stable_baselines3.common.env_util import make_vec_env
+
+from imitation.algorithms import adversarial, bc
+from imitation.data import rollout
+from imitation.util import logger
+
+from intersim.envs.intersimple import NRasterized
+
+from gail.discriminator import CnnDiscriminator
+
+model_name = 'gail_image'
+
+# %%
+# Load pickled test demonstrations.
+with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedAgent51w36h36mppx2.pkl", "rb") as f:
+    # This is a list of `imitation.data.types.Trajectory`, where
+    # every instance contains observations and actions for a single expert
+    # demonstration.
+    trajectories = pickle.load(f)
+
+# %%
+# Convert List[types.Trajectory] to an instance of `imitation.data.types.Transitions`.
+# This is a more general dataclass containing unordered
+# (observation, actions, next_observation) transitions.
+transitions = rollout.flatten_trajectories(trajectories)
+
+venv = make_vec_env(NRasterized, n_envs=2, env_kwargs={'agent': 51, 'width': 36, 'height': 36, 'm_per_px': 2})
+
+tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
+tempdir_path = pathlib.Path(tempdir.name)
+print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
+
+# Train GAIL on expert data.
+# GAIL, and AIRL also accept as `expert_data` any Pytorch-style DataLoader that
+# iterates over dictionaries containing observations, actions, and next_observations.
+logger.configure(tempdir_path / "GAIL/")
+gail_trainer = adversarial.GAIL(
+    venv,
+    expert_data=transitions,
+    expert_batch_size=32,
+    #n_disc_updates_per_round=2048,
+    discrim_kwargs={'discrim_net': CnnDiscriminator(venv)},
+    gen_algo=sb3.PPO("CnnPolicy", venv, verbose=1, n_steps=1024),
+    allow_variable_horizon=True,
+)
+gail_trainer.train(total_timesteps=100000)
+gail_trainer.gen_algo.save(model_name)
+
+#del gail_trainer
+
+# %%
+model = sb3.PPO.load(model_name)
+
+env = NRasterized(agent=51, width=36, height=36, m_per_px=2)
+
+obs = env.reset()
+while True:
+    action, _states = model.predict(obs)
+    obs, rewards, done, info = env.step(action)
+    env.render(mode='post')
+    if done:
+        break
+
+env.close(filestr='render/'+model_name)

scratch/etienne/intersimple/gail_options_image.py

@@ -0,0 +1,297 @@
+# %%
+from gail.discriminator import CnnDiscriminator
+from imitation.algorithms import adversarial
+import stable_baselines3
+import torch.utils.data
+import numpy as np
+from intersim.envs.intersimple import NRasterized
+import itertools
+from torch.distributions import Categorical
+import gym
+import torch
+import pickle
+import imitation.data.rollout as rollout
+import tempfile
+import pathlib
+from imitation.util import logger
+from stable_baselines3.common.env_util import make_vec_env
+
+model_name = 'gail_options_image'
+env_settings = {'agent': 51, 'width': 36, 'height': 36, 'm_per_px': 2}
+
+ALL_OPTIONS = [(v,t) for v in [0,2,4,6,8] for t in [5, 10, 20]] # option 0 is safe fallback
+
+class OptionsCnnPolicy(stable_baselines3.common.policies.ActorCriticCnnPolicy):
+
+    def __init__(self, observation_space, *args, **kwargs):
+        super().__init__(observation_space['obs'], *args, **kwargs)
+
+    def _prior_distribution(self, s):
+        latent_pi, latent_vf, latent_sde = self._get_latent(s)
+        distribution = self._get_action_dist_from_latent(latent_pi, latent_sde)
+        values = self.value_net(latent_vf)
+        return values, distribution.distribution
+
+    def predict(self, obs):
+        s, m = obs['obs'], obs['mask']
+        values, prior = self._prior_distribution(s)
+        posterior = Categorical(prior.probs * m)
+        ch = posterior.sample()
+        return ch, values, posterior.log_prob(ch)
+
+    def evaluate_actions(self, obs, ch):
+        s, m = obs['obs'], obs['mask']
+        values, prior = self._prior_distribution(s)
+        posterior = Categorical(prior.probs * m)
+        return values, posterior.log_prob(ch), posterior.entropy() # additional values used by PPO.train
+
+def available_actions(env):
+    """Return mask of available actions given current `env` state."""
+    valid = np.array([feasible(env, generate_plan(env, i), i) for i in range(len(ALL_OPTIONS))])
+    return valid
+
+def target_velocity_plan(current_v: float, target_v: float, t: int, dt: float):
+    """Smoothly target a velocity in a given number of steps"""
+    # for now, constant acceleration
+    a = (target_v - current_v)  / (t * dt)
+    return a*np.ones((t,))
+
+def generate_plan(env, i):
+    """Generate input profile for high-level action `i`."""
+    assert i < len(ALL_OPTIONS), "Invalid option index {i}"
+    target_v, t = ALL_OPTIONS[i]
+    current_v = env._env.state[env._agent, 1].item() # extract from env
+    plan = target_velocity_plan(current_v, target_v, t, env._env._dt)
+    assert len(plan) == t, "incorrect plan length"
+    return plan
+
+def check_future_collisions_fast(env, actions):
+    """Checks whether `env._agent` would collide with other agents assuming `actions` as input.
+    
+    Vehicles are (over-)approximated by single circles.
+
+    Args:
+        env (gym.Env): current environment state
+        actions (list of torch.Tensor): list of B (T, nv, adims) T-length action profiles
+    Returns:
+        feasible (torch.Tensor): tensor of shape (B,) indicating whether the respective action profiles are collision-free
+    """
+    B, (T, nv, _) = len(actions), actions[0].shape
+
+    states = torch.stack(env._env.propagate_action_profile(actions), axis=0)
+    assert states.shape == (B, T, nv, 5)
+
+    distance = ((states[:, :, :, :2] - states[:, :, env._agent:env._agent+1, :2])**2).sum(-1).sqrt()
+    distance = torch.where(distance.isnan(), np.inf*torch.ones_like(distance), distance) # only collide with spawned agents
+    distance[:, :, env._agent] = np.inf # cannot collide with itself
+    assert distance.shape == (B, T, nv)
+
+    radius = (env._env._lengths**2 + env._env._widths**2).sqrt() / 2
+    min_distance = radius[env._agent] + radius
+    min_distance = min_distance.unsqueeze(0).unsqueeze(0)
+    assert min_distance.shape == (1, 1, nv)
+
+    return (distance > min_distance).all(-1).all(-1)
+
+def feasible(env, plan, ch):
+    """Check if input profile is feasible given current `env` state. Action `ch=0` is safe fallback."""
+    
+    # zero pad plan - Take (T,) np plan and convert it to (T, nv, 1) torch.Tensor
+    full_plan = torch.zeros(len(plan), env._env._nv, 1)
+    full_plan[:, env._agent, 0] = torch.tensor(plan)
+    valid = check_future_collisions_fast(env, [full_plan]) # check_future_collisions_fast takes in B-list and outputs (B,) bool tensor
+    return ch == 0 or valid.item()
+
+def sample(env, generator, discriminator, level: str):
+    """
+        Sample low-level (state, action, next_state) tuples for discriminator training or 
+        high-level (state, action, reward) tuples for generator training.
+    """
+    done = True
+    while True:
+        episode_start = False
+        if done:
+            s = env.reset()
+            m = available_actions(env)
+            done = False
+            episode_start = True
+
+        obs = {'obs': s, 'mask': m}
+        ch, value, log_prob = generator.policy.predict({
+            'obs': torch.tensor(s).unsqueeze(0).to(generator.policy.device),
+            'mask': torch.tensor(m).unsqueeze(0).to(generator.policy.device),
+        })
+        plan = list(map(float, generate_plan(env, ch)))
+
+        assert not done
+        assert plan
+        assert feasible(env, plan, ch), f'Infeasible hl action {ch}'
+
+        r = 0
+        discount = 1
+        while not done and plan and feasible(env, plan, ch):
+            a, plan = env._normalize(plan[0]), plan[1:]
+            if level == 'high':
+                r += discount * discriminator.discrim_net.discriminator(
+                    torch.tensor(s).unsqueeze(0).to(discriminator.discrim_net.device()),
+                    torch.tensor([[a]]).to(discriminator.discrim_net.device()),
+                )
+                discount *= env.discount
+
+            nexts, _, done, _ = env.step(a)
+            m = available_actions(env)
+
+            if level == 'low':
+                yield {
+                    'obs': s,
+                    'next_obs': nexts,
+                    'acts': np.array((a,)),
+                    'dones': np.array(done),
+                }
+            s = nexts
+
+        if level == 'high':
+            yield {
+                'obs': obs,
+                'option': ch,
+                'reward': r.detach(),
+                'episode_start': episode_start,
+                'value': value.detach(),
+                'log_prob': log_prob.detach(),
+                'done': done,
+            }
+
+def flatten_transitions(transitions):
+    return {
+        'obs': np.stack(list(t['obs'] for t in transitions), axis=0),
+        'next_obs': np.stack(list(t['next_obs'] for t in transitions), axis=0),
+        'acts': np.stack(list(t['acts'] for t in transitions), axis=0),
+        'dones': np.stack(list(t['dones'] for t in transitions), axis=0),
+    }
+
+def train_discriminator(env, generator, discriminator, num_samples):
+    transitions = list(itertools.islice(sample(env, generator, None, 'low'), num_samples))
+    generator_samples = flatten_transitions(transitions)
+    discriminator.train_disc(gen_samples=generator_samples)
+
+def train_generator(env, generator, discriminator, num_samples):
+    generator_samples = list(itertools.islice(sample(env, generator, discriminator, 'high'), num_samples+1))
+    
+    generator.rollout_buffer.reset()
+    for s in generator_samples[:-1]:
+        generator.rollout_buffer.add(
+            obs=s['obs'],
+            action=s['option'].cpu(),
+            reward=s['reward'].cpu(),
+            episode_start=s['episode_start'],
+            value=s['value'],
+            log_prob=s['log_prob'],
+        )
+    
+    generator.rollout_buffer.compute_returns_and_advantage(
+        last_values=generator_samples[-1]['value'],
+        dones=generator_samples[-1]['done'],
+    )
+    
+    generator.train()
+
+class OptionsEnv(gym.Wrapper):
+    
+    def __init__(self, env):
+        super().__init__(env)
+        num_hl_options = len(ALL_OPTIONS)
+        self.action_space = gym.spaces.Discrete(num_hl_options)
+        self.observation_space = gym.spaces.Dict({
+            'obs': env.observation_space,
+            'mask': gym.spaces.Box(low=0, high=1, shape=(num_hl_options,)),
+        })
+
+def train(expert_data, epochs=10, expert_batch_size=32, generator_steps=2048, discount=0.99):
+    env = NRasterized(**env_settings)
+    env.discount = discount
+
+    tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
+    tempdir_path = pathlib.Path(tempdir.name)
+    logger.configure(tempdir_path / "GAIL/")
+    print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
+
+    venv = make_vec_env(NRasterized, n_envs=1, env_kwargs=env_settings)
+    discriminator = adversarial.GAIL(
+        expert_data=expert_data,
+        expert_batch_size=expert_batch_size,
+        discrim_kwargs={'discrim_net': CnnDiscriminator(venv)},
+        venv=venv, # unused
+        gen_algo=stable_baselines3.PPO("CnnPolicy", venv), # unused
+    )
+
+    generator = stable_baselines3.PPO(
+        OptionsCnnPolicy,
+        OptionsEnv(env),
+        verbose=1,
+        n_steps=generator_steps,
+    )
+
+    # PPO.train requires logger as set up in
+    # PPO._setup_learn (called by PPO.learn)
+    generator._logger = stable_baselines3.common.utils.configure_logger(
+        generator.verbose,
+        generator.tensorboard_log,
+    )
+
+    for _ in range(epochs):
+        train_discriminator(env, generator, discriminator, num_samples=expert_batch_size)
+        train_generator(env, generator, discriminator, num_samples=generator_steps)
+    
+    return generator
+
+# %%
+if __name__ == '__main__':
+    # %%
+    with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedAgent51w36h36mppx2.pkl", "rb") as f:
+        trajectories = pickle.load(f)
+    transitions = rollout.flatten_trajectories(trajectories)
+    generator = train(transitions, generator_steps=200)
+
+    generator.save(model_name)
+
+    # %%
+    model = stable_baselines3.PPO.load(model_name)
+
+    env = NRasterized(**env_settings)
+
+    for transition in sample(env, generator, None, 'low'):
+        env.render()
+        if transition['dones']:
+            break
+
+    env.close(filestr='render/'+model_name)
+
+# %% Tests
+
+def test_ll_transitions_vs_expert_data():
+    with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedAgent51w36h36mppx2.pkl", "rb") as f:
+        expert_trajectories = pickle.load(f)
+    expert_transitions = rollout.flatten_trajectories(expert_trajectories)
+
+    env = NRasterized(agent=51, width=36, height=36, m_per_px=2)
+
+    gen_transitions = list(itertools.islice(sample(
+        env=NRasterized(**env_settings),
+        generator=stable_baselines3.PPO(
+            OptionsCnnPolicy,
+            OptionsEnv(env),
+            verbose=1,
+        ),
+        discriminator=None,
+        level='low'
+    ), 10))
+    gen_transitions = flatten_transitions(gen_transitions)
+
+    assert expert_transitions[:10].obs.shape == gen_transitions['obs'].shape
+    assert expert_transitions[:10].next_obs.shape == gen_transitions['next_obs'].shape
+    assert expert_transitions[:10].acts.shape == gen_transitions['acts'].shape
+    assert expert_transitions[:10].dones.shape == gen_transitions['dones'].shape
+
+
+def test_hl_transitions():
+    pass

scratch/etienne/intersimple/imitation_quickstart.py

@@ -0,0 +1,63 @@
+# %%
+import pathlib
+import pickle
+import tempfile
+
+import stable_baselines3 as sb3
+from stable_baselines3.common.env_util import make_vec_env
+
+from imitation.algorithms import adversarial, bc
+from imitation.data import rollout
+from imitation.util import logger
+
+from intersim.envs.intersimple import IntersimpleReward
+
+# Load pickled test demonstrations.
+with open("data/NormalizedIntersimpleExpert_IntersimpleRewardAgent51.pkl", "rb") as f:
+    # This is a list of `imitation.data.types.Trajectory`, where
+    # every instance contains observations and actions for a single expert
+    # demonstration.
+    trajectories = pickle.load(f)
+
+# %%
+# Convert List[types.Trajectory] to an instance of `imitation.data.types.Transitions`.
+# This is a more general dataclass containing unordered
+# (observation, actions, next_observation) transitions.
+transitions = rollout.flatten_trajectories(trajectories)
+
+venv = make_vec_env(IntersimpleReward, n_envs=2, env_kwargs={'agent': 51})
+
+tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
+tempdir_path = pathlib.Path(tempdir.name)
+print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
+
+# Train BC on expert data.
+# BC also accepts as `expert_data` any PyTorch-style DataLoader that iterates over
+# dictionaries containing observations and actions.
+logger.configure(tempdir_path / "BC/")
+bc_trainer = bc.BC(venv.observation_space, venv.action_space, expert_data=transitions)
+bc_trainer.train(n_epochs=1)
+
+# Train GAIL on expert data.
+# GAIL, and AIRL also accept as `expert_data` any Pytorch-style DataLoader that
+# iterates over dictionaries containing observations, actions, and next_observations.
+logger.configure(tempdir_path / "GAIL/")
+gail_trainer = adversarial.GAIL(
+    venv,
+    expert_data=transitions,
+    expert_batch_size=32,
+    gen_algo=sb3.PPO("MlpPolicy", venv, verbose=1, n_steps=1024),
+)
+gail_trainer.train(total_timesteps=2048)
+
+# Train AIRL on expert data.
+logger.configure(tempdir_path / "AIRL/")
+airl_trainer = adversarial.AIRL(
+    venv,
+    expert_data=transitions,
+    expert_batch_size=32,
+    gen_algo=sb3.PPO("MlpPolicy", venv, verbose=1, n_steps=1024),
+)
+airl_trainer.train(total_timesteps=2048)
+
+# %%

scratch/etienne/intersimple/ppo_const.py

@@ -0,0 +1,36 @@
+# %%
+from stable_baselines3 import PPO
+from stable_baselines3.common.env_util import make_vec_env
+from intersim.envs.intersimple import IntersimpleReward, speed_reward
+
+model_name = "ppo_const"
+
+env = IntersimpleReward(
+    agent=51,
+    #reward=speed_reward,
+)
+
+# %%
+model = PPO(
+    "MlpPolicy", env,
+    verbose=1,
+)
+model.learn(total_timesteps=100000)
+model.save(model_name)
+
+print('Done training.')
+
+del model # remove to demonstrate saving and loading
+
+# %%
+model = PPO.load(model_name)
+
+obs = env.reset()
+while True:
+    action, _states = model.predict(obs)
+    obs, rewards, done, info = env.step(action)
+    env.render(mode='post')
+    if done:
+        break
+
+env.close(filestr='render/'+model_name)

scratch/etienne/intersimple/ppo_const_collision.py

@@ -0,0 +1,41 @@
+# %%
+from stable_baselines3 import PPO
+from stable_baselines3.common.env_util import make_vec_env
+from intersim.envs.intersimple import ConstCollisionReward, IntersimpleFlatAgent
+
+model_name = "ppo_const_collision"
+
+class IntersimpleConstCollisionAgent(ConstCollisionReward, IntersimpleFlatAgent):
+    pass
+
+env = IntersimpleConstCollisionAgent(
+    agent=51,
+    speed_reward_weight=0.001,
+    collision_penalty=1000
+)
+
+# %%
+model = PPO(
+    "MlpPolicy", env,
+    learning_rate=3e-6,
+    verbose=1,
+)
+model.learn(total_timesteps=2e5)
+model.save(model_name)
+
+print('Done training.')
+
+del model # remove to demonstrate saving and loading
+
+# %%
+model = PPO.load(model_name)
+
+obs = env.reset()
+while True:
+    action, _states = model.predict(obs)
+    obs, rewards, done, info = env.step(action)
+    env.render(mode='post')
+    if done:
+        break
+
+env.close(filestr='render/'+model_name)

scratch/etienne/intersimple/ppo_const_image.py

@@ -0,0 +1,35 @@
+# %%
+from stable_baselines3 import PPO
+from stable_baselines3.common.env_util import make_vec_env
+from intersim.envs.intersimple import NRasterized
+
+model_name = "ppo_const_image"
+
+env = NRasterized(
+    agent=51,
+)
+
+# %%
+model = PPO(
+    "CnnPolicy", env,
+    verbose=1,
+)
+model.learn(total_timesteps=100000)
+model.save(model_name)
+
+print('Done training.')
+
+del model # remove to demonstrate saving and loading
+
+# %%
+model = PPO.load(model_name)
+
+obs = env.reset()
+while True:
+    action, _states = model.predict(obs)
+    obs, rewards, done, info = env.step(action)
+    env.render(mode='post')
+    if done:
+        break
+
+env.close(filestr='render/'+model_name)

scratch/etienne/intersimple/ppo_const_image_random.py

@@ -0,0 +1,33 @@
+# %%
+from stable_baselines3 import PPO
+from intersim.envs.intersimple import NRasterizedRandomAgent
+import functools
+
+model_name = "ppo_const_image_random"
+
+env = NRasterizedRandomAgent()
+
+# %%
+model = PPO(
+    "CnnPolicy", env,
+    verbose=1,
+)
+model.learn(total_timesteps=2e5)
+model.save(model_name)
+
+print('Done training.')
+
+del model # remove to demonstrate saving and loading
+
+# %%
+model = PPO.load(model_name)
+
+obs = env.reset()
+while True:
+    action, _states = model.predict(obs)
+    obs, rewards, done, info = env.step(action)
+    env.render(mode='post')
+    if done:
+        break
+
+env.close(filestr='render/'+model_name)

scratch/etienne/intersimple/ppo_intersimple_tspeed.py

@@ -0,0 +1,25 @@
+from stable_baselines3 import PPO
+from stable_baselines3.common.env_util import make_vec_env
+from intersim.envs.intersimple import IntersimpleTargetSpeed
+
+env = IntersimpleTargetSpeed()
+
+model = PPO("MlpPolicy", env, verbose=1)
+model.learn(total_timesteps=25000)
+model.save("ppo_intersimple")
+
+print('Done training.')
+
+del model # remove to demonstrate saving and loading
+
+model = PPO.load("ppo_intersimple")
+
+obs = env.reset()
+while True:
+    action, _states = model.predict(obs)
+    obs, rewards, done, info = env.step(action)
+    env.render(mode='post')
+    if done:
+        break
+
+env.close()

scratch/etienne/intersimple/ppo_speed.py

@@ -0,0 +1,46 @@
+# %%
+from stable_baselines3 import PPO
+from intersim.envs.intersimple import IntersimpleReward, speed_reward
+import functools
+
+model_name = "ppo_speed"
+
+#def reward(state, action, info):
+#    speed = state[2].item()
+#    r = speed if speed < 10 else (10 - 5 * (speed - 10))
+#    return 0.1 * r
+
+env = IntersimpleReward(
+    agent=51,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+)
+
+# %%
+model = PPO(
+    "MlpPolicy", env,
+    verbose=1,
+)
+model.learn(total_timesteps=100000)
+model.save(model_name)
+
+print('Done training.')
+
+del model # remove to demonstrate saving and loading
+
+# %%
+model = PPO.load(model_name)
+
+obs = env.reset()
+while True:
+    action, _states = model.predict(obs)
+    obs, rewards, done, info = env.step(action)
+    env.render(mode='post')
+    if done:
+        break
+
+env.close(filestr='render/'+model_name)
+
+# %%

scratch/etienne/intersimple/ppo_speed_image.py

@@ -0,0 +1,46 @@
+# %%
+from stable_baselines3 import PPO
+from intersim.envs.intersimple import NRasterized, speed_reward
+import functools
+
+model_name = "ppo_speed_image"
+
+#def reward(state, action, info):
+#    speed = state[2].item()
+#    r = speed if speed < 10 else (10 - 5 * (speed - 10))
+#    return 0.1 * r
+
+env = NRasterized(
+    agent=20,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+)
+
+# %%
+model = PPO(
+    "CnnPolicy", env,
+    verbose=1,
+)
+model.learn(total_timesteps=100000)
+model.save(model_name)
+
+print('Done training.')
+
+del model # remove to demonstrate saving and loading
+
+# %%
+model = PPO.load(model_name)
+
+obs = env.reset()
+while True:
+    action, _states = model.predict(obs)
+    obs, rewards, done, info = env.step(action)
+    env.render(mode='post')
+    if done:
+        break
+
+env.close(filestr='render/'+model_name)
+
+# %%

scratch/etienne/intersimple/ppo_speed_image_lowres.py

@@ -0,0 +1,49 @@
+# %%
+from stable_baselines3 import PPO
+from intersim.envs.intersimple import NRasterized, speed_reward
+import functools
+
+model_name = "ppo_speed_image_lowres"
+
+#def reward(state, action, info):
+#    speed = state[2].item()
+#    r = speed if speed < 10 else (10 - 5 * (speed - 10))
+#    return 0.1 * r
+
+env = NRasterized(
+    agent=51,
+    height=36,
+    width=36,
+    m_per_px=2,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+)
+
+# %%
+model = PPO(
+    "CnnPolicy", env,
+    verbose=1,
+)
+model.learn(total_timesteps=100000)
+model.save(model_name)
+
+print('Done training.')
+
+del model # remove to demonstrate saving and loading
+
+# %%
+model = PPO.load(model_name)
+
+obs = env.reset()
+while True:
+    action, _states = model.predict(obs)
+    obs, rewards, done, info = env.step(action)
+    env.render(mode='post')
+    if done:
+        break
+
+env.close(filestr='render/'+model_name)
+
+# %%

scratch/etienne/intersimple/ppo_speed_image_lowres_random.py

@@ -0,0 +1,42 @@
+# %%
+from stable_baselines3 import PPO
+from intersim.envs.intersimple import NRasterizedRandomAgent, speed_reward
+import functools
+
+model_name = "ppo_speed_image_lowres_random"
+
+env = NRasterizedRandomAgent(
+    height=36,
+    width=36,
+    m_per_px=2,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    )
+)
+
+# %%
+model = PPO(
+    "CnnPolicy", env,
+    verbose=1,
+    batch_size=2048,
+)
+model.learn(total_timesteps=2e5)
+model.save(model_name)
+
+print('Done training.')
+
+del model # remove to demonstrate saving and loading
+
+# %%
+model = PPO.load(model_name)
+
+obs = env.reset()
+while True:
+    action, _states = model.predict(obs)
+    obs, rewards, done, info = env.step(action)
+    env.render(mode='post')
+    if done:
+        break
+
+env.close(filestr='render/'+model_name)

scratch/etienne/intersimple/ppo_speed_image_random.py

@@ -0,0 +1,39 @@
+# %%
+from stable_baselines3 import PPO
+from intersim.envs.intersimple import NRasterizedRandomAgent, speed_reward
+import functools
+
+model_name = "ppo_speed_image_random"
+
+env = NRasterizedRandomAgent(
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    )
+)
+
+# %%
+model = PPO(
+    "CnnPolicy", env,
+    verbose=1,
+    batch_size=2048,
+)
+model.learn(total_timesteps=2e5)
+model.save(model_name)
+
+print('Done training.')
+
+del model # remove to demonstrate saving and loading
+
+# %%
+model = PPO.load(model_name)
+
+obs = env.reset()
+while True:
+    action, _states = model.predict(obs)
+    obs, rewards, done, info = env.step(action)
+    env.render(mode='post')
+    if done:
+        break
+
+env.close(filestr='render/'+model_name)

scratch/etienne/intersimple/ppo_speed_random.py

@@ -0,0 +1,43 @@
+# %%
+from stable_baselines3 import PPO
+from intersim.envs.intersimple import IntersimpleFlatRandomAgent, Reward, RewardVisualization, speed_reward
+import functools
+
+model_name = "ppo_speed_random"
+
+class IntersimpleRewardRandom(RewardVisualization, Reward, IntersimpleFlatRandomAgent):
+    """`IntersimpleFlatAgent` with rewards."""
+    pass
+
+env = IntersimpleRewardRandom(
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    )
+)
+
+# %%
+model = PPO(
+    "MlpPolicy", env,
+    verbose=1,
+    batch_size=2048,
+)
+model.learn(total_timesteps=2e5)
+model.save(model_name)
+
+print('Done training.')
+
+del model # remove to demonstrate saving and loading
+
+# %%
+model = PPO.load(model_name)
+
+obs = env.reset()
+while True:
+    action, _states = model.predict(obs)
+    obs, rewards, done, info = env.step(action)
+    env.render(mode='post')
+    if done:
+        break
+
+env.close(filestr='render/'+model_name)

scratch/etienne/intersimple/ppo_tspeed.py

@@ -0,0 +1,39 @@
+# %%
+from stable_baselines3 import PPO
+from stable_baselines3.common.env_util import make_vec_env
+from intersim.envs.intersimple import IntersimpleTargetSpeedAgent
+
+model_name = "ppo_tspeed"
+
+env = IntersimpleTargetSpeedAgent(
+    agent=51,
+    target_speed=10,
+    speed_penalty_weight=0.001,
+    collision_penalty=1000
+)
+
+# %%
+model = PPO(
+    "MlpPolicy", env,
+    learning_rate=3e-6,
+    verbose=1,
+)
+model.learn(total_timesteps=2e5)
+model.save(model_name)
+
+print('Done training.')
+
+del model # remove to demonstrate saving and loading
+
+# %%
+model = PPO.load(model_name)
+
+obs = env.reset()
+while True:
+    action, _states = model.predict(obs)
+    obs, rewards, done, info = env.step(action)
+    env.render(mode='post')
+    if done:
+        break
+
+env.close(filestr='render/'+model_name)

scratch/etienne/intersimple/ppo_tspeed_random.py

@@ -0,0 +1,31 @@
+# %%
+from stable_baselines3 import PPO
+from stable_baselines3.common.env_util import make_vec_env
+from intersim.envs.intersimple import IntersimpleTargetSpeedRandom
+
+model_name = "ppo_tspeed_random"
+
+# %%
+env = IntersimpleTargetSpeedRandom(target_speed=10)
+
+# %%
+model = PPO("MlpPolicy", env, verbose=1)
+model.learn(total_timesteps=250000)
+model.save(model_name)
+
+print('Done training.')
+
+del model # remove to demonstrate saving and loading
+
+# %%
+model = PPO.load(model_name)
+
+obs = env.reset()
+while True:
+    action, _states = model.predict(obs)
+    obs, rewards, done, info = env.step(action)
+    env.render(mode='post')
+    if done:
+        break
+
+env.close()

scratch/etienne/pillbox/intersim_demos.ipynb

@@ -0,0 +1,249 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "source": [
+    "%cd learners"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "/home/buehrle/dev/InteractionImitation/scratch/etienne/pillbox/learners\n"
+     ]
+    }
+   ],
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2"
+   ],
+   "outputs": [],
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "source": [
+    "import torch\n",
+    "import numpy as np"
+   ],
+   "outputs": [],
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "source": [
+    "# save expert demos to ../experts/Intersim/demos.npz\n",
+    "# make sure to split different experts up\n",
+    "\n",
+    "from intersim.envs.simulator import InteractionSimulator\n",
+    "from intersim.utils import get_map_path, get_svt, SVT_to_stateactions\n",
+    "import gym\n",
+    "from tqdm import tqdm\n",
+    "\n",
+    "def pillbox_demo(observations, actions, rewards):\n",
+    "    demo = {\n",
+    "        'env': 'intersim:intersim-v0',\n",
+    "        'num_trajs': len(observations),\n",
+    "        'mean_reward': rewards.mean(),\n",
+    "        'std_reward': rewards.std(),\n",
+    "    }\n",
+    "    demo.update({\n",
+    "        str(i): {\n",
+    "            'states': o,\n",
+    "            'actions': a,\n",
+    "        } for i, (o, a) in enumerate(zip(observations, actions))\n",
+    "    })\n",
+    "    return demo\n",
+    "\n",
+    "def intersim_expert_demos(loc, track):\n",
+    "    svt, svt_path = get_svt(loc, track)\n",
+    "    osm = get_map_path(loc)\n",
+    "    \n",
+    "    n_actors = svt.simstate.size(1)\n",
+    "    observations = []\n",
+    "    actions = [] ##\n",
+    "    #states, actions = SVT_to_stateactions(svt) ##\n",
+    "    rewards = []\n",
+    "    \n",
+    "    print('Simulating')\n",
+    "    env = gym.make('intersim:intersim-v0', svt=svt, map_path=osm)\n",
+    "    obs, info = env.reset()\n",
+    "    for s in tqdm(svt.simstate[1:]): ##\n",
+    "    #for a in actions: ##\n",
+    "        relative_state = torch.stack((\n",
+    "            obs['relative_state'][..., 0],\n",
+    "            obs['relative_state'][..., 1],\n",
+    "            (obs['relative_state'][..., 2]**2 + obs['relative_state'][..., 3]**2).sqrt(),\n",
+    "            obs['relative_state'][..., 4],\n",
+    "            obs['relative_state'][..., 5],\n",
+    "        ), -1)\n",
+    "        observations.append(torch.cat((\n",
+    "            obs['state'].unsqueeze(1),\n",
+    "            relative_state,\n",
+    "        ), 1))\n",
+    "        obs, r, done, info = env.step(env.target_state(s, mu=.01))\n",
+    "        #obs, r, done, info = env.step(a) ##\n",
+    "        actions.append(info['action_taken'])\n",
+    "        rewards.append(r)\n",
+    "        assert not done, 'Episode terminated during expert demonstration.'\n",
+    "\n",
+    "    _except_idx = lambda o, i: torch.cat((o[:i], o[i+1:]))\n",
+    "        \n",
+    "    # transpose to per-agent observations and actions\n",
+    "    print('Transposing')\n",
+    "    observations = [torch.stack([_except_idx(o[i], i+1) for o in observations]) for i in range(n_actors)]\n",
+    "    actions = [torch.stack([a[i] for a in actions]) for i in range(n_actors)]\n",
+    "    \n",
+    "    print('Trimming')\n",
+    "    # trim observations and actions to start/end of trajectory\n",
+    "    _alive = lambda o: (~o.isnan().all(2).all(1)).nonzero()\n",
+    "    _start = lambda o: _alive(o).min()\n",
+    "    _end = lambda o: _alive(o).max() + 1\n",
+    "    start_end = [(_start(obs), _end(obs)) for obs in observations]\n",
+    "    observations = [obs[start:end] for obs, (start, end) in zip(observations, start_end)]\n",
+    "    actions = [act[start:end] for act, (start, end) in zip(actions, start_end)]\n",
+    "    \n",
+    "    #print('Cropping')\n",
+    "    ## crop observations to max number of observations\n",
+    "    #max_num_obs = max([(~obs.isnan().all(2)).sum(1).max() for obs in observations])\n",
+    "    #observations = [obs[:, :max_num_obs] for obs in observations]\n",
+    "    \n",
+    "    observations = [o.numpy() for o in observations]\n",
+    "    actions = [a.numpy() for a in actions]\n",
+    "    rewards = np.array(rewards)\n",
+    "    \n",
+    "    return pillbox_demo(observations, actions, rewards)"
+   ],
+   "outputs": [],
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "source": [
+    "demos = intersim_expert_demos(loc=0, track=0)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "Simulating\n",
+      "Custom Vehicle Trajectory Paths\n",
+      "Map Path: datasets/maps/DR_USA_Roundabout_FT.osm\n",
+      "Environment Reset\n"
+     ]
+    },
+    {
+     "output_type": "stream",
+     "name": "stderr",
+     "text": [
+      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████| 3006/3006 [01:17<00:00, 38.87it/s]\n"
+     ]
+    },
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "Transposing\n",
+      "Trimming\n"
+     ]
+    }
+   ],
+   "metadata": {
+    "scrolled": true,
+    "tags": [
+     "outputPrepend"
+    ]
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "source": [
+    "demos['num_trajs']"
+   ],
+   "outputs": [
+    {
+     "output_type": "execute_result",
+     "data": {
+      "text/plain": [
+       "151"
+      ]
+     },
+     "metadata": {},
+     "execution_count": 6
+    }
+   ],
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "source": [
+    "demos['25']['states'].shape"
+   ],
+   "outputs": [
+    {
+     "output_type": "execute_result",
+     "data": {
+      "text/plain": [
+       "(71, 151, 5)"
+      ]
+     },
+     "metadata": {},
+     "execution_count": 7
+    }
+   ],
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "source": [
+    "np.savez('../experts/intersim:intersim-v0/demos.npz', **demos)"
+   ],
+   "outputs": [],
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "source": [],
+   "outputs": [],
+   "metadata": {}
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "name": "python3",
+   "display_name": "Python 3.7.5 64-bit ('.venv': venv)"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.7.5"
+  },
+  "interpreter": {
+   "hash": "56465d2ea10f338edb3d30adb010c5849fd826fffc543ba31360f3db8b47a703"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

scratch/etienne/pillbox/intersim_expert.ipynb

@@ -0,0 +1,168 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "source": [
+    "%cd learners"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "[Errno 2] No such file or directory: 'learners'\n",
+      "/home/buehrle/dev/InteractionImitation/scratch/etienne/pillbox/learners\n"
+     ]
+    }
+   ],
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "source": [
+    "import gym\n",
+    "from tqdm import tqdm\n",
+    "\n",
+    "def rollout(pi, max_steps=1000):\n",
+    "    env = gym.make('intersim:intersim-v0')\n",
+    "    env.reset() # obs = env.reset()\n",
+    "    obs, _, done, _ = env.step(0 * env.action_space.sample())\n",
+    "    \n",
+    "    _except = lambda o, i: torch.cat((o[:i], o[i+1:]))\n",
+    "    \n",
+    "    _relative_state_v = lambda obs: torch.stack((\n",
+    "        obs[..., 0],\n",
+    "        obs[..., 1],\n",
+    "        (obs[..., 2]**2 + obs[..., 3]**2).sqrt(),\n",
+    "        obs[..., 4],\n",
+    "        obs[..., 5],\n",
+    "    ), -1)\n",
+    "    \n",
+    "    for _ in tqdm(range(max_steps)):\n",
+    "        pi_obs = [\n",
+    "            torch.cat((e.unsqueeze(0), _relative_state_v(_except(o, i)))).unsqueeze(0)\n",
+    "            for i, (e, o) in enumerate(zip(obs['state'], obs['relative_state']))\n",
+    "        ]\n",
+    "        \n",
+    "        actions = [pi(o).squeeze() for o in pi_obs]\n",
+    "        actions = torch.stack(actions).unsqueeze(1)\n",
+    "        obs, _, done, _ = env.step(actions)\n",
+    "        env.render(mode='post')\n",
+    "        if done:\n",
+    "            break\n",
+    "    env.close()"
+   ],
+   "outputs": [],
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "source": [
+    "import torch\n",
+    "import numpy as np\n",
+    "\n",
+    "def expert(obs):\n",
+    "    ego = obs[:, 0]\n",
+    "    rel = obs[:, 1:]\n",
+    "    front = torch.stack((torch.cos(ego[:, 3]), torch.sin(ego[:, 3])), -1)\n",
+    "    left = torch.stack((-torch.sin(ego[:, 3]), torch.cos(ego[:, 3])), -1)\n",
+    "    df = (rel[:, :, :2] * front.unsqueeze(1)).sum(-1)\n",
+    "    dl = (rel[:, :, :2] * left.unsqueeze(1)).sum(-1)\n",
+    "\n",
+    "    df = torch.where(df.isnan(), np.inf * torch.ones_like(df), df)\n",
+    "    dl = torch.where(dl.isnan(), np.inf * torch.ones_like(dl), dl)\n",
+    "    rel = torch.where(rel.isnan(), np.inf * torch.ones_like(rel), rel)\n",
+    "\n",
+    "    # relative speed in direction of position difference vector\n",
+    "    vrel = rel[:, :, 2] * (rel[:, :, :2] * torch.stack((\n",
+    "        torch.cos(ego[:, 3].unsqueeze(1) + rel[:, :, 3]),\n",
+    "        torch.sin(ego[:, 3].unsqueeze(1) + rel[:, :, 3])),\n",
+    "        -1)).sum(-1)\n",
+    "    vrel = torch.where(vrel.isnan(), np.inf * torch.ones_like(vrel), vrel)\n",
+    "    vrel = torch.maximum(vrel, torch.zeros_like(vrel))\n",
+    "    \n",
+    "    alpha = torch.atan2(dl, df)\n",
+    "    d = (rel[:, :, :2] ** 2).sum(-1)\n",
+    "    attn = torch.exp(-torch.where(alpha > 0, 0.8*alpha, 1*alpha)**2 - 0.01 * d - 0.1*vrel) \n",
+    "    \n",
+    "    act = 10 - ego[:, 2] - 20 * attn.sum(-1)\n",
+    "    \n",
+    "    return act"
+   ],
+   "outputs": [],
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "source": [
+    "rollout(expert, max_steps=500)"
+   ],
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "Vehicle Trajectory Paths: /home/buehrle/dev/InteractionImitation/InteractionSimulator/datasets/trackfiles/DR_USA_Roundabout_FT/vehicle_tracks_000.csv\n",
+      "Map Path: /home/buehrle/dev/InteractionImitation/InteractionSimulator/datasets/maps/DR_USA_Roundabout_FT.osm\n",
+      "Environment Reset\n"
+     ]
+    },
+    {
+     "output_type": "stream",
+     "name": "stderr",
+     "text": [
+      "  0%|          | 0/500 [00:00<?, ?it/s]\n"
+     ]
+    },
+    {
+     "output_type": "error",
+     "ename": "RuntimeError",
+     "evalue": "torch.cat(): Sizes of tensors must match except in dimension 0. Got 5 and 6 in dimension 1 (The offending index is 1)",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
+      "\u001b[0;32m/tmp/ipykernel_4266/633354333.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mrollout\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mexpert\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmax_steps\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m500\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
+      "\u001b[0;32m/tmp/ipykernel_4266/1839273282.py\u001b[0m in \u001b[0;36mrollout\u001b[0;34m(pi, max_steps)\u001b[0m\n\u001b[1;32m     12\u001b[0m         pi_obs = [\n\u001b[1;32m     13\u001b[0m             \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0me\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0munsqueeze\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_except_self\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mo\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mi\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0munsqueeze\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 14\u001b[0;31m             \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0me\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mo\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32min\u001b[0m \u001b[0menumerate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mzip\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mobs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'state'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mobs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'relative_state'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     15\u001b[0m         ]\n\u001b[1;32m     16\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m/tmp/ipykernel_4266/1839273282.py\u001b[0m in \u001b[0;36m<listcomp>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m     12\u001b[0m         pi_obs = [\n\u001b[1;32m     13\u001b[0m             \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0me\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0munsqueeze\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_except_self\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mo\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mi\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0munsqueeze\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 14\u001b[0;31m             \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0me\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mo\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32min\u001b[0m \u001b[0menumerate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mzip\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mobs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'state'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mobs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'relative_state'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     15\u001b[0m         ]\n\u001b[1;32m     16\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;31mRuntimeError\u001b[0m: torch.cat(): Sizes of tensors must match except in dimension 0. Got 5 and 6 in dimension 1 (The offending index is 1)"
+     ]
+    }
+   ],
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "source": [],
+   "outputs": [],
+   "metadata": {}
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "name": "python3",
+   "display_name": "Python 3.7.5 64-bit ('.venv': venv)"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.7.5"
+  },
+  "interpreter": {
+   "hash": "56465d2ea10f338edb3d30adb010c5849fd826fffc543ba31360f3db8b47a703"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

scratch/etienne/pillbox/learners/adril.py

@@ -0,0 +1,140 @@
+import gym
+from gym import spaces
+from sklearn.neighbors import KDTree
+from scipy.stats import norm
+import numpy as np
+import warnings
+from abc import ABC, abstractmethod
+from typing import Dict, Generator, Optional, Union
+import torch as th
+
+try:
+    # Check memory used by replay buffer when possible
+    import psutil
+except ImportError:
+    psutil = None
+
+from stable_baselines3.common.preprocessing import get_action_dim, get_obs_shape
+from stable_baselines3.common.type_aliases import ReplayBufferSamples, RolloutBufferSamples
+from stable_baselines3.common.vec_env import VecNormalize
+from stable_baselines3.common.buffers import ReplayBuffer
+
+
+class AdRILWrapper(gym.Env):
+  metadata = {'render.modes': ['human']}
+
+  def __init__(self, base_env):
+    super(AdRILWrapper, self).__init__()
+    self.base_env = base_env
+    self.iter = 0
+    self.observation_space = self.base_env.observation_space
+    self.action_space = self.base_env.action_space
+    self.trajs = list()
+    self.num_trajs = 0
+    self.curr_state = None
+  def step(self, action):
+    next_obs, _, done, info = self.base_env.step(action)
+    reward = self.iter # Transformed by replay buffer
+    self.trajs.append((self.curr_state, action, next_obs, done))
+    if done:
+        self.num_trajs += 1
+    self.curr_state = next_obs
+    return next_obs, reward, done, info
+  def reset(self):
+    obs = self.base_env.reset()
+    self.curr_state = obs
+    return obs
+  def render(self, mode='human'):
+    self.base_env.render(mode=mode)
+  def close (self):
+    self.base_env.close()
+  def get_learner_trajs(self):
+    return self.trajs
+  def set_iter(self, k):
+    self.iter = k
+    
+class AdRILReplayBuffer(ReplayBuffer):
+    def __init__(
+        self,
+        buffer_size: int,
+        observation_space: spaces.Space,
+        action_space: spaces.Space,
+        device: Union[th.device, str] = "cpu",
+        n_envs: int = 1,
+        optimize_memory_usage: bool = False,
+        expert_data: dict = dict(),
+            N_expert: int = 0,
+        balanced: bool = True,
+    ):
+        super(AdRILReplayBuffer, self).__init__(buffer_size, observation_space, action_space, device, n_envs=n_envs, optimize_memory_usage=optimize_memory_usage)
+
+        self.expert_states = expert_data['obs']
+        self.expert_actions = expert_data['acts']
+        self.expert_next_states = expert_data['next_obs']
+        self.expert_dones = expert_data['dones']
+        n_expert = len(expert_data["obs"])
+        self.iter = 0
+        self.N_expert = N_expert
+        self.N_learner = 0
+        self.normalizer = 1
+        self.balanced = balanced
+
+    def set_iter(self, k):
+      self.iter = k
+      normalizer = 0
+      for i in range(0, k):
+        normalizer += 1 ** (-i) # written to support decaying learning rate
+      self.normalizer = normalizer
+
+    def set_n_learner(self, n):
+        self.N_learner = n
+
+    def _get_samples(self, batch_inds: np.ndarray, env: Optional[VecNormalize] = None) -> ReplayBufferSamples:
+        num_samples = len(batch_inds)
+        if self.balanced:
+          num_expert_samples = int(num_samples / 2)
+          batch_inds = batch_inds[:num_expert_samples]
+          expert_inds = np.random.randint(0, len(self.expert_states), size=num_expert_samples)
+          # balanced sampling
+          if self.optimize_memory_usage:
+              next_obs = self._normalize_obs(self.observations[(batch_inds + 1) % self.buffer_size, 0, :], env)
+          else:
+              next_obs = self._normalize_obs(self.next_observations[batch_inds, 0, :], env)
+          next_obs = np.concatenate((next_obs, self._normalize_obs(self.expert_next_states[expert_inds], env)), axis=0)
+          obs = self._normalize_obs(self.observations[batch_inds, 0, :], env)
+          obs = np.concatenate((obs, self._normalize_obs(self.expert_states[expert_inds], env)), axis=0)
+          actions = self.actions[batch_inds, 0, :]
+          actions = np.concatenate((actions, self.expert_actions[expert_inds].reshape(num_expert_samples, -1)), axis=0)
+          dones = self.dones[batch_inds]
+          dones = np.concatenate((dones, self.expert_dones[expert_inds].reshape(num_expert_samples, -1)), axis=0)
+          # AdRIL Rewards (indicator kernel)
+          mask1 = (self.rewards[batch_inds] >= 0).astype(np.float32)
+          mask2 = (self.rewards[batch_inds] < self.iter).astype(np.float32)
+          r1 = - (1. ** (-self.rewards[batch_inds])) * mask1 * mask2 # Past iter
+          r2 = np.zeros_like(self.rewards[batch_inds]) * mask1 * (1 - mask2) # current iter
+          r3 = -self.rewards[batch_inds] * (1 - mask1) # Expert
+          if self.iter > 0:
+              rewards = (r1 / self.N_learner) + r2 + r3
+          else:
+              rewards = r1 + r2 + r3
+          rewards = np.concatenate((rewards, np.ones_like(rewards) / self.N_expert), axis=0)
+        else:
+          if self.optimize_memory_usage:
+              next_obs = self._normalize_obs(self.observations[(batch_inds + 1) % self.buffer_size, 0, :], env)
+          else:
+              next_obs = self._normalize_obs(self.next_observations[batch_inds, 0, :], env)
+          obs = self._normalize_obs(self.observations[batch_inds, 0, :], env)
+          actions = self.actions[batch_inds, 0, :]
+          dones = self.dones[batch_inds]
+          # AdRIL Rewards (indicator kernel)
+          mask1 = (self.rewards[batch_inds] >= 0).astype(np.float32)
+          mask2 = (self.rewards[batch_inds] < self.iter).astype(np.float32)
+          r1 = - (1. ** (-self.rewards[batch_inds])) * mask1 * mask2 # Past iter
+          r2 = np.zeros_like(self.rewards[batch_inds]) * mask1 * (1 - mask2) # current iter
+          r3 = -self.rewards[batch_inds] * (1 - mask1) / self.N_expert # Expert
+          if self.iter > 0:
+              rewards = (r1 * 1. / self.N_learner) + r2 + r3
+          else:
+              rewards = r1 + r2 + r3
+        data = (obs, actions, next_obs, dones, rewards)
+        return ReplayBufferSamples(*tuple(map(self.to_torch, data)))

scratch/etienne/pillbox/learners/advil.py

@@ -0,0 +1,222 @@
+import numpy as np
+
+import torch
+import torch.autograd as autograd
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from gym.spaces import Discrete
+import gym
+from stable_baselines3.common.preprocessing import get_action_dim
+from tqdm import tqdm
+from torch.autograd import Variable
+from itertools import repeat
+from torch.autograd import grad as torch_grad
+from typing import List, Type
+import types
+
+# Infinite dataloader
+def repeater(data_loader):
+    for loader in repeat(data_loader):
+        for data in loader:
+            yield data
+
+def create_mlp(
+    input_dim: int, output_dim: int, net_arch: List[int], activation_fn: Type[nn.Module] = nn.ReLU) -> List[nn.Module]:
+
+    if len(net_arch) > 0:
+        modules = [nn.Linear(input_dim, net_arch[0]), activation_fn()]
+    else:
+        modules = []
+
+    for idx in range(len(net_arch) - 1):
+        modules.append(nn.Linear(net_arch[idx], net_arch[idx + 1]))
+        modules.append(activation_fn())
+
+    if output_dim > 0:
+        last_layer_dim = net_arch[-1] if len(net_arch) > 0 else input_dim
+        modules.append(nn.Linear(last_layer_dim, output_dim))
+    return modules
+
+def init_ortho(layer):
+    if type(layer) == nn.Linear:
+        nn.init.orthogonal_(layer.weight)
+
+        
+class AdVILPolicy(nn.Module):
+    def __init__(self, env, mean=None, std=None):
+        super(AdVILPolicy, self).__init__()
+        if isinstance(env.action_space, Discrete):
+            self.net_arch = [64, 64]
+            self.action_dim = env.action_space.n
+            self.discrete = True
+        else:
+            self.net_arch = [256, 256]
+            self.action_dim = int(np.prod(env.action_space.shape))
+            self.low = torch.as_tensor(env.action_space.low)
+            self.high = torch.as_tensor(env.action_space.high)
+            self.discrete = False
+        self.obs_dim = int(np.prod(env.observation_space.shape))
+        self.observation_space = env.observation_space
+        net = create_mlp(self.obs_dim, self.action_dim, self.net_arch, nn.ReLU)
+        if self.discrete:
+            net.append(nn.Softmax(dim=1))
+        self.net = nn.Sequential(*net)
+        self.net.apply(init_ortho)
+        if mean is not None and std is not None:
+            self.mean = mean
+            self.std = std
+            self.is_normalized = True
+        else:
+            self.is_normalized = False
+    def forward(self, obs):
+        action = self.net(obs)
+        return action
+    def predict(self, obs, state, mask, deterministic):
+        obs = obs.reshape((-1,) + (self.obs_dim,))
+        if self.is_normalized:
+            obs = (obs - self.mean) / self.std
+        obs = torch.as_tensor(obs)
+        with torch.no_grad():
+            actions = self.forward(obs)
+            if self.discrete:
+                actions = actions.argmax(dim=1).reshape(-1)
+            else:
+                actions = self.low + ((actions + 1.0) / 2.0) * (self.high - self.low)
+                actions = torch.max(torch.min(actions, self.high), self.low)
+            actions = actions.cpu().numpy()
+        return actions, state
+
+
+class AdVILDiscriminator(nn.Module):
+    def __init__(self, env):
+        super(AdVILDiscriminator, self).__init__()
+        if isinstance(env.action_space, Discrete):
+            self.net_arch = [64, 64]
+            self.action_dim = env.action_space.n
+        else:
+            self.net_arch = [256, 256]
+            self.action_dim = int(np.prod(env.action_space.shape))
+        self.obs_dim = int(np.prod(env.observation_space.shape))
+        net = create_mlp(self.obs_dim + self.action_dim, 1, self.net_arch, nn.ReLU)
+        self.net = nn.Sequential(*net)
+        self.net.apply(init_ortho)
+
+    def forward(self, inputs):
+        output = self.net(inputs)
+        return output.view(-1)
+
+def pi_update(obs, acts, pi, f, pi_opt, prog):
+    pi_opt.zero_grad()
+    obs_v = Variable(obs)
+    pi_acts = pi(obs_v)
+    #learner_sa = torch.cat((obs, pi_acts), axis=1)
+    f_learner = f(obs, acts)
+    pi_loss = f_learner.mean() + orthogonal_reg(pi) + 2e-1 * (pi_acts - acts).square().mean()
+    pi_loss.backward()
+    if prog > 0.1:
+        torch.nn.utils.clip_grad_norm(pi.parameters(), 40.0)
+    pi_opt.step()
+    return pi_loss.item(), (2e-1 * (pi_acts - acts).square().mean()).item()
+
+def orthogonal_reg(pi):
+    with torch.enable_grad():
+        reg = 1e-4
+        orth_loss = torch.zeros(1)
+        for name, param in pi.named_parameters():
+            if 'bias' not in name:
+                x = torch.mm(torch.t(param), param)
+                x = x * (1. - torch.eye(param.shape[-1]))
+                orth_loss = orth_loss + reg * (x.square().sum())
+        return orth_loss
+
+def f_update(obs, acts, pi, f, f_opt, prog):
+    obs_v = Variable(obs)
+    pi_acts = pi(obs_v)
+    #learner_sa = torch.cat((obs, pi_acts), axis=1)
+    #expert_sa = Variable(torch.cat((obs, acts), axis=1))
+    f_learner = f(obs, pi_acts)
+    f_expert = f(obs, acts)
+    #gp = gradient_penalty((obs, pi_acts), (obs, acts), f)
+    f_opt.zero_grad()
+    f_loss = f_expert.mean() - f_learner.mean()# + 10 * gp
+    f_loss.backward()
+    if prog > 0.1:
+        torch.nn.utils.clip_grad_norm(f.parameters(), 40.0)
+    f_opt.step()
+    return f_loss.item()
+
+def gradient_penalty(learner_sa, expert_sa, f):
+    batch_size = expert_sa[0].size()[0]
+
+    #alpha = torch.rand(batch_size, 1)
+    #alpha = alpha.expand_as(expert_sa)
+
+    salpha = torch.rand(batch_size, 1, 1)
+    salpha = salpha.expand_as(expert_sa[0])
+
+    aalpha = torch.rand(batch_size, 1)
+    aalpha = aalpha.expand_as(expert_sa[1])
+
+    #interpolated = alpha * expert_sa.data + (1 - alpha) * learner_sa.data
+    #interpolated = Variable(interpolated, requires_grad=True)
+    #f_interpolated = f(interpolated.float())
+
+    sinterpolated = salpha * expert_sa[0].data + (1 - salpha) * learner_sa[0].data
+    sinterpolated = Variable(sinterpolated, requires_grad=True)
+
+    ainterpolated = aalpha * expert_sa[1].data + (1 - aalpha) * learner_sa[1].data
+    ainterpolated = Variable(ainterpolated, requires_grad=True)
+    
+    f_interpolated = f(sinterpolated, ainterpolated)
+
+    #gradients = torch_grad(outputs=f_interpolated, inputs=interpolated,
+    #                       grad_outputs=torch.ones(f_interpolated.size()),
+    #                       create_graph=True, retain_graph=True)[0]
+    
+    sgradients = torch_grad(outputs=f_interpolated, inputs=sinterpolated,
+                           grad_outputs=torch.ones(f_interpolated.size()),
+                           create_graph=True, retain_graph=True)[0]
+
+    agradients = torch_grad(outputs=f_interpolated, inputs=ainterpolated,
+                           grad_outputs=torch.ones(f_interpolated.size()),
+                           create_graph=True, retain_graph=True)[0]
+
+    #gradients = gradients.view(batch_size, -1)
+    sgradients = sgradients.view(batch_size, -1)
+    agradients = agradients.view(batch_size, -1)
+    #norm = gradients.norm(2, dim=1).mean().item()
+    #gradients_norm = torch.sqrt(torch.sum(gradients ** 2, dim=1) + 1e-12)
+    gradients_norm = torch.sqrt(torch.sum(sgradients ** 2, dim=1) + torch.sum(agradients ** 2, dim=1) + 1e-12)
+    # 2 * |f'(x_0)|
+    return ((gradients_norm - 0.4) ** 2).mean()
+
+def advil_training(data_loader, env, iters=int(1e5), policy_class=AdVILPolicy, discriminator_class=AdVILDiscriminator, lr_pi=8e-6, lr_f=8e-4):
+    if not isinstance(env.action_space, Discrete):
+        low = torch.as_tensor(env.action_space.low)
+        high = torch.as_tensor(env.action_space.high)
+    if data_loader.dataset.is_normalized:
+        pi = policy_class(env, data_loader.dataset.mean, data_loader.dataset.std)
+    else:
+        pi = policy_class(env)
+    f = discriminator_class(env)
+    pi_opt = optim.Adam(pi.parameters(), lr=lr_pi)
+
+    last_loss = 0
+    f_opt = optim.Adam(f.parameters(), lr=lr_f)
+    data_loader = repeater(data_loader)
+    for t in tqdm(range(iters)):
+        data = next(data_loader)
+        obs = data['obs']
+        acts = data['acts']
+        #if isinstance(env.action_space, Discrete):
+        #    acts = nn.functional.one_hot(acts, env.action_space.n)
+        #else:
+        #    acts = (((acts - low) / (high - low)) * 2.0) - 1.0
+        pi_loss, mse_reg = pi_update(obs, acts, pi, f, pi_opt, t/iters)
+        f_loss = f_update(obs, acts, pi, f, f_opt, t/iters)
+        if t % 100 == 0:
+            print("pi loss:", pi_loss)
+            print("mse reg:", mse_reg)
+            print("f loss:", f_loss)
+    return pi

scratch/etienne/pillbox/learners/intersim_advil.py

@@ -0,0 +1,155 @@
+import torch
+import torch.nn as nn
+
+def unnormalize(val, mean, std):
+    val *= std or 1
+    val += mean or 0
+    return val
+
+def normalize(val, mean, std):
+    val -= mean or 0
+    val /= std or 1
+    return val
+
+class IntersimPolicy(nn.Module):
+    def __init__(self, env, mean=None, std=None):
+        # assert "intersim" in env.unwrapped.spec.id
+        super().__init__()
+
+        self._ego_encoder = nn.Sequential(
+            # in 5, out 5
+            nn.Linear(5, 10),
+            nn.ReLU(),
+            nn.Linear(10, 10),
+            nn.ReLU(),
+            nn.Linear(10, 10),
+            nn.ReLU(),
+            nn.Linear(10, 5),
+            nn.ReLU(),
+        )
+        self._state_encoder = nn.Sequential(
+            # in 5, out 5
+            nn.Linear(5, 10),
+            nn.ReLU(),
+            nn.Linear(10, 10),
+            nn.ReLU(),
+            nn.Linear(10, 10),
+            nn.ReLU(),
+            nn.Linear(10, 5),
+            nn.ReLU(),
+        )
+        self._deepset = lambda e: e.sum(-2)
+        self._action_decoder = nn.Sequential(
+            # in 5 + 5, out 1
+            nn.Linear(5 + 5, 10),
+            nn.ReLU(),
+            nn.Linear(10, 10),
+            nn.ReLU(),
+            nn.Linear(10, 10),
+            nn.ReLU(),
+            nn.Linear(10, 10),
+            nn.ReLU(),
+            nn.Linear(10, 1),
+        )
+    
+    def forward(self, obs):
+        # obs.shape = (batch=514, 1 + others=150, 5)
+        # act.shape = (batch=514, 1)
+        
+        ego = obs[:, 0]#.detach().clone()
+        rel = obs[:, 1:]#.detach().clone()
+        nan = rel.isnan().any(-1, keepdim=True)
+        rel = torch.where(nan, torch.zeros_like(rel), rel) # required because of https://github.com/pytorch/pytorch/issues/15506
+        
+        d = (rel[:, :, :2] ** 2).sum(-1).sqrt()
+        front = torch.stack((torch.cos(ego[:, 3]), torch.sin(ego[:, 3])), -1)
+        left = torch.stack((-torch.sin(ego[:, 3]), torch.cos(ego[:, 3])), -1)
+        df = (rel[:, :, :2] * front.unsqueeze(1)).sum(-1)
+        dl = (rel[:, :, :2] * left.unsqueeze(1)).sum(-1)
+        alpha = torch.atan2(dl, df)
+        
+        rel[:, :, 0] = d
+        rel[:, :, 1] = alpha
+
+        e = self._ego_encoder(ego)
+        x = self._state_encoder(rel)
+        x = torch.where(nan, torch.zeros_like(x), x)
+        x = self._deepset(x)
+        a = self._action_decoder(torch.cat((e, x), 1))
+        
+        return 10 * a
+
+    def predict(self, state, mask, deterministic):
+        #action_distribution = self.forward(obs)
+        #action = action_distribution.argmax()
+        #return action
+        return self.forward(obs)
+
+class IntersimDiscriminator(nn.Module):
+    def __init__(self, env):
+        # assert "intersim" in env.unwrapped.spec.id
+        super().__init__()
+        
+        self._ego_encoder = nn.Sequential(
+            # in 5, out 5
+            nn.Linear(5, 10),
+            nn.ReLU(),
+            nn.Linear(10, 10),
+            nn.ReLU(),
+            nn.Linear(10, 10),
+            nn.ReLU(),
+            nn.Linear(10, 5),
+            nn.ReLU(),
+        )
+        self._state_encoder = nn.Sequential(
+            # in 5, out 5
+            nn.Linear(5, 10),
+            nn.ReLU(),
+            nn.Linear(10, 10),
+            nn.ReLU(),
+            nn.Linear(10, 10),
+            nn.ReLU(),
+            nn.Linear(10, 5),
+            nn.ReLU(),
+        )
+        self._deepset = lambda e: e.sum(-2)
+        self._discriminator = nn.Sequential(
+            # in 5 + 5 + 1, out 1
+            nn.Linear(5 + 5 + 1, 10),
+            nn.ReLU(),
+            nn.Linear(10, 10),
+            nn.ReLU(),
+            nn.Linear(10, 10),
+            nn.ReLU(),
+            nn.Linear(10, 10),
+            nn.ReLU(),
+            nn.Linear(10, 1),
+        )
+    
+    def forward(self, obs, acts):
+        # obs.shape = (batch=514, 1 + others=150, 5)
+        # acts.shape = (batch=514, 1)
+        # val.shape = (batch=514,)
+
+        ego = obs[:, 0]
+        rel = obs[:, 1:]
+        nan = rel.isnan().any(-1, keepdim=True)
+        rel = torch.where(nan, torch.zeros_like(rel), rel) # required because of https://github.com/pytorch/pytorch/issues/15506
+        
+        d = (rel[:, :, :2] ** 2).sum(-1).sqrt()
+        front = torch.stack((torch.cos(ego[:, 3]), torch.sin(ego[:, 3])), -1)
+        left = torch.stack((-torch.sin(ego[:, 3]), torch.cos(ego[:, 3])), -1)
+        df = (rel[:, :, :2] * front.unsqueeze(1)).sum(-1)
+        dl = (rel[:, :, :2] * left.unsqueeze(1)).sum(-1)
+        alpha = torch.atan2(dl, df)
+        
+        rel[:, :, 0] = d
+        rel[:, :, 1] = alpha
+        
+        e = self._ego_encoder(ego)
+        x = self._state_encoder(rel)
+        x = torch.where(nan, torch.zeros_like(x), x)
+        x = self._deepset(x)
+        v = self._discriminator(torch.cat((e, x, acts), 1))
+        
+        return v.squeeze(1)

scratch/etienne/pillbox/learners/soft_q.py

@@ -0,0 +1,31 @@
+from typing import Any, Dict, List, Optional, Type
+
+import gym
+import torch as th
+from torch import nn
+
+from stable_baselines3.common.policies import BasePolicy, register_policy
+from stable_baselines3.common.torch_layers import BaseFeaturesExtractor, FlattenExtractor, NatureCNN, create_mlp
+from stable_baselines3.dqn.policies import DQNPolicy, QNetwork
+
+
+class SoftQNetwork(QNetwork):
+    def _predict(self, observation: th.Tensor, deterministic: bool = True) -> th.Tensor:
+        q_values = self.forward(observation)
+        probs = nn.functional.softmax(q_values * 10, dim=1)
+        m = th.distributions.Categorical(probs)
+        action = m.sample().reshape(-1) 
+        return action
+
+
+class SQLPolicy(DQNPolicy):
+    def make_q_net(self) -> SoftQNetwork:
+        # Make sure we always have separate networks for features extractors etc
+        net_args = self._update_features_extractor(
+            self.net_args, features_extractor=None)
+        return SoftQNetwork(**net_args).to(self.device)
+
+
+SoftMlpPolicy = SQLPolicy
+
+register_policy("SoftMlpPolicy", SoftMlpPolicy)

scratch/etienne/pillbox/learners/sqil.py

@@ -0,0 +1,61 @@
+import warnings
+from abc import ABC, abstractmethod
+from typing import Dict, Generator, Optional, Union
+
+import numpy as np
+import torch as th
+from gym import spaces
+
+try:
+    # Check memory used by replay buffer when possible
+    import psutil
+except ImportError:
+    psutil = None
+
+from stable_baselines3.common.preprocessing import get_action_dim, get_obs_shape
+from stable_baselines3.common.type_aliases import ReplayBufferSamples, RolloutBufferSamples
+from stable_baselines3.common.vec_env import VecNormalize
+from stable_baselines3.common.buffers import ReplayBuffer
+
+
+class SQILReplayBuffer(ReplayBuffer):
+    def __init__(
+        self,
+        buffer_size: int,
+        observation_space: spaces.Space,
+        action_space: spaces.Space,
+        device: Union[th.device, str] = "cpu",
+        n_envs: int = 1,
+        optimize_memory_usage: bool = False,
+        expert_data: dict = dict(),
+    ):
+        super(SQILReplayBuffer, self).__init__(buffer_size, observation_space, action_space, device, n_envs=n_envs, optimize_memory_usage=optimize_memory_usage)
+
+        self.expert_states = expert_data['obs']
+        self.expert_actions = expert_data['acts']
+        self.expert_next_states = expert_data['next_obs']
+        self.expert_dones = expert_data['dones']
+
+    def _get_samples(self, batch_inds: np.ndarray, env: Optional[VecNormalize] = None) -> ReplayBufferSamples:
+        num_samples = len(batch_inds)
+        num_expert_samples = int(num_samples / 2)
+        batch_inds = batch_inds[:num_expert_samples]
+        expert_inds = np.random.randint(0, len(self.expert_states), size=num_expert_samples)
+        # Balanced sampling
+        if self.optimize_memory_usage:
+            next_obs = self._normalize_obs(self.observations[(batch_inds + 1) % self.buffer_size, 0, :], env)
+        else:
+            next_obs = self._normalize_obs(self.next_observations[batch_inds, 0, :], env)
+        next_obs = np.concatenate((next_obs, self._normalize_obs(self.expert_next_states[expert_inds], env)), axis=0)
+        obs = self._normalize_obs(self.observations[batch_inds, 0, :], env)
+        obs = np.concatenate((obs, self._normalize_obs(self.expert_states[expert_inds], env)), axis=0)
+        actions = self.actions[batch_inds, 0, :]
+        actions = np.concatenate((actions, self.expert_actions[expert_inds].reshape(num_expert_samples, -1)), axis=0)        
+        dones = self.dones[batch_inds]
+        dones = np.concatenate((dones, self.expert_dones[expert_inds].reshape(num_expert_samples, -1)), axis=0)
+        # SQIL Rewards
+        rewards = self.rewards[batch_inds] * 0.
+        rewards = np.concatenate((rewards, np.ones_like(rewards)), axis=0)
+
+        data = (obs, actions, next_obs, dones, rewards)
+        return ReplayBufferSamples(*tuple(map(self.to_torch, data)))

scratch/etienne/pillbox/learners/train.py

@@ -0,0 +1,248 @@
+from imitation.algorithms import adversarial, bc
+from imitation.util import logger, util
+from stable_baselines3 import PPO, DQN, SAC
+from soft_q import SQLPolicy
+from sqil import SQILReplayBuffer
+from stable_baselines3.common import policies
+from stable_baselines3.common.evaluation import evaluate_policy
+from imitation.rewards import discrim_nets
+import numpy as np
+import argparse
+from utils import make_sa_dataloader, make_sads_dataloader, make_sa_dataset, linear_schedule
+from stable_baselines3.common.vec_env import DummyVecEnv, VecNormalize
+from adril import AdRILWrapper, AdRILReplayBuffer
+import os
+from gym.spaces import Discrete
+import gym
+from advil import advil_training
+from stable_baselines3.common.running_mean_std import RunningMeanStd
+
+from advil import AdVILPolicy, AdVILDiscriminator
+
+def train_bc(env, n=0):
+    venv = util.make_vec_env(env, n_envs=8)
+    if isinstance(venv.action_space, Discrete):
+        w = 64
+    else:
+        w = 256
+    for i in range(n):
+        mean_rewards = []
+        std_rewards = []
+        for num_trajs in range(0, 26, 5):
+            if num_trajs == 0:
+                expert_data = make_sa_dataloader(env, normalize=False)
+            else:
+                expert_data = make_sa_dataloader(env, max_trajs=num_trajs, normalize=False)
+            bc_trainer = bc.BC(venv.observation_space, venv.action_space, expert_data=expert_data,
+                               policy_class=policies.ActorCriticPolicy,
+                               ent_weight=0., l2_weight=0., policy_kwargs=dict(net_arch=[w, w]))
+            if num_trajs > 0:
+                bc_trainer.train(n_batches=int(5e5))
+
+            def get_policy(*args, **kwargs):
+                return bc_trainer.policy
+            model = PPO(get_policy, env, verbose=1)
+            model.save(os.path.join("learners", env,
+                                    "bc_{0}_{1}".format(i, num_trajs)))
+            mean_reward, std_reward = evaluate_policy(
+                    model, model.get_env(), n_eval_episodes=10)               
+            mean_rewards.append(mean_reward)
+            std_rewards.append(std_reward)
+            print("{0} Trajs: {1}".format(num_trajs, mean_reward))
+            np.savez(os.path.join("learners", env, "bc_rewards_{0}".format(
+                i)), means=mean_rewards, stds=std_rewards)
+
+
+def train_gail(env, n=0):
+    venv = util.make_vec_env(env, n_envs=8)
+    if isinstance(venv.action_space, Discrete):
+        w = 64
+    else:
+        w = 256
+    expert_data = make_sads_dataloader(env, max_trajs=5)
+    logger.configure(os.path.join("learners", "GAIL"))
+
+    for i in range(n):
+        discrim_net = discrim_nets.ActObsMLP(
+                action_space=venv.action_space,
+                observation_space=venv.observation_space,
+                hid_sizes=(w, w),
+                )
+        gail_trainer = adversarial.GAIL(venv, expert_data=expert_data, expert_batch_size=32,
+                                        gen_algo=PPO("MlpPolicy", venv, verbose=1, n_steps=1024,
+                                                     policy_kwargs=dict(net_arch=[w, w])),
+                                                     discrim_kwargs={'discrim_net': discrim_net})
+        mean_rewards = []
+        std_rewards = []
+        for train_steps in range(20):
+            if train_steps > 0:
+                if 'Bullet' in env:
+                    gail_trainer.train(total_timesteps=25000)
+                else:
+                    gail_trainer.train(total_timesteps=16384)
+
+            def get_policy(*args, **kwargs):
+                return gail_trainer.gen_algo.policy
+            model = PPO(get_policy, env, verbose=1)
+            mean_reward, std_reward = evaluate_policy(
+                model, model.env, n_eval_episodes=10)
+            mean_rewards.append(mean_reward)
+            std_rewards.append(std_reward)
+            print("{0} Steps: {1}".format(train_steps, mean_reward))
+            np.savez(os.path.join("learners", env, "gail_rewards_{0}".format(i)),
+                     means=mean_rewards, stds=std_rewards)
+
+
+def train_sqil(env, n=0):
+    venv = gym.make(env)
+    expert_data = make_sa_dataset(env, max_trajs=5)
+
+    for i in range(n):
+        if isinstance(venv.action_space, Discrete):
+            model = DQN(SQLPolicy, venv, verbose=1, policy_kwargs=dict(net_arch=[64, 64]), learning_starts=1)
+        else:
+            model = SAC('MlpPolicy', venv, verbose=1, policy_kwargs=dict(net_arch=[256, 256]), ent_coef='auto',
+                        learning_rate=linear_schedule(7.3e-4), train_freq=64, gradient_steps=64, gamma=0.98, tau=0.02)
+
+        model.replay_buffer = SQILReplayBuffer(model.buffer_size, model.observation_space,
+                                               model.action_space, model.device, 1,
+                                               model.optimize_memory_usage, expert_data=expert_data)
+        mean_rewards = []
+        std_rewards = []
+        for train_steps in range(20):
+            if train_steps > 0:
+                if 'Bullet' in env:
+                    model.learn(total_timesteps=25000, log_interval=1)
+                else:
+                    model.learn(total_timesteps=16384, log_interval=1)
+            mean_reward, std_reward = evaluate_policy(
+                model, model.env, n_eval_episodes=10)
+            mean_rewards.append(mean_reward)
+            std_rewards.append(std_reward)
+            print("{0} Steps: {1}".format(train_steps, mean_reward))
+            np.savez(os.path.join("learners", env, "sqil_rewards_{0}".format(i)),
+                     means=mean_rewards, stds=std_rewards)
+
+
+def train_adril(env, n=0, balanced=False):
+    num_trajs = 20
+    expert_data = make_sa_dataset(env, max_trajs=num_trajs)
+    n_expert = len(expert_data["obs"])
+    expert_sa = np.concatenate((expert_data["obs"], np.reshape(expert_data["acts"], (n_expert, -1))), axis=1)
+
+    for i in range(0, n):
+        venv = AdRILWrapper(gym.make(env))
+        mean_rewards = []
+        std_rewards = []
+        # Create model
+        if isinstance(venv.action_space, Discrete):
+            model = DQN(SQLPolicy, venv, verbose=1, policy_kwargs=dict(net_arch=[64, 64]), learning_starts=1)
+        else:
+            model = SAC('MlpPolicy', venv, verbose=1, policy_kwargs=dict(net_arch=[256, 256]), ent_coef='auto',
+                        learning_rate=linear_schedule(7.3e-4), train_freq=64, gradient_steps=64, gamma=0.98, tau=0.02)
+        model.replay_buffer = AdRILReplayBuffer(model.buffer_size, model.observation_space,
+                                               model.action_space, model.device, 1,
+                                               model.optimize_memory_usage, expert_data=expert_data, N_expert=num_trajs,
+                                               balanced=balanced)
+        if not balanced:
+            for j in range(len(expert_sa)):
+                obs = expert_data["obs"][j]
+                act = expert_data["acts"][j]
+                next_obs = expert_data["next_obs"][j]
+                done = expert_data["dones"][j]
+                model.replay_buffer.add(obs, next_obs, act, -1, done)    
+        for train_steps in range(400):
+            # Train policy
+            if train_steps > 0:
+                if 'Bullet' in env:
+                    model.learn(total_timesteps=1250, log_interval=1000)
+                else:
+                    model.learn(total_timesteps=25000, log_interval=1000)
+                if train_steps % 1 == 0: # written to support more complex update schemes
+                    model.replay_buffer.set_iter(train_steps)
+                    model.replay_buffer.set_n_learner(venv.num_trajs)
+            
+            # Evaluate policy
+            if train_steps % 20 == 0:
+                model.set_env(gym.make(env))
+                mean_reward, std_reward = evaluate_policy(
+                    model, model.env, n_eval_episodes=10)
+                mean_rewards.append(mean_reward)
+                std_rewards.append(std_reward)
+                print("{0} Steps: {1}".format(int(train_steps * 1250), mean_reward))
+                np.savez(os.path.join("learners", env, "adril_rewards_{0}".format(i)),
+                        means=mean_rewards, stds=std_rewards)
+            # Update env
+            if train_steps > 0:
+                if train_steps % 1  == 0:
+                    venv.set_iter(train_steps + 1)
+            model.set_env(venv)
+
+
+def train_advil(env, policy_class=AdVILPolicy, discriminator_class=AdVILDiscriminator,
+                iters=int(1e5), lr_pi=8e-6, lr_f=8e-4):
+    venv = gym.make(env)
+    expert_data = make_sa_dataloader(
+        env,
+        normalize=False,
+        batch_size=1024,
+    )
+    pi = advil_training(
+        expert_data,
+        venv,
+        iters=iters,
+        policy_class=policy_class,
+        discriminator_class=discriminator_class, 
+        lr_pi=lr_pi,
+        lr_f=lr_f,
+    )
+    return pi
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Train expert policies.')
+    parser.add_argument(
+        '-a', '--algo', choices=['bc', 'gail', 'sqil', 'adril', 'advil', 'all'], required=True)
+    parser.add_argument('-e', '--env', choices=['cartpole', 'lunarlander', 'acrobot', 'pendulum', 'halfcheetah', 'walker', 'hopper', 'ant'],
+                        required=True)
+    parser.add_argument('-n', '--num_runs', required=False)
+    args = parser.parse_args()
+    if args.env == "cartpole":
+        envname = 'CartPole-v1'
+    elif args.env == "lunarlander":
+        envname = 'LunarLander-v2'
+    elif args.env == "acrobot":
+        envname = 'Acrobot-v1'
+    elif args.env == "pendulum":
+        envname = 'Pendulum-v0'
+    elif args.env == "halfcheetah":
+        envname = 'HalfCheetahBulletEnv-v0'
+    elif args.env == "walker":
+        envname = 'Walker2DBulletEnv-v0'
+    elif args.env == "hopper":
+        envname = 'HopperBulletEnv-v0'
+    elif args.env == "ant":
+        envname = 'AntBulletEnv-v0'
+    else:
+        print("ERROR: unsupported env.")
+    if args.num_runs is not None and args.num_runs.isdigit():
+        num_runs = int(args.num_runs)
+    else:
+        num_runs = 1
+    if args.algo == 'bc':
+        train_bc(envname, num_runs)
+    elif args.algo == 'gail':
+        train_gail(envname, num_runs)
+    elif args.algo == 'sqil':
+        train_sqil(envname, num_runs)
+    elif args.algo == 'adril':
+        train_adril(envname, num_runs)
+    elif args.algo == 'advil':
+        train_advil(envname, num_runs)
+    elif args.algo == 'all':
+        train_bc(envname, num_runs)
+        train_gail(envname, num_runs)
+        train_sqil(envname, num_runs)
+        train_adril(envname, num_runs)
+        train_advil(envname, num_runs)
+    else:
+        print("ERROR: unsupported algorithm")

scratch/etienne/pillbox/learners/utils.py

@@ -0,0 +1,129 @@
+import numpy as np
+import torch
+from torch.utils.data import Dataset, DataLoader
+from itertools import chain
+from typing import Callable, Union, Type, Optional, Dict, Any
+
+# From https://github.com/DLR-RM/rl-baselines3-zoo/blob/8ea4f4a87afa548832ca17e575b351ec5928c1b0/utils/utils.py
+def linear_schedule(initial_value: Union[float, str]) -> Callable[[float], float]:
+    """
+    Linear learning rate schedule.
+    :param initial_value: (float or str)
+    :return: (function)
+    """
+    if isinstance(initial_value, str):
+        initial_value = float(initial_value)
+
+    def func(progress_remaining: float) -> float:
+        """
+        Progress will decrease from 1 (beginning) to 0
+        :param progress_remaining: (float)
+        :return: (float)
+        """
+        return progress_remaining * initial_value
+
+    return func
+
+class SADataset(torch.utils.data.Dataset):
+    def __init__(self, obs, acts, normalize):
+        if normalize:
+            obs = np.array(obs)
+            self.mean = obs.mean(axis=0)
+            self.std = obs.std(axis=0) + 1e-3
+            obs = (obs - self.mean) / (self.std)
+            self.is_normalized = True
+        else:
+            self.is_normalized = False
+        self.obs = torch.tensor(obs)
+        self.acts = torch.tensor(acts)
+
+    def __len__(self):
+        return len(self.obs)
+
+    def __getitem__(self, idx):
+        if torch.is_tensor(idx):
+            idx = idx.tolist()
+        obs = self.obs[idx]
+        acts = self.acts[idx]
+        sample = {'obs': obs, 'acts': acts}
+        return sample
+
+def make_sa_dataloader(envname, max_trajs=None, normalize=False, batch_size=32):
+    demos = np.load(
+        "../experts/{0}/demos.npz".format(envname), allow_pickle=True)
+    num_trajs = demos["num_trajs"]
+    if max_trajs is None:
+        max_trajs = num_trajs
+    obs = []
+    acts = []
+    for traj in range(min(max_trajs, num_trajs)):
+        obs.extend(demos[str(traj)].item()['states'])
+        acts.extend(demos[str(traj)].item()['actions'])
+    dataset = SADataset(obs, acts, normalize)
+    dataloader = DataLoader(dataset, batch_size=batch_size,
+                             shuffle=True, num_workers=0)
+    return dataloader
+
+class SADSDataset(torch.utils.data.Dataset):
+    def __init__(self, obs, acts, next_obs, traj_lens):
+        self.obs = torch.tensor(obs)
+        self.acts = torch.tensor(acts)
+        self.next_obs = torch.tensor(next_obs)
+        dones = [[False for _ in range(l - 2)] + [True] for l in traj_lens]
+        self.dones = torch.tensor(list(chain.from_iterable(dones)))
+
+    def __len__(self):
+        return len(self.obs)
+
+    def __getitem__(self, idx):
+        if torch.is_tensor(idx):
+            idx = idx.tolist()
+        obs = self.obs[idx]
+        acts = self.acts[idx]
+        next_obs = self.next_obs[idx]
+        dones = self.dones[idx]
+        sample = {'obs': obs, 'acts': acts,
+                  'next_obs': next_obs, 'dones': dones}
+        return sample
+
+def make_sads_dataloader(envname, max_trajs=None):
+    demos = np.load(
+        "./experts/{0}/demos.npz".format(envname), allow_pickle=True)
+    num_trajs = demos["num_trajs"]
+    if max_trajs is None:
+        max_trajs = num_trajs
+    obs = []
+    next_obs = []
+    acts = []
+    lens = []
+    for traj in range(min(max_trajs, num_trajs)):
+        obs.extend(demos[str(traj)].item()['states'][:-1])
+        next_obs.extend(demos[str(traj)].item()['states'][1:])
+        acts.extend(demos[str(traj)].item()['actions'][:-1])
+        lens.append(len(demos[str(traj)].item()['states']))
+    dataset = SADSDataset(obs, acts, next_obs, lens)
+    dataloader = DataLoader(dataset, batch_size=32,
+                            shuffle=False, num_workers=0, drop_last=True)
+    return dataloader
+
+def make_sa_dataset(envname, max_trajs=None):
+    demos = np.load("../pillbox/experts/{0}/demos.npz".format(envname), allow_pickle=True)
+    num_trajs = demos["num_trajs"]
+    if max_trajs is None:
+        max_trajs = num_trajs
+    expert_states = []
+    expert_actions = []
+    expert_next_states = []
+    expert_dones = []
+    for traj in range(min(max_trajs, num_trajs)):
+        expert_states.extend(demos[str(traj)].item()['states'][:-1])
+        expert_next_states.extend(demos[str(traj)].item()['states'][1:])
+        expert_actions.extend(demos[str(traj)].item()['actions'][:-1])
+        l = len(demos[str(traj)].item()['states'])
+        expert_dones.extend([False for _ in range(l - 2)] + [True])
+    expert_data = dict()
+    expert_data['obs'] = np.array(expert_states)
+    expert_data['acts'] = np.array(expert_actions)
+    expert_data['next_obs'] = np.array(expert_next_states)
+    expert_data['dones'] = np.array(expert_dones)
+    return expert_data

scratch/etienne/pillbox/requirements.txt

@@ -0,0 +1,9 @@
+gym
+numpy
+psutil
+scikit_learn
+scipy
+stable_baselines3
+torch
+tqdm
+imitation

scratch/johannes/normalization.py

@@ -0,0 +1,73 @@
+import torch
+from torch import nn
+
+from sklearn import preprocessing
+
+class Normalization(nn.Module):
+    def __init__(self, X):
+        super(Normalization, self).__init__()
+        self.fit(X)
+
+    def fit(self, X):
+        raise NotImplementedError('Please implement fit()')
+
+    def transform(self, X):
+        raise NotImplementedError('Please implement transform()')
+
+    def inverse_transform(self, X):
+        raise NotImplementedError('Please implement inverse_transform()')
+
+    def forward(self, X):
+        return self.transform(X)
+
+class SciKitNormalization(Normalization):
+    def __init__(self, tf, X):
+        self.tf = tf
+        super(SciKitNormalization, self).__init__(X)
+
+    def fit(self, X):
+        self.tf.fit(X)
+
+    def transform(self, X):
+        return torch.tensor(self.tf.transform(X), dtype=torch.float)
+
+    def inverse_transform(self, X):
+        return torch.tensor(self.tf.inverse_transform(X), dtype=torch.float)
+
+class SciKitStandardization(SciKitNormalization):
+    def __init__(self, X):
+        super(SciKitStandardization, self).__init__(preprocessing.StandardScaler(), X)
+
+class SciKitMinMaxScaler(SciKitNormalization):
+    def __init__(self, X):
+        super(SciKitMinMaxScaler, self).__init__(preprocessing.MinMaxScaler(), X)
+
+
+ns = 5
+na = 1
+n_batch = 1000
+
+state = torch.rand(n_batch, ns)
+action = torch.rand(n_batch, na)
+
+s_tf = SciKitStandardization(state)
+a_tf = SciKitMinMaxScaler(action)
+
+print(torch.linalg.norm(s_tf.inverse_transform(s_tf(state)) - state))
+print(torch.linalg.norm(a_tf.inverse_transform(a_tf(action)) - action))
+
+
+
+
+# class Foo:
+#     def __init__(self):
+#         return None
+#     def baz(self):
+#         print("Foo.baz()")
+
+# class Bar(Foo):
+#     def __init__(self):
+#         return None
+
+# bar = Bar()
+# bar.baz()

setup.py

@@ -1,19 +0,0 @@
-from setuptools import setup, find_packages
-
-setup(
-    name='interimit',
-    version='0.0.1',
-    url='https://github.com/sisl/InteractionImitation/',
-    author='Arec Jamgochian',
-    author_email='jamgochian95@gmail.com',
-    description='Imitation Learning on INTERACTION Dataset',
-    packages=find_packages(),
-    install_requires=[
-        'tikzplotlib',
-        'torch',
-        'pytest',
-        'json5',
-        'gym',
-        'intersim @ git+https://github.com/sisl/InteractionSimulator',
-    ],
-)

src/__init__.py

@@ -0,0 +1,3 @@
+from src.expert_data import generate_expert_data, load_expert_data
+from src.data_utils import InteractionDatasetSingleAgent
+from src.metrics import metrics

src/bc/__init__.py

@@ -0,0 +1 @@
+from src.bc.bc import BehaviorCloningPolicy, train, bc_config

src/bc/bc.py

@@ -0,0 +1,191 @@
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+import pickle
+from torch.utils.tensorboard import SummaryWriter
+
+from src.policies import IntersimStateNet, IntersimPolicy, generate_transforms
+from src.util.nn_training import optimizer_factory
+from tqdm import tqdm
+import json5
+from ray import tune
+
+def bc_config(ray_config):
+    config = {
+        'ego_encoder': {'input_dim': 5, 'hidden_n': 0, 'hidden_dim':0, 'output_dim': 0},
+        'deepsets': {
+            'input_dim': 6,
+            'phi': {
+                'hidden_n': ray_config['deepsets_phi_hidden_n'],
+                'hidden_dim': ray_config['deepsets_phi_hidden_dim']
+                },
+            'latent_dim': ray_config['deepsets_latent_dim'],
+            'rho': {
+                'hidden_n': ray_config['deepsets_rho_hidden_n'],
+                'hidden_dim': ray_config['deepsets_rho_hidden_dim']
+                },
+            'output_dim': ray_config['deepsets_output_dim']
+        },
+        'path_encoder': {'input_dim': 40, 'hidden_n': 0, 'hidden_dim': 0, 'output_dim': 0},
+        'head': {
+            'input_dim': 0, # computed in constructor
+            'hidden_n': ray_config['head_hidden_n'],
+            'hidden_dim': ray_config['head_hidden_dim'],
+            'output_dim': 1, # number of outputs e.g. number of actions, or just one
+            'final_activation': ray_config['head_final_activation'],
+        },
+        'optim': {
+            'optimizer':'adam',
+            'lr':ray_config['lr'],
+            'weight_decay':ray_config['weight_decay']
+        },
+        'train_epochs': 40,
+        'train_batch_size': ray_config['train_batch_size'],
+        'loss': ray_config['loss'],
+
+    }
+    return config
+
+class BehaviorCloningPolicy(IntersimPolicy):
+    """
+    Class for (continuous) behavior cloning policy
+    """
+    
+    def __init__(self, config: dict, transforms: dict):
+        """
+        Initialize BehaviorCloningPolicy
+        Args:
+            config (dict): configuration file to initialize IntersimDeepSetsNet with
+            transforms (dict): dictionary of transforms to apply to different fields
+        """
+        super(BehaviorCloningPolicy, self).__init__(config, transforms)
+        self._policy = IntersimStateNet(config)
+        
+    @classmethod
+    def load_model(cls, filestr: str, config: dict = None):
+        """
+        Load a model from a file prefix
+        Args:
+            config (dict): configuration dict to set up model
+            filestr (str): string prefix to load model from
+        Returns
+            model (BehaviorCloningPolicy): loaded model
+        """
+        if not config:
+            with open(filestr+'_config.json', 'r') as cfg:
+                config = json5.load(cfg)
+        transforms = pickle.load(open(filestr+'_transforms.pkl', 'rb'))
+        model = cls(config, transforms=transforms)
+        model._policy.load_state_dict(torch.load(filestr+'_model.pt'))
+        return model
+    
+    def eval(self):
+        self._policy.eval()
+
+    def parameters(self):
+        return self._policy.parameters()
+
+    def save_model(self, filestr, save_config=True, save_transforms=True):
+        """
+        Save transforms and state_dict to a location specificed by filestr
+        Args:
+            filestr (str): string prefix to save model to
+            save_config (bool): whether to save the config file (as a json)
+            save_transforms (bool): whether to save transforms (as a pickle)
+        """
+        if save_config:
+            with open(filestr+'_config.json', 'w') as cfg:
+                json5.dump(self._config, cfg)
+        if save_transforms:
+            pickle.dump(self._transforms, open(filestr+'_transforms.pkl', 'wb'))
+        torch.save(self._policy.state_dict(), filestr+'_model.pt')
+
+def train(config, policy, train_dataset, cv_dataset, filestr, **kwargs):
+    
+    using_ray = kwargs.get('ray', False)
+    if using_ray:
+        print('using ray')
+
+    # hyperparams
+    loss_type = config['loss']
+    train_epochs = config['train_epochs']
+    train_batch_size = config['train_batch_size']
+
+    cv_every = 1
+    print_epoch_every = 1000
+    print_cv_every = 5
+    checkpoint_every = 100
+    cv_batch_size = 256 # doesn't matter
+
+    # training and testing dataloaders
+    training_loader = DataLoader(train_dataset, batch_size=train_batch_size, shuffle=True)
+    cv_loader = DataLoader(cv_dataset, batch_size=cv_batch_size, shuffle=True)
+
+    # change policy dtype
+    policy.policy = policy.policy.type(train_dataset[0]['state']['ego_state'].dtype)
+
+    # generate loss function, optimizer
+    cv_loss_fn = nn.MSELoss(reduction='sum')
+    if loss_type == 'huber':
+        loss_fn = nn.HuberLoss(reduction='sum') 
+    elif loss_type == 'mse':
+        loss_fn = nn.MSELoss(reduction='sum')
+    else:
+        raise NotImplementedError
+    optimizer = optimizer_factory(config['optim'], policy.parameters())
+
+    # generate tensorboard writer
+    if not using_ray:
+        writer = SummaryWriter(filestr)
+    
+    for i in tqdm(range(train_epochs)):
+        
+        # save model checkpoints
+        if i % checkpoint_every == 0:
+            policy.save_model(filestr + '_epoch%04i'%(i) )
+
+        # train
+        epoch_loss = 0
+        for (batch_idx, batch) in enumerate(training_loader):
+            
+            # sample mini-batch and run through policy
+            pred_action = policy(batch['state']) 
+            loss = loss_fn(pred_action, batch['action'])
+
+            # compute loss and step optimizer
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+            epoch_loss += loss.item() / len(train_dataset)
+        
+        # if i % print_epoch_every == 0:
+        #     print('Epoch: {}, Training Loss: {}'.format(i, epoch_loss))
+
+                # measure cv loss
+
+        if i % cv_every == 0:
+            with torch.no_grad():
+                cv_loss = 0.
+                for (batch_idx, batch) in enumerate(cv_loader):
+                    pred_action = policy(batch['state']) 
+                    loss = cv_loss_fn(pred_action, batch['action'])
+                    cv_loss += loss.item() / len(cv_dataset)
+            
+            
+        # Write epoch loss
+        if using_ray:
+            if i % cv_every == 0:
+                tune.report(training_loss=epoch_loss, cv_loss=cv_loss, training_iteration=i+1)
+            else:
+                tune.report(training_loss=epoch_loss, training_iteration=i+1)
+        else:
+            writer.add_scalar('training loss',epoch_loss, i)
+            if i % cv_every == 0:
+                writer.add_scalar('cv loss', cv_loss, i)
+
+        # if i % print_cv_every == 0:    
+        #     print('Epoch: {}, CV Loss: {}'.format(i, cv_loss))
+
+
+    policy.save_model(filestr)

src/data_utils.py

@@ -0,0 +1,96 @@
+import torch
+from torch.utils.data import Dataset
+import numpy as np
+from src.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], dtype=torch.float32):
+        """
+        Args:
+            output_dir (string): Directory with all the images.
+            loc (int): location index
+            tracks (list[int]): track indices
+        """
+        self.output_dir = output_dir
+        self.loc = loc
+        self.tracks = tracks
+        self.dtype = dtype
+        self.keys = ['ego_state', 'relative_state', 'path', 'action', 'next_ego_state', 'next_relative_state', 'next_path']
+        self._load_dataset()
+
+    def _load_dataset(self):
+        """
+        Load the full datasets ahead of time
+        """
+        self.raw_data = {key:[] for key in self.keys}
+        max_nv = 0
+        for track in self.tracks:
+            try:
+                data = 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
+            max_nv = max(max_nv, data['relative_state'].shape[1])
+            for key in self.keys:
+                self.raw_data[key].append(data[key])
+
+        # 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, dtype=self.dtype) * np.nan
+            self.raw_data['relative_state'][i] = torch.cat((self.raw_data['relative_state'][i], pad), dim=1)
+            self.raw_data['next_relative_state'][i] = torch.cat((self.raw_data['next_relative_state'][i], pad), dim=1)
+
+        # cat lists
+        for key in self.keys:
+            self.raw_data[key] = torch.cat(self.raw_data[key]).type(self.dtype)
+
+        # mandate equal length
+        lengths = [len(self.raw_data[key]) for key in self.keys]
+        assert min(lengths) == max(lengths), 'dataset lengths unequal'
+        
+    def __len__(self):
+        return len(self.raw_data['ego_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 (dict): state dictionary with the following entries:
+                    ego_state (torch.tensor): (B, 5) raw state
+                    relative_state (torch.tensor): (B, max_nv, d) relative state (padded with nans)
+                    path (torch.tensor): (B, P, 2) tensor of P future path x and y positions
+                action (torch.tensor): (B, 1) actions taken from each state
+                next_stat (dict): next state dictionary with the following entries:
+                    ego_state (torch.tensor): (B, 5) raw next state
+                    relative_state (torch.tensor): (B, max_nv, d) next relative state (padded with nans)
+                    path (torch.tensor): (B, P, 2) tensor of P future next path x and y positions
+        """
+        #sample = {key:self.raw_data[key][idx] for key in self.keys}
+        sample = {
+            'state':{
+                'ego_state':self.raw_data['ego_state'][idx],
+                'relative_state':self.raw_data['relative_state'][idx],
+                'path':self.raw_data['path'][idx]
+            },
+            'action':self.raw_data['action'][idx],
+            'next_state':{
+                'ego_state':self.raw_data['next_ego_state'][idx],
+                'relative_state':self.raw_data['next_relative_state'][idx],
+                'path':self.raw_data['next_path'][idx]},
+        }
+        return sample

src/expert_data.py

@@ -0,0 +1,210 @@
+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
+from intersim.graphs import ConeVisibilityGraph
+import os
+opj = os.path.join
+
+def generate_expert_data(path: str='expert_data', loc: int = 0, track:int = 0, 
+                            mask_relstate: bool = False, regularize_actions: bool = False, 
+                            **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
+        mask_relstate (bool): whether to mask the relative states from the cone visibility graph
+        regularize_actions (bool): whether to regularize the action selection
+        kwargs: arguments for environment instantiation
+    """
+    
+    action_reg = 0.002 if regularize_actions else 0
+    
+    if not os.path.isdir(path):
+        os.makedirs(path)
+    filestr = opj(path,intersim.LOCATIONS[loc]+'_track%03i'%(track))
+    
+    svt, svt_path = get_svt(loc=loc, track=track) #base='InteractionSimulator'
+    osm = get_map_path(loc=loc)
+    print('SVT path: {}'.format(svt_path))
+    print('Map path: {}'.format(osm))
+    states, actions = SVT_to_stateactions(svt)
+
+    # animate from environment
+    if mask_relstate:
+        cvg = ConeVisibilityGraph(r=20, half_angle=120)
+        env = gym.make('intersim:intersim-v0', svt=svt, map_path=osm, 
+            min_acc=-np.inf, max_acc=np.inf, graph=cvg, mask_relstate=True, **kwargs)
+    else:
+        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)
+        if len(norms)>0:
+            max_devs.append(norms.max())
+
+        # propagate environment
+        ob, r, done, info = env.step(env.target_state(svt.simstate[i+1], mu=action_reg))
+        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)
+    actions = torch.stack(actions_taken)
+
+    # save observations and actions
+    pickle.dump(obs,open(filestr+'_raw_observations.pkl', 'wb'))
+    torch.save(actions, filestr+'_raw_actions.pt')
+    process_expert_observations(obs, actions, filestr)
+
+def process_expert_observations(obs, actions, filestr, remove_outliers=True, dtype=torch.float32):
+    """
+    Process the expert observations and save them as torch tensors
+    Args:
+        obs (list[dict]): lost of observations
+        actions (torch.Tensor): (T, nv, a) tensor of actions
+        filestr (str): base filename with which to save out observation tensors
+        remove_outliers (bool): whether to remove datapoints with acceleration above or below 5 m/s/s
+        dtype (torch.Type): type to convert data to
+    """
+    keys = ['ego_state', 'relative_state', 'path', 'action', 'next_ego_state', 'next_relative_state', 'next_path']
+    data = {key:[] for key in keys}
+    assert len(obs) == len(actions), 'non-matching action and observation lengths'
+    T = len(obs)
+    max_nv = 0
+    for t in range(T-1):
+        nni = ~torch.isnan(obs[t]['state'][:,0]) & ~torch.isnan(obs[t+1]['state'][:,0]) 
+        max_nv = max(max_nv,nni.count_nonzero())
+
+        # state
+        data['ego_state'].append(obs[t]['state'][nni])
+        data['relative_state'].append(obs[t]['relative_state'].index_select(0, 
+            nni.nonzero()[:,0]).index_select(1, nni.nonzero()[:,0]))
+        data['path'].append(torch.stack((obs[t]['paths'][0][nni], obs[t]['paths'][1][nni]), dim=-1))
+
+        # action
+        data['action'].append(actions[t][nni])
+
+        # next state
+        data['next_ego_state'].append(obs[t+1]['state'][nni])
+        data['next_relative_state'].append(obs[t+1]['relative_state'].index_select(0, 
+            nni.nonzero()[:,0]).index_select(1, nni.nonzero()[:,0]))
+        data['next_path'].append(torch.stack((obs[t+1]['paths'][0][nni], obs[t+1]['paths'][1][nni]), dim=-1))
+
+    
+
+    # pad second dimension of relative state
+    for i in range(len(data['relative_state'])):
+        nv1, nv2, d = data['relative_state'][i].shape
+        pad = torch.zeros(nv1, max_nv-nv2, d, dtype=dtype) * np.nan
+        data['relative_state'][i] = torch.cat((data['relative_state'][i], pad), dim=1)
+        data['next_relative_state'][i] = torch.cat((data['next_relative_state'][i], pad), dim=1)
+
+    # cat lists 
+    for key in keys:
+        data[key] = torch.cat(data[key]).type(dtype)
+
+    if remove_outliers:
+        non_outlier_indices = torch.nonzero(torch.abs(data['action'][:,0]) < 5)
+        for key in keys:
+            data[key] = data[key][non_outlier_indices[:,0]]    
+
+    # mandate equal length
+    lengths = [len(data[key]) for key in keys]
+    assert min(lengths) == max(lengths), 'dataset lengths unequal'
+
+    # save out data
+    for key in keys:
+        torch.save(data[key], filestr+'_'+key+'.pt')
+
+def load_expert_data(path='expert_data', loc: int = 0, track:int = 0):
+    """
+    Load expert data from processed files.
+    Args:
+        path (str): directory to save data
+        loc (int): location index
+        track (int): track index
+    Returns:
+        data (dict[torch.Tensor]): dict of data
+    """
+    # load observations and actions
+    filestr = opj(path, intersim.LOCATIONS[loc]+'_track%03i'%(track))
+    data = {}
+    keys = ['ego_state', 'relative_state', 'path', 'action', 'next_ego_state', 'next_relative_state', 'next_path']
+    for key in keys:
+        data[key] = torch.load(filestr+'_'+key+'.pt')
+    return data
+
+def load_expert_data_raw(path='expert_data', loc: int = 0, track:int = 0):
+    """
+    Load expert data from raw 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+'_raw_observations.pkl', 'rb'))
+    actions = torch.load(filestr+'_raw_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)')
+    parser.add_argument('--all-tracks', action='store_true',
+                       help='whether to process all tracks at location')
+    parser.add_argument('--graph', action='store_true',
+                       help='whether to mask the relative states based on a ConeVisibilityGraph')
+    parser.add_argument('--reg', action='store_true',
+                       help='whether to regularize actions in the action targeter')
+    parser.add_argument('-o', default='./expert_data', type=str,
+                       help='output folder')
+    args = parser.parse_args()
+
+    kwargs = {
+        'loc':args.loc, 
+        'track': args.track,
+        'path':args.o, 
+        'mask_relstate':args.graph, 
+        'regularize_actions': args.reg
+    }
+
+    if args.all_tracks:
+        for i in range(intersim.MAX_TRACKS):
+            kwargs['track'] = i
+            generate_expert_data(**kwargs)
+    else:
+        generate_expert_data(**kwargs)

src/main.py

@@ -0,0 +1,118 @@
+import os
+import torch
+import gym
+import intersim
+import numpy as np
+from tqdm import tqdm
+from torch.utils.tensorboard import SummaryWriter
+
+from src import InteractionDatasetSingleAgent, metrics
+from intersim.utils import get_map_path, get_svt
+from src.policies.policy import generate_transforms
+
+def basestr(**kwargs):
+    """
+    Return base prefix for all files relating to a certain experiment
+    Args:
+        kwargs (dict): keyword arguments sent to main training loop
+    Returns:
+        basestr (str): prefix
+    """
+    return 'base'
+
+def main(config, method='bc', train=False, test=False, loc=0, datadir='./expert_data', filestr='', **kwargs):
+    """
+    Main loop for training and testing different imitation models
+    Args:
+        config (dict): configuration dictionary for model
+        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
+        datadir (str): path to expert data
+        kwargs (dict): remaining kwargs for training loop
+    """
+    # get/set seed
+    seed = kwargs.get('seed',0)
+    torch.manual_seed(seed)
+
+    # method-based training
+    if method=='bc':
+        from src import bc
+        policy_class = bc.BehaviorCloningPolicy
+        train_fn = bc.train
+    elif method=='vd':
+        from src import value_dice
+        policy_class = value_dice.ValueDicePolicy
+        train_fn = value_dice.train
+    else:
+        raise NotImplementedError("Method {} not implemented".format(method))
+    
+    # default train / cv / test split datasets 
+    if train:
+
+        # make policy, train and test datasets, and send to 
+        train_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=kwargs['train_tracks'])
+        # generate transform from train_dataset
+        transforms = generate_transforms(train_dataset)
+        policy = policy_class(config, transforms)
+        cv_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=kwargs['cv_tracks'])
+        train_fn(config, policy, train_dataset, cv_dataset, filestr, **kwargs)
+    
+    if test:
+
+        # load policy 
+        policy = policy_class.load_model(filestr, config)
+        policy.eval()
+
+        # simulate policy
+        simulate_policy(policy, loc=loc, track=kwargs['test_tracks'][0], filestr=filestr, nframes=kwargs['nframes'], graph=kwargs['graph'])
+
+        # run test metrics
+        test_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=kwargs['test_tracks'])
+        writer = SummaryWriter(filestr)
+        info = metrics(filestr, test_dataset, policy)
+        for k, m in info.items():
+            writer.add_scalar('test/{}'.format(k), m, 0)
+
+
+def simulate_policy(policy, loc=0, track=0, filestr='', nframes=float('inf'), graph=None):
+    """
+    Simulate a trained policy
+    Args:
+        policy: the policy to simulate, which should return action directly
+        loc (int): location index to test policy
+        track (int): track to test policy
+        filestr (str): path prefix to save simulation to
+    """
+    # animate from environment
+    basepath = os.path.abspath('./InteractionSimulator')
+    svt, svt_path = get_svt(base=basepath, loc=loc, track=track)
+    osm = get_map_path(base=basepath, loc=loc)
+    if graph:
+        env = gym.make('intersim:intersim-v0', svt=svt, map_path=osm, 
+            min_acc=-np.inf, max_acc=np.inf, graph=graph, mask_relstate=True)
+    else:
+        env = gym.make('intersim:intersim-v0', svt=svt, map_path=osm, 
+            min_acc=-np.inf, max_acc=np.inf)
+    # env = gym.make('intersim:intersim-v0', loc=loc, track=track, 
+    #     min_acc=-np.inf, max_acc=np.inf)
+    
+    ob, _ = env.reset()
+    env.render()
+    done = False
+    i = 0
+    with tqdm(total=min(nframes, env._svt.Tind)) as pbar:
+        while not done and i < nframes:  
+            i += 1
+            
+            # get action
+            action = policy(ob)
+
+            # propagate environment
+            ob, r, done, info = env.step(action)
+            env.render()
+
+            pbar.update()
+
+    env.close(filestr=filestr+'_sim')

src/metrics.py

@@ -0,0 +1,251 @@
+import torch
+import pickle
+import numpy as np
+import matplotlib.pyplot as plt
+from torch.utils.data import DataLoader
+from intersim import collisions
+
+def metrics(filestr: str, test_dataset, policy):
+    """
+    Calculate metrics using a) base filestring to a simulation, and b) the test dataset and learned policy
+    Args:
+        filestr (str): base string to outputs of a simulation
+        test_dataset: a dataset held for testing
+        policy: policy
+    Returns:
+        info (dict): metrics in a dictionary
+    """
+    info = {}
+
+    # compute metrics using either
+    # a) simulation files that were saved under the trained policy with prefix 'policy'
+    # b) applying the policy to observations in the test dataset
+
+    # load simulated trajectory
+    states = torch.load(filestr + '_sim_states.pt').detach()
+    lengths = torch.load(filestr + '_sim_lengths.pt').detach()
+    widths = torch.load(filestr + '_sim_widths.pt').detach()
+    xpoly = torch.load(filestr + '_sim_xpoly.pt').detach()
+    ypoly = torch.load(filestr + '_sim_ypoly.pt').detach()
+
+    # count collisions (from function in intersim.collisions)
+    n_collisions = collisions.count_collisions_trajectory(states, lengths, widths)
+    info['n_collisions'] = n_collisions
+
+    # calculate average velocity
+    avg_v = average_velocity(states)
+    info['average_velocity'] = avg_v
+
+    # convert policy dtype between float32 and float64
+    policy.policy = policy.policy.type(test_dataset[0]['state']['ego_state'].dtype)
+    
+    # generate actions in test dataset
+    true_actions, pred_actions = [], []
+    true_velocities = []
+    test_loader = DataLoader(test_dataset, batch_size=1024)
+    with torch.no_grad():
+        for (batch_idx, batch) in enumerate(test_loader):
+            pred_actions.append(policy(batch['state']))
+            true_actions.append(batch['action'])
+            true_velocities.append(batch['state']['ego_state'][:,2])
+
+    true_actions, pred_actions = torch.cat(true_actions,dim=0), torch.cat(pred_actions, dim=0)
+    visualize_distribution(true_actions[:,0], pred_actions[:,0], filestr+'_action_viz') 
+
+    # calculate divergence between acceleration distributions
+    acceleration_divergence = divergence(pred_actions, true_actions, type='js')
+    info['acceleration_divergence'] = acceleration_divergence
+
+    # calculate divergence between velocity distributions
+    sim_velocities = states[:,:,2]
+    sim_velocities = sim_velocities[~torch.isnan(sim_velocities)].flatten()
+    true_velocities = torch.cat(true_velocities, dim=0)
+    velocity_divergence = divergence(sim_velocities, true_velocities, type='js')
+    info['velocity_divergence'] = velocity_divergence
+
+    return info
+
+
+def visualize_distribution(true, pred, filestr):
+    """
+    Visualize two distributions
+    Args:
+        true (torch.tensor): (n,)-sized true distribution
+        pred (torch.tensor): (m,)-sized pred distribution
+        filestr (str): string to save figure to
+    """
+    nni1 = ~torch.isnan(true)
+    nni2 = ~torch.isnan(pred)
+    plt.figure()
+    plt.hist(true[nni1].numpy(), density=True, bins=20)
+    plt.hist(pred[nni2].numpy(), density=True, bins=20)
+    plt.legend(['True', 'Predicted'])
+    plt.savefig(filestr+'.png')
+    
+def average_velocity(states):
+    """
+    Compute average of average velocity over all vehicles.
+    Args:
+        states (torch.tensor): (T,nv,5) vehicle states where T is the number of time steps and nv the number of vehicles
+    Returns
+        avg_v (float): average velocity
+    """
+    velocities = states[:,:,2]
+    # average velocity per vehicle
+    vehicle_avg_v = nanmean(velocities, dim=0)
+    arg_v = nanmean(vehicle_avg_v)
+    return arg_v
+
+def divergence(p, q, type='js', n_components=-1):
+    """
+    Calculate a divergence between p and q
+    Args:
+        p (torch.tensor): (n) samples from p
+        q (torch.tensor): (m) samples from q
+        type (str): divergence to use
+            'kl': Kullback-Leibler divergence KL(p||q)
+            'js': Jensen-Shannon divergence (symmetric KLD)
+        n_components (int): method to use to compute kl divergence
+            n_components < 0: approximate samples with histogram density
+            n_components == 0: approximate samples by Gaussian distributions and compute analytically
+            n_components > 0: approximate samples as Gaussian mixture models with n_components components
+    Returns:
+        d (float): approximate divergence
+    """
+    if type == 'js':
+        # Use histogram binning to discretize sampled distributions
+        p_hist = np.histogram(p, bins='auto', density=True)
+        q_hist = np.histogram(q, bins='auto', density=True)
+        m = torch.cat([p, q], dim=0)
+        m_weights = torch.cat([torch.full_like(p, 1./len(p)), torch.full_like(q, 1./len(q))], dim=0)
+        m_bins = np.sort(np.concatenate([p_hist[1], q_hist[1]]))
+        m_hist = np.histogram(m, bins=m_bins, density=True, weights=m_weights)
+        d = .5 * kl_histogram(p, p_hist, m_hist) + .5 * kl_histogram(q, q_hist, m_hist)
+        return d
+    elif type == 'kl':
+        if n_components < 0:
+            # Use histogram binning to discretize sampled distributions
+            p_hist = np.histogram(p, bins='auto', density=True)
+            q_hist = np.histogram(q, bins='auto', density=True)
+            d = kl_histogram(p, p_hist, q_hist)
+            return d
+        elif n_components == 0:
+            # Assume p and q to be Gaussian
+            pm = torch.mean(p)
+            qm = torch.mean(q)
+            pv = torch.var(p)
+            qv = torch.var(q)
+            d = kl_normal(pm, pv, qm, qv).item()
+            return d
+        else:
+            from sklearn.mixture import GaussianMixture
+            p = p.unsqueeze(-1)
+            q = q.unsqueeze(-1)
+            p_gmm = GaussianMixture(n_components=n_components).fit(p)
+            q_gmm = GaussianMixture(n_components=n_components).fit(q)
+            px = p_gmm.score_samples(p)
+            qx = q_gmm.score_samples(p)
+            d = np.mean(px - qx).item()
+            return d
+    else:
+        raise NotImplementedError("Please implement divergence for type '{}'".format(type))
+
+def kl_histogram(p_sample, p_hist, q_hist):
+    """
+    Calculate the kl divergence between p and q based on a histogram representation
+    Args:
+        p_sample (torch.tensor): (n) samples from p
+        p_hist (tuple): result of np.histogram(density=True) for samples from p
+        q_hist (tuple): result of np.histogram(density=True) for samples from q
+    Returns:
+        d (float): approximate KL divergence
+    """
+    p_density, p_edges = p_hist
+    q_density, q_edges = q_hist
+    px = evaluate_histogram(p_sample, p_density, p_edges)
+    qx = evaluate_histogram(p_sample, q_density, q_edges)
+    p_supp = ~np.isclose(px, 0.0)
+    q_supp = ~np.isclose(qx, 0.0)
+    if np.any(np.logical_and(p_supp, ~q_supp)):
+        # if not support(p) subset support(q)
+        return np.inf
+    elif ~np.any(p_supp):
+        # if p is zero everywhere
+        return 0.
+    d = np.mean(np.log(px[p_supp] / qx[p_supp]))
+    return d
+
+
+def kl_normal(pm, pv, qm, qv):
+    """
+    Computes the elem-wise KL divergence between two normal distributions KL(p || q) and
+    sum over the last dimension
+
+    Args:
+        pm: tensor: (batch, dim): p mean
+        pv: tensor: (batch, dim): p variance
+        qm: tensor: (batch, dim): q mean
+        qv: tensor: (batch, dim): q variance
+
+    Return:
+        kl: tensor: (batch,): kl between each sample
+    """
+    element_wise = 0.5 * (torch.log(qv) - torch.log(pv) + pv / qv + (pm - qm).pow(2) / qv - 1)
+    kl = element_wise.sum(-1)
+    return kl
+
+
+def kl_cat(q, log_q, log_p):
+    """
+    Computes the KL divergence between two categorical distributions
+
+    Args:
+        q: tensor: (batch, dim): Categorical distribution parameters
+        log_q: tensor: (batch, dim): Log of q
+        log_p: tensor: (batch, dim): Log of p
+
+    Return:
+        kl: tensor: (batch,) kl between each sample
+    """
+    element_wise = (q * (log_q - log_p))
+    kl = element_wise.sum(-1)
+    return kl
+
+
+def nanmean(v, *args, inplace=False, **kwargs):
+    """
+    Calculate mean over not nan entries
+
+    To be added to torch as torch.nanmean in the next release
+    https://github.com/pytorch/pytorch/issues/61474, https://github.com/pytorch/pytorch/issues/21987
+
+    Args:
+        v (torch.tensor): arbitrary tensor
+    Returns:
+        result (torch.tensor): mean over non nan elements
+    """
+    if not inplace:
+        v = v.clone()
+    is_nan = torch.isnan(v)
+    v[is_nan] = 0
+    result = v.sum(*args, **kwargs) / (~is_nan).float().sum(*args, **kwargs)
+    return result
+
+
+def evaluate_histogram(x, hist, bin_edges):
+    """
+    Evaluate a histogram 
+    Args:
+        x (array) : points at which to evaluate the histogram
+        hist (array): histogram values in terms of number of occurrences or probability
+        bin_edges (array): edges of histogram bins
+            e.g. from hist, bin_edges = np.histogram(p, bins='auto', density=True)
+    Return:
+        r: tensor: (batch,) kl between each sample
+    """
+    idx = np.digitize(x, bin_edges)
+    mask = np.logical_and(np.less(0, idx), np.less(idx, len(bin_edges)))
+    r = np.zeros_like(x)
+    r[mask] = hist[idx[mask] - 1]
+    return r
+

src/nets/deepsets.py

@@ -1,7 +1,7 @@
 import torch
 from torch import nn
 
-from interimit.nets.util import parse_functional
+from src.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):
@@ -17,10 +17,11 @@ class DeepSetsModule(nn.Module):
         """
         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
+        self.latent_dim = latent_dim
+        self.phi = Phi(self.input_dim, phi_hidden_n, phi_hidden_dim, self.latent_dim)
+        self.rho = Phi(self.latent_dim, rho_hidden_n, rho_hidden_dim, output_dim)
+        self.output_dim = self.rho.output_dim
+        self.pooling = torch.sum
 
     @staticmethod
     def from_config(config):
@@ -54,18 +55,22 @@ class DeepSetsModule(nn.Module):
     def forward(self, x):
         """
         Args:
-            x (torch.tensor): (batch_size, dynamic_size, input_dim)
+            x (torch.tensor): ([B, ]max_nv, d)
         Returns:
-            y (torch.tensor): (batch_size, output_dim)
+            y (torch.tensor): ([B, ]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)
+        # mask for selecting only those batches and vehicles where all relative states are not nan
+        # shape (B, max_nv)
+        notnan_mask = torch.all(~torch.isnan(x), dim=-1)
+        # create zero tensor of shape (B, max_nv, latent_dim) to store phi evaluations in
+        latent = torch.zeros([*x.shape[:-1], self.latent_dim], dtype=x.dtype)
+        # evaluate phi for all not NaN entries
+        # x[batch_dynamic_mask] has shape (notnan_mask.sum(), input_dim)
+        latent[notnan_mask] = self.phi(x[notnan_mask])
+
+        # sum over relative state dimension
+        latent = self.pooling(latent, dim=-2)
+        
         # apply rho network
         y = self.rho(latent)
         return y
@@ -85,15 +90,16 @@ class Phi(nn.Module):
         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)
+        if hidden_n > 0:
+            self.layers = nn.ModuleList([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))
+        else:
+            self.layers = nn.ModuleList([nn.Identity()])
+            self.output_dim = self.input_dim
         self.activation = nn.functional.relu
-        self.final_activation = final_activation if final_activation else self.activation
+        self.final_activation = final_activation if final_activation else lambda x: x
 
     def forward(self, x):
         for layer in self.layers[:-1]:

src/nets/util.py

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

src/policies/__init__.py

@@ -0,0 +1 @@
+from src.policies.policy import IntersimPolicy, IntersimStateNet, IntersimStateActionNet, generate_transforms

src/policies/policy.py

@@ -0,0 +1,164 @@
+
+import torch
+from torch import nn
+
+from src.nets.deepsets import DeepSetsModule, Phi
+from src.util.transform import MinMaxScaler
+
+class IntersimStateNet(nn.Module):
+    def __init__(self, config):
+        """
+        Args:
+            config (dict): dictionary for configuring the deep sets policy
+        """
+        super(IntersimStateNet, self).__init__()
+        ego_config = config['ego_encoder']
+        deepsets_config = config['deepsets']
+        pathnet_config = config['path_encoder']
+
+        self.ego_net = Phi.from_config(ego_config)
+        self.deepsets_net = DeepSetsModule.from_config(deepsets_config)
+        self.path_net = Phi.from_config(pathnet_config)
+
+        cat_dim = self.ego_net.output_dim + self.deepsets_net.output_dim + self.path_net.output_dim
+        # head has number of concatenated features as input
+        head_config = config['head']
+        head_config["input_dim"] = cat_dim
+        self.head = Phi.from_config(head_config)
+
+    def forward(self, sample):
+        """
+        Args:
+            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 (torch.tensor): (B, P, 2) tensor of P future path x and y positions
+                action (torch.tensor): (B, 1) actions taken from each state
+        Returns:
+            x (torch.tensor): (head_output_dim,) output of common head network
+        """
+        ego = self.ego_net(sample["ego_state"])
+        relative = self.deepsets_net(sample["relative_state"])
+        path = self.path_net(sample["path"].reshape((sample["path"].shape[0], -1)))
+        x = torch.cat([ego, relative, path], dim=-1)
+        x = self.head(x)
+        return x
+
+
+class IntersimStateActionNet(nn.Module):
+    def __init__(self, config):
+        """
+        Args:
+            config (dict): dictionary for configuring the deep sets policy
+        """
+        super(IntersimStateActionNet, self).__init__()
+        ego_config = config['ego_encoder']
+        deepsets_config = config['deepsets']
+        pathnet_config = config['path_encoder']
+
+        self.ego_net = Phi.from_config(ego_config)
+        self.deepsets_net = DeepSetsModule.from_config(deepsets_config)
+        self.path_net = Phi.from_config(pathnet_config)
+        self.action_dim = config["action_dim"]
+
+        cat_dim = self.ego_net.output_dim + self.deepsets_net.output_dim + self.path_net.output_dim + self.action_dim
+        # head has number of concatenated features as input
+        head_config = config['head']
+        head_config["input_dim"] = cat_dim
+        self.head = Phi.from_config(head_config)
+
+    def forward(self, sample):
+        """
+        Args:
+            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
+        Returns:
+            x (torch.tensor): (head_output_dim,) output of common head network
+        """
+        ego = self.ego_net(sample["ego_state"])
+        relative = self.deepsets_net(sample["relative_state"])
+        action = sample["action"]
+        path = self.path_net(sample["path"].reshape((sample["path"].shape[0], -1)))
+        x = torch.cat([ego, relative, path, action], dim=-1)
+        x = self.head(x)
+        return x
+
+
+class IntersimPolicy():
+    """
+    Base class for intersim policies
+    """
+    def __init__(self, config, transforms):
+        super(IntersimPolicy, self).__init__()
+        self._config = config
+        self._transforms = transforms
+
+    @property
+    def transforms(self):
+        return self._transforms
+    
+    @transforms.setter
+    def transforms(self, transforms):
+        self._transforms=transforms
+
+    @property 
+    def policy(self):
+        return self._policy
+    
+    @policy.setter
+    def policy(self, policy):
+        self._policy = policy
+
+    def transform_observation(self, ob):
+        # run observation through transforms
+        transformed_ob = {}
+        for key in ['ego_state', 'relative_state', 'path', 'action']:             
+            if key in self._transforms.keys() and key in ob.keys():
+                transformed_ob[key] = self._transforms[key].transform(ob[key])
+        return transformed_ob
+
+    def __call__(self, ob):
+
+        if 'ego_state' in ob.keys():
+            # extract state from dataloader samples  
+            pass
+        else:
+            # extract state from observation (using simulator)
+            ob['ego_state'] = ob['state']
+            ob['path'] = torch.stack(ob['paths'],dim=-1)
+
+        ob = self.transform_observation(ob)
+
+        # run transformed state through model
+        action = self._policy(ob)
+        assert action.ndim == 2, 'action has incorrect shape'
+
+        # untransform action
+        if 'action' in self._transforms.keys():
+            action = self._transforms['action'].inverse_transform(action)
+        return action
+
+
+def generate_transforms(dataset):
+    """
+    Generate transform dictionary from dataset
+    Args:
+        dataset (Dataset): dataset of demo observations and actions
+    """
+    transforms = {
+        'action': MinMaxScaler(),
+        'ego_state': MinMaxScaler(),
+        'relative_state': MinMaxScaler(reduce_dim=2),
+        'path': MinMaxScaler(reduce_dim=2),
+    }
+    for key in transforms.keys():
+        if key == 'action':
+            transforms[key].fit(dataset[:][key])
+        else:
+            transforms[key].fit(dataset[:]['state'][key])
+
+    return transforms

src/util/nn_training.py

@@ -0,0 +1,11 @@
+import torch
+
+def optimizer_factory(config, parameters):
+    optimizer_type = config['optimizer']
+    learning_rate = config['lr']
+    weight_decay = config['weight_decay']
+    if optimizer_type == 'adam':  
+        optimizer = torch.optim.Adam(parameters, lr=learning_rate, weight_decay=weight_decay)
+    else:
+        raise NotImplementedError
+    return optimizer

src/util/transform.py

@@ -0,0 +1,140 @@
+import torch
+from torch import nn
+import numpy as np
+from sklearn import preprocessing
+
+class Transform(nn.Module):
+    """
+    Base class to normalize observations and actions for network.
+    """
+
+    def __init__(self):
+        super(Transform, self).__init__()
+        # self.fit(X)
+
+    def fit(self, X):
+        """
+        Fit transformer to X
+        Args:
+            X (torch.tensor): (B, N) tensor of B data points with N features
+        """
+        raise NotImplementedError('Please implement fit()')
+
+    def transform(self, X):
+        """
+        Transform X. fit() has to be called first
+        Args:
+            X (torch.tensor): (B, N) tensor where N has to be the same as during fit()
+        """
+        raise NotImplementedError('Please implement transform()')
+
+    def inverse_transform(self, X):
+        """
+        Inverse transformation
+        Args:
+            X (torch.tensor): (B, N) tensor
+        """
+        raise NotImplementedError('Please implement inverse_transform()')
+
+    def forward(self, X):
+        return self.transform(X)
+
+class MinMaxScaler(Transform):
+    """
+    Scale tensor so each feature is in [0, 1]
+    """
+    def __init__(self, reduce_dim:int=None):
+        """
+        Initialize SciKitTransform
+        Args:
+            reduce_dim (int): dimension to start calculating featues from
+                e.g. with reduce_dim=2, (A, B, C, D, E) will be reshaped to (A*B, C*D*E)
+        """
+        self.reduce_dim = reduce_dim
+        super(MinMaxScaler, self).__init__()
+
+    def fit(self, X):
+        nd = X.ndim
+        if self.reduce_dim:
+            self.nfeatures = int(torch.tensor(X.shape[self.reduce_dim:]).prod())
+        else:
+            assert nd==2, 'Invalid ndim'
+            self.nfeatures = X.shape[1]
+        
+        X = X.reshape((-1,self.nfeatures))
+        nans = torch.isnan(X)
+        X[nans] = float('inf')
+        self.min = X.min(0,keepdims=True)[0]
+
+        X[nans] = -float('inf')
+        self.span = X.max(0,keepdims=True)[0] - self.min
+
+        X[nans] = np.nan
+
+    def transform(self, X):
+        
+        assert hasattr(self, 'min') and hasattr(self, 'span'), 'Model not yet fit'
+        shape = X.shape
+        X = X.reshape((-1,self.nfeatures))
+        t = (X - self.min) / self.span
+        return t.reshape(shape)
+
+    def inverse_transform(self, X):
+
+        assert hasattr(self, 'min') and hasattr(self, 'span'), 'Model not yet fit'
+        shape = X.shape
+        X = X.reshape((-1,self.nfeatures))
+        it =  X * self.span + self.min
+        return it.reshape(shape)
+
+
+class SciKitTransform(Transform):
+    """
+    Wrappers around scikit-learn transforms
+    """
+    def __init__(self, tf, reduce_dim:int=None):
+        """
+        Initialize SciKitTransform
+        Args:
+            tf: transform
+            reduce_dim (int): dimension to start calculating featues from
+                e.g. with reduce_dim=2, (A, B, C, D, E) will be reshaped to (A*B, C*D*E)
+        """
+        self.tf = tf
+        self.reduce_dim = reduce_dim
+        super(SciKitTransform, self).__init__()
+
+    def fit(self, X):
+        nd = X.ndim
+        if self.reduce_dim:
+            self.nfeatures = int(torch.tensor(X.shape[self.reduce_dim:]).prod())
+        else:
+            assert nd==2, 'Invalid ndim'
+            self.nfeatures = X.shape[1]
+
+        self.tf.fit(X.reshape((-1,self.nfeatures)))
+
+    def transform(self, X):
+        shape = X.shape
+        t = torch.tensor(self.tf.transform(X.reshape((-1,self.nfeatures))), dtype=torch.float)
+        return t.reshape(shape)
+
+    def inverse_transform(self, X):
+        shape = X.shape
+        it =  torch.tensor(self.tf.inverse_transform(X.reshape((-1,self.nfeatures))), dtype=torch.float)
+        return it.reshape(shape)
+
+class SciKitStandardScaler(SciKitTransform):
+    """
+    Wrapper around scikit-learn's StandardScaler for standardizing each feature individually.
+    """
+    def __init__(self, **kwargs):
+        super(SciKitStandardScaler, self).__init__(preprocessing.StandardScaler(), **kwargs)
+
+class SciKitMinMaxScaler(SciKitTransform):
+    """
+    Wrapper around scikit-learn's MinMaxScaler for scaling features to [0, 1] individually.
+    """
+    def __init__(self, **kwargs):
+        super(SciKitMinMaxScaler, self).__init__(preprocessing.MinMaxScaler(), **kwargs)
+

src/value_dice/__init__.py

@@ -0,0 +1 @@
+from src.value_dice.value_dice import ValueDicePolicy, train, vd_config

src/value_dice/value_dice.py

@@ -0,0 +1,300 @@
+import numpy as np
+
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+from torch.nn.utils import clip_grad_norm_
+import pickle
+import itertools
+from torch.utils.tensorboard import SummaryWriter
+
+from src.policies import IntersimStateNet, IntersimStateActionNet, IntersimPolicy, generate_transforms
+from src.util.transform import MinMaxScaler
+from src.util.nn_training import optimizer_factory
+from tqdm import tqdm
+import json5
+from ray import tune
+
+def vd_config(ray_config):
+    config = {
+        'policy_net': {
+            'ego_encoder': {'input_dim': 5, 'hidden_n': 0, 'hidden_dim':0, 'output_dim': 0},
+            'deepsets': {
+                'input_dim': 6,
+                'phi': {
+                    'hidden_n': ray_config['deepsets_phi_hidden_n'],
+                    'hidden_dim': ray_config['deepsets_phi_hidden_dim']
+                    },
+                'latent_dim': ray_config['deepsets_latent_dim'],
+                'rho': {
+                    'hidden_n': ray_config['deepsets_rho_hidden_n'],
+                    'hidden_dim': ray_config['deepsets_rho_hidden_dim']
+                    },
+                'output_dim': ray_config['deepsets_output_dim']
+            },
+            'path_encoder': {'input_dim': 40, 'hidden_n': 0, 'hidden_dim': 0, 'output_dim': 0},
+            'head': {
+                'input_dim': 0, # computed in constructor
+                'hidden_n': ray_config['head_hidden_n'],
+                'hidden_dim': ray_config['head_hidden_dim'],
+                'output_dim': 1, # number of outputs e.g. number of actions, or just one
+                'final_activation': ray_config['head_final_activation'],
+            },
+        },
+        'value_net': {
+            'ego_encoder': {'input_dim': 5, 'hidden_n': 0, 'hidden_dim':0, 'output_dim': 0},
+            'deepsets': {
+                'input_dim': 6,
+                'phi': {
+                    'hidden_n': ray_config['deepsets_phi_hidden_n'],
+                    'hidden_dim': ray_config['deepsets_phi_hidden_dim']
+                    },
+                'latent_dim': ray_config['deepsets_latent_dim'],
+                'rho': {
+                    'hidden_n': ray_config['deepsets_rho_hidden_n'],
+                    'hidden_dim': ray_config['deepsets_rho_hidden_dim']
+                    },
+                'output_dim': ray_config['deepsets_output_dim']
+            },
+            'path_encoder': {'input_dim': 40, 'hidden_n': 0, 'hidden_dim': 0, 'output_dim': 0},
+            'action_dim': 1,
+            'head': {
+                'input_dim': 0, # computed in constructor
+                'hidden_n': ray_config['head_hidden_n'],
+                'hidden_dim': ray_config['head_hidden_dim'],
+                'output_dim': 1, # number of outputs e.g. number of actions, or just one
+                'final_activation': ray_config['head_final_activation'],
+            },
+        },
+        'policy_optim': {
+            'optimizer':'adam',
+            'lr':ray_config['policy_lr'],
+            'weight_decay':ray_config['policy_weight_decay']
+        },
+        'value_optim': {
+            'optimizer':'adam',
+            'lr':ray_config['value_lr'],
+            'weight_decay':ray_config['value_weight_decay']
+        },
+        'train_epochs': 40,
+        'train_batch_size': ray_config['train_batch_size'],
+        'discount': ray_config['discount'],
+        'clip_grad_norm': ray_config['clip_grad_norm'],
+    }
+    return config
+
+class ValueDicePolicy(IntersimPolicy):
+    """
+    Class for value dice policy
+    """
+    
+    def __init__(self, config: dict, transforms: dict):
+        """
+        Initialize ValueDicePolicy
+        Args:
+            config (dict): configuration file to initialize IntersimDeepSetsNet with
+            transforms (dict): dictionary of transforms to apply to different fields
+        """
+        super(ValueDicePolicy, self).__init__(config, transforms)
+        self._policy = IntersimStateNet(config['policy_net'])
+        self._value = IntersimStateActionNet(config['value_net'])
+        
+    @property 
+    def value(self):
+        return self._value
+    
+    @value.setter
+    def value(self, value):
+        self._value = value
+
+    @classmethod
+    def load_model(cls, filestr: str, config: dict = None):
+        """
+        Load a model from a file prefix
+        Args:
+            config (dict): configuration dict to set up model
+            filestr (str): string prefix to load model from
+        Returns
+            model (BehaviorCloningPolicy): loaded model
+        """
+        if not config:
+            with open(filestr+'_config.json', 'r') as cfg:
+                config = json5.load(cfg)
+        transforms = pickle.load(open(filestr+'_transforms.pkl', 'rb'))
+        model = cls(config, transforms=transforms)
+        model._policy.load_state_dict(torch.load(filestr+'_policy.pt'))
+        model._value.load_state_dict(torch.load(filestr+'_value.pt'))
+        return model
+    
+    def parameters(self):
+        return itertools.chain(self._policy.parameters(), self._value.parameters())
+
+    def policy_parameters(self):
+        return self.policy.parameters()
+
+    def value_parameters(self):
+        return self.value.parameters()
+
+    def eval(self):
+        self.policy.eval()
+        self.value.eval()
+
+    def save_model(self, filestr, save_config=True, save_transforms=True):
+        """
+        Save transforms and state_dict to a location specificed by filestr
+        Args:
+            filestr (str): string prefix to save model to
+            save_config (bool): whether to save the config file (as a json)
+            save_transforms (bool): whether to save transforms (as a pickle)
+        """
+        if save_config:
+            with open(filestr+'_config.json', 'w') as cfg:
+                json5.dump(self._config, cfg)
+        if save_transforms:
+            pickle.dump(self._transforms, open(filestr+'_transforms.pkl', 'wb'))
+        torch.save(self._policy.state_dict(), filestr+'_policy.pt')
+        torch.save(self._value.state_dict(), filestr+'_value.pt')
+
+
+
+def train(config, policy, train_dataset, cv_dataset, filestr, **kwargs):
+    
+    using_ray = kwargs.get('ray', False)
+    if using_ray:
+        print('using ray')
+
+    # hyperparams
+    train_epochs = config['train_epochs']
+    train_batch_size = config['train_batch_size']
+    discount = config['discount']
+    clip_grad_norm = config['clip_grad_norm']
+
+    cv_every = 1
+    print_epoch_every = 1000
+    print_cv_every = 5
+    checkpoint_every = 100
+    cv_batch_size = 256 # doesn't matter
+
+    # training and testing dataloaders
+    training_loader = DataLoader(train_dataset, batch_size=train_batch_size, shuffle=True)
+    cv_loader = DataLoader(cv_dataset, batch_size=cv_batch_size, shuffle=True)
+
+    # change policy dtype
+    dtype = train_dataset[0]['state']['ego_state'].dtype
+    policy.policy = policy.policy.type(dtype)
+    policy.value = policy.value.type(dtype)
+
+    # define loss function
+    def f_value_dice_loss(batch):
+        # get s, a, s', s_0 from batch
+        state = batch['state']
+        action = batch['action']
+        next_state = batch['next_state']
+        initial_state = state
+
+        # append action to state batches
+        #   use expert action for s
+        state['action'] = action
+        #   run s' and s_0 through policy
+        initial_state['action'] = policy(initial_state)
+        next_state['action'] = policy(next_state)
+
+        # transform state and action before inputting to value network
+        # (for the policy network this is done in policy.__call__() )
+        state = policy.transform_observation(state)
+        initial_state = policy.transform_observation(initial_state)
+        next_state = policy.transform_observation(next_state)
+
+        # evaluate value network
+        value = policy.value(state)
+        value_init = policy.value(initial_state)
+        value_next = policy.value(next_state)
+        
+        # linear loss
+        linear_loss = (1 - discount) * torch.mean(value_init)
+
+        # nonlinear loss
+        value_diff = value - discount * value_next
+        nonlinear_loss = torch.logsumexp(value_diff, dim=0) - np.log(len(value_diff))
+
+        loss = nonlinear_loss - linear_loss
+        return loss
+
+
+    policy_optimizer = optimizer_factory(config['policy_optim'], policy.policy_parameters())
+    value_optimizer = optimizer_factory(config['value_optim'], policy.value_parameters())
+
+    # generate tensorboard writer
+    if not using_ray:
+        writer = SummaryWriter(filestr)
+    
+    for i in tqdm(range(train_epochs)):
+        
+        # save model checkpoints
+        if i % checkpoint_every == 0:
+            policy.save_model(filestr + '_epoch%04i'%(i) )
+
+        # train
+        epoch_loss = 0
+        for (batch_idx, batch) in enumerate(training_loader):
+
+            loss = f_value_dice_loss(batch)
+
+            # In original implementation policy is regularized with orthogonal regularization, 
+            # value with L2 regularization on gradients
+            policy_loss = -loss
+            value_loss = loss
+
+            # # compute loss and step optimizer
+            # policy_optimizer.zero_grad()
+            # value_optimizer.zero_grad()
+            # policy_loss.backward(retain_graph=True)
+            # value_loss.backward()
+
+            # clip_grad_norm_(policy.policy.parameters(), clip_grad_norm)
+            # clip_grad_norm_(policy.value.parameters(), clip_grad_norm)
+
+            # policy_optimizer.step()
+            # value_optimizer.step()
+
+            if batch_idx % 2 == 0:
+                policy_optimizer.zero_grad()
+                policy_loss.backward()
+                clip_grad_norm_(policy.policy.parameters(), clip_grad_norm)
+                policy_optimizer.step()
+            else:
+                value_optimizer.zero_grad()
+                value_loss.backward()
+                clip_grad_norm_(policy.value.parameters(), clip_grad_norm)
+                value_optimizer.step()
+
+            epoch_loss += loss.item() / len(train_dataset)
+        
+        if i % print_epoch_every == 0:
+            print('Epoch: {}, Training Loss: {}'.format(i, epoch_loss))
+
+        # measure cv loss
+        if i % cv_every == 0:
+            with torch.no_grad():
+                cv_loss = 0.
+                for (batch_idx, batch) in enumerate(cv_loader):
+                    loss = f_value_dice_loss(batch)
+                    cv_loss += loss.item() / len(cv_dataset)
+            
+            
+        # Write epoch loss
+        if using_ray:
+            if i % cv_every == 0:
+                tune.report(training_loss=epoch_loss, cv_loss=cv_loss, training_iteration=i+1)
+            else:
+                tune.report(training_loss=epoch_loss, training_iteration=i+1)
+        else:
+            writer.add_scalar('training loss',epoch_loss, i)
+            if i % cv_every == 0:
+                writer.add_scalar('cv loss', cv_loss, i)
+
+        if i % print_cv_every == 0:    
+            print('Epoch: {}, CV Loss: {}'.format(i, cv_loss))
+
+
+    policy.save_model(filestr)

tests/nets/test_deepsets.py

@@ -1,7 +1,8 @@
 import torch
 import random
-from interimit.nets import deepsets as ds
+from src.nets import deepsets as ds
 import copy
+import numpy as np
 
 ds_config = {
     "input_dim": 5,
@@ -24,7 +25,11 @@ def test_constructor():
     phi_config["output_dim"] = 2
     phi_config["final_activation"] = "sigmoid"
     phi = ds.Phi.from_config(phi_config)
-    assert phi.final_activation == torch.nn.functional.sigmoid
+    assert phi.final_activation == torch.sigmoid
+
+    phi_config["final_activation"] = "relu"
+    phi = ds.Phi.from_config(phi_config)
+    assert phi.final_activation == torch.nn.functional.relu
 
 def test_phi():
     input_dim = 5
@@ -36,38 +41,75 @@ def test_phi():
     y = phi(torch.rand(input_dim))
     y = phi(torch.rand(7,7,7,input_dim))
 
+    assert len(list(phi.parameters())) > 0
+
 def test_deepsets():
     m = ds.DeepSetsModule.from_config(ds_config)
 
     input_dim = ds_config["input_dim"]
-    n_dynamic = random.randint(5, 15)
-    x = torch.rand(n_dynamic, input_dim)
+    B = 50
+    max_V = 10
+    batch = []
+    for i in range(B):
+        if i==0: # ensure that there is an example with no vehicles
+            n_dynamic = 0
+        elif i==1: # and one with full vehicles
+            n_dynamic = max_V
+        else:
+            n_dynamic = random.randint(1, max_V)
+        x = torch.rand(n_dynamic, input_dim)
+        n_nan = max_V - n_dynamic
+        x = torch.cat([x, torch.zeros(n_nan, input_dim) * np.nan])
+        assert x.shape == torch.Size([max_V, input_dim])
+        batch.append(x)
+    batch = torch.stack(batch)
+    assert batch.shape == torch.Size([B, max_V, input_dim])
 
-    n_batch = 20
-    x = x.unsqueeze(0).expand(n_batch, n_dynamic, input_dim)
+    y = m(batch)
+    assert y.shape == torch.Size([B, ds_config["output_dim"]])
+    assert torch.isnan(y).sum() == 0
 
-    y = m(x)
-    assert y.shape == torch.Size([n_batch, ds_config["output_dim"]])
-
-    for i in range(n_batch):
-        assert torch.allclose(y[i], y[0])
+    for i in range(B):
+        y = m(batch[i])
+        assert y.shape == torch.Size([ds_config["output_dim"]])
+        assert torch.isnan(y).sum() == 0
 
 def test_deepsets_computation():
-    n_dynamic = random.randint(5,15)
-    n_batch = 7
     input_dim = 5
-    output_dim = 3
-    x = torch.rand(n_dynamic, input_dim)
-    x = x.unsqueeze(0).expand(n_batch, n_dynamic, input_dim)
-    assert x.shape == torch.Size([n_batch, n_dynamic, input_dim])
+    latent_dim = 8
 
-    phi = torch.nn.Linear(input_dim, output_dim)
+    B = 20
+    max_V = 10
+    batch = []
+    for i in range(B):
+        if i==0:
+            n_dynamic = 0
+        else:
+            n_dynamic = random.randint(1, max_V)
+        x = torch.rand(n_dynamic, input_dim)
+        n_nan = max_V - n_dynamic
+        x = torch.cat([x, torch.zeros(n_nan, input_dim) * np.nan])
+        assert x.shape == torch.Size([max_V, input_dim])
+        batch.append(x)
+    batch = torch.stack(batch)
+    assert batch.shape == torch.Size([B, max_V, input_dim])
+    x = batch
 
-    y = torch.stack(tuple(phi(instance) for instance in x.unbind(-2)), dim=-2)
-    assert y.shape == torch.Size([n_batch, n_dynamic, output_dim])
+    ### create phi
+    phi = ds.Phi(input_dim, 1, 10, latent_dim)
 
-    y = y.sum(dim=-2)
-    assert y.shape == torch.Size([n_batch, output_dim])
+    max_nv = x.shape[-2]
+    input_mask = ~torch.isnan(x)
+    batch_dynamic_mask = torch.all(input_mask, dim=-1)
+    assert batch_dynamic_mask.shape == x.shape[:-1]
+    batch_mask = torch.all(batch_dynamic_mask, dim=-1)
+    assert batch_mask.shape == x.shape[:-2]
 
-    for i in range(n_batch):
-        assert torch.allclose(y[i], y[0])
+    batch_dims = x.shape[:-2]
+    latent = torch.zeros([*batch_dims, max_nv, latent_dim])
+    latent[batch_dynamic_mask] = phi(x[batch_dynamic_mask])
+    assert x[batch_dynamic_mask].shape == torch.Size([batch_dynamic_mask.sum(), input_dim])
+    assert phi(x[batch_dynamic_mask]).shape == torch.Size([batch_dynamic_mask.sum(), latent_dim])
+
+    latent = latent.sum(dim=-2)
+    assert latent.shape == torch.Size([B, latent_dim])

tests/policies/test_forward_pass.py

@@ -0,0 +1,21 @@
+import torch
+import pickle
+import json5
+from src.bc import BehaviorCloningPolicy
+config_path = "config/networks.json5"
+
+with open(config_path, 'r') as cfg:
+    config = json5.load(cfg)
+
+filestr = 'tests/policies/base_'
+
+
+def test_forward_pass():
+    batch = pickle.load(open(filestr + '_test_batch.pkl', 'rb'))
+    batch['relative_state'] = batch['relative_state'].float()
+    policy = BehaviorCloningPolicy.load_model(config, filestr)
+    policy.policy = policy.policy.type(batch['state'].dtype)
+    action = policy(batch)
+    i = torch.where(action.isnan())[0]
+    for key in batch.keys():
+        print(batch[key][i])

tests/policies/test_policies.py

@@ -1,5 +1,5 @@
 import torch
-from interimit.policies.policy import DeepSetsPolicy
+from src.policies.policy import DeepSetsPolicy
 import json5
 
 config_path = "config/networks.json5"
@@ -7,14 +7,30 @@ with open(config_path, 'r') as cfg:
     config = json5.load(cfg)
 
 def test_deepsets_policy():
-    module = DeepSetsPolicy(config["ego_state"], config["deepsets"], config["path_encoder"], config["head"])
+    module = DeepSetsPolicy(config)
     
+    B = 25
     ns = 5 
     nv = 7
     npath = 20
 
-    ego_state = torch.rand(ns)
-    relative_state = torch.rand(nv, ns)
-    path = torch.rand(npath, 2)
-    
-    module(ego_state, relative_state, path)
+    ego_state = torch.rand(B, ns)
+    relative_state = torch.rand(B, nv, ns)
+    path_x = torch.rand(B, npath)
+    path_y = torch.rand(B, npath)
+
+    sample = {
+        "state": ego_state[0],
+        "relative_state": relative_state[0],
+        "path_x": path_x[0],
+        "path_y": path_y[0],
+    }
+    module(sample)
+
+    sample = {
+        "state": ego_state,
+        "relative_state": relative_state,
+        "path_x": path_x,
+        "path_y": path_y,
+    }
+    module(sample)

tests/test_metrics.py

@@ -0,0 +1,79 @@
+import torch
+import numpy as np
+
+from src import metrics
+from src.metrics import divergence, evaluate_histogram
+from sklearn.model_selection import GridSearchCV
+from sklearn.neighbors import KernelDensity
+
+def test_kl_divergence():
+    p = torch.randn(1000000)
+    q = 1.0 + 2.0 * torch.randn(1000000)
+
+    d1 = divergence(p, q, type='kl', n_components=0)
+    print(d1)
+    assert isinstance(d1, float)
+    d2 = divergence(p, q, type='kl', n_components=1)
+    print(d2)
+    assert isinstance(d2, float)
+    assert np.isclose(d1, d2, atol=1e-5)
+    d3 = divergence(p, q, type='kl', n_components=3)
+    print(d3)
+    assert isinstance(d3, float)
+    d4 = divergence(p, q, type='kl', n_components=-1)
+    assert isinstance(d4, float)
+    d5 = divergence(q, p, type='kl', n_components=-1)
+    assert isinstance(d5, float)
+
+    p = 1000 * torch.randn(1000)
+    q = torch.randn(1000)
+    d6 = divergence(p, q, type='kl', n_components=-1)
+    assert ~np.isfinite(d6)
+
+def test_evaluate_histogram():
+    N = 10000
+    p = torch.randn(N)
+    q = 1.0 + 2.0 * torch.randn(2*N)
+    
+    p_hist, p_edges = np.histogram(p.unsqueeze(-1), bins='auto', density=True)
+    q_hist, q_edges = np.histogram(q.unsqueeze(-1), bins='auto', density=True)
+    px = evaluate_histogram(p, p_hist, p_edges)
+    assert px.shape == p.shape
+    qx = evaluate_histogram(q, p_hist, p_edges)
+    assert qx.shape == q.shape
+    px = evaluate_histogram(p, q_hist, q_edges)
+    assert px.shape == p.shape
+    qx = evaluate_histogram(q, q_hist, q_edges)
+    assert qx.shape == q.shape
+
+
+def test_js_divergence():
+    N = 1000
+    p = torch.randn(N)
+    q = 1.0 + 2.0 * torch.randn(2*N)
+
+    d1 = divergence(p, q, type='js')
+    d2 = divergence(q, p, type='js')
+    assert d1 == d2
+
+    p = 1000 * torch.randn(1000)
+    q = torch.randn(1000)
+    d1 = divergence(p, q, type='js')
+    d2 = divergence(q, p, type='js')
+    assert d1 == d2
+    assert np.isfinite(d1)
+
+
+if __name__ == '__main__':
+    p = torch.randn(10)
+    q = 1.0 + 2.0 * torch.randn(10)
+
+    p = p.unsqueeze(-1)
+    q = q.unsqueeze(-1)
+    p_hist, p_edges = np.histogram(p.unsqueeze(-1), bins='auto', density=True)
+    q_hist, q_edges = np.histogram(q.unsqueeze(-1), bins='auto', density=True)
+    # px = p_hist[np.digitize(p, p_edges) - 1]
+
+    # qx = q_hist[np.digitize(p, q_edges) - 1]
+    px = evaluate_histogram(p, p_hist, p_edges)
+    qx = evaluate_histogram(q, p_hist, p_edges)