Gaussian Multivariate Bayesian State Inference

Using RxInfer to estimate the hidden state vector

Aerospace Industry
Bayesian Inference
Active Inference
RxInfer
Julia
Author

Kobus Esterhuysen

Published

October 11, 2024

Modified

November 4, 2026

Back to Blog |  LearnableLoopAI.com |  Portfolio of Projects |  LinkedIn


Multivariate Bayesian State estimation is performed using the RxInfer Julia package. The hidden states are a vector of three Gaussians.

versioninfo() ## Julia version
Julia Version 1.10.5
Commit 6f3fdf7b362 (2024-08-27 14:19 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 12 × Intel(R) Core(TM) i7-8700B CPU @ 3.20GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-15.0.7 (ORCJIT, skylake)
Threads: 1 default, 0 interactive, 1 GC (on 12 virtual cores)
import Pkg
Pkg.add(Pkg.PackageSpec(;name="RxInfer"))
Pkg.add("Plots")
Pkg.add("StableRNGs")
Pkg.add("LaTeXStrings")
Pkg.add("Distributions")

using RxInfer, Random, Plots, StableRNGs, LaTeXStrings, Distributions
   Resolving package versions...
  No Changes to `/workspaces/GaussianMvStateInference/Project.toml`
  No Changes to `/workspaces/GaussianMvStateInference/Manifest.toml`
   Resolving package versions...
  No Changes to `/workspaces/GaussianMvStateInference/Project.toml`
  No Changes to `/workspaces/GaussianMvStateInference/Manifest.toml`
   Resolving package versions...
  No Changes to `/workspaces/GaussianMvStateInference/Project.toml`
  No Changes to `/workspaces/GaussianMvStateInference/Manifest.toml`
   Resolving package versions...
  No Changes to `/workspaces/GaussianMvStateInference/Project.toml`
  No Changes to `/workspaces/GaussianMvStateInference/Manifest.toml`
   Resolving package versions...
  No Changes to `/workspaces/GaussianMvStateInference/Project.toml`
  No Changes to `/workspaces/GaussianMvStateInference/Manifest.toml`
Pkg.status()
Status `/workspaces/GaussianMvStateInference/Project.toml`
  [6e4b80f9] BenchmarkTools v1.8.0
⌃ [31c24e10] Distributions v0.25.131
  [b964fa9f] LaTeXStrings v1.4.1
⌃ [91a5bcdd] Plots v1.41.6
⌃ [86711068] RxInfer v3.7.1
⌃ [860ef19b] StableRNGs v1.0.4
Info Packages marked with ⌃ have new versions available and may be upgradable.

Multivariate Linear Gaussian State Space Model

Working notation (FoAI Spec v1.23, with local choices for this notebook):

Object Form
Continuous state / observation italic bold arrays \(\boldsymbol{x}^{(t)}\), \(\boldsymbol{y}^{(t)}\) (D33)
Component (1-based here) \(x^{(t)[n]}\), \(n\in\{1,2,3\}\) — local exception to Spec zero-indexing; matches Julia
Classical LTI / engineering matrices bold upright sans-serif \boldsymbol{\mathsf{\ldots}}: \(\boldsymbol{\mathsf{A}},\boldsymbol{\mathsf{C}},\boldsymbol{\mathsf{Q}},\boldsymbol{\mathsf{R}},\boldsymbol{\mathsf{\Phi}}\) (D40–D41); not AIF \(\boldsymbol{A}\)–\(\boldsymbol{D}\); not light \mathsf{A}
Code mark for those matrices ordinary _ (_A, _Q, …); no separate sans code face (D41)
Trajectories (code) Spec ː (triangular colon), not _seq: _xˣː, _yː, _xː
Process (true) quantities asterisk before indices: \(\boldsymbol{x}^{\ast(t)}\), \(\boldsymbol{\mathsf{A}}^{\ast}\) (D36)
Biased generative model \(\breve{P}\) (D39) — unused in this notebook; do not put breve on matrices
Family parameters semicolon: \(\mathcal{N}(\boldsymbol{x}^{(t)};\boldsymbol{\mathsf{A}}\boldsymbol{x}^{(t-1)},\boldsymbol{\mathsf{Q}})\) (D34)

A multivariate Multivariate Linear Gaussian State Space Model (LGSSM) can be described with the equations:

\[\begin{aligned} p(\boldsymbol{x}^{(t)}\mid\boldsymbol{x}^{(t - 1)}) & = \mathcal{N}(\boldsymbol{x}^{(t)};\boldsymbol{\mathsf{A}}\, \boldsymbol{x}^{(t - 1)}, \boldsymbol{\mathsf{Q}}),\\ p(\boldsymbol{y}^{(t)}\mid\boldsymbol{x}^{(t)}) & = \mathcal{N}(\boldsymbol{y}^{(t)};\boldsymbol{\mathsf{C}}\, \boldsymbol{x}^{(t)}, \boldsymbol{\mathsf{R}}), \end{aligned}\]

where \(\boldsymbol{x}^{(t)}\) are hidden states, \(\boldsymbol{y}^{(t)}\) are noisy observations, \(\boldsymbol{\mathsf{A}}\), \(\boldsymbol{\mathsf{C}}\) are state transition and observation matrices, \(\boldsymbol{\mathsf{Q}}\) and \(\boldsymbol{\mathsf{R}}\) are state transition noise and observation noise covariance matrices (bold upright sans-serif \boldsymbol{\mathsf{\ldots}}: classical LTI / engineering matrices, Spec D40–D41 v1.23; not AIF \(\boldsymbol{A}\)–\(\boldsymbol{D}\)).

To make things more interesting, we will use a state space model that is subject to rotation in 3 dimensions, i.e. along the x, y, and z axes. This is a common use case for flying aircraft that rotate relative to a ground-based frame of coordinates. Once we have the transition and observation matrices for rotation around all 3 axes, we will use these matrices for the final state space model. Eventually, we will perform bayesian multivariate inference of the random vector \(\boldsymbol{x}^{(t)}\). Note that this state space model does not capture the complete dynamics of a rotating object - only the location is included.

  1. State Space Model without Rotation (3D)

\[ \boldsymbol{x}^{(t)} = \boldsymbol{\mathsf{A}}\, \boldsymbol{x}^{(t-1)} \] \[ \boldsymbol{y}^{(t)} = \boldsymbol{\mathsf{C}}\, \boldsymbol{x}^{(t)} \]

where: \[ \boldsymbol{x}^{(t)} = \begin{bmatrix} x^{(t)[1]} \\ x^{(t)[2]} \\ x^{(t)[3]} \end{bmatrix} \] and these entries are the location coordinates along the spatial \(x\), \(y\), and \(z\) axes.

\[ \boldsymbol{y}^{(t)} = \begin{bmatrix} y^{(t)[1]} \\ y^{(t)[2]} \\ y^{(t)[3]} \end{bmatrix} \]

  1. Rotation Matrices for Each Axis
  • Rotation around the x-axis by \(\alpha\):

\[ \boldsymbol{\mathsf{\Phi}}_x(\alpha) = \begin{bmatrix} 1 & 0 & 0 \\ 0 & \cos(\alpha) & -\sin(\alpha) \\ 0 & \sin(\alpha) & \cos(\alpha) \end{bmatrix} \]

  • Rotation around the y-axis by \(\beta\):

\[ \boldsymbol{\mathsf{\Phi}}_y(\beta) = \begin{bmatrix} \cos(\beta) & 0 & \sin(\beta) \\ 0 & 1 & 0 \\ -\sin(\beta) & 0 & \cos(\beta) \end{bmatrix} \]

  • Rotation around the z-axis by \(\gamma\):

\[ \boldsymbol{\mathsf{\Phi}}_z(\gamma) = \begin{bmatrix} \cos(\gamma) & -\sin(\gamma) & 0 \\ \sin(\gamma) & \cos(\gamma) & 0 \\ 0 & 0 & 1 \end{bmatrix} \]

  1. Combined Rotation Matrix in 3D

\[ \boldsymbol{\mathsf{\Phi}}(\alpha, \beta, \gamma) = \boldsymbol{\mathsf{\Phi}}_z(\gamma) \boldsymbol{\mathsf{\Phi}}_y(\beta) \boldsymbol{\mathsf{\Phi}}_x(\alpha) \]

\[ \boldsymbol{\mathsf{\Phi}}(\alpha, \beta, \gamma) = \begin{bmatrix} \cos(\gamma)\cos(\beta) & \cos(\gamma)\sin(\beta)\sin(\alpha) - \sin(\gamma)\cos(\alpha) & \cos(\gamma)\sin(\beta)\cos(\alpha) + \sin(\gamma)\sin(\alpha) \\ \sin(\gamma)\cos(\beta) & \sin(\gamma)\sin(\beta)\sin(\alpha) + \cos(\gamma)\cos(\alpha) & \sin(\gamma)\sin(\beta)\cos(\alpha) - \cos(\gamma)\sin(\alpha) \\ -\sin(\beta) & \cos(\beta)\sin(\alpha) & \cos(\beta)\cos(\alpha) \end{bmatrix} \]

  1. State Space Model with Rotation (3D)

\[ \boldsymbol{x}^{(t)} = \boldsymbol{\mathsf{A}}\, \boldsymbol{\mathsf{\Phi}}(\alpha, \beta, \gamma) \boldsymbol{x}^{(t-1)} = \boldsymbol{\mathsf{\Phi}}(\alpha, \beta, \gamma) \boldsymbol{x}^{(t-1)} = \boldsymbol{\mathsf{A}}\, \boldsymbol{x}^{(t-1)} \] (when the non-rotation factor \(\boldsymbol{\mathsf{A}}\) is the identity, so the effective transition is \(\boldsymbol{\mathsf{A}}\triangleq\boldsymbol{\mathsf{\Phi}}(\alpha,\beta,\gamma)\))

and the output equation:

\[ \boldsymbol{y}^{(t)} = \boldsymbol{\mathsf{C}}\, \boldsymbol{\mathsf{\Phi}}(\alpha, \beta, \gamma) \boldsymbol{x}^{(t)} = \boldsymbol{\mathsf{\Phi}}(\alpha, \beta, \gamma) \boldsymbol{x}^{(t)} = \boldsymbol{\mathsf{C}}\, \boldsymbol{x}^{(t)} \] (when the non-rotation factor \(\boldsymbol{\mathsf{C}}\) is the identity, so the effective observation matrix is \(\boldsymbol{\mathsf{C}}\triangleq\boldsymbol{\mathsf{\Phi}}(\alpha,\beta,\gamma)\)).

seed = 777
rng = MersenneTwister(seed)
_xˣ₀ = [ 1.0, 2.0, 3.0 ]
αˣ = 0.025 ## radians
βˣ = 0.03 ## radians
γˣ = 0.02 ## radians
_Φˣ_x = [1        0         0;
         0        cos(αˣ)  -sin(αˣ);
         0        sin(αˣ)  cos(αˣ)]
_Φˣ_y = [cos(βˣ)  0        sin(βˣ);
         0         1        0;
         -sin(βˣ) 0        cos(βˣ)]
_Φˣ_z = [cos(γˣ)  -sin(γˣ) 0;
         sin(γˣ)  cos(γˣ)  0;
         0        0        1]
_Φˣ = _Φˣ_x * _Φˣ_y * _Φˣ_z
_Aˣ = [1 0 0;
       0 1 0;
       0 0 1]*_Φˣ
_Cˣ = [1 0 0;
       0 1 0;
       0 0 1]*_Φˣ
_Qˣ = diageye(3)
_Rˣ = 20.0 .* diageye(3)
T = 200; ## number of observations
_Φˣ
3×3 Matrix{Float64}:
  0.99935    -0.0199897   0.0299955
  0.0207421   0.999473   -0.0249861
 -0.0294802   0.0255921   0.999238
_Aˣ
3×3 Matrix{Float64}:
  0.99935    -0.0199897   0.0299955
  0.0207421   0.999473   -0.0249861
 -0.0294802   0.0255921   0.999238
_Cˣ
3×3 Matrix{Float64}:
  0.99935    -0.0199897   0.0299955
  0.0207421   0.999473   -0.0249861
 -0.0294802   0.0255921   0.999238
_Qˣ
3×3 Matrix{Float64}:
 1.0  0.0  0.0
 0.0  1.0  0.0
 0.0  0.0  1.0
_Rˣ
3×3 Matrix{Float64}:
 20.0   0.0   0.0
  0.0  20.0   0.0
  0.0   0.0  20.0

The Generative Process

Next, we will generate some synthetic data.

Provision function (\(f_p\))

The provision function, provides another covariate vector. Because this is a sequential system, the provision function defines the transition between the previous state and the current state. This special case of the provision function is known as a transition function and it returns a provision/pre-state:

\[\boldsymbol{p}^{\ast(t)} = f_p(\boldsymbol{x}^{\ast(t-1)}) = f_{\mathsf{A}}(\boldsymbol{x}^{\ast(t-1)}) = \boldsymbol{\mathsf{A}}^{\ast} \boldsymbol{x}^{\ast(t-1)}\]

When we add system noise the next state is produced:

\[\boldsymbol{x}^{\ast(t)} \sim \mathcal{N}(\boldsymbol{p}^{\ast(t)}, \boldsymbol{\mathsf{Q}}^{\ast})\]

The process asterisk \(^{\ast}\) marks true (generative-process) quantities; it precedes any indices (e.g. \(\boldsymbol{x}^{\ast(t)}\)).

## provision function, provides another covariate vector
function fₚ(; _Aˣ, _xˣₜ₋₁)
    return _Aˣ * _xˣₜ₋₁
end
fₚ(_Aˣ=_Aˣ, _xˣₜ₋₁=_xˣ₀)
3-element Vector{Float64}:
 1.049357295076791
 1.9447288135798293
 3.019417014740016

Response function (\(f_r\))

The response function, provides the response to the covariate vector, called the response: \[\boldsymbol{r}^{\ast(t)} = f_r(\boldsymbol{x}^{\ast(t)}) = f_{\mathsf{C}}(\boldsymbol{x}^{\ast(t)}) = \boldsymbol{\mathsf{C}}^{\ast} \boldsymbol{x}^{\ast(t)}\]

After combining with observation noise the observation is produced:

\[\boldsymbol{y}^{(t)} \sim \mathcal{N}(\boldsymbol{r}^{\ast(t)}, \boldsymbol{\mathsf{R}}^{\ast})\]

The process asterisk \(^{\ast}\) marks true (generative-process) quantities; it precedes any indices (e.g. \(\boldsymbol{x}^{\ast(t)}\)). Observations \(\boldsymbol{y}^{(t)}\) carry no asterisk.

## response function, provides the response to a covariate vector
function fᵣ(; _Cˣ, _xˣₜ)
    return _Cˣ * _xˣₜ
end
fᵣ(; _Cˣ=_Cˣ, _xˣₜ=_xˣ₀)
3-element Vector{Float64}:
 1.049357295076791
 1.9447288135798293
 3.019417014740016
## Data comes from either a simulation/lab (sim|lab) OR from the field (fld)
## Data are handled either in batches (batch) OR online as individual points (point)
## Batch data accumulates either
    ## along the depth/examples dimension/axis (into the screen/page), OR
        ## typical for supervised & unsupervised learning
    ## along the time dimension/axis (down the screen page)
        ## typical for sequential decision learning (reinforcement learning & active inference)
function sim_batch_data(rng, T, _Aˣ, _Cˣ, _Qˣ, _Rˣ) ## simulated batch data
    _xˣₜ₋₁ = _xˣ₀
    _pˣː = Vector{Vector{Float64}}(undef, T)
    _xˣː = Vector{Vector{Float64}}(undef, T)
    _rˣː = Vector{Vector{Float64}}(undef, T)
    _yː = Vector{Vector{Float64}}(undef, T)
    for t in 1:T
        ## _pˣː[t] = _Aˣ * _xˣₜ₋₁
        _pˣː[t] = fₚ(_Aˣ=_Aˣ, _xˣₜ₋₁=_xˣₜ₋₁)
        _xˣː[t] = rand(rng, MvNormal(_pˣː[t], _Qˣ))

        ## _rˣː[t] = _Cˣ * _xˣː[t]
        _rˣː[t] = fᵣ(_Cˣ=_Cˣ, _xˣₜ=_xˣː[t])
        _yː[t] = rand(rng, MvNormal(_rˣː[t], _Rˣ))

        _xˣₜ₋₁ = _xˣː[t]
    end
    return _xˣː, _yː
end
sim_batch_data (generic function with 1 method)
_xˣː, _yː = sim_batch_data(rng, T, _Aˣ, _Cˣ, _Qˣ, _Rˣ);
_xˣː
200-element Vector{Vector{Float64}}:
 [0.816809072166705, 1.7217016515618577, 1.7266544980425234]
 [0.8462870749788376, 1.7866860719285167, 2.284274407952872]
 [0.9240522716813285, 1.418720769965815, 0.8247663268344083]
 [1.5074917335918394, 0.7066355754400989, 0.8491162345690277]
 [0.004147904602415986, 0.08005081422059745, 1.3484301531004605]
 [-0.40918536446570625, 0.5108956198778928, 1.2731758674141656]
 [-0.29479303139789226, 0.15609530935162597, 2.960739501191896]
 [-0.8665382984624709, 0.16374219670328205, 3.088073982225083]
 [-0.23988074147935623, 1.1424146037019616, 2.1747840274034957]
 [-1.0003584001330097, 2.852304474516476, 1.415656287370715]
 ⋮
 [-2.537308148940686, 15.591255092640719, 31.074458843049115]
 [-2.1058359114778606, 13.06175160006155, 32.669454354693556]
 [-0.0028260097871883882, 11.648552923438668, 33.356051855748966]
 [0.5719478763216154, 10.351997913061277, 34.50540423895058]
 [0.8630023982128899, 9.221911298703859, 35.250394798844084]
 [1.529652990431234, 10.229648860494272, 35.60678023193549]
 [0.24115939793664687, 10.256648422027853, 36.4859960288678]
 [1.521104148818912, 8.028052604901172, 36.928855425490276]
 [2.9986191444712427, 8.130049619613066, 35.814649123738945]
_yː
200-element Vector{Vector{Float64}}:
 [-3.169472568477776, -3.0112527573196353, -7.009918788993621]
 [-0.19736298654548257, -0.9315849230039661, 4.794843936719561]
 [-4.288405615039999, -5.762609514878842, -8.122024186064301]
 [2.8069893736881704, 1.7982263833951335, 7.44462480949098]
 [4.438067867580971, 0.19944421500499462, 9.017867444981764]
 [6.993569706363068, -3.1249164236964346, 0.9734242907700942]
 [-3.1965703217997126, 3.819146546268818, 2.698419965079931]
 [-1.1740176313445088, 0.04818588007655261, -4.817597171181123]
 [8.627593419582793, -0.6712097565032804, 8.148547430458786]
 [-9.742187368874497, 0.8006839079242203, 3.5438536430998697]
 ⋮
 [-1.6344882281183137, 20.829889441875068, 23.274840442528472]
 [1.2481039038512645, 7.542168851587531, 30.462872006215893]
 [7.131212683295738, 15.018484882402486, 41.0097343948868]
 [4.849334373916948, 7.283164402494926, 36.43111168355113]
 [2.6856851138600195, 16.090514888838772, 39.147352264172525]
 [-4.122825845486582, 12.324991477304, 37.39387349737904]
 [-2.4770139986439768, 7.1290241560389, 33.335583694008776]
 [-3.6565065833463866, 8.006062912322305, 36.2493283324553]
 [-0.6635946016021235, 2.6901737288236562, 41.04715457107512]

Let’s visualize the synthesized dataset. Lines represent the hidden states that need to be estimated/inferred. We only have acces to noisy observations which are represented as dots.

_p = plot(title="Hidden states with noisy observations", legendfontsize=12)

_p = plot!(_p, getindex.(_xˣː, 1), label="Hidden Signal " * L"x^{\ast[1]}", color=:red)
_p = scatter!(_p, getindex.(_yː, 1), label=false, markersize=2, color=:red)

_p = plot!(_p, getindex.(_xˣː, 2), label="Hidden Signal " * L"x^{\ast[2]}", color=:green)
_p = scatter!(_p, getindex.(_yː, 2), label=false, markersize=2, color=:green)

_p = plot!(_p, getindex.(_xˣː, 3), label="Hidden Signal " * L"x^{\ast[3]}", color=:blue)
_p = scatter!(_p, getindex.(_yː, 3), label=false, markersize=2, color=:blue)

plot(_p)

The Generative Model

We now use RxInfer:

@model function rotate_ssm(_yː, _x₀, _A, _C, _Q, _R)
    _x_prior ~ MvNormalMeanCovariance(mean(_x₀), cov(_x₀))
    _xₜ₋₁ = _x_prior
    for t in 1:length(_yː)
        _xː[t] ~ MvNormalMeanCovariance(_A*_xₜ₋₁, _Q) ## `_xː` trajectory of hidden states (Spec ː)
        ##- s[t] ~ MvNormalMeanCovariance(f(B= B, sₜ₋₁= sₜ₋₁), Q)
        _yː[t] ~ MvNormalMeanCovariance(_C*_xː[t], _R) ## `_yː` trajectory of clamped observations
        _xₜ₋₁ = _xː[t]
    end
end

Specify a prior for the initial hidden state:

_xˣ₀ = MvNormalMeanCovariance(zeros(3), 10.0*diageye(3));

Perform inference:

## We assume the process matrices _Aˣ, _Cˣ, _Qˣ, _Rˣ are known (not inferred).
## The ˣ mark is the process asterisk (Spec D36), not a tilde.
result = infer(
    model=       rotate_ssm(_x₀=_xˣ₀, _A=_Aˣ, _C=_Cˣ, _Q=_Qˣ, _R=_Rˣ),
    data=        (_yː = _yː,),
    free_energy= true
);

Extract the results:

result.posteriors
Dict{Symbol, Any} with 2 entries:
  :_xː      => MvNormalWeightedMeanPrecision{Float64, Vector{Float64}, Matrix{F…
  :_x_prior => MvNormalWeightedMeanPrecision(…
xmarginals  = result.posteriors[:_xː];

Visualize:

_p = plot(title="Estimated states from noisy observations", legendfontsize=12)
_p = plot!(_p, getindex.(_xˣː, 1), label="Hidden Signal " * L"x^{\ast[1]}", color=:red, linestyle=:dash)
_p = plot!(_p, getindex.(_xˣː, 2), label="Hidden Signal " * L"x^{\ast[2]}", color=:green, linestyle=:dash)
_p = plot!(_p, getindex.(_xˣː, 3), label="Hidden Signal " * L"x^{\ast[3]}", color=:blue, linestyle=:dash)

_p = plot!(_p, getindex.(mean.(xmarginals), 1), ribbon=getindex.(var.(xmarginals), 1) .|> sqrt, fillalpha=0.5, label="Estimated Signal " * L"x^{[1]}", color=:pink)
_p = plot!(_p, getindex.(mean.(xmarginals), 2), ribbon=getindex.(var.(xmarginals), 2) .|> sqrt, fillalpha=0.5, label="Estimated Signal " * L"x^{[2]}", color=:lightgreen)
_p = plot!(_p, getindex.(mean.(xmarginals), 3), ribbon=getindex.(var.(xmarginals), 3) .|> sqrt, fillalpha=0.5, label="Estimated Signal " * L"x^{[3]}", color=:lightblue)
plot(_p)

As we can see from our plot, estimated signal resembles closely to the real hidden states with small variance. We maybe also interested in the value for minus log evidence:

## given the analytical solution, the free energy will be equal to the negative log evidence
logevidence = -result.free_energy; 
logevidence
1-element Vector{Float64}:
 -1814.862291590095