initial commit

Author: sjgershm

Collins18_data.csv

@@ -0,0 +1,8 @@
+reward-complexity
+====
+
+Code for reproducing the analyses reported in the paper "Origin of perseveration in the trade-off between reward and complexity".
+
+Some of the code requires the mfit package: https://github.com/sjgershm/mfit
+
+Questions? Contact Sam Gershman ([email protected]).

README.md

@@ -0,0 +1,67 @@
+function results = analyze_collins(data)
+    
+    % Analyze Collins (2018) data.
+    
+    if nargin < 1
+        data = load_data('collins18');
+    end
+    
+    beta = linspace(0.1,15,30);
+    
+    for s = 1:length(data)
+        B = unique(data(s).learningblock);
+        cond = zeros(length(B),1);
+        R_data =zeros(length(B),1);
+        V_data =zeros(length(B),1);
+        for b = 1:length(B)
+            ix = data(s).learningblock==B(b) & data(s).phase==0;
+            stim = data(s).stim(ix)';
+            c = data(s).corchoice(ix);
+            choice = data(s).choice(ix)';
+            lowerx = 1; upperx = max(stim); lowery = 1; uppery = max(stim);
+            descriptor = [lowerx,upperx,upperx-lowerx;lowery,uppery,uppery-lowery];
+            R_data(b) = information(stim,choice,descriptor);
+            V_data(b) = mean(data(s).cor(ix));
+            
+            S = unique(stim);
+            Q = zeros(length(S),3);
+            Ps = zeros(1,length(S));
+            for i = 1:length(S)
+                ii = stim==S(i);
+                Ps(i) = mean(ii);
+                a = c(ii); a = a(1);
+                Q(i,a) = 1;
+            end
+            
+            [R(b,:),V(b,:)] = blahut_arimoto(Ps,Q,beta);
+            
+            if length(S)==3
+                cond(b) = 1;
+            else
+                cond(b) = 2;
+            end
+            
+            ix = data(s).learningblock==B(b) & data(s).phase==1;
+            stim = data(s).stim(ix)';
+            choice = data(s).choice(ix)';
+            try
+                R_test(b) = information(stim,choice);
+                V_test(b) = mean(data(s).cor(ix));
+            catch
+                R_test(b) = nan;
+                V_test(b) = nan;
+            end
+        end
+        
+        for c = 1:2
+            results.R(s,:,c) = nanmean(R(cond==c,:));
+            results.V(s,:,c) = nanmean(V(cond==c,:));
+            results.R_data(s,c) = nanmean(R_data(cond==c));
+            results.V_data(s,c) = nanmean(V_data(cond==c));
+            results.V_test(s,c) = nanmean(V_test(cond==c));
+            results.R_test(s,c) = nanmean(R_test(cond==c));
+        end
+        
+        clear R V
+        
+    end

analyze_collins.m

@@ -0,0 +1,34 @@
+function results = analyze_steyvers(data)
+    
+    % Analyze Steyvers et al. (2019) data.
+    
+    if nargin < 1
+        data = load_data('steyvers19');
+    end
+    
+    beta = linspace(1.5,5,30);
+    
+    lowerx = 1; upperx = 32; lowery = 1; uppery = 4;
+    descriptor = [lowerx,upperx,upperx-lowerx;lowery,uppery,uppery-lowery];
+    
+    for s = 1:length(data)
+        results.R_data(s,1) = information(data(s).state',data(s).action',descriptor);
+        results.V_data(s,1) = mean(data(s).reward);
+        for state = 1:32
+            Ps(s,state) = mean(data(s).state==state);
+        end
+    end
+    
+    Ps = mean(Ps);
+    [X, Y, Z] = ind2sub([4 4 2],1:32);
+    Q = zeros(32,4);
+    for i = 1:32
+        if Z(i)==1
+            a = X(i);
+        else
+            a = Y(i);
+        end
+        Q(i,a) = 1;
+    end
+    [results.R,results.V] = blahut_arimoto(Ps,Q,beta);
+    results.Q = Q; results.Ps = Ps;

analyze_steyvers.m

@@ -0,0 +1,39 @@
+function [R,V,Pa] = blahut_arimoto(Ps,Q,b)
+    
+    % Blahut-Arimoto algorithm applied to the reward-complexity trade-off.
+    %
+    % USAGE: [R,V,Pa] = blahut_arimoto(Ps,Q,[b])
+    %
+    % INPUTS:
+    %   Ps - [1 x S] state probabilities, where S is the number of states
+    %   Q - [S x A] expected reward, where A is the number of actions
+    %   b (optional) - vector of trade-off parameters. Default: linspace(0.1,15,30)
+    %
+    % OUTPUTS:
+    %   R - [K x 1] channel capacity values, where K is the length of b
+    %   V - [K x 1] average reward values
+    %   Pa - [K x A] marginal action policy
+    %
+    % Sam Gershman, Jan 2020
+    
+    A = size(Q,2);
+    nIter = 50;
+    if nargin < 3; b = linspace(0.1,15,30); end
+    R = zeros(length(b),1); V = zeros(length(b),1); Pa = zeros(length(b),A);
+    q = ones(1,A)./A;
+    
+    for j = 1:length(b)
+        F = b(j).*Q;
+        v0 = mean(Q(:));
+        for i = 1:nIter
+            logP = log(q) + F;
+            Z = logsumexp(logP,2);
+            Psa = exp(logP - Z);
+            q = Ps*Psa;
+            v = sum(Ps*(Psa.*Q));
+            if abs(v-v0) < 0.001; break; else v0 = v; end
+        end
+        Pa(j,:) = q;
+        V(j) = v;
+        R(j) = b(j)*v - Ps*Z;
+    end

blahut_arimoto.m

