can't get posteriors on similar model to mountain car #382
-
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:
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 |
Beta Was this translation helpful? Give feedback.
Replies: 3 comments 1 reply
-
Hey! I'll check it out in the next few days and will give you a detailed response. 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 |
Beta Was this translation helpful? Give feedback.
-
What exactly are you expecting to see? As far as I can tell The Also |
Beta Was this translation helpful? Give feedback.
-
Perhaps you want to get a predictive distribution over your observed variables? You can try to use
Read more about it here and also the general documentation. |
Beta Was this translation helpful? Give feedback.
Perhaps you want to get a predictive distribution over your observed variables? You can try to use
Read more about it here and also the general documentation.