%EKsimpleNeocortexForPosting.m - 
%Simplified neocortex example performing learning of path from start on
%left to goal on right. In control section below, 
%Neocortex=1 sets neocortex model.  
%SARSA=1 sets correct SARSA.  
%Neocortex=0, SARSA=0 gives actor critic using value function.
%- If set TotalTests= large value, withgraphics=0, get multiple tests, with a plot of action
%values for each test and a final plot of rewards received/optimal.
%- If set TotalTests=1, withgraphics=1, can then see the movements of the
%agents during a single process of learning the task.
%EXAMPLE OF STARTING CONTROLS:
%Neocortex=1; SARSA=1; SA=1;  %SA sets it to form representations using state and following action.
%Tsteps=500;  %DURATION OF TIME OF SIMULATION=Tsteps - Tsteps=800 for 3 by 5, 2400 works for 3 by 8.
%TotalTests=10; %NUMBER of tests for observing average time to optimal reward receipt.
%withgraphics=0; withendgraphics=1; %withgraphics=1 runs slower, but shows
%location and action-value function while running
%If remove central blockade, get reasonable learning, but not optimal in
%all cases.

%Script by Michael E. Hasselmo, Department of Psychology Center for Memory and Brain,
%Boston University, 2 Cummington St., Boston, MA 02215.
%hasselmo@bu.edu, http://people.bu.edu/hasselmo/

%***************************
close all; clear  %Clears the screen of previous windows.
%***************************
%CHOOSE NEOCORTEX (Neocortex=1) VS. REINFORCEMENT LEARNING (Neocortex=0).
Neocortex=1; SARSA=1; SA=1;  %SA sets it to form representations using state and following action.
Tsteps=500;  %DURATION OF TIME OF SIMULATION=Tsteps - Tsteps=800 for 3 by 5, 2400 works for 3 by 8.
TotalTests=1; %NUMBER of tests for observing average time to optimal reward receipt.
withgraphics=1; withendgraphics=1; %withgraphics=1 runs slower, but shows location and action-value function while running
%*******************
%ENVIRONMENT
%Size of environment - sv=vertical state coordinate, sh=horizontal state coordinate
vz=3; hz=5;
%Array for creating barriers in the environment.
B=zeros(vz,hz); %B(2,3)=1; %B(2,2)=1; %B(1,4)=1; B(2,4)=1; B(3,2)=1; %B(4,2)=1; B(2,3)=1;
%DEFINE START LOCATION OF VIRTUAL RAT AND REWARD/GOAL LOCATION.
stv=ceil(vz/2); sth=1; gv=ceil(vz/2); gh=hz;
%PUT AGENT INTO START LOCATION. "rat" appears as a black square at sv,sh
sv=stv; sh=sth; ov=sv; oh=sh;
%DEFINE SIZE OF STATE AND ACTION VECTORS ON BASIS OF ENVIRONMENT AND ACTIONS.
%"No" is number of output/actions 4 for up, down, left, right. "n" is number of states/locations.
No=4; n=vz*hz; N=No*n;
%**********************************
%GRAPHICS - INITIAL DEFINITION OF FIGURES OF NETWORK FUNCTION.
%The following is used to plot location, value function (V) and action value (Q).
xL=zeros(vz,hz); xC=zeros(vz,hz); xQ=zeros(vz*3,hz*3);
%Graphics showing location of agent in environment.
xL=64*(1-0.4*B); figure('Position',[360 40 300 300]);
figh=gcf; axh=gca; colormap(gray); image(xL)
set(figh,'Name','Location (state)','numbertitle','off');
%Graphics showing value function of each state.
xC=B; xC=64*(1-xC); figure('Position',[40 40 300 300]);
figk=gcf; axk=gca; colormap(gray); set(gca,'Drawmode','fast'); image(xC)
set(figk,'Name','StateVal(V) or Wib','numbertitle','off');
%Graphics showing action-value function Q plotted by state/location
xQ=64*(1-xQ); figure('Position',[360 440 300 300]);
figj=gcf; axj=gca; colormap(gray); set(gca,'Drawmode','fast'); image(xQ)
set(figj,'Name','Plot of State-Action-Value (Q)','numbertitle','off');
for m=1:hz, XTick(m)=3*m+0.5; YTick(m)=3*m+0.5; end %This last creates grid arrays for fig.

