def test_fpi_inference(self):
num_obs = [2, 4]
num_states = [2, 2]
num_control = [2, 2]
A = utils.random_A_matrix(num_obs, num_states)
B = utils.random_B_matrix(num_states, num_control)
C = utils.obj_array_zeros([num_ob for num_ob in num_obs])
C[1][0] = 1.0
C[1][1] = -2.0
agent = Agent(A=A, B=B, C=C, control_fac_idx=[1])
o, s = [0, 2], [0, 0]
qx = agent.infer_states(o)
agent.infer_policies()
action = agent.sample_action()
self.assertEqual(len(action), len(num_control))
def test_mmp_active_inference(self):
"""
Tests to make sure whole active inference loop works (with various past and future
inference/policy horizons).
@TODO: Need to check this against a MATLAB output, where
the sequence of all observations / actions / generative model
parameters are used (with deterministic action sampling and
pre-determined generative process outputs - i.e. no effects of action)
"""
num_obs = [3, 2]
num_states = [4, 3]
num_control = [1, 3]
A = utils.random_A_matrix(num_obs, num_states)
B = utils.random_B_matrix(num_states, num_control)
C = utils.obj_array_zeros(num_obs)
C[1][0] = 1.0
C[1][1] = -2.0
agent = Agent(A=A,
B=B,
C=C,
control_fac_idx=[1],
inference_algo="MMP",
policy_len=2,
inference_horizon=3)
T = 10
for t in range(T):
o = [
np.random.randint(num_ob) for num_ob in num_obs
] # just randomly generate observations at each timestep, no generative process
qx = agent.infer_states(o)
agent.infer_policies()
action = agent.sample_action()
print(agent.prev_actions)
print(agent.prev_obs)
def test_mmp_inference(self):
num_obs = [2, 4]
num_states = [2, 2]
num_control = [2, 2]
A = utils.random_A_matrix(num_obs, num_states)
B = utils.random_B_matrix(num_states, num_control)
C = utils.obj_array_zeros(num_obs)
C[1][0] = 1.0
C[1][1] = -2.0
agent = Agent(A=A,
B=B,
C=C,
control_fac_idx=[1],
inference_algo="MMP",
policy_len=5,
inference_horizon=1)
o = [0, 2]
qx = agent.infer_states(o)
print(qx[0].shape)
print(qx[1].shape)
def run_mmp(lh_seq,
B,
policy,
prev_actions=None,
prior=None,
num_iter=10,
grad_descent=False,
tau=0.25,
last_timestep=False,
save_vfe_seq=False):
"""
Marginal message passing scheme for updating posterior beliefs about multi-factor hidden states over time,
conditioned on a particular policy.
Parameters:
--------------
`lh_seq`[numpy object array]:
Likelihoods of hidden state factors given a sequence of observations over time. This is logged beforehand
`B`[numpy object array]:
Transition likelihood of the generative model, mapping from hidden states at T to hidden states at T+1. One B matrix per modality (e.g. `B[f]` corresponds to f-th factor's B matrix)
This is used in inference to compute the 'forward' and 'backward' messages conveyed between beliefs about temporally-adjacent timepoints.
`policy` [2-D numpy.ndarray]:
Matrix of shape (policy_len, num_control_factors) that indicates the indices of each action (control state index) upon timestep t and control_factor f in the element `policy[t,f]` for a given policy.
`prev_actions` [None or 2-D numpy.ndarray]:
If provided, should be a matrix of previous actions of shape (infer_len, num_control_factors) taht indicates the indices of each action (control state index) taken in the past (up until the current timestep).
`prior`[None or numpy object array]:
If provided, this a numpy object array with one sub-array per hidden state factor, that stores the prior beliefs about initial states (at t = 0, relative to `infer_len`).
`num_iter`[Int]:
Number of variational iterations
`grad_descent` [Bool]:
Flag for whether to use gradient descent (predictive coding style)
`tau` [Float]:
Decay constant for use in `grad_descent` version
`last_timestep` [Bool]:
Flag for whether we are at the last timestep of belief updating
`save_vfe_seq` [Bool]:
Flag for whether to save the sequence of variational free energies over time (for this policy). If `False`, then VFE is integrated across time/iterations.
Returns:
--------------
`qs_seq`[list]: the sequence of beliefs about the different hidden state factors over time, one multi-factor posterior belief per timestep in `infer_len`
`F`[Float or list, depending on setting of save_vfe_seq]
"""
# window
past_len = len(lh_seq)
future_len = policy.shape[0]
if last_timestep:
infer_len = past_len + future_len - 1
else:
infer_len = past_len + future_len
future_cutoff = past_len + future_len - 2
# dimensions
_, num_states, _, num_factors = get_model_dimensions(A=None, B=B)
B = to_arr_of_arr(B)
# beliefs
qs_seq = obj_array(infer_len)
for t in range(infer_len):
qs_seq[t] = obj_array_uniform(num_states)
# last message
qs_T = obj_array_zeros(num_states)
# prior
if prior is None:
prior = obj_array_uniform(num_states)
# transposed transition
trans_B = obj_array(num_factors)
for f in range(num_factors):
trans_B[f] = spm_norm(np.swapaxes(B[f], 0, 1))
# full policy
if prev_actions is None:
prev_actions = np.zeros((past_len, policy.shape[1]))
policy = np.vstack((prev_actions, policy))
# initialise variational free energy of policy (accumulated over time)
if save_vfe_seq:
F = []
F.append(0.0)
else:
F = 0.0
for itr in range(num_iter):
for t in range(infer_len):
for f in range(num_factors):
# likelihood
if t < past_len:
lnA = spm_log(spm_dot(lh_seq[t], qs_seq[t], [f]))
else:
lnA = np.zeros(num_states[f])
# past message
if t == 0:
lnB_past = spm_log(prior[f])
else:
past_msg = B[f][:, :, int(policy[t - 1,
f])].dot(qs_seq[t - 1][f])
lnB_past = spm_log(past_msg)
# future message
if t >= future_cutoff:
lnB_future = qs_T[f]
else:
future_msg = trans_B[f][:, :, int(policy[t, f])].dot(
qs_seq[t + 1][f])
lnB_future = spm_log(future_msg)
# inference
if grad_descent:
lnqs = spm_log(qs_seq[t][f])
coeff = 1 if (t >= future_cutoff) else 2
err = (coeff * lnA + lnB_past + lnB_future) - coeff * lnqs
err -= err.mean()
lnqs = lnqs + tau * err
qs_seq[t][f] = softmax(lnqs)
if (t == 0) or (t == (infer_len - 1)):
F += +0.5 * lnqs.dot(0.5 * err)
else:
F += lnqs.dot(
0.5 * (err - (num_factors - 1) * lnA / num_factors)
) # @NOTE: not sure why Karl does this in SPM_MDP_VB_X, we should look into this
else:
qs_seq[t][f] = softmax(lnA + lnB_past + lnB_future)
if not grad_descent:
if save_vfe_seq:
if t < past_len:
F.append(
F[-1] +
calc_free_energy(qs_seq[t],
prior,
num_factors,
likelihood=spm_log(lh_seq[t]))[0])
else:
F.append(
F[-1] +
calc_free_energy(qs_seq[t], prior, num_factors)[0])
else:
if t < past_len:
F += calc_free_energy(qs_seq[t],
prior,
num_factors,
likelihood=spm_log(lh_seq[t]))
else:
F += calc_free_energy(qs_seq[t], prior, num_factors)
return qs_seq, F
from pymdp.agent import Agent
from pymdp.core import utils
from pymdp.core.maths import softmax
from pymdp.distributions import Categorical, Dirichlet
import copy
obs_names = ["state_observation", "reward", "decision_proprioceptive"]
state_names = ["reward_level", "decision_state"]
action_names = ["uncontrolled", "decision_state"]
num_obs = [3, 3, 3]
num_states = [2, 3]
num_modalities = len(num_obs)
num_factors = len(num_states)
A = utils.obj_array_zeros([[o] + num_states for _, o in enumerate(num_obs)])
A[0][:, :, 0] = np.ones((num_obs[0], num_states[0])) / num_obs[0]
A[0][:, :, 1] = np.ones((num_obs[0], num_states[0])) / num_obs[0]
A[0][:, :, 2] = np.array([[0.8, 0.2], [0.0, 0.0], [0.2, 0.8]])
A[1][2, :, 0] = np.ones(num_states[0])
A[1][0:2, :, 1] = softmax(
np.eye(num_obs[1] - 1)
) # bandit statistics (mapping between reward-state (first hidden state factor) and rewards (Good vs Bad))
A[1][2, :, 2] = np.ones(num_states[0])
# establish a proprioceptive mapping that determines how the agent perceives its own `decision_state`
A[2][0, :, 0] = 1.0
A[2][1, :, 1] = 1.0
A[2][2, :, 2] = 1.0
def update_posterior_policies_mmp(
qs_seq_pi,
A,
B,
C,
policies,
use_utility=True,
use_states_info_gain=True,
use_param_info_gain=False,
prior=None,
pA=None,
pB=None,
F=None,
E=None,
gamma=16.0,
return_numpy=True,
):
"""
`qs_seq_pi`: numpy object array that stores posterior marginals beliefs over hidden states for each policy.
The structure is nested as policies --> timesteps --> hidden state factors. So qs_seq_pi[p_idx][t][f] is the belief about factor `f` at time `t`, under policy `p_idx`
`A`: numpy object array that stores likelihood mappings for each modality.
`B`: numpy object array that stores transition matrices (possibly action-conditioned) for each hidden state factor
`policies`: numpy object array that stores each (potentially-multifactorial) policy in `policies[p_idx]`. Shape of `policies[p_idx]` is `(num_timesteps, num_factors)`
`use_utility`: Boolean that determines whether expected utility should be incorporated into computation of EFE (default: `True`)
`use_states_info_gain`: Boolean that determines whether state epistemic value (info gain about hidden states) should be incorporated into computation of EFE (default: `True`)
`use_param_info_gain`: Boolean that determines whether parameter epistemic value (info gain about generative model parameters) should be incorporated into computation of EFE (default: `False`)
`prior`: numpy object array that stores priors over hidden states - this matters when computing the first value of the parameter info gain for the Dirichlet parameters over B
`pA`: numpy object array that stores Dirichlet priors over likelihood mappings (one per modality)
`pB`: numpy object array that stores Dirichlet priors over transition mappings (one per hidden state factor)
`F` : 1D numpy array that stores variational free energy of each policy
`E` : 1D numpy array that stores prior probability each policy (e.g. 'habits')
`gamma`: Float that encodes the precision over policies
`return_numpy`: Boolean that determines whether output should be a numpy array or an instance of the Categorical class (default: `True`)
"""
A = utils.to_numpy(A)
B = utils.to_numpy(B)
num_obs, num_states, num_modalities, num_factors = utils.get_model_dimensions(
A, B)
horizon = len(qs_seq_pi[0])
num_policies = len(qs_seq_pi)
# initialise`qo_seq` as object arrays to initially populate `qo_seq_pi`
qo_seq = utils.obj_array(horizon)
for t in range(horizon):
qo_seq[t] = utils.obj_array_zeros(num_obs)
# initialise expected observations
qo_seq_pi = utils.obj_array(num_policies)
for p_idx in range(num_policies):
# qo_seq_pi[p_idx] = copy.deepcopy(obs_over_time)
qo_seq_pi[p_idx] = qo_seq
efe = np.zeros(num_policies)
if F is None:
F = np.zeros(num_policies)
if E is None:
E = np.zeros(num_policies)
for p_idx, policy in enumerate(policies):
qs_seq_pi_i = qs_seq_pi[p_idx]
for t in range(horizon):
qo_pi_t = get_expected_obs(qs_seq_pi_i[t], A)
qo_seq_pi[p_idx][t] = qo_pi_t
if use_utility:
efe[p_idx] += calc_expected_utility(qo_seq_pi[p_idx][t], C)
if use_states_info_gain:
efe[p_idx] += calc_states_info_gain(A, qs_seq_pi_i[t])
if use_param_info_gain:
if pA is not None:
efe[p_idx] += calc_pA_info_gain(pA, qo_seq_pi[p_idx][t],
qs_seq_pi_i[t])
if pB is not None:
if t > 0:
efe[p_idx] += calc_pB_info_gain(
pB, qs_seq_pi_i[t], qs_seq_pi_i[t - 1], policy)
else:
if prior is not None:
efe[p_idx] += calc_pB_info_gain(
pB, qs_seq_pi_i[t], prior, policy)
q_pi = softmax(efe * gamma - F + E)
if return_numpy:
q_pi = q_pi / q_pi.sum(axis=0)
else:
q_pi = utils.to_categorical(q_pi)
q_pi.normalize()
return q_pi, efe