Add basic setup.py

.gitignore

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

README.md

@@ -24,26 +24,8 @@ 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 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
+You can then load the experts actions and observations using
 ```
 from src import expert_data
 observations, actions = expert_data.load_expert_data(loc = [LOCNUM], track = [TRACKNUM])

config/networks.json5

@@ -1,42 +1,34 @@
 {
-    ego_encoder: {
+    ego_state: {
         input_dim: 5, // number of state vars
-        hidden_n: 0,
+        hidden_n: 1,
         hidden_dim: 5,
         output_dim: 5
     },
     deepsets: {
-        input_dim: 6, // number of relative state vars for others
+        input_dim: 5, // number of relative state vars for others
         phi: {
-            hidden_n: 2,
+            hidden_n: 1,
             hidden_dim: 20,
         },
         latent_dim: 20,
         rho: {
-            hidden_n: 2,
+            hidden_n: 1,
             hidden_dim: 10,
         },
         output_dim: 10
     },
     path_encoder: {
         input_dim: 40, // 2 * path length for (x,y) coordinates
-        hidden_n: 0,
+        hidden_n: 2,
         hidden_dim: 20,
         output_dim: 10,
     },
     head: {
         input_dim: 0, // computed in policy constructor
-        hidden_n: 3,
+        hidden_n: 1,
         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

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

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

@@ -1,9 +0,0 @@
-#!/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

@@ -1,5 +0,0 @@
-#!/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

generate_demos.sh

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

interimit/__init__.py

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

interimit/data_utils.py

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

interimit/expert_data.py

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

interimit/nets/deepsets.py

@@ -1,7 +1,7 @@
 import torch
 from torch import nn
 
-from src.nets.util import parse_functional
+from interimit.nets.util import parse_functional
 
 class DeepSetsModule(nn.Module):
     def __init__(self, input_dim, phi_hidden_n, phi_hidden_dim, latent_dim, rho_hidden_n, rho_hidden_dim, output_dim):
@@ -17,11 +17,10 @@ class DeepSetsModule(nn.Module):
         """
         super(DeepSetsModule, self).__init__()
         self.input_dim = input_dim
-        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
+        self.output_dim = output_dim
+        self.phi = Phi(self.input_dim, phi_hidden_n, phi_hidden_dim, latent_dim)
+        self.rho = Phi(latent_dim, rho_hidden_n, rho_hidden_dim, self.output_dim)
+        self.pooling = torch.sum  # torch.max # torch.mean
 
     @staticmethod
     def from_config(config):
@@ -55,22 +54,18 @@ class DeepSetsModule(nn.Module):
     def forward(self, x):
         """
         Args:
-            x (torch.tensor): ([B, ]max_nv, d)
+            x (torch.tensor): (batch_size, dynamic_size, input_dim)
         Returns:
-            y (torch.tensor): ([B, ]output_dim)
+            y (torch.tensor): (batch_size, output_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)
-        
+        # use negative dynamic_dim since batch dimensions are inserted at the front
+        dynamic_dim = -2
+        # iterate over dynamic dimension to apply phi to every instance
+        latent = tuple(self.phi(instance) for instance in x.unbind(dynamic_dim))
+        # stack outputs of phi
+        latent = torch.stack(latent, dim=dynamic_dim)
+        # apply pooling function to reduce dynamic dimension
+        latent = self.pooling(latent, dim=dynamic_dim)
         # apply rho network
         y = self.rho(latent)
         return y
@@ -90,16 +85,15 @@ class Phi(nn.Module):
         super(Phi, self).__init__()
         self.input_dim = input_dim
         self.output_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.layers = [nn.Linear(self.input_dim, hidden_dim)]
+        for _ in range(hidden_n - 1):
+            self.layers.append(nn.Linear(hidden_dim, hidden_dim))
+        self.layers.append(nn.Linear(hidden_dim, self.output_dim))
+        # self.in_layer = nn.Linear(input_dim, hidden_dim)
+        # self.hidden_layers = [nn.Linear(hidden_dim, hidden_dim) for _ in range(hidden_n - 1)]
+        # self.out_layer = nn.Linear(hidden_dim, output_dim)
         self.activation = nn.functional.relu
-        self.final_activation = final_activation if final_activation else lambda x: x
+        self.final_activation = final_activation if final_activation else self.activation
 
     def forward(self, x):
         for layer in self.layers[:-1]:

interimit/nets/util.py

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

interimit/policies/policy.py

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

requirements.txt

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

scratch/arec/intersimple/commands.txt

@@ -1,22 +0,0 @@
-python -m render_options --model_name='gail_options_image_mid_wcollision' --env='NRasterizedRoute' --options=True --width=36 --height=36 --m_per_px=2 --agent=50 --stop_on_collision=False
-
-import torch, os
-from src.data import load_experts
-folder = 'expert_data/DR_USA_Roundabout_FT/track0000'
-single_agent = os.path.join(folder, 'expert.pkl')
-multi_agent = os.path.join(folder,'joint_expert_states.pt')
-multi_agent_actions = os.path.join(folder,'joint_expert_actions.pt')
-demonstrations = load_experts([single_agent], flatten=False)
-demonstrations[0].__dict__.keys()
-len(demonstrations[0].obs)
-single_agent_lengths = [len(demonstration.obs) for demonstration in demonstrations]
-states = torch.load(multi_agent)
-actions = torch.load(multi_agent_actions)
-multi_agent_lengths = [sum(~torch.isnan(states[:,i,0])).item() for i in range(states.shape[1])]
-
-single_agent_actions = [demonstration.acts for demonstration in demonstrations]
-multi_agent_actions = [actions[~torch.isnan(actions[:,i,0])] for i in range(actions.shape[1])]
-
-import pickle
-with open(single_agent, "rb") as f:
-    new_trajectories = pickle.load(f)

scratch/arec/intersimple/data/expert.py

@@ -1,145 +0,0 @@
-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(BasePolicy):
-
-    def __init__(self, experts):
-        self._experts = [e() for e in experts]
-
-    def forward(self, *args, **kwargs):
-        raise NotImplementedError()
-
-    def _predict(self, *args, **kwargs):
-        raise NotImplementedError()
-    
-    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 forward(self, *args, **kwargs):
-        raise NotImplementedError()
-
-    def _predict(self, *args, **kwargs):
-        raise NotImplementedError()
-
-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='NRasterizedIncrementingAgent', path=None, min_timesteps=None, min_episodes=None, video=False, env_args={}, policy_args={}):
-    """Rollout and save expert demos.
-    
-    Usage:
-        python -m intersimple.expert <flags>
-    Args: 
-        expert (class): class of expert
-        env (class): class of env intersim.envs.intersimple
-        path (str): path to store output
-        min_timesteps (int): min number of timesteps for call to rollout.rollout_and_save
-        min_episodes (int): min number of episodes for call to rollout.rollout_and_save
-        video (bool): whether to save a video of the expert until a single environment instantiation stops
-        env_args (dict): dictionary of kwargs when instantiating environment class
-        policy_args (dict): dictionary of kwargs when instantiating Expert policy
-    """
-    
-    Env = intersim.envs.intersimple.__dict__[env]
-    Expert = globals()[expert]
-
-    env = Env(**env_args)
-    info_env = RolloutInfoWrapper(env) # getting rollout info (dictionary) from environment
-    venv = DummyVecEnv([lambda: info_env]) # making a DummyVecEnv with a list of a function that when called returns the rollout info
-
-    policy = Expert(env, **policy_args) # instantiate an expert policy from specified class with instantiated environment and policy kwargs
-    venv_policy = DummyVecEnvPolicy([lambda: policy]) # make a DummyVecEnvPolicy with a list of a function that when called returns the Expert policy
-
-    if min_timesteps is None and min_episodes is None:
-        min_episodes = env.nv # one episode per vehicle being controlled in environment (hopefully an incrementing agent environment)
-
-    if video:
-        save_video(env, policy)
-    
-    path = path or (policy.__class__.__name__ + '_' + env.__class__.__name__ + '.pkl')
-    suntil = rollout.make_sample_until(
-            min_timesteps=min_timesteps,
-            min_episodes=min_episodes,
-        )
-    rollout.rollout_and_save(
-        path=path,
-        policy=venv_policy,
-        venv=venv,
-        sample_until=suntil
-    )
-
-if __name__ == '__main__':
-    import fire
-    fire.Fire(demonstrations)

scratch/arec/intersimple/data/generate.sh

@@ -1,9 +0,0 @@
-#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'
-# python -m expert --env=NRasterizedRandomAgent --min_timesteps=10000 --env_args='{width:36,height:36,m_per_px:2}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001N10000_NRasterizedRandomAgentw36h36mppx2.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'
-#python -m expert --env=NRasterized --min_timesteps=3000 --video --env_args='{width:36,height:36,m_per_px:2}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001_NRasterizedRandomAgentw36h36mppx2.pkl'
-#python -m expert --env=NRasterizedIncrementingAgent --env_args='{width:36,height:36,m_per_px:2}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001_NRasterizedIncrementingAgentw36h36mppx2.pkl'
-python -m process_all_experts --env_args='{width:36,height:36,m_per_px:2}' --policy_args='{mu:0.001}'

scratch/arec/intersimple/data/load_experts.py

@@ -1,23 +0,0 @@
-import pickle
-import imitation.data.rollout as rollout
-from tqdm import tqdm
-
-def load_experts(expert_files=[]):
-    """
-    Load expert trajectories from files and combine their transitions into a single RB
-
-    Args:
-        expert_files (list): list of expert file strings
-    Returns:
-        transitions (list): list of combined expert episode transitions 
-    """
-    transitions = []
-    for file in tqdm(expert_files):
-        with open(file, "rb") as f:
-            trajectories = pickle.load(f)
-        transitions = transitions + rollout.flatten_trajectories(trajectories)
-    return transitions
-
-if __name__=='__main__':
-    import fire
-    fire.Fire(load_experts)

scratch/arec/intersimple/data/process_all_experts.py

@@ -1,48 +0,0 @@
-import tqdm
-import expert
-import copy
-import os
-import intersim
-from tqdm import tqdm
-
-def process_all_experts(filename='expert.pkl',env_args={}, policy_args={}):
-    """
-    Process all experts in the Interaction Dataset
-    For now, using NormalizedIntersimpleExpert with NRasterizedIncrementingAgent environment
-
-    Args:
-        filename (str): name for track file
-        env_args (dict): default environment kwargs
-        policy_args (dict): default policy kwargs
-    """
-    I, J = len(intersim.LOCATIONS), intersim.MAX_TRACKS
-    pbar = tqdm(total=I*J)
-    for loc in range(I):
-        for track in range(J):
-            
-            it_env_args = copy.deepcopy(env_args)
-            it_env_args.update({
-                'loc':loc,
-                'track':track,
-            })
-            out_folder = os.path.join(intersim.LOCATIONS[loc], 'track%04i'%(track))
-            if not os.path.isdir(out_folder):
-                os.makedirs(out_folder) 
-            it_path = os.path.join(out_folder,filename)
-
-            expert.demonstrations(
-                expert='NormalizedIntersimpleExpert', 
-                env='NRasterizedIncrementingAgent', 
-                path=it_path, 
-                env_args=it_env_args, 
-                policy_args=policy_args,
-                )
-            pbar.update(1)
-    pbar.close()
-    
-
-if __name__=='__main__':
-    import fire
-    fire.Fire(process_all_experts)
-
-

scratch/arec/intersimple/gail/discriminator.py

@@ -1,101 +0,0 @@
-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)
-        assert sa.ndim == 4
-        return self.cnn(sa).squeeze(1)
-
-class CnnDiscriminatorFlatAction(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
-
-        self.cnn = torch.nn.Sequential(
-            torch.nn.Conv2d(in_channels, 32, kernel_size=(8, 8), stride=(4, 4)), # in_channels -> 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(128), # 28224 -> 128
-        )
-        self.decoder = torch.nn.Sequential(
-            torch.nn.LazyLinear(64), #128 + 2 -> 64
-            torch.nn.ReLU(),
-            torch.nn.LazyLinear(64), #64 -> 64
-            torch.nn.ReLU(),
-            torch.nn.LazyLinear(1) #64 -> 1
-        )
-    
-    @staticmethod
-    def _concatenate(state, action):
-        b, s= state.shape
-        b, a = action.shape
-        sa = torch.cat((state, action), -1)
-        return sa
-
-    def forward(self, state, action):
-        s = self.cnn(state.float())
-        sa = self._concatenate(s, action)
-        assert sa.ndim == 2
-        return self.decoder(sa).squeeze(1)
-
-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)
-        assert sa.ndim == 2
-        return self.mlp(sa).squeeze(1)

scratch/arec/intersimple/gail/test_discriminator.py

@@ -1,45 +0,0 @@
-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/arec/intersimple/gail_image_multiagent_nocollision.py

@@ -1,70 +0,0 @@
-# %%
-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 CnnDiscriminatorFlatAction
-
-model_name = 'gail_image_multiagent_nocollision'
-
-# %%
-# Load pickled test demonstrations.
-with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedRandomAgentw36h36mppx2.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={'stop_on_collision':False, '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': CnnDiscriminatorFlatAction(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(stop_on_collision=False, 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/arec/intersimple/gail_image_singleagent_nocollision.py

@@ -1,70 +0,0 @@
-# %%
-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_singleagent_nocollision'
-
-# %%
-# 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, 'stop_on_collision':False, '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, stop_on_collision=False)
-
-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/arec/intersimple/gail_options_image.py

@@ -1,172 +0,0 @@
-# %%
-import sys
-sys.path.append('../../../')
-from src.discriminator import CnnDiscriminator, CnnDiscriminatorFlatAction
-from src.policies import OptionsCnnPolicy
-from src.util import render_env
-from src.data import load_experts
-from src.gail.options import OptionsEnv, LLOptions, HLOptions, RenderOptions
-from src.gail.train import train_discriminator, train_generator
-
-from imitation.algorithms import adversarial
-from imitation.util import logger
-import imitation.data.rollout as rollout
-
-import stable_baselines3
-from stable_baselines3.common.env_util import make_vec_env
-
-import torch
-import torch.utils.data
-import numpy as np
-import itertools
-import gym
-import pickle
-import tempfile
-import pathlib
-from tqdm import tqdm
-
-from intersim.envs.intersimple import NRasterized, NRasterizedRoute, NRasterizedRandomAgent, NRasterizedIncrementingAgent, NRasterizedRouteRandomAgent
-
-ALL_OPTIONS = [(v,t) for v in [0,2,4,6,8] for t in [5, 10]] # option 0 is safe fallback
-
-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(expert_data, env_class=NRasterizedRouteRandomAgent, env_settings={}, 
-        epochs=10, discrim_batch_size=32, generator_steps=2048, discount=0.99):
-    """
-    Args: 
-        expert_data: list of transitions
-        env_class: environment class
-        env_settings: environment settings
-        epochs: number of epochs to train for
-        discrim_batch_size: discriminator batch size
-        generator_steps: number of steps taken in generator
-        discount: discount factor
-    Returns:
-        generator (stable_baselines3.PPO): options policy
-    """
-    env = env_class(**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(env_class, n_envs=1, env_kwargs=env_settings)
-    discriminator = adversarial.GAIL(
-        expert_data=expert_data,
-        expert_batch_size=discrim_batch_size,
-        discrim_kwargs={'discrim_net': CnnDiscriminatorFlatAction(venv)},
-        #discrim_kwargs={'discrim_net': CnnDiscriminator(venv)},
-        venv=venv, # unused
-        gen_algo=stable_baselines3.PPO("CnnPolicy", venv), # unused
-    )
-
-    generator = stable_baselines3.PPO(
-        OptionsCnnPolicy,
-        OptionsEnv(env, options=ALL_OPTIONS),
-        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 tqdm(range(epochs)):
-        train_discriminator(LLOptions(env, options=ALL_OPTIONS), generator, discriminator, num_samples=discrim_batch_size)
-        train_generator(HLOptions(env, options=ALL_OPTIONS), generator, discriminator, num_samples=generator_steps)
-    
-    return generator
-
-# %%
-if __name__ == '__main__':
-    # %%
-    model_name = 'gail_options_image_mid_wcollision'
-    env_class = NRasterizedRouteRandomAgent
-    env_settings = {'width': 36, 'height': 36, 'm_per_px': 2, 'stop_on_collision': False}
-    
-    #env_class = NRasterized
-    #env_settings = {'agent': 51, 'width': 36, 'height': 36, 'm_per_px': 2}
-    files = ['../../../expert_data/DR_USA_Roundabout_FT/track%04i/expert.pkl'%(i) for i in range(5)]
-    transitions=load_experts(files)
-
-    generator = train(
-        transitions,
-        env_class=env_class,
-        env_settings=env_settings,
-        epochs=2, 
-        discrim_batch_size=256, 
-        generator_steps=10,#256, 
-        discount=0.99
-        )
-
-    generator.save(model_name)
-
-    # Render
-    render_settings = {'width': 36, 'height': 36, 'm_per_px': 2, 'agent':51, 'stop_on_collision': False}
-    render_env(model_name=model_name, env='NRasterizedRoute', options=True, options_list=ALL_OPTIONS,
-        **render_settings)
-
-
-# %% Tests
-
-def test_ll_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 = LLOptions(NRasterized(agent=51, width=36, height=36, m_per_px=2))
-
-    gen_transitions = list(itertools.islice(env.sample_ll(
-        policy=stable_baselines3.PPO(
-            OptionsCnnPolicy,
-            OptionsEnv(env),
-            verbose=1,
-        )
-    ), 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_ll_states():
-    env = NRasterized()
-    policy = stable_baselines3.PPO(
-        OptionsCnnPolicy,
-        OptionsEnv(env),
-        verbose=1,
-    )
-    llenv = LLOptions(env)
-    transitions = list(itertools.islice(llenv.sample_ll(policy=policy), 100))
-
-    env2 = NRasterized()
-    s2 = env2.reset()
-    for i, t in enumerate(transitions):
-        assert i == 0 or np.array_equal(t['obs'], transitions[i-1]['next_obs'])
-        assert np.array_equal(t['obs'], s2)
-        assert t['acts'].shape == (1,)
-
-        nexts2, _, done2, _ = env2.step(t['acts'])
-        assert np.array_equal(t['next_obs'], nexts2)
-        assert np.array_equal(t['dones'], done2)
-
-        if done2:
-            break
-
-        s2 = nexts2
-
-def test_hl_transitions():
-    pass

scratch/arec/intersimple/gail_options_scratch.py

@@ -1,559 +0,0 @@
-# %%
-from gail.discriminator import CnnDiscriminator, CnnDiscriminatorFlatAction
-from imitation.algorithms import adversarial
-import stable_baselines3
-import torch.utils.data
-import numpy as np
-from intersim.envs.intersimple import NRasterized, NRasterizedRandomAgent
-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
-from tqdm import tqdm
-
-import logging 
-logging.basicConfig(level=logging.DEBUG)
-
-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):
-    """
-    Class for high-level options policy (generator)
-    """
-    def __init__(self, observation_space, *args, **kwargs):
-        super().__init__(observation_space['obs'], *args, **kwargs)
-
-    def _prior_distribution(self, s):
-        """
-        Return prior distribution over high-level options (before masking)
-        Args:
-            s (torch.tensor): observation
-        Returns:
-            values (torch.tensor): values from critic
-            dist (torch.distributions): prior distribution over actions
-        """
-        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):
-        """
-        Will mask invalid states before making action selections
-        Args:
-            obs: dict with keys:
-                obs (torch.tensor): (B,o)  true observations
-                mask (torch.tensor): (B,m) mask over valid actions
-        Returns:
-            ch (torch.tensor): (B,a) sampled actions
-            values (torch.tensor): (B,) predicted value at observation
-            log_probs (torch.tensor): (B,) log probabilities of selected actions
-        """
-        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):
-        """
-        Evaluate particular actions
-        Args:
-            obs: dict with keys:
-                obs (torch.tensor): (B,o) true observations
-                mask (torch.tensor): (B,m) masks over valid actions
-            ch (torch.tensor): (B,a) selected actions
-        Returns:
-            values (torch.tensor): (B,) predicted value at observation
-            log_probs (torch.tensor): (B,) log probabilities of selected actions
-            ent (torch.tensor): (B,) entropy of each distribution over actions
-        """
-        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
-
-class OptionsEnv(gym.Wrapper):
-    """
-    Wrap an intersimple environment with an options generator
-    """
-    def __init__(self, env, *args, **kwargs):
-        """
-        Initialize wrapped environment and set high-level action and observation spaces
-        """
-        super().__init__(env, *args, **kwargs)
-        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 _after_choice(self):
-        pass
-    
-    def _after_step(self):
-        pass
-
-    def _transitions(self):
-        raise NotImplementedError('Use `LLOptions` or `HLOptions` for sampling.')
-    
-    def sample(self, generator):
-        """
-        yield transitions using a generator
-        Args:
-            generator (sb3.PPO)
-        Yields:
-
-        """
-        self.done = True
-        while True:
-            self.episode_start = False
-
-            if self.done:
-                # reset environment
-                self.s = self.env.reset()
-                self.m = available_actions(self.env)
-                self.done = False
-                self.episode_start = True
-
-            # set the action, the value of the start state, and the logprob of the action
-            # according to the current environment state and mask
-            self.ch, self.value, self.log_prob = generator.policy.predict({
-                'obs': torch.tensor(self.s).unsqueeze(0).to(generator.policy.device),
-                'mask': torch.tensor(self.m).unsqueeze(0).to(generator.policy.device),
-            })
-
-            # store a float list of actions to take given the option selected in the environment
-            self.plan = list(map(float, generate_plan(self.env, self.ch)))
-
-            # run whatever _after_choice might dictate in a child class
-            self._after_choice()
-
-            # some checks
-            assert not self.done
-            assert self.plan
-            assert feasible(self.env, self.plan, self.ch)
-
-            # execute the option so long as the episode isn't complete and the plan is still feasible
-            while not self.done and self.plan and feasible(self.env, self.plan, self.ch):
-                
-                # pop first action
-                self.a, self.plan = self.plan[0], self.plan[1:]
-                
-                # normalize action ??
-                self.a = self.env._normalize(self.a)
-                
-                # step through environment
-                self.nexts, _, self.done, _ = self.env.step(self.a)
-                self.nextm = available_actions(self.env)
-
-                # run whatever _after_step might dictate in child class
-                self._after_step()
-
-                # update state and mask to current
-                self.s = self.nexts
-                self.m = self.nextm
-
-            # transitions yielded from self._transitions() functions specied in child classes
-            yield from self._transitions()
-
-            ### NOTE: only yields after a full option has been executed / exited
-
-class LLOptions(OptionsEnv):
-    """Sample low-level (state, action) tuples for discriminator training."""
-    
-    def __init__(self, *args, **kwargs):
-        """
-        LLOption uses the true LL observations
-        """
-        super().__init__(*args, **kwargs)
-        # overwrite observation space to just output obs directly
-        self.observation_space = self.observation_space['obs'] 
-
-    def _after_choice(self):
-        """
-        After each option choice, initialize/reset the transition buffer
-        """
-        self._transition_buffer = []
-
-    def _after_step(self):
-        """
-        After each ll action, append s, s', a, done to transition buffer
-        """
-        self._transition_buffer.append({
-            'obs': self.s,
-            'next_obs': self.nexts,
-            'acts': np.array((self.a,)),
-            'dones': np.array(self.done),
-        })
-
-    def _transitions(self):
-        """
-        Yield from the transition buffer
-        """
-        yield from self._transition_buffer
-    
-    def sample_ll(self, policy):
-        """
-        Args:
-            policy
-        Returns:
-            gen: iterable which samples low-level transitions from the environment
-        """
-        return self.sample(policy)
-
-class HLOptions(OptionsEnv):
-    """Sample high-level (state, action, reward) tuples for generator training."""
-    
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-    def _after_choice(self):
-        """
-        After an option selection, initialize total reward and number of steps
-        """
-        self.r = 0
-        self.steps = 0
-
-    def _after_step(self):
-        """
-        After each low-level action, add the discounted discriminated reward score (given a discriminator)
-        """
-        self.r += self.discount**self.steps * self.discriminator.discrim_net.reward_train(
-            state=torch.tensor(self.s).unsqueeze(0).to(self.discriminator.discrim_net.device()),
-            action=torch.tensor([[self.a]]).to(self.discriminator.discrim_net.device()),
-            next_state=torch.tensor(self.s).unsqueeze(0).to(self.discriminator.discrim_net.device()), # unused
-            done=torch.tensor(self.done).unsqueeze(0).to(self.discriminator.discrim_net.device()), # unused
-        )
-        self.steps += 1
-    
-    def _transitions(self):
-        """
-        Yield a single dictionary per high-level selected action
-        Fields:
-            obs: high-level state and mask at selection
-            action: chosen high-level action
-            reward: accumulated option reward
-            episode_start: whether the action was chosen at the episode start
-            value: the value estimate from the starting state
-            log_prob: the log_prob of the selected action from the starting state
-            done: whether the episode has ended
-
-        """
-        yield {
-            'obs': {'obs': self.s, 'mask': self.m},
-            'action': self.ch,
-            'reward': self.r.detach(),
-            'episode_start': self.episode_start,
-            'value': self.value.detach(),
-            'log_prob': self.log_prob.detach(),
-            'done': self.done,
-        }
-    
-    def sample_hl(self, policy, discriminator):
-        """
-        Args:
-            policy
-            discriminator: function with which to score rewards
-        Returns:
-            gen: iterable which samples high-level transitions from the environment
-        """
-        self.discriminator = discriminator
-        return self.sample(policy)
-
-class RenderOptions(LLOptions):
-
-    def _after_step(self):
-        """
-        Render the environment after each low-level step
-        """
-        super()._after_step()
-        self.env.render()
-    
-    def close(self, *args, **kwargs):
-        """
-        On 'close', close the environment
-        """
-        self.env.close(*args, **kwargs)
-
-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 check_future_collisions_circles(env, actions, n_circles:int=2):
-    """Checks whether `env._agent` would collide with other agents assuming `actions` as input.
-    
-    Vehicles are (over-)approximated by multiple 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
-    """
-    assert n_circles >= 2
-    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)
-    centers = states[:, :, :, :2]
-    psi = states[:, :, :, 3]
-    lon = torch.stack([psi.cos(), psi.sin()],dim=-1) # (B, T, nv, 2)
-
-    # offset between [-env._env.lengths+env._env.widths/2, env._env.lengths/2-env._env.widths/2]
-    back = (-env._env._lengths/2+env._env._widths/2).unsqueeze(-1) # (nv, 1)
-    length = (env._env._lengths-env._env._widths).unsqueeze(-1) # (nv, 1)
-    diff_d =  back + length*(torch.arange(n_circles)/(n_circles-1)).unsqueeze(0)  # (nv, n_circles) 
-    assert diff_d.shape == (nv, n_circles)
-    
-    offsets = diff_d[None, None, :, :, None] *  lon[:, :, :, None, :]
-    assert offsets.shape == (B, T, nv, n_circles, 2)
-
-    expanded_centers=centers.unsqueeze(-2) + offsets #(B, T, nv, n_circles, 2)
-    assert expanded_centers.shape == (B, T, nv, n_circles, 2)
-    agent_centers = expanded_centers[:,:,env._agent:env._agent+1,:,:] #(B, T, 1, n_circles, 2)
-    ds = expanded_centers.reshape((B, T, nv*n_circles, 1, 2)) - agent_centers #(B, T, nv*nc,1, 2) - (B, T, 1, nc, 2) = (B, T, nv*nc, nc, 2) 
-
-    distance = (ds**2).sum(-1).sqrt().reshape((B, T, nv, n_circles, n_circles)) # (B, T, nv, nc, nc)
-    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, n_circles, n_circles)
-
-    radius = env._env._widths*np.sqrt(2) / 2
-    min_distance = radius[env._agent] + radius
-    min_distance = min_distance[None, None, :, None, None]
-    assert min_distance.shape == (1, 1, nv, 1, 1)
-
-    return (distance > min_distance).all(-1).all(-1).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
-    valid = check_future_collisions_circles(env, [full_plan])
-    return ch == 0 or valid.item()
-
-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(env.sample_ll(generator), 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(env.sample_hl(generator, discriminator), num_samples+1))
-    
-    generator.rollout_buffer.reset()
-    for s in generator_samples[:-1]:
-        generator.rollout_buffer.add(
-            obs=s['obs'],
-            action=s['action'].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()
-
-def train(expert_data, env_class=NRasterizedRandomAgent, env_settings={}, epochs=10, discrim_batch_size=32, generator_steps=2048, discount=0.99):
-    """
-    Args: 
-        expert_data: list of transitions
-        env_class: environment class
-        env_settings: environment settings
-        epochs: number of epochs to train for
-        discrim_batch_size: discriminator batch size
-        generator_steps: number of steps taken in generator
-        discount: discount factor
-    Returns:
-        generator (stable_baselines3.PPO): options policy
-    """
-    env = env_class(**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(env_class, n_envs=1, env_kwargs=env_settings)
-    discriminator = adversarial.GAIL(
-        expert_data=expert_data,
-        expert_batch_size=discrim_batch_size,
-        discrim_kwargs={'discrim_net': CnnDiscriminatorFlatAction(venv)},
-        #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 tqdm(range(epochs)):
-        train_discriminator(LLOptions(env), generator, discriminator, num_samples=discrim_batch_size)
-        train_generator(HLOptions(env), generator, discriminator, num_samples=generator_steps)
-    
-    return generator
-
-# %%
-if __name__ == '__main__':
-    # %%
-    model_name = 'gail_options_image'
-    env_class = NRasterizedRandomAgent
-    env_settings = {'width': 36, 'height': 36, 'm_per_px': 2}
-
-    with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedIncrementingAgentw36h36mppx2.pkl", "rb") as f:
-        trajectories = pickle.load(f)
-    #import pdb 
-    #pdb.set_trace()
-    transitions = rollout.flatten_trajectories(trajectories)
-    generator = train(
-        transitions,
-        env_class=env_class,
-        env_settings=env_settings,
-        epochs=2, 
-        discrim_batch_size=32, 
-        generator_steps=2048, 
-        discount=0.99
-        )
-
-    generator.save(model_name) # save ppo sb3 generator class
-
-    # %%
-    model = stable_baselines3.PPO.load(model_name) # not actually used
-
-    env = RenderOptions(NRasterizedRandomAgent(**env_settings))
-    for s in env.sample_ll(generator):
-        if s['dones']:
-            break
-
-    env.close(filestr='render/'+model_name)
-
-# %% Tests
-
-def test_ll_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 = LLOptions(NRasterized(agent=51, width=36, height=36, m_per_px=2))
-
-    gen_transitions = list(itertools.islice(env.sample_ll(
-        policy=stable_baselines3.PPO(
-            OptionsCnnPolicy,
-            OptionsEnv(env),
-            verbose=1,
-        )
-    ), 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_ll_states():
-    env = NRasterized()
-    policy = stable_baselines3.PPO(
-        OptionsCnnPolicy,
-        OptionsEnv(env),
-        verbose=1,
-    )
-    llenv = LLOptions(env)
-    transitions = list(itertools.islice(llenv.sample_ll(policy=policy), 100))
-
-    env2 = NRasterized()
-    s2 = env2.reset()
-    for i, t in enumerate(transitions):
-        assert i == 0 or np.array_equal(t['obs'], transitions[i-1]['next_obs'])
-        assert np.array_equal(t['obs'], s2)
-        assert t['acts'].shape == (1,)
-
-        nexts2, _, done2, _ = env2.step(t['acts'])
-        assert np.array_equal(t['next_obs'], nexts2)
-        assert np.array_equal(t['dones'], done2)
-
-        if done2:
-            break
-
-        s2 = nexts2
-
-def test_hl_transitions():
-    pass

scratch/arec/intersimple/options_gail.py

@@ -1,510 +0,0 @@
-# %%
-from gail.discriminator import CnnDiscriminator, CnnDiscriminatorFlatAction
-from imitation.algorithms import adversarial
-import stable_baselines3
-import torch.utils.data
-import numpy as np
-from intersim.envs.intersimple import NRasterized, NRasterizedRandomAgent
-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
-from tqdm import tqdm
-
-import logging 
-logging.basicConfig(level=logging.DEBUG)
-
-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):
-    """
-    Class for high-level options policy (generator)
-    """
-    def __init__(self, observation_space, *args, **kwargs):
-        super().__init__(observation_space['obs'], *args, **kwargs)
-
-    def _prior_distribution(self, s):
-        """
-        Return prior distribution over high-level options (before masking)
-        Args:
-            s (torch.tensor): observation
-        Returns:
-            values (torch.tensor): values from critic
-            dist (torch.distributions): prior distribution over actions
-        """
-        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):
-        """
-        Will mask invalid states before making action selections
-        Args:
-            obs: dict with keys:
-                obs (torch.tensor): (B,o)  true observations
-                mask (torch.tensor): (B,m) mask over valid actions
-        Returns:
-            ch (torch.tensor): (B,a) sampled actions
-            values (torch.tensor): (B,) predicted value at observation
-            log_probs (torch.tensor): (B,) log probabilities of selected actions
-        """
-        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):
-        """
-        Evaluate particular actions
-        Args:
-            obs: dict with keys:
-                obs (torch.tensor): (B,o) true observations
-                mask (torch.tensor): (B,m) masks over valid actions
-            ch (torch.tensor): (B,a) selected actions
-        Returns:
-            values (torch.tensor): (B,) predicted value at observation
-            log_probs (torch.tensor): (B,) log probabilities of selected actions
-            ent (torch.tensor): (B,) entropy of each distribution over actions
-        """
-        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
-
-class OptionsEnv(gym.Wrapper):
-    """
-    Wrap an intersimple environment with an options generator
-    """
-    def __init__(self, env, render=False, *args, **kwargs):
-        """
-        Initialize wrapped environment and set high-level action and observation spaces
-        """
-        super().__init__(env, *args, **kwargs)
-        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,)),
-        })
-        self._hl_transition_buffer = []
-        self._ll_transition_buffer = []
-        self.render=render
-
-    def _after_option_choice(self):
-        """
-        After initial option choice, 
-        """
-        self._hl_r = 0
-        self._hl_steps = 0
-    
-    def _after_step(self):
-        """
-        After each step, add the ll transition to the appropriate buffer, add to reward, add to steps, and possibly render
-        """
-
-        self._ll_transition_buffer.append({
-            'obs': self.s,
-            'next_obs': self.nexts,
-            'acts': np.array((self.a,)),
-            'dones': np.array(self.done),
-        })  
-        self.r += self.discount**self.steps * self.discriminator.discrim_net.reward_train(
-            state=torch.tensor(self.s).unsqueeze(0).to(self.discriminator.discrim_net.device()),
-            action=torch.tensor([[self.a]]).to(self.discriminator.discrim_net.device()),
-            next_state=torch.tensor(self.s).unsqueeze(0).to(self.discriminator.discrim_net.device()), # unused
-            done=torch.tensor(self.done).unsqueeze(0).to(self.discriminator.discrim_net.device()), # unused
-        )
-        self.steps += 1
-        if self.render:
-            self.env.render()
-    
-    def _after_option(self):
-        """
-        After each low-level action, add the discounted discriminated reward score (given a discriminator)
-        """
-        self._hl_transition_buffer.append({
-            'obs': {'obs': self.os, 'mask': self.m},
-            'action': self.ch,
-            'reward': self.r.detach(),
-            'episode_start': self.episode_start,
-            'value': self.value.detach(),
-            'log_prob': self.log_prob.detach(),
-            'done': self.done,
-            })
-
-    def close(self, *args, **kwargs):
-        """
-        On 'close', close the environment
-        """
-        self.env.close(*args, **kwargs)
-    
-    def sample(self, generator, controller):
-        """
-        yield transitions using a generator
-        Args:
-            generator (sb3.PPO)
-            controller (str): 'high' or 'low' to yield from proper buffer
-        Yields:
-
-        """
-        self.done = True
-        # DO I WANT TO EMPTY THE BUFFERS??? Probs naw
-        while True:
-            
-            # yield from buffers to empty what was stored previously
-            if controller = 'high':
-                yield from self._hl_transition_buffer
-            elif controller == 'low':
-                yield from self._ll_transition_buffer
-            else:
-                raise('Improper buffer')
-            
-            self.episode_start = False
-            if self.done:
-                # reset environment
-                self.s = self.env.reset()
-                self.done = False
-                self.episode_start = True
-
-            self.os = self.s.copy() # option start state
-            self.m = available_actions(self.env)
-
-            # set the action, the value of the start state, and the logprob of the action
-            # according to the current environment state and mask
-            self.ch, self.value, self.log_prob = generator.policy.predict({
-                'obs': torch.tensor(self.os).unsqueeze(0).to(generator.policy.device),
-                'mask': torch.tensor(self.m).unsqueeze(0).to(generator.policy.device),
-            })
-
-            # store a float list of actions to take given the option selected in the environment
-            self.plan = list(map(float, generate_plan(self.env, self.ch)))
-
-            # run whatever _after_choice might dictate in a child class
-            self._after_option_choice()
-
-            # some checks
-            assert not self.done
-            assert self.plan
-            assert feasible(self.env, self.plan, self.ch)
-
-            # execute the option so long as the episode isn't complete and the plan is still feasible
-            while not self.done and self.plan and feasible(self.env, self.plan, self.ch):
-                
-                # pop first action
-                self.a, self.plan = self.plan[0], self.plan[1:]
-                
-                # normalize action ??
-                self.a = self.env._normalize(self.a)
-                
-                # step through environment
-                self.nexts, _, self.done, _ = self.env.step(self.a)
-
-                # run whatever _after_step might dictate in child class
-                self._after_step()
-
-                # update state and mask to current
-                self.s = self.nexts
-            
-            # run whatever to do after option
-            self._after_option()
-
-    def sample_ll(self, policy):
-        """
-        Not quite sure how this works????
-        Why would you do this over LLOptions.sample(policy)
-        """
-        return self.sample(policy, 'low')
-
-    def sample_hl(self, policy, discriminator):
-        """
-        Args:
-            policy
-            discriminator: function with which to score rewards
-        Returns:
-            gen: an which samples high-level transitions from the environment
-        """
-        self.discriminator = discriminator
-        return self.sample(policy)
-
-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 check_future_collisions_circles(env, actions, n_circles:int=2):
-    """Checks whether `env._agent` would collide with other agents assuming `actions` as input.
-    
-    Vehicles are (over-)approximated by multiple 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
-    """
-    assert n_circles >= 2
-    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)
-    centers = states[:, :, :, :2]
-    psi = states[:, :, :, 3]
-    lon = torch.stack([psi.cos(), psi.sin()],dim=-1) # (B, T, nv, 2)
-
-    # offset between [-env._env.lengths+env._env.widths/2, env._env.lengths/2-env._env.widths/2]
-    back = (-env._env._lengths/2+env._env._widths/2).unsqueeze(-1) # (nv, 1)
-    length = (env._env._lengths-env._env._widths).unsqueeze(-1) # (nv, 1)
-    diff_d =  back + length*(torch.arange(n_circles)/(n_circles-1)).unsqueeze(0)  # (nv, n_circles) 
-    assert diff_d.shape == (nv, n_circles)
-    
-    offsets = diff_d[None, None, :, :, None] *  lon[:, :, :, None, :]
-    assert offsets.shape == (B, T, nv, n_circles, 2)
-
-    expanded_centers=centers.unsqueeze(-2) + offsets #(B, T, nv, n_circles, 2)
-    assert expanded_centers.shape == (B, T, nv, n_circles, 2)
-    agent_centers = expanded_centers[:,:,env._agent:env._agent+1,:,:] #(B, T, 1, n_circles, 2)
-    ds = expanded_centers.reshape((B, T, nv*n_circles, 1, 2)) - agent_centers #(B, T, nv*nc,1, 2) - (B, T, 1, nc, 2) = (B, T, nv*nc, nc, 2) 
-
-    distance = (ds**2).sum(-1).sqrt().reshape((B, T, nv, n_circles, n_circles)) # (B, T, nv, nc, nc)
-    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, n_circles, n_circles)
-
-    radius = env._env._widths*np.sqrt(2) / 2
-    min_distance = radius[env._agent] + radius
-    min_distance = min_distance[None, None, :, None, None]
-    assert min_distance.shape == (1, 1, nv, 1, 1)
-
-    return (distance > min_distance).all(-1).all(-1).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
-    valid = check_future_collisions_circles(env, [full_plan])
-    return ch == 0 or valid.item()
-
-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(env.sample_ll(generator), 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(env.sample_hl(generator, discriminator), num_samples+1))
-    
-    generator.rollout_buffer.reset()
-    for s in generator_samples[:-1]:
-        generator.rollout_buffer.add(
-            obs=s['obs'],
-            action=s['action'].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()
-
-def train(expert_data, env_class=NRasterizedRandomAgent, env_settings={}, epochs=10, discrim_batch_size=32, generator_steps=2048, discount=0.99):
-    """
-    Args: 
-        expert_data: list of transitions
-        env_class: environment class
-        env_settings: environment settings
-        epochs: number of epochs to train for
-        discrim_batch_size: discriminator batch size
-        generator_steps: number of steps taken in generator
-        discount: discount factor
-    Returns:
-        generator (stable_baselines3.PPO): options policy
-    """
-    env = env_class(**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(env_class, n_envs=1, env_kwargs=env_settings)
-    discriminator = adversarial.GAIL(
-        expert_data=expert_data,
-        expert_batch_size=discrim_batch_size,
-        discrim_kwargs={'discrim_net': CnnDiscriminatorFlatAction(venv)},
-        #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 tqdm(range(epochs)):
-        train_discriminator(LLOptions(env), generator, discriminator, num_samples=discrim_batch_size)
-        train_generator(HLOptions(env), generator, discriminator, num_samples=generator_steps)
-    
-    return generator
-
-# %%
-if __name__ == '__main__':
-    # %%
-    model_name = 'gail_options_image'
-    env_class = NRasterizedRandomAgent
-    env_settings = {'width': 36, 'height': 36, 'm_per_px': 2}
-
-    with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedIncrementingAgentw36h36mppx2.pkl", "rb") as f:
-        trajectories = pickle.load(f)
-    #import pdb 
-    #pdb.set_trace()
-    transitions = rollout.flatten_trajectories(trajectories)
-    generator = train(
-        transitions,
-        env_class=env_class,
-        env_settings=env_settings,
-        epochs=2, 
-        discrim_batch_size=32, 
-        generator_steps=2048, 
-        discount=0.99
-        )
-
-    generator.save(model_name) # save ppo sb3 generator class
-
-    # %%
-    model = stable_baselines3.PPO.load(model_name) # not actually used
-
-    env = OptionsGail(NRasterizedRandomAgent(**env_settings), render=True)
-    for s in env.sample_ll(generator):
-        if s['dones']:
-            break
-
-    env.close(filestr='render/'+model_name)
-
-# %% Tests
-
-def test_ll_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 = LLOptions(NRasterized(agent=51, width=36, height=36, m_per_px=2))
-
-    gen_transitions = list(itertools.islice(env.sample_ll(
-        policy=stable_baselines3.PPO(
-            OptionsCnnPolicy,
-            OptionsEnv(env),
-            verbose=1,
-        )
-    ), 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_ll_states():
-    env = NRasterized()
-    policy = stable_baselines3.PPO(
-        OptionsCnnPolicy,
-        OptionsEnv(env),
-        verbose=1,
-    )
-    llenv = LLOptions(env)
-    transitions = list(itertools.islice(llenv.sample_ll(policy=policy), 100))
-
-    env2 = NRasterized()
-    s2 = env2.reset()
-    for i, t in enumerate(transitions):
-        assert i == 0 or np.array_equal(t['obs'], transitions[i-1]['next_obs'])
-        assert np.array_equal(t['obs'], s2)
-        assert t['acts'].shape == (1,)
-
-        nexts2, _, done2, _ = env2.step(t['acts'])
-        assert np.array_equal(t['next_obs'], nexts2)
-        assert np.array_equal(t['dones'], done2)
-
-        if done2:
-            break
-
-        s2 = nexts2
-
-def test_hl_transitions():
-    pass

scratch/arec/intersimple/plan.txt

@@ -1,55 +0,0 @@
-Environment 
--- each 'environment' follows a single roundabout and track id (recording of that roundabout)
--- on reset, the environment we will use changes the vehicle to control while having the other agents follow their true data (expert controller)
----- Note this can be problematic as it can lead to vehicles behind you crashing into you
-
-TRAINING
----------
-1. Load pre-trained massive set of transitions
--- For all roundabouts
-    -- For all tracks
-        -- For all vehicles
-            -- For all valid timesteps
-                -- Rasterized state (incl. path), action
-
-2. HGAIL
--- For each epoch
-    -- INSTANTIATE A NEW ENVIRONMENT (Roundabout + Track) w/ randomized agent, from set of all expert environments
-        -- Train discriminator off training data + yielded low-level transitions in replay buffer
-        -- Train generator off yielded high-level transitions + summed low-level discriminator rewards
-
-TESTING
-----------
-1. Save average vehicle velocities for all expert vehicles (loop roundabout + track + vehicle, average over time)
-
-2. Run test suite for: expert, BC, GAIL, RAIL, HGAIL, (and hopefully HRAIL)
--- For all roundabouts, tracks
-    -- Get expert velocities for track
-    -- Simulate incrementing agent environment (e.g. on reset, agent +=1)
-        -- Store low-level true joint states, actions, and controlled vehicle index
-    -- Per-vehicle statistics (v_all, v_mean, v_shortfall, a_all, jerk_all, n_collisions, T) 
--- Aggregate statistics + joint 
-
-Problems
------------
-Should train without stopping for collisions, however when doing so, end up with policy that always takes decelerate option
--- It seems safe at the start of each vehicles sim, but actually it isn't since a car will spawn and hit it
-Solutions:
-    -- Hold cars from spawning if their spawn location is full
-    -- Start simulations a few seconds later (after cars clear their spawn places) <- Preferred
-    
-Test could run indefinitely if stop_on_collision is off
-Solution:
-    -- Set maximum episode length in intersimple
-
-
-
-
-
-
-
-
-Save massive set of transition raw states beforehand (1 from training, but with raw states)
-# -- For all roundabouts, tracks
-#     -- For all vehicles, steps
-#         -- Raw vehicle state, action

scratch/arec/intersimple/render_env_from_model.py

@@ -1,58 +0,0 @@
-
-import stable_baselines3 as sb3
-from intersim.envs.intersimple import NRasterized
-
-
-def render_env(model_name='gail_image_multiagent_nocollision', agent=51, environment=NRasterized):
-    """
-    Render a video from an model, agent, and environment
-    Args:
-        model_name (str): name of the model
-        agent (int): agent to start the video from
-        environment (gym.Env): gym environment class to render environment on
-    """
-
-    model = sb3.PPO.load(model_name)
-
-    env = environment(stop_on_collision=False, width=36, height=36, m_per_px=2, agent=agent)
-
-    obs = env.reset()
-    i=0
-    while True and i < 600:
-        i+=1
-        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+'_agent%i'%(agent)) 
-
-def render_options_env(model_name='gail_image_multiagent_nocollision', agent=51, environment=NRasterized):
-    """
-    Render a video from an model, agent, and environment
-    Args:
-        model_name (str): name of the model
-        agent (int): agent to start the video from
-        environment (gym.Env): gym environment class to render environment on
-    """
-
-    model = sb3.PPO.load(model_name)
-
-    env = environment(stop_on_collision=False, width=36, height=36, m_per_px=2, agent=agent)
-
-    obs = env.reset()
-    i=0
-    while True and i < 600:
-        i+=1
-        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+'_agent%i'%(agent)) 
-
-if __name__ == '__main__':
-    import fire
-    fire.Fire(render_env)

scratch/arec/intersimple/render_options.py

@@ -1,11 +0,0 @@
-import sys
-sys.path.append('../../../')
-from src.util import render_env
-ALL_OPTIONS = [(v,t) for v in [0,2,4,6,8] for t in [5, 10]]
-
-def render_wrapper(**kwargs):
-    render_env(**kwargs, options_list=ALL_OPTIONS)
-
-if __name__=='__main__':
-    import fire
-    fire.Fire(render_wrapper)

scratch/arec/intersimple/test_model.py

@@ -1,148 +0,0 @@
-from tqdm import tqdm
-from copy import deepcopy
-import stable_baselines3 as sb3
-import intersim
-
-ALL_OPTIONS = [(v,t) for v in [0,2,4,6,8] for t in [5, 10]] # option 0 is safe fallback
-
-def load_model(model_path:str, method:str):
-    """
-    Load a model given a path and the method
-    
-    Args:
-        model_path (str): the path to the model
-        method (str): the method for the model
-    Returns:
-        model: the action model
-        is_heir (bool): whether the method is heirarchial
-    """
-    model = None
-    is_heir = False
-    if method == 'expert':
-        raise NotImplementedError
-    elif method == 'bc':
-        raise NotImplementedError
-    elif method == 'gail':
-        raise NotImplementedError
-    elif method == 'rail':
-        raise NotImplementedError
-    elif method == 'hgail':
-        is_heir = True
-        model = sb3.PPO.load(model_path)
-    elif method == 'hrail':
-        is_heir = True
-        raise NotImplementedError
-    else:
-        raise NotImplementedError
-    return model, is_heir
-
-def load_expert_states(roundabout, track):
-    """
-    Load expert states from roundabout/track info
-    Args:
-        roundabout (str): roundabout name
-        track (str): track id
-    Returns:
-        states (torch.tensor): (T+1, nv, 5) expert states for track file
-        actions (torch.tensor): (T, nv, 1) expert actions for track file
-    """
-    state_path = '../../../expert_data/%s/track%04i/joint_expert_states.pt'%(roundabout, track)] #FIXME when moving
-    action_path = '../../../expert_data/%s/track%04i/joint_expert_actions.pt'%(roundabout, track)] #FIXME when moving
-    states = torch.load(path)
-    actions = torch.load(path)
-    # nanify actions where vehicle's don't exist
-    import pdb
-    pdb.set_trace()
-    return states, actions
-
-def test_model(
-    locations=[(0,0)], 
-    model_name='gail_image_multiagent_nocollision', 
-    env='NRasterizedRouteIncrementingAgent', 
-    method='expert',
-    options_list=ALL_OPTIONS, 
-    **env_kwargs):
-    """
-    Test a particular model at different locations/tracks
-
-    Args:
-        locations (list of tuples): list of (roundabout, track) integer pairs
-        model_name (str): name of model to test
-        env (str): environment class
-        method (str): method (expert, bc, gail, rail, hgail, hrail)
-        options_list (list): list of options
-    """
-
-    # load policy
-    policy, is_heir = load_model(model_name, method)
-
-    # iterate through vehicles
-    all_vehicle_infos = []
-    for i, location in tqdm(enumerate(locations)):
-        
-        # add roundabout and track to environent 
-        roundabout, track = location
-        iround = intersim.LOCATIONS.index(roundabout)
-        it_env_kwargs = deepcopy(env_kwargs)
-        loc_kwargs = {
-            'loc':iround,
-            'track':track
-        }
-        it_env_kwargs.update(loc_kwargs)
-        
-        # load expert states and get average velocities
-        expert_states, expert_actions = load_expert_states(roundabout, track)
-        expert_vavg = torch.nanmean(expert_states[:,:,3], dim=-1)
-
-        # initialize environment
-        if not is_heir:
-            Env = src.options.envs.__dict__[env]
-        else: 
-            Env = intersim.envs.intersimple.__dict__[env]
-        env = Env(**env_kwargs)
-        s = env.reset()
-
-        # Iterate through every vehicle and time
-        vehicle_infos, done = [], False
-        for iv in range(env.nv):
-            v_number = env.agent
-            i_vehicle_infos = {'s':[], 'a':[], 'it':[]}
-            while not done:
-                a = policy(s)
-                sp, r, done, info   = env.step(a)
-                i_vehicle_infos['s'].append(env._env.state) # FIX
-                i_vehicle_infos['a'].append(a) 
-                i_vehicle_infos['it'].append(env._env.it) # FIX
-            i_vehicle_info.update({
-                'vehicle_id': env.agent,
-                'n_steps': len(i_vehicle_infos['a']),
-                'T': len(i_vehicle_infos['a'])*env._env.dt, # FIX
-                'n_collisions': collision.check(i_vehicle_infos['s'], env._env.lengths. env._env.widths), # FIX
-                'expert_vavg': expert_vavg[env.agent]
-            })
-            vehicle_infos.append(i_vehicle_info)
-            env.reset()
-
-        all_vehicle_infos.append({
-            'loc': location,
-            'track': track,
-            'stats': vehicle_infos
-            })
-        env.close()
-
-    # print and save model-specific metrics
-    outfolder = 'test_metrics'
-    print_and_save(all_vehicle_infos, method, model, outfolder)
-
-def print_and_save(stats, method, model, outfolder):
-    """
-    Print and save stats
-    """
-    pass
-
-def load_compare():
-    pass
-
-if __name__=='__main__':
-    import fire
-    fire.Fire()

scratch/etienne/intersimple/airl_flat.py

@@ -1,64 +0,0 @@
-# %%
-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

@@ -1,59 +0,0 @@
-# %%
-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

@@ -1,138 +0,0 @@
-from intersim.envs.intersimple import Intersimple, InfoFilter
-from stable_baselines3.common.policies import BasePolicy
-import gym
-from intersim.envs.intersimple import *
-from gail.envs import *
-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):
-        # RandomLocation mixin re-initializes the intersim sub-env
-        self._intersim_expert._intersim = self._intersimple._env
-        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(BasePolicy):
-
-    def __init__(self, experts):
-        self._experts = [e() for e in experts]
-
-    def forward(self, *args, **kwargs):
-        raise NotImplementedError()
-
-    def _predict(self, *args, **kwargs):
-        raise NotImplementedError()
-    
-    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 forward(self, *args, **kwargs):
-        raise NotImplementedError()
-
-    def _predict(self, *args, **kwargs):
-        raise NotImplementedError()
-
-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 = globals()[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')
-    include_infos = isinstance(env, InfoFilter)
-
-    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,
-        ),
-        exclude_infos=not include_infos,
-    )
-
-if __name__ == '__main__':
-    import fire
-    fire.Fire(demonstrations)

scratch/etienne/intersimple/data/generate.sh

@@ -1,15 +0,0 @@
-#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'
-#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'
-#python -m expert --env=NRasterized --min_timesteps=3000 --video --env_args='{width:36,height:36,m_per_px:2}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001_NRasterizedRandomAgentw36h36mppx2.pkl'
-#python -m expert --env=NRasterizedRandomAgent --min_timesteps=200 --env_args='{width:36,height:36,m_per_px:2}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001_NRasterizedRandomAgentw36h36mppx2.pkl'
-#python -m expert --env=NRasterizedRandomAgent --min_timesteps=10000 --env_args='{width:36,height:36,m_per_px:2}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001N10000_NRasterizedRandomAgentw36h36mppx2.pkl'
-#python -m expert --env=NRasterizedRouteRandomAgent --min_timesteps=10000 --env_args='{width:70,height:70,m_per_px:1}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001N10000_NRasterizedRouteRandomAgentw70h70mppx1.pkl'
-#python -m expert --env=NRasterizedRouteRandomAgentLocation --min_timesteps=100000 --env_args='{width:70,height:70,m_per_px:1}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001N100000_NRasterizedRouteRandomAgentLocationw70h70mppx1.pkl'
-#python -m expert --env=NRasterizedRouteRandomAgentLocation --min_timesteps=100000 --env_args='{width:70,height:70,m_per_px:1,map_color:128}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001N100000_NRasterizedRouteRandomAgentLocationw70h70mppx1mapc128.pkl'
-#python -m expert --env=NRasterizedRouteSpeedRandomAgentLocation --min_timesteps=10000 --env_args='{width:70,height:70,m_per_px:1,map_color:128,mu:0.001}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001N10000_NRasterizedRouteSpeedRandomAgentLocationw70h70mppx1mapc128mu.001.pkl'
-#python -m data.expert --env=NRasterizedRouteSpeedRandomAgentLocation --min_timesteps=10000 --env_args='{width:70,height:70,m_per_px:1,map_color:128,mu:0.001,skip_frames:5}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001N10000_NRasterizedRouteSpeedRandomAgentLocationw70h70mppx1mapc128mu.001skip5.pkl'
-#python -m data.expert --env=TLNRasterizedRouteRandomAgentLocation --min_timesteps=100000 --env_args='{width:70,height:70,m_per_px:1,mu:0.001,random_skip:True,max_episode_steps:50}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001N100000_TLNRasterizedRouteRandomAgentLocationw70h70mppx1mu.001rskips50.pkl'
-python -m data.expert --env=TLNRasterizedRouteRandomAgentLocation --min_timesteps=50000 --env_args='{width:70,height:70,m_per_px:1,mu:0.001,random_skip:True,max_episode_steps:50}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001N50000_TLNRasterizedRouteRandomAgentLocationw70h70mppx1mu.001rskips50.pkl'

scratch/etienne/intersimple/gail/envs.py

@@ -1,46 +0,0 @@
-import gym
-from gym.wrappers.time_limit import TimeLimit
-import numpy as np
-from intersim.envs.intersimple import NRasterizedRouteRandomAgentLocation, RandomLocation, RandomAgent, RewardVisualization, Reward, \
-    ImageObservationAnimation, RasterizedRoute, NObservations, RasterizedObservation, \
-    NormalizedActionSpace, ActionVisualization, InteractionSimulatorMarkerViz, ImitationCompat, Intersimple
-
-class RasterizedSpeed:
-
-    def __init__(self, max_speed=12, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        channels, height, width = self.observation_space.shape
-        self.observation_space = gym.spaces.Box(
-            low=0,
-            high=255,
-            shape=(channels+1, height, width),
-            dtype=np.uint8
-        )
-        self._max_speed = max_speed
-
-    def _simple_obs(self, intersim_obs, intersim_info):
-        img = super()._simple_obs(intersim_obs, intersim_info)
-
-        ego_speed = intersim_obs['state'][self._agent, 2]
-        scaled_speed = (255 * ego_speed) // self._max_speed
-        speed_layer = scaled_speed * np.ones_like(img[:1], dtype=np.uint8)
-        speed_layer = speed_layer.clamp(0, 255)
-        
-        obs = np.concatenate((img, speed_layer), axis=0)
-        return obs
-
-class NRasterizedRouteSpeedRandomAgentLocation(RandomLocation, RandomAgent, RewardVisualization,
-        Reward, ImageObservationAnimation, RasterizedRoute, NObservations, RasterizedSpeed, RasterizedObservation,
-        NormalizedActionSpace, ActionVisualization, InteractionSimulatorMarkerViz, ImitationCompat, Intersimple):
-    pass
-
-class TransparentTimeLimit(TimeLimit):
-
-    def __getattr__(self, name):
-        return getattr(self.env, name)
-    
-    def close(self, *args, **kwargs):
-        return self.env.close(*args, **kwargs)
-
-def TLNRasterizedRouteRandomAgentLocation(max_episode_steps, *args, **kwargs):
-    return TransparentTimeLimit(NRasterizedRouteRandomAgentLocation(*args, **kwargs), max_episode_steps=max_episode_steps)

scratch/etienne/intersimple/gail/options2.py

@@ -1,127 +0,0 @@
-import gym
-import torch
-from src.util.collisions import feasible
-import numpy as np
-from collections import deque
-
-def imitation_discriminator(discriminator):
-    return lambda obs, action, next_obs, done: discriminator.discrim_net.predict_reward_train(
-        state=torch.tensor(obs).unsqueeze(0).to(discriminator.discrim_net.device()),
-        action=torch.tensor([[action]]).to(discriminator.discrim_net.device()),
-        next_state=torch.tensor(next_obs).unsqueeze(0).to(discriminator.discrim_net.device()), # unused
-        done=torch.tensor(done).unsqueeze(0).to(discriminator.discrim_net.device()), # unused
-    ).item()
-
-class OptionsEnv(gym.Wrapper):
-
-    def __init__(self, env, options, discriminator, discount, ll_buffer, *args, **kwargs):
-        super().__init__(env, *args, **kwargs)
-
-        self.options = options
-        num_hl_options = len(self.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,)),
-        })
-
-        self.discriminator = discriminator
-        self.discount = discount
-        self.ll_buffer = ll_buffer
-    
-    @staticmethod
-    def _hl_observation(obs, mask):
-        return {
-            'obs': obs,
-            'mask': mask,
-        }
-    
-    def reset(self):
-        self.done = False
-        self.obs = self.env.reset()
-        self.m = available_actions(self.env, self.options)
-        return self._hl_observation(self.obs, self.m)
-    
-    def _ll_step(self, action):
-        return self.env.step(action)
-    
-    def step(self, action):
-        assert self.m[action]
-        assert not self.done
-
-        plan = list(map(float, generate_plan(self.env, action, self.options)))
-        reward = 0
-        steps = 0
-
-        while not self.done and plan and \
-            (feasible(self.env, safety_plan(self.env, plan)) or self.m.sum() == 1):
-            
-            a, plan = plan[0], plan[1:]
-            a = self.env._normalize(a)
-
-            next_obs, _, self.done, info = self._ll_step(a)
-
-            reward += self.discount**steps * self.discriminator(self.obs, a, next_obs, self.done)
-
-            self.ll_buffer.append({
-                'obs': self.obs,
-                'next_obs': next_obs,
-                'acts': np.array((a,)),
-                'dones': np.array(self.done),
-            })
-
-            steps += 1
-            self.obs = next_obs
-
-        self.m = available_actions(self.env, self.options)
-
-        return self._hl_observation(self.obs, self.m), reward, self.done, info
-
-class RenderOptions(OptionsEnv):
-
-    def __init__(self, env, options, *args, **kwargs):
-        super().__init__(env, options, discriminator=lambda s, a, n, d: 0, discount=1, ll_buffer=deque(maxlen=0), *args, **kwargs)
-
-    def _ll_step(self, action):
-        out = super()._ll_step(action)
-        self.env.render(mode='post')
-        return out
-    
-    def close(self, *args, **kwargs):
-        self.env.close(*args, **kwargs)
-
-def safety_plan(env, plan):
-    return np.concatenate((plan, np.array(5 * [env._env._min_acc])), axis=0)
-
-def available_actions(env, options):
-    """Return mask of available actions given current `env` state.
-    Action 0 is considered safe fallback.
-    """
-    plans = [generate_plan(env, i, options) for i, _ in enumerate(options)]
-    # is emergency braking still possible?
-    plans = list(map(lambda p: safety_plan(env, p), plans))
-
-    T = max(len(p) for p in plans)
-    plans = [np.pad(p, ((0, T-len(p)),), constant_values=np.nan) for p in plans]
-    plans = np.stack(plans, axis=0)
-    
-    valid = feasible(env, plans)
-    if not valid.any():
-        valid[0] = True
-
-    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, options):
-    """Generate input profile for high-level action `i`."""
-    assert i < len(options), "Invalid option index {i}"
-    target_v, t = 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

scratch/etienne/intersimple/gail_flat.py

@@ -1,70 +0,0 @@
-# %%
-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
-
-from gail.discriminator import MlpDiscriminator
-
-model_name = 'gail_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 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=150,
-    n_disc_updates_per_round=32,
-    discrim_kwargs={'discrim_net': MlpDiscriminator()},
-    gen_algo=sb3.PPO("MlpPolicy", venv, verbose=1, n_steps=4530),
-    allow_variable_horizon=True,
-)
-gail_trainer.train(total_timesteps=400000)
-gail_trainer.gen_algo.save(model_name)
-
-#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_flat_ray.py

@@ -1,115 +0,0 @@
-# %%
-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, 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,
-        allow_variable_horizon=True,
-    )
-
-    def callback(epoch):
-        print("callback")
-        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, return_episode_rewards=True)
-        tune.report(
-            reward=np.mean(episode_rewards),
-            length=np.mean(episode_lengths),
-            training_iteration=epoch,
-        )
-
-        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=40000, callback=callback)
-
-analysis = tune.run(
-    training_function,
-    config = {
-        'expert_batch_size': tune.randint(1, 22), #220,
-        'n_disc_updates_per_round': tune.randint(2, 100), #16,
-        'n_steps': tune.randint(1, 10000), #4096,
-    },
-    resources_per_trial={
-        'cpu': 1,
-        # 'gpu': 1,
-    },
-    local_dir='ray',
-    num_samples=10,
-)
-
-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

@@ -1,70 +0,0 @@
-# %%
-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_image_multiagent_nocollision.py

@@ -1,70 +0,0 @@
-# %%
-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 CnnDiscriminatorFlatAction
-
-model_name = 'gail_image_multiagent_nocollision'
-
-# %%
-# Load pickled test demonstrations.
-with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedRandomAgentw36h36mppx2.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={'stop_on_collision':False, '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': CnnDiscriminatorFlatAction(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(stop_on_collision=False, 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_image_random.py

@@ -1,79 +0,0 @@
-# %%
-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 NRasterizedRandomAgent, IntersimpleReward, speed_reward
-import functools
-from stable_baselines3.common.evaluation import evaluate_policy
-
-from gail.discriminator import CnnDiscriminator
-
-model_name = 'gail_image_random'
-
-# %%
-# Load pickled test demonstrations.
-with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedRandomAgentw36h36mppx2.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)
-
-env_kwargs = {'width': 36, 'height': 36, 'm_per_px': 2}
-venv = make_vec_env(NRasterizedRandomAgent, n_envs=2, env_kwargs=env_kwargs)
-
-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/")
-generator = sb3.PPO("CnnPolicy", venv, verbose=1, n_steps=1024)
-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=generator,
-    allow_variable_horizon=True,
-)
-def callback(round):
-    eval_env = NRasterizedRandomAgent(reward=functools.partial(speed_reward, collision_penalty=0.), **env_kwargs)
-    #sync_envs_normalization(self.training_env, self.eval_env)
-    episode_rewards, episode_lengths = evaluate_policy(generator, eval_env, return_episode_rewards=True)
-
-gail_trainer.train(total_timesteps=100000, callback=callback)
-gail_trainer.gen_algo.save(model_name)
-
-#del gail_trainer
-
-# %%
-model = sb3.PPO.load(model_name)
-
-env = NRasterizedRandomAgent(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_image_random_ray.py

@@ -1,171 +0,0 @@
-# %%
-import pathlib
-import pickle
-import tempfile
-import os
-import random
-import numpy as np
-import torch
-
-# set up ray tune
-import ray
-from ray import tune
-from ray.tune import Analysis, ExperimentAnalysis
-from ray.tune.schedulers import ASHAScheduler
-from ray.tune.suggest.hyperopt import HyperOptSearch
-from ray.tune.suggest import ConcurrencyLimiter
-
-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 NRasterizedRandomAgent, IntersimpleReward, speed_reward, NRasterized, NRasterizedRandomAgentVerbose
-import functools
-from stable_baselines3.common.evaluation import evaluate_policy
-from gym.wrappers import TimeLimit
-
-from gail.discriminator import CnnDiscriminator
-
-model_name = 'gail_image_random_ray'
-env_kwargs={'width': 36, 'height': 36, 'm_per_px': 2}
-
-# %%
-
-import argparse
-parser = argparse.ArgumentParser()
-parser.add_argument("--outdir", help="result directory", default='ray')
-parser.add_argument("--test", help="test run", default=False, action="store_true")
-args = parser.parse_args()
-outdir = args.outdir
-
-# %%
-# Load pickled test demonstrations.
-with open("data/NormalizedIntersimpleExpertMu.001N10000_NRasterizedRandomAgentw36h36mppx2.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)
-# Store transitions in shared ray memory
-ray_transitions = ray.put(transitions)
-
-# %%
-venv = make_vec_env(NRasterizedRandomAgent, n_envs=2, env_kwargs=env_kwargs)
-
-tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
-tempdir_path = pathlib.Path(tempdir.name)
-print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
-logger.configure(tempdir_path / "GAIL/")
-
-def get_ray_config(test=False):
-    if test:
-        return {
-            'expert_batch_size': 2,
-            'ppo_n_steps': 2,
-            'ppo_batch_size': 2,
-            'ppo_n_epochs': 1,
-            'total_timesteps': 10,
-        }
-    else:
-        return {
-            'expert_batch_size': tune.choice([2**x for x in range(6,10)]),
-            'ppo_n_steps': tune.choice([2048, 3072, 4096]),
-            'ppo_batch_size': tune.choice([2**x for x in range(9,13)]),
-            'ppo_n_epochs': tune.choice([6,10]),
-            'total_timesteps': 400_000,
-        }
-
-
-def ray_train(config, checkpoint_dir=None):
-    # 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.
-
-    discriminator = CnnDiscriminator(venv)
-    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(
-            "CnnPolicy", venv, verbose=0, 
-            n_steps=config["ppo_n_steps"], 
-            batch_size=config["ppo_batch_size"], 
-            n_epochs=config["ppo_n_epochs"]
-        )
-    gail_trainer = adversarial.GAIL(
-        venv,
-        expert_data=ray.get(ray_transitions),
-        expert_batch_size=config["expert_batch_size"],
-        #n_disc_updates_per_round=2048,
-        discrim_kwargs={'discrim_net': discriminator},
-        gen_algo=generator,
-        allow_variable_horizon=True,
-    )
-    def callback(round):
-        # eval_env = NRasterized(agent=51, reward=functools.partial(speed_reward, collision_penalty=0.), **env_kwargs)
-        eval_env = TimeLimit(NRasterizedRandomAgent(reward=functools.partial(speed_reward, collision_penalty=0.), **env_kwargs), max_episode_steps=1000)
-        episode_rewards, episode_lengths = evaluate_policy(generator, eval_env, return_episode_rewards=True)
-        tune.report(
-            reward=np.mean(episode_rewards),
-            length=np.mean(episode_lengths),
-            training_iteration=round,
-        )
-        with tune.checkpoint_dir(step=round) 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=config['total_timesteps'], callback=callback)
-
-
-ray_config = get_ray_config(args.test)
-search = HyperOptSearch(ray_config, metric='length', mode="max",)
-search = ConcurrencyLimiter(search, max_concurrent=10)
-custom_scheduler = ASHAScheduler(time_attr='training_iteration', metric='length', mode="max", grace_period=15)
-
-analysis = tune.run(
-    ray_train, 
-    # config=ray_config,
-    search_alg=search,
-    scheduler=custom_scheduler,
-    local_dir=outdir,
-    resources_per_trial={"cpu":10, "gpu": 0.2},
-    num_samples=1 if args.test else 100,
-)
-
-del analysis
-
-# %%
-# outdir = "ray/ray_train_2021-09-20_13-33-50/ray_train_f06785b0_33_expert_batch_size=128,ppo_batch_size=1024,ppo_n_epochs=6,ppo_n_steps=2048,total_timesteps=400000_2021-09-20_15-52-05"
-
-# %%
-analysis = Analysis(outdir, default_metric="length", default_mode="max")
-filepath = analysis.get_best_logdir()
-print("Best ray experiment:", filepath)
-config = analysis.get_best_config()
-print("Best config:", config)
-
-# %%
-
-model = sb3.PPO.load(os.path.join(analysis.get_last_checkpoint(), 'gen_checkpoint'))
-
-# env = NRasterized(agent=51, **env_kwargs)
-env = TimeLimit(NRasterizedRandomAgent(reward=functools.partial(speed_reward, collision_penalty=0.), **env_kwargs), max_episode_steps=1000)
-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.env.close(filestr='render/'+model_name)
-# %%
-

scratch/etienne/intersimple/gail_image_singleagent_nocollision.py

@@ -1,70 +0,0 @@
-# %%
-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_singleagent_nocollision'
-
-# %%
-# 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, 'stop_on_collision':False, '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, stop_on_collision=False)
-
-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

@@ -1,90 +0,0 @@
-# %%
-import sys
-sys.path.append('../../../')
-
-from src.discriminator import CnnDiscriminator, CnnDiscriminatorFlatAction
-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
-from tqdm import tqdm
-from src.policies.options import OptionsCnnPolicy
-from src.gail.options import OptionsEnv, LLOptions, HLOptions, RenderOptions
-from src.gail.train import train_discriminator, train_generator
-
-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]] # option 0 is safe fallback
-
-def train(expert_data, epochs=20, expert_batch_size=32, generator_steps=1024, 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': CnnDiscriminatorFlatAction(venv)},
-        #discrim_kwargs={'discrim_net': CnnDiscriminator(venv)},
-        venv=venv, # unused
-        gen_algo=stable_baselines3.PPO("CnnPolicy", venv), # unused
-    )
-
-    generator = stable_baselines3.PPO(
-        OptionsCnnPolicy,
-        OptionsEnv(env, options=ALL_OPTIONS),
-        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 tqdm(range(epochs)):
-        train_discriminator(LLOptions(env, options=ALL_OPTIONS), generator, discriminator, num_samples=expert_batch_size)
-        train_generator(HLOptions(env, options=ALL_OPTIONS), 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.save(model_name)
-
-    # %%
-    model = stable_baselines3.PPO.load(model_name)
-
-    env = RenderOptions(NRasterized(**env_settings), options=ALL_OPTIONS)
-
-    for s in env.sample_ll(model):
-        if s['dones']:
-            break
-
-    env.close(filestr='render/'+model_name)

scratch/etienne/intersimple/gail_options_image_alltracks.py

@@ -1,410 +0,0 @@
-# %%
-import sys
-sys.path.append('../../../')
-from src.discriminator import CnnDiscriminator, CnnDiscriminatorFlatAction
-from src.policies import OptionsCnnPolicy
-from src.util import feasible
-from src.data import load_experts
-
-from imitation.algorithms import adversarial
-from imitation.util import logger
-import imitation.data.rollout as rollout
-
-import stable_baselines3
-from stable_baselines3.common.env_util import make_vec_env
-
-import torch
-import torch.utils.data
-import numpy as np
-import itertools
-import gym
-import pickle
-import tempfile
-import pathlib
-from tqdm import tqdm
-
-from intersim.envs.intersimple import NRasterized, NRasterizedRandomAgent, NRasterizedIncrementingAgent
-
-ALL_OPTIONS = [(v,t) for v in [0,2,4,6,8] for t in [5, 10, 20]] # option 0 is safe fallback
-
-class OptionsEnv(gym.Wrapper):
-    """
-    Wrap an intersimple environment with an options generator
-    """
-    def __init__(self, env, *args, **kwargs):
-        """
-        Initialize wrapped environment and set high-level action and observation spaces
-        """
-        super().__init__(env, *args, **kwargs)
-        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 _after_choice(self):
-        pass
-    
-    def _after_step(self):
-        pass
-
-    def _transitions(self):
-        raise NotImplementedError('Use `LLOptions` or `HLOptions` for sampling.')
-    
-    def sample(self, generator):
-        """
-        yield transitions using a generator
-        Args:
-            generator (sb3.PPO)
-        Yields:
-
-        """
-        self.done = True
-        while True:
-            self.episode_start = False
-            if self.done:
-                self.s = self.env.reset()
-                self.done = False
-                self.episode_start = True
-
-            self.m = available_actions(self.env)
-            self.ch, self.value, self.log_prob = generator.policy.predict({
-                'obs': torch.tensor(self.s).unsqueeze(0).to(generator.policy.device),
-                'mask': torch.tensor(self.m).unsqueeze(0).to(generator.policy.device),
-            })
-            self.plan = list(map(float, generate_plan(self.env, self.ch)))
-
-            self._after_choice()
-
-            assert not self.done
-            assert self.plan
-            #assert feasible(self.env, self.plan, self.ch)
-
-            while not self.done and self.plan and feasible(self.env, self.plan, self.ch):
-                self.a, self.plan = self.plan[0], self.plan[1:]
-                self.a = self.env._normalize(self.a)
-                self.nexts, _, self.done, _ = self.env.step(self.a)
-
-                self._after_step()
-
-                self.s = self.nexts
-            
-            yield from self._transitions()
-
-class LLOptions(OptionsEnv):
-    """Sample low-level (state, action) tuples for discriminator training."""
-    
-    def __init__(self, *args, **kwargs):
-        """
-        LLOption uses the true LL observations
-        """
-        super().__init__(*args, **kwargs)
-        # overwrite observation space to just output obs directly
-        self.observation_space = self.observation_space['obs']
-
-    def _after_choice(self):
-        """
-        After each option choice, initialize/reset the transition buffer
-        """
-        self._transition_buffer = []
-
-    def _after_step(self):
-        """
-        After each ll action, append s, s', a, done to transition buffer
-        """
-        self._transition_buffer.append({
-            'obs': self.s,
-            'next_obs': self.nexts,
-            'acts': np.array((self.a,)),
-            'dones': np.array(self.done),
-        })
-
-    def _transitions(self):
-        """
-        Yield from the transition buffer
-        """
-        yield from self._transition_buffer
-    
-    def sample_ll(self, policy):
-        """
-        Args:
-            policy
-        Returns:
-            gen: iterable which samples low-level transitions from the environment
-        """
-        return self.sample(policy)
-
-class HLOptions(OptionsEnv):
-    """Sample high-level (state, action, reward) tuples for generator training."""
-    
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-    def _after_choice(self):
-        """
-        After an option selection, initialize total reward and number of steps
-        """
-        self.obs = {'obs': np.copy(self.s), 'mask': np.copy(self.m)}
-        self.r = 0
-        self.steps = 0
-
-    def _after_step(self):
-        """
-        After each low-level action, add the discounted discriminated reward score (given a discriminator)
-        """
-        self.r += self.discount**self.steps * self.discriminator.discrim_net.reward_train(
-            state=torch.tensor(self.s).unsqueeze(0).to(self.discriminator.discrim_net.device()),
-            action=torch.tensor([[self.a]]).to(self.discriminator.discrim_net.device()),
-            next_state=torch.tensor(self.s).unsqueeze(0).to(self.discriminator.discrim_net.device()), # unused
-            done=torch.tensor(self.done).unsqueeze(0).to(self.discriminator.discrim_net.device()), # unused
-        )
-        self.steps += 1
-    
-    def _transitions(self):
-        """
-        Yield a single dictionary per high-level selected action
-        Fields:
-            obs: high-level state and mask at selection
-            action: chosen high-level action
-            reward: accumulated option reward
-            episode_start: whether the action was chosen at the episode start
-            value: the value estimate from the starting state
-            log_prob: the log_prob of the selected action from the starting state
-            done: whether the episode has ended
-
-        """
-        yield {
-            'obs': self.obs,
-            'action': self.ch,
-            'reward': self.r.detach(),
-            'episode_start': self.episode_start,
-            'value': self.value.detach(),
-            'log_prob': self.log_prob.detach(),
-            'done': self.done,
-        }
-    
-    def sample_hl(self, policy, discriminator):
-        """
-        Args:
-            policy
-            discriminator: function with which to score rewards
-        Returns:
-            gen: iterable which samples high-level transitions from the environment
-        """
-        self.discriminator = discriminator
-        return self.sample(policy)
-
-class RenderOptions(LLOptions):
-
-    def _after_step(self):
-        """
-        Render the environment after each low-level step
-        """
-        super()._after_step()
-        self.env.render()
-    
-    def close(self, *args, **kwargs):
-        """
-        On 'close', close the environment
-        """
-        self.env.close(*args, **kwargs)
-
-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`.
-
-    Args:
-        env (gym.Env): current environment state
-        i (int): high-level action `i`
-    Returns:
-        plan (np.array): length T array of acceleration values
-    """
-    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 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(env.sample_ll(generator), 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(env.sample_hl(generator, discriminator), num_samples+1))
-    
-    generator.rollout_buffer.reset()
-    for s in generator_samples[:-1]:
-        generator.rollout_buffer.add(
-            obs=s['obs'],
-            action=s['action'].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()
-
-def train(expert_data, env_class=NRasterizedRandomAgent, env_settings={}, epochs=10, discrim_batch_size=32, generator_steps=2048, discount=0.99):
-    """
-    Args: 
-        expert_data: list of transitions
-        env_class: environment class
-        env_settings: environment settings
-        epochs: number of epochs to train for
-        discrim_batch_size: discriminator batch size
-        generator_steps: number of steps taken in generator
-        discount: discount factor
-    Returns:
-        generator (stable_baselines3.PPO): options policy
-    """
-    env = env_class(**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(env_class, n_envs=1, env_kwargs=env_settings)
-    discriminator = adversarial.GAIL(
-        expert_data=expert_data,
-        expert_batch_size=discrim_batch_size,
-        discrim_kwargs={'discrim_net': CnnDiscriminatorFlatAction(venv)},
-        #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 tqdm(range(epochs)):
-        train_discriminator(LLOptions(env), generator, discriminator, num_samples=discrim_batch_size)
-        train_generator(HLOptions(env), generator, discriminator, num_samples=generator_steps)
-    
-    return generator
-
-# %%
-if __name__ == '__main__':
-    # %%
-    model_name = 'gail_options_image'
-    env_class = NRasterizedRandomAgent
-    env_settings = {'width': 36, 'height': 36, 'm_per_px': 2}
-    
-    #env_class = NRasterized
-    #env_settings = {'agent': 51, 'width': 36, 'height': 36, 'm_per_px': 2}
-    files = ['../../../expert_data/DR_USA_Roundabout_FT/track%04i/expert.pkl'%(i) for i in range(5)]
-    transitions=load_experts(files)
-
-    generator = train(
-        transitions,
-        env_class=env_class,
-        env_settings=env_settings,
-        epochs=10, 
-        discrim_batch_size=32, 
-        generator_steps=2048, 
-        discount=0.99
-        )
-
-    generator.save(model_name)
-
-    # %%
-    model = stable_baselines3.PPO.load(model_name)
-
-    env = RenderOptions(NRasterizedRandomAgent(**env_args))
-
-    for s in env.sample_ll(model):
-        if s['dones']:
-            break
-
-    env.close(filestr='render/'+model_name)
-
-# %% Tests
-
-def test_ll_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 = LLOptions(NRasterized(agent=51, width=36, height=36, m_per_px=2))
-
-    gen_transitions = list(itertools.islice(env.sample_ll(
-        policy=stable_baselines3.PPO(
-            OptionsCnnPolicy,
-            OptionsEnv(env),
-            verbose=1,
-        )
-    ), 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_ll_states():
-    env = NRasterized()
-    policy = stable_baselines3.PPO(
-        OptionsCnnPolicy,
-        OptionsEnv(env),
-        verbose=1,
-    )
-    llenv = LLOptions(env)
-    transitions = list(itertools.islice(llenv.sample_ll(policy=policy), 100))
-
-    env2 = NRasterized()
-    s2 = env2.reset()
-    for i, t in enumerate(transitions):
-        assert i == 0 or np.array_equal(t['obs'], transitions[i-1]['next_obs'])
-        assert np.array_equal(t['obs'], s2)
-        assert t['acts'].shape == (1,)
-
-        nexts2, _, done2, _ = env2.step(t['acts'])
-        assert np.array_equal(t['next_obs'], nexts2)
-        assert np.array_equal(t['dones'], done2)
-
-        if done2:
-            break
-
-        s2 = nexts2
-
-def test_hl_transitions():
-    pass

scratch/etienne/intersimple/gail_options_image_random.py

@@ -1,91 +0,0 @@
-# %%
-import sys
-sys.path.append('../../../')
-
-from src.discriminator import CnnDiscriminatorFlatAction
-from imitation.algorithms import adversarial
-import stable_baselines3
-import torch.utils.data
-import numpy as np
-from intersim.envs.intersimple import NRasterizedRouteRandomAgent
-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
-from tqdm import tqdm
-from src.policies.options import OptionsCnnPolicy
-from src.gail.options import OptionsEnv, LLOptions, HLOptions, RenderOptions
-from src.gail.train import train_discriminator, train_generator
-
-model_name = 'gail_options_image_random'
-env_settings = {'width': 70, 'height': 70, 'm_per_px': 1}
-
-ALL_OPTIONS = [(v,t) for v in [0,2,4,6,8] for t in [5, 10, 20]] # option 0 is safe fallback
-
-def train(expert_data, epochs=100, expert_batch_size=64, generator_steps=1024, discount=0.99):
-    env = NRasterizedRouteRandomAgent(**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(NRasterizedRouteRandomAgent, n_envs=1, env_kwargs=env_settings)
-    discriminator = adversarial.GAIL(
-        expert_data=expert_data,
-        expert_batch_size=expert_batch_size,
-        discrim_kwargs={'discrim_net': CnnDiscriminatorFlatAction(venv)},
-        #discrim_kwargs={'discrim_net': CnnDiscriminator(venv)},
-        venv=venv, # unused
-        gen_algo=stable_baselines3.PPO("CnnPolicy", venv), # unused
-    )
-
-    generator = stable_baselines3.PPO(
-        OptionsCnnPolicy,
-        OptionsEnv(env, options=ALL_OPTIONS),
-        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 tqdm(range(epochs)):
-        train_discriminator(LLOptions(env, options=ALL_OPTIONS), generator, discriminator, num_samples=expert_batch_size)
-        train_generator(HLOptions(env, options=ALL_OPTIONS), generator, discriminator, num_samples=generator_steps)
-        generator.save(model_name)
-    
-    return generator
-
-def video(model_name, env):
-    model = stable_baselines3.PPO.load(model_name)
-    env = RenderOptions(env, options=ALL_OPTIONS)
-    for s in env.sample_ll(model):
-        if s['dones']:
-            break
-    env.close(filestr='render/'+model_name)
-
-def evaluate():
-    video(
-        model_name=model_name,
-        env=NRasterizedRouteRandomAgent(**env_settings)
-    )
-
-# %%
-if __name__ == '__main__':
-
-    with open("data/NormalizedIntersimpleExpertMu.001N10000_NRasterizedRouteRandomAgentw70h70mppx1.pkl", "rb") as f:
-        trajectories = pickle.load(f)
-    transitions = rollout.flatten_trajectories(trajectories)
-    train(transitions)

scratch/etienne/intersimple/gail_options_image_random_location.py

@@ -1,144 +0,0 @@
-# %%
-from collections import deque
-import sys
-sys.path.append('../../../')
-
-from src.discriminator import CnnDiscriminator, CnnDiscriminatorFlatAction
-from imitation.algorithms import adversarial
-import stable_baselines3
-import pickle
-import imitation.data.rollout as rollout
-import tempfile
-import pathlib
-from imitation.util import logger
-from tqdm import tqdm
-from src.policies.options import OptionsCnnPolicy
-from src.gail.train import flatten_transitions
-from gail.options2 import OptionsEnv, RenderOptions, imitation_discriminator
-from gail.envs import TLNRasterizedRouteRandomAgentLocation
-from stable_baselines3.common.vec_env.dummy_vec_env import DummyVecEnv
-from stable_baselines3.common.env_util import make_vec_env
-import torch
-import numpy as np
-
-model_name = 'gail_options_image_random_location'
-env_settings = {'width': 70, 'height': 70, 'm_per_px': 1, 'mu': 0.001, 'random_skip': True, 'max_episode_steps': 200}
-
-ALL_OPTIONS = [(v,t) for v in [0,2,4,8,10] for t in [5, 10, 20]] # option 0 is safe fallback
-
-class NoisyDiscriminator(CnnDiscriminatorFlatAction):
-
-    def __init__(self, *args, std=0.0, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.std = std
-
-    def forward(self, state, action):
-        noise = self.std * torch.randn(*action.shape, device=action.device)
-        return super().forward(state, action + noise)
-
-class LLBuffer(deque):
-
-    def sample(self, n):
-        assert n <= self.maxlen, f'Sample size of {n} exceeds buffer capacity of {self.maxlen}'
-        assert n <= len(self), f'Sample size of {n} exceeds buffer size of {len(self)}'
-        ind = np.random.randint(len(self), size=n)
-        return list(self[i] for i in ind)
-
-def train(
-        expert_data,
-        expert_batch_size=4096,
-        discriminator_updates_per_round=20,
-        generator_steps=1024,
-        generator_batch_size=1024,
-        generator_total_steps=8192,
-        generator_updates_per_round=10,
-        discount=1.0,
-        epochs=200,
-    ):
-    env = TLNRasterizedRouteRandomAgentLocation(**env_settings)
-
-    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 = DummyVecEnv([lambda: env])
-    discriminator = adversarial.GAIL(
-        expert_data=expert_data,
-        expert_batch_size=expert_batch_size,
-        #discrim_kwargs={'discrim_net': NoisyDiscriminator(venv, std=0.25)},
-        disc_opt_cls=torch.optim.RMSprop,
-        disc_opt_kwargs={'lr': 0.0001, 'weight_decay': 0.003},
-        discrim_kwargs={'discrim_net': CnnDiscriminator(venv)},
-        venv=venv, # unused
-        gen_algo=stable_baselines3.PPO("CnnPolicy", venv), # unused
-    )
-
-    ll_buffer = LLBuffer(maxlen=expert_batch_size*10)
-
-    options_env = make_vec_env(
-        OptionsEnv,
-        n_envs=1,
-        #vec_env_cls=SubprocVecEnv,
-        env_kwargs={
-            'env': env,
-            'options': ALL_OPTIONS,
-            'discriminator': imitation_discriminator(discriminator),
-            'discount': discount,
-            'll_buffer': ll_buffer,
-        }
-    )
-
-    generator = stable_baselines3.PPO(
-        OptionsCnnPolicy,
-        options_env,
-        verbose=1,
-        batch_size=generator_batch_size,
-        n_steps=generator_steps,
-        n_epochs=generator_updates_per_round,
-        gamma=1.0,
-        learning_rate=1e-4,
-    )
-
-    for _ in tqdm(range(epochs)):
-        ll_buffer.clear()
-
-        # train generator
-        generator.learn(total_timesteps=generator_total_steps)
-
-        # train discriminator
-        for _ in range(discriminator_updates_per_round):
-            generator_samples = ll_buffer.sample(expert_batch_size)
-            generator_samples = flatten_transitions(generator_samples)
-            discriminator.train_disc(gen_samples=generator_samples)
-
-        generator.save(model_name)
-    
-    return generator
-
-def video(model_name, env):
-    model = stable_baselines3.PPO.load(model_name)
-
-    done = False
-    obs = env.reset()
-    while not done:
-        action, _ = model.predict(obs)
-        obs, _, done, _ = env.step(action)
-
-    env.close(filestr='render/'+model_name)
-
-def evaluate():
-    video_settings = { **env_settings, 'random_skip': False, 'max_episode_steps': 200 }
-    env = TLNRasterizedRouteRandomAgentLocation(**video_settings)
-    env = RenderOptions(env, options=ALL_OPTIONS)
-    video(
-        model_name=model_name,
-        env=env
-    )
-
-# %%
-if __name__ == '__main__':
-    with open("data/NormalizedIntersimpleExpertMu.001N50000_TLNRasterizedRouteRandomAgentLocationw70h70mppx1mu.001rskips50.pkl", "rb") as f:
-        trajectories = pickle.load(f)
-    transitions = rollout.flatten_trajectories(trajectories)
-    train(transitions)

scratch/etienne/intersimple/imitation_quickstart.py

@@ -1,63 +0,0 @@
-# %%
-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

@@ -1,36 +0,0 @@
-# %%
-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

@@ -1,41 +0,0 @@
-# %%
-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

@@ -1,35 +0,0 @@
-# %%
-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

@@ -1,33 +0,0 @@
-# %%
-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

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

@@ -1,46 +0,0 @@
-# %%
-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

@@ -1,46 +0,0 @@
-# %%
-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

@@ -1,49 +0,0 @@
-# %%
-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

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

@@ -1,39 +0,0 @@
-# %%
-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

@@ -1,43 +0,0 @@
-# %%
-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

@@ -1,39 +0,0 @@
-# %%
-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

@@ -1,31 +0,0 @@
-# %%
-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/intersimple/render_env_from_model.py

@@ -1,33 +0,0 @@
-
-import stable_baselines3 as sb3
-from intersim.envs.intersimple import NRasterized
-
-
-def render_env(model_name='gail_image_multiagent_nocollision', agent=51, environment=NRasterized):
-    """
-    Render a video from an model, agent, and environment
-    Args:
-        model_name (str): name of the model
-        agent (int): agent to start the video from
-        environment (gym.Env): gym environment class to render environment on
-    """
-
-    model = sb3.PPO.load(model_name)
-
-    env = environment(stop_on_collision=False, width=36, height=36, m_per_px=2, agent=agent)
-
-    obs = env.reset()
-    i=0
-    while True and i < 600:
-        i+=1
-        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+'_agent%i'%(agent)) 
-
-if __name__ == '__main__':
-    import fire
-    fire.Fire(render_env)

scratch/etienne/intersimple/train_discrim.py

@@ -1,71 +0,0 @@
-# %%
-import sys
-sys.path.append('../../../')
-
-import pickle
-import imitation.data.rollout as rollout
-import imitation.data.types as types
-import torch
-from gail.envs import TLNRasterizedRouteRandomAgentLocation
-import tempfile
-import pathlib
-from imitation.util import logger
-from stable_baselines3.common.vec_env.dummy_vec_env import DummyVecEnv
-from imitation.algorithms import adversarial
-from src.discriminator import CnnDiscriminator
-import stable_baselines3
-from tqdm import tqdm
-
-with open("data/NormalizedIntersimpleExpertMu.001N50000_TLNRasterizedRouteRandomAgentLocationw70h70mppx1mu.001rskips50.pkl", "rb") as f:
-    trajectories = pickle.load(f)
-transitions = rollout.flatten_trajectories(trajectories)
-
-# %%
-env_settings = {'width': 70, 'height': 70, 'm_per_px': 1, 'mu': 0.001, 'random_skip': True, 'max_episode_steps': 200}
-env = TLNRasterizedRouteRandomAgentLocation(**env_settings)
-
-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}/.")
-
-expert_batch_size = 4096
-
-venv = DummyVecEnv([lambda: env])
-discriminator = adversarial.GAIL(
-    expert_data=transitions,
-    expert_batch_size=expert_batch_size,
-    #discrim_kwargs={'discrim_net': NoisyDiscriminator(venv, std=0.25)},
-    disc_opt_cls=torch.optim.RMSprop,
-    disc_opt_kwargs={'lr': 0.0001, 'weight_decay': 0.003},
-    discrim_kwargs={'discrim_net': CnnDiscriminator(venv)},
-    venv=venv, # unused
-    gen_algo=stable_baselines3.PPO("CnnPolicy", venv), # unused
-)
-
-expert_data_loader = torch.utils.data.DataLoader(
-    transitions,
-    batch_size=expert_batch_size,
-    collate_fn=types.transitions_collate_fn,
-    shuffle=True,
-    drop_last=True,
-)
-
-gen_data_loader = torch.utils.data.DataLoader(
-    transitions,
-    batch_size=expert_batch_size,
-    collate_fn=types.transitions_collate_fn,
-    shuffle=True,
-    drop_last=True,
-)
-
-# %%
-epochs = 1000
-for i in tqdm(range(epochs)):
-    for expert_samples, gen_samples in zip(expert_data_loader, gen_data_loader):
-        # randomly corrupt actions
-        gen_samples['acts'] = -1 + 2 * torch.rand(*gen_samples['acts'].shape)
-
-        discriminator.train_disc(expert_samples=expert_samples, gen_samples=gen_samples)
-
-    torch.save(discriminator.discrim_net.state_dict(), 'train_discrim.pt')

scratch/etienne/pillbox/intersim_demos.ipynb

@@ -1,249 +0,0 @@
-{
- "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

@@ -1,168 +0,0 @@
-{
- "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

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

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

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

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

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

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

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

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

scratch/johannes/evaluation.py

@@ -1,90 +0,0 @@
-
-def evaluate_policy_simple(
-    model,
-    env: gym.Env,
-    n_eval_episodes: int = 10,
-    deterministic: bool = True,
-    render: bool = False,
-    callback = None,
-    reward_threshold = None,
-    return_episode_rewards: bool = False,
-    warn: bool = True,
-):
-    """
-    Runs policy for ``n_eval_episodes`` episodes and returns average reward.
-    If a vector env is passed in, this divides the episodes to evaluate onto the
-    different elements of the vector env. This static division of work is done to
-    remove bias. See https://github.com/DLR-RM/stable-baselines3/issues/402 for more
-    details and discussion.
-
-    .. note::
-        If environment has not been wrapped with ``Monitor`` wrapper, reward and
-        episode lengths are counted as it appears with ``env.step`` calls. If
-        the environment contains wrappers that modify rewards or episode lengths
-        (e.g. reward scaling, early episode reset), these will affect the evaluation
-        results as well. You can avoid this by wrapping environment with ``Monitor``
-        wrapper before anything else.
-
-    :param model: The RL agent you want to evaluate.
-    :param env: The gym environment or ``VecEnv`` environment.
-    :param n_eval_episodes: Number of episode to evaluate the agent
-    :param deterministic: Whether to use deterministic or stochastic actions
-    :param render: Whether to render the environment or not
-    :param callback: callback function to do additional checks,
-        called after each step. Gets locals() and globals() passed as parameters.
-    :param reward_threshold: Minimum expected reward per episode,
-        this will raise an error if the performance is not met
-    :param return_episode_rewards: If True, a list of rewards and episode lengths
-        per episode will be returned instead of the mean.
-    :param warn: If True (default), warns user about lack of a Monitor wrapper in the
-        evaluation environment.
-    :return: Mean reward per episode, std of reward per episode.
-        Returns ([float], [int]) when ``return_episode_rewards`` is True, first
-        list containing per-episode rewards and second containing per-episode lengths
-        (in number of steps).
-    """
-    episode_rewards = []
-    episode_lengths = []
-
-    episode_counts = 0
-
-    current_rewards = 0
-    current_lengths = 0
-    observations = env.reset()
-    states = None
-    while (episode_counts < n_eval_episodes):
-        actions, states = model.predict(observations, state=states, deterministic=deterministic)
-        observations, rewards, dones, infos = env.step(actions)
-        print(env._env.t)
-        current_rewards += rewards
-        current_lengths += 1
-
-        # unpack values so that the callback can access the local variables
-        reward = rewards
-        done = dones
-        info = infos
-        if info['collision']:
-            print("COLLISION")
-
-        if callback is not None:
-            callback(locals(), globals())
-
-        if dones:
-            episode_rewards.append(current_rewards)
-            episode_lengths.append(current_lengths)
-            episode_counts += 1
-            current_rewards = 0
-            current_lengths = 0
-            if states is not None:
-                states *= 0
-
-        if render:
-            env.render()
-
-    mean_reward = np.mean(episode_rewards)
-    std_reward = np.std(episode_rewards)
-    if reward_threshold is not None:
-        assert mean_reward > reward_threshold, "Mean reward below threshold: " f"{mean_reward:.2f} < {reward_threshold:.2f}"
-    if return_episode_rewards:
-        return episode_rewards, episode_lengths
-    return mean_reward, std_reward

scratch/johannes/intersimple/gail_options_image.py

@@ -1,101 +0,0 @@
-# %%
-import sys
-sys.path.append('../../../')
-
-from src.discriminator import CnnDiscriminator, CnnDiscriminatorFlatAction
-from imitation.algorithms import adversarial
-import stable_baselines3
-import torch.utils.data
-import numpy as np
-from intersim.envs.intersimple import NRasterized, speed_reward
-import itertools
-import functools
-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
-from tqdm import tqdm
-from src.policies.options import OptionsCnnPolicy
-from src.gail.options import OptionsEnv, LLOptions, HLOptions, RenderOptions
-from src.gail.train import train_discriminator, train_generator
-from src.evaluation.evaluation import Evaluation
-from torch.utils.tensorboard import SummaryWriter
-
-
-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]] # option 0 is safe fallback
-
-def train(expert_data, epochs=20, expert_batch_size=32, generator_steps=1024, 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': CnnDiscriminatorFlatAction(venv)},
-        #discrim_kwargs={'discrim_net': CnnDiscriminator(venv)},
-        venv=venv, # unused
-        gen_algo=stable_baselines3.PPO("CnnPolicy", venv), # unused
-    )
-
-    generator = stable_baselines3.PPO(
-        OptionsCnnPolicy,
-        OptionsEnv(env, options=ALL_OPTIONS),
-        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,
-    )
-
-    filestr = os.path.join('out', model_name)
-    writer = SummaryWriter(filestr)
-    ev = Evaluation(filestr, env, expert_data, n_eval_episodes=100)
-    for epoch in tqdm(range(epochs)):
-        train_discriminator(LLOptions(env, options=ALL_OPTIONS), generator, discriminator, num_samples=expert_batch_size)
-        train_generator(HLOptions(env, options=ALL_OPTIONS), generator, discriminator, num_samples=generator_steps)
-
-        metrics = ev.evaluate(epoch, generator, discriminator)
-        for metric, value in metrics.items():
-            writer.add_scalar(metric, value, epoch)
-    
-    return generator
-
-# %%
-if __name__ == '__main__':
-    # %%
- 
-    with open("scratch/etienne/intersimple/data/NormalizedIntersimpleExpertMu.001_NRasterizedInfoAgent51w36h36mppx2.pkl", "rb") as f:
-        trajectories = pickle.load(f)
-    transitions = rollout.flatten_trajectories(trajectories)
-    generator = train(transitions)
-
-    generator.save(model_name)
-
-    # %%
-    model = stable_baselines3.PPO.load(model_name)
-
-    env = RenderOptions(NRasterized(**env_settings), options=ALL_OPTIONS)
-
-    for s in env.sample_ll(model):
-        if s['dones']:
-            break
-
-    env.close(filestr='render/'+model_name)

scratch/johannes/intersimple/gail_options_image_random.py

@@ -1,101 +0,0 @@
-# %%
-import sys
-sys.path.append('../../../')
-
-from src.discriminator import CnnDiscriminatorFlatAction
-from imitation.algorithms import adversarial
-import stable_baselines3
-import torch.utils.data
-import numpy as np
-from intersim.envs.intersimple import NRasterizedRandomAgent
-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
-from tqdm import tqdm
-from src.policies.options import OptionsCnnPolicy
-from src.gail.options import OptionsEnv, LLOptions, HLOptions, RenderOptions
-from src.gail.train import train_discriminator, train_generator
-from src.evaluation.evaluation import Evaluation
-from torch.utils.tensorboard import SummaryWriter
-import os 
-
-model_name = 'gail_options_image_random'
-env_settings = {'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
-
-def train(expert_data, epochs=100, expert_batch_size=16, generator_steps=16, discount=0.99):
-    env = NRasterizedRandomAgent(**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(NRasterizedRandomAgent, n_envs=1, env_kwargs=env_settings)
-    discriminator = adversarial.GAIL(
-        expert_data=expert_data,
-        expert_batch_size=expert_batch_size,
-        discrim_kwargs={'discrim_net': CnnDiscriminatorFlatAction(venv)},
-        #discrim_kwargs={'discrim_net': CnnDiscriminator(venv)},
-        venv=venv, # unused
-        gen_algo=stable_baselines3.PPO("CnnPolicy", venv), # unused
-    )
-
-    generator = stable_baselines3.PPO(
-        OptionsCnnPolicy,
-        OptionsEnv(env, options=ALL_OPTIONS),
-        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,
-    )
-    
-    filestr = os.path.join('out', model_name)
-    writer = SummaryWriter(filestr)
-    ev = Evaluation(filestr, env, expert_data, n_eval_episodes=100)
-    for epoch in tqdm(range(epochs)):
-        train_discriminator(LLOptions(env, options=ALL_OPTIONS), generator, discriminator, num_samples=expert_batch_size)
-        train_generator(HLOptions(env, options=ALL_OPTIONS), generator, discriminator, num_samples=generator_steps)
-        generator.save(model_name)
-    
-        metrics = ev.evaluate(epoch, generator, discriminator)
-        for metric, value in metrics.items():
-            writer.add_scalar(metric, value, epoch)
-    
-    return generator
-
-def video(model_name, env):
-    model = stable_baselines3.PPO.load(model_name)
-    env = RenderOptions(env, options=ALL_OPTIONS)
-    for s in env.sample_ll(model):
-        if s['dones']:
-            break
-    env.close(filestr='render/'+model_name)
-
-def evaluate():
-    video(
-        model_name=model_name,
-        env=NRasterizedRandomAgent(**env_settings)
-    )
-
-# %%
-if __name__ == '__main__':
-
-    with open("scratch/etienne/intersimple/data/NormalizedIntersimpleExpertMu.001N10000_NRasterizedRandomAgentInfow36h36mppx2.pkl", "rb") as f:
-        trajectories = pickle.load(f)
-    transitions = rollout.flatten_trajectories(trajectories)
-    train(transitions)

scratch/johannes/normalization.py

@@ -1,73 +0,0 @@
-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()

scratch/johannes/raytune_simple.py

@@ -1,69 +0,0 @@
-"""This example demonstrates basic Ray Tune random search and grid search."""
-import time
-
-import ray
-from ray import tune
-
-
-def evaluation_fn(step, width, height):
-    time.sleep(0.1)
-    return (0.1 + width * step / 100)**(-1) + height * 0.1
-
-def easy_objective(config):
-    # Hyperparameters
-    width, height = config["width"], config["height"]
-
-    mydata = ray.get(ray_data)
-    print(mydata)
-
-    for step in range(config["steps"]):
-        # Iterative training function - can be any arbitrary training procedure
-        intermediate_score = evaluation_fn(step, width, height)
-        # Feed the score back back to Tune.
-        tune.report(iterations=step, mean_loss=intermediate_score)
-
-
-if __name__ == "__main__":
-    import argparse
-
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "--smoke-test", action="store_true", help="Finish quickly for testing")
-    parser.add_argument(
-        "--server-address",
-        type=str,
-        default=None,
-        required=False,
-        help="The address of server to connect to if using "
-        "Ray Client.")
-    args, _ = parser.parse_known_args()
-    if args.server_address is not None:
-        ray.init(f"ray://{args.server_address}")
-    else:
-        ray.init(configure_logging=False)
-
-    # This will do a grid search over the `activation` parameter. This means
-    # that each of the two values (`relu` and `tanh`) will be sampled once
-    # for each sample (`num_samples`). We end up with 2 * 50 = 100 samples.
-    # The `width` and `height` parameters are sampled randomly.
-    # `steps` is a constant parameter.
-
-    import numpy as np
-    N = 3
-    data = np.random.rand(N,N,N)
-    ray_data = ray.put(data)
-
-
-    analysis = tune.run(
-        easy_objective,
-        metric="mean_loss",
-        mode="min",
-        num_samples=5 if args.smoke_test else 50,
-        config={
-            "steps": 5 if args.smoke_test else 100,
-            "width": tune.uniform(0, 20),
-            "height": tune.uniform(-100, 100),
-            "activation": tune.grid_search(["relu", "tanh"])
-        })
-
-    print("Best hyperparameters found were: ", analysis.best_config)

setup.py

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

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

src/bc/__init__.py

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

src/bc/bc.py

@@ -1,191 +0,0 @@
-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/__init__.py

@@ -1 +0,0 @@
-from src.data.expert import single_agent_expert, single_agent_demonstrations, multi_agent_demonstrations, load_experts, process_experts 

src/data/data_utils.py

@@ -1,96 +0,0 @@
-import torch
-from torch.utils.data import Dataset
-import numpy as np
-from src.data.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/data/expert.py

@@ -1,241 +0,0 @@
-import intersim
-from intersim.envs.intersimple import Intersimple
-from stable_baselines3.common.policies import BasePolicy
-import gym
-import intersim.envs.intersimple
-import pickle
-from tqdm import tqdm
-import copy
-import os
-import numpy as np
-
-from src.util.rollout import rollout_and_save, flatten_trajectories, make_sample_until
-
-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
-
-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 load_experts(expert_files, flatten=True):
-    """
-    Load expert trajectories from files and combine their transitions into a single RB
-
-    Args:
-        expert_files (list): list of expert file strings
-        flatten (bool): whether to flatten trajectory info
-    Returns:
-        transitions (list): list of combined expert episode transitions 
-    """
-    transitions = []
-    for file in tqdm(expert_files):
-        with open(file, "rb") as f:
-            new_trajectories = pickle.load(f)
-        transitions += new_trajectories
-    if flatten:
-        transitions = flatten_trajectories(transitions)
-    return transitions
-
-def single_agent_expert(expert='NormalizedIntersimpleExpert', 
-                        env='NRasterizedRouteIncrementingAgent', 
-                        env_args={}, policy_args={}, **kwargs):
-    """
-    Args: 
-        expert (class): class of expert
-        env (class): class of env intersim.envs.intersimple
-        env_args (dict): dictionary of kwargs when instantiating environment class
-        policy_args (dict): dictionary of kwargs when instantiating Expert policy
-        path (str): path to store output
-        min_timesteps (int): min number of timesteps for call to rollout.rollout_and_save
-        min_episodes (int): min number of episodes for call to rollout.rollout_and_save
-        video (bool): whether to save a video of the expert until a single environment instantiation stops
-    """
-    Env = intersim.envs.intersimple.__dict__[env]
-    Expert = globals()[expert]
-    env = Env(**env_args)
-    policy = Expert(env, **policy_args)
-    single_agent_demonstrations(env, policy, **kwargs)
-
-def single_agent_demonstrations(env, policy, 
-                                path=None, min_timesteps=None, 
-                                min_episodes=None, video=False,
-                                env_args={}, policy_args={}):
-    """Rollout and save expert demos.
-    
-    Usage:
-        python -m intersimple.expert <flags>
-    Args: 
-        env (class): intersimple environment
-        policy (BasePolicy): intersimple policy
-        path (str): path to store output
-        min_timesteps (int): min number of timesteps for call to rollout.rollout_and_save
-        min_episodes (int): min number of episodes for call to rollout.rollout_and_save
-        video (bool): whether to save a video of the expert until a single environment instantiation stops
-        env_args (dict): dictionary of kwargs when instantiating environment class
-        policy_args (dict): dictionary of kwargs when instantiating Expert policy
-    """
-    
-    if min_timesteps is None and min_episodes is None:
-        min_episodes = env.nv # one episode per vehicle being controlled in environment (hopefully an incrementing agent environment)
-
-    if video:
-        save_video(env, policy)
-
-    path = path or (policy.__class__.__name__ + '_' + env.__class__.__name__ + '.pkl')
-    suntil = make_sample_until(min_timesteps=min_timesteps,min_episodes=min_episodes)
-    
-    rollout_and_save(
-        path=path,
-        policy=policy,
-        env=env,
-        sample_until=suntil
-    )
-
-def multi_agent_demonstrations(expert='IntersimExpert',path=None, env_args={}, policy_args={}):
-    """
-    Run and save the `intersim' multiagent environment demonstration
-    
-    Args: 
-        expert (class): class of multi-agent expert
-        path (str): path to store output data
-        env_args (dict): dictionary of kwargs when instantiating environment class
-        policy_args (dict): dictionary of kwargs when instantiating Expert policy
-    """
-    if path is None:
-        raise('No path specified')
-
-    env = gym.make('intersim:intersim-v0',**env_args)
-    Expert = globals()[expert] 
-    policy = Expert(env, **policy_args)
-
-    s, done = env.reset(), False
-    env.render(mode='file')
-    while not done:
-        _,_,done,_ = env.step(policy.predict()[0])
-        env.render(mode='file')
-    env.close(filestr=path)
-
-def process_experts(filename:str='expert.pkl', 
-                    locs:list=None, 
-                    tracks:list=None,
-                    env_class:str='NRasterizedRouteIncrementingAgent',
-                    env_args:dict={'width':36,'height':36,'m_per_px':2},
-                    expert_class:str='NormalizedIntersimpleExpert',  
-                    expert_args:dict={'mu':0.001}):
-    """
-    Process all experts in the Interaction Dataset
-    For now, using NormalizedIntersimpleExpert with NRasterizedRouteIncrementingAgent environment
-
-    Args:
-        filename (str): name for track file
-        locs (list): list of location ids
-        tracks (list): list of track numbers
-        env_class (str): class of environment
-        env_args (dict): default environment kwargs
-        expert_class (str): class of expert
-        expert_args (dict): default expert kwargs
-    """
-    locs = locs or intersim.LOCATIONS
-    tracks = tracks or range(intersim.MAX_TRACKS)
-    pbar = tqdm(total=len(locs)*len(tracks))
-    for loc in locs:
-        for track in tracks:
-            
-            iloc = intersim.LOCATIONS.index(loc)
-
-            it_env_args = copy.deepcopy(env_args)
-            env_loc_args = {
-                'loc':iloc,
-                'track':track,
-            }
-            it_env_args.update(env_loc_args)
-            out_folder = os.path.join('expert_data',loc, 'track%04i'%(track))
-            if not os.path.isdir(out_folder):
-                os.makedirs(out_folder) 
-            it_path = os.path.join(out_folder,filename)
-
-            # Multi-Agent demonstrations
-            it_ma_path = os.path.join(out_folder,'joint_expert')
-            multi_agent_demonstrations(
-                expert='IntersimExpert',
-                path=it_ma_path,
-                env_args=env_loc_args,
-                policy_args=expert_args
-            )
-            # Single-Agent POV Demonstrations
-            single_agent_expert(
-                expert=expert_class, 
-                env=env_class, 
-                path=it_path, 
-                env_args=it_env_args, 
-                policy_args=expert_args,
-                )
-            
-            pbar.update(1)
-    pbar.close()
-    
-if __name__=='__main__':
-    import fire
-    fire.Fire(process_experts)
-