Blame for scripts/train_cdpr_wm.py

"""Train the Stringman SkyJEPA (cablepilot/model/learned.py) on task-free data, then test 1: open-loop prediction.
(armpilot's scripts/train_arm_wm.py on the cable robot.)

    python scripts/train_cdpr_wm.py --data data/cdpr_v0_playroom --out checkpoints/cdpr_wm_v0.pt
    python scripts/train_cdpr_wm.py --eval-only --out checkpoints/cdpr_wm_v0.pt

Stage 1 (task-free JEPA): predict the next 2 s of state latents from 1 s of history and the commands,
loss = MSE(pred, target) + 0.02 SIGReg(pred), as SkyJEPA. Stage 2: the prober on frozen latents, loss = weighted MSE of
the integrated states (normalised) against the recorded ones, the end of the horizon weighted most. The prober's last
layer starts at zero: training starts exactly at the hand-written backbone and learns only the correction.

Test 1 (docs/MILESTONES.md M1, the bar set before training): on held-out windows, after 20 steps (2.0 s): grasp point
error <= 2 cm median (also on payload windows), pole tilt error <= 2 deg median; next to the nominal oracle (MuJoCo, no
eta, no payload) and the backbone alone (nothing learned).
"""
from __future__ import annotations

import argparse
import json
import math
import sys
import time
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from cablepilot.model.learned import JEPA, Normalizer, Prober, features, integrate  # noqa: E402

H, L = 10, 20
L_SWING_DATA = 0.3089  # the data's pendulum (fw, Stringman hand); load() sets it from the manifest's hand
D_G = -0.5147          # grasp point down the pole; load() sets it from the data's manifest (SO-101 hand: -0.472)
# the machine the data came from (manifest "machine"; Stringman when absent) and its backbone priors: the person lift's
# swing runs over a 2.4x range of lengths (empty harness .. a person), its spools ramp at 0.8 m/s^2 (P1, 2026-10-12)
MACHINE, OMEGA_SPAN, A_MAX0 = "configs/stringman.yaml", 0.3, 0.25
# Stage-2 error scales per state dim, in METRES / radians: an error of one scale costs 1. In normalised units (v0) a 5 cm
# position error cost ~0.003 next to the swing and rate terms, and the grasp-point error at 2 s ROSE during stage 2
# (53 -> 68 mm) while the loss fell; near-still windows came out at 80 mm against 33 mm for "integrate the command"
# (2026-10-08). p 2 cm, v 5 cm/s, s 0.01 (~0.6 deg), s_dot 0.1 /s, wrist 0.1 rad, wrist rate 0.5, closure 0.05, its rate 0.5.
E_SCALE = torch.tensor([0.02, 0.02, 0.02, 0.05, 0.05, 0.05, 0.01, 0.01, 0.1, 0.1, 0.1, 0.5, 0.05, 0.5])


class SIGReg(nn.Module):
    """LeJEPA's Sketched Isotropic Gaussian Regularisation (copied from armpilot, from LeDrone/skyjepa)."""

    def __init__(self, n_directions=64, n_knots=17):
        super().__init__()
        nodes, weights = np.polynomial.hermite.hermgauss(n_knots)
        self.n = n_directions
        self.register_buffer("nodes", torch.tensor(nodes, dtype=torch.float32))
        self.register_buffer("weights", torch.tensor(weights, dtype=torch.float32))
        self.register_buffer("gauss_cf", torch.tensor(np.exp(-nodes ** 2 / 2.0), dtype=torch.float32))

    def forward(self, lat):                      # (T, B, D)
        T, B, D = lat.shape
        u = torch.randn(self.n, D)
        u = u / (u.norm(dim=1, keepdim=True) + 1e-8)
        stat = lat.new_zeros(())
        for t in range(T):
            arg = (lat[t] @ u.t()).unsqueeze(-1) * self.nodes.view(1, 1, -1)
            C, S = torch.cos(arg).mean(0), torch.sin(arg).mean(0)
            stat = stat + B * (((C - self.gauss_cf) ** 2 + S ** 2) * self.weights).sum(1).mean()
        return stat / T


