Feng_Simple_Advice_Model_CMG.m

% sim = 1;
% 
% % generative process (from task structure)
% 
% task.num_blocks = 2; % number of trials
% task.num_trials = 30; % number of trials
%
% % whether loss is -40 or -80
% task.block_type(1) = "SL"; %small loss
% task.block_type(2) = "LL"; %large loss
% 
% %true context probabilities by block
% task.true_p_right(1,1:15) = .9;
% task.true_p_right(1,16:30) = .5;
% 
% task.true_p_right(2,1:15) = .5;
% task.true_p_right(2,16:30) = .9;
% 
% %true hint accuracies by block
% task.true_p_a(1,1:15) = .6;
% task.true_p_a(1,16:30) = .8;
% 
% task.true_p_a(2,1:15) = .6;
% task.true_p_a(2,16:30) = .8;
% 
% % parameters
% 
% %initial context and trust priors
% params.p_right = .5;
% params.p_right = .9;
% 
% %learning and forgetting rates
%
% comment out to allow a/d specific types below
% 
% %params.omega = .2; 
% params.eta = .5;
% 
% % comment out to further allow win/loss specific types below
% 
% %params.omega_a = .2; 
% params.omega_d = .4;
% % params.eta_a = .5;
% % params.eta_d = .5;
% 
% % win/loss specific params (only used based on commenting out above)
% 
% params.omega_a_win = .1; params.eta_a_win = .4;
% params.omega_a_loss = .3; params.eta_a_loss = .6;
% 
% params.omega_d_win = .4; params.eta_d_win = .7;
% params.omega_d_loss = .5; params.eta_d_loss = .8;
% 
% %explore-exploit weights
% params.reward_value = 1;
% params.l_loss_value = 1;
% params.state_exploration = 1;
% params.parameter_exploration = 1;
% params.inv_temp = 4;
%
% %data
% observations.hints = [0 2 1;
%                       1 0 1]; %observations.hints(block,trial) 
% observations.rewards = [1 1 2;
%                         2 1 1];%observations.rewards(block,trial)
% choices(:,:,1) = [3 0;
%                   1 3;
%                   1 2];
% choices(:,:,2) = [1 2;
%                   3 0;
%                   1 2];
% if sim == 0
%     task.num_blocks = size(observations.rewards,1); % number of blocks
%     task.num_trials = size(observations.rewards,2); % number of trials
% end

