minimal-openvla-qwen05b/minimal_openvla_qwen05b.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "64ab06e2",
   "metadata": {},
   "source": [
    "# Minimal OpenVLA with Qwen2-0.5B — End-to-End Reproduction\n",
    "\n",
    "This notebook reproduces the complete minimal experiment from environment setup through action generation.\n",
    "\n",
    "> **Scope:** this is an OpenVLA-style minimal reproduction. It preserves the central `image + instruction → autoregressive action tokens → continuous action` loop, but replaces Stanford OpenVLA's fused DINOv2+SigLIP and Llama-2 7B stack with LLaVA-OneVision SigLIP + Qwen2-0.5B.\n",
    "\n",
    "## Table of contents\n",
    "\n",
    "1. [Paper and references](#1-paper-and-references)\n",
    "2. [Architecture](#2-architecture)\n",
    "3. [Runtime configuration](#3-runtime-configuration)\n",
    "4. [Environment installation](#4-environment-installation)\n",
    "5. [Source and weight download](#5-source-and-weight-download)\n",
    "6. [Dataset preparation](#6-dataset-preparation)\n",
    "7. [Model loading](#7-model-loading)\n",
    "8. [Minimal overfitting](#8-minimal-overfitting)\n",
    "9. [Evaluation and action generation](#9-evaluation-and-action-generation)\n",
    "10. [Plots and verified result](#10-plots-and-verified-result)\n",
    "11. [Limitations](#11-limitations)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3dc42f5f",
   "metadata": {},
   "source": [
    "## 1. Paper and references\n",
    "\n",
    "- **OpenVLA paper:** [OpenVLA: An Open-Source Vision-Language-Action Model](https://arxiv.org/abs/2406.09246)\n",
    "- **Project page:** [openvla.github.io](https://openvla.github.io/)\n",
    "- **Official source:** [openvla/openvla](https://github.com/openvla/openvla)\n",
    "- **0.5B VLM checkpoint:** [lmms-lab/llava-onevision-qwen2-0.5b-ov](https://huggingface.co/lmms-lab/llava-onevision-qwen2-0.5b-ov)\n",
    "- **LLaVA/OpenVLA adaptation used here:** [IP127000/openVLA-Qwen2-0.5B](https://github.com/IP127000/openVLA-Qwen2-0.5B)\n",
    "\n",
    "The original OpenVLA predicts discretized robot actions with a pretrained VLM. Its characteristic architecture uses fused DINOv2 and SigLIP image features, a two-layer projector, and Llama-2 7B. Actions are normalized, discretized into language-model tokens, trained with action-only cross entropy, and decoded back into continuous robot commands.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e04afff",
   "metadata": {},
   "source": [
    "## 2. Architecture\n",
    "\n",
    "```text\n",
    "+Original OpenVLA\n",
    "+image ─┬─ DINOv2 ─┐\n",
    "+       └─ SigLIP ─┴─ concat → 2-layer MLP → Llama-2 7B → action tokens\n",
    "+instruction ───────────────────────────────────────────────┘\n",
    "+\n",
    "+This minimal reproduction\n",
    "+image → SigLIP → LLaVA projector ─┐\n",
    "+                                  ├→ Qwen2-0.5B → 7 action tokens\n",
    "+instruction ──────────────────────┘\n",
    "+\n",
    "+7 action tokens → 256-bin decoder → normalized [Δx, Δy, Δz, Rx, Ry, Rz, gripper]\n",
    "+```\n",
    "+\n",
    "+Only the multimodal projector, token embeddings, LM head, and final Qwen block are trained. The experiment is designed to verify the learning path on a few samples, not robotics generalization.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec3624ea",
   "metadata": {},
   "source": [
    "## 3. Runtime configuration\n",
    "\n",
    "The defaults reproduce the reported 100-step experiment. Environment variables allow a fast smoke test or local assets:\n",
    "\n",
    "- `OPENVLA_STEPS=3` — shorten training for validation.\n",
    "- `OPENVLA_INSTALL_DEPS=0` — skip package installation in a prepared environment.\n",
    "- `OPENVLA_SKIP_DOWNLOAD=1` — use the Hugging Face cache/offline mode.\n",
    "- `OPENVLA_MODEL_PATH=/path/to/model` — load a local model directory.\n",
    "- `OPENVLA_LLAVA_REPO=/path/to/qwen2-openvla` — use a local source checkout.\n",
    "- `OPENVLA_DATA_JSON=/path/to/data.json` — use local examples instead of the synthetic dataset.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "768bd9ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, sys, json, subprocess\n",
    "from pathlib import Path\n",
    "\n",
    "REPO_ROOT = Path.cwd().resolve()\n",
    "if not (REPO_ROOT / \"minimal_openvla_qwen05b.py\").exists() and (REPO_ROOT.parent / \"minimal_openvla_qwen05b.py\").exists():\n",
    "    REPO_ROOT = REPO_ROOT.parent\n",
    "\n",
    "STEPS = int(os.getenv(\"OPENVLA_STEPS\", \"100\"))\n",
    "LEARNING_RATE = float(os.getenv(\"OPENVLA_LR\", \"5e-5\"))\n",
    "INSTALL_DEPS = os.getenv(\"OPENVLA_INSTALL_DEPS\", \"1\") == \"1\"\n",
    "SKIP_DOWNLOAD = os.getenv(\"OPENVLA_SKIP_DOWNLOAD\", \"0\") == \"1\"\n",
    "MODEL_ID = \"lmms-lab/llava-onevision-qwen2-0.5b-ov\"\n",
    "VISION_ID = \"google/siglip-so400m-patch14-384\"\n",
    "LOCAL_MODEL_PATH = os.getenv(\"OPENVLA_MODEL_PATH\")\n",
    "LOCAL_LLAVA_REPO = os.getenv(\"OPENVLA_LLAVA_REPO\")\n",
    "LOCAL_DATA_JSON = os.getenv(\"OPENVLA_DATA_JSON\")\n",
    "\n",
    "print({\n",
    "    \"repo_root\": str(REPO_ROOT), \"steps\": STEPS, \"install_deps\": INSTALL_DEPS,\n",
    "    \"skip_download\": SKIP_DOWNLOAD, \"local_model\": LOCAL_MODEL_PATH,\n",
    "    \"local_data\": LOCAL_DATA_JSON,\n",
    "})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89bc2b1b",
   "metadata": {},
   "source": [
    "## 4. Environment installation\n",
    "\n",
    "This cell installs the tested runtime into the current Jupyter kernel. On a clean machine, create and select a Python 3.10–3.12 virtual environment before opening the notebook. Re-running the cell is safe because pip skips satisfied packages.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf7ea935",
   "metadata": {},
   "outputs": [],
   "source": [
    "if INSTALL_DEPS:\n",
    "    subprocess.check_call([\n",
    "        sys.executable, \"-m\", \"pip\", \"install\", \"-r\", str(REPO_ROOT / \"requirements.txt\")\n",
    "    ])\n",
    "else:\n",
    "    print(\"Dependency installation skipped by OPENVLA_INSTALL_DEPS=0\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f67285d5",
   "metadata": {},
   "source": [
    "## 5. Source and weight download\n",
    "\n",
    "The model depends on the LLaVA adaptation source and a separate SigLIP visual tower. By default this cell clones/downloads them automatically. To use local assets, set the path environment variables in the configuration section.\n",
    "\n",
    "No access token is stored in this notebook. Public downloads normally work without one; if needed, set `HF_TOKEN` in the shell before launching Jupyter.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "509ee736",
   "metadata": {},
   "outputs": [],
   "source": [
    "if LOCAL_LLAVA_REPO:\n",
    "    llava_repo = Path(LOCAL_LLAVA_REPO).expanduser().resolve()\n",
    "else:\n",
    "    sibling = REPO_ROOT.parent / \"qwen2-openvla\"\n",
    "    nested = REPO_ROOT / \"qwen2-openvla\"\n",
    "    llava_repo = sibling if sibling.exists() else nested\n",
    "    if not llava_repo.exists():\n",
    "        subprocess.check_call([\n",
    "            \"git\", \"clone\", \"--depth\", \"1\",\n",
    "            \"https://github.com/IP127000/openVLA-Qwen2-0.5B.git\", str(llava_repo)\n",
    "        ])\n",
    "os.environ[\"OPENVLA_LLAVA_REPO\"] = str(llava_repo)\n",
    "print(\"LLaVA source:\", llava_repo)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f4208bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import snapshot_download\n",
    "\n",
    "if LOCAL_MODEL_PATH:\n",
    "    model_source = str(Path(LOCAL_MODEL_PATH).expanduser().resolve())\n",
    "elif SKIP_DOWNLOAD:\n",
    "    model_source = MODEL_ID\n",
    "    os.environ[\"HF_HUB_OFFLINE\"] = \"1\"\n",
    "else:\n",
    "    model_source = snapshot_download(MODEL_ID)\n",
    "    snapshot_download(VISION_ID)\n",
    "\n",
    "print(\"Model source:\", model_source)\n",
    "print(\"Weights are downloaded into the Hugging Face cache unless a local path is supplied.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25639a5d",
   "metadata": {},
   "source": [
    "## 6. Dataset preparation\n",
    "\n",
    "The default dataset contains three generated 224×224 images with colored squares and seven-dimensional normalized actions. To use local data, provide a JSON file with this schema:\n",
    "\n",
    "```json\n",
    "[\n",
    "  {\n",
    "    \"image\": \"/absolute/or/json-relative/image.png\",\n",
    "    \"instruction\": \"move to the object\",\n",
    "    \"action\": [-0.2, 0.1, 0.3, 0, 0, 0, 1]\n",
    "  }\n",
    "]\n",
    "```\n",
    "\n",
    "Actions must already be normalized to `[-1, 1]`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f7d12fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "import importlib.util\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "\n",
    "spec = importlib.util.spec_from_file_location(\"minimal_vla\", REPO_ROOT / \"minimal_openvla_qwen05b.py\")\n",
    "minimal_vla = importlib.util.module_from_spec(spec)\n",
    "sys.modules[spec.name] = minimal_vla\n",
    "spec.loader.exec_module(minimal_vla)\n",
    "\n",
    "if LOCAL_DATA_JSON:\n",
    "    data_path = Path(LOCAL_DATA_JSON).expanduser().resolve()\n",
    "    records = json.loads(data_path.read_text())\n",
    "    samples = []\n",
    "    for record in records:\n",
    "        image_path = Path(record[\"image\"])\n",
    "        if not image_path.is_absolute():\n",
    "            image_path = data_path.parent / image_path\n",
    "        samples.append(minimal_vla.ToySample(\n",
    "            Image.open(image_path).convert(\"RGB\"),\n",
    "            record[\"instruction\"],\n",
    "            np.asarray(record[\"action\"], dtype=np.float32),\n",
    "        ))\n",
    "else:\n",
    "    samples = minimal_vla.make_toy_samples()\n",
    "\n",
    "print(\"Number of samples:\", len(samples))\n",
    "for sample in samples:\n",
    "    print(sample.instruction, sample.action.tolist())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e98d111",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "fig, axes = plt.subplots(1, len(samples), figsize=(4 * len(samples), 3))\n",
    "if len(samples) == 1:\n",
    "    axes = [axes]\n",
    "for ax, sample in zip(axes, samples):\n",
    "    ax.imshow(sample.image)\n",
    "    ax.set_title(sample.instruction)\n",
    "    ax.axis(\"off\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f508a8c4",
   "metadata": {},
   "source": [
    "## 7. Model loading\n",
    "\n",
    "CUDA uses BF16 when available and otherwise FP16. Apple Silicon uses MPS with FP32 training because real FP16 AdamW testing produced NaNs. Most model parameters remain frozen.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ffe8283",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from llava.mm_utils import get_model_name_from_path\n",
    "from llava.model.builder import load_pretrained_model\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    device = torch.device(\"cuda\")\n",
    "    dtype_name = \"bfloat16\" if torch.cuda.is_bf16_supported() else \"float16\"\n",
    "    device_map = {\"\": 0}\n",
    "elif torch.backends.mps.is_available():\n",
    "    device = torch.device(\"mps\")\n",
    "    dtype_name = \"float16\"\n",
    "    device_map = \"auto\"\n",
    "else:\n",
    "    raise RuntimeError(\"A CUDA GPU or Apple-Silicon MPS device is required\")\n",
    "\n",
    "tokenizer, model, image_processor, _ = load_pretrained_model(\n",
    "    model_source, None, get_model_name_from_path(model_source),\n",
    "    device_map=device_map, torch_dtype=dtype_name, attn_implementation=\"sdpa\"\n",
    ")\n",
    "if device.type == \"mps\":\n",
    "    model.to(device=device, dtype=torch.float32)\n",
    "\n",
    "action_tokenizer = minimal_vla.ActionTokenizer(tokenizer, bins=256)\n",
    "model.resize_token_embeddings(len(tokenizer))\n",
    "model.config.use_cache = False\n",
    "trainable = minimal_vla.set_trainable_subset(model)\n",
    "print(\"device:\", device)\n",
    "print(\"trainable parameters:\", f\"{sum(p.numel() for p in trainable):,}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4eb7f49",
   "metadata": {},
   "source": [
    "## 8. Minimal overfitting\n",
    "\n",
    "Only action-token positions contribute to the loss. The default 100-step run takes a few minutes after weights are cached. Set `OPENVLA_STEPS=3` before launching Jupyter for a quick end-to-end smoke test.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e37333ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "import itertools, time\n",
    "\n",
    "torch.manual_seed(0)\n",
    "np.random.seed(0)\n",
    "optimizer = torch.optim.AdamW(trainable, lr=LEARNING_RATE, weight_decay=0.0)\n",
    "initial_loss, initial_accuracy = minimal_vla.teacher_forced_metrics(\n",
    "    model, tokenizer, action_tokenizer, samples, image_processor, device\n",
    ")\n",
    "print(f\"before: loss={initial_loss:.4f}, action_token_accuracy={initial_accuracy:.3f}\")\n",
    "\n",
    "history = [(0, initial_loss)]\n",
    "started = time.time()\n",
    "model.train()\n",
    "for step, sample in zip(range(1, STEPS + 1), itertools.cycle(samples)):\n",
    "    input_ids, labels = minimal_vla.prepare_example(tokenizer, action_tokenizer, sample, device)\n",
    "    images, image_sizes = minimal_vla.prepare_image(sample.image, image_processor, model, device)\n",
    "    output = model(\n",
    "        input_ids=input_ids, labels=labels, images=images,\n",
    "        image_sizes=image_sizes, modalities=[\"image\"], use_cache=False,\n",
    "    )\n",
    "    output.loss.backward()\n",
    "    torch.nn.utils.clip_grad_norm_(trainable, 1.0)\n",
    "    optimizer.step()\n",
    "    optimizer.zero_grad(set_to_none=True)\n",
    "    if step == 1 or step % max(1, STEPS // 10) == 0:\n",
    "        value = float(output.loss.detach())\n",
    "        history.append((step, value))\n",
    "        print(f\"step={step:04d} train_loss={value:.4f}\")\n",
    "\n",
    "print(f\"elapsed: {time.time() - started:.1f}s\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea0d10ff",
   "metadata": {},
   "source": [
    "## 9. Evaluation and action generation\n",
    "\n",
    "Teacher-forced accuracy measures individual action tokens. The second part performs explicit constrained autoregressive decoding: every generated position is restricted to the 256 action tokens and then converted back to a continuous value.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e151e39e",
   "metadata": {},
   "outputs": [],
   "source": [
    "final_loss, final_accuracy = minimal_vla.teacher_forced_metrics(\n",
    "    model, tokenizer, action_tokenizer, samples, image_processor, device\n",
    ")\n",
    "results = []\n",
    "for sample in samples:\n",
    "    prediction = minimal_vla.generate_action(\n",
    "        model, tokenizer, action_tokenizer, sample, image_processor, device\n",
    "    )\n",
    "    results.append((sample.instruction, sample.action.copy(), prediction))\n",
    "\n",
    "print(f\"after: loss={final_loss:.4f}, action_token_accuracy={final_accuracy:.3f}\")\n",
    "for instruction, target, prediction in results:\n",
    "    print(f\"{instruction}\\n  target={target.round(3)}\\n  pred  ={prediction.round(3)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd93da25",
   "metadata": {},
   "source": [
    "## 10. Plots and verified result\n",
    "\n",
    "The plot below is generated from the current notebook run. The repository's verified 100-step Apple M5 run produced:\n",
    "\n",
    "- loss: **15.2092 → 0.2834**;\n",
    "- action-token accuracy: **0% → 90.5%**;\n",
    "- red and blue examples closely matched their targets;\n",
    "- the green example was confused with the blue example.\n",
    "\n",
    "The mismatch is kept visible rather than presented as a perfect result.\n",
    "\n",
    "\n",
    "![Verified 100-step training result](assets/training-results.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd8f0218",
   "metadata": {},
   "outputs": [],
   "source": [
    "steps_seen, losses_seen = zip(*history)\n",
    "fig = plt.figure(figsize=(14, 8), constrained_layout=True)\n",
    "gs = fig.add_gridspec(2, len(results), height_ratios=[1.1, 1])\n",
    "ax = fig.add_subplot(gs[0, :])\n",
    "ax.plot(steps_seen, losses_seen, marker=\"o\", linewidth=2.5)\n",
    "ax.scatter([STEPS], [final_loss], label=f\"final eval loss={final_loss:.4f}\")\n",
    "ax.set(title=\"Minimal OpenVLA overfitting\", xlabel=\"optimization step\", ylabel=\"action loss\")\n",
    "ax.legend()\n",
    "\n",
    "names = [\"dx\", \"dy\", \"dz\", \"Rx\", \"Ry\", \"Rz\", \"grip\"]\n",
    "x = np.arange(7); width = 0.36\n",
    "for idx, (instruction, target, prediction) in enumerate(results):\n",
    "    axis = fig.add_subplot(gs[1, idx])\n",
    "    axis.bar(x - width/2, target, width, label=\"target\")\n",
    "    axis.bar(x + width/2, prediction, width, label=\"prediction\")\n",
    "    axis.set_title(instruction)\n",
    "    axis.set_xticks(x, names, rotation=35)\n",
    "    axis.set_ylim(-1.15, 1.15)\n",
    "    if idx == 0:\n",
    "        axis.set_ylabel(\"normalized action\")\n",
    "        axis.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "844c0926",
   "metadata": {},
   "source": [
    "## 11. Limitations\n",
    "\n",
    "- This is not the exact OpenVLA architecture: DINOv2 and Llama-2 7B are absent.\n",
    "- The synthetic dataset verifies memorization, not robot-policy generalization.\n",
    "- Local JSON actions must be normalized in advance.\n",
    "- Per-dataset `q01/q99` normalization and robot-specific action unnormalization are not implemented.\n",
    "- The full SigLIP vision tower remains large even though the language backbone is only 0.5B.\n",
    "- A successful few-shot overfit run validates the software path, not real-world safety or control quality.\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10+"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}