scripts/train_hf_space.py

"""Training script for HuggingFace Spaces (GPU runtime).

Downloads dataset from HF Hub, trains SmolVLA, pushes checkpoint back to Hub.

Usage on HF Spaces (L4 GPU):
    python train_hf_space.py \
        --dataset-repo venaychawda/elrobot-pickplace \
        --output-repo venaychawda/elrobot-smolvla-pickplace \
        --steps 5000 --batch-size 32
"""

from __future__ import annotations

import argparse
import math
import time
from functools import partial
from pathlib import Path

import torch
from huggingface_hub import HfApi, snapshot_download, create_repo
from torch.utils.data import DataLoader

from smolvla import SmolVLAPolicy
from smolvla.dataset import PickAndPlaceDataset, collate_samples
from smolvla.normalize import normalize_action, normalize_state
from smolvla.stats import compute_stats, save_stats


IMAGE_KEYS = ("observation.images.cam0",)
STATE_DIM = 8
ACTION_DIM = 8


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser()
    p.add_argument("--dataset-repo", required=True, help="HF dataset repo (e.g. venaychawda/elrobot-pickplace)")
    p.add_argument("--output-repo", required=True, help="HF model repo to push checkpoint (e.g. venaychawda/elrobot-smolvla-pickplace)")
    p.add_argument("--base-checkpoint", type=str, default="LBST/t01_pick_and_place")
    p.add_argument("--steps", type=int, default=5000)
    p.add_argument("--batch-size", type=int, default=32)
    p.add_argument("--lr", type=float, default=1e-4)
    p.add_argument("--weight-decay", type=float, default=1e-10)
    p.add_argument("--warmup-steps", type=int, default=500)
    p.add_argument("--decay-steps", type=int, default=10000)
    p.add_argument("--decay-lr", type=float, default=2.5e-6)
    p.add_argument("--grad-clip", type=float, default=10.0)
    p.add_argument("--save-every", type=int, default=500)
    p.add_argument("--log-every", type=int, default=20)
    p.add_argument("--num-workers", type=int, default=2)
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--private", action="store_true")
    return p.parse_args()


def make_loader(ds, batch_size: int, shuffle: bool, num_workers: int) -> DataLoader:
    return DataLoader(
        ds,
        batch_size=batch_size,
        shuffle=shuffle,
        num_workers=num_workers,
        collate_fn=partial(collate_samples, image_keys=IMAGE_KEYS),
        pin_memory=True,
        drop_last=shuffle,
        persistent_workers=num_workers > 0,
    )


def move_batch(batch: dict, device: torch.device) -> dict:
    return {
        k: (v.to(device, non_blocking=True) if isinstance(v, torch.Tensor) else v)
        for k, v in batch.items()
    }


def build_train_batch(batch: dict, policy, stats, device) -> dict:
    batch = move_batch(batch, device)
    batch["observation.state"] = normalize_state(batch["observation.state"], stats)
    batch["action"] = normalize_action(batch["action"], stats)
    tokens, mask = policy.tokenize_task(batch.pop("task"), device=device)
    batch["observation.language.tokens"] = tokens
    batch["observation.language.attention_mask"] = mask
    return batch


def lr_lambda(step, nominal_warmup, nominal_decay, total_steps, decay_lr, peak_lr):
    if total_steps < nominal_decay:
        scale = total_steps / nominal_decay
        warmup = max(int(nominal_warmup * scale), 1)
        decay = total_steps
    else:
        warmup = max(nominal_warmup, 1)
        decay = nominal_decay
    if step < warmup:
        return (step + 1) / warmup
    s = min(step - warmup, decay - warmup)
    cos = 0.5 * (1 + math.cos(math.pi * s / max(decay - warmup, 1)))
    alpha = decay_lr / peak_lr
    return (1 - alpha) * cos + alpha


