`waymax.config`#

Configs for Waymax Environments.

Module Contents#

Classes#

`CoordinateFrame`	Coordinate system used for data.
`DataFormat`	Data format for serialized records.
`DatasetConfig`	Configs for dataset/dataloader.
`ObservationConfig`	Config for the built-in Waymax Observation function.
`MetricsConfig`	Config for the built-in Waymax Metrics functions.
`LinearCombinationRewardConfig`	Config listing all metrics and their corresponding weights.
`ObjectType`	Types of objects that can be controlled by Waymax.
`SimAgentType`	Types of sim agents that can be used by Waymax.
`SimAgentConfig`	Config for sim agents.
`EnvironmentConfig`	Configs for Waymax Environments.
`WaymaxConfig`	Top level config for Waymax.

Attributes#

`WOD_1_0_0_TRAINING`
`WOD_1_0_0_VALIDATION`
`WOD_1_1_0_TRAINING`
`WOD_1_1_0_VALIDATION`
`WOD_1_2_0_TEST`

class waymax.config.CoordinateFrame(*args, **kwds)#

Bases: enum.Enum

Coordinate system used for data.

SDC = 'SDC'#

OBJECT = 'OBJECT'#

GLOBAL = 'GLOBAL'#

class waymax.config.DataFormat(*args, **kwds)#

Bases: enum.Enum

Data format for serialized records.

TFRECORD = 'TFRECORD'#

class waymax.config.DatasetConfig#

Configs for dataset/dataloader.

path#: Path/pattern for data, supporting sharded files with @.

data_format#: The format of data in path, string from womd_dataloader.DataFormat.

repeat#: Number of times to repeat the dataset. Set None to repeat infinitely.

batch_dims#: Batch dimensions in tuple format. Can be empty as ().

shuffle_seed#: Random seed. Set None to disable shuffle.

shuffle_buffer_size#: Buffer size for shuffling.

num_shards#: Number of shards for parallel loading, no effect on data returned.

deterministic#: Whether to use deterministic parallel processing.

include_sdc_paths#: Whether to include all valid future paths for SDC according to roadgraph connectivity from its starting position.

aggregate_timesteps#: Whether to aggregate keys from tf examples, need to set to True for parsing SimulatorState afterwards.

max_num_rg_points#: Max number of roadgraph points in data.

max_num_objects#: Max number of total objects in the scene. Set None to include all objects from the data.

num_paths#: Number of roadgraph traversal paths. Must be specified if include_sdc_paths is True.

num_points_per_path#: Number of points per roadgraph traversal path. Must be specified if include_sdc_paths is True.

drop_remainder#: Argument for tf.data.Dataset.batch. Set True to drop remainder if the last batch does not contain enough examples. Note training should not be affected since it is looping over all data for multiple epochs. For evaluation, it should be set to False to include all examples.

tf_data_service_address#: A string or a tuple indicating how to connect to the tf.data service. Used as the service argument for tf.data.experimental.service.distribute in dataloader.

distributed#: If True, the generated data batch will have an extra leading axis corresponding to the number of available devices. This is needed when the returned data is consumed by a pmap function.

batch_by_scenario#: If True, one example in a returned batch is the entire scenario containing all objects; if False, the dataset will treat individual object trajectories as a training example rather than an entire scenario.

path: str#

data_format: DataFormat#

repeat: int | None#

batch_dims: tuple[int, Ellipsis] = ()#

shuffle_seed: int | None#

shuffle_buffer_size: int = 1000#

num_shards: int = 4#

deterministic: bool = True#

include_sdc_paths: bool = False#

aggregate_timesteps: bool = True#

max_num_rg_points: int = 30000#

max_num_objects: int | None#

num_paths: int | None#

num_points_per_path: int | None#

drop_remainder: bool = True#

tf_data_service_address: str | None#

distributed: bool = False#

batch_by_scenario: bool = True#

class waymax.config.ObservationConfig#

Config for the built-in Waymax Observation function.

num_steps#: Number of trajectory history steps for observation.

