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