def load(data):
    """One dataset directory, or several joined with commas: (X, A, P, M, R, dt). M (N, 4, 3) mounts and R (N, 4) runs
    per trajectory (shards from a fixed room without them: the scenario's)."""
    dirs = [Path(d) for d in str(data).split(",")]
    parts = []
    for d in dirs:
        shards = sorted(d.glob("shard_*.npz"))
        X, A, P, M, R = [], [], [], [], []
        man = json.loads((d / "manifest.json").read_text())
        global D_G, L_SWING_DATA, MACHINE, OMEGA_SPAN, A_MAX0
        D_G = float(man.get("dG", D_G))
        from cablepilot.sim.scene import load_config as _lc
        if man.get("machine"):  # another machine (P1: the person lift): its own pendulum and spools
            MACHINE = man["machine"]
            mc = _lc(machine=MACHINE)
            L_SWING_DATA = float(mc["body"]["swing_length_m"])
            A_MAX0 = float(mc["spool"]["line_accel_mps2"])
            OMEGA_SPAN = 1.0 if mc.get("hand") == "harness" else 0.3
        elif man.get("hand"):   # the hand's pendulum, to start the backbone's swing from (Prober l_swing)
            L_SWING_DATA = float(_lc(hand=man["hand"])["body"]["swing_length_m"])
        for f in shards:
            z = np.load(f)
            X.append(z["states"])
            A.append(z["actions"])
            P.append(z["payload"])
            if "mounts" in z:
                M.append(z["mounts"])
                R.append(z["runs"])
            else:
                from cablepilot.sim.scene import load_config, mounts
                c = load_config(machine=MACHINE, scenario=man["scenario"])
                M.append(np.repeat(mounts(c)[None].astype(np.float32), len(z["states"]), 0))
                R.append(np.repeat(np.asarray(c["room"]["indirect_run_m"], np.float32)[None], len(z["states"]), 0))
        parts.append([np.concatenate(v) for v in (X, A, P, M, R)] + [float(np.load(shards[0])["dt"])])
    if len(parts) == 1:
        return tuple(parts[0])
    k = min(len(p[0]) for p in parts)   # interleave so that the held-out tail holds trajectories of every set
    return tuple(np.stack([p[j][:k] for p in parts], 1).reshape(-1, *parts[0][j].shape[1:]) for j in range(5)) + (parts[0][5],)


def windows(X, A, idx, rng, B, M, P=None):
    i = rng.choice(idx, B)
    t = rng.integers(0, A.shape[1] - (H + L) + 1, B)
    ar = np.arange(H + 1 + L)
    out = [torch.as_tensor(X[i[:, None], t[:, None] + ar]), torch.as_tensor(A[i[:, None], t[:, None] + ar[:-1]]),
           torch.as_tensor(M[i])]
    if P is not None:
        out += [P[i[:, None], t[:, None] + ar], i]
    return out


def lr_at(step, total, peak, final=1e-4, warm=200):
    if step < warm:
        return peak * step / warm
    p = (step - warm) / max(1, total - warm)
    return final + (peak - final) * 0.5 * (1 + math.cos(math.pi * p))


def grasp_point(x):
    s = x[..., 6:8]
    z = np.sqrt(np.clip(1 - (s ** 2).sum(-1, keepdims=True), 1e-6, None))
    return x[..., 0:3] + D_G * np.concatenate([s, z], -1)


