MUSE: Multimodal Uncertainty Quantification of State Estimation

UnCal-Flight Dataset is an indoor flight dataset recorded on a quadrotor with a ZED2i stereo camera, its onboard IMU, and Vicon motion capture ground truth. It contains 145 trajectories spanning ten 2D/3D shapes (Rectangle, Line, Circle, 3D Line, Vertical Fig 8, Figure 8, Clover, Complex, 3D Fig 8, Spiral) under two yaw regimes:

• Yaw-constant — heading fixed while translating.
• Yaw-forward — heading facing forward.

Each shape is flown at four speed levels (0–3) and under combinations of two lighting conditions (bright / dark) and two scene conditions (clear / obs, the latter with dynamic obstacles in the workspace).

Download the full raw rosbag dataset here: https://uofi.box.com/s/xtr8srlbliwxuw6espbqqo8e5vdv2erw.

The dataset ships as ROS1 .bag files (one bag per run). Each bag contains:

Filenames follow {trajectory}_{speed}_{lighting}_{scene}.bag, e.g. 2dComplex_2_dark_clear.bag is the 2dComplex shape at speed level 2, recorded in dark lighting with a clear scene.

The Box share is laid out as:

UnCal-Flight-Dataset/
├── config/                        # sensor calibrations
│   ├── cam-body.yaml              # T_cam0_body  body 
│   ├── body-gt.yaml               # T_body_gt    mocap-marker ↔ body (z-up)
│   ├── cam-imu.yaml               # T_cam_imu    camera intrinsics (Kalibr)
│   └── imu.yaml                   # IMU noise model
└── bags/
    ├── yaw_constant/             
    │   ├── 2dComplex/
    │   ├── 3dFig8/
    │   ├── ...
    └── yaw_forward/               
        └── (same 10 trajectory folders)

Each trajectory folder holds the per-run bags.

scripts/bag_to_dataset.py reads a single .bag and writes the per-sequence directory consumed by UncertaintyDataset (see §2.2): four CSVs (IMU, ZED VIO raw, ZED VIO aligned, mocap GT), the stereo PNGs under rogx_zed_left/ and rogx_zed_right/, and imgs/timestamps.npy. rogx_zed_vo_pose_aligned.csv is the raw VIO right-multiplied by T_body_gt from config/body-gt.yaml so it tracks the mocap-marker pose in the (z-up) VIO world; this is the file the trainer loads. The raw rogx_zed_vo_pose.csv is kept alongside it for debugging / reproducibility. CSV column names match muse/datasets/uncal_flight_headers.py and %time is the integer-nanosecond header.stamp so PNG filenames and CSV times share one clock.

pip install rosbags

python scripts/bag_to_dataset.py \
    --bag PATH/TO/UnCal-Flight-Dataset/yaw_constant/2dComplex/2dComplex_2_dark_clear.bag \
    --out uncal_flight \
    --split validation \
    --regime yaw_constant

This writes uncal_flight/yaw_constant/validation/2dComplex/2dComplex_2_dark_clear/. --regime is optional — omit it for a flat <out>/<split>/<group>/<seq>/ layout. Override --group / --seq if the auto-derived names (trajectory token from the bag stem and the full stem, respectively) don't match what you want.

For other datasets (EuRoC, …), edit the five *_TOPIC constants and decode_image() at the top of the script — the rest of the conversion is generic.

This repo also includes reference implementation of MUSE on learned pose-uncertainty quantification. The model takes pre-extracted image features, IMU readings, and the VIO estimate, and predicts a 6-DoF correction with a covariance over the residual.

We tested with Python 3.10, PyTorch 2.0.1+cu118 (older torch is fine too), and mamba_ssm 2.2.2. Newer torch 2.4 also works.

conda create -n muse python=3.10 -y
conda activate muse

# 1. PyTorch + CUDA (pick the wheel matching your driver)
pip install torch==2.0.1+cu118 torchvision==0.15.2+cu118 \
    --index-url https://download.pytorch.org/whl/cu118

# 2. Mamba SSM (must be the wheel built for your torch + CUDA combo)
pip install "https://github.com/state-spaces/mamba/releases/download/v2.2.2/mamba_ssm-2.2.2+cu118torch2.0cxx11abiFALSE-cp310-cp310-linux_x86_64.whl"

# 3. The rest
pip install "numpy<2" omegaconf hydra-core einops transformers==4.44.0 \
    pandas matplotlib scipy tqdm scikit-learn pyyaml \
    natsort numpy-quaternion opencv-python-headless

# 4. Project (editable install)
pip install -e .

# 5. (optional) wandb — only if you set wandb.enable=True in the config
pip install wandb

uncal_flight (or any compatible dataset) is expected as:

<DATASET_ROOT>/
  train/
    <group_a>/
      <sequence_1>/
        rogx_zed_imu.csv          # IMU (timestamp, la_xyz, av_xyz)
        rogx_zed_vo_pose.csv      # ZED VIO pose (with 36-element covariance)
        rogx_mocap_pose.csv       # Mocap GT pose
        imgs/
          images.npy              # stacked grayscale frames (T, H, W) or (T, C, H, W)
          timestamps.npy          # (T,) ns
        features/
          img_feature.npy         # (T, 256) precomputed image features
      <sequence_2>/
        ...
    <group_b>/
      ...
  validation/
    <group_x>/
      <sequence_1>/
        ...

