def test_mmp_v2(self):
array_path = os.path.join(os.getcwd(), DATA_PATH + "mmp_b.mat")
mat_contents = loadmat(file_name=array_path)
A = mat_contents["A"][0]
B = mat_contents["B"][0]
prev_obs = mat_contents["obs_idx"].astype("int64")
policy = mat_contents["policy"].astype("int64") - 1
curr_t = mat_contents["t"][0, 0].astype("int64") - 1
t_horizon = mat_contents["t_horizon"][0, 0].astype("int64")
prev_actions = mat_contents["previous_actions"].astype("int64") - 1
result_spm = mat_contents["qs"][0]
_ = mat_contents["likelihoods"][0]
num_obs, num_states, _, num_factors = get_model_dimensions(A, B)
prev_obs = convert_observation_array(
prev_obs[:, max(0, curr_t - t_horizon):(curr_t + 1)], num_obs)
ll_seq = get_joint_likelihood_seq(A, prev_obs, num_states)
qs_seq = run_mmp_v2(A,
B,
ll_seq,
policy,
prev_actions,
num_iter=5,
grad_descent=True)
result_pymdp = qs_seq[-1]
for f in range(num_factors):
self.assertTrue(
np.isclose(result_spm[f].squeeze(), result_pymdp[f]).all())
def test_mmp_d(self):
""" Testing our SPM-ified version of `run_MMP` with
2 hidden state factors & 2 outcome modalities, at the final
timestep of the generative process (boundary condition test)
@NOTE: mmp_d.mat test has issues with the prediction errors. But the future messages are
totally fine (even at the last timestep of variational iteration."""
array_path = os.path.join(os.getcwd(), DATA_PATH + "mmp_d.mat")
mat_contents = loadmat(file_name=array_path)
A = mat_contents["A"][0]
B = mat_contents["B"][0]
prev_obs = mat_contents["obs_idx"].astype("int64")
policy = mat_contents["policy"].astype("int64") - 1
curr_t = mat_contents["t"][0, 0].astype("int64") - 1
t_horizon = mat_contents["t_horizon"][0, 0].astype("int64")
prev_actions = mat_contents["previous_actions"].astype("int64") - 1
result_spm = mat_contents["qs"][0]
likelihoods = mat_contents["likelihoods"][0]
num_obs, num_states, _, num_factors = get_model_dimensions(A, B)
prev_obs = convert_observation_array(
prev_obs[:, max(0, curr_t - t_horizon):(curr_t + 1)], num_obs)
prev_actions = prev_actions[(max(0, curr_t - t_horizon) - 1):, :]
prior = np.empty(num_factors, dtype=object)
for f in range(num_factors):
uniform = np.ones(num_states[f]) / num_states[f]
prior[f] = B[f][:, :, prev_actions[0, f]].dot(uniform)
lh_seq = get_joint_likelihood_seq(A, prev_obs, num_states)
qs_seq, _ = run_mmp(lh_seq,
B,
policy,
prev_actions[1:],
prior=prior,
num_iter=5,
grad_descent=True,
last_timestep=True)
result_pymdp = qs_seq[-1]
for f in range(num_factors):
self.assertTrue(
np.isclose(result_spm[f].squeeze(), result_pymdp[f]).all())
def test_mmp_a(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
"""
array_path = os.path.join(os.getcwd(), DATA_PATH + "mmp_a.mat")
mat_contents = loadmat(file_name=array_path)
A = mat_contents["A"][0]
B = mat_contents["B"][0]
prev_obs = mat_contents["obs_idx"].astype("int64")
policy = mat_contents["policy"].astype("int64") - 1
curr_t = mat_contents["t"][0, 0].astype("int64") - 1
t_horizon = mat_contents["t_horizon"][0, 0].astype("int64")
prev_actions = mat_contents["previous_actions"].astype("int64") - 1
result_spm = mat_contents["qs"][0]
likelihoods = mat_contents["likelihoods"][0]
num_obs, num_states, _, num_factors = get_model_dimensions(A, B)
prev_obs = convert_observation_array(
prev_obs[:, max(0, curr_t - t_horizon):(curr_t + 1)], num_obs)
prev_actions = prev_actions[(max(0, curr_t - t_horizon) - 1):, :]
prior = np.empty(num_factors, dtype=object)
for f in range(num_factors):
uniform = np.ones(num_states[f]) / num_states[f]
prior[f] = B[f][:, :, prev_actions[0, f]].dot(uniform)
ll_seq = get_joint_likelihood_seq(A, prev_obs, num_states)
qs_seq = run_mmp_v2(A,
B,
ll_seq,
policy,
prev_actions[1:],
prior=prior,
num_iter=5,
grad_descent=True)
result_pymdp = qs_seq[-1]
for f in range(num_factors):
self.assertTrue(
np.isclose(result_spm[f].squeeze(), result_pymdp[f]).all())
def test_mmp_c(self):
""" Testing our SPM-ified version of `run_MMP` with
2 hidden state factors & 2 outcome modalities, at the very first
timestep of the generative process (boundary condition test). So there
are no previous actions"""
array_path = os.path.join(os.getcwd(), DATA_PATH + "mmp_c.mat")
mat_contents = loadmat(file_name=array_path)
A = mat_contents["A"][0]
B = mat_contents["B"][0]
prev_obs = mat_contents["obs_idx"].astype("int64")
policy = mat_contents["policy"].astype("int64") - 1
curr_t = mat_contents["t"][0, 0].astype("int64") - 1
t_horizon = mat_contents["t_horizon"][0, 0].astype("int64")
# prev_actions = mat_contents["previous_actions"].astype("int64") - 1
result_spm = mat_contents["qs"][0]
likelihoods = mat_contents["likelihoods"][0]
num_obs, num_states, _, num_factors = get_model_dimensions(A, B)
prev_obs = convert_observation_array(
prev_obs[:, max(0, curr_t - t_horizon):(curr_t + 1)], num_obs)
# prev_actions = prev_actions[(max(0, curr_t - t_horizon)) :, :]
lh_seq = get_joint_likelihood_seq(A, prev_obs, num_states)
qs_seq, _ = run_mmp(lh_seq,
B,
policy,
prev_actions=None,
prior=None,
num_iter=5,
grad_descent=True)
result_pymdp = qs_seq[-1]
for f in range(num_factors):
self.assertTrue(
np.isclose(result_spm[f].squeeze(), result_pymdp[f]).all())
def test_mmp_fixedpoints(self):
array_path = os.path.join(os.getcwd(), DATA_PATH + "mmp_a.mat")
mat_contents = loadmat(file_name=array_path)
A = mat_contents["A"][0]
B = mat_contents["B"][0]
prev_obs = mat_contents["obs_idx"].astype("int64")
policy = mat_contents["policy"].astype("int64") - 1
curr_t = mat_contents["t"][0, 0].astype("int64") - 1
t_horizon = mat_contents["t_horizon"][0, 0].astype("int64")
prev_actions = mat_contents["previous_actions"].astype("int64") - 1
result_spm = mat_contents["qs"][0]
likelihoods = mat_contents["likelihoods"][0]
num_obs, num_states, _, num_factors = get_model_dimensions(A, B)
prev_obs = convert_observation_array(
prev_obs[:, max(0, curr_t - t_horizon):(curr_t + 1)], num_obs)
prev_actions = prev_actions[(max(0, curr_t - t_horizon) - 1):, :]
prior = np.empty(num_factors, dtype=object)
for f in range(num_factors):
uniform = np.ones(num_states[f]) / num_states[f]
prior[f] = B[f][:, :, prev_actions[0, f]].dot(uniform)
lh_seq = get_joint_likelihood_seq(A, prev_obs, num_states)
qs_seq, F = run_mmp(lh_seq,
B,
policy,
prev_actions[1:],
prior=prior,
num_iter=5,
grad_descent=False,
save_vfe_seq=True)
self.assertTrue((np.diff(np.array(F)) < 0).all())
return obs
def rand_controls(num_controls):
if type(num_controls) is int:
num_controls = [num_controls]
controls = np.zeros(len(num_controls))
for i in range(len(num_controls)):
controls[i] = np.random.randint(num_controls[i])
return controls
if __name__ == "__main__":
past_len = 4
future_len = 4
num_states = [8, 12, 13]
num_controls = [12, 3, 16]
num_obs = [12, 14, 6]
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)])
policy = np.array([rand_controls(num_controls) for _ in range(future_len)])
ll_seq = get_joint_likelihood_seq(A, prev_obs, num_states)
qs_seq = run_mmp_v2(A, B, ll_seq, policy, grad_descent=True)
for t, qs in enumerate(qs_seq):
print(f"Step {t} shape {[el.shape for el in qs]}")
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