Add verified minimal OpenVLA Qwen2-0.5B reproduction

.gitignore

@@ -0,0 +1,13 @@
+.DS_Store
+.venv*/
+__pycache__/
+*.py[cod]
+.pytest_cache/
+.mypy_cache/
+output/
+models/
+*.pt
+*.pth
+*.bin
+*.safetensors
+.env

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.

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.

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"<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

@@ -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