How to implement a two-step task

[katiacracco]

Hi,

I am trying to use this package to implement an agent in a t-maze that has to execute two actions in a single trial. Then in the next trials it has to begin from the initial location, having in memory the inference of the precedent trials.
In particular, I tried to implement the behaviour described in the attached paper.
Which approach do you suggest to use?

Thank you very much for your availability!

Uploading Friston_Deeply Felt Affect_ The Emergence of Valence in Deep Active Inference.pdf…

[conorheins]

Hi @katiacracco,

I am not aware of the details of the implementation or the specific generative model used in paper you shared, but from what you described, I would go about constructing a series of trials as follows (roughly in code below):

n_trials = 10 # number of total trials or epochs
n_timesteps = 2 # number of timesteps per trial

from pymdp.agent import Agent
from pymdp import utils

A, B, C, ... = create_A(params...), create_B(params...), create_C(params...), ... # create the generative model
D = utils.obj_array(1) # instantiate a hidden state prior with one state factor
D[0] = np.ones(2) / 2. # make a flat ('50-50') prior over the two contexts, for the single hidden state factor
pD = utils.dirichlet_like(D, scale = 1.0) # create a Dirichlet prior `pD` over the Categorical hidden state prior `D`

my_agent = Agent(A = ..., B = ..., C = ..., pD = pD, inference_algo = "MMP", use_BMA = True, policy_sep_prior = False, inference_horizon = 2, save_belief_hist = True) # instantiate the agent
env = MyEnvironment(params...) # instantiate your environment

for trial in range(n_trials):

   # run a single trial
   obs = env.reset()
   for t in range(n_timesteps):
        my_agent.infer_states(obs)
        my_agent.infer_policies()
        chosen_action = my_agent.sample_action()
        obs = env.step(action)
        
  # at the end of each trial, use the agent's beliefs about how the world was at the beginning of each trial (`t = 0`) to update its prior over hidden states, via the Dirichlet prior `pD`
  my_agent.update_D()
  my_agent.reset()

What this does is allows the agent to accumulate hidden state beliefs across trials, such that its beliefs about the hidden states at the first timestep of every trial (D), will be learned over time through Dirichlet hyperparameter learning (specifically, updating the entries of the pD array using the learning function learning.update_state_prior_dirichlet()), that itself is a function of hidden state beliefs that the agent optimizes through each trial. Note that the use of marginal message passing (inference_algo = "MMP") is quite important in this context, because you want the agent to entertain a belief about what, by the end of the trial, it believes the state of the world was at the beginning of each trial. In the POMDP / active inference world we call these "post-dictive beliefs" anecdotally.

I hope this helps, let me know if you have any more questions!

[katiacracco]

Hello,
thank you very much for your answer!

I tried to implement the code you suggested.

I observed that when I use the VANILLA algorithm and call my_agent.infer_states(obs), it returns an array qs composed of 2 other arrays: one with 4 elements and one with 2 elements. In fact, I have two hidden states, one of four levels and one of two.

Now that I use the MMP algorithm, when I call my_agent.infer_states(obs), it returns an array qs composed of 10 arrays, each with other 4 arrays containing two arrays, one with 4 elements and one with 2 elements (as the hidden states). I understand that I have 10 arrays as many as the policies I have. However I don't understand the 4 arrays in each of these 10 main arrays. I thought they would be as many as the timepoints, but I have only two timepoints.

Could you clarify how the inference is computed with the MMP algorithm?

Thank you!

[conorheins]

This is a good and subtle point.

The reason there are 4, and not 2, timepoint-specified arrays within each policy-conditioned posterior (each of your qs_pi[pi_i] arrays), is because the agent holds both backward ("postdictive" i.e. retrospective) beliefs about the previous timepoints, as well as forward (predictive, i.e. anticipated) beliefs about the future timepoints. From hearing that you have 4 timepoints total, I assume in your case agent.policy_len = 2 -- is that correct? So in this case, at the final timepoint of the trial, the agent has beliefs about hidden states, under all policies, for two timepoints in the future (t = 2 and t = 3, two "impossible" timepoints into the future, since your trials only last 2 timepoints), as well as two timepoints in the present/past (the current timepoint and the states at t = 0).

