# Realising Synthetic Active Inference Agents,
Part I: Epistemic Objectives and Graphical Specification Language

Magnus Koudahl  BIASLab, Department of Electrical Engineering, Eindhoven
University of Technology, The Netherlands  VERSES Research Lab, Los Angeles,
CA, USA, 90016  Thijs van de Laar  BIASLab, Department of Electrical
Engineering, Eindhoven University of Technology, The Netherlands  Bert de
Vries  BIASLab, Department of Electrical Engineering, Eindhoven University of
Technology, The Netherlands  GN Hearing, JF Kennedylaan 2, 5612 AB Eindhoven,
The Netherlands

###### Abstract

The Free Energy Principle (FEP) is a theoretical framework for describing how
(intelligent) systems self-organise into coherent, stable structures by
minimising a free energy functional. Active Inference (AIF) is a corollary of
the FEP that specifically details how systems that are able to plan for the
future (agents) function by minimising particular free energy functionals that
incorporate information seeking components. This paper is the first in a
series of two where we derive a synthetic version of AIF on free form factor
graphs. The present paper focuses on deriving a local version of the free
energy functionals used for AIF. This enables us to construct a version of AIF
which applies to arbitrary graphical models and interfaces with prior work on
message passing algorithms. The resulting messages are derived in our
companion paper. We also identify a gap in the graphical notation used for
factor graphs. While factor graphs are great at expressing a generative model,
they have so far been unable to specify the full optimisation problem
including constraints. To solve this problem we develop Constrained Forney-
style Factor Graph (CFFG) notation which permits a fully graphical description
of variational inference objectives. We then proceed to show how CFFG’s can be
used to reconstruct prior algorithms for AIF as well as derive new ones. The
latter is demonstrated by deriving an algorithm that permits direct policy
inference for AIF agents, circumventing a long standing scaling issue that has
so far hindered the application of AIF in industrial settings. We demonstrate
our algorithm on the classic T-maze task and show that it reproduces the
information seeking behaviour that is a hallmark feature of AIF.

##  1 Introduction

Active Inference (AIF) is an emerging framework for modelling intelligent
agents interacting with an environment. Originating in the field of
computational neuroscience, it has since been spread to numerous other fields
such as modern machine learning. At its core, AIF relies on variational
inference techniques to minimise a free energy functional. A key
differentiator of AIF compared to other approaches is the use of custom free
energy functionals such as Expected and Generalised free energies (EFE and
GFE, respectively). These functionals are specifically constructed so as to
elicit epistemic, information seeking behaviour when used to infer actions.

Optimisation of these functionals have so far relied on custom algorithms that
require evaluating very large search trees which has rendered upscaling of AIF
difficult. Many recent works, such as Branching Time AIF [4] and Sophisticated
AIF [8] have investigated algorithmic ways to prune the search tree in order
to solve this problem.

This paper is the first in a series of two parts where we take a different
approach and formulate a variation of the GFE optimisation problem using a
custom Lagrangian derived from constrained Bethe free energy (BFE) on factor
graphs. Using variational calculus we then derive a custom message passing
algorithm that can directly solve for fixed points of local GFE terms,
removing the need for a search tree. This allows us to construct a purely
synthetic approach to AIF which is motivated from the point of view of
engineering rather than neurobiology. From here on, we abbreviate "Synthetic
Active Inference" as SAIF.

SAIF applies to arbitrary graph topologies and interfaces with generic message
passing algorithms which allows for scaling up of AIF using off-the-shelf
tools. We accomplish this by constructing a node-local GFE on generic factor
graphs.

The present paper is structured as follows: In section 2, we review relevant
background material concerning Forney-style Factor Graphs (FFG’s) and Bethe
free energy (BFE). In section 3 we formalise what we mean by "epistemics" and
construct an objective that is local to a single node on an FFG and possesses
an epistemic, information-seeking drive. This objective turns out to be a
local version of the GFE which we review in section 3.2.

Once we start modifying the free energy functional, we recognise a problem
with current FFG notation. FFG’s visualise generative models but fail to
display a significant part of the optimisation problem, namely, the
variational distribution and the functional to be optimised. To remedy this,
we develop the Constrained FFG (CFFG) graphical notation in section 5 as a
method for visualising both the variational distribution and any adaptations
to the free energy functional that are needed for AIF. These tools form the
basis of the update rules derived in our companion paper [16].

With the CFFG notation in hand, in section 6 we then proceed to demonstrate
how to recover prior algorithms for AIF as message passing on a CFFG. AIF is
often described as message passing on a probabilistic graphical model, see [5,
7, 14, 27] for examples. However, this relationship has not been properly
formalised before, in part because adequate notation has been lacking. Using
CFFG’s it is straightforward to accurately write down this relation. Further,
due to the modular nature of CFFG’s it becomes easy to devise extensions to
prior AIF algorithms that can be implemented using off-the-shelf message
passing tools.

Finally, section 7 demonstrates a new algorithm for policy inference using
SAIF that scales linearly in the planning horizon, providing a solution to a
long standing barrier for scaling AIF to larger models and more complex tasks.

##  2 The Lagrangian approach to message massing

In this section we will review the Bethe free energy (BFE) along with Forney-
style factor graphs (FFG’s). These concepts will form the foundation from
which we will build towards local epistemic, objectives.

As a free energy functional, the BFE is unique because stationary points of
the BFE correspond to solutions of the Belief Propagation (BP) algorithm [21,
26, 29] which provides exact inference on tree-structured graphs.

Furthermore, by adding constraints to the BFE using Lagrange multipliers, one
can form a custom Lagrangian for an inference problem. Taking the first
variation of this Lagrangian, one can then solve for stationary points and
obtain message passing algorithms that solve the desired problem. Prior work
[26, 29] have shown that adding additional constraints to the BFE allows for
deriving a host of different message passing algorithms including Variational
Message Passing (VMP) [28] and Expectation Propagation (EP) [19] among others.
We refer interested readers to [26] for a comprehensive overview of this
technique and how different choices of constraints can lead to different
algorithms.

Constraints are specified at the level of nodes and edges, meaning this
procedure can produce hybrid message passing algorithms, foreshadowing the
approach we are going to take for deriving a local, epistemic objective for
AIF.

###  2.1 Bethe free energy and Forney-style factor graphs

Throughout the remainder of the paper we will use FFG’s to visualise
probabilistic models models. In Section 5 we extend this notation to
additionally allow for specifying constraints on the variational optimisation
objective.

Following [26] we define an FFG as a graph
$\mathcal{G}=(\mathcal{V},\mathcal{E})$ with nodes $\mathcal{V}$ and edges
$\mathcal{E}\subseteq\mathcal{V}\times\mathcal{V}$. For a node
$a\in\mathcal{V}$ we denote the connected edges by $\mathcal{E}(a)$. Similarly
for an edge $i\in\mathcal{E}$, we denote the connected nodes by
$\mathcal{V}(i)$.

An FFG can be used to represent a factorised function (model) over variables
$\bm{s}$, as

$\displaystyle f(\bm{s})=\prod_{a\in\mathcal{V}}f_{a}(\bm{s}_{a})\,,$ (1)

where $\bm{s}_{a}$ collects the argument variables of the factor $f_{a}$.
Throughout this paper we will use cursive bold font to denote collections of
variables. In the corresponding FFG, the factor $f_{a}$ is denoted by a square
node and connected edges $\mathcal{E}(a)$ represent the argument variables
$\bm{s}_{a}$.

$f_{b}$$f_{c}$$f_{a}$$f_{d}$$f_{e}$$s_{1}$$s_{2}$$s_{3}$$s_{4}$ Figure 1:
Example of an FFG

As an example, we can consider the FFG shown in Fig. 1 which corresponds to
the model

$\displaystyle f(s_{1},s_{2},s_{3},s_{4})$
$\displaystyle=f_{a}(s_{1})f_{b}(s_{1},s_{2},s_{3})f_{c}(s_{2},s_{4})f_{d}(s_{3})f_{e}(s_{4})$
(2)

In Fig. 1, the vertex set is $\mathcal{V}=\\{a,\dots,e\\}$ and the edge set is
$\mathcal{E}=\\{1,\dots,4\\}$. As an example, the neighbouring edges of the
node $c$ are given by $\mathcal{E}(c)=\\{2,4\\}$. In this way, FFG’s allow for
a simple visualisation of the factorisation properties of a high-dimensional
function.

The problem we will focus on concerns minimisation of a free energy functional
over a generative model. More formally, given a model (1) and a "variational"
distribution $q(\bm{s})$, the Variational free energy (VFE) is defined as

$\displaystyle F[q]\triangleq\int
q(\bm{s})\log\frac{q(\bm{s})}{f(\bm{s})}\operatorname{d}\\!{\bm{s}}\,.$ (3)

Variational inference concerns minimising this functional, leading to the
solution

$\displaystyle q^{*}=\arg\min_{q\in\mathcal{Q}}F[q]\,,$ (4)

with $\mathcal{Q}$ denoting the admissible family of functions $q$. The
optimised VFE upper-bounds the negative log-evidence (surprisal), as

$\displaystyle F[q^{*}]=\underbrace{\int
q^{*}\\!(\bm{s})\log\frac{q^{*}\\!(\bm{s})}{p(\bm{s})}\operatorname{d}\\!{\bm{s}}}_{\text{Posterior
divergence}}\underbrace{-\log Z}_{\text{Surprisal}}\,,$ (5)

with $Z=\int f(\bm{s})\differential\bm{s}$ is the model evidence and the exact
posterior is given by $p(\bm{s})=f(\bm{s})/Z$.

The BFE applies the Bethe assumption to the factorisation of $q$ which yields
an objective that decomposes into a sum of local free energy terms, each local
to a node on the corresponding FFG. Each node local free energy will include
entropy terms from all connected edges. Since an edge is connected to (at
most) two nodes, that means the corresponding entropy term would be counted
twice. To prevent overcounting of the edge entropies, the BFE includes
additional terms entropy terms that cancel out overcounted terms.

Under the Bethe approximation $q(\bm{s})$ is given by

$\displaystyle
q(\bm{s})=\prod_{a\in\mathcal{V}}q_{a}(\bm{s}_{a})\prod_{i\in\mathcal{E}}q_{i}(s_{i})^{1-d_{i}}\,$
(6)

with $d_{i}$ the degree of edge $i$. As an example, on the FFG shown in Fig. 1
this would correspond to a variational distribution of the form

$\displaystyle q(s_{1},\dots,s_{4})$
$\displaystyle=\frac{q_{a}(s_{1})q_{b}(s_{1},s_{2},s_{3})q_{c}(s_{2},s_{4})q_{d}(s_{3})q_{e}(s_{4})}{q_{1}(s_{1})q_{2}(s_{2})q_{3}(s_{3})q_{4}(s_{4})}$
(7)

where we see terms for the edges in the denominator and for the nodes in the
numerator. With this definition, the free energy factorises over the FFG as

$\displaystyle F[q]=\sum_{a\in\mathcal{V}}\underbrace{\int
q_{a}(\bm{s}_{a})\log\frac{q_{a}(\bm{s}_{a})}{f_{a}(\bm{s}_{a})}\operatorname{d}\\!{\bm{s}_{a}}}_{F[q_{a}]}+\sum_{i\in\mathcal{E}}(1-d_{i})\underbrace{\int
q_{i}(s_{i})\log\frac{1}{q_{i}(s_{i})}\operatorname{d}\\!{s_{i}}}_{H[q_{i}]}\,.$
(8)

Eq. (8) defines the BFE. Note that $F$ defines a free energy functional which
can be either local or global depending on its arguments. More specifically,
$F[q]$ defines the free energy for the entire model, while $F[q_{a}]$ defines
a node local (to the node $a$) BFE contribution of the same functional form.

Under optimisation of the BFE - solving Eq. (4) - the admissible set of
functions $\mathcal{Q}$ enforces consistent normalisation and marginalisation
of the node and edge local distributions, such that

$\displaystyle\int q_{i}(s_{i})\operatorname{d}\\!{s_{i}}=1$
$\displaystyle\text{ for all }i\in\mathcal{E}$ (9a) $\displaystyle\int
q_{a}(\bm{s}_{a})\operatorname{d}\\!{\bm{s}_{a}}=1$ $\displaystyle\text{ for
all }a\in\mathcal{V}$ (9b) $\displaystyle\int
q_{a}(\bm{s}_{a})\operatorname{d}\\!{\bm{s}_{a\setminus i}}=q_{i}(s_{i})$
$\displaystyle\text{ for all }a\in\mathcal{V},i\in\mathcal{E}(a)\,.$ (9c)

Because we always assume the constraints given by Eq. (9) to be in effect, we
will omit the subscript on individual $q$’s moving forwards and instead let
the arguments determine which marginal we are referring to, for example
writing $q(s_{a})$ instead of $q_{a}(s_{a})$.

A core aspect of FFG’s is that they allow for easy visualisation of message
passing algorithms. Message passing algorithms are a family of distributed
inference algorithms with the common trait that they can be viewed as messages
flowing on an FFG.

As a general rule, to infer a marginal for a variable, messages are passed on
the graph toward the associated edge for that variable. Multiplication of
colliding (forward and backward) messages on an edge yields the desired
posterior marginal. We denote a message on an FFG by arrows pointing in the
direction that the message flows.

$f_{b}$$f_{c}$$f_{a}$$f_{d}$$f_{e}$$s_{1}$$\rightarrow$$s_{2}$$\rightarrow\leftarrow$$s_{3}$$\uparrow$$s_{4}$$\uparrow$
Figure 2: Example of messages flowing on an FFG

We show an example in Fig. 2 of inferring a posterior marginal for the
variable $s_{2}$. In this example, messages are flowing on the FFG of Fig. (1)
towards the variable $s_{2}$ where we see two arrows colliding. The exact form
of the individual messages depend on the algorithm being used.

##  3 Defining epistemic objectives

Now we move on to the central topic of AIF, namely agents that interact with
the world they inhabit. Under the heading of AIF, an agent entails a
generative model of its environment and is engaged in the process of achieving
future goals (specified by a target distribution) through actions. This task
can be cast as a process of free energy minimization [6, 8, 10]. A natural
question to ask is then, what should this free energy functional look like and
why? Here we wish to highlight a core features of AIF that sets it apart from
other approaches: Systematic information gathering by targeted exploration of
an environment. Whatever the form of our free energy functional, it should
lead to agents that possess an exploratory drive consistent with AIF.

While it is tempting to default to a standard variational free energy (VFE),
prior work [25] has shown that directly optimising BFE or VFE when inferring a
sequence of actions (which we will refer to as a _policy_) does not lead to
agents that systematically explore their environment. Instead, directly
inferring a policy by minimising a VFE/BFE leads to KL-control [13] 111Other
works will sometimes use $\pi$ as a symbol to denote a policy. This means that
VFE/BFE is not the correct choice when we desire an agent that actively
samples its environment with the explicit purpose of gathering information.

Instead a hallmark feature of AIF is the use of alternative functionals in
place of either the VFE or the BFE, specifically for inferring policies. The
goal of these alternative functionals is often specifically to induce an
epistemic, explorative term that drives AIF agents to seek out information.

Epistemics, epistemic behaviour or "foraging for information" are commonly
used terms in the AIF literature and related fields. While we have used the
term colloquially until this point, we now clarify formally how we use the
term in the present paper and how it relates to the objective functionals we
consider. A core problem is that epistemics is most often defined in terms of
the behaviour of agents ("what does my agent _do_?") rather than from a
mathematical point of view. Prior work on this point includes [11, 18].

We take the view that epistemics arise from the optimisation of either a
mutual information (MI) term or a bound thereon. The MI (between two variables
$x$ and $z$) is defined as [17]

$\displaystyle I(x,z)$ $\displaystyle=\iint
p(x,z)\log\frac{p(x,z)}{p(x)p(z)}\differential x\differential z\,.$ (10)

To gain an intuition for why maximising MI leads to agents that seek out
information, we can rewrite MI as

$\displaystyle I(x,z)$ $\displaystyle=\iint p(x,z)\log\frac{p(z\mid
x)}{p(z)}\differential x\differential z$ $\displaystyle=H(z)-H(z\mid
x)=H(x)-H(x\mid z)$ (11)

Since MI is symmetric in its arguments, Eq. (3) can equally well be written in
terms of $x$ rather than $z$. Eq. (3) shows that MI decomposes as the
difference between the marginal entropy $z$ and the expected entropy of $z$
conditional on $x$.

If we let $z$ denote an internal state of an agent and $x$ an observation and
allow our agent to choose $x$ \- for instance through acting on an environment
- we can see why maximising Eq. (10) biases the agent towards seeking our
observations that reduce entropy in $z$. Maximising Eq. (10) means the agent
will prefer observations that provide useful (in the sense of reducing
uncertainty) information about its internal states $z$. For this reason MI is
also known as Information Gain.

Actually computing Eq. (10) is often intractable and in practice a bound is
often optimised instead.

###  3.1 Constructing a local epistemic objective

At this point we have seen that the BFE is defined over arbitrary FFG’s, yet
does not lead to epistemic behaviour. On the other hand, maximising MI leads
to the types of epistemic behaviour we desire, yet is not distributed like the
BFE. The question becomes whether there is a way to merge the two and obtain a
distributed functional - like the BFE - that includes an epistemic term?

We will now show how to construct such a functional. Our starting point will
be the BFE given by Eq. (8). We will focus on a single node $a$ and its
associated local energy term and partition the incoming edges into two sets,
$\bm{x}$ and $\bm{z}$. This gives the node local free energy

$\displaystyle F[q_{a}]$ $\displaystyle=\iint
q(\bm{x},\bm{z})\log\frac{q(\bm{x},\bm{z})}{f(\bm{x},\bm{z})}\differential\bm{x}\differential\bm{z}$
(12)

Now we need to add on an MI term to induce epistemics. Since we are minimising
our free energy functional and want to maximise MI, we augment the free energy
with a negative MI term as

$\displaystyle G[q_{a}]$ $\displaystyle=\overbrace{\iint
q(\bm{x},\bm{z})\log\frac{q(\bm{x},\bm{z})}{f(\bm{x},\bm{z})}\differential\bm{x}\differential\bm{z}}^{\text{Variational
free energy}}+\overbrace{\iint
q(\bm{x},\bm{z})\log\frac{q(\bm{x})q(\bm{z})}{q(\bm{x},\bm{z})}\differential\bm{x}\differential\bm{z}}^{\text{Negative
Mutual Information}}$ (13a) $\displaystyle=\iint
q(\bm{x},\bm{z})\log\frac{\cancel{q(\bm{x},\bm{z})}}{f(\bm{x},\bm{z})}\differential\bm{x}\differential\bm{z}+\iint
q(\bm{x},\bm{z})\log\frac{q(\bm{x})q(\bm{z})}{\cancel{q(\bm{x},\bm{z})}}\differential\bm{x}\differential\bm{z}$
(13b) $\displaystyle=\iint
q(\bm{x},\bm{z})\log\frac{q(\bm{x})q(\bm{z})}{f(\bm{x},\bm{z})}\differential\bm{x}\differential\bm{z}$
(13c)

which we can recognise as a node-local GFE [20]. In section 3.2 we review the
results of [20] to expand upon this statement. Eq. (13a) provides a
straightforward explanation of the kinds of behaviour that we can expect out
of agents optimising a GFE. Namely, minimising BFE with a goal prior
corresponds to performing KL-control [13] while the MI term adds an epistemic,
information seeking component. Viewed in this way, there is nothing mysterious
about the kind of objective optimised by AIF agents: It is simply the sum of
two well known and established objectives that are each widely used within the
control and reinforcement learning communities.

###  3.2 Generalised free energy

The objective derived in Eq. (13) is a node local version of the GFE
originally introduced by [20]. In this section we review the GFE as
constructed by [20] in order to relate our construction to prior work on
designing AIF functionals. In section 6 we show how to reconstruct the exact
method of [20] using the tools we develop in this paper. Prior to [20], the
functional of choice was the Expected free energy (EFE) [5, 7]. [20]
identified some issues with the EFE and proposed the GFE as a possible
solution.

We show how to reconstruct the original EFE-based algorithm of [7] using our
local objective in section 6. We also provide a detailed description of the
EFE in Appendix. A since it is still a popular choice for designing AIF
agents.

A core issue with the EFE is that it is strictly limited to planning over
future timesteps. This means that AIF agents that utilise the EFE functional
need to maintain two separate models: One for inferring policies (using EFE)
and one for state inference (using VFE/BFE) that updates as observations
become available. The key advantage of EFE is that it induces the epistemic
drive that we desire from AIF agents.

The goal of [20] was to extend upon the EFE by introducing a functional that
could induce similar epistemic behaviour when used to infer policies while at
the same time reducing to a VFE when dealing with past data points. In this
way an agent would no longer have to maintain two separate models and could
instead utilise only one.

The GFE as introduced by [20] is tied to a specific choice of generative
model. The model introduced by [20] is given by

$\displaystyle p(\bm{\mathrm{x}},\bm{\mathrm{z}}\mid\bm{\mathrm{\hat{u}}})$
$\displaystyle\propto
p(\bm{\mathrm{z}}_{0})\prod_{k=1}^{T}p(\bm{\mathrm{x}}_{k}|\bm{\mathrm{z}}_{k})p(\bm{\mathrm{z}}_{k}|\bm{\mathrm{\hat{u}}}_{k},\bm{\mathrm{z}}_{k-1})\tilde{p}(\bm{\mathrm{x}}_{k})$
(14)

where $\bm{\mathrm{x}}$ denotes observations, $\bm{\mathrm{z}}$ denotes latent
states and $\bm{\mathrm{\hat{u}}}$ denotes a fixed policy. Throughout this
paper we will use $t$ to refer to the current time step in a given model and
denote fixed values by a hat, here exemplified by $\bm{\mathrm{\hat{u}}}$.
Further we also use $\bm{\mathrm{z}}$ to denote vectors.

Given a current time step $t$, we have that for future timesteps $k>t$,
$\tilde{p}(\bm{\mathrm{x}}_{k})$ defines a goal prior over desired future
observations. For past time steps $k\leq t$, we instead have
$\tilde{p}(\bm{\mathrm{x}}_{k})=1$ which makes it uninformative 222[20] writes
this as the prior being flat to achieve a similar effect. Writing the model in
this way allows for a single model for both perception (integrating past data
points) and action (inferring policies) since both past and future time steps
are included.

GFE is defined as

$\displaystyle G[q;\hat{\bm{u}}]$ $\displaystyle=\sum_{k=1}^{T}\iint
q(\bm{\mathrm{x}}_{k}|\bm{\mathrm{z}}_{k})q(\bm{\mathrm{z}}_{k}|\bm{\mathrm{\hat{u}}}_{k})\log\frac{q(\bm{\mathrm{x}}_{k}|\bm{\mathrm{\hat{u}}}_{k})q(\bm{\mathrm{z}}_{k}|\bm{\mathrm{\hat{u}}}_{k})}{\tilde{p}(\bm{\mathrm{x}}_{k})p(\bm{\mathrm{x}}_{k},\bm{\mathrm{z}}_{k}|\bm{\mathrm{\hat{u}}}_{k})}\differential\bm{\mathrm{x}}_{k}\differential\bm{\mathrm{z}}_{k}$
(15)

where

$\displaystyle q(\bm{\mathrm{x}}_{k}|\bm{\mathrm{z}}_{k})$
$\displaystyle=\begin{cases}\delta(\bm{\mathrm{x}}_{k}-\bm{\mathrm{\hat{x}}}_{k})&\text{if
}k\leq t\\\ p(\bm{\mathrm{x}}_{k}|\bm{\mathrm{z}}_{k})&\text{if
}k>t\end{cases}\,.$ (16)

and $\bm{\mathrm{\hat{x}}}_{k}$ denotes the observed data point at timestep
$k$. $p(\bm{\mathrm{x}}_{k},\bm{\mathrm{z}}_{k}\mid\bm{\mathrm{\hat{u}}}_{k})$
can be found recursively by Bayesian smoothing, see [15, 24] for details. To
see how the GFE reduces to a VFE when data is available, we refer to Appendix
B.

The GFE introduced by [20] improves upon prior work utilising EFE but still
has some issues. The first lies in tieing it to the model definition in Eq.
(14). Being committed to a model specification apriori severely limits what
the GFE can be applied to since not all problems are going to fit the model
specification. Additionally a more subtle issue lies in the commitment to a
hard time horizon. If $t$ denotes the current timestep, at some point time
will advance to the point $t>T$. When this happens Eq. (14) loses the capacity
to plan and becomes static since all observations are clamped by virtue of Eq.
(16). A further complication arises from $\hat{\bm{u}}$ being a fixed
parameter and not a random variable. Being a fixed parameter means it is not
possible to perform inference for $\hat{\bm{u}}$ which in turn makes scaling
difficult. Moving to a fully local version of the GFE instead means we can
construct a new synthetic approach to AIF that addresses these issues.

##  4 SAIF - Synthetic Active Inference

Armed with the node local GFE derived in Eq. (13), we can construct a
Lagrangian for AIF. The goal is to adapt the Lagrangian approach to message
passing sketched in section 2 to derive a message passing algorithm that
optimises local GFE in order to have a distributed inference procedure that
incorporates epistemic terms. Naively applying the method of [26] to Eq. (13)
does not yield useful results because the numerator differs from the term we
take the expectation with respect to. To obtain a useful solution we need that

$\displaystyle q(\bm{x},\bm{z})$ $\displaystyle=q(\bm{x}\mid\bm{z})q(\bm{z})$
(17)

and add the original assumption in [20], given by Eq. (16)

$\displaystyle q(\bm{x}\mid\bm{z})$ $\displaystyle\triangleq
p(\bm{x}\mid\bm{z})\,.$ (18)

With these additional assumptions we can obtain meaningful solutions and
derive a message passing algorithm that optimises a local GFE and induces
epistemic behaviour which we demonstrate in section7. The detailed derivation
of these results can be found in our companion paper [16]. For practical AIF
modelling, we provide the relevant messages in Fig. 21.

At this point we will instead show a different way to arrive at the local GFE.
This approach is more "mechanical" but has the advantage that it can more
easily be written in terms of constraints on the BFE. Our starting point will
once again be a free energy term local to the node $a$

$\displaystyle F[q_{a}]$ $\displaystyle=\int
q(\bm{s}_{a})\log\frac{q(\bm{s}_{a})}{f(\bm{s}_{a})}\differential\bm{s}_{a}$
(19)

To turn Eq. (19) in to a local GFE we need to perform two steps. The first is
to enforce a mean field factorisation

$\displaystyle q(\bm{s}_{a})=\prod_{i\in\mathcal{E}(a)}q(s_{i})$ (20)

The second is to change the expectation to obtain

$\displaystyle\int
q(\bm{s}_{a})\log\frac{q(\bm{s}_{a})}{f(\bm{s}_{a})}\differential\bm{s}_{a}\implies\int
p(s_{i}|\bm{s}_{a\backslash i})q(\bm{s}_{a\backslash
i})\log\frac{q(\bm{s}_{a})}{f(\bm{s}_{a})}\differential\bm{s}_{a}$ (21)

where $i\in\mathcal{E}(a)$. This partitions the connected variables into two
sets: A set of variables where the expectation is modified and any remainder
that is not modified. $p(s_{i}\mid\bm{s}_{a\backslash i})$ denotes a
conditional probability distribution. We write $p$ here instead of $f$ to
emphasise that the conditional needs to be normalised in order for us to be
able to take the expectation.

We refer to this move as a _P-substitution_. Once the mean field factorisation
is enforced, we can recognise Eq. (21) as a node local GFE $G[q_{a}]$.

As an example of constructing a node local GFE with this approach, we can
consider a node with two connected variables, $\\{\bm{x},\bm{z}\\}$.

The local free energy becomes

$\displaystyle F[q_{a}]$ $\displaystyle=\iint
q(\bm{x},\bm{z})\log\frac{q(\bm{x},\bm{z})}{p(\bm{x},\bm{z})}\differential\bm{x}\differential\bm{z}$
(22)

Now we apply a mean field factorisation

$\displaystyle F[q_{a}]$ $\displaystyle=\iint
q(\bm{x})q(\bm{z})\log\frac{q(\bm{x})q(\bm{z})}{p(\bm{x},\bm{z})}\differential\bm{x}\differential\bm{z}$
(23)

And finally perform P-substitution to obtain a node local GFE

$\displaystyle\iint{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}q(\bm{x})}q(\bm{z})\log\frac{q(\bm{x})q(\bm{z})}{p(\bm{x},\bm{z})}\differential\bm{x}\differential\bm{z}$
$\displaystyle\implies\iint{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}p(\bm{x}\mid\bm{z})}q(\bm{z})\log\frac{q(\bm{x})q(\bm{z})}{p(\bm{x},\bm{z})}\differential\bm{x}\differential\bm{z}$
(24)

We denote the set of nodes for which we want to perform P-substitution with
$\mathcal{P}\subseteq\mathcal{V}$. When performing P-substitution, we are
replacing a local _variational_ free energy $F[q_{a}]$ with a local
_generalised_ free energy $\mathcal{G}[q_{a}]$. Armed with P-substitution, we
can now write the simplest instance of a Lagrangian for Active Inference

$\displaystyle\mathcal{L}[q]$
$\displaystyle=\underbrace{\sum_{a\in\mathcal{P}}G[q_{a}]}_{\begin{subarray}{c}\text{P-substituted
subgraph}\\\ \text{w / naive mean
field}\end{subarray}}+\underbrace{\sum_{b\in\mathcal{V}\backslash\mathcal{P}}F[q_{b}]}_{\begin{subarray}{c}\text{Node
local}\\\ \text{free
energies}\end{subarray}}+\sum_{i\in\mathcal{E}}(1-d_{i})\underbrace{H[q(s_{i})]}_{\text{Edge
entropy}}$
$\displaystyle\quad\quad+\underbrace{\sum_{a\in\mathcal{V}}\sum_{i\in\mathcal{E}}\int\lambda_{ia}(s_{i})\Bigg{[}q(s_{i})-\int
q(\bm{s}_{a})\differential s_{a\backslash i}\Bigg{]}\differential
s_{i}}_{\text{Marginalisation}}$
$\displaystyle\quad\quad+\underbrace{\sum_{a\in\mathcal{V}}\lambda_{a}\Bigg{[}\int
q(\bm{s}_{a})\differential\bm{s}_{a}-1\Bigg{]}}_{\text{Normalisation of node
marginals}}+\underbrace{\sum_{i\in\mathcal{E}}\lambda_{i}\Bigg{[}\int
q(s_{i})\differential s_{i}-1\Bigg{]}}_{\text{Normalisation of edge
marginals}}\,.$ (25)

With the Active Inference Lagrangian in hand, we can now solve for stationary
points using variational calculus and obtain message passing algorithms for
Synthetic Active Inference. The key insight is that messages flowing out of
$\mathcal{P}$ derived from stationary points of (4) will correspond to
stationary points of the local GFE rather than BFE, meaning the result will
include an epistemic component. This paves the way for a localised, synthetic
version of AIF that applies to arbitrary graph structures, does not suffer
from scaling issues as the planning horizon increases and can be solved
efficiently and asynchronously using message passing.

##  5 Constrained Forney-style factor graphs

While FFG’s are a useful tool for writing down generative models, we have by
now established the importance of knowing the exact functional to be
minimised. This requires specifying not just the model $f$ but also the family
$\mathcal{Q}$ through constraints and any potential P-substitutions. This is
important if we want to be able to succintly specify not just the model but
also the exact inference problem we aim to solve.

We will now develop just such a new notation for writing constraints directly
as part of the FFG. We refer to FFG’s with added constraints specification as
Constrained FFG’s or CFFG’s.

Figure 3: An example FFG.

Fig 3 shows a model comprised of five edges and six nodes. FFG’s traditionally
represent the model $f$ using squares connected by lines as shown in Fig. 3.
The squares represent factors and the connections between them represent
variables. Connecting an edge to a square node indicates that the variable on
that edge is an argument of the factor it is connected to.

The notation for CFFG’s adhere to similar principles when specifying $f$.
However, we augment the FFG with circular beads to indicate the constraints
that define our family $\mathcal{Q}$. Each factor of the variational
distribution in $q$ will correspond to a bead and the position of a bead
indicates to which marginal it refers - a bead on an edge denotes an edge
marginal $q(s_{i})$ and a bead inside a node denotes a node marginal
$q(\bm{s}_{a})$. An empty bead will denote the default normalisation
constraints while a connection between beads indicates marginalisation
constraints following Eq. (9). These beads form the basic building blocks of
our notation.

To write the objective corresponding to the model in Fig. 3 given the default
constraints of Eq. (9), we add beads for every term and extend edges through
the node boundary to connect variables that are under marginalisation
constraints as shown in Fig. 4.

Figure 4: Example CFFG with normalisation and marginalisation constraints.

###  5.1 Factorisation constraints

We will now extend our notation with the most common types of constraints used
for defining $\mathcal{Q}$. A common choice is factorisation of the
variational distribution with the most well-known example being the naive mean
field approximation. Under a naive mean field factorisation all marginals are
considered independent. Formally this means we enforce

$\displaystyle q(\bm{s}_{a})$
$\displaystyle=\prod_{i\in\mathcal{E}(a)}q(s_{i})\,.$ (26)

To write Eq. (26) on a CFFG we need to replace the joint node marginal with
the product of adjacent edge marginals.

$f_{a}$$f_{a}$ Figure 5: Changing a local joint factorisation to a naive mean
field assumption on a CFFG.

To do this we can replace the bead indicating the joint marginal with a bead
for each edge marginal in Eq. (26) as shown in Fig. 5.

The naive mean field is the strongest factorisation possible. It is possible
to utilise less aggressive factorisations by appealing to a structured mean
field approximation instead. The structured mean field constraint takes the
form

$\displaystyle q(\bm{s}_{a})=\prod_{n\in l(a)}q^{n}(\bm{s}_{a}^{n})$ (27)

where $l(a)$ denotes a set of one or more edges connected to the node $a$ such
that each element in $\mathcal{E}(a)$ can only appear in $l(a)$ once [26]. For
example if $\mathcal{E}(a)=\\{i,j,k\\}$ corresponding to variables
$\\{x,y,z\\}$, we can factorise $q(x,y,z)$ as $q(x)q(y,z)$ or $q(z)q(x,y)$ but
not as $q(x,y)q(y,z)$ since $y$ appears twice. The naive mean field is a
special case of the structured mean field where every variable appears only by
itself.

To write a structured mean field factorisation on a CFFG we can apply a
similar logic and replace the single bead denoting the joint with beads that
match the structure of $l(a)$. Each set of variables that are factorised
together correspond to a single bead connected to the edges in the set that
factor together.

$f_{a}$$f_{a}$$f_{a}$ Figure 6: Changing a local joint factorisation to
structured mean field on a CFFG.

Fig. 6 shows two example factorisations. The first option factorises the four
incoming edges into two sets of two while the second partitions the incoming
edges into two sets of one and a single set of two. Using these principles it
is possible to specify complex factorisation constraints as part of the CFFG
by augmenting each node on the original FFG.

