
%% Lab on Reinforcement Learning
% OS 2017

%% Rescora Wagner model  
% After figure 9.1 in Dayan and Abbott textbook
% Pairing stimulus u with reward r
epsilon = .05;   % epsilon
w = 0;
wVec = w;
v = w;
for i=1:200
  u = 1;
  if i<=100   % stimulus paired with reward
     r = 1;
  else
     r = 0;   % stimulus not paired with reward
  end
  w = w + epsilon .* (r-v) .* u;
  v = w;
  wVec = [wVec w];
end

figure(1);
h=plot(0:200, wVec, 'ko');
set(h, 'markersize',6)
set(h,'markerfacecolor','k')
set(gca, 'fontsize', 18);
xlabel('trial number');
ylabel('w');

%% To Do:
% Write a version of the above that for each
% stimulus has 50 percent chance of being
% paired with a reward. Call the resulting
% weight vector wVec2 and plot it alongside
% wVec above (you can use hold on for the figure).  

% To Do:
% Try different learning rates (epsilon values)


%% Temporal Difference (TD) learning. 
% After figure 9.2 in Dayan and Abbott.
% Using w to make predictions 
epsilon = .2;      % learning rate
tStim = 100;       % time of stimulus
tDelayedReward = 100;
timeSteps = 250;   % time steps per trial
numTrials = 1000;  % number of trials


u = zeros(timeSteps,1);  % stimulus
w = zeros(timeSteps,1);
r = zeros(timeSteps,1);  % reward
v = zeros(timeSteps,1);  % expected future reward

u(tStim) = 1;
r(tStim + tDelayedReward + [-1, 0, 1]) = [.5, 1, 0.5];
delta = zeros(timeSteps,numTrials);  % prediction error over trials and time steps
for trial = 1:numTrials
    v(1) = 0;
    for t = 1 : timeSteps - 1
        % compute prediction based on current w
        v(t + 1) = 0;
        for Tau = 0 : timeSteps -1
            if t + 1 - Tau > 1  % we assume that u(t)  for t < 1 is zero
                v(t + 1) = v(t + 1) + w(Tau + 1)*u(t + 1  - Tau);
                % recall w(Tau + 1) to provide the right indexing
            end
        end
        % compute delta based on TD
        delta(t, trial) = r(t) + v(t + 1) - v(t);      % at time t, update the previous
        % update w based on delta
        for Tau = 0 : timeSteps -1
            if t - Tau > 1  % we assume that u(t)  for t < 1 is zero
                w(Tau + 1) = w(Tau + 1) + epsilon*delta(t, trial)*u(t - Tau);
            end
        end
    end

end
figure(2);
rotate3d on
h=surf(1:numTrials,1:timeSteps,delta)
axis('ij');
ylabel('t'); xlabel('trials'); zlabel('delta');

% To Do: How does delta change over trials? Take slices at
% early trials and late trials, and plot the vectors.
% What are the changes over time signifying in terms of what is learned?
% What happens when you change r -- for instance to have uniform reward values
% between 181:220, or to just have one value at time 150?
% What happens when you change the time of the stimulus?
% How does the expected future reward (v) change over trials?
% You can define a matrix vMat = zeros(timeSteps,numTrials)
% and update it every trial, and then look at slices through 
% the matrix.
% What are the changes in v over time signifying in terms of what is learned?

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%% Beta distribution and Bayes 
% Simple example after NYU Math Tools for Neuroscience tutorial of flipping coins
% Here we consider reward versus no reward.

priorReward = 2;
priorNone = 2;
range = 101;
thetas = linspace(0,1,range);
priorA = betapdf(thetas, priorReward+1, priorNone+1);

% The Beta distribution is:
% thetas.^priorReward .* (1-thetas).^priorNone
% (and normalized to sum to 1).
% Check: This is the same as:
priorB = (thetas.^priorReward .* (1-thetas).^priorNone)/beta(priorReward+1, priorNone+1);
%prior = prior./sum(prior);

figure(3);
subplot(3,1,1);
plot(thetas(:), priorA(:))
title('prior ')
% If there is new data (likelihood) or observation from
% next state
likeReward = 6;
likeNone = 2;
likelihood = betapdf(thetas, likeReward+1, likeNone+1);
%likelihood = likelihood./sum(likelihood);

subplot(3,1,2);
plot(thetas, likelihood)
title('likelihood')

%%
% Bayes
posteriorA = likelihood .* priorB;
%posterior = posterior./sum(posterior);

% This is proportional to (check):
posteriorB = betapdf(thetas, 1 + priorReward + likeReward, 1+ priorNone + likeNone);
%posterior = posterior./sum(posterior);

subplot(3,1,3)
plot(thetas, posteriorB)
xlabel('theta')
title('posterior')

%% We can look at some of these summary statistics:
[tmp, ml_theta] = max(likelihood);
ml_theta = thetas(ml_theta);
% maximum a posteriori (MAP)
[tmp, map_theta] = max(posteriorB);
map_theta = thetas(map_theta);
% Expectation (mean)
mean_theta = dot(thetas , posteriorB/sum(posteriorB));
hold on;
stem_height = 4;
plot(ml_theta,stem_height, '*r')
line([ml_theta,ml_theta],[0,stem_height],'LineStyle','--','color','r')
plot(map_theta,stem_height, '*m')
line([map_theta,map_theta],[0,stem_height],'LineStyle','--','color','m')
plot(mean_theta,stem_height, '*b')
line([mean_theta,mean_theta],[0,stem_height],'LineStyle','--','color','b')
hold off;

%% Try changing the prior and likelihood. How does the posterior
%  relate to the prior and likelihood?



