Free Energy Projective Simulation (FEPS): Active inference with interpretability

Abstract

In the last decade, the free energy principle (FEP) and active inference (AIF) have achieved many successes connecting conceptual models of learning and cognition to mathematical models of perception and action. This effort is driven by a multidisciplinary interest in understanding aspects of self-organizing complex adaptive systems, including elements of agency. Various reinforcement learning (RL) models performing active inference have been proposed and trained on standard RL tasks using deep neural networks. Recent work has focused on improving such agents’ performance in complex environments by incorporating the latest machine learning techniques. In this paper, we build upon these techniques. Within the constraints imposed by the FEP and AIF, we attempt to model agents in an interpretable way without deep neural networks by introducing Free Energy Projective Simulation (FEPS). Using internal rewards only, FEPS agents build a representation of their partially observable environments with which they interact. Following AIF, the policy to achieve a given task is derived from this world model by minimizing the expected free energy. Leveraging the interpretability of the model, techniques are introduced to deal with long-term goals and reduce prediction errors caused by erroneous hidden state estimation. We test the FEPS model on two RL environments inspired from behavioral biology: a timed response task and a navigation task in a partially observable grid. Our results show that FEPS agents fully resolve the ambiguity of both environments by appropriately contextualizing their observations based on prediction accuracy only. In addition, they infer optimal policies flexibly for any target observation in the environment.

Introduction

A key challenge in both cognitive science and artificial intelligence is understanding cognitive processes in biological systems and applying this knowledge to the development of artificial agents. In this work, we develop an interpretable artificial agent which integrates key aspects of both reinforcement learning (RL) [1] and active inference [2–6]. In RL, an external reward signal is typically used to guide an agent’s behavior while in active inference no such reward signal exists. Instead, agents follow an intrinsic motivation which is rooted in the free energy principle (FEP) [7, 8].

Central to the FEP is the idea that adaptive systems can be modeled as performing an approximate form of Bayesian inference. The outcome of this process minimizes a quantity, called variational free energy (VFE), which lends its name to the FEP. The FEP encompasses various paradigms that align with the goal of Bayesian inference, such as predictive processing [9] and the Bayesian brain hypothesis [10].

One way to implement the FEP is through active inference which involves a planning method for action selection based on an internal model of the environment, referred to as the world model, along with a preference distribution, which encodes the agent’s desired states (e.g., maintaining a certain body temperature). According to the FEP, it is—in principle—always possible to explain the observed behavior of living systems through active inference [11]. The Free Energy Principle has been applied to a wide range of domains, including neuroscience [11–13], psychology [14, 15], behavior biology [4, 16, 17], and machine learning [7, 18–21].

Active inference has recently gained traction in theoretical discussions [5, 22–25] leading to several successful algorithmic implementations [6, 7, 14, 18–21, 26, 27]. Existing implementations are based on methods developed in the context of model-based reinforcement learning and rely mostly on neural networks to represent the world model [6, 7, 14, 18, 21, 26, 28, 29]. As an alternative to neural networks, this work proposes implementing active inference using Projective Simulation (PS) [30].

In PS, memory is organized as a directed graph, which—unlike neural networks—is designed to be interpretable such that individual vertices carry semantic information [31, 32]. The agent’s cognitive processes are modelled as (1) a random walk through the graph (deliberation), and (2) the updating of transition probabilities along the graph’s edges (learning). The primary motivation for using PS over neural networks is its inherent interpretability.

While so far PS has mostly been employed in the model-free RL framework, this work extends PS to the active inference framework by proposing and testing an agent model called Free Energy Projective Simulation (FEPS). FEPS is an active inference model which uses PS for both the world model and the action policy, combining existing and novel methods for training these components. Key features of FEPS include an internal reward signal derived from the prediction accuracy, and the use of two distinct preference distributions—one for learning the world model and another for achieving a specific goal. For the latter purpose, we propose new heuristics to estimate the long-term value of belief states in sight of the goal.

In the following sections, we first introduce active inference (Sect 2) and PS (Sect 3), before presenting the architecture (Sect 4) and algorithmic features (Sect 5) of the FEPS model. We then test FEPS through numerical simulations in biologically inspired environments, focusing on timed responses and navigation (Sect 6). Finally, we conclude with a discussion of future research directions (Sect 7).

1 Related work

FEPS is a learning model that performs an approximate form of Bayesian inference, guided by internal rewards generated by the model itself. It thus falls within the broad category of reinforcement learning models: agents that interact with their environment to acquire data and select actions aimed at maximizing a utility function. More specifically, FEPS agents build upon model-based reinforcement learning (MBRL) since they are equipped with a world model that they use to plan actions. Using an active inference approach to interact with and learn about the environment, FEPS is also intended to model some cognitive processes involving cognitive maps to support adaptive behavior, drawing on existing techniques in this field. In the remainder of this section, we elaborate on how our model relates to the aforementioned methods and what it contributes beyond them.

Model-based Reinforcement learning. agents use a model of the environment to make predictions about outcomes of actions and the resulting expected utility. This not only helps agents to learn from sparse rewards [33], but also makes the training more sample efficient by using the world model as a generative model [34–38]. In particular, some successful approaches such as Dreamer encode sensory inputs into latent space, while a recurrent neural network captures the dynamics in latent space [37, 39, 40]. In spite of their almost human-like performance in some complex environments [39, 40], the world models they learn remain inaccessible to users. Interpretability is therefore a key challenge in order to use MBRL models to reason about cognitive processes underlying and using cognitive maps. While FEPS also encodes observations into latent space and simulates dynamics on the latter, it remains interpretable in three ways: 1) the world model, encoded on a graph-structured memory, is readable by users, 2) the agent’s decision-making process, modeled as a random walk on a graph, is traceable, and 3) the learning process amounts to strengthening associations between nodes on the memory.

RL with internal rewards: Finding efficient exploration strategies constitutes an essential problem in RL [41]. Curiosity mechanisms have been proposed [42], which often rely on a world model. In particular, the policy can be chosen to maximize both internal reward signals emitted by the agent itself, and external goal-oriented rewards distributed by the environment [43, 44]. Popular formulations of internal rewards include information-theoretical quantities, such as prediction errors [45], epistemic uncertainty [38], empowerment [46] or equivalently [47], mutual information [48, 49]. FEPS agents require no external rewards from the environment: thanks to active inference, their actions are internally motivated by information gain, ensuring exploration, and rewards are generated internally from the accuracy of predictions to reinforce relevant transitions in the world model.

Learning with active inference: Leveraging the free energy as a learning signal–such as a loss function [50] or a reward [51]– for RL agents with active inference is advantageous in that it establishes a trade-off between exploration and exploitation, and allows for two types of exploration: goal-directed and epistemic [7, 52]. Internal motivation in active inference is rooted in a preference prior that describes the desirability of states, akin to a utility function in RL, and that can be either given or learned from experts [51], by predicting rewards [53], or calculated iteratively from experience [18, 20, 54]. Most RL applications of active inference model the agents’ distributions with deep neural networks [7, 19, 20, 50, 55–61], thereby inheriting their lack of interpretability. The contribution of FEPS is threefold. 1) We provide an iterative heuristics to calculate the preference prior from the world model, allowing the agent to adapt rapidly to target changes. 2) FEPS is interpretable. 3) Free energy is not used as a learning signal but it instead integrates world model and preferences to plan for actions, such that the agent’s representation of the environment directly influences its behavior.

Cognitive maps: Cognitive maps [62] are mental representations of spatial and abstract spaces [63, 64], learned from experience and supported by episodic memory [65–67]. Clone-structured cognitive graphs (CSCG) [68–72] have recently emerged as biologically-plausible candidates to model cognitive maps: clones of sensory signals, organized on a graph, are allocated to specific paths or contexts in the partially observable environment, thanks to the sequences of observations collected. FEPS makes use of CSCG to structure the memory of the agent and performs Bayesian inference to formulate beliefs about hidden states, as in [72, 73] for example. Relying on CSCG, FEPS uses active inference to add a model of behavior based on the cognitive map of the agent. Instead of the expectation-maximization algorithm, the reinforcement of associations between successive states update the cognitive map of the agent.

Note that there are also parallels between energy-based techniques, such as Boltzmann machines and simulated annealing, and aspects of FEPS. Both use physics-inspired loss functions like the variational free energy to approximate complex probability distributions on graphs. However, the Free Energy Principle has a broader aim: to describe complex adaptive systems. In this framework, the minimization of the VFE is a tool to measure adaptation in terms of uncertainty reduction, and it remains agnostic of any method to implement it. In FEPS, the FEP further constrains the agent’s architecture and its interactions with external degrees of freedom. A final conceptual distinction is that, unlike energy-based graph techniques, the vertices of the graph are treated as fragments of the agent’s experiences in FEPS. The graph thus represents the agent’s memory, which is used to guide decisions.

2 Active inference

The following presentation of active inference is divided into two parts. First, we outline how the world model is constructed and updated in response to sensations (perceptual inference). Next, we explain how planning and action selection are modeled (active inference).

