scripts/train.py

#!/usr/bin/env python
"""Stage-2+ training loop: weighted multi-dataset mixture + accelerate.

Thin replacement for lerobot-train adding:
- weighted sampling across LeRobotDatasets (per-dataset embodiment ids)
- staleness augmentation (Stage 2b)
- spatial-distillation aux loss (Stage 3)

Usage:
    accelerate launch scripts/train.py --config configs/stage2_mixture.yaml
"""

from __future__ import annotations

import argparse
import math
import time
from pathlib import Path

import torch
import yaml


def make_policy(cfg: dict):
    """Canonical-schema policy: two fixed camera slots, padded state/action."""
    from lerobot.configs import FeatureType, PolicyFeature
    from tinyvla.configuration_tinyvla import TinyVLAConfig
    from tinyvla.modeling_tinyvla import TinyVLAPolicy

    pcfg = TinyVLAConfig(**cfg.get("policy", {}))
    s = pcfg.image_size
    pcfg.input_features = {
        "observation.images.cam0": PolicyFeature(type=FeatureType.VISUAL, shape=(3, s, s)),
        "observation.images.cam1": PolicyFeature(type=FeatureType.VISUAL, shape=(3, s, s)),
        "observation.state": PolicyFeature(type=FeatureType.STATE, shape=(pcfg.max_state_dim,)),
    }
    pcfg.output_features = {
        "action": PolicyFeature(type=FeatureType.ACTION, shape=(pcfg.max_action_dim,)),
    }
    pcfg.validate_features()
    return TinyVLAPolicy(pcfg), pcfg


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", type=Path, required=True)
    args = parser.parse_args()
    cfg = yaml.safe_load(args.config.read_text())

    from accelerate import Accelerator
    from lerobot.datasets.lerobot_dataset import LeRobotDataset

    accelerator = Accelerator(mixed_precision=cfg.get("mixed_precision", "bf16"))

    # ---- datasets ------------------------------------------------------
    # spec forms:
    #   {repo_id, weight, root?, episodes?, revision?}          — one dataset
    #   {root_glob, weight, embodiment_group?}                  — local converted dirs,
    #       weight is split across matches proportionally to episode count
    datasets, weights, names, embodiment_ids = [], [], [], []
    chunk = cfg["policy"]["chunk_size"]

    def add(src, w, name, emb_id):
        datasets.append(src)
        weights.append(w)
        names.append(name)
        embodiment_ids.append(emb_id)
        accelerator.print(
            f"dataset[{len(datasets)-1}] {name}: eps={src.ds.num_episodes} w={w:.4f} emb={emb_id}"
        )

    policy, pcfg = make_policy(cfg)

    from tinyvla.data.mixture import CanonicalSource, WeightedMixtureDataset

    from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata

    def make_ds(repo_id, root=None, episodes=None, revision=None):
        # delta_timestamps must be set at construction (it feeds DatasetReader),
        # and needs fps — read metadata first
        meta = LeRobotDatasetMetadata(repo_id, root=root, revision=revision)
        return LeRobotDataset(
            repo_id,
            root=root,
            episodes=episodes,
            revision=revision,
            delta_timestamps={"action": [t / meta.fps for t in range(chunk)]},
            video_backend="torchcodec",
        )

    labels_dir = cfg.get("spatial_labels_dir")

    def wrap(ds, emb_id):
        store = None
        if labels_dir:
            from tinyvla.data.spatial_labels import SpatialLabelStore

            store = SpatialLabelStore(labels_dir, ds.repo_id.split("/")[-1])
            if len(store) == 0:
                store = None
        return CanonicalSource(
            ds,
            embodiment_id=emb_id,
            image_size=pcfg.image_size,
            max_state_dim=pcfg.max_state_dim,
            max_action_dim=pcfg.max_action_dim,
            staleness_max_s=pcfg.staleness_max_s,
            # staleness_prob switched on at cfg["staleness_start_step"]
            spatial_labels=store,
        )

    next_emb = 0
    for spec in cfg["datasets"]:
        if "root_glob" in spec:
            roots = sorted(Path(p) for p in __import__("glob").glob(spec["root_glob"]))
            subs = [
                make_ds(r.name, root=r)
                for r in roots
                if (r / "meta" / "info.json").exists()
            ]
            total_eps = sum(d.num_episodes for d in subs) or 1
            for d in subs:
                add(wrap(d, next_emb), spec["weight"] * d.num_episodes / total_eps, d.root.name, next_emb)
        else:
            emb = spec.get("embodiment_id", next_emb)
            ds = make_ds(
                spec["repo_id"],
                root=spec.get("root"),
                episodes=list(range(spec["episodes"])) if spec.get("episodes") else None,
                revision=spec.get("revision"),
            )
            add(wrap(ds, emb), spec["weight"], spec["repo_id"], emb)
        next_emb += 1

    mixture = WeightedMixtureDataset(datasets, weights)
    loader = torch.utils.data.DataLoader(
        mixture,
        batch_size=cfg["batch_size"],
        num_workers=cfg.get("num_workers", 8),
        pin_memory=True,
        persistent_workers=True,
        drop_last=True,
    )

    # tokenizer for task strings (per-source normalization already done in CanonicalSource)
    from transformers import AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained(pcfg.lm_model_name)

    backbone_params = [
        p for n, p in policy.named_parameters() if p.requires_grad and "semantic.vlm" in n
    ]
    head_params = [
        p for n, p in policy.named_parameters() if p.requires_grad and "semantic.vlm" not in n
    ]
    groups = [{"params": head_params, "lr": cfg["lr"]}]
    if backbone_params:
        groups.append({"params": backbone_params, "lr": cfg["lr"] * cfg.get("backbone_lr_mult", 0.1)})
        accelerator.print(f"backbone group: {sum(p.numel() for p in backbone_params)/1e6:.1f}M params at {cfg.get('backbone_lr_mult', 0.1)}x lr")
    opt = torch.optim.AdamW(groups, betas=(0.9, 0.95), weight_decay=1e-10)
    steps = cfg["steps"]
    warmup = cfg.get("warmup_steps", 1000)

    def lr_lambda(s):
        if s < warmup:
            return s / warmup
        p = (s - warmup) / max(1, steps - warmup)
        return 0.025 + 0.975 * 0.5 * (1 + math.cos(math.pi * p))

    sched = torch.optim.lr_scheduler.LambdaLR(opt, lr_lambda)

    policy, opt, loader, sched = accelerator.prepare(policy, opt, loader, sched)

    out_dir = Path(cfg["output_dir"])
    out_dir.mkdir(parents=True, exist_ok=True)
    if cfg.get("wandb") and accelerator.is_main_process:
        import wandb

        wandb.init(project=cfg["wandb"], config=cfg)

    step, t0 = 0, time.time()
    if cfg.get("resume_from"):
        from safetensors.torch import load_file

        sd = load_file(Path(cfg["resume_from"]) / "model.safetensors")
        missing, unexpected = accelerator.unwrap_model(policy).load_state_dict(sd, strict=False)
        step = int(cfg.get("resume_step", 0))
        for _ in range(step):
            sched.step()  # fast-forward LR schedule
        accelerator.print(f"resumed from {cfg['resume_from']} at step {step} "
                          f"(missing {len(missing)}, unexpected {len(unexpected)})")
    grad_accum = cfg.get("grad_accum", 1)
    staleness_start = cfg.get("staleness_start_step")
    staleness_on = False
    data_iter = iter(loader)
    while step < steps:
        if staleness_start is not None and not staleness_on and step >= staleness_start:
            # persistent workers hold dataset copies — rebuild the loader
            for src in datasets:
                src.staleness_prob = cfg.get("staleness_prob", 0.5)
            del data_iter
            loader = torch.utils.data.DataLoader(
                mixture,
                batch_size=cfg["batch_size"],
                num_workers=cfg.get("num_workers", 8),
                pin_memory=True,
                persistent_workers=True,
                drop_last=True,
            )
            data_iter = iter(loader)
            staleness_on = True
            accelerator.print(f"staleness augmentation ON at step {step}")
        opt.zero_grad()
        for _ in range(grad_accum):
            try:
                batch = next(data_iter)
            except StopIteration:
                data_iter = iter(loader)
                batch = next(data_iter)
            tok = tokenizer(
                list(batch.pop("task")),
                padding=True,
                truncation=True,
                max_length=pcfg.tokenizer_max_length,
                return_tensors="pt",
            )
            batch["observation.language.tokens"] = tok["input_ids"]
            batch["observation.language.attention_mask"] = tok["attention_mask"].bool()
            batch = {
                k: v.to(accelerator.device, non_blocking=True) if torch.is_tensor(v) else v
                for k, v in batch.items()
            }
            loss, info = policy(batch)
            accelerator.backward(loss / grad_accum)
        accelerator.clip_grad_norm_(policy.parameters(), cfg.get("grad_clip", 10.0))
        opt.step()
        sched.step()
        step += 1

        if step % cfg.get("log_freq", 50) == 0:
            it_s = cfg.get("log_freq", 50) / (time.time() - t0)
            t0 = time.time()
            accelerator.print(f"step {step}/{steps} loss {info['loss']:.4f} {it_s:.2f} it/s")
            if cfg.get("wandb") and accelerator.is_main_process:
                wandb.log({"loss": info["loss"], "lr": sched.get_last_lr()[0]}, step=step)

        if step % cfg.get("save_freq", 2000) == 0 and accelerator.is_main_process:
            accelerator.unwrap_model(policy).save_pretrained(out_dir / f"step_{step}")

    if accelerator.is_main_process:
        accelerator.unwrap_model(policy).save_pretrained(out_dir / "final")


if __name__ == "__main__":
    main()