can't get posteriors on similar model to mountain car

[adammanoogian]

Hello,

I'm working off of the mountain car example, and cannot understand why my model doesn't show the posteriors for all of the variables I've input into the model.

For the mountain car, I can see posteriors are available for

Inference result: Inference results:

  Posteriors       | available for (s_t_min, u_h_k, s, u_s_sum, u, s_g_k, x)
  
  And for my example, I can only see 
  
  Inference results:
  Posteriors       | available for (Fa_movement, u)
  Free Energy:     | Real[630.081]

Can you help me understand the process here and what's going wrong in the infer function, or any way to debug it?

Here's my code:

#actinf groundup


#import Pkg; Pkg.activate("testEnv3"); Pkg.instantiate();
using RxInfer
using Plots, GraphPlot, Cairo, Distributions
using RxInfer.ReactiveMP: getrecent, messageout

#get task env
mutable struct NoTransitionEnvironment
    num_trials::Int
    trial_data::Vector{Float64}
    switch_freq::Int
    obs::Vector{Float64}
end

env = NoTransitionEnvironment(120, [], 20, [.25, .50, .75]) #hardcoding of environment here!!!

function create_gen_process(; env)
    context_index = 1
    #make a vector of those clusters with switch info and noise
    for i in 1:env.num_trials
        #check if switch condition
        if i % env.switch_freq == 0 
            context_index = context_index < length(env.obs) ? context_index += 1 : 1
        end
        trial_t = env.obs[context_index]
        #random walk
        walk_size = .05 + .05 * randn() #hardcoded noise !
        env.obs = clamp.([obs + walk_size * rand([-1,1]) for obs in env.obs], 0.0, 1.0)

        push!(env.trial_data, trial_t)
        #env.trial_data = trial_t
    end
    return env
end
create_gen_process(; env)

#include("utils.jl")

function create_physics(noise::Float64)
    Fa = (a::Float64) -> begin
        return a + randn() * noise
    end
    return Fa
end

@meta function delta_meta()
    # approximate the state transition function using the Unscented transform
    Fa() -> Unscented()
end

function create_world(; env, init_state::Float64)
    #init first state
    next_state = init_state

    execute = (a_t::Float64, timestep::Int) -> begin
        # Compute next state- for now a passive environment
        #action has no effect on next state

        next_state = env.trial_data[timestep]
        return next_state        
    end
    observe = () -> begin
        return next_state
    end
    return execute, observe
end

#create the model
@model function heli_task_draft(
    m_u, V_u,
    m_x, V_x,
    Fa, x0,
    T)

    println("Model heli_task_draft called with inputs: m_u=$m_u, V_u=$V_u, m_x=$m_x, V_x=$V_x, Fa=$Fa, T=$T")
    
    #start loop
    if T > 1 #multiple timesteps

        m_x_last = x0

        for k in 1:T
            println("Loop count: $k")
            println("Before inference: m_x = ", m_x[:], ",V_x = ", V_x[:], ", m_u = ", m_u[:], ", V_u = ", V_u[:], ", m_x_last = ", m_x_last)
            u[k] ~ Normal(mean = m_u[k], variance = V_u[k])
            Fa_movement ~ Fa(u[k]) where { meta = DeltaMeta(method = UnscentedTransform()) } #check usage of Fa here
            # u_sum[k] ~ m_x_last + Fa_movement
            # m_x[k] ~ Normal(mean = u_sum[k], variance = V_x[k])
            m_x[k] ~ Normal(mean = m_x_last + Fa_movement, variance = V_x[k])
            m_x_last = m_x[k]
            println("After inference at step $k: u = ", u[:], ", Fa_movement = ", Fa_movement[:], ", m_x = ", m_x[:], ", m_u = ", m_u[:], ", m_x_last = ", m_x_last)

        end

    else #single timestep
        u ~ Normal(mean=m_u, variance=V_u)
        u_sum ~ Fa(m_x + u)
        m_x ~ Normal(mean=u_sum, variance=V_x)
    end

    println("Inference loop concluded")

    return (m_x, V_x, m_u, V_u)
end

#action - inference loop
#init agent and setup compute, act, slide functions
#compute takes in the prev action and observation, updates the model and runs inference
#act returns the action to take based on the updated inference
#slide moves the model forward in time

