feat/kitscenes-local-training-scripts

[FLagbusted]

Summary

Adds local training and evaluation scripts for the KITScenes grid search requested in #168. The scripts allow contributors to run exploratory single-GPU experiments without the full cluster training pipeline, while keeping local disk usage bounded by downloading and removing one scene at a time.

What's in here

• build_shard_split.py — builds the custom 585/65 scene split requested in Federated Model Training #168
• rotating_dataset.py — downloads, extracts, yields, and deletes one scene at a time via PinnedKITScenesDownloader
• train_one_combo.py — Bezier training with bfloat16 autocast
• train_one_combo_fm.py — Flow Matching training via compute_planner_loss (Fix/ part of #115 compute planner loss #172)
• eval_checkpoint.py — open-loop ADE/FDE evaluation at 3 s and 6.4 s
• constant_velocity_baseline.py — no-perception reference baseline using the last observed acceleration/curvature
  @gcordova10 thanks for this.

Epoch-0 results

Configuration: swin_v2_tiny + residual, custom 585-scene training / 65-scene validation split, 4,049 validation samples.

Run	ADE @ 3 s	FDE @ 3 s	ADE @ 6.4 s	FDE @ 6.4 s	Comfort violation rate
Constant-velocity baseline	0.883 m	2.662 m	—	—	—
Bezier epoch 0	3.087 m	8.086 m	10.789 m	27.907 m	0.362
Flow Matching epoch 0	2.572 m	6.594 m	8.764 m	22.341 m	1.000

The constant-velocity baseline was evaluated first on the same validation split before the learned models.
and 1 is expected for epoch 0 for Flow matching's comfort violation rate and which improves over furthur epochs.

At epoch 0, both learned models remain worse than the constant-velocity baseline. Flow Matching improves over Bezier on all reported ADE/FDE horizons, but still does not pass the baseline gate. Further training epochs are in progress.

Validation used the custom 65-scene split from the 585/65 split. 18 scenes were skipped because they had fewer than the required 129 ego poses, leaving 4,049 usable samples. These results are not directly comparable with runs using the frozen train_il validation scope.

This is a custom 585/65 split, not the frozen training/validation manifest used by the standard train_il pipeline, so these numbers are not directly comparable with results reported using --validation_scope full.

Hardware / training configuration

Experiments were run on an NVIDIA GeForce RTX 4060 8 GB with swin_v2_tiny + residual.

To fit the models on the 8 GB GPU, the runs use a reduced BEV grid and gradient checkpointing. cross_attn fusion requires a smaller BEV configuration (--bev_h 60 --bev_w 60, ≤4096 tokens).

Checkpoints

I will upload the Bezier and Flow Matching checkpoints from the 585-scene training split to Drive and add the links here later. These are the epoch-0 model weights produced by the local single-GPU training runs reported in this PR. Evaluation was performed on the corresponding 65-scene validation split defined by the same custom split.

Further participation

Additional training epochs are still in progress. More contributors with local GPUs are encouraged to run additional epochs/configurations and report results using the same split and evaluation procedure.

Relates-to: #168, #172

[FLagbusted]

for now I’ve uploaded the Bezier and Flow Matching checkpoints from the 585-scene training split to Drive. You can access them here:

https://drive.google.com/drive/folders/11dElZ0g_BDk7DELG_2WLt19krHGMOWCb?usp=sharing

These are the epoch-0 model weights produced by the local single-GPU training runs reported in this PR. Please note: the associated evaluation uses the 65-scene validation split from the same custom 585/65 split, not the frozen train_il validation manifest.

Continuing Training:
If anyone wants to test the pipeline on their own local setup right now and continue from where these left off, you can easily use these weights to kick off epoch 1 without having to start from scratch. I haven't wired a formal --resume flag into the scripts for this PR yet, but you can quickly hack it into train_one_combo_fm.py or train_one_combo.py by loading the states right before the training loop:

[gcordova10]

Read through the scripts — the rotating loader is careful work, and the cross_attn token
guard is the kind of check that saves someone an afternoon. Three small things.

1. Evaluation runs with the route switched off, which is upstream of the epoch-0 table.
eval_checkpoint.py calls the model with route_mask but no route_valid:

out = model(
    camera_tiles, map_context, visual_history, egomotion_history,
    route_mask=route_mask,
    projection=PinholeProjection(camera_params),
    geometry_type="pinhole",
    mode="infer",
)

ReactiveE2E gates the route with default=False when the flag is absent
(Model/model_components/reactive_e2e.py:193-203 and :218-223), so the two route channels
are multiplied by zero for every sample. map_valid defaults to True, so the map still
gets through — only the route is dropped. Both training scripts do pass
route_valid=route_valid, so the model is trained with the route and scored without it,
which makes any ADE from this path pessimistic. That is worth ruling out before reading
"both learned models remain worse than the baseline" as a statement about the models.
Passing the flag through at eval should close the gap.

2. The baseline is named for a different baseline. The README describes it correctly —
"holds the last observed (accel, curvature)" — but the filename and the printed label say
constant velocity, and the table above reports 0.883 under that name. Those are two
different references: constant velocity is accel and curvature at zero. On the frozen
40-scene holdout at dataset_version v3.3 they come out at 1.096 and 0.843 ADE@3s
respectively — different split from yours, so not the same numbers, but the gap between the
two is the point. Holding the last action is the harder bar of the two, so the label
undersells the result.

3. The retry in the baseline script has no success guard. rotating_dataset.py already
does this right:

if not download_success:
    print(f"Skipping scene {scene_id} due to persistent network failures.")
    continue

constant_velocity_baseline.py retries three times without the equivalent check, so a
network failure falls through and the scene is counted as skipped (too short). Its own
Skipped scenes line would then be wrong, and the ADE it prints would cover fewer samples
than intended with nothing in the output saying so. The eval path does not have this
problem — rotating_dataset.py is where the guard already lives.

[riita10069]

Local execution scripts are a very important initiative.
I want to align them with the consistency of other existing scripts.
I want to create a standard for how local execution should be done. I believe it should be executed with flyte run --local.
When different people train using different scripts, the results obtained will vary. My opinion is that it should be unified.

#187 migit be similar..?