Home

Here you find the coding material for the workshop on reinforcement learning and sequential sampling modeling, presented at CMAH 2021 (9 and 10 Sept 2021).

Find the source code here.

Program:

• Sequential sampling
  1. Diffusion model
  2. Race diffusion model
• Combining RL with sequential sampling
  1. RL + Diffusion model
• Examples
  1. Fitting the RLDM on Fontanesi and colleagues (2019) data with hierarchical Bayesian estimation
  2. Fitting the DM on Ratcliff and Rouder (1998) data with MLE (group activity)

Before starting

Check the installation of the following packages:

require(tidyverse)
require(dfoptim)
require(rtdists)
require(rstan)
require(bayesplot)

And download the rwald code.r at this OSF page, that was attached to the Tillman paper on the race diffusion model.

Stan models

Here is a summary of the stan models we used throughout the tutorials:

DM without across-trial variability

data {
    int<lower=1> N;                                                // number of data items
    int<lower=-1,upper=1> accuracy[N];                   // accuracy (-1, 1)
    real<lower=0> rt[N];                                         // rt
    real<lower=0, upper=1> starting_point;             // starting point diffusion model not to estimate
}
parameters {
    real drift;
    real threshold;
    real ndt;
}
transformed parameters {
    real drift_ll[N];                                          // trial-by-trial drift rate for likelihood (incorporates accuracy)
    real drift_t[N];                                           // trial-by-trial drift rate for predictions
    real<lower=0> threshold_t[N];                        // trial-by-trial threshold
    real<lower=0> ndt_t[N];                                // trial-by-trial ndt

    real transf_threshold;
    real transf_ndt;
    transf_threshold = log(1 + exp(threshold));
    transf_ndt = log(1 + exp(ndt));

    for (n in 1:N) {
        drift_t[n] = drift;
        drift_ll[n] = drift_t[n]*accuracy[n];
        threshold_t[n] = transf_threshold;
        ndt_t[n] = transf_ndt;
    }
}
model {
    drift ~ normal(1, 3);
    threshold ~ normal(0, 3);
    ndt ~ normal(-1, 1);
    rt ~ wiener(threshold_t, ndt_t, starting_point, drift_ll);
}
generated quantities {
    vector[N] log_lik;
    {for (n in 1:N) {
        log_lik[n] = wiener_lpdf(rt[n] | threshold_t[n], ndt_t[n], starting_point, drift_ll[n]);
    }
}
}

DM with across-trial variability in the drift-rate

data {
    int<lower=1> N;                                              // number of data items
    int<lower=-1,upper=1> accuracy[N];                 // accuracy (-1, 1)
    real<lower=0> rt[N];                                       // rt
    real<lower=0, upper=1> starting_point;           // starting point diffusion model not to estimate
}
parameters {
    real drift_trialmu;                     // group mean drift-rate 
    real<lower=0> drift_trialsd;            // group SD drift-rate
    real z_drift_trial[N];                  // single trial drift-rate estimates
    real threshold;                          
    real ndt;
}
transformed parameters {
    real drift_ll[N];                                       // trial-by-trial drift rate for likelihood (incorporates accuracy)
    real drift_t[N];                                        // trial-by-trial drift rate for predictions
    real<lower=0> threshold_t[N];                     // trial-by-trial threshold
    real<lower=0> ndt_t[N];                             // trial-by-trial ndt

    real transf_threshold;
    real transf_ndt;
    real<lower=0> transf_drift_trialsd;

    transf_threshold = log(1 + exp(threshold));
    transf_ndt = log(1 + exp(ndt));
    transf_drift_trialsd = drift_trialsd;

    for (n in 1:N) {
        drift_t[n] = drift_trialmu + z_drift_trial[n]*drift_trialsd;
        drift_ll[n] = drift_t[n]*accuracy[n];
        threshold_t[n] = transf_threshold;
        ndt_t[n] = transf_ndt;
    }
}
model {
    drift_trialmu ~ normal(0, 3);
    threshold ~ normal(0, 5);
    ndt ~ normal(0, 1);

    drift_trialsd ~ normal(0, 3);
    z_drift_trial ~ normal(0, 1);

    rt ~ wiener(threshold_t, ndt_t, starting_point, drift_ll);
}
generated quantities {
    vector[N] log_lik;

    {for (n in 1:N) {
        log_lik[n] = wiener_lpdf(rt[n] | threshold_t[n], ndt_t[n], starting_point, drift_ll[n]);
    }
}
}

Hierarchical DM with separate drift-rates and thresholds per condition

