#!/usr/bin/env python3 """CoT + BeliefToken fast-SFT on Qwen3-VL-4B-Instruct (skip pretrain, domain-adapt only). Pipeline: 1. Load Qwen3-VL-4B processor + model (bf16). 2. Add 4 special tokens: <|BELIEF|>, , <|ALERT|>, <|OBSERVE|>, <|SILENT|> (resizes embedding + lm_head; new rows initialized with mean of existing). 3. Freeze vision tower, LoRA on LLM (q/k/v/o/gate/up/down_proj). 4. Train on data/vla_cot_belief/train500_belief.jsonl — assistant target = scene + threat + <|BELIEF|> <|ACTION|> . At belief-extraction time (separate script), we teacher-force the prefix + CoT up through "<|BELIEF|>" and read hidden_states[-1] at that position. Run: python -m training.VLA.train_cot_belief \ --cot_jsonl data/vla_cot_belief/train500_belief.jsonl \ --video_dir nexar-collision-prediction/train \ --out_dir checkpoints/VLA/qwen3vl4b_cot_belief \ --epochs 5 --batch_size 1 --grad_accum 4 --lr 2e-4 """ from __future__ import annotations import argparse import json import math import sys from functools import partial from pathlib import Path import torch from torch.optim import AdamW from torch.utils.data import DataLoader from tqdm import tqdm sys.path.insert(0, str(Path(__file__).resolve().parents[2])) from peft import LoraConfig, get_peft_model from transformers import AutoProcessor, AutoModelForImageTextToText from transformers.optimization import get_cosine_schedule_with_warmup from training.VLA.cot_belief_dataset import ( CoTBeliefDataset, collate_fn, ALL_SPECIAL, ) def parse_args(): ap = argparse.ArgumentParser() ap.add_argument("--model_name", default="PROJECT_ROOT/models/Qwen3-VL-4B-Instruct") ap.add_argument("--cot_jsonl", required=True) ap.add_argument("--video_dir", required=True) ap.add_argument("--out_dir", required=True) ap.add_argument("--lora_r", type=int, default=32) ap.add_argument("--lora_alpha", type=int, default=16) ap.add_argument("--lora_dropout", type=float, default=0.05) ap.add_argument("--lr", type=float, default=2e-4) ap.add_argument("--epochs", type=int, default=5) ap.add_argument("--batch_size", type=int, default=1) ap.add_argument("--grad_accum", type=int, default=4) ap.add_argument("--warmup_ratio", type=float, default=0.05) ap.add_argument("--n_frames", type=int, default=8) ap.add_argument("--resize_short", type=int, default=336) ap.add_argument("--max_len", type=int, default=3072) ap.add_argument("--max_samples", type=int, default=0, help="If >0, truncate dataset for smoke-test") ap.add_argument("--log_every", type=int, default=10) ap.add_argument("--save_every_epoch", action="store_true") ap.add_argument("--seed", type=int, default=0) ap.add_argument("--resume", type=str, default="", help="Path to existing PEFT adapter dir to warm-start from") ap.add_argument("--per_frame", action="store_true", help="Per-frame POMDP target (requires belief.actions_per_frame)") ap.add_argument("--state_conditional", action="store_true", help="VLAlert-X Stage A: emit state-specific phrases inside " "<|BELIEF|> blocks (forces state-distinguishing belief)") return ap.parse_args() def add_special_tokens_and_resize(processor, model): """Add the 4 belief/action special tokens; resize embeddings; init new rows.""" tok = processor.tokenizer before = len(tok) added = tok.add_special_tokens({"additional_special_tokens": ALL_SPECIAL}) after = len(tok) print(f"[tokens] vocab {before} -> {after} ({added} new)") if added == 0: return # already present model.resize_token_embeddings(after) emb = model.get_input_embeddings() with torch.no_grad(): mean_vec = emb.weight[:before].mean(dim=0) for tok_str in ALL_SPECIAL: tid = tok.convert_tokens_to_ids(tok_str) emb.weight[tid] = mean_vec + 0.01 * torch.randn_like(mean_vec) # Qwen3-VL ties input/output embeddings; get_output_embeddings may still # return a separate Linear — handle both cases. out_emb = model.get_output_embeddings() if out_emb is not None and out_emb.weight.data_ptr() != emb.weight.data_ptr(): with torch.no_grad(): mean_out = out_emb.weight[:before].mean(dim=0) for tok_str in ALL_SPECIAL: tid = tok.convert_tokens_to_ids(tok_str) out_emb.weight[tid] = mean_out + 0.01 * torch.randn_like(mean_out) def main(): args = parse_args() torch.manual_seed(args.seed) out_dir = Path(args.out_dir) out_dir.mkdir(parents=True, exist_ok=True) print(f"[train] loading processor/model from {args.model_name}") processor = AutoProcessor.from_pretrained(args.model_name, trust_remote_code=True) model = AutoModelForImageTextToText.from_pretrained( args.model_name, torch_dtype=torch.bfloat16, trust_remote_code=True, attn_implementation="sdpa", ) # 1) Inject special tokens + resize. add_special_tokens_and_resize(processor, model) # 2) Freeze vision tower. for attr in ("visual", "vision_tower"): if hasattr(model, attr): for p in getattr(model, attr).parameters(): p.requires_grad = False # 3) Enable input_require_grads + gradient checkpointing. if hasattr(model, "enable_input_require_grads"): model.enable_input_require_grads() model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) # 4) LoRA on LLM — warm-start from an existing adapter if --resume is set. if args.resume: from peft import PeftModel print(f"[resume] loading PEFT adapter from {args.resume}") model = PeftModel.from_pretrained(model, args.resume, is_trainable=True) else: lora_cfg = LoraConfig( r=args.lora_r, lora_alpha=args.lora_alpha, lora_dropout=args.lora_dropout, target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], bias="none", task_type="CAUSAL_LM", # Train the new embedding rows + lm_head rows too (cheap, ~5 rows). modules_to_save=["embed_tokens", "lm_head"], ) model = get_peft_model(model, lora_cfg) model.print_trainable_parameters() model.to("cuda") model.config.use_cache = False # 5) Dataset. ds = CoTBeliefDataset( jsonl_path=args.cot_jsonl, video_dir=args.video_dir, processor=processor, n_frames=args.n_frames, resize_short=args.resize_short, max_len=args.max_len, per_frame=args.per_frame, state_conditional=args.state_conditional, ) if args.state_conditional: print("[stage-A] state_conditional=True — <|BELIEF|> blocks " "will contain state-specific phrases.") print(f"[train] dataset size = {len(ds)}") if args.max_samples > 0 and len(ds) > args.max_samples: from torch.utils.data import Subset ds = Subset(ds, list(range(args.max_samples))) print(f"[smoke] truncated to {len(ds)}") if len(ds) == 0: raise SystemExit("empty dataset — check --cot_jsonl") pad_id = processor.tokenizer.pad_token_id or processor.tokenizer.eos_token_id dl = DataLoader( ds, batch_size=args.batch_size, shuffle=True, num_workers=0, collate_fn=partial(collate_fn, pad_token_id=pad_id), pin_memory=True, ) # 6) Optim. trainable = [p for p in model.parameters() if p.requires_grad] opt = AdamW(trainable, lr=args.lr, betas=(0.9, 0.95), weight_decay=0.0) total_steps = math.ceil(len(dl) * args.epochs / args.grad_accum) warmup_steps = max(1, int(total_steps * args.warmup_ratio)) sched = get_cosine_schedule_with_warmup(opt, warmup_steps, total_steps) print(f"[train] total_updates={total_steps} warmup={warmup_steps} lr={args.lr}") # 7) Loop. global_step = 0 model.train() for epoch in range(args.epochs): pbar = tqdm(enumerate(dl), total=len(dl), desc=f"ep{epoch}", ncols=80, leave=True) running = 0.0 running_n = 0 for step, batch in pbar: input_ids = batch["input_ids"].to("cuda", non_blocking=True) attn = batch["attention_mask"].to("cuda", non_blocking=True) labels = batch["labels"].to("cuda", non_blocking=True) pix = batch["pixel_values"].to("cuda", dtype=torch.bfloat16, non_blocking=True) grid = batch["image_grid_thw"].to("cuda", non_blocking=True) out = model( input_ids=input_ids, attention_mask=attn, labels=labels, pixel_values=pix, image_grid_thw=grid, ) loss = out.loss / args.grad_accum loss.backward() running += out.loss.detach().float().item() running_n += 1 if (step + 1) % args.grad_accum == 0 or (step + 1) == len(dl): torch.nn.utils.clip_grad_norm_(trainable, 1.0) opt.step() sched.step() opt.zero_grad(set_to_none=True) global_step += 1 if global_step % args.log_every == 0: pbar.set_postfix(loss=running / max(1, running_n), lr=sched.get_last_lr()[0]) running, running_n = 0.0, 0 if args.save_every_epoch or epoch == args.epochs - 1: ep_dir = out_dir / f"epoch_{epoch}" ep_dir.mkdir(parents=True, exist_ok=True) model.save_pretrained(ep_dir) processor.save_pretrained(ep_dir) with (ep_dir / "train_args.json").open("w") as f: json.dump(vars(args), f, indent=2) print(f"[train] saved -> {ep_dir}") # Final "best" final = out_dir / "best" final.mkdir(parents=True, exist_ok=True) model.save_pretrained(final) processor.save_pretrained(final) with (final / "train_args.json").open("w") as f: json.dump(vars(args), f, indent=2) with (final / "belief_tokens.json").open("w") as f: json.dump({"special_tokens": ALL_SPECIAL, "belief_open": "<|BELIEF|>", "belief_close": "", "actions": ["<|ALERT|>", "<|OBSERVE|>", "<|SILENT|>"]}, f, indent=2) print(f"[train] done. final -> {final}") if __name__ == "__main__": main()