HELLO CYBERNETICS

深層学習、機械学習、強化学習、信号処理、制御工学、量子計算などをテーマに扱っていきます

JuliaでExtended Kalman Filterを書いてみる

 

 

follow us in feedly

はじめに

今回は拡張カルマンフィルタをお題にJuliaの練習をしてみようと思いました。 色々Julia的な書き方をしようと頑張ったのですが、もしもあとで状態遷移関数を渡しさえすればヤコビ行列を計算する関数を保持してくれる保持してくれるオブジェクトを自動微分で作れたら…とか考えてたら、最終的に構造体に関数を持たせるという、思いっきりクラス-メソッド的な実装になってしまったのですが、致し方なし。

とりあえず備忘録として、いつかjuliaをもっとちゃんと書けるようになったときに、こんな躓き方していたなぁと、自分で振り替えれるように記事にします。

このコードを書く上ではPython Roboticsのリポジトリを大いに参考にしました(というか問題設定はそのまま借りて、実装だけしなおした感じ)。

github.com

拡張カルマンフィルタのモジュールを作成する

状態空間モデルは

$$ \begin{align} x[t+1] = F(x[t], u[t]) \\ y[t+1] = H(x[t+1]) \end{align} $$

という形式で二本立ての構造を取ります。ロボットの世界では特に1本目の式がロボットの移動を記述する方程式であり、2本目がセンサーによる観測を記述する方程式になります。というわけで、今回はPythonRoboticsを題材にしているということでOdModelSnsModel をそれぞれ1本目の式と2本目の式を表すために作るとしましょう。

次に状態空間モデルの抽象型を作っておきます。 別に今回は意味はないのですが、Pythonで言うところの抽象クラスに相当します。抽象型は型の階層構造の出発点になり、構造を整理するために利用します。 今回は拡張カルマンフィルタを扱うので、対象となる状態空間モデルはガウス型になります。なのでGaussStateSpaceModelという構造体を作っておくことにします。

#ExtendedKalmanFilter.jl
using LinearAlgebra

struct OdModel
    F::Function
    ∂F::Function
end

struct SnsModel
    H::Function
    ∂H::Function
end

abstract type StateSpaceModel end

struct GaussStateSpaceModel <: StateSpaceModel
    odmodel::OdModel
    snsmodel::SnsModel
    Q::Matrix{Float64}
    R::Matrix{Float64}
end

状態空間モデルに対しては、状態発展と時間発展が非線形であったとしても、それを各時刻において都度テイラー展開で線形近似することで、一応線形なモデルとして扱うことができます。また線形ガウスモデルまで落とし込めれば、観測から状態の推定量がカルマンフィルタによって得られます。このテイラー展開とカルマンフィルタを組み合わせた手法を拡張カルマンフィルタというのでした。

この推定の手続きは状態方程式を用いた予測と、観測方程式を用いた更新の2つによって実施されます。

function predict(m::GaussStateSpaceModel, xEst, PEst, u)
    #  Predict
    xPred = m.odmodel.F(xEst, u)
    ∂F = m.odmodel.∂F(xEst, u)
    PPred = ∂F*PEst*∂F' + Q
    return xPred, PPred
end

function update(m::GaussStateSpaceModel, xPred, PPred, z)
    #  Update
    ∂H = m.snsmodel.∂H()
    zPred = m.snsmodel.H(xPred)
    y = z - zPred
    S = ∂H*PPred*∂H' + R
    K = PPred*∂H'*inv(S)
    xEst = xPred + K*y
    # Matrix{::T}(I(K)) create KxK indentity matrix of T type  
    Iₓ = Matrix{typeof(xEst[1])}(I(size(xEst,1)))
    PEst = (Iₓ - K*∂H)*PPred
    return (xEst, PEst)
end

function filt_est(m::GaussStateSpaceModel, xEst, PEst, z, u)
    xPred, PPred = predict(m, xEst, PEst, u)
    return update(m, xPred, PPred, z)
end

このように書いておけば、あとは状態方程式と観測方程式、およびその線形近似関数(ヤコビ行列)を作ってしまえば各々の問題に適用することができます。

