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)