training/VLA/train_cot_belief.py

#!/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|>, </|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|> </|BELIEF|>.

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": "</|BELIEF|>",
                   "actions": ["<|ALERT|>", "<|OBSERVE|>", "<|SILENT|>"]}, f, indent=2)
    print(f"[train] done. final -> {final}")


if __name__ == "__main__":
    main()