Add RACE2 ablation A1 timing-only H20 40k EMA (weights, overlay, logs, config, source)

.gitattributes

@@ -40,3 +40,4 @@ acot_h10_40k/params/ocdbt.process_0/d/508e05f18e9bd423e6b0fa51e86e91b4 filter=lf
 acot_h10_40k/params/ocdbt.process_0/d/6717b350a38fef3e1f333634f8ffb1e9 filter=lfs diff=lfs merge=lfs -text
 acot_h10_40k/params/ocdbt.process_0/d/873cd1adcf38ed97d5dc1d23fd516e11 filter=lfs diff=lfs merge=lfs -text
 acot_h10_40k/params/ocdbt.process_0/d/f1e4a83571097cb44f24769987d96dd0 filter=lfs diff=lfs merge=lfs -text
+abl_timing_only_h20_40k/logs/train.log filter=lfs diff=lfs merge=lfs -text

abl_timing_only_h20_40k/README.md

@@ -0,0 +1,34 @@
+# RACE2 ablation A1 — timing loss only, no conditioning (VLABench, H=20, 40k EMA)
+
+Config `pi05_vlabench_race2abl_timing_only_h20` (model class `Race2TimingOnlyPI0`), experiment `abl_timing_only_h20_s42`.
+Full RACE2 recipe with `e_trans` / `field_mod` / `tau_gate` removed: training loss = flow-matching (no modulation)
++ 0.05 · timing BCE; no auxiliary base-FM loss, no teacher forcing. Serving is the plain pi0.5 sampler (no auxiliary pass).
+
+| | |
+|---|---|
+| Init | official VLABench pi0.5 teacher (`teacher_official/`, SHA256 `f6762a23…`) |
+| Training | 40,000 steps, seed 42, global batch 64 = 4 GPUs × micro 16 (no accumulation), warmup 500 → 5e-5 → cosine → 5e-6, AdamW (0.9, 0.95, wd 1e-10, clip 1.0), EMA 0.999, VLM (PaliGemma + SigLIP) frozen |
+| Data | `SeonghoonYu/vlabench_primitive_ft_lerobot_224` rev `5772f03c…` (5,000 ep), changepoints `vlabench_cps.npz` SHA256 `6e31b274…` |
+| Hardware | 4 × RTX 6000 Ada, 2026-09-14 19:04 → 09-15 13:12 UTC (18 h 08 m, ~1.6 s/step) |
+| Final train metrics (step 39,900) | fm 0.039, bce 0.137, peak±1 0.72, pres_acc 0.81 |
+
+## Evaluation (40k EMA, horizon = replan = 20, seed 7, tracks 1/2/3/4/6, 50 ep/task, 2,460 ep, physics errors excluded: 24)
+
+| Model | SR | IS | PS |
+|---|---:|---:|---:|
+| PlainFT H20 | 28.04 | 61.86 | 44.77 |
+| RACE2 full H20 (`h20_40k/`) | 32.86 | 61.48 | 47.84 |
+| **A1 timing-only H20 (this)** | **32.62** | **62.77** | **48.46** |
+
+Per track SR (A1): T1 52.41 · T2 28.71 · T3 25.88 · T4 23.56 · T6 32.54. Full details in `logs/eval_summary.json`.
+
+## Files
+
+- `model.safetensors` — full serving model: 603 frozen teacher VLM tensors + 225 EMA tensors (action expert, projections, RACE timing head). SHA256 in `SHA256SUMS`.
+- `overlay/ema.safetensors`, `overlay/metadata.pt` — the slim trainable-only EMA checkpoint written by the trainer (exactly what was evaluated).
+- `assets/vlabench/vlabench_ft_primitive/norm_stats.json` — normalization statistics (identical to the teacher's).
+- `training_config.json` — readable copy of the training config from `metadata.pt`.
+- `logs/train.log`, `logs/history.jsonl`, `logs/status.json`, `logs/commands.json`, `logs/launch_validation.json` — training driver records; `logs/eval.log`, `logs/eval_summary.json` — pooled evaluation; `logs/smoke_timing_only_h20.log` — pre-flight smoke.
+- `source/` — the ablation model, config and trainer files used (apply onto the OpenPI fork `Seonghoon-Yu/Libero` @ `43bc120d` with the `h25_support/openpi_overlay` config; also in the VLABench repo `race_eval/code`).
+
+Load with `--policy.config pi05_vlabench_race2abl_timing_only_h20 --policy.dir <this dir>` in the matching fork; the horizon comes from the config. Optimizer/resume state is not included.

abl_timing_only_h20_40k/SHA256SUMS

@@ -0,0 +1,2 @@
+9be111ce1db525928ec24a20fb102ce7632d9e6f81becd1b732328cf42cc9874  /root/vlabench/.cache/hf_stage_abl_timing_only_h20_40k/model.safetensors
+add103644143aa4cf09d8bacdd4ccf5987bedcd9e7fd4e1d68f1c463e62b5155  /root/vlabench/.cache/hf_stage_abl_timing_only_h20_40k/overlay/ema.safetensors

abl_timing_only_h20_40k/assets/vlabench/vlabench_ft_primitive/norm_stats.json

@@ -0,0 +1,280 @@
+{
+  "norm_stats": {
+    "state": {
+      "mean": [
+        0.02967369742691517,
+        0.39555251598358154,
+        0.25719010829925537,
+        0.4673101305961609,
+        -0.37718722224235535,
+        -0.11985507607460022,
+        0.44498270750045776,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.1581457555294037,
+        0.13765743374824524,
+        0.1370578110218048,
+        2.873211145401001,
+        0.625612199306488,
+        1.7478904724121094,
+        0.49696388840675354,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        -0.2880732418596744,
+        0.16642616012468936,
+        0.016219121195189656,
+        -3.1415927410125732,
+        -1.5606922270774841,
+        -3.1378225915908815,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.33706892974972724,
+        0.7168794604014606,
+        0.4705753956319765,
+        3.1403361039161677,
+        0.029222321224212644,
+        3.1353093175888063,
+        0.9998,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    },
+    "actions": {
+      "mean": [
+        -0.00008998554403660819,
+        0.004090703558176756,
+        0.0006838542176410556,
+        0.5465309023857117,
+        -0.004646406974643469,
+        -0.019831420853734016,
+        0.5484480857849121,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.012184085324406624,
+        0.013969186693429947,
+        0.021805619820952415,
+        3.0321147441864014,
+        0.028920045122504234,
+        0.7395531535148621,
+        0.497647225856781,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        -0.01572922980785374,
+        -0.011965612983703577,
+        -0.012324783468246447,
+        -6.28066982460022,
+        -0.07970612406730648,
+        -1.807010852813721,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.014919160246849028,
+        0.016237698739767104,
+        0.020379505300521905,
+        6.280669824600221,
+        0.06261572251319891,
+        1.7844429779052735,
+        0.9998,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    }
+  }
+}

abl_timing_only_h20_40k/logs/commands.json

@@ -0,0 +1,59 @@
+{
+  "method": "race2abl_timing_only",
+  "config": "pi05_vlabench_race2abl_timing_only_h20",
+  "exp_name": "abl_timing_only_h20_s42",
+  "horizon": 20,
+  "micro_batch_size": 16,
+  "batch_size": 64,
+  "world_size": 4,
+  "per_gpu_batch_size": 16,
+  "accumulation_steps": 1,
+  "gpus": [
+    0,
+    1,
+    2,
+    3
+  ],
+  "train_command": [
+    "/root/vlabench/runtime/openpi/.venv/bin/torchrun",
+    "--standalone",
+    "--nnodes=1",
+    "--nproc_per_node=4",
+    "scripts/race2_ablation/train_race2abl_timing_only_vlabench.py",
+    "pi05_vlabench_race2abl_timing_only_h20",
+    "--exp_name",
+    "abl_timing_only_h20_s42",
+    "--model.micro-batch-size",
+    "16"
+  ],
+  "eval_command": [
+    "/root/vlabench/runtime/openpi/.venv/bin/python",
+    "scripts/vlabench/eval_pool.py",
+    "--method",
+    "race2abl_timing_only",
+    "--horizon",
+    "20",
+    "--exp-name",
+    "abl_timing_only_h20_s42",
+    "--step",
+    "40000",
+    "--gpus",
+    "0,1,2,3",
+    "--jobs-per-gpu",
+    "2",
+    "--groups-per-job",
+    "2",
+    "--n-episode",
+    "50",
+    "--openpi-root",
+    "/root/vlabench/runtime/openpi",
+    "--vlabench-root",
+    "/root/vlabench/VLABench",
+    "--sim-python",
+    "/root/vlabench/runtime/sim-env/bin/python"
+  ],
+  "checkpoint": "/root/vlabench/runtime/openpi/checkpoints/pi05_vlabench_race2abl_timing_only_h20/abl_timing_only_h20_s42/40000/ema.safetensors",
+  "log_dir": "/root/vlabench/runtime/openpi/logs/h20/race2abl_timing_only_abl_timing_only_h20_s42",
+  "openpi_root": "/root/vlabench/runtime/openpi",
+  "setup_root": "/root/vlabench/race_eval/logs/h25_g4_setup"
+}

abl_timing_only_h20_40k/logs/eval.log

@@ -0,0 +1 @@
+[2026-09-15 16:07:19] POOL_EVALUATION_COMPLETE: /root/vlabench/runtime/openpi/evaluate_results/vlabench_h20/race2abl_timing_only_h20_abl_timing_only_h20_s42_40000_ema_replan20_n50_seed7/summary.json

abl_timing_only_h20_40k/logs/eval_summary.json

@@ -0,0 +1,1246 @@
+{
+  "status": "complete",
+  "spec": {
+    "config": "pi05_vlabench_race2abl_timing_only_h20",
+    "exp_name": "abl_timing_only_h20_s42",
+    "step": 40000,
+    "checkpoint": "/root/vlabench/runtime/openpi/checkpoints/pi05_vlabench_race2abl_timing_only_h20/abl_timing_only_h20_s42/40000",
+    "ema_sha256": "add103644143aa4cf09d8bacdd4ccf5987bedcd9e7fd4e1d68f1c463e62b5155",
+    "horizon": 20,
+    "replan_steps": 20,
+    "seed": 7,
+    "policy_seed": 7,
+    "policy_noise_rng": "independent_track_servers",
+    "requested_episodes": 50,
+    "tracks": {
+      "track_1_in_distribution": {
+        "select_painting": 50,
+        "select_book": 50,
+        "select_drink": 50,
+        "select_chemistry_tube": 50,
+        "select_poker": 50,
+        "select_mahjong": 50,
+        "select_toy": 50,
+        "select_fruit": 50,
+        "add_condiment": 50,
+        "insert_flower": 50
+      },
+      "track_2_cross_category": {
+        "select_painting": 50,
+        "select_fruit": 50,
+        "select_toy": 50,
+        "select_poker": 50,
+        "select_mahjong": 50,
+        "insert_flower": 10,
+        "select_chemistry_tube": 50,
+        "add_condiment": 50,
+        "select_drink": 50,
+        "select_book": 50
+      },
+      "track_3_common_sense": {
+        "select_painting": 50,
+        "select_book": 50,
+        "select_drink": 50,
+        "select_chemistry_tube": 50,
+        "select_toy": 50,
+        "select_fruit": 50,
+        "add_condiment": 50,
+        "insert_flower": 50,
+        "select_nth_largest_poker": 50,
+        "select_unique_type_mahjong": 50
+      },
+      "track_4_semantic_instruction": {
+        "select_painting": 50,
+        "select_book": 50,
+        "select_drink": 50,
+        "select_chemistry_tube": 50,
+        "select_poker": 50,
+        "select_mahjong": 50,
+        "select_toy": 50,
+        "select_fruit": 50,
+        "add_condiment": 50,
+        "insert_flower": 50
+      },
+      "track_6_unseen_texture": {
+        "select_painting": 50,
+        "select_book": 50,
+        "select_drink": 50,
+        "select_chemistry_tube": 50,
+        "select_poker": 50,
+        "select_mahjong": 50,
+        "select_toy": 50,
+        "select_fruit": 50,
+        "add_condiment": 50,
+        "insert_flower": 50
+      }
+    },
+    "track_sha256": {
+      "track_1_in_distribution": "61b07440c6f4272705bea0c8d721bc196e744232f48408075613cc9b6f7c8c04",
+      "track_2_cross_category": "cf004383de22d994f83126ca30a315cbfc6a92fd0dbb8d25fa4c185dd75a21fb",
+      "track_3_common_sense": "2aaa298580e74bec13555a6950b95241edf6c5264b498eeaf0e919837952ec6d",
+      "track_4_semantic_instruction": "9aff9dfc778ae35f299dce4c3aa07afbdc5275847490abd8037f3f5bdd9de411",
+      "track_6_unseen_texture": "249c1bca941b13267571b7641a91bcfbaffbc9c2b6b210d9ee2d2513324aa817"
+    },
+    "vlabench_root": "/root/vlabench/VLABench",
+    "execution_mode": "pool",
+    "pool_protocol": {
+      "name": "track_pool_v1",
+      "policy_noise_rng": "independent_track_servers",
+      "policy_seed": 7,
+      "jobs_per_gpu": 2,
+      "groups_per_job": 2,
+      "note": "Each track has an independent seed-7 policy server with a shared noise RNG among its task clients. Client ordering and partial resumes can change draws; this is not bit-identical to a single-server or four-shard evaluation."
+    },
+    "gpus": [
+      0,
+      1,
+      2,
+      3
+    ],
+    "env_limits": {
+      "TORCHINDUCTOR_COMPILE_THREADS": "2",
+      "OMP_NUM_THREADS": "1",
+      "OPENBLAS_NUM_THREADS": "1",
+      "MKL_NUM_THREADS": "1"
+    }
+  },
+  "tracks": {
+    "track_1_in_distribution": {
+      "status": "complete",
+      "track": "track_1_in_distribution",
+      "n_tasks": 10,
+      "episodes_per_task": {
+        "select_painting": 50,
+        "select_book": 50,
+        "select_drink": 50,
+        "select_chemistry_tube": 50,
+        "select_poker": 50,
+        "select_mahjong": 50,
+        "select_toy": 50,
+        "select_fruit": 50,
+        "add_condiment": 50,
+        "insert_flower": 50
+      },
+      "total_episodes": 500,
+      "per_task": {
+        "select_painting": {
+          "success_rate": 0.28,
+          "intention_score": 0.6,
+          "progress_score": 0.28
+        },
+        "select_book": {
+          "success_rate": 0.5918367346938775,
+          "intention_score": 0.8979591836734694,
+          "progress_score": 0.7244897959183674
+        },
+        "select_drink": {
+          "success_rate": 0.58,
+          "intention_score": 0.96,
+          "progress_score": 0.68
+        },
+        "select_chemistry_tube": {
+          "success_rate": 0.46938775510204084,
+          "intention_score": 0.0,
+          "progress_score": 0.6938775510204082
+        },
+        "select_poker": {
+          "success_rate": 0.56,
+          "intention_score": 0.8,
+          "progress_score": 0.7200000000000002
+        },
+        "select_mahjong": {
+          "success_rate": 0.5,
+          "intention_score": 0.8333333333333334,
+          "progress_score": 0.5
+        },
+        "select_toy": {
+          "success_rate": 0.56,
+          "intention_score": 0.22,
+          "progress_score": 0.77
+        },
+        "select_fruit": {
+          "success_rate": 0.7,
+          "intention_score": 0.98,
+          "progress_score": 0.84
+        },
+        "add_condiment": {
+          "success_rate": 0.64,
+          "intention_score": 1.0,
+          "progress_score": 0.7933333333333336
+        },
+        "insert_flower": {
+          "success_rate": 0.36,
+          "intention_score": 0.98,
+          "progress_score": 0.67
+        }
+      },
+      "avg_success_rate": 0.5241224489795918,
+      "avg_intention_score": 0.7271292517006802,
+      "avg_progress_score": 0.6671700680272109,
+      "error_policy": "exclude_physics_errors_v1",
+      "per_task_episode_accounting": {
+        "select_painting": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 50,
+          "excluded_episodes": 0,
+          "metric_denominator": 50,
+          "excluded_episode_ids": []
+        },
+        "select_book": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 49,
+          "excluded_episodes": 1,
+          "metric_denominator": 49,
+          "excluded_episode_ids": [
+            48
+          ]
+        },
+        "select_drink": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 50,
+          "excluded_episodes": 0,
+          "metric_denominator": 50,
+          "excluded_episode_ids": []
+        },
+        "select_chemistry_tube": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 49,
+          "excluded_episodes": 1,
+          "metric_denominator": 49,
+          "excluded_episode_ids": [
+            13
+          ]
+        },
+        "select_poker": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 50,
+          "excluded_episodes": 0,
+          "metric_denominator": 50,
+          "excluded_episode_ids": []
+        },
+        "select_mahjong": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 48,
+          "excluded_episodes": 2,
+          "metric_denominator": 48,
+          "excluded_episode_ids": [
+            11,
+            34
+          ]
+        },
+        "select_toy": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 50,
+          "excluded_episodes": 0,
+          "metric_denominator": 50,
+          "excluded_episode_ids": []
+        },
+        "select_fruit": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 50,
+          "excluded_episodes": 0,
+          "metric_denominator": 50,
+          "excluded_episode_ids": []
+        },
+        "add_condiment": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 50,
+          "excluded_episodes": 0,
+          "metric_denominator": 50,
+          "excluded_episode_ids": []
+        },
+        "insert_flower": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 50,
+          "excluded_episodes": 0,
+          "metric_denominator": 50,
+          "excluded_episode_ids": []
+        }
+      },
+      "valid_episodes": 496,
+      "excluded_episodes": 4
+    },
+    "track_2_cross_category": {
+      "status": "complete",
+      "track": "track_2_cross_category",
+      "n_tasks": 10,
+      "episodes_per_task": {
+        "select_painting": 50,
+        "select_fruit": 50,
+        "select_toy": 50,
+        "select_poker": 50,
+        "select_mahjong": 50,
+        "insert_flower": 10,
+        "select_chemistry_tube": 50,
+        "add_condiment": 50,
+        "select_drink": 50,
+        "select_book": 50
+      },
+      "total_episodes": 460,
+      "per_task": {
+        "select_painting": {
+          "success_rate": 0.24,
+          "intention_score": 0.48,
+          "progress_score": 0.24
+        },
+        "select_fruit": {
+          "success_rate": 0.8,
+          "intention_score": 0.96,
+          "progress_score": 0.87
+        },
+        "select_toy": {
+          "success_rate": 0.54,
+          "intention_score": 0.2,
+          "progress_score": 0.69
+        },
+        "select_poker": {
+          "success_rate": 0.54,
+          "intention_score": 0.78,
+          "progress_score": 0.68
+        },
+        "select_mahjong": {
+          "success_rate": 0.391304347826087,
+          "intention_score": 0.7608695652173914,
+          "progress_score": 0.43478260869565216
+        },
+        "insert_flower": {
+          "success_rate": 0.1,
+          "intention_score": 1.0,
+          "progress_score": 0.45
+        },
+        "select_chemistry_tube": {
+          "success_rate": 0.06,
+          "intention_score": 0.02,
+          "progress_score": 0.17
+        },
+        "add_condiment": {
+          "success_rate": 0.04,
+          "intention_score": 0.32,
+          "progress_score": 0.16666666666666663
+        },
+        "select_drink": {
+          "success_rate": 0.1,
+          "intention_score": 0.62,
+          "progress_score": 0.26
+        },
+        "select_book": {
+          "success_rate": 0.06,
+          "intention_score": 0.14,
+          "progress_score": 0.09
+        }
+      },
+      "avg_success_rate": 0.28713043478260875,
+      "avg_intention_score": 0.5280869565217391,
+      "avg_progress_score": 0.4051449275362319,
+      "error_policy": "exclude_physics_errors_v1",
+      "per_task_episode_accounting": {
+        "select_painting": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 50,
+          "excluded_episodes": 0,
+          "metric_denominator": 50,
+          "excluded_episode_ids": []
+        },
+        "select_fruit": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 50,
+          "excluded_episodes": 0,
+          "metric_denominator": 50,
+          "excluded_episode_ids": []
+        },
+        "select_toy": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 50,
+          "excluded_episodes": 0,
+          "metric_denominator": 50,
+          "excluded_episode_ids": []
+        },
+        "select_poker": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 50,
+          "excluded_episodes": 0,
+          "metric_denominator": 50,
+          "excluded_episode_ids": []
+        },
+        "select_mahjong": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 46,
+          "excluded_episodes": 4,
+          "metric_denominator": 46,
+          "excluded_episode_ids": [
+            1,
+            22,
+            25,
+            49
+          ]
+        },
+        "insert_flower": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 10,
+          "valid_episodes": 10,
+          "excluded_episodes": 0,
+          "metric_denominator": 10,
+          "excluded_episode_ids": []
+        },
+        "select_chemistry_tube": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 50,
+          "excluded_episodes": 0,
+          "metric_denominator": 50,
+          "excluded_episode_ids": []
+        },
+        "add_condiment": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 50,
+          "excluded_episodes": 0,
+          "metric_denominator": 50,
+          "excluded_episode_ids": []
+        },
+        "select_drink": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 50,
+          "excluded_episodes": 0,
+          "metric_denominator": 50,
+          "excluded_episode_ids": []
+        },
+        "select_book": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 50,
+          "excluded_episodes": 0,
+          "metric_denominator": 50,
+          "excluded_episode_ids": []
+        }
+      },
+      "valid_episodes": 456,
+      "excluded_episodes": 4
+    },
+    "track_3_common_sense": {
+      "status": "complete",
+      "track": "track_3_common_sense",
+      "n_tasks": 10,
+      "episodes_per_task": {
+        "select_painting": 50,
+        "select_book": 50,
+        "select_drink": 50,
+        "select_chemistry_tube": 50,
+        "select_toy": 50,
+        "select_fruit": 50,
+        "add_condiment": 50,
+        "insert_flower": 50,
+        "select_nth_largest_poker": 50,
+        "select_unique_type_mahjong": 50
+      },
+      "total_episodes": 500,
+      "per_task": {
+        "select_painting": {
+          "success_rate": 0.24,
+          "intention_score": 0.5,
+          "progress_score": 0.24
+        },
+        "select_book": {
+          "success_rate": 0.3617021276595745,
+          "intention_score": 0.7659574468085106,
+          "progress_score": 0.5319148936170213
+        },
+        "select_drink": {
+          "success_rate": 0.1,
+          "intention_score": 0.78,
+          "progress_score": 0.24
+        },
+        "select_chemistry_tube": {
+          "success_rate": 0.54,
+          "intention_score": 0.0,
+          "progress_score": 0.76
+        },
+        "select_toy": {
+          "success_rate": 0.24,
+          "intention_score": 0.28,
+          "progress_score": 0.53
+        },
+        "select_fruit": {
+          "success_rate": 0.62,
+          "intention_score": 0.94,
+          "progress_score": 0.78
+        },
+        "add_condiment": {
+          "success_rate": 0.1,
+          "intention_score": 0.7,
+          "progress_score": 0.2800000000000001
+        },
+        "insert_flower": {
+          "success_rate": 0.12,
+          "intention_score": 0.98,
+          "progress_score": 0.55
+        },
+        "select_nth_largest_poker": {
+          "success_rate": 0.1,
+          "intention_score": 0.22,
+          "progress_score": 0.12666666666666668
+        },
+        "select_unique_type_mahjong": {
+          "success_rate": 0.16666666666666666,
+          "intention_score": 0.5416666666666666,
+          "progress_score": 0.21875
+        }
+      },
+      "avg_success_rate": 0.2588368794326241,
+      "avg_intention_score": 0.5707624113475177,
+      "avg_progress_score": 0.42573315602836886,
+      "error_policy": "exclude_physics_errors_v1",
+      "per_task_episode_accounting": {
+        "select_painting": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 50,
+          "excluded_episodes": 0,
+          "metric_denominator": 50,
+          "excluded_episode_ids": []
+        },
+        "select_book": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 47,
+          "excluded_episodes": 3,
+          "metric_denominator": 47,
+          "excluded_episode_ids": [
+            16,
+            30,
+            40
+          ]
+        },
+        "select_drink": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 50,
+          "excluded_episodes": 0,
+          "metric_denominator": 50,
+          "excluded_episode_ids": []
+        },
+        "select_chemistry_tube": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 50,
+          "excluded_episodes": 0,
+          "metric_denominator": 50,
+          "excluded_episode_ids": []
+        },
+        "select_toy": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 50,
+          "excluded_episodes": 0,
+          "metric_denominator": 50,
+          "excluded_episode_ids": []
+        },
+        "select_fruit": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 50,
+          "excluded_episodes": 0,
+          "metric_denominator": 50,
+          "excluded_episode_ids": []
+        },
+        "add_condiment": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 50,
+          "excluded_episodes": 0,
+          "metric_denominator": 50,
+          "excluded_episode_ids": []
+        },
+        "insert_flower": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 50,
+          "excluded_episodes": 0,
+          "metric_denominator": 50,
+          "excluded_episode_ids": []
+        },
+        "select_nth_largest_poker": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 50,
+          "excluded_episodes": 0,
+          "metric_denominator": 50,
+          "excluded_episode_ids": []
+        },
+        "select_unique_type_mahjong": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 48,
+          "excluded_episodes": 2,
+          "metric_denominator": 48,
+          "excluded_episode_ids": [
+            3,
+            28
+          ]
+        }
+      },
+      "valid_episodes": 495,
+      "excluded_episodes": 5
+    },
+    "track_4_semantic_instruction": {
+      "status": "complete",
+      "track": "track_4_semantic_instruction",
+      "n_tasks": 10,
+      "episodes_per_task": {
+        "select_painting": 50,
+        "select_book": 50,
+        "select_drink": 50,
+        "select_chemistry_tube": 50,
+        "select_poker": 50,
+        "select_mahjong": 50,
+        "select_toy": 50,
+        "select_fruit": 50,
+        "add_condiment": 50,
+        "insert_flower": 50
+      },
+      "total_episodes": 500,
+      "per_task": {
+        "select_painting": {
+          "success_rate": 0.18,
+          "intention_score": 0.56,
+          "progress_score": 0.18
+        },
+        "select_book": {
+          "success_rate": 0.375,
+          "intention_score": 0.7083333333333334,
+          "progress_score": 0.4791666666666667
+        },
+        "select_drink": {
+          "success_rate": 0.18,
+          "intention_score": 0.86,
+          "progress_score": 0.33
+        },
+        "select_chemistry_tube": {
+          "success_rate": 0.24,
+          "intention_score": 0.0,
+          "progress_score": 0.53
+        },
+        "select_poker": {
+          "success_rate": 0.3,
+          "intention_score": 0.78,
+          "progress_score": 0.38666666666666666
+        },
+        "select_mahjong": {
+          "success_rate": 0.2608695652173913,
+          "intention_score": 0.8043478260869565,
+          "progress_score": 0.34782608695652173
+        },
+        "select_toy": {
+          "success_rate": 0.16,
+          "intention_score": 0.28,
+          "progress_score": 0.5
+        },
+        "select_fruit": {
+          "success_rate": 0.4,
+          "intention_score": 0.92,
+          "progress_score": 0.65
+        },
+        "add_condiment": {
+          "success_rate": 0.18,
+          "intention_score": 0.7,
+          "progress_score": 0.33333333333333326
+        },
+        "insert_flower": {
+          "success_rate": 0.08,
+          "intention_score": 0.98,
+          "progress_score": 0.5
+        }
+      },
+      "avg_success_rate": 0.23558695652173914,
+      "avg_intention_score": 0.6592681159420289,
+      "avg_progress_score": 0.4236992753623189,
+      "error_policy": "exclude_physics_errors_v1",
+      "per_task_episode_accounting": {
+        "select_painting": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 50,
+          "excluded_episodes": 0,
+          "metric_denominator": 50,
+          "excluded_episode_ids": []
+        },
+        "select_book": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 48,
+          "excluded_episodes": 2,
+          "metric_denominator": 48,
+          "excluded_episode_ids": [
+            2,
+            5
+          ]
+        },
+        "select_drink": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 50,
+          "excluded_episodes": 0,
+          "metric_denominator": 50,
+          "excluded_episode_ids": []
+        },
+        "select_chemistry_tube": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 50,
+          "excluded_episodes": 0,
+          "metric_denominator": 50,
+          "excluded_episode_ids": []
+        },
+        "select_poker": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 50,
+          "excluded_episodes": 0,
+          "metric_denominator": 50,
+          "excluded_episode_ids": []
+        },
+        "select_mahjong": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 46,
+          "excluded_episodes": 4,
+          "metric_denominator": 46,
+          "excluded_episode_ids": [
+            7,
+            19,
+            20,
+            39
+          ]
+        },
+        "select_toy": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 50,
+          "excluded_episodes": 0,
+          "metric_denominator": 50,
+          "excluded_episode_ids": []
+        },
+        "select_fruit": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 50,
+          "excluded_episodes": 0,
+          "metric_denominator": 50,
+          "excluded_episode_ids": []
+        },
+        "add_condiment": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 50,
+          "excluded_episodes": 0,
+          "metric_denominator": 50,
+          "excluded_episode_ids": []
+        },
+        "insert_flower": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 50,
+          "excluded_episodes": 0,
+          "metric_denominator": 50,
+          "excluded_episode_ids": []
+        }
+      },
+      "valid_episodes": 494,
+      "excluded_episodes": 6
+    },
+    "track_6_unseen_texture": {
+      "status": "complete",
+      "track": "track_6_unseen_texture",
+      "n_tasks": 10,
+      "episodes_per_task": {
+        "select_painting": 50,
+        "select_book": 50,
+        "select_drink": 50,
+        "select_chemistry_tube": 50,
+        "select_poker": 50,
+        "select_mahjong": 50,
+        "select_toy": 50,
+        "select_fruit": 50,
+        "add_condiment": 50,
+        "insert_flower": 50
+      },
+      "total_episodes": 500,
+      "per_task": {
+        "select_painting": {
+          "success_rate": 0.24,
+          "intention_score": 0.5,
+          "progress_score": 0.24
+        },
+        "select_book": {
+          "success_rate": 0.3333333333333333,
+          "intention_score": 0.7083333333333334,
+          "progress_score": 0.5
+        },
+        "select_drink": {
+          "success_rate": 0.34,
+          "intention_score": 0.92,
+          "progress_score": 0.57
+        },
+        "select_chemistry_tube": {
+          "success_rate": 0.4,
+          "intention_score": 0.02,
+          "progress_score": 0.62
+        },
+        "select_poker": {
+          "success_rate": 0.38,
+          "intention_score": 0.64,
+          "progress_score": 0.4266666666666667
+        },
+        "select_mahjong": {
+          "success_rate": 0.3404255319148936,
+          "intention_score": 0.7659574468085106,
+          "progress_score": 0.3829787234042553
+        },
+        "select_toy": {
+          "success_rate": 0.3,
+          "intention_score": 0.14,
+          "progress_score": 0.6
+        },
+        "select_fruit": {
+          "success_rate": 0.48,
+          "intention_score": 0.94,
+          "progress_score": 0.66
+        },
+        "add_condiment": {
+          "success_rate": 0.38,
+          "intention_score": 0.98,
+          "progress_score": 0.6533333333333335
+        },
+        "insert_flower": {
+          "success_rate": 0.06,
+          "intention_score": 0.92,
+          "progress_score": 0.36
+        }
+      },
+      "avg_success_rate": 0.3253758865248227,
+      "avg_intention_score": 0.6534290780141845,
+      "avg_progress_score": 0.5012978723404256,
+      "error_policy": "exclude_physics_errors_v1",
+      "per_task_episode_accounting": {
+        "select_painting": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 50,
+          "excluded_episodes": 0,
+          "metric_denominator": 50,
+          "excluded_episode_ids": []
+        },
+        "select_book": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 48,
+          "excluded_episodes": 2,
+          "metric_denominator": 48,
+          "excluded_episode_ids": [
+            8,
+            9
+          ]
+        },
+        "select_drink": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 50,
+          "excluded_episodes": 0,
+          "metric_denominator": 50,
+          "excluded_episode_ids": []
+        },
+        "select_chemistry_tube": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 50,
+          "excluded_episodes": 0,
+          "metric_denominator": 50,
+          "excluded_episode_ids": []
+        },
+        "select_poker": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 50,
+          "excluded_episodes": 0,
+          "metric_denominator": 50,
+          "excluded_episode_ids": []
+        },
+        "select_mahjong": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 47,
+          "excluded_episodes": 3,
+          "metric_denominator": 47,
+          "excluded_episode_ids": [
+            24,
+            25,
+            41
+          ]
+        },
+        "select_toy": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 50,
+          "excluded_episodes": 0,
+          "metric_denominator": 50,
+          "excluded_episode_ids": []
+        },
+        "select_fruit": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 50,
+          "excluded_episodes": 0,
+          "metric_denominator": 50,
+          "excluded_episode_ids": []
+        },
+        "add_condiment": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 50,
+          "excluded_episodes": 0,
+          "metric_denominator": 50,
+          "excluded_episode_ids": []
+        },
+        "insert_flower": {
+          "error_policy": "exclude_physics_errors_v1",
+          "attempted_episodes": 50,
+          "valid_episodes": 50,
+          "excluded_episodes": 0,
+          "metric_denominator": 50,
+          "excluded_episode_ids": []
+        }
+      },
+      "valid_episodes": 495,
+      "excluded_episodes": 5
+    }
+  },
+  "n_tracks": 5,
+  "n_tasks": 50,
+  "total_episodes": 2460,
+  "shards": {
+    "track_1_in_distribution": {
+      "gpu": 0,
+      "output": "shards/track_1/race2abl_timing_only_h20_abl_timing_only_h20_s42_40000_ema_replan20_n50_seed7_subset_9d9b3e45",
+      "summary": "shards/track_1/race2abl_timing_only_h20_abl_timing_only_h20_s42_40000_ema_replan20_n50_seed7_subset_9d9b3e45/summary.json",
+      "completed_tasks": "shards/track_1/race2abl_timing_only_h20_abl_timing_only_h20_s42_40000_ema_replan20_n50_seed7_subset_9d9b3e45/track_1_in_distribution/completed_tasks",
+      "job": {
+        "gpu": 0,
+        "pid": 713409,
+        "command": [
+          "/root/vlabench/runtime/openpi/.venv/bin/python",
+          "/root/vlabench/runtime/openpi/scripts/vlabench/eval_h25.py",
+          "--pool-worker",
+          "--method",
+          "race2abl_timing_only",
+          "--horizon",
+          "20",
+          "--exp-name",
+          "abl_timing_only_h20_s42",
+          "--step",
+          "40000",
+          "--gpu",
+          "0",
+          "--n-episode",
+          "50",
+          "--groups",
+          "2",
+          "--tracks",
+          "1",
+          "--replan-steps",
+          "20",
+          "--openpi-root",
+          "/root/vlabench/runtime/openpi",
+          "--vlabench-root",
+          "/root/vlabench/VLABench",
+          "--checkpoint-root",
+          "/root/vlabench/runtime/openpi/checkpoints",
+          "--output-root",
+          "/root/vlabench/runtime/openpi/evaluate_results/vlabench_h20/race2abl_timing_only_h20_abl_timing_only_h20_s42_40000_ema_replan20_n50_seed7/shards/track_1",
+          "--policy-python",
+          "/root/vlabench/runtime/openpi/.venv/bin/python",
+          "--sim-python",
+          "/root/vlabench/runtime/sim-env/bin/python",
+          "--startup-timeout",
+          "1200.0",
+          "--warmup-timeout",
+          "3600.0",
+          "--task-timeout",
+          "21600.0"
+        ],
+        "status": "complete",
+        "started_at": 1789477934.447253,
+        "env_limits": {
+          "TORCHINDUCTOR_COMPILE_THREADS": "2",
+          "OMP_NUM_THREADS": "1",
+          "OPENBLAS_NUM_THREADS": "1",
+          "MKL_NUM_THREADS": "1"
+        },
+        "finished_at": 1789487279.4960856
+      }
+    },
+    "track_2_cross_category": {
+      "gpu": 1,
+      "output": "shards/track_2/race2abl_timing_only_h20_abl_timing_only_h20_s42_40000_ema_replan20_n50_seed7_subset_baa99750",
+      "summary": "shards/track_2/race2abl_timing_only_h20_abl_timing_only_h20_s42_40000_ema_replan20_n50_seed7_subset_baa99750/summary.json",
+      "completed_tasks": "shards/track_2/race2abl_timing_only_h20_abl_timing_only_h20_s42_40000_ema_replan20_n50_seed7_subset_baa99750/track_2_cross_category/completed_tasks",
+      "job": {
+        "gpu": 1,
+        "pid": 713410,
+        "command": [
+          "/root/vlabench/runtime/openpi/.venv/bin/python",
+          "/root/vlabench/runtime/openpi/scripts/vlabench/eval_h25.py",
+          "--pool-worker",
+          "--method",
+          "race2abl_timing_only",
+          "--horizon",
+          "20",
+          "--exp-name",
+          "abl_timing_only_h20_s42",
+          "--step",
+          "40000",
+          "--gpu",
+          "1",
+          "--n-episode",
+          "50",
+          "--groups",
+          "2",
+          "--tracks",
+          "2",
+          "--replan-steps",
+          "20",
+          "--openpi-root",
+          "/root/vlabench/runtime/openpi",
+          "--vlabench-root",
+          "/root/vlabench/VLABench",
+          "--checkpoint-root",
+          "/root/vlabench/runtime/openpi/checkpoints",
+          "--output-root",
+          "/root/vlabench/runtime/openpi/evaluate_results/vlabench_h20/race2abl_timing_only_h20_abl_timing_only_h20_s42_40000_ema_replan20_n50_seed7/shards/track_2",
+          "--policy-python",
+          "/root/vlabench/runtime/openpi/.venv/bin/python",
+          "--sim-python",
+          "/root/vlabench/runtime/sim-env/bin/python",
+          "--startup-timeout",
+          "1200.0",
+          "--warmup-timeout",
+          "3600.0",
+          "--task-timeout",
+          "21600.0"
+        ],
+        "status": "complete",
+        "started_at": 1789477934.4488165,
+        "env_limits": {
+          "TORCHINDUCTOR_COMPILE_THREADS": "2",
+          "OMP_NUM_THREADS": "1",
+          "OPENBLAS_NUM_THREADS": "1",
+          "MKL_NUM_THREADS": "1"
+        },
+        "finished_at": 1789487215.1072192
+      }
+    },
+    "track_3_common_sense": {
+      "gpu": 2,
+      "output": "shards/track_3/race2abl_timing_only_h20_abl_timing_only_h20_s42_40000_ema_replan20_n50_seed7_subset_21b01caf",
+      "summary": "shards/track_3/race2abl_timing_only_h20_abl_timing_only_h20_s42_40000_ema_replan20_n50_seed7_subset_21b01caf/summary.json",
+      "completed_tasks": "shards/track_3/race2abl_timing_only_h20_abl_timing_only_h20_s42_40000_ema_replan20_n50_seed7_subset_21b01caf/track_3_common_sense/completed_tasks",
+      "job": {
+        "gpu": 2,
+        "pid": 713411,
+        "command": [
+          "/root/vlabench/runtime/openpi/.venv/bin/python",
+          "/root/vlabench/runtime/openpi/scripts/vlabench/eval_h25.py",
+          "--pool-worker",
+          "--method",
+          "race2abl_timing_only",
+          "--horizon",
+          "20",
+          "--exp-name",
+          "abl_timing_only_h20_s42",
+          "--step",
+          "40000",
+          "--gpu",
+          "2",
+          "--n-episode",
+          "50",
+          "--groups",
+          "2",
+          "--tracks",
+          "3",
+          "--replan-steps",
+          "20",
+          "--openpi-root",
+          "/root/vlabench/runtime/openpi",
+          "--vlabench-root",
+          "/root/vlabench/VLABench",
+          "--checkpoint-root",
+          "/root/vlabench/runtime/openpi/checkpoints",
+          "--output-root",
+          "/root/vlabench/runtime/openpi/evaluate_results/vlabench_h20/race2abl_timing_only_h20_abl_timing_only_h20_s42_40000_ema_replan20_n50_seed7/shards/track_3",
+          "--policy-python",
+          "/root/vlabench/runtime/openpi/.venv/bin/python",
+          "--sim-python",
+          "/root/vlabench/runtime/sim-env/bin/python",
+          "--startup-timeout",
+          "1200.0",
+          "--warmup-timeout",
+          "3600.0",
+          "--task-timeout",
+          "21600.0"
+        ],
+        "status": "complete",
+        "started_at": 1789477934.450404,
+        "env_limits": {
+          "TORCHINDUCTOR_COMPILE_THREADS": "2",
+          "OMP_NUM_THREADS": "1",
+          "OPENBLAS_NUM_THREADS": "1",
+          "MKL_NUM_THREADS": "1"
+        },
+        "finished_at": 1789488167.2772944
+      }
+    },
+    "track_4_semantic_instruction": {
+      "gpu": 3,
+      "output": "shards/track_4/race2abl_timing_only_h20_abl_timing_only_h20_s42_40000_ema_replan20_n50_seed7_subset_5fae63be",
+      "summary": "shards/track_4/race2abl_timing_only_h20_abl_timing_only_h20_s42_40000_ema_replan20_n50_seed7_subset_5fae63be/summary.json",
+      "completed_tasks": "shards/track_4/race2abl_timing_only_h20_abl_timing_only_h20_s42_40000_ema_replan20_n50_seed7_subset_5fae63be/track_4_semantic_instruction/completed_tasks",
+      "job": {
+        "gpu": 3,
+        "pid": 713412,
+        "command": [
+          "/root/vlabench/runtime/openpi/.venv/bin/python",
+          "/root/vlabench/runtime/openpi/scripts/vlabench/eval_h25.py",
+          "--pool-worker",
+          "--method",
+          "race2abl_timing_only",
+          "--horizon",
+          "20",
+          "--exp-name",
+          "abl_timing_only_h20_s42",
+          "--step",
+          "40000",
+          "--gpu",
+          "3",
+          "--n-episode",
+          "50",
+          "--groups",
+          "2",
+          "--tracks",
+          "4",
+          "--replan-steps",
+          "20",
+          "--openpi-root",
+          "/root/vlabench/runtime/openpi",
+          "--vlabench-root",
+          "/root/vlabench/VLABench",
+          "--checkpoint-root",
+          "/root/vlabench/runtime/openpi/checkpoints",
+          "--output-root",
+          "/root/vlabench/runtime/openpi/evaluate_results/vlabench_h20/race2abl_timing_only_h20_abl_timing_only_h20_s42_40000_ema_replan20_n50_seed7/shards/track_4",
+          "--policy-python",
+          "/root/vlabench/runtime/openpi/.venv/bin/python",
+          "--sim-python",
+          "/root/vlabench/runtime/sim-env/bin/python",
+          "--startup-timeout",
+          "1200.0",
+          "--warmup-timeout",
+          "3600.0",
+          "--task-timeout",
+          "21600.0"
+        ],
+        "status": "complete",
+        "started_at": 1789477934.452021,
+        "env_limits": {
+          "TORCHINDUCTOR_COMPILE_THREADS": "2",
+          "OMP_NUM_THREADS": "1",
+          "OPENBLAS_NUM_THREADS": "1",
+          "MKL_NUM_THREADS": "1"
+        },
+        "finished_at": 1789487807.8258038
+      }
+    },
+    "track_6_unseen_texture": {
+      "gpu": 0,
+      "output": "shards/track_6/race2abl_timing_only_h20_abl_timing_only_h20_s42_40000_ema_replan20_n50_seed7_subset_498302e6",
+      "summary": "shards/track_6/race2abl_timing_only_h20_abl_timing_only_h20_s42_40000_ema_replan20_n50_seed7_subset_498302e6/summary.json",
+      "completed_tasks": "shards/track_6/race2abl_timing_only_h20_abl_timing_only_h20_s42_40000_ema_replan20_n50_seed7_subset_498302e6/track_6_unseen_texture/completed_tasks",
+      "job": {
+        "gpu": 0,
+        "pid": 713413,
+        "command": [
+          "/root/vlabench/runtime/openpi/.venv/bin/python",
+          "/root/vlabench/runtime/openpi/scripts/vlabench/eval_h25.py",
+          "--pool-worker",
+          "--method",
+          "race2abl_timing_only",
+          "--horizon",
+          "20",
+          "--exp-name",
+          "abl_timing_only_h20_s42",
+          "--step",
+          "40000",
+          "--gpu",
+          "0",
+          "--n-episode",
+          "50",
+          "--groups",
+          "2",
+          "--tracks",
+          "6",
+          "--replan-steps",
+          "20",
+          "--openpi-root",
+          "/root/vlabench/runtime/openpi",
+          "--vlabench-root",
+          "/root/vlabench/VLABench",
+          "--checkpoint-root",
+          "/root/vlabench/runtime/openpi/checkpoints",
+          "--output-root",
+          "/root/vlabench/runtime/openpi/evaluate_results/vlabench_h20/race2abl_timing_only_h20_abl_timing_only_h20_s42_40000_ema_replan20_n50_seed7/shards/track_6",
+          "--policy-python",
+          "/root/vlabench/runtime/openpi/.venv/bin/python",
+          "--sim-python",
+          "/root/vlabench/runtime/sim-env/bin/python",
+          "--startup-timeout",
+          "1200.0",
+          "--warmup-timeout",
+          "3600.0",
+          "--task-timeout",
+          "21600.0"
+        ],
+        "status": "complete",
+        "started_at": 1789477934.4536605,
+        "env_limits": {
+          "TORCHINDUCTOR_COMPILE_THREADS": "2",
+          "OMP_NUM_THREADS": "1",
+          "OPENBLAS_NUM_THREADS": "1",
+          "MKL_NUM_THREADS": "1"
+        },
+        "finished_at": 1789488436.9035316
+      }
+    }
+  },
+  "pool_manifest": "pool_manifest.json",
+  "error_policy": "exclude_physics_errors_v1",
+  "valid_episodes": 2436,
+  "excluded_episodes": 24,
+  "aggregation": "Equal weight for five tracks; equal task weight within each track. Recorded PhysicsError episodes are excluded from each task's metric denominator; ordinary unsuccessful rollouts remain included.",
+  "avg_success_rate": 0.3262105212482773,
+  "avg_intention_score": 0.62773516270523,
+  "avg_progress_score": 0.4846090598589113
+}

