Model & belief

The generative model the agent assumes, and the Gaussian belief it carries over the hidden state.

LinearGaussianModel dataclass

LinearGaussianModel(
    dynamics_matrix: ArrayLike,
    *,
    observation_matrix: ArrayLike,
    dynamics_noise: ArrayLike,
    observation_noise: ArrayLike,
    prior: Belief,
    control_matrix: ArrayLike | None = None,
    observation_model: ObservationModel | None = None,
    dynamics_noise_model: DynamicsNoise | None = None,
    structure: ModelStructure | None = None,
)

A linear-Gaussian state-space model — the agent's generative model.

The agent's assumed story for how a hidden state evolves and produces observations, under linear maps and Gaussian noise::

next_state  = dynamics @ state         + control @ action + dynamics noise
observation = observation_matrix @ state                  + observation noise

The noise terms are zero-mean Gaussians with covariances dynamics_noise and observation_noise; the initial state is drawn from prior.

Parameters are role-named rather than using the traditional control-theory letters, to avoid the letter collision with discrete active inference (pymdp), where the same letters mean different things. The "also known as" column lists the terms other backgrounds use, so readers can still find the right field. (Letters survive as .A/.B/.C/.Q/.R aliases for backend use.)

role name	letter	meaning	shape	also known as
dynamics_matrix	A	state -> next state	(n,n)	state-transition
control_matrix	B	action -> state (optional)	(n,p)	input/control
observation_matrix	C	state -> reading	(m,n)	measurement/emission
dynamics_noise	Q	dynamics-noise covariance	(n,n)	process noise
observation_noise	R	reading-noise covariance	(m,m)	measurement noise
prior	--	initial belief over state	n-D	Belief / D (pymdp)

Dimensions: n = state, m = observation, p = action. A model with no control_matrix is a pure filtering (tracking) model.

Everything after dynamics_matrix is keyword-only. Two of the four matrices are maps and two are covariances, all of the same rank, and only the covariances are content-checked. A transposed pair therefore constructs in silence whenever the maps are square and symmetric. Naming them at the call site rules that out.

Three optional fields (all default None → the plain fixed-matrix model) extend it: observation_model (an ObservationModel for state-dependent sensing R(x)), dynamics_noise_model (a DynamicsNoise for state-dependent process noise Q(x)), and structure (a ModelStructure declaring the factor / Markov-blanket partition).

Source code in src/cpomdp/types.py

def __init__(
    self,
    dynamics_matrix: ArrayLike,
    *,
    observation_matrix: ArrayLike,
    dynamics_noise: ArrayLike,
    observation_noise: ArrayLike,
    prior: Belief,
    control_matrix: ArrayLike | None = None,
    observation_model: ObservationModel | None = None,
    dynamics_noise_model: DynamicsNoise | None = None,
    structure: ModelStructure | None = None,
) -> None:
    object.__setattr__(
        self, "dynamics_matrix", jnp.asarray(dynamics_matrix, dtype=float)
    )
    object.__setattr__(
        self, "observation_matrix", jnp.asarray(observation_matrix, dtype=float)
    )
    object.__setattr__(
        self, "dynamics_noise", jnp.asarray(dynamics_noise, dtype=float)
    )
    object.__setattr__(
        self, "observation_noise", jnp.asarray(observation_noise, dtype=float)
    )
    object.__setattr__(self, "prior", prior)
    object.__setattr__(
        self,
        "control_matrix",
        None
        if control_matrix is None
        else jnp.asarray(control_matrix, dtype=float),
    )
    object.__setattr__(self, "observation_model", observation_model)
    object.__setattr__(self, "dynamics_noise_model", dynamics_noise_model)
    object.__setattr__(self, "structure", structure)
    self._validate()

n_states property

n_states: int

Dimension of the hidden state (n).

n_observations property

n_observations: int

Dimension of an observation (m).

n_controls property

n_controls: int

Dimension of an action (p); 0 if the model has no control.

A property

A: Float64[Array, 'n n']

A: the state-transition matrix (alias of dynamics_matrix).

B property

B: Float64[Array, 'n p'] | None

B: the control matrix (alias of control_matrix); None if unset.

C property

C: Float64[Array, 'm n']

C: the observation matrix (alias of observation_matrix).

Q property

Q: Float64[Array, 'n n']

Q: the process-noise covariance (alias of dynamics_noise).

R property

R: Float64[Array, 'm m']

R: the observation-noise covariance (alias of observation_noise).

tree_flatten

tree_flatten() -> tuple[
    tuple[_ModelLeaf, ...], ModelStructure | None
]

Leaves for JAX: every matrix plus the prior belief; structure is aux.

control_matrix, observation_model and dynamics_noise_model are included as (possibly None) children; an uncontrolled / fixed-sensor / fixed-Q model contributes no leaf there and the None is restored on rebuild. A non-None observation_model/dynamics_noise_model is itself a pytree and recurses into its own leaves. structure (declarative metadata, no array leaves) rides in the static aux_data, so two models differing only in structure are different pytrees and a jit keyed on the model re-specialises when it changes.