The final situation we need to consider is the case when a single node has a
variable or very high number of incoming edges. An example could be a Gaussian
Mixture Model with a variable number of mixture components. On a CFFG we
indicate variable or large numbers of identical edges by drawing two of the
relevant edges and separating them with dots ($\cdots$) as shown in Fig. 7. To
indicate factorisation constraints we can write either a joint or a naive mean
field factorisation between the two edges, letting the dots denote that a
similar factorisation applies to the remaining edges.

$f_{a}$$f_{a}$ Figure 7: Mean field (left) and joint (right) factorisation
constraints for variable numbers of edges on a CFFG.

###  5.2 Form constraints

We will now extend CFFG notation with constraints on the functional form of
nodes and edges. Form constraints are used to enforce a particular form for a
local marginal on either an edge or a node. For an edge $s_{i}$ they enforce

$\displaystyle\int q(\bm{s}_{a})\differential s_{a\backslash i}$
$\displaystyle=q(s_{i})=g(s_{i})$ (28)

where $g(s_{i})$ denotes the functional form we are constraining the edge
marginal $s_{i}$ to take. Form constraints on node marginals take the form

$\displaystyle q(\bm{s}_{a})$ $\displaystyle=g(\bm{s}_{a})$ (29)

Conventionally FFG’s denote the form of a factor by a symbol inside the node.
We adopt a similar convention to denote form constraints on $q$ by adding
symbols within the corresponding beads. For instance we can indicate a
Gaussian form constraint on an edge as shown in Fig 8

$\mathcal{N}$ Figure 8: Notation for enforcing a Gaussian form constraint on
an edge.

Note that this is not dependent on the form of the neighbouring factors. This
is a subtle point as it allows us to write approximations into the
specification of $\mathcal{Q}$. As an example the unconstrained marginal in
Fig 8 might be bimodal or highly skewed but by adding a form constraint, we
are enforcing a Gaussian approximation. Outside of a few special cases,
enforcing form constraints on edges is rarely done in practice since the
functional form of $q$ most often follows from optimisation [26].

A special case of form constraints is the case of dangling edges (edges that
are not terminated by a factor node). Technically these would not warrant a
bead since they would not appear explicitly in the BFE due to having degree
$1$. Intuitively this means that the edge marginal is only counted once and we
therefore do not need to correct for overcounting. However without a bead
there is nowhere to annotate a form constraint which is problematic.

The solution for CFFG notation is to simply draw the bead anyway, in case a
form constraint is needed. This is formally equivalent to terminating the
dangling edge by a factor node with the node function $f_{a}(\bm{s}_{a})=1$.
Terminating the edge in this way means the edge in question now has degree $2$
and therefore warrants a bead. This is always a valid move since
multiplication by $1$ does not change the underlying function [26].

We can denote form constraints on node marginals in the same manner as edge
marginals. We show an example in Fig 9 where we enforce a Gaussian form
constraint on one node marginal and a Wishart on the other. Again it is
important to note that these are constraints on $q$ and not part of the
underlying model specification $f$.

$\mathcal{N}$$\mathcal{W}$ Figure 9: Notation for enforcing form constraints
on nodes.

Two kinds of form constraints warrant extra attention: $\delta$-constraints
and moment matching. We will now deal with these in turn.

###  5.3 $\delta$-constraints and data points

$\delta$-constraints are the most commonly used form constraints because they
allow us to incorporate data points into a model. A $\delta$-constraint on an
edge defines the function $g(s_{i})$ in Eq. (28) to be

$\displaystyle g(s_{i})$ $\displaystyle=\delta(s_{i}-\hat{s}_{i})$ (30)

What makes the $\delta$-constraint special is that $\hat{s}_{i}$ can either be
a known value or a parameter to optimise [26]. In the case where $\hat{s}_{i}$
is known, it commonly corresponds to a data point. We will refer to this case
as a data constraint and denote it with a filled circle as shown in Fig 10

$\delta$$\delta$ Figure 10: Terminating and non-terminating notation for data
constraints.

Data constraints are special because they denote observations. They also block
any information flow across the edge in question [26]. Because they block
information flow, CFFG notation optionally allows data constraints to
terminate edges.

Here we wish to raise a subtle point about prior FFG notation. Previous work
has used small black _squares_ to denote data constraints following [22]. In
keeping with our convention, a small black square on a CFFG denotes a
$\delta$-distributed variable in the model $f$ rather than the variational
distribution $q$. Being able to differentiate data constrained variables in
$q$ and apriori fixed parameter of the model $f$ allows us to be explicit
about what actually constitutes a data point for the inference problem at hand
[3].

In the case where $\hat{s}_{j}$ is not known, it can be treated as a parameter
to be optimised. We refer to this case as a $\delta$-constraint or a pointmass
constraint and notate it with an unfilled circle as shown in Fig 11

$\delta$ Figure 11: Notation for $\delta$-constraints.

Unlike data constraints, the $\delta$-constraint allows messages to pass and
is therefore not allowed to terminate an edge. Optimising the value of
$\hat{s}_{i}$ under a $\delta$-constraint leads to EM as message passing [26].

###  5.4 Moment matching constraints

Moment matching constraints are special in that they replace the hard
marginalisation constraints of Eq. (8) with constraints of the form

$\displaystyle q(s_{i})=\int q(\bm{s}_{a})\differential s_{a\backslash i}$
$\displaystyle\implies\int
q(\bm{s}_{a})T_{i}(s_{i})\differential\bm{s}_{a}=\int
q(s_{i})T_{i}(s_{i})\differential s_{i}$ (31)

where $T_{i}(s_{i})$ are the sufficient statistics of an exponential family
distribution. This move loosens the marginalisation constraint by instead only
requiring that the moments in question align. When taking the first variation
and solving, one obtains the Expectation Propagation algorithm.

For notational purposes, moment matching constraints are unique in that the
involve both and edge- and a node-marginal. That means the effects are not
localised to a single bead. To indicate which beads are involved, we replace
the solid lines betwen them with dashed lines instead.

We denote moment matching constraints by an $\mathbb{E}$ inside the
corresponding edge-bead as shown in Fig 12. Choosing the edge-bead over the
node-bead is an arbitrary decision made mainly for convenience.

$\mathbb{E}$$\mathbb{E}$ Figure 12: Notation for moment matching with a single
sided (left) and double sided (right) node/edge pairs.

The left side of Fig. 12 shows notation for constraining a single node/edge
pair by moment matching. If both nodes connected to an edge are under moment
matching constraints, the double sided notation on the right of Fig. 12
applies.

Given the modular nature of CFFG notation it is easy to compose different
local constraints to accurate specify a Lagrangian and by extension an
inference problem. Adding custom marginal constraints to a CFFG is also
straightforward as it simply requires defining the meaning of a symbol inside
a bead.

###  5.5 P-substitution on CFFG’s

The final piece needed to represent the Active Inference Lagrangian on a CFFG
is P-substitution. Being able to represent LAIF on a CFFG is the reason for
constructing the local GFE using a mean field factorisation and
P-substitution. This construction is much more amenable to the tools we have
developed so far as we will now demonstrate.

Recall that P-substitution involves substituting part of the model $p$ for $q$
in the expectation only. To write P-substitution on a CFFG, the logical
notation is therefore to replace a circle with a square. Fig. 13 shows an
example of adding a P-substitution to a mean field factorised node marginal

$f_{a}$$x$$y$$z$$f_{a}$$x$$y$$z$ Figure 13: P-substitution on a CFFG with
naive mean field factorisation.

The square notation for P-substitution on CFFG’s imply a conditioning of the
P-substituted variable on all other connected variables that are not
P-substituted. For example, in Fig. 13 the P-substitution changes local VFE to
a GFE by

$\displaystyle\iiint{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}q(y)}q(x)q(z)\log\frac{q(y)q(x)q(z)}{p(y,x,z)}\differential
y\differential x\differential z$
$\displaystyle\quad\quad\implies\iiint{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}p(y\mid
x,z)}q(x)q(z)\log\frac{q(y)q(x)q(z)}{p(y,x,z)}\differential y\differential
x\differential z$ (32)

Here the P-substituted variable is $y$ and the remainder are $\\{x,z\\}$.

###  5.6 CFFG compression

The value of CFFG notation is measured by how much it aids other researchers
and practitioners in expressing their ideas accurately and succinctly. We
envision two main groups for whom CFFG’s might be of particular interest. The
first group is comprised of mathematical researchers working on constrained
free energy optimisation on FFG’s. For this group we expect that the notation
developed so far will be both useful and practically applicable since work is
often focused on the intricacies of performing local optimisation. Commonly an
FFG in this tradition is small but a very high level of accuracy is desired in
order to be mathematically rigorous.

However there is a second group composed of applied researchers for whom the
challenge is to accurately specify a larger inference problem and its solution
in order to solve an auxiliary goal - for instance controlling a drone,
transmitting a coded message or simulating some phenomenon using AIF. For this
group of researchers in particular, CFFG notation as described so far might be
too verbose and the overhead of using it may not outweigh the benefits gained.
To this end we now complete CFFG notation by a mandatory compression step.

The compression step is designed to remove redundant information by enforcing
an emphasis on _deviations from a default BFE_. By default, we mean no
constraints other than normalisation and marginalisation and with a joint
factorisation around every node. Recall that default normalisation constraints
are denoted by empty, round beads and marginalisation by connected lines. A
joint factorisation means all incoming edges are connected and the node only
has a single, internal bead.

To provide a recipe for compressing a CFFG, we will need the concept of a
_bead chain_. A bead chain is simply a series of beads connected by edges. In
the following recipe, a bead will only be summarized as part of a chain if it
contains no additional information, meaning if it is round and empty. To
compress a CFFG, we follow a series of four steps:

  1. 1.

Summarize every bead chain by their terminating beads.

  2. 2.

For nodes with no factorisation constraints, remove empty, internal beads.

  3. 3.

For nodes with no factorisation constraints, remove all internal edges.

  4. 4.

For all factor nodes, push remaining internal beads to the border of the
corresponding factor node.

After performing these steps, we are left with a compressed version of the
original CFFG where each node can be much smaller and more concise. To
exemplify, we apply the recipe to the CFFG in Fig. 14.

$f_{a}$$f_{b}$$f_{c}$$f_{d}$
$f_{e}$$\delta$$=$$f_{f}$$f_{g}$$\mathbb{E}$$\delta$ Figure 14: Initial CFFG
before compression.

Fig. 14 is a complicated graph with loops, dangling edges and multiple
different factorisations in play. As a result, Fig. 14 covers a lot of special
cases that one might encounter on a CFFG. We will now apply the steps in
sequence, starting by removing empty beads on chains. This removes most of the
beads in the inner loop as shown in Fig. 15

$f_{a}$$f_{b}$$f_{c}$$f_{d}$
$f_{e}$$\delta$$=$$f_{f}$$f_{g}$$\mathbb{E}$$\delta$$f_{a}$$f_{b}$$f_{c}$$f_{d}$
$f_{e}$$\delta$$=$$f_{f}$$f_{g}$$\mathbb{E}$$\delta$

Figure 15: Step 1: Removing empty beads from chains. Affected beads are
highlighted in  red.

Note that the dangling edge extending from $f_{a}$ is treated as terminated by
a bead and the bead on the edge is removed. Next we remove any internal beads
for nodes with no factorisation constraints. We show this step in Fig. 16

$f_{a}$$f_{b}$$f_{c}$$f_{d}$
$f_{e}$$\delta$$=$$f_{f}$$f_{g}$$\mathbb{E}$$\delta$$f_{a}$$f_{b}$$f_{c}$$f_{d}$
$f_{e}$$\delta$$=$$f_{f}$$f_{g}$$\mathbb{E}$$\delta$

Figure 16: Step 2. Removing beads on nodes with default factorisation.
Affected nodes are highlighted in  red.

The next step is to remove any internal edge extensions for nodes with no
factorisation constraints. We demonstrate this step in Fig. 17. Note that the
internal edge of the node $f_{d}$ does not get cancelled since one of the
connected edges is a pointmass in the model and therefore not present in $q$,
meaning $f_{d}$ does not have a default factorisation.

$f_{a}$$f_{b}$$f_{c}$$f_{d}$
$f_{e}$$\delta$$=$$f_{f}$$f_{g}$$\mathbb{E}$$\delta$$f_{a}$$f_{b}$$f_{c}$$f_{d}$
$f_{e}$$\delta$$=$$f_{f}$$f_{g}$$\mathbb{E}$$\delta$

Figure 17: Step 3. Removing internal edges. Affected edges are highlighted in
red.

Finally we can push the any remaining internal beads to their node border,
illustrated in Fig. 18. This step allows for writing the CFFG much more
compactly, as we can see in Fig. 19.

$f_{a}$$f_{b}$$f_{c}$$f_{d}$
$f_{e}$$\delta$$=$$f_{f}$$f_{g}$$\mathbb{E}$$\delta$$f_{a}$$f_{b}$$f_{c}$$f_{d}$
$f_{e}$$\delta$$=$$f_{f}$$f_{g}$$\mathbb{E}$$\delta$

Figure 18: Step 4. Pushing beads to node borders. Affected beads are
highlighted in  red.

At this point we have removed most of the beads and edges and still retain
most of the relevant information around all nodes. What is left is only what
deviates from a default BFE specification.

The goal here is as stated initially to make it easier to work with CFFG’s for
larger models which necessitates working with smaller nodes.

$f_{a}$$f_{b}$$f_{c}$$f_{d}$ $f_{e}$$\delta$ $=$
$f_{f}$$f_{g}$$\mathbb{E}$$\delta$$f_{a}$$f_{b}$$f_{c}$$f_{d}$
$f_{e}$$\delta$$=$$f_{f}$$f_{g}$$\mathbb{E}$$\delta$

Figure 19: compressed allows for much more compact CFFG’s which in turn are
better suited for larger inference problems.

From Fig. 19 it is clear that using compressed allows for much more compact
CFFG’s which in turn are more amenable to larger CFFG’s. This is especially
true in the absence of constraints and P-substitutions where the CFFG and the
underlying FFG are identical since compressed only denotes deviations from the
default. We see this with nodes $f_{c}$, $f_{f}$, $f_{g}$ and $=$ in Fig. 19.

##  6 Classical AIF and the original GFE algorithm as special cases of SAIF

An integral part of working with message passing algorithms is the choice of
schedule - the order in which messages are passed. Many message passing
algorithms are iterative in nature and can therefore be sensitive to
scheduling. SAIF is also an iterative message passing algorithm and might
consequently be sensitive to the choice of schedule.

A particularly interesting observation is that we can recover both the
classical AIF planning algorithm of [7] and the scheme proposed by [20] as
special cases of SAIF by carefully choosing the schedule and performing model
comparison. This result extends upon prior work by [14] who reinterpreted the
classical algorithm through the lens of message passing on an FFG.

###  6.1 Example: CFFG and message updates for a discrete observation model
with goals

$T$$\mathcal{C}at$$\bm{\mathrm{x}}$$\bm{\mathrm{z}}$$\bm{\mathrm{A}}$$\bm{\mathrm{c}}$
Figure 20: CFFG of composite node for SAIF on discrete state spaces,
reproduced from [16].

In order to both demonstrate SAIF and recover prior work, we need to derive
the required messages. We refer to our companion paper [16] for the detailed
derivations and summarise the results here. For the remaining paragraphs, we
will use $\bm{\mathrm{x}}$ to denote that $x$ is vector-valued and
$\bm{\mathrm{A}}$ to denote that $A$ is either a matrix or a tensor. For the
discrete case we work with a composite node corresponding to the factor

$\displaystyle p(\bm{\mathrm{x}}\mid\bm{\mathrm{A}},\bm{\mathrm{z}})$
$\displaystyle=\mathcal{C}at(\bm{\mathrm{x}}\mid\bm{\mathrm{A}}\bm{\mathrm{z}})$
(33a) $\displaystyle\tilde{p}(\bm{\mathrm{x}}\mid\bm{\mathrm{c}})$
$\displaystyle=\mathcal{C}at(\bm{\mathrm{x}}\mid\bm{\mathrm{c}})\,.$ (33b)

We show the corresponding CFFG in Fig. 20 where we indicate the necessary
factorisation and P-substitution. In the parlance of AIF Eq. 33a defines a
model composed of a discrete state transition ($T$) with transition matrix
$\bm{\mathrm{A}}$ and a categorical goal prior ($\mathcal{C}at$) with
parameter vector $\bm{\mathrm{c}}$. Before stating the messages, we define the
vector

$\displaystyle\bm{\mathrm{h}}(\bm{\mathrm{A}})$
$\displaystyle=-\mathrm{diag}(\bm{\mathrm{A}}^{T}\log\bm{\mathrm{A}})$ (34)

This term is often denoted H in other works [5, 20]. We instead opt for using
lowercase since the term is a vector and for making the dependence on
$\bm{\mathrm{A}}$ explicit. In the following section, we will sometimes use an
$\overline{\mathrm{overbar}}$ to denote expectations, such that
$\mathbb{E}_{q(x)}[g(x)]=\overline{g(x)}$. Additionally, we define

$\displaystyle\bm{\mathrm{\upxi}}(\bm{\mathrm{A}})$
$\displaystyle=\bm{\mathrm{A}}^{\operatorname{T}}\\!\left(\overline{\log\bm{\mathrm{c}}}-\log(\overline{\bm{\mathrm{A}}}\overline{\bm{\mathrm{z}}})\right)-\bm{\mathrm{h}}(\bm{\mathrm{A}})$
(35a) $\displaystyle\bm{\mathrm{\uprho}}$
$\displaystyle=\overline{\bm{\mathrm{A}}}^{\operatorname{T}}\\!\left(\overline{\log\bm{\mathrm{c}}}-\log(\overline{\bm{\mathrm{A}}}\overline{\bm{\mathrm{z}}})\right)-\overline{\bm{\mathrm{h}}(\bm{\mathrm{A}})}\,.$
(35b)