def train(args):
    X, A, P, M, R, dt = load(args.data)
    n = len(X)
    rng = np.random.default_rng(0)
    test_idx, train_idx = np.arange(n - args.test, n), np.arange(0, n - args.test)
    with torch.no_grad():
        sub = features(torch.as_tensor(X[train_idx[:2000], ::5]), torch.as_tensor(M[train_idx[:2000]])[:, None]).reshape(-1, 28).numpy()
    asub = A[train_idx[:2000]].reshape(-1, 5)
    norm = Normalizer(sub.mean(0), sub.std(0) + 1e-3, asub.mean(0), asub.std(0) + 1e-3)
    cfg = {"latent": 24, "action_latent": 12, "hidden": 48, "prober_hidden": 64, "H": H, "L": L, "dt": dt,
           "data": str(args.data), "trajectories_train": int(len(train_idx)), "l_swing": L_SWING_DATA,
           "omega_span": OMEGA_SPAN, "a_max0": A_MAX0, "machine": MACHINE}
    torch.manual_seed(0)
    jepa, prober, sig = JEPA(cfg["latent"], cfg["action_latent"], cfg["hidden"]), Prober(cfg["latent"], cfg["prober_hidden"], cfg["l_swing"], cfg["omega_span"], cfg["a_max0"]), SIGReg()
    n_j, n_p = sum(p.numel() for p in jepa.parameters()), sum(p.numel() for p in prober.parameters())
    print(f"{n} trajectories ({len(train_idx)} train, {len(test_idx)} test), dt {dt * 1000:.0f} ms; "
          f"JEPA {n_j} params, prober {n_p}", flush=True)
    t0 = time.perf_counter()
    opt = torch.optim.Adam(jepa.parameters(), lr=args.lr, weight_decay=1e-5)
    if args.stage1_from:                             # reuse stage 1 (frozen latents) from an earlier checkpoint
        old = torch.load(args.stage1_from, weights_only=False)
        jepa.load_state_dict(old["jepa"])
        norm = Normalizer(**old["norm"])
        print(f"stage 1 from {args.stage1_from}", flush=True)
    for step in range(0 if args.stage1_from else args.steps1):
        for g in opt.param_groups:
            g["lr"] = lr_at(step, args.steps1, args.lr)
        xs, acts, Mw = windows(X, A, train_idx, rng, args.batch, M)
        pred, tgt = jepa(norm.s(features(xs, Mw[:, None])), norm.a(acts), H)
        l_pred = ((pred - tgt) ** 2).mean()
        l_sig = sig(pred.transpose(0, 1)[::4])
        loss = l_pred + 0.02 * l_sig
        opt.zero_grad()
        loss.backward()
        nn.utils.clip_grad_norm_(jepa.parameters(), 0.5)
        opt.step()
        if step % 250 == 0 or step == args.steps1 - 1:
            std = pred.detach().reshape(-1, pred.shape[-1]).std(0).min().item()
            print(f"  stage 1  step {step:5d}  L_pred {l_pred.item():.4f}  L_sig {l_sig.item():.3f}  "
                  f"latent min std {std:.3f}  {time.perf_counter() - t0:.0f} s", flush=True)
    t1 = time.perf_counter() - t0
    jepa.eval()
    for p in jepa.parameters():
        p.requires_grad_(False)
    opt = torch.optim.Adam(prober.parameters(), lr=args.lr2, weight_decay=1e-5)
    tw = torch.linspace(0.3, 1.0, L)[None, :, None]
    for step in range(args.steps2):
        for g in opt.param_groups:
            g["lr"] = lr_at(step, args.steps2, args.lr2)
        xs, acts, Mw = windows(X, A, train_idx, rng, args.batch, M)
        with torch.no_grad():
            pred, _ = jepa(norm.s(features(xs, Mw[:, None])), norm.a(acts), H)
        out = integrate(xs[:, H], acts[:, H:], acts[:, H - 1], pred, prober, dt, norm, Mw)
        err = nn.functional.huber_loss((out - xs[:, H + 1:]) / E_SCALE, torch.zeros_like(out), reduction="none", delta=3.0)
        loss = (err * tw).mean()
        opt.zero_grad()
        loss.backward()
        nn.utils.clip_grad_norm_(prober.parameters(), 0.5)
        opt.step()
        if step % 250 == 0 or step == args.steps2 - 1:
            g_err = np.linalg.norm(grasp_point(out[:, -1].detach().numpy()) - grasp_point(xs[:, -1].numpy()), axis=-1)
            print(f"  stage 2  step {step:5d}  loss {loss.item():.4f}  grasp-point error at 2 s median "
                  f"{1000 * np.median(g_err):.1f} mm  {time.perf_counter() - t0:.0f} s", flush=True)
    ckpt = {"config": cfg, "jepa": jepa.state_dict(), "prober": prober.state_dict(), "norm": norm.state_dict(),
            "train_s": time.perf_counter() - t0, "stage1_s": t1, "params": {"jepa": n_j, "prober": n_p}}
    Path(args.out).parent.mkdir(parents=True, exist_ok=True)
    torch.save(ckpt, args.out)
    print(f"saved {args.out}  ({time.perf_counter() - t0:.0f} s)")
    return X, A, P, M, R, test_idx


