How to Select Specific Scenarios in Waymax Based on WAYMO OPEN MOTION DATASET Methodology?

[Jynxzzz]

I am exploring the utilization of Waymax for simulating scenarios based on the Waymo's data. I recently came across a paper titled "Large Scale Interactive Motion Forecasting for Autonomous Driving: The WAYMO OPEN MOTION DATASET", where the authors outlined a method to mine specific pairwise interaction scenarios from a large dataset. They mentioned:
"we specifically mined for the following pairwise interaction scenarios: merges, lane changes, unprotected turns, intersection left turns, intersection right turns, pedestrian-vehicle interactions, cyclist-vehicle interactions, interactions with close proximity, and interactions with high accelerations.”

I am intrigued to understand how I could adopt a similar approach within Waymax to select specific scenarios like intersections or merging scenarios. Is there a recommended method or tool within Waymax that facilitates the selection or annotation of these specific interaction scenarios?

I am particularly interested in automating the process of scenario selection based on certain criteria as outlined in the paper, and I'm looking for any guidance on how to approach this within the Waymax environment.

Here's the excerpt from the paper for reference:
"we specifically mined for the following pairwise interaction scenarios: merges, lane changes, unprotected turns, intersection left turns, intersection right turns, pedestrian-vehicle interactions, cyclist-vehicle interactions, interactions with close proximity, and interactions with high accelerations.”

Thank you!

[justinjfu]

My understanding of this question is that you would like to filter the dataset elements by some criteria, and then write them back to disk?

You could try the following strategy:

1. Use the dataloader to read TF examples as raw dictionaries (https://github.com/waymo-research/waymax/blob/main/waymax/dataloader/dataloader_utils.py#L63)
2. Write a filter function to match the criteria of your choice. Inside the filter function you could use the Waymax datastructures if you find those easier to work with, but it's optional. To convert dictionaries to Waymax datastructures use the provided factory functions in: https://github.com/waymo-research/waymax/blob/main/waymax/dataloader/womd_factories.py#L135
3. Convert the results back to TF record format and write to disk by following this example: https://www.tensorflow.org/tutorials/load_data/tfrecord.

[ameesh-shah]

Hi Justin,

Thanks for your helpful comment. I have a similar question as the original author of this thread; namely, I'd like to find dataset elements that are of a specific type of scenario. For example, I'd like to go through each element of the dataset, and pick all the ones that take place in four-way intersections, and use those.

I followed the steps you outlined in your response, but found that the raw dictionaries that get produced from instantiating a tf_examples_dataset are difficult to parse; even loading them into a SimulatorState doesn't provide any insight as to how to determine whether a given element takes place on a highway, or an intersection, or so forth.

Are there particular fields or attributes of these elements that I should be looking for in order to determine which dataset element corresponds to which scenario?

If you need any further clarification from me, don't hesitate to let me know and I'll provide more details.

Thanks!