data {
    int<lower=1> N;                                                     // number of data items
    int<lower=1> L;                                                     // number of levels
    int<lower=1> C;                                                                       // number of conditions
    int<lower=1, upper=L> participant[N];                             // level (participant)
    int<lower=1, upper=C> condition[N];                                 // condition
    int<lower=-1,upper=1> accuracy[N];                                  // accuracy (-1, 1)
    real<lower=0> rt[N];                                                            // rt
    real<lower=0, upper=1> starting_point;                          // starting point diffusion model not to estimate
}
parameters {
    real mu_drift[C];
    real mu_threshold[C];
    real mu_ndt;

    real<lower=0> sd_drift[C];
    real<lower=0> sd_threshold[C];
    real<lower=0> sd_ndt;

    vector[C] z_drift[L];
    vector[C] z_threshold[L];
    real z_ndt[L];
}
transformed parameters {
    real drift_ll[N];                                     // trial-by-trial drift rate for likelihood (incorporates accuracy)
    real drift_t[N];                                      // trial-by-trial drift rate for predictions
    real<lower=0> threshold_t[N];                   // trial-by-trial threshold
    real<lower=0> ndt_t[N];                           // trial-by-trial ndt

    vector[C] drift_sbj[L];
    vector[C] threshold_sbj[L];
    real<lower=0> ndt_sbj[L];
    
    for (l in 1:L) {
        for (c in 1:C) {
            drift_sbj[l][c] = mu_drift[c] + z_drift[l][c]*sd_drift[c];
            threshold_sbj[l][c] = log(1 + exp(mu_threshold[c] + z_threshold[l][c]*sd_threshold[c]));
        }
        ndt_sbj[l] = log(1 + exp(mu_ndt + z_ndt[l]*sd_ndt));
    }

    for (n in 1:N) {
        drift_t[n] = drift_sbj[participant[n]][condition[n]];
        drift_ll[n] = drift_t[n]*accuracy[n];
        threshold_t[n] = threshold_sbj[participant[n]][condition[n]];
        ndt_t[n] = ndt_sbj[participant[n]];
    }
}
model {
    mu_drift ~ normal(0, 3);
    mu_threshold ~ normal(0, 5);
    mu_ndt ~ normal(-1, 3);

    sd_drift ~ normal(0, 3);
    sd_threshold ~ normal(0, 5);
    sd_ndt ~ normal(0, 1);

    for (l in 1:L) {
        z_drift[l] ~ normal(0, 1);
        z_threshold[l] ~ normal(0, 1);
        z_ndt[l] ~ normal(0, 1);
    }

    rt ~ wiener(threshold_t, ndt_t, starting_point, drift_ll);
}
generated quantities {
    vector[N] log_lik;

    {for (n in 1:N) {
        log_lik[n] = wiener_lpdf(rt[n] | threshold_t[n], ndt_t[n], starting_point, drift_ll[n]);
    }
}
}

Race diffusion model

functions {
     real race_pdf(real t, real b, real v){
          real pdf;
          pdf = b/sqrt(2 * pi() * pow(t, 3)) * exp(-pow(v*t-b, 2) / (2*t));
          return pdf;
     }

     real race_cdf(real t, real b, real v){
          real cdf;
          cdf = Phi((v*t-b)/sqrt(t)) + exp(2*v*b) * Phi(-(v*t+b)/sqrt(t));
          return cdf;
     }

     real race_lpdf(matrix RT, vector  ndt, vector b, vector drift_cor, vector drift_inc){

          real t;
          vector[rows(RT)] prob;
          real cdf;
          real pdf;
          real out;

          for (i in 1:rows(RT)){
               t = RT[i,1] - ndt[i];
               if(t > 0){
                  if(RT[i,2] == 1){
                    pdf = race_pdf(t, b[i], drift_cor[i]);
                    cdf = 1 - race_cdf(t, b[i], drift_inc[i]);
                  }
                  else{
                    pdf = race_pdf(t, b[i], drift_inc[i]);
                    cdf = 1 - race_cdf(t, b[i], drift_cor[i]);
                  }
                  prob[i] = pdf*cdf;

                if(prob[i] < 1e-10){
                    prob[i] = 1e-10;
                }
               }
               else{
                    prob[i] = 1e-10;
               }
          }
          out = sum(log(prob));
          return out;
     }
}
data {
    int<lower=1> N;                                          // number of data items
    int<lower=1,upper=2> choices[N];                 // choices, coded as 1 and 2
    real<lower=0> rt[N];                                   // rt
}
transformed data {
    matrix [N, 2] RT;
    for (n in 1:N){
        RT[n, 1] = rt[n];
        RT[n, 2] = choices[n];
    }
}
parameters {
    real ndt;
    real threshold;
    real drift1;
    real drift2;
}
transformed parameters {
    vector<lower=0> [N] drift1_t;                    // trial-by-trial drift rate option 1
    vector<lower=0> [N] drift2_t;                    // trial-by-trial drift rate option 2
    vector<lower=0> [N] threshold_t;                 // trial-by-trial threshold
    vector<lower=0> [N] ndt_t;                         // trial-by-trial ndt

    real<lower=0> transf_drift1;
    real<lower=0> transf_drift2;
    real<lower=0> transf_threshold;
    real<lower=0> transf_ndt;
    transf_drift1 = log(1 + exp(drift1));
    transf_drift2 = log(1 + exp(drift2));
    transf_threshold = log(1 + exp(threshold));
    transf_ndt = log(1 + exp(ndt));

    for (n in 1:N) {
        drift1_t[n] = transf_drift1;
        drift2_t[n] = transf_drift2;
        threshold_t[n] = transf_threshold;
        ndt_t[n] = transf_ndt;
    }
}
model {
    ndt ~  normal(-1, 3);
    threshold ~ normal(0, 3);
    drift1 ~ normal(1, 5);
    drift2 ~ normal(1, 5);
    RT ~ race(ndt_t, threshold_t, drift1_t, drift2_t);
}
generated quantities {
    vector[N] log_lik;
    {
    for (n in 1:N){
        log_lik[n] = race_lpdf(block(RT, n, 1, 1, 2)| segment(ndt_t, n, 1), segment(threshold_t, n, 1), segment(drift1_t, n, 1), segment(drift2_t, n, 1));
    }
    }
}

