docs: add mnist rgm notebook

[Arun-Niranjan]

Initial attempt post v1 release to implement scale free active inference MNIST classifier as per this paper and this demo script in SPM.

I could not achieve the accuracy levels of ~95% mentioned in the paper (this implementation gets to ~87% in the test set), and there are quite a few deviations between the paper and the demo script, or underspecified methods in general.

Of particular note:

• the SPM demo additionally includes asymptotic forgetting controlled with an eta hyperparameter
• the SPM demo does joint optimisation of A and Q during parameter learning with expectation maximisation, which we don't have in pymdp - we have the infer_parameters(...) method with a pre-specified learning rate.

Because this PR required quite a lot of extra infrastructure to get going, I recommend checking out the branch, running the mnist_classifier.ipynb notebook and then clicking back into the various functions to get a better idea of the implementation

Disclaimer: The work on this branch was done with heavy use of Claude code, especially for trying to compare against the SPM demo and also to make siginificant performance improvements with JAX optimisations

THIS PR SHOULD NOT BE MERGED AS IT IS. I would like to get some broad feedback on direction and algorithmic choices, and once we can get decent accuracy we can decide on:

• how much of this infrastructure (if any) is worth moving into pymdp core package
• how much should live in code next to the notebook to demonstrate to users how they can do scale free active inference in this way
• how to extend to the dove animation or atari games (I have simply not considered actions yet)
• how to have at least some of this notebook tested (it's currently not suitable for CI because it will need the MNIST dataset to run anything useful)

Then I will split code from this branch into sensible chunks that can be properly reviewed.

[conorheins]

Thanks for this @Arun-Niranjan and sorry for the late review.

I pulled this down and ran the training path on the full 10k images. I got roughly 83% on the first 1,000 official test images and 86.3% on the SPM-style training[10000:11000] slice. Structure learning and generation also work, so I think this is a useful and promising direction.

My main thoughts would be: maybe don't spend the next phase tuning performance to 95.1% quite yet. The SPM reference fits its SVD bases and bin ranges on all 11,000 images before splitting 10k/1k, and its “test” set is drawn from MNIST’s training split. Your implementation fits the basis from only the 130 structure exemplars and evaluates the official test set, so it's not entirely surprising that the reported accuracies are not comparable.

I’d suggest defining two modes:

1. An exact SPM-parity mode, including the original protocol, for validating the port.
2. A clean benchmark mode with train-only preprocessing and the complete official test set.

Before further optimisation, I think we should validate several algorithmic details against small exported SPM fixtures:

• soft hierarchical likelihood messages: the current pymdp categorical-observation calculation differs from SPM for non-one-hot beliefs (in pymdp the log happens inside the expectation $ \sum_o q(o) \log p(o|s)$ , whereas in SPM it happens outside i.e.$ \log \sum_o q(o) p(o|s) $
• the mutual-information-based gate, which SPM applies per likelihood mapping rather than once across a whole level like you do here;
• full-spectrum SVD thresholding followed by the 16-mode truncation;
• symmetric zero-centred bins;
• explicit masking of padded child outcomes;
• total hierarchical ELBO, so you can reproduce the paper’s accuracy-versus-coverage result, as done here.

I’d keep the implementation under examples/experimental for now and add small deterministic unit/golden tests around those primitives. The expensive notebook can be a nightly/manual integration test. Once parity is demonstrated, we can decide which pieces -- hierarchical message passing, conjugate updates, and the MI gate -- are mature enough to promote into pymdp proper.

Actions or the Dove/Atari extensions do not need to be in scope for this PR; establishing a correct static MNIST RGM is already a substantial and useful milestone.

[conorheins]

from what i can tell, this is summing across every patch and modality, whereas SPM gates each likelihood mapping independently. With $\beta = 512$, this could materially change what gets learned

[conorheins]

it looks like you first truncate to the top nksingular values here before calculating the 'retention denominator', whereas SPM thresholds using the full spectrum (more than nk vlaues) and then only caps them post-threshold at 16: https://github.com/spm/spm/blob/23a9359a60bb13eec207591e238ec24b367f42ad/spm_svd.m#L35-L82

[conorheins]

SPM uses [-max(abs(u)), +max(abs(u))]

[conorheins]

maybe i'm confused here, but doesn't this conflate the fitting of the SVD frnot-end with the fitting of the structuer learning algorithm? I thought they should be separate sets of inputs --one dataset for fitting the front end and then 13-per-class exemplars for fast structure learning itself, using the separately-learned SVD basis frontend

[conorheins]

The code from here seems to suggest that probability can leak into nonexistent states.
The batched JAX representation pads every local mapping to the largest child-state vocabulary. Concentration is then added to those padded rows, while only hidden-state masks (not observation mask) are consistently carried through inference and generation, see also

pymdp/examples/renormalising_generative_models/rgm_hierarchy.py

Lines 370 to 385 in 2618449

	# --- Phase A: initialize learnable agents at ALL levels ---
	for lv_idx, level in enumerate(self.levels):
	is_cls = lv_idx == len(self.levels) - 1
	concentration = concentration_cls if is_cls else concentration_lower
	pA = [A_m + concentration for A_m in level.agent.A]
	updated_agent = _build_agent(
	list(level.agent.A),
	level.agent.batch_size,
	level.agent.num_states[0],
	level.agent.num_modalities,
	level.valid_mask,
	learn_A=True,
	pA=pA,
	categorical_obs=level.agent.categorical_obs,
	)
	self.levels[lv_idx] = RGMLevel(updated_agent, level.valid_mask, level.stats)

[conorheins]

i don't think this matches the code path the demo actually runs: the 16-iteration qa = pa; qa += cross(O,Q) scheme lives in spm_backwards, which is only entered when OPTIONS.B=1 -- and the MNIST demo sets OPTIONS.B = 0. spm_VBX itself is explicitly non-iterative (one belief-propagation pass with A held fixed), so the demo's update is a single inference under the prior A followed by a single Dirichlet accumulation. iterating Q against an A refit to the same image is self-reinforcing and will sharpen the counts towards one-hot relative to SPM -- worth trying num_iter=1 and comparing before attributing the accuracy gap elsewhere

[conorheins]

refined_soft[lv_idx] is a posterior that has already absorbed this level's observations (the refinement pass re-ran infer_states on the same obs list), and the E-step in _interleaved_em then multiplies the same observations in again -- so the evidence gets counted twice. in SPM the prior for the learning posterior is the top-down message alone (mdp.D{f} = spm_dot(A, Q_parent) folded into the child's D), with the observation entering once. i think passing the _top_down_D_hierarchical / _top_down_D_from_cls output here instead of the refined posterior would be closer to the reference

[conorheins]

the outcome side of the cross(O, Q) accumulation here (and obs_soft_cls = [child_beliefs] below for the classification level) comes from the unsupervised bottom-up pass. in SPM the supervised one-hot D conditions the whole downward sweep before inference happens, so the child posteriors that come back up as the outcomes O are label-conditioned at every level by the time qa += cross(O,Q) runs. early in training this matters: SPM's counts get pulled toward the correct class's exemplar states, whereas these reinforce whatever the unsupervised parse said (only the Q side is label-corrected). using the refined (post-top-down) beliefs as the observation side would be closer to the reference