adding functions to process expert data across locations and tracks in intersimple environment

src/data/__init__.py

@@ -0,0 +1 @@
+from src.data.expert_trajectories import demonstrations, load_experts, process_experts 

src/data/data_utils.py

@@ -0,0 +1,96 @@
+import torch
+from torch.utils.data import Dataset
+import numpy as np
+from src.expert_data import load_expert_data
+import os
+opj = os.path.join
+
+class InteractionDatasetMultiAgent(Dataset):
+    """
+    Class to handle getting full multi-agent observations and actions
+    """
+    pass
+
+class InteractionDatasetSingleAgent(Dataset):
+    """Class to load states and actions for individual agents."""
+
+    def __init__(self, output_dir='expert_data', loc:int = 0, tracks:list = [0], dtype=torch.float32):
+        """
+        Args:
+            output_dir (string): Directory with all the images.
+            loc (int): location index
+            tracks (list[int]): track indices
+        """
+        self.output_dir = output_dir
+        self.loc = loc
+        self.tracks = tracks
+        self.dtype = dtype
+        self.keys = ['ego_state', 'relative_state', 'path', 'action', 'next_ego_state', 'next_relative_state', 'next_path']
+        self._load_dataset()
+
+    def _load_dataset(self):
+        """
+        Load the full datasets ahead of time
+        """
+        self.raw_data = {key:[] for key in self.keys}
+        max_nv = 0
+        for track in self.tracks:
+            try:
+                data = load_expert_data(path=self.output_dir, loc=self.loc, track=track)
+                print('Loaded location {} track {}'.format(self.loc,track))
+            except:
+                print('Failed to load location {} track {}'.format(self.loc,track))
+                continue
+            max_nv = max(max_nv, data['relative_state'].shape[1])
+            for key in self.keys:
+                self.raw_data[key].append(data[key])
+
+        # pad second dimension of relative state
+        for i in range(len(self.raw_data['relative_state'])):
+            nv1, nv2, d = self.raw_data['relative_state'][i].shape
+            pad = torch.zeros(nv1, max_nv-nv2, d, dtype=self.dtype) * np.nan
+            self.raw_data['relative_state'][i] = torch.cat((self.raw_data['relative_state'][i], pad), dim=1)
+            self.raw_data['next_relative_state'][i] = torch.cat((self.raw_data['next_relative_state'][i], pad), dim=1)
+
+        # cat lists
+        for key in self.keys:
+            self.raw_data[key] = torch.cat(self.raw_data[key]).type(self.dtype)
+
+        # mandate equal length
+        lengths = [len(self.raw_data[key]) for key in self.keys]
+        assert min(lengths) == max(lengths), 'dataset lengths unequal'
+        
+    def __len__(self):
+        return len(self.raw_data['ego_state'])
+
+    def __getitem__(self, idx):
+        """ 
+        Sample from the dataset
+        Args:
+            idx: index or indices of B samples
+        Returns:
+            sample (dict): sample dictionary with the following entries:
+                state (dict): state dictionary with the following entries:
+                    ego_state (torch.tensor): (B, 5) raw state
+                    relative_state (torch.tensor): (B, max_nv, d) relative state (padded with nans)
+                    path (torch.tensor): (B, P, 2) tensor of P future path x and y positions
+                action (torch.tensor): (B, 1) actions taken from each state
+                next_stat (dict): next state dictionary with the following entries:
+                    ego_state (torch.tensor): (B, 5) raw next state
+                    relative_state (torch.tensor): (B, max_nv, d) next relative state (padded with nans)
+                    path (torch.tensor): (B, P, 2) tensor of P future next path x and y positions
+        """
+        #sample = {key:self.raw_data[key][idx] for key in self.keys}
+        sample = {
+            'state':{
+                'ego_state':self.raw_data['ego_state'][idx],
+                'relative_state':self.raw_data['relative_state'][idx],
+                'path':self.raw_data['path'][idx]
+            },
+            'action':self.raw_data['action'][idx],
+            'next_state':{
+                'ego_state':self.raw_data['next_ego_state'][idx],
+                'relative_state':self.raw_data['next_relative_state'][idx],
+                'path':self.raw_data['next_path'][idx]},
+        }
+        return sample

src/data/expert.py

@@ -0,0 +1,214 @@
+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 imitation.data.rollout as rollout
+from stable_baselines3.common.vec_env.dummy_vec_env import DummyVecEnv
+from imitation.data.wrappers import RolloutInfoWrapper
+import copy
+import os
+
+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 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
+
+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
+    )
+
+def process_experts(filename:str='expert.pkl', 
+                    locs:list=None, 
+                    tracks:list=None,
+                    env_class:str='NRasterizedIncrementingAgent',
+                    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 NRasterizedIncrementingAgent 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)
+            it_env_args.update({
+                'loc':iloc,
+                'track':track,
+            })
+            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)
+
+            demonstrations(
+                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)
+

src/data/expert_data.py

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