trajdata

repository·main·Indexed 19 days ago

A unified interface for accessing multiple human trajectory forecasting datasets, including nuScenes, Waymo, Lyft, nuPlan, View-of-Delft, and INTERACTION. Version 1.4.0 provides a consistent API for data loading, timestep interpolation for multi-dataset training, vector map access via MapAPI, and simulation capabilities using SimulationScene.

Tokens: 17.6K
Snippets: 58
Records: 63
Agent score: 64%

What's inside trajdata

How timestep interpolation works with multiple datasets
main
When loading multiple datasets simultaneously (e.g., ["nusc_mini", "eupeds_eth"]), they may have different underlying annotation frequencies (dt).
To ensure all batched data shares the same frequency, use the desired_dt argument in the UnifiedDataset constructor. trajdata performs timestep interpolation to this common frequency.
Constraint: Interpolation can only be performed to integer multiples of the original data annotation frequency. For example, if nuScenes has dt=0.5s and ETH BIWI has dt=0.4s, you can interpolate both to a common desired_dt=0.1s.
```
dataset = UnifiedDataset(
    desired_data=["nusc_mini", "eupeds_eth"],
    data_dirs={
        "nusc_mini": "~/datasets/nuScenes",
        "eupeds_eth": "~/datasets/eth_ucy_peds"
    },
    desired_dt=0.1,
)
```
https://github.com/nvlabs/trajdata/blob/main/README.md

Setup Waymo Open Motion Dataset

main

Download Waymo Open Motion Dataset v1.1 following the instructions on the dataset website.

Note: You do not need to download all parts. Only download the necessary directories from the Google Cloud Bucket (e.g., validation for the validation dataset).

The directory structure must match the following:

/path/to/waymo/
            ├── training/
            |   ├── training.tfrecord-00000-of-01000
            |   ├── training.tfrecord-00001-of-01000
            |   └── ...
            ├── validation/
            |   ├── validation.tfrecord-00000-of-00150
            |   ├── validation.tfrecord-00001-of-00150
            |   └── ...
            └── testing/
            |   ├── testing.tfrecord-00000-of-00150
            |   ├── testing.tfrecord-00001-of-00150
            |   └── ...

Setup ETH/UCY Pedestrians dataset

main

The ETH/UCY Pedestrians raw data can be sourced from various online locations. This data loader assumes the data was sourced from the Trajectron-plus-plus GitHub repository.

The directory structure must match the following:

/path/to/eth_ucy/
            ├── biwi_eth.txt
            ├── biwi_hotel.txt
            ├── crowds_zara01.txt
            ├── crowds_zara02.txt
            ├── crowds_zara03.txt
            ├── students001.txt
            ├── students003.txt
            └── uni_examples.txt

Setup nuScenes dataset
main
Download the nuScenes dataset following the nuscenes-devkit instructions.
Note: You only need to download the annotations; raw sensor data (radar, camera, lidar, etc.) is not required.
The directory structure must match the following:
```
/path/to/nuScenes/
            ├── maps/
            ├── samples/
            ├── sweeps/
            ├── v1.0-mini/
            ├── v1.0-test/
            └── v1.0-trainval/
```
https://github.com/nvlabs/trajdata/blob/main/DATASETS.md
Install trajdata
main
Install the core package via PyPI:
```
pip install trajdata
```
To use specific datasets that require additional devkits or dependencies (like nuScenes, Lyft, Waymo, INTERACTION, or View-of-Delft), install the corresponding extras:
- nuScenes: pip install "trajdata[nusc]"
- Lyft: pip install "trajdata[lyft]"
- Waymo: pip install "trajdata[waymo]"
- INTERACTION: pip install "trajdata[interaction]"
- View-of-Delft: pip install "trajdata[vod]"
- All supported extras: pip install "trajdata[nusc,lyft,waymo,interaction,vod]"
After installation, you must manually download the raw datasets. Refer to DATASETS.md for instructions on how to structure your dataset folders.
https://github.com/nvlabs/trajdata/blob/main/README.md

Setup INTERACTION Dataset

main

Download the INTERACTION Dataset following the instructions on the dataset website.

The directory structure must match the following (supports both interaction_single and interaction_multi formats):

/path/to/interaction_single/
            ├── maps/
            |   ├── DR_CHN_Merging_ZS0.osm
            |   ├── DR_CHN_Merging_ZS0.osm_xy
            |   └── ...
            ├── test_conditional-single-agent/
            |   ├── DR_CHN_Merging_ZS0_obs.csv
            |   ├── DR_CHN_Merging_ZS2_obs.csv
            |   └── ...
            ├── test_single-agent/
            |   ├── DR_CHN_Merging_ZS0_obs.csv
            |   ├── DR_CHN_Merging_ZS2_obs.csv
            |   └── ...
            └── train/
            |   ├── DR_CHN_Merging_ZS0_train.csv
            |   ├── DR_CHN_Merging_ZS2_train.csv
            |   └── ...
            └── val/
                ├── DR_CHN_Merging_ZS0_val.csv
                ├── DR_CHN_Merging_ZS2_val.csv
                └── ...

/path/to/interaction_multi/
            ├── maps/
            |   ├── DR_CHN_Merging_ZS0.osm
            |   ├── DR_CHN_Merging_ZS0.osm_xy
            |   └── ...
            ├── test_conditional-multi-agent/
            |   ├── DR_CHN_Merging_ZS0_obs.csv
            |   ├── DR_CHN_Merging_ZS2_obs.csv
            |   └── ...
            ├── test_multi-agent/
            |   ├── DR_CHN_Merging_ZS0_obs.csv
            |   ├── DR_CHN_Merging_ZS2_obs.csv
            |   └── ...
            └── train/
            |   ├── DR_CHN_Merging_ZS0_train.csv
            |   ├── DR_CHN_Merging_ZS2_train.csv
            |   └── ...
            └── val/
                ├── DR_CHN_Merging_ZS0_val.csv
                ├── DR_CHN_Merging_ZS2_val.csv
                └── ...

