PyAIF (pyaif-toolkit)

PyAIF is an active-inference toolkit for Python. It constructs agents with factorised discrete hidden states and either categorical or continuous observations. It provides single-step inference, deep temporal inference, policy evaluation, action selection, and parameter-learning hooks.

Documentation · PyPI · Source code

Citation

If you use PyAIF in your research, please cite:

D. A. Warnakulasuriya, J. Plosila and H. Haghbayan, “Towards Self-Supervised Intent Recognition in Human-Robot Collaboration using Active Inference,” 2026 11th International Conference on Control and Robotics Engineering (ICCRE), Kyoto, Japan, 2026, pp. 191–198, doi:10.1109/ICCRE69951.2026.11593576.

Demonstration

Active Inference POMDP implementation on a Franka Emika Panda Robot — ASL Lab UTU.

Learning under uncertainty

The agent minimizes variational free energy (VFE) as it improves its generative model. When the task context changes, prediction errors produce a temporary spike in VFE surprise. Continued learning then refines the model for the new context, reducing surprise again. The animation plots VFE surprise (the smoothed deviation from an adaptive VFE baseline), rather than raw cumulative model evidence.

Version 0.2 adds domain-independent continuous-observation inference. Research applications and domain-specific likelihood construction remain in separate repositories.

Installation

PyAIF requires Python 3.9 or newer.

Install the published distribution:

python -m pip install pyaif-toolkit

The distribution is named pyaif-toolkit on PyPI, while the Python import package remains PyAIF.

To install from a source checkout:

python -m pip install .

For development:

python -m pip install -e ".[dev]"
pytest

Quick start

The public API separates the state-transition model, observation model, and inference algorithm.

import numpy as np

from PyAIF import (
    ActiveInfAgent,
    CategoricalLikelihood,
    GenerativeModel,
    ShallowInference,
)


def object_array(*values):
    result = np.empty(len(values), dtype=object)
    for index, value in enumerate(values):
        result[index] = np.asarray(value, dtype=float)
    return result


# A[o, s]: categorical observation likelihood.
A = object_array(
    np.array(
        [
            [0.95, 0.05],
            [0.05, 0.95],
        ]
    )
)

# B[s_next, s_previous, action]: controlled transitions.
B_factor = np.zeros((2, 2, 2))
B_factor[:, :, 0] = np.eye(2)
B_factor[:, :, 1] = np.fliplr(np.eye(2))
B = object_array(B_factor)

# D[s]: initial-state prior. C[o]: outcome preferences.
D = object_array(np.array([0.5, 0.5]))
C = object_array(np.zeros(2))

model = GenerativeModel(
    B=B,
    D=D,
    controls_dim=[2],
    controllable_factors=[0],
)
likelihood = CategoricalLikelihood(
    A=A,
    preferences=C,
    modality_dependencies=[[0]],
)
agent = ActiveInfAgent(
    model=model,
    likelihood=likelihood,
    inference=ShallowInference(),
    action_selection="deterministic",
)

agent.reset()
agent.observe([0])
agent.infer_states()
expected_free_energy, _ = agent.infer_policies()
action = agent.select_action()

Use DeepTemporalInference(horizon=...) for policy-dependent beliefs over a fixed temporal window. For continuous control, choose either RecedingHorizonInference(horizon=...) for policy-conditioned temporal message passing or FilteredRecedingHorizonInference(horizon=...) for the conventional robotics split: filter the current state once, then predict and score future states under each policy. Deep preferences have shape (number_of_outcomes, horizon).

Continuous observations

The likelihood component selects observation semantics; the inference component independently selects the temporal algorithm:

likelihood = ContinuousLikelihood(
    likelihood_fn=density,  # density(value, modality) -> state tensor
    observation_grids=[rssi_grid],
    log_preferences={0: rssi_log_preferences},
    modality_dependencies=[[0, 1]],
)

agent = ActiveInfAgent(
    model=model,
    likelihood=likelihood,
    inference=ShallowInference(),  # or a deep/receding inference configuration
)

For an existing domain model that implements likelihoods, get_o_grid, and log_preferences, use ContinuousLikelihood.from_model(...). PyAIF then integrates densities on the supplied grids for preference and information-gain evaluation. Small state spaces are enumerated exactly; large ones use reproducible Monte Carlo sampling controlled by policy_samples, exact_state_limit, and random_seed.

