Hi all!
I have an aversive learning task, in which participant is watching two stimuli (one at a time). One is associated with shocks (30% of the times) and the other is neve paired with shock. We measure SCR as a way to assess if aversive learning had occurred.
Now, I try to build a Bayesian model to assess the learning rate (Rescorla Wagner’s alpha). To make sure the model is reasonable, I have created a simulated data based on the same assumptions.
Code of creating the simulated data:
shockVec = np.zeros(len(scrVec), dtype=np.int32) # vector to capture shock (1=yes, 0=no)
stimVec = np.zeros(len(scrVec), dtype=np.int32) # vector to capture stimulus (1=CS+, 2= CS-)
# build stim vector and shock vector based on real trial
for i, cond in enumerate(scrTwo['Condition'].values):
#print(i)
if cond=='CSplusUS1':
shockVec[i]= 1
stimVec[i] = 1
else:
shockVec[i] = 0
if cond=='CSminus1':
stimVec[i] = 0
else:
stimVec[i] = 1
print(shockVec.shape)
print(stimVec.shape)
def simulateSCR(alpha, stimVec, shockVec, intercept, slope):
scrSim = np.zeros(len(stimVec))
scrCSp = 0.5
scrCSm = 0.5
# set intercept and slopes
for i,(s,t) in enumerate(zip(stimVec,shockVec)):
if s==1:
pe = shockVec[i] - scrCSp # prediction error
scrCSp = scrCSp + alpha*pe
scrSim[i] = scrCSp
if s==0:
pe = shockVec[i] - scrCSm # prediction error
scrCSm = scrCSm + alpha*pe
scrSim[i] = scrCSm
# add intercept and slope
scrSim[i] = scrSim[i] + np.random.normal(0,1) # add noise #
scrSim[i] = slope*scrSim[i]
return scrSim
# generate 10 subjects with different alphas
alphalist = []
interceptList = []
slopeList = []
subjects = np.empty([30,10]) # create an empty matrix of trials X subjects
for i in np.arange(10):
print(i)
alpha = np.random.beta(a=1,b=1)
intercept = np.random.normal(0,1)
slope = np.random.normal(0,1)
subjects[:,i] = simulateSCR(alpha, stimVec, shockVec, intercept, slope)
alphalist.append(alpha)
interceptList.append(intercept)
slopeList.append(slope)
n_trials=30
trials, subj = np.meshgrid(range(n_trials), range(n_subj))
trials = tt.as_tensor_variable(trials.T)
subj = tt.as_tensor_variable(subj.T)
stim =np.reshape([stimVec]*n_subj, (n_subj,30)).T # transform to matrix trials x subject
shock = np.reshape([shockVec]*n_subj, (n_subj,30)).T
# turn to tensores
stim = tt.as_tensor_variable(stim)
shock = tt.as_tensor_variable(shock)
No, I know the alpha of each subjects of course.
Thanks to @Maria and @ricardoV94 (see here, and here) I was able to strat building the model.
Because SCR is a vector that includes responses for both reinforced and non reinforced, I have tweaked Ricardo’s code a bit.
The code of the model is as follows:
def update_Q(stim, shock,
Qs,vec,
alpha, n_subj):
"""
This function updates the Q table according to the RL update rule.
It will be called by theano.scan to do so recursevely, given the observed data and the alpha parameter
This could have been replaced be the following lamba expression in the theano.scan fn argument:
fn=lamba action, reward, Qs, alpha: tt.set_subtensor(Qs[action], Qs[action] + alpha * (reward - Qs[action]))
"""
PE = shock - Qs[tt.arange(n_subj), stim]
Qs = tt.set_subtensor(Qs[tt.arange(n_subj),stim], Qs[tt.arange(n_subj),stim] + alpha * PE)
# in order to get a vector of expected outcome (dependent on the stimulus presentes [CS+, CS-]
# we us if statement (switch in theano)
vec = tt.set_subtensor(vec[tt.arange(n_subj),0], (tt.switch(tt.eq(stim,1),
Qs[tt.arange(n_subj),1], Qs[tt.arange(n_subj),0])))
return Qs, vec
with pm.Model() as mB:
alpha = pm.Beta('alpha', 1,1, shape=n_subj)
beta = pm.Normal('beta',0, 1, shape=n_subj)
eps = pm.HalfNormal('eps', 10)
#betas = tt.tile(tt.repeat(beta,1), n_trials).reshape([n_trials, n_subj]) # Q_sub.shape
Qs = 0.5 * tt.ones((n_subj,2), dtype='float64') # set values for boths stimuli (CS+, CS-)
vec0 = 0.5 * tt.ones((n_subj,1), dtype='float64') # vector to save the relevant stimulus's expactation
[Qs,vec], updates = theano.scan(
fn=update_Q,
sequences=[stim, shock],
outputs_info=[Qs, vec0],
non_sequences=[alpha, n_subj])
vec = tt.concatenate([[vec0], vec[:-1]], axis=0) # add first value, remove last
vec_ = vec[trials,subj,0]* beta
# change to vector
vec_reshape = vec_.T.reshape([n_trials * n_subj])#vec_.T.reshape([30*n_subj])
scrVec = subjects.T.flatten()
scrs = pm.Normal('scrs',vec_reshape, eps, observed=scrVec)
trB = pm.sample(target_accept=.9, chains=4, cores=8, return_inferencedata=True)
Unfortunately, the model is far from recovering the real alpha of each participant. I suspect I’ve messed up with the shapes but was unable to figure it out (theano is hard…).
Any help would be much appreciated.
Thanks!