abl_timing_only_h20_40k/logs/history.jsonl

@@ -0,0 +1,5 @@
+{"method": "race2abl_timing_only", "config": "pi05_vlabench_race2abl_timing_only_h20", "exp_name": "abl_timing_only_h20_s42", "horizon": 20, "micro_batch_size": 16, "batch_size": 64, "world_size": 4, "per_gpu_batch_size": 16, "accumulation_steps": 1, "gpus": [0, 1, 2, 3], "train_command": ["/root/vlabench/runtime/openpi/.venv/bin/torchrun", "--standalone", "--nnodes=1", "--nproc_per_node=4", "scripts/race2_ablation/train_race2abl_timing_only_vlabench.py", "pi05_vlabench_race2abl_timing_only_h20", "--exp_name", "abl_timing_only_h20_s42", "--model.micro-batch-size", "16"], "eval_command": ["/root/vlabench/runtime/openpi/.venv/bin/python", "scripts/vlabench/eval_pool.py", "--method", "race2abl_timing_only", "--horizon", "20", "--exp-name", "abl_timing_only_h20_s42", "--step", "40000", "--gpus", "0,1,2,3", "--jobs-per-gpu", "2", "--groups-per-job", "2", "--n-episode", "50", "--openpi-root", "/root/vlabench/runtime/openpi", "--vlabench-root", "/root/vlabench/VLABench", "--sim-python", "/root/vlabench/runtime/sim-env/bin/python"], "checkpoint": "/root/vlabench/runtime/openpi/checkpoints/pi05_vlabench_race2abl_timing_only_h20/abl_timing_only_h20_s42/40000/ema.safetensors", "log_dir": "/root/vlabench/runtime/openpi/logs/h20/race2abl_timing_only_abl_timing_only_h20_s42", "openpi_root": "/root/vlabench/runtime/openpi", "setup_root": "/root/vlabench/race_eval/logs/h25_g4_setup", "stage": "launch_validated", "time": "2026-09-14T19:04:28.566790+00:00", "pid": 659463, "launch_validation": "/root/vlabench/runtime/openpi/logs/h20/race2abl_timing_only_abl_timing_only_h20_s42/launch_validation.json"}
+{"method": "race2abl_timing_only", "config": "pi05_vlabench_race2abl_timing_only_h20", "exp_name": "abl_timing_only_h20_s42", "horizon": 20, "micro_batch_size": 16, "batch_size": 64, "world_size": 4, "per_gpu_batch_size": 16, "accumulation_steps": 1, "gpus": [0, 1, 2, 3], "train_command": ["/root/vlabench/runtime/openpi/.venv/bin/torchrun", "--standalone", "--nnodes=1", "--nproc_per_node=4", "scripts/race2_ablation/train_race2abl_timing_only_vlabench.py", "pi05_vlabench_race2abl_timing_only_h20", "--exp_name", "abl_timing_only_h20_s42", "--model.micro-batch-size", "16"], "eval_command": ["/root/vlabench/runtime/openpi/.venv/bin/python", "scripts/vlabench/eval_pool.py", "--method", "race2abl_timing_only", "--horizon", "20", "--exp-name", "abl_timing_only_h20_s42", "--step", "40000", "--gpus", "0,1,2,3", "--jobs-per-gpu", "2", "--groups-per-job", "2", "--n-episode", "50", "--openpi-root", "/root/vlabench/runtime/openpi", "--vlabench-root", "/root/vlabench/VLABench", "--sim-python", "/root/vlabench/runtime/sim-env/bin/python"], "checkpoint": "/root/vlabench/runtime/openpi/checkpoints/pi05_vlabench_race2abl_timing_only_h20/abl_timing_only_h20_s42/40000/ema.safetensors", "log_dir": "/root/vlabench/runtime/openpi/logs/h20/race2abl_timing_only_abl_timing_only_h20_s42", "openpi_root": "/root/vlabench/runtime/openpi", "setup_root": "/root/vlabench/race_eval/logs/h25_g4_setup", "stage": "training", "time": "2026-09-14T19:04:28.568297+00:00", "pid": 659463, "child_pid": 659473, "log": "/root/vlabench/runtime/openpi/logs/h20/race2abl_timing_only_abl_timing_only_h20_s42/train.log"}
+{"method": "race2abl_timing_only", "config": "pi05_vlabench_race2abl_timing_only_h20", "exp_name": "abl_timing_only_h20_s42", "horizon": 20, "micro_batch_size": 16, "batch_size": 64, "world_size": 4, "per_gpu_batch_size": 16, "accumulation_steps": 1, "gpus": [0, 1, 2, 3], "train_command": ["/root/vlabench/runtime/openpi/.venv/bin/torchrun", "--standalone", "--nnodes=1", "--nproc_per_node=4", "scripts/race2_ablation/train_race2abl_timing_only_vlabench.py", "pi05_vlabench_race2abl_timing_only_h20", "--exp_name", "abl_timing_only_h20_s42", "--model.micro-batch-size", "16"], "eval_command": ["/root/vlabench/runtime/openpi/.venv/bin/python", "scripts/vlabench/eval_pool.py", "--method", "race2abl_timing_only", "--horizon", "20", "--exp-name", "abl_timing_only_h20_s42", "--step", "40000", "--gpus", "0,1,2,3", "--jobs-per-gpu", "2", "--groups-per-job", "2", "--n-episode", "50", "--openpi-root", "/root/vlabench/runtime/openpi", "--vlabench-root", "/root/vlabench/VLABench", "--sim-python", "/root/vlabench/runtime/sim-env/bin/python"], "checkpoint": "/root/vlabench/runtime/openpi/checkpoints/pi05_vlabench_race2abl_timing_only_h20/abl_timing_only_h20_s42/40000/ema.safetensors", "log_dir": "/root/vlabench/runtime/openpi/logs/h20/race2abl_timing_only_abl_timing_only_h20_s42", "openpi_root": "/root/vlabench/runtime/openpi", "setup_root": "/root/vlabench/race_eval/logs/h25_g4_setup", "stage": "trained", "time": "2026-09-15T13:12:13.170535+00:00", "pid": 659463, "checkpoint_validation": {"step": 40000, "config": "pi05_vlabench_race2abl_timing_only_h20", "ema": "/root/vlabench/runtime/openpi/checkpoints/pi05_vlabench_race2abl_timing_only_h20/abl_timing_only_h20_s42/40000/ema.safetensors", "n_tensors": 225, "horizon": 20, "micro_batch_size": 16}}
+{"method": "race2abl_timing_only", "config": "pi05_vlabench_race2abl_timing_only_h20", "exp_name": "abl_timing_only_h20_s42", "horizon": 20, "micro_batch_size": 16, "batch_size": 64, "world_size": 4, "per_gpu_batch_size": 16, "accumulation_steps": 1, "gpus": [0, 1, 2, 3], "train_command": ["/root/vlabench/runtime/openpi/.venv/bin/torchrun", "--standalone", "--nnodes=1", "--nproc_per_node=4", "scripts/race2_ablation/train_race2abl_timing_only_vlabench.py", "pi05_vlabench_race2abl_timing_only_h20", "--exp_name", "abl_timing_only_h20_s42", "--model.micro-batch-size", "16"], "eval_command": ["/root/vlabench/runtime/openpi/.venv/bin/python", "scripts/vlabench/eval_pool.py", "--method", "race2abl_timing_only", "--horizon", "20", "--exp-name", "abl_timing_only_h20_s42", "--step", "40000", "--gpus", "0,1,2,3", "--jobs-per-gpu", "2", "--groups-per-job", "2", "--n-episode", "50", "--openpi-root", "/root/vlabench/runtime/openpi", "--vlabench-root", "/root/vlabench/VLABench", "--sim-python", "/root/vlabench/runtime/sim-env/bin/python"], "checkpoint": "/root/vlabench/runtime/openpi/checkpoints/pi05_vlabench_race2abl_timing_only_h20/abl_timing_only_h20_s42/40000/ema.safetensors", "log_dir": "/root/vlabench/runtime/openpi/logs/h20/race2abl_timing_only_abl_timing_only_h20_s42", "openpi_root": "/root/vlabench/runtime/openpi", "setup_root": "/root/vlabench/race_eval/logs/h25_g4_setup", "stage": "evaluating", "time": "2026-09-15T13:12:13.175344+00:00", "pid": 659463, "child_pid": 713408, "log": "/root/vlabench/runtime/openpi/logs/h20/race2abl_timing_only_abl_timing_only_h20_s42/eval.log"}
+{"method": "race2abl_timing_only", "config": "pi05_vlabench_race2abl_timing_only_h20", "exp_name": "abl_timing_only_h20_s42", "horizon": 20, "micro_batch_size": 16, "batch_size": 64, "world_size": 4, "per_gpu_batch_size": 16, "accumulation_steps": 1, "gpus": [0, 1, 2, 3], "train_command": ["/root/vlabench/runtime/openpi/.venv/bin/torchrun", "--standalone", "--nnodes=1", "--nproc_per_node=4", "scripts/race2_ablation/train_race2abl_timing_only_vlabench.py", "pi05_vlabench_race2abl_timing_only_h20", "--exp_name", "abl_timing_only_h20_s42", "--model.micro-batch-size", "16"], "eval_command": ["/root/vlabench/runtime/openpi/.venv/bin/python", "scripts/vlabench/eval_pool.py", "--method", "race2abl_timing_only", "--horizon", "20", "--exp-name", "abl_timing_only_h20_s42", "--step", "40000", "--gpus", "0,1,2,3", "--jobs-per-gpu", "2", "--groups-per-job", "2", "--n-episode", "50", "--openpi-root", "/root/vlabench/runtime/openpi", "--vlabench-root", "/root/vlabench/VLABench", "--sim-python", "/root/vlabench/runtime/sim-env/bin/python"], "checkpoint": "/root/vlabench/runtime/openpi/checkpoints/pi05_vlabench_race2abl_timing_only_h20/abl_timing_only_h20_s42/40000/ema.safetensors", "log_dir": "/root/vlabench/runtime/openpi/logs/h20/race2abl_timing_only_abl_timing_only_h20_s42", "openpi_root": "/root/vlabench/runtime/openpi", "setup_root": "/root/vlabench/race_eval/logs/h25_g4_setup", "stage": "complete", "time": "2026-09-15T16:07:19.098081+00:00", "pid": 659463}