Juliaはモジュールを作る手続きがシステム化されており、だれでも簡単にパッケージを作成できます。 今回はそれらに触れていると話がそれてしまうため、include("ExtendedKalmanFilter.jl") を使って他のソースコードにまるまる上記のコードを(実質的に)コピペする形式で実装してしまいます。

拡張カルマンフィルタを使ってみる

一気にコードを貼ってしまいます。

まずすでに述べたとおり include("ExtendedKalmanFilter.jl") によって別ファイルの ExtendedKalmanFilter.jl の中身があたかもそこに書かれているかのように扱うことができます。要するに一行でコピペ相当になるということです。

次に、Juliaでは素性の不明なグローバル変数は非常にパフォーマンスを悪くします。素性が悪いというのは、あとで値が変更されたり、型が決まっていなかったりするような状態を指します。 最も簡単にこの問題を解決する方法は const をつけてしまうことです。こうすることでJuliaがこの変数は一切今後変更されることは無いということで、安心してメモリを確定させてくれます。

using PyCall
@pyimport matplotlib.pyplot as plt
using Plots
using ProgressBars
include("ExtendedKalmanFilter.jl")

# Plot params
const SHOW_ANIMATION = true
const SAVE_FIG = false
# Simulation parameter
const INPUT_NOISE = diagm([1.0, deg2rad(30.0)]).^2
const GPS_NOISE = diagm([0.5, 0.5]).^2
const DT = 0.1  # time tick [s]
const SIM_TIME = 50.0  # simulation time [s]

function calc_input()
    v = 1.5  # [m/s]
    yawrate = 0.15  # [rad/s]
    u = [v; yawrate]
    return u
end

# Covariance for EKF simulation
const Q = diagm([
    0.1,  # variance of location on x-axis
    0.1,  # variance of location on y-axis
    deg2rad(1.0),  # variance of yaw angle
    1.0  # variance of velocity
]).^2  # predict state covariance
const R = diagm([1.0, 1.0]).^2

function state_step(x, u)
    F = [1.0 0.0 0.0 0.0;
         0.0 1.0 0.0 0.0;
         0.0 0.0 1.0 0.0;
         0.0 0.0 0.0 1.0]
    B = [DT*cos(x[3,1]) 0.0;
         DT*sin(x[3,1]) 0.0;
         0.0             DT;
         1.0            0.0]
    x = F*x + B*u
    return x
end

function obs_step(x)
    H = [1.0 0.0 0.0 0.0;
         0.0 1.0 0.0 0.0]
    z = H*x
    return z
end

function jacob_f(x, u)
    yaw = x[3,1]
    v = u[1,1]
    ∂F = [1.0 0.0 -DT*v*sin(yaw) DT*cos(yaw);
           0.0 1.0  DT*v*cos(yaw) DT*sin(yaw);
           0.0 0.0            1.0         0.0;
           0.0 0.0            0.0         1.0]
    return ∂F
end

function jacob_h()
    # Jacobian of simulate_onestep Model
    ∂H = [1.0 0.0 0.0 0.0;
           0.0 1.0 0.0 0.0]
    return ∂H
end

function simlate_onestep(m::GaussStateSpaceModel, xEst, PEst, xTrue, xDR)
    u = calc_input()
    xTrue = m.odmodel.F(xTrue, u)
    z = m.snsmodel.H(xTrue) + GPS_NOISE * randn(2)
    ud = u + INPUT_NOISE * randn(2)
    xDR = m.odmodel.F(xDR, ud)
    xEst, PEst = filt_est(m, xEst, PEst, z, ud)
    return xEst, PEst, xTrue, xDR, z
end

# utilities
function makerot(angle)
    [cos(angle) -sin(angle); sin(angle) cos(angle)] 
end

function flatten_fn(x)
    collect(Iterators.flatten(x))
end