def tilt_angle_err(a, b):
    def axis(x):
        s = x[..., 6:8]
        return np.concatenate([s, np.sqrt(np.clip(1 - (s ** 2).sum(-1, keepdims=True), 1e-6, None))], -1)
    c = np.clip((axis(a) * axis(b)).sum(-1), -1, 1)
    return np.degrees(np.arccos(c))


def test1(args, data=None, test_idx=None):
    from cablepilot.model.learned import LearnedCDPRModel
    from cablepilot.model.oracle import OracleWorldModel
    from cablepilot.sim.scene import load_config
    if data is None:
        data = load(args.data)
        test_idx = np.arange(len(data[0]) - args.test, len(data[0]))
    X, A, P, M, R, dt = data
    manifest = json.loads((Path(str(args.data).split(",")[0]) / "manifest.json").read_text())
    cfg = load_config(machine=manifest.get("machine", "configs/stringman.yaml"), scenario=manifest["scenario"],
                      hand=manifest.get("hand"))
    person = cfg.get("hand") == "harness"
    ck = torch.load(args.out, weights_only=False)
    model = LearnedCDPRModel(ck, dt=ck["config"]["dt"])
    rng = np.random.default_rng(123)
    xs, acts, Mw, pay, traj = windows(X, A, test_idx, rng, args.n_eval, M, P)
    xs, acts = xs.numpy(), acts.numpy()
    payload_w = (pay > 1e-3).any(1)
    event_w = np.ptp(pay, axis=1) > 1.0                       # the person's weight coming on / going off / shifting
    t0 = time.perf_counter()
    learned = []
    for b in range(len(xs)):
        model.set_room(M[traj[b]])
        learned.append(model.predict(xs[b, :H + 1], acts[b, :H], acts[b:b + 1, H:]).numpy())
    learned = np.concatenate(learned)
    t_l = (time.perf_counter() - t0) / len(xs)
    # the person lift's stiff lines need the 2 ms step; its oracle is told 75 kg (the range's middle), never the truth
    orc = OracleWorldModel(cfg, dt=ck["config"]["dt"], timestep=0.002 if person else 0.01, nthread=4,
                           payload_kg=75.0 if person else 0.0)
    t0 = time.perf_counter()
    oracle, last = [], None
    for b in range(len(xs)):
        j = traj[b]
        if last is None or not np.allclose(M[j], last):       # the oracle in this window's room
            orc.set_room({"mounts_xyz": M[j].tolist(), "indirect_run_m": R[j].tolist()})
            last = M[j]
        oracle.append(orc.rollout(xs[b, H][None], acts[b, H - 1][None], acts[b:b + 1, H:]).numpy())
    oracle = np.concatenate(oracle)
    t_o = (time.perf_counter() - t0) / len(xs)
    fresh = Prober(ck["config"]["latent"], ck["config"]["prober_hidden"], ck["config"].get("l_swing", 0.3089),
                   ck["config"].get("omega_span", 0.3), ck["config"].get("a_max0", 0.25))   # the hand-written backbone alone
    with torch.no_grad():
        bb = integrate(torch.as_tensor(xs[:, H]), torch.as_tensor(acts[:, H:]), torch.as_tensor(acts[:, H - 1]),
                       torch.zeros(len(xs), L, ck["config"]["latent"]), fresh, ck["config"]["dt"], model.norm,
                       torch.as_tensor(M[traj])).numpy()
    truth = xs[:, -1]
    res = {}
    for name, pr in (("learned", learned[:, -1]), ("oracle", oracle[:, -1]), ("backbone", bb[:, -1])):
        e = np.linalg.norm(grasp_point(pr) - grasp_point(truth), axis=-1) * 1000
        ea = tilt_angle_err(pr, truth)
        eh = np.linalg.norm(pr[:, 0:3] - truth[:, 0:3], axis=-1) * 1000
        res[name] = {"g_med_mm": float(np.median(e)), "g_p90_mm": float(np.percentile(e, 90)),
                     "g_payload_med_mm": float(np.median(e[payload_w])), "tilt_med_deg": float(np.median(ea)),
                     "tilt_p90_deg": float(np.percentile(ea, 90)),
                     "hook_med_mm": float(np.median(eh)), "hook_event_med_mm": float(np.median(eh[event_w])) if event_w.any() else None,
                     "tilt_event_med_deg": float(np.median(ea[event_w])) if event_w.any() else None}
    print(f"\nTEST 1 · {args.data} · open loop, {len(xs)} held-out windows ({payload_w.sum()} with payload on or changing), "
          f"{len(set(map(tuple, M[traj].reshape(len(traj), -1).round(3))))} rooms, after {L} steps ({L * ck['config']['dt']:.1f} s)")
    print("(backbone = hand-written physics alone, nothing learned; oracle = MuJoCo in the same room, nominal physics, no payload)")
    print(f"{'model':9s} {'grasp pt median':>15s} {'p90':>8s} {'payload median':>15s} {'tilt median':>12s} {'p90':>7s}")
    for k, r in res.items():
        print(f"{k:9s} {r['g_med_mm']:12.1f} mm {r['g_p90_mm']:5.1f} mm {r['g_payload_med_mm']:12.1f} mm "
              f"{r['tilt_med_deg']:9.2f} deg {r['tilt_p90_deg']:4.2f}")
    r = res["learned"]
    if person:   # P1 (docs/MILESTONES.md): hook <= 20 mm, also on attach / set-down windows; swing <= 1.5 deg
        f1 = lambda v: "n/a" if v is None else f"{v:.1f}"
        print(f"\nP1 · {event_w.sum()} windows where the person's weight changes. At 2 s:")
        for k, v in res.items():
            print(f"  {k:9s} hook {v['hook_med_mm']:.1f} mm (events {f1(v['hook_event_med_mm'])}), swing "
                  f"{v['tilt_med_deg']:.2f} deg (events {f1(v['tilt_event_med_deg'])}), person COM {v['g_med_mm']:.1f} mm")
        ok = (r["hook_med_mm"] <= 20 and (r["hook_event_med_mm"] is None or r["hook_event_med_mm"] <= 20)
              and r["tilt_med_deg"] <= 1.5)
        print(f"P1 model bar (hook <= 20 mm, also on events; swing <= 1.5 deg): {'YES' if ok else 'NO'}")
        res["weigh"] = weigh_test(args, cfg, test_idx)
    else:
        ok = r["g_med_mm"] <= 20 and r["g_payload_med_mm"] <= 20 and r["tilt_med_deg"] <= 2
        print(f"pass (<= 20 mm, <= 20 mm on payload, <= 2 deg): {'YES' if ok else 'NO'}")
    print(f"learned {t_l * 1000:.1f} ms per window (S = 1), oracle {t_o * 1000:.1f} ms")
    Path(args.out).with_suffix(f".test1.{Path(str(args.data).split(',')[0]).name}.json").write_text(json.dumps(res, indent=1))


