Add EFE Minimization via Message Passing notebook

[meditans]

Add a tutorial recreation of the paper "A Message Passing Realization of Expected Free Energy Minimization" by Wouter Nuijten, Mykola Lukashchuk, Thijs van de Laar and Bert de Vries

@wouterwln this is the Stochastic Maze example from the paper. Thank you again for all your insights, I'd welcome some feedback when you get the time!

[bvdmitri]

Thanks @meditans ! I'm ok with the example, @wouterwln is currently a bit busy, when he is available he will also check and then we can merge :)

[meditans]

Thank you guys, and no problem at all, it's not urgent on my part!

[FraserP117]

Excellent stuff @meditans! Having just quickly experimented with your notebook, I suspect the following function, from the observation tensor section:

function observation_tensor(grid :: Grid)
    obss = zeros((25, 25))
    for source in grid.indices
        if source in grid.noisy_observations
            observations = 0.2 * exact(source) + 0.8 * linear_combination(adjacent(source, example_grid))
        else
            observations = exact(source)
        end

        for (observation, weight) in observations
            obss[to_linear(grid, observation), to_linear(grid, source)] = weight
        end
    end
    return obss
end

Should, perhaps be amended to:

function observation_tensor(grid :: Grid)
    n = length(grid.indices)
    obss = zeros(n, n) # now generic 
    for source in grid.indices
        if source in grid.noisy_observations
            observations = 0.2 * exact(source) + 0.8 * linear_combination(adjacent(source, grid)) # evaluates the actual argument!
        else
            observations = exact(source)
        end

        for (observation, weight) in observations
            obss[to_linear(grid, observation), to_linear(grid, source)] = weight
        end
    end
    return obss
end

This is a very minor point.

[meditans]

Good catch @FraserP117, fixed!

[wouterwln]

Hey @meditans ! Thanks so much for putting this together. This is a really clean reproduction of the message-passing approach for Expected Free Energy. The way you’ve broken down the factor graph implementation for the stochastic maze is super clear and will be a huge help for users.

I have a few small "pedantic" requests and some phrasing tweaks before we can get this merged into the official examples:

1. Tone & Documentation Style

Since this is going to live on the official package documentation, I’m a bit self-conscious about the "insightful paper" bit in the intro! Could we pivot the phrasing to be a bit more neutral/objective?

• Suggestion: "This notebook reproduces the methodology from 'A Message Passing Realization of Expected Free Energy Minimization'..."
• Pro-tip: It might be worth running the prose through an LLM just to give it a final "documentation-style" polish to ensure it matches the rest of the site’s voice.

2. Mathematical Notation ($x \sim$ vs. $x =$)

In the section describing the priors, you currently use the equality sign:

• u_t = Cat(...)
• x_t = Cat(...)

Technically, $x_t$ is a random variable distributed as a categorical, rather than being equal to the distribution object itself. Could you swap those to the "distributed as" tilde symbol ($\sim$)?

• $u_t \sim \text{Cat}(\dots)$
• $x_t \sim \text{Cat}(\dots)$

3. Factor Graph Visuals & Equality Nodes

I noticed the figure for the factor graph is missing some equality nodes. For the sake of formal correctness, it would be great to get those added in.

• Tooling Note: I actually use Ipe to draw all the factor graphs in the original paper. If you want the notebook figure to perfectly match the paper's aesthetic and formal notation, I’d recommend giving Ipe a shot!

Aside from those minor points, the logic is spot on. I like the visualizations on the transition tensors to build intuition.

Let me know when you’ve had a chance to update these, and I’ll get this merged!

[meditans]

Thank you very much for the review, @wouterwln!

1. I reworked the tone of the introduction, making it more formal. LLMs flagged another passage,
   ...the planning part contains the key insight of the paper: two carefully chosen priors are added...,
   which in their opinion could have been ...the planning part introduces two priors added to the model... but I find that bland, so kept my version, that focuses the attention of the reader on the important point (imo). But if you prefer, I can modify that one also.
2. Thank you for catching that, oversight on my part!
3. I have fixed the graph, adding the missing equality nodes, and also adding closed edges for the parts in which the observation and the previous action enter the model. Being used to typing categorical diagrams, I bit the bullet and redid the diagram in fletcher (basically tikz for typst), as I think this will allow me to generate this kind of diagrams mostly automatically from now on.

Tell me if there's anything else you want me to change

[wouterwln]

Thanks for the quick update! Looks good to me now. Let's see what CI has to say about it and if all checks pass, I'll merge it.

[bvdmitri]

@wouterwln I think CI is fine despite it failing. The failing example is the one that uses OPENAI_KEY and it is not available for forked repositories

[bvdmitri]

I will merge at my responsibility, should pass on the main branch with the proper env variables