def reset(self, init_qs=None):
self.curr_timestep = 1
if init_qs is None:
if self.inference_algo == 'VANILLA':
self.qs = utils.obj_array_uniform(self.n_states)
else: # in the case you're doing MMP (i.e. you have an inference_horizon > 1), we have to account for policy- and timestep-conditioned posterior beliefs
self.qs = utils.obj_array(len(self.policies))
for p_i, _ in enumerate(self.policies):
self.qs[p_i] = utils.obj_array(
self.inference_horizon + self.policy_len +
1) # + 1 to include belief about current timestep
# self.qs[p_i][0] = copy.deepcopy(self.D) # initialize the very first belief of the inference_horizon as the prior over initial hidden states
self.qs[p_i][0] = utils.obj_array_uniform(self.n_states)
first_belief = utils.obj_array(len(self.policies))
for p_i, _ in enumerate(self.policies):
first_belief[p_i] = copy.deepcopy(self.D)
if self.edge_handling_params['policy_sep_prior']:
self.set_latest_beliefs(last_belief=first_belief)
else:
self.set_latest_beliefs(last_belief=self.D)
else:
if isinstance(init_qs, Categorical):
self.qs = init_qs
else:
self.qs = Categorical(values=init_qs)
return self.qs
def rand_dist_states(num_states):
if type(num_states) is int:
num_states = [num_states]
states = obj_array(len(num_states))
for i in range(len(num_states)):
states[i] = norm_dist(np.random.rand(num_states[i]))
return states
def average_states_over_policies(qs_pi, q_pi):
"""
Parameters
----------
`qs_seq_pi` - marginal posteriors over hidden states, per policy, at the current time point
`q_pi` - posterior beliefs about policies - (num_policies x 1) numpy 1D array
Returns:
---------
`qs_bma` - marginal posterior over hidden states for the current timepoint, averaged across policies according to their posterior probability given by `q_pi`
"""
q_pi = utils.to_numpy(q_pi)
num_factors = len(qs_pi[0]) # get the number of hidden state factors using the shape of the first-policy-conditioned posterior
num_states = [qs_f.shape[0] for qs_f in qs_pi[0]] # get the dimensionalities of each hidden state factor
qs_bma = utils.obj_array(num_factors)
for f in range(num_factors):
qs_bma[f] = np.zeros(num_states[f])
for p_idx, policy_weight in enumerate(q_pi):
for f in range(num_factors):
qs_bma[f] += qs_pi[p_idx][f] * policy_weight
return qs_bma
def rand_onehot_obs(num_obs):
if type(num_obs) is int:
num_obs = [num_obs]
obs = obj_array(len(num_obs))
for i in range(len(num_obs)):
ob = np.random.randint(num_obs[i])
obs[i] = onehot(ob, num_obs[i])
return obs
def get_future_qs(self):
"""
This method only gets the last `policy_len` timesteps of each policy-conditioned belief
over hidden states. This is a step of pre-processing that needs to be done before computing
the expected free energy of policies. We do this to avoid computing the expected free energy of
policies using ('post-dictive') beliefs about hidden states in the past
"""
future_qs_seq = utils.obj_array(len(self.qs))
for p_idx in range(len(self.qs)):
future_qs_seq[p_idx] = self.qs[p_idx][
-self.
policy_len:] # this grabs only the last `policy_len` beliefs about hidden states, under each policy
return future_qs_seq
def set_latest_beliefs(self, last_belief=None):
"""
This method sets the 'last' belief before the inference horizon. In the case that the inference
horizon reaches back to the first timestep of the simulation, then the `latest_belief` is
identical to the first belief / the prior (`self.D`).
"""
if last_belief is None:
last_belief = utils.obj_array(len(self.policies))
for p_i, _ in enumerate(self.policies):
last_belief[p_i] = copy.deepcopy(self.qs[p_i][0])
if self.edge_handling_params['use_BMA'] and (
self.curr_timestep - self.inference_horizon > 0):
self.latest_belief = inference.average_states_over_policies(
last_belief, self.q_pi
) # average the earliest marginals together using posterior over policies (`self.q_pi`)
else:
self.latest_belief = last_belief
return self.latest_belief
def test_update_posterior_states_v2(self):
"""
Testing our SPM-ified version of `run_MMP` with
1 hidden state factor & 1 outcome modality, at a random fixed
timestep during the generative process
"""
past_len = 3
future_len = 4
num_policies = 5
num_states = [6, 7, 8]
num_controls = [9, 10, 11]
num_obs = [12, 13, 14]
num_modalities = len(num_obs)
A = random_A_matrix(num_obs, num_states)
B = random_B_matrix(num_states, num_controls)
prev_obs = [rand_onehot_obs(num_obs) for _ in range(past_len)]
prev_actions = np.array([rand_controls(num_controls) for _ in range(past_len)])
policies = [
np.array([rand_controls(num_controls) for _ in range(future_len)])
for _ in range(num_policies)
]
prior = rand_dist_states(num_states)
qs_seq_pi, VFE_policies = update_posterior_states_v2(A, B, prev_obs, policies, prev_actions, prior=prior, policy_sep_prior = False)
qs_seq_pi_future = utils.obj_array(num_policies)
for p_idx in range(num_policies):
qs_seq_pi_future[p_idx] = qs_seq_pi[p_idx][(1 + past_len) :]
# create C matrix
# horizon = len(qs_seq_pi_future[0])
# C = utils.obj_array(horizon)
# for t in range(horizon):
# C[t] = utils.obj_array(num_modalities)
# for g in range(num_modalities):
# C[t][g] = np.ones(num_obs[g])
C = utils.obj_array_uniform(num_obs)
q_pi, efe = update_posterior_policies_mmp(
qs_seq_pi_future,
A,
B,
C,
policies,
use_utility=True,
use_states_info_gain=True,
use_param_info_gain=False,
prior = None,
pA=None,
pB=None,
F = VFE_policies,
E = None,
gamma=16.0,
return_numpy=True,
)
qs_pi_curr_t = utils.obj_array(num_policies)
for p_idx in range(num_policies):
qs_pi_curr_t[p_idx] = qs_seq_pi[p_idx][past_len]
qs_bma = average_states_over_policies(qs_pi_curr_t, q_pi) # Bayesian model average of hidden states across policies
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
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
control_fac_idx = [1]
B = utils.obj_array(num_factors)
for f, ns in enumerate(num_states):
B[f] = np.eye(ns)
if f in control_fac_idx:
B[f] = B[f].reshape(ns, ns, 1)
B[f] = np.tile(B[f], (1, 1, ns))
B[f] = B[f].transpose(1, 2, 0)
else:
B[f] = B[f].reshape(ns, ns, 1)
C = utils.obj_array_zeros([num_ob for num_ob in num_obs])
C[1][0] = 1.0 # put a 'reward' over first observation
C[1][1] = -2.0 # put a 'punishment' over first observation
# this implies that C[1][2] is 'neutral'
agent = Agent(A=A, B=B, C=C, control_fac_idx=[1])
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
def update_posterior_states_v2(
A,
B,
prev_obs,
policies,
prev_actions=None,
prior=None,
return_numpy=True,
policy_sep_prior = True,
**kwargs,
):
"""
Update posterior over hidden states using marginal message passing
"""
# safe convert to numpy
A = utils.to_numpy(A)
num_obs, num_states, num_modalities, num_factors = utils.get_model_dimensions(A, B)
A = utils.to_arr_of_arr(A)
B = utils.to_arr_of_arr(B)
prev_obs = utils.process_observation_seq(prev_obs, num_modalities, num_obs)
if prior is not None:
if policy_sep_prior:
for p_idx, policy in enumerate(policies):
prior[p_idx] = utils.process_prior(prior[p_idx], num_factors)
else:
prior = utils.process_prior(prior, num_factors)
lh_seq = get_joint_likelihood_seq(A, prev_obs, num_states)
if prev_actions is not None:
prev_actions = np.stack(prev_actions,0)
qs_seq_pi = utils.obj_array(len(policies))
F = np.zeros(len(policies)) # variational free energy of policies
if policy_sep_prior:
for p_idx, policy in enumerate(policies):
# get sequence and the free energy for policy
qs_seq_pi[p_idx], F[p_idx] = run_mmp(
lh_seq,
B,
policy,
prev_actions=prev_actions,
prior=prior[p_idx],
**kwargs
)
else:
for p_idx, policy in enumerate(policies):
# get sequence and the free energy for policy
qs_seq_pi[p_idx], F[p_idx] = run_mmp(
lh_seq,
B,
policy,
prev_actions=prev_actions,
prior=prior,
**kwargs
)
return qs_seq_pi, F