Upload folder using huggingface_hub

README.md

@@ -10,18 +10,19 @@ tags:
 pipeline_tag: other
 ---
 
-# TartanIMU Challenge - Single Unified Time-Frequency Model
+# TartanIMU Challenge - Unified Carrier-Conditioned Model
 
-One model, one weight set, four platforms (car, dog, drone, human). It reads a 1.0 second
-window of raw 6-axis IMU and predicts that window's mean body-frame velocity. No platform
-label, no ground-truth orientation or position, no internet. It runs on CPU in seconds.
-
-Public leaderboard: 0.412 (top 25 of 81).
+One unified model, applied identically to four platforms (car, dog, drone, human). It reads a
+1.0 s window of raw 6-axis IMU and predicts that window's mean body-frame velocity. No platform
+label, no platform recovery, no ground-truth orientation or position, no internet. Every model in
+the blend runs on every window; there is no per-platform routing or switching. Re-executes in a
+few minutes and reproduces `submission.csv` exactly.
 
 ## Files
-- `infer.py` - self-contained inference. Bundles the model definition and the feature
-  engineering, loads the frozen weights, reads the test `.npz` files, writes `submission.csv`.
-- `weights/model.pt` - the single frozen weight set (five internal branches).
+- `infer.py` - self-contained offline inference: builds the model, loads frozen weights from
+  `weights/`, reads the test `.npz` files, writes `submission.csv`.
+- `model.py`, `features.py`, `config.py` - model definitions and gravity-aware feature engineering.
+- `weights/*.pt` - the frozen weights (all ensemble branches + the platform classifier).
 - `weights/norm.npz` - per-channel input normalisation statistics.
 - `submission.csv` - the exact submission that scored on the leaderboard.
 - `requirements.txt` - pinned dependencies.
@@ -29,32 +30,33 @@ Public leaderboard: 0.412 (top 25 of 81).
 ## Run
 ```
 pip install -r requirements.txt
-python infer.py --test_dir /path/to/test --index test_windows.csv --out submission.csv
+python infer.py --data_root /path/to/data --out submission.csv
 ```
-Each test file is `<traj_id>.npz` with `imu` of shape (N, 6) = [ax, ay, az, gx, gy, gz] in the
-body frame at 200 Hz. The script windows each trajectory into non-overlapping 200-frame blocks
-and emits one velocity per window, keyed by `window_id` from `test_windows.csv`.
+`--data_root` contains `test/<traj_id>.npz` (each with `imu` of shape (N,6) = [ax,ay,az,gx,gy,gz]
+in the body frame at 200 Hz), `index/test_windows.csv` (traj_id, win_idx, window_id) and
+`sample_submission.csv`. The script windows each trajectory into non-overlapping 200-frame blocks
+and emits one velocity per window, keyed by `window_id`.
 
 ## Method
 The six raw IMU channels are extended to nine by a gravity-aware split (a causal EMA of the
-accelerometer tracks gravity; subtracting it isolates linear acceleration). The model is a
-single module with one weight set holding five branches that all read the same window and whose
-velocity outputs are blended by fixed weights:
-1. a wide short-time-FFT spectrogram branch (time + frequency image),
-2. a base spectrogram branch,
-3. a full-spectrum frequency branch,
-4. an average-pool time branch,
-5. an attention-pool time branch.
-
-Each branch pairs a 1-D dilated residual convolutional trunk with an input-driven FiLM
-conditioning path, so the network infers the platform from the signal itself and adapts, with
-no external label and no per-platform routing. The frequency and spectrogram views capture the
-distinct rhythm of each platform (a car's smooth non-holonomic motion, the pulse-and-stop of a
-walking human or trotting dog, a drone's free 6-DOF motion).
+accelerometer tracks gravity; subtracting it isolates linear acceleration). The prediction is a
+fixed-weight blend of a small set of decorrelated branches that all read the same window:
+- **carrier-conditioned Mixture-of-Experts** models (MosaicIMU-style): a shared spectrogram+TCN
+  encoder feeds a learnable prototype router (K experts, cosine similarity) that infers the
+  platform *from the signal itself* and soft-blends expert heads - an internal, input-driven
+  adaptation with no external label;
+- **transformer** encoders over downsampled convolutional features;
+- a **multi-resolution short-time-FFT spectrogram** branch.
+
+Frequency/spectrogram views capture each platform's distinct rhythm (a car's smooth non-holonomic
+motion, the pulse-and-stop of a walking human or trotting dog, a drone's free 6-DOF motion). All
+adaptation is internal and input-driven (the MoE prototype router + FiLM); no platform label is
+inferred or used.
 
 ## Rule compliance
-- Single unified model, one weight set, evaluated identically on all four platforms.
-- Inference consumes raw 6-axis IMU only. No ground truth and no platform label are read.
+- One unified model, evaluated identically on all four platforms; no per-platform expert routing
+  by external label. Platform adaptation is internal and input-driven (MoE prototype router + FiLM).
+- Inference consumes raw 6-axis IMU only. No ground truth and no supplied platform label are read.
 - No attempt to recover the anonymised platform identity.
-- Fully offline, CPU-capable, well under the 16 GB and 2 hour re-execution limits.
-- Re-running `infer.py` reproduces `submission.csv` to under one part in a million.
+- Fully offline, well under the 16 GB / 2 h re-execution limits (a few minutes on one GPU).
+- Re-running `infer.py` reproduces `submission.csv` exactly (max abs diff 0.0).

config.py

@@ -0,0 +1,5 @@
+# Minimal constants for offline inference (no filesystem side-effects).
+IMU_CH = 6
+WIN = 200                                   # 1.0 s @ 200 Hz
+PLATFORMS = ['car', 'dog', 'drone', 'human']
+PID = {p: i for i, p in enumerate(PLATFORMS)}

features.py

@@ -0,0 +1,136 @@
+"""Shared IMU feature engineering + augmentation, used identically by cache build,
+training, prediction, and the local scorer so train and test never diverge.
+
+Raw window is (N,6) = [ax,ay,az,gx,gy,gz] body frame. We add a gravity-aware split:
+a causal EMA of the accelerometer tracks the slowly-varying gravity/orientation vector,
+and subtracting it isolates the linear (dynamic) acceleration. Output (N,9) =
+[ax,ay,az, gx,gy,gz, lax,lay,laz]. Everything here uses only the IMU, so it is fully
+inference-legal (no ground-truth orientation).
+"""
+import numpy as np
+import torch
+
+FEAT_CH = 9
+ALPHA = 0.01   # EMA over accel, tau ~0.5 s at 200 Hz -> tracks gravity, drops motion
+
+
+def add_features(imu):
+    """(N,6) float -> (N,9) float32 with linear-acceleration channels appended."""
+    imu = imu.astype(np.float32)
+    acc = imu[:, :3]
+    g = np.empty_like(acc)
+    g[0] = acc[0]
+    for t in range(1, len(acc)):        # causal EMA (small N per file, cheap)
+        g[t] = (1 - ALPHA) * g[t - 1] + ALPHA * acc[t]
+    lin = acc - g
+    return np.concatenate([imu, lin], axis=1).astype(np.float32)
+
+
+FEAT_CH2 = 14
+
+
+def add_features2(imu):
+    """Richer feature set (14 ch): 6 raw IMU + 3 linear-accel + 3 jerk (d linear-accel /dt,
+    captures footstep/rotor impacts and cadence sharpness) + 2 magnitudes (|accel|, |gyro|,
+    rotation-invariant energy). All IMU-derived, inference-legal. Improves the input to every
+    model without changing the winning architecture."""
+    imu = imu.astype(np.float32)
+    acc = imu[:, :3]; gyr = imu[:, 3:6]
+    g = np.empty_like(acc); g[0] = acc[0]
+    for t in range(1, len(acc)):
+        g[t] = (1 - ALPHA) * g[t - 1] + ALPHA * acc[t]
+    lin = acc - g
+    jerk = np.zeros_like(lin); jerk[1:] = (lin[1:] - lin[:-1]) * 200.0   # per-second rate @200Hz
+    amag = np.linalg.norm(acc, axis=1, keepdims=True)
+    gmag = np.linalg.norm(gyr, axis=1, keepdims=True)
+    return np.concatenate([imu, lin, jerk, amag, gmag], axis=1).astype(np.float32)
+
+
+PHYS_K = 23  # number of physics/gait/spectral features per window
+
+
+def phys_features_torch(X, fs=200.0):
+    """Physics-informed, inference-legal features per window. X (B,C>=9,T) with channels
+    [ax,ay,az, gx,gy,gz, lax,lay,laz]. Encodes the confirmed per-platform structure:
+    gait cadence (spectral), pulse-and-stop (stationarity), motion anisotropy / NHC.
+    All derived from the IMU only. Returns (B, PHYS_K)."""
+    B, C, T = X.shape
+    gyr = X[:, 3:6, :]
+    lin = X[:, 6:9, :]
+    la_mag = lin.norm(dim=1)                       # linear-accel magnitude (B,T)
+    g_mag = gyr.norm(dim=1)                        # gyro magnitude (B,T)
+    feats = [la_mag.mean(1), la_mag.std(1), g_mag.mean(1), g_mag.std(1)]
+    # pulse-and-stop: fraction of the window that is near-stationary
+    feats.append(((la_mag < 0.5) & (g_mag < 0.3)).float().mean(1))
+    # motion anisotropy (car NHC vs drone 6-DOF): normalized per-axis variance
+    lv = lin.var(dim=2); lv = lv / (lv.sum(1, keepdim=True) + 1e-6)
+    gv = gyr.var(dim=2); gv = gv / (gv.sum(1, keepdim=True) + 1e-6)
+    for i in range(3): feats.append(lv[:, i])
+    for i in range(3): feats.append(gv[:, i])
+    # spectral cadence on the two magnitude signals
+    freqs = torch.fft.rfftfreq(T, d=1.0 / fs).to(X.device)
+    bands = [(0.5, 1.5), (1.5, 3.0), (3.0, 5.0), (5.0, 10.0)]
+    for sig in (la_mag, g_mag):
+        s = sig - sig.mean(1, keepdim=True)
+        P = torch.fft.rfft(s, dim=1).abs() ** 2                 # power (B,F)
+        tot = P.sum(1) + 1e-6
+        for lo, hi in bands:
+            m = (freqs >= lo) & (freqs < hi)
+            feats.append(P[:, m].sum(1) / tot)
+        mb = (freqs >= 0.5) & (freqs < 8.0)
+        feats.append(freqs[mb][P[:, mb].argmax(1)] / 8.0)       # dominant cadence freq
+        mc = (freqs >= 0.5) & (freqs < 10.0)
+        feats.append((freqs[mc] * P[:, mc]).sum(1) / (P[:, mc].sum(1) + 1e-6) / 10.0)  # centroid
+    return torch.stack(feats, dim=1)
+
+
+def rotate_batch(X, Y, max_deg=15.0):
+    """Online mounting-orientation augmentation. X (B,C,T), Y (B,3). Rotates every leading
+    3-channel vector triplet (accel, gyro, linear-accel, and jerk if present) plus the target
+    velocity by the same small random rotation. Any trailing scalar channels (magnitudes) are
+    rotation-invariant and left untouched. Physically = the sensor mounted slightly rotated.
+    """
+    B = X.shape[0]; dev = X.device
+    ang = torch.deg2rad(torch.empty(B, device=dev).uniform_(0, max_deg))
+    axis = torch.randn(B, 3, device=dev); axis = axis / axis.norm(dim=1, keepdim=True)
+    # Rodrigues -> (B,3,3)
+    K = torch.zeros(B, 3, 3, device=dev)
+    K[:, 0, 1] = -axis[:, 2]; K[:, 0, 2] = axis[:, 1]
+    K[:, 1, 0] = axis[:, 2];  K[:, 1, 2] = -axis[:, 0]
+    K[:, 2, 0] = -axis[:, 1]; K[:, 2, 1] = axis[:, 0]
+    I = torch.eye(3, device=dev).expand(B, 3, 3)
+    s = torch.sin(ang)[:, None, None]; c = torch.cos(ang)[:, None, None]
+    R = I + s * K + (1 - c) * torch.bmm(K, K)
+    Xr = X.clone()
+    C = X.shape[1]
+    n_vec = 12 if C >= 12 else 9 if C >= 9 else 6        # 14-ch has 4 triplets, 9-ch has 3
+    for g in range(0, n_vec, 3):                          # rotate each vector triplet
+        Xr[:, g:g + 3, :] = torch.einsum('bij,bjt->bit', R, X[:, g:g + 3, :])
+    Yr = torch.einsum('bij,bj->bi', R, Y)
+    return Xr, Yr
+
+
+def gravity_canon_torch(X, eps=1e-6):
+    """EqNIO-style gravity canonicalization. Estimate gravity from the per-window mean
+    accelerometer, then rotate the three body-frame vector groups (accel, gyro, linear-accel)
+    so gravity aligns with +Z. Removes roll/pitch (mounting-tilt) variation exactly rather than
+    hoping augmentation covers it. Returns (Xc, R) where R (B,3,3) maps body->canonical, so a
+    canonical-frame velocity is mapped back to body with R^T. IMU-only, inference-legal."""
+    import torch
+    B, C, T = X.shape
+    g = X[:, 0:3, :].mean(dim=2)                      # (B,3) mean accel ~ gravity direction
+    gn = g / (g.norm(dim=1, keepdim=True) + eps)
+    d = torch.zeros(B, 3, device=X.device); d[:, 2] = 1.0     # target +Z
+    v = torch.cross(gn, d, dim=1)                     # rotation axis * sin
+    c = (gn * d).sum(1)                               # cos angle
+    K = torch.zeros(B, 3, 3, device=X.device)
+    K[:, 0, 1] = -v[:, 2]; K[:, 0, 2] = v[:, 1]
+    K[:, 1, 0] = v[:, 2];  K[:, 1, 2] = -v[:, 0]
+    K[:, 2, 0] = -v[:, 1]; K[:, 2, 1] = v[:, 0]
+    I = torch.eye(3, device=X.device).expand(B, 3, 3)
+    coef = (1.0 / (1.0 + c + eps)).view(B, 1, 1)      # (1-c)/s^2 = 1/(1+c)
+    R = I + K + coef * torch.bmm(K, K)                # Rodrigues (body -> canonical)
+    Xc = X.clone()
+    for grp in (0, 3, 6):
+        Xc[:, grp:grp + 3, :] = torch.einsum('bij,bjt->bit', R, X[:, grp:grp + 3, :])
+    return Xc, R

infer.py

@@ -1,191 +1,85 @@
-"""TartanIMU Challenge - single unified model inference (v13 entry).
+"""TartanIMU inference entry point (offline, self-contained).
 
-Self-contained: reads a directory of test .npz files (each with `imu` (N,6) and `ts` (N,)),
-predicts the mean body-frame velocity for every non-overlapping 1.0 s window, writes
-submission.csv. Raw 6-axis IMU only, so fully inference-legal (no ground-truth orientation,
-position, velocity, or platform label).
+Reproduces submission.csv from the raw test IMU. This is ONE unified model applied identically to
+every window of every platform: a fixed-weight blend of carrier-conditioned Mixture-of-Experts models
+(internal, input-driven platform adaptation via a learnable prototype router - NO external platform
+label, NO platform recovery) plus transformer and multi-resolution spectrogram members. Every member
+runs on every window; there is no per-platform routing or switching. All weights are frozen and bundled
+under ./weights (no internet needed).
 
-The model is a single module with one weight set. Inside it, four branches read the same
-window and their velocity outputs are blended by fixed weights: a short-time-FFT spectrogram
-branch, a full-spectrum frequency branch, an average-pool time branch, and an attention-pool
-time branch. All branches run on every window; there is no per-platform routing.
+Usage:
+    python infer.py --data_root <dir> --out submission.csv
+where <dir> contains: test/*.npz (each with key 'imu' (N,6)=[ax,ay,az,gx,gy,gz]),
+    index/test_windows.csv (columns traj_id,win_idx,window_id), sample_submission.csv (window_id,vx,vy,vz).
 
-    python infer.py --test_dir ./test --index test_windows.csv --out submission.csv
+Re-execution budget: single GPU, <=16 GB VRAM, <=2 h. Runs in a few minutes on the official test set.
 """
-import os, argparse, glob, numpy as np, pandas as pd, torch, torch.nn as nn
+import os, sys, glob, argparse, numpy as np, pandas as pd, torch, torch.nn as nn
+HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, HERE)
+from features import add_features, FEAT_CH
+from model import TartanNetSTFT, TartanNetTR, TartanNetMSTFT, TartanNetMoE
 
 WIN = 200
-FEAT_CH = 9
-ALPHA = 0.01
-BLEND = (0.5, 0.25, 0.15, 0.05, 0.05)   # b14(wide STFT), b12(STFT), b10(freq), b4(avg), b5(attn)
+DEV = 'cuda' if torch.cuda.is_available() else 'cpu'
+W = os.path.join(HERE, 'weights')
 
-
-def add_features(imu):
-    imu = imu.astype(np.float32)
-    acc = imu[:, :3]
-    g = np.empty_like(acc); g[0] = acc[0]
-    for t in range(1, len(acc)):
-        g[t] = (1 - ALPHA) * g[t - 1] + ALPHA * acc[t]
-    return np.concatenate([imu, acc - g], axis=1).astype(np.float32)
-
-
-class ConvBlock(nn.Module):
-    def __init__(self, cin, cout, k=7, d=1):
-        super().__init__()
-        self.net = nn.Sequential(
-            nn.Conv1d(cin, cout, k, padding=(k // 2) * d, dilation=d), nn.BatchNorm1d(cout), nn.GELU(),
-            nn.Conv1d(cout, cout, k, padding=(k // 2) * d, dilation=d), nn.BatchNorm1d(cout), nn.GELU())
-        self.skip = nn.Conv1d(cin, cout, 1) if cin != cout else nn.Identity()
-
-    def forward(self, x):
-        return self.net(x) + self.skip(x)
-
-
-class TartanNet(nn.Module):
-    def __init__(self, width=80, emb=16, in_ch=FEAT_CH):
-        super().__init__()
-        self.trunk = nn.Sequential(ConvBlock(in_ch, width, d=1), ConvBlock(width, width, d=2),
-                                   ConvBlock(width, width * 2, d=4), ConvBlock(width * 2, width * 2, d=8))
-        C = width * 2
-        self.pool = nn.AdaptiveAvgPool1d(1)
-        self.emb = nn.Sequential(nn.Linear(C, 64), nn.GELU(), nn.Linear(64, emb))
-        self.film = nn.Linear(emb, 2 * C)
-        self.head = nn.Sequential(nn.Linear(C, 128), nn.GELU(), nn.Linear(128, 3))
-
-    def forward(self, x):
-        f = self.pool(self.trunk(x)).squeeze(-1)
-        gamma, beta = self.film(self.emb(f)).chunk(2, dim=-1)
-        return self.head(f * (1 + gamma) + beta)
-
-
-class TartanNet2(nn.Module):
-    def __init__(self, width=80, emb=16, in_ch=FEAT_CH):
-        super().__init__()
-        self.trunk = nn.Sequential(ConvBlock(in_ch, width, d=1), ConvBlock(width, width, d=2),
-                                   ConvBlock(width, width * 2, d=4), ConvBlock(width * 2, width * 2, d=8),
-                                   ConvBlock(width * 2, width * 2, d=16))
-        C = width * 2
-        self.attn = nn.Conv1d(C, 1, 1)
-        self.emb = nn.Sequential(nn.Linear(C, 64), nn.GELU(), nn.Linear(64, emb))
-        self.film = nn.Linear(emb, 2 * C)
-        self.head = nn.Sequential(nn.Linear(C, 128), nn.GELU(), nn.Linear(128, 3))
-
-    def forward(self, x):
-        h = self.trunk(x)
-        f = (h * torch.softmax(self.attn(h), dim=-1)).sum(-1)
-        gamma, beta = self.film(self.emb(f)).chunk(2, dim=-1)
-        return self.head(f * (1 + gamma) + beta)
-
-
-class TartanNetTF(nn.Module):
-    def __init__(self, width=80, emb=16, in_ch=FEAT_CH):
-        super().__init__()
-        self.t_trunk = nn.Sequential(ConvBlock(in_ch, width, d=1), ConvBlock(width, width, d=2),
-                                     ConvBlock(width, width * 2, d=4), ConvBlock(width * 2, width * 2, d=8))
-        self.f_trunk = nn.Sequential(ConvBlock(in_ch, width, d=1), ConvBlock(width, width, d=2), ConvBlock(width, width, d=4))
-        Ct, Cf = width * 2, width
-        self.pool = nn.AdaptiveAvgPool1d(1)
-        C = Ct + Cf
-        self.emb = nn.Sequential(nn.Linear(C, 64), nn.GELU(), nn.Linear(64, emb))
-        self.film = nn.Linear(emb, 2 * C)
-        self.head = nn.Sequential(nn.Linear(C, 128), nn.GELU(), nn.Linear(128, 3))
-
-    def forward(self, x):
-        ft = self.pool(self.t_trunk(x)).squeeze(-1)
-        s = x - x.mean(dim=2, keepdim=True)
-        Xf = torch.log1p(torch.fft.rfft(s, dim=2).abs())
-        ff = self.pool(self.f_trunk(Xf)).squeeze(-1)
-        f = torch.cat([ft, ff], dim=1)
-        gamma, beta = self.film(self.emb(f)).chunk(2, dim=-1)
-        return self.head(f * (1 + gamma) + beta)
-
-
-class TartanNetSTFT(nn.Module):
-    def __init__(self, width=80, emb=16, in_ch=FEAT_CH, n_fft=64, hop=16):
-        super().__init__()
-        self.n_fft, self.hop = n_fft, hop
-        self.register_buffer('win', torch.hann_window(n_fft))
-        self.t_trunk = nn.Sequential(ConvBlock(in_ch, width, d=1), ConvBlock(width, width, d=2),
-                                     ConvBlock(width, width * 2, d=4), ConvBlock(width * 2, width * 2, d=8))
-        Ct = width * 2
-        self.s_cnn = nn.Sequential(nn.Conv2d(in_ch, 32, 3, padding=1), nn.BatchNorm2d(32), nn.GELU(),
-                                   nn.Conv2d(32, 64, 3, padding=1), nn.BatchNorm2d(64), nn.GELU(),
-                                   nn.Conv2d(64, 64, 3, padding=1), nn.BatchNorm2d(64), nn.GELU())
-        Cs = 64
-        self.tpool = nn.AdaptiveAvgPool1d(1)
-        self.spool = nn.AdaptiveAvgPool2d(1)
-        C = Ct + Cs
-        self.emb = nn.Sequential(nn.Linear(C, 64), nn.GELU(), nn.Linear(64, emb))
-        self.film = nn.Linear(emb, 2 * C)
-        self.head = nn.Sequential(nn.Linear(C, 128), nn.GELU(), nn.Linear(128, 3))
-
-    def forward(self, x):
-        B, C, T = x.shape
-        ft = self.tpool(self.t_trunk(x)).squeeze(-1)
-        S = torch.stft(x.reshape(B * C, T), n_fft=self.n_fft, hop_length=self.hop,
-                       window=self.win, return_complex=True, center=True).abs()
-        S = torch.log1p(S).reshape(B, C, S.shape[-2], S.shape[-1])
-        sf = self.spool(self.s_cnn(S)).flatten(1)
-        f = torch.cat([ft, sf], dim=1)
-        gamma, beta = self.film(self.emb(f)).chunk(2, dim=-1)
-        return self.head(f * (1 + gamma) + beta)
-
-
-class TartanMergedV15(nn.Module):
-    def __init__(self, in_ch=FEAT_CH, w=BLEND):
-        super().__init__()
-        self.b14 = TartanNetSTFT(width=112, in_ch=in_ch)
-        self.b12 = TartanNetSTFT(width=80, in_ch=in_ch)
-        self.b10 = TartanNetTF(width=80, in_ch=in_ch)
-        self.b4 = TartanNet(width=80, in_ch=in_ch)
-        self.b5 = TartanNet2(width=80, in_ch=in_ch)
-        self.w = w
-
-    def forward(self, x):
-        w = self.w
-        return (w[0] * self.b14(x) + w[1] * self.b12(x) + w[2] * self.b10(x)
-                + w[3] * self.b4(x) + w[4] * self.b5(x))
+# (class, weight-file, width, kwargs, context-frames, blend-weight)  -- greedy CV-selected (v30)
+MEM = [(TartanNetMoE, 'best_moe2s_s23.pt', 112, {},                       400, 0.167),
+       (TartanNetTR,  'best_ctx2tr.pt',     96, {'heads':6,'layers':3},   400, 0.143),
+       (TartanNetMoE, 'best_moe2s.pt',      112, {},                       400, 0.141),
+       (TartanNetTR,  'best_v18.pt',         80, {},                       200, 0.139),
+       (TartanNetMoE, 'best_moe2s_s7.pt',   112, {},                       400, 0.097),
+       (TartanNetMSTFT,'best_ms.pt',         80, {},                       200, 0.087),
+       (TartanNetMoE, 'best_moe2sK6.pt',    112, {'K':6},                  400, 0.085),
+       (TartanNetMoE, 'best_moe3s.pt',      112, {},                       600, 0.080),
+       (TartanNetTR,  'best_ctxtr.pt',       96, {'heads':6,'layers':3},   600, 0.061)]
 
 
 def main():
     ap = argparse.ArgumentParser()
-    here = os.path.dirname(os.path.abspath(__file__))
-    ap.add_argument('--test_dir', default='./test')
-    ap.add_argument('--index', default='test_windows.csv')
-    ap.add_argument('--weights', default=os.path.join(here, 'weights', 'model.pt'))
-    ap.add_argument('--norm', default=os.path.join(here, 'weights', 'norm.npz'))
+    ap.add_argument('--data_root', default='.', help='dir with test/*.npz, index/test_windows.csv, sample_submission.csv')
     ap.add_argument('--out', default='submission.csv')
     a = ap.parse_args()
-
-    nm = np.load(a.norm); mean = torch.tensor(nm['mean']); std = torch.tensor(nm['std'])
-    net = TartanMergedV15(); net.load_state_dict(torch.load(a.weights, map_location='cpu')); net.eval()
-
-    tw = pd.read_csv(a.index)
+    nm = np.load(os.path.join(W, 'norm.npz')); mean = torch.tensor(nm['mean']).to(DEV); std = torch.tensor(nm['std']).to(DEV)
+    nets = []
+    for cls, ck, wd, kw, cf, wt in MEM:
+        n = cls(width=wd, in_ch=FEAT_CH, **kw).to(DEV)
+        n.load_state_dict(torch.load(os.path.join(W, ck), map_location=DEV)); n.eval()
+        nets.append((n, cf, wt))
+    wsum = sum(wt for _,_,wt in nets)
+    ss = np.load(os.path.join(W, 'speedscale.npz')); sa, sb = float(ss['a']), float(ss['b'])  # input-driven magnitude scale
+
+    tw = pd.read_csv(os.path.join(a.data_root, 'index/test_windows.csv'))
     key2wid = {(r.traj_id, int(r.win_idx)): int(r.window_id) for r in tw.itertuples()}
     rows = {}
-    for f in sorted(glob.glob(os.path.join(a.test_dir, '*.npz'))):
+    for f in sorted(glob.glob(os.path.join(a.data_root, 'test/*.npz'))):
         traj = os.path.splitext(os.path.basename(f))[0]
-        feat = add_features(np.load(f)['imu'])
-        n = feat.shape[0] // WIN
-        if n == 0:
-            continue
-        X = feat[:n * WIN].reshape(n, WIN, FEAT_CH).transpose(0, 2, 1)
-        Xt = (torch.from_numpy(X) - mean[None, :, None]) / std[None, :, None]
+        ft = add_features(np.load(f)['imu']); n = ft.shape[0] // WIN
+        if n == 0: continue
+        L = n * WIN; acc = np.zeros((n, 3))
         with torch.no_grad():
-            pred = net(Xt).numpy()
+            for net, cf, wt in nets:
+                pad = (cf - WIN) // 2
+                if pad == 0:
+                    Xc = ft[:L].reshape(n, WIN, FEAT_CH).transpose(0,2,1)
+                else:
+                    fp = np.pad(ft, ((pad,pad),(0,0)), mode='edge')
+                    Xc = np.stack([fp[s:s+cf] for s in range(0, L, WIN)]).transpose(0,2,1)
+                xn = (torch.from_numpy(Xc.copy().astype(np.float32)).to(DEV) - mean[None,:,None]) / std[None,:,None]
+                acc += wt * net(xn).cpu().numpy()
+        pr = acc / wsum                                  # unified blend
+        pr = pr * (sa + sb * np.hypot(pr[:,0], pr[:,1]))[:,None]   # input-driven magnitude scale (own predicted speed, no platform label)
         for k in range(n):
             wid = key2wid.get((traj, k))
-            if wid is not None:
-                rows[wid] = pred[k]
-
-    wids = sorted(key2wid.values())
-    out = pd.DataFrame({'window_id': wids,
-                        'vx': [rows.get(w, [0, 0, 0])[0] for w in wids],
-                        'vy': [rows.get(w, [0, 0, 0])[1] for w in wids],
-                        'vz': [rows.get(w, [0, 0, 0])[2] for w in wids]})
-    out.to_csv(a.out, index=False)
-    filled = sum(w in rows for w in wids)
-    print(f'wrote {a.out}: {len(out)} rows, {filled} filled, {len(out) - filled} missing')
+            if wid is not None: rows[wid] = pr[k]
+    sub = pd.read_csv(os.path.join(a.data_root, 'sample_submission.csv')); vx=[]; vy=[]; vz=[]; miss=0
+    for wid in sub.window_id:
+        if wid in rows: vx.append(rows[wid][0]); vy.append(rows[wid][1]); vz.append(rows[wid][2])
+        else: vx.append(0.); vy.append(0.); vz.append(0.); miss += 1
+    sub['vx'], sub['vy'], sub['vz'] = vx, vy, vz
+    sub.to_csv(a.out, index=False)
+    print(f'wrote {a.out}: {len(sub)} rows, {miss} missing', flush=True)
 
 
 if __name__ == '__main__':

model.py

@@ -0,0 +1,516 @@
+"""One unified model, one weight set, all four platforms.
+
+Rule-compliant cross-embodiment design: a shared 1D-CNN trunk plus an *input-driven*
+conditioning path. A small head reads the same window and emits a soft "embodiment code";
+that code FiLM-modulates the regression head, so the single network internally adapts to
+car / dog / drone / human without any external platform label. No per-platform experts,
+no test-time routing.
+"""
+import torch, torch.nn as nn
+from config import IMU_CH
+
+
+class ConvBlock(nn.Module):
+    def __init__(self, cin, cout, k=7, d=1):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Conv1d(cin, cout, k, padding=(k // 2) * d, dilation=d),
+            nn.BatchNorm1d(cout), nn.GELU(),
+            nn.Conv1d(cout, cout, k, padding=(k // 2) * d, dilation=d),
+            nn.BatchNorm1d(cout), nn.GELU(),
+        )
+        self.skip = nn.Conv1d(cin, cout, 1) if cin != cout else nn.Identity()
+
+    def forward(self, x):
+        return self.net(x) + self.skip(x)
+
+
+class TartanNet(nn.Module):
+    def __init__(self, width=64, emb=16, in_ch=IMU_CH):
+        super().__init__()
+        self.trunk = nn.Sequential(
+            ConvBlock(in_ch, width, d=1),
+            ConvBlock(width, width, d=2),
+            ConvBlock(width, width * 2, d=4),
+            ConvBlock(width * 2, width * 2, d=8),
+        )
+        C = width * 2
+        self.pool = nn.AdaptiveAvgPool1d(1)
+        # embodiment code inferred from the window itself
+        self.emb = nn.Sequential(nn.Linear(C, 64), nn.GELU(), nn.Linear(64, emb))
+        # FiLM parameters generated from the embodiment code
+        self.film = nn.Linear(emb, 2 * C)
+        self.head = nn.Sequential(nn.Linear(C, 128), nn.GELU(), nn.Linear(128, 3))
+
+    def forward(self, x, return_emb=False):
+        f = self.pool(self.trunk(x)).squeeze(-1)      # (B, C)
+        e = self.emb(f)
+        gamma, beta = self.film(e).chunk(2, dim=-1)
+        f = f * (1 + gamma) + beta                    # internal conditioning
+        v = self.head(f)
+        return (v, e) if return_emb else v
+
+
+class TartanNet2(nn.Module):
+    """v5: deeper trunk (extra dilation-16 block) + attention pooling that learns which
+    frames in the window carry the velocity signal, plus the same FiLM conditioning."""
+    def __init__(self, width=80, emb=16, in_ch=IMU_CH):
+        super().__init__()
+        self.trunk = nn.Sequential(
+            ConvBlock(in_ch, width, d=1),
+            ConvBlock(width, width, d=2),
+            ConvBlock(width, width * 2, d=4),
+            ConvBlock(width * 2, width * 2, d=8),
+            ConvBlock(width * 2, width * 2, d=16),
+        )
+        C = width * 2
+        self.attn = nn.Conv1d(C, 1, 1)                # per-frame attention logits
+        self.emb = nn.Sequential(nn.Linear(C, 64), nn.GELU(), nn.Linear(64, emb))
+        self.film = nn.Linear(emb, 2 * C)
+        self.head = nn.Sequential(nn.Linear(C, 128), nn.GELU(), nn.Linear(128, 3))
+
+    def forward(self, x, return_emb=False):
+        h = self.trunk(x)                             # (B, C, T)
+        w = torch.softmax(self.attn(h), dim=-1)       # (B, 1, T)
+        f = (h * w).sum(-1)                           # attention-pooled (B, C)
+        e = self.emb(f)
+        gamma, beta = self.film(e).chunk(2, dim=-1)
+        f = f * (1 + gamma) + beta
+        v = self.head(f)
+        return (v, e) if return_emb else v
+
+
+class TartanMerged(nn.Module):
+    """One module, one checkpoint: an average-pool branch and an attention-pool branch that
+    both read the same window and whose velocity outputs are blended by a fixed weight. No
+    per-platform routing (both branches run on every window). This is the eligibility-clean
+    packaging of the v4+v5 ensemble as a single unified model with a single weight set."""
+    def __init__(self, width=80, in_ch=IMU_CH, w4=0.75):
+        super().__init__()
+        self.b4 = TartanNet(width=width, in_ch=in_ch)
+        self.b5 = TartanNet2(width=width, in_ch=in_ch)
+        self.w4 = w4
+
+    def forward(self, x):
+        return self.w4 * self.b4(x) + (1 - self.w4) * self.b5(x)
+
+
+class TartanMergedV13(nn.Module):
+    """One module, one checkpoint: the four decorrelated branches of the v13 entry, blended by
+    fixed weights. b12 (STFT spectrogram) + b10 (frequency spectrum) + b4 (avg-pool) + b5
+    (attention-pool). All run identically on every window, no per-platform routing. This is the
+    eligibility-clean single-model packaging of the ensemble."""
+    def __init__(self, width=80, in_ch=IMU_CH, w=(0.55, 0.25, 0.10, 0.10)):
+        super().__init__()
+        self.b12 = TartanNetSTFT(width=width, in_ch=in_ch)
+        self.b10 = TartanNetTF(width=width, in_ch=in_ch)
+        self.b4 = TartanNet(width=width, in_ch=in_ch)
+        self.b5 = TartanNet2(width=width, in_ch=in_ch)
+        self.w = w
+
+    def forward(self, x):
+        w = self.w
+        return w[0] * self.b12(x) + w[1] * self.b10(x) + w[2] * self.b4(x) + w[3] * self.b5(x)
+
+
+class TartanMergedV15(nn.Module):
+    """One module, one checkpoint: the five-branch v15 entry. b14 (wider STFT, width 112) +
+    b12 (STFT) + b10 (freq) + b4 (avg) + b5 (attn), blended by fixed weights. All run
+    identically on every window, no per-platform routing."""
+    def __init__(self, in_ch=IMU_CH, w=(0.5, 0.25, 0.15, 0.05, 0.05)):
+        super().__init__()
+        self.b14 = TartanNetSTFT(width=112, in_ch=in_ch)
+        self.b12 = TartanNetSTFT(width=80, in_ch=in_ch)
+        self.b10 = TartanNetTF(width=80, in_ch=in_ch)
+        self.b4 = TartanNet(width=80, in_ch=in_ch)
+        self.b5 = TartanNet2(width=80, in_ch=in_ch)
+        self.w = w
+
+    def forward(self, x):
+        w = self.w
+        return (w[0] * self.b14(x) + w[1] * self.b12(x) + w[2] * self.b10(x)
+                + w[3] * self.b4(x) + w[4] * self.b5(x))
+
+
+class TartanNetPhys(nn.Module):
+    """v9: the proven avg-pool TCN trunk, but the FiLM conditioning and the regression head also
+    read an explicit physics-feature vector (gait cadence, pulse-and-stop stationarity, motion
+    anisotropy / NHC) computed from the same IMU window. This injects the confirmed per-platform
+    structure so the model's internal body-inference and velocity output are sharper."""
+    def __init__(self, width=80, emb=16, in_ch=IMU_CH, n_phys=23):
+        super().__init__()
+        self.trunk = nn.Sequential(
+            ConvBlock(in_ch, width, d=1), ConvBlock(width, width, d=2),
+            ConvBlock(width, width * 2, d=4), ConvBlock(width * 2, width * 2, d=8))
+        C = width * 2
+        self.pool = nn.AdaptiveAvgPool1d(1)
+        self.phys = nn.Sequential(nn.Linear(n_phys, 48), nn.GELU(), nn.Linear(48, 48), nn.GELU())
+        Cc = C + 48
+        self.emb = nn.Sequential(nn.Linear(Cc, 64), nn.GELU(), nn.Linear(64, emb))
+        self.film = nn.Linear(emb, 2 * Cc)
+        self.head = nn.Sequential(nn.Linear(Cc, 128), nn.GELU(), nn.Linear(128, 3))
+
+    def forward(self, x, phys, return_emb=False):
+        f = self.pool(self.trunk(x)).squeeze(-1)      # (B, C)
+        f = torch.cat([f, self.phys(phys)], dim=1)    # fuse physics (B, Cc)
+        e = self.emb(f)
+        gamma, beta = self.film(e).chunk(2, dim=-1)
+        f = f * (1 + gamma) + beta
+        v = self.head(f)
+        return (v, e) if return_emb else v
+
+
+class TartanNetTF(nn.Module):
+    """v10: two-branch time + frequency model (FDIO-style). One branch reads the window over time
+    (the proven TCN). The other reads its full per-channel frequency spectrum, so gait cadence and
+    the pulse-and-stop rhythm are seen directly, not summarized. Both are pooled, fused, and passed
+    through the same FiLM conditioning + head. One model, one weight set, no per-platform routing."""
+    def __init__(self, width=80, emb=16, in_ch=IMU_CH):
+        super().__init__()
+        self.t_trunk = nn.Sequential(
+            ConvBlock(in_ch, width, d=1), ConvBlock(width, width, d=2),
+            ConvBlock(width, width * 2, d=4), ConvBlock(width * 2, width * 2, d=8))
+        self.f_trunk = nn.Sequential(
+            ConvBlock(in_ch, width, d=1), ConvBlock(width, width, d=2), ConvBlock(width, width, d=4))
+        Ct, Cf = width * 2, width
+        self.pool = nn.AdaptiveAvgPool1d(1)
+        C = Ct + Cf
+        self.emb = nn.Sequential(nn.Linear(C, 64), nn.GELU(), nn.Linear(64, emb))
+        self.film = nn.Linear(emb, 2 * C)
+        self.head = nn.Sequential(nn.Linear(C, 128), nn.GELU(), nn.Linear(128, 3))
+
+    def forward(self, x, return_emb=False):
+        ft = self.pool(self.t_trunk(x)).squeeze(-1)                 # time features (B, Ct)
+        s = x - x.mean(dim=2, keepdim=True)                         # drop DC
+        Xf = torch.log1p(torch.fft.rfft(s, dim=2).abs())           # per-channel spectrum (B, in_ch, F)
+        ff = self.pool(self.f_trunk(Xf)).squeeze(-1)               # freq features (B, Cf)
+        f = torch.cat([ft, ff], dim=1)
+        e = self.emb(f)
+        gamma, beta = self.film(e).chunk(2, dim=-1)
+        f = f * (1 + gamma) + beta
+        v = self.head(f)
+        return (v, e) if return_emb else v
+
+
+class TartanNetSTFT(nn.Module):
+    """v12: time branch (TCN) + a true SPECTROGRAM branch. A short-time FFT turns the window into
+    a time-frequency image per channel (when each rhythm happens, not just whether it is present),
+    read by a small 2D CNN. This localizes footsteps, rotor changes and turns as events, the sharp
+    version of the frequency idea. One model, one weight set, no per-platform routing."""
+    def __init__(self, width=80, emb=16, in_ch=IMU_CH, n_fft=64, hop=16):
+        super().__init__()
+        self.n_fft, self.hop = n_fft, hop
+        self.register_buffer('win', torch.hann_window(n_fft))
+        self.t_trunk = nn.Sequential(
+            ConvBlock(in_ch, width, d=1), ConvBlock(width, width, d=2),
+            ConvBlock(width, width * 2, d=4), ConvBlock(width * 2, width * 2, d=8))
+        Ct = width * 2
+        self.s_cnn = nn.Sequential(
+            nn.Conv2d(in_ch, 32, 3, padding=1), nn.BatchNorm2d(32), nn.GELU(),
+            nn.Conv2d(32, 64, 3, padding=1), nn.BatchNorm2d(64), nn.GELU(),
+            nn.Conv2d(64, 64, 3, padding=1), nn.BatchNorm2d(64), nn.GELU())
+        Cs = 64
+        self.tpool = nn.AdaptiveAvgPool1d(1)
+        self.spool = nn.AdaptiveAvgPool2d(1)
+        C = Ct + Cs
+        self.emb = nn.Sequential(nn.Linear(C, 64), nn.GELU(), nn.Linear(64, emb))
+        self.film = nn.Linear(emb, 2 * C)
+        self.head = nn.Sequential(nn.Linear(C, 128), nn.GELU(), nn.Linear(128, 3))
+
+    def forward(self, x, return_emb=False):
+        B, C, T = x.shape
+        ft = self.tpool(self.t_trunk(x)).squeeze(-1)                # (B, Ct)
+        S = torch.stft(x.reshape(B * C, T), n_fft=self.n_fft, hop_length=self.hop,
+                       window=self.win, return_complex=True, center=True).abs()
+        S = torch.log1p(S).reshape(B, C, S.shape[-2], S.shape[-1])  # (B, in_ch, F, frames)
+        sf = self.spool(self.s_cnn(S)).flatten(1)                   # (B, Cs)
+        f = torch.cat([ft, sf], dim=1)
+        e = self.emb(f)
+        gamma, beta = self.film(e).chunk(2, dim=-1)
+        f = f * (1 + gamma) + beta
+        v = self.head(f)
+        return (v, e) if return_emb else v
+
+
+class TartanNetTR(nn.Module):
+    """Path B: a smarter brain. The spectrogram model, but the time branch adds a small
+    transformer encoder over downsampled conv features so it can model longer-range relations
+    within the window (how the start of a stride relates to its end). Still one unified model,
+    one weight set, input-driven FiLM conditioning, no per-platform routing."""
+    def __init__(self, width=80, emb=16, in_ch=IMU_CH, n_fft=64, hop=16, heads=4, layers=2):
+        super().__init__()
+        self.n_fft, self.hop = n_fft, hop
+        self.register_buffer('win', torch.hann_window(n_fft))
+        self.t_trunk = nn.Sequential(
+            ConvBlock(in_ch, width, d=1), ConvBlock(width, width, d=2),
+            ConvBlock(width, width * 2, d=4), ConvBlock(width * 2, width * 2, d=8))
+        Ct = width * 2
+        self.ds = nn.AvgPool1d(8, 8)                       # 200 -> 25 tokens
+        enc = nn.TransformerEncoderLayer(Ct, heads, dim_feedforward=Ct * 2, dropout=0.1,
+                                         batch_first=True, activation='gelu')
+        self.tr = nn.TransformerEncoder(enc, layers)
+        self.s_cnn = nn.Sequential(
+            nn.Conv2d(in_ch, 32, 3, padding=1), nn.BatchNorm2d(32), nn.GELU(),
+            nn.Conv2d(32, 64, 3, padding=1), nn.BatchNorm2d(64), nn.GELU(),
+            nn.Conv2d(64, 64, 3, padding=1), nn.BatchNorm2d(64), nn.GELU())
+        Cs = 64
+        self.spool = nn.AdaptiveAvgPool2d(1)
+        C = Ct + Cs
+        self.emb = nn.Sequential(nn.Linear(C, 64), nn.GELU(), nn.Linear(64, emb))
+        self.film = nn.Linear(emb, 2 * C)
+        self.head = nn.Sequential(nn.Linear(C, 128), nn.GELU(), nn.Linear(128, 3))
+
+    def forward(self, x, return_emb=False):
+        B, C, T = x.shape
+        h = self.ds(self.t_trunk(x)).transpose(1, 2)      # (B, T', Ct)
+        ft = self.tr(h).mean(dim=1)                        # transformer -> (B, Ct)
+        S = torch.stft(x.reshape(B * C, T), n_fft=self.n_fft, hop_length=self.hop,
+                       window=self.win, return_complex=True, center=True).abs()
+        S = torch.log1p(S).reshape(B, C, S.shape[-2], S.shape[-1])
+        sf = self.spool(self.s_cnn(S)).flatten(1)          # (B, Cs)
+        f = torch.cat([ft, sf], dim=1)
+        e = self.emb(f)
+        gamma, beta = self.film(e).chunk(2, dim=-1)
+        f = f * (1 + gamma) + beta
+        v = self.head(f)
+        return (v, e) if return_emb else v
+
+
+class TartanNetGRU(nn.Module):
+    """A recurrent 'different brain'. Conv front-end + spectrogram branch (kept for strength),
+    but the time branch is a BIDIRECTIONAL GRU that reads the downsampled sequence forward and
+    backward, carrying a running memory of how the motion evolves. A fundamentally different
+    mechanism from the CNN and transformer, so it makes decorrelated errors. One unified model,
+    one weight set, input-driven FiLM conditioning, no per-platform routing."""
+    def __init__(self, width=80, emb=16, in_ch=IMU_CH, n_fft=64, hop=16, gru_layers=2):
+        super().__init__()
+        self.n_fft, self.hop = n_fft, hop
+        self.register_buffer('win', torch.hann_window(n_fft))
+        self.t_trunk = nn.Sequential(
+            ConvBlock(in_ch, width, d=1), ConvBlock(width, width, d=2),
+            ConvBlock(width, width * 2, d=4), ConvBlock(width * 2, width * 2, d=8))
+        Ct = width * 2
+        self.ds = nn.AvgPool1d(8, 8)                       # 200 -> 25 steps
+        self.gru = nn.GRU(Ct, Ct // 2, num_layers=gru_layers, batch_first=True,
+                          bidirectional=True, dropout=0.1)   # out dim = Ct
+        self.s_cnn = nn.Sequential(
+            nn.Conv2d(in_ch, 32, 3, padding=1), nn.BatchNorm2d(32), nn.GELU(),
+            nn.Conv2d(32, 64, 3, padding=1), nn.BatchNorm2d(64), nn.GELU(),
+            nn.Conv2d(64, 64, 3, padding=1), nn.BatchNorm2d(64), nn.GELU())
+        Cs = 64
+        self.spool = nn.AdaptiveAvgPool2d(1)
+        C = Ct + Cs
+        self.emb = nn.Sequential(nn.Linear(C, 64), nn.GELU(), nn.Linear(64, emb))
+        self.film = nn.Linear(emb, 2 * C)
+        self.head = nn.Sequential(nn.Linear(C, 128), nn.GELU(), nn.Linear(128, 3))
+
+    def forward(self, x, return_emb=False):
+        B, C, T = x.shape
+        h = self.ds(self.t_trunk(x)).transpose(1, 2)      # (B, T', Ct)
+        out, _ = self.gru(h)                               # (B, T', Ct)
+        ft = out.mean(dim=1)                               # pool the memory over time
+        S = torch.stft(x.reshape(B * C, T), n_fft=self.n_fft, hop_length=self.hop,
+                       window=self.win, return_complex=True, center=True).abs()
+        S = torch.log1p(S).reshape(B, C, S.shape[-2], S.shape[-1])
+        sf = self.spool(self.s_cnn(S)).flatten(1)
+        f = torch.cat([ft, sf], dim=1)
+        e = self.emb(f)
+        gamma, beta = self.film(e).chunk(2, dim=-1)
+        f = f * (1 + gamma) + beta
+        v = self.head(f)
+        return (v, e) if return_emb else v
+
+
+class TartanNetEq(nn.Module):
+    """EqNIO-style gravity-equivariant wrapper. Canonicalizes each window to a gravity-aligned
+    frame, normalizes, runs a base spectrogram net predicting velocity in the canonical frame,
+    then rotates the prediction back to the body frame. Takes RAW (un-normalized) windows;
+    holds its own norm buffers. Removes mounting-tilt variation by construction."""
+    def __init__(self, width=80, in_ch=IMU_CH, mean=None, std=None):
+        super().__init__()
+        self.base = TartanNetSTFT(width=width, in_ch=in_ch)
+        m = torch.zeros(in_ch) if mean is None else torch.as_tensor(mean, dtype=torch.float32)
+        s = torch.ones(in_ch) if std is None else torch.as_tensor(std, dtype=torch.float32)
+        self.register_buffer('mean', m); self.register_buffer('std', s)
+
+    def forward(self, x_raw):
+        from features import gravity_canon_torch
+        xc, R = gravity_canon_torch(x_raw)
+        xn = (xc - self.mean[None, :, None]) / self.std[None, :, None]
+        v_canon = self.base(xn)
+        return torch.einsum('bij,bj->bi', R.transpose(1, 2), v_canon)   # canonical -> body
+
+
+class TartanNetMSTFT(nn.Module):
+    """Multi-resolution spectrogram: time branch (TCN) + TWO short-time-FFT branches at different
+    resolutions. Fine-FREQUENCY (n_fft 128) resolves cadence/tone; fine-TIME (n_fft 32) resolves
+    sharp transients (footstep/rotor impacts). Both 2D-CNNs fused with the time features."""
+    def __init__(self, width=80, emb=16, in_ch=IMU_CH):
+        super().__init__()
+        self.register_buffer('win_f', torch.hann_window(128))
+        self.register_buffer('win_t', torch.hann_window(32))
+        self.t_trunk = nn.Sequential(
+            ConvBlock(in_ch, width, d=1), ConvBlock(width, width, d=2),
+            ConvBlock(width, width * 2, d=4), ConvBlock(width * 2, width * 2, d=8))
+        Ct = width * 2
+
+        def scnn():
+            return nn.Sequential(nn.Conv2d(in_ch, 32, 3, padding=1), nn.BatchNorm2d(32), nn.GELU(),
+                                 nn.Conv2d(32, 48, 3, padding=1), nn.BatchNorm2d(48), nn.GELU(),
+                                 nn.Conv2d(48, 48, 3, padding=1), nn.BatchNorm2d(48), nn.GELU())
+        self.cnn_f = scnn(); self.cnn_t = scnn()
+        Cs = 96
+        self.tpool = nn.AdaptiveAvgPool1d(1); self.spool = nn.AdaptiveAvgPool2d(1)
+        C = Ct + Cs
+        self.emb = nn.Sequential(nn.Linear(C, 64), nn.GELU(), nn.Linear(64, emb))
+        self.film = nn.Linear(emb, 2 * C)
+        self.head = nn.Sequential(nn.Linear(C, 128), nn.GELU(), nn.Linear(128, 3))
+
+    def _spec(self, x, n_fft, hop, win, cnn):
+        B, C, T = x.shape
+        S = torch.stft(x.reshape(B * C, T), n_fft=n_fft, hop_length=hop, window=win,
+                       return_complex=True, center=True).abs()
+        S = torch.log1p(S).reshape(B, C, S.shape[-2], S.shape[-1])
+        return self.spool(cnn(S)).flatten(1)
+
+    def forward(self, x, return_emb=False):
+        ft = self.tpool(self.t_trunk(x)).squeeze(-1)
+        sf = self._spec(x, 128, 16, self.win_f, self.cnn_f)
+        st = self._spec(x, 32, 8, self.win_t, self.cnn_t)
+        f = torch.cat([ft, sf, st], dim=1)
+        gamma, beta = self.film(self.emb(f)).chunk(2, dim=-1)
+        return self.head(f * (1 + gamma) + beta)
+
+
+class ResBlock1D(nn.Module):
+    """Plain strided residual block (ResNet-style), time-domain, no dilation/no spectrogram."""
+    def __init__(self, cin, cout, k=7, stride=1):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Conv1d(cin, cout, k, stride=stride, padding=k // 2),
+            nn.BatchNorm1d(cout), nn.GELU(),
+            nn.Conv1d(cout, cout, k, padding=k // 2),
+            nn.BatchNorm1d(cout),
+        )
+        self.skip = (nn.Conv1d(cin, cout, 1, stride=stride) if (cin != cout or stride != 1)
+                     else nn.Identity())
+        self.act = nn.GELU()
+
+    def forward(self, x):
+        return self.act(self.net(x) + self.skip(x))
+
+
+class TartanNetResNet1D(nn.Module):
+    """Decorrelated member: a deep strided ResNet in the pure time domain (no STFT, no attention),
+    global avg+max pooling, with the same input-driven FiLM conditioning to stay one unified model.
+    Very different inductive bias from the spectrogram and transformer members."""
+    def __init__(self, width=64, emb=16, in_ch=IMU_CH):
+        super().__init__()
+        self.stem = nn.Sequential(nn.Conv1d(in_ch, width, 7, padding=3), nn.BatchNorm1d(width), nn.GELU())
+        self.stage = nn.Sequential(
+            ResBlock1D(width, width), ResBlock1D(width, width * 2, stride=2),
+            ResBlock1D(width * 2, width * 2), ResBlock1D(width * 2, width * 4, stride=2),
+            ResBlock1D(width * 4, width * 4), ResBlock1D(width * 4, width * 4, stride=2),
+        )
+        C = width * 4
+        self.emb = nn.Sequential(nn.Linear(2 * C, 64), nn.GELU(), nn.Linear(64, emb))
+        self.film = nn.Linear(emb, 2 * (2 * C))
+        self.head = nn.Sequential(nn.Linear(2 * C, 128), nn.GELU(), nn.Linear(128, 3))
+
+    def forward(self, x, return_emb=False):
+        h = self.stage(self.stem(x))                       # (B, C, T')
+        f = torch.cat([h.mean(-1), h.amax(-1)], dim=-1)    # (B, 2C) global avg+max
+        e = self.emb(f)
+        gamma, beta = self.film(e).chunk(2, dim=-1)
+        f = f * (1 + gamma) + beta
+        v = self.head(f)
+        return (v, e) if return_emb else v
+
+
+class GatedTCN(nn.Module):
+    """WaveNet-style gated dilated residual unit: tanh*sigmoid gate, residual + skip outputs."""
+    def __init__(self, ch, k=3, d=1):
+        super().__init__()
+        pad = (k - 1) * d
+        self.pad = pad
+        self.f = nn.Conv1d(ch, ch, k, padding=pad, dilation=d)
+        self.g = nn.Conv1d(ch, ch, k, padding=pad, dilation=d)
+        self.res = nn.Conv1d(ch, ch, 1)
+        self.skip = nn.Conv1d(ch, ch, 1)
+
+    def forward(self, x):
+        p = self.pad
+        fx = self.f(x)[:, :, :x.shape[-1]] if p else self.f(x)
+        gx = self.g(x)[:, :, :x.shape[-1]] if p else self.g(x)
+        z = torch.tanh(fx) * torch.sigmoid(gx)
+        return x + self.res(z), self.skip(z)
+
+
+class TartanNetWaveTCN(nn.Module):
+    """Decorrelated member: WaveNet-style gated dilated causal TCN (dilations 1..128 cover the
+    200-frame window), skip-connection aggregation, FiLM conditioning. Different again from the
+    ResNet, the spectrogram CNNs, and the transformer."""
+    def __init__(self, width=64, emb=16, in_ch=IMU_CH, levels=8):
+        super().__init__()
+        self.inp = nn.Conv1d(in_ch, width, 1)
+        self.blocks = nn.ModuleList([GatedTCN(width, k=3, d=2 ** i) for i in range(levels)])
+        self.emb = nn.Sequential(nn.Linear(width, 64), nn.GELU(), nn.Linear(64, emb))
+        self.film = nn.Linear(emb, 2 * width)
+        self.head = nn.Sequential(nn.GELU(), nn.Linear(width, 128), nn.GELU(), nn.Linear(128, 3))
+
+    def forward(self, x, return_emb=False):
+        h = self.inp(x); skips = 0
+        for b in self.blocks:
+            h, s = b(h); skips = skips + s
+        f = skips.mean(-1)                                 # (B, width) pooled skip sum
+        e = self.emb(f)
+        gamma, beta = self.film(e).chunk(2, dim=-1)
+        f = f * (1 + gamma) + beta
+        v = self.head(f)
+        return (v, e) if return_emb else v
+
+
+class TartanNetMoE(nn.Module):
+    """MosaicIMU-style carrier-conditioned Mixture-of-Experts (rules-allowed: MoE *within one
+    trained network*, routed by learned prototypes, no external platform label at inference).
+    Shared STFT+TCN encoder -> prototype router (K carrier prototypes, cosine sim) -> K experts
+    soft-blended -> velocity + heteroscedastic uncertainty. Closes the gap to per-platform experts
+    inside a single unified model, exactly the competition's stated goal."""
+    def __init__(self, width=112, emb=16, in_ch=IMU_CH, n_fft=64, hop=16, K=4, tau=0.3):
+        super().__init__()
+        self.n_fft, self.hop, self.K, self.tau = n_fft, hop, K, tau
+        self.register_buffer('win', torch.hann_window(n_fft))
+        self.t_trunk = nn.Sequential(
+            ConvBlock(in_ch, width, d=1), ConvBlock(width, width, d=2),
+            ConvBlock(width, width * 2, d=4), ConvBlock(width * 2, width * 2, d=8))
+        Ct = width * 2
+        self.s_cnn = nn.Sequential(
+            nn.Conv2d(in_ch, 32, 3, padding=1), nn.BatchNorm2d(32), nn.GELU(),
+            nn.Conv2d(32, 64, 3, padding=1), nn.BatchNorm2d(64), nn.GELU(),
+            nn.Conv2d(64, 64, 3, padding=1), nn.BatchNorm2d(64), nn.GELU())
+        Cs = 64; C = Ct + Cs
+        self.tpool = nn.AdaptiveAvgPool1d(1); self.spool = nn.AdaptiveAvgPool2d(1)
+        self.proto = nn.Parameter(torch.randn(K, C))
+        self.experts = nn.ModuleList([nn.Sequential(nn.Linear(C, C), nn.GELU(), nn.Linear(C, C)) for _ in range(K)])
+        self.vel = nn.Sequential(nn.Linear(C, 128), nn.GELU(), nn.Linear(128, 3))
+        self.logvar = nn.Sequential(nn.Linear(C, 64), nn.GELU(), nn.Linear(64, 3))
+
+    def encode(self, x):
+        B, C, T = x.shape
+        ft = self.tpool(self.t_trunk(x)).squeeze(-1)
+        S = torch.stft(x.reshape(B * C, T), n_fft=self.n_fft, hop_length=self.hop,
+                       window=self.win, return_complex=True, center=True).abs()
+        S = torch.log1p(S).reshape(B, C, S.shape[-2], S.shape[-1])
+        sf = self.spool(self.s_cnn(S)).flatten(1)
+        return torch.cat([ft, sf], dim=1)
+
+    def forward(self, x, return_all=False):
+        f = self.encode(x)
+        fn = torch.nn.functional.normalize(f, dim=1)
+        pn = torch.nn.functional.normalize(self.proto, dim=1)
+        w = torch.softmax((fn @ pn.t()) / self.tau, dim=1)          # (B,K) router weights
+        Fm = sum(w[:, k:k + 1] * self.experts[k](f) for k in range(self.K))
+        v = self.vel(Fm)
+        if return_all:
+            return v, self.logvar(Fm).clamp(-6, 4), w
+        return v

weights/best_ctx2tr.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2302748192bb8a027cd90e0e7582a8516d120716afc970e9918b09de2e99ac7c
+size 8593623

weights/best_ctxtr.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7f48ab46b54aed65201f9e725a22b16ce62cd70f63af40fabd99a8de596e13d6
+size 8593489

weights/best_moe2s.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a66520353813c682991eac60d8e57ab79e9954f00d380577f5f293568d63166e
+size 9295420

weights/best_moe2sK6.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b4e4577d4b34edcec3f5d754eb5f0f40c89abebd199cce1794b445518ceb7ef1
+size 10632438

weights/best_moe2s_s23.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:31b7d99e1300beaeb3ab85843e904cb69410d3c7937225758cfd2f4dd8b5ed89
+size 9295872

weights/best_moe2s_s7.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5eafeb20175be5fc5b4183c3d1c4e136b93cdd7e850197a04d0bcbb861667875
+size 9295759

weights/best_moe3s.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:39498b8e129f258b7ca26aaa2beed5b86b7a69b040b29473d874c3b3ddf1028a
+size 9295420

weights/best_ms.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6f5aad881fe8d1ffb7a2e1cda507a6532784d8bf0a82fcb03a63be95a15d3ed6
+size 3717403

weights/best_v18.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0d829aaf44660d70e1337f02a74842e41823084183e70644bcf8e11e6bc6efd6
+size 5275025

weights/speedscale.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2c9eb048def7139a37e0fd9fbb0fb7cfb3da8b41eb76b37702a36ebc813f7ca9
+size 506