PyMDP and RxInfer: The Active Inference Software Stack

Series: Active Inference Applied | Part: 5 of 10

There's a gap between understanding the theory and actually building something that works. You can grasp the free energy principle, follow the math of expected free energy, even understand how message passing implements inference—and still have no idea where to start coding.

This is where theory meets silicon. This is the software stack.

Active inference isn't one framework—it's an ecosystem of tools built by different communities for different purposes. Some prioritize ease of learning. Others optimize for speed. Some target research exploration. Others aim at production deployment. The landscape is fragmented but maturing fast.

Two libraries dominate the current landscape: PyMDP for discrete state spaces and RxInfer.jl for continuous inference at scale. Understanding what each does well—and where each falls short—is essential for anyone trying to move from equations to implementations.

The Discrete World: PyMDP

If you're learning active inference and want to see it work today, you start with PyMDP.

PyMDP is a Python library designed for discrete active inference—agents operating in environments with finite, countable states. Think grid worlds, simple navigation tasks, classification problems. It's the workhorse of the Active Inference Institute and the tool behind most pedagogical demonstrations.

What PyMDP Gets Right

Accessibility. PyMDP was built to teach. The API is clean, the tutorials are extensive, and the barrier to entry is low. If you understand Python and have worked through the generative model construction process, you can have a working agent in an afternoon.

Research-ready flexibility. Discrete models let you specify everything explicitly—state spaces, observation models, transition matrices, preferences. This makes PyMDP ideal for hypothesis testing. Want to see how an agent behaves when you change its prior beliefs about state transitions? Adjust the B matrix and re-run. Want to test different precision weightings? Modify the C vector. Everything is inspectable, modifiable, interpretable.

Integration with the community. PyMDP's development is tightly coupled to the research output of Karl Friston's lab and the broader Active Inference Institute network. New theoretical developments often get implemented in PyMDP within months. The library tracks the cutting edge.

Where PyMDP Struggles

Scalability. Discrete state spaces explode combinatorially. A 10x10 grid with 4 possible actions is manageable. A high-dimensional continuous control problem is not. PyMDP simply wasn't designed for the kinds of environments where modern reinforcement learning excels—robotics, continuous control, pixel-based observations.

Speed. PyMDP is written in Python and optimized for clarity, not performance. Inference is exact but slow. For research demonstrations, this is fine. For real-time robotics or large-scale simulations, it's a bottleneck.

Production readiness. PyMDP is a research library. It's not hardened for deployment. No one is running PyMDP agents in production systems managing real-world infrastructure. That's not what it was built for.

When to Use PyMDP

Use PyMDP when:

• You're learning active inference and need to see the concepts work concretely
• You're doing research on discrete decision-making problems
• You want full control over the generative model specification
• Interpretability matters more than speed
• Your state space is small enough to enumerate explicitly

PyMDP is the pedagogical standard and the research prototype tool. It's not the answer for scaling to complex real-world systems, but it's the best place to start.

The Continuous Frontier: RxInfer.jl

If PyMDP is the classroom, RxInfer is the workshop where production systems get built.

RxInfer.jl is a Julia library for reactive message passing on probabilistic graphical models. It's not only for active inference—it's a general-purpose Bayesian inference engine—but it's become the de facto choice for researchers and engineers building active inference agents in continuous, high-dimensional spaces.

What RxInfer Gets Right

Speed. Julia's just-in-time compilation gives RxInfer performance within striking distance of C++. Inference that takes minutes in PyMDP can run in milliseconds in RxInfer. This isn't a 2x improvement—it's orders of magnitude. That difference is what makes real-time robotics feasible.

Scalability to continuous domains. RxInfer implements variational message passing, which makes approximate inference on continuous state spaces tractable. Instead of enumerating all possible states, RxInfer uses probability distributions (often Gaussians) as messages, updating beliefs through efficient linear algebra. This is how you move from toy grid worlds to robots navigating real space.

