report.md

# Decoder-at-Layer Sweep for OpenVLA-OFT — Technical Report

**Question studied:** when OpenVLA-OFT decodes robot actions from its final LLM layer,
is the final layer's advantage due to *better features* or to *feature–decoder
compatibility*? The answer determines whether the model's depth can be cut at inference.

**Summary of findings:** attaching and training the real action head at intermediate
layers recovers 94–100/100 closed-loop success from layer 16 upward (half the network);
an 8.4M-parameter adapter in front of the *frozen original* head recovers most of it;
and tuning as few as 4 transformer blocks recovers full baseline success from layer 12
— i.e. with 51% of the deployed parameters removed. The final layer holds no privileged
action information; its advantage is a small, learnable compatibility transformation.

**Setup:** OpenVLA-OFT 7B with the L1 regression action head
(`moojink/openvla-7b-oft-finetuned-libero-spatial`), LIBERO-Spatial,
`libero_spatial_no_noops` data, seed 7. Offline evaluation: 1,491 action-chunk samples
(first 10 episodes per task). Readout training: 20 disjoint episodes per task.
Closed-loop evaluation: 10 trials × 10 tasks per configuration, fixed benchmark initial
states. Hardware: A100-SXM4-40GB, torch 2.2.0+cu121 (exact pins: `env/pip_freeze.txt`).
All analysis code: `code_snapshot/`.

## 1. Experimental-stack validation

| Check | Result |
|---|---|
| Model + headless-simulator smoke tests | pass |
| Offline final-layer L1 (1,491 samples) | **0.008120** (reference value 0.0081, tolerance ±5e-4) |
| Live bitwise invariant (instrumented layer-32 decode == stock output) | 0.0 on every sample of every run |
| Stock predictions vs reference artifacts (`baseline_reference/`) | **bitwise identical (max dev 0.0)** |
| Control: layer-16 linear probe retrained on the fresh extraction | L1 0.1733 vs reference 0.1724 (tolerance 0.01) |

The evaluation stack reproduces the reference predictions exactly, so differences
reported below are attributable to the interventions.
**Replicate:** `python code_snapshot/gate2_check.py`, `python code_snapshot/gate3_probe_check.py`;
recomputed independently by `code_snapshot/verify_results.py` (checks 1–2).

## 2. Main result: success vs attach layer

The real `L1RegressionActionHead` (151M params) is fresh-initialized per layer and
trained on cached action-token hidden states with the LLM frozen (lr 3e-4, 60 epochs,
cosine schedule, early stopping on a 10%-episode validation split; the recipe was tuned
at layer 32 only and then applied unchanged to all layers).

| Attach layer | 1 | 2 | 6 | 12 | 16 | 20 | 24 | 28 | 30 | 31 | 32 |
|---|---|---|---|---|---|---|---|---|---|---|---|
| Offline eval L1 | .205 | .189 | .168 | .153 | .129 | .095 | .048 | .024 | .019 | .017 | .013 |
| **Rollout success /100** | **11** | **24** | **22** | **77** | **97** | **94** | **97** | **99** | **100** | **99** | **100** |

Reference points: full model with the original head 98/100; a 200k-parameter linear
probe at layer 16 achieves 77/100; the original head reused *frozen* at layer 31
achieves 70/100 (probe and reference records in `baseline_reference/`).

**Findings:**

1. **Reusing the pretrained head at intermediate layers fails for compatibility
   reasons, not information reasons.** With a per-layer retrained head, offline L1 is
   strictly monotonic in depth (0.205 → 0.013); the error spike that the frozen
   pretrained head shows mid-network (peaking at layer 24) vanishes. Closed-loop,
   layer 24 rolls out at 97/100.
2. **No cliff at the final layer.** Retrained heads at layers 16–32 all land at
   94–100/100. **Verdict: the final-layer advantage is a compatibility artifact, fully
   repaired by retraining the head at the attach layer.**
3. **The information boundary for static readouts sits between layers 12 and 16**
   (77 → 97). Below it the curve collapses (22/100 at layer 6; the 11–24 floor reflects
   tasks solvable by trivial policies). Early exit at half depth costs ~1 point of
   success (97 vs the 98 full-model reference; the retrained head at layer 32 itself
   reaches 100).
4. **Readout capacity matters closed-loop:** at layer 16, the real head reaches 97
   vs the linear probe's 77 on identical training data.
5. **Offline action metrics do not predict closed-loop success** — the offline L1
   ordering across layers 16–24 does not match the rollout plateau, and §4 shows a
   starker case. All accept/reject conclusions here rest on rollouts only.

**Representation analysis (linear CKA, action-token positions, all 33 layers):**
adjacent-layer similarity is 0.95–0.99 through the mid-network, with three landmarks: a
large early reorganization at layers 2–3 (CKA 0.11–0.17), a dip at 27–28 (0.94), and
0.92 at 31→32 (the final entry includes the last RMSNorm). Similarity to the final
layer rises smoothly from 0.33 (layer 3) to 0.91 (layer 28+): representations converge
gradually — consistent with the smooth success curve and inconsistent with any single
"decision layer". Values: `cka_results.json`.

**Replicate:** every number above recomputes from the raw artifacts in this repository
via `python code_snapshot/verify_results.py --results-dir . --ref-dir baseline_reference --skip-cka`
(exit 0 = pass). Sweep table: `sweep_results.json`; per-episode records:
`rollout_head_layer*.json`.

## 3. Compatibility mechanism: adapters into the frozen original head

