openvla-oft-libero-early-exit

OpenVLA-OFT Early Exit: Decoder-at-Layer Study (LIBERO-Spatial)

Trained readouts and complete evaluation artifacts for a depth-compression study of OpenVLA-OFT 7B on LIBERO-Spatial. Core question: is the final LLM layer's advantage for action decoding better features or feature–decoder compatibility?

Answer: compatibility. Retraining the action head at intermediate layers recovers near-full closed-loop performance from half the network depth, and tuning a few attention blocks recovers full performance from layer 12 of 32.

Headline results (closed-loop rollout success, 10 trials × 10 tasks, seed 7)

Retrained L1 action head (151M params, fresh init per layer, LLM frozen, trained on 20 episodes/task):

Reference: full model with original head = 98/100 (matches the published result).

Additional arms:

Key findings:

1. No cliff at the final layer — the prior diagnostic observation that only layers 31–32 work with the pretrained head is a head-compatibility artifact, fully repaired by retraining the head (or even an 8.4M adapter in front of the frozen original head).
2. The information boundary sits between layers 12 and 16; below it, no static readout of frozen features passes ~80/100 — but tuning as little as one attention block below layer 12 jumps success to 93, and four blocks reach baseline parity (98).
3. Offline action-prediction metrics do not predict closed-loop success (e.g. the unfreeze arm moves offline L1 by 7% while rollout success moves 77 → 98).

Full write-up: report.md / report.pdf.

Repository layout

heads/                head_layer{n}.pt (fresh-init L1RegressionActionHead per layer)
                      + per-sample eval predictions head_pred_eval_layer{n}.npy
adapters/             adapter{A|B}_layer{n}.pt (A: pairs with frozen original head)
unfreeze/             unfreeze_k{1,2,4}.pt (tuned blocks 12-k..11 + head, attach@12)
rollout_*.json        per-episode rollout records (all 20 configurations)
rollout_*.queries.npz per-query executed vs stock action chunks (drift analysis)
*_results.json        offline metrics (heads, adapters, unfreeze, CKA, sweep summary)
offline_check/        fresh offline sweep used for stack validation
hidden/               extraction metadata + normalized GT labels (raw caches excluded, see below)
code_snapshot/        all analysis code (training, rollouts, verification, figure, this study)
env/                  setup_env.sh + pip_freeze.txt + conda_env.yml (exact environment)
logs/                 complete run logs, from driver install to final rollout
verify: code_snapshot/verify_results.py

Verifying every reported number (no GPU needed)

pip install numpy                # the only requirement for this check
python code_snapshot/verify_results.py --results-dir . --ref-dir baseline_reference --skip-cka
# expected final line: ALL CHECKS PASSED  (exit code 0)

(The min-model provenance check below additionally needs torch, safetensors, and — for --check-base — huggingface_hub. Both commands were tested on a fresh clone of this repo exactly as written.)

Fully self-contained: dataset_statistics.json and the baseline-study reference artifacts (baseline_reference/) are bundled. See REPRODUCING.md for the three verification tiers (recompute-only → re-rollout with these checkpoints → full re-derivation from public sources).

Reproducing / regenerating

• Environment: env/setup_env.sh (mirrors the exact build; pins in env/pip_freeze.txt).
• Base model: moojink/openvla-7b-oft-finetuned-libero-spatial; data: openvla/modified_libero_rlds (libero_spatial_no_noops). Both public.
• The raw hidden-state caches (~60 GB) are excluded; they regenerate deterministically with code_snapshot/extract_hidden2.py and code_snapshot/extract_fullseq.py (single A100 pass each; built-in gates verify correctness — the study machine reproduced the base model's predictions bitwise).
• Head/adapter/unfreeze training: train_heads2.py, train_adapters2.py, train_unfreeze.py; rollouts: rollout_eval2.py, rollout_unfreeze.py.

min_model/ — the minimum-parameter deployment bundle

min_model/model.safetensors is the assembled exit-at-layer-12 model: vision encoders + projector + embeddings + LLM blocks 0–11 (blocks 8–11 are the k=4 tuned copies) + proprio projector + L1 action head. Blocks 12–31, the final RMSNorm, and the LM head are removed.

Forward pass: standard OpenVLA-OFT pipeline, stopping at block 11 and decoding the action-token hidden states with action_head (bidirectional LLM attention — see code_snapshot/unfreeze_common.py; the parity rollout used exactly these weights via code_snapshot/rollout_unfreeze.py, record: rollout_unfreeze_k4.json).

Provenance is fully verifiable: min_model/manifest.json records, per tensor, its SHA256 and source — either a named tensor of the public base checkpoint or a key of unfreeze/unfreeze_k4.pt (whose offline predictions and 98/100 rollout are covered by the main verify script). Check it:

python code_snapshot/verify_min_model.py --bundle min_model --unfreeze-ckpt unfreeze/unfreeze_k4.pt
# expected: MIN-MODEL BUNDLE VERIFIED (exit 0); add --check-base to also
# byte-compare every frozen tensor against a fresh download of the base repo

Using a checkpoint

import torch
from prismatic.models.action_heads import L1RegressionActionHead  # openvla-oft repo

head = L1RegressionActionHead(input_dim=4096, hidden_dim=4096, action_dim=7)
head.load_state_dict(torch.load("heads/head_layer16.pt"))
# feed hidden_states[16][:, action_token_positions, :]  ->  (B, 8, 7) normalized actions

Note: OpenVLA-OFT uses bidirectional attention in the LLM (not causal). If you re-run decoder blocks standalone (as the unfreeze checkpoints require), follow code_snapshot/unfreeze_common.py.

License and attribution

This release is licensed under the Llama 2 Community License (see LICENSE): the OpenVLA backbone derives from Llama-2-7B, and unfreeze/unfreeze_k*.pt contain modified copies of Llama-2-derived transformer blocks.

Llama 2 is licensed under the LLAMA 2 Community License, Copyright © Meta Platforms, Inc. All Rights Reserved.

Built on OpenVLA-OFT (Kim et al.) and the LIBERO benchmark. Base checkpoint and dataset by the OpenVLA team.