Schema-Based Hierarchical Active Inference

• Schema-Based Hierarchical Active Inference (S-HAI) is a computational framework that merges abstract schemas, predictive processing, and active inference to enable rapid generalization and flexible task adaptation.
• The framework uses a hierarchical generative model that binds abstract goals to concrete sensorimotor states via variational free energy minimization and Bayesian learning.
• S-HAI demonstrates behavioral efficiency in spatial navigation tasks and replicates neural coding motifs in frontal and hippocampal circuits, providing robust insights into adaptive planning.

Schema-Based Hierarchical Active Inference (S-HAI) is a computational framework for rapid generalization, flexible task adaptation, and hierarchical abstraction in behavioral and neural systems. S-HAI formalizes the concept of schemas—abstract relational structures capturing invariant patterns across experiences—and integrates them with predictive processing and active inference. The approach enables the binding of abstract goals or task-structures to novel sensory and spatial contexts via a multi-layer generative model, supports fast learning and generalization, and accounts for key neural correlates in frontal and hippocampal circuits (Maele et al., 26 Jan 2026), with categorical formalizations via polynomial functors and monoidal bicategories (Smithe, 2022).

1. Model Architecture and Mathematical Foundations

S-HAI employs a hierarchical generative model comprising at least two levels:

• Level 1 ("navigation"): Encodes agent's discrete spatial location $s^1_t \in \{1,\dots,N\}$, local observations $o^1_t$, and navigation actions $a^1_t$ (e.g., Up, Down, Left, Right).
• Level 2 ("schema/task"): Represents latent abstract phases $s^2_t \in \{A, B, C, D, \dots\}$, observations $o^2_t$ (signaling goal achievement), and schema-level actions $a^2_t$ ("next goal in sequence").

Key to S-HAI is the grounding likelihood $A^{2\to1}$: a mapping from high-level schema observations to low-level spatial states, implemented as a categorical distribution parameterized by Dirichlet priors. This enables abstract goals (e.g., "visit goal A") to be grounded in specific spatial locations within each context.

The generative model at each hierarchy is a factored POMDP: $$P(\hat{s}^i, \hat{o}^i, \hat{a}^i) = P(s^i_0) \prod_{t=1}^T P(o^i_t|s^i_t) P(s^i_t|s^i_{t-1}, a^i_{t-1}) P(a^i_{t-1})$$ Parameters (observation likelihood $A^i$, transition $B^i$) and policy priors are learned and updated online, with actions selected via minimization of expected free energy and an inductive cost (goal-distance penalty).

Categorically, S-HAI can be formalized via polynomial functors (schemas) and open dynamical systems (coalgebras) (Smithe, 2022). Hierarchical composition utilizes monoidal bicategories, and approximate inference doctrines act as functors from statistical (probabilistic) models to hierarchical active inference systems, ensuring compositionality and convergence.

2. Inference, Learning Algorithms, and Planning Procedures

Both levels of S-HAI employ variational free energy minimization for inference: $$\mathcal{F}^i = \mathbb{E}_{Q^i}[ \ln Q^i(\hat{s}^i) - \ln P(\hat{s}^i, o^i | a^i) ] = \sum_{t} \mathrm{D_{KL}}(q^i(s^i_t) \| P(s^i_t|s^i_{t-1},a^i_{t-1})) - \mathbb{E}_{q^i} \ln P(o^i_t|s^i_t)$$ Posterior updates proceed via time-factorized mean-field inference and softmaxed belief propagation, augmented with message passing between levels:

• Bottom-up: Posterior beliefs about low-level states $s^1_t$ are propagated upward, transformed via $A^{2\to1}$ into likelihoods for schema-level inference, especially on reward events.
• Top-down: High-level policy $\pi^2$ sets priors for the navigation level, either by informing goal preferences or through inductive cost maps based on latent graph distances.

Policy selection leverages minimization of expected free energy for each candidate policy over a planning horizon, combining epistemic value (information gain) and expected utility (reward). The policy posterior is

$$Q(\pi^i) \propto \exp(-\gamma [ \mathcal{G}^i(\pi^i) + \mathcal{H}^i(\pi^i) ])$$

where $\gamma$ is a temperature parameter, and $\mathcal{H}^i$ encodes costs for deviations from goal priors.

Learning of grounding likelihoods $A^{2\to1}$, schema transitions $B^2$, and context-specific mapping is performed online via conjugate Bayesian updates (Dirichlet) and mixture models (Dirichlet processes for multiple schemas or groundings).

