train_floorplan_grpo.py

"""
FloorplanVLM GRPO Training (Stage 2) - Run after SFT

Loads the SFT model and applies geometric reward-based RL.
Reward: R = 0.1·R_val + 0.5·R_ext + α·0.4·R_int (FloorplanVLM Eq. 9)

Usage:
    pip install shapely  # + same deps as SFT script
    python train_floorplan_grpo.py

Requires: SFT model already pushed to HUB (default: manitocross/floorplan-vlm-sft)
"""

import os, json, re, math, torch, numpy as np
from PIL import Image, ImageDraw
from datasets import Dataset
from transformers import AutoProcessor, TrainerCallback
from trl import GRPOTrainer, GRPOConfig
from peft import LoraConfig
from shapely.geometry import Polygon, LineString
from shapely.ops import unary_union

# ══════════════════════════════════════════════════════════════════════════════
# CONFIGURATION
# ══════════════════════════════════════════════════════════════════════════════
SFT_MODEL_ID = "manitocross/floorplan-vlm-sft"   # your SFT model from Stage 1
HUB_MODEL_ID = "manitocross/floorplan-vlm-grpo"
OUTPUT_DIR = "./floorplan-vlm-grpo"

# Uses same CubiCasa5K data directory as SFT script
DATA_DIR = "./cubicasa_data"

NUM_EPOCHS = 1
BATCH_SIZE = 1
GRAD_ACCUM = 4
LEARNING_RATE = 1e-6
NUM_GENERATIONS = 4     # G: completions per prompt
MAX_COMPLETION_LENGTH = 4096
KL_COEF = 0.01
PUSH_TO_HUB = True
# ══════════════════════════════════════════════════════════════════════════════

SYSTEM_PROMPT = "You are a floor plan vectorization expert. Extract wall, door, window geometry from floor plan images into structured JSON. Output ONLY valid JSON."
USER_PROMPT = "Vectorize this floor plan into structured JSON with all walls, doors, windows, and rooms."


# ── Geometric Helpers ────────────────────────────────────────────────────────

def extract_json(text):
    if isinstance(text, list):
        text = text[0].get("content", "") if text else ""
    text = text.strip()
    try: return json.loads(text)
    except: pass
    m = re.search(r'\{[\s\S]*\}', text)
    if m:
        try: return json.loads(m.group())
        except: pass
    return None


def walls_to_polygon(walls):
    if not walls or len(walls) < 3: return None
    try:
        polys = []
        for w in walls:
            s, e = w.get('start',[0,0]), w.get('end',[0,0])
            t = max(w.get('thickness',10), 1)
            polys.append(LineString([s, e]).buffer(t/2, cap_style=2))
        combined = unary_union(polys)
        return combined.convex_hull if not combined.is_empty else None
    except: return None


def poly_iou(p1, p2):
    if p1 is None or p2 is None: return 0.0
    try:
        if not p1.is_valid: p1 = p1.buffer(0)
        if not p2.is_valid: p2 = p2.buffer(0)
        inter = p1.intersection(p2).area
        union = p1.union(p2).area
        return inter / union if union > 0 else 0.0
    except: return 0.0


# ── Reward Function ──────────────────────────────────────────────────────────

