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

Hidden Markov Model with Control (Part 1)

1 BUSINESS UNDERSTANDING

This is the first part of the Hidden Markov Model with Control series, built on the Hidden Markov Model RxInfer example. Eventually this work will be used in a more elaborate scheduling project making use of active inference.

Text conventions. The names of model variables are written in italic lower case, such as whereabouts, and their values in double quotes, such as “bedroom”. Section 4.1 describes these conventions, and Appendix A gives their motivation.

In this part, we will simply be tracking a Roomba as it moves throughout a 3-room apartment consisting of a bedroom, a living room, and a bathroom. The agent infers the whereabouts at each time step and learns its matrices \(\boldsymbol{A}^{(0)}\) and \(\boldsymbol{B}^{(0)}\). We also swapped the roles of these matrices compared to the original RxInfer example, which follows the engineering convention of using \(A\) for state transitions. Here, as in the active inference literature, \(\boldsymbol{A}\) is the observation (likelihood) matrix and \(\boldsymbol{B}\) the state transition matrix. There is no control yet.

Terms used in this series

Roomba. A vacuum-cleaning robot that moves among the rooms of a small apartment.
Apartment. Three rooms: a bedroom, a living room, and a bathroom. There is no door between the bedroom and the living room.
Tick. One discrete time step of the simulation, \(t = 0, \ldots, T\).
Generative process. The environment that produces the true whereabouts and the observed surface.
Generative model. The agent’s internal model of that process, used for inference.
Hidden Markov Model. A Markov chain on a hidden state (here the whereabouts) that emits an observation (here the surface) at each tick.
Active inference. Inferring hidden states, and later choosing actions, by minimizing a free-energy functional.

Assumptions

The numbers in \(\boldsymbol{B}^{\ast(0)}\) and \(\boldsymbol{A}^{\ast(0)}\) are assumed, not measured: they invent a layout and a fallible camera so that a closed simulation exists. The Roomba is assumed always to start in the bedroom. Section 4.3.4 states the numbers; Appendix B records why they were chosen.

Structure of the code

This notebook uses the series protocol of van de Laar and de Vries (2019).

create_envir(...; seed) → (execute, observe, state)
- execute(û) carries out the chosen action for one tick. In this part the action is ignored: the Roomba moves on its own.
- observe() returns the surface observation of this tick, \(\hat{\mathbf{o}}^{(t)}\).
- state() returns the true whereabouts, \(\hat{\mathbf{s}}^{\ast(t)}\), for evaluation only.
create_agent(...) → (act, future, compute, slide)
- act() the chosen action. Not yet.
- future() predicted beliefs under a policy. Not yet.
- compute(ô) perception: one infer over the whole observation sequence, returning \(Q\big(s^{(0:T,0)}\big)\), \(Q\big(\boldsymbol{A}^{(0)}\big)\) and \(Q\big(\boldsymbol{B}^{(0)}\big)\).
- slide() move one tick forward. Not yet.

The environment draws all randomness from MersenneTwister(seed). In this part compute is called once on \(\hat{\mathbf{o}}^{(0:T)}\), not once per tick.

versioninfo() ## Julia version

