Package: trajdata
trajdata: A Unified Interface to Multiple Human Trajectory DatasetsAnnouncementsSept 2023:Our paper about trajdatahas been accepted to the NeurIPS 2023 Datasets and Benchmarks Track!InstallationThe easiest way to install trajdata is through PyPI withpipinstalltrajdataIn case you would also like to use datasets such as nuScenes, Lyft Level 5, or Waymo Open Motion Dataset (which require their own devkits to access raw data or additional package dependencies), the following will also install the respective devkits and/or package dependencies.# For nuScenespipinstall"trajdata[nusc]"# For Lyftpipinstall"trajdata[lyft]"# For Waymopipinstall"trajdata[waymo]"# For INTERACTIONpipinstall"trajdata[interaction]"# Allpipinstall"trajdata[nusc,lyft,waymo,interaction]"Then, download the raw datasets (nuScenes, Lyft Level 5, ETH/UCY, etc.) in case you do not already have them. For more information about how to structure dataset folders/files, please seeDATASETS.md.Package Developer InstallationFirst, in whichever environment you would like to use (conda, venv, ...), make sure to install all required dependencies withpip install -r requirements.txtThen, install trajdata itself in editable mode withpip install -e .Data Preprocessing [Optional]The dataloader operates via a two-stage process, visualized below.While optional, we recommend first preprocessing data into a canonical format. Take a look at theexamples/preprocess_data.pyscript for an example script that does this. Data preprocessing will execute the first part of the diagram above and create data caches for each specified dataset.Note: Explicitly preprocessing datasets like this is not necessary; the dataloader will always internally check if there exists a cache for any requested data and will create one if not.Data LoadingAt a minimum, batches of data for training/evaluation/etc can be loaded the following way:importosfromtorch.utils.dataimportDataLoaderfromtrajdataimportAgentBatch,UnifiedDataset# See below for a list of already-supported datasets and splits.dataset=UnifiedDataset(desired_data=["nusc_mini"],data_dirs={# Remember to change this to match your filesystem!"nusc_mini":"~/datasets/nuScenes"},)dataloader=DataLoader(dataset,batch_size=64,shuffle=True,collate_fn=dataset.get_collate_fn(),num_workers=os.cpu_count(),# This can be set to 0 for single-threaded loading, if desired.)batch:AgentBatchforbatchindataloader:# Train/evaluate/etc.passFor a more comprehensive example, please seeexamples/batch_example.py.For more information on all of the possibleUnifiedDatasetconstructor arguments, please seesrc/trajdata/dataset.py.Supported DatasetsCurrently, the dataloader supports interfacing with the following datasets:DatasetIDSplitsLocationsDescriptiondtMapsnuScenes Train/TrainVal/Valnusc_trainvaltrain,train_val,valboston,singaporenuScenes prediction challenge training/validation/test splits (500/200/150 scenes)0.5s (2Hz):white_check_mark:nuScenes Testnusc_testtestboston,singaporenuScenes test split, no annotations (150 scenes)0.5s (2Hz):white_check_mark:nuScenes Mininusc_minimini_train,mini_valboston,singaporenuScenes mini training/validation splits (8/2 scenes)0.5s (2Hz):white_check_mark:nuPlan Trainnuplan_trainN/Aboston,singapore,pittsburgh,las_vegasnuPlan training split (947.42 GB)0.05s (20Hz):white_check_mark:nuPlan Validationnuplan_valN/Aboston,singapore,pittsburgh,las_vegasnuPlan validation split (90.30 GB)0.05s (20Hz):white_check_mark:nuPlan Testnuplan_testN/Aboston,singapore,pittsburgh,las_vegasnuPlan testing split (89.33 GB)0.05s (20Hz):white_check_mark:nuPlan Mininuplan_minimini_train,mini_val,mini_testboston,singapore,pittsburgh,las_vegasnuPlan mini training/validation/test splits (942/197/224 scenes, 7.96 GB)0.05s (20Hz):white_check_mark:Waymo Open Motion Trainingwaymo_traintrainN/AWaymo Open Motion Datasettrainingsplit0.1s (10Hz):white_check_mark:Waymo Open Motion Validationwaymo_valvalN/AWaymo Open Motion Datasetvalidationsplit0.1s (10Hz):white_check_mark:Waymo Open Motion Testingwaymo_testtestN/AWaymo Open Motion Datasettestingsplit0.1s (10Hz):white_check_mark:Lyft Level 5 Trainlyft_traintrainpalo_altoLyft Level 5 training data - part 1/2 (8.4 GB)0.1s (10Hz):white_check_mark:Lyft Level 5 Train Fulllyft_train_fulltrainpalo_altoLyft Level 5 training data - part 2/2 (70 GB)0.1s (10Hz):white_check_mark:Lyft Level 5 Validationlyft_valvalpalo_altoLyft Level 5 validation data (8.2 GB)0.1s (10Hz):white_check_mark:Lyft Level 5 Samplelyft_samplemini_train,mini_valpalo_altoLyft Level 5 sample data (100 scenes, randomly split 80/20 for training/validation)0.1s (10Hz):white_check_mark:INTERACTION Dataset Single-Agentinteraction_singletrain,val,test,test_conditionalusa,china,germany,bulgariaSingle-agent split of the INTERACTION Dataset (where the goal is to predict one target agents' future motion)0.1s (10Hz):white_check_mark:INTERACTION Dataset Multi-Agentinteraction_multitrain,val,test,test_conditionalusa,china,germany,bulgariaMulti-agent split of the INTERACTION Dataset (where the goal is to jointly predict multiple agents' future motion)0.1s (10Hz):white_check_mark:ETH - Univeupeds_ethtrain,val,train_loo,val_loo,test_loozurichThe ETH (University) scene from the ETH BIWI Walking Pedestrians dataset0.4s (2.5Hz)ETH - Hoteleupeds_hoteltrain,val,train_loo,val_loo,test_loozurichThe Hotel scene from the ETH BIWI Walking Pedestrians dataset0.4s (2.5Hz)UCY - Univeupeds_univtrain,val,train_loo,val_loo,test_loocyprusThe University scene from the UCY Pedestrians dataset0.4s (2.5Hz)UCY - Zara1eupeds_zara1train,val,train_loo,val_loo,test_loocyprusThe Zara1 scene from the UCY Pedestrians dataset0.4s (2.5Hz)UCY - Zara2eupeds_zara2train,val,train_loo,val_loo,test_loocyprusThe Zara2 scene from the UCY Pedestrians dataset0.4s (2.5Hz)Stanford Drone Datasetsddtrain,val,teststanfordStanford Drone Dataset (60 scenes, randomly split 42/9/9 (70%/15%/15%) for training/validation/test)0.0333...s (30Hz)Adding New DatasetsThe code that interfaces the original datasets (dealing with their unique formats) can be found insrc/trajdata/dataset_specific.To add a new dataset, one needs to:Create a new folder undersrc/trajdata/dataset_specificwhich will contain all the code specific to a particular dataset (e.g., for extracting data into our canonical format). In particular, there must be:An__init__.pyfile.A file that defines a subclass ofRawDatasetand implements some of its functions. Reference implementations can be found in thenusc/nusc_dataset.py,lyft/lyft_dataset.py, andeth_ucy_peds/eupeds_dataset.pyfiles.Add a subclass ofNamedTupletosrc/trajdata/dataset_specific/scene_records.pywhich contains the minimal set of information sufficient to describe a scene. This "scene record" will be used in conjunction with the raw dataset class above and relates to how scenes are stored and efficiently accessed later.Add a section to theDATASETS.mdfile which outlines how users should store the raw dataset locally.Add a section tosrc/trajdata/utils/env_utils.pywhich allows users to get the raw dataset via its name, and specify if the dataset is a good candidate for parallel processing (e.g., does its native dataset object have a large memory footprint which might not allow it to be loaded in multiple processes, such as nuScenes?) and if it has maps.ExamplesPlease see theexamples/folder for more examples, below are just a few demonstrations of core capabilities.Multiple DatasetsThe following will load data from both the nuScenes mini dataset as well as the ETH - University scene from the ETH BIWI Walking Pedestrians dataset.dataset=UnifiedDataset(desired_data=["nusc_mini","eupeds_eth"],data_dirs={# Remember to change this to match your filesystem!"nusc_mini":"~/datasets/nuScenes","eupeds_eth":"~/datasets/eth_ucy_peds"},desired_dt=0.1,# Please see the note below about common dt!)Note: Be careful about loading multiple datasets without an associateddesired_dtargument; many datasets do not share the same underlying data annotation frequency. To address this, we've implemented timestep interpolation to a common frequency which will ensure that all batched data shares the same dt. Interpolation can only be performed to integer multiples of the original data annotation frequency. For example, nuScenes'dt=0.5and the ETH BIWI dataset'sdt=0.4can be interpolated to a commondesired_dt=0.1.Map APItrajdataalso provides an API to access the raw vector map information from datasets that provide it.frompathlibimportPathfromtrajdataimportMapAPI,VectorMapcache_path=Path("~/.unified_data_cache").expanduser()map_api=MapAPI(cache_path)vector_map:VectorMap=map_api.get_map("nusc_mini:boston-seaport")Simulation InterfaceOne additional feature of trajdata is that it can be used to initialize simulations from real data and track resulting agent motion, metrics, etc.At a minimum, a simulation can be initialized and stepped through as follows (also present inexamples/simple_sim_example.py):fromtypingimportDict# Just for type annotationsimportnumpyasnpfromtrajdataimportAgentBatch,UnifiedDatasetfromtrajdata.data_structures.scene_metadataimportScene# Just for type annotationsfromtrajdata.simulationimportSimulationScene# See below for a list of already-supported datasets and splits.dataset=UnifiedDataset(desired_data=["nusc_mini"],data_dirs={# Remember to change this to match your filesystem!"nusc_mini":"~/datasets/nuScenes",},)desired_scene:Scene=dataset.get_scene(scene_idx=0)sim_scene=SimulationScene(env_name="nusc_mini_sim",scene_name="sim_scene",scene=desired_scene,dataset=dataset,init_timestep=0,freeze_agents=True,)obs:AgentBatch=sim_scene.reset()fortinrange(1,sim_scene.scene.length_timesteps):new_xyh_dict:Dict[str,np.ndarray]=dict()# Everything inside the forloop just sets# agents' next states to their current ones.foridx,agent_nameinenumerate(obs.agent_name):curr_yaw=obs.curr_agent_state[idx,-1]curr_pos=obs.curr_agent_state[idx,:2]next_state=np.zeros((3,))next_state[:2]=curr_posnext_state[2]=curr_yawnew_xyh_dict[agent_name]=next_stateobs=sim_scene.step(new_xyh_dict)examples/sim_example.pycontains a more comprehensive example which initializes a simulation from a scene in the nuScenes mini dataset, steps through it by replaying agents' GT motions, and computes metrics based on scene statistics (e.g., displacement error from the original GT data, velocity/acceleration/jerk histograms).CitationIf you use this software, please cite it as follows:@Inproceedings{ivanovic2023trajdata,
  author = {Ivanovic, Boris and Song, Guanyu and Gilitschenski, Igor and Pavone, Marco},
  title = {{trajdata}: A Unified Interface to Multiple Human Trajectory Datasets},
  booktitle = {{Proceedings of the Neural Information Processing Systems (NeurIPS) Track on Datasets and Benchmarks}},
  month = dec,
  year = {2023},
  address = {New Orleans, USA},
  url = {https://arxiv.org/abs/2307.13924}
}TODOCreate a method like finalize() which writes all the batch information to a TFRecord/WebDataset/some other format which is (very) fast to read from for higher epoch training.