commit 7c3180828aae655898b661a6798fa427e91c5d8c Author: seeker Date: Sun Jul 12 23:26:48 2026 +0800 Add verified minimal OpenVLA Qwen2-0.5B reproduction diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..d66d565 --- /dev/null +++ b/.gitignore @@ -0,0 +1,13 @@ +.DS_Store +.venv*/ +__pycache__/ +*.py[cod] +.pytest_cache/ +.mypy_cache/ +output/ +models/ +*.pt +*.pth +*.bin +*.safetensors +.env diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..696450f --- /dev/null +++ b/LICENSE @@ -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. diff --git a/README.md b/README.md new file mode 100644 index 0000000..cf929ce --- /dev/null +++ b/README.md @@ -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. diff --git a/assets/training-results.png b/assets/training-results.png new file mode 100644 index 0000000..ed13ea2 Binary files /dev/null and b/assets/training-results.png differ diff --git a/minimal_openvla_qwen05b.py b/minimal_openvla_qwen05b.py new file mode 100644 index 0000000..f14e76b --- /dev/null +++ b/minimal_openvla_qwen05b.py @@ -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"" for i in range(bins)] + + # The local model loader currently creates .... + # 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() diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..35b1c8e --- /dev/null +++ b/requirements.txt @@ -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