Experimental code for comparing the new thrml library with action-less POMDPs (HMMs) optimizing via iterative VFE minimization in pymdp.
This repository contains experimental code that compares two Python libraries for probabilistic discrete state inference: THRML and pymdp. The code implements a minimal single-factor POMDP test case (3 hidden states, 2 observations, 5 timesteps) and validates that THRML's sampling-based inference can approximate pymdp's variational inference prior to further testing. This project is purely experimental and makes no assumption that the libraries should necessarily be applied to the same problems nor applied for the same purposes.
THRML is a JAX-based library that uses block Gibbs sampling to generate samples from posterior distributions over discrete graphical models, while pymdp uses fixed-point iteration to compute exact marginal distributions analytically. The key difference is that THRML performs joint inference over all timesteps simultaneously (smoothing), computing P(s_t | o_0:T-1), whereas pymdp performs sequential filtering, computing P(s_t | o_0:t) one observation at a time. The validation tests confirm that when both methods have observed all data, THRML's empirical marginals (estimated from 1,000-10,000 samples) match pymdp's exact posteriors with 1-2% error at the final timestep, though earlier timesteps show larger discrepancies due to the filtering vs. smoothing distinction.
The code includes five validation tests: (1) matrix convention verification to ensure both libraries use equivalent probability matrices despite different indexing conventions, (2) energy function verification to confirm THRML's factor construction correctly encodes the POMDP's joint distribution, (3) multiple trial runs to measure sampling stability across different random seeds, (4) increased sample size tests to verify error decreases with more samples, and (5) timestep-by-timestep comparison documenting the expected divergence between filtering and smoothing. Next steps for experimentation include extending to multi-factor state spaces (e.g., num_states = [3, 2]), implementing actual policy evaluation beyond the dummy action used here, testing longer time horizons (T > 10), comparing computational performance between THRML's GPU-accelerated sampling and pymdp's CPU-based or GPU-based variational inference, and exploring whether THRML's ability to handle more complex posterior distributions (which pymdp's mean-field approximation may struggle with) provides practical advantages for specific Active Inference problems. This code serves as a minimal proof-of-concept demonstrating that THRML's categorical sampling machinery can represent discrete POMDP models, not as a production-ready implementation.
See also: