Assistance With a Basic Active Inference Implementation

[FraserP117]

Hello BIASlab,

My name is Fraser, I'm a Computer Science Honours student at the University of Western Australia and I'm conducting my thesis on Active Inference. I fell in love with Active Inference shortly after watching the demonstration of ForneyLab.jl on the Active Inference YouTube channel. I’ve come here at the suggestion of Dmitry Bagaev.

Now I’m intending to use RxInfer.jl for implementing several Active Inference agents (along with python’s pyro package). I’m over the moon that something like RxInfer exists, thank you very much. While my thesis is on Active Inference, the whole topic is still very new to me and I'm having some difficulty with implementing my first few models. My only experience with implementing Active Inference thus far has been in extremely contrived scenarios with the pymdp package. I’ve walked through the initial tutorials on the “Getting Started” and “Model Specification” pages on the official site: https://biaslab.github.io/RxInfer.jl/stable/manuals/getting-started/.

I'm reaching out for any advice on how to use RxInfer to specify a basic Active Inference agent to solve the canonical "Bayesian Thermostat" problem. I am not attempting to use generalised coordinates – at least not yet – and so I’ve formulated the problem like this:

**"The agent can move left or right on the interval [0.0, 100.0].

For all x in the interval, T(x) = T_0/(1+x^2) is the temperature at
position x. T_0 is the temperature at x = 0.0, let T_0 = 100.0.

The agent can only observe the temperature at its current position and
it must use Active Inference to infer the position that will yield an
observed temperature as close as possible to 4.0 (an arbitrary preference).  

The agent maintains an initial, uniform prior over the optimal location for soliciting the preferred observation, and the agent begins at x = 1.0 (for simplicity)."**

Enclosed is my first attempt to implement just such an approach with RxInfer. It won’t run without error at this stage, I only attach it to demonstrate how I have gone about framing an implementation of a solution to this problem. Please forgive my implementation, indeed, the form of the “custom_inference” function was lifted from the getting started page. 

I wouldn't be reaching out with such a “forward” query if I was not feeling a sense of mounting dread, at thus far having failed to accomplish this embarrassingly simple task in the last two weeks. I feel that being able to implement a solution to this problem, will set me up to be able to conquer the more sophisticated problems I have set for myself, in the course of my thesis.

Absolutely any comments and suggestions from anyone here would be very greatly appreciated.

Thank you for your time and attention.

Fraser

Here is my attempted implementation:

using RxInfer

# the generative model: params is a one-element vector containing the 
# temperature at the origin and x is the current position of the agent
@model function generative_model(params, x)
    T_0 = params[1]
    T_obs ~ NormalMeanVariance(T_0 / (1.0 + x^2), 0.1)
    return T_obs
end

function custom_inference(x, params)
    # yield the model and the observed temperature
    model, (T_obs,) = generative_model(params, x)

    # a reference for the future posterior marginal:
    pos_inf = nothing

    # `getmarginal` function returns an observable of future posterior marginal updates
    # We use `Rocket.jl` API to subscribe on this observable
    # As soon as posterior marginal update is available we just save it in `pos_inf`
    subscription = subscribe!(getmarginal(T_obs), (pos) -> pos_inf = pos)

    # `update!` function passes data to our data inputs
    update!(T_obs, params)

    # It is always a good practice to unsubscribe and to
    # free computer resources held by the subscription
    unsubscribe!(subscription)

    # Here we return our resulting posterior marginal
    return pos_inf
end

# the preferred temperature observation
goal_temp = 4.0

# The temperature at the origin: T_0
params = [100.0]

# set a uniform pirior over the temperature:
prior(params) = pdf(Uniform(0, 100), params[1])

# initial position of the agent:
x = 1.0

# initial observation:
T_obs = rand(
    NormalMeanVariance(
        generative_model(params, x)[1],
        0.1
    )
)

position_beliefs = [1.0] # agent believes it begins at x = 1.0
observed_temps = [T_obs] # obviously, the temperature is fully observable at any step

while abs(T_obs - goal_temp) > 1e-6
    # infer the agent's current position, given the observation of the temperature.
    position_estimated = custom_inference(x, params)

    # infer the next position via free energy minimisation
    x = argmin(
        x -> expected_free_energy(
            generative_model, prior, position_beliefs, [x]
        ), 0.0, 100.0
    )

    # observe the temperature at the new positon
    T_obs = rand(
        NormalMeanVariance(
            generative_model(params, x)[1],
            0.1
        ),
    )

    push!(position_beliefs, x)
end

println("Reached goal temperature at position $x after $(length(position_beliefs)-1) steps.")

[bvdmitri]

Hey @FraserP117 , thank you for the question. @Chengfeng-Jia implemented the "Bayesian thermostat" experiment in the biaslab/RxAgents repository and opened a PR. You can find it here: biaslab/RxAgent-Zoo#3. The PR is going to be merged in the main branch soon, but if you want you can already check the experiment in its separate branch.

Please check the experiment first as it uses modern capabilities of the RxInfer package, e.g. the inference function. There is no need to write your custom_inference function. RxInfer also does not have the expected_free_energy so I'm not sure where is it coming from. The @model specification must have at least one observation (datavar function), otherwise RxInfer cannot run inference.

[bvdmitri]

https://github.com/biaslab/RxAgents/blob/main/Thermostat/Active%20Inference%20Bayesian%20thermostat.ipynb

[FraserP117]

Thank you very much Dmitry, this has been very helpful!