admin/minimal-openvla-qwen05b
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add verified minimal OpenVLA Qwen2-0.5B reproduction

Browse Source
This commit is contained in:
seeker
2026-07-12 23:26:48 +08:00
commit 7c3180828a
6 changed files with 489 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+13
View File
@@ -0,0 +1,13 @@
.DS_Store
.venv*/
__pycache__/
*.py[cod]
.pytest_cache/
.mypy_cache/
output/
models/
*.pt
*.pth
*.bin
*.safetensors
.env
LICENSE
+21
View File
@@ -0,0 +1,21 @@
MIT License
Copyright (c) 2026 IP127000
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
README.md
+142
View File
@@ -0,0 +1,142 @@
# Minimal OpenVLA with Qwen2-0.5B VLM
A single-file, experimentally verified reproduction of the core OpenVLA training loop using the Hugging Face checkpoint [`lmms-lab/llava-onevision-qwen2-0.5b-ov`](https://huggingface.co/lmms-lab/llava-onevision-qwen2-0.5b-ov).
This is intentionally a **minimal OpenVLA-style experiment**, not a reproduction of the complete Stanford OpenVLA architecture. It demonstrates:
- image and language instruction input;
- a pretrained SigLIP visual encoder and multimodal projector;
- a Qwen2-0.5B language-model backbone;
- 256-bin discretization of normalized continuous actions;
- autoregressive prediction of seven robot action tokens;
- action-only cross-entropy supervision;
- overfitting and constrained action decoding on three synthetic examples.
![Training loss and action comparison](assets/training-results.png)
## Verified result
The experiment was run on an Apple M5 GPU through PyTorch MPS for 100 optimizer steps.
| Metric | Before training | After 100 steps |
| --- | ---: | ---: |
| Mean action loss | 15.2092 | 0.2834 |
| Teacher-forced action-token accuracy | 0.0% | 90.5% |
Two of the three autoregressive action predictions closely matched their training targets:
```text
red target = [-0.750, -0.500, 0.200, 0, 0, 0, 1]
red pred = [-0.746, -0.496, 0.199, 0.004, 0.004, 0.004, 0.996]
blue target = [0.650, 0.600, -0.150, 0, 0, 0, 1]
blue pred = [0.652, 0.598, -0.152, 0.004, 0.004, 0.004, 0.996]
```
The green example was confused with the blue example. This limitation is reported deliberately: the run proves that the complete multimodal training and decoding path works and strongly overfits, but the three-example run does not reach 100% sequence memorization.
## Architecture
```text
synthetic image ─→ SigLIP vision tower ─→ LLaVA projector ─┐
├→ Qwen2-0.5B → 7 action tokens
language instruction ──────────────────────────────────────┘
7 action tokens → 256-bin decoder → normalized 7-D continuous action
```
Compared with the Stanford OpenVLA paper, this experiment replaces the fused DINOv2+SigLIP encoder and Llama-2 7B backbone with the smaller LLaVA-OneVision SigLIP + Qwen2-0.5B stack. It preserves the central VLA mechanism rather than the exact model architecture or Open-X training recipe.
## Repository layout
```text
.
├── minimal_openvla_qwen05b.py # complete training and evaluation experiment
├── requirements.txt # verified Python dependencies
└── assets/training-results.png # plots from the verified 100-step run
```
The script expects a local checkout of [`IP127000/openVLA-Qwen2-0.5B`](https://github.com/IP127000/openVLA-Qwen2-0.5B) next to it:
```text
parent-directory/
├── minimal-openvla-qwen05b/
└── qwen2-openvla/
```
## Installation
Python 3.10-3.12 is recommended.
```bash
git clone https://github.com/IP127000/openVLA-Qwen2-0.5B.git qwen2-openvla
git clone https://github.com/IP127000/minimal-openvla-qwen05b.git
python3 -m venv .venv-openvla
source .venv-openvla/bin/activate
pip install --upgrade pip
pip install -r minimal-openvla-qwen05b/requirements.txt
```
The Hugging Face model is public. A token is normally unnecessary. If your environment requires authentication, pass it through the environment rather than putting it in source code:
```bash
export HF_TOKEN="your-token"
```
## Run
From the parent directory containing both repositories:
```bash
PYTORCH_ENABLE_MPS_FALLBACK=1 \
python minimal-openvla-qwen05b/minimal_openvla_qwen05b.py \
--steps 100 \
--learning-rate 5e-5 \
--save-dir minimal-openvla-qwen05b/output/minimal-openvla-overfit
```
On CUDA, `PYTORCH_ENABLE_MPS_FALLBACK` can be omitted.
The first run downloads approximately 1.8GB of LLaVA-Qwen weights plus the separate SigLIP vision tower. Later runs use the Hugging Face cache.
## What is trained
To keep memory and runtime manageable, the script freezes most parameters and trains:
- the multimodal projector;
- input token embeddings;
- the language-model output head;
- the final Qwen transformer block.
On Apple Silicon the model is trained in FP32 because FP16 AdamW produced NaNs during real testing. CUDA uses BF16 when supported and otherwise FP16.
## Synthetic dataset
The script creates three deterministic 224×224 images containing red, green, or blue squares. Each image/instruction pair is assigned a normalized seven-dimensional action:
```text
[Δx, Δy, Δz, Rx, Ry, Rz, gripper]
```
This dataset is only intended to verify that the multimodal model can learn and decode action tokens. It does not measure robotics performance or generalization.
## Known limitations
- Uses SigLIP only, not OpenVLA's fused DINOv2+SigLIP visual encoder.
- Uses Qwen2-0.5B instead of Llama-2 7B.
- Uses synthetic images rather than Open-X/RLDS demonstrations.
- Assumes actions are already normalized to `[-1, 1]`.
- Does not implement per-dataset quantile normalization or robot-specific unnormalization.
- The verified three-sample run reached 90.5%, not 100%, action-token accuracy.
## References
- [OpenVLA paper](https://arxiv.org/abs/2406.09246)
- [Official OpenVLA repository](https://github.com/openvla/openvla)
- [LLaVA-OneVision Qwen2-0.5B checkpoint](https://huggingface.co/lmms-lab/llava-onevision-qwen2-0.5b-ov)
- [openVLA-Qwen2-0.5B source](https://github.com/IP127000/openVLA-Qwen2-0.5B)
## License
MIT. Model checkpoints and upstream components remain subject to their respective licenses.
assets/training-results.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 139 KiB

minimal_openvla_qwen05b.py
+301
View File
@@ -0,0 +1,301 @@
#!/usr/bin/env python3
"""Minimal OpenVLA-style overfitting experiment on a 0.5B Qwen VLM.
This intentionally small experiment uses the Hugging Face checkpoint
``lmms-lab/llava-onevision-qwen2-0.5b-ov`` and the local LLaVA-OneVision
implementation under ``qwen2-openvla``. It demonstrates the core VLA loop:
image + language instruction -> autoregressive action tokens -> 7-D action
The script generates a few deterministic synthetic samples, discretizes actions
into 256 bins, fine-tunes a small subset of the VLM, and reports teacher-forced
action-token accuracy plus constrained autoregressive predictions.
Authentication is read from HF_TOKEN/HUGGING_FACE_HUB_TOKEN. Never put a token
in this file.
"""
from __future__ import annotations
import argparse
import itertools
import os
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable
import numpy as np
import torch
from PIL import Image, ImageDraw
ROOT = Path(__file__).resolve().parent
LLAVA_REPO = ROOT / "qwen2-openvla"
if not LLAVA_REPO.exists():
raise FileNotFoundError(
f"Expected the LLaVA implementation at {LLAVA_REPO}. "
"Clone https://github.com/IP127000/openVLA-Qwen2-0.5B there first."
)
sys.path.insert(0, str(LLAVA_REPO))
from llava.constants import DEFAULT_IMAGE_TOKEN, IGNORE_INDEX, IMAGE_TOKEN_INDEX # noqa: E402
from llava.mm_utils import get_model_name_from_path, process_images, tokenizer_image_token # noqa: E402
from llava.model.builder import load_pretrained_model # noqa: E402
@dataclass(frozen=True)
class ToySample:
image: Image.Image
instruction: str
action: np.ndarray
class ActionTokenizer:
"""OpenVLA-style uniform discretization over normalized actions in [-1, 1]."""
def __init__(self, tokenizer, bins: int = 256) -> None:
self.tokenizer = tokenizer
self.n_bins = bins
self.edges = np.linspace(-1.0, 1.0, bins + 1, dtype=np.float32)
self.centers = (self.edges[:-1] + self.edges[1:]) / 2.0
self.tokens = [f"<ACTION_{i}>" for i in range(bins)]
# The local model loader currently creates <ACTION_0>...<ACTION_270>.
# add_tokens is idempotent, so this also works if that behavior changes.
tokenizer.add_tokens(self.tokens)
self.token_ids = tokenizer.convert_tokens_to_ids(self.tokens)
if len(set(self.token_ids)) != bins:
raise RuntimeError("Action tokens were not assigned distinct tokenizer IDs")
def encode(self, action: np.ndarray) -> list[int]:
action = np.asarray(action, dtype=np.float32)
indices = np.digitize(np.clip(action, -1.0, 1.0), self.edges) - 1
indices = np.clip(indices, 0, self.n_bins - 1)
return [self.token_ids[int(i)] for i in indices]
def decode(self, token_ids: Iterable[int]) -> np.ndarray:
id_to_bin = {token_id: i for i, token_id in enumerate(self.token_ids)}
try:
indices = [id_to_bin[int(token_id)] for token_id in token_ids]
except KeyError as exc:
raise ValueError(f"Generated non-action token ID: {exc.args[0]}") from exc
return self.centers[np.asarray(indices)]
def make_toy_samples() -> list[ToySample]:
specs = [
((220, 50, 50), (45, 65), "move to the red square", [-0.75, -0.50, 0.20, 0.0, 0.0, 0.0, 1.0]),
((50, 190, 70), (130, 70), "move to the green square", [0.10, -0.45, 0.35, 0.0, 0.0, 0.0, -1.0]),
((55, 90, 220), (185, 175), "move to the blue square", [0.65, 0.60, -0.15, 0.0, 0.0, 0.0, 1.0]),
]
samples: list[ToySample] = []
for color, center, instruction, action in specs:
image = Image.new("RGB", (224, 224), "white")
draw = ImageDraw.Draw(image)
x, y = center
draw.rectangle((x - 24, y - 24, x + 24, y + 24), fill=color)
samples.append(ToySample(image, instruction, np.asarray(action, dtype=np.float32)))
return samples
def build_prompt(tokenizer, instruction: str) -> str:
user_text = f"{DEFAULT_IMAGE_TOKEN}\nWhat action should the robot take to {instruction}?"
if getattr(tokenizer, "chat_template", None):
return tokenizer.apply_chat_template(
[{"role": "user", "content": user_text}],
tokenize=False,
add_generation_prompt=True,
)
return user_text + "\nASSISTANT:"
def prepare_example(tokenizer, action_tokenizer: ActionTokenizer, sample: ToySample, device: torch.device):
prompt = build_prompt(tokenizer, sample.instruction)
prompt_ids = tokenizer_image_token(
prompt, tokenizer, image_token_index=IMAGE_TOKEN_INDEX, return_tensors="pt"
)
action_ids = torch.tensor(action_tokenizer.encode(sample.action), dtype=torch.long)
input_ids = torch.cat([prompt_ids, action_ids]).unsqueeze(0).to(device)
labels = torch.full_like(input_ids, IGNORE_INDEX)
labels[:, -len(action_ids) :] = action_ids.to(device)
return input_ids, labels
def prepare_image(image, image_processor, model, device: torch.device):
pixels = process_images([image], image_processor, model.config)
dtype = next(model.parameters()).dtype
if isinstance(pixels, torch.Tensor):
pixels = [item.to(device=device, dtype=dtype) for item in pixels]
else:
pixels = [item.to(device=device, dtype=dtype) for item in pixels]
return pixels, [image.size]
def set_trainable_subset(model) -> list[torch.nn.Parameter]:
"""Train projector, token embeddings/head, and the final Qwen block."""
model.requires_grad_(False)
modules = [model.get_model().mm_projector, model.get_model().embed_tokens, model.lm_head]
layers = getattr(model.get_model(), "layers", None)
if layers is not None and len(layers):
modules.append(layers[-1])
for module in modules:
module.requires_grad_(True)
return [parameter for parameter in model.parameters() if parameter.requires_grad]
@torch.no_grad()
def teacher_forced_metrics(model, tokenizer, action_tokenizer, samples, image_processor, device):
model.eval()
losses, correct, total = [], 0, 0
for sample in samples:
input_ids, labels = prepare_example(tokenizer, action_tokenizer, sample, device)
images, image_sizes = prepare_image(sample.image, image_processor, model, device)
output = model(
input_ids=input_ids,
labels=labels,
images=images,
image_sizes=image_sizes,
modalities=["image"],
)
losses.append(float(output.loss))
# Image patch embeddings expand the multimodal sequence, so output
# logits are longer than the original text labels. Action tokens remain
# the final tokens; the preceding logits predict those actions.
action_count = len(sample.action)
action_logits = output.logits[:, -(action_count + 1) : -1]
action_labels = labels[:, -action_count:]
predicted = action_logits.argmax(dim=-1)
correct += int((predicted == action_labels).sum())
total += action_count
return float(np.mean(losses)), correct / max(total, 1)
@torch.no_grad()
def generate_action(model, tokenizer, action_tokenizer, sample, image_processor, device):
input_ids = tokenizer_image_token(
build_prompt(tokenizer, sample.instruction),
tokenizer,
image_token_index=IMAGE_TOKEN_INDEX,
return_tensors="pt",
).unsqueeze(0).to(device)
allowed = torch.tensor(action_tokenizer.token_ids, device=device)
generated_ids: list[int] = []
# Manual greedy decoding avoids a LLaVA-OneVision generation incompatibility
# where prefix_allowed_tokens_fn receives an empty synthetic input_ids tensor
# after the wrapper converts multimodal inputs to inputs_embeds.
for _ in range(len(sample.action)):
images, image_sizes = prepare_image(sample.image, image_processor, model, device)
output = model(
input_ids=input_ids,
images=images,
image_sizes=image_sizes,
modalities=["image"],
use_cache=False,
)
allowed_logits = output.logits[0, -1, allowed]
next_id = int(allowed[allowed_logits.argmax()].item())
generated_ids.append(next_id)
input_ids = torch.cat([input_ids, torch.tensor([[next_id]], device=device)], dim=1)
return action_tokenizer.decode(generated_ids)
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--model-id", default="lmms-lab/llava-onevision-qwen2-0.5b-ov")
parser.add_argument("--steps", type=int, default=100)
parser.add_argument("--learning-rate", type=float, default=5e-5)
parser.add_argument("--log-every", type=int, default=10)
parser.add_argument("--save-dir", type=Path)
return parser.parse_args()
def main() -> None:
args = parse_args()
torch.manual_seed(0)
np.random.seed(0)
if torch.cuda.is_available():
device = torch.device("cuda")
dtype_name = "bfloat16" if torch.cuda.is_bf16_supported() else "float16"
load_device_map = {"": device.index or 0}
elif torch.backends.mps.is_available():
device = torch.device("mps")
dtype_name = "float16"
# The bundled LLaVA loader hard-codes a CUDA move whenever device_map
# is not "auto". Let it load automatically, then move to MPS below.
load_device_map = "auto"
else:
raise RuntimeError("A CUDA GPU or Apple-Silicon MPS device is required.")
# huggingface_hub/transformers automatically consume HF_TOKEN. We only
# check for it so private/gated downloads fail with a useful message.
if not (os.getenv("HF_TOKEN") or os.getenv("HUGGING_FACE_HUB_TOKEN")):
print("Note: HF_TOKEN is unset; continuing because this checkpoint may be public.")
model_name = get_model_name_from_path(args.model_id)
tokenizer, model, image_processor, _ = load_pretrained_model(
args.model_id,
model_base=None,
model_name=model_name,
device_map=load_device_map,
torch_dtype=dtype_name,
attn_implementation="sdpa",
)
if device.type == "mps":
# FP16 AdamW on MPS overflowed after one real optimizer step. The model
# fits in unified memory at FP32 and is stable enough for this tiny test.
model.to(device=device, dtype=torch.float32)
action_tokenizer = ActionTokenizer(tokenizer, bins=256)
model.resize_token_embeddings(len(tokenizer))
model.config.use_cache = False
samples = make_toy_samples()
trainable = set_trainable_subset(model)
optimizer = torch.optim.AdamW(trainable, lr=args.learning_rate, weight_decay=0.0)
initial_loss, initial_accuracy = teacher_forced_metrics(
model, tokenizer, action_tokenizer, samples, image_processor, device
)
print(f"before: loss={initial_loss:.4f}, action_token_accuracy={initial_accuracy:.3f}")
model.train()
for step, sample in zip(range(1, args.steps + 1), itertools.cycle(samples)):
input_ids, labels = prepare_example(tokenizer, action_tokenizer, sample, device)
images, image_sizes = prepare_image(sample.image, image_processor, model, device)
output = model(
input_ids=input_ids,
labels=labels,
images=images,
image_sizes=image_sizes,
modalities=["image"],
use_cache=False,
)
output.loss.backward()
torch.nn.utils.clip_grad_norm_(trainable, 1.0)
optimizer.step()
optimizer.zero_grad(set_to_none=True)
if step == 1 or step % args.log_every == 0:
print(f"step={step:04d} train_loss={float(output.loss.detach()):.4f}")
final_loss, final_accuracy = teacher_forced_metrics(
model, tokenizer, action_tokenizer, samples, image_processor, device
)
print(f"after: loss={final_loss:.4f}, action_token_accuracy={final_accuracy:.3f}")
for sample in samples:
predicted = generate_action(model, tokenizer, action_tokenizer, sample, image_processor, device)
print(f"{sample.instruction}\n target={sample.action.round(3)}\n pred ={predicted.round(3)}")
if args.save_dir:
args.save_dir.mkdir(parents=True, exist_ok=True)
trainable_names = {name for name, parameter in model.named_parameters() if parameter.requires_grad}
torch.save(
{name: value.detach().cpu() for name, value in model.state_dict().items() if name in trainable_names},
args.save_dir / "minimal_openvla_trainable.pt",
)
tokenizer.save_pretrained(args.save_dir)
print(f"Saved trainable state and tokenizer to {args.save_dir}")
if __name__ == "__main__":
main()
requirements.txt
+12
View File
@@ -0,0 +1,12 @@
torch>=2.5
torchvision>=0.20
transformers==4.46.3
accelerate==1.1.1
timm==1.0.11
einops>=0.8
einops-exts==0.0.4
huggingface-hub>=0.26,<1.0
numpy>=1.26
pillow>=10
sentencepiece>=0.2
protobuf>=4