VTA_stan.csv (47.0 KB)
Hi everyone,
I have been trying to work on writing a custom stan model for modelling data from an Instrumental learning task. I have attached both the stan model and sample data as a csv file. I will describe the task here for context.
The task has 3 trial types, reward loss and neutral. One pair of fractal images was associated with each potential outcome type; win, loss or neutral. Associations between outcomes and fractal shapes were randomized across subjects. Win trials had the possible outcomes ‘you win’ or ‘nothing’ (no change) and loss trials had the possible outcomes ‘you lost’ or ‘nothing’. A neutral condition with outcomes ‘look’ or ‘nothing was included whereby no change occurred.
I have attached the sample data from this experiment. I am using cmstandR in R.
This is what the sample data looks like. The attached csv file should have the data for all 14 subjects.
Subjid – subject id
Cond – condition i.e. trial type. 1 is win trial, 2 is loss trial, 3 is netural trial
Decisions – will be 1 or 2 as there are only two options to choose from in each trial type
Outcomes – 1 is a win in the win trial, -1 is a loss in the loss trial. 0 Is nothing across all the trial types.
An illustrative figure of the experiment has been included.
I am looking to extract trial by trial prediction errors across trial types as I want to correlate to neuroimaging data. I am also calculating learning rates and sensitivities for each participant for each trial type.
The way I have written my model is to check what the condition of each trial is and then calculate the parameters such as learning rate and sensitivities as well as the prediction errors per trial.
My main question is, is my model written well for the given experiment and data? Should I rescore the outcome data of the trials that would be different from the current condition? Am I calculating the prediction errors and other model parameters correctly?
The model does run etc but I want to see if I have conceptually messed up somewhere in the stan code.
Thanks a lot!
data {
int<lower=1> N; // # of subjects
int<lower=1> T; // max # of trials
array[N] int<lower=1, upper=T> Tsubj; // # of trials within subjects
array[N, T] int<lower=-1, upper=3> cond; // trial condition
array[N, T] int<lower=-1, upper=2> decisions; // decisions
array[N, T] int<lower=-1, upper=1> outcomes; // outcomes
array[N, T] int<lower=-1, upper=1> new_block; // new block indicator to reset Q-values
}
transformed data {
array[3, 2] real initV;
initV = rep_array(0.0, 3, 2);
}
parameters {
// Declare all parameters as vectors for vectorizing
// Hyper(group)-parameters
vector[6] mu_pr;
vector<lower=0>[6] sigma;
// Subject-level raw parameters (for Matt trick)
vector[N] Arew_pr;
vector[N] Apun_pr;
vector[N] Aneu_pr;
vector[N] Brew_pr;
vector[N] Bpun_pr;
vector[N] Bneu_pr;
}
transformed parameters {
// Transform subject-level raw parameters
vector<lower=0, upper=1>[N] Arew;
vector<lower=0, upper=1>[N] Apun;
vector<lower=0, upper=1>[N] Aneu;
vector<lower=0>[N] Brew;
vector<lower=0>[N] Bpun;
vector<lower=0>[N] Bneu;
for (i in 1:N) {
Arew[i] = Phi_approx(mu_pr[1] + sigma[1] * Arew_pr[i]);
Apun[i] = Phi_approx(mu_pr[2] + sigma[2] * Apun_pr[i]);
Aneu[i] = Phi_approx(mu_pr[3] + sigma[3] * Aneu_pr[i]);
Brew[i] = exp(mu_pr[4] + sigma[4] * Brew_pr[i]);
Bpun[i] = exp(mu_pr[5] + sigma[5] * Bpun_pr[i]);
Bneu[i] = exp(mu_pr[6] + sigma[6] * Bneu_pr[i]);
}
}
model {
// Hyperparameters
mu_pr ~ normal(0, 1);
sigma ~ normal(0, 0.2);
// individual parameters
Arew_pr ~ normal(0, 1.0);
Apun_pr ~ normal(0, 1.0);
Aneu_pr ~ normal(0, 1.0);
Brew_pr ~ normal(0, 1.0);
Bpun_pr ~ normal(0, 1.0);
Bneu_pr ~ normal(0, 1.0);
for (i in 1:N) {
// Define values
array[3, 2] real Q; // expected value table
real A; // trial-level learning rate
real B; // trial-level sensitivity
real PE; // prediction error
for (t in 1:Tsubj[i]) {
// Initialize values
if (new_block[i, t] == 1) {
Q = initV;
}
// Which state/stimuli?
if (cond[i, t] == 1) {
A = Arew[i];
B = Brew[i];
} else if (cond[i, t] == 2) {
A = Apun[i];
B = Bpun[i];
} else {
A = Aneu[i];
B = Bneu[i];
}
// decisions
decisions[i, t] ~ categorical(softmax(to_vector(Q[cond[i, t], ])));
// Prediction error signals
PE = B * outcomes[i, t] - Q[cond[i, t], decisions[i, t]];
// Learning
Q[cond[i, t], decisions[i, t]] += A * PE;
}
}
}
generated quantities {
// For group level parameters
real<lower=0, upper=1> mu_Arew;
real<lower=0, upper=1> mu_Apun;
real<lower=0, upper=1> mu_Aneu;
real<lower=0> mu_Brew;
real<lower=0> mu_Bpun;
real<lower=0> mu_Bneu;
// For log likelihood calculation
real log_lik[N];
// For posterior predictive check
matrix[N, T] y_pred;
// For prediction errors
matrix[N, T] pred_error;
// Set all posterior predictions to 0 (avoids NULL values)
y_pred = rep_matrix(0.0, N, T);
pred_error = rep_matrix(0.0, N, T);
mu_Arew = Phi_approx(mu_pr[1]);
mu_Apun = Phi_approx(mu_pr[2]);
mu_Aneu = Phi_approx(mu_pr[3]);
mu_Brew = exp(mu_pr[4]);
mu_Bpun = exp(mu_pr[5]);
mu_Bneu = exp(mu_pr[6]);
for (i in 1:N) {
// Define values
array[3, 2] real Q; // expected value table
real A; // trial-level learning rate
real B; // trial-level sensitivity
real PE; // prediction error
for (t in 1:Tsubj[i]) {
// Initialize values
if (new_block[i, t] == 1) {
Q = initV;
}
// Which state/stimuli?
if (cond[i, t] == 1) {
A = Arew[i];
B = Brew[i];
} else if (cond[i, t] == 2) {
A = Apun[i];
B = Bpun[i];
} else {
A = Aneu[i];
B = Bneu[i];
}
// compute log likelihood of current trial
log_lik[i] += categorical_lpmf(decisions[i, t] | softmax(to_vector(Q[cond[i, t], ])));
// generate posterior prediction for current trial
y_pred[i, t] = categorical_rng(softmax(to_vector(Q[cond[i, t],])));
// Prediction error signals
PE = B * outcomes[i, t] - Q[cond[i, t], decisions[i, t]];
// Store prediction error
pred_error[i, t] = PE;
// Learning
Q[cond[i, t], decisions[i, t]] += A * PE;
}
}
}
VTA_stan.csv (47.0 KB)
