How to limit the number of scenarios in a data_iter

[daphnecor]

Description

Is it possible to set the number of scenarios to a small number, such as 1?

import dataclasses
from waymax import config as _config
from waymax import dataloader

# Create the data iterator
config = dataclasses.replace(_config.WOD_1_0_0_VALIDATION, max_num_objects=max_num_objects)
data_iter = dataloader.simulator_state_generator(config=config)

# Get a scene (this takes ~ 80 seconds for the first time, < 1 seconds after that )
scenario = next(data_iter)

Motivation

It would be helpful to this because calling next(data_iter) with the default settings for the first time takes approximately 80 seconds. I assume this is due to the iterator being jitted.

Speeding this up would facilitate faster debugging/testing.

What I tried

I went through the config.py file and code in the dataloader folder but couldn't find a way to reduce the number of loaded scenes. Would it be possible to include something like max_num_rg_points for the number of files?

WOD_1_0_0_TRAINING = DatasetConfig(
 path='gs://waymo_open_dataset_motion_v_1_0_0/uncompressed/tf_example/training/training_tfexample.tfrecord@1000',
    max_num_rg_points=20000,
    data_format=DataFormat.TFRECORD,
)

Thank you for considering this issue and, of course, for releasing waymax!