abl_timing_only_h20_40k/logs/launch_validation.json

@@ -0,0 +1,153 @@
+{
+  "status": "verified",
+  "method": "race2abl_timing_only",
+  "config": "pi05_vlabench_race2abl_timing_only_h20",
+  "exp_name": "abl_timing_only_h20_s42",
+  "horizon": 20,
+  "micro_batch_size": 16,
+  "batch_size": 64,
+  "world_size": 4,
+  "per_gpu_batch_size": 16,
+  "accumulation_steps": 1,
+  "gpus": [
+    0,
+    1,
+    2,
+    3
+  ],
+  "train_command": [
+    "/root/vlabench/runtime/openpi/.venv/bin/torchrun",
+    "--standalone",
+    "--nnodes=1",
+    "--nproc_per_node=4",
+    "scripts/race2_ablation/train_race2abl_timing_only_vlabench.py",
+    "pi05_vlabench_race2abl_timing_only_h20",
+    "--exp_name",
+    "abl_timing_only_h20_s42",
+    "--model.micro-batch-size",
+    "16"
+  ],
+  "eval_command": [
+    "/root/vlabench/runtime/openpi/.venv/bin/python",
+    "scripts/vlabench/eval_pool.py",
+    "--method",
+    "race2abl_timing_only",
+    "--horizon",
+    "20",
+    "--exp-name",
+    "abl_timing_only_h20_s42",
+    "--step",
+    "40000",
+    "--gpus",
+    "0,1,2,3",
+    "--jobs-per-gpu",
+    "2",
+    "--groups-per-job",
+    "2",
+    "--n-episode",
+    "50",
+    "--openpi-root",
+    "/root/vlabench/runtime/openpi",
+    "--vlabench-root",
+    "/root/vlabench/VLABench",
+    "--sim-python",
+    "/root/vlabench/runtime/sim-env/bin/python"
+  ],
+  "checkpoint": "/root/vlabench/runtime/openpi/checkpoints/pi05_vlabench_race2abl_timing_only_h20/abl_timing_only_h20_s42/40000/ema.safetensors",
+  "log_dir": "/root/vlabench/runtime/openpi/logs/h20/race2abl_timing_only_abl_timing_only_h20_s42",
+  "openpi_root": "/root/vlabench/runtime/openpi",
+  "setup_root": "/root/vlabench/race_eval/logs/h25_g4_setup",
+  "setup_reports": {
+    "dataset_validation.json": {
+      "path": "/root/vlabench/race_eval/logs/h25_g4_setup/dataset_validation.json",
+      "sha256": "d3ebe82d7a50cbc0ca22fd6e6360d0fd673700180ec7078a08b59d6c94877838",
+      "status": "verified"
+    },
+    "teacher_validation.json": {
+      "path": "/root/vlabench/race_eval/logs/h25_g4_setup/teacher_validation.json",
+      "sha256": "09ad2c05835dd5922b84d412bfcb657b6434c8eaa2af70cbc52c46f3860f6be6",
+      "status": "verified"
+    },
+    "cache_ready.json": {
+      "path": "/root/vlabench/race_eval/logs/h25_g4_setup/cache_ready.json",
+      "sha256": "41d61eb1996d17ba7bba60b0b37fd1c5b9efcb59edc4ee934dea91c8480b2c84",
+      "status": "verified"
+    }
+  },
+  "recipe": {
+    "name": "pi05_vlabench_race2abl_timing_only_h20",
+    "exp_name": "abl_timing_only_h20_s42",
+    "pytorch_weight_path": "checkpoints/pi05_vlabench_official/pytorch",
+    "seed": 42,
+    "batch_size": 64,
+    "pytorch_training_precision": "bfloat16",
+    "optimizer": {
+      "b1": 0.9,
+      "b2": 0.95,
+      "eps": 1e-08,
+      "weight_decay": 1e-10,
+      "clip_gradient_norm": 1.0
+    },
+    "lr_schedule": {
+      "warmup_steps": 500,
+      "peak_lr": 5e-05,
+      "decay_steps": 39500,
+      "decay_lr": 5e-06
+    },
+    "ema_decay": 0.999,
+    "model": {
+      "action_dim": 32,
+      "action_horizon": 20,
+      "max_token_len": 200,
+      "dtype": "bfloat16",
+      "paligemma_variant": "gemma_2b",
+      "action_expert_variant": "gemma_300m",
+      "pi05": true,
+      "discrete_state_input": false,
+      "gemma_depth": 18,
+      "sigma": 1.0,
+      "support": 3.0,
+      "jitter_prob": 0.5,
+      "lambda_timing": 0.05,
+      "lambda_base_fm": 0.0,
+      "read_tau": 1.0,
+      "n_tau": 10,
+      "tau_gate_init": 0.0,
+      "tau_gate_lr_mult": 20.0,
+      "timing_proj_dim": 256,
+      "timing_num_heads": 4,
+      "fast_serving": 0,
+      "kv_layer": -1
+    },
+    "data": {
+      "repo_id": "SeonghoonYu/vlabench_primitive_ft_lerobot_224",
+      "assets": {
+        "assets_dir": "./assets/pi05_vlabench",
+        "asset_id": "vlabench/vlabench_ft_primitive"
+      },
+      "phase_labels_path": null,
+      "phase_label_horizon": null,
+      "cps_path": "assets/gauss_timing/vlabench_cps.npz",
+      "gauss_sigma": 1.0,
+      "gauss_support": 3.0,
+      "base_config": {
+        "repo_id": null,
+        "asset_id": null,
+        "use_quantile_norm": false,
+        "action_sequence_keys": [
+          "actions"
+        ],
+        "prompt_from_task": true
+      }
+    }
+  },
+  "gpu_uuids": [
+    "GPU-4feb92e1-28ff-042a-fea0-1a3883fa3012",
+    "GPU-75266c29-c27f-b973-931f-918bed539110",
+    "GPU-330794ad-628b-fa4d-a46d-6ca53a33d45d",
+    "GPU-cca7b1a0-5e7c-6c6e-e23f-84ff32cca083"
+  ],
+  "hf_lerobot_home": "/root/vlabench/.cache/huggingface/lerobot",
+  "hf_datasets_cache": "/root/vlabench/.cache/huggingface/datasets",
+  "openpi_data_home": "/root/vlabench/.cache/openpi"
+}