function [results] = Feng_Simple_Advice_Model_CMG(task, params)

    % observations.hints = 0 is no hint, 1 is left hint, 2 is right hint
    % observations.rewards(trial) 1 is win, 2 is loss
    % choices : 1 is advisor, 2 is left, 3 is right
    task.num_trials = 30;
    task.num_blocks = 12;
    observations.hints = nan(1,task.num_trials);
    observations.rewards = nan(1,task.num_trials);
    choices = nan(task.num_trials,2);
    
    
    params.p_right = .5;
    single_omega = 0;
    single_eta = 0;
    field = fieldnames(params);

    % if single learning rate or forget_rate
    for i = 1:length(field)
        if strcmp(field{i},'omega')
            params.omega_d_win = params.omega;
            params.omega_d_loss = params.omega;
            params.omega_a_win = params.omega;
            params.omega_a_loss = params.omega;
            params.omega_d = params.omega;
            params.omega_a = params.omega;
            single_omega = 1;
        elseif strcmp(field{i},'eta')
            params.eta_d_win = params.eta;
            params.eta_d_loss = params.eta;
            params.eta_a_win = params.eta;
            params.eta_a_loss = params.eta;
            params.eta_d = params.eta;
            params.eta_a = params.eta;
            single_eta = 1;
        end
    end
    
    % win loss same learnin rate or forget rate
    for i = 1:length(field)
        if strcmp(field{i},'omega_d') & single_omega ~= 1
            params.omega_d_win = params.omega_d;
            params.omega_d_loss = params.omega_d;
        elseif strcmp(field{i},'eta_d') & single_eta ~= 1
            params.eta_d_win = params.eta_d;
            params.eta_d_loss = params.eta_d;
        elseif strcmp(field{i},'omega_a') & single_omega ~= 1
            params.omega_a_win = params.omega_a;
            params.omega_a_loss = params.omega_a;
        elseif strcmp(field{i},'eta_a') & single_eta ~= 1
            params.eta_a_win = params.eta_a;
            params.eta_a_loss = params.eta_a;
        end
    end

    p_win = 1;
    a_floor = 1;
    context_floor = 1;

    
    
    hint_outcomes(1,1:task.num_trials) = 0;

    for trial = 1:task.num_trials
        hint_outcome_vector(:,trial) = [0 0]';
        dir_context(:,:,trial) = zeros(2,1);
        p_context(:,:,trial) = zeros(2,1);
        pp_context(:,:,trial) = zeros(2,1);
        true_context(:,:,trial) = zeros(2,1);
        a(:,:,trial) = zeros(2,2);
        true_A(:,:,trial) = zeros(2,2);
        A{2}(:,:,trial) = zeros(2,2);
        action_probs(:,:,trial) = zeros(3,2);
        Q(:,:,trial) = zeros(3,1);
        epistemic_value(:,:,trial) = zeros(3,1);
        pragmatic_value(:,:,trial) = zeros(3,1);
        novelty_value(:,:,trial) = zeros(3,1);
        p_o_win(:,:,trial) = zeros(2,3);
        % sample if left or right win
        true_context_vector(trial) = find(rand < cumsum([1-task.true_p_right(1,trial) task.true_p_right(1,trial)]'),1)-1;
    end
 

    d_0 = [1-params.p_right params.p_right]'*context_floor;
    a_0 =  [params.p_a 1-params.p_a;         % "try left"
            1-params.p_a params.p_a]*a_floor;% "try right"   

    actions = zeros(task.num_trials,2);
    
    % reward value distribution
    if task.block_type(1)== "LL"
        R =  spm_softmax([params.reward_value+eps -params.l_loss_value-eps]');
    else
        R =  spm_softmax([params.reward_value+eps -eps]');
    end


    for trial = 1:task.num_trials
        tp = 1; % initial timepoint

        % priors
        if trial == 1
            dir_context(:,:,trial) = d_0;
        end 

        p_context(:,:,trial) = spm_norm(dir_context(:,:,trial)); % choice left or right?
        true_context(:,:,trial) = [1-true_context_vector(trial) true_context_vector(trial)]';
            
        % likelihood mapping

        %hint (concentration parameters)
        if trial == 1
            a(:,:,trial) = a_0;
        end

        % p(hint|context)
        A{1}(:,:,trial) = spm_norm(a(:,:,trial));
        true_A(:,:,trial) =  [task.true_p_a(1,trial)   1-task.true_p_a(1,trial); % "try left"
                                 1-task.true_p_a(1,trial) task.true_p_a(1,trial)];  % "try right"
        
        %left better
        A{2}(:,:,1) = [p_win   1-p_win; % win
                       1-p_win p_win];  % lose
        
        %right better
        A{2}(:,:,2) = [1-p_win p_win;   % win
                       p_win   1-p_win];% lose
   
    
    % compute information gain
    % a_sums{1}(:,:,trial) = [sum(a(:,1,trial)) sum(a(:,2,trial));
    %                         sum(a(:,1,trial)) sum(a(:,2,trial))];
    % 
    % info_gain = .5*((a(:,:,trial).^-1)-(a_sums{1}(:,:,trial).^-1));

% compute action values (negative EFE)
        for option = 1:3
            if option == 1
                p_o_hint(:,trial) = A{1}(:,:,1)*p_context(:,:,trial);
                true_p_o_hint(:,trial) = true_A(:,:,1)*true_context(:,:,trial);
                % novelty_value(option,tp,trial) = .5*dot(A{1}(:,1,trial),info_gain(:,1)) + .5*dot(A{1}(:,2,trial),info_gain(:,1));
                % novelty_value(option,tp,trial) = (sum(sum(a(:,:,trial))))^-1;
                a_sums{1}(:,:,trial) = [sum(a(:,1,trial)) sum(a(:,2,trial)); sum(a(:,1,trial)) sum(a(:,2,trial))];
                info_gain = (a(:,:,trial).^-1) - (a_sums{1}(:,:,trial).^-1);
                %marginalize over context state factor (i.e. left better or
                %right better)
                novelty_for_each_observation = info_gain(:,1)*p_context(1,:,trial) + info_gain(:,2)*p_context(2,:,trial);
                novelty_value(option,tp,trial) = sum(novelty_for_each_observation);
                epistemic_value(option,tp,trial) = G_epistemic_value(A{1}(:,:,trial),p_context(:,:,trial));
                pragmatic_value(option,tp,trial) = 0;
            elseif option == 2 
                p_o_win(:,option,trial) = A{2}(:,:,1)*p_context(:,:,trial);
                true_p_o_win(:,option,trial) = A{2}(:,:,1)*true_context(:,:,trial);
                novelty_value(option,tp,trial) = 0;
                epistemic_value(option,tp,trial) = 0;
                pragmatic_value(option,tp,trial) = dot(p_o_win(:,option,trial),R);
            elseif option == 3 
                p_o_win(:,option,trial) = A{2}(:,:,2)*p_context(:,:,trial);
                true_p_o_win(:,option,trial) = A{2}(:,:,2)*true_context(:,:,trial);
                novelty_value(option,tp,trial) = 0;
                epistemic_value(option,tp,trial) = 0;
                pragmatic_value(option,tp,trial) = dot(p_o_win(:,option,trial),R);
            end
            Q(option, tp,trial) = params.state_exploration*epistemic_value(option,tp,trial) + pragmatic_value(option,tp,trial) + params.parameter_exploration*novelty_value(option,tp,trial);
        end
    
            
        % compute action probabilities
        action_probs(:,tp,trial) = spm_softmax(params.inv_temp*Q(:,tp,trial))';

        % select actions
        % note that 1 corresponds to choosing advisor, 2 corresponds to
        % choosing left bandit, 3 corresponds to choosing right bandit.
        
        actions(trial,tp) = find(rand < cumsum(action_probs(:,tp,trial)'),1);
   
        
        % if first action was choosing a bandit, only update context state
        % vector
        if actions(trial,1) ~= 1
   
         
            hint_outcomes(trial) = 0;
        
            % hint_outcome_vector(:,trial) = [0 0]';
    
            % get reward outcome
         
            reward_outcomes(trial) = find(rand < cumsum(true_p_o_win(:,actions(trial,tp),trial)'),1);
        
            if actions(trial,tp) == 3 && reward_outcomes(trial) == 1
                ppp_context(:,trial) = [0 1]';
            elseif actions(trial,tp) == 2 && reward_outcomes(trial) == 2
                ppp_context(:,trial) = [0 1]';
            elseif actions(trial,tp) == 2 && reward_outcomes(trial) == 1
                ppp_context(:,trial) = [1 0]';
            elseif actions(trial,tp) == 3 && reward_outcomes(trial) == 2
                ppp_context(:,trial) = [1 0]';
            end

        % if first action was choosing advisor, update likelihood matrices
        % before picking pandit
        elseif actions(trial,1) == 1
            tp = 2; % second time point

            % get hint outcome
       
            hint_outcomes(trial) = find(rand < cumsum(true_p_o_hint(:,trial)'),1);
           

            hint_outcome_vector(hint_outcomes(trial),trial) = 1;

            % state belief update
        
            pp_context(:,:,trial) = p_context(:,:,trial).*A{1}(:,hint_outcomes(trial),trial)...
                            /sum(p_context(:,:,trial).*A{1}(:,hint_outcomes(trial),trial),1);


            for option = 2:3
                Q(1, tp,trial) = eps;
                if option == 2 
                    p_o_win(:,option,trial) = A{2}(:,:,1)*pp_context(:,:,trial);
                    novelty_value(option,tp,trial) = 0;
                    epistemic_value(option,tp,trial) = 0;
                    pragmatic_value(option,tp,trial) = dot(p_o_win(:,option,trial),R);
                elseif option == 3 
                    p_o_win(:,option,trial) = A{2}(:,:,2)*pp_context(:,:,trial);
                    novelty_value(option,tp,trial) = 0;
                    epistemic_value(option,tp,trial) = 0;
                    pragmatic_value(option,tp,trial) = dot(p_o_win(:,option,trial),R);
                end
                Q(option, tp,trial) = params.state_exploration*epistemic_value(option,tp,trial) + pragmatic_value(option,tp,trial) + params.parameter_exploration*novelty_value(option,tp,trial);
            end
    
            % compute action probabilities
            action_probs(:,tp,trial) = [0; spm_softmax(params.inv_temp*Q(2:3,tp,trial))]';

            % select actions
        
            actions(trial,tp) = find(rand < cumsum(action_probs(:,tp,trial)'),1);
            
    
            % get reward outcome
            
            reward_outcomes(trial) = find(rand < cumsum(true_p_o_win(:,actions(trial,tp),trial)'),1);
               
    
            if actions(trial,tp) == 3 && reward_outcomes(trial) == 1
                ppp_context(:,trial) = [0 1]';
            elseif actions(trial,tp) == 2 && reward_outcomes(trial) == 2
                ppp_context(:,trial) = [0 1]';
            elseif actions(trial,tp) == 2 && reward_outcomes(trial) == 1
                ppp_context(:,trial) = [1 0]';
            elseif actions(trial,tp) == 3 && reward_outcomes(trial) == 2
                ppp_context(:,trial) = [1 0]';
            end  
    
        end   

        if reward_outcomes(trial) == 1
            if actions(trial,1) == 1 
                % forgetting part
                a(:,:,trial+1) = (a(:,:,trial) - a_0)*(1-params.omega_a_win) + a_0;
                % learning part
                a(:,:,trial+1) = a(:,:,trial+1) + params.eta_a_win*(ppp_context(:,trial)*hint_outcome_vector(:,trial)')';
            else
                a(:,:,trial+1) = a(:,:,trial);
            end
    
                % forgetting part
                dir_context(:,:,trial+1) = (dir_context(:,:,trial) - d_0)*(1-params.omega_d_win) + d_0;
                % learning part
                dir_context(:,:,trial+1) = dir_context(:,:,trial+1) + params.eta_d_win*ppp_context(:,trial);
        elseif reward_outcomes(trial) == 2
            if actions(trial,1) == 1 
                % forgetting part
                a(:,:,trial+1) = (a(:,:,trial) - a_0)*(1-params.omega_a_loss) + a_0;
                % learning part
                a(:,:,trial+1) = a(:,:,trial+1) + params.eta_a_loss*(ppp_context(:,trial)*hint_outcome_vector(:,trial)')';
            else
                a(:,:,trial+1) = a(:,:,trial);
            end
    
            % forgetting part
            dir_context(:,:,trial+1) = (dir_context(:,:,trial) - d_0)*(1-params.omega_d_loss) + d_0;
            % learning part
            dir_context(:,:,trial+1) = dir_context(:,:,trial+1) + params.eta_d_loss*ppp_context(:,trial);
        end
    end

results.observations.hints(1,:) = hint_outcomes;
results.observations.rewards(1,:) = reward_outcomes;
results.choices(:,:,1) = actions(:,:);
results.R = R;


results.input.task = task;
results.input.params = params;
results.input.observations = observations;
results.input.choices = choices;
results.input.sim = 1;



results.blockwise.action_probs = action_probs;
results.blockwise.actions = actions;
results.blockwise.true_context = true_context;
results.blockwise.hint_outcomes = hint_outcomes;
results.blockwise.hint_outcome_vector = hint_outcome_vector;
results.blockwise.reward_outcomes = reward_outcomes;
results.blockwise.action_values_Q = Q;
results.blockwise.epistemic_value = epistemic_value;
results.blockwise.pragmatic_value = pragmatic_value;
results.blockwise.state_priors_d = dir_context;
results.blockwise.norm_priors_d = p_context;
results.blockwise.norm_posteriors_d_t2 = pp_context;
results.blockwise.norm_posteriors_d_final = ppp_context;
results.blockwise.trust_priors_a = a;
results.blockwise.norm_trust_priors_a = A{1};




end

 % actions
 % reward_outcome
 % true_context
 % hint_outcome

 %action_probs
 


% epistemic value term (Bayesian surprise) in expected free energy 
function G = G_epistemic_value(A,s)
    
% A   - likelihood array (probability of outcomes given causes)
% s   - probability density of causes


% probability distribution over the hidden causes: i.e., Q(s)

qx = spm_cross(s); % this is the outer product of the posterior over states
                   % calculated with respect to itself

% accumulate expectation of entropy: i.e., E[lnP(o|s)]
G     = 0;
qo    = 0;
for i = find(qx > exp(-16))'
    % probability over outcomes for this combination of causes
    po   = 1;
    for g = 1:numel(A)
        %po = spm_cross(po,A{g}(:,i));
        po = spm_cross(po,A(:,i));
    end
    po = po(:);
    qo = qo + qx(i)*po;
    G  = G  + qx(i)*po'*nat_log(po);
end

% subtract entropy of expectations: i.e., E[lnQ(o)]
G  = G - qo'*nat_log(qo);
    
end 

function y = nat_log(x)
y = log(x+exp(-16));
end 

function A  = spm_norm(A)
% normalisation of a probability transition matrix (columns)
%--------------------------------------------------------------------------
A           = bsxfun(@rdivide,A,sum(A,1));
A(isnan(A)) = 1/size(A,1);
end