It depends on what you'd like to do with the posterior beliefs, but if you would like to truncate the posterior beliefs down to just the agent's beliefs about the past / present, you could do something like the following:

T = 2 # total length of a trial
n_policies = len(my_agent.policies) # number of policies, assume `my_agent` is your instance of the Agent() class
qs_pi_truncated = utils.obj_array(n_policies) # we are initializing a new qs_pi array to contain just the beliefs you're interested in

start_tstep = max(0, my_agent.current_timestep - my_agent.inference_horizon) # this is the first timestep of the agent's inference horizon, from the perspective of the current timestep (lower bounded by beginning of trial)
end_tstep = min(agent.current_timestep + agent.policy_len, T) # this is the last timestep of the agent's inference horizon, from the perspective of the current timestep (upper bounded by the absolute end of the trial)

for p in range(n_policies):
        qs_pi_truncated[p] = qs_pi[p][start_tstep:end_tstep] # extract just the timesteps you want from the original array. Those indices should get just the forward and backwards beliefs from the perspective of the current timestep 

[katiacracco]

Thank you for your answer, now it is clearer how the qs_pi[pi_i] arrays are constructed.

I have another question. In my model I have a rat in a tmaze that starts from the center and in one trial (in which I have two timepoints) it can do two actions. It can ask a cue (bottom arm) and then go to the arm with the reward (left or right arm), or it can go directly in one arm (without knowing if there will be a reward or a shock) and then get stuck there also for the second time point.
In my model I have 64 trials: in the first 32 the reward is in the left arm and in the last 32 the reward is in the right arm.
The behaviour I expect to have it would be:

• the rat ask the cue and then go to the left arm - for about 10 trials
• the rat improves its confidence so go directly into the left arm - until trial 32, or a few more
• the rat understand that the reward is not anymore on the left so ask for the cue and then go to the right arm - for about something more than 10 trials
• the rat improves again its confidence and go directly on the right arm - until the end

I implemented the scheme proposed in your previous answer but in that case the rat asks for the cue only on the first trial and then it is right away confident about going on the left. Then ask for the cue at trial 33 (because it received a shock) and from trial 34 go directly in the right arm. I think that the prior relies excessively on the posterior of the precedent trial and I would like to smooth the confidence improvement.

Now my question is how I would have to manage the prior update? Does the pymdp library expect to handle this behaviour?

Thank you for your suggestions, I really appreciate your help!

[conorheins]

Hi @katiacracco,
So from your description, it seems that they are learning "too quickly" -- is that a fair characterization? Basically, after a single trial they already update D significantly using their posterior from that single trial, with the consequence being that at the very next trial, they are already over-confident and don't need the cue anymore.

In order to achieve your desired behaviour as described: "the rat ask the cue and then go to the left arm - for about 10 trials", I would recommend you try a couple things:

• decrease the learning rate (lr_pD) or analogously, increase the initial scale or "pseudocount" of pD (the initial dirichlet prior over the prior p(s) over the two contexts).
• Try changing the relative magnitude of the (log) probabilities assigned to the different outcomes in the C vector (specifically, the modality-specific C vector corresponding to the reward outcome observation). If you make the agent more risk-averse (a lower probability or more negative log probability for the entry of C corresponding to observing the "shock" outcome), then even if they believe the right (respectively left) arm is better, they will still not "risk" visiting it until they are sufficiently confident. Note however, this strategy is only possible if the agent also believes that the reward /shock contingencies are probabilistic, given the context -- this would be encoded in the corresponding slices of the reward-modality-specific A array that encodes the probability of reward vs. shock, given the context identity (Left Better vs RIght Better). If you put any uncertainty into that mapping (e.g. 80% reward, 20% shock in the Right Arm when RIght Better is the context), then that will suddenly give you the ability to modulate the agent's risk-averseness using the C vector. I encourage you to write down the equation for the expected utility component of the (negative) expected free energy, if you would like to understand why.

I hope this helps! Thank you for formulating your issue so clearly.

Best,
Conor

[katiacracco]

Hi Conor,

Thank you very much for your availability! You were really helpful.

Have a good evening,
Katia