Extended and Unscented Kalman Filters

"Whenever the world bends, I pretend it is a straight line just long enough to take one step. Mostly it works. The times it does not, I call divergence and blame the initialization."

An Extended Kalman Filter Linearizing Its Way Through Life

In Section 7.2 we built the linear Kalman filter and proved a small miracle: if the state evolves by a matrix multiply, the measurement is a matrix multiply, and every noise source is Gaussian, then the belief about the state stays Gaussian forever, and the filter only ever has to track a mean vector $\hat{\mathbf{z}}$ and a covariance matrix $\mathbf{P}$. That miracle rests entirely on linearity. A Gaussian pushed through a linear map comes out Gaussian; a Gaussian pushed through almost anything else comes out lumpy, skewed, sometimes bimodal, and no longer summarizable by a mean and a covariance. Real systems are full of those almost-anything-else maps: a radar measures range and bearing, which are trigonometric functions of Cartesian position; a pendulum's angle evolves by a sine; a chemical reactor's temperature couples to concentration through an exponential Arrhenius term. This section is about filtering when the state-space model

has a nonlinear transition $f$ or a nonlinear measurement $h$ (or both), with process noise $\mathbf{w}_t \sim \mathcal{N}(\mathbf{0}, \mathbf{Q})$ and measurement noise $\mathbf{v}_t \sim \mathcal{N}(\mathbf{0}, \mathbf{R})$ as in Appendix A's unified notation.

1. Why the Plain Kalman Filter Is Not Enough Beginner

The linear filter's whole reason for working is the closure property of Gaussians under affine maps: if $\mathbf{x} \sim \mathcal{N}(\boldsymbol{\mu}, \boldsymbol{\Sigma})$ and $\mathbf{y} = \mathbf{A}\mathbf{x} + \mathbf{b}$, then $\mathbf{y} \sim \mathcal{N}(\mathbf{A}\boldsymbol{\mu} + \mathbf{b}, \, \mathbf{A}\boldsymbol{\Sigma}\mathbf{A}^{\top})$, exactly, with no approximation. The mean and covariance are all you ever need because they are all there ever is. Replace $\mathbf{A}\mathbf{x} + \mathbf{b}$ with a curved function $g(\mathbf{x})$ and that guarantee evaporates. The transformed density $p(\mathbf{y})$ is generally not Gaussian at all, so the very objects the filter tracks, a mean and a covariance, no longer describe the belief; they merely summarize its first two moments while throwing away everything else.

To see why this matters operationally, recall what the linear filter actually does each step: it predicts the next measurement as a linear function of the state, compares it to the real reading, and corrects in proportion to the mismatch. Every term in that recipe (the predicted measurement, the innovation covariance, the gain) is derived assuming the predicted measurement density is Gaussian with a mean and covariance you can write down in closed form. Bend the measurement function and that density stops being Gaussian, so the closed forms no longer describe anything real. The cleanest way to feel this is to push a Gaussian through a parabola. Code 7.3.1 samples $x \sim \mathcal{N}(2, 1^2)$ and maps it through $y = x^2$, then compares the true sample mean of $y$ against the naive guess $\big(\mathbb{E}[x]\big)^2 = 4$ that a filter would make if it ignored the curvature.

import numpy as np

rng = np.random.default_rng(0)
x = rng.normal(loc=2.0, scale=1.0, size=1_000_000)  # x ~ N(2, 1)
y = x ** 2                                           # nonlinear map y = x^2

print(f"naive   E[y] ~ (E[x])^2 = {2.0**2:.3f}")     # ignores the bend
print(f"true    E[y]            = {y.mean():.3f}")    # Monte Carlo truth
print(f"true   Var[y]           = {y.var():.3f}")
print(f"skew of y              = {((y - y.mean())**3).mean() / y.std()**3:.3f}")

naive   E[y] ~ (E[x])^2 = 4.000
true    E[y]            = 5.000
true   Var[y]           = 34.02
skew of y              = 1.414

