saeedrmd/trajectory-prediction-waymo

Waymo Trajectory Prediction

Dataset Description

This dataset contains preprocessed trajectory prediction samples for autonomous driving research, formatted for use with DiscoBench's TrajectoryPrediction task.

• —Original Dataset: Waymo Open Motion Dataset
• —Number of Samples: 850
• —Format: Pickle files with numpy arrays
• —Task: Multi-modal trajectory prediction

Dataset Structure

Each sample is a pickle file containing:

• —obj_trajs (32, 21, 2): Past trajectories of surrounding agents (2.1s @ 10Hz)
• —obj_trajs_mask (32, 21): Validity mask for past trajectories
• —map_polylines (128, 20, 2): Map polylines (lanes, boundaries)
• —map_polylines_mask (128, 20): Validity mask for map polylines
• —center_gt_trajs (60, 2): Ground truth future trajectory (6s @ 10Hz)
• —center_gt_trajs_mask (60,): Validity mask for future trajectory
• —center_gt_final_valid_idx: Last valid timestep index
• —track_index_to_predict: Index of the track to predict

Usage

Download Dataset

from huggingface_hub import hf_hub_download
import pickle

# Download a single sample
file_path = hf_hub_download(
    repo_id="saeedrmd/trajectory-prediction-waymo",
    filename="sample_0000.pkl",
    repo_type="dataset"
)

# Load the sample
with open(file_path, 'rb') as f:
    data = pickle.load(f)

print("Agent trajectories:", data['obj_trajs'].shape)
print("Ground truth:", data['center_gt_trajs'].shape)

Download All Samples

from huggingface_hub import snapshot_download

# Download entire dataset
dataset_path = snapshot_download(
    repo_id="saeedrmd/trajectory-prediction-waymo",
    repo_type="dataset"
)

print(f"Dataset downloaded to: {dataset_path}")

Use with DiscoBench

1. 1.Download the dataset using the code above
2. 2.Place the samples in your DiscoBench cache directory:

   cp {dataset_path}/*.pkl /path/to/DiscoBench/cache/Waymo Trajectory Prediction/

1. 1.Update your task configuration to point to this cache directory

Dataset Statistics

• —Number of samples: 850
• —Average valid agents per sample: ~31
• —Average valid polylines per sample: ~128
• —Average valid future timesteps: ~60

Coordinate System

All trajectories and map features are in the focal agent's coordinate frame:

• —Origin: Focal agent's position at current timestep
• —Orientation: Aligned with focal agent's heading
• —Units: Meters

Citation

If you use this dataset, please cite:

@dataset{trajectory_prediction_discobench,
  title={Waymo Trajectory Prediction},
  author={DiscoBench Team},
  year={2025},
  publisher={Hugging Face},
  howpublished={\url{https://huggingface.co/datasets/saeedrmd/trajectory-prediction-waymo}}
}

License

Apache 2.0

Original Dataset Citations