# pymdp: active inference for Markov decision processes

> pymdp is an MIT licensed Python package for simulating active inference agents in Markov decision process environments, with the low level mathematics ported from the SPM implementation in MATLAB and validated against it, a JAX-first API, and a paper published in the Journal of Open Source Software. The judgement is that the validation claim is the strongest part of the project, and that the dependency pinning tells you more about the state of the ecosystem than the version number does.

**infer-actively/pymdp** — A Python implementation of active inference for Markov Decision Processes

- Repository: https://github.com/infer-actively/pymdp
- Stars: 746 · Forks: 144
- Language: Python
- License: MIT
- Published: 2026-10-21 · Updated: 2026-10-21 · Language: en
- Canonical page: https://hysenlabs.com/projects/infer-actively-pymdp

## An active inference package with a peer reviewed paper

pymdp is a Python package for simulating active inference agents in Markov decision process environments, and the first thing the readme does is point at the paper.

The companion paper was published in the Journal of Open Source Software as pymdp: A Python library for active inference in discrete state spaces, with a DOI, and a longer tutorial-style arxiv version exists for a more in-depth introduction to the package and to the mathematics of active inference in Markov decision processes. The citation block is right there in the readme, and the request is that you cite it if you use the package in your work or research.

The paper is from 2022 and lists a wide group of authors, Conor Heins, Beren Millidge, Daphne Demekas, Brennan Klein, Karl Friston, Iain D. Couzin and Alexander Tschantz, which is a good indication of where the method comes from.

The package is hosted on the infer-actively GitHub organisation, which was built with the intention of hosting open source active inference and free-energy-principle related software, so the repository is one of several rather than a one-off. The project describes itself as active, it is published on PyPI as inferactively-pymdp, the documentation is on Read the Docs, and the licence is MIT.

The releases are in the 1.0 series, with 1.0.2 on 2026-05-14, 1.0.3 on 2026-06-15 and 1.0.4 on 2026-08-23, and the last push to the repository was on 2026-09-29.

## The mathematics is a port of SPM, and that is the point

One sentence in the readme is the strongest claim the project makes, and it is about provenance rather than about features.

Most of the low-level mathematical operations are NumPy ports of their equivalent functions from the SPM implementation in MATLAB, and the readme adds that most of these functions have been benchmarked and validated against their SPM counterparts.

That matters for a research package. Active inference implementations are notoriously easy to get subtly wrong, because the update rules are a stack of matrix operations where a transpose or a normalisation mistake produces results that look plausible and are wrong. Porting from a reference implementation that the field already uses, and validating the port against it, is a much stronger position than writing the equations from a paper and hoping.

The package has grown past a straight port since then, most visibly in its move to JAX, but the SPM comparison remains the reference point the readme points to, and it is the reason to trust the numbers the package produces.

The organisation around the code is visible in the dependency list: equinox for the neural network modules, mctx for the Monte Carlo tree search used in some policy routines, jaxtyping for shape checking, opt_einsum for the contraction work, and multimethod for the dispatch pattern that lets the same agent class run under different backends.

## The JAX-first API, and the guide for people coming from NumPy

The quick start section shows the shape of the modern API, and it is a small, readable program.

The example imports JAX for arrays and random keys, imports the utils module and the Agent class, and then splits a PRNG key into three:

```python
from jax import numpy as jnp, random as jr
from pymdp import utils
from pymdp.agent import Agent

key = jr.PRNGKey(0)
keys = jr.split(key, 3)
```

The arrays are then built with helper functions, a random A array for the observation model, a random B array for the transition model, and a uniform C array for the preferences. The agent is constructed from those three and a batch size, and the loop is three calls: infer_states for the posterior over hidden states, infer_policies for the posterior over policies, and sample_action for the choice.

```python
qs, info = agent.infer_states(observation, empirical_prior=agent.D, return_info=True)
vfe = info["vfe"]
q_pi, neg_efe = agent.infer_policies(qs)
```

The variational free energy comes back from the inference call as a diagnostic, which is a nice touch for a package whose users are researchers who want to see the numbers.

For anyone arriving from the older interface, there is a dedicated migration guide notebook, and the readme recommends three notebooks in order: the JAX quickstart, the NumPy or legacy to JAX migration guide, and a guide to the rollout() active inference loop. All three can be run on Google Colab as well as locally.

## The JAX version cap is a story about numpyro

The dependency list in pyproject.toml contains the most interesting comment in the whole repository, and it explains a version pin that would otherwise look arbitrary.

The JAX and JAXlib requirements are both capped below 0.11.1, and the comment says why. JAX and numpyro are upper-bounded until numpyro supports JAX 0.11.1, which moves closed-over constants from jaxpr.constvars into jaxpr.invars. numpyro's eval_provenance zips jaxpr.invars against the explicit keyword arguments only, so it trips a length assertion. The path to that is pybefit to numpyro, and the issue is numbered 447.

So the cap is not conservatism about JAX itself. It is one upstream package, numpyro, not having caught up with a change in how JAX represents closed-over values, and that breaking a length assertion in numpyro's provenance evaluation. A user of pymdp who never touches numpyro still inherits the cap, because the dependency chain through pybefit pulls numpyro in.