%***********************************
%Compute performance measure over time.
cTest=zeros(Tsteps+10,1);

for Mult=1:TotalTests;
   Mult
    
%***********************************
%SET NEOCORTEX VECTORS AND MATRICES.
if Neocortex==1
%a is sensory input, oa buffers a(t-1), g is spread from goal c is
%gated spread for guiding output, goal is goal representation.
a=zeros(1,N); oa=zeros(1,N); c=zeros(1,N); oc=zeros(1,N); g=zeros(1,N); Rew=0; goal=zeros(1,N);
o=zeros(1,No); Out=4; oldOut=4; gm=0.9999; %Both TD learn and Neocortex use output vector o.
if SA==0; for m=1:No; goal((m-1)*n + ((gv-1)*hz+gh))=gm; end; end  %A-S version.
if SA==1; for m=1:No; goal(((gv-1)*hz+gh-1)*No +m)=gm; end; end  %S-A version.
%Set neocortex weight matrix - W for spread from goal, Wo for spread to output.
W=zeros(N,N); eta=0.2; %eta=rate of synaptic modification
for m=1:No; for m2=1:n; Wo(m,(m2-1)*No+m)=1; end; end;
%SET NEOCORTEX GRAPHICS VECTORS FOR ACTIVITY. For viewing of vectors over all time steps at end of simulation.
cA=zeros(1,N); coldA=zeros(1,N); cG=zeros(1,N); cC=zeros(1,N); coC=zeros(1,N);
end %End Neocortex==1 conditional.

%***************
%SET TD LEARNING VECTORS (LOOK-UP TABLES). V=value function, Q=action-value function
if Neocortex==0; a=zeros(1,N); cA=zeros(1,N); cQ=zeros(1,N); V=zeros(1,n); dV=zeros(1,n);
Q=zeros(1,N); dQ=zeros(1,N); Rew=0; o=zeros(1,No); Out=4; oldOut=4; alf=0.2; end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%MAIN - Main loop.  This part of the script runs the agent ("virtual rat") through the environment.
for steps = 1:Tsteps
%LOGIC FOR COMPUTING REWARD PRESENTATION.
%Compute finding of reward if state=goal location.
if(sv==gv & sh==gh); Rew=1; end
if (ov==gv & oh==gh); sv=stv; sh=sth; Rew=0;   %Commenting this out can allow local loops to form!
    if Neocortex==1; a=zeros(1,N); o=zeros(1,No); end; end;  %Prevents learning to output at goal location.

%CREATE INDEX FOR STATE si=state index and oi= old state index (Index for action/output appears later)
si=((sv-1)*hz+(sh-1))*No;
oi=((ov-1)*hz+(oh-1))*No;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%Main retrieval for NEOCORTEX.
%Retrieval comes first.  Encoding comes later after computation of Output.
if Neocortex==1;
%COMPUTE THE INPUT "a"
if SA==0; a((Out-1)*n+(sv-1)*hz+sh)=1; end %Action-State based on proprioception of OUTPUT/ACTION and next state obtained sv,sh
if SA==1; a(si+1:si+No)=1; end %SA uses states combined with the following action.

%RETRIEVAL SPREAD
%Rsteps=6; for ret=0:Rsteps; g= goal + (W*g')'; %end; %Take away the "end" here - put GOAL CLAMPING into same loop.
%if SA==1; for m=1:No; g(((gv-1)*hz+gh-1)*No +m)=gm; end; end; end %Add "end" here to Clamp goal units to goal magnitude "gm".
g=((inv(eye(N) - W))*goal')';  %This replaces iterative loop= sumk(Wgg^k*goal). Works well with clipping below, but does NOT work
%with postsynaptic gated decay because stability is lost due to no clamping of goal activity.
c=a.*g; %Used in plotting xQ and in computing output below.

%GRAPHICS FOR NEOCORTEX "ACTION VALUES".  This command shows the output function oq=Wo*c due to backward retrieval
 oq=(Wo*c')';
   xQ(sv*3-2,sh*3-1)=64*(1-oq(1)); xQ(sv*3,sh*3-1)=64*(1-oq(2));
   xQ(sv*3-1,sh*3-2)=64*(1-oq(3)); xQ(sv*3-1,sh*3)=64*(1-oq(4)); %Scaled by 0.3.
end %End neocortex==1 conditional for RETRIEVAL.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%BEFORE UPDATE OUTPUT SHOW CURRENT LOCATION (state) OF AGENT (rat) IN ENVIRONMENT.
%The if function shuts off graphics to speed up simulation for longer runs.
if withgraphics==1 %| Tsteps <= 1000 | ((Tsteps > 1000) & ((steps > Tsteps - 100) | steps < 100));
 xL=64*(1-0.4*B); xL(sv,sh)=0; axes(axh); image(xL,'EraseMode','none'); drawnow; end

%If withgraphics then show the output function oq=xQ (neocortex=1) or the action-value Q (neocortex=0) on each step
if withgraphics==1 | steps==Tsteps
axes(axj); image(xQ,'EraseMode','none'); set(axj,'XTick',XTick, 'YTick',YTick);
set(axj,'XGrid','on', 'YGrid','on','ZGrid','off','GridLineStyle','-'); drawnow; end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%OUTPUT - COMPUTATION OF OUTPUT/ACTION.
%COMPUTE VECTORS WHICH GUIDE BEHAVIOR ON BASIS OF NEOCORTEX RETRIEVAL OR Q VALUES.
rando=[sprand(1,4,1)];  %Creates a 4 element random vector with values between 0 and 1.
if sprand(1,1,1) > 1.0; o=rando; %This line sets random exploration even with prior learning if value not 1.0.

elseif Neocortex==1; o=0.001*rando + (Wo*c')'; %Same as c(si+1:si+No), because Wo just translates elements into o.
elseif Neocortex==0; o=0.001*rando + Q(si+1:si+No); end;

%Make sure output "o" won't push the rat into barrier or off field.
    if sv<=1 | B(sv-1,sh)==1; o(1,1)=-100; end
    if sv>=vz | B(sv+1,sh)==1; o(1,2)=-100; end
    if sh<=1 | B(sv,sh-1)==1; o(1,3)=-100; end
    if sh>=hz | B(sv,sh+1)==1; o(1,4)=-100; end

%CHOOSE THE MAXIMUM OUTPUT/ACTION IN THE OUTPUT VECTOR.
%The following line finds the max of the output vector for output
   oldOut=Out;
   [C,Out]=max(o);
   ov=sv; oh=sh; o=zeros(1,No); %Before altering sv and sh, save the ov and oh values.
 %Change location on basis of largest output value (detected by "max" function above).
 if Out==1; sv=sv-1; sh=sh; elseif Out==2; sv=sv+1; elseif Out==3; sh=sh-1; sv=sv; elseif Out==4; sh=sh+1; sv=sv; end
 o(1,Out)=1;  %update 4 element output vector.

%COMPUTE LEARNING FOR NEOCORTEX
if Neocortex==1;
if SA==1; a=zeros(1,N); a(si+Out)=1; end  %Set specific input to state*No+action.
W=W+eta*(oa'*(a.*g)); %This version with oa learns in about 250 steps for 3by5 with block in middle B(2,3)=1
%for i=1:N;for j=1:N; dW(i,j)=eta*(oa(i)*c(j) - oa(i)*W(i,j)); end; end; W=W+dW; %About 400 steps. 
%Note that testing stability with postsynaptic gated decay requires commenting out the clipping at 0.9 below.

W(find(W > 0.9))=0.9; %Set max of matrix W. Works much better with max near 1.0 (mathematically only need less than 1.0).
if withendgraphics==1; cA=vertcat(cA,a); cC=vertcat(cC,c); coC=vertcat(coC,oc); cG=vertcat(cG,g); end %Update graphics vectors for withendgraphics.
%Buffer previous value of "a" as "oa" for next cycle of encoding, then zero out vectors.
oa=a; a=zeros(1,N); o=zeros(1,No); g=zeros(1,N);
end %End Neocortex==1 conditional.

%%%%%%%%%%%%%%%%%%%%%%%%%%
%TDLEARN COMPUTATIONS FOR TDLEARNING OF VALUE AND ACTION-VALUES TDlearning (Neocortex==0) CONDITIONAL
if Neocortex==0;

%These graphics show the value function V of each state (not the Q)
xC(ov,oh)=64*(1-2*V((ov-1)*hz+oh));
if withgraphics==1 | steps==Tsteps
    %This command makes the new image of the value function on each step
    axes(axk); image(xC,'EraseMode','none'); drawnow; end
%Update the graphics file for plotting the action-value function Q.
if Out==1; xQ(ov*3-2,oh*3-1)=64*(1-Q(oi+oldOut));
elseif Out==2; xQ(ov*3,oh*3-1)=64*(1-Q(oi+oldOut));
elseif Out==3; xQ(ov*3-1,oh*3-2)=64*(1-Q(oi+oldOut));
elseif Out==4; xQ(ov*3-1,oh*3)=64*(1-Q(oi+oldOut)); end
if SARSA==0;
%SARSA==0 is actor-critic, which uses TD error to update action value.
dV((ov-1)*hz+oh)=V((sv-1)*hz+sh)+Rew-V((ov-1)*hz+oh);
V((ov-1)*hz+oh)=V((ov-1)*hz+oh)+alf*dV((ov-1)*hz+oh);
Q(oi+oldOut)=Q(oi+oldOut)+alf*dV((ov-1)*hz+oh);
elseif SARSA==1; %This is SARSA.
dQ(oi+oldOut)=Q(si+Out)+Rew-Q(oi+oldOut);
Q(oi+oldOut)=Q(oi+oldOut)+alf*dQ(oi+oldOut);
end %End the SARSA conditional for Q learning.
    a=zeros(1,N); a(oi+oldOut)=1; cA=vertcat(cA,a); cQ=vertcat(cQ,Q);
end %End the Neocortex==0 conditional segment.

    cTest(steps)=cTest(steps)+ Rew;
    end %End multiple loop testing of performance.

%End the whole MAIN LOOP (for steps=1:Tsteps).
end

%SHOW END GRAPHICS SUMMARIZING ACTIVITY DURING SIMULATION.
if withendgraphics==1
if Neocortex==1
%Show the activity of pop g with W at end.
cG=vertcat(cG,W); xB=64*(1-cG); figure('Position',[80 40 300 300]); h=gcf;
set(h,'Name','g with W at end'); set(h,'numbertitle','off'); colormap(gray); image(xB)
%Show the activity of pop c with Wo at end.
cC=vertcat(cC,Wo); xB=64*(1-cC); figure('Position',[680 40 300 300]); h=gcf;
set(h,'Name','"c" vector with Wo at end'); set(h,'numbertitle','off'); colormap(gray); image(xB)
%Show the activity of pop oc with Wo at end.
coC=vertcat(coC,Wo); xB=64*(1-coC); figure('Position',[380 40 300 300]); h=gcf;
set(h,'Name','"oc" vector with Wo at end'); set(h,'numbertitle','off'); colormap(gray); image(xB)
end %End the Neocortex==1 conditional
%Show the activity of pop "a"
xB=64*(1-cA); figure('Position',[980 40 300 300]); h=gcf;
set(h,'Name','"a" vector all steps'); set(h,'numbertitle','off'); colormap(gray); image(xB)
if Neocortex==0; xB=64*(1-cQ); figure('Position',[680 40 300 300]); h=gcf;
set(h,'Name','"Q" vector all steps'); set(h,'numbertitle','off'); colormap(gray); image(xB); end
end %End final withgraphics==1 conditional.
%Note that need to use optimalSteps to determine number of rewards received relative to maximal rate.
optimalSteps=hz;
cTest2=zeros(ceil(Tsteps/optimalSteps), 1);

for m=0:floor((Tsteps-1)/optimalSteps)
    for m2=1:(optimalSteps)
    cTest2(m+1,1)=cTest2(m+1,1)+cTest(m*(optimalSteps)+m2,1);
    end
end

figure('Position',[500 500 400 200]);
h=gcf;
set(h,'Name','Approach to optimal function');
colormap(gray);
plot(cTest2/TotalTests, '-k')
xlabel('trial')
ylabel('Rewards received/optimal')
title('Average # rewards/optimal','FontSize',12)