Modularity through factor graphs. RxInfer represents generative models as factor graphs—nodes for variables, edges for dependencies, factors for local probability functions. This representation makes it easy to compose complex models from reusable components. You define the structure once, and RxInfer handles the inference automatically. Change the graph, and the inference updates accordingly.

Reactive programming model. RxInfer is built on reactive streams—data flows through the graph as observations arrive, and inference happens continuously. This maps naturally to real-world sensing and acting. An agent doesn't wait for a full batch of data—it updates its beliefs incrementally as new information comes in. This is essential for real-time systems.

Where RxInfer Struggles

Learning curve. Julia is less familiar than Python to most researchers. The reactive programming paradigm takes time to internalize. Factor graphs require thinking differently about model specification. RxInfer is powerful, but the barrier to entry is real.

Documentation gaps. RxInfer's documentation is improving but still lags behind PyMDP. The examples are often advanced, assuming familiarity with probabilistic programming and message passing. For newcomers, this can be daunting.

Active inference is not the primary focus. RxInfer is a general Bayesian inference library that supports active inference but isn't designed around it. You can build active inference agents in RxInfer, and many do, but it requires stitching together components yourself. There's no turnkey "active inference agent" class the way PyMDP provides one.

When to Use RxInfer

Use RxInfer when:

• You need real-time performance for robotics or control
• Your state space is continuous or high-dimensional
• You're building production systems, not research prototypes
• You're comfortable with Julia and reactive programming
• You need to integrate inference with streaming sensor data

RxInfer is where you go when PyMDP stops scaling. It's the production-grade inference engine for serious applications.

The Supporting Cast: Other Tools in the Ecosystem

PyMDP and RxInfer dominate, but they're not alone. The active inference ecosystem includes several other libraries, each solving specific problems.

SPM and the MATLAB Legacy

Statistical Parametric Mapping (SPM) is where active inference originated. Karl Friston's group at University College London built the original implementations in MATLAB as part of SPM's DEM (Dynamic Expectation Maximization) toolbox. This code is still actively maintained and used in neuroscience research.

SPM implementations are rigorous and well-tested, but they inherit MATLAB's limitations—licensing costs, slower execution, less integration with modern machine learning workflows. For neuroscience applications analyzing fMRI data with active inference models, SPM is irreplaceable. For building autonomous agents, it's largely been superseded.

pymdp-gym: Bridging to Reinforcement Learning Environments

Active inference needs environments to act in. Reinforcement learning has OpenAI Gym, a standardized API for environments. pymdp-gym provides adapters that let PyMDP agents interact with Gym environments, enabling direct comparisons between active inference and RL approaches.

This matters for research. If you want to demonstrate that an active inference agent outperforms an RL agent on some task, you need both to operate in the same environment under the same conditions. pymdp-gym makes that possible.

Genal: Factor Graphs for Active Inference

Genal is a Julia library explicitly designed for active inference using factor graphs. It's younger and less mature than RxInfer, but it focuses specifically on the active inference use case rather than general Bayesian inference.

Genal's advantage is conceptual clarity—it maps directly to the active inference literature. Its disadvantage is maturity—the ecosystem around RxInfer is larger, the performance optimizations more developed, the community more active. Genal is worth watching, but for now, RxInfer is the safer bet for production work.

ForneyLab: The RxInfer Predecessor

ForneyLab was the original Julia library for message passing on factor graphs. RxInfer is its successor, redesigned from scratch with better performance and a cleaner API. ForneyLab still exists and is maintained, but new projects should start with RxInfer.

Understanding this lineage matters because much of the early active inference literature in Julia references ForneyLab. If you're reading papers from 2018-2020, they're often using ForneyLab code. The concepts translate directly to RxInfer, but the syntax has changed.

The Practical Question: Which Tool for Which Job?

If you're trying to decide which library to learn, the decision tree is straightforward:

If you're learning active inference for the first time: Start with PyMDP. Work through the tutorials. Build a simple grid-world agent. Understand the generative model specification process. Get intuition for how beliefs update and actions get selected. Once you grasp the concepts in discrete space, you can generalize to continuous.

If you're doing research on discrete decision-making: Stick with PyMDP. The flexibility, interpretability, and community support make it ideal for hypothesis testing and concept development.

If you need real-time performance or continuous control: Move to RxInfer. Accept the steeper learning curve. Invest time in understanding factor graphs and variational message passing. The performance payoff is essential for robotics, control systems, and production deployments.

If you're in neuroscience analyzing brain data: Use SPM. It's what the field uses, and it integrates with the rest of the neuroimaging analysis pipeline.

If you're comparing active inference to reinforcement learning: Use PyMDP with pymdp-gym adapters. This gives you a fair comparison on standardized benchmarks.

There's no universal solution. The right tool depends on what you're building and why.

What's Missing: The Production Gap

Here's the hard truth: none of these tools are quite production-ready for large-scale autonomous systems.

PyMDP is too slow and too limited in scope. RxInfer is fast enough but requires significant engineering to wrap into a full agent architecture. SPM is tied to MATLAB. The gym adapters are research tools, not deployment infrastructure.

What's missing is the equivalent of what PyTorch or TensorFlow became for deep learning—a mature, battle-tested, production-grade framework that handles the full pipeline from model specification to deployment, with support for GPU acceleration, distributed training, model serving, monitoring, and iteration.

The active inference ecosystem hasn't reached that level of maturity yet. The tools exist to build such a system, but no one has packaged them into a turnkey solution.

This gap is closing. Commercial interest is growing. Robotics applications are driving demand for performant implementations. Open-source contributors are filling gaps. But we're not there yet.

If you want to deploy active inference in production today, expect to do significant custom engineering. You'll be integrating RxInfer (or writing your own inference engine), wrapping it in application-specific control logic, handling edge cases, and building monitoring infrastructure. It's feasible—several research groups and startups are doing it—but it's not plug-and-play.

The Convergence Ahead: Active Inference Meets Deep Learning

The most interesting developments aren't happening within any single library—they're happening at the intersections.

Researchers are integrating active inference with neural network function approximators. Instead of specifying transition and observation models explicitly, they're learning them from data using deep learning. Instead of exact inference over discrete states, they're using amortized variational inference with neural networks as encoders and decoders.

This hybrid approach combines the structured reasoning of active inference with the representation learning of deep learning. It's how you get agents that can operate in high-dimensional perceptual spaces (like pixels) while maintaining the interpretability and theoretical grounding of active inference.

Libraries like Pyro (Python) and Turing.jl (Julia) provide the probabilistic programming infrastructure to support this integration. These aren't active inference libraries per se, but they're being used to build active inference agents that leverage deep learning components.

The future isn't PyMDP or RxInfer or neural networks. It's all three, integrated. Structured probabilistic models for planning and reasoning, learned representations for perception, and efficient inference engines to tie it all together.

Further Reading

For those ready to dive into implementation:

• PyMDP Documentation and Tutorials: https://pymdp-rtd.readthedocs.io/ — Start here for discrete active inference
• RxInfer.jl Documentation: https://rxinfer.ml/ — The primary resource for continuous inference
• Heins, C., Millidge, B., et al. (2022). "pymdp: A Python library for active inference in discrete state spaces." Journal of Open Source Software. — The official PyMDP paper
• Bagaev, D., & de Vries, B. (2023). "Reactive Message Passing for Scalable Bayesian Inference." Scientific Programming. — Technical foundations of RxInfer
• Friston, K., et al. (2017). "Active Inference: A Process Theory." Neural Computation. — Still the canonical reference for understanding what you're implementing

This is Part 5 of the Active Inference Applied series, exploring practical implementation of active inference for AI systems and robotics.

Previous: Message Passing and Belief Propagation in Active Inference
Next: Active Inference Agents vs Reinforcement Learning: A Comparison