Learning in Hybrid Active Inference Models (2024)

In the recurrent-only (ro) formulation of the rSLDS (see Fig.2), the discrete latent states $z_{t}\in\{1,2,...,K\}$ are generated as a function of the continuous latents $x_{t}\in\mathbb{R}^{M}$ and the control input $u_{t}\in\mathbb{R}^{N}$ (specified by some controller) via a softmax regression model,

whereby $W_{x}\in\mathbb{R}^{K\times M}$ and $W_{u}\in\mathbb{R}^{K\times N}$ are weight matrices and $r$ is a bias of size $\mathbb{R}^{K}$. The continuous latent states $x_{t}$ evolve according to a linear dynamical system indexed by the current discrete state $z_{t}$.

$A_{z_{t}}$ is the state transition matrix, which defines how the state $x_{t}$ evolves in the absence of input. $B_{z_{t}}$ is the control matrix which defines how external inputs influence the state of the system while $b_{z_{t}}$ is an offset vector. At each time-step $t$, we observe an observation $y_{t}\in\mathbb{R}^{M}$ produced by a simple linear-Gaussian emission model with an identity matrix $C_{z_{t}}$. Both the dynamics of the continuous latents and the observations are perturbed by zero-mean Gaussian noise with covariance matrices of $Q_{z_{t}}$ and $S_{z_{t}}$ respectively.

Where $\mathbb{s}=\{s_{1},...,s_{T}\}$ and $\mathbb{o}=\{o_{1},...,o_{T}\}$ are the states and observations being evaluated under a particular policy or sequence of actions, $\pi=a_{1:T}$. The integration of rewards in the inference procedure is achieved by biasing the agent’s generative model with an optimistic prior over observing desired outcomes $\tilde{p}(o)$. Action selection then involves converting this into a probability distribution over the set of policies and sampling from this distribution accordingly.

In order to create approximate plans at the discrete level, we derive a high-level planner based on a re-description of the discrete latent states found by the rSLDS by approximately ‘integrating out’ the continuous variables and the continuous prior. This process involves calculating the expected free energy ($\mathcal{G}$) for a continuous controller to drive the system from one mode to another. Importantly, the structure of the lifted discrete state transition model has been constrained by the polyhedral partition of the continuous state space extracted from the parameters of the rSLDS ¹¹1For a visualisation of this partitioning of the state space, see Fig.4(a): invalid transitions are assigned zero probability while valid transitions are assigned a high probability. In order to generate the possible transitions from the rSLDS, we calculate the set of active constraints for each region from the softmax representation, $p(z\mid x)=\sigma(Wx+b)$. Specifically, to check that the region $i$ is adjacent to region $j$, we verify the solution using a linear program,

where $(x_{lb},x_{ub})$ are bounds chosen to reflect realistic values for the problem. This ensures we only lift transitions to the discrete model if they are possible. After integration, we are left with a discrete MDP which contains averaged information about all of the underlying continuous quantities. This includes information about the transitions that the structure of the task space allows, and their corresponding approximate control costs (see 0.A.2). Note that after each batch update of the rSLDS parameters, this discrete planner must be refitted accordingly.

The lifted discrete generative model has all the components of a standard POMDP in the active inference framework:

along with prior over policies $P(\pi)=Cat(E)$, and preference distribution $\tilde{P}(o_{t})=Cat(C)$. Specifically our lifted $P(\pi)$ reflects the approximate control costs of each continuous transition and $\tilde{P}(o_{t})$ reflects the reward available in each mode. We assume an identity mapping between states and observation meaning the state information gain term in Eq.4 collapses into a maximum entropy regulariser, while we maintain Dirichlet priors over the transition parameters $B$, facilitating directed exploration. Due to conjugate structure Bayesian updates amount to a simple count-based update of the Dirichlet parameters [25]. At each time step, the discrete planner selects a policy by sampling from the following distribution:

The policy is then communicated to the continuous controller. Specifically, the first action of the selected policy is a requested transition $i\rightarrow j$ and is translated into a continuous control prior $\tilde{p}(x)\sim N(x_{j},\Sigma_{j})$ via the following link function,

Continuous closed-loop control is handled by a set of continuous active inference controllers. For controlling the transition from mode $i$ to mode $j$ ($x_{i}$ to $x_{j}$), the objective of the controller is to minimise the following (discrete-time) expected free energy functional ³³3As shown in [20] linear state space models preclude state information gain terms leaving the simplified form seen here.:

Where $S$ is the finite time horizon and the quadratic terms derive from Gaussian preferences about the final state $\tilde{p}_{j}(x_{S})\sim N(x_{j},Q_{f}^{-1})$ and time invariant control input prior $p(u_{t})\sim N(0,R^{-1})$ (0.A.4). Importantly we design the control priors such that the controller only provides solutions within the environments given constraints (for further discussion, see Sec.5). The approximate closed-loop solution to each of these sub-problems is computed offline each time the rSLDS is refitted (see 0.A.3) using the updated parameters of the linear dynamical systems, allowing for fast discrete-only planning when online.