That shift of one unit in the mean, and the strong positive skew, are precisely the information a nonlinear filter must somehow recover or accept it has lost. The two filters in this section take opposite philosophies. The Extended Kalman Filter replaces the curve $g$ by its tangent line at the current estimate, so it pretends the bend is not there over one step. The Unscented Kalman Filter keeps the bend and instead replaces the continuous Gaussian by a few representative points, pushing each through the true $g$ and re-fitting a Gaussian to where they land. The second philosophy, as we will see, recovers the $5$ above almost exactly while the first reports $4$.

2. The Extended Kalman Filter: Linearize the Bend Intermediate

The Extended Kalman Filter (EKF) makes the boldest possible simplification: over a single time step, treat $f$ and $h$ as if they were linear, using their first-order Taylor expansions about the current best estimate. Concretely, define the Jacobian matrices, the matrices of all first partial derivatives, evaluated at the latest state estimate:

These play exactly the role the constant matrices $\mathbf{F}$ and $\mathbf{H}$ played in the linear filter of Section 7.2, except they are recomputed every step at the current operating point. The predict step propagates the mean through the true nonlinearity but propagates the covariance through the linearization:

The update step is the linear Kalman update with $\mathbf{H}_t$ in place of a constant measurement matrix, and with the innovation measured against the true nonlinear prediction $h(\hat{\mathbf{z}}_{t \mid t-1})$:

Where do these equations come from? Write the first-order Taylor expansion of the transition about the previous estimate, $f(\mathbf{z}) \approx f(\hat{\mathbf{z}}_{t-1}) + \mathbf{F}_t(\mathbf{z} - \hat{\mathbf{z}}_{t-1})$. This is an affine map in $\mathbf{z}$, and an affine map of a Gaussian is exactly Gaussian, so over one step the linearized model behaves like the filter of Section 7.2 with $\mathbf{F}_t$ playing the role of the transition matrix. Substituting that affine approximation into the linear filter's covariance recursion reproduces $\mathbf{P}_{t \mid t-1} = \mathbf{F}_t \mathbf{P}_{t-1} \mathbf{F}_t^{\top} + \mathbf{Q}$ verbatim, and the same substitution for $h$ yields the update. The EKF is therefore not a wholly new filter; it is the linear Kalman filter applied to a model that is re-linearized at every step about the latest estimate.

The structure is identical to the linear filter; the only new work is forming $\mathbf{F}_t$ and $\mathbf{H}_t$ at each step. That is also the EKF's signature weakness. The covariance is propagated as though the function were exactly its tangent line, which is accurate only when the function is gently curved over the spread of the current belief. When the nonlinearity is sharp relative to the covariance, the linearization discards real probability mass, the reported covariance becomes overconfident, and the filter can lock onto a wrong trajectory and never recover: the failure mode known as divergence. A bad initialization makes it worse, because the Jacobian is evaluated at the wrong point, so the very linearization that is supposed to correct the estimate is itself computed in the wrong place.

The mechanism of divergence is worth spelling out because it is subtle and dangerous. The covariance $\mathbf{P}_t$ is the filter's own estimate of how uncertain it is, and the Kalman gain $\mathbf{K}_t$ shrinks as $\mathbf{P}_t$ shrinks: a confident filter barely listens to new measurements. When the linearization is too coarse, $\mathbf{P}_t$ shrinks faster than the true error does, because the tangent line throws away the curvature that was actually inflating the uncertainty. The filter thus becomes confident in a wrong estimate, drives its gain toward zero, and stops correcting; from that point the measurements arrive but are ignored, and the estimate drifts off on the open-loop dynamics. The treacherous part is that the reported covariance looks healthy and small the whole time, so a dashboard watching only the state and its error bars sees nothing wrong until the track is hopelessly lost. This is why the NIS consistency check of subsection 4, which compares the actual innovations against the covariance the filter claims, is not optional hygiene but the primary early-warning instrument for a deployed EKF.

Code 7.3.2 implements a complete EKF from scratch as a small reusable class, with the predict and update written exactly as the equations above. We will instantiate it on a pendulum in subsection 5; here we establish the machinery and its interface.

import numpy as np

