Enable fast WOSAC evals on large datasets with resampling

[daphne-cornelisse]

Description

Problem

When evaluating WOSAC on large datasets, the current implementation loads all target scenes into memory simultaneously, leading to excessive memory consumption and potential out-of-memory errors.

Solution

This PR refactors the evaluation pipeline to iterate through datasets incrementally, loading and processing scenes one at a time rather than all at once.

Usage

• Configure the sampling parameters like this:

[eval]
; Number of scenarios to process in each batch
wosac_batch_size = 10
; Target number of unique scenarios perform evaluation in
wosac_target_scenarios = 50
; Total number of scenarios to sample from
wosac_scenario_pool_size = 1000

• Run:

puffer eval puffer_drive --eval.wosac-realism-eval True

>>> Running WOSAC realism evaluation with training dataset. 

Collecting rollout 14/32...████████████████████████████                                                | 40/100 [00:10<00:16,  3.73%/s, n=20, batch=4]

...

Implentation

• Refactor resample scenes logic into a separate function
• Use the util function in training and eval code
• Simplify sampling logic: num_maps -> num_agents mapping is done under the hood in WOSAC eval so that users can think in scene units (easier for WOSAC eval purposes)

[WaelDLZ]

Nice !