
%% Lab on Reinforcement Learning
% OS 2020 updated

%% 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.  

w = 0;
wVec2 = w;
v = w;
for i=1:200
  u = 1;
  therand = rand;
  if therand >= .5
     r = 1;
  else
     r = 0;
  end
  w = w + epsilon .* (r-v) .* u;
  v = w;
  wVec2 = [wVec2 w];
end

figure(1);
hold on;
h = plot(0:200, wVec2, 'sb');
set(h, 'markersize',6)
hold off;


% 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;
%if 0
%r(181:220)=rand(size(r(181:220)));
%end
%if 0
%    r(191:210)= 1/10;
%end
delta = zeros(timeSteps,numTrials);  % prediction error over trials and time steps
vMat = zeros(timeSteps,numTrials);   % expected future reward over trials and time steps
vDiff = zeros(timeSteps,numTrials);
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)); % v expected future reward is updating every trial
     end
     vMat(t-1,trial) = v(t-1);
     vDiff(t-1,trial) = v(t)-v(t-1);
  end
end
figure(2);
rotate3d on
h=surf(1:numTrials,1:timeSteps,delta)
axis('ij');
ylabel('t'); xlabel('trials'); zlabel('delta');
set(h,'facealpha',0.2)
set(h,'edgealpha',0.2)

%% 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?

%% before learning:
plot(1:timeSteps,delta(:,1))
%% after learning
plot(1:timeSteps,delta(:,1000))

%r = zeros(timeSteps,1);  % reward
%r(191:210)= 1/10;

%% before
plot(1:timeSteps,vMat(:,1))
%% after
plot(1:timeSteps,vMat(:,1000))

%% before
plot(1:timeSteps,vDiff(:,1))
%% after
plot(1:timeSteps,vDiff(:,1000))