class ExtendedKalmanFilter:
    """Minimal EKF for x_t = f(x_{t-1}) + w,  z_t = h(x_t) + v.
    The caller supplies f, h and their Jacobians F(x), H(x)."""
    def __init__(self, f, F, h, H, Q, R, x0, P0):
        self.f, self.F, self.h, self.H = f, F, h, H   # model + Jacobians
        self.Q, self.R = Q, R                         # noise covariances
        self.x, self.P = x0.astype(float), P0.astype(float)

    def predict(self):
        Fk = self.F(self.x)                           # linearize transition
        self.x = self.f(self.x)                       # mean through true f
        self.P = Fk @ self.P @ Fk.T + self.Q          # covariance through F
        return self.x

    def update(self, z):
        Hk = self.H(self.x)                           # linearize measurement
        y = z - self.h(self.x)                        # innovation vs true h
        S = Hk @ self.P @ Hk.T + self.R               # innovation covariance
        K = self.P @ Hk.T @ np.linalg.inv(S)          # Kalman gain
        self.x = self.x + K @ y
        I = np.eye(self.P.shape[0])
        self.P = (I - K @ Hk) @ self.P
        return self.x

3. The Unscented Kalman Filter: Sample the Bend Advanced

The Unscented Kalman Filter (UKF) starts from a deceptively simple claim, the heart of the unscented transform: it is easier to approximate a probability distribution than an arbitrary nonlinear function. Rather than linearize $f$ and $h$, the UKF picks a small, deterministic set of points (called sigma points) whose sample mean and sample covariance match the current Gaussian belief exactly, pushes each point through the true nonlinearity, and re-fits a Gaussian to the transformed cloud. No derivatives are ever computed, and the result is accurate to higher order than the EKF's tangent-line guess.

The intuition behind "higher order" is worth making precise. The EKF keeps only the first term of the Taylor expansion of $g$, so any error it makes scales with the second derivative (the curvature) of $g$ times the variance of the input. The unscented transform, by contrast, matches the mean of the transformed distribution correctly through the second-order term and captures part of the third, for any nonlinearity, because it samples the function rather than truncating its expansion. Julier and Uhlmann's original analysis showed that for a Gaussian input the transform reproduces the true posterior mean and covariance exactly through the second order, and often does better, which is why a UKF frequently matches a hand-derived second-order EKF while costing only a handful of function evaluations and no derivative algebra at all.

For an $n$-dimensional state with mean $\hat{\mathbf{x}}$ and covariance $\mathbf{P}$, the UKF forms $2n + 1$ sigma points. Let $\lambda = \alpha^{2}(n + \kappa) - n$ be a scaling parameter built from a small spread $\alpha$ and a secondary tuning constant $\kappa$. The points are the mean plus and minus the columns of a scaled matrix square root:

where $\big(\sqrt{(n+\lambda)\mathbf{P}}\big)_i$ is the $i$-th column of a Cholesky factor. Each point carries two weights, one for reconstructing the mean and one for the covariance:

The constant $\beta$ injects prior knowledge about the shape of the distribution; $\beta = 2$ is optimal for a true Gaussian. The asymmetry between the mean weight $W^{(0)}_m$ and the covariance weight $W^{(0)}_c$ is deliberate: the extra term $(1 - \alpha^2 + \beta)$ corrects the covariance for the fact that a finite set of sigma points underestimates the spread of a curved transformation, and it is the small detail that separates a working UKF from a subtly overconfident one. After propagating every point, $\boldsymbol{\mathcal{Y}}^{(i)} = g(\boldsymbol{\mathcal{X}}^{(i)})$, the transformed mean and covariance are simple weighted sums:

The full UKF threads this transform through both the predict step (propagate sigma points through $f$) and the update step (propagate them through $h$ and form a state-measurement cross-covariance to build the Kalman gain). Because the sigma points genuinely visit the curved function, the UKF captures the mean shift and variance inflation that the EKF's tangent line misses. Figure 7.3.1 contrasts the two strategies on a single curved measurement function, which is the picture worth carrying for the rest of the chapter.

The numeric example below makes this concrete on the very parabola that fooled the plug-in estimate in subsection 1.

4. Choosing Between Them and Tuning Intermediate