Setup Stanford Drone Dataset

main

The Stanford Drone Dataset can be found on Kaggle (e.g., this space-optimized version).

Note: You only need to download the annotations; the videos are not required.

The directory structure must match the following:

/path/to/sdd/
            ├── bookstore/
            |   ├── video0
            |   |   ├── annotations.txt
            |   |   └── reference.jpg
            |   ├── video1
            |   |   ├── annotations.txt
            |   |   └── reference.jpg
            |   └── ...
            ├── coupa/
            |   ├── video0
            |   |   ├── annotations.txt
            |   |   └── reference.jpg
            |   ├── video1
            |   |   ├── annotations.txt
            |   |   └── reference.jpg
            |   └── ...
            └── ...

Setup Argoverse 2 Motion Forecasting dataset

main

Download the Argoverse 2 Motion Forecasting dataset from the official website.

The directory structure must match the following:

/path/to/av2mf/
            ├── train/
            |   ├── 0000b0f9-99f9-4a1f-a231-5be9e4c523f7/
            |   |   ├── log_map_archive_0000b0f9-99f9-4a1f-a231-5be9e4c523f7.json
            |   |   └── scenario_0000b0f9-99f9-4a1f-a231-5be9e4c523f7.parquet
            |   ├── 0000b6ab-e100-4f6b-aee8-b520b57c0530/
            |   |   ├── log_map_archive_0000b6ab-e100-4f6b-aee8-b520b57c0530.json
            |   |   └── scenario_0000b6ab-e100-4f6b-aee8-b520b57c0530.parquet
            |   └── ...
            ├── val/
            |   ├── 00010486-9a07-48ae-b493-cf4545855937/
            |   |   ├── log_map_archive_00010486-9a07-48ae-b493-cf4545855937.json
            |   |   └── scenario_00010486-9a07-48ae-b493-cf4545855937.parquet
            |   └── ...
            └── test/
            |   ├── 0000b329-f890-4c2b-93f2-7e2413d4ca5b/
            |   |   ├── log_map_archive_0000b329-f890-4c2b-93f2-7e2413d4ca5b.json
            |   |   └── scenario_0000b329-f890-4c2b-93f2-7e2413d4ca5b.parquet
            |   └── ...

Initialize and step through a simulation using SimulationScene

main

The SimulationScene class allows you to initialize a simulation from real trajectory data and step through it to track agent motion and metrics.

To use it, you must first create a UnifiedDataset pointing to your local data directories, retrieve a Scene object using get_scene(), and then instantiate a SimulationScene.

Key parameters for SimulationScene:

env_name: A string identifier for the environment.
scene_name: A string identifier for the scene.
scene: The Scene object retrieved from the dataset.
dataset: The UnifiedDataset instance.
init_timestep: The starting timestep for the simulation.
freeze_agents: A boolean indicating whether to freeze agent states.

After calling .reset(), you can iterate through the scene length using .step(new_xyh_dict), where new_xyh_dict is a dictionary mapping agent names to their next state (typically [x, y, heading]).

from typing import Dict
import numpy as np
from trajdata import AgentBatch, UnifiedDataset
from trajdata.data_structures.scene_metadata import Scene
from trajdata.simulation import SimulationScene

# 1. Setup Dataset
dataset = UnifiedDataset(
    desired_data=["nusc_mini"],
    data_dirs={
        "nusc_mini": "~/datasets/nuScenes",
    },
)

# 2. Get a specific scene
desired_scene: Scene = dataset.get_scene(scene_idx=0)

# 3. Initialize Simulation
sim_scene = SimulationScene(
    env_name="nusc_mini_sim",
    scene_name="sim_scene",
    scene=desired_scene,
    dataset=dataset,
    init_timestep=0,
    freeze_agents=True,
)

# 4. Run Simulation Loop
obs: AgentBatch = sim_scene.reset()
for t in range(1, sim_scene.scene.length_timesteps):
    new_xyh_dict: Dict[str, np.ndarray] = dict()

    # Define next states for agents
    for idx, agent_name in enumerate(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_pos
        next_state[2] = curr_yaw
        new_xyh_dict[agent_name] = next_state

    # Step the simulation
    obs = sim_scene.step(new_xyh_dict)

Setup View-of-Delft dataset
main
To use the View-of-Delft Prediction dataset, download it following the instructions in the official devkit README. The directory structure must match the following:
```
/path/to/VoD/
            ├── maps/
            ├── v1.0-test/
            └── v1.0-trainval/
```
https://github.com/nvlabs/trajdata/blob/main/DATASETS.md
Install trajdata as a package developer
main
To install trajdata in editable mode for development:
1. Install all required dependencies:
```
pip install -r requirements.txt
```
1. Install the package in editable mode:
```
pip install -e .
```
```
pip install -r requirements.txt
pip install -e .
```
https://github.com/nvlabs/trajdata/blob/main/README.md

Setup Lyft Level 5 dataset

main

Download the Lyft Level 5 dataset following the instructions on the dataset website.

Note: You only need to download the necessary .zarr files (e.g., sample.zarr for the small sample dataset).

The directory structure must match the following:

/path/to/lyft/
            ├── LICENSE
            ├── aerial_map
            ├── feedback.txt
            ├── meta.json
            ├── scenes/
            |   ├── sample.zarr
            |   ├── train.zarr
            |   └── ...
            └── semantic_map/
                └── semantic_map.pb