Fix traffic light timestamp desync when loading WOMD scenarios (#60)

[Abhishek21g]

Summary

Fixes waymax#60 — traffic light states can appear out of sync with logged agent motion in ~10% of WOMD v1.1/v1.2 scenarios (also reported in waymo-open-dataset#853).

Root cause: traffic_lights_from_womd_dict couples TL states to agents by timestep index only, ignoring traffic_light_state/*/timestamp_micros. When per-step TL timestamps diverge from the SDC agent timeline (often by one 100 ms step), visualization and index-based lookup show the wrong phase (e.g. agent moving while TL shows STOP).

Fix: New traffic_light_alignment module resamples TL states onto the SDC reference timeline when timestamps differ (integer offset search, then nearest-timestamp fallback). Hooked from simulator_state_from_womd_dict when time_key='all'.

Repro scenario

Shard	Index
training_tfexample.tfrecord-00008-of-01000	4

(Requires WOMD v1.2.1 motion download — not in repo testdata.)

Test plan

• python -m pytest waymax/dataloader/traffic_light_alignment_test.py -v — 7 passed
• python -m pytest waymax/dataloader/womd_factories_test.py -v — 23 passed (no regressions)
• python tools/validate_traffic_lights.py --tfrecord waymax/dataloader/testdata/tfrecord_with_routes --max-scenarios 3 — PASS (0 timestamp mismatches on bundled shard)

pip install -e .
python -m pytest waymax/dataloader/traffic_light_alignment_test.py -v
python tools/validate_traffic_lights.py --tfrecord <WOMD_SHARD> --scenario-index 4 --report tl_report.json

Design notes

• No-op when SDC and TL timestamp_micros already match (common case).
• Does not change TrafficLights dataclass shape — alignment runs at load time only.
• tools/validate_traffic_lights.py flags scenarios by timestamp_mismatch_count; kinematic heuristics only run when mismatches exist to limit false positives from CAUTION phases.

Related

Companion open-source validation tool (not part of this PR): waymax-sim-receipts — traffic-light consistency doctor.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.