Hybrid extensions utilize recurrent switching linear dynamical systems (rSLDS) for continuous state/event decomposition, with a high-level planner learning discrete modes/schemas over rSLDS partitioned state-space and a low-level controller implementing LQR-based continuous active inference (Collis et al., 2024).

3. Behavioral and Simulation Results

S-HAI has been validated in spatial navigation multi-goal tasks (e.g., ABCD cycle, alternation/ABCB variant, and novel/repeating block mixtures) (Maele et al., 26 Jan 2026):

• Generalization to novel layouts: S-HAI achieves rapid reacquisition (approaching optimal trial length $\sim32\pm2$ steps) in novel mazes after brief adaptation ($<2000$ steps), outperforming baseline hierarchical active inference agents (HAI-20, HAI-40), which exhibit slower learning and plateau at suboptimal rates.
• Schema reuse vs. accommodation: Block-specific mixtures allow the agent to reuse groundings for repeated layouts and instantiate new schema components for genuinely novel contexts, with model selection governed by Bayesian surprise (log-likelihood dips).
• Ambiguity resolution via clones: Alternation tasks (e.g., two goals labeled "B") require clone-structured graphs at the schema level for goal disambiguation; S-HAI-2C agents resolve these with near-optimal performance.
• Continuous control tasks: Hybrid S-HAI agents with rSLDS decomposition realize rapid exploration and reward acquisition (100% state coverage, reward found in $<3$ episodes), outperforming RL baselines (Collis et al., 2024).

4. Neural Correlates and Representational Insights

S-HAI reproduces multiple distinct neural coding motifs observed in rodent medial prefrontal cortex and hippocampal circuits during schema-dependent navigation tasks (Maele et al., 26 Jan 2026):

• Goal-progress cells: Latent expected inductive cost ramps as the agent traverses from start to goal, matching spatially-tuned ramping cells.
• Goal-identity cells: Abstract schema-phase posteriors yield units tuned to specific goals, invariant across maze layouts.
• Conjunctive goal-location cells: Grounding likelihood parameters encode cells firing only for particular goal-location combinations, reflecting flexible remapping.
• Place-like codes: Navigation-level posteriors support place field-like neural populations.

These latent activations, binned by spatial and task phase, closely align with empirical neural activity, supporting the hypothesis that hierarchical active inference mechanisms underpin observed neural coding in frontal and hippocampal areas.

5. Categorical and Compositional Generalization

From a categorical perspective, S-HAI defines schemas as polynomial functors and organizes hierarchical inference as a monoidal bicategory:

• Schemas = polynomial functors: Encapsulate configuration signatures and input/output types.
• Open dynamical systems = coalgebras: Capture state-dependent processing on schema interfaces.
• Hierarchical composition: Constructed via monoidal and categorical composition, yielding arbitrarily deep/complex inference architectures.
• Approximate inference doctrines: Functor from probabilistic models to hierarchical dynamical systems; canonical examples include Laplace and Hebb-Laplace doctrines for Gaussian channels and affine parameterizations.

Correctness and local convergence are guaranteed under standard convexity assumptions, with functorial mappings ensuring compositionality of inference dynamics (Smithe, 2022).

6. Theoretical Implications and Extensions

S-HAI mechanistically implements schema assimilation and accommodation in line with Piagetian cognitive theory—schemas are learned as abstract graphs and flexibly grounded via context-specific mappings. Active inference principles (free energy minimization, inductive cost, epistemic value) yield perceptual inference, model learning, and adaptive planning.

Extensions include:

• Mixture models: Nonparametric mixtures over both schema-level structures and grounding likelihoods for flexible adaptation to evolving environments.
• Higher-order domains: Applicability to narrative, language templates, and sequential decision tasks.
• Neural circuit mapping: Direct links to specific cortical and hippocampal pathways; experimental perturbations predict loss of remapping (grounding disruption) or step-sequence inference (schema disruption).

A plausible implication is that hierarchical predictive processing underlies rapid abstraction and generalization capacities in biological and artificial agents.

7. Software Implementation and Reproducibility

A reference Python/TensorFlow implementation is available (pymdp + S-HAI modules; https://github.com/toonvdm/grounding-schemas) (Maele et al., 26 Jan 2026), supporting generative model specification, message passing, variational updates (Smith et al. 2022), expected free energy policy selection (Friston et al. 2016), Bayesian learning, and schema/grounding mixture expansions. This enables full replication of both theoretical and empirical results described above.