Learned Memory Attenuation in Sage-Husa Kalman Filters for Robust UAV State Estimation

This is the official implementation of the N-Deep Recurrent Sage-Husa Kalman Filter (NDR-SHKF). Built in JAX/Equinox, this repository provides a fully differentiable filtering framework that replaces the static scalar forgetting factor of the classical Sage-Husa filter with a learned, vector-valued memory attenuation policy.

Authors: Kenan Majewski, Marcin Żugaj

TL;DR | Installation | Usage | Citation

TL;DR

Traditional Kalman filters are limited by rigid, stationary noise assumptions that prevent effective adaptation to complex dynamic environments. The NDR-SHKF resolves this by employing a hierarchical recurrent neural network to act as a dynamic hyper-controller for the Sage Husa filter. By replacing static parameters with a learned, vector-valued memory attenuation policy, this differentiable architecture bridges the gap between classical Bayesian stability and deep learning, enabling robust state estimation during sensor outages where conventional adaptive estimators lose convergence.

Installation

We use uv for fast Python package management.

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

# Install dependencies
uv sync

# Activate virtual environment
source .venv/bin/activate

# Install dev dependencies (includes ruff, pyright)
uv sync --extra dev

# (Optional) Install with GPU support
uv sync --extra gpu

Usage

Quick Start

Run the complete pipeline with a single command:

# Train, benchmark, and plot (both simulated and UAV data)
ndrshkf run

# Or step-by-step:
ndrshkf train --epochs 2000
ndrshkf benchmark --runs 1000
ndrshkf plot

Training

Train a new model with custom hyperparameters:

# Train on both simulated and UAV data (default)
ndrshkf train

# Train only on simulated data (Lorenz/Rössler)
ndrshkf train --sim

# Train only on UAV data
ndrshkf train --uav

# Custom epochs
ndrshkf train --epochs 2000

# Load config from file
ndrshkf train --config config.json

Benchmarking

Evaluate model performance against standard EKF and SHKF:

# Run benchmark on both simulated and UAV data (default)
ndrshkf benchmark

# Benchmark only simulated data
ndrshkf benchmark --sim

# Benchmark only UAV data
ndrshkf benchmark --uav

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

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

Visualization

Generate plots from benchmark results:

# Generate all plots (both simulated and UAV)
ndrshkf plot

# Plot only simulated results
ndrshkf plot --sim

# Plot only UAV results
ndrshkf plot --uav

Full Pipeline

Run complete pipeline with training, benchmarking, and plotting:

# Run everything (both sim and UAV)
ndrshkf run

# Run only simulated pipeline
ndrshkf run --sim

# Run only UAV pipeline
ndrshkf run --uav

# Skip training, only benchmark and plot
ndrshkf run --no-train

Configuration

You can save and load configuration files:

from ndrshkf import Config, UAVConfig

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

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

Then use with CLI:

ndrshkf train --config config.json
ndrshkf benchmark --config uav_config.json

Scripts

Ablation Study

Ablation across three axes: input representation (whitened, raw, whitened_no_log, nis), recurrent depth (1, 3, 5), and auxiliary decoder weight (lambda_aux=1.0, lambda_aux=0.1, lambda_aux=0.01, lambda_aux=0.0).

# Train all ablation variants (all axes, 100 seeds each)
uv run python scripts/ablation.py train

# Train a specific axis only
uv run python scripts/ablation.py train input
uv run python scripts/ablation.py train depth
uv run python scripts/ablation.py train decoder

# Benchmark trained models (10,000 Monte Carlo runs per seed)
uv run python scripts/ablation.py benchmark
uv run python scripts/ablation.py benchmark depth

# Generate IQR box plots and violin plots
uv run python scripts/ablation.py plot
uv run python scripts/ablation.py plot input

Cross-Dynamics Transfer Benchmark

Evaluates a UAV-trained model on low-dimensional chaotic systems (Lorenz, Rössler) using semantic input mapping with zero-fill.

# Run with default UAV model
uv run python scripts/cross_dynamics_benchmark.py

# Custom model and settings
uv run python scripts/cross_dynamics_benchmark.py --model models/uav.eqx --runs 10000 --seed 43

# Save results
uv run python scripts/cross_dynamics_benchmark.py --save results/cross_dynamics.npz

Latency Benchmark

Times the full loop over UAV data, comparing EKF vs NDR-SHKF per-step latency.

uv run python scripts/single_step_uav.py

Project Structure

ndrshkf/
├── __init__.py         # Package exports
├── cli.py              # Command-line interface
├── config.py           # Configuration management (Config, UAVConfig)
├── model.py            # NDR-SHKF neural network architecture
├── dynamics.py         # Dynamical systems (Lorenz, Rössler, UAV)
├── train.py            # Training logic (simulated and UAV data)
├── benchmark.py        # Monte Carlo evaluation
├── plot.py             # Visualization utilities
├── utils.py            # Model I/O and data loading helpers
└── filters/
    ├── ekf.py           # Extended Kalman Filter
    ├── shkf.py          # Sage-Husa Kalman Filter
    ├── kalmannet.py     # KalmanNet Filter
    ├── enn.py           # Elman Neural Network Adaptive Kalman Filter
    ├── lekf.py          # Learned Extended Kalman Filter
    ├── vbakf.py         # Variational Bayesian Adaptive Kalman Filter 
    └── utils.py         # Filter utilities
scripts/
├── ablation.py                  # Ablation studies
├── cross_dynamics_benchmark.py  # Cross-dynamics transfer benchmark
└── single_step_uav.py           # Single-step UAV evaluation

Citation

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

@misc{majewski2026learnedmemoryattenuationsagehusa,
      title={Learned Memory Attenuation in Sage-Husa Kalman Filters for Robust UAV State Estimation}, 
      author={Kenan Majewski and Marcin Żugaj},
      year={2026},
      eprint={2605.18704},
      archivePrefix={arXiv},
      primaryClass={eess.SP},
      url={https://arxiv.org/abs/2605.18704}, 
}

License

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