This notebook looks at a technique for simultaneous localization (finding the position of a robot) and mapping (finding the positions of any obstacles), abbreviated as SLAM. In this model, the probability distribution for the robot's trajectory $x_{1:t}$ is represented with a set of weighted particles. Let the weight for particle $i$ be denoted $w^i$, and let $W$ denote the sum of all the weights. Each particle also keeps track of the conditional distribution of obstacle locations $m$ given $x_1, \dotsc x_t$. Even though each particle $i$ should be thought of as describing a full potential trajectory $x_1^i, \dotsc x_t^i$, we will only need to store the position at the most recent time $x_t^i$.

using LinearAlgebra, StaticArrays, LogExpFunctions, Distributions, ArraysOfArrays, FillArrays, Plots

const Dist = Tuple{SVector{2, Float64}, SMatrix{2,2, Float64}}

Tuple{SVector{2, Float64}, SMatrix{2, 2, Float64}}

sample(d::Dist) = rand(Distributions.MvNormal(d[1], d[2]))

sample (generic function with 1 method)

struct Particle
    weight::Float64 # stored in log scale
    x::SVector{2, Float64}
    m::Vector{Dist}
end

const k = 6; # Number of obstacles

function init_particles(n)
    d = (zeros(SVector{2}), SMatrix{2,2}(10.0I))
    p = Particle(0.0, zeros(SVector{2}), fill(d, k))
    fill(p, n)
end;

const n = 100; # Number of particles

particles = init_particles(n);

The particles $x^i$ at time $t$ will always be samples from a proposal distribution $q_t(x^i)$. The weights $w^i$ will be importance ratios of the form $C \pi_t(x^i)/q_t(x^i)$, where $\pi_t(x^i)$ is our target distribution and $C$ is a constant. Given these particles, we can approximate $E_{X_t \sim \pi_t} f(X_t)$ for any function $f$ by computing

$$ \frac{1}{W}\sum_{i=1}^n w^i f(x_t^i) \approx E_{x_t^i \sim q} \frac{\pi(x_t^i)f(x_t^i)}{q(x_t^i)}=\int \pi(x_t^i) dx_t^i = E_{x_t^i \sim \pi} f(x_t^i) $$

In our case, the target distribution $\pi_{t+1}(x^i)$ at time $t+1$ will be the conditional distribution $p(x^i_{1:t+1} | z_{1:t+1})$, and the proposal distribution $q_{t+1}(x^i) = \pi_t(x^i_{1:t})p(x_{t+1} | x_t, z_{1:t+1})$. By Bayes' rule

$$ \pi_{t+1}(x^i) \propto \pi_t(x^i) p(x_{t+1} | x_t) p(z_{t+1} | x_{t+1}, z_{1:t}) $$

Similarly,

$$ \begin{align*} q_{t+1}(x^i) &\propto \frac{\pi_t(x^i) p(x_{t+1}, z_{t+1} | x_t, z_{1:t})}{p(z_{t+1} | x_t)} \\ &= \frac{\pi_t(x^i) p(x_{t+1} | x_t) p(z_{t+1} | x_{t+1}, z_{1:t})}{p(z_{t+1} | x_t, z_{1:t})} \end{align*} $$

This means the ratio simplifies to $p(z_{t+1} | x_t, z_{1:t})$.

At each step $t$ in time, our algorithm will do the following: 1. Resample the particles in proportion to their importance weights. This effectively turns our samples from $q_t$ into samples from $\pi_t$. 2. Extend each particle, turning the samples from $\pi_t$ into samples from $q_{t+1}$. This requires sampling from $p(x_{t+1}^i | x_t^i, z_{1:t})$, computing the importance weight $p(z_{t+1} | x_T, z_{1:t})$, and updating the associated distribution $p(m | x_{1:t}, z_{1:t})$.

At $t=0$, the proposal and target distributions are the same so we don't need to resample.

