""" Qwen2.5-VL-3B perception + high-level planner (v2). Uses the locally-cached Qwen/Qwen2.5-VL-3B-Instruct weights (hf_cache/) as the System-2 perception+reasoning layer: given the scene image and the operator's instruction it returns an ordered high-level plan that the learned low-level action model (models/policy_learned.py) then executes. The VLM decides WHICH cubes to move and the relation (into the box / stacking order / onto which cube); the policy computes the exact place geometry and the trained action model produces the joint trajectory. It loads the 3B model in 4-bit (fits the ~6 GB laptop GPU via bitsandbytes) and is queried once per instruction (deliberative, ~1-2 s), not in the control loop. If transformers / a GPU / the weights are unavailable it degrades to the same keyword parser the policy already uses, so nothing breaks. Output (consumed by _LearnedRX1Policy): {"plan": {"type":"stack"|"sort", "colors":[...]} | None, "queue": [ {"color":c, "kind":"box"|"on", "on":c2}, ... ] } """ from __future__ import annotations import json import logging import os import re from pathlib import Path from typing import Optional import numpy as np log = logging.getLogger(__name__) _MODEL_ID = "Qwen/Qwen2.5-VL-3B-Instruct" _COLORS = ("red", "blue", "green", "yellow") _SYS = ( "You are the perception and planning module of a tabletop robot. You see an " "image of coloured cubes (red, blue, green, yellow) and an open box, and an " "instruction. Reply with ONLY a JSON object, no prose:\n" '{"task":"stack"|"sort"|"pick"|"place","colors":[ordered colours],' '"onto":,"into_box":}\n' "Rules: 'stack' builds a tower in colours order (first = base); 'sort'/'tidy'/" "'clear' = put every cube into the box; 'pick/place X [on Y|in box]' = a single " "move. Only use colours actually present in the image." ) def _keyword_plan(instruction: str) -> dict: """Deterministic fallback identical to the policy's built-in grammar.""" t = (instruction or "").lower() colors = [c for c in _COLORS if c in t] on = (" on " in t or "onto" in t or "on top" in t) if any(k in t for k in ("stack", "tower", "pile")): if len(colors) >= 2 and on: return {"plan": None, "queue": [{"color": colors[0], "kind": "on", "on": colors[1]}]} return {"plan": {"type": "stack", "colors": colors or None}, "queue": []} if any(k in t for k in ("sort", "tidy", "organi", "collect", "clean", "clear", "put away")): return {"plan": {"type": "sort", "colors": colors or None}, "queue": []} if len(colors) >= 2 and on: return {"plan": None, "queue": [{"color": colors[0], "kind": "on", "on": colors[1]}]} return {"plan": None, "queue": [{"color": colors[0] if colors else None, "kind": "box"}]} def _to_policy_plan(vlm: dict) -> dict: """Map the VLM's JSON to the policy's {plan, queue} structure.""" task = (vlm.get("task") or "").lower() colors = [c for c in (vlm.get("colors") or []) if c in _COLORS] onto = vlm.get("onto") if vlm.get("onto") in _COLORS else None if task == "stack" and not (onto and colors): return {"plan": {"type": "stack", "colors": colors or None}, "queue": []} if task == "sort": return {"plan": {"type": "sort", "colors": colors or None}, "queue": []} if onto and colors: return {"plan": None, "queue": [{"color": colors[0], "kind": "on", "on": onto}]} return {"plan": None, "queue": [{"color": colors[0] if colors else None, "kind": "box"}]} class VLMPlanner: def __init__(self, model_id: str = _MODEL_ID, device: str = "cuda", hf_home: Optional[str] = None, enabled: bool = True): self._ok = False self._model = self._proc = None if not enabled: log.info("[VLMPlanner] disabled — using keyword fallback"); return # point HF at the bundled cache so it loads offline hf_home = hf_home or str(Path(__file__).resolve().parents[1] / "hf_cache") os.environ.setdefault("HF_HOME", hf_home) os.environ.setdefault("HF_HUB_OFFLINE", "1") try: import torch from transformers import (Qwen2_5_VLForConditionalGeneration, AutoProcessor, BitsAndBytesConfig) quant = None if device == "cuda" and torch.cuda.is_available(): quant = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16) self._model = Qwen2_5_VLForConditionalGeneration.from_pretrained( model_id, torch_dtype="auto", device_map="auto" if quant else None, quantization_config=quant) self._proc = AutoProcessor.from_pretrained(model_id) self._torch = torch self._ok = True log.info("[VLMPlanner] Qwen2.5-VL ready (%s, 4-bit=%s)", device, quant is not None) except Exception as e: log.warning("[VLMPlanner] load failed (%s) — keyword fallback", e) def plan(self, image: Optional[np.ndarray], instruction: str) -> dict: """Return {plan, queue} for the instruction; VLM if available else keyword.""" if not self._ok or image is None: return _keyword_plan(instruction) try: from PIL import Image pil = Image.fromarray(np.asarray(image).astype(np.uint8)) messages = [{"role": "user", "content": [ {"type": "image"}, {"type": "text", "text": _SYS + f"\n\nInstruction: {instruction}\nJSON:"}]}] text = self._proc.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) inputs = self._proc(text=[text], images=[pil], return_tensors="pt").to(self._model.device) with self._torch.no_grad(): out = self._model.generate(**inputs, max_new_tokens=128, do_sample=False) dec = self._proc.batch_decode(out[:, inputs.input_ids.shape[1]:], skip_special_tokens=True)[0] m = re.search(r"\{.*\}", dec, re.S) vlm = json.loads(m.group(0)) if m else {} log.info("[VLMPlanner] '%s' -> %s", instruction, vlm) return _to_policy_plan(vlm) except Exception as e: log.warning("[VLMPlanner] inference failed (%s) — keyword fallback", e) return _keyword_plan(instruction)