Julia Version 1.11.6
Commit 9615af0f269 (2025-07-09 12:58 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
  LLVM: libLLVM-16.0.6 (ORCJIT, skylake)
Threads: 1 default, 0 interactive, 1 GC (on 12 virtual cores)

import Pkg
# Pkg.add(Pkg.PackageSpec(;name="RxInfer", version="3.6.0"))
Pkg.add(Pkg.PackageSpec(;name="RxInfer"))
Pkg.add(Pkg.PackageSpec(;name="Random"))
# Pkg.add(Pkg.PackageSpec(;name="Plots", version="1.40.1"))
Pkg.add(Pkg.PackageSpec(;name="Plots"))
Pkg.add(Pkg.PackageSpec(;name="XLSX"))
Pkg.add(Pkg.PackageSpec(;name="DataFrames"))
Pkg.add(Pkg.PackageSpec(;name="BenchmarkTools"))
Pkg.add(Pkg.PackageSpec(;name="Distributions"))
Pkg.add(Pkg.PackageSpec(;name="OffsetArrays"))

using RxInfer, Random, BenchmarkTools, Distributions, Plots, XLSX, DataFrames, OffsetArrays

   Resolving package versions...
  No Changes to `~/.julia/environments/v1.11/Project.toml`
  No Changes to `~/.julia/environments/v1.11/Manifest.toml`
   Resolving package versions...
  No Changes to `~/.julia/environments/v1.11/Project.toml`
  No Changes to `~/.julia/environments/v1.11/Manifest.toml`
   Resolving package versions...
  No Changes to `~/.julia/environments/v1.11/Project.toml`
  No Changes to `~/.julia/environments/v1.11/Manifest.toml`
   Resolving package versions...
  No Changes to `~/.julia/environments/v1.11/Project.toml`
  No Changes to `~/.julia/environments/v1.11/Manifest.toml`
   Resolving package versions...
  No Changes to `~/.julia/environments/v1.11/Project.toml`
  No Changes to `~/.julia/environments/v1.11/Manifest.toml`
   Resolving package versions...
  No Changes to `~/.julia/environments/v1.11/Project.toml`
  No Changes to `~/.julia/environments/v1.11/Manifest.toml`
   Resolving package versions...
  No Changes to `~/.julia/environments/v1.11/Project.toml`
  No Changes to `~/.julia/environments/v1.11/Manifest.toml`
   Resolving package versions...
  No Changes to `~/.julia/environments/v1.11/Project.toml`
  No Changes to `~/.julia/environments/v1.11/Manifest.toml`

Pkg.status()

Status `~/.julia/environments/v1.11/Project.toml`
  [6e4b80f9] BenchmarkTools v1.8.0
  [a93c6f00] DataFrames v1.8.2
  [31c24e10] Distributions v0.25.132
  [6fe1bfb0] OffsetArrays v1.17.0
  [91a5bcdd] Plots v1.41.7
  [86711068] RxInfer v6.1.0
  [fdbf4ff8] XLSX v0.13.2
  [9a3f8284] Random v1.11.0

## tuple0(x...): build a 0-based tuple (x^{(0)}, ..., x^{(N)}), mirroring the math.
##   - each argument becomes one entry: tuple0(v) wraps v; it never re-indexes v itself
##   - the index range 0:N is derived from the number of arguments
##   - entries may differ in shape, e.g. likelihood matrices for different modalities
## Returns an OffsetVector (mutable), not a Julia Tuple.
tuple0(x...) = OffsetVector(collect(x), 0:length(x)-1)

tuple0 (generic function with 1 method)

Why tuple0() rather than OffsetVector() directly?

It prevents a level mix-up. With OffsetVector, it is easy to offset the wrong level:

    OffsetVector([1.0, 0.0, 0.0], 0:2)    ## the one-hot vector itself, re-indexed 0:2 (wrong level)
    OffsetVector([[1.0, 0.0, 0.0]], 0:0)  ## a tuple containing the one-hot vector (intended)
    tuple0([1.0, 0.0, 0.0])               ## the same as the second line, with no way to get it wrong

The first line is valid code but means something else: the one-hot vector $\hat{\boldsymbol{s}}^{\ast(0,0)}$ itself, now indexed from 0, rather than the tuple $\hat{\mathbf{s}}^{\ast(0)} = \big(\hat{\boldsymbol{s}}^{\ast(0,0)}\big)$. Nothing would complain, and later code that indexes `[0]` would get a single number instead of a vector. `tuple0` always wraps its arguments as *entries*, so this mistake cannot happen.

The index range is computed automatically. With OffsetVector, the range is written by hand and has to agree with the number of entries: adding a second modality to LAˣ while forgetting to change 0:0 to 0:1 gives an error at best. tuple0 derives the range from the number of arguments, so adding or removing an entry is a one-line change.
It reads like the math. tuple0(A0, A1) looks like \(\big(\boldsymbol{A}^{(0)}, \boldsymbol{A}^{(1)}\big)\): entries separated by commas, nothing else. OffsetVector([A0, A1], 0:1) adds square brackets and a range that have nothing to do with the math.
It states the intent. OffsetVector says “an array with shifted indices”, which could be anything. tuple0 says “a tuple of factors or modalities, as in the math”, the specific concept from (10.33). And since all tuples are built the same way, a search for tuple0( finds every one of them.
There is one place to change things. Changing how tuples work, e.g. switching to an immutable 0-based type, or adding a check that every column of every \(\boldsymbol{A}\) and \(\boldsymbol{B}\) sums to 1, means changing one line instead of every place a tuple is built.
Entries of different shapes just work. collect builds a vector of whatever is passed in, so tuple0(A0, A1) with a \(3\times 3\) and a \(2\times 3\) matrix gives a vector of matrices, which is exactly what a tuple of likelihood matrices with different numbers of observations needs.

2 DATA UNDERSTANDING

The data are simulated. One tick is one step of the Roomba. A run covers \(T + 1 = 600\) ticks, \(t = 0, \ldots, T\) with \(T = 599\).

The recorded quantities are three kinds of one-hot encoded vectors:

Observations, which the agent receives: the observation tuple \(\hat{\mathbf{o}}^{(t)} = \big(\hat{\boldsymbol{o}}^{(t,0)}\big)\), with one modality, surface, whose values are “hardwood”, “carpet”, “tiles”.
Actions, which the agent would choose. There are none in this part: the Roomba is not steered.
True states, which the agent never sees: the state tuple \(\hat{\mathbf{s}}^{\ast(t)} = \big(\hat{\boldsymbol{s}}^{\ast(t,0)}\big)\), with one factor, whereabouts, whose values are “bedroom”, “living room”, “bathroom”.

In a later project the observations would come from a real camera. Here they are drawn from \(\boldsymbol{A}^{\ast(0)}\).

The agent works with the observations only, and with the actions it chose itself. The true states are kept for evaluation: comparing the inferred whereabouts with the room the Roomba was actually in, and comparing the learned \(\boldsymbol{A}^{(0)}\) and \(\boldsymbol{B}^{(0)}\) with \(\boldsymbol{A}^{\ast(0)}\) and \(\boldsymbol{B}^{\ast(0)}\).

3 DATA PREPARATION

The simulation writes the trajectories directly, indexed from 0:

Lôː: observations \(\hat{\mathbf{o}}^{(0:T)}\), \(t = 0, \ldots, T\)
Lŝˣː: true states \(\hat{\mathbf{s}}^{\ast(0:T)}\), for evaluation only
Lsː: the agent’s marginal beliefs, once inference has run (held inside the RxInfer result as sː)
Lûː: actions, \(t = 0, \ldots, T-1\). Not used in this part.

The single state factor is whereabouts (\(n = 0\)). The single observation modality is surface (\(m = 0\)). The 0-based helper tuple0 is defined above, before section 2. Plots in this part use the existing style; the step_x / step_y helpers of later notebooks are not needed yet.

4 MODELING

4.1 Narrative

The typeface shows what kind of object a symbol is, and an index produces the typeface of its result. They are shown for states \(s\); observations \(o\) follow the same pattern.

italic lower case: the name of a model variable. In this notebook: whereabouts (the state factor) and surface (the observation modality).
“double quotes”: a value of a variable, matching the label in the code: “bedroom”, “living room”, “bathroom”; “hardwood”, “carpet”, “tiles”.
bold: low-level headings and labels in lists and tables, such as Inference.

Italics and bold are not used for emphasis. Appendix A gives the motivation.

Indices. Round brackets mark a member of an ordered collection: time \(t\), factors \(n\), modalities \(m\). Square brackets mark a position along an array: an entry (every index fixed, plain) or a slice (an index free, still bold). Time comes first when both are present. A lone numeral on a random variable is time, \(s^{(0)}\); on a time-invariant array it is a member, \(\boldsymbol{D}^{(0)}\).

\(t\): time step, \(t = 0, \ldots, T\)
\(s^{(t,n)}\): categorical random variable for state factor \(n\) at time \(t\). Its values are labels, e.g. \(\{\text{bedroom}, \text{living room}, \text{bathroom}\}\). It is not bold, because a single categorical random variable has no structure.
\(\mathbf{s}^{(t)} = \big(s^{(t,0)}, \ldots, s^{(t,N)}\big)\): tuple of the random variables of all state factors at time \(t\)

Fonts. Bold marks anything with structure, and the style of bold tells you which kind:

upright bold (\(\mathbf{s}\), \(\mathbf{A}\)): a tuple, an ordered collection whose entries may differ in size, such as one entry per state factor or observation modality
italic bold (\(\boldsymbol{s}\), \(\boldsymbol{A}\)): a single vector or matrix, such as a one-hot vector, a probability vector, a belief, or a likelihood matrix

For example, \(\mathbf{A} = \big(\boldsymbol{A}^{(0)}, \ldots, \boldsymbol{A}^{(M)}\big)\) is the tuple of likelihood matrices, and \(\boldsymbol{A}^{(m)}\) is the matrix for observation modality \(m\). Likewise, \(\mathbf{B} = \big(\boldsymbol{B}^{(0)}, \ldots, \boldsymbol{B}^{(N)}\big)\) and \(\mathbf{D} = \big(\boldsymbol{D}^{(0)}, \ldots, \boldsymbol{D}^{(N)}\big)\) collect one array per state factor. Upright bold is only used for Latin letters, because Greek letters such as \(\pi\) do not render in upright bold.

Marks. A hat marks a realized value: something that has actually happened. The realized category is plain, \(\hat{s}^{\ast(t,n)}\) or \(\hat{o}^{(t,0)}\); its one-hot encoding is bold, \(\hat{\boldsymbol{s}}^{\ast(t,n)}\) or \(\hat{\boldsymbol{o}}^{(t,0)}\). In code, ô is the category and _ô is the one-hot vector. A bold symbol without a hat is a distribution: a probability vector over possible values.

\(\hat{\boldsymbol{s}}^{(t,n)}\): realized value of \(s^{(t,n)}\), one-hot encoded as a vector, e.g. \([1,0,0]^{\top}\)
\(\hat{\mathbf{s}}^{(t)} = \big(\hat{\boldsymbol{s}}^{(t,0)}, \ldots, \hat{\boldsymbol{s}}^{(t,N)}\big)\): tuple of the realized values of all state factors at time \(t\)
\(\hat{\mathbf{s}}^{(0:T)} = \big(\hat{\mathbf{s}}^{(0)}, \ldots, \hat{\mathbf{s}}^{(T)}\big)\): trajectory (sequence over time) of realized states
\(\boldsymbol{s}^{(t,n)}\), without a hat: the agent’s belief about state factor \(n\) at time \(t\), a probability vector with one entry per value of \(s^{(t,n)}\)
\(\boldsymbol{B}^{\ast(0)}\,\hat{\boldsymbol{s}}^{\ast(t-1,0)}\): a probability vector (here, over the next state), not a realized value

The agent never observes the states, so its state variables never carry a hat: it has random variables \(s^{(t,n)}\) and beliefs \(\boldsymbol{s}^{(t,n)}\). The observations are different, because the agent does receive them. Each observation therefore appears in three forms, which share the same index but are different objects:

\(o^{(t,0)}\): the categorical random variable for the surface at time \(t\), with values \(\{\text{hardwood}, \text{carpet}, \text{tiles}\}\)
\(\hat{\boldsymbol{o}}^{(t,0)}\): the observation the agent actually received, one-hot encoded. If the camera reads hardwood at time \(t\), then \(\hat{\boldsymbol{o}}^{(t,0)} = [1,0,0]^{\top}\), and the tuple of all modalities is \(\hat{\mathbf{o}}^{(t)} = \big([1,0,0]^{\top}\big)\).
\(\boldsymbol{o}^{(t,0)}\), without a hat: an expected observation, i.e. a probability vector the agent computes from its beliefs, \(\boldsymbol{o}^{(t,0)} = \boldsymbol{A}^{(0)}\,\boldsymbol{s}^{(t,0)}\)

For example, suppose the agent believes the Roomba is most likely in the bedroom, \(\boldsymbol{s}^{(t,0)} = [0.7, 0.2, 0.1]^{\top}\), and its likelihood matrix \(\boldsymbol{A}^{(0)}\) has the same values as \(\boldsymbol{A}^{\ast(0)}\). Then it expects

\[ \boldsymbol{o}^{(t,0)} = \boldsymbol{A}^{(0)}\,\boldsymbol{s}^{(t,0)} = \begin{bmatrix} 0.9 & 0.05 & 0.05 \\ 0.05 & 0.9 & 0.05 \\ 0.05 & 0.05 & 0.9 \end{bmatrix} \begin{bmatrix} 0.7 \\ 0.2 \\ 0.1 \end{bmatrix} = \begin{bmatrix} 0.645 \\ 0.22 \\ 0.135 \end{bmatrix} \]

that is, hardwood with probability 0.645, carpet 0.22, and tiles 0.135. If the camera then reads carpet, the received observation is \(\hat{\boldsymbol{o}}^{(t,0)} = [0,1,0]^{\top}\): a single definite outcome, which the agent uses to update its belief \(\boldsymbol{s}^{(t,0)}\) toward the living room.

So, for observations the agent has both hatted and unhatted forms: what it expects (\(\boldsymbol{o}\)) and what it gets (\(\hat{\boldsymbol{o}}\)). For states, it only ever has the unhatted forms: the random variable \(s^{(t,n)}\) and the belief \(\boldsymbol{s}^{(t,n)}\). Only the environment has \(\hat{\boldsymbol{s}}^{\ast(t,n)}\), the room the Roomba is really in.

The asterisk \({}^{\ast}\) marks what belongs to the environment (the generative process), e.g. \(\hat{\mathbf{s}}^{\ast(t)}\) and \(\boldsymbol{B}^{\ast(0)}\). Symbols without it belong to the agent (the generative model), e.g. its own matrices \(\boldsymbol{A}^{(m)}\) and \(\boldsymbol{B}^{(n)}\), which it learns so that they approximate \(\boldsymbol{A}^{\ast(m)}\) and \(\boldsymbol{B}^{\ast(n)}\). The observation \(\hat{\mathbf{o}}^{(t)}\) has no asterisk because it is shared: the environment generates it and the agent receives it.

Code symbols. The code mirrors the math as closely as identifiers allow:

L prefix: upright bold (\(\mathbf{s}\), \(\mathbf{A}\)), a tuple. The L suggests verticality (upright).
_ prefix: italic bold (\(\boldsymbol{s}\), \(\boldsymbol{A}\)), a single vector or matrix that is not part of a tuple. The _ suggests the underline of bold. A leading _ if and only if the math symbol is bold.
hat: a Unicode circumflex on the letter, e.g. ŝ for \(\hat{s}\), ô for \(\hat{o}\); with _ when the math is bold (_ô for \(\hat{\boldsymbol{o}}\))
asterisk \({}^{\ast}\): a superscript ˣ (typed \^x then Tab), before the time index in the math: \(\hat{\boldsymbol{s}}^{\ast(t)}\)
time index \((t)\): a subscript in the name: ₜ for \(t\), ₜ₋₁ for \(t-1\), ₀ for \(t=0\)
trajectory \((0{:}T)\): a ː suffix (triangular colon), e.g. Lôː for \(\hat{\mathbf{o}}^{(0:T)}\)
member index: ı…ı (dotless i, typed \imath then Tab), e.g. _Aˣı0ı for \(\boldsymbol{A}^{\ast(0)}\). This notebook keeps the Julia tuple index LAˣ[0] for the same object; _Aˣı0ı is the named-member equivalent.
array position: I…I, e.g. GIpτI for the entry \(G^{[p,\tau]}\) (not used in this part)
tuple index in Julia: real 0-based indexing, e.g. LAˣ[0]. Tuples are built with tuple0(...) (an OffsetVector indexed from 0).

Indexing and the tear-off mnemonic. Only tuples have names; their entries are written by indexing. Read a factor or modality index ([n] or [m]) as tearing the vertical line off the L, turning it into _: LAˣ[0] reads as _Aˣ[0], i.e. italic bold \(\boldsymbol{A}^{\ast(0)}\), just as indexing into a tuple turns the upright \(\mathbf{A}^{\ast}\) into the italic \(\boldsymbol{A}^{\ast(0)}\) in the math. Three details:

A time index into a trajectory does not tear off the line: Lŝˣː[t] is still a tuple, \(\hat{\mathbf{s}}^{\ast(t)}\). Only the factor index that follows tears it off: Lŝˣː[t][0] is \(\hat{\boldsymbol{s}}^{\ast(t,0)}\).
The mnemonic applies wherever the brackets appear, on either side of an assignment: LAˣ[0] = … assigns to \(\boldsymbol{A}^{\ast(0)}\).
Each further index changes the font of the result as it would in the math: LAˣ[0][1,2] is a single number, an entry of the matrix \(\boldsymbol{A}^{\ast(0)}\), so it is no longer bold at all.

Since each tuple has exactly one name, there is nothing to keep in sync: when a tuple gets a new value, a single assignment updates it, e.g. Lŝˣₜ = tuple0(...).

Math	Code	Meaning
\(\mathbf{A}^{\ast}\), \(\mathbf{B}^{\ast}\), \(\mathbf{D}^{\ast}\)	`LAˣ`, `LBˣ`, `LDˣ`	environment’s tuples of arrays
\(\boldsymbol{A}^{\ast(0)}\), \(\boldsymbol{B}^{\ast(0)}\), \(\boldsymbol{D}^{\ast(0)}\)	`LAˣ[0]`, `LBˣ[0]`, `LDˣ[0]`	environment’s arrays for surface / whereabouts 0; named-member form `_Aˣı0ı`, `_Bˣı0ı`, `_Dˣı0ı`
\(\mathbf{A}\), \(\mathbf{B}\), \(\mathbf{D}\)	`LA`, `LB`, `LD`	agent’s tuples of arrays
\(\boldsymbol{A}^{(0)}\), \(\boldsymbol{B}^{(0)}\), \(\boldsymbol{D}^{(0)}\)	`LA[0]`, `LB[0]`, `LD[0]`	agent’s arrays for surface / whereabouts 0
\(\hat{\mathbf{s}}^{\ast(0)}\)	`Lŝˣ₀`	true initial state tuple
\(\hat{\mathbf{s}}^{\ast(t)}\), \(\hat{\mathbf{s}}^{\ast(t-1)}\)	`Lŝˣₜ`, `Lŝˣₜ₋₁`	true state tuple at \(t\), \(t-1\)
\(\hat{s}^{\ast(t,0)}\), \(\hat{\boldsymbol{s}}^{\ast(t,0)}\)	`ŝ`, `Lŝˣₜ[0]`	true whereabouts at \(t\): category, and one-hot
\(\hat{\mathbf{s}}^{\ast(0:T)}\)	`Lŝˣː`	trajectory of true states
\(\hat{\mathbf{s}}^{\ast(t)}\), \(\hat{\boldsymbol{s}}^{\ast(t,0)}\) (from the trajectory)	`Lŝˣː[t]`, `Lŝˣː[t][0]`	entry \(t\) of the trajectory, and its Room factor
\(\hat{\mathbf{o}}^{(t)}\)	`Lôₜ`	received observation tuple at \(t\)
\(\hat{\boldsymbol{o}}^{(t,0)}\)	`Lôₜ[0]`	received surface at \(t\) (one-hot)
\(\hat{\mathbf{o}}^{(0:T)}\)	`Lôː`	trajectory of received observations
\(\boldsymbol{B}^{\ast(0)}\,\hat{\boldsymbol{s}}^{\ast(t-1,0)}\)	`LBˣ[0] * Lŝˣₜ₋₁[0]`	probability vector over the next room
\(\boldsymbol{s}^{(t,n)}\)	`Lsₜ[n]`	agent’s belief about factor \(n\) at \(t\) (probability vector)
\(t\), \(T\)	`t`, `T`	time step, last time step

The code stores the agent’s beliefs, not its random variables. So the tuple Lsₜ holds the beliefs \(\big(\boldsymbol{s}^{(t,0)}, \ldots, \boldsymbol{s}^{(t,N)}\big)\), while in the math, \(\mathbf{s}^{(t)}\) denotes the tuple of random variables.

What is new in this part. The original RxInfer Hidden Markov Model example is wrapped in create_envir / create_agent. The agent tracks the whereabouts and learns \(\boldsymbol{A}^{(0)}\) and \(\boldsymbol{B}^{(0)}\). There is no control yet: act, future and slide are left for later parts.

Appendix B records the modeling choices and assumed numbers, so that they are not undone by accident.

In Python, ˣ is normalized to a plain x, so LAˣ and LAx are the same name; never use LAx for anything else.

4.2 Core Elements

Metrics.

tracking: the fraction of ticks at which the most probable whereabouts under \(Q\big(s^{(t,0)}\big)\) equals the true whereabouts \(\hat{s}^{\ast(t,0)}\). Higher is better. Symbol \(J_{\text{track}}\).
layout: how close the posterior mean of \(\boldsymbol{B}^{(0)}\) is to \(\boldsymbol{B}^{\ast(0)}\) (which rooms connect).
camera: how close the posterior mean of \(\boldsymbol{A}^{(0)}\) is to \(\boldsymbol{A}^{\ast(0)}\) (how often the camera misreads).

There is no baseline comparison in this part: the agent is not choosing actions, and a single seed is used to follow the beliefs in detail. A good agent should recover the apartment layout and the camera’s error rates, and keep the inferred whereabouts close to the true one.

Decisions. There are no control decisions yet. The Roomba is not steered; we only track it.

Sources of uncertainty.

In the environment: state transitions are stochastic, captured by \(\mathbf{B}^{\ast}\) with member \(\boldsymbol{B}^{\ast(n)}\) for state factor \(n\). The camera sometimes misidentifies the floor, captured by \(\mathbf{A}^{\ast}\) with member \(\boldsymbol{A}^{\ast(m)}\) for observation modality \(m\). The initial state adds no uncertainty: the Roomba always starts in the bedroom.
On the agent’s side: the agent never knows the true whereabouts, only its belief \(\boldsymbol{s}^{(t,0)}\). It does not know \(\mathbf{A}^{\ast}\) and \(\mathbf{B}^{\ast}\); it learns its own \(\mathbf{A}\) and \(\mathbf{B}\) from the observations. It also does not know the starting room, so it begins with a uniform prior over the three rooms.

4.3 System-Under-Steer / Environment / Generative Process

First, before tracking the Roomba’s movements using RxInfer, we describe the environment it moves in, i.e. the generative process that produces the data. Since we have a discrete set of rooms in the apartment, we can use a categorical distribution to represent the Roomba’s position. There are three rooms in the apartment, meaning we need three states in our categorical distribution. At time \(t\), the Roomba’s true position is the realized state tuple \[ \hat{\mathbf{s}}^{\ast(t)} = \big(\hat{\boldsymbol{s}}^{\ast(t,0)}\big) \] whose entry \(\hat{\boldsymbol{s}}^{\ast(t,0)}\) is the one-hot encoded value of a single state factor, factor 0 (whereabouts), with possible values

\(s^{\ast(t,0)} \in \{\text{bedroom}, \text{living room}, \text{bathroom}\}\)

The Roomba’s starting room is given by the initial state distribution \(\mathbf{D}^{\ast} = (\boldsymbol{D}^{\ast(0)})\). In our simulation, the Roomba always starts in the bedroom, so \(\boldsymbol{D}^{\ast(0)} = [1,0,0]^{\top}\). Some rooms are more accessible than others, meaning the Roomba is more likely to move between these rooms - for example, there is no door directly between the bedroom and the living room, so the Roomba has to pass through the bathroom. The true transition probabilities are captured by \(\mathbf{B}^{\ast} = (\boldsymbol{B}^{\ast(0)})\), with transition matrix \(\boldsymbol{B}^{\ast(0)} \in [0,1]^{3\times 3}\) for this single state factor.

Our Roomba is equipped with a small camera that tracks the surface it is moving over. This camera provides the observations, since there are

hardwood floors in the bedroom
carpet in the living room, and
tiles in the bathroom.

However, this method is not foolproof, and sometimes the Roomba will mistake one floor type for another. Don’t be too hard on the little guy, it’s just a Roomba after all.

At time \(t\), the received observation tuple is \(\hat{\mathbf{o}}^{(t)} = \big(\hat{\boldsymbol{o}}^{(t,0)}\big)\), whose entry \(\hat{\boldsymbol{o}}^{(t,0)}\) is the one-hot encoded value of a single observation modality, modality 0 (surface), with possible values

\(o^{(t,0)} \in \{\text{hardwood}, \text{carpet}, \text{tiles}\}\)

The true mapping from the Roomba’s position to the observations is \(\mathbf{A}^{\ast} = (\boldsymbol{A}^{\ast(0)})\), with likelihood matrix \(\boldsymbol{A}^{\ast(0)} \in [0,1]^{3\times 3}\) for this single observation modality. \(\boldsymbol{A}^{\ast(0)}\) also encodes the probability that the Roomba will make a mistake and get the wrong observation. This leaves us with the following generative process specification:

\[ \mathcal{E} \;\triangleq\; \begin{cases} \hat{\boldsymbol{s}}^{\ast(0,0)} \sim \mathcal{Cat}\big(\boldsymbol{D}^{\ast(0)}\big) & \text{INITIAL STATE}\\[2pt] \hat{\boldsymbol{s}}^{\ast(t,0)} \sim \mathcal{Cat}\big(\boldsymbol{B}^{\ast(0)}\,\hat{\boldsymbol{s}}^{\ast(t-1,0)}\big), \quad t>0 & \text{STATE TRANSITION}\\[2pt] \hat{\boldsymbol{o}}^{(t,0)} \sim \mathcal{Cat}\big(\boldsymbol{A}^{\ast(0)}\,\hat{\boldsymbol{s}}^{\ast(t,0)}\big) & \text{OBSERVATION GENERATION} \end{cases} \]

Because \(\hat{\boldsymbol{s}}^{\ast(t-1,0)}\) and \(\hat{\boldsymbol{s}}^{\ast(t,0)}\) are one-hot, the products \(\boldsymbol{B}^{\ast(0)}\,\hat{\boldsymbol{s}}^{\ast(t-1,0)}\) and \(\boldsymbol{A}^{\ast(0)}\,\hat{\boldsymbol{s}}^{\ast(t,0)}\) simply pick out the column for the current room: a probability vector over next rooms and over textures, respectively.

This process has the structure of a Hidden Markov Model or HMM for short: the room is a hidden Markov chain, and the camera emits noisy observations of it. The agent never sees \(\boldsymbol{A}^{\ast(0)}\) or \(\boldsymbol{B}^{\ast(0)}\) directly. Its goal is to learn its own matrices \(\boldsymbol{A}^{(0)}\) and \(\boldsymbol{B}^{(0)}\) from the observations, so that they approximate the true ones and it can track where the Roomba is.

Lŝˣ₀ = tuple0([1.0, 0.0, 0.0])    ## ŝ^{*(0)} = (ŝ^{*(0,0)}): initial state, Roomba starts in the bedroom

1-element OffsetArray(::Vector{Vector{Float64}}, 0:0) with eltype Vector{Float64} with indices 0:0:
 [1.0, 0.0, 0.0]

4.3.1 State and Observation variables

The state variables represent what we need to know. The true state of the Roomba at time \(t\) is given by a (compound) tuple \(\hat{\mathbf{s}}^{\ast(t)}\) whose components are the one-hot encoded state factors (here a single factor):

\[ \hat{\mathbf{s}}^{\ast(t)} = \big(\hat{\boldsymbol{s}}^{\ast(t,0)}\big) \]

where the state factors are:

factor 0 (whereabouts):
- \(s^{\ast(t,0)} \in \{\text{bedroom}, \text{living room}, \text{bathroom}\}\),
- one-hot encoded as \(\hat{\boldsymbol{s}}^{\ast(t,0)} \in \big\{[1,0,0]^{\top}, [0,1,0]^{\top}, [0,0,1]^{\top}\big\}\)

The observation variables represent what we can sense. Similarly, the (compound) observation tuple received by the Roomba at time \(t\) is given by

\[ \hat{\mathbf{o}}^{(t)} = \big(\hat{\boldsymbol{o}}^{(t,0)}\big) \]

where the observation modalities are:

modality 0 (surface):
- \(o^{(t,0)} \in \{\text{hardwood}, \text{carpet}, \text{tiles}\}\),
- one-hot encoded as \(\hat{\boldsymbol{o}}^{(t,0)} \in \big\{[1,0,0]^{\top}, [0,1,0]^{\top}, [0,0,1]^{\top}\big\}\)

4.3.2 Decision variables

There are no decisions to be made (yet). We are simply observing the behavior of the Roomba, i.e. there is no control/steering applied to the environment.

4.3.3 Exogenous forces

There are no exogenous force variables (yet). We may add some in followup parts, for example, having certain doors closed at random.

4.3.4 State Transition and Observation Generation functions

Starting from \(\hat{\boldsymbol{s}}^{\ast(0,0)} \sim \mathcal{Cat}\big(\boldsymbol{D}^{\ast(0)}\big)\), with \(\boldsymbol{D}^{\ast(0)} = [1,0,0]^{\top}\) (the bedroom), the environment evolves according to

\[ \begin{aligned} \hat{\boldsymbol{s}}^{\ast(t,0)} &\sim \mathcal{Cat}\big(\boldsymbol{B}^{\ast(0)}\,\hat{\boldsymbol{s}}^{\ast(t-1,0)}\big), \quad t>0 && \text{STATE TRANSITION}\\ \hat{\boldsymbol{o}}^{(t,0)} &\sim \mathcal{Cat}\big(\boldsymbol{A}^{\ast(0)}\,\hat{\boldsymbol{s}}^{\ast(t,0)}\big) && \text{OBSERVATION GENERATION} \end{aligned} \]

State Transition model \(\boldsymbol{B}^{\ast(0)}\) (whereabouts)

Each column is a probability distribution over the next room, given the previous room, so each column sums to 1.

next room \(s^{\ast(t,0)}\) previous room \(s^{\ast(t-1,0)}\)	bedroom	living room	bathroom
bedroom	0.9	0.0	0.1
living room	0.0	0.9	0.1
bathroom	0.1	0.1	0.8

The zeros mean there is no door between the bedroom and the living room: the Roomba has to pass through the bathroom.

Observation Generation model \(\boldsymbol{A}^{\ast(0)}\) (surface)

Each column is a probability distribution over the observed texture, given the room the Roomba is in, so each column sums to 1.

observation \(o^{(t,0)}\) room \(s^{\ast(t,0)}\)	bedroom	living room	bathroom
hardwood	0.9	0.05	0.05
carpet	0.05	0.9	0.05
tiles	0.05	0.05	0.9

The off-diagonal entries are the camera’s mistakes: in each room, it misreads the floor with probability 0.1, split equally between the two other textures.

LBˣ = tuple0(
    [0.9  0.0  0.1;    ## rows: next room (bedroom, living room, bathroom)
     0.0  0.9  0.1;    ## columns: previous room (bedroom, living room, bathroom)
     0.1  0.1  0.8],
)
LBˣ[0]    ## the Room transition matrix, B^{*(0)}

3×3 Matrix{Float64}:
 0.9  0.0  0.1
 0.0  0.9  0.1
 0.1  0.1  0.8

Observation Generation model \(\boldsymbol{A}^{\ast(0)}\) (surface)

Each column is a probability distribution over the observed texture, given the room the Roomba is in, so each column sums to 1.

observation \(o^{(t,0)}\) room \(s^{\ast(t,0)}\)	bedroom	living room	bathroom
hardwood	0.9	0.05	0.05
carpet	0.05	0.9	0.05
tiles	0.05	0.05	0.9

The off-diagonal entries are the camera’s mistakes: in each room, it misreads the floor with probability 0.1, split equally between the two other textures.

LAˣ = tuple0(
    [0.9  0.05 0.05;    ## rows: hardwood, carpet, tiles
     0.05 0.9  0.05;    ## columns: bedroom, living room, bathroom
     0.05 0.05 0.9],
)
LAˣ[0]    ## the surface likelihood matrix, A^{*(0)}

3×3 Matrix{Float64}:
 0.9   0.05  0.05
 0.05  0.9   0.05
 0.05  0.05  0.9

4.3.5 Objective function

The designer judges tracking with knowledge of the true whereabouts:

\[ J_{\text{track}} = \frac{1}{T+1}\sum_{t=0}^{T} \mathbb{1}\big[\arg\max_i\, Q\big(s^{(t,0)}=i\big) = \hat{s}^{\ast(t,0)}\big] \]

Higher is better. The layout and camera checks in 4.2 are secondary: they compare posterior means with \(\boldsymbol{B}^{\ast(0)}\) and \(\boldsymbol{A}^{\ast(0)}\).

Inference itself minimizes the Bethe free energy (BFE), or the generalized free energy (GFE), which RxInfer handles.

4.3.6 Implementation of the System-Under-Steer / Environment / Generative Process

## returns a one-hot encoding of a random sample from a categorical distribution
function rand_1hot_vec(rng, distribution::Categorical)
    K = ncategories(distribution)
    sample = zeros(K)
    drawn_category = rand(rng, distribution)
    sample[drawn_category] = 1.0
    return sample
end

rand_1hot_vec (generic function with 1 method)

function create_envir(; LAˣ, LBˣ, Lŝˣ₀, seed=42)
    rng = MersenneTwister(seed)

    ## t = 0
    Lŝˣₜ = Lŝˣ₀                                              ## ŝ^{*(0)}
    Loₜ  = tuple0(LAˣ[0] * Lŝˣₜ[0])                          ## A^{*(0)} ŝ^{*(0,0)}
    Lôₜ  = tuple0(rand_1hot_vec(rng, Categorical(Loₜ[0])))   ## ô^{(0,0)} ~ Cat(...)

    execute = () -> begin
        Lŝˣₜ₋₁ = Lŝˣₜ
        Lsˣₜ = tuple0(LBˣ[0] * Lŝˣₜ₋₁[0])                         ## B^{*(0)} ŝ^{*(t-1,0)}
        Lŝˣₜ = tuple0(rand_1hot_vec(rng, Categorical(Lsˣₜ[0])))   ## ŝ^{*(t,0)} ~ Cat(...)
        Loₜ  = tuple0(LAˣ[0] * Lŝˣₜ[0])                          ## A^{*(0)} ŝ^{*(t,0)}
        Lôₜ  = tuple0(rand_1hot_vec(rng, Categorical(Loₜ[0])))   ## ô^{(t,0)} ~ Cat(...)
        nothing
    end

    observe = () -> Lôₜ     ## ô^{(t)}
    state   = () -> Lŝˣₜ    ## ŝ^{*(t)}, true state (for evaluation only)

    return (execute, observe, state)
end

create_envir (generic function with 1 method)

In order to generate data to mimic the observations of the Roomba, we need to specify three things: the initial state (i.e., in which room the Roomba starts), the actual transition probabilities between the states \(\boldsymbol{B}^{\ast(0)}\) (i.e., how likely the Roomba is to move from one room to another), and the observation distribution \(\boldsymbol{A}^{\ast(0)}\) (i.e., what type of texture the Roomba will observe in each room). We can then use these specifications to generate observations from the generative process, which has the structure of a hidden Markov model (HMM).

To generate our observation data, we’ll follow these steps: 1. Assume an initial state for the Roomba, \(\hat{\boldsymbol{s}}^{\ast(0,0)}\). For example, we can start the Roomba in the bedroom. 2. Determine the observation in the starting room by drawing from the categorical distribution \(\boldsymbol{A}^{\ast(0)}\,\hat{\boldsymbol{s}}^{\ast(0,0)}\), i.e. the column of \(\boldsymbol{A}^{\ast(0)}\) for that room. 3. Determine where the Roomba goes next by drawing from the categorical distribution \(\boldsymbol{B}^{\ast(0)}\,\hat{\boldsymbol{s}}^{\ast(t-1,0)}\), i.e. the column of \(\boldsymbol{B}^{\ast(0)}\) for the current room. 4. Determine the observation in the new room by drawing from \(\boldsymbol{A}^{\ast(0)}\,\hat{\boldsymbol{s}}^{\ast(t,0)}\). 5. Repeat steps 3-4 for as many time steps as we want.

We will generate 600 data points (\(t = 0, \ldots, T\) with \(T = 599\)) to simulate 600 ticks of the Roomba moving through the apartment. Lôː will contain the Roomba’s observations of the floor it’s on at each tick, \(\hat{\mathbf{o}}^{(0:T)}\), and Lŝˣː will contain the room the Roomba was actually in, \(\hat{\mathbf{s}}^{\ast(0:T)}\).

The following code implements this process and generates our observation data:

T = 599                           ## 600 ticks: t = 0, ..., T
Lŝˣ₀ = tuple0([1.0, 0.0, 0.0])    ## ŝ^{*(0)}: Roomba starts in the bedroom
(execute_sim, observe_sim, state_sim) = create_envir(; LAˣ=LAˣ, LBˣ=LBˣ, Lŝˣ₀=Lŝˣ₀)

Lŝˣː = OffsetVector(Vector{Any}(undef, T+1), 0:T)    ## ŝ^{*(0:T)}: true states
Lôː  = OffsetVector(Vector{Any}(undef, T+1), 0:T)    ## ô^{(0:T)}: observations
for t = 0:T
    t > 0 && execute_sim()     ## move and observe (t = 0 is already set up)
    Lŝˣː[t] = state_sim()
    Lôː[t]  = observe_sim()
end

[Lŝˣ[0] for Lŝˣ in Lŝˣː]    ## ŝ^{*(t,0)} for t = 0, ..., T: the Room factor

600-element OffsetArray(::Vector{Vector{Float64}}, 0:599) with eltype Vector{Float64} with indices 0:599:
 [1.0, 0.0, 0.0]
 [1.0, 0.0, 0.0]
 [1.0, 0.0, 0.0]
 [1.0, 0.0, 0.0]
 [1.0, 0.0, 0.0]
 [1.0, 0.0, 0.0]
 [1.0, 0.0, 0.0]
 [1.0, 0.0, 0.0]
 [1.0, 0.0, 0.0]
 [1.0, 0.0, 0.0]
 ⋮
 [0.0, 1.0, 0.0]
 [0.0, 1.0, 0.0]
 [0.0, 1.0, 0.0]
 [0.0, 1.0, 0.0]
 [0.0, 0.0, 1.0]
 [0.0, 0.0, 1.0]
 [0.0, 0.0, 1.0]
 [0.0, 0.0, 1.0]
 [0.0, 0.0, 1.0]

[Lô[0] for Lô in Lôː]    ## ô^{(t,0)} for t = 0, ..., T: the surface

600-element OffsetArray(::Vector{Vector{Float64}}, 0:599) with eltype Vector{Float64} with indices 0:599:
 [1.0, 0.0, 0.0]
 [1.0, 0.0, 0.0]
 [1.0, 0.0, 0.0]
 [1.0, 0.0, 0.0]
 [1.0, 0.0, 0.0]
 [1.0, 0.0, 0.0]
 [1.0, 0.0, 0.0]
 [1.0, 0.0, 0.0]
 [1.0, 0.0, 0.0]
 [1.0, 0.0, 0.0]
 ⋮
 [0.0, 1.0, 0.0]
 [0.0, 1.0, 0.0]
 [0.0, 1.0, 0.0]
 [0.0, 0.0, 1.0]
 [0.0, 0.0, 1.0]
 [0.0, 0.0, 1.0]
 [0.0, 0.0, 1.0]
 [0.0, 0.0, 1.0]
 [1.0, 0.0, 0.0]

rooms = [argmax(Lŝˣ[0]) for Lŝˣ in parent(Lŝˣː)]    ## room index (1, 2, 3) of ŝ^{*(t,0)} for t = 0, ..., T
plot(title="True States of the Simulated Roomba")
plot!(
    0:T, rooms, linetype=:steppost, linestyle=:solid, leg=:right,
    label="true room", xlabel="Time t", yticks=(1:3, ["Bedroom", "Living room", "Bathroom"]))

textures = [argmax(Lô[0]) for Lô in parent(Lôː)]    ## texture index (1, 2, 3) of ô^{(t,0)} for t = 0, ..., T
plot(title="Observations of the Simulated Roomba")
plot!(
    0:T, textures, linetype=:steppost, linestyle=:solid, leg=:right,
    label="observed texture", xlabel="Time t", yticks=(1:3, ["Hardwood", "Carpet", "Tiles"]))

plot(title="True whereabouts and observed surface")
plot!(0:T, rooms, linetype=:steppost, label="true room",
      yticks=(1:3, ["Bedroom / Hardwood", "Living room / Carpet", "Bathroom / Tiles"]))
plot!(0:T, textures, linetype=:steppost, linestyle=:dash, label="observed texture",
      xlabel="Time t", leg=:right)

mistakes = count(rooms .!= textures)
println("camera mistakes: $mistakes of $(T+1) ticks ($(round(100*mistakes/(T+1), digits=1))%)")

camera mistakes: 57 of 600 ticks (9.5%)

4.4 Uncertainty Model

The uncertainty in the environment is captured in the state transition matrix \(\boldsymbol{B}^{\ast(0)}\) (factor 0, Room) and the observation generation matrix \(\boldsymbol{A}^{\ast(0)}\) (modality 0, Texture). Their entries are the probabilities of the Roomba’s moves and of the camera’s readings. The initial state adds no uncertainty, because the Roomba always starts in the bedroom.

The agent’s uncertainty about \(\mathbf{A}^{\ast}\) and \(\mathbf{B}^{\ast}\) themselves is not part of the environment. It is captured later, in the agent’s model, by priors over its own tuples of matrices \(\mathbf{A}\) and \(\mathbf{B}\).

4.5 Agent / Generative Model

4.5.1 State and Observation variables

On the agent’s side, the state of the environment at time \(t\) is not known: the agent represents it by a (compound) tuple of categorical random variables, one per state factor (here a single factor):

\[ \mathbf{s}^{(t)} = \big(s^{(t,0)}\big) \]

where the state factors are:

factor 0 (whereabouts):
- \(s^{(t,0)} \in \{\text{bedroom}, \text{living room}, \text{bathroom}\}\)

The agent never observes \(s^{(t,0)}\) directly. What it holds instead is a belief, the probability vector \(\boldsymbol{s}^{(t,0)} \in [0,1]^{3}\), with one probability per room.

The initial state prior is given by

\[ P_D\big(s^{(0,0)}\big) = \mathcal{Cat}\big(s^{(0,0)};\, \boldsymbol{D}^{(0)}\big) \]

where \(\boldsymbol{D}^{(0)} \in [0,1]^{3}\) parameterizes the categorical distribution of \(s^{(0,0)}\). Unlike the environment, the agent does not know that the Roomba starts in the bedroom, so it uses the uniform prior \(\boldsymbol{D}^{(0)} = [\tfrac{1}{3}, \tfrac{1}{3}, \tfrac{1}{3}]^{\top}\).

The (compound) observation tuple received by the agent at time \(t\) is the same one the environment generates:

\[ \hat{\mathbf{o}}^{(t)} = \big(\hat{\boldsymbol{o}}^{(t,0)}\big) \]

where the observation modalities are:

modality 0 (surface):
- \(o^{(t,0)} \in \{\text{hardwood}, \text{carpet}, \text{tiles}\}\),
- one-hot encoded as \(\hat{\boldsymbol{o}}^{(t,0)} \in \big\{[1,0,0]^{\top}, [0,1,0]^{\top}, [0,0,1]^{\top}\big\}\)

4.5.2 Decision variables

There will be no decision variables (yet).

4.5.3 Preference variables / Setpoints

There will be no preference variables (yet).

4.5.4 State Transition and Observation Generation functions

The state transition function is given by

\[ P\big(s^{(t,0)} \mid s^{(t-1,0)}, \boldsymbol{B}^{(0)}\big) = \mathcal{Cat}\big(s^{(t,0)};\, \boldsymbol{B}^{(0)}[\bullet, s^{(t-1,0)}]\big) \]

where \(\boldsymbol{B}^{(0)}[\bullet, s^{(t-1,0)}]\), the column of \(\boldsymbol{B}^{(0)}\) for the previous room, parameterizes the categorical distribution of \(s^{(t,0)}\).

The observation generation function is given by

\[ P\big(o^{(t,0)} \mid s^{(t,0)}, \boldsymbol{A}^{(0)}\big) = \mathcal{Cat}\big(o^{(t,0)};\, \boldsymbol{A}^{(0)}[\bullet, s^{(t,0)}]\big) \]

where \(\boldsymbol{A}^{(0)}[\bullet, s^{(t,0)}]\), the column of \(\boldsymbol{A}^{(0)}\) for the current room, parameterizes the categorical distribution of \(o^{(t,0)}\).

An entry of \(\boldsymbol{A}^{(0)}\) is the probability of a specific observation given a specific state, and each column of \(\boldsymbol{A}^{(0)}\) is a categorical distribution over observations. The same holds for \(\boldsymbol{B}^{(0)}\), with the next state in place of the observation. If the state were known, its column could be selected by multiplying with the one-hot encoding of that state. The agent, however, does not know the state; it only has its belief \(\boldsymbol{s}^{(t,0)}\). Multiplying with the belief gives a weighted average of the columns:

\(\boldsymbol{B}^{(0)}\,\boldsymbol{s}^{(t-1,0)}\) is the agent’s predicted belief about the next room,
\(\boldsymbol{A}^{(0)}\,\boldsymbol{s}^{(t,0)}\) is the agent’s expected observation.

Unlike the environment’s \(\boldsymbol{A}^{\ast(0)}\) and \(\boldsymbol{B}^{\ast(0)}\), the agent’s matrices \(\boldsymbol{A}^{(0)}\) and \(\boldsymbol{B}^{(0)}\) are not known in advance: they are learned from the observations. This is why they appear as conditions in the two distributions above.

4.5.5 Implementation of the Agent / Generative Model / Internal Model

We start by specifying a probabilistic model for the agent that describes the agent’s internal beliefs about the external dynamics of the environment. The generative model is defined as follows:

\[ \begin{aligned} P\big(o^{(0:T,0)}, s^{(0:T,0)}, \boldsymbol{A}^{(0)}, \boldsymbol{B}^{(0)}\big) &= \underbrace{P_{\check{a}}\big(\boldsymbol{A}^{(0)}\big)\, P_{\check{b}}\big(\boldsymbol{B}^{(0)}\big)}_{\text{parameter prior}} \,\underbrace{P_D\big(s^{(0,0)}\big)}_{\substack{\text{state}\\ \text{prior}}} \prod_{t=0}^{T} \underbrace{P_A\big(o^{(t,0)} \mid s^{(t,0)}\big)}_{\substack{\text{observation-generation}\\ \text{model}}} \prod_{t=1}^{T} \underbrace{P_B\big(s^{(t,0)} \mid s^{(t-1,0)}\big)}_{\substack{\text{state-transition}\\ \text{model}}} \\[6pt] &= P_{\check{a}}\big(\boldsymbol{A}^{(0)}\big)\, P_{\check{b}}\big(\boldsymbol{B}^{(0)}\big) \, P_D\big(s^{(0,0)}\big) \prod_{t=0}^{T} P\big(o^{(t,0)} \mid s^{(t,0)}, \boldsymbol{A}^{(0)}\big) \prod_{t=1}^{T} P\big(s^{(t,0)} \mid s^{(t-1,0)}, \boldsymbol{B}^{(0)}\big) \end{aligned} \]

where \(\check{\boldsymbol{a}}^{(0)}\) and \(\check{\boldsymbol{b}}^{(0)}\) are the concentration parameters of the DirichletCollection priors on \(\boldsymbol{A}^{(0)}\) and \(\boldsymbol{B}^{(0)}\), and \(\boldsymbol{D}^{(0)}\) is the agent’s uniform initial state prior from 4.5.1.

The two lines say the same thing in two notations. In the first, the subscripts \(A\) and \(B\) mark which matrix parameterizes each distribution. In the second, \(\boldsymbol{A}^{(0)}\) and \(\boldsymbol{B}^{(0)}\) appear as explicit conditions: because the agent learns them, they are random variables with their own priors, not fixed numbers. \(\boldsymbol{D}^{(0)}\) is not learned, so it remains a subscript in both lines.

Now it is time to build our model. As mentioned earlier, we will use categorical distributions for the states and observations. To learn the agent’s matrices \(\boldsymbol{A}^{(0)}\) and \(\boldsymbol{B}^{(0)}\), we give them DirichletCollection priors: a Dirichlet distribution over each column, with concentration parameters \(\check{\boldsymbol{a}}^{(0)}\) and \(\check{\boldsymbol{b}}^{(0)}\):

\[ P_{\check{a}}\big(\boldsymbol{A}^{(0)}\big) = \mathcal{MatDir}\big(\boldsymbol{A}^{(0)};\, \check{\boldsymbol{a}}^{(0)}\big), \qquad P_{\check{b}}\big(\boldsymbol{B}^{(0)}\big) = \mathcal{MatDir}\big(\boldsymbol{B}^{(0)};\, \check{\boldsymbol{b}}^{(0)}\big) \]

The concentration parameters act as pseudo-counts: imagined observations the agent starts out with, before it has seen any data.

For \(\boldsymbol{B}^{(0)}\), since we have no a priori idea how the Roomba is going to move, we fill \(\check{\boldsymbol{b}}^{(0)}\) with ones. This makes every transition equally likely on average, while leaving the prior very uncertain, so the data will quickly overrule it. Remember that this prior gets updated once we start learning, so it’s fine if our initial guess is not quite accurate.

As for the observations, we have good reason to trust our Roomba’s camera. To represent this, we put large values on the diagonal of \(\check{\boldsymbol{a}}^{(0)}\), i.e. many pseudo-counts for correct readings. However, we also acknowledge that the camera is not infallible, so we put small positive values on the off-diagonal entries, which allow for the occasional misreading.

Since we will use variational inference, we also have to specify inference constraints. We will use a structured variational approximation to the true posterior distribution, in which the variational posterior over the trajectory of states is decoupled from the posteriors over the two matrices:

\[ Q\big(s^{(0:T,0)}, \boldsymbol{A}^{(0)}, \boldsymbol{B}^{(0)}\big) = Q\big(s^{(0:T,0)}\big)\, Q\big(\boldsymbol{A}^{(0)}\big)\, Q\big(\boldsymbol{B}^{(0)}\big) \]

The states keep their dependencies on each other over time, which is what makes the approximation structured rather than fully factorized. Decoupling the states from the matrices is what makes inference tractable. Let’s build the model!

@model function hidden_markov_model(oː, T)
    _B ~ DirichletCollection(ones(3, 3))          ## P_b̌(B^{(0)}), b̌^{(0)} = all ones
    _A ~ DirichletCollection([10.0  1.0  1.0;
                           1.0 10.0  1.0;
                           1.0  1.0 10.0])        ## P_ǎ(A^{(0)})

    sː[1] ~ Categorical(fill(1.0/3.0, 3))        ## t = 0: P_D(s^{(0,0)}), uniform
    oː[1] ~ DiscreteTransition(sː[1], _A)                ## t = 0: P(o^{(0,0)} | s^{(0,0)}, A^{(0)})
    for t in 1:T
        sː[t+1] ~ DiscreteTransition(sː[t], _B)          ## P(s^{(t,0)} | s^{(t-1,0)}, B^{(0)})
        oː[t+1] ~ DiscreteTransition(sː[t+1], _A)        ## P(o^{(t,0)} | s^{(t,0)}, A^{(0)})
    end
end

@constraints function hidden_markov_model_constraints()
    q(sː, _A, _B) = q(sː)q(_A)q(_B)
end

hidden_markov_model_constraints (generic function with 1 method)

Notes on the model code

Names: the model mirrors the agent’s math.

sː and oː have no hats and no prefix. Inside the model, these are the agent’s random variables, \(s^{(0:T,0)}\) and \(o^{(0:T,0)}\), the unbold symbols on the left-hand side of the generative model. The data fills in the observations, but in the model, oː is still the variable being observed.
_A and _B are single italic-bold matrices, \(\boldsymbol{A}^{(0)}\) and \(\boldsymbol{B}^{(0)}\). A model has no tuples, and tuple0 indexing would not work inside @model anyway, so the member index \((0)\) is implicit in the names. That is harmless with one factor and one modality. With more, names such as _Aı0ı and _Aı1ı would be used.
The comments map each line of the model to the corresponding factor of the generative model.

Structure: the observation at \(t = 0\).

The uniform prior applies to \(s^{(0,0)}\) itself, the room is also observed at \(t = 0\), and there are \(T + 1 = 600\) observations, matching the generative model and the simulated data. In the code:

The prior goes directly on sː[1] (that is, \(t = 0\)), and the first observation is attached to it.
The loop variable t is the math’s \(t\). The storage index is t+1 only because RxInfer’s vectors start at 1, so sː[t+1] holds \(s^{(t,0)}\).
The constraints factorize as q(sː)q(_A)q(_B), i.e. \(Q\big(s^{(0:T,0)}\big)\, Q\big(\boldsymbol{A}^{(0)}\big)\, Q\big(\boldsymbol{B}^{(0)}\big)\).

The data is passed to inference under the model’s name for the observations: data = (oː = Lôː_data,).

The model uses RxInfer’s DiscreteTransition node. Older RxInfer versions called it Transition.

Next, we define the agent and the time-stepping procedure.

function create_agent(; Lôː_data)
    T = length(Lôː_data) - 1    ## observations at t = 0, ..., T

    compute = () -> begin
        imarginals = @initialization begin
            q(_A) = vague(DirichletCollection, (3, 3))
            q(_B) = vague(DirichletCollection, (3, 3))
            q(sː) = vague(Categorical, 3)
        end
        infer(
            model          = hidden_markov_model(T=T),
            data           = (oː = Lôː_data,),
            constraints    = hidden_markov_model_constraints(),
            initialization = imarginals,
            returnvars     = (_A = KeepLast(), _B = KeepLast(), sː = KeepLast()),
            iterations     = 20,
            free_energy    = true,
        )
    end

    return compute
end

create_agent (generic function with 1 method)

4.5.6 Expected free energy and action selection

There is no expected free energy and no action selection in this part. The agent does not choose a control: it only infers \(Q\big(s^{(0:T,0)}\big)\), \(Q\big(\boldsymbol{A}^{(0)}\big)\) and \(Q\big(\boldsymbol{B}^{(0)}\big)\) from the observations. The per-step array \(\boldsymbol{G}\), its sum \(\boldsymbol{G}_{\Sigma}\), and the policy posterior \(Q(\pi)=\sigma(-\boldsymbol{G}_{\Sigma})\) belong to later parts that add act and a preference.

4.6 Agent Evaluation

Now it’s time to perform inference and find out where the Roomba went in our absence. Did it remember to clean the bathroom?

We’ll be using variational inference, which improves its estimates over a number of iterations (here 20). It therefore needs initial marginals as a starting point. RxInfer makes this easy with the vague function, which provides an uninformative guess. If you have better ideas, you can try a different initial guess and see what happens.

Inference gives us three things: the posterior over the Roomba’s room at every time step, \(Q\big(s^{(0:T,0)}\big)\), and the posteriors over the agent’s matrices, \(Q\big(\boldsymbol{A}^{(0)}\big)\) and \(Q\big(\boldsymbol{B}^{(0)}\big)\). Since we’re only interested in the final estimates, we keep only the result of the last iteration (KeepLast()), which still contains the posterior for every time step. We also track the free energy, which should decrease over the iterations as the estimates improve. Let’s start the inference process!

4.6.1 Evaluate with simulated data

4.6.1.1 Naive approach {N/A}

4.6.1.2 Active inference approach

### Simulation parameters
## Time steps t = 0, ..., T (600 ticks)
T = 599

## Initial state: the Roomba starts in the bedroom
Lŝˣ₀ = tuple0([1.0, 0.0, 0.0])    ## ŝ^{*(0)} = (ŝ^{*(0,0)})

1-element OffsetArray(::Vector{Vector{Float64}}, 0:0) with eltype Vector{Float64} with indices 0:0:
 [1.0, 0.0, 0.0]

### OVERRIDES

Let’s take a look at our results. If we’re successful, we should have a good idea about the actual layout of the apartment, i.e. which rooms are connected (a good posterior over \(\boldsymbol{B}^{(0)}\)), and about how often the Roomba’s camera misreads the floor (a good posterior over \(\boldsymbol{A}^{(0)}\)). Since we simulated the data ourselves, we can check this directly: the agent’s estimates should come close to the true \(\boldsymbol{B}^{\ast(0)}\) and \(\boldsymbol{A}^{\ast(0)}\). Let’s see if it worked.

In the plot above, look for the whereabouts sequence staying in a room for long stretches, with brief visits to the bathroom as the only way between the bedroom and the living room, and for surface readings that mostly match that room but are occasionally wrong.

## Environment: generate t = 0, ..., T
(execute_ai, observe_ai, state_ai) = create_envir(; LAˣ=LAˣ, LBˣ=LBˣ, Lŝˣ₀=Lŝˣ₀)
Lŝˣː = OffsetVector(Vector{Any}(undef, T+1), 0:T)    ## ŝ^{*(0:T)}: true states
Lôː  = OffsetVector(Vector{Any}(undef, T+1), 0:T)    ## ô^{(0:T)}: observations
for t = 0:T
    t > 0 && execute_ai()      ## move and observe (t = 0 is already set up)
    Lŝˣː[t] = state_ai()
    Lôː[t]  = observe_ai()
end

## Agent: infer the Roomba's rooms and learn A^{(0)}, B^{(0)} from all observations at once
Lôː_data   = [Lô[0] for Lô in parent(Lôː)]
compute_ai = create_agent(; Lôː_data=Lôː_data)
result     = compute_ai()

Inference results:
  Posteriors       | available for (_A, _B, sː)
  Free Energy:     | [692.251, 690.14, 688.281, 686.54, 684.685, 682.29, 678.45, 670.811, 651.831, 596.963, 490.151, 436.201, 425.596, 422.306, 421.243, 420.876, 420.731, 420.67, 420.644, 420.633]

[Lŝˣ[0] for Lŝˣ in Lŝˣː]    ## ŝ^{*(t,0)} for t = 0, ..., T: the true Room factor
## Lŝˣː

600-element OffsetArray(::Vector{Vector{Float64}}, 0:599) with eltype Vector{Float64} with indices 0:599:
 [1.0, 0.0, 0.0]
 [1.0, 0.0, 0.0]
 [1.0, 0.0, 0.0]
 [1.0, 0.0, 0.0]
 [1.0, 0.0, 0.0]
 [1.0, 0.0, 0.0]
 [1.0, 0.0, 0.0]
 [1.0, 0.0, 0.0]
 [1.0, 0.0, 0.0]
 [1.0, 0.0, 0.0]
 ⋮
 [0.0, 1.0, 0.0]
 [0.0, 1.0, 0.0]
 [0.0, 1.0, 0.0]
 [0.0, 1.0, 0.0]
 [0.0, 0.0, 1.0]
 [0.0, 0.0, 1.0]
 [0.0, 0.0, 1.0]
 [0.0, 0.0, 1.0]
 [0.0, 0.0, 1.0]

[Lô[0] for Lô in Lôː]    ## ô^{(t,0)} for t = 0, ..., T: the received surface

600-element OffsetArray(::Vector{Vector{Float64}}, 0:599) with eltype Vector{Float64} with indices 0:599:
 [1.0, 0.0, 0.0]
 [1.0, 0.0, 0.0]
 [1.0, 0.0, 0.0]
 [1.0, 0.0, 0.0]
 [1.0, 0.0, 0.0]
 [1.0, 0.0, 0.0]
 [1.0, 0.0, 0.0]
 [1.0, 0.0, 0.0]
 [1.0, 0.0, 0.0]
 [1.0, 0.0, 0.0]
 ⋮
 [0.0, 1.0, 0.0]
 [0.0, 1.0, 0.0]
 [0.0, 1.0, 0.0]
 [0.0, 0.0, 1.0]
 [0.0, 0.0, 1.0]
 [0.0, 0.0, 1.0]
 [0.0, 0.0, 1.0]
 [0.0, 0.0, 1.0]
 [1.0, 0.0, 0.0]

result.posteriors

Dict{Symbol, Any} with 3 entries:
  :_A => DirichletCollection{Float64, 2, Matrix{Float64}}(…
  :_B => DirichletCollection{Float64, 2, Matrix{Float64}}(…
  :sː => Categorical{Float64, Vector{Float64}}[Categorical{Float64, Vector{Floa…

println("Posterior mean of B^{(0)} (Room transitions):")
_B̄ = mean(result.posteriors[:_B])    ## B̄^{(0)} = E_Q[B^{(0)}], the agent's estimate of B^{*(0)}

Posterior mean of B^{(0)} (Room transitions):

3×3 Matrix{Float64}:
 0.874632   0.00434098  0.0923197
 0.0147468  0.925755    0.0841563
 0.110621   0.0699038   0.823524

This compares well with:

_Bˣı0ı = [
    0.9  0.0  0.1;
    0.0  0.9  0.1;
    0.1  0.1  0.8]

println("Posterior mean of A^{(0)} (surface likelihood):")
_Ā = mean(result.posteriors[:_A])    ## Ā^{(0)} = E_Q[A^{(0)}], the agent's estimate of A^{*(0)}

Posterior mean of A^{(0)} (surface likelihood):

3×3 Matrix{Float64}:
 0.901596   0.0304646  0.113377
 0.078098   0.939542   0.0289036
 0.0203056  0.0299935  0.857719

This compares well with:

_Aˣı0ı = [
    0.9  0.05 0.05;
    0.05 0.9  0.05;
    0.05 0.05 0.9]

Finally, we can check whether we were successful in keeping tabs on our Roomba’s whereabouts: we compare the agent’s inferred room at each time step, the most probable room under \(Q\big(s^{(t,0)}\big)\), with the room the Roomba was actually in, \(\hat{s}^{\ast(t,0)}\). We can also check whether the inference has converged by looking at the free energy: it should decrease over the iterations and then level off.

In the plots above, look for the posterior mean of \(\boldsymbol{B}^{(0)}\) recovering the missing bedroom–living room door (the zeros), and for \(\boldsymbol{A}^{(0)}\) recovering a strong diagonal with small off-diagonal camera errors.

rooms          = [argmax(Lŝˣ[0]) for Lŝˣ in parent(Lŝˣː)]                ## true room of ŝ^{*(t,0)}, t = 0, ..., T
inferred_rooms = argmax.(ReactiveMP.probvec.(result.posteriors[:sː]))   ## most probable room under Q(s^{(t,0)})
accuracy       = mean(inferred_rooms .== rooms)

p1 = scatter(0:T, rooms,
    title="Simulation results (inferred room correct $(round(100*accuracy, digits=1))% of the time)",
    label="True room",
    legend=:right,
    xlabel="Time t",
    yticks=(1:3, ["Bedroom", "Living room", "Bathroom"]),
    markercolor=:cyan,
    markersize=6,
    markerstrokewidth=0,
    alpha=.5,
    size=(900,550)
)
scatter!(p1, 0:T, inferred_rooms,
    label="Inferred room",
    markercolor=:green,
    markersize=3,
    markerstrokewidth=0,
    alpha=.5
)
p2 = plot(result.free_energy,
    label="Bethe free energy",
    xlabel="Iteration"
)
plot(p1, p2, layout=@layout([ a; b ]))

4.7 Conclusion and outlook

The agent observed only the surface of the floor and inferred the hidden whereabouts at each of \(T+1 = 600\) ticks. It also learned \(\boldsymbol{A}^{(0)}\) and \(\boldsymbol{B}^{(0)}\) from that single sequence. It chose no actions.

Three ingredients shaped what happened:

A hidden Markov structure. The whereabouts is a Markov chain; the camera emits a noisy surface.
Priors that match what is known. \(\check{\boldsymbol{b}}^{(0)}\) is uninformative, because the layout is unknown; \(\check{\boldsymbol{a}}^{(0)}\) trusts the camera but allows mistakes.
Structured variational inference. The trajectory of rooms stays coupled in time; the two matrices are decoupled from it, which keeps inference tractable.

This part asked whether wrapping the original RxInfer example in create_envir / create_agent still recovers the apartment and the Roomba’s path. It does, on this seed.

The learned layout matches the apartment. The posterior mean of \(\boldsymbol{B}^{(0)}\) is close to \(\boldsymbol{B}^{\ast(0)}\): the bedroom and the living room do not connect, and the bathroom is the passage. The posterior mean of \(\boldsymbol{A}^{(0)}\) is close to \(\boldsymbol{A}^{\ast(0)}\).

Caveats. These numbers come from one seed and from assumed transition and camera probabilities, not from a measured Roomba. A single run can be misleading; later parts that compare strategies should average over seeds. There is no control yet, so there is no expected free energy and no baseline policy.

Possible next steps:

Control. Add act, future and slide, and a preference over surface or whereabouts.
More factors. Doors that open and close, or a second modality.
Online inference. Call compute once per tick, rather than once on the whole sequence.

The next part can start from this environment and this agent, and add a decision.

APPENDIX A: Text conventions and their motivation

Three typographic styles, each with one role:

italic lower case names a model variable: whereabouts, surface. The italics mark the identifier, not a first mention. Names stay lower case, also at the start of a sentence; start with an article when needed.
“double quotes” name a value, matching the label in the code: “bedroom”, “hardwood”. They are used for nothing else.
bold marks a low-level heading (Inference.) or a label in a list or table. A low-level heading starts with a word, not with a name, so that the name can stay italic and lower case.

backticks mark code identifiers, such as LAˣ[0].

Italics, bold, underlining and capitals are not used for emphasis. New domain terms go in Terms used in this series, or are defined in the sentence, without italics.

The earlier rule that round brackets meant time only is superseded. Members of a collection — time, factors, modalities — take round brackets, \(s^{(t,n)}\), \(\boldsymbol{A}^{(m)}\). Positions along an array take square brackets, \(G^{[p,\tau]}\). Julia tuple indexing LAˣ[0] is kept; the named-member equivalent is _Aˣı0ı.

These choices keep each style doing one job, so a search for "bedroom" or for whereabouts finds the same thing in the text and in the code.

APPENDIX B: Modeling choices and assumptions, and their motivation

Assumed numbers, not measurements. The entries of \(\boldsymbol{B}^{\ast(0)}\) and \(\boldsymbol{A}^{\ast(0)}\) in 4.3.4 are assumed, so that a closed simulation exists. They can be replaced when real movement and camera-error rates arrive.

Always start in the bedroom. \(\boldsymbol{D}^{\ast(0)} = [1,0,0]^{\top}\) removes initial-state uncertainty from the environment, so that what the agent has to learn is the layout and the camera, not the starting room. The agent itself is not told this: its \(\boldsymbol{D}^{(0)}\) is uniform.

No door between bedroom and living room. The zeros in \(\boldsymbol{B}^{\ast(0)}\) force a passage through the bathroom. That makes the learned layout checkable by eye, and it is the structure the original RxInfer example used.

Camera mostly right. The 0.9 diagonal of \(\boldsymbol{A}^{\ast(0)}\) and the large diagonal of \(\check{\boldsymbol{a}}^{(0)}\) say that surface is informative but not perfect. A noiseless camera would make the whereabouts almost observed; a very noisy one would make 600 ticks too few.

One factor, one modality. This part is the tracking wrap of the original example. Control factors, preferences, and expected free energy are left out on purpose.

One inference on the whole sequence. compute calls infer once on \(\hat{\mathbf{o}}^{(0:T)}\), with 20 variational iterations. That is batch smoothing, not filtering. Online compute / slide is left for later parts.

One seed, no baselines. This part is for following one run. Comparing strategies over seeds is left until there are decisions to compare.

\(A\) is likelihood, \(B\) is transition. The original RxInfer example uses the engineering convention. The swap aligns the notebook with the active inference chapters and with the rest of the series.