Buckets:
| """cwm3/train/pretrain_u_trunk_aide.py — self-supervised pretraining of the U TRUNK on AIDE (claim C1, E10 (b)). | |
| The trunk is exactly the survivor DriverBranch of U0 kin_can_noK_clip8_ln (models.branches.DriverBranch: in_ln LayerNorm(12) -> GRU(18 -> 128) | |
| ⊕ mode embedding; built with build_branch("U", seed, hid 128, hidden 128, dropout 0, cabin none, in_ln True) so the seed-k initialisation equals | |
| the scratch baseline's). Task: from the AIDE kin11 history predict the next 1 s of the same 11 channels at the BATON K grid (5 steps, +0.2 s | |
| .. +1.0 s), as a residual to K_0 (the last valid history step) — the form of the BATON residual kinematics head. | |
| Inputs follow train/batch.prep_U's [B,40,18] layout bit for bit: [11 standardised Kh (clip ±8, NaN -> 0), valid bit, 5 CAN channels, Δt]: | |
| * history = the 2 s (20 steps) ending at the anchor, right-aligned in the 40-step window; earlier steps carry valid 0 and zeros | |
| (--hist-steps); Δt = -(39-k)*0.1 as in BATON; | |
| * the 5 CAN channels are ZERO (AIDE has no CAN; 0 = the standardised mean state) — there is no separate CAN valid bit in the contract; | |
| * mode = 2 ("unknown", the third embedding row; BATON uses 0 / 1) — the mode embedding rows the fine-tune uses keep their seed init | |
| (up to AdamW decay); | |
| * standardisation = AIDE TRAIN kmu / ksd stored in aide_kin11.npz (deviation: BATON standardises with its own stats v2; the input LayerNorm | |
| of the survivor recipe makes the trunk insensitive to the per-channel scale, docs/roadmap_2026-09-11.md §三). | |
| Anchors: every step j of a clip with >= --min-hist valid-eligible steps of history (j >= min_hist - 1) and >= 1 future step inside the clip; | |
| future targets are masked where the step is outside the clip or invalid. Fit on the AIDE train split, early stop on val (test unused). | |
| Loss (--loss laplace | l1): masked mean Laplace NLL |y - mu| / b + log b with a learned per-cell log-scale (l1: |y - mu| only). | |
| Recipe: AdamW lr 1e-3, wd 1e-4, batch 512, <= 60 epochs, patience 6 on the val loss, grad clip 1, bf16 autocast on CUDA. | |
| Output results/runs/U0_aidepre_trunk_s<k>/: trunk.pt in the train/train_u0_from_init.py format ({"gru", "mode", "in_ln"} state_dicts + | |
| "pretrain_head", epoch, val_loss, seed, recipe, model, stats (AIDE kmu/ksd), source, channels), history.jsonl, BEST, val_metrics.json | |
| (per-channel / per-horizon MAE in raw units vs the copy-K0 and constant-mean predictors), run json. No BATON file, no DMS field. | |
| GPU etiquette: no lock (coordinator 2026-09-10), threads <= 4, VRAM / RAM / trainer-count checks; --cpu runs the same on the host. | |
| python train/pretrain_u_trunk_aide.py --seed 0 [--epochs E] [--cpu] [--out DIR] | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import os | |
| import re | |
| import sys | |
| import time | |
| from pathlib import Path | |
| os.environ.setdefault("OMP_NUM_THREADS", "4") | |
| os.environ.setdefault("MKL_NUM_THREADS", "4") | |
| import numpy as np # noqa: E402 | |
| CWM3 = Path(__file__).resolve().parents[1] | |
| for p in (CWM3, CWM3 / "train", CWM3 / "common", CWM3 / "data"): | |
| sys.path.insert(0, str(p)) | |
| import batch as bt # noqa: E402 (X_CLIP / K_CLIP, HIST, DK — the BATON input contract) | |
| import cwm3_dataset as ds # noqa: E402 (KIN_CHANNELS: the pack channel order) | |
| from models import build_branch, params_by_module # noqa: E402 | |
| from models.branches import flat_head # noqa: E402 | |
| from runmeta import write_run_json, sha256_file # noqa: E402 | |
| CACHE = CWM3 / "aide" / "results" / "cache" / "aide_kin11.npz" | |
| HIST, DK = bt.HIST, bt.DK | |
| N_FUT, FUT_STRIDE = 5, 2 # +0.2 .. +1.0 s at 10 Hz | |
| MODE_UNKNOWN = 2 | |
| RECIPE = dict(batch=512, max_epochs=60, patience=6, lr=1e-3, weight_decay=1e-4, grad_clip=1.0, amp=True, hist_steps=20, min_hist=10, loss="laplace", | |
| kclip=float(bt.K_CLIP)) | |
| MODEL = dict(hid=128, hidden=128, dropout=0.0, cabin="none", in_ln=True) | |
| MIN_FREE_VRAM_GB, MIN_AVAIL_RAM_GB, MAX_TRAINERS = 8.0, 14.0, 2 | |
| def log(msg, out): | |
| line = f"[{time.strftime('%H:%M:%S')}] {msg}" | |
| print(line, flush=True) | |
| with open(Path(out) / "train.log", "a") as f: | |
| f.write(line + "\n") | |
| def live_trainers(): | |
| me = os.getpid() | |
| n = [] | |
| for pid in os.listdir("/proc"): | |
| if not pid.isdigit() or int(pid) == me: | |
| continue | |
| try: | |
| cmd = open(f"/proc/{pid}/cmdline", "rb").read().replace(b"\0", b" ").decode(errors="ignore") | |
| except OSError: | |
| continue | |
| if (re.search(r"train_cwm3(_joint)?\.py|train_u0_from_init\.py|dms_teacher_pretrain\.py|pretrain_u_trunk_aide\.py", cmd) and "pgrep" not in cmd | |
| and re.match(r"\S*python[\d.]*\s", cmd)): # the interpreter process only (not the `bash -c` shell that launched it) | |
| n.append((int(pid), cmd[:120])) | |
| return n | |
| def check_resources(out, cpu): | |
| avail_gb = float([l for l in open("/proc/meminfo") if l.startswith("MemAvailable")][0].split()[1]) / 2**20 | |
| alive = live_trainers() | |
| msg = f"resources: RAM available {avail_gb:.1f} GB, trainers alive {len(alive)}" | |
| if not cpu: | |
| import torch | |
| free_gb = torch.cuda.mem_get_info()[0] / 2**30 | |
| msg += f", VRAM free {free_gb:.1f} GB" | |
| assert free_gb >= MIN_FREE_VRAM_GB, f"free VRAM {free_gb:.1f} GB < {MIN_FREE_VRAM_GB}" | |
| log(msg, out) | |
| assert avail_gb >= MIN_AVAIL_RAM_GB, f"available RAM {avail_gb:.1f} GB < {MIN_AVAIL_RAM_GB}" | |
| assert len(alive) < MAX_TRAINERS, f"{len(alive)} trainers alive (limit {MAX_TRAINERS}): {[c for _, c in alive]}" | |
| def build_samples(z, splits, hist_steps, min_hist): | |
| """AIDE kin11 -> anchors. -> dict(Kh f32[S,40,11] NaN-padded raw, Khv bool[S,40], Y f32[S,5,11] raw, Yv bool[S,5], clip i32[S], j i32[S]).""" | |
| K, V, n_steps, split = z["K"], z["valid"], z["n_steps"], np.array([str(s) for s in z["split"]]) | |
| rows = np.flatnonzero(np.isin(split, splits)) | |
| Kh, Khv, Y, Yv, ci, jj = [], [], [], [], [], [] | |
| for i in rows: | |
| n = int(n_steps[i]) | |
| for j in range(min_hist - 1, n - FUT_STRIDE): | |
| lo = j - hist_steps + 1 | |
| h = np.full((HIST, DK), np.nan, np.float32) | |
| hv = np.zeros(HIST, bool) | |
| src = np.arange(max(lo, 0), j + 1) | |
| dst = src - (j - HIST + 1) | |
| h[dst] = K[i, src] | |
| hv[dst] = V[i, src] | |
| if hv.sum() < min_hist // 2 or not hv[-1]: # need the anchor step (K_0) and at least half the minimal history | |
| continue | |
| f = j + FUT_STRIDE * np.arange(1, N_FUT + 1) | |
| ok = f < n | |
| y = np.full((N_FUT, DK), np.nan, np.float32) | |
| yv = np.zeros(N_FUT, bool) | |
| y[ok] = K[i, f[ok]] | |
| yv[ok] = V[i, f[ok]] | |
| if not yv.any(): | |
| continue | |
| Kh.append(h); Khv.append(hv); Y.append(y); Yv.append(yv); ci.append(i); jj.append(j) | |
| return {"Kh": np.stack(Kh), "Khv": np.stack(Khv), "Y": np.stack(Y), "Yv": np.stack(Yv), "clip": np.array(ci, np.int32), "j": np.array(jj, np.int32)} | |
| def to_inputs(S, kmu, ksd, kclip, device): | |
| """samples -> tensors: Kx [S,40,18] (prep_U layout, CAN zero), mode [S], K0std [S,11], Ystd [S,5,11], Ym [S,5,11].""" | |
| import torch | |
| Kh = torch.as_tensor(S["Kh"], device=device) | |
| kv = torch.as_tensor(S["Khv"], device=device) | |
| mu, sd = torch.as_tensor(kmu, device=device), torch.as_tensor(ksd, device=device) | |
| Khs = torch.nan_to_num(torch.clamp((Kh - mu) / sd, -kclip, kclip)) * kv[..., None].float() | |
| dt = (-(HIST - 1 - torch.arange(HIST, device=device).float()) * 0.1)[None, :, None].expand(len(Kh), -1, -1) | |
| Xs = torch.zeros(len(Kh), HIST, 5, device=device) # no CAN on AIDE | |
| Kx = torch.cat([Khs, kv[..., None].float(), Xs, dt], -1) | |
| assert Kx.shape[-1] == 18 | |
| Y = torch.as_tensor(S["Y"], device=device) | |
| Yv = torch.as_tensor(S["Yv"], device=device) | |
| Ym = Yv[..., None] & torch.isfinite(Y) | |
| Ystd = torch.nan_to_num(torch.clamp((Y - mu) / sd, -kclip, kclip)) * Ym.float() | |
| K0std = Khs[:, -1, :] # anchor step always valid by construction | |
| return {"Kx": Kx, "mode": torch.full((len(Kh),), MODE_UNKNOWN, dtype=torch.long, device=device), "K0std": K0std, "Ystd": Ystd, "Ym": Ym} | |
| class Pretrainer: | |
| def __init__(self, seed, device): | |
| import torch | |
| self.base = build_branch("U", seed, **MODEL).to(device) # torch.manual_seed(seed) inside -> trunk init == U0 seed k (in_ln, gru, mode) | |
| torch.manual_seed(seed + 1) | |
| self.head = flat_head(MODEL["hid"] + 16, MODEL["hidden"], N_FUT * DK * 2, 0.0).to(device) | |
| def parameters(self): | |
| return list(self.base.in_ln.parameters()) + list(self.base.gru.parameters()) + list(self.base.mode.parameters()) + list(self.head.parameters()) | |
| def train(self): | |
| self.base.train(); self.head.train() | |
| def eval(self): | |
| self.base.eval(); self.head.eval() | |
| def forward(self, Kx, mode): | |
| import torch | |
| kx = torch.cat([self.base.in_ln(Kx[..., :self.base.N_KIN_IN]), Kx[..., self.base.N_KIN_IN:]], -1) # DriverBranch.forward's trunk path | |
| _, h = self.base.gru(kx) | |
| z = torch.cat([h[-1], self.base.mode(mode)], 1) | |
| o = self.head(z).float().view(-1, N_FUT, DK, 2) | |
| return o[..., 0], o[..., 1].clamp(-6, 4) # d_mu (residual on K0std), log_b | |
| def trunk_state(self): | |
| sd = lambda m: {k: v.detach().cpu() for k, v in m.state_dict().items()} | |
| return {"in_ln": sd(self.base.in_ln), "gru": sd(self.base.gru), "mode": sd(self.base.mode), "pretrain_head": sd(self.head)} | |
| def loss_fn(d_mu, log_b, T, kind): | |
| import torch | |
| mu = T["K0std"][:, None, :] + d_mu | |
| m = T["Ym"].float() | |
| err = (T["Ystd"] - mu).abs() | |
| per = err / torch.exp(log_b) + log_b if kind == "laplace" else err | |
| return (per * m).sum() / m.sum().clamp_min(1.0), err | |
| def evaluate(model, T, kind, ksd, bs=4096): | |
| """-> val loss, MAE raw units [5,11] (model), copy-K0 baseline, constant-mean baseline.""" | |
| import torch | |
| model.eval() | |
| tot, n = 0.0, 0.0 | |
| e_m = torch.zeros(N_FUT, DK, device=T["Kx"].device); e_c = torch.zeros_like(e_m); e_0 = torch.zeros_like(e_m); w = torch.zeros_like(e_m) | |
| sd = torch.as_tensor(ksd, device=T["Kx"].device) | |
| with torch.no_grad(): | |
| for i in range(0, len(T["Kx"]), bs): | |
| sl = slice(i, i + bs) | |
| Tb = {k: v[sl] for k, v in T.items()} | |
| d_mu, log_b = model.forward(Tb["Kx"], Tb["mode"]) | |
| l, err = loss_fn(d_mu, log_b, Tb, kind) | |
| m = Tb["Ym"].float() | |
| tot += float(l) * float(m.sum()); n += float(m.sum()) | |
| e_m += (err * m).sum(0) * sd | |
| e_c += ((Tb["Ystd"] - Tb["K0std"][:, None, :]).abs() * m).sum(0) * sd | |
| e_0 += (Tb["Ystd"].abs() * m).sum(0) * sd # constant predictor = the standardised mean (0) | |
| w += m.sum(0) | |
| w = w.clamp_min(1.0) | |
| return tot / max(n, 1.0), (e_m / w).cpu().numpy(), (e_c / w).cpu().numpy(), (e_0 / w).cpu().numpy() | |
| def main(): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--seed", type=int, default=0) | |
| ap.add_argument("--epochs", type=int, default=0) | |
| ap.add_argument("--out", default="") | |
| ap.add_argument("--cpu", action="store_true") | |
| ap.add_argument("--loss", default=RECIPE["loss"], choices=["laplace", "l1"]) | |
| ap.add_argument("--hist-steps", type=int, default=RECIPE["hist_steps"]) | |
| ap.add_argument("--fit-splits", default="train") | |
| ap.add_argument("--val-split", default="val") | |
| a = ap.parse_args() | |
| import torch | |
| torch.set_num_threads(4) | |
| out = Path(a.out) if a.out else CWM3 / "results" / "runs" / f"U0_aidepre_trunk_s{a.seed}" | |
| out.mkdir(parents=True, exist_ok=True) | |
| if (out / "TRAIN_DONE").exists(): | |
| print(f"{out} already done") | |
| return | |
| t0 = time.time() | |
| rec = dict(RECIPE, loss=a.loss, hist_steps=a.hist_steps, max_epochs=a.epochs or RECIPE["max_epochs"]) | |
| log(f"=== pretrain_u_trunk_aide seed {a.seed} out={out} recipe={rec} model={MODEL} ===", out) | |
| z = np.load(CACHE, allow_pickle=True) | |
| assert [str(c) for c in z["channels"]] == list(ds.KIN_CHANNELS), "AIDE kin11 channel order != BATON pack contract" | |
| kmu, ksd = z["kmu"].astype(np.float32), z["ksd"].astype(np.float32) | |
| S_tr = build_samples(z, a.fit_splits.split(","), rec["hist_steps"], rec["min_hist"]) | |
| S_va = build_samples(z, [a.val_split], rec["hist_steps"], rec["min_hist"]) | |
| log(f"samples: fit {len(S_tr['clip']):,} anchors / {len(np.unique(S_tr['clip']))} clips ({a.fit_splits}); val {len(S_va['clip']):,} / {len(np.unique(S_va['clip']))} ({a.val_split}); " | |
| f"future steps valid {S_tr['Yv'].mean():.3f}; history steps valid {S_tr['Khv'][:, -rec['hist_steps']:].mean():.3f}", out) | |
| device = "cpu" if a.cpu else "cuda" | |
| check_resources(out, a.cpu) | |
| amp = device == "cuda" and rec["amp"] | |
| T_tr, T_va = to_inputs(S_tr, kmu, ksd, rec["kclip"], device), to_inputs(S_va, kmu, ksd, rec["kclip"], device) | |
| assert (T_tr["Kx"][..., 12:17] == 0).all(), "CAN channels must be zero on AIDE" | |
| model = Pretrainer(a.seed, device) | |
| n_par = {"in_ln": sum(p.numel() for p in model.base.in_ln.parameters()), "gru": sum(p.numel() for p in model.base.gru.parameters()), | |
| "mode": sum(p.numel() for p in model.base.mode.parameters()), "pretrain_head": sum(p.numel() for p in model.head.parameters())} | |
| opt = torch.optim.AdamW(model.parameters(), lr=rec["lr"], weight_decay=rec["weight_decay"]) | |
| log(f"model: trunk params {n_par}; device {device} amp {amp}", out) | |
| N, bs = len(T_tr["Kx"]), rec["batch"] | |
| best, best_ep, patience = 1e18, -1, 0 | |
| for ep in range(rec["max_epochs"]): | |
| t_ep = time.time() | |
| rng = np.random.default_rng(a.seed * 1000 + ep) | |
| order = torch.as_tensor(rng.permutation(N), device=device) | |
| model.train() | |
| tot, nb = 0.0, 0 | |
| for i in range(0, N, bs): | |
| idx = order[i:i + bs] | |
| Tb = {k: v[idx] for k, v in T_tr.items()} | |
| with torch.autocast("cuda", dtype=torch.bfloat16, enabled=amp): | |
| d_mu, log_b = model.forward(Tb["Kx"], Tb["mode"]) | |
| loss, _ = loss_fn(d_mu, log_b, Tb, rec["loss"]) | |
| opt.zero_grad(set_to_none=True) | |
| loss.backward() | |
| torch.nn.utils.clip_grad_norm_(model.parameters(), rec["grad_clip"]) | |
| opt.step() | |
| tot += float(loss.detach()); nb += 1 | |
| vl, mae, mae_copy, mae_const = evaluate(model, T_va, rec["loss"], ksd) | |
| improved = vl < best - 1e-5 | |
| if improved: | |
| best, best_ep, patience = vl, ep, 0 | |
| torch.save({**model.trunk_state(), "epoch": ep, "val_loss": vl, "seed": a.seed, "recipe": rec, "model": MODEL, | |
| "stats": {"kmu": kmu, "ksd": ksd, "stats_source": "aide_kin11.npz AIDE train split"}, "source": "train/pretrain_u_trunk_aide.py", | |
| "channels": [str(c) for c in z["channels"]], "dms_fields_used": [], "cache_sha256": sha256_file(CACHE), | |
| "val_metrics": {"mae_model": mae.tolist(), "mae_copy_k0": mae_copy.tolist(), "mae_const": mae_const.tolist()}}, out / "trunk.pt") | |
| (out / "BEST").write_text(f"trunk.pt epoch {ep}\n") | |
| else: | |
| patience += 1 | |
| h = {"epoch": ep, "train_loss": tot / max(nb, 1), "val_loss": vl, "best": best, "best_epoch": best_ep, "mae_model_mean": float(mae.mean()), | |
| "mae_copy_k0_mean": float(mae_copy.mean()), "mae_const_mean": float(mae_const.mean()), "t_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())} | |
| with open(out / "history.jsonl", "a") as f: | |
| f.write(json.dumps(h) + "\n") | |
| log(f"ep {ep}: train {h['train_loss']:.4f} val {vl:.4f} best {best:.4f}@{best_ep} MAE model {h['mae_model_mean']:.4f} copy-K0 {h['mae_copy_k0_mean']:.4f} " | |
| f"const {h['mae_const_mean']:.4f} ({time.time() - t_ep:.0f}s)", out) | |
| if patience >= rec["patience"]: | |
| break | |
| zb = torch.load(out / "trunk.pt", map_location="cpu", weights_only=False) | |
| vm = zb["val_metrics"] | |
| chans = [str(c) for c in z["channels"]] | |
| (out / "val_metrics.json").write_text(json.dumps({"best_epoch": zb["epoch"], "val_loss": zb["val_loss"], "channels": chans, "horizons_s": [0.2 * k for k in range(1, N_FUT + 1)], | |
| **vm, "per_channel_mean_over_horizons": {c: {"model": float(np.mean(vm["mae_model"], 0)[j]), "copy_k0": float(np.mean(vm["mae_copy_k0"], 0)[j]), | |
| "const": float(np.mean(vm["mae_const"], 0)[j])} for j, c in enumerate(chans)}}, indent=1)) | |
| (out / "TRAIN_DONE").write_text(json.dumps({"best_epoch": zb["epoch"], "best": zb["val_loss"], "epochs_run": ep + 1, "params": n_par})) | |
| write_run_json(out, f"pretrain_u_trunk_aide_s{a.seed}", [Path(__file__), CWM3 / "models" / "branches.py", CWM3 / "train" / "batch.py"], [CACHE], seed=a.seed, t_start=t0, | |
| extra={"recipe": rec, "model": MODEL, "params": n_par, "n_fit": int(len(S_tr["clip"])), "n_val": int(len(S_va["clip"])), "fit_splits": a.fit_splits, "val_split": a.val_split, | |
| "best_epoch": zb["epoch"], "val_loss": zb["val_loss"], "epochs_run": ep + 1, "device": device, "dms_fields_used": [], | |
| "inputs_note": "AIDE kin11 only (no BATON file); CAN channels zero; mode 2 (unknown); AIDE-train standardisation", "gpu_lock": "none (--no-lock convention)"}) | |
| log(f"done: best epoch {zb['epoch']} val {zb['val_loss']:.4f}; trunk.pt sha256 {sha256_file(out / 'trunk.pt')[:16]} ({time.time() - t0:.0f}s)", out) | |
| print("PRETRAIN_AIDE_DONE", flush=True) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 17.9 kB
- Xet hash:
- cd9d214b4cf5325cb0556d9aa1888fecdf5f166131c1c0b292c3296c9a1950f3
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.