Optimal Control To motivate our approximate hierarchical decomposition, we adopt the optimal control framework,specifically we consider discrete time state space dynamics of the form:

with known initial condition $x_{0}$, and noise $\eta_{t}$ drawn from some time invariant distribution $\eta_{t}\sim D$, where we assume $f$ to be $p(x_{t+1}\mid x_{t},u_{t})$ and is a valid probability density throughout. We use $c_{t}:X\times U\rightarrow\mathbb{R}$ for the control cost function at time $t$ and let $\mathbb{U}$ be the set of admissible (non-anticipative, continuous) feedback control laws, possibly restricted by affine constraints. The optimal control law for the finite horizon problem is given as:

PWA Optimal Control The fact we do not have access to the true dynamical system $f$ motivates the use of a piecewise affine (PWA) approximation. Also known as hybrid systems:

Where $\mathbb{H}=\{H_{i}:i\in[K]\}$ is a polyhedral partition of the space $X\times U$. In the case of a quadratic cost function, it can be shown the optimal control law for such a system is peicewise linear. Further there exist many completeness (universal approximation) type theorems for peicewise linear approximations implying if the original system is controllable, there will exist a peicewise affine approximation through which the system is still controllable [3, 6].

Relationship to rSLDS We perform a canonical decomposition of the control objective $J$ in terms of the components or modes of the system. By slight abuse of notation $[x_{t}=i]:=[(x_{t},u_{t})\in H_{i}]$ represent the Iverson bracket.

Now let $z_{t}$ be the random variable on $[K]$ induced by $Z_{t}=i$ if $[x_{t}=i]$ we can rewrite the above more concisely as,

which is just the expectation under a recurrent dynamical system with deterministic switches. Later (see 0.A.5), we exploit the non-deterministic switches of rSLDS in order to drive exploration. Eq.21 demonstrates the global problem can be partitioned solving problems within each region (inner expectation), and a global discrete problem which decides which sequence of regions to visit. In the next section, we introduce a new set of variables which allows us to approximately decouple the problems.

Decoupling from clock time Let the random variable $(\zeta_{s})_{s>0}$ record the transitions of $(z_{t})_{t>0}$ i.e. let

be the sequence of first exit times, then $\zeta$ is given by $\zeta_{s}=z_{\tau_{s}}$. With these variables in hand, we frame a small section of the global problem as a first exit problem.

Low level problems Consider the first exit problem for exiting region $i$ and entering $j$ defined by:

where $\partial H_{ij}$ is the boundary $H_{i}\bigcap H_{j}$. Due, to convexity of the polyhedral partition, the full objective admits the decomposition in terms of these subproblems,

Ideally, we would like to simply solve all possible subproblems $\{J_{ij}^{*}(x):i,j\in[K]\times[K]\}$ and then find a sequence of discrete states, $\zeta(1),...,\zeta(S)$, which minimises the sum of the sub-costs, however notice each sub-cost depends on the starting state, and further this is determined by the final state of the previous problem. A pure separation into discrete and continuous problems is not possible without a simplifying assumption.

High level problem we let $J_{ij}^{*}=\min_{\pi}\int_{x_{0}}J_{ij}(\pi,x_{0})p(x_{0})$ be the average cost of each low-level problem. We form a Markov decision process by introducing abstract actions $a\in[K]$:

and let $p_{\pi_{d}}$ be the associated distribution over trajectories induced by some discrete state feedback policy, along with the discrete state action cost $c_{d}(a=j,\zeta=i)=J_{ij}^{*}$ we may write the high level problem:

Our overall approximate control law is then given by choosing the action of the continuous controller $\pi_{ij}(x)$ suggested by the discrete policy $\pi_{d}(i(x))$, or more concisely, $\pi(x)=\pi_{i(x),\pi_{d}^{*}\circ i(x)}(x)$, where $i$ is calculates the discrete label (MAP estimate) for the continuous state $x$. In the next sections we describe the methods used to solve the high and low level problems.

Rather than solve the first-exit problem directly, we formulate an approximate problem by finding trajectories that end at specific ‘control priors’ (see 0.A.6).Recall the low level problem given by:

In order to approximate this problem with one solvable by a finite horizon LQR controller, we adopt a fixed goal state, $x^{*}\in H_{j}$. Imposing costs $c_{t}(x_{t},u_{t})=u_{t}^{T}Ru_{t}$ and $c_{S}(x_{S},u_{S})=(x-x^{*})Q_{f}(x-x^{*})$. Formally we solve,

by integrating the discrete Ricatti equation backwards. Numerically, we found optimising over different time horizons made little difference to the solution, so we opted to instead specify a fixed horizon (hyperparameter). These solutions are recomputed offline every time the linear system matrices change.