def main() -> None:
    args = parse_args()
    torch.manual_seed(args.seed)

    if not torch.cuda.is_available():
        raise SystemExit("CUDA required. Use a GPU-enabled HF Space (L4/A10G/T4).")
    device = torch.device("cuda")
    print(f"GPU: {torch.cuda.get_device_name()}")

    # Download dataset from HF Hub
    print(f"\nDownloading dataset from {args.dataset_repo}...")
    dataset_path = Path(snapshot_download(args.dataset_repo, repo_type="dataset"))
    data_dir = dataset_path / "data"
    parquets = sorted(data_dir.glob("*.parquet"))
    if not parquets:
        raise SystemExit(f"No parquet files in {data_dir}")
    print(f"Found {len(parquets)} parquet files")

    # Load dataset
    train_ds = PickAndPlaceDataset(parquets, image_keys=IMAGE_KEYS)
    print(f"Episodes: {train_ds.total_episodes}  Frames: {len(train_ds)}")

    # Compute stats
    output_dir = Path("checkpoints/elrobot-run")
    output_dir.mkdir(parents=True, exist_ok=True)
    stats = compute_stats(train_ds._state, train_ds._action)
    stats_path = output_dir / "stats.safetensors"
    save_stats(stats, stats_path)
    print(f"\nStats saved -> {stats_path}")
    for k, v in stats.items():
        print(f"  {k:12s}  {v.tolist()}")
    stats = {k: v.to(device) for k, v in stats.items()}

    # Load base model
    print(f"\nLoading base checkpoint {args.base_checkpoint}...")
    policy = SmolVLAPolicy.from_pretrained(
        args.base_checkpoint,
        config_overrides={
            "image_keys": list(IMAGE_KEYS),
            "state_dim": STATE_DIM,
            "action_dim": ACTION_DIM,
            "load_vlm_weights": False,
            "empty_cameras": 0,
        },
        strict=False,
    ).to(device)
    policy.train()

    trainable = [p for p in policy.parameters() if p.requires_grad]
    n_train = sum(p.numel() for p in trainable)
    n_total = sum(p.numel() for p in policy.parameters())
    print(f"Trainable params: {n_train:,} / {n_total:,} ({100*n_train/n_total:.1f}%)")

    opt = torch.optim.AdamW(trainable, lr=args.lr, betas=(0.9, 0.95), weight_decay=args.weight_decay)
    sched = torch.optim.lr_scheduler.LambdaLR(
        opt, lr_lambda=partial(
            lr_lambda,
            nominal_warmup=args.warmup_steps,
            nominal_decay=args.decay_steps,
            total_steps=args.steps,
            decay_lr=args.decay_lr,
            peak_lr=args.lr,
        )
    )

    train_loader = make_loader(train_ds, args.batch_size, shuffle=True, num_workers=args.num_workers)
    print(f"Batches/epoch: {len(train_loader)}")
    print(f"\n{'='*60}")
    print(f"Starting training: {args.steps} steps, batch_size={args.batch_size}")
    print(f"{'='*60}\n")

    step = 0
    t0 = time.time()
    running_loss = 0.0
    running_n = 0

    while step < args.steps:
        for raw_batch in train_loader:
            if step >= args.steps:
                break
            batch = build_train_batch(raw_batch, policy, stats, device)
            loss, info = policy.forward(batch)

            opt.zero_grad(set_to_none=True)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(trainable, args.grad_clip)
            opt.step()
            sched.step()

            running_loss += info["loss"]
            running_n += 1

            if (step + 1) % args.log_every == 0:
                dt = time.time() - t0
                avg = running_loss / running_n
                lr = sched.get_last_lr()[0]
                print(f"step {step+1:>6}/{args.steps}  loss {avg:.4f}  lr {lr:.2e}  "
                      f"({(step+1) / max(dt, 1e-9):.2f} step/s)")
                running_loss, running_n = 0.0, 0

            if (step + 1) % args.save_every == 0:
                out = output_dir / f"step-{step+1:06d}"
                policy.save_pretrained(out)
                (out / "stats.safetensors").write_bytes(stats_path.read_bytes())
                print(f"  [save] {out}")

            step += 1

    # Save final checkpoint
    final = output_dir / "final"
    policy.save_pretrained(final)
    (final / "stats.safetensors").write_bytes(stats_path.read_bytes())
    print(f"\nTraining complete. Final checkpoint: {final}")

    # Push to HF Hub
    print(f"\nPushing checkpoint to {args.output_repo}...")
    create_repo(args.output_repo, exist_ok=True, private=args.private)
    api = HfApi()
    api.upload_folder(
        repo_id=args.output_repo,
        folder_path=str(final),
        path_in_repo=".",
    )
    print(f"Checkpoint pushed: https://huggingface.co/{args.output_repo}")
    print(f"\nTotal training time: {(time.time()-t0)/60:.1f} min")


if __name__ == "__main__":
    main()