Skip to content

Repository files navigation

TartanIMU

A Light "Foundation Model" for Inertial Positioning in Robotics
CVPR 2025

CVPR 2025 paper TartanIMU project website Live TartanIMU inference TartanIMU model weights on Hugging Face TartanIMU Kaggle challenge TartanIMU challenge setup guide

TartanIMU drone inertial odometry demo

Overview

TartanIMU learns a shared inertial representation across ground vehicles, quadrupeds, drones, and humans. Given 6-axis accelerometer and gyroscope measurements, it predicts 3D body-frame velocity for inertial positioning.

Pretrain
Pretrain
Generalize
Generalize
Adapt
Adapt
Deploy
Deploy
100+ hours 4 platforms 36% 200 FPS
Multi-platform training data Car, quadruped, drone, human Reported ATE improvement Reported online adaptation speed

The released implementation provides a ResNet-LSTM multi-head foundation model, pretrained inference, configurable training and evaluation, and an IROS 2026 challenge starter kit.

Note

This public release includes the LSTM-based Foundation_Model. The Transformer registration is retained for compatibility, but its core is not included. Selecting model_name: Transformer raises NotImplementedError.

Quick Start

Requirements: Python 3.10+ and PyTorch 2.0+.

Install

git clone https://github.com/superxslam/TartanIMU.git
cd TartanIMU

# Library and inference dependencies
pip install -e .

# Add experiment tracking for the training CLI
pip install -e ".[logging]"

# Add linting and tests for development
pip install -e ".[logging,dev]"

Confirm that the released model is available:

python -c "from tartan_imu.model.registry import available; print(available())"
# ['Foundation_Model', 'Transformer']

Run Pretrained Inference

Download the released configuration and weights from the TartanIMU model repository:

pip install huggingface_hub
huggingface-cli download Tartan-IMU/TartanIMU \
  --local-dir ./tartanimu_weights

Run inference on one trajectory:

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

See example/minimal_example.py for a complete fine-tuning and evaluation example.

Train and Evaluate

main_net.py is the entry point for both training and evaluation. Experiment behavior is defined by a YAML configuration.

Run the small end-to-end smoke experiment:

WANDB_MODE=disabled CUDA_VISIBLE_DEVICES=0 \
  python main_net.py \
  --config ./config/datasets/tartanimu/tartan_imu_multihead_smoke.yaml

Train on one or more GPUs:

# Single GPU
CUDA_VISIBLE_DEVICES=0 \
  python main_net.py --config <path/to/config.yaml>

# Multi-GPU: also set train.use_multi_gpu: True in the YAML
CUDA_VISIBLE_DEVICES=0,1,2,3 \
  python main_net.py --config <path/to/config.yaml>

Evaluate a checkpoint:

CUDA_VISIBLE_DEVICES=0 \
  python main_net.py \
  --config <path/to/config.yaml> \
  --checkpoint <path/to/checkpoints/best_model.pt>

Command-Line Options

Option Purpose
--config, --yaml Experiment YAML path
--checkpoint Load weights for evaluation or warm start
--resume_from Resume model, optimizer, scheduler, and AMP state
--exp_name Override the experiment and output name
--pdb Use single-process debug mode and disable W&B

Configuration

Dataset configurations live under config/datasets/. Their model.model_yaml field points to a model definition such as config/resnet_lstm_multihead.yaml.

The release includes two TartanIMU experiment configs:

Config Purpose
tartan_imu_dataset.yaml Full car, drone, dog, and human training/evaluation
tartan_imu_multihead_smoke.yaml Short single-GPU smoke run used by the example and tests
Section Key settings
data Dataset reader, platform paths, split names, and sample rates
model Registered model name, model YAML, and prediction targets
train Output directory, epochs, AMP, and multi-GPU behavior

Training artifacts are written to train.out_dir. Checkpoints are stored in <train.out_dir>/checkpoints/.

Data

Library Dataset Format

Each configured dataset root contains trajectory files grouped by split:

