Bayesian Filtering and Belief States › Code Implementation
World Models HandbookAll chapters Bayesian Filtering and Belief States › Code ImplementationOne section · Sections in this chapter8 / 11 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
We will implement three filters on two test problems: a Kalman filter on a linear-Gaussian tracking task, then an EKF and a particle filter on a nonlinear bearing-only tracking task, and finally a small neural network trained to perform the belief update itself. The progression is deliberate. We start where the analytic solution is optimal under the known model, so we can compare implementation output with simulated state truth. We then move to a problem where the analytic filter is only approximate, and finally remove the equations from the learning procedure to preview amortized inference in modern world models.
A Linear-Gaussian Tracking Problem
Our first testbed is a noisy double integrator: a cart whose state is position and velocity, driven by random accelerations, observed through a noisy position sensor. The Kalman filter is the model-specified analytic baseline for this known linear-Gaussian system. The state components are coupled through time, and that coupling lets the filter infer unobserved velocity from position observations.
import numpy as np
rng = np.random.default_rng(42)
## State: [position, velocity]. Constant-velocity motion driven by acceleration.
A = np.array([[1.0, 1.0], [0.0, 1.0]])
B = np.array([[0.5], [1.0]])
Q = np.diag([0.02, 0.05]) # process noise covariance
H = np.array([[1.0, 0.0]]) # we observe position only
R = np.array([[4.0]]) # position sensor variance (std = 2 m)
## This diagonal Q is a simple stand-in, not the correlated covariance obtained
## by discretizing a single continuous-time acceleration disturbance.
T = 80
actions = rng.normal(scale=1.0, size=(T, 1))
truth = np.zeros((T + 1, 2))
obs = np.zeros((T, 1))
for t in range(T):
noise = rng.multivariate_normal(np.zeros(2), Q)
truth[t + 1] = A @ truth[t] + B @ actions[t] + noise
obs[t] = H @ truth[t + 1] + rng.normal(scale=float(np.sqrt(R[0, 0])))The simulation produces paired transitions and noisy position readings. Every trajectory starts from the known state ; the default Kalman initialization uses that same mean but a deliberately broad covariance . The sensor sees position with a two-meter standard deviation, while velocity must be inferred from correlations across time. Comparing the result with raw measurement differences checks only that the implementation beats a noisy one-step baseline. It does not isolate model-based fusion from generic position smoothing.
Running the Kalman Filter
The filter itself is compact. We write it generically so the same function serves as a baseline later. Note how directly the code mirrors the equations: each line corresponds to one formula from the Filter Equations section, in the same order, which is a good habit when implementing anything mathematical, since it makes review and debugging a matter of matching lines to symbols.
def kalman_filter(A, B, H, Q, R, obs, actions, x0=None, P0=None):
n = A.shape[0]
x = np.zeros(n) if x0 is None else x0.copy()
P = np.eye(n) * 10.0 if P0 is None else P0.copy()
means, covs = [], []
for t in range(len(obs)):
# Predict: push belief through the dynamics
x = A @ x + B @ actions[t]
P = A @ P @ A.T + Q
# Correct: fold in the measurement
innovation = obs[t] - H @ x
S = H @ P @ H.T + R
K = P @ H.T @ np.linalg.inv(S)
x = x + K @ innovation
# The compact covariance update is adequate here; ill-conditioned
# production filters often use the Joseph form for numerical stability.
P = (np.eye(n) - K @ H) @ P
means.append(x.copy())
covs.append(P.copy())
return np.array(means), np.array(covs)
linear_x0 = np.zeros(2)
linear_P0 = np.eye(2) * 10.0
estimates, covariances = kalman_filter(
A, B, H, Q, R, obs, actions, x0=linear_x0, P0=linear_P0
)
assert T == 80 and len(obs) == len(actions) == T
assert np.allclose(linear_x0, np.zeros(2))
assert np.allclose(linear_P0, np.eye(2) * 10.0)
assert estimates.shape == (T, 2) and covariances.shape == (T, 2, 2)
rmse_pos = np.sqrt(np.mean((estimates[:, 0] - truth[1:, 0]) ** 2))
rmse_vel = np.sqrt(np.mean((estimates[:, 1] - truth[1:, 1]) ** 2))
raw_diff_rmse = np.sqrt(
np.mean((obs[1:, 0] - obs[:-1, 0] - (truth[2:, 0] - truth[1:-1, 0])) ** 2)
)
std_pos = np.sqrt(covariances[:, 0, 0])
prior_means = np.vstack([linear_x0, estimates[:-1]])
predicted_means = prior_means @ A.T + actions @ B.T
innovation_residuals = obs[:, 0] - (predicted_means @ H.T)[:, 0]
prior_covariances = np.concatenate(
[linear_P0[None, ...], covariances[:-1]], axis=0
)
predicted_covariances = np.array([A @ P @ A.T + Q for P in prior_covariances])
innovation_variances = np.array(
[(H @ P @ H.T + R).item() for P in predicted_covariances]
)
normalized_innovations = innovation_residuals / np.sqrt(innovation_variances)
innovations = normalized_innovationsPosition RMSE (filter): 1.044 m Velocity RMSE (filter): 0.378 m/s Velocity RMSE (measurement differencing): 2.885 m/s
The filter's velocity estimate beats naive measurement differencing substantially in this run. With the unit sampling period used here, the difference of two independent sensor errors has variance , so this is deliberately a weak baseline. The filter uses the dynamic model and full observation history instead of amplifying sensor noise through subtraction, but this comparison alone does not show how much of the gain comes from model-based fusion rather than generic smoothing. Now let us look at the estimate itself.
The shaded band reports the filter's own confidence: early on, when the filter knows little, the band is wide; during the initial transient it tightens and then approaches a nonzero steady-state width set jointly by , , , and . The plotted half-width is and the full band width is . The band is not decoration, it is the filter reporting its own confidence, and downstream planners can consume it. A controller that knows when its state estimate is unreliable can act cautiously, slow down, or request more sensing, and a controller fed only point estimates has no way to make that distinction. This is the first concrete appearance of uncertainty quality as a property distinct from accuracy: two filters with identical RMSE can differ enormously in how honestly their bands reflect their actual error.
Next we examine the internal quantities that drive this behavior: the posterior variance and the innovation magnitude.
Both diagnostics matter in practice. Converging variance confirms the filter reaches an information steady state, where each step sheds as much uncertainty as the process noise injects. The second panel shows the signed innovation divided by its predicted standard deviation, not the raw residual. Under a correctly specified linear-Gaussian model, these signed normalized innovations are standard normal and uncorrelated over time. In the run above, the filter deliberately starts from a broadened prior ( against a deterministic initial state), so the first few normalized innovations are underscaled until that initial transient passes; thereafter they follow the standard-normal reference, and unusually large magnitudes or persistent temporal structure provide a principled model-health alarm. Learned systems can also compare predictions with outcomes, but the point-estimate GRU below does not report a predictive variance, so its benchmark checks estimation error rather than uncertainty calibration.
Bearing-Only Tracking with an EKF
Now we move to a nonlinear problem. A sensor at the origin measures only the bearing to a moving target, not its range. The observation function is , so one measurement constrains the target to a ray. Our simulated target follows an offset circle, which gives a more varied bearing trace than a circle centered on the sensor. It does not make absolute scale observable to the filters used here: scaling every position and velocity by the same positive factor leaves all bearings unchanged, and neither filter encodes the circle center or radius. Their Cartesian trajectories are therefore anchored by the initial belief alone; the constant-velocity transition is scale-equivariant, so it carries whatever scale the prior sets without contributing independent scale information. The truth uses circular kinematics, while both filters deliberately predict with a constant-velocity transition; their approximation error therefore includes dynamics mismatch as well as, for the EKF, local linearization.
## State: [x, y, vx, vy]. The target's circle is offset from the beacon.
T_nl = 60
radius, omega = 30.0, 0.08
circle_center = np.array([8.0, 0.0])
beacon_xy = np.zeros(2)
truth_nl = np.zeros((T_nl + 1, 4))
for t in range(T_nl + 1):
theta = omega * t
truth_nl[t] = [
circle_center[0] + radius * np.cos(theta),
circle_center[1] + radius * np.sin(theta),
-radius * omega * np.sin(theta),
radius * omega * np.cos(theta),
]
bearing_noise_std = 0.05 # radians, about 3 degrees
rng_bearing = np.random.default_rng(7)
truth_relative_to_beacon = truth_nl[1:, :2] - beacon_xy
bearings = np.arctan2(
truth_relative_to_beacon[:, 1], truth_relative_to_beacon[:, 0]
) + rng_bearing.normal(scale=bearing_noise_std, size=T_nl)
assert len(bearings) == T_nl == len(truth_nl) - 1
assert np.min(np.linalg.norm(truth_nl[:, :2] - beacon_xy, axis=1)) > 20.0The EKF needs the Jacobian of the observation function. For the partial-derivative formulas away from the origin are and . The derivatives depend on the current estimate, so the linearization moves with the belief. The Jacobian is undefined at the sensor origin; a production implementation must detect that singular geometry rather than divide by zero. This example stays away from the origin, but its bearing trace crosses the principal-angle branch cut near the negative -axis. The wrapped innovation below is therefore essential: the rational formulas remain finite along the cut, but the observation function itself is discontinuous there.
F = np.array(
[[1, 0, 1, 0], [0, 1, 0, 1], [0, 0, 1, 0], [0, 0, 0, 1]], dtype=float
)
Q_nl = np.diag([0.05, 0.05, 0.01, 0.01])
R_nl = np.array([[bearing_noise_std**2]])
def ekf_bearing(bearings, x0, P0):
x = x0.copy()
P = P0.copy()
ests = []
for z in bearings:
# Predict with a constant-velocity approximation to the circular truth.
# Q_nl absorbs some, but not all, of this deliberate model mismatch.
x = F @ x
P = F @ P @ F.T + Q_nl
# Jacobian of atan2(y, x) at the predicted mean
px, py = x[:2] - beacon_xy
s = px**2 + py**2
if s < 1e-12:
raise ValueError(
"bearing Jacobian is undefined at the sensor origin"
)
H_jac = np.array([[-py / s, px / s, 0.0, 0.0]])
# Wrap across the atan2 branch cut before applying the correction.
innov = np.arctan2(
np.sin(z - np.arctan2(py, px)), np.cos(z - np.arctan2(py, px))
)
S = H_jac @ P @ H_jac.T + R_nl
K = P @ H_jac.T @ np.linalg.inv(S)
x = x + K @ np.array([innov])
P = (np.eye(4) - K @ H_jac) @ P
ests.append(x.copy())
return np.array(ests)
x0_guess = np.array([20.0, 20.0, 0.0, 0.0]) # deliberately wrong initial guess
P0 = np.diag([100.0, 100.0, 10.0, 10.0])
ekf_estimates = ekf_bearing(bearings, x0_guess, P0)
ekf_rmse = np.sqrt(np.mean((ekf_estimates[:, :2] - truth_nl[1:, :2]) ** 2))EKF position RMSE: 5.821 m
Two implementation details deserve comment. First, the innovation is wrapped into the principal interval using the sine/cosine trick, because angles live on a circle. Under the code's ordering, measurement minus prediction, a prediction of and a measurement of give a raw difference of , which wraps to , not or . At an exact antipodal difference of , floating-point signs can choose either equivalent endpoint. Ignoring wrapping is a classic EKF bug. The deeper lesson generalizes: whenever a state or measurement lives on a manifold rather than flat Euclidean space, angles on circles, orientations on spheres, poses on SE(3), the arithmetic of averaging and differencing must respect the geometry, and naive subtraction silently corrupts the filter. Second, we start the filter with a deliberately wrong initial guess and broad covariance. Early measurements can pull the estimate toward a bearing-consistent trajectory, while the prior continues to anchor absolute scale, which the scale-equivariant constant-velocity model merely carries forward. A broad covariance correctly represents that the initial Cartesian state is uncertain; it does not manufacture range information absent from the observations.
A Particle Filter for the Same Problem
The particle filter handles the same problem without a single derivative. Each particle is a full four-dimensional state hypothesis; the bearing likelihood reweights them each step. Where the EKF maintains one Gaussian hypothesis and bends it locally, the particle filter maintains thousands of discrete hypotheses and lets the data decide which survive. This flexibility does not guarantee a better estimate. The bootstrap proposal below samples only from the constant-velocity transition model, while each of the bearing observations has noise standard deviation radians and constrains angle but not absolute scale. The initialization corresponds to ten-meter marginal position standard deviations and m/s marginal velocity standard deviations, broad enough to include very different ranges and speeds. With that broad prior and circular truth that violates the proposal dynamics, resampling can eliminate range hypotheses that the bearing stream cannot restore on its own. Their survival depends on prior support and the assumed transition model.
N_particles = 4000
assert len(bearings) == T_nl == len(truth_nl) - 1
assert T_nl == 60
assert np.isclose(bearing_noise_std, 0.05)
assert np.allclose(np.diag(P0), [100.0, 100.0, 10.0, 10.0])
def systematic_resample(weights, rng):
if not np.isclose(weights.sum(), 1.0):
raise ValueError("systematic_resample requires normalized weights")
n = len(weights)
positions = (rng.random() + np.arange(n)) / n
cumulative = np.cumsum(weights)
cumulative[-1] = 1.0 # guard against floating-point roundoff
return np.searchsorted(cumulative, positions)
try:
systematic_resample(np.array([0.2, 0.2]), np.random.default_rng(0))
except ValueError:
pass
else:
raise AssertionError("non-normalized resampling weights must be rejected")
def particle_filter_bearing(bearings, x0, P0, n_particles, rng):
particles = rng.multivariate_normal(x0, P0, size=n_particles)
weights = np.ones(n_particles) / n_particles
ests, ess_trace, resampled_trace = [], [], []
for z in bearings:
# Propagate through dynamics with process noise
particles = particles @ F.T + rng.multivariate_normal(
np.zeros(4), Q_nl, size=n_particles
)
# Reweight by bearing likelihood (wrapped Gaussian)
particle_relative_to_beacon = particles[:, :2] - beacon_xy
pred_bearings = np.arctan2(
particle_relative_to_beacon[:, 1], particle_relative_to_beacon[:, 0]
)
diff = np.arctan2(np.sin(z - pred_bearings), np.cos(z - pred_bearings))
log_w = (
np.log(weights + np.finfo(float).tiny)
- 0.5 * (diff / bearing_noise_std) ** 2
)
log_w -= log_w.max()
weights = np.exp(log_w)
weights /= weights.sum()
ess_trace.append(1.0 / np.sum(weights**2))
# Estimate the weighted posterior before any resampling noise is added.
ests.append(weights @ particles)
# Resample only when degeneracy threatens
should_resample = ess_trace[-1] < n_particles / 2
resampled_trace.append(should_resample)
if should_resample:
idx = systematic_resample(weights, rng)
particles = particles[idx]
weights = np.ones(n_particles) / n_particles
return np.array(ests), np.array(ess_trace), np.array(resampled_trace)
pf_seeds = [17, 29, 41, 53, 71, 83, 97, 107]
pf_runs = [
particle_filter_bearing(
bearings, x0_guess, P0, N_particles, np.random.default_rng(seed)
)
for seed in pf_seeds
]
pf_rmse_runs = np.array(
[
np.sqrt(np.mean((est[:, :2] - truth_nl[1:, :2]) ** 2))
for est, _, _ in pf_runs
]
)
median_rmse = np.median(pf_rmse_runs)
representative_run_index = np.argmin(np.abs(pf_rmse_runs - median_rmse))
pf_estimates, pf_ess, pf_resampled = pf_runs[representative_run_index]
pf_rmse = pf_rmse_runs[representative_run_index]
pf_q25, pf_q75 = np.quantile(pf_rmse_runs, [0.25, 0.75])
pf_worse_than_ekf_in_this_sequence = bool(pf_rmse > ekf_rmse)
assert np.array_equal(pf_resampled, pf_ess < N_particles / 2)EKF position RMSE: 5.821 m Representative run closest to median particle-filter RMSE across eight random seeds: 8.054 m PF RMSE interquartile range across seeds: [6.387, 17.623] m
The filter reports the weighted posterior mean as its point estimate. In a multimodal posterior that mean can lie between modes; a mode estimate or the full particle set may be more useful to a downstream decision maker.
The printed particle trajectory is the representative run closest to the median particle-filter RMSE across eight random seeds; if two runs are equally close, np.argmin selects the earliest seed in the listed order. The interquartile range summarizes sensitivity to randomness in particle initialization, propagation, and resampling. For this executed configuration, pf_worse_than_ekf_in_this_sequence is true: the representative particle-filter run is worse than the EKF. That comparison is a reported outcome, not a build invariant; changing the numerical stack, seeds, particle budget, or proposal may reverse it. The code suggests two plausible contributors rather than proving a single cause: an angle-only likelihood cannot distinguish scale-equivalent trajectories, and the constant-velocity proposal is misspecified for circular motion. The comparison does not imply that sample-based beliefs are intrinsically inferior. It shows that a flexible representation still needs an informative proposal, an adequate particle budget, and dynamics that place mass near the truth.
Each downward slope is weight degeneracy in action: likelihood reweighting concentrates mass on fewer particles. Each vertical jump is a resampling event restoring uniform weights, not necessarily restoring state diversity. Effective sample size measures concentration of weights before resampling; it does not prove that the surviving particles cover the correct range or motion mode. If ESS repeatedly collapses, inspect trajectory error and particle diversity together rather than treating the post-resampling jump as evidence that the state estimate recovered.
Finally, the trajectory comparison:
The trajectory plot makes the representative run's RMSE concrete. Bearing residuals can remain small along a plausible ray even when Cartesian position error is large, so evaluate the state quantities a downstream controller needs rather than only the likelihood used for filtering. Observation-informed proposals, motion constraints, or conditional analytic updates can be more effective than increasing particle count alone when the proposal is structurally weak.
In this reproducible sequence, the EKF and the representative particle run settle on different Cartesian scales even though both receive the same bearings. That contrast is not evidence that the angles identify absolute range. It shows how the prior, transition approximation, and finite-sample proposal select among scale-equivalent bearing explanations.
Learning a Belief Update with a GRU
Classical filters require known models and hand-derived update rules. A learned alternative trains a recurrent network to consume observations and controls and output a state estimate, amortizing an approximation to the composed filtering recursion into its weights. Here “amortize” means paying an upfront training cost so that each later update uses one recurrent forward step.
We return to the linear-Gaussian cart, generate training trajectories, and train a gated recurrent unit (GRU) to map to an estimate of . In the zero-based arrays below, row (t) pairs control (u[:,t]) for the transition from (x_t) to (x_{t+1}) with measurement (z[:,t]) of (x_{t+1}); the target is therefore (x[:,t+1]). This is the same convention with array indices shifted by one. The controls are random exogenous inputs, not actions selected by a policy. The network never sees the Kalman equations. Because position and velocity have different scales, we standardize inputs and residual targets using training-set statistics. Position uses a residual target, , and evaluation reconstructs the position by adding back. The linear head itself is ordinary; the residual structure lies in this target parameterization. The learned update remains tied to its training distribution, which here contains only trajectories generated by the specified linear-Gaussian simulator and random exogenous controls. Every training and test episode starts from the deterministic state . The held-out Kalman baseline uses that correct zero initial mean but retains a broad covariance , so this benchmark does not test generalization to unknown or nonzero initial states.
import torch
import torch.nn as nn
## The setup cell seeds Python and NumPy; this framework-specific seed makes the
## GRU initialization and minibatch draws reproducible as well.
torch.manual_seed(0)
torch.set_num_threads(1)
def simulate_batch(n_seq, T, rng):
"""Episodes using the same A, B, H, Q, and R as the earlier simulator."""
x = np.zeros((n_seq, T + 1, 2))
u = rng.normal(scale=1.0, size=(n_seq, T, 1))
z = np.zeros((n_seq, T, 1))
q_chol = np.linalg.cholesky(Q)
for t in range(T):
noise = rng.normal(size=(n_seq, 2)) @ q_chol.T
x[:, t + 1] = x[:, t] @ A.T + u[:, t] @ B.T + noise
z[:, t] = x[:, t + 1] @ H.T + rng.normal(
scale=float(np.sqrt(R[0, 0])), size=(n_seq, 1)
)
return x, u, z
class BeliefRNN(nn.Module):
def __init__(self, hidden=32):
super().__init__()
self.gru = nn.GRU(input_size=2, hidden_size=hidden, batch_first=True)
self.head = nn.Linear(hidden, 2)
def forward(self, z, u):
inp = torch.cat([z, u], dim=-1) # (batch, T, 2)
out, _ = self.gru(inp)
return self.head(out) # (batch, T, 2) state estimates
train_x, train_u, train_z = simulate_batch(128, T, np.random.default_rng(101))
test_x, test_u, test_z = simulate_batch(64, T, np.random.default_rng(202))
## Learn a position correction relative to the observation, plus velocity.
train_targets = np.stack(
[train_x[:, 1:, 0] - train_z[..., 0], train_x[:, 1:, 1]], axis=-1
)
z_mean = train_z.mean(axis=(0, 1), keepdims=True)
z_scale = np.maximum(train_z.std(axis=(0, 1), keepdims=True), 1e-6)
u_mean = train_u.mean(axis=(0, 1), keepdims=True)
u_scale = np.maximum(train_u.std(axis=(0, 1), keepdims=True), 1e-6)
target_mean = train_targets.mean(axis=(0, 1), keepdims=True)
target_scale = np.maximum(train_targets.std(axis=(0, 1), keepdims=True), 1e-6)
model = BeliefRNN(hidden=32)
opt = torch.optim.Adam(model.parameters(), lr=3e-3)
loss_fn = nn.MSELoss()
zx = torch.tensor((train_z - z_mean) / z_scale, dtype=torch.float32)
ux = torch.tensor((train_u - u_mean) / u_scale, dtype=torch.float32)
yx = torch.tensor(
(train_targets - target_mean) / target_scale, dtype=torch.float32
)
assert np.allclose(zx.numpy().mean(axis=(0, 1)), 0.0, atol=1e-5)
assert np.allclose(ux.numpy().mean(axis=(0, 1)), 0.0, atol=1e-5)
assert np.allclose(yx.numpy().mean(axis=(0, 1)), 0.0, atol=1e-5)
training_steps = 1200
minibatch_size = 32
## This compact demonstration deliberately uses no validation split, early
## stopping, or hyperparameter search; those omissions are discussed below.
for step in range(training_steps):
idx = torch.randint(0, len(zx), (minibatch_size,))
pred = model(zx[idx], ux[idx])
loss = loss_fn(pred, yx[idx])
opt.zero_grad()
loss.backward()
opt.step()Now evaluate against the analytic Kalman filter on held-out trajectories. The evaluation standardizes test observations and controls with training-set statistics, de-standardizes the predicted residual targets, adds the observation back to reconstruct position, and supplies the same held-out controls to the Kalman filter. This is still a deliberately asymmetric benchmark: the Kalman filter receives the exact data-generating matrices and noise covariances , , , , and , while the GRU must learn its update from 128 simulated training trajectories. The GRU also receives injected structure through its residual target parameterization. The comparison tells us whether the learned update improves on the raw sensor and how far it remains from the model-specified analytic filter. It does not isolate amortization alone, because the gap includes both model knowledge given to the Kalman filter and structural information given to the GRU.
model.eval()
with torch.no_grad():
gru_pred_normalized = model(
torch.tensor((test_z - z_mean) / z_scale, dtype=torch.float32),
torch.tensor((test_u - u_mean) / u_scale, dtype=torch.float32),
).numpy()
gru_features = gru_pred_normalized * target_scale + target_mean
gru_pred = np.empty_like(test_x[:, 1:])
gru_pred[..., 0] = test_z[..., 0] + gru_features[..., 0]
gru_pred[..., 1] = gru_features[..., 1]
assert np.allclose(gru_pred[..., 0] - test_z[..., 0], gru_features[..., 0])
gru_rmse_pos = np.sqrt(np.mean((gru_pred[..., 0] - test_x[:, 1:, 0]) ** 2))
gru_rmse_vel = np.sqrt(np.mean((gru_pred[..., 1] - test_x[:, 1:, 1]) ** 2))
zero_velocity_rmse = np.sqrt(np.mean(test_x[:, 1:, 1] ** 2))
measurement_rmse_pos = np.sqrt(
np.mean((test_z[..., 0] - test_x[:, 1:, 0]) ** 2)
)
kf_estimates_test = []
for i in range(len(test_z)):
est_i, _ = kalman_filter(A, B, H, Q, R, test_z[i], test_u[i])
kf_estimates_test.append(est_i)
kf_estimates_test = np.stack(kf_estimates_test)
## Use pooled RMSE for both methods so aggregation does not favor either one.
kf_mean_pos = np.sqrt(
np.mean((kf_estimates_test[..., 0] - test_x[:, 1:, 0]) ** 2)
)
kf_mean_vel = np.sqrt(
np.mean((kf_estimates_test[..., 1] - test_x[:, 1:, 1]) ** 2)
)
kf_err_curve = np.mean(
np.abs(kf_estimates_test[..., 0] - test_x[:, 1:, 0]), axis=0
)
gru_err_curve = np.mean(np.abs(gru_pred[..., 0] - test_x[:, 1:, 0]), axis=0)
gru_t_axis = np.arange(1, T + 1)
assert len(gru_t_axis) == len(kf_err_curve) == len(gru_err_curve) == T
assert len(train_x) == 128 and len(test_x) == 64
assert training_steps == 1200 and minibatch_size == 32
gru_benchmark_observations = {
"position ordering in this held-out set": (
kf_mean_pos < gru_rmse_pos < measurement_rmse_pos
),
"velocity ordering in this held-out set": (
kf_mean_vel < gru_rmse_vel < zero_velocity_rmse
),
"Kalman has lower time-averaged position error in this held-out set": (
np.mean(kf_err_curve) < np.mean(gru_err_curve)
),
}
assert len(gru_benchmark_observations) == 3Position RMSE: Kalman 1.323 m | GRU 1.850 m Velocity RMSE: Kalman 0.520 m/s | GRU 0.655 m/s Position RMSE of raw measurements: 2.042 m Velocity RMSE of zero-velocity baseline: 6.802 m/s
In this executed benchmark, the residual-target GRU improves on the raw position measurement and beats the zero-velocity baseline, but it does not match the analytic Kalman filter on either state component. This implementation uses standardization and an observation skip connection at reconstruction time; because we do not ablate either choice, the comparison does not isolate their individual contributions. What the learned update buys is flexibility: the same recurrent architecture can be trained where an analytic filter is unavailable, provided its training distribution covers the deployment regime. Recurrent state-space models such as PlaNet and Dreamer (Hafner et al., 2019; 2020) extend this idea with probabilistic objectives and latent uncertainty.
Beyond aggregate RMSE, it is worth checking how errors change over a trajectory. The next figure uses per-timestep mean absolute error across held-out sequences, not pooled RMSE, because its purpose is to show temporal shape. Its vertical values therefore should not be compared numerically with the preceding RMSE bars.
In this executed benchmark, the Kalman curve has lower mean absolute error after averaging over time; the claim is not that it is lower at every individual timestep. That outcome is consistent with the Kalman filter receiving the exact linear-Gaussian model, but the benchmark does not isolate the causes of the gap. It combines known versus learned model information, finite training data, residual-target parameterization, network capacity, optimization, and the GRU's point-estimate output.
The executable checks used in the comparisons are collected below. Three structural invariants fail the build: bearing/state alignment, the particle resampling trigger, and the shared error-curve time axis. Seven seeded numerical observations are reported separately rather than asserted because they can change across numerical stacks: four checks cover sensor distance, the seed-42 linear comparison, innovation scale, and the representative particle run, while three come from the 64-sequence GRU benchmark. This integration cell deliberately depends on the preceding notebook cells; the explicit name check turns accidental reordering into an informative failure.
required_result_names = {
"rmse_vel",
"raw_diff_rmse",
"normalized_innovations",
"bearings",
"truth_nl",
"beacon_xy",
"pf_resampled",
"pf_ess",
"N_particles",
"pf_worse_than_ekf_in_this_sequence",
"gru_benchmark_observations",
"gru_t_axis",
"kf_err_curve",
"gru_err_curve",
"T",
}
missing_result_names = sorted(required_result_names.difference(globals()))
assert not missing_result_names, (
"Run the chapter cells in order before consolidation; missing: "
f"{missing_result_names}"
)
structural_checks = {
"bearing observations align with next states": len(bearings)
== len(truth_nl) - 1,
"particle filter uses the stated ESS trigger": np.array_equal(
pf_resampled, pf_ess < N_particles / 2
),
"error curves share the stated time axis": (
len(gru_t_axis) == len(kf_err_curve) == len(gru_err_curve) == T
),
}
benchmark_observations = {
"bearing trajectory stays more than 20 m from the sensor": np.min(
np.linalg.norm(truth_nl[:, :2] - beacon_xy, axis=1)
)
> 20.0,
"Kalman beats differencing on the seed-42 linear trajectory": (
rmse_vel < raw_diff_rmse
),
"normalized innovation scale is plausible on that trajectory": (
0.5 < np.std(normalized_innovations) < 1.5
),
"representative PF is worse than EKF on this bearing sequence": (
pf_worse_than_ekf_in_this_sequence
),
**gru_benchmark_observations,
}
assert len(structural_checks) == 3
assert len(benchmark_observations) == 7
assert all(structural_checks.values()), structural_checks
print(
"Executed structural checks passed: "
f"{sum(structural_checks.values())}/{len(structural_checks)}"
)
print(
"Seeded benchmark observations true in this run: "
f"{sum(benchmark_observations.values())}/{len(benchmark_observations)}"
)Executed structural checks passed: 3/3 Seeded benchmark observations true in this run: 7/7