Questions on JointMarginalStorage in EFEasVFE

[meditans]

Hey @wouterwln I'd like to ask a question about the EFEasVFE implementation. Let's take for example the TMaze model; in it we have something like:

        loc_marginalstorage = JointMarginalStorage(Contingency(ones(size(location_transition_tensor))))
        loc_marginalcomponent = JointMarginalMetaComponent(loc_marginalstorage, 1, 3)
        u[t] ~ Exploration(reward_observation) where {meta=JointMarginalMeta([loc_marginalcomponent])}
        location[t] ~ DiscreteTransition(previous_location, location_transition_tensor, u[t]) where {meta=loc_marginalstorage}

Could you help me understand what the JointMarginalStorage and JointMarginalMetaComponent do and why they are needed? I'm particularly interested in what problem they are solving and the steps you used to come up with the implementation.

[wouterwln]

Hey @meditans , thanks for checking out the paper! In the paper it is stated that for every iteration of the variational inference procedure, we recompute the epistemic priors. The JointMarginalStorage and JointMarginalMeta objects are used to funnel the joint marginals around transition and observation nodes back into the Exploration and Ambiguity nodes such that all information is available in these nodes at the end of every iteration of message passing.

The reason JointMarginalMetaComponent is here is because I want to be able to pass multiple JointMarginalStorage's into a JointMarginalMeta, for example in the Minigrid experiments when we want to compute ambiguity over multiple observations.

[meditans]

Thank you @wouterwln, this helped getting me unblocked. Basically the gist, looking at the StochasticMaze example, is that the equations for the empirical priors are:

but in the code we write:

        u[t] ~ Exploration(y_current) where {meta=JointMarginalMeta([state_marginalcomponent])}
        s[t] ~ Ambiguity(y_current) where {meta=JointMarginalMeta([observation_marginalcomponent])}

and the JointMarginalMeta captures the information about the state in the first case (for u[t]) and about the observations in the second case (for s[t]).

Now I have a followup question on an aspect that still confuses me: consider this snippet from the StochasticMaze example again

for t in 1:T
    state_marginalstorage = JointMarginalStorage(Contingency(ones(n_states, n_states, n_actions)))
    state_marginalcomponent = JointMarginalMetaComponent(state_marginalstorage, 1, 3)
    u[t] ~ Exploration(y_current) where {meta=JointMarginalMeta([state_marginalcomponent])}
    s[t] ~ DiscreteTransition(s_prev, B, u[t]) where {meta=state_marginalstorage}
    ...
end

Here, the variables state_marginalstorage and state_marginalcomponent are initialized to constants. How do they mutate in the course of the run? It doesn't seem to be the Exploration node, that's just computing the entropies, so by exclusion I think it could be the DiscreteTransition node with that meta. Could you elucidate?

[wouterwln]

In the rules that compute the joint marginals (@marginalrule), this storage is updated with the updated joint marginal distribution around the transition node. Then that can be used by the Exploration node in the message it sends. The constant initialization is just indeed an initialization which will be updated over the course of the variational inference procedure.

[meditans]

Thank you, I was confused by the fact that the lines for state_marginalstorage and state_marginalcomponent

for t in 1:T
    state_marginalstorage = JointMarginalStorage(Contingency(ones(n_states, n_states, n_actions)))
    state_marginalcomponent = JointMarginalMetaComponent(state_marginalstorage, 1, 3)
    ...
end

are in the inner for loop; that initially tricked me into thinking they are executed every time, nullifying the result of the update. But of course this is just done once by the macro at model construction time, while the updated happen during the actual inference of the model, so they are never reset.