CunYv/tartanimu-challenge-unified

TartanIMU Challenge — single unified model (team Tong Joye)

Fine-tuned TartanIMU foundation model for the IROS 2026 TartanIMU Challenge: Multi-Platform Inertial Odometry. This repository contains one model with one shared set of weights; it reproduces the Kaggle submission.

What the model does

Input: a 1.0 s window of 6-axis IMU (accelerometer + gyroscope, gravity retained) sampled at 200 Hz. Output: the mean 3-D body-frame velocity (vx, vy, vz) in m/s for that window, for any of the four embodiments (car, quadruped, drone, handheld).

Architecture: the released TartanIMU Foundation_Model trunk (ResNet + LSTM, 5.72 M parameters) with its four per-embodiment heads, plus a router (128 -> 128 -> 4, ~17 k parameters) evaluated on the shared trunk features. The exported velocity is the router-weighted soft combination of the four heads, followed by a small bidirectional GRU refiner (hidden 64, one layer, ~25 k parameters) that consumes the whole predicted velocity sequence and adds a zero-initialised residual. There is no external platform classifier and no separately selected expert network.

On the official full 89-sequence test set (per-sequence scoring service): TartanIMU Score 0.37177 (macro AVE 0.312 m/s, macro ATE20 0.912 m). Submission MD5 64ce896ea931a310d814ff6eb59022fc.

Files

File	Description
checkpoint.pt	fine-tuned weights (single model, 5.74 M parameters)
router_model.py	model definition (trunk + heads + internal router)
tartan_imu/	vendored copy of the official Apache-2.0 TartanIMU package (see TARTANIMU_LICENSE), so inference needs no external clone
infer.py	end-to-end inference: test .npz files -> submission CSV
config/unified.yaml, config/resnet_lstm_multihead.yaml	model configuration

Environment

python >= 3.10
torch >= 2.0
numpy
pyyaml
tartan_imu  # the official package: pip install -e ".[logging]" from github.com/superxslam/TartanIMU

Reproduce the submission

The official tartan_imu package is vendored in this repository, so nothing has to be cloned or installed beyond torch, numpy and pyyaml.

python infer.py \
  --test-root  /path/to/test \
  --windows    /path/to/index/test_windows.csv \
  --out        submission.csv \
  --checkpoint checkpoint.pt \
  --config     config/unified.yaml \
  --seq-len    1024

This writes a window_id,vx,vy,vz CSV with 30,644 rows (the same rows as sample_submission.csv). Inference takes ~2 minutes on an RTX 4060 Laptop GPU and runs one deterministic forward pass per 512-window chunk; no test-time augmentation, no ensembling and no checkpoint averaging are used.

Training summary

Initialised from the organizers' released unified checkpoint and fine-tuned for 150 epochs on all labelled challenge data (train + val: 105,645 windows / 475 trajectories). Loss: masked mean Euclidean velocity error with metric-normalised weights w_i ∝ 1 / (windows in trajectory × trajectories of platform), further modulated by each window's rotational dynamics (1 + alpha z, gyroscope magnitude, alpha = 1), plus a router cross-entropy term (weight 0.3) and a differentiable ATE20 surrogate (weight 0.05). AdamW, lr 1e-4 (trunk/heads) and 5e-4 (router), cosine schedule with 3-epoch warm-up, batch 16 sequences of 32 windows, bf16 autocast. Windows are decimated 200 Hz -> 40 Hz and channels ordered [gyro | accel], exactly as the released model expects.

Candidate checkpoints were scored with the organizers' per-sequence scoring service and this one was selected because it has the best official full-test score of all our candidates; the labelled validation split was folded into training, so it cannot be used for selection.

Compliance declaration

1. Single model with one shared set of weights — yes (this checkpoint).
2. Platform classification / routing / specialisation at inference — an internal learned router on shared trunk features (explicitly allowed by the challenge rules); no separately trained expert models.
3. Ensembling / checkpoint averaging / test-time augmentation — none.
4. Test-set leakage — no test label was ever seen by any training procedure, and no pseudo-labelling, transductive training or test-input adaptation was used. Architecture and hyper-parameters were chosen on the labelled val split; the final checkpoint was selected among our trained candidates by the organizers' official per-sequence scoring service (selection therefore used the official evaluation, which we state explicitly).
5. Development vs submitted system — identical apart from the selected epoch.