@@ -0,0 +1,59 @@
+function [results, bms_results] = fit_models_collins(data)
+    
+    % Fit Collins (2018) data. Requires mfit package.
+    
+    for m = 1:2
+        disp(['... fitting model ',num2str(m)]);
+        
+        switch m
+            
+            case 1
+                
+                param(1) = struct('name','b1','logpdf',@(x) 0);
+                param(2) = struct('name','b2','logpdf',@(x) 0);
+                param(3) = struct('name','sticky','logpdf',@(x) 0);
+                fun = @lik_collins;
+                
+            case 2
+                
+                param(1) = struct('name','b1','logpdf',@(x) 0);
+                param(2) = struct('name','b2','logpdf',@(x) 0);
+                fun = @lik_collins;
+                
+        end
+        
+        results(m) = mfit_optimize(fun,param,data);
+        clear param
+    end
+    
+    % Bayesian model selection
+    if nargout > 1
+        bms_results = mfit_bms(results,1);
+    end
+    
+end
+
+function lik = lik_collins(x,data)
+    
+    B = x(1:2);
+    if length(x) > 2
+        sticky = x(3);
+    else
+        sticky = 1;
+    end
+    
+    lik = 0;
+    
+    for t = 1:size(data.Q,1)
+        a = data.choice(t);
+        if a > 0
+            if data.ns(t)==3
+                b = B(1);
+            else
+                b = B(2);
+            end
+            d = b*data.Q(t,:) + sticky*data.logPa(t,:);
+            lik = lik + d(a) - logsumexp(d,2);
+        end
+    end
+end

fit_models_collins.m

@@ -0,0 +1,51 @@
+function [results, bms_results] = fit_models_steyvers(data,models)
+    
+    % Fit models to Steyvers data. Requires mfit package.
+    
+    if nargin < 2; models = 1:2; end
+    
+    for m = models
+        disp(['... fitting model ',num2str(m)]);
+        
+        switch m
+            
+            case 1
+                
+                param(1) = struct('name','b','logpdf',@(x) 0);
+                param(2) = struct('name','sticky','logpdf',@(x) 0);
+                fun = @lik_steyvers;
+                
+            case 2
+                
+                param(1) = struct('name','b','logpdf',@(x) 0);
+                fun = @lik_steyvers;
+                
+        end
+        
+        results(m) = mfit_optimize(fun,param,data);
+        clear param
+    end
+    
+    % Bayesian model selection
+    if nargout > 1
+        bms_results = mfit_bms(results,1);
+    end
+    
+end
+
+function lik = lik_steyvers(x,data)
+    
+    b = x(1);
+    if length(x) > 1
+        sticky = x(2);
+    else
+        sticky = 1;
+    end
+    
+    d = b*data.Q + sticky*data.logPa;
+    lik = -sum(logsumexp(d,2));
+    
+    for t = 1:data.N
+        lik = lik + d(t,data.action(t));
+    end
+end

fit_models_steyvers.m

@@ -0,0 +1,70 @@
+function data = load_data(dataset)
+    
+    % Load data sets.
+    %
+    % USAGE: data = load_data(dataset)
+    %
+    % INPUTS:
+    %   dataset - 'collins18' or 'steyvers19'
+    
+    switch dataset
+        
+        case 'collins18'
+            
+            T = {'ID' 'learningblock' 'trial' 'ns' 'stim' 'iter' 'corchoice' 'choice' 'cor' 'rt' 'pcor' 'delay' 'phase'};
+            x = csvread('Collins18_data.csv',1);
+            S = unique(x(:,1));
+            for s = 1:length(S)
+                ix = x(:,1)==S(s);
+                for j = 1:length(T)
+                    data(s).(T{j}) = x(ix,j);
+                end
+                
+                trials = find(data(s).phase==0);
+                data(s).N = length(trials);
+                data(s).C = 3;
+                for t = trials'
+                    if t==1 || data(s).learningblock(t)~=data(s).learningblock(t-1)
+                        Q = zeros(data(s).ns(t),3) + 0.5;
+                        n = zeros(data(s).ns(t),3) + 0.7;
+                        Pa = ones(1,3)/3;
+                    end
+                    data(s).Q(t,:) = Q(data(s).stim(t),:);
+                    data(s).logPa(t,:) = safelog(Pa);
+                    if data(s).choice(t)>0
+                        n(data(s).stim(t),data(s).choice(t)) = n(data(s).stim(t),data(s).choice(t)) + 1;
+                        lr = 1./n(data(s).stim(t),data(s).choice(t));
+                        Q(data(s).stim(t),data(s).choice(t)) = Q(data(s).stim(t),data(s).choice(t)) + lr*(data(s).cor(t)-Q(data(s).stim(t),data(s).choice(t)));
+                        Pa = n; Pa = Pa./sum(Pa(:)); Pa = sum(Pa);
+                    end
+                end
+            end
+            
+        case 'steyvers19'
+            
+            load steyvers19_data.mat
+            
+            [X, Y, Z] = ind2sub([4 4 2],1:32);
+            Q = zeros(32,4);
+            for i = 1:32
+                if Z(i)==1
+                    a = X(i);
+                else
+                    a = Y(i);
+                end
+                Q(i,a) = 1;
+            end
+            
+            for s = 1:length(data)
+                data(s).N = length(data(s).state);
+                data(s).C = 4;
+                A = zeros(data(s).N,4);
+                for t = 1:length(data(s).state)
+                    A(t,data(s).action(t)) = 1;
+                    data(s).Q(t,:) = Q(data(s).state(t),:);
+                end
+                for i = 1:size(A,2); A(:,i) = eps + smooth(A(:,i)); end
+                data(s).logPa = log(A./sum(A,2));
+            end
+            
+    end