Add training script

train_drone_vla.py

@@ -0,0 +1,189 @@
+#!/usr/bin/env python3
+"""
+🚁 Drone VLA Training: Qwen3.5-0.8B + LoRA + Action Head → [Vx, Vy, Vz, Yaw_rate]
+
+Usage:
+    python train_drone_vla.py --dataset_repo YOUR_USER/drone-nav-data --hub_model_id YOUR_USER/drone-vla --num_epochs 10
+"""
+import os, sys, json, math, argparse, numpy as np
+from pathlib import Path
+import torch, torch.nn as nn, torch.nn.functional as F
+from torch.utils.data import DataLoader, Dataset as TorchDataset
+from transformers import AutoProcessor, AutoModelForImageTextToText
+from peft import LoraConfig, get_peft_model, TaskType
+from datasets import load_from_disk, load_dataset
+from PIL import Image
+try:
+    import trackio; HAS_TRACKIO = True
+except ImportError:
+    HAS_TRACKIO = False
+
+class ActionHead(nn.Module):
+    def __init__(self, hidden_size=1024, action_dim=4, hidden_dim=256):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(hidden_size, hidden_dim), nn.GELU(), nn.Dropout(0.1),
+            nn.Linear(hidden_dim, hidden_dim), nn.GELU(), nn.Dropout(0.1),
+            nn.Linear(hidden_dim, action_dim), nn.Tanh())
+    def forward(self, x): return self.net(x)
+
+class DroneVLA(nn.Module):
+    def __init__(self, model_name="Qwen/Qwen3.5-0.8B", action_dim=4, lora_r=16, lora_alpha=32, freeze_vision=True):
+        super().__init__()
+        print(f"Loading {model_name}...")
+        self.backbone = AutoModelForImageTextToText.from_pretrained(model_name, torch_dtype=torch.bfloat16, trust_remote_code=True)
+        hs = getattr(self.backbone.config, 'hidden_size', None) or getattr(self.backbone.config, 'text_config', type('',(),{'hidden_size':1024})).hidden_size
+        print(f"Hidden size: {hs}")
+        exclude = ["lm_head"]
+        if freeze_vision:
+            for name, _ in self.backbone.named_modules():
+                if any(v in name.lower() for v in ['visual','vision','vit','patch_embed']) and name.count('.') <= 1:
+                    exclude.append(name)
+            exclude = list(set(exclude)) or ["visual", "lm_head"]
+        lora_config = LoraConfig(r=lora_r, lora_alpha=lora_alpha, lora_dropout=0.05, bias="none",
+            task_type=TaskType.CAUSAL_LM, target_modules="all-linear", exclude_modules=exclude)
+        self.backbone = get_peft_model(self.backbone, lora_config)
+        self.backbone.print_trainable_parameters()
+        self.action_head = ActionHead(hidden_size=hs, action_dim=action_dim)
+        self.hidden_size = hs
+
+    def forward(self, input_ids, attention_mask, pixel_values=None, image_grid_thw=None, action_targets=None, **kw):
+        bkw = {'input_ids': input_ids, 'attention_mask': attention_mask, 'output_hidden_states': True, 'return_dict': True}
+        if pixel_values is not None: bkw['pixel_values'] = pixel_values
+        if image_grid_thw is not None: bkw['image_grid_thw'] = image_grid_thw
+        out = self.backbone(**bkw)
+        h = out.hidden_states[-1]
+        if attention_mask is not None:
+            sl = attention_mask.sum(dim=1) - 1
+            feat = h[torch.arange(h.shape[0], device=h.device), sl]
+        else:
+            feat = h[:, -1, :]
+        pred = self.action_head(feat.float())
+        if action_targets is not None:
+            return F.l1_loss(pred, action_targets.float()), pred
+        return pred
+
+    def save_pretrained(self, save_dir):
+        os.makedirs(save_dir, exist_ok=True)
+        self.backbone.save_pretrained(save_dir)
+        torch.save(self.action_head.state_dict(), os.path.join(save_dir, "action_head.pt"))
+        with open(os.path.join(save_dir, "vla_config.json"), 'w') as f:
+            json.dump({'hidden_size': self.hidden_size, 'action_dim': 4, 'action_labels': ['Vx','Vy','Vz','Yaw_rate']}, f)
+
+class VLADataset(TorchDataset):
+    def __init__(self, hf_dataset, processor, max_length=512):
+        self.dataset, self.processor, self.max_length = hf_dataset, processor, max_length
+    def __len__(self): return len(self.dataset)
+    def __getitem__(self, idx):
+        s = self.dataset[idx]
+        img = s['image'] if isinstance(s['image'], Image.Image) else Image.fromarray(np.array(s['image']))
+        img = img.convert('RGB')
+        action = torch.tensor(s['action'], dtype=torch.float32)
+        inst = s.get('language_instruction','navigate forward') or 'navigate forward'
+        msgs = [{"role":"user","content":[{"type":"image","image":img},{"type":"text","text":inst}]}]
+        try:
+            text = self.processor.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
+            inputs = self.processor(text=[text], images=[img], padding='max_length', max_length=self.max_length, truncation=True, return_tensors="pt")
+        except:
+            inputs = self.processor(text=[inst], images=[img], padding='max_length', max_length=self.max_length, truncation=True, return_tensors="pt")
+        result = {k: v.squeeze(0) for k, v in inputs.items()}
+        result['action_targets'] = action
+        return result
+
+def collate_fn(batch):
+    keys = batch[0].keys()
+    c = {}
+    for key in keys:
+        vals = [b[key] for b in batch]
+        if isinstance(vals[0], torch.Tensor):
+            if vals[0].dim() >= 1 and key in ['input_ids','attention_mask']:
+                ml = max(v.shape[0] for v in vals)
+                c[key] = torch.stack([F.pad(v,(0,ml-v.shape[0]),value=0) for v in vals])
+            else:
+                try: c[key] = torch.stack(vals)
+                except: c[key] = vals
+        else: c[key] = vals
+    return c
+
+def train(args):
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print(f"Device: {device}")
+    ds = load_dataset(args.dataset_repo, split='train') if args.dataset_repo else load_from_disk(args.dataset_path)
+    print(f"Dataset: {len(ds)} rows")
+    processor = AutoProcessor.from_pretrained(args.model_name, trust_remote_code=True, min_pixels=224*224, max_pixels=512*512)
+    if processor.tokenizer.pad_token is None: processor.tokenizer.pad_token = processor.tokenizer.eos_token
+    model = DroneVLA(args.model_name, 4, args.lora_r, args.lora_alpha, args.freeze_vision).to(device)
+    loader = DataLoader(VLADataset(ds, processor, args.max_length), batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, collate_fn=collate_fn, pin_memory=device.type=='cuda', drop_last=True)
+    lora_p = [p for n,p in model.backbone.named_parameters() if p.requires_grad]
+    head_p = list(model.action_head.parameters())
+    opt = torch.optim.AdamW([{'params':lora_p,'lr':args.lr_lora,'weight_decay':0.01},{'params':head_p,'lr':args.lr_head}], betas=(0.9,0.95))
+    total = len(loader)*args.num_epochs; warmup = min(100, total//10)
+    sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: s/max(warmup,1) if s<warmup else 0.1+0.9*(1+math.cos(math.pi*(s-warmup)/max(total-warmup,1)))/2)
+    
+    if HAS_TRACKIO and args.hub_model_id:
+        try:
+            trackio.init(project="drone-vla", name=args.hub_model_id.split('/')[-1], space_id=f"{args.hub_model_id.split('/')[0]}/drone-vla-trackio")
+        except: pass
+    
+    print(f"\nTraining: {args.num_epochs} epochs, batch={args.batch_size}, steps={total}")
+    step, best = 0, float('inf')
+    for epoch in range(args.num_epochs):
+        model.train(); el, es = 0, 0
+        for batch in loader:
+            ids = batch['input_ids'].to(device); mask = batch['attention_mask'].to(device); tgt = batch['action_targets'].to(device)
+            pv = batch.get('pixel_values')
+            if pv is not None: pv = (torch.stack(pv) if isinstance(pv,list) else pv).to(device, dtype=torch.bfloat16)
+            igt = batch.get('image_grid_thw')
+            if igt is not None: igt = (torch.stack(igt) if isinstance(igt,list) else igt).to(device)
+            loss, pred = model(ids, mask, pv, igt, tgt)
+            loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            opt.step(); sched.step(); opt.zero_grad()
+            lv = loss.item(); el += lv; es += 1; step += 1
+            with torch.no_grad(): de = (pred-tgt).abs().mean(0)
+            if step % args.log_every == 0 or step == 1:
+                print(f"Step {step:5d} | Ep {epoch+1}/{args.num_epochs} | Loss: {lv:.4f} | Avg: {el/es:.4f} | Vx={de[0]:.4f} Vy={de[1]:.4f} Vz={de[2]:.4f} Yaw={de[3]:.4f}")
+                if HAS_TRACKIO:
+                    try: trackio.log({"train/loss":lv,"train/avg_loss":el/es,"action_error/Vx":de[0].item(),"action_error/Vy":de[1].item(),"action_error/Vz":de[2].item(),"action_error/Yaw":de[3].item()})
+                    except: pass
+            if step % args.save_every == 0:
+                model.save_pretrained(f"{args.output_dir}/ckpt-{step}")
+                if lv < best: best=lv; model.save_pretrained(f"{args.output_dir}/best"); print(f"  ⭐ best: {best:.4f}")
+        print(f"\nEpoch {epoch+1} done | Avg: {el/max(es,1):.4f}\n")
+    
+    final = f"{args.output_dir}/final"; model.save_pretrained(final)
+    if args.hub_model_id:
+        try:
+            model.backbone.push_to_hub(args.hub_model_id)
+            from huggingface_hub import HfApi; api = HfApi()
+            api.upload_file(path_or_fileobj=f"{final}/action_head.pt", path_in_repo="action_head.pt", repo_id=args.hub_model_id)
+            api.upload_file(path_or_fileobj=f"{final}/vla_config.json", path_in_repo="vla_config.json", repo_id=args.hub_model_id)
+            print(f"✅ https://huggingface.co/{args.hub_model_id}")
+        except Exception as e: print(f"Push failed: {e}")
+    if HAS_TRACKIO:
+        try: trackio.finish()
+        except: pass
+    print(f"Done! Best loss: {best:.4f}")
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument('--dataset_repo', type=str, default=None)
+    p.add_argument('--dataset_path', type=str, default=None)
+    p.add_argument('--model_name', type=str, default='Qwen/Qwen3.5-0.8B')
+    p.add_argument('--lora_r', type=int, default=16)
+    p.add_argument('--lora_alpha', type=int, default=32)
+    p.add_argument('--freeze_vision', action='store_true', default=True)
+    p.add_argument('--max_length', type=int, default=512)
+    p.add_argument('--num_epochs', type=int, default=10)
+    p.add_argument('--batch_size', type=int, default=4)
+    p.add_argument('--lr_lora', type=float, default=2e-4)
+    p.add_argument('--lr_head', type=float, default=1e-3)
+    p.add_argument('--num_workers', type=int, default=2)
+    p.add_argument('--log_every', type=int, default=10)
+    p.add_argument('--save_every', type=int, default=500)
+    p.add_argument('--output_dir', type=str, default='./checkpoints')
+    p.add_argument('--hub_model_id', type=str, default=None)
+    a = p.parse_args()
+    if not a.dataset_repo and not a.dataset_path: print("Need --dataset_repo or --dataset_path"); sys.exit(1)
+    train(a)
+
+if __name__ == '__main__': main()