The optional learning_fn, preference_learning_fn, and parameter_information_gain_fn callbacks keep domain-specific updates outside the reusable agent core.

Performance and parallel policies

Deep temporal inference automatically batches independent policies into NumPy tensor operations. This reduces Python-loop overhead while preserving policy-by-policy inference results. Multi-factor categorical likelihoods use optimized contraction paths; likelihood entropy and packed posterior batches are reused across the future-time calculation.

Policy evaluation can also use bounded worker threads:

inference = DeepTemporalInference(
    horizon=3,
    message_passing_iterations=16,
    policy_workers=4,
)

policy_workers=1 is the default and is usually fastest for small categorical models because batched NumPy operations already process all policies efficiently. Values greater than one split policy batches across threads without copying the agent's arrays. Benchmark representative workloads before increasing the worker count; threading is most useful when each policy contains sufficiently large tensor contractions. ShallowInference accepts the same policy_workers option for concurrent policy scoring.

See the reproducible filtered receding-horizon CPU benchmark for comparisons among PyAIF temporal modes, classic NumPy pymdp, and current JAX pymdp. The original benchmark is retained for historical comparison. A separate resolution-scaling profile breaks policy evaluation into rollout, EFE contractions, posterior updating, and Bayesian model averaging.

Agent lifecycle

The supported component-based lifecycle is:

agent.reset(trial=trial)
agent.observe(observation, time_step=t)
agent.infer_states()
agent.infer_policies()
action = agent.select_action()
agent.learn()  # only when parameter learning is enabled

The older matrix-heavy constructor and methods such as choose_action() and perform_learning() remain available while examples migrate, but new projects should use the component constructor and lifecycle above.

Model structure

• GenerativeModel owns state transitions (B), initial-state priors (D), control dimensions, controllable factors, and optional policies.
• CategoricalLikelihood owns observation likelihoods (A), preferences (C), and modality-to-factor dependencies.
• ContinuousLikelihood owns density evaluation, observation grids, log preferences, dependencies, and optional domain learning hooks.
• ShallowInference performs single-step factorised inference.
• DeepTemporalInference performs marginal message passing over a fixed window.
• RecedingHorizonInference replans a full window at every observation and uses policy-conditioned marginal message passing, including backward future-to-present messages.
• FilteredRecedingHorizonInference filters the current state once, then rolls that shared posterior forward under every policy. This estimator/planner separation is usually the more natural default for online robotics. It carries the shared current posterior directly and skips policy-averaged future-state diagnostics by default; set average_future_states=True when those diagnostics are needed. Both receding modes select the first action and support uninterrupted control. Horizon 3 means the current state plus two future states. See the lifecycle and migration notes and the runnable temporal-message-passing and filter-then-plan examples. Parameter learning is not yet supported in either receding mode.
• PyAIF.learning contains reusable categorical updates for A, B, C, D, and E.

See Model shapes and Public API for details.

Examples

• examples/quickstart_discrete.py: minimal categorical agent.
• examples/quickstart_continuous.py: minimal Gaussian continuous-observation agent.
• examples/quickstart_receding.py: receding-horizon temporal message passing.
• examples/quickstart_filtered_receding.py: robotics-style current-state filtering followed by predictive policy rollout.
• examples/learning_under_uncertainty/: deep temporal parameter-learning experiments under epistemic and aleatoric uncertainty.

Model-selection and bounded-rationality experiments are maintained separately from the reusable PyAIF package.

The automated regression suite covers the reusable component API and executes one trial of each parameter-learning experiment.

Run the minimal component examples with:

python examples/quickstart_discrete.py
python examples/quickstart_continuous.py

Version policy

• 0.1.x: categorical observations and discrete parameter learning.
• 0.2.x: continuous-observation likelihood components for discrete hidden states.

Behavioral changes are protected with numerical regression tests. Research experiments may evolve independently from the packaged API.

License

PyAIF is distributed under the BSD 3-Clause License. Copyright © 2026 Diluna A. Warnakulasuriya.

This software license does not automatically apply to papers, datasets, figures, trained models, or other research artifacts.

Development

ruff check PyAIF tests
ruff format --check PyAIF tests
pytest
python -m build
python -m twine check dist/*

See CONTRIBUTING.md for the pull-request and release process.