def floorplan_reward(completions, **kwargs):
    """Combined: 0.1·R_val + 0.5·R_ext + α·0.4·R_int"""
    gt_jsons = kwargs.get("json_gt", [])
    rewards = []

    for c, gt_str in zip(completions, gt_jsons):
        text = c[0]["content"] if isinstance(c, list) else c
        pred = extract_json(text)
        if pred is None:
            rewards.append(0.0); continue
        try:
            gt = json.loads(gt_str) if isinstance(gt_str, str) else gt_str
        except:
            rewards.append(0.0); continue

        # R_val
        r_val = 0.0
        has_walls = "walls" in pred and isinstance(pred["walls"], list) and len(pred["walls"]) > 0
        if has_walls:
            valid = sum(1 for w in pred["walls"]
                        if all(k in w for k in ["id","start","end","thickness"])
                        and isinstance(w.get("start"), list) and len(w.get("start",[])) == 2)
            r_val = 0.3 + 0.5 * (valid / max(len(pred["walls"]), 1))
            if "rooms" in pred and isinstance(pred["rooms"], list):
                wids = {w.get("id") for w in pred["walls"]}
                vr = sum(1 for r in pred["rooms"]
                         if "label" in r and "walls" in r
                         and all(wid in wids for wid in r.get("walls",[])))
                r_val += 0.2 * (vr / max(len(pred["rooms"]), 1))

        # R_ext
        pred_poly = walls_to_polygon(pred.get("walls", []))
        gt_poly = walls_to_polygon(gt.get("walls", []))
        r_ext = poly_iou(pred_poly, gt_poly)

        # Alpha gating (FloorplanVLM Eq. 8)
        if r_ext < 0.3: alpha = 0.1
        elif r_ext < 0.7: alpha = 0.1 + 0.9 * (r_ext - 0.3) / 0.4
        else: alpha = 1.0

        # R_int (room label overlap)
        r_int = 0.0
        pred_rooms = pred.get("rooms", [])
        gt_rooms = gt.get("rooms", [])
        if pred_rooms and gt_rooms:
            pl = set(r.get("label","") for r in pred_rooms)
            gl = set(r.get("label","") for r in gt_rooms)
            overlap = len(pl & gl)
            total = len(pl | gl)
            r_int = overlap / total if total > 0 else 0.0

        total = 0.1 * min(r_val, 1.0) + 0.5 * r_ext + alpha * 0.4 * r_int
        rewards.append(float(total))

    return rewards


# ── Dataset (reuses SFT data dir) ───────────────────────────────────────────
# You'll need to adapt this to load your data (same CubiCasa5K directory).
# For GRPO format: prompt-only (no completion), plus metadata for reward.

def build_grpo_dataset(data_dir, max_samples=100):
    """Build GRPO dataset — loads pre-converted JSON annotations."""
    # Look for the annotation file from SFT stage
    ann_path = os.path.join(data_dir, "annotations.json")
    if not os.path.exists(ann_path):
        print(f"  No pre-converted annotations at {ann_path}")
        print(f"  Run train_floorplan_vlm.py (SFT) first to generate data.")
        print(f"  Creating synthetic GRPO dataset as fallback...")
        return create_synthetic_grpo(max_samples or 20)

    with open(ann_path) as f:
        annotations = json.load(f)

    if max_samples:
        annotations = annotations[:max_samples]

    records = []
    for ann in annotations:
        img_path = ann.get("image_path")
        if not img_path or not os.path.exists(img_path):
            continue
        img = Image.open(img_path).convert("RGB")
        records.append({
            "prompt": [
                {"role":"system","content":[{"type":"text","text":SYSTEM_PROMPT}]},
                {"role":"user","content":[{"type":"image"},{"type":"text","text":USER_PROMPT}]},
            ],
            "images": [img],
            "json_gt": ann["json_annotation"],
        })

    print(f"  GRPO dataset: {len(records)} samples")
    return Dataset.from_list(records)