<dataset_root>/
|-- train/
|-- val/
`-- test/

The included AirLab reader expects synchronized arrays in each .npz file:

Key Description
retargetted_ts Timestamps
retargetted_imu Accelerometer and gyroscope measurements
retargetted_pos Ground-truth position
retargetted_quat Ground-truth orientation in xyzw order

Available readers:

  • AirLab
  • Humanoid
  • HumanoidPostProcessed
  • HumanoidPostProcessedCached

Select a reader with data.dataset. To support another format, add a module under tartan_imu/dataloader/ and register it in tartan_imu/utils/registry.py.

Challenge Dataset Format

The TartanIMU Challenge uses a separate window-level format with 1-second, 200-frame IMU windows. It does not use the retargetted_* schema above.

train/<platform>/*.npz
val/<platform>/*.npz
test/test_*.npz
index/*_windows.csv
sample_submission.csv

The test split is anonymized and contains neither pose nor platform labels. See the challenge starter guide for the complete schema, submission commands, and evaluation protocol.

IROS 2026 TartanIMU Challenge

The challenge evaluates one shared model across car, dog/legged, drone, and human motion. Given a 1-second IMU window, the model predicts mean body-frame velocity (vx, vy, vz).

Resource Purpose
Kaggle competition Join the challenge and check the current schedule, rules, submissions, and leaderboard
Challenge setup guide Follow the complete data, training, evaluation, and submission workflow
Challenge dataset Access the multi-platform training and validation data (Hugging Face access may be required)
Foundation model Download the released unified configuration and model weights
Live model demo Explore the reference models interactively

Submissions are scored with the TartanIMU Score, a dimensionless combination of 60 % per-window Absolute Velocity Error (AVE, m/s) and 40 % 20-meter segment Absolute Trajectory Error (ATE20, m):

TartanIMU Score = 0.6 * (AVE / 0.7356384388)  +  0.4 * (ATE20 / 3.1160277267)

Both components are macro-averaged so that all four platforms weigh equally, and each is normalized by the value the all-zero submission reaches on the test set, which makes the score dimensionless and pins an all-zero submission to exactly 1.000. Lower is better; the released baseline scores 0.637 on the public split.

File Purpose
starter/starter.ipynb Data-to-submission walkthrough
starter/baseline_submission.py Valid zero or constant baseline
starter/tartanimu_submission.py Released model inference
starter/kaggle_metric_tartanimu_score.py Leaderboard metric for validation

Predictions must come from one model with one shared set of weights. Platform-specific internal routing is allowed, but four separately selected expert models are not.

The released weights and full model card are available at Tartan-IMU/TartanIMU. Review the model card for artifact-specific terms and known limitations.

Repository Structure

TartanIMU/
|-- tartan_imu/
|   |-- config/          # Configuration loading and object construction
|   |-- dataloader/      # Dataset readers
|   |-- evaluation/      # Metrics and trajectory analysis
|   |-- model/
|   |   |-- backbones/   # Registered model builders
|   |   |-- common/      # Shared blocks, losses, and helpers
|   |   `-- lstm/        # Released ResNet-LSTM model
|   |-- training/        # Trainer, checkpoints, and plots
|   `-- utils/           # Logging, constants, and registries
|-- config/              # Model and experiment YAML files
|-- doc/                 # README media
|-- example/
|   |-- inference_example.py
|   `-- minimal_example.py
|-- starter/             # Challenge metric, baselines, and notebook
|-- tests/unit/          # Unit and characterization tests
|-- tools/               # Dataset, analysis, and plotting utilities
|-- main_net.py
|-- train.py
`-- test.py

Development

pip install -e ".[logging,dev]"
ruff check tartan_imu/
pytest tests/unit -q

See CONTRIBUTING.md for code style, naming, testing, checkpoint compatibility, and pull request guidelines.

Analysis and plotting helpers are available under tools/, including drift_analysis.py, plot_2d_traj.py, and gen_experiment_entry.py.

Citation

If TartanIMU supports your research, please cite the CVPR 2025 paper:

@inproceedings{zhao2025tartan,
  title={Tartan IMU: A Light Foundation Model for Inertial Positioning in Robotics},
  author={Zhao, Shibo and Zhou, Sifan and Blanchard, Raphael and Qiu, Yuheng and Wang, Wenshan and Scherer, Sebastian},
  booktitle={Proceedings of the Computer Vision and Pattern Recognition Conference},
  pages={22520--22529},
  year={2025}
}

License

TartanIMU is released under the Apache License 2.0.

Copyright 2026 Shibo Zhao.

About

First Attempt on IMU "Foundation Model"

Resources

Contributing

Stars

76 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages