Add basic setup.py

interimit/__init__.py

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

interimit/data_utils.py

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

interimit/expert_data.py

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

interimit/nets/deepsets.py

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

interimit/nets/util.py

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

interimit/policies/policy.py

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