abl_timing_only_h20_40k/logs/smoke_timing_only_h20.log

@@ -0,0 +1,26 @@
+Race2TimingOnlyPI0: conditioning OFF | lambda_base_fm 0.0 lambda_timing 0.05 | head params=790,017 | dtype=bfloat16
+WARNING:root:
+The dataset you requested (SeonghoonYu/vlabench_primitive_ft_lerobot_224) is in 2.0 format.
+While current version of LeRobot is backward-compatible with it, the version of your dataset still uses global
+stats instead of per-episode stats. Update your dataset stats to the new format using this command:
+```
+python lerobot/common/datasets/v21/convert_dataset_v20_to_v21.py --repo-id=SeonghoonYu/vlabench_primitive_ft_lerobot_224
+```
+
+If you encounter a problem, contact LeRobot maintainers on [Discord](https://discord.com/invite/s3KuuzsPFb)
+or open an [issue on GitHub](https://github.com/huggingface/lerobot/issues/new/choose).
+
+WARNING:root:
+The dataset you requested (SeonghoonYu/vlabench_primitive_ft_lerobot_224) is in 2.0 format.
+While current version of LeRobot is backward-compatible with it, the version of your dataset still uses global
+stats instead of per-episode stats. Update your dataset stats to the new format using this command:
+```
+python lerobot/common/datasets/v21/convert_dataset_v20_to_v21.py --repo-id=SeonghoonYu/vlabench_primitive_ft_lerobot_224
+```
+
+If you encounter a problem, contact LeRobot maintainers on [Discord](https://discord.com/invite/s3KuuzsPFb)
+or open an [issue on GitHub](https://github.com/huggingface/lerobot/issues/new/choose).
+
+train: loss 0.4261 fm 0.3879 base_fm 0.00000 bce 0.7641 | grad race_head 1.02e+00 expert 1.31e+01 field 0.00e+00 VLM 0.00e+00
+serve: |a|max 6.812 finite=True | determinism 0.00e+00 | parity vs reference 6.68e-03 | per-call 191 ms (first 246 ms)
+SMOKE_OK

abl_timing_only_h20_40k/logs/status.json

@@ -0,0 +1,62 @@
+{
+  "method": "race2abl_timing_only",
+  "config": "pi05_vlabench_race2abl_timing_only_h20",
+  "exp_name": "abl_timing_only_h20_s42",
+  "horizon": 20,
+  "micro_batch_size": 16,
+  "batch_size": 64,
+  "world_size": 4,
+  "per_gpu_batch_size": 16,
+  "accumulation_steps": 1,
+  "gpus": [
+    0,
+    1,
+    2,
+    3
+  ],
+  "train_command": [
+    "/root/vlabench/runtime/openpi/.venv/bin/torchrun",
+    "--standalone",
+    "--nnodes=1",
+    "--nproc_per_node=4",
+    "scripts/race2_ablation/train_race2abl_timing_only_vlabench.py",
+    "pi05_vlabench_race2abl_timing_only_h20",
+    "--exp_name",
+    "abl_timing_only_h20_s42",
+    "--model.micro-batch-size",
+    "16"
+  ],
+  "eval_command": [
+    "/root/vlabench/runtime/openpi/.venv/bin/python",
+    "scripts/vlabench/eval_pool.py",
+    "--method",
+    "race2abl_timing_only",
+    "--horizon",
+    "20",
+    "--exp-name",
+    "abl_timing_only_h20_s42",
+    "--step",
+    "40000",
+    "--gpus",
+    "0,1,2,3",
+    "--jobs-per-gpu",
+    "2",
+    "--groups-per-job",
+    "2",
+    "--n-episode",
+    "50",
+    "--openpi-root",
+    "/root/vlabench/runtime/openpi",
+    "--vlabench-root",
+    "/root/vlabench/VLABench",
+    "--sim-python",
+    "/root/vlabench/runtime/sim-env/bin/python"
+  ],
+  "checkpoint": "/root/vlabench/runtime/openpi/checkpoints/pi05_vlabench_race2abl_timing_only_h20/abl_timing_only_h20_s42/40000/ema.safetensors",
+  "log_dir": "/root/vlabench/runtime/openpi/logs/h20/race2abl_timing_only_abl_timing_only_h20_s42",
+  "openpi_root": "/root/vlabench/runtime/openpi",
+  "setup_root": "/root/vlabench/race_eval/logs/h25_g4_setup",
+  "stage": "complete",
+  "time": "2026-09-15T16:07:19.098081+00:00",
+  "pid": 659463
+}

abl_timing_only_h20_40k/logs/train.log

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5dfd1c2b27e0066952db43e252bfebf0a4fb255725b7c5b7af023dee2dc9e9cf
+size 58618421

abl_timing_only_h20_40k/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9be111ce1db525928ec24a20fb102ce7632d9e6f81becd1b732328cf42cc9874
+size 7474154628

abl_timing_only_h20_40k/overlay/ema.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:add103644143aa4cf09d8bacdd4ccf5987bedcd9e7fd4e1d68f1c463e62b5155
+size 1627379860

abl_timing_only_h20_40k/overlay/metadata.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a800acd1b2546db5f6ec4d56f6e7d14dd6cb49bcade22b3fb817ab277e1521ea
+size 3699

abl_timing_only_h20_40k/provenance.json

@@ -0,0 +1,18 @@
+{
+  "method": "RACE2 ablation A1 (timing loss only, no conditioning)",
+  "config": "pi05_vlabench_race2abl_timing_only_h20",
+  "model_class": "Race2TimingOnlyPI0",
+  "exp_name": "abl_timing_only_h20_s42",
+  "step": 40000, "weights": "EMA", "ema_decay": 0.999, "action_horizon": 20, "training_seed": 42,
+  "world_size": 4, "micro_batch_size": 16, "accumulation_steps": 1, "global_batch": 64,
+  "ema_overlay_sha256": "add103644143aa4cf09d8bacdd4ccf5987bedcd9e7fd4e1d68f1c463e62b5155",
+  "materialized_model_sha256": "9be111ce1db525928ec24a20fb102ce7632d9e6f81becd1b732328cf42cc9874",
+  "materialization": {"frozen_tensors": 603, "ema_tensors": 225, "total_tensors": 828},
+  "teacher_sha256": "f6762a23e4374bcbd8f3a7297f159b02d253a335f125eb93595d431d37607172",
+  "dataset": {"repo_id": "SeonghoonYu/vlabench_primitive_ft_lerobot_224", "revision": "5772f03c3af8df38ee75d8f71cab0075563ed12a"},
+  "changepoints_sha256": "6e31b27469d69fa65f2e3c41bcbf016b0f8b426048ca2d88d32367d1aa5aa4fc",
+  "openpi_base_repository": "https://github.com/Seonghoon-Yu/Libero", "openpi_base_revision": "43bc120d863a2a67fafd3ee2464029ac37040231",
+  "trained": "2026-09-14T19:04Z to 2026-09-15T13:12Z", "evaluated": "2026-09-15T13:12Z to 2026-09-15T16:07Z",
+  "evaluation": {"replan_steps": 20, "seed": 7, "tracks": [1, 2, 3, 4, 6], "n_episode": 2460, "physics_errors_excluded": 24,
+                 "success_rate_pct": 32.62, "intention_score_pct": 62.77, "progress_score_pct": 48.46}
+}

abl_timing_only_h20_40k/source/scripts/race2_ablation/train_race2abl_timing_only_vlabench.py

@@ -0,0 +1,447 @@
+"""RACE2 ablation trainer: A1 — timing loss only, NO conditioning (e_trans/field_mod/tau_gate absent, plain FM, pi0.5 serving).
+
+Runs ONLY pi05_vlabench_race2abl_timing_only_h{H} presets (model Race2TimingOnlyPI0). Direct use:
+
+  cd /root/vlabench/runtime/openpi && source /root/vlabench/race_eval/logs/h30_h35_setup/env.sh
+  .venv/bin/torchrun --standalone --nnodes=1 --nproc_per_node=2 \
+      scripts/race2_ablation/train_race2abl_timing_only_vlabench.py pi05_vlabench_race2abl_timing_only_h30 \
+      --exp_name timing_only_h30_s42 --model.micro-batch-size 32 [--resume]
+
+or through the driver (train -> final-EMA check -> pooled eval):
+  .venv/bin/python scripts/vlabench/run_h25_job.py --method race2abl_timing_only --horizon 30 --gpus 0,1 \
+      --exp-name timing_only_h30_s42 --micro-batch-size 32
+
+
+Copy of scripts/race2/train_race2_vlabench.py (same recipe machinery: fp32 master
+weights, EMA, cosine LR, slim overlay checkpoints) with the model checks relaxed to the
+ablation classes in openpi.models_pytorch.race2_ablation:
+  Race2TimingOnlyPI0  pi05_vlabench_race2abl_timing_only_h{H}  timing loss only, no field
+  Race2NoAuxPI0       pi05_vlabench_race2abl_noaux_h{H}        full minus auxiliary loss
+The full-method trainer is left untouched.
+
+Original header follows.
+--------------------------------------------------------------------------------------
+Training entrypoint for RACE v2 on VLABench (Race2PI0).
+
+Standalone counterpart of scripts/race/train_race.py for the v2 timing head
+(cross-attention to the cached VLM values -> self-attention among the action slots ->
+linear -> sigmoid; openpi.models_pytorch.race2). Same data pipeline, objective and
+VLABench recipe machinery as the v1 trainer; kept as a separate file so the two methods
+never share (or break) each other's code.
+
+Recipe machinery (added for the VLABench campaign, usable for any RACE config):
+  * fp32 MASTER WEIGHTS for the trainable set. The model itself stays bf16 for compute,
+    exactly as before, but AdamW steps fp32 copies that are written back after every
+    step. Without this, the cosine tail (lr 5e-6 -> ~2.5e-4 relative updates) is rounded
+    away by the bf16 parameter spacing (2^-8).
+  * EMA of the master weights (config.ema_decay). The decay is warmed up as
+    min(decay, (1 + step) / (10 + step)) so the early 1k checkpoints are not dominated
+    by the initial weights.
+  * warmup + cosine LR straight from config.lr_schedule (CosineDecaySchedule).
+  * SLIM CHECKPOINTS. Only the trainable tensors are written (the VLM prefix is frozen
+    and identical to the teacher):
+        <step>/ema.safetensors    EMA weights, model dtype  (~1.6 GB)   every save
+        <step>/raw.safetensors    non-EMA weights            (~1.6 GB)   every RAW_EVERY-th save
+        <step>/train_state.pt     fp32 master + EMA + optimizer (~13 GB) newest step only
+        <step>/metadata.pt, <step>/assets/<asset_id>/norm_stats.json
+    A full serving checkpoint (teacher VLM + overlay) is materialised on demand with
+    scripts/race/materialize_ckpt.py (config-driven, shared with v1).
+
+Usage:
+  torchrun --standalone --nnodes=1 --nproc_per_node=<n> \
+      scripts/race2_ablation/train_race2abl_timing_only_vlabench.py pi05_vlabench_race2abl_timing_only_h30 --exp_name timing_only_h30_s42 [--resume]
+"""
+
+import contextlib
+import dataclasses
+import logging
+import os
+import pathlib
+import shutil
+import sys
+import time
+
+sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent))  # scripts/ for train_pytorch
+
+import numpy as np
+import safetensors.torch
+import torch
+import torch.distributed as dist
+import tqdm
+
+import jax
+
+import openpi.models.race2_ablation_config as race2_ablation_config  # noqa: F401  (config registry)
+import openpi.models.model as _model
+import openpi.shared.normalize as _normalize
+import openpi.training.config as _config
+import openpi.training.data_loader as _data
+import openpi.training.gauss_timing_data as gauss_data
+
+import train_pytorch as base
+from vlabench.checkpoint_validation import validate_resume
+
+# module name fragments that are NEW in this model (absent from the teacher ckpt)
+FRESH_OK = ("race_head", "e_trans", "field_mod", "tau_gate")
+ABLATION_MODELS = ("Race2TimingOnlyPI0",)
+# a raw (non-EMA) overlay is written on every RAW_EVERY-th save (5 x 1k = every 5k steps)
+RAW_EVERY = 5
+TRAIN_STATE = "train_state.pt"
+
+
+def load_teacher(model, config):
+    model_path = os.path.join(config.pytorch_weight_path, "model.safetensors")
+    logging.info(f"Loading teacher weights from: {model_path}")
+    missing, unexpected = safetensors.torch.load_model(model, model_path, strict=False)
+    if unexpected:
+        raise ValueError(f"Unexpected keys in teacher checkpoint: {sorted(unexpected)[:5]}")
+    stale = [k for k in missing if not any(m in k for m in FRESH_OK)]
+    if stale:
+        raise ValueError(f"Teacher checkpoint missing expected policy keys: {sorted(stale)[:5]}")
+    logging.info(f"Loaded teacher; {len(missing)} new-module tensors keep fresh init")
+
+
+def make_loader(config, world_size):
+    """(observation, actions, b) triples; mirrors data_loader.create_torch_data_loader
+    but keeps the extra 'b' key instead of dropping it at Observation.from_dict."""
+    data_config = config.data.create(config.assets_dirs, config.model)
+    dataset = _data.create_torch_dataset(data_config, config.model.action_horizon, config.model)
+    dataset = _data.transform_dataset(dataset, data_config)
+    sampler = None
+    local_bs = config.batch_size
+    if dist.is_initialized():
+        sampler = torch.utils.data.distributed.DistributedSampler(
+            dataset, num_replicas=dist.get_world_size(), rank=dist.get_rank(), shuffle=True, drop_last=True
+        )
+        local_bs = config.batch_size // world_size
+    loader = _data.TorchDataLoader(
+        dataset,
+        local_batch_size=local_bs,
+        sharding=None,
+        shuffle=(sampler is None),
+        sampler=sampler,
+        num_batches=None,
+        num_workers=config.num_workers,
+        seed=config.seed,
+        framework="pytorch",
+    )
+    return loader, sampler, data_config
+
+
+class MasterState:
+    """fp32 master copies + EMA of the trainable parameters.
+
+    Every rank holds an identical copy: DDP averages the bf16 grads, so identical AdamW
+    steps run on identical fp32 masters. `sync_to_model` writes the masters back into the
+    live (bf16 / fp32) parameters after each optimizer step."""
+
+    def __init__(self, named_params, ema_decay: float | None, ema_warmup: bool = True):
+        self.names = [n for n, _ in named_params]
+        self.params = [p for _, p in named_params]
+        self.master = []
+        for p in self.params:
+            m = p.detach().clone().float().requires_grad_(True)
+            m.grad = torch.zeros_like(m)
+            self.master.append(m)
+        self.ema = [m.detach().clone() for m in self.master]
+        self.decay = 0.0 if ema_decay is None else float(ema_decay)
+        self.ema_warmup = ema_warmup
+
+    def grads_to_master(self):
+        for p, m in zip(self.params, self.master):
+            if p.grad is None:
+                m.grad.zero_()
+            else:
+                m.grad.copy_(p.grad)
+
+    @torch.no_grad()
+    def sync_to_model(self):
+        for p, m in zip(self.params, self.master):
+            p.copy_(m)
+
+    @torch.no_grad()
+    def update_ema(self, step: int) -> float:
+        d = self.decay
+        if self.ema_warmup:
+            d = min(d, (1.0 + step) / (10.0 + step))
+        torch._foreach_lerp_(self.ema, self.master, 1.0 - d)
+        return d
+
+    @torch.no_grad()
+    def overlay(self, tensors) -> dict[str, torch.Tensor]:
+        """Trainable tensors in the model's own dtypes (bf16 expert, fp32 norms/head), on CPU."""
+        return {n: t.detach().to(p.dtype).cpu().contiguous() for n, t, p in zip(self.names, tensors, self.params)}
+
+    @torch.no_grad()
+    def state_dict(self):
+        return {
+            "names": list(self.names),
+            "master": [m.detach().cpu() for m in self.master],
+            "ema": [e.detach().cpu() for e in self.ema],
+        }
+
+    @torch.no_grad()
+    def load_state_dict(self, sd):
+        if list(sd["names"]) != self.names:
+            raise ValueError("train_state parameter set differs from the model's trainable set")
+        for m, v in zip(self.master, sd["master"]):
+            m.copy_(v)
+        for e, v in zip(self.ema, sd["ema"]):
+            e.copy_(v)
+
+
+def latest_resumable_step(ckpt_dir: pathlib.Path):
+    if not ckpt_dir.exists():
+        return None
+    steps = [
+        int(d.name) for d in ckpt_dir.iterdir()
+        if d.is_dir() and d.name.isdigit() and (d / TRAIN_STATE).is_file()
+    ]
+    return max(steps) if steps else None
+
+
+def save_checkpoint(config, step, ms: MasterState, opt, data_config, is_main):
+    if not is_main or step == 0 or step % config.save_interval != 0:
+        return
+    final = config.checkpoint_dir / str(step)
+    tmp = config.checkpoint_dir / f"tmp_{step}"
+    shutil.rmtree(tmp, ignore_errors=True)
+    tmp.mkdir(parents=True)
+    t0 = time.time()
+    # evaluation weights only from save_overlays_from_step on; the resume state is always written
+    overlays = step >= int(getattr(config, "save_overlays_from_step", 0) or 0)
+    if overlays:
+        safetensors.torch.save_file(ms.overlay(ms.ema), str(tmp / "ema.safetensors"))
+        if (step // config.save_interval) % RAW_EVERY == 0 or step == config.num_train_steps:
+            safetensors.torch.save_file(ms.overlay(ms.master), str(tmp / "raw.safetensors"))
+        norm_stats = data_config.norm_stats
+        if norm_stats is not None and data_config.asset_id is not None:
+            _normalize.save(tmp / "assets" / data_config.asset_id, norm_stats)
+    torch.save(
+        {
+            "global_step": step,
+            "config": dataclasses.asdict(config),
+            "timestamp": time.time(),
+            "ema_decay": ms.decay,
+            "format": "trainable-only overlays; materialise with scripts/race/materialize_ckpt.py",
+        },
+        tmp / "metadata.pt",
+    )
+    torch.save({"step": step, **ms.state_dict(), "optimizer": opt.state_dict()}, tmp / TRAIN_STATE)
+    if final.exists():
+        shutil.rmtree(final)
+    tmp.rename(final)
+    logging.info(f"Saved checkpoint {final} ({'overlays + ' if overlays else ''}resume state, {time.time() - t0:.0f}s)")
+    # keep the full resume state (fp32 master + EMA + optimizer, ~13 GB) for the newest step only;
+    # a stale step directory that holds nothing but its resume state is removed altogether
+    for d in config.checkpoint_dir.iterdir():
+        if d.is_dir() and d.name.isdigit() and int(d.name) < step and (d / TRAIN_STATE).is_file():
+            (d / TRAIN_STATE).unlink()
+            if not (d / "ema.safetensors").is_file() and not (d / "raw.safetensors").is_file():
+                shutil.rmtree(d)
+                logging.info(f"Pruned resume-only checkpoint dir: {d}")
+            else:
+                logging.info(f"Pruned stale train state: {d / TRAIN_STATE}")
+
+
+def train_loop(config: _config.TrainConfig):
+    use_ddp, local_rank, device = base.setup_ddp()
+    is_main = (not use_ddp) or (dist.get_rank() == 0)
+    base.set_seed(config.seed, local_rank)
+    world_size = dist.get_world_size() if use_ddp else 1
+    logging.info(f"Batch size per GPU: {config.batch_size // world_size} (total: {config.batch_size})")
+
+    data_loader, sampler, data_config = make_loader(config, world_size)
+
+    model_cfg = config.model
+    object.__setattr__(model_cfg, "dtype", config.pytorch_training_precision or "float32")
+    model = model_cfg.build_model().to(device)
+    if type(model).__name__ not in ABLATION_MODELS:
+        raise ValueError(f"config built {type(model).__name__}, expected one of {ABLATION_MODELS}")
+    logging.info(f"RACE2 ablation model: {type(model).__name__} | conditioning "
+                 f"{'ON' if hasattr(model, 'field_mod') else 'OFF'} | lambda_base_fm {model_cfg.lambda_base_fm} "
+                 f"| lambda_timing {model_cfg.lambda_timing}")
+    logging.info(f"RACE v2 head: {type(model.race_head).__name__} "
+                 f"({sum(p.numel() for p in model.race_head.parameters()):,} params), kv_layer {model.kv_layer}")
+
+    resume_step = latest_resumable_step(config.checkpoint_dir) if config.resume else None
+    if resume_step is not None:
+        validate_resume(config, config.checkpoint_dir / str(resume_step), resume_step)
+    load_teacher(model, config)  # frozen VLM prefix always comes from the teacher
+    n_train, n_frozen = model.freeze_vlm()
+    logging.info(f"trainable {n_train/1e6:.1f}M / frozen {n_frozen/1e6:.0f}M (VLM prefix)")
+    # bias-free guarantee for the b=0 => zero-modulation claim (conditioned ablations only)
+    for lin in getattr(model, "field_mod", []):
+        assert lin.bias is None, "field_mod projections must be bias-free"
+    if not hasattr(model, "field_mod"):
+        assert not hasattr(model, "e_trans") and not hasattr(model, "tau_gate"), "timing-only model must carry no field tensors"
+    model.gradient_checkpointing_enable()
+
+    raw = model
+    if use_ddp:
+        model = torch.nn.parallel.DistributedDataParallel(
+            model, device_ids=[device.index] if device.type == "cuda" else None,
+            find_unused_parameters=True, gradient_as_bucket_view=True,
+        )
+
+    ms = MasterState([(n, p) for n, p in raw.named_parameters() if p.requires_grad], config.ema_decay)
+    gate_mult = float(getattr(model_cfg, "tau_gate_lr_mult", 1.0))
+    base_m = [m for n, m in zip(ms.names, ms.master) if "tau_gate" not in n]
+    gate_m = [m for n, m in zip(ms.names, ms.master) if "tau_gate" in n]
+    opt = torch.optim.AdamW(
+        [{"params": base_m, "lr_scale": 1.0},
+         {"params": gate_m, "lr_scale": gate_mult}],
+        lr=config.lr_schedule.peak_lr, betas=(config.optimizer.b1, config.optimizer.b2),
+        eps=config.optimizer.eps, weight_decay=config.optimizer.weight_decay,
+    )
+    if gate_m:
+        logging.info(f"tau_gate param group: {sum(p.numel() for p in gate_m)} scalars, lr x{gate_mult}")
+    logging.info(f"fp32 master weights + EMA(decay={ms.decay}, warmup) over {len(ms.master)} tensors")
+
+    warmup, peak, decay, end = (config.lr_schedule.warmup_steps, config.lr_schedule.peak_lr,
+                                config.lr_schedule.decay_steps, config.lr_schedule.decay_lr)
+
+    def lr_at(step):
+        if step < warmup:
+            return peak * (step + 1) / max(1, warmup)
+        t = min(1.0, (step - warmup) / max(1, decay))
+        return end + 0.5 * (peak - end) * (1 + np.cos(np.pi * t))
+
+    logging.info(f"lr: warmup {warmup} -> peak {peak:.1e} -> cosine over {decay} -> {end:.1e}")
+
+    global_step = 0
+    if resume_step is not None:
+        path = config.checkpoint_dir / str(resume_step) / TRAIN_STATE
+        logging.info(f"Resuming from {path}")
+        sd = torch.load(path, map_location="cpu", weights_only=False)
+        ms.load_state_dict(sd)
+        opt.load_state_dict(sd["optimizer"])
+        ms.sync_to_model()
+        global_step = int(sd["step"])
+        if global_step != resume_step:
+            raise ValueError(f"Checkpoint directory step {resume_step} differs from saved step {global_step}")
+        del sd
+        logging.info(f"Resumed at step {global_step}")
+    elif config.resume:
+        logging.info("--resume given but no resumable checkpoint found; starting from the teacher")
+    raw._step = global_step
+
+    per_gpu = config.batch_size // world_size
+    micro = getattr(model_cfg, "micro_batch_size", 0) or per_gpu
+    if per_gpu % micro != 0:
+        raise ValueError(f"per-GPU batch {per_gpu} not divisible by micro_batch_size {micro}")
+    accum = per_gpu // micro
+    if accum > 1:
+        logging.info(f"grad accumulation: micro {micro} x {accum} = per-GPU {per_gpu}")
+
+    infos = []
+    pbar = tqdm.tqdm(total=config.num_train_steps, initial=global_step, desc="Training",
+                     disable=not is_main, dynamic_ncols=True)
+    start = time.time()
+    epoch = global_step  # distinct sampler seeds across resumes
+    while global_step < config.num_train_steps:
+        # iterate the underlying torch loader per epoch so DistributedSampler
+        # reshuffles (TorchDataLoader's infinite __iter__ never calls set_epoch)
+        if sampler is not None:
+            sampler.set_epoch(epoch)
+        epoch += 1
+        for batch in data_loader.torch_loader:
+            if global_step >= config.num_train_steps:
+                break
+            batch = jax.tree.map(torch.as_tensor, batch)
+            b_star = batch["b"]
+            observation = _model.Observation.from_dict(batch)
+            actions = batch["actions"]
+            observation = jax.tree.map(lambda x: x.to(device), observation)
+            actions = actions.to(torch.float32).to(device)
+            b_star = b_star.to(torch.float32).to(device)
+            lr = lr_at(global_step)
+            for pg in opt.param_groups:
+                pg["lr"] = lr * pg.get("lr_scale", 1.0)
+            model.zero_grad(set_to_none=True)
+            loss_val = 0.0
+            metrics_acc = {}
+            for a in range(accum):
+                sl = slice(a * micro, (a + 1) * micro)
+                obs_m = jax.tree.map(lambda x: x[sl], observation)
+                # skip the DDP all-reduce on all but the last micro-step
+                ctx = model.no_sync() if (use_ddp and a < accum - 1) else contextlib.nullcontext()
+                with ctx:
+                    loss = model(obs_m, actions[sl], b_star=b_star[sl]) / accum
+                    loss.backward()
+                loss_val += loss.item()
+                for k, v in raw.last_metrics.items():
+                    metrics_acc[k] = metrics_acc.get(k, 0.0) + v / accum
+            raw.last_metrics = metrics_acc
+            ms.grads_to_master()
+            gn = torch.nn.utils.clip_grad_norm_(
+                ms.master, max_norm=config.optimizer.clip_gradient_norm, error_if_nonfinite=True
+            )
+            opt.step()
+            ms.sync_to_model()
+            ema_d = ms.update_ema(global_step)
+            if is_main:
+                m = {"loss": loss_val, "grad_norm": float(gn), **raw.last_metrics}
+                infos.append(m)
+                if global_step % config.log_interval == 0:
+                    avg = {k: sum(i.get(k, 0) for i in infos) / len(infos) for k in infos[0]}
+                    extra = " ".join(f"{k}={avg[k]:.4f}" for k in avg if k not in ("loss", "grad_norm"))
+                    mem = torch.cuda.max_memory_allocated(device) / 2**30 if device.type == "cuda" else 0.0
+                    logging.info(f"step={global_step} loss={avg['loss']:.4f} {extra} "
+                                 f"lr={lr:.2e} ema={ema_d:.4f} gn={avg['grad_norm']:.2f} "
+                                 f"mem={mem:.1f}G t={time.time()-start:.1f}s")
+                    infos = []
+                    start = time.time()
+                last = raw.last_metrics
+                pbar.update(1)
+                pbar.set_postfix({
+                    "loss": f"{loss_val:.3f}",
+                    "fm": f"{last.get('fm_loss', float('nan')):.3f}",
+                    "bfm": f"{last.get('base_fm', float('nan')):.3f}",
+                    "bce": f"{last.get('bce', float('nan')):.3f}",
+                    "pm1": f"{last.get('peak_pm1', float('nan')):.2f}",
+                    "pres": f"{last.get('pres_acc', float('nan')):.2f}",
+                    "b+": f"{last.get('b_pos', float('nan')):.2f}",
+                    "b-": f"{last.get('b_neg', float('nan')):.2f}",
+                    "g1": f"{last.get('g_first', float('nan')):.2f}",
+                    "g10": f"{last.get('g_last', float('nan')):.2f}",
+                    "dmod": f"{last.get('dmod', float('nan')):.2f}",
+                    "lr": f"{lr:.1e}",
+                })
+            global_step += 1
+            raw._step = global_step
+            save_checkpoint(config, global_step, ms, opt, data_config, is_main)
+            if use_ddp:
+                dist.barrier()  # rank 0 may have spent a minute writing the checkpoint
+    if is_main:
+        pbar.close()
+
+
+def main():
+    base.init_logging()
+    config = _config.cli()
+    if type(config.model) is not race2_ablation_config.Race2TimingOnlyConfig:
+        raise ValueError(f"this trainer runs race2abl_timing_only presets only (Race2TimingOnlyConfig); got {type(config.model).__name__}")
+    if "race2abl_timing_only_h" not in config.name:
+        raise ValueError(f"expected a pi05_vlabench_race2abl_timing_only_h<H> config, got {config.name!r}")
+    if config.resume and latest_resumable_step(config.checkpoint_dir) is None:
+        raise FileNotFoundError(f"No resumable checkpoint in {config.checkpoint_dir}")
+    if not config.resume and config.checkpoint_dir.exists():
+        saved_steps = [p for p in config.checkpoint_dir.iterdir() if p.is_dir() and p.name.isdigit()]
+        if saved_steps:
+            raise FileExistsError(
+                f"Checkpoints already exist in {config.checkpoint_dir}; use --resume or a new --exp_name"
+            )
+    cps = gauss_data.load_changepoints(config.data.cps_path)
+    n_cp = sum(len(v) for v in cps.values())
+    logging.info(f"changepoints: {config.data.cps_path} — {len(cps)} episodes, "
+                 f"{n_cp} CPs ({n_cp / len(cps):.2f}/episode)")
+    try:
+        train_loop(config)
+    finally:
+        if dist.is_initialized():
+            dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()

abl_timing_only_h20_40k/source/src/openpi/models/race2_ablation_config.py

@@ -0,0 +1,58 @@
+"""Configs for the RACE2 ablations (openpi.models_pytorch.race2_ablation).
+
+Both keep every RacePi0Config knob of the full method; lambda_base_fm is pinned to 0.
+  Race2TimingOnlyConfig -> Race2TimingOnlyPI0 (A1: timing loss only, no conditioning)
+  Race2NoAuxConfig      -> Race2NoAuxPI0      (A3: full method without the auxiliary loss)
+  Race2FmOnlyConfig     -> Race2FmOnlyPI0     (A2: conditioning by predicted b_hat, FM loss only)
+"""
+
+import dataclasses
+
+from openpi.models import race2_config
+
+
+@dataclasses.dataclass(frozen=True)
+class Race2TimingOnlyConfig(race2_config.Race2Pi0Config):
+    lambda_base_fm: float = 0.0
+
+    def __post_init__(self):
+        super().__post_init__()
+        if float(self.lambda_base_fm) != 0.0:
+            raise ValueError("Race2TimingOnlyConfig: lambda_base_fm must be 0")
+
+    def build_model(self):
+        from openpi.models_pytorch.race2_ablation.model import Race2TimingOnlyPI0
+
+        return Race2TimingOnlyPI0(config=self)
+
+
+@dataclasses.dataclass(frozen=True)
+class Race2NoAuxConfig(race2_config.Race2Pi0Config):
+    lambda_base_fm: float = 0.0
+
+    def __post_init__(self):
+        super().__post_init__()
+        if float(self.lambda_base_fm) != 0.0:
+            raise ValueError("Race2NoAuxConfig: lambda_base_fm must be 0")
+
+    def build_model(self):
+        from openpi.models_pytorch.race2_ablation.model import Race2NoAuxPI0
+
+        return Race2NoAuxPI0(config=self)
+
+
+@dataclasses.dataclass(frozen=True)
+class Race2FmOnlyConfig(race2_config.Race2Pi0Config):
+    lambda_base_fm: float = 0.0
+    lambda_timing: float = 0.0
+    jitter_prob: float = 0.0  # unused: the field comes from b_hat, never from b*
+
+    def __post_init__(self):
+        super().__post_init__()
+        if float(self.lambda_base_fm) != 0.0 or float(self.lambda_timing) != 0.0:
+            raise ValueError("Race2FmOnlyConfig: lambda_base_fm and lambda_timing must be 0")
+
+    def build_model(self):
+        from openpi.models_pytorch.race2_ablation.model import Race2FmOnlyPI0
+
+        return Race2FmOnlyPI0(config=self)

abl_timing_only_h20_40k/source/src/openpi/models_pytorch/race2_ablation/__init__.py

@@ -0,0 +1,3 @@
+from openpi.models_pytorch.race2_ablation.model import Race2FmOnlyPI0, Race2NoAuxPI0, Race2TimingOnlyPI0
+
+__all__ = ["Race2FmOnlyPI0", "Race2NoAuxPI0", "Race2TimingOnlyPI0"]

abl_timing_only_h20_40k/source/src/openpi/models_pytorch/race2_ablation/model.py

@@ -0,0 +1,240 @@
+"""RACE2 ablations (VLABench). Separate subclasses so the full method's code is untouched.
+
+Full RACE2 (openpi.models_pytorch.race.model.RacePI0.forward), per training step:
+    aux pass   tau = read_tau, unmodulated  -> timing head -> bce(b*)        [timing loss]
+                                            -> action_out_proj -> base_fm     [auxiliary loss]
+    FM pass    tau ~ U, modulated by alpha(tau) * jitter(b*) * M(e_trans)    [fm_loss]
+    total = fm_loss + lambda_base_fm * base_fm + lambda_timing * bce
+Serving: aux pass -> b_hat -> hook modulation of the denoising passes.
+
+A1  Race2TimingOnlyPI0  "timing loss only, no conditioning"
+    The head is trained with the timing BCE and its gradient reaches the action expert
+    through the auxiliary pass (hidden_read is NOT detached, as in the full method), but
+    there is no field: e_trans / field_mod / tau_gate are removed, the FM pass is the
+    plain pi0.5 flow-matching loss and serving is pi0.5's own sampler. This isolates the
+    effect of timing supervision as a multi-task signal.
+        total = fm_loss(unmodulated) + lambda_timing * bce
+
+A3  Race2NoAuxPI0       "full method without the auxiliary (base FM) loss"
+    Identical model, conditioning and serving; only the base FM term is dropped
+    (lambda_base_fm must be 0, which RacePI0.forward already honours exactly).
+        total = fm_loss(modulated, teacher-forced b*) + lambda_timing * bce
+
+A2  Race2FmOnlyPI0      "conditioning driven end-to-end by the FM loss"
+    Same parameters, conditioning sites and serving path as the full method, but the
+    field is built from the head's own prediction b_hat instead of the teacher-forced
+    b* (no jitter), and the only loss is the flow-matching loss. The timing head,
+    alpha(tau), e_trans and field_mod are therefore trained solely through the FM
+    gradient that flows back through the modulation; the timing BCE and the base FM
+    term are computed under no_grad for monitoring only.
+        total = fm_loss(modulated by alpha(tau) * b_hat * M(e_trans))
+"""
+
+import torch
+import torch.nn.functional as F  # noqa: N812
+
+from openpi.models_pytorch.pi0_pytorch import PI0Pytorch
+from openpi.models_pytorch.race2.model import Race2PI0
+
+
+class Race2TimingOnlyPI0(Race2PI0):
+    """A1: RACE2 timing head + BCE, no field modulation anywhere."""
+
+    def __init__(self, config):
+        super().__init__(config)
+        if float(config.lambda_base_fm) != 0.0:
+            raise ValueError("Race2TimingOnlyPI0 has no auxiliary loss; set lambda_base_fm=0")
+        # No conditioning: drop the field parameters so checkpoints carry no dead weights
+        # and nothing can modulate the expert. The norm hooks installed by RacePI0 stay
+        # inert because _mod_buf is never allocated.
+        del self.e_trans
+        del self.field_mod
+        del self.tau_gate
+        self._mod_consts = None
+        self._mod_buf = None
+
+    # ---- conditioning is structurally absent -----------------------------------
+    @torch.compiler.disable
+    def _refresh_mod_consts(self):
+        self._mod_consts = None
+        return None
+
+    def _prepare_mod_buf(self, bsize, horizon, device):  # never used; keep hooks inert
+        self._mod_buf = None
+
+    def _field(self, b, tau):
+        raise RuntimeError("Race2TimingOnlyPI0 has no field modulation")
+
+    # ---- training --------------------------------------------------------------
+    def forward(self, observation, actions, b_star=None, noise=None, time=None):
+        cfg = self.config
+        if b_star is None:
+            raise ValueError("Race2TimingOnlyPI0.forward requires b_star (B,H)")
+        images, img_masks, lang_tokens, lang_masks, state = self._preprocess_observation(observation, train=True)
+        device = actions.device
+        bsize = actions.shape[0]
+        b_star = b_star.to(device=device, dtype=torch.float32)
+
+        if noise is None:
+            noise = self.sample_noise(actions.shape, device)
+        u_t = noise - actions
+        use_bf16 = (
+            self.paligemma_with_expert.paligemma.language_model.layers[0].self_attn.q_proj.weight.dtype
+            == torch.bfloat16
+        )
+        prefix_pad_masks, prefix_keys, prefix_values, _prefix_hidden = self._prefix_h(
+            images, img_masks, lang_tokens, lang_masks, use_bf16
+        )
+
+        # ---- auxiliary pass (tau = read_tau, unmodulated) -> timing head only.
+        # hidden_read is not detached: the BCE gradient reaches the action expert.
+        read_time = torch.full((bsize,), cfg.read_tau, device=device)
+        x_read = cfg.read_tau * noise + (1 - cfg.read_tau) * actions
+        hidden_read = self._run_expert(
+            state, x_read, read_time, prefix_pad_masks, prefix_keys, prefix_values, field=None
+        )
+        logits = self._timing_logits(hidden_read, prefix_values[self.kv_layer], prefix_pad_masks)
+        timing_bce = F.binary_cross_entropy_with_logits(logits, b_star)
+
+        # ---- plain flow-matching pass: no field, no teacher forcing
+        if time is None:
+            time = self.sample_time(bsize, device)
+        x_t = time[:, None, None] * noise + (1 - time[:, None, None]) * actions
+        hidden = self._run_expert(state, x_t, time, prefix_pad_masks, prefix_keys, prefix_values, field=None)
+        fm_loss = F.mse_loss(u_t, self.action_out_proj(hidden))
+
+        total = fm_loss + cfg.lambda_timing * timing_bce
+
+        metrics = {"fm_loss": fm_loss.detach(), "bce": timing_bce.detach(),
+                   "base_fm": torch.zeros((), device=device)}
+        with torch.no_grad():
+            b_hat = torch.sigmoid(logits)
+            pos = b_star.max(dim=1).values >= 0.5
+            z = torch.tensor(0.0, device=device)
+            if pos.any():
+                peak_err = (b_hat[pos].argmax(dim=1) - b_star[pos].argmax(dim=1)).abs()
+                metrics["peak_pm1"] = (peak_err <= 1).float().mean()
+                metrics["b_pos"] = b_hat[pos].max(dim=1).values.mean()
+            else:
+                metrics["peak_pm1"] = z
+                metrics["b_pos"] = z
+            metrics["b_neg"] = b_hat[~pos].max(dim=1).values.mean() if (~pos).any() else z
+            metrics["pres_acc"] = ((b_hat.max(dim=1).values >= 0.5) == pos).float().mean()
+
+        self._step += 1
+        self.last_metrics = {k: float(v) for k, v in metrics.items()}
+        return total
+
+    # ---- serving: pi0.5's own sampler (hooks inert, no auxiliary pass) --------------
+    @torch.no_grad()
+    def sample_actions(self, device, observation, noise=None, num_steps=10):
+        return PI0Pytorch.sample_actions(self, device, observation, noise=noise, num_steps=num_steps)
+
+    @torch.no_grad()
+    def _sample_actions_reference(self, device, observation, noise=None, num_steps=10):
+        """Hand-rolled unmodulated rollout, parity target for tests only."""
+        cfg = self.config
+        bsize = observation.state.shape[0]
+        if noise is None:
+            noise = self.sample_noise((bsize, cfg.action_horizon, cfg.action_dim), device)
+        images, img_masks, lang_tokens, lang_masks, state = self._preprocess_observation(observation, train=False)
+        ppm, keys, values, _ = self._prefix_h(images, img_masks, lang_tokens, lang_masks, use_bf16=False)
+        dt = -1.0 / num_steps
+        x_t = noise.clone()
+        t = 1.0
+        for _ in range(num_steps):
+            tt = torch.full((bsize,), t, dtype=torch.float32, device=device)
+            hidden = self._run_expert(state, x_t, tt, ppm, keys, values, field=None)
+            x_t = x_t + dt * self.action_out_proj(hidden)
+            t += dt
+        return x_t
+
+
+class Race2FmOnlyPI0(Race2PI0):
+    """A2: RACE2 conditioning with b_hat in the loop, FM loss only (no timing / aux loss)."""
+
+    def __init__(self, config):
+        super().__init__(config)
+        if float(config.lambda_base_fm) != 0.0 or float(config.lambda_timing) != 0.0:
+            raise ValueError("Race2FmOnlyPI0 trains with the FM loss only; set lambda_base_fm=lambda_timing=0")
+
+    def forward(self, observation, actions, b_star=None, noise=None, time=None):
+        cfg = self.config
+        if b_star is None:
+            raise ValueError("Race2FmOnlyPI0.forward requires b_star (B,H) for monitoring metrics")
+        images, img_masks, lang_tokens, lang_masks, state = self._preprocess_observation(observation, train=True)
+        device = actions.device
+        bsize = actions.shape[0]
+        b_star = b_star.to(device=device, dtype=torch.float32)
+
+        if noise is None:
+            noise = self.sample_noise(actions.shape, device)
+        u_t = noise - actions
+        use_bf16 = (
+            self.paligemma_with_expert.paligemma.language_model.layers[0].self_attn.q_proj.weight.dtype
+            == torch.bfloat16
+        )
+        prefix_pad_masks, prefix_keys, prefix_values, _prefix_hidden = self._prefix_h(
+            images, img_masks, lang_tokens, lang_masks, use_bf16
+        )
+
+        # ---- auxiliary pass (tau = read_tau, unmodulated) -> head -> b_hat (kept in the graph)
+        read_time = torch.full((bsize,), cfg.read_tau, device=device)
+        x_read = cfg.read_tau * noise + (1 - cfg.read_tau) * actions
+        hidden_read = self._run_expert(
+            state, x_read, read_time, prefix_pad_masks, prefix_keys, prefix_values, field=None
+        )
+        logits = self._timing_logits(hidden_read, prefix_values[self.kv_layer], prefix_pad_masks)
+        b_hat = torch.sigmoid(logits)
+
+        # ---- modulated FM pass: field from the PREDICTED b_hat (no teacher forcing, no jitter)
+        if time is None:
+            time = self.sample_time(bsize, device)
+        x_t = time[:, None, None] * noise + (1 - time[:, None, None]) * actions
+        field = self._field(b_hat, time)
+        hidden = self._run_expert(state, x_t, time, prefix_pad_masks, prefix_keys, prefix_values, field=field)
+        fm_loss = F.mse_loss(u_t, self.action_out_proj(hidden))
+
+        total = fm_loss
+
+        metrics = {"fm_loss": fm_loss.detach()}
+        with torch.no_grad():
+            # monitoring only: how the FM-trained head relates to the changepoint target
+            metrics["bce"] = F.binary_cross_entropy_with_logits(logits.detach(), b_star)
+            metrics["base_fm"] = F.mse_loss(u_t, self.action_out_proj(hidden_read.detach()))
+            pos = b_star.max(dim=1).values >= 0.5
+            z = torch.tensor(0.0, device=device)
+            if pos.any():
+                peak_err = (b_hat[pos].argmax(dim=1) - b_star[pos].argmax(dim=1)).abs()
+                metrics["peak_pm1"] = (peak_err <= 1).float().mean()
+                metrics["b_pos"] = b_hat[pos].max(dim=1).values.mean()
+            else:
+                metrics["peak_pm1"] = z
+                metrics["b_pos"] = z
+            metrics["b_neg"] = b_hat[~pos].max(dim=1).values.mean() if (~pos).any() else z
+            metrics["pres_acc"] = ((b_hat.max(dim=1).values >= 0.5) == pos).float().mean()
+            metrics["b_hat_mean"] = b_hat.mean()
+            g = torch.sigmoid(self.tau_gate)
+            for i in range(cfg.n_tau):
+                metrics[f"alpha_{i}"] = g[i].detach()
+            metrics["g_first"] = g[0].detach()
+            metrics["g_last"] = g[-1].detach()
+            v = self.e_trans.float()
+            mods = torch.stack([lin(v) for lin in self.field_mod])
+            s_p, b_p, _ = mods.chunk(3, dim=-1)
+            metrics["dmod"] = (s_p.norm(dim=-1) + b_p.norm(dim=-1)).mean()
+
+        self._step += 1
+        self.last_metrics = {k: float(v) for k, v in metrics.items()}
+        return total
+
+
+class Race2NoAuxPI0(Race2PI0):
+    """A3: full RACE2 minus the auxiliary base-FM loss. RacePI0.forward skips the base FM
+    term entirely when lambda_base_fm == 0, so training, conditioning, serving and the
+    checkpoint layout are inherited unchanged; this class only pins that setting."""
+
+    def __init__(self, config):
+        super().__init__(config)
+        if float(config.lambda_base_fm) != 0.0:
+            raise ValueError("Race2NoAuxPI0 drops the auxiliary loss; set lambda_base_fm=0")

abl_timing_only_h20_40k/training_config.json

@@ -0,0 +1,107 @@
+{
+  "global_step": 40000,
+  "config": {
+    "name": "pi05_vlabench_race2abl_timing_only_h20",
+    "project_name": "openpi",
+    "exp_name": "abl_timing_only_h20_s42",
+    "model": {
+      "action_dim": 32,
+      "action_horizon": 20,
+      "max_token_len": 200,
+      "dtype": "bfloat16",
+      "paligemma_variant": "gemma_2b",
+      "action_expert_variant": "gemma_300m",
+      "pi05": true,
+      "discrete_state_input": false,
+      "gemma_depth": 18,
+      "sigma": 1.0,
+      "support": 3.0,
+      "jitter_prob": 0.5,
+      "lambda_timing": 0.05,
+      "lambda_base_fm": 0.0,
+      "read_tau": 1.0,
+      "n_tau": 10,
+      "tau_gate_init": 0.0,
+      "tau_gate_lr_mult": 20.0,
+      "timing_proj_dim": 256,
+      "timing_num_heads": 4,
+      "micro_batch_size": 16,
+      "fast_serving": 0,
+      "kv_layer": -1
+    },
+    "weight_loader": {},
+    "pytorch_weight_path": "checkpoints/pi05_vlabench_official/pytorch",
+    "pytorch_weight_path_teacher": null,
+    "pytorch_training_precision": "bfloat16",
+    "lr_schedule": {
+      "warmup_steps": 500,
+      "peak_lr": 5e-05,
+      "decay_steps": 39500,
+      "decay_lr": 5e-06
+    },
+    "optimizer": {
+      "b1": 0.9,
+      "b2": 0.95,
+      "eps": 1e-08,
+      "weight_decay": 1e-10,
+      "clip_gradient_norm": 1.0
+    },
+    "ema_decay": 0.999,
+    "freeze_filter": "Nothing()",
+    "data": {
+      "repo_id": "SeonghoonYu/vlabench_primitive_ft_lerobot_224",
+      "assets": {
+        "assets_dir": "./assets/pi05_vlabench",
+        "asset_id": "vlabench/vlabench_ft_primitive"
+      },
+      "base_config": {
+        "repo_id": null,
+        "asset_id": null,
+        "norm_stats": null,
+        "repack_transforms": {
+          "inputs": [],
+          "outputs": []
+        },
+        "data_transforms": {
+          "inputs": [],
+          "outputs": []
+        },
+        "model_transforms": {
+          "inputs": [],
+          "outputs": []
+        },
+        "use_quantile_norm": false,
+        "action_sequence_keys": [
+          "actions"
+        ],
+        "prompt_from_task": true,
+        "rlds_data_dir": null,
+        "action_space": null,
+        "datasets": []
+      },
+      "phase_labels_path": null,
+      "phase_label_horizon": null,
+      "cps_path": "assets/gauss_timing/vlabench_cps.npz",
+      "gauss_sigma": 1.0,
+      "gauss_support": 3.0
+    },
+    "assets_base_dir": "./assets",
+    "checkpoint_base_dir": "./checkpoints",
+    "seed": 42,
+    "batch_size": 64,
+    "num_workers": 8,
+    "num_train_steps": 40000,
+    "log_interval": 100,
+    "save_interval": 1000,
+    "save_overlays_from_step": 20000,
+    "keep_period": 5000,
+    "overwrite": false,
+    "resume": false,
+    "wandb_enabled": false,
+    "policy_metadata": null,
+    "fsdp_devices": 1
+  },
+  "ema_decay": 0.999,
+  "timestamp": 1789477910.1500757,
+  "format": "trainable-only overlay; materialised with scripts/race/materialize_ckpt.py"
+}