
%% Lab on Reinforcement Learning
% OS 2016

%% 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.
epsilon = .2;      % learning rate
tStim = 100;       % time of stimulus
timeSteps = 250;   % time steps per trial
numTrials = 1000;  % number of trials

u = zeros(timeSteps,1);  % stimulus
r = zeros(timeSteps,1);  % reward
v = zeros(timeSteps,1);  % expected future reward
u(tStim) = 1;
r(199) = .5;
r(200) = 1;
r(201) = .5;
delta = zeros(timeSteps,numTrials);  % prediction error over trials and time steps
for trial = 1:numTrials
  for t=2:timeSteps
     delta(t-1,trial) = r(t-1) + v(t) - v(t-1);      % at time t, update the previous
     if t>=100
        v(t-1) = v(t-1) + epsilon.*(delta(t-1,trial));
     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);
prior = betapdf(thetas, priorReward+1, priorNone+1);
prior = prior./sum(prior);

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

figure(3);
subplot(3,1,1);
plot(thetas, prior)
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
posterior = likelihood .* prior;
posterior = posterior./sum(posterior);

% This is the same as:
posterior = betapdf(thetas, 1 + priorReward + likeReward, 1+ priorNone + likeNone);
posterior = posterior./sum(posterior);

subplot(3,1,3)
plot(thetas, posterior)
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(posterior);
map_theta = thetas(map_theta);
% Expectation (mean)
mean_theta = dot(thetas , posterior);
hold on;
plot(ml_theta,0.04, '*r')
line([ml_theta,ml_theta],[0,.04],'LineStyle','--','color','r')
plot(map_theta,0.04, '*m')
line([map_theta,map_theta],[0,.04],'LineStyle','--','color','m')
plot(mean_theta,0.04, '*b')
line([mean_theta,mean_theta],[0,.04],'LineStyle','--','color','b')
hold off;

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



