Frank-Insert-Wipe: AdaSpeVLA checkpoints for Franka insert_tube and wipe_whiteboard (H=16)

AdaSpeVLA is a pi0.5 (OpenPI, JAX) variant. It predicts cubic B-spline action chunks with adaptive knots, plus a timing head that decides when to query the policy again. This repo holds two variants that differ only in the generator that produces the 11 control points:

• flow_head (<task>/, 2026-11-12): a 20.4 M FlowPolicy trained from scratch.
• pi05_expert (<task>_expert/, 2026-11-14): pi0.5's pretrained Gemma-300M action expert (427.9 M with adaRMS, from pi05_base) runs on the 11 control-point tokens, as pi0.5 does on its action tokens.

Everything else (data, norm stats, H=16, K=11, labels, knot head, timing head, 40k + 10k steps) is the same. The training code for both is in code/.

These are params-only Orbax checkpoints in the OpenPI layout. The optimizer train_state is not included.

Folder	Training stage	Config name (with TASK=<task>)	Use
<task>/timing-from40k-step-9999/	Timing head, 10k steps, from Action step 39999 (everything except freq_head frozen)	<task>_adaspe_timing_from40k	Deployment: full model with the trained timing head
<task>/action-step-39999/	Action stage, 40k steps (EMA params)	<task>_adaspe_action	Action generator / knot head only. Its timing head is untrained.
<task>_expert/timing-from40k-step-9999/	Same, pi05_expert generator	<task>_adaspe_expert_timing_from40k	Deployment (expert variant)
<task>_expert/action-step-39999/	Same, pi05_expert generator	<task>_adaspe_expert_action	Action generator / knot head only (expert variant)
<task>/timing-from40k-lr3e-4-step-9999/	Timing head only, retrained 2026-11-14 with peak LR 3e-4 (same 10k steps; see below)	<task>_adaspe_timing_from40k_lr3e-4	Deployment candidate
<task>/timing-from40k-lr1e-3-step-9999/	Timing head only, retrained 2026-11-14 with peak LR 1e-3 (same 10k steps; see below)	<task>_adaspe_timing_from40k_lr1e-3	Deployment candidate
<task>_expert/timing-from40k-lr3e-4-step-9999/	Timing head only, retrained 2026-11-14 with peak LR 3e-4 (same 10k steps; see below)	<task>_adaspe_expert_timing_from40k_lr3e-4	Deployment candidate (expert variant)
<task>_expert/timing-from40k-lr1e-3-step-9999/	Timing head only, retrained 2026-11-14 with peak LR 1e-3 (same 10k steps; see below)	<task>_adaspe_expert_timing_from40k_lr1e-3	Deployment candidate (expert variant)

<task> is insert_tube or wipe_whiteboard. Each folder has params/, assets/dddddsds31341/<dataset>/norm_stats.json (the normalization stats used in training) and _CHECKPOINT_METADATA. SHA256SUMS lists every file.

Loading

You need the AdaSpeVLA OpenPI plugin, included here as code/adaspe_train (run with TASK=<task>; it registers these config names). It runs on OpenPI at the commit in code/OPENPI_COMMIT. Setup is in code/README.md, and code/adaspe_train/EXPERT_VARIANT.md documents the expert variant. The folders use the standard OpenPI checkpoint layout (params/ + assets/).

from huggingface_hub import snapshot_download
root = snapshot_download("cognaclee/Frank-Insert-Wipe", allow_patterns=["insert_tube/timing-from40k-step-9999/*"])
# higher-LR timing head: "insert_tube/timing-from40k-lr3e-4-step-9999/*" with config insert_tube_adaspe_timing_from40k_lr3e-4
ckpt_dir = f"{root}/insert_tube/timing-from40k-step-9999"   # with config insert_tube_adaspe_timing_from40k

Data

Task	Dataset (revision)	Size	pi0.5 baseline (same norm stats)
insert_tube	dddddsds31341/insert-tube-lerobot (70bcd1f86ed35fc93a0422a564472355242a51ab)	102 episodes / 16,403 frames	dddddsds31341/Franka-Insert-Tube-pi05-step-20000
wipe_whiteboard	dddddsds31341/wipe-whiteboard-lerobot (2c8cfa8ddef1f83cb09ca7881a0944cd637cd6d0)	112 episodes / 29,795 frames	dddddsds31341/Franka-Wipe-Whiteboard-pi05-step-20000

Two 128×128 cameras, 19-D state, 7-D action, 10 Hz.

Recipe

• Action horizon H=16 (1.5 s), 11 cubic B-spline control points, query interval 0.2–1.5 s.
• Gripper events for the timing and knot labels come from the gripper command (the measured width lags by about 1.9 s).
• Knot density is tanh(motion_variation / e_crit), with e_crit calibrated per task at H=16: insert_tube 11.98, wipe_whiteboard 8.96.
• Initialized from pi05_base. Batch 32, Action 40k steps (LR decay over 30k), then timing 10k steps.

Training loss

Mean training loss over the last 1k–2k steps of each stage.

Task	Variant	Action loss, steps 38k–40k	Timing loss (log-interval MSE), steps 9k–10k
insert_tube	flow_head	0.086	0.206
insert_tube	pi05_expert	0.030	0.255
wipe_whiteboard	flow_head	0.118	0.196
wipe_whiteboard	pi05_expert	0.040	0.271

Action loss over training (mean per 2k-step window):

Steps	insert_tube flow_head	insert_tube expert	wipe flow_head	wipe expert
0–2k	0.624	0.334	0.597	0.295
8–10k	0.238	0.121	0.238	0.121
18–20k	0.156	0.067	0.169	0.083
28–30k	0.105	0.038	0.130	0.049
38–40k	0.086	0.030	0.118	0.040

Notes:

• The two variants use different flow-matching conventions (flow_head: t ~ U(0,1); pi05_expert: pi0.5's Beta(1.5,1)), so the action losses are not exactly the same metric.
• The expert's timing head has a higher loss. It reads a pooled observation feature that, in the expert variant, is trained only by the knot loss (see code/adaspe_train/EXPERT_VARIANT.md).
• These are training losses. No held-out or real-robot evaluation results are included here.

Timing head retrained with a higher learning rate (2026-11-14)

The original timing heads (timing-from40k-step-9999) used OpenPI's default schedule: peak LR 2.5e-5 with a cosine decay over 30k steps. The timing stage only runs 10k steps, so the LR stayed at about 2.0–2.5e-5, which is low for freq_head, a freshly initialized 2-layer MLP (633 → 128 → 1, ~81k params) and the only trainable part. The loss was still falling at 10k.

The timing-from40k-lr<LR>-step-9999 folders retrain only freq_head, with the same data, labels, 10k steps, batch 32 and 1k warmup, a cosine decay to LR/10 over the 10k steps, and peak LR 3e-4 or 1e-3. They start from the same Action step-39999 params, and everything except freq_head is frozen. The predicted actions, knots and B-spline are therefore identical to timing-from40k-step-9999. Only the query interval (and so execute_steps) changes.

Timing loss (log-interval MSE, training loss, mean over steps 9k–10k):

Task	Variant	LR 2.5e-5 (original)	LR 3e-4	LR 1e-3
insert_tube	flow_head	0.206	0.157	0.091
insert_tube	pi05_expert	0.255	0.182	0.100
wipe_whiteboard	flow_head	0.196	0.158	0.112
wipe_whiteboard	pi05_expert	0.271	0.214	0.153

These are training losses. With the higher LR, part of the drop may come from fitting the training frames; a held-out comparison has not been done yet. The configs are configs.TIMINGS_LR in code/adaspe_train/configs.py, and code/adaspe_train/slurm/timing_lr.slurm runs one of them.