function step(particles::Vector{Particle}, z)
    new_particles = propose.(particles, Ref(z))
    ix = resample(softmax([p.weight for p in new_particles]))
    resampled = new_particles[ix]
    Particle[Particle(0.0, p.x, p.m) for p in resampled]
end;

Resampling

If w is our vector of weights, we want to sample $k$ elements of w in proportion to their values. The most obvious way to do this is with inverse transform sampling. Let b = cumsum(w). For each random sample, we can pick a uniform $u_i$ between $0$ and $W$, which corresponds to searchsortedfirst(b, u_i).

To lower the variance of this approach, we can use stratified sampling. Instead of letting each $u_i$ come from the full interval $[0,W]$, we can restrict each to a separate interval of size $W/n$.

An even lower variance approach is to sample a single uniform number in $[0,W]$ and shift it to each of the stratified intervals. This is called systematic resampling.

function resample(w::Vector{Float64})
    u = rand()
    ui = [(u + i-1) / length(w) for i in 1:length(w)]
    bins = cumsum(w)
    [searchsortedfirst(bins, i) for i in ui]
end;

Proposing

We will assume a simple dynamics model where $p(x_{t+1} | x_t) = \mathcal{N}(x_t + \delta, \Sigma_A)$. For the likelihood, we observe noisy distances to obstacles, so that an observation at time $t$ is distributed $p(z_t | x_t, m) = \mathcal{N}(m - x_t, \Sigma_B)$.

To find the posterior $p(x_{t+1} | x_t, z_{t+1})$, it will be easier to work with the general case where $A \sim \mathcal{N}(\mu_A, \Sigma_A)$, $B \sim \mathcal{N}(\mu_B, \Sigma_B)$, and $Z \sim \mathcal{N}(B - A, \Sigma_Z)$. Marginalizing out $B$ from the definition of $Z$ gives $Z \sim \mathcal{N}(\mu_B - A, \Sigma_B + \Sigma_Z)$. Let $\Sigma_B + \Sigma_Z = \Sigma_{BZ}$ and $A \circ B \coloneqq \text{Trace}(A^TB)$. By Bayes' rule:

$$ \begin{align*} \log p(A | Z) &= \log p(A) + \log p(Z | A) + c \\ &= -\frac{1}{2}(A - \mu_A)^T \Sigma_A^{-1} (A - \mu_A) -\frac{1}{2}(Z - \mu_B + A)^T (\Sigma_{BZ})^{-1} (Z - \mu_B + A) +c \\ &= -\frac{1}{2} AA^T \circ \Sigma_A^{-1} + A^T \Sigma_A^{-1} \mu_A -\frac{1}{2} AA^T \circ \Sigma_{BZ}^{-1} + A^T \Sigma_{BZ}^{-1} (\mu_B - Z) \\ &= -\frac{1}{2} AA^T \circ (\Sigma_A^{-1} + \Sigma_{BZ}^{-1}) + A^T (\Sigma_A^{-1} \mu_A + \Sigma_{BZ}^{-1} (\mu_B - Z)) \end{align*} \\ $$

This shows that the posterior is Gaussian with natural parameters $(\Sigma_A^{-1} \mu_A + \Sigma_{BZ}^{-1} (\mu_B - Z),\, \Sigma_A^{-1} + \Sigma_{BZ}^{-1})$.

For FastSLAM, $\mu_A = x_t + \delta$, so $p(x_{t+1} | x_t, z_{t+1})$ has natural parameters

$$ (\Sigma_A^{-1} (x_t + \delta) + (\Sigma_B + \Sigma_Z)^{-1} (\mu_B - z_{t+1}),\, \Sigma_A^{-1} + (\Sigma_B + \Sigma_Z)^{-1}) $$

Similarly, to update the obstacle distribution upon seeing $x_{t+1}$ and $z_{t+1}$:

