210 lines
8.0 KiB
Python
210 lines
8.0 KiB
Python
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) |