uncal_flight_headers.py defines the CSV column names and frame conventions for each pose source (gt, ekf, msckf, zed).

We feed MUSE the output of frozen pretrained encoders applied to each raw sensor stream:

• IMU → RoninLSTM, a BilinearLSTM in the RoNIN style. Computed on-the-fly during dataset loading (UncertaintyDataset runs the encoder once per chunk).
• Image → ImageEncoderSuperPoint, built on the MagicLeap SuperPoint v1 checkpoint.

Running SuperPoint at every training step is expensive, so we precompute image features once and let training read (T, 256) tensors from disk. This is optional — without it, the data loader can run SuperPoint inline at the cost of slower training.

scripts/extract_features.py reads <timestamp_ns>.png frames from one or two camera directories, converts each to grayscale, runs SuperPoint, and writes:

• <out>/img_feature.npy — (T, 256) float32
• <out>/timestamps.npy — (T,) int64 ns
• <out>/preview_cam0.png (and preview_cam1.png if stereo) — optional, with --save_preview, the actual grayscale fed to the encoder so you can sanity-check the preprocessing.

Stereo input fuses the two views (SuperPoint(L) + SuperPoint(R) → concat 512 → Linear(512, 256)); mono input uses the SuperPoint avg-pool output directly. Either way the per-frame feature is 256-D.

# uncal_flight: all sequences under <root>[/<regime>]/{train,validation}/<group>/<seq>/
python scripts/extract_features.py --root PATH/TO/uncal_flight --save_preview

# uncal_flight: a single sequence
python scripts/extract_features.py \
    --seq PATH/TO/uncal_flight/yaw_constant/validation/2dComplex/2dComplex_2_dark_clear

# Generic stereo (any dataset; PNG filenames must be <timestamp_ns>.png)
python scripts/extract_features.py \
    --cam0 PATH/TO/cam0 --cam1 PATH/TO/cam1 \
    --out  PATH/TO/features --save_preview

# Generic mono (omit --cam1)
python scripts/extract_features.py \
    --cam0 PATH/TO/cam0 --out PATH/TO/features

Useful flags: --batch_size 16 (default 8), --device cpu, --superpoint_ckpt PATH (or set SUPERPOINT_CKPT_PATH; default is the bundled muse/models/encoders/superpoint/superpoint_v1.pth).

python scripts/train_uncertainty.py \
    dataset.path=PATH/TO/uncal_flight \
    train.pretrained_path=null \
    wandb.enable=False

Two-stage training is supported. A typical run is:

# Stage 1 — train the mean head only
python scripts/train_uncertainty.py \
    dataset.path=$DATA \
    model.decoder.predict=mean_only_error \
    train.freeze_mean=False

# Stage 2 — freeze the trained mean and learn the covariance head
python scripts/train_uncertainty.py \
    dataset.path=$DATA \
    model.decoder.predict=error \
    train.freeze_mean=True \
    train.pretrained_path=_ckpt/uncal_flight_mean_only_error/<RUN_ID>/model_val.pt

Best checkpoints (model_val.pt) and per-chunk validation dumps are written to _ckpt/<wandb.group>/<run-timestamp>/ whenever validation loss improves.

python scripts/validate_model.py \
    --config configs/uncal_flight.yaml \
    --weights _ckpt/<group>/<run>/model_val.pt \
    --out outputs/eval/

Each validation chunk is dumped under outputs/eval/dumped_output/<seq>/<chunk_idx>/:

muse/
  datasets/                 # data loading + chunking
    uncertainty_dataset.py
    dataset_utils.py
    uncal_flight_headers.py
  losses/                   # SE(3) ops + DICE / NLL losses
    lie_algebra.py
    losses.py
  metrics/                  # RMSE / geodesic / NLL
    metrics.py
  models/
    encoders/
      ronin_lstm.py         # IMU encoder (BilinearLSTM, RoNIN-style)
      superpoint.py         # image feature encoder
    uncertainty/
      model.py              # main UncertaintyModel (Mamba backbone + heads)
  trainers/
    trainer.py              # base train loop + checkpoint dumping
    uncertainty_trainer.py  # UncertaintyTrainer with NLL/DICE loss
  visualization/
    trainer_visualizer.py   # per-batch plots + reliability/ECE
    video_generator.py      # 2D/3D trajectory animation w/ uncertainty ellipses
  utils/
    transformation.py
scripts/
  extract_features.py       # SuperPoint preprocessing (raw PNGs -> img_feature.npy)
  train_uncertainty.py      # Hydra entry point for training
  validate_model.py         # standalone evaluation script
configs/
  uncal_flight.yaml         # example training config

@article{kim2026muse,
  author  = {Kim, Minkyung and Che, Henry and Chandaka, Bhargav and Pramuanpornsatid, Bhumsitt and Yang, Chengyu and Cheng, Sheng and Wang, Xiaofeng and Hovakimyan, Naira and Wang, Shenlong},
  title   = {MUSE: Multimodal Uncertainty Quantification of State Estimation},
  journal = {ICRA},
  year    = {2026},
}