def weigh_test(args, cfg, test_idx):
    """P1's weighing bar: 1 s after every attach ends in the held-out trajectories, the lines' weighing
    (perception/weigh.py: mean vertical pull over that second / g - the rig's nominal mass) within 10 % of the person."""
    from cablepilot.perception.weigh import rig_mass, weigh_kg
    shards = sorted(Path(str(args.data).split(",")[0]).glob("shard_*.npz"))
    S, TEN, PAY, MO = [], [], [], []
    for f in shards:
        z = np.load(f)
        S.append(z["states"]), TEN.append(z["tensions"]), PAY.append(z["payload"]), MO.append(z["mounts"])
    S, TEN, PAY, MO = (np.concatenate(v) for v in (S, TEN, PAY, MO))
    rig, errs = rig_mass(cfg), []
    for j in test_idx:
        p = PAY[j]
        k = 1
        while k < len(p):
            if p[k - 1] < 1.0 <= p[k]:                       # an attach starts
                e = k
                while e + 1 < len(p) and abs(p[e + 1] - p[e]) > 1e-3:
                    e += 1                                    # ... and ends at e
                # an attach that ENDS with a person on (2026-10-12: one attach in the held-out set was cut off by a
                # set-down back to 0 kg, and the first run divided by that 0)
                if e + 10 < len(p) and p[e] >= 1.0 and abs(p[e + 10] - p[e]) < 1e-3:
                    w = slice(e + 1, e + 11)                  # the second after it ends
                    est = weigh_kg(TEN[j, w], S[j, w, 0:3], MO[j], rig)
                    errs.append(abs(est - p[e + 10]) / p[e + 10])
                k = e + 1
            k += 1
    errs = np.array(errs) * 100
    if not len(errs):
        print("weighing: no complete attach in the held-out set")
        return None
    r = {"attaches": int(len(errs)), "median_pct": float(np.median(errs)), "p90_pct": float(np.percentile(errs, 90)),
         "max_pct": float(errs.max()), "within_10pct": float((errs <= 10).mean())}
    print(f"P1 weighing bar (within 10 % 1 s after an attach ends): {r['attaches']} attaches, median {r['median_pct']:.1f} %, "
          f"p90 {r['p90_pct']:.1f} %, max {r['max_pct']:.1f} %, within 10 %: {100 * r['within_10pct']:.0f} %  "
          f"{'YES' if r['within_10pct'] == 1.0 else 'NO'}")
    return r


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", default=str(ROOT / "data/cdpr_v0_playroom"))
    ap.add_argument("--out", default=str(ROOT / "checkpoints/cdpr_wm_v0.pt"))
    ap.add_argument("--steps1", type=int, default=3000)
    ap.add_argument("--steps2", type=int, default=3000)
    ap.add_argument("--batch", type=int, default=256)
    ap.add_argument("--lr", type=float, default=3e-3)
    ap.add_argument("--test", type=int, default=500, help="last N trajectories held out")
    ap.add_argument("--n-eval", type=int, default=500)
    ap.add_argument("--test-all", action="store_true", help="eval-only: every trajectory of --data is held out (an unseen room)")
    ap.add_argument("--eval-only", action="store_true")
    ap.add_argument("--stage1-from", default=None, help="checkpoint whose stage 1 (JEPA, normaliser) is reused")
    ap.add_argument("--lr2", type=float, default=1e-3)
    args = ap.parse_args()
    if args.eval_only:
        if args.test_all:
            d = load(args.data)
            return test1(args, d, np.arange(len(d[0])))
        return test1(args)
    X, A, P, M, R, test_idx = train(args)
    test1(args, (X, A, P, M, R, float(np.load(sorted(Path(str(args.data).split(",")[0]).glob("shard_*.npz"))[0])["dt"])), test_idx)


if __name__ == "__main__":
    main()