Add verified minimal OpenVLA Qwen2-0.5B reproduction
This commit is contained in:
seeker
@@ -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()
|
||||
Reference in New Issue
Block a user