During perceptual inference, through interactions with the environment, the agent develops the world model [37, 38], sometimes referred to as a generative model. The world model [6, 11] is described as a partially observable Markov decision process (POMDP) ($ℬ,𝒮,𝒜,T,L,R$), with belief states $ℬ$, observations $𝒮$, and actions $𝒜$. The transition function T assigns probabilities to transitions between belief states given some action, while L, sometimes referred to as the emission function, defines the probability that some belief state $b\in ℬ$ emits a sensory signal $s\in 𝒮$. While sensory states are observed and shared between the world model and the environment, belief states are part of the world model only and support the encoding of generally different unobservable hidden states in the environment.

In what follows, the time-indexed random variables B_t, S_t, and A_t represent beliefs, observations, and actions at a given time t, with their possible values denoted by the corresponding lower-case letters. The transition function T between times t–1 and t corresponds to the conditional distribution $p(B_t|B_{t-1},A_{t-1})$, while the emission function L at time t is the likelihood $p(S_t|B_t)$. Finally, rewards R are used to learn the transition and emission functions. In this work, rewards are internal (generated by the agent itself) [43] but in general they could also be external (given by the environment). It is further assumed that action A_t is selected based on the current belief state B_t, as represented by the conditional distribution $\pi (A_t|B_t)$, which defines the policy. Combining these elements, the joint distribution of the world model up to time t can be decomposed as follows:

$$p(B_{0:t},A_{0:t-1},S_{0:t})=p(B_0,S_0){\prod }_{t^{\prime }=1}^t\pi (A_{t^{\prime }-1}|B_{t^{\prime }-1})p(S_{t^{\prime }}|B_{t^{\prime }})p(B_{t^{\prime }}|B_{t^{\prime }-1},A_{t^{\prime }-1}),$$ (1)

where the index notation $B_{0:t}=(B_0,B_1,\dots ,B_t)$ is used for sequences of random variables, and $p(B_0,S_0)$ denotes the initial distribution over beliefs and observations. Eq 1 defines the relations between random variables: while the distribution over observation and an action is specified by the current belief state, both the previous belief state and the previous action determine the current belief state.

The transition function $p(B_t|B_{t-1},A_{t-1})$ plays a crucial role in active inference because it has to be learned by the agent. Following the idea of active inference, learning involves updating $p(B_t|B_{t-1},A_{t-1})$ through variational inference, an approximate version of Bayesian inference. Variational inference simplifies Bayesian inference by restricting the set of possible posterior transition functions to a family $\{Q_{\varphi }(B_t|B_{t-1},A_{t-1}){\}}_{\varphi }$, parametrized by some variable ϕ, thus reducing computational complexity. When a new observation $s_t^{\text{env}}$ is received from the environment, the approximate posterior is obtained as the solution to an optimization problem which minimizes the VFE, here conditioned on the specific values b_t−1 and a_t−1 for the random variables B_t−1 and A_t−1 in the previous step:

$$ℱ= {𝔻}_{\mathrm{K}\mathrm{L}}[q_{\varphi }(B_t|b_{t-1},a_{t-1})||p(B_t|b_{t-1},a_{t-1})]+{𝔼}_{b_t~q_{\varphi }(B_t|b_{t-1},a_{t-1})}[-\log (p(s_t^{\text{env}}|b_t)],$$ (2)

where ${𝔼}_{b_t~q_{\varphi }(B_t|b_{t-1},a_{t-1})}[f(b_t)]={\sum }_{b_t}{q}_{\varphi }(b_t|b_{t-1},a_{t-1})f(b_t)$ represents the expectation value over belief states distributed according to the posterior distribution of some function f of b_t. Minimizing the VFE balances two effects: while the first term penalizes drastic changes in the distribution, the second promotes accuracy in the model to predict observations coming from the environment. The Kullback-Leibler distance ${𝔻}_{\mathrm{K}\mathrm{L}}(•)$ represents the dissimilarity between the posterior and prior distributions for the transition function, $q_{\varphi }(B_t|b_{t-1},a_{t-1})$ and $p(B_t|b_{t-1},a_{t-1})$ respectively. The second term is the surprise of the current observation. In order to calculate this surprise, an expectation value is calculated from the posterior distribution $q_{\varphi }(B_t|b_{t-1},a_{t-1})$ over belief states. Realizing that the first term is either larger or equal to zero, the VFE is an upper bound on the surprise raised by the current input from the environment. By updating its model and minimizing the variational free energy, the agent therefore minimizes its surprise. Typically, the prior is replaced by the posterior after collecting a number of observations and implementing the corresponding updates on the posterior.

So far, the model acquisition has been described in terms of perceptual inference: the system has no control over its actions yet, and constructs the internal model based solely on its observations. The reduction of the free energy describes the adaptation of the system to its environment. The distribution over actions, or policy $\pi (A_t|B_t)$, can take any form and is not yet optimized over.

A FEPS agent performs active inference [2] by exploiting the world model that has been learned through perceptual inference in order to plan sequences of future actions. Using the current world model, the agent estimates the future free energy for each of its actions based on the current transition function and implements the one action that minimizes it. This estimate of the free energy for future states does not have a unique form [74] and different formulations can lead to different trade-offs between exploration and exploitation. Here, we use the most common one, denoted in the literature as the Expected Free Energy (EFE):

$${𝒢}_{b_{t-1}}[a]={𝔼}_{b_t,s_t~p(B_t,S_t|b_{t-1},a)}[\log p(b_t|b_{t-1},a)-\log \text{pref}(s_t,b_t|b_{t-1},a)],$$ (3)

$$=-𝐇[B_t|b_{t-1},a]+{𝔼}_{s_t,b_t~p(S_t,B_t|b_{t-1},a)}[{𝒮}^{\text{pref}}(s_t,b_t|b_{t-1},a)]$$ (4)

where $p(B_t,S_t|b_{t-1},a)$ refers to the world model over a single time step, and the surprise of getting outcomes s_t and b_t is denoted ${𝒮}^{\text{pref}}(s_t,b_t|b_{t-1},a)=-\log \text{pref}(s_t,b_t|b_{t-1},a)$. $𝐇(Y|x)=-{\sum }_{y}p(Y=y|X=x)\log p(Y=y|X=x)$ stands for the conditional entropy of the random variable Y conditioned on a specific value X = x. It corresponds to the expected surprise over belief states and is therefore always positive. $\text{pref}(S_t,B_t|b_{t-1},a)$ is a distribution over belief states and observations.

The EFE relies on two measures to determine actions: how much uncertainty is expected, and how useful is the action in order to fulfill a goal. In an active inference setting, an action’s utility refers to its expected capacity to meet some preferences. The first term in Eq 4 is the negative entropy over belief states, related to the expected information gain about the transition function: the larger the entropy, the larger the gain. Maximizing this entropy minimizes the expected free energy and leads to explorative behaviors [74]. Finite entropy can have two causes: (1) the transition to the next state is still uncertain in the world model, or (2) the transition to the next state in the environment is stochastic. In the second case, if entropy were used alone, agents could fall in the so-called curiosity trap [45] by always choosing actions that lead to random outcomes in the environment and inherently result in high entropy values. The second term reflects the utility gained from taking the action under consideration. When minimized, the transitions are expected to match the preferences to a greater extent. Therefore, minimizing the EFE maximizes the utility an agent expects from its action. For this purpose, a new biased distribution, the preference distribution $\text{pref}(S_t,B_t|b_{t-1},a)$, encodes the desirability of some states and observations. The larger the preference for a state, the larger the probability in the preference distribution, and the smaller the associated surprise. This second term encourages exploitation of the world model in order to fulfill the preferences. In a biological agent, these preferences could be of genetic origin (a preference for homeostatic states for example), socially learned (a preference for some type of songs for birds of certain areas, that would be different for the same species in a different place), acquired or externally given. As a result, they can encode a set of states and observations that are favorable to the survival of the agent. As a modeling choice, preferences cover both belief states and observations, and are conditioned on the previous state and action.

For the purpose of learning with an artificial agent, the preferences can naturally encode a task to fulfill as a preferred target state or observation. Choosing the action that minimizes the EFE, the agent selects the transition in its world model that will bring it as close as possible to its preferred states according to its world model. Furthermore, by using the full world model over long periods of time, say from 0 to t, as in Eq 1, the long-term EFE can be calculated to project t steps ahead in the world model and plan for sequences of t actions. Finding the optimal sequence of actions can rely on tree search with pruning [75], or calculating a long-term, discounted, expected free energy quantity [28, 54, 76] for example.

To summarize, according to the free energy principle, an agent that adapts to its surroundings can be modeled as learning a world model of its environment. In this process, the variational free energy is minimized, reflecting the reduction in the surprise the agent experiences about new sensory events. Active inference prescribes a method to plan actions in order to reduce this surprise. In other words, actions are chosen to consolidate the world model. An action is chosen if it minimizes the expected free energy, or equivalently, if it is expected to lead to transitions associated with high uncertainty and high utility. As noticed already in [74], this can be counter-intuitive at first glance, since minimizing the VFE and the EFE respectively minimizes and maximizes uncertainty. This paradox can be resolved by realizing that in active inference, actions target areas in the world model whose outcomes are least predictable in order learn about them and to receive less surprising outcomes in the future, thereby minimizing the VFE in the long run.

3 Projective simulation

Projective Simulation (PS) [30] is a model for embodied reinforcement learning and agency inspired from physics that performs associative learning on a memory encoded as a graph. It is composed of a network of vertices, denoted clips, with a defined structure that gives each clip a semantic meaning, that can be assigned from the start or acquired progressively through past experiences. For example, a clip may represent a sensory state the agent’s sensors are capable of receiving, or it can inherit supplementary semantics from past experiences and reinforcements, through associations to other clips in the graph memory. Directed edges between clips are weighted and can be modified dynamically to learn and adapt to the environment. In particular, PS allows the simulation of percept-action loops to represent previous interactions with the environment and the associated decision process. The resulting graph is the Episodic and Compositional Memory (ECM). When a clip is visited, either because it is currently experienced, or because it is used for deliberation, it is excited. After a first clip was excited, deliberation takes place as a traceable random walk in the ECM originating from this initial clip. It ends when a decoupling criterion is met, which leads to an action on the environment.

For the purpose of solving different environments while retaining interpretability, the ECM can adopt different structures. In order to imitate a percept-action loop, the ECMs are often structured as bipartite graphs, where a first layer contains percepts and the second is composed of actions [77, 78]. In this case, the trained ECM is directly relatable to a policy. For more complex environments, or in order to extract abstract concepts from the percepts, an intermediate layer can be added between the percept and action layers [31, 79]. The ECM can either contain a fixed number of clips, or it can dynamically add and erase some of them when needed [79]. To consider composite percepts and actions, the ECM graph can be replaced by a hypergraph, where each hyperedge connects a set of clips to another set [80].

Each directed edge is equipped with at least one attribute to track learning and allow the agent to react adaptively to the environment. h-values increase as the edges are rewarded. They encode the strength of associations between clips that are useful to fulfill some task and record the learning of the agent. The probability of a transition associated with an edge with h-value h_ij connecting two clips, $c_i\to c_j$, is inferred from the h-values:

$$p(c_j|c_i)=\frac{h_{ij}}{{\sum }_k{h}_{ik}}.$$ (5)

Alternatively, a softmax function can also be implemented to enhance the differences between probabilities, especially in large ECMs. Upon receiving a new percept, the agent deliberates by taking a random walk through the ECM, using the probabilities defined from the h-values to sample a new edge.

As the agent modeled with PS interacts with its environment, it receives rewards that are distributed over the different edges, which changes the corresponding h-values. Specifically, the h-value of the edge $c_i\to c_j$ is updated as follows:

$$h_{ij}^{t+1}=h_{ij}^t-\gamma (h_{ij}^t-h_{ij}^0)+R$$ (6)

where γ is the forgetting parameter, $h_{ij}^0$ is the initial h-value for $c_i\to c_j$ and R is the reward. When the reward is positive, the h-value of the corresponding edge is increased accordingly. If an edge is not visited or does not receive a positive reward, the corresponding h-value decreases back to their initial value thanks to the forgetting mechanism in the second term. In order to accept continuous percepts and to unlock generalization on some task, neural networks have been used to update the h-values in some cases. Training was then implemented by minimizing a loss function [81].

Projective simulation has been tested in multiple tasks, ranging from the standard RL toy environments [77, 78], robotics [82], to animal behavior simulations [83, 84] and quantum computations [85, 86]. Extension of the model include modeling the ECM with a quantum photonic circuit [32] and considering composite concepts in the form of multiple joint excited clips using hypergraphs [80].

4 The free energy projective simulation agent

We combine Projective Simulation, with Active Inference, following the free energy principle’s framework. A FEPS agent is a model-based Projective Simulation agent, where the world model is an ECM with a clone-structured architecture [68, 70]. Consequently, clone clips inherit the semantics of the unique sensory state they relate to, and context creates distinctions between hidden states that emit the same observations. As in the FEP, the agent does not need external rewards. Instead, prediction accuracy, weighted with confidence, is used as a reinforcement signal. The world model is directly exploited by the agent to set the edges’ h-values in the policy with active inference.

4.1 Architecture of the agent

To mimic a system described by the FEP, the agent is composed of two structures: the world model and the policy, each represented by separate graphs with vertices corresponding to random variables, and edges that can be sampled to perform a random walk, as in Fig 1. Consider the agent can perceive N_S sensory states, has N_B possible belief states and a repertoire of N_A actions. Each state is supported by one vertex in a graph. The world model’s vertices can either support belief states or sensory states. The world model is the representation of the environment (see Eq 1) required by active inference. For reasons that will become clear shortly, the ECM of a FEPS agent is made of all vertices, that we call clips, that support belief states, and edges that represent transitions between such clips. A belief state is then formally defined by the excitation configuration of clips at one step of the deliberation. We limit the number of excitations in the ECM at any given time to one. In this case, the excitation configuration on the clips is analogous, from a Bayesian perspective, to a belief– a probability distribution– over clips, where a single clip in the distribution is associated with a probability of 1. Such a vertex deserves the name "clip" because it receives an excitation when the corresponding representation of an event is revisited in order to make a decision. The policy covers all possible conditioned responses coming from any belief state, given the repertoire of actions of the agent.

Fig 1. Architecture and training of an FEPS agent.

a) Architecture of a FEPS agent, with four sensory states (squares) and two possible actions (diamonds). The agent has two main components: the world model and the policy. The world model is composed of vertices representing observations (squares) while clone clips represent all values a belief state can take (circles). As in a clone-structured graph, each clone clip b relates to exactly one observation s and the emission function $p(s|b)$ is deterministic. The clone clips, together with the set of edges between them, form an ECM. A belief state, circled in purple, is designated by an excited clone clip. The weighted edges in the ECM encode the transition function and are trainable with reinforcement: there is one set of edges per action (light and dark turquoise arrows). The belief state in the ECM is an input to the policy, where the probability of sampling an action is a function of the EFE. In turn, the action that was selected determines the edge set to sample from in the world model in order to make a prediction for the next belief state and observation. b) Training of the world model of a FEPS agent. The agent interacts with the environment by receiving observations and implementing actions. When an action a_t is chosen, a corresponding edge $b_t\stackrel{a_t}{\to }{b}_{t+1}$ is sampled in the world model, from the current to the next belief state, conditioned on the action. The observation s_t + 1 associated with the next belief state is the prediction for the next sensory state. Simultaneously, the action is applied to the environment and creates a transition in the hidden states of the environment, $e_t\stackrel{a_t}{\to }{e}_{t+1}$ (bottom, green rectangle). This transition is perceived by the agent through the observation $s_{t+1}^{\mathrm{e}\mathrm{n}\mathrm{v}}$. Finally, the weights of the edges are updated. The reinforcement of an edge is proportional to the number of correct predictions it enabled in a row, as depicted with the thickness of the arrows in the world model. When the agent makes an incorrect prediction (the purple arrow), the reinforcements are applied to the edges that contributed to the trajectory. The last, incorrect, edge is not reinforced.

In the world model, two sets of edges can be traversed at different times during the deliberation: we denote them emission and transition edges respectively. They aim at predicting and explaining sensory signals received from the environment.

In the first set, emission edges $b\to s$ relate belief states and sensory states, modeling the latter as parts of larger, possibly contextualized, hidden states by using clone-structured Hidden Markov Models (HMM) [68–70]. Each clip is bound to a single observation and is denoted a clone clip. A single edge in the emission set carries a non-zero probability for each clip, as shown in Fig 1a. Consequently, this set of edges defines a deterministic likelihood $p(s_t|b_t)={\delta }_{s_t,s(b_t)}$ in the world model in Eq 1 and the agent remains initially agnostic of the dynamical structure of the hidden states $\{e_{\tau }{\}}_{\tau =0}^t$ in the environment. Meanwhile, a clone clip can readily be interpreted as an augmented description of a specific observation. In particular, the additional information can relate to the cause or the context of the observation, such as the previous belief state and previous action, for example. Sampling some sensory state ${\hat{s}}_{t+1}$ amounts to predicting the next observation: if it turns out to coincide with the actual observation perceived from the environment, a reward will be distributed to the edges contributing to the random walk that led to this prediction. For clarity, we denote predicted states with hatted low case letters. We choose to associate each observation to a fixed number $N_{\text{clones}}$ of clone clips, such that $N_B=N_S\times N_{\text{clones}}$. This approach works remarkably well for navigation problems [69, 70]. Transferring the dynamics learned on some set of sensory states to another set is also possible by keeping the transition function unchanged, but redistributing clone clips to the sensory states in the new set [70].

The purpose of the set of transition edges $b_t\stackrel{a_t}{\to }{b}_{t+1}$ in the world model is to encode the presumed dynamics in the environment as transitions between belief states, conditioned on actions. They are represented as edges between clone clips in Fig 1a. In contrast to other sets of edges, transition edges are endorsed with attributes such as h-values that enable learning with reinforcements (as in Eq 6 for example). Therefore, clone clips together with transition edges constitute an ECM for the FEPS agent. From a given clone clip, for each action, a set of edges points to the next possible estimated belief states. The h-values of those edges indicate how certain the agent is that taking a particular action from the current belief state will lead to any of the clone clips in the future. There can be at most N_B edges in each such group of transition edges. The full set of transition edges – $N_{\text{A}}·(N_{\text{S}}·N_{\text{clones}}){ }^2$ edges in total – defines the transition function $p(B_{t+1}|B_t,A_t)$ in the world model in Eq 1 and corresponds to the trained part of the model. Before reinforcement, this distribution is referred to as the prior. The posterior is the updated version of the transition function, after the agent distributed rewards to the relevant transition edges. The posterior is labeled $q(B_{t+1}|B_t,A_t)$.

