TartanIMU / README.md
YiZhaoJasper's picture
Fix metric labelling: leaderboard metric is macro 20m-segment ATE, not window-RMSE; add challenge-metric tables + platform_id off-by-one note
adc5165 verified
|
Raw
History Blame Contribute Delete
10.9 kB
metadata
license: mit
library_name: tartan_imu
pipeline_tag: other
tags:
  - inertial-odometry
  - imu
  - robotics
  - velocity-estimation
  - state-estimation
  - tartanimu
  - lstm
language:
  - en
model-index:
  - name: TartanIMU-unified
    results:
      - task:
          type: inertial-odometry
          name: Body-frame velocity regression (Macro-ATE)
        metrics:
          - type: ATE
            name: >-
              Challenge leaderboard metric — macro 20 m-segment ATE (m, v4 test,
              all)
            value: 1.261
          - type: ATE
            name: Macro-ATE all-4 (m, 20 m seg, test.py on tartan_layout test)
            value: 2.327
          - type: ATE
            name: Unified drone ATE (m, drone_v3 test, 20 m seg, test.py)
            value: 3.945
          - type: RMSE
            name: >-
              Offline diagnostic — macro window-RMSE (v3 test, all; NOT the
              leaderboard metric)
            value: 0.664

TartanIMU — Multi-Platform Inertial Foundation Model

TartanIMU is a light foundation model for neural inertial positioning. Given a 6-axis IMU stream (3-axis accelerometer + 3-axis gyroscope), the network regresses per-window body-frame velocity; position is obtained by integrating velocity with the platform orientation. It is the model released alongside the IROS 2026 Tartan IMU Challenge.

This repository ships one multi-platform unified checkpoint (4 motion-type heads: car / dog / drone / human) plus four single-platform expert checkpoints for reproducing the per-platform ablations.

v3 (2026-07-18). The drone platform is now drone_v3 — EuRoC MAV + UZH-FPV + Blackbird, hover → 20 m/s real flight — replacing the low-dynamic indoor-mocap drone_v2. unified.pt and expert_drone.pt are retrained on it. Headline change: the car/dog multi-task collapse of previous unified releases is gone (see Results). car/dog/human expert checkpoints unchanged.

wdfix2 (2026-07-24). unified.pt now exempts the per-motion-type head parameters from weight decay. The previous unified.pt had a dead output channel — drone forward-velocity (vx) was stuck at constant zero — because weight decay is applied every step while each head only receives gradients on its own platform's samples, driving weakly-supervised head channels to zero. The fix revives drone-vx (dense-forward correlation up to 0.80 vs ground truth, was 0). On the offline window-RMSE diagnostic it goes 0.741 → 0.664, and on test.py 20 m-segment macro ATE 2.52 → 2.33. On the challenge leaderboard metric the gain is concentrated in drone (2.47 → 2.21 m) and the four-platform macro is a wash (1.23 → 1.26 m) — see Results. Expert checkpoints unchanged.

  • Backbone: ResNet-1D feature trunk + 2-layer LSTM (hidden 128) → per-motion-type OutputHead
  • Model name (code): Foundation_Model
  • Input: 6-axis IMU, raw 200 Hz, resampled to 40 Hz network input, 1.0 s windows
  • Output: 3-D body-frame velocity (vx, vy, vz) per window
  • Coordinate frame: honest body frame with gravity present (use_local_coord: true) — no ground-truth orientation is used to de-rotate the IMU input (no leakage)
  • License: MIT © Shibo Zhao
  • Code: https://github.com/superxslam/TartanIMU

Files

File Description
checkpoints/unified.pt Flagship 4-head unified model (car/dog/drone/human), drone_v3 + wdfix2 (heads exempt from weight decay)
checkpoints/expert_car.pt Single-platform expert — car
checkpoints/expert_dog.pt Single-platform expert — dog (legged)
checkpoints/expert_drone.pt Single-platform expert — drone (drone_v3: EuRoC+UZH-FPV+Blackbird)
checkpoints/expert_human.pt Single-platform expert — human
config/resnet_lstm_multihead.yaml Network definition (model.model_yaml)
config/unified.yaml Experiment config used to train the unified model (drone_v3 paths + per-platform speed gates)
inference_example.py Self-contained: checkpoint → body-frame velocity

Each checkpoint is a torch.save dict with a model_state_dict key.

Results

All numbers on the challenge test split, LSTM Foundation_Model, use_local_coord: true. ⚠️ drone numbers are on the drone_v3 test (real flight, up to 20 m/s) and must not be compared with older drone_v2-era numbers (that task was far easier).

1. Challenge leaderboard metric — macro 20 m-segment ATE

This is the metric the Kaggle leaderboard reports: per-window body-frame velocities are integrated with ground-truth orientation inside each ~20 m segment of the ground-truth path, SE(3)-aligned (Umeyama, no scale), scored by RMS position error, then averaged over segments → trajectories → platforms. Metres, lower is better. Evaluated on the anonymized v4 challenge test split (89 trajectories / 30,644 windows).

Split Previous unified.pt This unified.pt (wdfix2)
all 1.230 1.261
public 1.535 1.587
private 1.029 1.048

Per-platform (all split):

