Autoregressive action pretraining shapes trunk features that a frozen flow decoder can read

bijou stage-2 report · generated 2026-08-05T12:43+01:00 · controlled two-arm phase: steps 0–2,500 (both arms stopped at the 2,500 save boundary — the separation was ~9× the probe noise floor, and the GPUs were reallocated to continuing the winner; see "status" below). Conclusions remain provisional until the 40k frozen-panel eval of the continuation run.

TL;DR. We pretrained our Gemma-4 E2B trunk for 100k steps with an autoregressive action-token objective, froze it, and trained an identical flow-matching expert twice: once on the stock trunk, once on the AR-pretrained one. Same expert, same init, same seed, same data order — only the trunk weights differ. At step 2,500 the AR-trunk arm reads step 2,500: stock 13.41 vs AR-trunk 10.71 (Δ -2.69, -20%) on the holdout probe, ~8× the probe's noise floor. If this holds to 40k, "pretrain the trunk autoregressively, then mount a flow head" becomes our default recipe.

Why we ran this

Bijou trains vision-language-action models for SO-100/101 arms on crowd-sourced teleoperation data: a Gemma-4 E2B backbone encodes (images, instruction, state) into a prefix, and an action decoder predicts 50-step joint-space chunks. Two decoder families coexist in the codebase: a flow-matching expert (a ~0.5B cross-attention velocity network over continuous chunks, trunk frozen) and an autoregressive decoder-only path (the full backbone predicts discretized action tokens — a FAST-style DCT codec — plus auxiliary language tasks, trunk live).

The AR path owns our offline leaderboard: its best checkpoint scores 5.64 chunk MAE on the community holdout vs 10.88 for a state-copy baseline, and its first-action error beats the copy baseline — the initial-placement measure that used to be the wall. But two observations kept nagging:

Separately, mechanism probes said the AR pretrain attacks exactly what limits generalization: cross-rig error is a frame-dependent level mis-estimation (the model mis-localizes the working point visually), and the AR runs were the first whose first-action error beat the copy baseline. Hypothesis: the AR objective shapes the trunk's features; the flow head's deployment behavior could inherit that for free. This experiment isolates the trunk contribution.

Experimental design

Two arms, identical in everything except trunk weights:

A — stockB — AR-pretrained
trunk (frozen, 15-layer prefix) google/gemma-4-e2b-it rcond ar_backbone pretrain, step 100k
prompt state tokeninert (zero-init, frozen) the pretrain's trained projection, frozen
flow expert (474.4M, trained) identical: adaRMS time conditioning, bidirectional self-attention, h1024/8 heads/4096 MLP, cross-attention to trunk layers 4/9/14 (4-4-8 schedule), same init, same seed
data & schedule identical: 880 datasets (878 community + 2 rig), 18.7M frames, fps-30/≤2-cam filters, 10% episode holdout; B96 × DDP2, LR 1e-4 cosine to 40k, warmup 500, clip 10, seed 0

Details that took care to get right:

Results

The cleanest view is the training loss: both arms share the expert init, the seed and the exact data order, so at every step the two curves score the same batches with the same head — the vertical gap is purely what the trunk's features give the expert to work with:

training loss

The holdout probe tells the same story where it counts — the separation is immediate and stable, the AR-trunk arm reaching any given probe value roughly 2× faster in steps and sitting ~2–3 MAE below at matched steps, against a probe noise floor of ±0.3:

holdout chunk MAE, both arms paired gap

For provenance, the pretrain that produced the inherited trunk (different frame set and objective — not comparable in absolute terms to the curves above):

AR pretrain curve
Read the numbers with these caveats. (1) The controlled comparison covers steps 0–2,500 of a planned 40,000 — we stopped early on a clear separation, which means the late-training picture (does stock eventually catch up?) is unmeasured by design. (2) One seed per arm. (3) The in-run probe is 256 frames on a new frame set (rig data joined the mix) — the arms compare to each other, not to older ledger numbers. (4) adaRMS starts as an identity velocity field, so both arms were still in their fast-improvement regime; the paired gap was stable-to-widening over the measured range, not closing.

