yangwinnie/tartanimu-unified

TartanIMU unified model: wide-window TCN + bidirectional GRU trajectory head

Single unified model for the TartanIMU Challenge (multi-platform inertial odometry: car / dog / drone / human from raw 6-axis IMU only), IROS 2026 workshop Beyond Exteroception.

• Parameters: 3.37M
• Source experiment: EXP-212-rep-s10tv (epoch 30)
• Input: 2.5 s of raw 200 Hz IMU (6 ch) per scored window, mean/std-normalised with the frozen per-channel statistics shipped in model.pt
• Architecture: dilated residual TCN encoder (64/96/128/192 ch, attention + mean + std pooling) -> projection -> 2-layer bidirectional GRU (2x256) -> per-window body-frame velocity (+ auxiliary platform logits, unused at inference)
• Inference: per trajectory, chunks of 32 windows with 16 overlap so the recurrent head sees neighbouring context; overlapping predictions are averaged. Inference uses the single scored window placement (no test-time augmentation): the final submission line is single-view. fp16 autocast on CUDA.

Run inference (offline, single GPU, no internet)

pip install -r requirements.txt
python infer.py --test-root <dir of test_*.npz> \
    --windows test_windows.csv --out submission.csv

--test-root is scanned recursively for *.npz; each trajectory is matched to the traj_id column of the windows CSV by file stem. Output is a CSV with header window_id,vx,vy,vz, one row per window, sorted by window_id.

Device: CUDA is used when torch.cuda.is_available(), otherwise the code falls back to CPU automatically (no flag needed). Runtime on the 89-trajectory test set: ~1 min on one A10 (<2 GB VRAM); ~8 min on 8 CPU cores. Both paths are far inside the 2 h budget. Prediction values differ between the two paths at the 1e-3 m/s level (fp16 autocast vs fp32, reduction order); the official score differs by <2e-5, i.e. the submitted score reproduces on either device.

Training summary (see the technical report for details)

Trained on the released train+val split only (475 real trajectories) plus physics-exact time-warp augmentations of the training data (per-trajectory warp factor s ~ U(0.80, 1.50): gyro' = sgyro(st), acc' = s^2acc(st) - (s^2-1)R(st)^T*[0,0,+g], velocity labels scaled by s; generated solely from the released training split). Objective: vector-Huber velocity loss (beta 0.25) + auxiliary platform cross-entropy (0.05) + IMU-strapdown consistency (0.2) + differentiable ATE20 trajectory loss (0.05), 30 epochs, AdamW + OneCycle, from a TCN encoder warm-started on the per-window task. No external data, no test-time calibration tables, deterministic (fixed seed).

Provenance

The TCN trunk follows the public Kaggle kernel by nomannic19 (ts-tcmpio-unified-temporal-resnet), itself building on the official TartanIMU release. Everything in this repository is trained from scratch on the challenge data.