function plot_covariance_ellipse(xEst, PEst)  # pragma: no cover
    Pxy = PEst[1:2, 1:2]
    eig = eigen(Pxy)
    eigval = eig.values
    eigvec = eig.vectors
    if eigval[1] >= eigval[2]
        bigind = 1
        smallind = 2
    else
        bigind = 2
        smallind = 1
    end
    t = 0.0:0.1:2*pi+0.1
    a = sqrt(eigval[bigind])
    b = sqrt(eigval[smallind])
    x = a*cos.(t)
    y = b*sin.(t)
    angle = atan(eigvec[2, bigind], eigvec[1, bigind])
    rot = makerot(angle)
    fx = rot*[x y]'
    px = flatten_fn((fx[1, :] .+ xEst[1, 1])')
    py = flatten_fn((fx[2, :] .+ xEst[2, 1])')
    return px, py
end

function save_plot!(hz, hxTrue, hxDR, hxEst, px, py, anim)
    plot(hz[1, :], hz[2, :], label="obs", aspectratio=1.0, linecolor=:green, linestyle=:dot)
    plot!(hxTrue[1, :], hxTrue[2, :], label="true", linecolor=:blue)
    plot!(hxDR[1, :], hxDR[2, :], label="od + noise", linecolor=:black)
    plot!(hxEst[1, :], hxEst[2, :], label="est mean", linecolor=:red)
    plot!(px, py, label="est cov", linecolor=:red)
    frame(anim)     
end

function show_plot(hz, hxTrue, hxDR, hxEst, px, py)
    plt.cla()
    plt.plot(hz[1, :], hz[2, :], "-go", alpha=0.3)
    plt.plot(hxTrue[1, :], hxTrue[2, :], "-b")
    plt.plot(hxDR[1, :], hxDR[2, :], "-k")
    plt.plot(hxEst[1, :], hxEst[2, :], "-r")
    plt.plot(px, py, "--r")
    plt.legend(["obs", "true", "od + noise", "est kf", "cov"])
    plt.axis("equal")
    plt.grid(true)
    plt.pause(0.001)
end

function main(show_animation, save_plot, sim_time)

    od = OdModel(state_step, jacob_f)
    sns = SnsModel(obs_step, jacob_h)
    model = GaussStateSpaceModel(od, sns, Q, R)

    time = 0.0:DT:sim_time
    xEst::Vector{Float64} = zeros(4)
    xTrue::Vector{Float64} = zeros(4)
    PEst::Matrix{Float64} = Matrix{Float64}(I(4))
    xDR::Vector{Float64} = zeros(4)  # Dead reckoning

    # history
    hxEst = xEst
    hxTrue = xTrue
    hxDR = xTrue
    hz = zeros(2)
    anim = Animation()

    for t = ProgressBar(time)
        xEst, PEst, xTrue, xDR, z = simlate_onestep(model, xEst, PEst, xTrue, xDR)
        # store data history
        hxEst = hcat(hxEst, xEst)
        hxDR = hcat(hxDR, xDR)
        hxTrue = hcat(hxTrue, xTrue)
        hz = hcat(hz, z)
        px, py = plot_covariance_ellipse(xEst, PEst)

        if save_plot
            save_plot!(hz, hxTrue, hxDR, hxEst, px, py, anim)
        end
        
        if show_animation
            show_plot(hz, hxTrue, hxDR, hxEst, px, py)
        end

    end
    if SAVE_FIG
        gif(anim, "ekf.gif", fps=1/DT)
    end
end

main(SHOW_ANIMATION, SAVE_FIG, SIM_TIME)

備忘録なので特にくどくど解説はしません(し、そもそもまともなコードなのかもわかっていない)。 ひとまず上記の SAVE_FIG = true によって保存された gif をここに貼っておきます。

f:id:s0sem0y:20210502014305g:plain

ちなみに描画を実施すると、JuliaよりもPythonの方が速いです。 描画無しで数字だけ計算していくとJitの時間を除けばJuliaの方が3倍速く終わります。今回程度のシミュレーションだとJitがボトルネックになって、やはりPythonの方が速くなります。 パッとシミュレーションしてみたいだけならPythonの方が手軽で良いかもしれませんね。一方で、アルゴリズムが固まって、ガッツリ計算をし続ける場合はJuliaの方が良いということになるでしょうか(ただ、Pythonにもjax.numpyがいてforiloop速いですよね)。

GPUでの計算が主戦場になると、どの言語でも関係ない気がするので(CUDAを何で叩くかの違いなので)なかなか難しい立ち位置かもしれません。