Robust Unscented Kalman Filtering via Recurrent Meta-Adaptation of Sigma-Point Weights

This is the official implementation of the Meta-Adaptive Unscented Kalman Filter (MA-UKF). Built in JAX/Equinox, this repository provides a fully differentiable filtering framework that uses a recurrent neural policy to dynamically adapt sigma-point weights in real-time for robust nonlinear state estimation.

Authors: Kenan Majewski, Michał Modzelewski, Marcin Żugaj, Piotr Lichota

TL;DR | Installation | Usage | Citation

TL;DR

The standard Unscented Kalman Filter (UKF) relies on static scaling parameters, limiting its robustness against unmodeled maneuvers and heavy-tailed sensor noise. The Meta-Adaptive UKF (MA-UKF) solves this by reformulating sigma-point weight synthesis as a continuous control problem. By using a small neural network to continuously analyze measurement residuals, the MA-UKF dynamically adapts its mean and covariance weights at every time step. Trained end-to-end, this lightweight approach effectively filters out extreme radar glint and seamlessly generalizes to unseen, out-of-distribution trajectories with negligible computational overhead.

Installation

We use uv for fast Python package management.

# Clone the repository
git clone https://github.com/kenanmajewski/maukf
cd maukf

# Install dependencies
uv sync

# Activate virtual environment
source .venv/bin/activate

# Install dev dependencies
uv sync --extra dev

Usage

Quick Start

Run the complete pipeline with a single command:

# Train, benchmark, and plot
maukf run

# Or step-by-step:
maukf train --epochs 1800
maukf benchmark --runs 1000
maukf plot

Training

Train a new model with custom hyperparameters:

# Basic training
maukf train

# Custom epochs and learning rate
maukf train --epochs 2000 --learning-rate 1e-4

# Save to custom path
maukf train --output model.eqx

Benchmarking

Evaluate model performance against standard UKF and IMM-UKF:

# Run benchmark with default settings
maukf benchmark

# Custom number of Monte Carlo runs
maukf benchmark --runs 1000 --model model.eqx

# Save results to custom path
maukf benchmark --output results.npz

Visualization

Generate plots from benchmark results:

# Generate all plots
maukf plot

# Plot specific trajectory
maukf plot --data results.npz --index 8

UKF Parameter Search

Perform hyperparameter search for the standard UKF:

# Run parameter search 
maukf search

# Custom parameters
maukf search --trials 200 --runs 200
# Custom output path
maukf search --output search_results.json

Configuration

You can save and load configuration files:

from maukf import Config

# Save current config
config = Config(n_epochs=2000)
config.save("config.json")

# Load config
config = Config.load("config.json")

Then use with CLI:

maukf train --config config.json

Project Structure

maukf/
├── __init__.py             # Package exports
├── benchmark.py            # Monte Carlo evaluation
├── cli.py                  # Command-line interface
├── config.py               # Configuration management
├── dynamics.py             # Target motion models
├── filters.py              # UKF and IMM-UKF implementations
├── model.py                # MAUKF neural network
├── parameter_search.py     # Search for UKF parameters
├── plot.py                 # Visualization utilities
├── train.py                # Training loop
└── utils.py                # Model I/O helpers

Citation

If you use this code in your research, please cite:

@misc{majewski2026robustunscentedkalmanfiltering,
      title={Robust Unscented Kalman Filtering via Recurrent Meta-Adaptation of Sigma-Point Weights}, 
      author={Kenan Majewski and Michał Modzelewski and Marcin Żugaj and Piotr Lichota},
      year={2026},
      eprint={2603.04360},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2603.04360}, 
}

License

This project is licensed under the MIT License - see the LICENSE file for details.