{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Kalman Filter\n",
"\n",
"Kalman filters are linear models for state estimation of dynamic systems [1]. They have been the de facto standard in many robotics and tracking/prediction applications because they are well suited for systems with uncertainty about an observable dynamic process. They use a \"observe, predict, correct\" paradigm to extract information from an otherwise noisy signal. In Pyro, we can build differentiable Kalman filters with learnable parameters using the `pyro.contrib.tracking` [library](http://docs.pyro.ai/en/dev/contrib.tracking.html#module-pyro.contrib.tracking.extended_kalman_filter)\n",
"\n",
"## Dynamic process\n",
"\n",
"To start, consider this simple motion model:\n",
"\n",
"$$ X_{k+1} = FX_k + \\mathbf{W}_k $$\n",
"$$ \\mathbf{Z}_k = HX_k + \\mathbf{V}_k $$\n",
"\n",
"where $k$ is the state, $X$ is the signal estimate, $Z_k$ is the observed value at timestep $k$, $\\mathbf{W}_k$ and $\\mathbf{V}_k$ are independent noise processes (ie $\\mathbb{E}[w_k v_j^T] = 0$ for all $j, k$) which we'll approximate as Gaussians. Note that the state transitions are linear."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Kalman Update\n",
"At each time step, we perform a prediction for the mean and covariance:\n",
"$$ \\hat{X}_k = F\\hat{X}_{k-1}$$\n",
"$$\\hat{P}_k = FP_{k-1}F^T + Q$$\n",
"\n",
"and a correction for the measurement:\n",
"\n",
"$$ K_k = \\hat{P}_k H^T(H\\hat{P}_k H^T + R)^{-1}$$\n",
"$$ X_k = \\hat{X}_k + K_k(z_k - H\\hat{X}_k)$$\n",
"$$ P_k = (I-K_k H)\\hat{P}_k$$\n",
"\n",
"where $X$ is the position estimate, $P$ is the covariance matrix, $K$ is the Kalman Gain, and $Q$ and $R$ are covariance matrices.\n",
"\n",
"For an in-depth derivation, see \\[2\\]"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Nonlinear Estimation: Extended Kalman Filter\n",
"\n",
"What if our system is non-linear, eg in GPS navigation? Consider the following non-linear system:\n",
"\n",
"$$ X_{k+1} = \\mathbf{f}(X_k) + \\mathbf{W}_k $$\n",
"$$ \\mathbf{Z}_k = \\mathbf{h}(X_k) + \\mathbf{V}_k $$\n",
"\n",
"Notice that $\\mathbf{f}$ and $\\mathbf{h}$ are now (smooth) non-linear functions.\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The Extended Kalman Filter (EKF) attacks this problem by using a local linearization of the Kalman filter via a [Taylors Series expansion](https://en.wikipedia.org/wiki/Taylor_series).\n",
"\n",
"$$ f(X_k, k) \\approx f(x_k^R, k) + \\mathbf{H}_k(X_k - x_k^R) + \\cdots$$\n",
"\n",
"where $\\mathbf{H}_k$ is the Jacobian matrix at time $k$, $x_k^R$ is the previous optimal estimate, and we ignore the higher order terms. At each time step, we compute a Jacobian conditioned the previous predictions (this computation is handled by Pyro under the hood), and use the result to perform a prediction and update.\n",
"\n",
"Omitting the derivations, the modification to the above predictions are now:\n",
"$$ \\hat{X}_k \\approx \\mathbf{f}(X_{k-1}^R)$$\n",
"$$ \\hat{P}_k = \\mathbf{H}_\\mathbf{f}(X_{k-1})P_{k-1}\\mathbf{H}_\\mathbf{f}^T(X_{k-1}) + Q$$\n",
"\n",
"and the updates are now:\n",
"\n",
"$$ X_k \\approx \\hat{X}_k + K_k\\big(z_k - \\mathbf{h}(\\hat{X}_k)\\big)$$\n",
"$$ K_k = \\hat{P}_k \\mathbf{H}_\\mathbf{h}(\\hat{X}_k) \\Big(\\mathbf{H}_\\mathbf{h}(\\hat{X}_k)\\hat{P}_k \\mathbf{H}_\\mathbf{h}(\\hat{X}_k) + R_k\\Big)^{-1} $$\n",
"$$ P_k = \\big(I - K_k \\mathbf{H}_\\mathbf{h}(\\hat{X}_k)\\big)\\hat{P}_K$$\n",
"\n",
"In Pyro, all we need to do is create an `EKFState` object and use its `predict` and `update` methods. Pyro will do exact inference to compute the innovations and we will use SVI to learn a MAP estimate of the position and measurement covariances.\n",
"\n",
"As an example, let's look at an object moving at near-constant velocity in 2-D in a discrete time space over 100 time steps."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"import math\n",
"\n",
"import torch\n",
"import pyro\n",
"import pyro.distributions as dist\n",
"from pyro.infer.autoguide import AutoDelta\n",
"from pyro.optim import Adam\n",
"from pyro.infer import SVI, Trace_ELBO, config_enumerate\n",
"from pyro.contrib.tracking.extended_kalman_filter import EKFState\n",
"from pyro.contrib.tracking.distributions import EKFDistribution\n",
"from pyro.contrib.tracking.dynamic_models import NcvContinuous\n",
"from pyro.contrib.tracking.measurements import PositionMeasurement\n",
"\n",
"smoke_test = ('CI' in os.environ)\n",
"assert pyro.__version__.startswith('1.8.3')"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"dt = 1e-2\n",
"num_frames = 10\n",
"dim = 4\n",
"\n",
"# Continuous model\n",
"ncv = NcvContinuous(dim, 2.0)\n",
"\n",
"# Truth trajectory\n",
"xs_truth = torch.zeros(num_frames, dim)\n",
"# initial direction\n",
"theta0_truth = 0.0\n",
"# initial state\n",
"with torch.no_grad():\n",
" xs_truth[0, :] = torch.tensor([0.0, 0.0, math.cos(theta0_truth), math.sin(theta0_truth)])\n",
" for frame_num in range(1, num_frames):\n",
" # sample independent process noise\n",
" dx = pyro.sample('process_noise_{}'.format(frame_num), ncv.process_noise_dist(dt))\n",
" xs_truth[frame_num, :] = ncv(xs_truth[frame_num-1, :], dt=dt) + dx"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Next, let's specify the measurements. Notice that we only measure the positions of the particle."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Measurements\n",
"measurements = []\n",
"mean = torch.zeros(2)\n",
"# no correlations\n",
"cov = 1e-5 * torch.eye(2)\n",
"with torch.no_grad():\n",
" # sample independent measurement noise\n",
" dzs = pyro.sample('dzs', dist.MultivariateNormal(mean, cov).expand((num_frames,)))\n",
" # compute measurement means\n",
" zs = xs_truth[:, :2] + dzs"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We'll use a [Delta autoguide](http://docs.pyro.ai/en/dev/infer.autoguide.html#autodelta) to learn MAP estimates of the position and measurement covariances. The `EKFDistribution` computes the joint log density of all of the EKF states given a tensor of sequential measurements."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def model(data):\n",
" # a HalfNormal can be used here as well\n",
" R = pyro.sample('pv_cov', dist.HalfCauchy(2e-6)) * torch.eye(4)\n",
" Q = pyro.sample('measurement_cov', dist.HalfCauchy(1e-6)) * torch.eye(2)\n",
" # observe the measurements\n",
" pyro.sample('track_{}'.format(i), EKFDistribution(xs_truth[0], R, ncv,\n",
" Q, time_steps=num_frames),\n",
" obs=data)\n",
" \n",
"guide = AutoDelta(model) # MAP estimation"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"optim = pyro.optim.Adam({'lr': 2e-2})\n",
"svi = SVI(model, guide, optim, loss=Trace_ELBO(retain_graph=True))\n",
"\n",
"pyro.set_rng_seed(0)\n",
"pyro.clear_param_store()\n",
"\n",
"for i in range(250 if not smoke_test else 2):\n",
" loss = svi.step(zs)\n",
" if not i % 10:\n",
" print('loss: ', loss)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# retrieve states for visualization\n",
"R = guide()['pv_cov'] * torch.eye(4)\n",
"Q = guide()['measurement_cov'] * torch.eye(2)\n",
"ekf_dist = EKFDistribution(xs_truth[0], R, ncv, Q, time_steps=num_frames)\n",
"states= ekf_dist.filter_states(zs)"
]
},
{
"cell_type": "raw",
"metadata": {
"raw_mimetype": "text/html"
},
"source": [
"

\n",
" \n",
" \n",
" \n",
" \n",
"

\n",
" \n",
"

\n",
" \n",
" Figure 1:True track and EKF prediction with error. \n",
"

\n",
"

"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## References\n",
"\n",
"\\[1\\] Kalman, R. E. *A New Approach to Linear Filtering and Prediction Problems.* 1960\n",
"\n",
"\\[2\\] Welch, Greg, and Bishop, Gary. *An Introduction to the Kalman Filter.* 2006.\n"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.6.10"
}
},
"nbformat": 4,
"nbformat_minor": 2
}