#!/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()