pymdp

Hello, welcome. It is April 9th, 2026 and we're in active inference model stream 007.3 on pymdp 1.0, the JAX first release of pymdp with Connor Hines and other guests. So, we will be hearing about pymdp release and discussing some topics about applications of active inference and the package. So, thank you to Connor and all the developers for their work. And Connor to you, thank you. Cool. Yeah, thanks a lot, Daniel. Um yeah, happy to be with you all. Thanks for attending. Uh yeah, so I'm really excited to share this release with you. Um this is the 100 release of pymdp. Um this is an active inference package that we first started working on back in 2019 and it's matured and changed a lot over the years. So, this um talk uh this presentation is about the 100 release, which is an upgrade from the last release, which I think was almost 3 years ago or maybe a little over than 3 years ago. Um and the kind of biggest headline change to highlight here is the fact that we rewrote the entire back end of pymdp in JAX. Um so, now it's a kind of a JAX first, JAX only package. And because [clears throat] of that rewrite and motivating that rewrite we have a bunch of new uh capabilities uh for active inference agents in in discrete state spaces. So, first of all, just because it's in JAX, it's really trivial to now put everything on GPUs and TPUs. Um that's just comes from having everything in JAX. So, now you can run, you know, whole groups, hundreds, thousands of active inference agents uh paralyzed on GPUs. You're only constrained basically by your hardware memory, whether it's GPU or TPU. Um also, because it's in JAX, which is natively fully differentiable, you can kind of hook up arbitrary differentiable modules to the front end or the back end or in the middle of any pymdp agent. So, that means you can pass gradients through the whole system and do like end-to-end learning sort of uh workflows, which is cool. Um and you know, it it it I would say um significantly increases the sort of things that you can do with discrete pymdp agents. So, you can imagine hooking up a pymdp agent to say images or audio or more high-dimensional uh input data than normally you couldn't because of the kind of poor scaling features you get with fully uh discrete categorical inputs. Um another big uh benefit of JAX is the just-in-time compilation uh feature. So, that allows you to basically run things much faster uh cuz the kind of com- computation graph of the full active inference process is being compiled ahead of time. And then you can actually just execute that in this lower uh accelerated linear algebra language XLA that that JAX kind of writes and dispatches to. So, um what that means in practice is that now pymdp can be used like a lot of these other JAX-based RL frameworks where you can get really high throughput simulations on arbitrary JAX-based reinforcement learning environments. Um so, you can kind of JIT or just-in-time compile the full agent-environment loop and that allows you to run things like orders of magnitude than you would be able to in the old um un-JITed NumPy back end. And so, beside all those I kind of just core capabilities that are unlocked by putting everything in JAX, there's also just loads of new features. Um so, new kind of algorithms that have been developed in the last couple years in the active inference community, like sophisticated inference, inductive inference, um those are all now uh in integrated into the new pymdp. So, those just allow you uh to have smarter agents in the end. So, you're actually using better planning algorithms, better inference algorithms. Um another nice thing and actually this is what originally motivated me and Dimitri Markovik, who's on this call, I believe. What originally motivated us to write the JAX back end, this was like, I think going back to 2022 even, was the ability to fit pymdp models or or active inference pomdps to experimental data from from, for instance, experimental um…