roadgraph_top_k#: Number of closest roadgraph elements to get for observation.

coordinate_frame#: What coordinate frame the observation is using.

num_steps: int = 1#

roadgraph_top_k: int = 2000#

coordinate_frame: CoordinateFrame#

class waymax.config.MetricsConfig#

Config for the built-in Waymax Metrics functions.

metrics_to_run#: A list of metric names to run. Available metrics are: log_divergence, overlap, offroad, sdc_wrongway, sdc_off_route, sdc_progression, kinematic_infeasibility. Additional custom metrics can be registered with metric_factory.register_metric.

metrics_to_run: tuple[str, Ellipsis] = ('log_divergence', 'overlap', 'offroad')#

class waymax.config.LinearCombinationRewardConfig#

Config listing all metrics and their corresponding weights.

rewards#: Dictionary of metric names to floats indicating the weight of each metric to create a reward of a linear combination.

rewards: dict[str, float]#

class waymax.config.ObjectType(*args, **kwds)#

Bases: enum.Enum

Types of objects that can be controlled by Waymax.

SDC = 'SDC'#

MODELED = 'MODELED'#

VALID = 'VALID'#

NON_SDC = 'NON_SDC'#

class waymax.config.SimAgentType(*args, **kwds)#

Bases: enum.Enum

Types of sim agents that can be used by Waymax.

IDM = 'IDM'#

class waymax.config.SimAgentConfig#

Config for sim agents.

agent_type#: What sim agent algorithm to use.

controlled_objects#: Which objects the sim agent should control.

agent_type: SimAgentType#

controlled_objects: ObjectType#

class waymax.config.EnvironmentConfig#

Configs for Waymax Environments.

max_num_objects#: Max number of objects in the scene, should be consistent with DatasetConfig.max_num_objects.

init_steps#: Initial/warmup steps taken for the environment. For WOMD, this includes 10 warmup steps and 1 for the current step

controlled_object#: What type of objects are controlled.

compute_reward#: Whether to compute the reward. Computing the reward is expensive and unnecessary for cases like rollout data generation.

allow_new_objects_after_warmup#: Whether to allow new objects to appear after environment warmup. If this is set to False, all non-controlled objects that are invalid in the log at timestep t=warmup will be forever invalid in the simulation. This means that objects that appear out of occlusion after t=warmup will still be invalid as if they never appeared. If this is set to True, agents will be able to appear in the future simulated state if they appeared in the future logged state. Note when set to True, the environment could expect users to control objects even before their first appearance, users can either ignore or simply provide any invalid action for those objects.

observation#: Optional config for the provided observation function found at waymax/data/observation.py.

metrics#: Specifies which built-in Waymax metrics to run when calling env.metrics(…).

rewards#: Specifies the metrics and weights to create a reward as a linear combination of metrics.

sim_agents#: Configurations for sim agents used to control non user-controlled objects. Sim agents are applied in the order of that they are specified (if multiple sim agents control the same object, only the last sim agent will be applied for that object).

max_num_objects: int = 128#

init_steps: int = 11#

controlled_object: ObjectType#

compute_reward: bool = True#

allow_new_objects_after_warmup: bool = True#

observation: ObservationConfig | None#

metrics: MetricsConfig#

rewards: LinearCombinationRewardConfig#

sim_agents: Sequence[SimAgentConfig] | None#

__post_init__()#

class waymax.config.WaymaxConfig#

Top level config for Waymax.

data_config#: Data related configurations, including how to parse and load the data.

env_config#: Configurations about the environment itself, observation, and metrics.

data_config: DatasetConfig#

env_config: EnvironmentConfig#

__post_init__()#

waymax.config.WOD_1_0_0_TRAINING#

waymax.config.WOD_1_0_0_VALIDATION#

waymax.config.WOD_1_1_0_TRAINING#

waymax.config.WOD_1_1_0_VALIDATION#

waymax.config.WOD_1_2_0_TEST#

waymax.config#

Module Contents#

Classes#

Attributes#

`waymax.config`#