EM-KalmanNet: Learned Expectation-Maximization for Adaptive Tracking in Partially Known, Block-Wise Time-Varying State-Space Models

Abstract

State estimation in partially known state space (SS) models is challenging when the dynamics or observation model varies across short data blocks. Classical model-based approaches, such as the expectation-maximization (EM) Kalman filter, jointly recover the latent states and the unknown model parameters, but rely on linear-Gaussian assumptions that should accurately describe the system and require numerous forward-backward passes. Consequently, their performance and computational efficiency may deteriorate under complex and non-stationary real-world conditions. On the other hand, learned Kalman smoothers are robust to model mismatch yet cannot adapt at inference to unseen model variations without labeled data. In this work, we propose EM-KalmanNet, an AI-aided tracking algorithm for adaptive smoothing in blockwise time-varying SS models. The method unfolds a fixed, small number of EM-like iterations into a trainable architecture: a parameter-aware RTSNet implements a learned E-step conditioned on the current model-parameter estimate, while a lightweight M-Net implements a learned M-step that updates the state-transition or the observation matrix using empirical moments, residuals, and gradient-related statistics. The two modules are shared across the unfolded iterations and are trained offline via a dedicated three-stage procedure. During deployment, the per-block parameter estimate is propagated between consecutive blocks, enabling observation-driven adaptation without labeled online data or knowledge of the noise statistics. Experiments involving linear and nonlinear models, Gaussian and non-Gaussian noise, Lorenz attractor tracking, and acoustic source localization demonstrate that EM-KalmanNet consistently outperforms model-based and data-driven benchmarks while substantially reducing inference latency relative to the EM-KF.