CommonRoad-Geometric
Our Python framework for graph-based autonomous driving research provides a user-friendly and fully customizable data processing pipeline for extracting PyTorch-based graph datasets from traffic scenarios.
The spatiotemporal heterogeneous graph representations provided by CommonRoad-Geometric explicitly models current and historic vehicle interactions in a map-aware manner. Being based on the well-established PyTorch-Geometric framework for graph-based deep learning, our framework allows researchers to leverage cutting-edge graph neural networks to model the complex interaction effects among a varying number of traffic participants and the underlying road infrastructure, which can then be used for various downstream learning tasks such as traffic prediction and motion planning.

pip install git+https://github.com/CommonRoad/crgeo.git
Motivating example:
simulation = SumoSimulation(
"data/osm_crawled/DEU_Munich_1-100.xml"
)
traffic_extractor = TrafficExtractor(
simulation=simulation,
options=TrafficExtractorOptions(
edge_drawer=VoronoiEdgeDrawer(dist_threshold=30.0),
#feature_computers=...
)
)
with simulation:
for time_step, scenario in simulation(to_time_step=200):
data = traffic_extractor.extract(TrafficExtractionParams(
index=time_step,
))
print(data)
print(data.vehicle.velocity)
>> CommonRoadData(scenario_id=ZAM_TMP-0_0, t=0, dt=0.04)
>> - VirtualAttributesNodeStorage('vehicle')
>> x: Tensor (torch.Size([1, 19]), torch.float32)
>> velocity: idx (0, 2)
>> acceleration: idx (2, 4)
>> ...
>> ...
>> - VirtualAttributesEdgeStorage(('vehicle', 'to', 'vehicle'))
>> edge_index: Tensor (torch.Size([2, 0]), torch.int64)
>> edge_attr: Tensor (torch.Size([0, 10]), torch.float32)
>> same_lanelet: idx (0, 1)
>> distance: idx (1, 2)
>> ...
>> ...
>> - VirtualAttributesNodeStorage('lanelet')
>> ...
>> tensor([[6.21., 0.], [5.12, 0.], [7.45, 0.], ...])