Reinforcement learning diffusion model

data {
    int<lower=1> N;                                           // number of data items
    int<lower=1> K;                                           // number of options
    vector[N] f_cor;                                          // feedback correct option
    vector[N] f_inc;                                          // feedback incorrect option
    vector[N] trial;                                          // trial number
    int<lower=1, upper=K> cor_option[N];            // correct option
    int<lower=1, upper=K> inc_option[N];            // incorrect option
    int<lower=-1,upper=1> accuracy[N];              // accuracy (-1, 1)
    real<lower=0> rt[N];                                    // rt

    real initial_value;                                     // intial value for learning in the first block

    real<lower=0, upper=1> starting_point;      // starting point diffusion model not to estimate
}
transformed data {
    vector[K] Q0;
    Q0 = rep_vector(initial_value, K);        // initialize Q values
}
parameters {
    real alpha;                               // learning rate parameter
    real drift_scaling;                       // parameter that scales the Q difference that defines the drift-rate
    real threshold;
    real ndt;
}
transformed parameters {
    real drift_ll[N];                                       // trial-by-trial drift rate for likelihood (incorporates accuracy)
    real drift_t[N];                                        // trial-by-trial drift rate for predictions
    real<lower=0> threshold_t[N];                     // trial-by-trial threshold
    real<lower=0> ndt_t[N];                             // trial-by-trial ndt

    vector[K] Q;                                              // Qs values

    real transf_alpha;
    real transf_drift_scaling;
    real transf_threshold;
    real transf_ndt;

    transf_alpha = Phi(alpha);          // learning rate is scaled using cumulative density function of standard normal
    transf_drift_scaling = log(1 + exp(drift_scaling));
    transf_threshold = log(1 + exp(threshold));
    transf_ndt = log(1 + exp(ndt));

    Q = Q0;
    for (n in 1:N) {
        drift_t[n] = transf_drift_scaling*(Q[cor_option[n]] - Q[inc_option[n]]);
        drift_ll[n] = drift_t[n]*accuracy[n];
        threshold_t[n] = transf_threshold;
        ndt_t[n] = transf_ndt;
        
        Q[cor_option[n]] = Q[cor_option[n]] + transf_alpha*(f_cor[n] - Q[cor_option[n]]);
        Q[inc_option[n]] = Q[inc_option[n]] + transf_alpha*(f_inc[n] - Q[inc_option[n]]);
    }
}
model {
    alpha ~ normal(0, .8);
    drift_scaling ~ normal(0, 5);
    threshold ~ normal(0, 5);
    ndt ~ normal(-1, 3);

    rt ~ wiener(threshold_t, ndt_t, starting_point, drift_ll);
}
generated quantities {
    vector[N] log_lik;

    {for (n in 1:N) {
        log_lik[n] = wiener_lpdf(rt[n] | threshold_t[n], ndt_t[n], starting_point, drift_ll[n]);
    }
    }
}