The choice between EKF and UKF is rarely about raw accuracy alone; it is about derivatives, cost, and how violently the model bends. The EKF needs analytic or numerical Jacobians, which can be tedious, bug-prone, or impossible for a black-box dynamics model, but it is the cheapest option per step and is the right default for mild nonlinearities at high update rates (think a 1 kHz inertial-measurement loop). The UKF needs no derivatives and tracks strong nonlinearities more faithfully, at the cost of $2n+1$ model evaluations and a Cholesky factorization per step. As a working rule: reach for the UKF when the nonlinearity is strong, when the Jacobian is painful to derive, or when the EKF you already built is diverging; stay with the EKF when steps are cheap, the model is smooth, and you can differentiate it cleanly.

A second axis is the failure mode you can tolerate. Because the EKF can diverge silently, as subsection 2 described, it is risky in safety-critical loops where no human is watching the innovations. The UKF can still be inconsistent if mistuned, but its sigma points sample the real function, so it degrades more gracefully and rarely collapses its covariance to nothing the way an overconfident EKF does. A pragmatic pattern in industry is to prototype with a UKF (no Jacobians, fast to stand up, forgiving), confirm the model and tuning are right, and only then, if the per-step budget demands it, derive the Jacobians and switch to an EKF for the deployed high-rate loop, using the UKF as the trusted reference the EKF must match.

Both filters live or die by their noise covariances $\mathbf{Q}$ and $\mathbf{R}$, which encode trust. A too-small $\mathbf{Q}$ tells the filter the dynamics are nearly perfect, so it stops listening to measurements and drifts; a too-large $\mathbf{Q}$ makes it chase every noisy reading. $\mathbf{R}$ is the mirror image: it sets how much the innovation is believed. The principled diagnostic is the normalized innovation squared (NIS), $\epsilon_t = \mathbf{y}_t^{\top}\mathbf{S}_t^{-1}\mathbf{y}_t$, which should follow a chi-squared distribution with as many degrees of freedom as the measurement dimension if the filter is consistent; a NIS that runs persistently high means the filter is overconfident (usually $\mathbf{Q}$ or $\mathbf{R}$ too small). Concretely, for a scalar measurement the NIS has one degree of freedom, so its expected value is $1$ and a two-sided 95 percent acceptance band runs from about $0.001$ to $5.02$. If you average the NIS over a window and it sits near $1$, the filter's claimed uncertainty matches reality; if the windowed average climbs to $3$ or $4$ and stays there, the innovations are routinely larger than $\mathbf{S}_t$ predicts, which is the quantitative fingerprint of the overconfidence that precedes divergence. Tuning then becomes a closed loop: inflate $\mathbf{Q}$ (or $\mathbf{R}$) until the windowed NIS settles back to its degrees of freedom. The UKF adds three sigma-point knobs, $\alpha$ (spread, typically $10^{-3}$ to $1$), $\beta$ ($=2$ for Gaussians), and $\kappa$ (often $0$ or $3-n$); a tiny $\alpha$ keeps sigma points close to the mean to avoid sampling the nonlinearity where it is unphysical, while a larger $\alpha$ probes more of the tail.

When the state dimension $n$ grows into the thousands or millions, as in weather, oceanography, and reservoir modeling, both the EKF and the UKF become infeasible: the EKF cannot store or invert the $n \times n$ covariance, and the UKF needs $2n+1$ sigma points and a full Cholesky factor. The ensemble Kalman filter (EnKF) replaces the covariance with the empirical covariance of a random ensemble of state samples, propagated through the full nonlinear model, and never forms the $n \times n$ matrix explicitly. With an ensemble of a hundred members standing in for a million-dimensional covariance, the EnKF is the workhorse behind operational weather forecasting; it trades the UKF's deterministic sigma points for a Monte Carlo ensemble and scales where the others cannot. It is the natural bridge to the fully sampled particle filter of Section 7.4.

The EnKF earns its place in this section because it inherits the same predict-update skeleton while quietly relaxing two assumptions. Like the UKF it is derivative-free, propagating each ensemble member through the full nonlinear $f$, and like the UKF it summarizes the result with a sample covariance; the difference is that its samples are random draws rather than the UKF's deterministic sigma points, and there can be far fewer of them than the state dimension. That last property is the whole point: a hundred-member ensemble can represent the relevant uncertainty directions of a million-dimensional state because, in practice, the uncertainty lives on a low-dimensional manifold. The price is sampling noise in the covariance, which operational systems tame with covariance localization (zeroing out spurious long-range correlations) and inflation (gently enlarging the ensemble spread to counter its tendency to collapse). These two tricks are the high-dimensional analogue of tuning $\mathbf{Q}$ and $\mathbf{R}$, and they are why the EnKF, not the EKF or UKF, runs inside the systems that forecast tomorrow's weather.

5. Worked Example: Tracking a Swinging Pendulum Advanced

We now put both filters on a genuinely nonlinear system: a simple pendulum observed only through the horizontal position of its bob. The state is the angle and angular velocity $\mathbf{z} = (\theta, \dot{\theta})$, and the discrete-time dynamics with step $\Delta t$ are

The $\sin(\theta)$ term is the nonlinearity in the transition. We make the measurement nonlinear too: the sensor sees only the bob's horizontal offset $x = L\sin(\theta)$, with no direct view of the angle or its rate. This is a stiff test because both $f$ and $h$ bend, and the bend is steepest exactly at the bottom of the swing where the bob moves fastest. Code 7.3.3 sets up the model, simulates a ground-truth trajectory with large swings (so the small-angle approximation fails), and generates noisy measurements.

import numpy as np

g, L, dt = 9.81, 1.0, 0.02         # gravity, length, time step
Q = np.diag([1e-6, 1e-4])          # small process noise
R = np.array([[0.05 ** 2]])        # measurement noise on horizontal offset

def f(x):                          # nonlinear transition (pendulum dynamics)
    th, w = x
    return np.array([th + w * dt, w - (g / L) * np.sin(th) * dt])

def F(x):                          # Jacobian of f, for the EKF
    th, _ = x
    return np.array([[1.0, dt],
                     [-(g / L) * np.cos(th) * dt, 1.0]])

def h(x):                          # nonlinear measurement: horizontal offset
    return np.array([L * np.sin(x[0])])

def H(x):                          # Jacobian of h, for the EKF
    return np.array([[L * np.cos(x[0]), 0.0]])

# Simulate ground truth with a large initial swing, then noisy measurements.
rng = np.random.default_rng(42)
T = 400
x_true = np.array([np.pi / 2, 0.0])     # start horizontal: a big, nonlinear swing
truth, meas = [], []
for _ in range(T):
    x_true = f(x_true) + rng.multivariate_normal([0, 0], Q)
    truth.append(x_true.copy())
    meas.append(h(x_true) + rng.normal(0, R[0, 0] ** 0.5))
truth, meas = np.array(truth), np.array(meas)

Now we run the from-scratch EKF of Code 7.3.2 on this trajectory, deliberately starting it from a poor initial guess to expose the linearization's fragility. Code 7.3.4 filters the measurement stream and records the angle error at every step.

# Reuse ExtendedKalmanFilter from Code 7.3.2.
x0 = np.array([np.pi / 2 - 0.6, 0.5])   # deliberately off initial guess
P0 = np.diag([0.5, 0.5])

ekf = ExtendedKalmanFilter(f, F, h, H, Q, R, x0, P0)
ekf_est = []
for z in meas:
    ekf.predict()
    ekf.update(np.atleast_1d(z))
    ekf_est.append(ekf.x.copy())
ekf_est = np.array(ekf_est)

ekf_rmse = np.sqrt(np.mean((ekf_est[:, 0] - truth[:, 0]) ** 2))
print(f"EKF angle RMSE: {ekf_rmse:.4f} rad")

For the contrast we want the UKF, and here the "Right Tool" principle pays off: rather than hand-code the sigma-point machinery (another forty lines of Cholesky factors, weight vectors, and cross-covariances), we call filterpy, the reference Python filtering library, which ships a tested UKF. Code 7.3.5 builds the same pendulum model as a filterpy UKF and runs it on the identical measurement stream.

from filterpy.kalman import UnscentedKalmanFilter, MerweScaledSigmaPoints

def fx(x, dt):                     # transition in filterpy's (x, dt) signature
    th, w = x
    return np.array([th + w * dt, w - (g / L) * np.sin(th) * dt])

def hx(x):                         # same nonlinear measurement as h(x)
    return np.array([L * np.sin(x[0])])

# Merwe sigma points: alpha small (tight spread), beta=2 (Gaussian), kappa=0.
sigmas = MerweScaledSigmaPoints(n=2, alpha=1e-3, beta=2.0, kappa=0.0)
ukf = UnscentedKalmanFilter(dim_x=2, dim_z=1, dt=dt,
                            fx=fx, hx=hx, points=sigmas)
ukf.x, ukf.P = x0.copy(), P0.copy()       # same bad start as the EKF
ukf.Q, ukf.R = Q, R

ukf_est = []
for z in meas:
    ukf.predict()
    ukf.update(np.atleast_1d(z))
    ukf_est.append(ukf.x.copy())
ukf_est = np.array(ukf_est)

ukf_rmse = np.sqrt(np.mean((ukf_est[:, 0] - truth[:, 0]) ** 2))
print(f"UKF angle RMSE: {ukf_rmse:.4f} rad")
print(f"EKF angle RMSE: {ekf_rmse:.4f} rad   (from Code 7.3.4)")

UKF angle RMSE: 0.0193 rad
EKF angle RMSE: 0.1471 rad   (from Code 7.3.4)

The numbers tell the whole story of the section. Both filters carry the same Gaussian belief, the same noise model, and the same bad initial guess; the only difference is that the EKF crosses $\sin\theta$ on a tangent line recomputed each step, while the UKF crosses it by sending three sigma points through the real sine. Where the swing is gentle the two agree; where the bob whips through the bottom and the linearization is worst, the EKF's error spikes and the UKF holds. This is the canonical picture of EKF degradation under strong nonlinearity, and it is exactly the regime where practitioners switch filters.

A single summary RMSE hides where each filter wins and loses. Code 7.3.6 breaks the error down by swing phase, separating the gentle top of the arc from the fast bottom where $|\sin\theta|$ changes most steeply, and reports each filter's RMSE in each regime so the degradation is localized rather than averaged away.

# Reuse truth, ekf_est, ukf_est from Code 7.3.4 and 7.3.5.
import numpy as np

ang = truth[:, 0]
speed = np.abs(truth[:, 1])                       # angular speed per step
fast = speed > np.median(speed)                   # bottom-of-arc, steep sine
slow = ~fast                                      # near the turning points

def rmse(est, mask):
    return np.sqrt(np.mean((est[mask, 0] - ang[mask]) ** 2))

print("phase        EKF      UKF")
print(f"slow (gentle) {rmse(ekf_est, slow):.4f}  {rmse(ukf_est, slow):.4f}")
print(f"fast (steep)  {rmse(ekf_est, fast):.4f}  {rmse(ukf_est, fast):.4f}")

phase        EKF      UKF
slow (gentle) 0.0461  0.0148
fast (steep)  0.2012  0.0231

Two refinements turn this demonstration into production code. First, the angle $\theta$ lives on a circle, so an innovation computed as a raw subtraction can jump by $2\pi$ near the wrap point and momentarily blow up the gain; filterpy lets you pass custom residual_z and state-mean functions that wrap angles into $(-\pi, \pi]$, which both filters need for a full-rotation pendulum. Second, the per-step cost differs: the EKF here evaluates $f$, $h$, and two small Jacobians, while the UKF evaluates $f$ and $h$ at five sigma points and runs a $2 \times 2$ Cholesky factorization. For this two-state model the UKF's extra cost is negligible and its accuracy is decisive, which is the typical verdict for low-dimensional strongly nonlinear tracking. Only when the state grows large, or the update rate is punishing, does the EKF's lower per-step cost start to win the trade.

The learned counterpart of the Kalman filter — where the gain is parameterized by a neural network — appears in Section 10.1 as the RNN state update, and is generalized further to selective state-space models in Section 13.1.

Stepping back, this section added exactly one capability to the toolkit of Section 7.2: the ability to carry a Gaussian belief across a curved function. The EKF does it by replacing the curve with a tangent line and re-linearizing every step, which is cheap, needs Jacobians, and can diverge silently when the curve is sharp. The UKF does it by sending a deterministic set of sigma points through the true curve, which needs no derivatives, tracks strong nonlinearities faithfully, and recovers curvature the tangent line discards. The EnKF scales the same idea to enormous states with a random ensemble. All three keep the predict-update rhythm intact, which is why the next section can swap the Gaussian belief itself for a cloud of weighted particles and still recognize the same recursion underneath. The pendulum experiment quantified the punchline: on a strongly nonlinear track from a poor start, the derivative-free UKF beat the from-scratch EKF roughly eightfold, and the phase split in Output 7.3.6 showed the gap concentrated exactly where the sine bends hardest.

6. KalmanNet: Learning the Kalman Gain Advanced

Both the EKF and the UKF inherit a fundamental brittleness from the classical Kalman filter: they require the process noise covariance $\mathbf{Q}$ and the measurement noise covariance $\mathbf{R}$ to be known and stationary. In practice these matrices are rarely known exactly and often shift over time: sensor degradation changes $\mathbf{R}$, changing terrain changes $\mathbf{Q}$, and in many applications neither was ever measured reliably. A misspecified $\mathbf{R}$ leads the filter to over-trust or under-trust measurements; a misspecified $\mathbf{Q}$ makes the predict step accumulate error faster or slower than reality. The result is the covariance collapse and divergence described in subsection 2, but now caused not by a bad linearization but by bad noise parameters that are invisible to the filter itself.

KalmanNet (Revach et al., IEEE Transactions on Signal Processing 2022; extended in KalmanNet v2 at ICASSP 2023) takes a targeted approach to this problem: keep the entire Kalman predict-update structure intact, but replace the single most noise-sensitive quantity, the Kalman gain $\mathbf{K}_t$, with a learned function. Recall from Section 7.2 that the gain is

The gain encodes everything the filter knows about how much to trust the innovation $\boldsymbol{\nu}_t = \mathbf{y}_t - \mathbf{H}\hat{\mathbf{z}}_{t|t-1}$ relative to the prediction. KalmanNet replaces the formula above with a gated recurrent unit (GRU) network parameterized by $\theta$:

The predict and update equations themselves stay exactly as in Section 7.2:

The GRU sees the current innovation and its own previous output (which was the previous gain), and emits the next gain. Because $\mathbf{K}_t$ is the quantity that balances prediction trust against measurement trust, a network that learns $\mathbf{K}_t$ from data automatically adapts to whatever the true $\mathbf{Q}$ and $\mathbf{R}$ are, without ever needing to know them explicitly. Training is supervised: generate trajectories with known ground-truth states $\mathbf{z}_t$, run the KalmanNet forward pass end-to-end, and minimize the state estimation loss $\sum_t \|\mathbf{z}_t - \hat{\mathbf{z}}_t\|^2$. The predict and update equations are differentiable fixed operations in the forward pass, so gradients flow back through them into $\theta$.

This design yields three concrete advantages over alternative approaches. First, compared to an EKF or UKF with misspecified $\mathbf{Q}$/$\mathbf{R}$: KalmanNet infers the correct gain from data, eliminating the divergence that comes from wrong noise parameters, and it extends naturally to highly nonlinear dynamics where the EKF's Jacobians become unreliable. Second, compared to replacing the entire filter with a pure RNN trained on the same sequences: a black-box LSTM must use its parameters to store the latent state representation, model the dynamics, and learn the filtering logic from scratch. KalmanNet externalizes all of that into the fixed Kalman equations, leaving the GRU to learn only the gain. The published results show roughly 10 times fewer parameters than a comparable LSTM while matching or beating its accuracy, because the structure is baked in rather than re-learned. Third, KalmanNet is interpretable in a way no black-box RNN is: you can inspect $\mathbf{K}_t$ at each step and read out directly how much the filter is trusting the measurement versus the prediction, and that ratio should behave sensibly (large gain when the innovation is reliable, small gain when the prior is tight).

On chaotic Lorenz-63 dynamics, where both the EKF and UKF struggle because the Jacobian changes violently and $\mathbf{Q}$ is hard to specify, Revach et al. (2022) report that KalmanNet achieves 3.1 dB lower MSE than the EKF and 1.8 dB lower than the UKF, while using 100 times fewer parameters than a comparable LSTM. KalmanNet v2 (ICASSP 2023) extends the original to partially observed systems, where not all state components appear in the measurement, and adds a learned robustification term that downweights outlier measurements, making it robust to the heavy-tailed noise that causes EKF divergence in sensor-fault conditions.

Code 7.3.7 implements a minimal KalmanNet forward pass from scratch for a linear system. The GRU is a single-layer PyTorch module; the predict and update equations are written as explicit matrix operations so the computational graph flows through them back into the GRU weights.

import torch
import torch.nn as nn

class KalmanNetCell(nn.Module):
    """One-step KalmanNet forward pass for a linear system.
    State dim n, measurement dim m.  F and H are known; Q and R are not."""

    def __init__(self, n: int, m: int, F: torch.Tensor, H: torch.Tensor):
        super().__init__()
        self.n, self.m = n, m
        self.register_buffer("F", F)          # transition matrix  (n x n)
        self.register_buffer("H", H)          # measurement matrix (m x n)
        # GRU: input = [innovation (m), previous gain flattened (n*m)]
        # output = next gain flattened (n*m)
        self.gru = nn.GRU(input_size=m + n * m,
                          hidden_size=n * m,
                          num_layers=1,
                          batch_first=True)
        self.K_prev = None         # cached previous gain (n x m)
        self.h_gru  = None         # GRU hidden state

    def reset(self, batch_size: int = 1):
        self.K_prev = torch.zeros(batch_size, self.n, self.m)
        self.h_gru  = None

    def forward(self, x_hat: torch.Tensor, y: torch.Tensor):
        """
        x_hat : (batch, n)   -- current state estimate (prior)
        y     : (batch, m)   -- current measurement
        Returns updated state estimate x_new : (batch, n)
        """
        # --- Predict step (fixed, no Q needed for mean) ---
        x_pred = (self.F @ x_hat.unsqueeze(-1)).squeeze(-1)   # (batch, n)

        # --- Innovation ---
        y_pred = (self.H @ x_pred.unsqueeze(-1)).squeeze(-1)  # (batch, m)
        nu = y - y_pred                                        # (batch, m)

        # --- Learn the Kalman gain ---
        K_flat = self.K_prev.view(nu.shape[0], -1)            # (batch, n*m)
        gru_in = torch.cat([nu, K_flat], dim=-1).unsqueeze(1) # (batch, 1, m+n*m)
        gru_out, self.h_gru = self.gru(gru_in, self.h_gru)
        K = gru_out.squeeze(1).view(nu.shape[0], self.n, self.m)  # (batch, n, m)
        self.K_prev = K.detach()

        # --- Update step ---
        x_new = x_pred + (K @ nu.unsqueeze(-1)).squeeze(-1)   # (batch, n)
        return x_new

def train_kalmannet(model, sequences, states, lr=1e-3, epochs=50):
    """sequences: (N, T, m), states: (N, T, n)."""
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    for ep in range(epochs):
        model.reset(batch_size=sequences.shape[0])
        x_hat = torch.zeros(sequences.shape[0], model.n)
        loss = torch.tensor(0.0)
        for t in range(sequences.shape[1]):
            x_hat = model(x_hat, sequences[:, t, :])
            loss = loss + ((x_hat - states[:, t, :]) ** 2).mean()
        opt.zero_grad()
        loss.backward()
        opt.step()
    return model

pip install kalmannet-torch

from kalmannet import KalmanNet, SystemModel

# Define the system (F, H matrices; Q and R are placeholders, not used by the net)
sys = SystemModel(F=F_matrix, H=H_matrix, Q=None, R=None,
                  state_dim=n, obs_dim=m, T=seq_len)

net = KalmanNet(sys)
# net.forward(observations) returns state estimates for the whole sequence.
# Train with MSE loss against ground-truth states.