JackLiu0406/b1k-checkpoints

BEHAVIOR-1K DA3 — K/V-split spatial VLA checkpoints

Two checkpoints from the DA3 spatial-reasoning line on the 2025 BEHAVIOR-1K winner (PiBehavior, a pi0.5 JAX fork): gemma_2b VLM + gemma_300m action expert + a DA3 spatial cross-attention branch, finetuned on the single task clean_up_your_desk.

In these checkpoints the DA3 spatial injection is load-bearing — the model depends on the bank (zeroing or degenerating it is very costly) — achieved by a K/V-split redesign of the spatial bank (below). See the honest specificity caveat under "Geometry verdict".

What's here

folder	run	step	action_loss	recipe in one line
kvsplit_spatretrain_step19999/	spatretrain	19999 (final)	~0.103 (best)	warm-start from the K/V-split base; spatial frozen 0–2k → ramp to 5e-4 peak at 10k → decay (geom anneals to 1/10 peak, core to 1/5); base held low at 2e-5
kvsplit_spatretrain_step18000/	spatretrain	18000	~0.115	(earlier snapshot of the same run)
kvsplit_spatretrain_step12000/	spatretrain	12000	0.162	(earlier snapshot)
kvsplit_strongbase_step15000/	strongbase	15000 (final)	~0.18	fresh from the 50-task meta ckpt; base ramps gradually 0→1e-4; spatial 3.5e-4 from step 0

Each folder contains params/ (bf16 weights) and assets/ (norm stats). Optimizer state (train_state/) is intentionally omitted — these are for eval/inference, not resume.

Headline: the original K/V-split run plateaued at action_loss ≈ 0.19 (base throttled to 1/4 peak). spatretrain — aggressively re-training the spatial branch on top of the load-bearing checkpoint — broke that plateau to 0.14. The lever that mattered was spatial training, not a hotter base: strongbase's full 1e-4 base only reached ~0.18.

Architecture — the K/V split (why geometry is now used)

The earlier additive bank (spatial = fused + pos_emb + ray_emb + view_emb) let input-independent "where" terms flood the attention value stream; under broad attention they averaged into a batch-constant bank that the model learned to use as a bias (wrong/shuffled geometry cost nothing).

The K/V split separates address from payload in the perceiver cross-attention:

• values (payload) = DA3 latents + a depth encoding — the only thing that can flow into the bank
• keys (address) = pos_emb + ray_emb(Plücker-6) + view_emb — added to the keys only, so they steer which patches each query reads but are structurally excluded from the output

A constant can now route but can never enter the bank, so the bank carries only per-sample content. Two more pieces: bank-centering (subtract the batch mean → removes any residual constant DOF) and depth-dropout 0.5 (zero the explicit depth encoding for half the training samples, forcing the DA3 features to carry geometry redundantly).

Geometry verdict (measured on the K/V-split base, GPU features, BS=16)

• T1 per-sample: cross-sample cosine ≈ 0 (was +0.98 in the constant-collapse era); ‖b−mean‖/‖b‖ = 1.0
• T2 feature-driven: swapping a sample's DA3 features changes its bank by 0.4–1.3 (≈ complete)
• T3 load-bearing (robust): zeroing the injection costs +40–54% action_loss and a degenerate/mean bank +340–530% — consistent across many batches. The model genuinely depends on the bank.
• T3 sample-specificity (NOT robust — honest caveat): shuffling geometry across samples costs only ~0% ± 6% over 6 batches (single-batch draws range −7% to +21% — earlier "+21%" was one lucky batch). So the model reads the pathway but does not reliably use which per-sample geometry it gets — a large generic component remains. Judge specificity with multi-batch shuffle-damage, never one draw. (This motivates the next arch: cross-view 3D fusion + perceiver locality to make it discriminative.)

How to load

These require the PiBehavior + openpi code with the DA3 spatial branch, built with the K/V-split flags (the arch differs from pre-split checkpoints — loading with the wrong flags will mismatch):

USE_DA3_FULL=1 DA3_KV_SPLIT=1 DA3_BANK_CENTER=1 DA3_SCALE=1.0
DA3_QK_NORM=1 DA3_PERC_NORM_FINAL=1 DA3_PERC_NORM_OUT=0 DA3_POS_EMB_SCALE=0.25
DA3_LOGIT_GAIN_INIT=3.0 DA3_INJ_GAIN_MAX=8.0 DA3_PERC_GAIN_INIT=3.0 DA3_PERC_GAIN_MAX=8.0

Point the weight loader at <folder>/params.

How to verify the geometry (specificity check)

verify/_da3_spatial_specificity.py runs T1/T2/T3. Use GPU-extracted DA3 features (CPU features mismatch the trained model) and BS ≥ 16 (BS=6 draws near-identical desk frames → the shuffle arm loses statistical power). Set the same DA3_KV_SPLIT=1 DA3_BANK_CENTER=1 flags so the rebuilt model matches the checkpoint. The decisive number is the shuffled-bank damage: >0 (ideally >10%) means the model is using per-sample geometry; ≈0 means it collapsed to a constant.