Status & the continuation run

Both arms were stopped at step 2,500 (a save boundary) on 2026-08-04, and arm B resumed from that checkpoint on 4×H100 (per-rank batch 96→48, so the effective batch is unchanged at 192 — no LR-scale seam) with one deliberate prompt change from step 2,500 onward: flow prompts now render the constant [generate|actions] bracket (--prompt-generate-bracket). The AR trunk shaped its conditioning and state-token positions with that bracket present in every pretraining prompt, so bracket-free prompts read those positions slightly off-manifold; the controlled phase shared the shift across arms (fair), and after the change the flow prompt distribution is a strict sub-distribution of the trunk's pretraining forms (element-by- element audit in the session notes). The 2,500-step seam (bracket + world size) is recorded in the run header; steps 0–2,500 above stand as the completed controlled experiment. The resumed run showed no adaptation shock at the seam (loss continued at ~0.15, grad norms <1).

The continuation in historical context

wandb: continuation vs the historical
adaRMS h1536 lineage

wandb workspace view (2026-08-05). The AR-trunk run spans three sessions under one name: the short yellow segment (controlled phase, 0–2.5k), teal (resumed on 4×H100 with the trunk-consistency bracket, 2.5k→40k), and salmon (the 40k→80k extension: fresh data seed + LR re-warmup at the 40k seam — the restart wobble is visible — finishing at 6.57). Solid magenta: the stock control arm, 0–2.5k. Blue → orange → dashed magenta: the historical stock-trunk flow lineage (adaRMS, h1536) across its three sessions, 0→100k.

Read coarsely — the lineage trained on a different data mix (fps-30 all-camera-count corpus, no rig data, pre-conditioning prompt format) and a WIDER expert (h1536 vs h1024), so absolute values do not transfer exactly. But the shape is hard to argue with: the AR-trunk run passes ~7.6 around step 21k — a value the stock-trunk lineage first reached near step 100k with more parameters — and its extension finishes at 6.57 by 80k, below anything the lineage reached in 100k. A ~5× step-efficiency gain for the flow family at matched quality, and new-best quality at matched-or-less compute.

Interpretation

Under the acuity probes, position information is sharpest at the vision-tower output and degrades through the stock LM stack — the text layers' use of visual tokens, not the vision encoder, is the bottleneck. The AR pretrain forces the trunk to carry action-relevant visual detail all the way to token prediction, and this experiment says a flow head can read that improvement through frozen features at layers 4/9/14. The practical consequence, if the gap holds: a two-stage recipe — pretrain the trunk once with the AR objective (expensive, data-hungry, memorization-prone, but it shapes features), then train flow heads on the frozen result (cheap, zero episode-fit gap, samples coherent trajectories, fine-tunes gracefully on small rig datasets). The stages decouple: trunk quality scales with corpus size while the deployment head stays replaceable.

What's next

Appendix: provenance

arms:      bijou_flow_stock_e2b_h1024_40k_ddp2  (wandb 61tqffoz)
           bijou_flow_artrunk_h1024_40k_ddp2  (wandb 7j7cia7l)
trunk:     bijou_arb_rcond_100k_ddp4/step_100000 (wandb uvgqd5ci; hub: mcobzarenco/bijou-checkpoints)
corpus:    community_curated_v0 (annotation stamp 9b796de, judge opus-5)
           + so101_pick_place_v2 + so101_pick_place_clean (rejudged opus-5)
code:      c9ce380 (ablation launcher: outputs/run_ablation_stage2.sh)
panel:     plans/holdout_curated_v0_k4l2.json (17,204 core + 8,596 labeled frames)
report:    outputs/make_stage2_report.py (rerun to refresh charts)