The final component of the agent guides its behavior. The policy $\pi (A_t|B_t)$ is modeled as a separate graph with two layers of N_B and N_A clips respectively. Given the clone clip corresponding to belief state b_t is excited in the world model, an action a_t is sampled to be applied to the environment, based on how much surprise it expects from this decision. Each edge $b_t\to a_t$ is weighted with the expected free energy ${𝒢}_{b_t}[a_t]$, that defines the policy.

4.2 Reinforcement with prediction accuracy

Each interaction step with the environment involves a deliberation over three states: (1) the next belief state is proposed, (2) the next sensory state is predicted and (3) an action is chosen. The agent excites a belief state b_t + 1 it believes it will transition to, given its current action a_t, by sampling a transition edge $b_t\stackrel{a_t}{\to }{b}_{t+1}$. From there, the agent makes a prediction ${\hat{s}}_{t+1}$ about the next sensory state. Meanwhile, an action a_t + 1 is selected in the policy and it is applied to the hidden state in the environment, that emits a new observation. The interaction step ends by comparing the predicted and perceived sensory states, ${\hat{s}}_{t+1}$ and $s_{t+1}^{\text{env}}$.

The world model is trained without external rewards, and reinforcement is instead based on matching predictions and observations. We call trajectory a sequence of transitions that led to correct predictions about sensory states. To record the trajectory, transition edges are equipped with a new attribute: the confidence, f. Initialized at zero, it increases for all transitions in the trajectory every time the prediction ${\hat{s}}_{t+1}$ and the actual sensory state $s_{t+1}^{\text{env}}$ coincide. The more subsequent predictions an edge enabled, the higher the confidence for that edge: it reflects the number of correct predictions the edge enabled until the end of the trajectory. Formally, a trajectory τ is a sequence of transitions whose predictions were confirmed by the observations from the environment. If at step n in the t-th trajectory ${\tau }_t$ the sensory prediction was accurate, confidence is enhanced for all edges $i\to j$ in ${\tau }_t$:

$$f_{ij}^{({\tau }_t),n}=\{\begin{array}{cc}{f}_{ij}^{({\tau }_t),n-1}+1\hfill & \text{if} b_i\to b_j\in {\tau }_t\hfill \\ 0\hfill & \text{otherwise.}\hfill \end{array}$$ (7)

When the prediction and observation do not match, the trajectory ${\tau }_t$ is interrupted, and the rewards are distributed to the transition edges’ h-values proportionally to the corresponding confidence:

$$h_{ij}^t=h_{ij}^{t-1}-\gamma (h_{ij}^{t-1}-h_{ij}^0)+f_{ij}^{({\tau }_t)}R$$ (8)

where $h_{ij}^{t-1}$ is the h-value at the end of the previous trajectory, $h_{ij}^0$ the initial h-value of the edge, and R scales the reinforcement of the edges. Confidence values are reinitialized at zero to start the next trajectory. This mechanism provides a built-in learning schedule such that the scale of the reinforcement signals grows progressively: rewards are initially small when trajectories are short, and they become larger when transitions are accurately captured in the model. For the world model to yield accurate predictions, at most each of the $N_A·(N_S·N_{\text{clones}}){ }^2$ transition edges must have been visited and reinforced appropriately. As a result, the computational cost of learning scales at least quadratically in the number of observations and clones per observation. During the deliberations, states are sampled from the prior ECM that did not receive the rewards yet, while the posterior ECM is updated with confidence and rewards at the end of the trajectory. It is equivalent to sampling states from the prior ECM, but updating the posterior ECM with the rewards R for all edges in the trajectory every time a prediction was verified by the observation in the environment. Metaphorically speaking, this mechanism is analogous to layers of snow accumulating in time on salient features. At the end of a trajectory, the snow is cleared away, bringing all salient points back at an equal level.

To complete the update of the FEPS agent, the policy is modified according to the EFE inferred from the new world model and can be adjusted to make behaviors more or less explorative. In particular, h-values of an edge $b_i\to a_j$ are set to the expected free energy in Eq 4 with a world model conditioned on b_i and a_j. Each h-value directly carries the surprise expected from traversing the corresponding edge. As in [6], the policy is defined using a softmax function:

$$\pi (a_j|b_i)=\text{softmax}(\zeta {𝒢}_{b_i}[a_j]).$$ (9)

where ζ is a (real-valued) scaling parameter and ${𝒢}_{b_i}[a_j]$ is the value of the EFE for action a_j coming from state b_i. In active inference, ζ is typically negative. When it becomes more negative, actions associated with small EFE receive a large probability. More specifically, looking at the decomposition of the EFE in Eq 4, actions associated with larger entropies, that is lower certainty, together with higher chances of landing on preferable states or observations, become more attractive during the deliberation. In contrast, if the scaling parameter is positive, large EFE yield large probabilities in the policy, and actions with high certainty but also less chances of meeting preferences are more likely to be sampled. In this case, the agent implements a non-explorative policy that is confined to a known region of the environment at the expense of reaching the preferred states. When $\zeta =0$, the policy is uniform, and the agent has no bias towards certainty nor utility.

5 Algorithmic features of the FEPS

The FEPS can be augmented with a number of techniques that take advantage of the world model and its interpretability. During learning, the internal model can be leveraged to identify transitions that are instrumental to gain information or to get closer to a preferred state. Furthermore, the performance of an agent in completing a task can be enhanced by evaluating the correct belief state accurately and quickly. Since the policy depends on the EFE, the preference distribution can be tuned according to the task: to seek information to complete the model, or to complete a given goal. Therefore, we propose to separate training into two phases, depending on how the preference distribution is constructed. We introduce a belief state estimation scheme that distributes belief states over multiple clone clips and eliminates those that are incompatible with new observations.

5.1 Leveraging preferences as a learning tool

So far, the preference distribution entering the EFE was not defined. One can optionally leverage it to define a goal in the environment, be it for the purpose of gaining information, or to solve an actual task. Therefore, we propose to separate learning into two tasks: model the environment and attain a goal in it. During the first phase, which we denote the exploration phase, the agent explores the environment without a prescribed goal. Instead, actions whose outcomes are expected to reduce prediction errors should be favorized. This phase spreads over multiple episodes and relies on interacting with the environment. In contrast to the exploration phase, the second phase is dedicated to learning to complete a given task. For this purpose, we designed an algorithm to infer a goal-oriented policy from the world model in a single step and without further information.

5.1.1 Seek information gain about the world model.

Before designating any task bound to the environment as a preference, we investigate whether the preference distribution can be used to incentivize actions that minimize prediction errors, according to the current world model. This is directly related to the minimization of the VFE. Specifically, preferences should encourage the agent to seek transitions the world model associates with certainty – or equivalently with high probabilities, irrespective of actions. Sequences of interaction steps with the environment guided by this preference distribution belong to the exploration phase. To reflect the preference for highly probable transitions in the world model regardless of the action chosen, the preference distribution is constructed as the marginal of the world model over actions:

$$\text{pref}(B_{t+1},S_{t+1}|B_t)={\sum }_a\pi (A_t=a|B_t)p(B_{t+1}|B_t,A_t=a)p(S_{t+1}|B_{t+1})$$ (10)

$$=p(B_{t+1},S_{t+1}|B_t).$$ (11)

Plugging this distribution into the expected free energy evaluated for an action a in Eq 4 results in the following:

$${𝒢}_{b_t}[a_t]={\sum }_{b_{t+1},s_{t+1}}p(b_{t+1},s_{t+1}|b_t,a_t)\log \frac{p(b_{t+1}|b_t,a_t)}{{\sum }_a\pi (a|b_t)p(b_{t+1}|b_t,a)p(s_{t+1}|b_{t+1})}$$ (12)

$$={𝔻}_{\mathrm{K}\mathrm{L}}[p(B_{t+1}|b_t,a_t)||p(B_{t+1}|b_t)]$$ (13)

$$=\mathbf{I}\mathbf{G}(B_{t+1},A_t=a_t),$$ (14)

where $\mathbf{I}\mathbf{G}(X,Y=y)={𝔻}_{\mathrm{K}\mathrm{L}}[p(X|Y=y)||p(X)]$ is the information gain about the random variable X when the value y for the second random variable Y is known. The dependency on the observations dropped from the first to second line thanks to the constraints the clone structure imposes on the emission function. A complete derivation of this formula is provided in S2 Appendix.

If we follow the conventional formulation of active inference as in Sect 2, the agent should increase the probability of sampling an action that minimizes this EFE. Doing so during the exploration phase, the agent would therefore seek actions it estimates will yield the lowest information gain about belief states. As a result, the agent would stay in a region of the environment where it predicts it will receive the least surprise, according to its world model. This situation is sometimes referred to as the Dark Room problem [87]: an agent that adapts by minimizing its surprise about observations would act to stay in a dark room instead of using actions to explore other places that may be more surprising, but also more favorable to its survival, because all observations there would be predictable.

There is, however, an easy solution to this problem for FEPS. In order to avoid the dark room problem and to select actions that are expected to improve the model of the environment, the scaling parameter ζ in Eq 9 can be set to a positive value. In this case, the larger the EFE associated to an action – and the estimated information gain about the next belief state, the larger the probability of this action in the policy. The scaling parameter ζ can be understood as a way to determine how greedy an agent is in its exploration, or how strongly the information gain associated with each action influences its behavior. We call exploration phase the interaction steps in which the agent samples its actions from such a policy.

5.1.2 Task-oriented behavior by inferring preferences on belief states.

At the end of the exploration phase, a task is designated by encoding the associated targets with a high probability in the preference distribution. From there, the agent can plan, that is sample a sequence of actions to follow to achieve the goal. While the target is identified as a sensory state for the FEPS, transitions that are useful to reach it are deduced from the world model. In our framework, updating the policy takes a single step and does not require further interaction with the environment.

Though active inference commonly determines the behavior by planning sequences of actions, it becomes expensive for distant horizons T_h. A sequence of T_h actions must be chosen out of $N_{\text{actions}}^{T_h}$ possible combinations, by evaluating the EFE in a space of $(N_B\times N_S){ }^{T_h}$ possible outcomes for each sequence. Methods such as habitual tree search or sophisticated inference have been developed to mitigate this scaling issue [19, 76]. Alternative approaches are presented in [6].

Instead of planning by evaluating the generative distribution over all possible future sequences of outcomes, we propose to encode the long-term value of a state directly into the preference distribution. Our scheme is reminiscent of iterative value estimation [88, 89] and successor representation [35, 90, 91], in that it estimates a value function provided some expectations about occupancies of states in the future, either acquired by experience with reinforcement for example, or by inverting a learned transition function. In contrast to searching a tree of future sequences of actions, this method does not rely on mental time traveling [92], to the extent that agents do not simulate possible future scenarios. Instead, they are “stuck in time", and infer preferences in one go from stochastic quantities stored in the world model, in contrast to [20]. Our method also departs from sophisticated inference [28, 76]: instead of bootstrapping expected free energies in time, we bootstrap preferences over belief states and calculate the EFE only once.

We model the preference distribution to factorize over sensory and belief states [93], and we condition it on the current belief state b_t:

$$\text{pref}(S_{t+1}=s,B_{t+1}=b|b_t)=\text{pref}(s) \text{pref}(b|b_t).$$ (15)

The first part, $\text{pref}(s)$, is an absolute preference distribution over sensory states, that is independent of where the agent believes it is in the world model. More specifically, if s  is the target state for an observation, a probability $\text{pref}(S_{t+1}=s^{*})=p^{*}$ is associated to it. All other observations are given a uniform distribution $\text{pref}(S_{t+1}\ne s^{*})=(1-p^{*})/(|𝒮|-1)$. Since the target observation is given by design in the absolute preference distribution, it plays a role analogous to that of a reward function in reinforcement learning. In addition, the policy derived from Eq 9 aims at maximizing the probability of receiving the target observation: the more probable a transition to the target observation, the lower the EFE and the larger the probability of taking the corresponding action, thereby maximizing rewards. The second part $\text{pref}(b|b_t)$ then reflects look-ahead preferences over belief states and how useful an agent estimates a transition to be in order to satisfy its absolute preferences. In other words, the utility of belief states over longer horizons is inferred from the value associated with the observations they can transition to. This way, even if the goal might be far away in the world model, the preference for the target observation propagates to intermediate belief states that contribute to reaching it. Metaphorically speaking, the preference for a target observation propagates to the belief states that are useful to get to that target observation. For example, consider an animal in a maze. The target is manifested with a high preference towards observing food. However, the preference distribution over the locations does not indicate how to reach the food. To remedy this, the agent infers the value of belief states in order to reach the target: if a transition to some location brings the animal closer to the food, it is assigned a higher probability in the preference distribution. This way, the preference distribution highlights the path of relevant actions to the target.

We propose a heuristic algorithm to estimate the look-ahead preference distribution $\text{pref}(B_{t+1}|b_t)$ over transitions in the world model. The algorithm can update the policy $K_{\text{pref}}$ times if needed, to refine the preference distribution. The initial policy ${\pi }^{(0)}(A_t|B_t)$ is uniform. During each update k, two quantities are calculated: (1) the look-ahead preference distribution results from the value of each transition, that estimates how useful a transition is to reach a target within a prediction horizon $T_h$, and (2) the policy ${\pi }^{(k+1)}$ is calculated with Eq 9 and the latest preference distribution.

To initialize each update step k, the policy ${\pi }^{(k)}$ is used together with the world model to evaluate how easy a belief state can be reached from the current state, in a distribution we denote reachability:

$$r^{(k)}(b_{t+1}|b_t)={\sum }_{a}p(b_{t+1}|b_t,a){\pi }^{(k)}(a|b_t).$$ (16)

The reachability of b_t + 1 coming from b_t is large if there exists transitions $b_t\stackrel{a}{\to }{b}_{t+1}$ associated with high probabilities in the world model and the corresponding actions have high chances to be sampled in the policy.

In addition, the initial value of a belief state equals the value of the observation it corresponds to:

$$v_0(b_t)={\sum }_s\text{pref}(s)p(s|b_t).$$ (17)

For the clone-structured model, this sum reduces to a single term. At this stage, the only belief states associated with a high value are those that represent a target observation in the absolute preference distribution.

Next, for each iteration $n\in [1,T_h]$, where T_h is a prediction horizon, the value of a belief state is increased if it can lead to transitions that are useful to reach the target within n steps in the environment:

$$v_n(b_{t+1})=max\{v_{n-1}(b_{t+1}), {max}_{b^{+}}\{{\beta }^{n-1}{r}^{(k)}(b^{+}|b_{t+1})v_{n-1}(b^{+})\}\},$$ (18)

for $0\le \beta \le 1$ some discount factor that makes the value of the state decrease with the number of steps between this state and the target. At each n, the value of a belief state b_t + 1 can either keep its previous value $v_{n-1}(b_{t+1})$, or it can take the discounted value of the best state b⁺ it can transition to. As a result, the value of a state can only increase. If n = 1, using this value function can point at the right decision while being one step away from s . However, it does not incentivize the correct action when starting further away from the goal. To mitigate this effect, a state can inherit the value of the states it can reach over larger time scales n > 1, thereby propagating the preference for the target to more distant but useful belief state states.

When the prediction horizon is reached, a transition $b_t\to b_{t+1}$ is associated with the following probability in the look-ahead preference distribution:

$${\text{pref}}^{(k)}(b_{t+1}|b_t)={\delta }_{b_{t+1}\in \mathrm{c}\mathrm{h}(b_t)} \frac{v_{n=T_h}(b_{t+1})}{{\sum }_{b_{t+1}^{\prime }\in \mathrm{c}\mathrm{h}(b_t)}{v}_{n=T_h}(b_{t+1}^{\prime })}$$ (19)

where the set $ch(b_t)=\{b^{+} | r^{(k)}(b^{+}|b_t)>\overline{r^{(k)}(B_{t+1}|b_t)}\}$ contains the children of b_t, or the states that are easily reachable from b_t. $\overline{r^{(k)}(B_{t+1}|b_t)}$ is the mean reachability of belief states coming from state b_t.

Finally, to conclude the k-th update step in the algorithm, the policy ${\pi }^{(k+1)}$ is calculated by using the preference distribution $\text{pref}(S_{t+1}){\text{pref}}^{(k)}(B_{t+1}|B_t)$ in the expected free energy, as in Eq 9. We show a possible look-ahead preference distribution in the world model in Fig 2. The computational time cost of this planning procedure scales linearly with the product of the prediction horizon T_h and the number of update iterations $K_{\mathrm{p}\mathrm{r}\mathrm{e}\mathrm{f}}$. Keeping the number of clones per observation constant, it also scales as $𝒪(N_B^3)$. This cost can easily become constant if the procedure is parallelized.

Fig 2. Estimation of belief states in superposition, after the world model has been trained.

To minimize its prediction error due to faulty belief state estimation, an agent considers multiple clone clips as candidate belief states simultaneously. For the initial observation, $s_1^{\mathrm{e}\mathrm{n}\mathrm{v}}$ (on the left), the agent includes all corresponding clone clips $\{b_t^i{\}}_{i=1}^3$ to its hypothesis, as depicted on the right. Conditioned on the chosen action, a clone clip is sampled for each candidate belief state to represent the next one. Finally, all clone clips that are incompatible with the observation from the environment are eliminated from the hypothesis. The clips that remain become the current candidate belief states. In the world model, the thickness of the arrows represents the look-ahead preferences: the larger the arrow, the more advantageous is the transition in order to reach the target observation, s₄ in this case.

5.2 Delineate belief states for the same observation

In spite of prescribing a method to sample actions, active inference does not include techniques to efficiently choose belief states when multiple of them could explain the current observation. In particular, models involving neural networks lack the interpretability to design suitable belief state selection methods. Therefore, letting the agent learn a world model with a clone-structured HMM has advantages beyond planning. We propose a technique to evaluate belief states in superposition, depicted in Fig 2. When placed in an environment and receiving its first observation, the agent makes an initial hypothesis about its belief state by distributing an excitation to any clone clip compatible with the observation, as shown in Fig 1c. At each step, an action is sampled from the policy for each excited clone clip. The resulting frequencies define a new distribution, from which an action is sampled before it is applied to the environment. For each compatible clone clip, the agent samples a new clip to represent the belief state it anticipates it would transition to if the clone clip under consideration stands for the correct belief state. The excitation jumps onto the new clip. After applying the action to the environment and receiving the resulting sensory signal, the agent takes away the excitation on any clip that does not match the current observation. Depending on the structure of the environment and the number of clones for each observation, the agent progressively narrows down its candidate belief states to a single clone clip, in spite of the initial uncertainty. This elimination process can be thought of as successively applying Bayes’ rule to the clone clips while keeping the posterior uniform as the agent receives more observations: thanks to the deterministic likelihood imposed by the clone structure of the ECM, the probability of any incompatible candidate clone clip becomes zero. In the event that the world model is imperfect and the agent has eliminated excitations on all clips, it starts its hypothesis over, and considers all the clone clips of its current, unpredicted observation as candidate belief states. This mechanism allows the agent to evolve in environments with ambiguous observations in the absence of more contextual information about its initial conditions. The computational cost of this method depends on the structure of the environment. More specifically, the maximum number of steps required to disambiguate hidden states is the length of the longest sequence of actions that would produce the same observation sequences when starting from those hidden states.

6 Numerical results

In this section, we present a numerical analysis of the model on environments inspired from behavioral psychology tasks, namely a timed response experiment in a Skinner box and a navigation task to forage for food. The parameters used for the simulations are given in S1 Table in the Appendix.

6.1 The timed response task

6.1.1 Learn short-term associations.

The timed response task is a minimal environment for an agent to learn to contextualize its observations with past states and actions when the sensory signals emitted by the environment are ambiguous. This environment simulates an animal standing in front of a door that can be opened with a lever. The goal is for the agent to learn a conditioned response and to press a lever at the right time in order to access food. The environment’s MDP is depicted in Fig 3. For this task, the observations combine two sensory inputs: $𝒮=$ { (light off, hungry), (light off, satiated), (light on, hungry)}. Since food can be consumed only when the light is off, the observation (light on, satiated) in excluded from the set. The actions are $𝒜=$ {wait, press the lever}. The environment is initialized in the (hidden) state E₀, that emits observation (light off, hungry). From there, the light turns on, regardless of the action taken by the agent. Once the light has turned on, the agent must learn to wait one step before pressing the lever. If it does so too early, it gets back to the initial state. If the agent activates the lever on time, the environment transitions to state $E_0^{\prime }$, where the light is off, but the agent is satiated.

Fig 3. MDP for the timed response environment.

This environment has four hidden states. The observations are compositional and contain information that are both external (light on or off) and internal (hungry or satiated) to the agent. Arrows correspond to the transitions the actions the agent can result in. In this environment, the agent can either wait or press a lever. In order to complete the task, the agent must reach $E_0^{\prime }$ and feel satiated. The only way to do so is to follow the actions marked with thicker arrows. The observation (light on, hungry) is called ambiguous because it can be emitted by two hidden states E₁ and E₂ that can only be distinguished with context.

For the simulations, we give each observation two clones. It makes it possible to accommodate enough candidate belief states to model up to two ambiguous hidden states that emit the same sensory signal, enabling the agent to adapt its policy to a one-step waiting time between the light turning on and the food being accessible for example. When observations are not ambiguous and can be emitted only by a single hidden state in the environment, some belief states become redundant. This redundancy can make the training more challenging to the extent that the agent has to find a convention and adapt its model accordingly before it can account for all transitions in the environment faithfully. The larger the waiting time n, the more clones might be necessary to learn. We train 100 agents for 4000 episodes of 80 steps in this environment and we test two scenarios. In the first, the agent is directly given the preference for the target observation (light off, satiated). The EFE is then scaled with a scaling parameter of $\zeta =-1$ in the policy in Eq 9. In the second scenario, we test whether the agent can learn more efficiently if it explores aimlessly the environment for the same number of episodes without preference for the target before adapting its policy to the task, with a scaling parameter of 0.

6.1.2 Simulation results.

The timed response environment is a minimal testbed for the FEPS, where some hidden states are not uniquely defined by the observation they emit, but also by the recent past. The agent must learn two types of belief states. While clones for (light off, hungry) and (light off, satiated) only support transitions that are independent of the actions the agent takes, the clones for (light on, hungry) must appropriately use information about the previous observation and action to be contextualized and distinguished. Some results are reported in Fig 4.

Fig 4. Training FEPS agents for the timed response task.

a) Evolution of the variational free energy (top) defined in Eq 2 and expected free energy as in Eq 4 (bottom) during the training, averaged over 100 agents and a time window of 100 episodes. At each step, the VFE depends on the specific belief states and actions that were sampled. Two types of training are compared: a first set, “task" in dark purple, learned the model with a preference to fulfill a task in the environment, while the second set, “explore", in green, experimented aimlessly in the environment with a uniform policy before switching to the task. Both were trained for 4000 steps before being tested on the task. The best and worst agents are represented with dashed and dotted lines respectively, and examples of individual agents were traced with transparent lines, full for task-oriented agents, and dashed for the exploring ones. When the VFE converges to its minimal value, the world model is precise enough for most belief states to make planning possible. As expected from the values chosen for the scaling parameter ζ, agents select actions that minimize and maximize the EFE for task-oriented and exploring agents respectively. The EFE of exploring agents plateaus quickly at the limit derived for S3 Appendix. b) World model learned by one of the agents trained on the task, where each circle is a belief state, whose observation is denoted by its colors and label. The numbers at the center of the circles are the clone indices for each clone clip. Arrows indicate the transitions learned in the world model, red for action “press the lever", and blue for “waiting". Dashed lines indicate that both actions lead to the same transition. The weight on the arrow indicates its probability in the world model. Stars mark transitions that were identified as useful to achieve the goal with a probability of 1 in the preference distribution. The policy is indicated by the thickness of the arrows, where a thick arrow corresponds to probabilities close to 1, and thinner close to 0.5.