Adapter = token-wise residual MLP, LayerNorm→4096→1024→GELU→4096 (+x), **8.4M params**
(1/18 of the head's size), same training recipe. Variant A trains ONLY the adapter and
feeds the **frozen original checkpoint head**; variant B trains adapter + fresh head.

| Rollout success /100 at layer | 12 | 16 | 30 |
|---|---|---|---|
| A: adapter + frozen original head | 77 | 92 | 98 |
| B: adapter + trained head | 81 | 96 | 99 |
| Retrained full head (151M, §2) | 77 | 97 | 100 |
| Offline eval L1 (A) | .149 | .119 | .014 |

**Findings:**

1. **The compatibility artifact is a small, learnable transformation.** An 8.4M
   projector suffices for the *untouched* original head to read layer-30 features at
   98/100 (= full-model reference) and layer-16 features at 92/100.
2. **At layer 12, readout capacity is irrelevant:** adapter+frozen head (77),
   adapter+trained head (81), and the 151M retrained head (77) are statistically
   indistinguishable (SE ≈ 4). No *static readout* of frozen layer-12 features exceeds
   ~80/100 — §4 shows this ceiling belongs to the frozen features, not the information.
3. Offline, the adapter variants match or beat the retrained head at every layer —
   another instance of offline metrics compressing real closed-loop differences.

**Replicate:** `adapter_results.json` + `rollout_adapter*_layer*.json`; recomputed by
the verify script. Training code: `code_snapshot/train_adapters2.py`.

## 4. Tuned attention blocks below layer 12

Attach at layer 12 and jointly train the head plus the last k decoder blocks below it
(blocks 12−k .. 11; block lr 2e-5, head lr 3e-4, 12 epochs, early stopping). Training
and deployment share an identical block-recompute path from a layer-8 full-sequence
cache, gated functionally: the recompute with original weights reproduces the layer-12
sweep result exactly (offline L1 0.1531) before any training.

| k tuned blocks | 0 (frozen, §2) | 1 | 2 | 4 |
|---|---|---|---|---|
| Trainable params | 151M (head) | +202M | +404M | +808M |
| Offline eval L1 | .153 | .142 | .142 | .140 |
| **Rollout success /100** | **77** | **93** | **95** | **98** |

**Findings:**

1. **The layer-12 ceiling is shallow.** One tuned block lifts closed-loop success
   77 → 93; four tuned blocks reach **98/100 = full-model reference** while exiting at
   layer 12 (20 of 32 blocks, the final norm, and the LM head removed).
2. Contrast with §3: at layer 12, more *readout* capacity never beat 81 — but letting
   *attention over the full sequence* adapt recovers everything. What limits layer 12
   is cross-token integration, not information in the action-token features.
3. **The sharpest offline↔rollout dissociation observed:** offline L1 improves 7%
   (0.153 → 0.142, flat in k) while rollout success moves 77 → 93 → 95 → 98. Offline
   metrics cannot rank these configurations at all.

**Replicate:** `unfreeze_results.json` + `rollout_unfreeze_k*.json`; checkpoints in
`unfreeze/`; training code `code_snapshot/train_unfreeze.py`, rollout code
`code_snapshot/rollout_unfreeze.py`. Note: the LLM uses **bidirectional** attention
(not causal) — `code_snapshot/unfreeze_common.py` documents the verified re-run
semantics and the bf16 kernel-noise bounds.

## 5. The assembled minimum-parameter model

`min_model/model.safetensors` packages the best size/accuracy point (§4, k=4): vision
encoders + projector + embeddings + LLM blocks 0–11 (8–11 tuned) + proprio projector +
action head.

| | Params | Size | LIBERO-Spatial success |
|---|---|---|---|
| Full OpenVLA-OFT stack | 7.709 B | ~15.5 GB | 98/100 |
| **min_model (exit@12, 4 tuned blocks)** | **3.772 B (48.9%)** | **7.9 GB** | **98/100** |

Parameter count, not latency, is the claim: the rollout above ran the full forward with
the tuned blocks applied on the side (`rollout_unfreeze_k4.json`); a truncated forward
implementation is required to realize wall-clock savings.

**Replicate/verify:** every tensor's SHA256 and source (public base checkpoint or
`unfreeze/unfreeze_k4.pt`) is recorded in `min_model/manifest.json`; check with
`python code_snapshot/verify_min_model.py --bundle min_model --unfreeze-ckpt unfreeze/unfreeze_k4.pt`
(add `--check-base` to byte-compare all frozen tensors against the base repo).

## 6. Included measurement data

Every rollout logged, per policy query, the executed normalized action chunk and the
stock (original head, layer 32) chunk (`rollout_*.queries.npz`). Mean |executed −
stock| grows with attach distance from the final layer (0.010 at layer 32 → 0.084 at
layer 16) and rises modestly within episodes (layer 16: 0.073 in the first third of an
episode → 0.087 in the last third) — bounded drift, consistent with the high success
rates. Complete run logs are in `logs/`.

## 7. Replication guide

Three tiers, detailed in `REPRODUCING.md`:

1. **Verify every reported number** (CPU, ~2 min): the two verify commands above.
2. **Re-run rollouts with the released checkpoints** (1 GPU, hours):
   `env/setup_env.sh` rebuilds the exact environment; `code_snapshot/rollout_eval2.py`
   / `rollout_unfreeze.py` re-run any configuration with fixed seeds and benchmark
   initial states.
3. **Re-derive everything from public sources** (~1 GPU-day): base checkpoint and
   dataset are public; extraction, training, and rollout scripts include built-in
   gates (bitwise stock-reproduction, probe-reproduction control, sweep-reproduction
   check) that re-validate the chain at each stage.