function initAgent(T::Int)
    # init variables
    x0 = 0.5
    Fa = create_physics(0.0)
    T = T #input
    m_u = zeros(T)
    V_u = ones(T)
    m_x = zeros(T)
    V_x = ones(T)

    #setup goals and priors
    m_u[end] = 0.0 #want to be no need to take actions at end
    V_x[end] = 1e-4 #goal to be more certain as time goes on
    m_x[end] = .50 #inductive bias / goal state 
    V_x[end] = 1e-4 #goal to be more certain as time goes on

    result = nothing

    compute = (action::Float64, observation::Float64) -> begin
        # update the model with the new observation
        m_u[1] = action 
        V_u[1] = 1e-4 #certainty on action

        m_x[1] = observation
        V_x[1] = 1e-4 #certainty on observation

        #smooth out priors from beginning to end
        # m_u[2:end-1] .= range(m_u[1], stop=m_u[end], length=T-2)
        # V_u[2:end-1] .= range(V_u[1], stop=V_u[end], length=T-2)
        # m_x[2:end-1] .= range(m_x[1], stop=m_x[end], length=T-2)
        # V_x[2:end-1] .= range(V_x[1], stop=V_x[end], length=T-2)

        # instantiate the model with the new action / obs
        model = heli_task_draft(Fa=Fa, T=T)
        data = Dict(:m_u => m_u, :V_u => V_u, :m_x => m_x, :V_x => V_x, :x0 => x0)
        # debugging code
        # data2 = (m_u = m_u, V_u = V_u, m_x = m_x, V_x = V_x, x0 = x0)

        imarginals = @initialization begin
            q(m_x) = m_x 
            q(V_x) = V_x
            q(m_u) = m_u
            q(V_u) = V_u
        end
        
        inferred_vars = (
            m_x = KeepLast(),
            V_x = KeepLast(),
            m_u = KeepLast(),
            V_u = KeepLast()
        )
        result = infer(model=model,
            meta=delta_meta(), 
            data=data, 
            #initialization = imarginals,
            free_energy=true, 
            #returnvars=inferred_vars
            )

        println("Inference result: ", result.posteriors)
        
        return result
    end

    act = () -> begin
        if result !== nothing
            return mode(result.posteriors[:u][2])
        else
            return 0.0 # Without inference result we return some 'random' action
            println("No inference result")
        end
    end

    # change to future states? 
    future = () -> begin 
        if result !== nothing 
            return getindex.(mode.(result.posteriors[:u]), 1)
        else
            return zeros(T)
        end
    end

    slide = () -> begin
        # # slide the model forward in time

        model  = RxInfer.getmodel(result.model)
        (s, )  = RxInfer.getreturnval(model)
        varref = RxInfer.getvarref(model, s) 
        variable = RxInfer.getvariable(varref)

        slide_msg_idx = 1# This index is model dependend
        (x0, v0) = mean_var(getrecent(messageout(variable[2], slide_msg_idx)))


        m_u[end] = 0.05 #want to be no need to take actions at end
        V_x[end] = 1e-4 #goal to be more certain as time goes on
        m_x[end] = .5 #inductive bias / goal state 
        V_x[end] = 1e-4 #goal to be more certain as time goes on
    
        
    end  

    return (compute, act, future, slide)
end

#params 
T = 5 #number of timesteps
action_noise = .05
#setup functions
rnd = rand(0.0:1.00)
(execute, observe) = create_world(env = env, init_state = rnd)
Fa = create_physics(action_noise)
(compute, act, future, slide) = initAgent(T)


# preallo  actions and Observations
agent_actions = Vector{Float64}(undef, env.num_trials)
agent_observations = Vector{Float64}(undef, env.num_trials)
agent_future = Vector{Vector{Float64}}(undef, env.num_trials) #will have to change to vector at some point
agent_inf_results = Vector{Any}(undef, env.num_trials)

#run the active inference loop

for i in 1:env.num_trials #tau vs Timerstep, got to cleanup
    agent_actions[i] = act() 
    agent_future[i] = future()
    execute(agent_actions[i], i)
    agent_observations[i] = observe()
    agent_inf_results[i] = compute(agent_actions[i], agent_observations[i])
    slide()
    println("Trial $i: Action = ", agent_actions[i], ", Observation = ", agent_observations[i])
end

plots = false  # Set to true to enable plotting, false to disable

if plots
    plot(title = "Prior updating")
    plot!(0:0.1:1.0, (x) -> pdf(NormalMeanVariance(m_u[1], V_u[1]), x), label = "prior", fill = 0, fillalpha = 0.2)
    plot!(0:0.1:1.0, (x) -> pdf(agent_inf_results[end].posteriors[:u][1], x), label = "posterior", fill = 0, fillalpha = 0.2)
    vline!([ .5 ], label = "data point")

    # Plot the true data, inferred m_u, and inferred Fa_action in one plot
    plot(title = "Inference results on state")
    m_u_inferred = [mean([mode(agent_inf_results[i].posteriors[:u][j]) for j in 1:T]) for i in 1:env.num_trials]
    Fa_movement_inferred = [mode(agent_inf_results[i].posteriors[:Fa_movement]) for i in 1:env.num_trials]
    plot!(1:env.num_trials, env.trial_data[1:env.num_trials], label = "True Data", lw = 2, color = :blue)
    plot!(1:env.num_trials, m_u_inferred, label = "Inferred m_u", lw = 2, color = :green)
    plot!(1:env.num_trials, Fa_movement_inferred, label = "Inferred Fa_movement", lw = 2, color = :purple)

    # Plot agent actions against the trial data
    plot(title = "Agent Actions vs Trial Data")
    plot!(1:env.num_trials, agent_actions, label = "Agent Actions", lw = 2, color = :green)
    plot!(1:env.num_trials, env.trial_data[1:env.num_trials], label = "Trial Data", lw = 2, color = :blue)



    #pe = choice_prediction_error(agent_actions, env.trial_data, v_u)


end

[abpolym]

Hey! I'll check it out in the next few days and will give you a detailed response.
Edit: I consulted Dmitry.

I think you might want to read up on the difference between (how to specify) latent variables and data variables in RxInfer?

See the coin toss example here as a minimal example where y is both provided data and a random variable in the generative model (similar to yourm_x).

[bvdmitri]

What exactly are you expecting to see? As far as I can tell Fa_movement and u are the only random variables in your model. All other variables are observed, so there are no posteriors over them.

The return statement in the model doesn't really specify which posteriors to compute, though it might be useful for certain scenarios, please read the model specification guide for more information here and here.

Also println statements in the model specifications aren't being executed during the inference, but rather during the model creation. It happens before the inference is running.

[bvdmitri]

Perhaps you want to get a predictive distribution over your observed variables? You can try to use

infer(
   ..., # other options
   predictvars = KeepLast(),
)

Read more about it here and also the general documentation.