Platform Previous unified.pt This unified.pt (wdfix2)
car 0.947 1.174
dog 0.550 0.640
drone 2.473 2.205
human 0.951 1.023

Read this honestly: the dead-channel fix does what it was meant to do — drone improves 10.8 % — but on this ruler the other three platforms give a little back, so the four-platform macro is a wash (1.23 → 1.26 m, within run-to-run noise). The fix is justified by the revived drone-vx channel and by the test.py ruler below, not by the challenge macro.

Reference points on the same metric (all / public / private): ground-truth velocities 0.152 / 0.162 / 0.147 (the metric's floor), per-platform mean velocity 3.496 / 3.154 / 3.870, all-zeros 3.116 / 3.058 / 3.103.

2. test.py evaluation — a different split and integration granularity

⚠️ Not comparable with the table above. These numbers come from the repository's test.py (20 m segments, IMU-rate integration) on the tartan_layout test split — a different test set from the challenge one.

Platform Expert Unified 4-head (this unified.pt, wdfix2)
car 1.704 1.666
dog 2.009 1.924
drone (v3 test) 3.921 3.945
human 1.823 1.773
macro (all 4) 2.364 2.327

On this ruler the unified 4-head model matches or beats every single-platform expert.

3. Offline window-RMSE diagnostic — not the leaderboard metric

Mean body-frame velocity error per 1 s window. Useful for spotting dead channels (that is how the drone-vx collapse was found), but the competition does not score submissions this way.

Split Previous unified.pt This unified.pt (wdfix2)
all 0.741 0.664
public 0.844 0.731
private 0.614 0.587

(Reference baselines: all-zeros 0.941, platform-mean 0.931. Per-platform window RMSE of this model, all split: car 0.270 / dog 0.263 / drone 1.961 / human 0.164.)

What changed vs the previous release: the earlier 4-head unified collapsed on car (7.6) and dog (5.6); replacing the near-static mocap drone data with drone_v3 removed that data-quality interference (near-static drone windows are feature-space neighbors of slow human motion, and the shared trunk paid for it). The wdfix2 update then fixed a dead output channel — the previous unified.pt predicted a constant-zero drone forward velocity (drone-vx) because weight decay drove the weakly-supervised head to zero; exempting the head parameters from weight decay revives it and drives the drone window RMSE from 2.26 to 1.96 (macro window-RMSE 0.741 → 0.664) and the drone challenge-metric ATE from 2.47 to 2.21 m. Aggressive UZH-FPV flight remains the hardest regime.

Quick start

git clone https://github.com/superxslam/TartanIMU
cd TartanIMU
pip install -e .

# download this model repo (weights + config)
huggingface-cli download Tartan-IMU/TartanIMU --local-dir ./tartanimu_weights

python inference_example.py \
  --config ./tartanimu_weights/config/unified.yaml \
  --model  ./tartanimu_weights/checkpoints/unified.pt \
  --npz    <path/to/trajectory.npz> \
  --motion_type human

See inference_example.py for a minimal load → predict loop. The model call is:

pred = model(imu_window, motion_type=label, predict_cov=False, compute_all_heads=False)
velocity_body = pred[motion_type]   # (B, 3) body-frame velocity

Motion-type IDs (this model's API): car=1, dog=2, drone=3, human=4.

⚠️ Off-by-one with the dataset. The challenge .npz files store platform_id 0-based (car=0, dog=1, drone=2, human=3), while the model's motion_type is 1-based as above. Add 1 when feeding a dataset platform_id into the model.

Intended use & limitations

Intended use. Body-frame velocity estimation from wearable / platform-mounted IMU for short-horizon inertial odometry, and as a starting point for the Tartan IMU Challenge.

Known limitations (be honest about these):

  1. Aggressive-flight observability. With honest body-frame IMU (gravity present, no external orientation), high-speed drone velocity is weakly observable from 1 s windows: at the window level this model reaches drone RMSE 1.96 (vs platform-mean baseline 2.41), and trajectory-level ATE on UZH-FPV racing remains several meters. This is the open problem the v3 benchmark deliberately exposes.
  2. Reproduction pitfall. The release dataloader historically gated windows at 5 m/s (max_v_norm) and clipped GT speed at 15 m/s. The shipped config/unified.yaml raises both to 25 m/s for drone; without that, the fast half of drone_v3 is silently discarded and none of these numbers reproduce.
  3. Slow-regime bias elsewhere. car/dog/human corpora remain dominated by moderate speeds; generalization to extreme regimes on those platforms is untested.

Training data

Challenge dataset v3 (car / dog / drone / human), de-duplicated with verified zero train/val/test content leakage (sha256 window fingerprints). Drone = drone_v3: EuRoC MAV (In Copyright – Non-Commercial Use Permitted), UZH-FPV (CC BY-NC-SA 3.0), Blackbird (MIT); 200 Hz, whole-recording splits, UZH same-flight IMU-stream pairs bound to one split. The drone-derived weights are therefore research / non-commercial use; cite Burri 2016 / Delmerico 2019 / Antonini 2020.

Citation

@misc{tartanimu2026,
  title  = {TartanIMU: A Light Foundation Model for Neural Inertial Positioning},
  author = {Zhao, Shibo and others},
  year   = {2026},
  note   = {IROS 2026 Tartan IMU Challenge}
}