Buckets:
| """cwm3/data/token_bank.py — per-route front (road) token bank from the dense V-JEPA2 cache + causal anchor gather indices. | |
| Source: s12_wm/cache/frontfeat/<uid>.npz {t f32[n], feats f16[n,1408], region f16[n,16,1408]} written by | |
| s12_wm/extract_crop.py --cam front. CAUSALITY RULE (derived from extract_crop.py, verified 2026-09-10): | |
| * `t` is the END time of each 2 s clip on a 0.5 s grid: clip_ts() returns k/2 for k in [ceil(2*max(2,t0)), floor(2*min(dur,t1))] | |
| (docstring "clip 覆盖 [ts-2, ts]"), and the 16 frames are decoded at frame index round((t + rel) * 20) with | |
| rel = np.linspace(-2.0, 0.0, 16) -> the window is [t-2.0, t], its LAST frame is exactly at t. Frame index / 20 is the | |
| frame-anchored CAN time axis of the CSVs (build_targets_s12 / canspec_v2: "CSV time_s == 帧号/20"), so bank time and | |
| anchor time share one axis. A bank row is causal for an anchor at time T iff t_row <= T (+1e-6). | |
| * Anchors sit on a 0.2 s grid, the bank on a 0.5 s grid: slot j (j = 0..15) of an anchor at T targets | |
| tau_j = floor(2T)/2 - (15 - j) * 0.5, i.e. the 16 most recent bank times <= T (span 7.5 s, "8 s @ 2 Hz"). All slots are | |
| therefore <= T by construction; the index builder asserts it and tests/test_bank_causality.py re-checks it on disk. | |
| The real lag tau_j - T is in [-7.9, 0] and is what the model receives as the Δt channel (cwm3_dataset.py). | |
| Bank layout (results/banks/front_tok.npz, uncompressed): | |
| tok f16[T_all,4,64] 4x4 region grid -> 2x2 average pool (4 spatial tokens x 1408) -> PCA-64 (fit on a manual-train sample) | |
| pooled256 f16[T_all,256] feats (clip-mean token, 1408) -> PCA-256 (fit on the same sample) | |
| t f32[T_all] clip END time (s) on the route's frame-anchored axis | |
| uid_index str[R] route uid per bank segment; route_off int64[R+1] segment offsets (rows of route r: [off[r], off[r+1])) | |
| Anchor gather: results/banks/anchor_index_{regime}_{split}.npy int32[N,16] (-1 = missing) + anchor_index_meta.jsonl. | |
| No openpilot DMS field is read. | |
| python token_bank.py --fit --build --index [--workers 8] | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import sys | |
| import time | |
| from multiprocessing import Pool | |
| from pathlib import Path | |
| import numpy as np | |
| CWM3 = Path(__file__).resolve().parents[1] | |
| EXPLORE = CWM3.parent | |
| S12R = EXPLORE / "s12_wm" / "results" | |
| FD = EXPLORE / "s12_wm" / "cache" / "frontfeat" | |
| PCA_DIR = CWM3 / "results" / "pca" | |
| BANK_DIR = CWM3 / "results" / "banks" | |
| sys.path.insert(0, str(CWM3 / "common")) | |
| from runmeta import write_run_json, sha256_file # noqa: E402 | |
| D_REG, K_TOK, K_POOL = 1408, 64, 256 | |
| N_SLOTS, SLOT_S = 16, 0.5 | |
| CLIP_SPAN = 2.0 | |
| CAUSALITY_RULE = ("bank row usable for anchor T iff t_row <= T + 1e-6, where t_row is the clip END time (window " | |
| "[t_row - 2.0, t_row]; extract_crop.py rel = linspace(-2, 0, 16), frames at round((t + rel) * 20)); slot j " | |
| "targets floor(2T)/2 - (15 - j) * 0.5") | |
| def pool_regions(region): | |
| """region [n,16,1408] (4x4 grid, row-major) -> [n,4,1408] 2x2 average pool (quadrants TL, TR, BL, BR).""" | |
| r = region.astype(np.float32).reshape(len(region), 2, 2, 2, 2, D_REG) # [n, gy_hi, gy_lo, gx_hi, gx_lo, D] | |
| return r.mean((2, 4)).reshape(len(region), 4, D_REG) | |
| def load_route(uid): | |
| with np.load(FD / f"{uid}.npz") as z: | |
| return z["t"].astype(np.float32), z["feats"], z["region"] | |
| def pca_fit(X, k): | |
| X = X.astype(np.float64) | |
| mu = X.mean(0) | |
| cov = (X - mu).T @ (X - mu) / len(X) | |
| evals, evecs = np.linalg.eigh(cov) | |
| order = np.argsort(evals)[::-1] | |
| evals, evecs = evals[order], evecs[:, order] | |
| return mu.astype(np.float32), evecs[:, :k].astype(np.float32), (evals[:k] / evals.sum()).astype(np.float32) | |
| # ----------------------------------------------------------------------------- fit | |
| def fit(n_tok, n_pool, seed, log): | |
| t0 = time.time() | |
| out_tok, out_pool = PCA_DIR / f"pca_tok_{K_TOK}.npz", PCA_DIR / f"pca_pooled_{K_POOL}.npz" | |
| if out_tok.exists() and out_pool.exists(): | |
| log("[fit] pca_tok / pca_pooled exist — skipped") | |
| return | |
| with np.load(S12R / "targets_manual_train.npz", allow_pickle=True) as z: | |
| uids = np.unique(z["uid"]) | |
| uids = [u for u in uids if (FD / f"{u}.npz").exists()] | |
| rng = np.random.default_rng(seed) | |
| rng.shuffle(uids) | |
| toks, pooled = [], [] | |
| n_t = n_p = 0 | |
| per_route = max(n_pool // 60, 500) # ~60 routes -> device / scene diversity | |
| for u in uids: | |
| t, feats, region = load_route(u) | |
| sel = np.sort(rng.choice(len(t), min(per_route, len(t)), replace=False)) | |
| toks.append(pool_regions(region[sel]).reshape(-1, D_REG)) | |
| pooled.append(feats[sel].astype(np.float32)) | |
| n_t += toks[-1].shape[0] | |
| n_p += len(sel) | |
| if n_t >= n_tok and n_p >= n_pool: | |
| break | |
| Xt = np.concatenate(toks)[:n_tok] | |
| Xp = np.concatenate(pooled)[:n_pool] | |
| mu_t, W_t, evr_t = pca_fit(Xt, K_TOK) | |
| mu_p, W_p, evr_p = pca_fit(Xp, K_POOL) | |
| stamp = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) | |
| np.savez(out_tok, mu=mu_t, W=W_t, evr=evr_t, n_fit=np.array(len(Xt)), seed=np.array(seed), created_utc=np.array(stamp), | |
| pooling=np.array("4x4 region grid -> 2x2 mean -> 4 tokens"), n_routes=np.array(len(toks))) | |
| np.savez(out_pool, mu=mu_p, W=W_p, evr=evr_p, n_fit=np.array(len(Xp)), seed=np.array(seed), created_utc=np.array(stamp), | |
| pooling=np.array("feats = clip-mean token (1408)"), n_routes=np.array(len(pooled))) | |
| with open(PCA_DIR / "pca_manifest.jsonl", "a") as f: | |
| for p, evr, n in ((out_tok, evr_t, len(Xt)), (out_pool, evr_p, len(Xp))): | |
| f.write(json.dumps({"file": p.name, "sha256": sha256_file(p), "stream": p.stem, "n_fit": n, "k": int(evr.shape[0]), | |
| "evr_sum": float(evr.sum()), "evr_first8": [float(v) for v in evr[:8]], "n_routes": len(toks), | |
| "seed": seed, "wall_s": round(time.time() - t0, 1)}) + "\n") | |
| log(f"[fit] tok: n={len(Xt)} evr64={evr_t.sum():.4f}; pooled: n={len(Xp)} evr256={evr_p.sum():.4f}; routes={len(toks)} {time.time() - t0:.0f}s") | |
| # ----------------------------------------------------------------------------- build | |
| _P = {} | |
| def _init_worker(): | |
| zt, zp = np.load(PCA_DIR / f"pca_tok_{K_TOK}.npz"), np.load(PCA_DIR / f"pca_pooled_{K_POOL}.npz") | |
| _P.update(mu_t=zt["mu"], W_t=zt["W"], mu_p=zp["mu"], W_p=zp["W"]) | |
| def _route_job(uid): | |
| t, feats, region = load_route(uid) | |
| n = len(t) | |
| tok = np.zeros((n, 4, K_TOK), np.float16) | |
| for i in range(0, n, 2048): | |
| r = pool_regions(region[i:i + 2048]) # [k,4,1408] | |
| tok[i:i + 2048] = ((r - _P["mu_t"]) @ _P["W_t"]).astype(np.float16) | |
| pooled = ((feats.astype(np.float32) - _P["mu_p"]) @ _P["W_p"]).astype(np.float16) | |
| assert np.all(np.diff(t) > 0), f"{uid}: bank times not strictly increasing" | |
| return uid, t, tok, pooled | |
| def build(workers, log): | |
| t0 = time.time() | |
| out = BANK_DIR / "front_tok.npz" | |
| if out.exists(): | |
| log("[build] front_tok.npz exists — skipped") | |
| return | |
| uids = sorted(p.stem for p in FD.glob("*.npz")) | |
| T, TOK, POOL, off, names = [], [], [], [0], [] | |
| with Pool(workers, initializer=_init_worker) as pool: | |
| for k, (uid, t, tok, pooled) in enumerate(pool.imap(_route_job, uids, chunksize=1)): | |
| names.append(uid) | |
| T.append(t) | |
| TOK.append(tok) | |
| POOL.append(pooled) | |
| off.append(off[-1] + len(t)) | |
| if k % 50 == 0: | |
| log(f"[build] {k + 1}/{len(uids)} routes, rows so far {off[-1]:,} ({time.time() - t0:.0f}s)") | |
| tok = np.concatenate(TOK) | |
| pooled = np.concatenate(POOL) | |
| t = np.concatenate(T) | |
| np.savez(out, tok=tok, pooled256=pooled, t=t, uid_index=np.array(names), route_off=np.array(off, np.int64), | |
| pca_tok_sha256=np.array(sha256_file(PCA_DIR / f"pca_tok_{K_TOK}.npz")), | |
| pca_pooled_sha256=np.array(sha256_file(PCA_DIR / f"pca_pooled_{K_POOL}.npz")), | |
| causality_rule=np.array(CAUSALITY_RULE), clip_span_s=np.array(CLIP_SPAN), stride_s=np.array(SLOT_S), | |
| source_dir=np.array(str(FD)), created_utc=np.array(time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()))) | |
| log(f"[build] front_tok.npz: routes={len(names)} T_all={len(t):,} tok={tok.nbytes / 2**30:.2f} GiB pooled={pooled.nbytes / 2**30:.2f} GiB {time.time() - t0:.0f}s") | |
| # ----------------------------------------------------------------------------- anchor index | |
| def anchor_index_for(bank_t, route_of, off, uid, t_anchor): | |
| """-> int32[N,16] bank rows for the 16 slots (or -1). Pure numpy; asserts t_bank[idx] <= t_anchor.""" | |
| N = len(t_anchor) | |
| idx = np.full((N, N_SLOTS), -1, np.int32) | |
| base = np.floor(t_anchor * 2 + 1e-6) / 2.0 # largest 0.5-grid time <= T | |
| order = np.argsort(uid, kind="stable") | |
| bounds = np.flatnonzero(np.r_[1, uid[order][1:] != uid[order][:-1], 1]) | |
| for a, b in zip(bounds[:-1], bounds[1:]): | |
| u = uid[order[a]] | |
| r = route_of.get(u) | |
| if r is None: | |
| continue | |
| lo, hi = off[r], off[r + 1] | |
| ts = bank_t[lo:hi].astype(np.float64) | |
| rows = order[a:b] | |
| for j in range(N_SLOTS): | |
| tau = base[rows] - (N_SLOTS - 1 - j) * SLOT_S | |
| k = np.searchsorted(ts, tau - 1e-3) | |
| k = np.clip(k, 0, len(ts) - 1) | |
| hit = np.abs(ts[k] - tau) < 1e-3 | |
| idx[rows[hit], j] = (lo + k[hit]).astype(np.int32) | |
| ok = idx >= 0 | |
| assert np.all(bank_t[idx[ok]].astype(np.float64) <= t_anchor[:, None].repeat(N_SLOTS, 1)[ok] + 1e-6), "non-causal gather" | |
| return idx | |
| def index(log): | |
| t0 = time.time() | |
| with np.load(BANK_DIR / "front_tok.npz", allow_pickle=True) as z: | |
| bank_t, names, off = z["t"], [str(u) for u in z["uid_index"]], z["route_off"] | |
| route_of = {u: i for i, u in enumerate(names)} | |
| for regime in ("manual", "engaged"): | |
| for split in ("train", "val", "test"): | |
| out = BANK_DIR / f"anchor_index_{regime}_{split}.npy" | |
| if out.exists(): | |
| log(f"[index] {out.name} exists — skipped") | |
| continue | |
| with np.load(S12R / f"targets_{regime}_{split}.npz", allow_pickle=True) as z: | |
| uid, t = z["uid"], z["t"].astype(np.float64) | |
| idx = anchor_index_for(bank_t, route_of, off, uid, t) | |
| np.save(out, idx) | |
| ok = idx >= 0 | |
| rec = {"file": out.name, "N": int(len(t)), "slot_coverage": [float(v) for v in ok.mean(0)], | |
| "rows_all16": float(ok.all(1).mean()), "rows_last_slot": float(ok[:, -1].mean()), "rows_none": float((~ok.any(1)).mean()), | |
| "max_lag_s": float(np.max(t[:, None] - bank_t[np.where(ok, idx, 0)].astype(np.float64), where=ok, initial=0.0)), | |
| "min_lag_last_slot_s": float(np.min((t - bank_t[np.where(ok[:, -1], idx[:, -1], 0)])[ok[:, -1]])) if ok[:, -1].any() else None, | |
| "causality_rule": CAUSALITY_RULE, "created_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())} | |
| with open(BANK_DIR / "anchor_index_meta.jsonl", "a") as f: | |
| f.write(json.dumps(rec) + "\n") | |
| log(f"[index] {out.name}: N={len(t):,} all16={rec['rows_all16']:.4f} last={rec['rows_last_slot']:.4f} none={rec['rows_none']:.4f} " | |
| f"max_lag={rec['max_lag_s']:.2f}s {time.time() - t0:.0f}s") | |
| def main(): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--fit", action="store_true") | |
| ap.add_argument("--build", action="store_true") | |
| ap.add_argument("--index", action="store_true") | |
| ap.add_argument("--n-tok", type=int, default=200_000) | |
| ap.add_argument("--n-pool", type=int, default=200_000) | |
| ap.add_argument("--seed", type=int, default=0) | |
| ap.add_argument("--workers", type=int, default=8) | |
| a = ap.parse_args() | |
| t_start = time.time() | |
| PCA_DIR.mkdir(parents=True, exist_ok=True) | |
| BANK_DIR.mkdir(parents=True, exist_ok=True) | |
| logp = BANK_DIR / "token_bank.log" | |
| def log(m): | |
| line = f"[{time.strftime('%H:%M:%S')}] {m}" | |
| print(line, flush=True) | |
| with open(logp, "a") as f: | |
| f.write(line + "\n") | |
| if a.fit: | |
| fit(a.n_tok, a.n_pool, a.seed, log) | |
| if a.build: | |
| build(a.workers, log) | |
| if a.index: | |
| index(log) | |
| write_run_json(BANK_DIR, "token_bank", [Path(__file__), EXPLORE / "s12_wm" / "extract_crop.py"], | |
| [p for p in (PCA_DIR / f"pca_tok_{K_TOK}.npz", PCA_DIR / f"pca_pooled_{K_POOL}.npz", BANK_DIR / "front_tok.npz") if p.exists()], | |
| seed=a.seed, t_start=t_start, input_listing=False, | |
| extra={"fit": a.fit, "build": a.build, "index": a.index, "n_routes_frontfeat": len(list(FD.glob("*.npz"))), | |
| "causality_rule": CAUSALITY_RULE, "dms_fields_used": []}) | |
| log("DONE") | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 13 kB
- Xet hash:
- 06c9958a94d1aa9145d9984cd3738fb0cb056bc3179a323ba0129bc180190704
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.