Add CPU inference script, update README with model details and perf stats

bd89217 verified 8 days ago

5.62 kB

license: mit
language:
  - en
library_name: transformers
pipeline_tag: reinforcement-learning
tags:
  - robotics
  - vla
  - vision-language-action
  - openvla
  - omnivla
  - robot
  - qwen
  - dinov2
  - siglip
datasets:
  - libero_90
  - cast
model-index:
  - name: openvla-micro
    results: []

OpenVLA-Micro

A drop-in replacement for OmniVLA 7B that runs ~14× faster with 0.997 action cosine similarity.

OpenVLA-Micro is a compact Vision-Language-Action model that replaces OmniVLA 7B (DINOv2-L + SigLIP-so400m + Llama-2 7B) with a much smaller stack (DINOv2-S/14 + SigLIP-B/16 + Qwen2.5 0.5B) while preserving compatibility with OmniVLA's pretrained action head through a learned hidden-state shim (896→4096).

Model Architecture

Component	Encoder	Output Dim
Vision (DINO)	`DINOv2-S/14` (facebook/dinov2-small)	256 tokens × 384d → MLP → 8704
Vision (SigLIP)	`SigLIP-B/16` (google/siglip-base-patch16-224)	196 tokens × 768d → MLP → 8704
Projector	Linear(8704→896) + GELU + Linear(896→896)	452 tokens × 896d
LLM	`Qwen2.5-0.5B` (24 layers, 896 hidden, 8 heads)	Variable × 896d
Shim	Linear(896→2048) + GELU + Linear(2048→4096)	32 action tokens × 4096d
Action Head	OmniVLA's pretrained head (unchanged)	8 chunks × 7-DoF

Total parameters: ~0.6B (vs 7B for OmniVLA/OpenVLA).

Performance

vs OmniVLA 7B (teacher)

Metric	Value	Note
Hidden state cosine	0.63	Last-layer HS at action positions
Action cosine	0.997	After OmniVLA action head
Action MSE	~0.001	Effectively identical predictions
Inference speed	~14× faster	0.5B vs 7B LLM

Trained on 1000 CAST episodes (17k steps) via hidden-state distillation. The shim was trained for ~21k steps (plateaued at ~8k).

vs OpenVLA 7B (original)

OpenVLA-Micro is distilled from OmniVLA, which itself fine-tuned OpenVLA with 32 action tokens and a modified action head. Direct comparison with the original OpenVLA is not apples-to-apples due to different action tokenization, but the ~14× speedup and near-lossless action quality relative to OmniVLA apply similarly vs OpenVLA.

Quick Start

from PIL import Image
from modeling_openvla_micro import OpenVLAMicro

model = OpenVLAMicro.from_pretrained("theguy21/openvla-micro", device="cuda")
model.eval()

image = Image.open("demo.jpg").convert("RGB")
action = model.predict_action(image, "pick up the red block")
print(action)  # [0.12, -0.03, 0.45, -0.01, 0.22, 0.08, -0.15]

CLI — GPU

python inference.py --image demo.jpg "pick up the red block"

CLI — CPU / Edge

# Standard CPU (~6GB RAM, 3-5 sec/step)
python inference_cpu.py --image demo.jpg "pick up the red block"

# Low-RAM CPU (~2.5GB RAM, requires bitsandbytes)
python inference_cpu.py --low-ram --image demo.jpg "pick up the red block"

As an OmniVLA drop-in replacement

Use OpenVLAMicroWrapper (from model_wrapper.py) to expose the same forward interface as OmniVLA's VLAForActionPrediction:

from model_wrapper import OpenVLAMicroWrapper
from modeling_openvla_micro import DinoSigLIPEncoder, CombinedProjector, ShimMLP

ckpt = torch.load("openvla-micro-distill.pt", map_location="cpu")
ve = DinoSigLIPEncoder()
ve.load_state_dict(ckpt["model"]["vision_backbone"])
# ... (see model_wrapper.py for full example)

output = vla(input_ids, attention_mask, pixel_values, labels=labels, output_hidden_states=True)
actions_hidden_states = extract_actions(output.hidden_states[-1], labels)
predicted_actions = omnivla_action_head.predict_action(actions_hidden_states, modality_id)

Architecture Diagram

Image (224×224)
  ├── DINOv2-S/14 → 256 patches × 384d → ShimMLP(384→8704)
  └── SigLIP-B/16 → 196 patches × 768d → ShimMLP(768→8704)
       └── Concat (452 tokens) → Linear(8704→896) → GELU → Linear(896→896)
              └── Qwen2.5 0.5B (24 layers, 896 hidden)
                     └── Hidden State Shim (896→2048→4096)
                            └── OmniVLA Action Head (pretrained, frozen)
                                   └── 8 chunks × 7-DoF actions

Files

File	Size	Description
`modeling_openvla_micro.py`	15 KB	Model definitions
`model_wrapper.py`	11 KB	OmniVLA-compatible interface
`inference.py`	1.5 KB	GPU/CPU CLI inference
`inference_cpu.py`	2 KB	Edge device inference (with low-RAM mode)
`train_shim.py`	15 KB	Reference shim training script
`config.json`	1.2 KB	Model configuration
`openvla-micro-merged.pt`	1.6 GB	Base checkpoint (no shim, 896-dim output)
`openvla-micro-distill.pt`	1.6 GB	Full checkpoint (with baked-in shim, 4096-dim)

Which checkpoint to use?

openvla-micro-distill.pt — recommended. Outputs 4096-dim hidden states that plug directly into OmniVLA's action head. One-step inference.
openvla-micro-merged.pt — base model only (896-dim). Use if you want to train your own shim or action head.

Requirements

torch>=2.0.0
torchvision>=0.15.0
transformers>=4.38.0
timm>=0.9.0
Pillow>=10.0.0
numpy>=1.24.0

For low-RAM CPU: bitsandbytes>=0.43.0

Training the Shim

python train_shim.py \
    --cache-dir ./teacher_cache \
    --data-dir ./dataset \
    --base-model openvla-micro-merged.pt \
    --teacher-dim 4096

See train_shim.py for full options. The script expects pre-cached teacher hidden states; adapt DistillDataset to your format.

License

MIT