Upload tinyvla2/scripts/train_groot_morph.py with huggingface_hub

tinyvla2/scripts/train_groot_morph.py

@@ -0,0 +1,153 @@
+#!/usr/bin/env python
+"""Convert the id-conditioned GR00T checkpoint to morphology conditioning.
+
+Why this run exists: the physical_ai_ft checkpoint is `conditioning: id`, and a
+learned id is a lookup table — it has no row for a robot that was never trained on,
+and none can be synthesised. Measured on held-out LeKiwi: every trained id scores
+1.18-1.33 on the absolute dims (worse than holding still) and a PURE RANDOM identity
+vector does just as well. Simply switching the config to `morph` and grafting
+tv2_C_scaled's morph encoder does not help either (1.30/1.11) — the expert spent 60k
+steps learning to read an id token and has never seen morph tokens.
+
+So the expert has to be taught to read the descriptor. Everything else is already
+trained, so this only has to re-route one conditioning slot: 16 physical numbers ->
+MLP -> 2 tokens, in place of the 1 id token.
+
+Data is streamed from the hub exactly as the original run did (no disk), with the
+six GR00T family descriptors from configs/groot_morphology.yaml.
+
+The payoff is the project's central claim, measured rather than argued: if a
+descriptor-conditioned model transfers to LeKiwi where the id-conditioned one cannot,
+that is the difference between a code you can look up and a code you can compute.
+"""
+
+from __future__ import annotations
+
+import argparse
+import time
+from pathlib import Path
+
+import torch
+import yaml
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--resume", default="/home/alexw/tinyvla/outputs/physical_ai_ft_morph")
+    ap.add_argument("--out", default="/home/alexw/tinyvla/outputs/groot_morph")
+    ap.add_argument("--spec", default="/home/alexw/tinyvla/configs/physical_ai_stream.yaml")
+    ap.add_argument("--descriptors", default="/home/alexw/tinyvla/configs/groot_morphology.yaml")
+    ap.add_argument("--steps", type=int, default=6000)
+    ap.add_argument("--batch-size", type=int, default=32)
+    ap.add_argument("--num-workers", type=int, default=6)
+    ap.add_argument("--lr", type=float, default=1.0e-4)
+    ap.add_argument("--backbone-lr-mult", type=float, default=0.1)
+    ap.add_argument("--warmup", type=int, default=200)
+    ap.add_argument("--save-freq", type=int, default=1000)
+    ap.add_argument("--log-freq", type=int, default=50)
+    args = ap.parse_args()
+
+    from safetensors.torch import load_file
+    from transformers import AutoTokenizer
+
+    from tinyvla.data.hub_stream import HubEpisodeStream
+    from tinyvla.modeling_tinyvla import TinyVLAPolicy
+    from tinyvla.modules.embodiment import MORPH_FIELDS
+
+    out = Path(args.out)
+    out.mkdir(parents=True, exist_ok=True)
+
+    pol = TinyVLAPolicy.from_pretrained(args.resume)
+    pol.load_state_dict(load_file(f"{args.resume}/model.safetensors"), strict=True)
+    cfg = pol.config
+    assert cfg.conditioning == "morph", f"expected morph conditioning, got {cfg.conditioning}"
+    pol = pol.cuda().train()
+    tok = AutoTokenizer.from_pretrained(cfg.lm_model_name)
+    print(f"resumed {args.resume} | conditioning={cfg.conditioning} "
+          f"max_state={cfg.max_state_dim} max_action={cfg.max_action_dim}", flush=True)
+
+    # same field scaling the rest of the project uses (see MORPH_FIELDS 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}
+    raw = yaml.safe_load(Path(args.descriptors).read_text())
+    morph = {k: torch.tensor([v.get(f, 0) * sc.get(f, 1) for f in MORPH_FIELDS], dtype=torch.float32)
+             for k, v in raw.items()}
+    print("descriptors:", list(morph))
+
+    specs = yaml.safe_load(Path(args.spec).read_text())["datasets"]
+    missing = {s.get("morph_key") for s in specs} - set(morph)
+    assert not missing, f"specs reference descriptors that do not exist: {missing}"
+
+    # Fix the sampler bug the original run shipped with: HubEpisodeStream draws a
+    # DATASET PER EPISODE, so a weight of sqrt(frames) is silently multiplied by
+    # episode length. r1_pro episodes are ~37x longer, which turned an intended 49%
+    # share into a realized 96.5% — and here it also starves the loader, because
+    # every draw pulls a ~190 MB two-camera episode and the first batch never
+    # arrives. Dividing by relative episode length restores the intended mixture.
+    EP_LEN = {"r1_pro": 37.0, "gr1": 1.0, "panda_single": 1.0,
+              "panda_bi_gripper": 1.0, "panda_bi_hand": 1.0, "g1": 0.6}
+    for sp in specs:
+        sp["weight"] = float(sp["weight"]) / EP_LEN.get(sp.get("morph_key"), 1.0)
+    from collections import defaultdict as _dd
+    _agg = _dd(float)
+    for sp in specs:
+        _agg[sp["morph_key"]] += sp["weight"]
+    _tot = sum(_agg.values())
+    print("reweighted family shares:", {k: round(v / _tot, 3) for k, v in sorted(_agg.items())})
+
+    stream = HubEpisodeStream(specs, chunk=cfg.chunk_size, image_size=cfg.image_size,
+                              max_state_dim=cfg.max_state_dim, max_action_dim=cfg.max_action_dim,
+                              shuffle_buffer=64, seed=1234, morph_descriptors=morph)
+    loader = torch.utils.data.DataLoader(
+        stream, batch_size=args.batch_size, num_workers=args.num_workers,
+        pin_memory=True, prefetch_factor=4, persistent_workers=True,
+        multiprocessing_context="spawn",  # fork deadlocks the hub HTTP pool (B200.md)
+    )
+
+    backbone, head = [], []
+    for n, p in pol.named_parameters():
+        if not p.requires_grad:
+            continue
+        (backbone if n.startswith("semantic.") else head).append(p)
+    opt = torch.optim.AdamW([
+        {"params": head, "lr": args.lr},
+        {"params": backbone, "lr": args.lr * args.backbone_lr_mult},
+    ], weight_decay=1e-4)
+    sched = torch.optim.lr_scheduler.LambdaLR(
+        opt, lambda s: min(1.0, (s + 1) / args.warmup))
+    scaler_dtype = torch.bfloat16
+    print(f"head {sum(p.numel() for p in head)/1e6:.1f}M | "
+          f"backbone {sum(p.numel() for p in backbone)/1e6:.1f}M at {args.backbone_lr_mult}x", flush=True)
+
+    step, t0 = 0, time.time()
+    for batch in loader:
+        t = tok(list(batch["task"]), padding="max_length", truncation=True,
+                max_length=48, return_tensors="pt")
+        b = {k: v.cuda(non_blocking=True) for k, v in batch.items()
+             if isinstance(v, torch.Tensor)}
+        b["observation.language.tokens"] = t["input_ids"].cuda()
+        b["observation.language.attention_mask"] = t["attention_mask"].bool().cuda()
+        with torch.autocast("cuda", scaler_dtype):
+            loss, info = pol.forward(b)
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_([p for g in opt.param_groups for p in g["params"]], 10.0)
+        opt.step()
+        sched.step()
+        opt.zero_grad(set_to_none=True)
+        step += 1
+        if step % args.log_freq == 0:
+            print(f"step {step}/{args.steps} loss {info['loss']:.4f} "
+                  f"{step/(time.time()-t0):.2f} it/s", flush=True)
+        if step % args.save_freq == 0 or step == args.steps:
+            d = out / (f"step_{step}" if step < args.steps else "final")
+            d.mkdir(parents=True, exist_ok=True)
+            pol.save_pretrained(str(d))
+            print(f"saved {d}", flush=True)
+        if step >= args.steps:
+            break
+    print("DONE", flush=True)
+
+
+if __name__ == "__main__":
+    main()