tinyvla2/scripts/eval_physical_ai.py

#!/usr/bin/env python
"""Evaluate the physical_ai_ft (GR00T) checkpoint — the run that was never measured.

Everything here exists because this checkpoint can produce a plausible-looking
number for the wrong reason. Three specific traps, each with a gate or control:

1. The published weights carry torch.compile's `expert._orig_mod.` prefix, and
   LeRobot loads with strict=False while FlowMatchingExpert zero-inits action_out.
   Loaded naively the model returns its own initial Gaussian noise. Gate: strict
   load of the rekeyed copy; Control B re-runs the broken path as a labelled row.

2. Action semantics are per-DIM, not per-family: r1_pro's `action[3:6]` IS
   `state[6:9]`, gr1 mixes 12 absolute dims with deltas, and 17 of gr1's 44 dims
   are constant. So the trivial baseline is searched per dim, in raw units, not
   assumed. Getting this wrong on r1_pro (hold-still 0.895 vs predict-mean 3.79)
   would manufacture a 4x win out of nothing.

3. The realized training mixture is ~97% r1_pro, because the stream samples a
   dataset per EPISODE and r1_pro episodes are ~37x longer. Per-family numbers are
   printed next to that table or they will be misread.

Usage:
    python scripts/eval_physical_ai.py [--preflight] [--datasets-per-family 8]
"""

from __future__ import annotations

import argparse
import json
import math
from collections import defaultdict
from pathlib import Path

import numpy as np
import torch
import yaml

CKPT = "/home/alexw/tinyvla/outputs/physical_ai_ft_fixed"
RAW_CKPT = "/home/alexw/tinyvla_data/b200/outputs/physical_ai_ft/final"
SPEC = "/home/alexw/tinyvla/configs/physical_ai_stream.yaml"
FAMILY = {10: "gr1", 11: "bi_panda_grip", 12: "bi_panda_hand",
          13: "single_panda", 14: "r1_pro", 15: "g1"}
DEAD_STD = 1e-4
EVAL_SEED = 20260830  # deliberately not the training seed (42)


# ----------------------------------------------------------------- calibration

def classify_dims(state_raw, action_raw):
    """Per action dim: ABS (tracks a state dim), DELTA, or DEAD (constant).

    Searched in RAW units and over lags 0..2, because the z-score hides exact
    identities and commanded targets lead the measured state.
    """
    T, A = action_raw.shape
    S = state_raw.shape[1]
    std_a = action_raw.std(0)
    kind = np.full(A, "delta", dtype=object)
    amap = np.full(A, -1, dtype=int)
    lag = np.zeros(A, dtype=int)

    dead = std_a < DEAD_STD
    kind[dead] = "dead"

    for j in range(A):
        if dead[j]:
            continue
        best = (np.inf, -1, 0)
        for lg in (0, 1, 2):
            a = action_raw[: T - lg, j]
            sl = state_raw[lg:, :]
            rms = np.sqrt(((sl - a[:, None]) ** 2).mean(0))
            i = int(np.argmin(rms))
            if rms[i] < best[0]:
                best = (float(rms[i]), i, lg)
        rms, i, lg = best
        if i < 0:
            continue
        a = action_raw[: T - lg, j]
        s = state_raw[lg:, i]
        if s.std() < 1e-9 or a.std() < 1e-9:
            continue
        corr = float(np.corrcoef(a, s)[0, 1])
        if corr > 0.95 and rms < 0.5 * std_a[j]:
            kind[j], amap[j], lag[j] = "abs", i, lg
    return kind, amap, lag, std_a


def displacement(A, kind, amap, ref):
    """(chunk, adim) actions -> displacement, per the dim's semantics.

    ABS   : residual against the hold-still reference  (trivial baseline = hold still)
    DELTA : integrated                                  (trivial baseline = zero motion)
    In both cases the trivial baseline maps to d == 0, which is what makes the
    ratio comparable across families.
    """
    d = np.array(A, dtype=np.float64, copy=True)
    is_abs = kind == "abs"
    is_del = kind == "delta"
    d[:, is_abs] = A[:, is_abs] - ref[is_abs][None, :]
    d[:, is_del] = np.cumsum(A[:, is_del], axis=0)
    return d


def chunk_err(pred, gt, kind, amap, ref, sel):
    """Mean-over-steps L2 of the displacement error, restricted to dims `sel`."""
    dp = displacement(pred, kind, amap, ref)[:, sel]
    dg = displacement(gt, kind, amap, ref)[:, sel]
    return float(np.linalg.norm(dp - dg, axis=1).mean()), float(np.linalg.norm(dg, axis=1).mean())


# --------------------------------------------------------------------- helpers

def make_batch(items, tok, cfg, device="cuda", state_override=None, image_override=None,
               task_override=None, emb_override=None, padding="max_length"):
    tasks = task_override if task_override is not None else [it["task"] for it in items]
    t = tok(tasks, padding=padding, truncation=True, max_length=48, return_tensors="pt")
    cam0 = torch.stack([it["observation.images.cam0"] for it in items])
    cam1 = torch.stack([it["observation.images.cam1"] for it in items])
    if image_override is not None:
        cam0, cam1 = image_override
    st = torch.stack([it["observation.state"] for it in items])
    if state_override is not None:
        st = state_override
    emb = torch.stack([it["embodiment_id"] for it in items])
    if emb_override is not None:
        emb = emb_override
    return {
        "observation.images.cam0": cam0.to(device),
        "observation.images.cam1": cam1.to(device),
        "observation.state": st.to(device),
        "observation.language.tokens": t["input_ids"].to(device),
        "observation.language.attention_mask": t["attention_mask"].bool().to(device),
        "embodiment_id": emb.to(device),
    }


def predict_mean(pol, batch, n_seeds: int):
    """Average the flow ODE over n_seeds draws of the initial noise.

    predict_action_chunk starts from torch.randn, and control C measured seed
    sensitivity S ~ 1.2 on this checkpoint: two draws disagree with each other
    about as much as either disagrees with ground truth. The expert itself is fine
    (flow loss L/L0 = 0.12) — the 10-step Euler integration simply does not
    contract the initial noise. Since the metric is L2, the conditional MEAN is the
    optimal point estimate, so averaging is the right estimator, not a fudge.
    """
    acc = None
    for k in range(n_seeds):
        torch.manual_seed(1000 + k)
        with torch.autocast("cuda", torch.bfloat16):
            p = pol.predict_action_chunk(batch).float()
        acc = p if acc is None else acc + p
    return (acc / n_seeds).cpu().numpy()


def bootstrap_ci(per_ep, n=2000, seed=0):
    """CI over EPISODES — frames inside an episode are not independent."""
    if len(per_ep) < 2:
        return (float("nan"), float("nan"))
    rng = np.random.default_rng(seed)
    a = np.asarray(per_ep, dtype=float)
    means = [a[rng.integers(0, len(a), len(a))].mean() for _ in range(n)]
    return (float(np.percentile(means, 2.5)), float(np.percentile(means, 97.5)))


def realized_mixture(specs, metas):
    """Sample share actually seen, given the stream draws a DATASET PER EPISODE."""
    rows = []
    for si, s in enumerate(specs):
        m = metas.get(si)
        if m is None:
            continue
        rows.append((s["embodiment_id"], s["weight"], m["samples_per_ep"], m["n_eps"]))
    tot_w = sum(r[1] for r in rows) or 1.0
    fam = defaultdict(lambda: [0.0, 0.0, 0])
    for emb, w, spe, n_eps in rows:
        share = (w / tot_w) * spe
        fam[emb][0] += share
        fam[emb][1] += w / tot_w
        fam[emb][2] += n_eps
    norm = sum(v[0] for v in fam.values()) or 1.0
    return {k: (v[0] / norm, v[1], v[2]) for k, v in fam.items()}


# ------------------------------------------------------------------------ main

@torch.no_grad()
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--preflight", action="store_true", help="gates and tables only")
    ap.add_argument("--datasets-per-family", type=int, default=8)
    ap.add_argument("--episodes", type=int, default=3)
    ap.add_argument("--starts", type=int, default=12)
    ap.add_argument("--seeds", type=int, default=4,
                    help="average the flow ODE over N noise draws (control C measured S~1.2)")
    ap.add_argument("--out", default="/home/alexw/tinyvla/outputs/physical_ai_eval.json")
    args = ap.parse_args()

    from safetensors.torch import load_file
    from transformers import AutoTokenizer

    from tinyvla.data.hub_stream import EvalEpisodeStream
    from tinyvla.modeling_tinyvla import TinyVLAPolicy

    specs_all = yaml.safe_load(open(SPEC))["datasets"]

    # ---------------------------------------------------------------- GATE 1
    pol = TinyVLAPolicy.from_pretrained(CKPT)
    sd = load_file(f"{CKPT}/model.safetensors")
    pol.load_state_dict(sd, strict=True)
    aow = pol.expert.action_out.weight.abs().mean().item()
    assert aow > 1e-3, "expert.action_out at zero init — expert never trained"
    cfg = pol.config
    pol = pol.cuda().eval()
    print(f"GATE  rekey {len(sd)}/{len(pol.state_dict())} strict=True OK | "
          f"action_out |w|={aow:.4f} | action_dim={pol.action_dim} "
          f"max_action={cfg.max_action_dim} max_state={cfg.max_state_dim}")
    tok = AutoTokenizer.from_pretrained(cfg.lm_model_name)

    # ------------------------------------------------- dataset selection
    rng = np.random.default_rng(EVAL_SEED)
    by_fam = defaultdict(list)
    for s in specs_all:
        by_fam[s["embodiment_id"]].append(s)
    chosen = []
    for emb in sorted(by_fam):
        pool = by_fam[emb]
        k = min(args.datasets_per_family if emb != 14 else 12, len(pool))
        idx = rng.choice(len(pool), size=k, replace=False)
        chosen += [pool[i] for i in sorted(idx)]
    print(f"selected {len(chosen)} datasets across {len(by_fam)} families "
          f"(seed {EVAL_SEED}, not the training seed)")

    stream = EvalEpisodeStream(chosen, 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=1)

    # ------------------------------------------------- GATE 2/3 + calibration
    metas, calib = {}, {}
    for si, s in enumerate(chosen):
        try:
            m = stream._load_meta(si)
            eps = m["episodes"]
            metas[si] = {"n_eps": len(eps),
                         "samples_per_ep": float(np.mean([e[1] for e in eps[:200]])) / m["stride"],
                         "stride": m["stride"], "adim": m["action_dim"],
                         "stats_src": m.get("stats_src", "?")}
            sraw, araw, _ = stream.episode_raw(si)
            kind, amap, lag, std_a = classify_dims(sraw, araw)
            calib[si] = dict(kind=kind, amap=amap, lag=lag, std_a=std_a)
            print(f"  [{si:>3}] {s['prefix'][:44] or s['repo_id'][-30:]:46} "
                  f"emb={s['embodiment_id']} adim={m['action_dim']:>3} "
                  f"abs={int((kind=='abs').sum()):>3} del={int((kind=='delta').sum()):>3} "
                  f"dead={int((kind=='dead').sum()):>3}", flush=True)
        except Exception as e:
            print(f"  [{si:>3}] {s.get('prefix','')[:44]:46} CALIB FAIL {type(e).__name__}: {str(e)[:60]}",
                  flush=True)

    mix = realized_mixture(chosen, metas)
    print("\nREALIZED MIXTURE (stream draws a dataset PER EPISODE, so share = weight x episode length)")
    for emb in sorted(mix):
        share, wshare, n_eps = mix[emb]
        print(f"  {FAMILY[emb]:14} realized {share*100:5.1f}%   intended(weight) {wshare*100:5.1f}%   "
              f"episodes {n_eps}")
    draws_8gpu = 8 * 256 * 60000 / max(1e-9, np.mean([m["samples_per_ep"] for m in metas.values()]))
    print("HOLD-OUT  P(episode never drawn), 8-GPU / 1-GPU scenario (no log disambiguates):")
    for emb in sorted(mix):
        share, _, n_eps = mix[emb]
        lam8 = draws_8gpu * share / max(1, n_eps)
        print(f"  {FAMILY[emb]:14} P(unseen) {math.exp(-lam8)*100:5.1f}% / {math.exp(-lam8/8)*100:5.1f}%")

    if args.preflight:
        return

    # ------------------------------------------------------------ main pass
    print("\n=== main pass ===", flush=True)
    results = defaultdict(lambda: defaultdict(list))
    per_ep_json = []
    subsample = []  # kept for the control block

    for si, s in enumerate(chosen):
        if si not in calib:
            continue
        c = calib[si]
        emb = s["embodiment_id"]
        m = stream._load_meta(si)
        eps = m["episodes"]
        n_ep = 1 if emb == 14 else args.episodes
        n_st = 24 if emb == 14 else args.starts
        pick = rng.choice(len(eps), size=min(n_ep, len(eps)), replace=False)
        for e_i in pick:
            ep_idx, length, task = eps[e_i]
            n_rows = length // m["stride"]
            hi = n_rows - cfg.chunk_size - 1
            if hi <= 0:
                continue
            if emb == 14:  # prefix-decode: starts confined to the episode head (labelled bias)
                hi = min(hi, 600)
            starts = np.unique(np.linspace(0, hi, n_st).astype(int))
            try:
                items = stream.episode_samples(si, ep_idx, task, starts)
            except Exception as ex:
                print(f"  ep fail {s.get('prefix','')[:36]} {type(ex).__name__}: {str(ex)[:60]}", flush=True)
                continue
            if not items:
                continue

            preds = []
            for k in range(0, len(items), 8):
                b = make_batch(items[k:k + 8], tok, cfg)
                preds.append(predict_mean(pol, b, args.seeds))
            pred = np.concatenate(preds, 0)

            ad = m["action_dim"]
            sm, ss = m["stats"]["observation.state"]
            am, as_ = m["stats"]["action"]
            kind, amap = c["kind"], c["amap"]
            for subset in ("abs", "delta"):
                sel = (kind == subset)
                if not sel.any():
                    continue
                e_l, f0_l, f1_l, f2_l = [], [], [], []
                for t_i, it in enumerate(items):
                    A = it["action"][:, :ad].numpy().astype(np.float64)
                    P = pred[t_i][:, :ad].astype(np.float64)
                    s_norm = it["observation.state"].numpy()
                    s_raw = s_norm[:len(sm)] * np.maximum(ss, 1e-6) + sm
                    ref = np.zeros(ad)
                    for j in range(ad):
                        if kind[j] == "abs" and amap[j] >= 0 and amap[j] < len(s_raw):
                            ref[j] = (s_raw[amap[j]] - am[j]) / max(as_[j], 1e-6)
                    err, floor_zero = chunk_err(P, A, kind, amap, ref, sel)
                    b0, _ = chunk_err(np.zeros_like(A), A, kind, amap, ref, sel)
                    b1, _ = chunk_err(np.tile(ref, (len(A), 1)), A, kind, amap, ref, sel)
                    b2, _ = chunk_err(np.tile(A[0], (len(A), 1)), A, kind, amap, ref, sel)
                    e_l.append(err); f0_l.append(b0); f1_l.append(b1); f2_l.append(b2)
                key = (emb, subset)
                results[key]["err"].append(np.mean(e_l))
                results[key]["b0"].append(np.mean(f0_l))
                results[key]["b1"].append(np.mean(f1_l))
                results[key]["b2"].append(np.mean(f2_l))
                results[key]["ndim"].append(int(sel.sum()))
                per_ep_json.append(dict(emb=int(emb), subset=subset, dataset=s.get("prefix", s["repo_id"]),
                                        episode=int(ep_idx), err=float(np.mean(e_l)),
                                        b0=float(np.mean(f0_l)), b1=float(np.mean(f1_l)),
                                        b2=float(np.mean(f2_l)), ndim=int(sel.sum())))
            if len(subsample) < 24:
                subsample.append((si, items[:4], c, m))
            print(f"  {FAMILY[emb]:14} {s.get('prefix','')[:34]:36} ep{ep_idx:<6} "
                  f"n={len(items)}", flush=True)

    # ------------------------------------------------------------- report
    print("\n=== RESULTS (displacement error / trivial floor; <1 beats the floor) ===")
    print(f"{'family':16}{'subset':7}{'nep':>4}{'dims':>5}{'err':>8}{'B0':>8}{'B1':>8}{'B2':>8}"
          f"{'ratio':>8}   95% CI (bootstrap over episodes)")
    summary = {}
    for (emb, subset), v in sorted(results.items()):
        err = np.mean(v["err"]); b0 = np.mean(v["b0"]); b1 = np.mean(v["b1"]); b2 = np.mean(v["b2"])
        floor = min(b0, b1)
        ratios = [e / max(min(x, y), 1e-9) for e, x, y in zip(v["err"], v["b0"], v["b1"])]
        lo, hi = bootstrap_ci(ratios)
        tag = "  [IN-TRAINING]" if emb == 14 else ""
        print(f"{FAMILY[emb]:16}{subset:7}{len(v['err']):>4}{int(np.mean(v['ndim'])):>5}"
              f"{err:>8.3f}{b0:>8.3f}{b1:>8.3f}{b2:>8.3f}{err/floor:>8.2f}   "
              f"[{lo:.2f}, {hi:.2f}]{tag}")
        summary[f"{FAMILY[emb]}/{subset}"] = dict(err=err, b0=b0, b1=b1, b2=b2,
                                                  ratio=err / floor, ci=[lo, hi],
                                                  n_ep=len(v["err"]))
        if abs(b2 - floor) / max(floor, 1e-9) < 0.1:
            print(f"{'':16}^^ oracle repeat-A0 ({b2:.3f}) is within 10% of the floor "
                  f"({floor:.3f}) — this subset carries almost no predictable signal")

    # ------------------------------------------------------------ controls
    print("\n=== CONTROLS ===", flush=True)
    ctrl = {}

    # C: ODE seed sensitivity
    si, items, c, m = subsample[0]
    b = make_batch(items, tok, cfg)
    torch.manual_seed(1)
    with torch.autocast("cuda", torch.bfloat16):
        p1 = pol.predict_action_chunk(b).float().cpu().numpy()
    torch.manual_seed(2)
    with torch.autocast("cuda", torch.bfloat16):
        p2 = pol.predict_action_chunk(b).float().cpu().numpy()
    gt = torch.stack([it["action"] for it in items]).numpy()[:, :, :p1.shape[-1]]
    S = float(np.abs(p1 - p2).mean() / max(np.abs(p1 - gt).mean(), 1e-9))
    a1 = predict_mean(pol, b, args.seeds)
    torch.manual_seed(77)
    a2 = predict_mean(pol, b, args.seeds)
    S_avg = float(np.abs(a1 - a2).mean() / max(np.abs(a1 - gt).mean(), 1e-9))
    ctrl["ode_seed_sensitivity"] = S
    ctrl["ode_seed_sensitivity_averaged"] = S_avg
    print(f"  C  ODE seed sensitivity: single draw S = {S:.3f}, "
          f"{args.seeds}-seed average S = {S_avg:.3f}   (headline uses the average)")

    # E: flow-matching training loss vs the trivial predictor
    e_rows = {}
    for si, items, c, m in subsample[:12]:
        emb = chosen[si]["embodiment_id"]
        b = make_batch(items, tok, cfg)
        b["action"] = torch.stack([it["action"] for it in items]).cuda()
        b["action_dim_mask"] = torch.stack([it["action_dim_mask"] for it in items]).cuda()
        b["action_is_pad"] = torch.stack([it["action_is_pad"] for it in items]).cuda()
        torch.manual_seed(0)
        with torch.autocast("cuda", torch.bfloat16):
            loss, _ = pol.forward(b)
        A = b["action"][:, :, : m["action_dim"]]
        l0 = float((torch.randn_like(A) - A).pow(2).mean())
        e_rows.setdefault(emb, []).append(float(loss) / l0)
    ctrl["flow_loss_ratio"] = {FAMILY[k]: float(np.mean(v)) for k, v in e_rows.items()}
    print("  E  flow loss L/L0 :", {k: round(v, 3) for k, v in ctrl["flow_loss_ratio"].items()},
          "  (<1 = expert learned something, independent of the ODE and the metric)")

    # D3/D4: in-distribution donor swap — is vision used at all, or is this a proprio regressor?
    def ratio_for(items, c, m, state_override=None, image_override=None, emb_override=None):
        b = make_batch(items, tok, cfg, state_override=state_override,
                       image_override=image_override, emb_override=emb_override)
        pr = predict_mean(pol, b, args.seeds)
        ad = m["action_dim"]
        sm, ss = m["stats"]["observation.state"]
        am, as_ = m["stats"]["action"]
        kind, amap = c["kind"], c["amap"]
        sel = (kind != "dead")
        num, den = [], []
        for t_i, it in enumerate(items):
            A = it["action"][:, :ad].numpy().astype(np.float64)
            P = pr[t_i][:, :ad].astype(np.float64)
            s_raw = it["observation.state"].numpy()[:len(sm)] * np.maximum(ss, 1e-6) + sm
            ref = np.zeros(ad)
            for j in range(ad):
                if kind[j] == "abs" and 0 <= amap[j] < len(s_raw):
                    ref[j] = (s_raw[amap[j]] - am[j]) / max(as_[j], 1e-6)
            e, f = chunk_err(P, A, kind, amap, ref, sel)
            b0, _ = chunk_err(np.zeros_like(A), A, kind, amap, ref, sel)
            b1, _ = chunk_err(np.tile(ref, (len(A), 1)), A, kind, amap, ref, sel)
            num.append(e); den.append(min(b0, b1))
        return float(np.mean(num) / max(np.mean(den), 1e-9))

    d_rows = defaultdict(dict)
    for pos, (si, items, c, m) in enumerate(subsample[:8]):
        emb = chosen[si]["embodiment_id"]
        donor = subsample[(pos + 1) % len(subsample)][1]
        n = min(len(items), len(donor))
        it_, dn_ = items[:n], donor[:n]
        base = ratio_for(it_, c, m)
        img = (torch.stack([d["observation.images.cam0"] for d in dn_]),
               torch.stack([d["observation.images.cam1"] for d in dn_]))
        d3 = ratio_for(it_, c, m, image_override=img)
        d4 = ratio_for(it_, c, m, state_override=torch.stack([d["observation.state"] for d in dn_]))
        d6 = ratio_for(it_, c, m, emb_override=torch.full((n,), 11 if emb != 11 else 10).long())
        d_rows[emb] = dict(base=base, d3_img=d3, d4_state=d4, d6_emb=d6)
    ctrl["ablation"] = {FAMILY[k]: {kk: round(vv, 3) for kk, vv in v.items()} for k, v in d_rows.items()}
    print("  D  donor-swap ablation (ratio; higher = input mattered):")
    for k, v in ctrl["ablation"].items():
        print(f"       {k:14} base {v['base']:.3f} | images {v['d3_img']:.3f} "
              f"| state {v['d4_state']:.3f} | wrong emb id {v['d6_emb']:.3f}")

    # B: the broken (un-rekeyed) path, as shipped
    broken = TinyVLAPolicy.from_pretrained(RAW_CKPT).cuda().eval()
    si, items, c, m = subsample[0]
    b = make_batch(items, tok, cfg)
    with torch.autocast("cuda", torch.bfloat16):
        pb = broken.predict_action_chunk(b).float().cpu().numpy()
    ad = m["action_dim"]
    A = torch.stack([it["action"] for it in items]).numpy()[:, :, :ad]
    ctrl["broken_err"] = float(np.linalg.norm(pb[:, :, :ad] - A, axis=2).mean())
    with torch.autocast("cuda", torch.bfloat16):
        pg = pol.predict_action_chunk(b).float().cpu().numpy()
    ctrl["fixed_err"] = float(np.linalg.norm(pg[:, :, :ad] - A, axis=2).mean())
    print(f"  B  un-rekeyed (as published) err {ctrl['broken_err']:.3f} vs rekeyed "
          f"{ctrl['fixed_err']:.3f}   (equal => the rename did not take effect)")
    del broken
    torch.cuda.empty_cache()

    # F: tokenizer padding fidelity (training used max_length=48)
    si, items, c, m = subsample[0]
    r_max = ratio_for(items, c, m)
    b_long = make_batch(items, tok, cfg, padding=True)
    with torch.autocast("cuda", torch.bfloat16):
        _ = pol.predict_action_chunk(b_long)
    ctrl["padding_note"] = "headline uses padding=max_length(48), as train_fast.Collate did"
    print(f"  F  padding: headline uses max_length=48 (training convention); ratio {r_max:.3f}")

    Path(args.out).write_text(json.dumps(
        {"summary": summary, "controls": ctrl, "per_episode": per_ep_json,
         "mixture": {FAMILY[k]: v[0] for k, v in mix.items()}}, indent=2))
    print(f"\nwrote {args.out}")


if __name__ == "__main__":
    main()