During training, regardless of the strategy that the agent uses to resolve the environment, all agents follow a similar learning pattern. The acquisition of the world model happens in stages, where the transition from one to the next is manifested in a steep increase in the length of the trajectories of correct predictions, or equivalently, as a steep decrease in the free energy, as in Fig 4a. First, agents quickly eliminate transitions that are impossible in the environment, leading to an initial drop in the VFE. For example, as in Fig 3, a direct transition from (light off, hungry) to (light off, satiated) is incompatible with the timed response task. During the second phase, the number of rewards the agent can collect is limited by the absence of convention on the context-sensitive belief states. Therefore, a plateau is observed on the VFE. A final sharp decrease in the VFE signals the adoption of a convention between clones to accurately disentangle ambiguous hidden states that cannot be told apart with observations alone. The EFE evolves as expected during the training: it decreases for the task-oriented agents, while it rapidly plateaus to its limit (see S3 Appendix) in with an exploration phase. However, in contrast to the VFE, the convergence of the EFE to its asymptote value does not reflect the fact that the model is good enough to make accurate predictions.

For most agents, training with or without a preference for the target in order to construct the world model does not influence the final model. As in Fig 4a, the two best agents, that is, those that converge to free energy values below 1 the earliest, share similar learning curves. The averaged behaviors remain fairly identical, except for the length of the second stage of the learning phase, where the agents have yet to adopt a successful convention. In particular, in the absence of aimless experimentation with the environment, this second phase can last longer, such that convergence is not attained for a fraction of the agents trained in this way.

After convergence of the free energy, or equivalently of the number of successful predictions, the model has collapsed on a single representation of the environment, where the ambiguous observation is contextualized by the previous belief state and action. For example, the observation (light on, hungry) is divided into two belief states with different uses, reflecting the two hidden states in the environment. The corresponding clone clips bear more information than the observation to which they are linked. In Fig 4b, clone 2 for (light on, hungry) is the first hidden state encountered when the light turns on, and clone 1 can only be accessed by waiting from clone 2. When more clones are provided than necessary, two possibilities can arise: either the agent uses all the clones as duplicates of the same hidden state in the environment (as for clones of (light off, hungry)), or a single clone is reachable (clone 2 of (light off, satiated)) and the other is never trained because it was not visited. When multiple clone clips participate equally in the representation of the observation, the actions that lead to it also split with equal probability (as derived in S3 Appendix).