def create_synthetic_grpo(n=20):
    records = []
    for i in range(n):
        size = 256
        img = Image.new('RGB', (size,size), 'white')
        d = ImageDraw.Draw(img)
        m = 30+i*3; wt=6; s=1024.0/size; mid=size//2+i*2
        d.rectangle([m,m,size-m,size-m], outline='black', width=wt)
        d.line([(m,mid),(size-m,mid)], fill='black', width=wt)

        jd = {"walls":[
            {"id":"wall_1","start":[round(m*s),round(m*s)],"end":[round((size-m)*s),round(m*s)],"thickness":round(wt*s),"curvature":0,"openings":[]},
            {"id":"wall_2","start":[round((size-m)*s),round(m*s)],"end":[round((size-m)*s),round((size-m)*s)],"thickness":round(wt*s),"curvature":0,"openings":[]},
            {"id":"wall_3","start":[round((size-m)*s),round((size-m)*s)],"end":[round(m*s),round((size-m)*s)],"thickness":round(wt*s),"curvature":0,"openings":[]},
            {"id":"wall_4","start":[round(m*s),round((size-m)*s)],"end":[round(m*s),round(m*s)],"thickness":round(wt*s),"curvature":0,"openings":[]},
            {"id":"wall_5","start":[round(m*s),round(mid*s)],"end":[round((size-m)*s),round(mid*s)],"thickness":round(wt*s),"curvature":0,"openings":[]},
        ],"rooms":[
            {"label":"bedroom","walls":["wall_1","wall_2","wall_5","wall_4"]},
            {"label":"living_room","walls":["wall_5","wall_2","wall_3","wall_4"]},
        ]}
        records.append({
            "prompt":[
                {"role":"system","content":[{"type":"text","text":SYSTEM_PROMPT}]},
                {"role":"user","content":[{"type":"image"},{"type":"text","text":USER_PROMPT}]},
            ],
            "images":[img],
            "json_gt": json.dumps(jd, separators=(',',':')),
        })
    return Dataset.from_list(records)


# ── Main ─────────────────────────────────────────────────────────────────────
def main():
    use_gpu = torch.cuda.is_available()
    print("="*64)
    print(f"  FloorplanVLM GRPO Training  ({'GPU' if use_gpu else 'CPU'})")
    print(f"  SFT Model : {SFT_MODEL_ID}")
    print(f"  Output    : {HUB_MODEL_ID}")
    print(f"  Generations: {NUM_GENERATIONS}, KL: {KL_COEF}")
    print("="*64)

    dataset = build_grpo_dataset(DATA_DIR)

    peft_config = LoraConfig(
        r=16, lora_alpha=32,
        target_modules=["q_proj","k_proj","v_proj","o_proj","gate_proj","up_proj","down_proj"],
        lora_dropout=0.05, bias="none", task_type="CAUSAL_LM",
    )

    grpo_config = GRPOConfig(
        output_dir=OUTPUT_DIR,
        num_train_epochs=NUM_EPOCHS,
        per_device_train_batch_size=BATCH_SIZE,
        gradient_accumulation_steps=GRAD_ACCUM,
        learning_rate=LEARNING_RATE,
        bf16=use_gpu, fp16=False,
        gradient_checkpointing=True,
        logging_steps=5, logging_first_step=True,
        logging_strategy="steps", disable_tqdm=True,
        save_steps=200, save_total_limit=2,
        num_generations=NUM_GENERATIONS,
        max_prompt_length=512,
        max_completion_length=MAX_COMPLETION_LENGTH,
        scale_rewards=True,
        beta=KL_COEF,
        temperature=0.7,
        push_to_hub=PUSH_TO_HUB,
        hub_model_id=HUB_MODEL_ID if PUSH_TO_HUB else None,
        report_to="none",
    )

    trainer = GRPOTrainer(
        model=SFT_MODEL_ID,
        reward_funcs=[floorplan_reward],
        args=grpo_config,
        train_dataset=dataset,
        peft_config=peft_config,
    )

    print("\nStarting GRPO training...")
    trainer.train()

    trainer.save_model(OUTPUT_DIR)
    if PUSH_TO_HUB:
        try:
            trainer.push_to_hub()
            print(f"\n✅ Model: https://huggingface.co/{HUB_MODEL_ID}")
        except Exception as e:
            print(f"Push failed: {e}")
            from huggingface_hub import HfApi
            HfApi().create_repo(HUB_MODEL_ID, exist_ok=True)
            HfApi().upload_folder(folder_path=OUTPUT_DIR, repo_id=HUB_MODEL_ID)

    print(f"\n✅ Done! Local: {OUTPUT_DIR}/")


if __name__ == "__main__":
    main()