tinyvla2/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)
        # LeRobot-backed sources count episodes; wds packs are flat sample lists
        size = f"eps={src.ds.num_episodes}" if hasattr(src, "ds") else f"n={len(src)}"
        accelerator.print(
            f"dataset[{len(datasets)-1}] {name}: {size} 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")
    action_space = pcfg.action_space  # "native" (A) or "canonical" (B, C)
    morph_yaml = cfg.get("morphology_descriptors")
    morph_map = {}
    if pcfg.conditioning == "morph" and morph_yaml:
        import torch as _t

        from tinyvla.modules.embodiment import MORPH_FIELDS
        raw = yaml.safe_load(Path(morph_yaml).read_text())
        # normalization applied here (see MORPH_FIELDS comment in embodiment.py)
        _sc = {"arm_dof": 0.1, "reach_m": 2, "gripper_width_m": 10, "num_cameras": 1 / 3,
               "control_hz": 1 / 30, "joint_lo_mean": 1 / 3.1416, "joint_hi_mean": 1 / 3.1416,
               "workspace_x": 2, "workspace_y": 2, "workspace_z": 2, "payload_kg": 0.2}
        for krobot, d in raw.items():
            vec = [d.get(f, 0) * _sc.get(f, 1) for f in MORPH_FIELDS]
            morph_map[krobot] = _t.tensor(vec, dtype=_t.float32)

    # natural-language robot descriptions for the Qwen prompt (slow path)
    prompt_map = {}
    if cfg.get("robot_prompts"):
        prompt_map = yaml.safe_load(Path(cfg["robot_prompts"]).read_text())

    # conditioning="morph_qwen": tokenize each robot's text description once,
    # fixed-length, for the shared-Qwen text-only morphology encoder
    morph_text_map = {}
    if pcfg.conditioning == "morph_qwen" and prompt_map:
        from transformers import AutoTokenizer as _Tok

        _tok0 = _Tok.from_pretrained(pcfg.lm_model_name)
        for krobot, text in prompt_map.items():
            t = _tok0([text], padding="max_length", truncation=True,
                      max_length=pcfg.morph_text_max_len, return_tensors="pt")
            morph_text_map[krobot] = (t["input_ids"][0], t["attention_mask"][0].bool())

    def wrap(ds, emb_id, morph_key=None):
        name = ds.repo_id.split("/")[-1]
        store = None
        if labels_dir:
            from tinyvla.data.spatial_labels import SpatialLabelStore

            store = SpatialLabelStore(labels_dir, name)
            if len(store) == 0:
                store = None
        canon_store, canon_stats = None, None
        if action_space == "canonical":
            from tinyvla.data.canonical import CanonicalChunkStore

            canon_store = CanonicalChunkStore(name, src_fps=ds.fps, chunk=chunk)
            canon_stats = canon_store.compute_stats()
        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,
            spatial_labels=store,
            action_space=action_space,
            canonical_store=canon_store,
            canonical_stats=canon_stats,
            morphology=morph_map.get(morph_key) if morph_key else None,
            # prepend-to-task text conditioning (separate experiment, didn't help)
            robot_prompt=prompt_map.get(morph_key) if (morph_key and pcfg.conditioning != "morph_qwen") else None,
            morph_text_ids=morph_text_map.get(morph_key, (None, None))[0] if morph_key else None,
            morph_text_mask=morph_text_map.get(morph_key, (None, None))[1] if morph_key else None,
            # demos are needed either by the MLP demo-encoder or by the rich-slow
            # LM sequence (vlm_native), so enable sampling for both
            n_support=pcfg.n_support if (pcfg.use_demo_conditioning or pcfg.vlm_native) else 0,
            support_other_task=pcfg.support_other_task,
        )

    next_emb = 0
    for spec in cfg["datasets"]:
        mkey = spec.get("morph_key")  # e.g. "so101"/"bridge"/"rt1" for variant C
        if "wds_root" in spec:
            # Prebuilt WebDataset pack (unitree / navigation). Actions are NATIVE
            # here (joint positions 7-28d, nav waypoints 3d), not canonical EE —
            # per-sample action_dim_mask is what lets one head serve both spaces.
            from tinyvla.data.wds_shards import WdsShardSource

            roots = sorted(
                q for q in Path(spec["wds_root"]).expanduser().glob(spec.get("glob", ""))
            ) if spec.get("glob") else [Path(spec["wds_root"]).expanduser()]
            roots = [r for r in roots if (r / "manifest.json").exists()]
            if not roots:
                raise FileNotFoundError(f"no wds packs under {spec['wds_root']} {spec.get('glob','')}")
            sizes = []
            srcs = []
            for r in roots:
                src = WdsShardSource(
                    r,
                    embodiment_id=spec.get("embodiment_id"),
                    chunk=chunk,
                    image_size=pcfg.image_size,
                    max_state_dim=pcfg.max_state_dim,
                    max_action_dim=pcfg.max_action_dim,
                    morphology=morph_map.get(mkey) if mkey else None,
                    robot_prompt=prompt_map.get(mkey) if (mkey and pcfg.conditioning != "morph_qwen") else None,
                    emit_latent_image=pcfg.staleness_prob > 0,
                )
                srcs.append(src)
                sizes.append(len(src))
            total = sum(sizes) or 1
            for src, n in zip(srcs, sizes):  # split the spec weight by sample count
                add(src, spec["weight"] * n / total, src.name, src.embodiment_id)
            next_emb += 1
            continue
        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, mkey), 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, mkey), spec["weight"], spec["repo_id"], emb)
        next_emb += 1

    mixture = WeightedMixtureDataset(datasets, weights, seed=cfg.get("seed", 42))
    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()