$$ \begin{align*} \log p(M | x_{1:t+1}, z_{1:t+1}) &= -\frac{1}{2} MM^T \circ (\Sigma_B^{-1} + \Sigma_Z^{-1}) + M^T (\Sigma_B^{-1} \mu_B + \Sigma_{Z}^{-1} (Z - x_{t+1})) + c \end{align*} $$

For the importance weight, note that $p(z_{t+1} | X_t, z_{1:t}) = \mathcal{N}(\mu_B - (x_t + \delta), \Sigma_B + \Sigma_Z)$.

begin
ΣA = Eye(2,2)
ΛA = inv(ΣA)
ΣZ =  (1/6) * Eye(2,2)
ΛZ = inv(ΣZ)
δ = SVector(2.,2.)
end;

function post_msg_x(m, z)
    ΣBZ = m[2] + ΣZ
    (ΣBZ \ (m[1] - z), inv(ΣBZ))
end;

post_msg_m(x, z) = (ΛZ \ (z - x), ΛZ)

post_msg_m (generic function with 1 method)

function propose(p, z)
    post_x_msgs = sum(post_msg_x.(p.m, z))
    xmean = (p.x + δ)
    prior_x_msg = (ΛA * xmean, ΛA)
    xpost = meanform(prior_x_msg .+ post_x_msgs)
    new_x = sample(xpost)
    w = sum(importance.(p.m, z, Ref(xmean)))
    post_m_msgs = post_msg_m.(Ref(new_x), z)
    mpost = meanform.(stdform.(p.m) .+ post_m_msgs)
    Particle(w, new_x, mpost)
end;

function importance(m, z, xmean)
    logpdf(Distributions.MvNormal(m[1] - xmean, m[2] + ΣZ), z)
end

importance (generic function with 1 method)

function stdform((mu, var))
    (var \ mu, inv(var))
end

stdform (generic function with 1 method)

function meanform((prec_mu, prec))
    (prec \ prec_mu, inv(prec))
end

meanform (generic function with 1 method)

Fake Data

Here we generate the unknown, true trajectory that the robot takes by following our dynamics model.

begin
RΣA = sqrt(ΣA)
true_xs = cumsum(δ .+ RΣA * randn(2, 100), dims=2)
end;

We'll also set up the true but unknown obstacle locations.

true_ms = 50 * randn(2, k);

begin
plot(true_xs[1, :], true_xs[2, :], label="robot path")
scatter!(true_ms[1,:], true_ms[2,:], label="obstacles")
end

svg

The displacement observations the robot sees along its trajectory look like this:

begin
RΣZ = sqrt(ΣZ)
true_zs = [(true_ms .- x) + RΣZ * randn(2, k) for x in nestedview(true_xs)]
z = flatview(ArrayOfSimilarArrays(true_zs));
end;

To show the direction of time, lines will get fainter with time.

timealpha = reverse(0.1 .+ ((1:size(z, 3))./size(z, 3)));

begin
plt = plot(z[1,1,:], z[2,1,:], label="z1", alpha=timealpha)
for i in 2:k
    plot!(z[1,i,:], z[2,i,:], label="z$i", alpha=timealpha)
end
plt
end

svg

Simulating the Algorithm

Base.:+(a::Tuple, b::Tuple) = a .+ b # So we can add distribution parameters

sz = map(SVector{2}, eachslice(z; dims=(3,2))); # Rewrites z as a matrix of SVec

steps = accumulate(step, eachslice(sz;dims=1); init=particles);

predicted = SVector{2, Float64}[mean(p-> p.x, s) for s in steps];

flat_pred = flatview(ArrayOfSimilarArrays(predicted));

begin
plot(true_xs[1, :], true_xs[2, :], label="robot path")
plot!(flat_pred[1,:], flat_pred[2,:], label="predictions", alpha=timealpha)
end

svg

Even though the robot's trajectory is noisy, the observations allow us to track it quite accurately.

Sam's Blog

Fast SLAM

Resampling

Proposing

Fake Data

Simulating the Algorithm