"""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()