Finally, for the agents that converged to an appropriate representation of the environment, the look-ahead preferences inferred from the world model result in an optimal policy when using the EFE. A visualization of the preferences and policy is provided in Fig 4. We chose a prediction horizon of 2 steps, and a single iteration was required to calculate appropriate look-ahead preferences for the transitions. The transitions with maximal preference were kept so that, coming from a belief state b_t, a single belief state b  is more preferable than the others. The agents were tested for 1000 rounds, each starting in E₀ in the environment. For a prediction horizon of a single step, the edge between clone 1 of (light on, hungry) and clone 2 of (light off, satiated) would be the only edge to carry a probability larger than other transitions in the preference distribution. Thanks to the look-ahead preferences (and the adoption of a convention between belief states), waiting between clones 1 and 2 of (light on, hungry) is also preferred. The policy that results from this preference distribution is optimal to solve the task, in spite of initially having no hints about the target prior to the last transition.

6.2 Navigation task in a partially observable grid

6.2.1 Long-term planning in an ambiguous environment with symmetry.

The FEPS is further challenged in a larger navigation task, where observations are shared among hidden states, and multiple sequences of actions can emit the same observations, due to the symmetry of the environment. In order to disentangle the hidden states, the agent must use long-term information about its past observations and actions to contextualize its current state in a way that is consistent across actions. In this environment, food is hidden in one position in a grid. Locations in a 3 by 3 grid world are associated with smell intensities, denoted by their integer intensities from 0 to 3 from the lower left to the upper right corners, according to their closeness to the food.

The world model is provided with 3 clones for each of the four observations, and the behavior repertoire comprises directional actions $𝒜=$ {go right, go left, go up, go down}. 30 agents were trained on different hyperparameters and preferences, for a total of 40000 episodes of 80 steps. The hyperparameters are provided in S1 Table in the Appendix. Since this environment is larger and more complex, we test different training techniques, by changing the preference distributions and varying the scaling parameter. In particular, we test two scenarios. In the first the agent is trained with preferences pointing at a target in the environment, while in the second, the preferences are identified with the marginal of the world model over actions, which incentivizes dissociating the effect of actions on the environment. While in the first case, the agent tries to directly solve the task, in the second, the agent explores first before it learns the task. For each scenario, different scaling parameters were implemented, ranging from -3 to 3 for the task-oriented training, and from 0 to 3 with an exploration phase. Note that choosing a scaling parameter of 0 while training either directly on the task or with a preceding exploration phase results in the same uniform policy.

6.3 Simulation results.

Training task-oriented and exploring agents. Compared to agents trained with a uniform policy, agents using policies inferred from the EFE can achieve longer trajectories. This is apparent in Fig 5a), where with a scaling parameter of 0, the trajectories of agents reach a length of 60 on average by the end of the training, whereas given a proper tuning, agents using the EFE to define the policy can predict up to 69 transitions, regardless of the training method. Adapting the scaling parameter to the preference distribution is crucial for the agents to learn the model and query actions that result in longer trajectories. For task-oriented agents, as expected from Eq 4, a negative scaling parameter at -3 is optimal, whereas +1 yields the best results for the exploring agents. For both training methods, setting it at +3 results in the shortest trajectories, that are 28 and 50 steps long, respectively. In the case of task-oriented agents, a possible explanation is that in this regime and when the preference distribution designates a target in the environment, the policy that is derived from the EFE in Eq 4 minimizes the drive of the agent to explore and to fulfill any preferences in the environment. Therefore, the agent has no incentive to go outside of a region of the environment where it can predict its observations, and does not try to learn the rest of it. In contrast, a scaling parameter $\zeta =+3$ makes exploring agents very greedy. It is then possible that previously explored regions are not visited often enough to reinforce the correct associations in the world model in the long-run. Finally, our simulation shows that for the best parameter settings, the task-oriented agents converge on average faster than the exploring ones: while 17500 episodes are needed to predict 65 steps for the former, the latter crosses this milestone at 28000 episodes.

Fig 5. Training results for the grid world environment.

a) Evolution of the length of the trajectories during the training, for different scaling parameters ranging from −3 to 3, and different preference distributions: the agent can either learn to complete the task from the start (“task”), or first explore the grid (“explore”). We represent the running averages over a time window of 1500 steps of the lengths of trajectories, averaged over 30 agents. These lengths depend on the specific belief states and actions sampled by the agents. b) Evolution of the variational and expected free energies during the training for the two best settings in a): the task-oriented preferences are paired with a scaling parameter of −3, and the exploration preferences with a parameter of +1. The thick lines represent the running average of energies over a time window of 1500 steps, averaged over all 30 agents, while the dotted and dashed lines stand for the best and worst agents, that is the agents whose VFE converge first and last, respectively, to a minimum. The transparent lines indicate the behavior of agents selected randomly: these lines are full for task-oriented agents and dashed for exploring agents. c) Comparison of the accuracy of the model to make predictions between the two best parameter settings in a). Blue and red colors indicate explorative and task-oriented agents respectively. Darker and lighter shades distinguish the belief state estimation methods, “sup.” designating superposition. The policies of the agents are uniform, such that the actions yielding most certain outcomes cannot be relied upon to validate predictions. Two belief state estimation techniques are tested. Belief state estimation samples a single clone clip at a time, whereas the evaluation of belief states in superposition allows multiple clone clips to represent candidate belief states simultaneously, as long as they produce predictions that are compatible with the next observation. Each individual agent is tested over 1000 trajectories of at most 80 steps.

Looking at the evolution of the free energy during the training in Fig 5b), one sees that in both cases, the VFE decreases as the length of trajectories increases. It looks as if agents minimize their free energy by maintaining a world model and improving their predictions about future sensory states. This matches the free energy principle. Comparing individual learning curves, exploring agents have more diverse behaviors than those trained on the task. We define the best and worst agents as those converging the earliest and the latest to a free energy minimum, respectively. The best exploring agents can converge faster than the task-oriented ones, whereas the worst exploring agent lands on a higher free energy by the end of the training. However, the free energy of the worst exploring agent either goes down or plateaus, but it does not rise again, as for the worst agent directly trained on the task. When the training scenarios influence the scale of the EFE, no significant variations around the average behavior is observed in the EFE to distinguish successful agents from those that did not converge. In particular, the EFE exploring agents converge very quickly to its asymptotic value, regardless of the accuracy of the world model.

In order to fairly compare the models obtained from each training method, note that whenever the length of trajectories or the free energy are optimized, two contributions are at play. On the one hand, training the world model to capture the environment more accurately decreases the uncertainty about the transitions following actions, and therefore improves the length of the trajectories and the free energy. But on the other hand, as long as the policy deviates from the uniform distribution, it influences how much uncertainty the agents seek. Therefore, in order to evaluate how much the world model helps making accurate predictions, we test the length of trajectories with a uniform policy for the agents trained in the two best parameter and preference settings. The trajectory lengths averaged over all agents trained with the respective parameters are provided in Fig 5c).

With the belief state estimation protocol where a single clone clip is excited at each step, trajectories achieved by FEPS agents are analogous in length for both types of model-learning strategies, either exploration- or task-oriented. More precisely, half of the agents are able to predict half of the maximum length for a trajectory. Similarly, extreme cases are comparable.

Estimating belief states in superposition, i.e. considering all clone clips that are compatible with the current trajectory at each step, doubles the length of trajectories for both training strategies. Indeed, agents of both types reach near optimal length. For explorative agents, the top 50% of the agents achieves the same performance, showing a slight improvement compared to task-oriented agents. We expect that this improvement will become greater when the environments reach larger scales, and when the target given to task-oriented agents become more ambiguous. In addition, the ability of the agents to adapt to different goals in the environment would also be affected.

The agent learns an interpretable map of the environment. After the free energy converges, the belief states in the world model no longer stand only for the observation they relate to, but also convey information about the path they can appear in. As an example, the world model of an agent trained directly on the task is represented in Fig 6a). Because the environment the agent was trained in is deterministic, actions tend to map onto a single belief state when the number of clones for an observation matches the number of hidden states in the environment.

Fig 6. Interpretability of the world model and robustness to reward reevaluation for the grid world environment.

a) Example of a world model learned by an agent directly trained on the task, with the target positioned in the top right corner of the grid. The circles represent the belief states as in Fig 1, numbered with clone indices, and colored according to the observation they relate to in the grid. The arrows stand for the transition probabilities: the thicker the arrow, the more the agent believes taking an action from a belief state will lead to the state the arrow points at. b) Mapping of the world model in a) onto the grid: each clone is associated with exactly one cell, and can be interpreted as a single, specific hidden state in the environment. Stars stand for the preferred transitions, and the arrows for the policy resulting from these preferences. c) Median number of steps from each initial position in the grid in order to reach the target, where belief states are estimated in superposition. Two targets are given to the agent and symbolized with a gray triangle: reach observation of 3 (top right, red) and 0 (bottom left, blue), each requiring opposite policies. When targets were swapped, no interaction with the environment was necessary to re-evaluate the value of the transitions and the resulting policy. The median time to target is compared to that of a random agent with uniform policy and tested with the same procedure.

In order to understand how the world model can be interpreted, let us look at observation 0 (blue in Fig 6a) and b)), that can be emitted by three different hidden states in the lower left corner of the grid. After the training, the role of the three clone clips splits based on what has been observed, including sequences of actions and observations that led to it, and those that were predicted from it. For example in Fig 6a) , going left or down from clone c₃(0) maps it onto itself, suggesting that it represents the lower left corner of the grid. Going right or up from c₃(0) predicts transitions to clones c₂(0) and c₁(0) respectively, making them candidates for the second positions on the first row and first column respectively. In this way, each clone clip can be associated to a single location in the grid, as shown in Fig 6b). As a result, training the world model in this environment amounts to assigning a single hidden state to each clone clip in a manner that is consistent across actions.

In particular, in spite of being agnostic of the spatial nature of the environment, the world model can be interpreted as a topological map of the grid, as pictured in Fig 6b). Each clone maps onto a single location of the grid, and this mapping is consistent for all actions. When too many clones were provided compared to the degeneracy of an observation, multiple clones are associated with a single location, as for observation 3 in the upper right corner of the grid. This consistent mapping suggests that the agent is able to contextualize the observations with a sequence of previous states and actions. Thanks to this mapping of clones onto single locations, the look-ahead preferences over belief states encode paths ending at the target. They steer the policy toward an optimal one.

Exploit the world model to complete tasks flexibly. Provided the target observation is signaled in the preference distribution, the agents can adapt their policy to reach it following near optimal trajectories. In Fig 6c), agents whose model successfully converged are tested to reach targets with observations 3 (top, red) and 0 (bottom, blue), requiring opposite policies to be attained. For both tasks, irrespective of the initial location of the agent, the performance is drastically improved compared to a random agent with uniform policy. In order to move the target, no other interaction with the environment was required than changing the preference distribution over sensory states, and no new observation was exchanged.

While the paths towards observation 0 are mostly optimal, there is an overhead by at most one step to reach observation 3. Since the policy is an optimal one, as shown in Fig 6b), we suggest that this is due to the estimation of the belief states. In particular, the small size of the grid prevents agents from narrowing down their hypothesis to a single belief state early enough to opt for the right policy at the boundary of the grid. This could explain why on average, agents take 1.5 steps to reach the target when being right under it: they have 0.5 probability of finding the wrong belief state, and therefore choose the wrong action 50% of the time. This eliminates the last erroneous state from the hypothesis and the agent behaves optimally in the next step.

7 Discussion and outlook

In this work, we propose an interpretable, minimal model for cognitive processes that combines associative learning and active inference to adapt to any environment, without any prior knowledge about it and independently of any goal tied to it. We develop the Free Energy Projective Simulation model, a physics-inspired model that seeks inspiration in current paradigms in cognitive sciences, namely active inference, that encompasses the Bayesian brain hypothesis (BBH) and fits more broadly into predictive processing. A FEPS agent is equipped with a memory structured as a graph, the ECM, where the agent encodes associations between events represented as clone clips. As in models of associative learning in the cerebellum [94–96], internal reinforcement signals based on the prediction accuracy of the agent reinforce associations between observations. Clones clips are the support for belief states and acquire contextual meaning as the agent collects longer sequences of correct predictions about its observations and improves its world model. The behavior of FEPS agents does not depend on any reinforcement and is fully determined with active inference, such that the policy optimizes the expected free energy estimated for each action, given the current world model. The resulting model is interpretable in three ways: 1) The world model is readable, 2) Deliberation is traceable and 3) Credit assignment during training is explainable. We perform numerical analysis of the agents in two environments inspired from behavioral psychology: a timed response environment, as well as a navigation task in a partially observable grid world.

Leveraging the interpretability of the model, we identified three capacities a FEPS agent requires in order to interact with an environment and to reach a goal while controlling it, that is, minimizing its surprise about future events. (1) The representation of the environment in the world model should be accurate enough in its predictions, i.e. it should be capable of predicting future sensory signals. (2) The agent should be able to choose the appropriate belief state for aliased observations. (3) For a given time horizon, the agent should be equipped with an efficient mechanism to plan its course of actions.

First, in order to build a complete and accurate representation of the environment, we propose to start learning with an exploration phase. It focuses on exploring the environment strategically rather than completing a task. For this, we provide a model-dependent preference distribution that evolves as the agent learns and we show that the corresponding expected free energy of an action is equal to the information gain about belief states resulting from this action. As a result, the behavior of the agent is motivated by the acquisition of information in its world model rather than by a target bound to the environment. Our numerical simulations showed that a model learned in such a way is able to predict longer sequences of observations when actions are selected at random, compared to the model trained directly on a task. Previous works usually involve either adopting a different definition for the expected free energy [74], or more commonly, adding terms to it to encourage explorative behaviors, [20], leveraging existing literature on curiosity mechanisms, for example [29, 43, 45, 97]. Alternatively, a count-based boredom mechanism [98] could be well-suited for FEPS, as it could be directly implemented on the edge attributes, such as the confidence, to deviate the agent from transitions it has already resolved.

Second, we design a simple procedure to progressively disambiguate belief states with sequences of observations. It can be used to interact with the environment after the model has been trained. Leveraging the clone structure imposed to the clips in the memory of the agent, belief states can be estimated in superposition. As predictions for each candidate belief state are validated or not by the observation collected from the environment, the agent can eliminate candidate belief states that are incompatible with the context provided by its current sensory state. We show in our numerical analysis that this simple mechanism allows to double the number of correct, consecutive predictions, regardless of the technique used to learn the world model.

Thirdly, we introduce a planning method based on the expected free energy. Instead of relying on a tree search or a simulation of future scenarios, we propose to encode the utility of a transition between two belief states in the preference distribution. For this, we factorize the preference distribution into an absolute preference distribution that designates a target among the possible sensory states, and a look-ahead preference distribution. The latter assigns a probability to transitions between belief states that is commensurate with its estimated utility in reaching the target within a given number of actions. Using the world model, the value for each transition is determined in an iterative manner. We tested this scheme in two numerical experiments and achieved optimal policies in both cases, as long as the world model was predictive enough. The closest model to our knowledge is Successor Representation [35, 90, 91], that has been hypothesized to account for so-called model-free learning in cognitive systems [90]. A major difference is that in Successor Representation, the value of a transition depends on the prediction error over the expected reward in the environment, whereas we assign value to a belief state via the probability of the associated sensory state in the absolute preference distribution. A limitation of our scheme, however, is that a target can only be encoded as a possibly ambiguous observation. Using reinforcement and a few interactions with the environment, the preference for a particular hidden state could be encoded in the look-ahead preference distribution in a hybrid scheme.

The numerical analysis of FEPS in larger environments has revealed that the scaling of the method requires modifications of the model in some cases, inheriting in this regard scalability issues from tabular RL methods. More specifically, the main obstacle in efficiently learning a world model is the number of edges to reinforce, because it scales quadratically with the number of clones, and therefore, with the number of distinguishable hidden states in the environment. Consequently, FEPS can learn a representation of a large environment with little to no ambiguity in the observations, but can fail to model a smaller one with few but very ambiguous observations. In addition, clone differentiation can become harder if many different paths in the environment can lead to the same observation sequence using the same actions. Including a mechanism to add and delete clones, such that the world model approximates an ε-transducer for the environment [99] on finite degrees of freedom, could help mitigate this scaling issue. Modifications to the decision-making and reinforcement algorithms could also include multiple edges at each step in order to speed up the learning process.

Conceptually, the FEPS framework fits into the field of NeuroAI [100], at the intersection of behavioral sciences, engineering and neurosciences. While Projective Simulation can be used to learn in an artificial environment, its vocation is to understand agency and the behavior of agents in the world. Furthermore, the FEPS attempts to give a biologically plausible account of learning and adaptive behavior, grounding internal computations in the active inference framework and the predictive processing paradigm. ECMs are potentially implementable on physical platforms and can be considered embodied structures underlying the memory of the agents. For example, a parallel can be drawn between the role the network of belief states plays for FEPS agents, and that of place cells and grid cells in the hippocampus [69, 70, 72]. Both integrate stimuli to create a contextualized representation of an event in order to make predictions about future stimuli. For the FEPS to provide a modeling platform of interest for cognitive and behavioral sciences, the next challenge is to implement learning on real-world tasks and in a fully embodied way, including the calculation of coincidences between predictions and observations, and the update of the associations between states.

Supporting information

S1 Appendix. Table of hyperparameters for the numerical analysis.

(PDF)

S2 Appendix. Derivation of the expected free energy in the exploration phase.

(PDF)

S3 Appendix. Analytical estimation of the limits of the expected free energy.

(PDF)

Data Availability

The code and data supporting this study are accessible at https://doi.org/10.5281/zenodo.15526876.

Funding Statement

This research was funded in whole, or in part, by the Austrian Science Fund (FWF) DOI 10.55776/F71 and 10.55776/WIT9503323. For the purpose of open access, the author has applied a CC BY public copyright license to any Author Accepted Manuscript version arising from this submission. We gratefully acknowledge support from the European Union (ERC Advanced Grant, QuantAI, No. 101055129). The views and opinions expressed in this article are however those of the author(s) only and do not necessarily reflect those of the European Union or the European Research Council - neither the European Union nor the granting authority can be held responsible for them. The funders had no role in study design, data collection and analysis, decision to publish, or preparation of the manuscript.