diff --git a/experiments/experiment.py b/experiments/experiment.py index 8f26d9d..17ae7c5 100644 --- a/experiments/experiment.py +++ b/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") diff --git a/src/bc/bc.py b/src/bc/bc.py index 687e476..b81512d 100644 --- a/src/bc/bc.py +++ b/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['path_y'] = ob['paths'][1] + ob['ego_state'] = ob['state'] + 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_y': MinMaxScaler(reduce_dim=2), + 'path': 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) diff --git a/src/data_utils.py b/src/data_utils.py index f6b5cba..b14ec3a 100644 --- a/src/data_utils.py +++ b/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']) - self.raw_data['relative_state'].append(data['relative_state']) - 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) + 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['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']) \ - == len(self.raw_data['action']) \ - == len(self.raw_data['path_x']) \ - == len(self.raw_data['path_y']), 'dataset lengths unequal' + 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['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 - 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 + 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 """ - keys = ['state', 'relative_state', 'path_x', 'path_y', 'action'] - sample = {key:self.raw_data[key][idx] for key in keys} + #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 \ No newline at end of file diff --git a/src/expert_data.py b/src/expert_data.py index 543d4a2..4b4b5ba 100644 --- a/src/expert_data.py +++ b/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) - osm = get_map_path(base='InteractionSimulator', loc=loc) + 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) @@ -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): - nni = ~torch.isnan(obs[t]['state'][:,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()) - 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_x'].append(obs[t]['paths'][0][nni]) - data['path_y'].append(obs[t]['paths'][1][nni]) + data['path'].append(torch.stack((obs[t]['paths'][0][nni], obs[t]['paths'][1][nni]), dim=-1)) - # cat lists - data['state'] = torch.cat(data['state']).type(dtype) - data['action'] = torch.cat(data['action']).type(dtype) - data['path_x'] = torch.cat(data['path_x']).type(dtype) - data['path_y'] = torch.cat(data['path_y']).type(dtype) + # 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['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']) \ - == len(data['action']) == len(data['path_x']) \ - == len(data['path_y']), 'dataset lengths unequal' + lengths = [len(data[key]) for key in keys] + assert min(lengths) == max(lengths), '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 diff --git a/src/main.py b/src/main.py index 943f8a7..c31ba6a 100644 --- a/src/main.py +++ b/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]) - cv_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=[3]) + train_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=kwargs['train_tracks']) + 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=track, filestr=filestr, nframes=kwargs['nframes'], graph=kwargs['graph']) + simulate_policy(policy, loc=loc, track=kwargs['test_tracks'][0], filestr=filestr, nframes=kwargs['nframes'], graph=kwargs['graph']) # run test metrics - test_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=[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(): diff --git a/src/metrics.py b/src/metrics.py index 07e0879..3b1028d 100644 --- a/src/metrics.py +++ b/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') diff --git a/src/policies/policy.py b/src/policies/policy.py index 524345e..53683f3 100644 --- a/src/policies/policy.py +++ b/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_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 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 = torch.cat([sample["path_x"], sample["path_y"]], dim=-1) - path = self.path_net(path) + path = self.path_net(sample["path"].reshape((sample["path"].shape[0], -1))) x = torch.cat([ego, relative, path], dim=-1) x = self.head(x) return x