making expert data save s, a, sp. making dataloader also load batches thisway. renaming state to ego_state. converting path_x and path_y to single path variable. making number of samples for ray an argument. adjusting metrics, policy, and other functions to be able to handle this

experiments/experiment.py

@@ -48,6 +48,8 @@ def parse_args():
         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,
@@ -64,6 +66,7 @@ def parse_args():
         '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)),
@@ -156,7 +159,7 @@ if __name__ == '__main__':
             local_dir=kwargs['outdir'],
             #resources_per_trial={"cpu": 2},
             time_budget_s=120*60,
-            num_samples=200,
+            num_samples=kwargs['nsamples'],
         )
     elif kwargs['ray'] and kwargs['test']:
         analysis = Analysis(kwargs['outdir'], default_metric="cv_loss", default_mode="min")

src/bc/bc.py

@@ -80,16 +80,16 @@ class BehaviorCloningPolicy():
 
     def __call__(self, ob):
 
-        if 'action' in ob.keys():
+        if 'ego_state' in ob.keys():
             # extract state from dataloader samples  
             pass
         else:
             # extract state from observation (using simulator)
-            ob['path_x'] = ob['paths'][0]
+            ob['ego_state'] = ob['state']
-            ob['path_y'] = ob['paths'][1]
+            ob['path'] = torch.stack(ob['paths'],dim=-1)
 
         # run observation through transforms
-        for key in ['state', 'relative_state', 'path_x', 'path_y']:             
+        for key in ['ego_state', 'relative_state', 'path']:             
             if key in self._transforms.keys():
                 ob[key] = self._transforms[key].transform(ob[key])
 
@@ -149,13 +149,15 @@ def generate_transforms(dataset):
     """
     transforms = {
         'action': MinMaxScaler(),
-        'state': MinMaxScaler(),
+        'ego_state': MinMaxScaler(),
         'relative_state': MinMaxScaler(reduce_dim=2),
-        'path_x': MinMaxScaler(reduce_dim=2),
+        'path': MinMaxScaler(reduce_dim=2),
-        'path_y': MinMaxScaler(reduce_dim=2),
     }
     for key in transforms.keys():
-        transforms[key].fit(dataset[:][key])
+        if key == 'action':
+            transforms[key].fit(dataset[:][key])
+        else:
+            transforms[key].fit(dataset[:]['state'][key])
 
     return transforms
 
@@ -190,7 +192,7 @@ def train(config, policy, train_dataset, cv_dataset, filestr, **kwargs):
     cv_loader = DataLoader(cv_dataset, batch_size=cv_batch_size, shuffle=True)
 
     # change policy dtype
-    policy.policy = policy.policy.type(train_dataset[0]['state'].dtype)
+    policy.policy = policy.policy.type(train_dataset[0]['state']['ego_state'].dtype)
 
     # generate loss function, optimizer
     cv_loss_fn = nn.MSELoss(reduction='sum')
@@ -220,7 +222,7 @@ def train(config, policy, train_dataset, cv_dataset, filestr, **kwargs):
         for (batch_idx, batch) in enumerate(training_loader):
             
             # sample mini-batch and run through policy
-            pred_action = policy(batch) 
+            pred_action = policy(batch['state']) 
             loss = loss_fn(pred_action, batch['action'])
 
             # compute loss and step optimizer
@@ -239,7 +241,7 @@ def train(config, policy, train_dataset, cv_dataset, filestr, **kwargs):
             with torch.no_grad():
                 cv_loss = 0.
                 for (batch_idx, batch) in enumerate(cv_loader):
-                    pred_action = policy(batch) 
+                    pred_action = policy(batch['state']) 
                     loss = cv_loss_fn(pred_action, batch['action'])
                     cv_loss += loss.item() / len(cv_dataset)
             

src/data_utils.py

@@ -25,13 +25,14 @@ class InteractionDatasetSingleAgent(Dataset):
         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 = {'state':[], 'relative_state':[], 'action':[], 'path_x':[], 'path_y':[]}
+        self.raw_data = {key:[] for key in self.keys}
         max_nv = 0
         for track in self.tracks:
             try:
@@ -41,33 +42,26 @@ class InteractionDatasetSingleAgent(Dataset):
                 print('Failed to load location {} track {}'.format(self.loc,track))
                 continue
             max_nv = max(max_nv, data['relative_state'].shape[1])
-            self.raw_data['state'].append(data['state'])
+            for key in self.keys:
-            self.raw_data['relative_state'].append(data['relative_state'])
+                self.raw_data[key].append(data[key])
-            self.raw_data['action'].append(data['action'])
-            self.raw_data['path_x'].append(data['path_x'])
-            self.raw_data['path_y'].append(data['path_y'])
-
-        # cat lists
-        self.raw_data['state'] = torch.cat(self.raw_data['state']).type(self.dtype)
-        self.raw_data['action'] = torch.cat(self.raw_data['action']).type(self.dtype)
-        self.raw_data['path_x'] = torch.cat(self.raw_data['path_x']).type(self.dtype)
-        self.raw_data['path_y'] = torch.cat(self.raw_data['path_y']).type(self.dtype)
 
         # 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['relative_state'] = torch.cat(self.raw_data['relative_state']).type(self.dtype)
+            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
-        assert len(self.raw_data['state']) == len(self.raw_data['relative_state']) \
+        lengths = [len(self.raw_data[key]) for key in self.keys]
-            == len(self.raw_data['action']) \
+        assert min(lengths) == max(lengths), 'dataset lengths unequal'
-            == len(self.raw_data['path_x']) \
-            == len(self.raw_data['path_y']), 'dataset lengths unequal'
         
     def __len__(self):
-        return len(self.raw_data['state'])
+        return len(self.raw_data['ego_state'])
 
     def __getitem__(self, idx):
         """ 