in order to make the expressions more concise. Where expectations cannot be
computed in closed form, we instead resort to Monte Carlo estimates. With this
notation in place, we can write the required messages as

$\mu_{\leavevmode\hbox to7.91pt{\vbox to7.91pt{\pgfpicture\makeatletter\hbox{\hskip 3.95439pt\lower-3.95439pt\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{ }\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{ }\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{ }\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{ }\nullfont\pgfsys@beginscope\pgfsys@invoke{ }\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{ }{ {{}}\hbox{\hbox{{\pgfsys@beginscope\pgfsys@invoke{ }\definecolor[named]{pgffillcolor}{rgb}{0.25390625,0.25390625,0.25390625}\pgfsys@color@rgb@fill{0.25390625}{0.25390625}{0.25390625}\pgfsys@invoke{ }{{}{{{}}}{{}}{}{}{}{}{}{}{}{}{}{\pgfsys@beginscope\pgfsys@invoke{ }\definecolor[named]{pgffillcolor}{rgb}{0.25390625,0.25390625,0.25390625}\pgfsys@color@rgb@fill{0.25390625}{0.25390625}{0.25390625}\pgfsys@invoke{ }{}\pgfsys@moveto{3.7544pt}{0.0pt}\pgfsys@curveto{3.7544pt}{2.07352pt}{2.07352pt}{3.7544pt}{0.0pt}{3.7544pt}\pgfsys@curveto{-2.07352pt}{3.7544pt}{-3.7544pt}{2.07352pt}{-3.7544pt}{0.0pt}\pgfsys@curveto{-3.7544pt}{-2.07352pt}{-2.07352pt}{-3.7544pt}{0.0pt}{-3.7544pt}\pgfsys@curveto{2.07352pt}{-3.7544pt}{3.7544pt}{-2.07352pt}{3.7544pt}{0.0pt}\pgfsys@closepath\pgfsys@moveto{0.0pt}{0.0pt}\pgfsys@fillstroke\pgfsys@invoke{ } \pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}{{{{}}\pgfsys@beginscope\pgfsys@invoke{ }\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{-2.84528pt}{-2.45pt}\pgfsys@invoke{ }\hbox{{\definecolor[named]{.}{rgb}{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@invoke{ }\pgfsys@color@gray@fill{1}\pgfsys@invoke{ }\hbox{\minipage[t]{5.69055pt}\leavevmode\hbox{}{\definecolor[named]{.}{rgb}{1,1,1}\color[rgb]{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@color@gray@fill{1}\tiny{1}}\endminipage} }}\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}} \pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}} } \pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope{{{}}}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope\hss}}\lxSVG@closescope\endpgfpicture}}}(\bm{\mathrm{c}})$ | $\propto\mathcal{D}ir(\bm{\mathrm{c}}|\overline{\bm{\mathrm{A}}}\overline{\bm{\mathrm{z}}}+\bm{\mathrm{1}})$
---|---
$\mu_{\leavevmode\hbox to7.91pt{\vbox to7.91pt{\pgfpicture\makeatletter\hbox{\hskip 3.95439pt\lower-3.95439pt\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{ }\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{ }\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{ }\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{ }\nullfont\pgfsys@beginscope\pgfsys@invoke{ }\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{ }{ {{}}\hbox{\hbox{{\pgfsys@beginscope\pgfsys@invoke{ }\definecolor[named]{pgffillcolor}{rgb}{0.25390625,0.25390625,0.25390625}\pgfsys@color@rgb@fill{0.25390625}{0.25390625}{0.25390625}\pgfsys@invoke{ }{{}{{{}}}{{}}{}{}{}{}{}{}{}{}{}{\pgfsys@beginscope\pgfsys@invoke{ }\definecolor[named]{pgffillcolor}{rgb}{0.25390625,0.25390625,0.25390625}\pgfsys@color@rgb@fill{0.25390625}{0.25390625}{0.25390625}\pgfsys@invoke{ }{}\pgfsys@moveto{3.7544pt}{0.0pt}\pgfsys@curveto{3.7544pt}{2.07352pt}{2.07352pt}{3.7544pt}{0.0pt}{3.7544pt}\pgfsys@curveto{-2.07352pt}{3.7544pt}{-3.7544pt}{2.07352pt}{-3.7544pt}{0.0pt}\pgfsys@curveto{-3.7544pt}{-2.07352pt}{-2.07352pt}{-3.7544pt}{0.0pt}{-3.7544pt}\pgfsys@curveto{2.07352pt}{-3.7544pt}{3.7544pt}{-2.07352pt}{3.7544pt}{0.0pt}\pgfsys@closepath\pgfsys@moveto{0.0pt}{0.0pt}\pgfsys@fillstroke\pgfsys@invoke{ } \pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}{{{{}}\pgfsys@beginscope\pgfsys@invoke{ }\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{-2.84528pt}{-2.45pt}\pgfsys@invoke{ }\hbox{{\definecolor[named]{.}{rgb}{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@invoke{ }\pgfsys@color@gray@fill{1}\pgfsys@invoke{ }\hbox{\minipage[t]{5.69055pt}\leavevmode\hbox{}{\definecolor[named]{.}{rgb}{1,1,1}\color[rgb]{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@color@gray@fill{1}\tiny{2}}\endminipage} }}\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}} \pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}} } \pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope{{{}}}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope\hss}}\lxSVG@closescope\endpgfpicture}}}(\bm{\mathrm{z}})$ | $\propto\mathcal{C}at(\bm{\mathrm{z}}|\sigma(\bm{\mathrm{\uprho}}))$
Solve: $\overline{\bm{\mathrm{z}}}$ | $\stackrel{{\scriptstyle!}}{{=}}\sigma\\!\left(\bm{\mathrm{\uprho}}(\overline{\bm{\mathrm{z}}})+\log\bm{\mathrm{d}}\right)$
$\mu_{\leavevmode\hbox to7.91pt{\vbox to7.91pt{\pgfpicture\makeatletter\hbox{\hskip 3.95439pt\lower-3.95439pt\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{ }\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{ }\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{ }\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{ }\nullfont\pgfsys@beginscope\pgfsys@invoke{ }\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{ }{ {{}}\hbox{\hbox{{\pgfsys@beginscope\pgfsys@invoke{ }\definecolor[named]{pgffillcolor}{rgb}{0.25390625,0.25390625,0.25390625}\pgfsys@color@rgb@fill{0.25390625}{0.25390625}{0.25390625}\pgfsys@invoke{ }{{}{{{}}}{{}}{}{}{}{}{}{}{}{}{}{\pgfsys@beginscope\pgfsys@invoke{ }\definecolor[named]{pgffillcolor}{rgb}{0.25390625,0.25390625,0.25390625}\pgfsys@color@rgb@fill{0.25390625}{0.25390625}{0.25390625}\pgfsys@invoke{ }{}\pgfsys@moveto{3.7544pt}{0.0pt}\pgfsys@curveto{3.7544pt}{2.07352pt}{2.07352pt}{3.7544pt}{0.0pt}{3.7544pt}\pgfsys@curveto{-2.07352pt}{3.7544pt}{-3.7544pt}{2.07352pt}{-3.7544pt}{0.0pt}\pgfsys@curveto{-3.7544pt}{-2.07352pt}{-2.07352pt}{-3.7544pt}{0.0pt}{-3.7544pt}\pgfsys@curveto{2.07352pt}{-3.7544pt}{3.7544pt}{-2.07352pt}{3.7544pt}{0.0pt}\pgfsys@closepath\pgfsys@moveto{0.0pt}{0.0pt}\pgfsys@fillstroke\pgfsys@invoke{ } \pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}{{{{}}\pgfsys@beginscope\pgfsys@invoke{ }\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{-2.84528pt}{-2.45pt}\pgfsys@invoke{ }\hbox{{\definecolor[named]{.}{rgb}{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@invoke{ }\pgfsys@color@gray@fill{1}\pgfsys@invoke{ }\hbox{\minipage[t]{5.69055pt}\leavevmode\hbox{}{\definecolor[named]{.}{rgb}{1,1,1}\color[rgb]{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@color@gray@fill{1}\tiny{2}}\endminipage} }}\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}} \pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}} } \pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope{{{}}}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope\hss}}\lxSVG@closescope\endpgfpicture}}}(\bm{\mathrm{z}})$ | $\propto\mathcal{C}at\\!\left({\bm{\mathrm{z}}|\sigma(\log\overline{\bm{\mathrm{z}}}^{*}-\log\bm{\mathrm{d}})}\right)$
$\log\mu_{\leavevmode\hbox to7.91pt{\vbox to7.91pt{\pgfpicture\makeatletter\hbox{\hskip 3.95439pt\lower-3.95439pt\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{ }\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{ }\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{ }\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{ }\nullfont\pgfsys@beginscope\pgfsys@invoke{ }\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{ }{ {{}}\hbox{\hbox{{\pgfsys@beginscope\pgfsys@invoke{ }\definecolor[named]{pgffillcolor}{rgb}{0.25390625,0.25390625,0.25390625}\pgfsys@color@rgb@fill{0.25390625}{0.25390625}{0.25390625}\pgfsys@invoke{ }{{}{{{}}}{{}}{}{}{}{}{}{}{}{}{}{\pgfsys@beginscope\pgfsys@invoke{ }\definecolor[named]{pgffillcolor}{rgb}{0.25390625,0.25390625,0.25390625}\pgfsys@color@rgb@fill{0.25390625}{0.25390625}{0.25390625}\pgfsys@invoke{ }{}\pgfsys@moveto{3.7544pt}{0.0pt}\pgfsys@curveto{3.7544pt}{2.07352pt}{2.07352pt}{3.7544pt}{0.0pt}{3.7544pt}\pgfsys@curveto{-2.07352pt}{3.7544pt}{-3.7544pt}{2.07352pt}{-3.7544pt}{0.0pt}\pgfsys@curveto{-3.7544pt}{-2.07352pt}{-2.07352pt}{-3.7544pt}{0.0pt}{-3.7544pt}\pgfsys@curveto{2.07352pt}{-3.7544pt}{3.7544pt}{-2.07352pt}{3.7544pt}{0.0pt}\pgfsys@closepath\pgfsys@moveto{0.0pt}{0.0pt}\pgfsys@fillstroke\pgfsys@invoke{ } \pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}{{{{}}\pgfsys@beginscope\pgfsys@invoke{ }\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{-2.84528pt}{-2.45pt}\pgfsys@invoke{ }\hbox{{\definecolor[named]{.}{rgb}{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@invoke{ }\pgfsys@color@gray@fill{1}\pgfsys@invoke{ }\hbox{\minipage[t]{5.69055pt}\leavevmode\hbox{}{\definecolor[named]{.}{rgb}{1,1,1}\color[rgb]{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@color@gray@fill{1}\tiny{3}}\endminipage} }}\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}} \pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}} } \pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope{{{}}}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope\hss}}\lxSVG@closescope\endpgfpicture}}}(\bm{\mathrm{A}})$ | $=\overline{\bm{\mathrm{z}}}^{\operatorname{T}}\bm{\mathrm{\upxi}}(\bm{\mathrm{A}})$
$U_{\bm{\mathrm{x}}}$ | $=-\overline{\bm{\mathrm{z}}}^{\operatorname{T}}\bm{\mathrm{\uprho}}$
Figure 21: Message updates for the discrete composite node, reproduced from
[16].

Fig. 21 shows the message updates towards all connected variables as well as
as the average energy term ($U_{x}$) for the composite node. $\bm{\mathrm{d}}$
denotes the parameters of the incoming message from the rest of the CFFG on
the edge $z$. Interestingly, the energy term $U_{x}$ corresponds exactly to
the expected free energy as used in standard AIF [5, 7].

Special attention needs to be given to the messages $\mu_{\leavevmode\hbox
to7.91pt{\vbox to7.91pt{\pgfpicture\makeatletter\hbox{\hskip
3.95439pt\lower-3.95439pt\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{
}\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{
}\nullfont\pgfsys@beginscope\pgfsys@invoke{ }\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{ }{
{{}}\hbox{\hbox{{\pgfsys@beginscope\pgfsys@invoke{
}\definecolor[named]{pgffillcolor}{rgb}{0.25390625,0.25390625,0.25390625}\pgfsys@color@rgb@fill{0.25390625}{0.25390625}{0.25390625}\pgfsys@invoke{
}{{}{{{}}}{{}}{}{}{}{}{}{}{}{}{}{\pgfsys@beginscope\pgfsys@invoke{
}\definecolor[named]{pgffillcolor}{rgb}{0.25390625,0.25390625,0.25390625}\pgfsys@color@rgb@fill{0.25390625}{0.25390625}{0.25390625}\pgfsys@invoke{
}{}\pgfsys@moveto{3.7544pt}{0.0pt}\pgfsys@curveto{3.7544pt}{2.07352pt}{2.07352pt}{3.7544pt}{0.0pt}{3.7544pt}\pgfsys@curveto{-2.07352pt}{3.7544pt}{-3.7544pt}{2.07352pt}{-3.7544pt}{0.0pt}\pgfsys@curveto{-3.7544pt}{-2.07352pt}{-2.07352pt}{-3.7544pt}{0.0pt}{-3.7544pt}\pgfsys@curveto{2.07352pt}{-3.7544pt}{3.7544pt}{-2.07352pt}{3.7544pt}{0.0pt}\pgfsys@closepath\pgfsys@moveto{0.0pt}{0.0pt}\pgfsys@fillstroke\pgfsys@invoke{
} \pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope}{{{{}}\pgfsys@beginscope\pgfsys@invoke{
}\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{-2.84528pt}{-2.45pt}\pgfsys@invoke{
}\hbox{{\definecolor[named]{.}{rgb}{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@invoke{
}\pgfsys@color@gray@fill{1}\pgfsys@invoke{
}\hbox{\minipage[t]{5.69055pt}\leavevmode\hbox{}{\definecolor[named]{.}{rgb}{1,1,1}\color[rgb]{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@color@gray@fill{1}\tiny{2}}\endminipage}
}}\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}}
\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}} }
\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope{{{}}}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope\hss}}\lxSVG@closescope\endpgfpicture}}}(\bm{\mathrm{z}})$
and $\mu_{\leavevmode\hbox to7.91pt{\vbox
to7.91pt{\pgfpicture\makeatletter\hbox{\hskip 3.95439pt\lower-3.95439pt\hbox
to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{
}\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{
}\nullfont\pgfsys@beginscope\pgfsys@invoke{ }\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{ }{
{{}}\hbox{\hbox{{\pgfsys@beginscope\pgfsys@invoke{
}\definecolor[named]{pgffillcolor}{rgb}{0.25390625,0.25390625,0.25390625}\pgfsys@color@rgb@fill{0.25390625}{0.25390625}{0.25390625}\pgfsys@invoke{
}{{}{{{}}}{{}}{}{}{}{}{}{}{}{}{}{\pgfsys@beginscope\pgfsys@invoke{
}\definecolor[named]{pgffillcolor}{rgb}{0.25390625,0.25390625,0.25390625}\pgfsys@color@rgb@fill{0.25390625}{0.25390625}{0.25390625}\pgfsys@invoke{
}{}\pgfsys@moveto{3.7544pt}{0.0pt}\pgfsys@curveto{3.7544pt}{2.07352pt}{2.07352pt}{3.7544pt}{0.0pt}{3.7544pt}\pgfsys@curveto{-2.07352pt}{3.7544pt}{-3.7544pt}{2.07352pt}{-3.7544pt}{0.0pt}\pgfsys@curveto{-3.7544pt}{-2.07352pt}{-2.07352pt}{-3.7544pt}{0.0pt}{-3.7544pt}\pgfsys@curveto{2.07352pt}{-3.7544pt}{3.7544pt}{-2.07352pt}{3.7544pt}{0.0pt}\pgfsys@closepath\pgfsys@moveto{0.0pt}{0.0pt}\pgfsys@fillstroke\pgfsys@invoke{
} \pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope}{{{{}}\pgfsys@beginscope\pgfsys@invoke{
}\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{-2.84528pt}{-2.45pt}\pgfsys@invoke{
}\hbox{{\definecolor[named]{.}{rgb}{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@invoke{
}\pgfsys@color@gray@fill{1}\pgfsys@invoke{
}\hbox{\minipage[t]{5.69055pt}\leavevmode\hbox{}{\definecolor[named]{.}{rgb}{1,1,1}\color[rgb]{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@color@gray@fill{1}\tiny{3}}\endminipage}
}}\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}}
\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}} }
\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope{{{}}}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope\hss}}\lxSVG@closescope\endpgfpicture}}}(\bm{\mathrm{A}})$.
While $\mu_{\leavevmode\hbox to7.91pt{\vbox
to7.91pt{\pgfpicture\makeatletter\hbox{\hskip 3.95439pt\lower-3.95439pt\hbox
to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{
}\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{
}\nullfont\pgfsys@beginscope\pgfsys@invoke{ }\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{ }{
{{}}\hbox{\hbox{{\pgfsys@beginscope\pgfsys@invoke{
}\definecolor[named]{pgffillcolor}{rgb}{0.25390625,0.25390625,0.25390625}\pgfsys@color@rgb@fill{0.25390625}{0.25390625}{0.25390625}\pgfsys@invoke{
}{{}{{{}}}{{}}{}{}{}{}{}{}{}{}{}{\pgfsys@beginscope\pgfsys@invoke{
}\definecolor[named]{pgffillcolor}{rgb}{0.25390625,0.25390625,0.25390625}\pgfsys@color@rgb@fill{0.25390625}{0.25390625}{0.25390625}\pgfsys@invoke{
}{}\pgfsys@moveto{3.7544pt}{0.0pt}\pgfsys@curveto{3.7544pt}{2.07352pt}{2.07352pt}{3.7544pt}{0.0pt}{3.7544pt}\pgfsys@curveto{-2.07352pt}{3.7544pt}{-3.7544pt}{2.07352pt}{-3.7544pt}{0.0pt}\pgfsys@curveto{-3.7544pt}{-2.07352pt}{-2.07352pt}{-3.7544pt}{0.0pt}{-3.7544pt}\pgfsys@curveto{2.07352pt}{-3.7544pt}{3.7544pt}{-2.07352pt}{3.7544pt}{0.0pt}\pgfsys@closepath\pgfsys@moveto{0.0pt}{0.0pt}\pgfsys@fillstroke\pgfsys@invoke{
} \pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope}{{{{}}\pgfsys@beginscope\pgfsys@invoke{
}\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{-2.84528pt}{-2.45pt}\pgfsys@invoke{
}\hbox{{\definecolor[named]{.}{rgb}{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@invoke{
}\pgfsys@color@gray@fill{1}\pgfsys@invoke{
}\hbox{\minipage[t]{5.69055pt}\leavevmode\hbox{}{\definecolor[named]{.}{rgb}{1,1,1}\color[rgb]{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@color@gray@fill{1}\tiny{2}}\endminipage}
}}\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}}
\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}} }
\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope{{{}}}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope\hss}}\lxSVG@closescope\endpgfpicture}}}(\bm{\mathrm{z}})$
can be solved for in closed form, in practice applying this result directly
can lead to unstable solutions that fluctuate between multiple extrema. For
this reason we opt for parameterising the message by
$\overline{\bm{\mathrm{z}}}$ and solving for parameters of the marginal
directly using Newtons method. Having found the optimum
$\overline{\bm{\mathrm{z}}}^{*}$ we can then substitute the result into the
message expression in order to obtain a stable solution.

The message $\mu_{\leavevmode\hbox to7.91pt{\vbox
to7.91pt{\pgfpicture\makeatletter\hbox{\hskip 3.95439pt\lower-3.95439pt\hbox
to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{
}\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{
}\nullfont\pgfsys@beginscope\pgfsys@invoke{ }\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{ }{
{{}}\hbox{\hbox{{\pgfsys@beginscope\pgfsys@invoke{
}\definecolor[named]{pgffillcolor}{rgb}{0.25390625,0.25390625,0.25390625}\pgfsys@color@rgb@fill{0.25390625}{0.25390625}{0.25390625}\pgfsys@invoke{
}{{}{{{}}}{{}}{}{}{}{}{}{}{}{}{}{\pgfsys@beginscope\pgfsys@invoke{
}\definecolor[named]{pgffillcolor}{rgb}{0.25390625,0.25390625,0.25390625}\pgfsys@color@rgb@fill{0.25390625}{0.25390625}{0.25390625}\pgfsys@invoke{
}{}\pgfsys@moveto{3.7544pt}{0.0pt}\pgfsys@curveto{3.7544pt}{2.07352pt}{2.07352pt}{3.7544pt}{0.0pt}{3.7544pt}\pgfsys@curveto{-2.07352pt}{3.7544pt}{-3.7544pt}{2.07352pt}{-3.7544pt}{0.0pt}\pgfsys@curveto{-3.7544pt}{-2.07352pt}{-2.07352pt}{-3.7544pt}{0.0pt}{-3.7544pt}\pgfsys@curveto{2.07352pt}{-3.7544pt}{3.7544pt}{-2.07352pt}{3.7544pt}{0.0pt}\pgfsys@closepath\pgfsys@moveto{0.0pt}{0.0pt}\pgfsys@fillstroke\pgfsys@invoke{
} \pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope}{{{{}}\pgfsys@beginscope\pgfsys@invoke{
}\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{-2.84528pt}{-2.45pt}\pgfsys@invoke{
}\hbox{{\definecolor[named]{.}{rgb}{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@invoke{
}\pgfsys@color@gray@fill{1}\pgfsys@invoke{
}\hbox{\minipage[t]{5.69055pt}\leavevmode\hbox{}{\definecolor[named]{.}{rgb}{1,1,1}\color[rgb]{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@color@gray@fill{1}\tiny{3}}\endminipage}
}}\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}}
\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}} }
\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope{{{}}}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope\hss}}\lxSVG@closescope\endpgfpicture}}}(\bm{\mathrm{A}})$
does not follow a nice exponential family distribution. To circumvent this
problem, we can pass on the logpdf directly. When we need to compute the
marginal $q(\bm{\mathrm{A}})$, we can then use sampling procedures to estimate
the necessary expectations [1]. Further details and full derivations can be
found in our companion paper [16].

###  6.2 Reconstructing classical AIF

With our new messages in hand and equipped with CFFG notation, we can now
restate prior work unambiguously, starting with the classical algorithm of
[7]. To reconstruct the algorithm of [7], we start by defining the generative
model. The generative model is a discrete partially observed Markov decision
process (POMDP) over future time steps given by

$\displaystyle
p(\underbrace{\bm{\mathrm{x}}_{t+1:T},\bm{\mathrm{z}}_{t:T}}_{\text{Future}}$
$\displaystyle\mid\underbrace{\bm{\mathrm{\hat{u}}}_{t+1:T}}_{\text{Policy}},\underbrace{\bm{\mathrm{x}}_{1:t},\bm{\mathrm{u}}_{1:t}}_{\text{Past}})$
$\displaystyle\propto\underbrace{p(\bm{\mathrm{z}}_{t}\mid\bm{\mathrm{x}}_{1:t},\bm{\mathrm{u}}_{1:t})}_{\text{State
Prior}}\prod_{k=t+1}^{T}\underbrace{p(\bm{\mathrm{x}}_{k}\mid\bm{\mathrm{z}}_{k})}_{\text{Likelihood}}\underbrace{p(\bm{\mathrm{z}}_{k}\mid\bm{\mathrm{z}}_{k-1},\bm{\mathrm{\hat{u}}}_{k})}_{\text{State
Transition}}\underbrace{\tilde{p}(\bm{\mathrm{x}}_{k})}_{\text{Goal prior}}$
(36)

where

$\displaystyle
p(\bm{\mathrm{z}}_{t}\mid\bm{\mathrm{x}}_{1:t},\bm{\mathrm{u}}_{1:t})$
$\displaystyle=\mathcal{C}at(\bm{\mathrm{z}}_{t}\mid\bm{\mathrm{d}})$ (37a)
$\displaystyle
p(\bm{\mathrm{z}}_{k}\mid\bm{\mathrm{z}}_{k-1},\bm{\mathrm{\hat{u}}}_{k})$
$\displaystyle=\mathcal{C}at(\bm{\mathrm{z}}_{k}\mid\bm{\mathrm{B}}_{\hat{u}_{k}}\bm{\mathrm{z}}_{k-1})$
(37b) $\displaystyle p(\bm{\mathrm{x}}_{k}\mid\bm{\mathrm{z}}_{k})$
$\displaystyle=\mathcal{C}at(\bm{\mathrm{x}}_{k}\mid\bm{\mathrm{A}}\bm{\mathrm{z}}_{k})$
(37c) $\displaystyle\tilde{p}(\bm{\mathrm{x}}_{k})$
$\displaystyle=\mathcal{C}at(\bm{\mathrm{x}}_{k}\mid\bm{\mathrm{c}}_{k})\,.$
(37d)

Here we let $\bm{\mathrm{x}}_{k}$ denote observations, $\bm{\mathrm{z}}_{k}$
latent states and $\bm{\mathrm{\hat{u}}}_{k}$ a fixed control, with the
subscript $k$ indicating the future time step in question. The initial state
$\bm{\mathrm{z}}_{t}$ represents the filtering solution given the agents
trajectory so far which we summarize in the parameter vector
$\bm{\mathrm{d}}$. Control signals correspond to particular transition
matrices $\bm{\mathrm{B}}_{\bm{\mathrm{\hat{u}}}_{k}}$ where we use the
subscript to emphasize that each $\bm{\mathrm{B}}$ matches a particular
control $\bm{\mathrm{\hat{u}}}_{k}$. The observation model is given by the
known matrix $\bm{\mathrm{A}}$. Note that Eq. (6.2) is not normalised as it
includes goal priors over $\bm{\mathrm{x}}$. The CFFG of (6.2) is shown in Fig
22.

T$\mathcal{C}at$ $=$ T T $=$ T $\mathcal{C}at$$\mathcal{C}at$
$\cdots$$\bm{\mathrm{z}}_{t}$$\bm{\mathrm{d}}$$\bm{\mathrm{B}}_{\hat{u}_{t+1}}$$\bm{\mathrm{x}}_{t+1}$$\bm{\mathrm{A}}$$\bm{\mathrm{c}}_{t+1}$$\bm{\mathrm{z}}_{t+1}$$\bm{\mathrm{B}}_{\hat{u}_{t+2}}$$\bm{\mathrm{x}}_{t+2}$$\bm{\mathrm{A}}$$\bm{\mathrm{c}}_{t+2}$$\bm{\mathrm{z}}_{t+2}$
Figure 22: CFFG of discrete POMDP as used for planning in standard AIF models.

where $T$ nodes denote a discrete state transition (multiplication of a
categorical variable by a transition matrix). Given a generative model with a
fixed set of controls, the next step is to compute the EFE [5, 7, 14, 15]. We
provide a brief description here and refer to Appendix A and [5, 14, 15, 23]
for more detailed descriptions. The EFE is given by

$\displaystyle G_{EFE}(\hat{\bm{\mathrm{u}}}_{t+1:T})$
$\displaystyle=\sum_{k=t}^{T}\iint
p(\bm{\mathrm{x}}_{k}\mid\bm{\mathrm{z}}_{k})q(\bm{\mathrm{z}}_{k}\mid\bm{\mathrm{\hat{u}}}_{k})\log\frac{q(\bm{\mathrm{z_{k}}}\mid\bm{\mathrm{\hat{u}}}_{k})}{p(\bm{\mathrm{x}}_{k},\bm{\mathrm{z}}_{k}\mid\bm{\mathrm{\hat{u}}}_{k})}\differential\bm{\mathrm{x}}_{k}\differential\bm{\mathrm{z}}_{k}$
(38)

With a slight abuse of notation, we can compute EFE by first applying
transition matrices to the latent state and generating predicted observations
as

$\displaystyle\bm{\mathrm{z}}_{k}$
$\displaystyle=\bm{\mathrm{B}}_{\hat{u}_{k}}\bm{\mathrm{z}}_{k-1}$ (39a)
$\displaystyle\bm{\mathrm{x}}_{k}$
$\displaystyle=\bm{\mathrm{A}}\bm{\mathrm{z}}_{k}\,.$ (39b)

In Eq. (39) we slightly abuse notation by having $\bm{\mathrm{z}}_{k}$ (resp.
$\bm{\mathrm{x}}_{k}$) refer to the prediction after applying the transition
matrix $\bm{\mathrm{B}}_{\hat{u}_{k}}$ (resp. $\bm{\mathrm{A}}$) instead of
the random variable as we have done elsewhere.

We can recognise these operations as performing a forwards message passing
sweep using belief propagation messages. With this choice of generative model,
the EFE of a policy $\bm{\mathrm{\hat{u}}}_{t+1:T}$ is found by by [5, Eq.
D.2-3].

$\displaystyle G_{EFE}(\hat{\bm{\mathrm{u}}}_{t+1:T})$
$\displaystyle=\sum_{k=t+1}^{T}-\mathrm{diag}(\bm{\mathrm{A}}^{T}\log\bm{\mathrm{A}})^{T}\bm{\mathrm{z}}_{k}+\bm{\mathrm{x}}_{k}^{T}(\log\bm{\mathrm{x}}_{k}-\log\bm{\mathrm{c}}_{k})$
(40)

where $\bm{\mathrm{c}}_{k}$ denotes the parameter vector of a goal prior at
the $k$’th time step. To select a policy we simply pick the sequence
$\bm{\mathrm{\hat{u}}}_{t+1:T}$ that results in the lowest numerical value
when solving Eq. (40).

To write this method on the CFFG in Fig. 22, we can simply add messages to the
CFFG and note that the sum of energy terms of the P-substituted composite
nodes in Fig. 21 matches Eq. (40). We show this result in Fig. 23.

T$\mathcal{C}at$ $=$ T T $=$ T $\mathcal{C}at$$\mathcal{C}at$
$\cdots$$\rightarrow$$\rightarrow$$\downarrow$$\rightarrow$$\downarrow$$\rightarrow$$\rightarrow$$\rightarrow$$\downarrow$$\rightarrow$$\downarrow$$\rightarrow$$\rightarrow$$\rightarrow$
Figure 23: Classical EFE computation on the corresponding CFFG.

Comparing different policies (choices of $\bm{\mathrm{\hat{u}}}_{t+1:T}$) and
computing the energy terms of the P-substituted composite nodes is then
exactly equal to the EFE-computation detailed in [7].

###  6.3 Reconstructing the original GFE method

We can further exemplify the capabilities of CFFG notation by recapitulating
the update rules given in the original GFE paper [20]. To recover the
procedure of [20], we need to extend the generative model to encompass past
observations as

$\displaystyle p(\bm{\mathrm{x}}_{1:T}$
$\displaystyle,\bm{\mathrm{z}}_{0:T}\mid\bm{\mathrm{\hat{u}}}_{1:T})$
$\displaystyle\propto
p(\bm{\mathrm{z}}_{0})\underbrace{\prod_{l=1}^{t}p(\bm{\mathrm{x}}_{l}\mid\bm{\mathrm{z}}_{l})p(\bm{\mathrm{z}}_{l}\mid\bm{\mathrm{z}}_{l-1},\bm{\mathrm{\hat{u}}}_{l})}_{\text{Past}}\underbrace{\prod_{k=t+1}^{T}p(\bm{\mathrm{x}}_{k}\mid\bm{\mathrm{z}}_{k})p(\bm{\mathrm{z}}_{k}\mid\bm{\mathrm{z}}_{k-1},\bm{\mathrm{\hat{u}}}_{k})\overbrace{\tilde{p}(\bm{\mathrm{x}}_{k})}^{\text{Goal
prior}}}_{\text{Future}}$

where

$\displaystyle p(\bm{\mathrm{z}}_{0})$
$\displaystyle=\mathcal{C}at(\bm{\mathrm{z}}_{0}\mid\bm{\mathrm{d}})$ (41a)
$\displaystyle
p(\bm{\mathrm{z}}_{t}\mid\bm{\mathrm{z}}_{t-1},\bm{\mathrm{\hat{u}}}_{t})$
$\displaystyle=\mathcal{C}at(\bm{\mathrm{z}}_{t}\mid\bm{\mathrm{B}}_{\bm{\mathrm{\hat{u}}}_{t}}\bm{\mathrm{z}}_{t-1})$
(41b) $\displaystyle p(\bm{\mathrm{x}}_{t}\mid\bm{\mathrm{z}}_{t})$
$\displaystyle=\mathcal{C}at(\bm{\mathrm{x}}_{t}\mid\bm{\mathrm{A}}\bm{\mathrm{z}}_{t})$
(41c) $\displaystyle\tilde{p}(\bm{\mathrm{x}}_{k})$
$\displaystyle=\mathcal{C}at(\bm{\mathrm{x}}_{k}\mid\bm{\mathrm{c}}_{k})\,.$
(41d)

For past time steps we add data constraints and for future time steps we
perform P-substitution. We also apply a naive mean field factorisation for
every node. The final ingredient we need is a schedule that includes both
forwards and backwards passes as hinted at in [14]. For $t=1$ the
corresponding CFFG is shown in Fig. 24 with messages out of the P-substituted
nodes highlighted in red. These messages are given by $\mu_{\leavevmode\hbox
to7.91pt{\vbox to7.91pt{\pgfpicture\makeatletter\hbox{\hskip
3.95439pt\lower-3.95439pt\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{
}\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{
}\nullfont\pgfsys@beginscope\pgfsys@invoke{ }\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{ }{
{{}}\hbox{\hbox{{\pgfsys@beginscope\pgfsys@invoke{
}\definecolor[named]{pgffillcolor}{rgb}{0.25390625,0.25390625,0.25390625}\pgfsys@color@rgb@fill{0.25390625}{0.25390625}{0.25390625}\pgfsys@invoke{
}{{}{{{}}}{{}}{}{}{}{}{}{}{}{}{}{\pgfsys@beginscope\pgfsys@invoke{
}\definecolor[named]{pgffillcolor}{rgb}{0.25390625,0.25390625,0.25390625}\pgfsys@color@rgb@fill{0.25390625}{0.25390625}{0.25390625}\pgfsys@invoke{
}{}\pgfsys@moveto{3.7544pt}{0.0pt}\pgfsys@curveto{3.7544pt}{2.07352pt}{2.07352pt}{3.7544pt}{0.0pt}{3.7544pt}\pgfsys@curveto{-2.07352pt}{3.7544pt}{-3.7544pt}{2.07352pt}{-3.7544pt}{0.0pt}\pgfsys@curveto{-3.7544pt}{-2.07352pt}{-2.07352pt}{-3.7544pt}{0.0pt}{-3.7544pt}\pgfsys@curveto{2.07352pt}{-3.7544pt}{3.7544pt}{-2.07352pt}{3.7544pt}{0.0pt}\pgfsys@closepath\pgfsys@moveto{0.0pt}{0.0pt}\pgfsys@fillstroke\pgfsys@invoke{
} \pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope}{{{{}}\pgfsys@beginscope\pgfsys@invoke{
}\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{-2.84528pt}{-2.45pt}\pgfsys@invoke{
}\hbox{{\definecolor[named]{.}{rgb}{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@invoke{
}\pgfsys@color@gray@fill{1}\pgfsys@invoke{
}\hbox{\minipage[t]{5.69055pt}\leavevmode\hbox{}{\definecolor[named]{.}{rgb}{1,1,1}\color[rgb]{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@color@gray@fill{1}\tiny{2}}\endminipage}
}}\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}}
\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}} }
\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope{{{}}}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope\hss}}\lxSVG@closescope\endpgfpicture}}}(\bm{\mathrm{z}})$ in
Fig. 21.

T$\mathcal{C}at$ $=$ T $\delta$T $=$ T $\mathcal{C}at$T $=$ T $\mathcal{C}at$
$\cdots$$\rightarrow$$\leftarrow$$\uparrow$$\downarrow$$\rightarrow$$\leftarrow$$\uparrow$$\uparrow$$\rightarrow$$\rightarrow$$\leftarrow$$\downarrow$$\rightarrow$$\leftarrow$$\downarrow$${\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\uparrow}$$\rightarrow$$\rightarrow$$\bm{\mathrm{z}}_{t}$$\rightarrow$$\leftarrow$$\downarrow$$\rightarrow$$\leftarrow$$\downarrow$${\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\uparrow}$$\rightarrow$$\rightarrow$$\rightarrow$$\leftarrow$
Figure 24: GFE computation on the CFFG.

With these choices, the update equations become identical to those of [20].
Indeed, careful inspection of the update rules given in [20] reveals the
component parts of the P-substituted message. However by using CFFG’s and
P-substitution, we can cast their results as message passing on generic graphs
which immediately generalises their results.

Once inference has converged, we note that [20] shows that GFE evaluates
identically to the EFE, meaning we can use the same model comparison procedure
to select between policies as we used for reconstructing the classical
algorithm. This shows how we can obtain the algorithm of [20] as a special
case of SAIF.

##  7 SAIF for policy inference

The tools presented in this paper are not limited to restating prior work.
Indeed, SAIF offers several advantages over prior methods, one of which is the
ability to directly infer a policy instead of relying on a post-hoc comparison
of energy terms. To demonstrate, we solve two variations of the classic T-maze
task [7]. This is a well-studied setting within the AIF literature and
therefore constitutes a good minimal benchmark. In the T-maze experiment, the
agent lives in a maze with four locations as depicted in Fig 25. The agent (☺)
starts in position 1 and knows that a reward is present at either position 2
or 3, but not which one. At position 4 is a cue that informs the agent which
arm contains the reward. The optimal action to take is therefore to first
visit 4 and learn which arm contains the reward before going to the rewarded
arm. Because this course of action requires delaying the reward, an agent
following a greedy policy behaves sub-optimally. The T-maze is therefore
considered a reasonable minimal example of the epistemic, information-seeking
behaviour that is a hallmark of AIF agents. We implemented our experiments in
the reactive message passing toolbox RxInfer [2]

1324☺ Figure 25: The T-maze environment

###  7.1 Model specification

The generative model for the T-maze is an adaptation of the discrete POMDP
used by [7, 20]. We assume the agent starts at some current timestep $t$ and
want to infer a policy up to a known time horizon $T$. The generative model is
then

$\displaystyle p(\bm{x},\bm{z},\bm{u})$
$\displaystyle\propto\underbrace{p(\bm{\mathrm{z}}_{t})}_{\begin{subarray}{c}\text{Initial}\\\
\text{state}\end{subarray}}\prod_{k=t+1}^{T}\underbrace{p(\bm{\mathrm{x}}_{k}|\bm{\mathrm{z}}_{k})}_{\begin{subarray}{c}\text{Observation}\\\
\text{model}\end{subarray}}\underbrace{p(\bm{\mathrm{z}}_{k}|\bm{\mathrm{z}}_{k-1},\bm{\mathrm{u}}_{k})}_{\begin{subarray}{c}\text{Transition}\\\
\text{model}\end{subarray}}\underbrace{p(\bm{\mathrm{u}}_{k})}_{\begin{subarray}{c}\text{Control}\\\
\text{prior}\end{subarray}}\underbrace{\tilde{p}(\bm{\mathrm{x}}_{k})}_{\begin{subarray}{c}\text{Goal}\\\
\text{prior}\end{subarray}}$ (42)

where

$\displaystyle p(\bm{\mathrm{z}}_{t})$
$\displaystyle=\mathcal{C}at(\bm{\mathrm{z}}_{t}|\bm{\mathrm{d}})$ (43a)
$\displaystyle p(\bm{\mathrm{x}}_{k}|\bm{\mathrm{z}}_{k})$
$\displaystyle=\mathcal{C}at(\bm{\mathrm{x}}_{k}|\bm{\mathrm{A}}\bm{\mathrm{z}}_{k})$
(43b) $\displaystyle
p(\bm{\mathrm{z}}_{k}|\bm{\mathrm{z}}_{k-1},\bm{\mathrm{u}}_{k})$
$\displaystyle=\prod_{n}\mathcal{C}at(\bm{\mathrm{z}}_{k}|\bm{\mathrm{B}}_{n}\bm{\mathrm{z}}_{k-1})^{\bm{\mathrm{u}}_{nk}}$
(43c) $\displaystyle p(\bm{\mathrm{u}}_{k})$
$\displaystyle=\mathcal{C}at(\bm{\mathrm{u}}_{k}|\bm{\mathrm{e}}_{k})$ (43d)
$\displaystyle\tilde{p}(\bm{\mathrm{x}}_{k})$
$\displaystyle=\mathcal{C}at(\bm{\mathrm{x}}_{k}|\bm{\mathrm{c}}_{k})\,.$
(43e)

$\bm{\mathrm{u}}_{k}$ is a one-hot encoded vector of length $n$. The notation
$\bm{\mathrm{u}}_{nk}$ picks the $n$’th entry of $\bm{\mathrm{u}}_{k}$. We
show the corresponding CFFG in Fig 26.

$\mathcal{C}at$$\mathcal{T}M$ = $\mathcal{T}M$ $\bm{\mathrm{B}}_{1}$
$\bm{\mathrm{B}}_{4}$$\cdots$ $\bm{\mathrm{B}}_{1}$
$\bm{\mathrm{B}}_{4}$$\cdots$ $T$ $T$ $\mathcal{C}at$$\mathcal{C}at$
$\mathcal{C}at$$\mathcal{C}at$
$\bm{\mathrm{d}}$$\bm{\mathrm{z}}_{t}$$\bm{\mathrm{z}}_{t+1}$$\bm{\mathrm{u}}_{t}$$\bm{\mathrm{u}}_{t+1}$$\bm{\mathrm{e}}_{t}$$\bm{\mathrm{e}}_{t+1}$$\bm{\mathrm{A}}_{t}$$\bm{\mathrm{A}}_{t+1}$$\bm{\mathrm{c}}_{t}$$\bm{\mathrm{c}}_{t+1}$
Figure 26: CFFG for the T-maze experiment

Eq. (43c) defines a mixture model over candidate transition matrices indexed
by $\bm{\mathrm{u}}_{k}$. We give the details of this node function and the
required messages in Appendix C. The message passing schedule is shown in Fig.
27 with GFE-based messages highlighted in red

$\mathcal{C}at$$\mathcal{T}M$ = $\mathcal{T}M$ $\uparrow$ $\uparrow$$\cdots$
$\uparrow$ $\uparrow$$\cdots$ $T$ $T$ $\mathcal{C}at$$\mathcal{C}at$
$\mathcal{C}at$$\mathcal{C}at$
$\rightarrow$$\rightarrow$$\leftarrow$$\rightarrow$$\leftarrow$$\rightarrow$$\leftarrow$$\rightarrow$$\uparrow$$\downarrow$$\uparrow$$\downarrow$$\rightarrow$$\rightarrow$${\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\uparrow}$$\downarrow$${\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\uparrow}$$\rightarrow$$\rightarrow$$\rightarrow$$\rightarrow$
Figure 27: Message passing schedule for the T-maze experiment

Following [7] we define the initial state and control prior as

$\displaystyle\bm{\mathrm{d}}$
$\displaystyle=(1,0,0,0)^{T}\otimes(0.5,0.5)^{T}$ (44a)
$\displaystyle\bm{\mathrm{e}}_{k}$ $\displaystyle=(0.25,0.25,0.25,0.25)\forall
k$ (44b)

with $\otimes$ denoting the Kronecker product. The transition mixture node
requires a set of candidate transition matrices. The T-maze utilises four
possible transitions, given below

$\displaystyle\bm{\mathrm{B}}_{1}=\begin{bmatrix}1&1&1&1\\\ 0&0&0&0\\\
0&0&0&0\\\
0&0&0&0\end{bmatrix}\otimes\bm{\mathrm{I}}_{2}\,,\bm{\mathrm{B}}_{2}=\begin{bmatrix}0&1&1&0\\\
1&0&0&1\\\ 0&0&0&0\\\ 0&0&0&0\end{bmatrix}\otimes\bm{\mathrm{I}}_{2}$
$\displaystyle\bm{\mathrm{B}}_{3}=\begin{bmatrix}0&1&1&0\\\ 0&0&0&0\\\
1&0&0&1\\\
0&0&0&0\end{bmatrix}\otimes\bm{\mathrm{I}}_{2}\,,\bm{\mathrm{B}}_{4}=\begin{bmatrix}0&1&1&0\\\
0&0&0&0\\\ 0&0&0&0\\\ 1&0&0&1\end{bmatrix}\otimes\bm{\mathrm{I}}_{2}$

where $\bm{\mathrm{I}}_{2}$ denotes a $2\times 2$ identity matrix. Note that
these differ slightly from the original implementation of [7]. In [7] invalid
transitions are represented by an identity mapping where we instead model
invalid transitions by sending the agent back to position 1. The likelihood
matrix $\bm{\mathrm{A}}$ is given by four blocks, corresponding to the
observation likelihood in each position.

$\displaystyle\bm{\mathrm{A}}=\begin{bmatrix}\bm{\mathrm{A}}_{1}&&&\\\
&\bm{\mathrm{A}}_{2}&&\\\ &&\bm{\mathrm{A}}_{3}&\\\
&&&\bm{\mathrm{A}}_{4}\end{bmatrix}\,,$ (45)

with everything outside the blocks being set to 0. The blocks are

$\displaystyle\bm{\mathrm{A}}_{1}=\begin{bmatrix}0.5&0.5\\\ 0.5&0.5\\\ 0&0\\\
0&0\end{bmatrix}\,,$ $\displaystyle\bm{\mathrm{A}}_{2}=\begin{bmatrix}0&0\\\
0&0\\\ \alpha&1-\alpha\\\ 1-\alpha&\alpha\end{bmatrix}$
$\displaystyle\bm{\mathrm{A}}_{3}=\begin{bmatrix}0&0\\\ 0&0\\\
1-\alpha&\alpha\\\ \alpha&1-\alpha\end{bmatrix}\,,$
$\displaystyle\bm{\mathrm{A}}_{4}=\begin{bmatrix}1&0\\\ 0&1\\\ 0&0\\\
0&0\end{bmatrix}\,,$ (46)

with $\alpha$ being the probability of observing a reward. The goal prior is
given by

$\displaystyle\bm{\mathrm{c}}_{k}$
$\displaystyle=\sigma\big{(}(0,0,c,-c)^{T}\otimes(1,1,1,1)^{T}\big{)}\forall
k$ (47)

with $c$ being the utility ascribed to a reward and $\sigma(\cdot)$ the
softmax function. Inference for the parts of the model not in $\mathcal{P}$
can be accomplished using belief propagation. We follow the experimental setup
of [7] and let $c=2,\alpha=0.9$. For inference we perform two iterations of
our message passing procedure and use 20 Newton steps for obtaining the
parameters of the outgoing message from the P-substituted nodes.

We show the results in Fig. 28. The number in each cell is the posterior
probability mass assigned to the corresponding action, with the most likely
actions highlighted in  red.

$0.25$$0.20$$0.20$$0.35$Controls at time step
1$0.13$$0.30$$0.30$$0.26$Controls at time step 2 Figure 28: Posterior controls
for the T-maze experiment

Fig. 28 shows an agent that initially prefers the epistemic action (move to
state 4) at time $t+1$ and subsequently exhibits a preference for either of
the potentially rewarding arms (indifferent between states 2 and 3). This
shows that SAIF is able to infer the optimal policy and that our approach can
reproduce prior results on the T-maze.

Since CFFG’s are inherently modular, they allow us to modify the inference
task without changing the model. To demonstrate, we now add pointmass
constraints to the control variables. This corresponds to selecting the MAP
estimate of the control posterior [26] and results in the CFFG shown in Fig.
29. The schedule and all messages remain the same as our previous experiment -
however now the agent will instead select the most likely course of action
instead of providing us with a posterior distribution.

$\mathcal{C}at$$\mathcal{T}M$ = $\mathcal{T}M$ $B_{1}$ $B_{4}$$\cdots$ $B_{1}$
$B_{4}$$\cdots$ $T$ $T$ $\mathcal{C}at$$\mathcal{C}at$
$\mathcal{C}at$$\mathcal{C}at$ $\delta$$\delta$ Figure 29: CFFG for the T-maze
model with additional pointmass constraints

Performing this experiment yields the policy shown in Fig. 30

$0.0$$0.0$$0.0$$1.0$Controls at time step 1$0.0$$0.0$$1.0$$0.0$Controls at
time step 2 Figure 30: Posterior controls for the T-maze experiment with
$\delta$-constraints

Fig. 30 once again shows that the agent is able to infer the optimal policy
for solving the task. For repeated runs, the agent will randomly select to
move to either position 2 or 3 at the second step due to minute differences in
the Monte Carlo estimates used for computing the messages. The pointmass
constraint obscures this since it forces the marginal to put all mass on the
MAP estimate.

The point of repeating the experiments with $\delta$-constraints is not to
show that the behaviour of the agent changes dramatically. Instead, the idea
is to demonstrate that CFFG’s allow for modular specification of AIF agents
which allows for adapting parts of the model without having to touch the rest.
In this case, the only parts of inference that changes are those involving the
control marginal. This means all messages out of the P-substituted composite
nodes are unaffected since the $\delta$-constraints are only applied to the
control marginals.

##  8 Conclusions

In this paper we have proposed a novel approach to AIF based on lagrangian
optimisation which we have named SAIF. We demonstrated SAIF on a classic
benchmark problem from the AIF literature and found that it inherits the
epistemic drive that is a hallmark feature of AIF. SAIF presents three main
advantages over previous algorithms for AIF.

Firstly, an advantage of SAIF is the computational efficiency afforded by
being able to pass backward messages instead of needing to perform forwards
rollouts for every policy. While SAIF is still an iterative procedure, the
computational complexity of each iteration scales linearly in the size of the
planning horizon $T$ instead of exponentially.

A second advantage is that SAIF allows for directly inferring posteriors over
control signals instead of relying on a model comparison step based on
EFE/GFE. This means that SAIF unifies inference for perception, learning, and
actions into a single procedure without any overhead - it all becomes part of
the same inference procedure. See our companion paper [16] for more details.

Thirdly, SAIF is inherently modular and consequently works for freely
definable CFFG’s, while prior work has focused mostly on specific generative
models. Extensions to hierarchical or heterarchical models are straightforward
and only require writing out the corresponding CFFG.

We have also introduced CFFG notation for writing down constraints and
P-subtitutions on an FFG. CFFG’s are useful not just for AIF but for
specifying free energy functionals in general. CFFG’s accomplish this through
a simple and intuitive graphical syntax. Our hope is that CFFG’s can become a
standard tool similar to FFG’s when it is desirable to write not just a model
$f$ but also a family of approximating distributions $\mathcal{Q}$.
Specifically in the context of AIF we have also introduced P-substitution as a
way to modify the underlying free energy functional. In doing so we have
formalised the relation between AIF and message passing on a CFFG, paving the
way for future developments.

In future work we plan to extend SAIF to more node constructions to further
open up the scope of available problems that can be attacked using AIF.

## Acknowledgements

This research was made possible by funding from GN Hearing A/S. This work is
part of the research programme Efficient Deep Learning with project number
P16-25 project 5, which is (partly) financed by the Netherlands Organisation
for Scientific Research (NWO).

We gratefully acknowledge stimulating discussions with the members of the
BIASlab research group at the Eindhoven University of Technology and members
of VERSES Research Lab, in particular Ismail Senoz, Bart van Erp, Dmitry
Bagaev, Karl Friston, Chris Buckley, Conor Heins and Tim Verbelen.

## References

  * [1] Semih Akbayrak, Ivan Bocharov and Bert Vries  “Extended variational message passing for automated approximate Bayesian inference”  In _Entropy_ 23.7 MDPI, 2021, pp. 815
  * [2] Dmitry Bagaev and Bert Vries  “Reactive message passing for scalable Bayesian inference”  In _arXiv preprint arXiv:2112.13251_ , 2021
  * [3] Ariel Caticha  “Entropic inference and the foundations of physics”  In _Brazilian Chapter of the International Society for Bayesian Analysis-ISBrA, Sao Paulo, Brazil_ , 2012
  * [4] Théophile Champion, Lancelot Da Costa, Howard Bowman and Marek Grześ  “Branching Time Active Inference: the theory and its generality” Number: arXiv:2111.11107 arXiv:2111.11107 [cs]  arXiv, 2022  DOI: 10.48550/arXiv.2111.11107
  * [5] Lancelot Da Costa et al.  “Active inference on discrete state-spaces: a synthesis” arXiv: 2001.07203  In _arXiv:2001.07203 [q-bio]_ , 2020  URL: http://arxiv.org/abs/2001.07203
  * [6] Karl Friston  “A free energy principle for a particular physics” arXiv: 1906.10184  In _arXiv:1906.10184 [q-bio]_ , 2019  URL: http://arxiv.org/abs/1906.10184
  * [7] Karl Friston et al.  “Active inference and epistemic value”  In _Cognitive Neuroscience_ 6.4, 2015, pp. 187–214  DOI: 10.1080/17588928.2015.1020053
  * [8] Karl Friston et al.  “Sophisticated Inference” arXiv: 2006.04120  In _arXiv:2006.04120 [cs, q-bio]_ , 2020  URL: http://arxiv.org/abs/2006.04120
  * [9] Karl J Friston, N Trujillo-Barreto and Jean Daunizeau  “DEM: a variational treatment of dynamic systems”  In _Neuroimage_ 41.3 Elsevier, 2008, pp. 849–885
  * [10] Karl J. Friston, Jean Daunizeau, James Kilner and Stefan J. Kiebel  “Action and behavior: a free-energy formulation”  In _Biological Cybernetics_ 102.3, 2010, pp. 227–260  DOI: 10.1007/s00422-010-0364-z
  * [11] Danijar Hafner et al.  “Action and Perception as Divergence Minimization” arXiv: 2009.01791  In _arXiv:2009.01791 [cs, math, stat]_ , 2020  URL: http://arxiv.org/abs/2009.01791
  * [12] Conor Heins et al.  “pymdp: A Python library for active inference in discrete state spaces”  In _arXiv preprint arXiv:2201.03904_ , 2022
  * [13] Hilbert J. Kappen, Vicenç Gómez and Manfred Opper  “Optimal control as a graphical model inference problem”  In _Machine Learning_ 87.2, 2012, pp. 159–182  DOI: 10.1007/s10994-012-5278-7
  * [14] Magnus Koudahl, Christopher L Buckley and Bert Vries  “A Message Passing Perspective on Planning Under Active Inference”  In _Active Inference: Third International Workshop, IWAI 2022, Grenoble, France, September 19, 2022, Revised Selected Papers_ , 2023, pp. 319–327  Springer
  * [15] Magnus T. Koudahl, Wouter M. Kouw and Bert Vries  “On Epistemics in Expected Free Energy for Linear Gaussian State Space Models”  In _Entropy_ 23.12 MDPI AG, 2021, pp. 1565  DOI: 10.3390/e23121565
  * [16] Thijs Laar, Magnus Koudahl and Bert De Vries  “Realising Synthetic Active Inference Agents, Part II: Variational Message Updates”  In _In preparation_
  * [17] David JC MacKay and David JC Mac Kay  “Information theory, inference and learning algorithms”  Cambridge university press, 2003
  * [18] Beren Millidge, Alexander Tschantz, Anil Seth and Christopher Buckley  “Understanding the Origin of Information-Seeking Exploration in Probabilistic Objectives for Control” arXiv: 2103.06859  In _arXiv:2103.06859 [cs]_ , 2021  URL: http://arxiv.org/abs/2103.06859
  * [19] Thomas P. Minka  “Expectation Propagation for approximate Bayesian inference”  In _CoRR_ abs/1301.2294, 2013  arXiv: http://arxiv.org/abs/1301.2294
  * [20] Thomas Parr and Karl J. Friston  “Generalised free energy and active inference”  In _Biological Cybernetics_ 113.5-6, 2019, pp. 495–513  DOI: 10.1007/s00422-019-00805-w
  * [21] Judea Pearl  “Reverend Bayes on inference engines: A distributed hierarchical approach”  In _Probabilistic and Causal Inference: The Works of Judea Pearl_ , 2022, pp. 129–138
  * [22] Christoph Reller  “State-space methods in statistical signal processing: New ideas and applications”  ETH Zurich, 2013
  * [23] Noor Sajid, Philip J. Ball and Karl J. Friston  “Active inference: demystified and compared” arXiv: 1909.10863  In _arXiv:1909.10863 [cs, q-bio]_ , 2020  URL: http://arxiv.org/abs/1909.10863
  * [24] Simo Sarkka  “Bayesian Filtering and Smoothing”  Cambridge: Cambridge University Press, 2013  DOI: 10.1017/CBO9781139344203
  * [25] Sarah Schwöbel, Stefan Kiebel and Dimitrije Marković  “Active Inference, Belief Propagation, and the Bethe Approximation”  In _Neural Computation_ 30.9, 2018, pp. 2530–2567  DOI: 10.1162/neco_a_01108
  * [26] İsmail Şenöz, Thijs Laar, Dmitry Bagaev and Bert Vries  “Variational Message Passing and Local Constraint Manipulation in Factor Graphs” Number: 7 Publisher: Multidisciplinary Digital Publishing Institute  In _Entropy_ 23.7, 2021, pp. 807  DOI: 10.3390/e23070807
  * [27] Ryan Smith, Karl Friston and Christopher Whyte  “A Step-by-Step Tutorial on Active Inference and its Application to Empirical Data”, 2021  DOI: 10.31234/osf.io/b4jm6
  * [28] John Winn, Christopher M Bishop and Tommi Jaakkola  “Variational message passing.”  In _Journal of Machine Learning Research_ 6.4, 2005
  * [29] Dan Zhang et al.  “Unifying Message Passing Algorithms Under the Framework of Constrained Bethe Free Energy Minimization” arXiv: 1703.10932  In _arXiv:1703.10932 [cs, math]_ , 2021  URL: http://arxiv.org/abs/1703.10932

##  Appendix A Expected free energy

This section will focus on the the most common alternative functional used for
AIF, the Expected Free Energy (EFE).

EFE is the standard choice for AIF models and is what is found in most AIF
focused software such as SPM [9] and PyMDP [12]. In this section we will cover
the details of how EFE is commonly treated.

EFE is defined specifically over future time steps and on a particular choice
of generative model. The model in question is a state space model (SSM) of the
form

$\displaystyle p(\bm{x},\bm{z}|\hat{\bm{u}})$
$\displaystyle=\underbrace{p(\bm{\mathrm{z}}_{t})}_{\begin{subarray}{c}\text{State}\\\
\text{prior}\end{subarray}}\prod_{k=t+1}^{T}\underbrace{p(\bm{\mathrm{x}}_{k}|\bm{\mathrm{z}}_{k})}_{\begin{subarray}{c}\text{Observation}\\\
\text{model}\end{subarray}}\underbrace{p(\bm{\mathrm{z}}_{k}|\bm{\mathrm{z}}_{k-1},\bm{\mathrm{\hat{u}}}_{k})}_{\begin{subarray}{c}\text{Transition}\\\
\text{model}\end{subarray}}$ (48)

where $\bm{\mathrm{\hat{u}}}_{k}$ denotes a particular control parameter which
is known apriori and fixed. Note that (48) does not include a prior over
actions $\bm{u}$ and is instead conditional on a fixed policy
$\hat{\bm{u}}=\bm{\mathrm{\hat{u}}}_{t+1:T}$. We use the $\hat{}$ notation on
$\bm{\mathrm{u}}$ to indicate that it is treated as a known value instead of a
random variable.

The EFE is evaluated as a function of a particular policy. The EFE is defined
as [7]

$\displaystyle G_{EFE}[q;\bm{\mathrm{\hat{u}}}]$
$\displaystyle=\sum_{k=t+1}^{T}\iint
p(\bm{\mathrm{x}}_{k}|\bm{\mathrm{z}}_{k})\underbrace{q(\bm{\mathrm{z}}_{k}|\bm{\mathrm{\hat{u}}}_{k})\log\frac{q(\bm{\mathrm{z}}_{k}|\bm{\mathrm{\hat{u}}}_{k})}{p(\bm{\mathrm{x}}_{k},\bm{\mathrm{z}}_{k}|\bm{\mathrm{\hat{u}}}_{k})}}_{\text{VFE
conditioned on
}\bm{\mathrm{\hat{u}}}_{k}}\differential\bm{\mathrm{x}}_{k}\differential\bm{\mathrm{z}}_{k}$
(49a) $\displaystyle=\sum_{k=t+1}^{T}\iint
p(\bm{\mathrm{x}}_{k}|\bm{\mathrm{z}}_{k})q(\bm{\mathrm{z}}_{k}|\bm{\mathrm{\hat{u}}}_{k})\log\frac{q(\bm{\mathrm{z}}_{k}|\bm{\mathrm{\hat{u}}}_{k})}{p(\bm{\mathrm{z}}_{k}|\bm{\mathrm{x}}_{k},\bm{\mathrm{\hat{u}}}_{k})}-\log
p(\bm{\mathrm{x}}_{k})\differential\bm{\mathrm{x}}_{k}\differential\bm{\mathrm{z}}_{k}$
(49b)

where $p(\bm{\mathrm{x}}_{k})$ denotes a goal prior over preferred
observations. Note that $p(\bm{\mathrm{x}}_{k})$ is not part of the generative
model in Eq. (48). To compute Eq. (49) we also need

$\displaystyle
p(\bm{\mathrm{x}}_{k},\bm{\mathrm{z}}_{k}\mid\bm{\mathrm{\hat{u}}}_{k})$
$\displaystyle=\int
p(\bm{\mathrm{x}}_{k}\mid\bm{\mathrm{z}}_{k})p(\bm{\mathrm{z}}_{k}\mid\bm{\mathrm{z}}_{k-1},\hat{u}_{k}){\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}q(\bm{\mathrm{z}}_{k-1})}\differential\bm{\mathrm{z}}_{k-1}$
(50)

meaning we use the  forward prediction from the previous time step to compute
$p(\bm{\mathrm{x}}_{k},\bm{\mathrm{z}}_{k}\mid\bm{\mathrm{\hat{u}}}_{k})$. If
we further we assume that

$\displaystyle
q(\bm{\mathrm{x}}_{k},\bm{\mathrm{z}}_{k}\mid\bm{\mathrm{\hat{u}}}_{k})$
$\displaystyle=p(\bm{\mathrm{x}}_{k}\mid\bm{\mathrm{z}}_{k})q(\bm{\mathrm{z}}_{k}\mid\bm{\mathrm{\hat{u}}}_{k})$
(51)

It can be shown [5, 7, 15, 23] that (49) can be decomposed into a bound on a
mutual information term and a cross-entropy loss between predicted
observations and the goal prior.

##  Appendix B VFE and GFE

In this section we walk through how the formulation of GFE given by [20]
reduces to the VFE when observations are available. To show how this comes
about, we note that [20] assumes that

$\displaystyle q(\bm{\mathrm{x}}_{k}|\bm{\mathrm{z}}_{k})$
$\displaystyle=\begin{cases}\delta(\bm{\mathrm{x}}_{k}-\bm{\mathrm{\hat{x}}}_{k})&\text{if
}k\leq t\\\ p(\bm{\mathrm{x}}_{k}|\bm{\mathrm{z}}_{k})&\text{if
}k>t\end{cases}\,.$ (52)

where $t$ is the current time step and $\bm{\mathrm{\hat{x}}}_{k}$ denotes the
observed data point at time step $k$. This is a problematic move since
$q(\bm{\mathrm{x}}_{k}\mid\bm{\mathrm{z}}_{k})$ is no longer a function of
$\bm{\mathrm{z}}_{k}$ for $k\leq t$. However if we do take Eq. (16) as valid
and further assume

$\displaystyle
q(\bm{\mathrm{x}}_{k},\bm{\mathrm{z}}_{k}\mid\bm{\mathrm{\hat{u}}}_{k})$
$\displaystyle=q(\bm{\mathrm{x}}_{k}\mid\bm{\mathrm{z}}_{k})q(\bm{\mathrm{z}}_{k}\mid\bm{\mathrm{\hat{u}}}_{k})$
(53a) $\displaystyle\tilde{p}(\bm{\mathrm{x}}_{k})$ $\displaystyle=1\text{ if
}k\leq t$ (53b)

and note that

$\displaystyle q(\bm{\mathrm{x}}_{k}|\bm{\mathrm{\hat{u}}}_{k})$
$\displaystyle=\int
q(\bm{\mathrm{x}}_{k}|\bm{\mathrm{z}}_{k})q(\bm{\mathrm{z}}_{k}|\bm{\mathrm{\hat{u}}}_{k})\differential\bm{\mathrm{z}}_{k}\,.$
(54)

We can plug this result into (15) and obtain for $k\leq t$

$\displaystyle\sum_{k=1}^{t}\iint
q(\bm{\mathrm{x}}_{k}|\bm{\mathrm{z}}_{k})q(\bm{\mathrm{z}}_{k}\mid\bm{\mathrm{\hat{u}}}_{k})\log\frac{q(\bm{\mathrm{x}}_{k}\mid\bm{\mathrm{\hat{u}}}_{k})q(\bm{\mathrm{z}}_{k}\mid\bm{\mathrm{\hat{u}}}_{k})}{p(\bm{\mathrm{x}}_{k},\bm{\mathrm{z}}_{k}\mid\bm{\mathrm{\hat{u}}}_{k})\tilde{p}(\bm{\mathrm{x}}_{k})}\differential\bm{\mathrm{x}}_{k}\differential\bm{\mathrm{z}}_{k}$
$\displaystyle=\sum_{k=1}^{t}\iint
q(\bm{\mathrm{x}}_{k}|\bm{\mathrm{z}}_{k})q(\bm{\mathrm{z}}_{k}\mid\hat{u}_{k})\log\frac{\int\left[q(\bm{\mathrm{x}}_{k}|\bm{\mathrm{z}}_{k})q(\bm{\mathrm{z}}_{k}\mid\bm{\mathrm{\hat{u}}}_{k})\right]\differential\bm{\mathrm{z}}_{k}q(\bm{\mathrm{z}}_{k}\mid\bm{\mathrm{\hat{u}}}_{k})}{p(\bm{\mathrm{x}}_{k},\bm{\mathrm{z}}_{k}\mid\bm{\mathrm{\hat{u}}}_{k})}\differential\bm{\mathrm{x}}_{k}\differential\bm{\mathrm{z}}_{k}$
(55a)

where the last line follows from Eq. (53b). Then by (16) we have

$\displaystyle=\sum_{k=1}^{t}\iint\delta(\bm{\mathrm{x}}_{k}-\bm{\mathrm{\hat{x}}}_{k})q(\bm{\mathrm{z}}_{k}\mid\bm{\mathrm{\hat{u}}}_{k})\log\frac{\int\left[\delta(\bm{\mathrm{x}}_{k}-\bm{\mathrm{\hat{x}}}_{k})q(\bm{\mathrm{z}}_{k}\mid\bm{\mathrm{\hat{u}}}_{k})\right]\differential\bm{\mathrm{z}}_{k}q(\bm{\mathrm{z}}_{k}\mid\bm{\mathrm{\hat{u}}}_{k})}{p(\bm{\mathrm{x}}_{k},\bm{\mathrm{z}}_{k}\mid\bm{\mathrm{\hat{u}}}_{k})}\differential\bm{\mathrm{x}}_{k}\differential\bm{\mathrm{z}}_{k}$
(56)

Now we pull $\delta(\bm{\mathrm{x}}_{k}-\hat{x}_{k})$ out of the integral and
solve to find

$\displaystyle=\sum_{k=1}^{t}\iint\delta(\bm{\mathrm{x}}_{k}-\bm{\mathrm{\hat{x}}}_{k})q(\bm{\mathrm{z}}_{k}\mid\bm{\mathrm{\hat{u}}}_{k})\log\frac{\delta(\bm{\mathrm{x}}_{k}-\bm{\mathrm{\hat{x}}}_{k})\overbrace{\Big{[}\int
q(\bm{\mathrm{z}}_{k}\mid\bm{\mathrm{\hat{u}}}_{k})\differential\bm{\mathrm{z}}_{k}\Big{]}}^{\text{Integrates
to
1}}q(\bm{\mathrm{z}}_{k}\mid\bm{\mathrm{\hat{u}}}_{k})}{p(\bm{\mathrm{x}}_{k},\bm{\mathrm{z}}_{k}\mid\bm{\mathrm{\hat{u}}}_{k})}\differential\bm{\mathrm{x}}_{k}\differential\bm{\mathrm{z}}_{k}$
(57a) $\displaystyle=\sum_{k=1}^{t}\int
q(\bm{\mathrm{z}}_{k}\mid\bm{\mathrm{\hat{u}}}_{k})\log\frac{q(\bm{\mathrm{z}}_{k}\mid\bm{\mathrm{\hat{u}}}_{k})}{p(\bm{\mathrm{\hat{x}}}_{k},\bm{\mathrm{z}}_{k}\mid\bm{\mathrm{\hat{u}}}_{k})}\differential\bm{\mathrm{z}}_{k}$
(57b)

which we can recognise as a VFE with data constraints.

##  Appendix C The transition mixture node

In this section we derive the message passing rules required for the
Transition Mixture node used in our experiments on policy inference. Before
doing so, we first establish some preliminary results for the categorical
distribution and the standard transition node. In this section we will also on
occasion use an $\overline{\text{overbar}}$ to indicate an expectation,
$\mathbb{E}_{q(x)}[g(x)]=\overline{g(x)}$ The categorical distribution is
given by

$\displaystyle\mathcal{C}at(\bm{\mathrm{x}}\mid\bm{\mathrm{z}})$
$\displaystyle=\prod_{i=1}^{I}\bm{\mathrm{z}}_{i}^{\bm{\mathrm{x}}_{i}}$ (58a)
$\displaystyle=\sum_{i=1}^{I}\bm{\mathrm{x}}_{i}\bm{\mathrm{z}}_{i}$ (58b)

where $\bm{\mathrm{x}}$ and $\bm{\mathrm{z}}$ are both one-hot encoded
vectors. The move Eq. (58a) to Eq. (58b) is possible only because
$\bm{\mathrm{x}}$ is one-hot encoded. Similarly we have for the standard
Transition node that

$\displaystyle\mathcal{C}at(\bm{\mathrm{x}}\mid\bm{\mathrm{A}}\bm{\mathrm{z}})$
$\displaystyle=\prod_{i=1}^{I}\prod_{j=1}^{I}\big{[}\bm{\mathrm{A}}_{ij}\big{]}^{\bm{\mathrm{x}}_{i}\bm{\mathrm{z}}_{j}}$
(59a)
$\displaystyle=\sum_{i=1}^{I}\sum_{j=1}^{I}\bm{\mathrm{x}}_{i}\bm{\mathrm{z}}_{j}\bm{\mathrm{A}}_{ij}$
(59b)

and again, since $\bm{\mathrm{x}}$ and $\bm{\mathrm{z}}$ are one-hot encoded

$\displaystyle\log\mathcal{C}at(\bm{\mathrm{x}}\mid\bm{\mathrm{A}}\bm{\mathrm{z}})$
$\displaystyle=\log[\prod_{i=1}^{I}\prod_{j=1}^{I}\big{[}\bm{\mathrm{A}}_{ij}\big{]}^{\bm{\mathrm{x}}_{i}\bm{\mathrm{z}}_{j}}\Bigg{]}$
(60a)
$\displaystyle=\sum_{i=1}^{I}\sum_{j=1}^{I}\log[\bm{\mathrm{A}}_{ij}^{\bm{\mathrm{x}}_{i}\bm{\mathrm{z}}_{j}}\Bigg{]}$
(60b)
$\displaystyle=\sum_{i=1}^{I}\sum_{j=1}^{I}\bm{\mathrm{x}}_{i}\bm{\mathrm{z}}_{j}\log\bm{\mathrm{A}}_{ij}$
(60c)

Now we are ready to start derivations for the transition mixture node. The
transition mixture node has the node function

$\displaystyle
f(\bm{\mathrm{x}},\bm{\mathrm{y}},\bm{\mathrm{z}},\bm{\mathrm{A}})$
$\displaystyle=\prod_{k}\mathcal{C}at\big{(}\bm{\mathrm{x}}\mid\bm{\mathrm{A}}_{k}\bm{\mathrm{z}}\big{)}^{\bm{\mathrm{y}}_{k}}$
(61a)
$\displaystyle=\prod_{k}\Bigg{[}\prod_{i}\prod_{j}\big{[}\bm{\mathrm{A}}_{k}\big{]}_{ij}^{\bm{\mathrm{z}}_{i}\bm{\mathrm{x}}_{j}}\Bigg{]}^{\bm{\mathrm{y}}_{k}}$
(61b)
$\displaystyle=\prod_{i,j,k}\bm{\mathrm{A}}_{ijk}^{\bm{\mathrm{z}}_{i}\bm{\mathrm{x}}_{j}\bm{\mathrm{y}}_{k}}$
(61c)

Where we use $i$ to index over columns, $j$ to index over rows and $k$ to
index over factors. $\bm{\mathrm{A}}$ is a 3-tensor that is normalized over
columns. In other words, each slice of $\bm{\mathrm{A}}$ corresponds to a
valid transition matrix and slices are indexed by $k$. We assume a structured
mean field factorisation such that

$\displaystyle
q(\bm{\mathrm{x}},\bm{\mathrm{y}},\bm{\mathrm{z}},\bm{\mathrm{A}})$
$\displaystyle=q(\bm{\mathrm{x}},\bm{\mathrm{y}},\bm{\mathrm{z}})\prod_{k}q(\bm{\mathrm{A}}_{k})$
(62)

The CFFG of the transition mixture node is shown in Fig. 31

$\mathcal{T}M$$z$$x$$y$$A_{1}$$A_{k}$ Figure 31: The transition mixture node

We now define

$\displaystyle\log
f_{\bm{\mathrm{A}}}(\bm{\mathrm{x}},\bm{\mathrm{y}},\bm{\mathrm{z}})$
$\displaystyle=\mathbb{E}_{q(\bm{\mathrm{A}})}\Big{[}\log
f(\bm{\mathrm{x}},\bm{\mathrm{y}},\bm{\mathrm{z}},\bm{\mathrm{A}})\Big{]}$
(63a)
$\displaystyle=\mathbb{E}_{q(\bm{\mathrm{A}})}\Big{[}\sum_{ijk}\bm{\mathrm{z}}_{i}\bm{\mathrm{x}}_{j}\bm{\mathrm{y}}_{k}\log\bm{\mathrm{A}}_{ijk}\Big{]}$
(63b)
$\displaystyle=\sum_{k}\bm{\mathrm{y}}_{k}\mathbb{E}_{q(\bm{\mathrm{A}}_{k})}\Big{[}\sum_{ij}\bm{\mathrm{z}}_{i}\bm{\mathrm{x}}_{j}\log\bm{\mathrm{A}}_{ijk}\Big{]}$
(63c)
$\displaystyle=\sum_{k}\bm{\mathrm{y}}_{k}\sum_{ij}\bm{\mathrm{z}}_{i}\bm{\mathrm{x}}_{j}\Big{[}\overline{\log\bm{\mathrm{A}}_{k}}\Big{]}_{ij}$
(63d)

which implies that

$\displaystyle f_{A}(\bm{\mathrm{x}},\bm{\mathrm{y}},\bm{\mathrm{z}})$
$\displaystyle=\exp\Bigg(\sum_{k}\bm{\mathrm{y}}_{k}\sum_{ijk}\bm{\mathrm{z}}_{i}\bm{\mathrm{x}}_{j}\Big{[}\overline{\log\bm{\mathrm{A}}_{k}}\Big{]}_{ij}\Bigg{missing})$
(64a)
$\displaystyle=\prod_{ijk}\exp\Big(\overline{\log\bm{\mathrm{A}}_{k}}\Big{missing})^{\bm{\mathrm{z}}_{i}\bm{\mathrm{x}}_{j}\bm{\mathrm{y}}_{k}}$
(64b)
$\displaystyle=\sum_{ijk}\bm{\mathrm{z}}_{i}\bm{\mathrm{x}}_{j}\bm{\mathrm{y}}_{k}\underbrace{\exp\Big(\overline{\log\bm{\mathrm{A}}_{k}}\Big{missing})_{ij}}_{\tilde{\bm{\mathrm{A}}}_{ijk}}$
(64c)

Now we are ready to derive the desired messages. The first message
$\nu(\cdot)$ will be towards $\bm{\mathrm{x}}$. If we use
$\mu(\bm{\mathrm{y}})$ to denote the incoming message on the edge
$\bm{\mathrm{y}}$ (similar for $\mu{\bm{\mathrm{z}}}$), then the message
towards $\bm{\mathrm{x}}$ is given by

$\displaystyle\nu(\bm{\mathrm{x}})$
$\displaystyle=\sum_{\bm{\mathrm{z}}}\sum_{\bm{\mathrm{y}}}\mu(\bm{\mathrm{z}})\mu(\bm{\mathrm{y}})\exp\Big(\mathbb{E}_{q(\bm{\mathrm{A}})}\log
f(\bm{\mathrm{x}},\bm{\mathrm{y}},\bm{\mathrm{z}},\bm{\mathrm{A}})\Big{missing})$
(65a)
$\displaystyle=\sum_{\bm{\mathrm{z}}}\sum_{\bm{\mathrm{y}}}\mu(\bm{\mathrm{z}})\mu(\bm{\mathrm{y}})f_{A}(\bm{\mathrm{x}},\bm{\mathrm{y}},\bm{\mathrm{z}})\Big{)}$
(65b)
$\displaystyle=\sum_{\bm{\mathrm{z}}}\sum_{\bm{\mathrm{y}}}\mu(\bm{\mathrm{z}})\mu(\bm{\mathrm{y}})\sum_{ijk}\bm{\mathrm{z}}_{i}\bm{\mathrm{x}}_{j}\bm{\mathrm{y}}_{k}\tilde{\bm{\mathrm{A}}}_{ijk}$
(65c)

Assuming the incoming messages are Categorically distributed, we can continue
as

$\displaystyle=\mathbb{E}_{\mathcal{C}at(\bm{\mathrm{z}}\mid\pi_{z})}\mathbb{E}_{\mathcal{C}at(\bm{\mathrm{y}}\mid\pi_{y})}\sum_{ijk}\bm{\mathrm{z}}_{i}\bm{\mathrm{x}}_{j}\bm{\mathrm{y}}_{k}\tilde{\bm{\mathrm{A}}}_{ijk}$
(66a)
$\displaystyle=\sum_{ijk}\pi_{zi}\pi_{yk}\bm{\mathrm{x}}_{j}\tilde{\bm{\mathrm{A}}}_{ijk}$
(66b)
$\displaystyle=\prod_{j}\Big{[}\sum_{ik}\pi_{zi}\pi_{yk}\tilde{\bm{\mathrm{A}}}_{ijk}\Big{]}^{\bm{\mathrm{x}}_{j}}$
(66c)

where the last line follows from $\bm{\mathrm{x}}$ being one-hot encoded. At
this point we can recognise the message

$\displaystyle\nu(\bm{\mathrm{x}})$
$\displaystyle\propto\mathcal{C}at(\bm{\mathrm{x}}\mid\bm{\mathrm{\rho}})\text{
where
}\bm{\mathrm{\rho}}_{j}=\frac{\sum_{ik}\pi_{zi}\pi_{yk}\tilde{\bm{\mathrm{A}}}_{ijk}}{\sum_{ijk}\pi_{zi}\pi_{yk}\tilde{\bm{\mathrm{A}}}_{ijk}}$
(67)

By symmetry, similar results hold for messages $\nu(\bm{\mathrm{z}})$ and
$\nu(\bm{\mathrm{y}})$. $\nu(\bm{\mathrm{z}})$ is given by

$\displaystyle\nu(\bm{\mathrm{z}})$
$\displaystyle=\sum_{x}\sum_{y}\mu(x)\mu(\bm{\mathrm{y}})f_{A}(\bm{\mathrm{x}},\bm{\mathrm{y}},\bm{\mathrm{z}})$
(68a)
$\displaystyle\propto\mathcal{C}at(\bm{\mathrm{z}}\mid\bm{\mathrm{\rho}})\text{
where
}\bm{\mathrm{\rho}}_{i}=\frac{\sum_{jk}\pi_{xj}\pi_{yk}\tilde{\bm{\mathrm{A}}}_{ijk}}{\sum_{ijk}\pi_{xj}\pi_{yk}\tilde{\bm{\mathrm{A}}}_{ijk}}\,.$
(68b)

Similarly, $\nu(\bm{\mathrm{y}})$ evaluates to

$\displaystyle\nu(\bm{\mathrm{y}})$
$\displaystyle=\sum_{x}\sum_{z}\mu(\bm{\mathrm{x}})\mu(\bm{\mathrm{z}})f_{A}(\bm{\mathrm{x}},\bm{\mathrm{y}},\bm{\mathrm{z}})$
(69a)
$\displaystyle\propto\mathcal{C}at(\bm{\mathrm{y}}\mid\bm{\mathrm{\rho}})\text{
where
}\bm{\mathrm{\rho}}_{k}=\frac{\sum_{ij}\pi_{x,j}\pi_{z,i}\tilde{\bm{\mathrm{A}}}_{ijk}}{\sum_{ijk}\pi_{x,j}\pi_{z,i}\tilde{\bm{\mathrm{A}}}_{ijk}}\,.$
(69b)

To compute the message towards the $n$’th candidate transition matrix
$\bm{\mathrm{A}}_{n}$, we will need to take an expectation with respect to
$q(\bm{\mathrm{x}},\bm{\mathrm{y}},\bm{\mathrm{z}})$. It is given by

$\displaystyle q(\bm{\mathrm{x}},\bm{\mathrm{y}},\bm{\mathrm{z}})$
$\displaystyle=\mu(\bm{\mathrm{y}})\mu(\bm{\mathrm{x}})\mu(\bm{\mathrm{z}})f_{A}(\bm{\mathrm{x}},\bm{\mathrm{y}},\bm{\mathrm{z}})$
(70a)
$\displaystyle=\prod_{ijk}\pi_{xj}^{\bm{\mathrm{x}}_{j}}\pi_{zi}^{\bm{\mathrm{z}}_{i}}\pi_{yk}^{\bm{\mathrm{y}}_{k}}\tilde{\bm{\mathrm{A}}}_{ijk}^{\bm{\mathrm{x}}_{j}\bm{\mathrm{z}}_{i}\bm{\mathrm{y}}_{k}}$
(70b)
$\displaystyle=\prod_{ijk}\Big{[}\pi_{xj}\pi_{zi}\pi_{yk}\tilde{\bm{\mathrm{A}}}_{ijk}\Big{]}^{\bm{\mathrm{x}}_{j}\bm{\mathrm{z}}_{i}\bm{\mathrm{y}}_{k}}$
(70c)
$\displaystyle\propto\mathcal{C}at(\bm{\mathrm{x}},\bm{\mathrm{y}},\bm{\mathrm{z}}\mid\bm{\mathrm{B}})$
(70d)

where $\bm{\mathrm{B}}$ is a three dimensional contingency tensor with entries

$\displaystyle\bm{\mathrm{B}}_{ijk}$
$\displaystyle=\frac{\pi_{xj}\pi_{zi}\pi_{yk}\tilde{\bm{\mathrm{A}}}_{ijk}}{\sum_{ijk}\pi_{xj}\pi_{zi}\pi_{yk}\tilde{\bm{\mathrm{A}}}_{ijk}}$
(71)

Now we can compute the message towards a transition matrix
$\bm{\mathrm{A}}_{n}$ as

$\displaystyle\log\nu(\bm{\mathrm{A}}_{n})$
$\displaystyle=\mathbb{E}_{q(\bm{\mathrm{A}}_{k\backslash
n})q(\bm{\mathrm{x}},\bm{\mathrm{y}},\bm{\mathrm{z}})}\Big{[}\log
f(\bm{\mathrm{x}},\bm{\mathrm{y}},\bm{\mathrm{z}},\bm{\mathrm{A}})\Big{]}$
(72a) $\displaystyle=\mathbb{E}_{q(\bm{\mathrm{A}}_{k\backslash
n})q(\bm{\mathrm{x}},\bm{\mathrm{y}},\bm{\mathrm{z}})}\Big{[}\sum_{ijk}\bm{\mathrm{x}}_{j}\bm{\mathrm{z}}_{i}\bm{\mathrm{y}}_{k}\log\bm{\mathrm{A}}_{ijk}\Big{]}$
(72b) $\displaystyle=\mathbb{E}_{q(\bm{\mathrm{A}}_{k\backslash
n})}\Big{[}\sum_{ijk}\bm{\mathrm{B}}_{ijk}\log\bm{\mathrm{A}}_{ijk}\Big{]}$
(72c) $\displaystyle=\underbrace{\sum_{k\backslash
n}\Big{[}\sum_{ijk}\bm{\mathrm{B}}_{ijk}\log\bm{\mathrm{A}}_{ijk}\Big{]}}_{\text{Constant
w.r.t
$\bm{\mathrm{A}}_{n}$}}+\sum_{ij}\bm{\mathrm{B}}_{ijn}\log\bm{\mathrm{A}}_{ijn}$
(72d) $\displaystyle\propto
tr\Big{(}\bm{\mathrm{B}}_{ijn}\log\bm{\mathrm{A}}_{ijn}\Big{)}$ (72e)

which implies that

$\displaystyle\nu(\bm{\mathrm{A}}_{n})$
$\displaystyle\propto\mathcal{D}ir(\bm{\mathrm{A}}_{n}\mid\bm{\mathrm{B}}_{n}+1)\,.$
(73)

We see that messages towards each component in $\bm{\mathrm{A}}$ are
distributed according to a Dirichlet distribution with parameters
$\bm{\mathrm{B}}_{n}+1$. The final expression we need to derive is the average
energy term for the transition mixture node. It is given by

$\displaystyle
U\big{[}f(\bm{\mathrm{x}},\bm{\mathrm{y}},\bm{\mathrm{z}},\bm{\mathrm{A}})\big{]}$
$\displaystyle=-\mathbb{E}_{q(\bm{\mathrm{A}})q(\bm{\mathrm{x}},\bm{\mathrm{y}},\bm{\mathrm{z}})}\Big{(}\log
f(\bm{\mathrm{x}},\bm{\mathrm{y}},\bm{\mathrm{z}},\bm{\mathrm{A}})\Big{)}$
(74a)
$\displaystyle=-\mathbb{E}_{q(\bm{\mathrm{x}},\bm{\mathrm{y}},\bm{\mathrm{z}})}\Big{(}f_{A}(\bm{\mathrm{x}},\bm{\mathrm{y}},\bm{\mathrm{z}})\Big{)}$
(74b)
$\displaystyle=-\mathbb{E}_{q(\bm{\mathrm{x}},\bm{\mathrm{y}},\bm{\mathrm{z}})}\Big{[}\sum_{ijk}\bm{\mathrm{x}}_{j}\bm{\mathrm{z}}_{i}\bm{\mathrm{y}}_{k}\big{[}\overline{\log\bm{\mathrm{A}}_{k}}\big{]}_{ij}\Big{)}$
(74c)
$\displaystyle=-\sum_{ijk}\bm{\mathrm{B}}_{ijk}\big{[}\overline{\log\bm{\mathrm{A}}_{k}}\big{]}_{ij}$
(74d)
$\displaystyle=-\sum_{k}tr\Big{(}\bm{\mathrm{B}}_{k}^{T}\overline{\log\bm{\mathrm{A}}_{k}}\Big{)}$
(74e)