The other pins in the same file are of a different kind and are just as informative. The dev group pins ruff to below 0.16 because 0.16.0 broadened ruff's default rule set, with a note to keep the bound in sync with the ruff specification in the lint workflow. And the test group keeps ipykernel below 7 while the documentation tooling requires a version above 6 and below 7.

Those three pins, one for a transitive break, one for a linter's defaults and one for a documentation conflict, are a good summary of what it costs to maintain a research package across a moving scientific Python stack.

## The mouse, the cues and epistemic chaining

The demonstration in the readme is the clearest description of what the package is for, and it is a behavioural one.

One of the defining features of active inference agents is the drive to maximise epistemic value, which is curiosity. Equipped with such a drive in environments with uncertain but disclosable hidden structure, active inference lets an agent learn about the environment while also maximising reward.

The simulation shows what the readme calls epistemic chaining. An agent, analogised to a mouse seeking food, forages for a chain of cues, each of which discloses the location of the subsequent cue in the chain, and the final cue reveals the location of a hidden reward. The readme compares it to behaviour chaining in operant conditioning, with the difference that each successive action in the sequence does not need to be learned through instrumental conditioning. An active inference agent forages the sequence of cues on its own, from an intrinsic desire to disclose information, and reaches the reward in the fewest number of moves possible.

That is the claim in one paragraph: the agent is not trained to follow a chain, it discovers the chain because resolving uncertainty is what it values.

The animations are produced by a script in the examples directory, examples/envs/chained_cue_navigation.py, and there is an associated JAX notebook in the documentation. The examples directory is the place to look for the rest of the code, and it is organised into advanced, api, envs, experimental, inductive_inference, inference_and_learning, learning, legacy, model_fitting and sparse.

## Packaging, documentation and the meaning of version 1.0

The packaging is conventional and the metadata is precise about the supported environment.

The build uses setuptools with version 77 or newer and the standard build_meta backend, the package is named inferactively-pymdp, and the version is 1.0.4. The description in the manifest is a Python package for solving Markov decision processes with active inference, and the classifiers call it production or stable, with the intended audience as developers and the topic as scientific and engineering artificial intelligence. Python 3.10 or newer is required.

The authors listed in the manifest are Conor Heins, Alexander Tschantz, Tim Verbelen and Dimitrije Markovic, and the licence is MIT with the licence file declared in the metadata.

The documentation is built with mkdocs, with a readthedocs configuration and a directory of documentation sources named for the tool, plus a separate notebook sanitising configuration for the test suite that runs the documentation examples.

The repository also carries the machinery of a maintained project: a pre-commit configuration, release-please for the versioning and changelog, a contributing guide, a change log, a paper directory for the manuscript and the citation material, and a scripts directory.

The last thing to say is about the version number itself. A 1.0 label from a research package means the API is considered stable for the users the maintainers have in mind, which is a much narrower claim than the same label on a database driver, and the 1.0.2, 1.0.3 and 1.0.4 releases across May, June and August 2026 read as a project in maintenance rather than in development.

## Conclusion

pymdp fits a researcher or an engineer who wants to build active inference agents in discrete state spaces and needs the mathematics to be defensible, because the core routines are ports of the SPM implementation and have been benchmarked and validated against it, and because the package is published with a peer reviewed paper and a citation to copy. It is the wrong choice if you want a stable numerical stack you can pin without reading the dependency comments, because the JAX version is capped below 0.11.1 for a reason that lives in a numpyro issue, and it is the wrong choice if your work is in continuous state spaces, which is outside what this package covers. Before you adopt it, read the dependency pin in pyproject.toml rather than trusting the version number, decide whether you want the JAX path or the legacy NumPy path, and run the three recommended notebooks first, since they are the intended entry point.

## FAQ

### What is pymdp?

It is a Python package for simulating active inference agents in Markov decision process environments, published as inferactively-pymdp on PyPI under the MIT licence. Its low level mathematical operations are NumPy ports of the SPM implementation in MATLAB, benchmarked and validated against their SPM counterparts.

### How do I install pymdp?

The readme recommends uv with an explicit virtual environment, running uv venv .venv, activating it, then uv pip install inferactively-pymdp. A plain pip install inferactively-pymdp also works. If building pygraphviz fails, install Graphviz and the development packages for it first.

### Why is the JAX version capped below 0.11.1?

The comment in pyproject.toml explains that numpyro does not yet support JAX 0.11.1, which moves closed-over constants from jaxpr.constvars into jaxpr.invars, so numpyro's eval_provenance trips a length assertion. The dependency path runs through pybefit to numpyro, and the upstream issue is numbered 447.

### What is epistemic chaining in pymdp?

It is the behaviour shown in the demonstration, where an agent forages for a chain of cues, each of which discloses the location of the next, because resolving uncertainty is what it values, and it reaches the hidden reward in the fewest moves without the chain being trained through instrumental conditioning.

## Sources

- [infer-actively/pymdp on GitHub](https://github.com/infer-actively/pymdp)
- [Issues](https://github.com/infer-actively/pymdp/issues)
- [License: MIT](https://github.com/infer-actively/pymdp/blob/main/LICENSE)
- [README](https://github.com/infer-actively/pymdp/blob/main/README.md)
- [Releases](https://github.com/infer-actively/pymdp/releases)

---

Hysen Labs editorial analysis, written from the project's own repository and release notes. Cite the canonical page: https://hysenlabs.com/projects/infer-actively-pymdp