huangfu/InteractionImitation
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Arec
2021-08-04 09:45:36 -07:00
parent 7ae01f73a2
commit f9729b0a9d
7 changed files with 89 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
experiments/experiment.py
View File

@@ -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")

24
src/bc/bc.py
View File

@@ -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)

57
src/data_utils.py
View File

@@ -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

49
src/expert_data.py
View File

@@ -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))
# 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))
# 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)
# 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

9
src/main.py
View File

@@ -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():

6
src/metrics.py
View File

@@ -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')

9
src/policies/policy.py
View File

@@ -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