@@ -76,12 +70,27 @@ class InteractionDatasetSingleAgent(Dataset):
             idx: index or indices of B samples
         Returns:
             sample (dict): sample dictionary with the following entries:
-                state (torch.tensor): (B, 5) raw state
+                state (dict): state dictionary with the following entries:
-                relative_state (torch.tensor): (B, max_nv, d) relative state (padded with nans)
+                    ego_state (torch.tensor): (B, 5) raw state
-                path_x (torch.tensor): (B, P) tensor of P future path x positions
+                    relative_state (torch.tensor): (B, max_nv, d) relative state (padded with nans)
-                path_y (torch.tensor): (B, P) tensor of P future path y positions
+                    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
         """
-        keys = ['state', 'relative_state', 'path_x', 'path_y', 'action']
+        #sample = {key:self.raw_data[key][idx] for key in self.keys}
-        sample = {key:self.raw_data[key][idx] for key in keys}
+        sample = {
+            'state':{
+                'ego_state':self.raw_data['ego_state'][idx],
+                'relative_state':self.raw_data['relative_state'][idx],
+                'path':self.raw_data['path'][idx]
+            },
+            'action':self.raw_data['action'][idx],
+            'next_state':{
+                'ego_state':self.raw_data['next_ego_state'][idx],
+                'relative_state':self.raw_data['next_relative_state'][idx],
+                'path':self.raw_data['next_path'][idx]},
+        }
         return sample

src/expert_data.py

@@ -31,8 +31,8 @@ def generate_expert_data(path: str='expert_data', loc: int = 0, track:int = 0,
         os.makedirs(path)
     filestr = opj(path,intersim.LOCATIONS[loc]+'_track%03i'%(track))
     
-    svt, svt_path = get_svt(base='InteractionSimulator', loc=loc, track=track)
+    svt, svt_path = get_svt(loc=loc, track=track) #base='InteractionSimulator'
-    osm = get_map_path(base='InteractionSimulator', loc=loc)
+    osm = get_map_path(loc=loc)
     print('SVT path: {}'.format(svt_path))
     print('Map path: {}'.format(osm))
     states, actions = SVT_to_stateactions(svt)
@@ -91,33 +91,42 @@ def process_expert_observations(obs, actions, filestr, remove_outliers=True, dty
         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 = ['state', 'action', 'relative_state', 'path_x', 'path_y']
+    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):
+    for t in range(T-1):
-        nni = ~torch.isnan(obs[t]['state'][:,0])
+        nni = ~torch.isnan(obs[t]['state'][:,0]) & ~torch.isnan(obs[t+1]['state'][:,0]) 
         max_nv = max(max_nv,nni.count_nonzero())
-        data['state'].append(obs[t]['state'][nni])
+
+        # 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['action'].append(actions[t][nni])
+        data['path'].append(torch.stack((obs[t]['paths'][0][nni], obs[t]['paths'][1][nni]), dim=-1))
-        data['path_x'].append(obs[t]['paths'][0][nni])
-        data['path_y'].append(obs[t]['paths'][1][nni])
 
-    # cat lists 
+        # action
-    data['state'] = torch.cat(data['state']).type(dtype)
+        data['action'].append(actions[t][nni])
-    data['action'] = torch.cat(data['action']).type(dtype)
+
-    data['path_x'] = torch.cat(data['path_x']).type(dtype)
+        # next state
-    data['path_y'] = torch.cat(data['path_y']).type(dtype)
+        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['relative_state'] = torch.cat(data['relative_state']).type(dtype)
+        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)
@@ -125,12 +134,11 @@ def process_expert_observations(obs, actions, filestr, remove_outliers=True, dty
             data[key] = data[key][non_outlier_indices[:,0]]    
 
     # mandate equal length
-    assert len(data['state']) == len(data['relative_state']) \
+    lengths = [len(data[key]) for key in keys]
-        == len(data['action']) == len(data['path_x']) \
+    assert min(lengths) == max(lengths), 'dataset lengths unequal'
-        == len(data['path_y']), 'dataset lengths unequal'
 
     # save out data
-    for key in data.keys():
+    for key in keys:
         torch.save(data[key], filestr+'_'+key+'.pt')
 
 def load_expert_data(path='expert_data', loc: int = 0, track:int = 0):
@@ -146,7 +154,8 @@ def load_expert_data(path='expert_data', loc: int = 0, track:int = 0):
     # load observations and actions
     filestr = opj(path, intersim.LOCATIONS[loc]+'_track%03i'%(track))
     data = {}
-    for key in ['state','action','relative_state','path_x','path_y']:
+    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
 

src/main.py

@@ -48,8 +48,8 @@ def main(config, method='bc', train=False, test=False, loc=0, datadir='./expert_
 
         # make policy, train and test datasets, and send to 
         policy = policy_class(config)
-        train_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=[0,1,2])
+        train_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=kwargs['train_tracks'])
-        cv_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=[3])
+        cv_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=kwargs['cv_tracks'])
         train_fn(config, policy, train_dataset, cv_dataset, filestr, **kwargs)
     
     if test:
@@ -59,11 +59,10 @@ def main(config, method='bc', train=False, test=False, loc=0, datadir='./expert_
         policy.eval()
 
         # simulate policy
-        track = 4
+        simulate_policy(policy, loc=loc, track=kwargs['test_tracks'][0], filestr=filestr, nframes=kwargs['nframes'], graph=kwargs['graph'])
-        simulate_policy(policy, loc=loc, track=track, filestr=filestr, nframes=kwargs['nframes'], graph=kwargs['graph'])
 
         # run test metrics
-        test_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=[track])
+        test_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=kwargs['test_tracks'])
         writer = SummaryWriter(filestr)
         info = metrics(filestr, test_dataset, policy)
         for k, m in info.items():

src/metrics.py

@@ -37,7 +37,7 @@ def metrics(filestr: str, test_dataset, policy):
     info['average_velocity'] = avg_v
 
     # convert policy dtype between float32 and float64
-    policy.policy = policy.policy.type(test_dataset[0]['state'].dtype)
+    policy.policy = policy.policy.type(test_dataset[0]['state']['ego_state'].dtype)
     
     # generate actions in test dataset
     true_actions, pred_actions = [], []
@@ -45,9 +45,9 @@ def metrics(filestr: str, test_dataset, policy):
     test_loader = DataLoader(test_dataset, batch_size=1024)
     with torch.no_grad():
         for (batch_idx, batch) in enumerate(test_loader):
-            pred_actions.append(policy(batch))
+            pred_actions.append(policy(batch['state']))
             true_actions.append(batch['action'])
-            true_velocities.append(batch['state'][:,2])
+            true_velocities.append(batch['state']['ego_state'][:,2])
 
     true_actions, pred_actions = torch.cat(true_actions,dim=0), torch.cat(pred_actions, dim=0)
     visualize_distribution(true_actions[:,0], pred_actions[:,0], filestr+'_action_viz') 

src/policies/policy.py

@@ -34,17 +34,14 @@ class DeepSetsPolicy(Policy, nn.Module):
             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 (torch.tensor): (B, P, 2) tensor of P future path x and y positions
-                path_y (torch.tensor): (B, P) tensor of P future path y positions
                 action (torch.tensor): (B, 1) actions taken from each state
         Returns:
             x (torch.tensor): (head_output_dim,) output of common head network
         """
-        ego = self.ego_net(sample["state"])
+        ego = self.ego_net(sample["ego_state"])
         relative = self.deepsets_net(sample["relative_state"])
-        # cat path_x, path_y to tensor of dim (B, 2*P)
+        path = self.path_net(sample["path"].reshape((sample["path"].shape[0], -1)))
-        path = torch.cat([sample["path_x"], sample["path_y"]], dim=-1)
-        path = self.path_net(path)
         x = torch.cat([ego, relative, path], dim=-1)
         x = self.head(x)
         return x