waymax.dataloader.womd_dataloader#

Waymax data loading functions.

WOMD represents the data we typically use for simulation in our environment. See https://waymo.com/open/data/motion/tfexample for definitions on many of the data fields mentioned in this file.

Module Contents#

Functions#

preprocess_serialized_womd_data(→ dict[str, ...)

Parses serialized tf example into tf Tensor dict.

preprocess_womd_example(→ dict[str, tensorflow.Tensor])

Preprocesses dict of tf tensors, keyed by str.

simulator_state_generator(...)

Wrapper for SimulatorState iterator.

Attributes#

T

AUTOTUNE

waymax.dataloader.womd_dataloader.T#
waymax.dataloader.womd_dataloader.AUTOTUNE#
waymax.dataloader.womd_dataloader.preprocess_serialized_womd_data(serialized: bytes, config: waymax.config.DatasetConfig) → dict[str, tensorflow.Tensor]#

Parses serialized tf example into tf Tensor dict.

waymax.dataloader.womd_dataloader.preprocess_womd_example(example: dict[str, tensorflow.Tensor], aggregate_timesteps: bool, max_num_objects: int | None = None) → dict[str, tensorflow.Tensor]#

Preprocesses dict of tf tensors, keyed by str.

waymax.dataloader.womd_dataloader.simulator_state_generator(config: waymax.config.DatasetConfig) → Iterator[waymax.datatypes.simulator_state.SimulatorState]#

Wrapper for SimulatorState iterator.

This is the high level api for Waymax data loading that takes Waymax data config and outputs generator of SimulatorState.

Parameters:

config – dataset config.

Returns:

A SimulatorState iterator.