# Demo series: Controllers & Agents

---

## Active Inference for controlling a driven damped harmonic oscillator

Wouter Kouw, 06-08-2021

### System

This project considers a [harmonic oscillator](https://en.wikipedia.org/wiki/Harmonic_oscillator#Driven_harmonic_oscillators). Consider $y(t)$ as observed displacement, $x(t)$ as the unobserved state, $u(t)$ as the driving force and $e(t)$ as measurement noise. The continuous-time dynamics of the system are:

$$\begin{align}
m \frac{d^2 x(t)}{dt^2} =&\ - c \frac{d x(t)}{dt} - k x(t) + u(t) \, ,
\end{align}$$
where 
$$
m     = \text{mass} \, , \quad
c     = \text{damping} \, , \quad
k     = \text{spring stiffness} \, ,
$$
constitute the physical parameters. 

This oscillator is observed with a noisy sensor:
$$\begin{align}
y(t) =&\ x(t) + e(t)
\end{align}$$

The measurement noise is a Wiener process, with $e(t) \sim \mathcal{N}(0, \xi^{-1}dt)$ and $\xi$ as precision parameter. Note that nearly all Gaussian distributions in this demo use [precision](https://en.wikipedia.org/wiki/Precision_%28statistics%29) instead of variance. By convention, $\mathcal{N}$ uses a mean-covariance notation, so we invert precision parameters / matrices.

#### First-order system in generalised coordinates

We will first adopt a more concise notation and then divide by the leading coefficient:

$$ x'' = \frac{-c}{m} x' + \frac{-k}{m} x + \frac{1}{m} u \, .$$

In this demo, we will make the following variable substitutions:

$$ \theta_1 = \frac{-k}{m} \, , \quad \theta_2 = \frac{-c}{m} \, , \quad \eta = \frac{1}{m}, .$$

With that, we cast the above system into a multi-variate first-order form:

$$ \underbrace{\begin{bmatrix} x' \\ x'' \end{bmatrix}}_{z'} = \begin{bmatrix} 0 & 1 \\ \theta_1 & \theta_2 \end{bmatrix}  \begin{bmatrix} x \\ x' \end{bmatrix} + \begin{bmatrix} 0 \\ \eta \end{bmatrix} u \, ,$$

and discretize with $z' = (z_{k+1} - z_k) / \Delta t$ giving:

$$\begin{align} 
\left( \begin{bmatrix} x_{k+1} \\ x'_{k+1} \end{bmatrix} - \begin{bmatrix} x_{k} \\ x'_{k} \end{bmatrix} \right) / \Delta t =&\ \begin{bmatrix} 0 & 1 \\ \theta_1 & \theta_2 \end{bmatrix} \begin{bmatrix} x_k \\ x'_k \end{bmatrix} + \begin{bmatrix} 0 \\ \eta \end{bmatrix} u_k \\
\begin{bmatrix} x_{k+1} \\ x'_{k+1} \end{bmatrix} =&\ \underbrace{\begin{bmatrix} 1 & \Delta t \\ \theta_1 \Delta t & \theta_2 \Delta t + 1 \end{bmatrix}}_{A(\theta)} \begin{bmatrix} x_k \\ x'_k \end{bmatrix} + \underbrace{\begin{bmatrix} 0 \\ \eta \Delta t \end{bmatrix}}_{B(\eta)} u_k \, .
\end{align}$$

This produces a compact state-space model formulation:

$$\begin{align} 
z_{k+1} =&\ A(\theta) z_k + B(\eta) u_k \\
y_k =&\ Cx_k + e_k \, ,
\end{align}$$

where $C = \begin{bmatrix} 1 & 0 \end{bmatrix}$ is a masking vector.


```julia
using Pkg
Pkg.activate(".")
Pkg.instantiate();
```

    [32m[1m  Activating[22m[39m environment at `C:\Users\kouww\Research\actinf-oscillator\Project.toml`



```julia
using Revise
using Plots
using Colors
pyplot();
```


```julia
# Sampling step size
Δt = 0.1

# Physical parameters
m = 8.0
c = 2.0
k = 4.0

# Pack parameters
phys_params = (m, c, k);
```


```julia
# Interesting variable substitutions
ζ = c/(2*sqrt(m*k))
ω0 = sqrt(k/m)
ω1 = ω0*sqrt(1 - ζ^2)

# Report damping ratio
println("Damping ratio ζ = "*string(ζ))

# Report undamped angular frequency
println("Undamped angular frequency ω_0 = "*string(ω0))

# Report underdamped angular frequency
println("Underdamped angular frequency ω_1 = "*string(ω1))
```

    Damping ratio ζ = 0.17677669529663687
    Undamped angular frequency ω_0 = 0.7071067811865476
    Underdamped angular frequency ω_1 = 0.6959705453537528



```julia
# Measurement noise precision obtained from data
ξ = 1e8;
```


```julia
# Utility matrices
S = [1. Δt; 0. 1.]
s = [0., Δt]

# State transition coefficients
θ1 = -k/m
θ2 = -c/m
θ_true = [θ1, θ2]

# Control coefficient
η_true = 1/m

# Pack substituted variables
subs_params = (θ_true, η_true)
```




    ([-0.5, -0.25], 0.125)




```julia
function sim_sys(input, z_kmin1, params)
   "Simulate dynamic system"
    
    # Unpack parameters
    θ, η = params
    
    A = [1. Δt; 0. 1.] + [0., Δt]*θ'
    B = [0., Δt]*η
    C = [1., 0.]'
    
    # State transition
    z_k = A*z_kmin1 + B*input
    
    # Generate observation
    y_k = C*z_k + sqrt(inv(ξ))*randn(1,)[1]
    
    return y_k, z_k    
end;
```


```julia
# Time 
T = 600
stop = T-300

# Input signal
input = [cos.(Δt .*(1:stop)* 2/π); zeros(length(stop+1:T),)]

# Preallocate arrays
states = zeros(T,2)
output = zeros(T,)

for k = 2:T
    output[k], states[k,:] = sim_sys(input[k], states[k-1,:], subs_params)
end

p101 = plot(1:T, input, linewidth=1, color="red", ylabel="input", label="")
vline!([stop], color="gray", label="")

p102 = plot(1:T, states[:,1], linewidth=1, color="blue", ylabel="position", label="states")
vline!([stop], color="gray", label="")

p103 = plot(1:T, states[:,2], linewidth=1, color="darkblue", ylabel="velocity", label="states")
vline!([stop], color="gray", label="")

p100 = plot(p101, p102, p103, layout=(3,1), size=(900,400))
```


```julia
savefig(p100, "figures/example-input-output_seq1.png")
```

## Model: Adaptive filter in generalised coordinates

In practice, a system is usually affected by external circumstances. We include unknown effects by adding a zero-mean Gaussian noise vector $v_k$:

$$v_k \sim \mathcal{N}\big(0, W(\tau)^{-1} \big) \, , \quad \text{with} \quad W(\tau)^{-1} = \tau \begin{bmatrix} \frac{\Delta t^3}{3} & \frac{\Delta t^2}{2} \\ \frac{\Delta t^2}{2} & \Delta t \end{bmatrix}^{-1} \, ,
$$

where $W$ represents the precision matrix, consisting of a single precision parameter $\tau$ scaling a structured noise precision matrix.

### Likelihood

Integrating out $v_k$ and $e_k$ lets us formulate the dynamics as a Bayesian state-space model:

$$\begin{align}
z_k \sim&\ \mathcal{N}(A(\theta) z_{k-1} + B(\eta) u_k , W(\tau)^{-1}) \\
y_k \sim&\ \mathcal{N}(c^{\top} z_k, \xi^{-1}) \, .
\end{align}$$

### Priors

We choose Gaussian priors for the unknown coefficients and Gamma priors for the unknown noise precisions:

$$\begin{align}
p(\theta) = \mathcal{N}(m^{0}_{\theta}, V^{0}_{\theta}) \, , \quad 
p(\eta) = \mathcal{N}(m^{0}_{\eta}, v^{0}_{\eta}) \, , \quad  
p(\tau)= \Gamma(a^{0}_\tau, b^{0}_\tau) \, , \quad
p(\xi)= \Gamma(a^{0}_\xi, b^{0}_\xi) \, .
\end{align}$$

We also need a prior distribution for the initial state: $ p(z_0) = \mathcal{N}(0, V_{z_0})$.


### Posteriors

We can apply the Bayesian filtering equations to obtain recursive posterior estimates:

$$\begin{align}\label{eq:filtering}
    \underbrace{p(z_k | x_{1:t})}_{\text{state posterior}} &= \frac{\overbrace{p(x_t | z_t)}^{\text{likelihood}}}{\underbrace{p(x_t | x_{1:t-1})}_{\text{evidence}} }  \overbrace{\iint \underbrace{p(z_t| z_{t-1},u_{t})}_{\text{state transition}} \underbrace{q(u_t)}_{\substack{\text{control} \\ \text{signal}}} \underbrace{p(z_{t-1} | x_{1:t-1})}_{\text{state prior}} \ \mathrm{d} z_{t-1} \mathrm{d} u_t}^{\text{prior predictive $p(z_t | x_{1:t-1})$}} \, ,
\end{align}$$


## Variational Bayesian inference

### Recognition model

The recognition model will follow the generative model:

$$\begin{align}
q(\theta) = \mathcal{N}(m_{\theta}, V_{\theta}) \ , \quad 
q(\eta) = \mathcal{N}(m_{\eta}, v_{\eta}) \ , \quad  
q(\tau)= \Gamma(a_\tau, b_\tau) \, , \quad
q(\xi)= \Gamma(a_\xi, b_\xi) \, .
\end{align}$$

For the states, we have

### Free Energy function



#### Computing expectations

We need to isolate the unknown variables as separate scalar or vector variables, which can be done by splitting the equation as follows:

$$ A(\theta) = \Big(\underbrace{\begin{bmatrix} 1 & \Delta t \\ 0 & 1 \end{bmatrix}}_{S} + \underbrace{\begin{bmatrix} 0 \\ \Delta t \end{bmatrix}}_{s} \begin{bmatrix} \theta_1 & \theta_2 \end{bmatrix} \Big) \, , \quad B(\eta) = \underbrace{\begin{bmatrix} 0 \\ \Delta t \end{bmatrix}}_{s} \eta \, ,$$




```julia
using LinearAlgebra
using LaTeXStrings
using ProgressMeter
using Optim
using ForneyLab
```


```julia
using GFX
```


```julia
include("util.jl");
```


```julia
graph1 = FactorGraph()

# Coefficients
@RV θ ~ GaussianMeanPrecision(placeholder(:m_θ, dims=(2,)), placeholder(:w_θ, dims=(2,2)))
@RV η ~ GaussianMeanPrecision(placeholder(:m_η), placeholder(:w_η))

# Noise precisions
@RV τ ~ Gamma(placeholder(:a_τ), placeholder(:b_τ))

# State prior
@RV z_kmin1 ~ GaussianMeanPrecision(placeholder(:m_kmin1, dims=(2,)), placeholder(:w_kmin1, dims=(2, 2)), id=:z_kmin1)

# Composite state transition
@RV z_k ~ GeneralisedFilterX(θ, z_kmin1, η, placeholder(:u), τ, Δt=Δt, id=:z_k)

# Likelihood
@RV y ~ GaussianMeanPrecision(dot([1., 0.], z_k), ξ, id=:y_k)
placeholder(y, :y);

# Specify recognition model
q = PosteriorFactorization(z_k, z_kmin1, θ, η, τ, ids=[:z_k, :z_kmin1, :θ, :η, :τ])

# Specify and compile message passing algorithm
algo = messagePassingAlgorithm([z_k, z_kmin1, θ, η, τ], q)
eval(Meta.parse(algorithmSourceCode(algo)));
# println(algorithmSourceCode(algo))
```

### Experiment 1: Online system identification from known inputs


```julia
# Time horizon
T = 2000
stop = T-500

# Design input signal
input = [cos.(Δt .*(1:stop).* 2/π); zeros(length(stop+1:T),)];
```


```julia
# Preallocation
states = zeros(T,2)
output = zeros(T,)

# Initialize constraints
constraints = Dict()

# Initialize marginals in factor graph
marginals = Dict()
marginals[:z_kmin1] = ProbabilityDistribution(Multivariate, GaussianMeanPrecision, m=zeros(2,), w=0.1*I(2))
marginals[:z_k] = ProbabilityDistribution(Multivariate, GaussianMeanPrecision, m=zeros(2,), w=0.1*I(2))
marginals[:θ] = ProbabilityDistribution(Multivariate, GaussianMeanPrecision, m=[-.5, -.2], w=0.1*I(2))
marginals[:η] = ProbabilityDistribution(Univariate, GaussianMeanPrecision, m=0.1, w=0.1)
marginals[:τ] = ProbabilityDistribution(Univariate, Gamma, a=1e0, b=1e-3)

# Number of variational updates
num_iterations = 10

# Track estimated states
est_states = (zeros(T,2), zeros(T,2))
est_coeffs = (zeros(T,3), zeros(T,3))
est_noises = (zeros(T,), zeros(T,))

@showprogress for k = 2:T
    
    # Execute input and observe output
    output[k], states[k,:] = sim_sys(input[k], states[k-1,:], subs_params)
    
    "Update model params"
    
    # Update constraints
    constraints = Dict(:y => output[k],
                       :u => input[k],
                       :m_kmin1 => mn(marginals[:z_k]),
                       :w_kmin1 => pc(marginals[:z_k]),
                       :m_θ => mn(marginals[:θ]),
                       :w_θ => pc(marginals[:θ]),
                       :m_η => mn(marginals[:η]),
                       :w_η => pc(marginals[:η]),
                       :a_τ => marginals[:τ].params[:a],
                       :b_τ => marginals[:τ].params[:b])

    # Iterate recognition factor updates
    for ii = 1:num_iterations
        stepz_k!(constraints, marginals)
        stepθ!(constraints, marginals)
        stepz_kmin1!(constraints, marginals)
        stepη!(constraints, marginals)
        stepτ!(constraints, marginals)
    end    
    
    # Store state estimates
    est_states[1][k,:] = mn(marginals[:z_kmin1])
    est_states[2][k,:] = diag(cv(marginals[:z_kmin1]))
    est_coeffs[1][k,1:2] = mn(marginals[:θ])
    est_coeffs[2][k,1:2] = diag(cv(marginals[:θ]))
    est_coeffs[1][k,3] = mn(marginals[:η])
    est_coeffs[2][k,3] = cv(marginals[:η])
    est_noises[1][k,1] = mn(marginals[:τ])
    est_noises[2][k,1] = var(marginals[:τ])
    
end
```

    [32mProgress: 100%|█████████████████████████████████████████| Time: 0:00:09[39m[39m[39m[39m



```julia
# Inputs
p201 = plot(2:T, input[2:T], color="red", ylabel="input", label="")
vline!([stop], color="gray", label="")

# Position
p202 = plot(2:T, states[2:T,1], color="blue", ylabel="position", label="true")
# scatter!(2:T, output[2:T], markersize=2, alpha=0.5, color="black", label="outputs")
plot!(2:T, est_states[1][2:T,1], color="purple", label="inferred")
plot!(2:T, est_states[1][2:T,1], ribbon=[sqrt.(est_states[2][2:T,1]), sqrt.(est_states[2][2:T,1])], alpha=0.5, color="purple", label="")
vline!([stop], color="gray", label="")

# Velocity
p203 = plot(2:T, states[2:T,2], color="blue", ylabel="velocity", label="true")
plot!(2:T, est_states[1][2:T,2], color="purple", label="inferred")
plot!(2:T, est_states[1][2:T,2], ribbon=[sqrt.(est_states[2][2:T,2]), sqrt.(est_states[2][2:T,2])], alpha=0.5, color="purple", label="")
vline!([stop], color="gray", label="")

plot(p201, p202, p203, layout=(3,1), size=(900,600))
```


```julia
# First state coefficient
p201 = plot(2:T, est_coeffs[1][2:T,1], color="purple1", ylabel="θ1", label="inferred")
plot!(2:T, est_coeffs[1][2:T,1], ribbon=[sqrt.(est_coeffs[2][2:T,1]), sqrt.(est_coeffs[2][2:T,1])], color="purple1", label="")
plot!(θ_true[1]*ones(T-1,), label="true", color="blue")
vline!([stop], color="gray", label="")

# Second state coefficient
p202 = plot(2:T, est_coeffs[1][2:T,2], color="darkviolet", ylabel="θ2", label="inferred")
plot!(2:T, est_coeffs[1][2:T,2], ribbon=[sqrt.(est_coeffs[2][2:T,2]), sqrt.(est_coeffs[2][2:T,2])], color="darkviolet", label="")
plot!(θ_true[2]*ones(T-1,), label="true", color="blue")
vline!([stop], color="gray", label="")

# Control coefficient
p203 = plot(2:T, est_coeffs[1][2:T,3], color="darkmagenta", ylabel="η", label="inferred")
plot!(2:T, est_coeffs[1][2:T,3], ribbon=[sqrt.(est_coeffs[2][2:T,3]), sqrt.(est_coeffs[2][2:T,3])], color="darkmagenta", label="")
plot!(η_true*ones(T-1,), label="true", color="blue")
vline!([stop], color="gray", label="")

plot(p201,p202,p203, layout=(3,1), size=(900,500))
```


```julia
# Process noise precision
p401 = plot(2:T, est_noises[1][2:T], color="purple", ylabel="noise precision [τ]", label="inferred", size=(900,200))
plot!(2:T, est_noises[1][2:T], ribbon=[sqrt.(est_noises[2][2:T]), sqrt.(est_noises[2][2:T])], color="purple", label="")
```

#### Validate identified system

We benchmark the identified system by computing simulation error on a validation set.


```julia
# Time horizon
T_val = 100

# Input signal
input_val = cos.(Δt .*(1:T_val).* 4/π);
```


```julia
# Preallocate arrays
states_val = zeros(T_val,2)
output_val = zeros(T_val,)

# Preallocate prediction arrays
preds = (zeros(T_val,2), zeros(T_val,2,2))

# Initial states
states_val[1,1] = 0.0
states_val[1,2] = 0.0

# Initial predictions
preds[1][1,:] = states_val[1,:]
preds[2][1,:,:] = 0.1*I(2)

# Convenience variables
EA = (S + s*mn(marginals[:θ])')
EB = s*mn(marginals[:η])

@showprogress for k = 2:T_val
    
    # Simulate forward
    output_val[k], states_val[k,:] = sim_sys(input_val[k], states_val[k-1,:], subs_params)
    
    # Posterior predictive
    preds[1][k,:] = EA*preds[1][k-1,:] + EB*input_val[k]
    preds[2][k,:,:] = EA*preds[2][k-1,:,:]*EA' + inv(mn(marginals[:τ]))*noisecov(Δt, dims=2)

end
```


```julia
# zoom = 3:min(100,T_val)
zoom = 1:T_val

# Input
p301 = plot(zoom, input_val[zoom], ylabel="input", label="", color="red")

# Simulated position
p302 = plot(zoom, states_val[zoom,1], ylabel="position", label="true", color="blue", size=(900,300), legend=:topleft)
plot!(zoom, preds[1][zoom,1], label="simulated", linewidth=2, color="purple")
# plot!(zoom, preds[1][zoom,1], ribbon=[sqrt.(preds[2][zoom,1,1]), sqrt.(preds[2][zoom,1,1])], label="", color="purple")

# Simulated velocity
p303 = plot(zoom, states_val[zoom,2], ylabel="velocity", label="true", xlabel="time [k]", color="blue", size=(900,300), legend=:topleft)
plot!(zoom, preds[1][zoom,2], label="simulated", linewidth=2, color="purple")
# plot!(zoom, preds[1][zoom,2], ribbon=[sqrt.(preds[2][zoom,2,2]), sqrt.(preds[2][zoom,2,2])], label="", color="purple")

p300 = plot(p301,p302,p303, layout=(3,1), size=(900,400))
```


```julia
savefig(p300, "figures/simulations.png")
```


```julia
# Compute prediction error
sq_pred_error = (preds[1] .- states_val).^2

# Simulation error
MSE_sim = mean(sq_pred_error, dims=1)

# Scatter error over time
p501 = plot(1:T_val, sq_pred_error[:,1], color="blue", ylabel="Sqd. error", label="position")
title!("MSE = "*string(MSE_sim[1]))
p502 = plot(1:T_val, sq_pred_error[:,2], color="blue", xlabel="time [k]", ylabel="Sqd. error", label="velocity")
title!("MSE = "*string(MSE_sim[2]))

p500 = plot(p501, p502, layout=(2,1), size=(900,300))
```


```julia
savefig(p500, "figures/known-input_sim-errors.png")
```

### Experiment 2: Online system identification with agent-regulated inputs


```julia
function sim_future(actions, state_k, params; Δt = 1.)
    "Expected Free Energy"
    
    # Time horizon
    T = length(actions)

    # Unpack parameters
    θ, η, τ = params
    
    # Model matrices
    A = [1. Δt; 0. 1.] + [0., Δt]*θ'
    B = [0., Δt]*η
    C = [1., 0.]'
    
    # Process noise
    Σ_z = inv(τ)*[Δt^3/3 Δt^2/2; Δt^2/2 Δt]
    
    # Initialize future states
    sim_states = (zeros(2,T), zeros(2,2,T))
    
    # Unpack prior state     
    sim_states[1][:,1] = mn(state_k)
    sim_states[2][:,:,1] = cv(state_k)
    
    for t in 2:T

        # Posterior predictive
        sim_states[1][:,t] = A*sim_states[1][:,t-1] + B*actions[t]
        sim_states[2][:,:,t] = A*sim_states[2][:,:,t-1]*A' + Σ_z
        
    end
    return sim_states
end;
```


```julia
function EFE(action, state_k, goal_state, params; Δt = 1., time_horizon=1)
    "Expected Free Energy"
    
    # Unpack goal state
    μ_star = goal_state[1]
    Σ_star = goal_state[2]
    
    # Unpack prior state     
    μ_tmin1 = mn(state_k)
    Σ_tmin1 = cv(state_k)

    # Unpack model parameters
    θ, η, τ = params

    # Process noise
    Σ_τ = inv(τ)*[Δt^3/3 Δt^2/2; Δt^2/2 Δt]
    
    # Model matrices
    A = [1. Δt; 0. 1.] + [0., Δt]*θ'
    B = [0., Δt]*η
    C = [1., 0.]'
    
    cumEFE = 0
    for t in 1:time_horizon
        
        # State transition
        μ_t = A*μ_tmin1 + B*action[t]
        Σ_t = A*Σ_tmin1*A' + Σ_τ

        # Predicted observation 
        μ_y = C*μ_t
        Σ_y = C*Σ_t*C' .+ inv(ξ)

        # Ambiguity represents conditional entropy of future observation given state
        ambiguity = 0.5(log2π + log(det(inv(ξ))) - 1)
        
        # Risk represents KL-divergence between marginal future observation and goal prior
        risk = 0.5*(log(det(Σ_star)/det(Σ_y)) + 1 + (μ_y - μ_star)'*inv(Σ_star)*(μ_y - μ_star) + tr(inv(Σ_star)*Σ_y))
        
        # Cumulate EFE over future time steps
        cumEFE += risk + ambiguity
        
        # Update previous state
        μ_tmin1 = μ_t
        Σ_tmin1 = Σ_t
        
    end
    return cumEFE
end;
```

#### Plot planned policy


```julia
# Planning time horizon
TT = 50

# Goal state (mean and var)
goal_state = (-.5, 1e-6)

# Beliefs
marginals = Dict()
marginals[:z_k] = ProbabilityDistribution(Multivariate, GaussianMeanPrecision, m=zeros(2,), w=100 .*I(2))
marginals[:θ] = ProbabilityDistribution(Multivariate, GaussianMeanPrecision, m=θ_true, w=0.1*I(2))
marginals[:η] = ProbabilityDistribution(Multivariate, GaussianMeanPrecision, m=η_true, w=0.1)
marginals[:τ] = ProbabilityDistribution(Univariate, Gamma, a=1e1, b=1e-2)

# Pack parameters estimated by model
model_params = (mn(marginals[:θ]), mn(marginals[:η]), mn(marginals[:τ]))

# Objective function
G(u_k) = EFE(u_k, marginals[:z_k], goal_state, model_params, Δt=Δt, time_horizon=TT)

# Miminimize EFE
results = optimize(G, -10, 10, zeros(TT,), Fminbox(LBFGS()); autodiff=:forward)
input = Optim.minimizer(results)

# Simulated future states given selected actions
sim_states = sim_future(input, marginals[:z_k], model_params, Δt=Δt)

# Actual future states given selected actions
states = zeros(2,TT)
output = zeros(TT,)
for k = 2:TT
    output[k], states[:,k] = sim_sys(input[k], states[:,k-1], subs_params)
end
```


```julia
# Show inferred policy
p101 = plot(1:TT, input, linewidth=1, color="red", ylabel="input", label="")

# Simulated future during planning
p102 = plot(1:TT, sim_states[1][1,:], linewidth=1, color="purple", ylabel="position", label="simulated")
plot!(1:TT, sim_states[1][1,:], ribbon=[sqrt.(sim_states[2][1,1,:]), sqrt.(sim_states[2][1,1,:])], label="", color="purple", legend=:topleft)
plot!(1:TT, states[1,:], linewidth=1, color="blue", label="actual")
plot!(1:TT, goal_state[1]*ones(TT+1,), color="green", label="goal")

# Plot all
p100 = plot(p101, p102, layout=(2,1), size=(900,400))
```


```julia
savefig(p100, "ActInf-GFX-planning.png")
```

#### Execute entire policy before planning again


```julia
# Trial time
T = 100

# Planning time horizon
TT = 50

# Generate planning points
plan_t = 2:TT:T;
```


```julia
# Goal state (mean and var)
goal_state = (-.5, 1e-6)

# Preallocation
input = zeros(T,)
states = zeros(T,2)
output = zeros(T,)

# Initialize constraints
constraints = Dict()

# Initialize marginals in factor graph
marginals = Dict()
marginals[:z_kmin1] = ProbabilityDistribution(Multivariate, GaussianMeanPrecision, m=zeros(2,), w=0.1*I(2))
marginals[:z_k] = ProbabilityDistribution(Multivariate, GaussianMeanPrecision, m=zeros(2,), w=0.1*I(2))
marginals[:θ] = ProbabilityDistribution(Multivariate, GaussianMeanPrecision, m=[-0.5, -0.2], w=0.1*I(2))
marginals[:η] = ProbabilityDistribution(Univariate, GaussianMeanPrecision, m=0.1, w=0.1)
marginals[:τ] = ProbabilityDistribution(Univariate, Gamma, a=1e1, b=1e-3)

# Track estimated states
est_states = (zeros(T,2), zeros(T,2))
est_coeffs = (zeros(T,3), zeros(T,3))
est_noises = (zeros(T,), zeros(T,))

@showprogress for k = 2:T
    
    "Plan actions"

    if k in plan_t
    
        # Pack parameters estimated by model
        model_params = (mn(marginals[:θ]), mn(marginals[:η]), mn(marginals[:τ]))

        # Objective function
        G(u_k) = EFE(u_k, marginals[:z_k], goal_state, model_params, Δt=Δt, time_horizon=TT)

        # Miminimize EFE
        results = optimize(G, -10, 10, zeros(TT,), Fminbox(LBFGS()); autodiff=:forward)
    
        # Add policy to input signal 
        if k+TT-1 > T
            input[k:T] = Optim.minimizer(results)[1:T-k+1]
        else
            input[k:k+TT-1] = Optim.minimizer(results)
        end      
    end
    
    "Execute control"
    
    # Execute input and observe output
    output[k], states[k,:] = sim_sys(input[k], states[k-1,:], subs_params)
    
    "Update model params"
    
    # Update constraints
    constraints = Dict(:y => output[k],
                       :u => input[k],
                       :m_kmin1 => mn(marginals[:z_k]),
                       :w_kmin1 => pc(marginals[:z_k]),
                       :m_θ => mn(marginals[:θ]),
                       :w_θ => pc(marginals[:θ]),
                       :m_η => mn(marginals[:η]),
                       :w_η => pc(marginals[:η]),
                       :a_τ => marginals[:τ].params[:a],
                       :b_τ => marginals[:τ].params[:b])

    # Iterate recognition factor updates
    for ii = 1:num_iterations
        stepz_k!(constraints, marginals)
        stepz_kmin1!(constraints, marginals)
        stepη!(constraints, marginals)
        stepθ!(constraints, marginals)
        stepτ!(constraints, marginals)
    end    
    
    # Store state estimates
    est_states[1][k,:] = mn(marginals[:z_k])
    est_states[2][k,:] = diag(cv(marginals[:z_k]))
    est_coeffs[1][k,1:2] = mn(marginals[:θ])
    est_coeffs[2][k,1:2] = diag(cv(marginals[:θ]))
    est_coeffs[1][k,3] = mn(marginals[:η])
    est_coeffs[2][k,3] = cv(marginals[:η])
    est_noises[1][k,1] = mn(marginals[:τ])
    est_noises[2][k,1] = var(marginals[:τ])
    
end
```

    [32mProgress: 100%|█████████████████████████████████████████| Time: 0:00:04[39m



```julia
p601 = scatter(2:T, input[2:T], color="red", label="inputs")
vline!(collect(plan_t), color="gray", label="plan times")
p602 = scatter(2:T, output[2:T], color="black", markersize=1, alpha=0.5, label="outputs")
plot!(2:T, est_states[1][2:T,1], color="purple", label="inferred")
plot!(2:T, est_states[1][2:T,1], ribbon=[est_states[2][2:T,1], est_states[2][2:T,1]], color="purple", label="")
plot!(2:T, goal_state[1]*ones(T-1,), color="green", label="goal state")
p600 = plot(p601, p602, layout=(2,1), size=(900,400))
```


```julia
# First state coefficient
p211 = plot(2:T, est_coeffs[1][2:T,1], color="purple1", ylabel="θ1", label="inferred")
plot!(2:T, est_coeffs[1][2:T,1], ribbon=[est_coeffs[2][2:T,1], est_coeffs[2][2:T,1]], color="purple1", label="")
plot!(θ_true[1]*ones(T-1,), label="true", color="blue")

# Second state coefficient
p212 = plot(2:T, est_coeffs[1][2:T,2], color="darkviolet", ylabel="θ2", label="inferred")
plot!(2:T, est_coeffs[1][2:T,2], ribbon=[est_coeffs[2][2:T,2], est_coeffs[2][2:T,2]], color="darkviolet", label="")
plot!(θ_true[2]*ones(T-1,), label="true", color="blue")

# Control coefficient
p213 = plot(2:T, est_coeffs[1][2:T,3], color="darkmagenta", ylabel="η", label="inferred")
plot!(2:T, est_coeffs[1][2:T,3], ribbon=[est_coeffs[2][2:T,3], est_coeffs[2][2:T,3]], color="darkmagenta", label="")
plot!(η_true*ones(T-1,), label="true", color="blue")

p210 = plot(p211,p212,p213, layout=(3,1), size=(900,500))
```

#### Execute first planned action only


```julia
# Trial time
T = 30

# Planning time horizon
TT = 30;
```


```julia
# Preallocation
input = zeros(T,)
states = zeros(T,2)
output = zeros(T,)

# Initialize constraints
constraints = Dict()

# Initialize marginals in factor graph
marginals = Dict()
marginals[:z_kmin1] = ProbabilityDistribution(Multivariate, GaussianMeanPrecision, m=zeros(2,), w=I(2))
marginals[:z_k] = ProbabilityDistribution(Multivariate, GaussianMeanPrecision, m=zeros(2,), w=I(2))
marginals[:θ] = ProbabilityDistribution(Multivariate, GaussianMeanPrecision, m=[-.5, -.2], w=I(2))
marginals[:η] = ProbabilityDistribution(Univariate, GaussianMeanPrecision, m=0.1, w=1)
marginals[:τ] = ProbabilityDistribution(Univariate, Gamma, a=1e-2, b=1e-1)

# Track estimated states
est_states = (zeros(T,2), zeros(T,2))
est_coeffs = (zeros(T,3), zeros(T,3))
est_noises = (zeros(T,2), zeros(T,2))

@showprogress for k = 2:T
    
    "Find control"

    # Pack parameters estimated by model
    model_params = (mn(marginals[:θ]), mn(marginals[:η]), mn(marginals[:τ]))

    # Objective function
    G(u_k) = EFE(u_k, marginals[:z_k], goal_state, model_params, Δt=Δt, time_horizon=TT)

    # Miminimize EFE
    results = optimize(G, -10, 10, zeros(TT,), Fminbox(LBFGS()); autodiff=:forward)

    # Execute first planned action only
    input[k] = Optim.minimizer(results)[1]
    
    "Execute control"
    
    # Execute input and observe output
    output[k], states[k,:] = sim_sys(input[k], states[k-1,:], subs_params)
    
    "Update model params"
    
    # Update constraints
    constraints = Dict(:y => output[k],
                       :u => input[k],
                       :m_kmin1 => mn(marginals[:z_k]),
                       :w_kmin1 => pc(marginals[:z_k]),
                       :m_θ => mn(marginals[:θ]),
                       :w_θ => pc(marginals[:θ]),
                       :m_η => mn(marginals[:η]),
                       :w_η => pc(marginals[:η]),
                       :a_τ => marginals[:τ].params[:a],
                       :b_τ => marginals[:τ].params[:b])

    # Iterate recognition factor updates
    for ii = 1:num_iterations
        stepz_k!(constraints, marginals)
        stepz_kmin1!(constraints, marginals)
        stepθ!(constraints, marginals)
        stepη!(constraints, marginals)
        stepτ!(constraints, marginals)
    end    
    
    # Store state estimates
    est_states[1][k,:] = mn(marginals[:z_k])
    est_states[2][k,:] = diag(cv(marginals[:z_k]))
    est_coeffs[1][k,1:2] = mn(marginals[:θ])
    est_coeffs[2][k,1:2] = diag(cv(marginals[:θ]))
    est_coeffs[1][k,3] = mn(marginals[:η])
    est_coeffs[2][k,3] = cv(marginals[:η])
    est_noises[1][k,1] = mn(marginals[:τ])
    est_noises[2][k,1] = var(marginals[:τ])
    
end
```

    [32mProgress: 100%|█████████████████████████████████████████| Time: 0:00:08[39m



```julia
p1 = plot(2:T, input[2:T], color="red", label="inputs")
p2 = plot(2:T, output[2:T], color="blue", label="outputs")
plot!(2:T, est_states[1][2:T,1], color="purple", label="inferred")
plot!(2:T, est_states[1][2:T,1], ribbon=[est_states[2][2:T,1], est_states[2][2:T,1]], color="purple", label="")
plot!(2:T, goal_state[1]*ones(T-1,), color="green", label="goal state")
plot(p1, p2, layout=(2,1), size=(900,300))
```


```julia
# First state coefficient
p211 = plot(2:T, est_coeffs[1][2:T,1], color="purple1", ylabel="θ1", label="inferred")
plot!(2:T, est_coeffs[1][2:T,1], ribbon=[est_coeffs[2][2:T,1], est_coeffs[2][2:T,1]], color="purple1", label="")
plot!(θ_true[1]*ones(T-1,), label="true", color="blue")

# Second state coefficient
p212 = plot(2:T, est_coeffs[1][2:T,2], color="darkviolet", ylabel="θ2", label="inferred")
plot!(2:T, est_coeffs[1][2:T,2], ribbon=[est_coeffs[2][2:T,2], est_coeffs[2][2:T,2]], color="darkviolet", label="")
plot!(θ_true[2]*ones(T-1,), label="true", color="blue")

# Control coefficient
p213 = plot(2:T, est_coeffs[1][2:T,3], color="darkmagenta", ylabel="η", label="inferred")
plot!(2:T, est_coeffs[1][2:T,3], ribbon=[est_coeffs[2][2:T,3], est_coeffs[2][2:T,3]], color="darkmagenta", label="")
plot!(η_true*ones(T-1,), label="true", color="blue")

p210 = plot(p211,p212,p213, layout=(3,1), size=(900,500))
```

#### Validation


```julia
# Time horizon
T_val = 100

# Input signal
inputs_val = cos.(Δt .*range(1, stop=T_val) /π);
```


```julia
# Preallocate arrays
states_val = zeros(T_val,2)
output_val = zeros(T_val,)

# Convenience variables
EA = (S + s*mn(marginals[:θ])')
EB = s*mn(marginals[:η])

# Preallocate prediction arrays
preds = (zeros(T_val,2), zeros(T_val,2,2))

@showprogress for k = 2:T_val
    
    # Simulate forward
    output_val[k], states_val[k,:] = sim_sys(input_val[k], states_val[k-1,:], subs_params)
    
    # Posterior predictive
    preds[1][k,:] = EA*preds[1][k-1,:] + EB*input_val[k]
    preds[2][k,:,:] = EA'*preds[2][k-1,:,:]*EA + inv(mn(marginals[:τ]))*noisecov(Δt, dims=2)

end
```


```julia
# zoom = 3:min(300,T_val)
zoom = 1:T_val

# Input
p301 = plot(zoom, input_val[zoom], ylabel="input", label="", color="red")

# Simulated position
p302 = plot(zoom, states_val[zoom,1], ylabel="position", label="true", color="blue", size=(900,300), legend=:topleft)
plot!(zoom, preds[1][zoom,1], label="simulated", linewidth=2, color="purple")
# plot!(zoom, preds[1][zoom,1], ribbon=[sqrt.(preds[2][zoom,1,1]), sqrt.(preds[2][zoom,1,1])], label="", color="purple")

# Simulated velocity
p303 = plot(zoom, states_val[zoom,2], ylabel="velocity", label="true", xlabel="time [k]", color="blue", size=(900,300), legend=:topleft)
plot!(zoom, preds[1][zoom,2], label="simulated", linewidth=2, color="purple")
# plot!(zoom, preds[1][zoom,2], ribbon=[sqrt.(preds[2][zoom,2,2]), sqrt.(preds[2][zoom,2,2])], label="", color="purple")

p300 = plot(p301,p302,p303, layout=(3,1), size=(900,600))
```


```julia
# Compute prediction error
sq_pred_error = (preds[1] .- states_val).^2

# Simulation error
MSE_sim = mean(sq_pred_error, dims=1)

# Scatter error over time
p601 = plot(1:T_val, sq_pred_error[:,1], color="blue", ylabel="Sqd. error", label="position")
title!("MSE = "*string(MSE_sim[1]))
p602 = plot(1:T_val, sq_pred_error[:,2], color="blue", xlabel="time [k]", ylabel="Sqd. error", label="velocity")
title!("MSE = "*string(MSE_sim[2]))

p600 = plot(p601, p602, layout=(2,1), size=(900,300))
```


```julia
savefig(p600, "figures/agent-input_sim-errors.png")
```


```julia

```

046e3359a300b5086c2bf499a6785413972a7b93
ipynb
Jupyter Notebook
demo_ActInf-GFX-oscillator.ipynb
wmkouw/actinf-oscillator
3cee2e5e49423f431c3a85eebbe17195ed800ebc
MIT
demo_ActInf-GFX-oscillator.ipynb
wmkouw/actinf-oscillator
3cee2e5e49423f431c3a85eebbe17195ed800ebc
MIT
demo_ActInf-GFX-oscillator.ipynb
wmkouw/actinf-oscillator
3cee2e5e49423f431c3a85eebbe17195ed800ebc
MIT
Qwen/Qwen-72B
1. YES
2. YES