Source code in src/cpomdp/types.py

def tree_flatten(self) -> tuple[tuple[_ModelLeaf, ...], ModelStructure | None]:
    """Leaves for JAX: every matrix plus the ``prior`` belief; ``structure`` is aux.

    ``control_matrix``, ``observation_model`` and ``dynamics_noise_model`` are
    included as (possibly
    ``None``) children; an uncontrolled / fixed-sensor / fixed-Q model contributes
    no leaf there and the ``None`` is restored on rebuild. A non-``None``
    ``observation_model``/``dynamics_noise_model`` is itself a pytree and recurses
    into its own
    leaves. ``structure`` (declarative metadata, no array leaves) rides in the
    static aux_data, so two models differing only in ``structure`` are different
    pytrees and a jit keyed on the model re-specialises when it changes.
    """
    children = (
        self.dynamics_matrix,
        self.observation_matrix,
        self.dynamics_noise,
        self.observation_noise,
        self.prior,
        self.control_matrix,
        self.observation_model,
        self.dynamics_noise_model,
    )
    return children, self.structure

tree_unflatten classmethod

tree_unflatten(
    aux_data: ModelStructure | None,
    children: tuple[_ModelLeaf, ...],
) -> LinearGaussianModel

Rebuild from leaves without validating — the leaves may be tracers.

aux_data is the static structure (or None), restored as-is.

Source code in src/cpomdp/types.py

@classmethod
def tree_unflatten(
    cls,
    aux_data: ModelStructure | None,
    children: tuple[_ModelLeaf, ...],
) -> "LinearGaussianModel":
    """Rebuild from leaves without validating — the leaves may be tracers.

    ``aux_data`` is the static ``structure`` (or ``None``), restored as-is.
    """
    obj = object.__new__(cls)
    fields = (
        "dynamics_matrix",
        "observation_matrix",
        "dynamics_noise",
        "observation_noise",
        "prior",
        "control_matrix",
        "observation_model",
        "dynamics_noise_model",
    )
    for name, value in zip(fields, children, strict=True):
        object.__setattr__(obj, name, value)
    object.__setattr__(obj, "structure", aux_data)
    return obj

Belief dataclass

Belief(mean: ArrayLike, cov: ArrayLike)

A Gaussian belief over a continuous state.

In active inference an agent never knows the hidden state directly — it holds a probability distribution over what the state might be. For the linear-Gaussian case that distribution is always a Gaussian, fully described by two things:

• mean -- the centre, the best single estimate. A 1-D vector of length n.
• cov -- the covariance, the uncertainty. An n x n matrix; its diagonal is the variance per state dimension, its off-diagonals the correlations between them.

Beliefs are immutable values: updating a belief produces a new Belief rather than mutating an existing one. Inputs are accepted as anything array-like (lists, tuples, arrays) and stored as float jax.Array.

A Belief is a registered JAX pytree (its leaves are mean and cov), so it passes through jit/vmap/grad as data. JAX rebuilds it from its leaves without re-running validation; the shape/symmetry checks fire only on direct construction, at the trust boundary. Positive-semi-definiteness is enforced at the trust boundary too, not here (see DECISIONS.md ADR-002).

Source code in src/cpomdp/types.py

def __init__(self, mean: ArrayLike, cov: ArrayLike) -> None:
    object.__setattr__(self, "mean", jnp.asarray(mean, dtype=float))
    object.__setattr__(self, "cov", jnp.asarray(cov, dtype=float))
    self._validate()

ndim property

ndim: int

Dimensionality of the state — the length of the mean vector.

tree_flatten

tree_flatten() -> tuple[
    tuple[Float64[Array, n], Float64[Array, "n n"]], None
]

Leaves for JAX: (mean, cov), no static aux data.

Source code in src/cpomdp/types.py

def tree_flatten(
    self,
) -> tuple[tuple[Float64[Array, "n"], Float64[Array, "n n"]], None]:
    """Leaves for JAX: ``(mean, cov)``, no static aux data."""
    return (self.mean, self.cov), None

tree_unflatten classmethod

tree_unflatten(
    aux_data: None,
    children: tuple[
        Float64[Array, n], Float64[Array, "n n"]
    ],
) -> Belief

Rebuild from leaves without validating — the leaves may be tracers.

Source code in src/cpomdp/types.py

@classmethod
def tree_unflatten(
    cls,
    aux_data: None,
    children: tuple[Float64[Array, "n"], Float64[Array, "n n"]],
) -> "Belief":
    """Rebuild from leaves without validating — the leaves may be tracers."""
    mean, cov = children
    obj = object.__new__(cls)
    object.__setattr__(obj, "mean", mean)
    object.__setattr__(obj, "cov", cov)
    return obj