%% Setup

n = 4; %state space is nxn grid, numbered columnwise from upper left

%transition probabilities, T(s,a,s')
T = zeros(n^2,4,n^2);
for i = 1:n
    for j=1:n-1
        T(i+n*j,1,i+n*(j-1)) = 1; %move left
        T(i+n*(j-1),2,i+n*j) = 1; %move right
        T((i-1)*n+j+1,3,(i-1)*n+j) = 1; %move up
        T((i-1)*n+j,4,(i-1)*n+j+1) = 1; %move down
    end
    T(i,1,i) = 1; %move left against boundary (no state change)
    T(n*(n-1)+i,2,n*(n-1)+i) = 1; %move right against boundary (no state change)
    T((i-1)*n+1,3,(i-1)*n+1) = 1; %move up against boundary (no state change)
    T(i*n,4,i*n) = 1; %move down against boundary (no state change)
end

%rewards (set to depend only on current state)
R = zeros(n^2); 
R(n+2) = 1; %positive reward near upper-left corner
R(n*(n-2)) = -1; %negative reward near lower-right corner

%% Simulate state-value learning with random actions

%model parameters
e = .2; %learning rate
g = .7; %discount factor

V = zeros(n^2,1); %state values
s = randi([1 n^2],1); %starting state
figure(1) %open graphics
while ishandle(1) %run until user closes figure
    a = randi([1 4],1); %random action
    snew = find(rand<cumsum(squeeze(T(s,a,:))),1,'first'); %new state sampled according to T
    V(s) = (1-e)*V(s) + e*(R(s)+g*V(snew)); %learn
    s = snew;
    %draw state space and state values
    figure(1),clf,hold on
    set(gca,'XTick',[],'YTick',[],'Xlim',[0 n],'YLim',[0 n])
    axis square; box
    for i=1:n-1
        plot([0 n],[i i],'k')
        plot([i i],[0 n],'k')
    end
    for i=1:n, for j=1:n
            greyscale = (nthroot(V((i-1)*n+j)*(1-g),3)+1)/2; 
            fill([i-1 i i i-1 i-1],[n-j n-j n-j+1 n-j+1 n-j],[1 1 1]*greyscale)
    end, end
    fill(ceil(snew/n)-[.6 .5 .4 .5 .6],n-mod(snew-1,n)-[.5 .6 .5 .4 .5],'b')
    drawnow
end

%% Simulate Q-learning

T(n+2,:,:) = 1/n^2; %random transport after entering 'good' state

%model parameters
e = .2; %learning rate
g = .7; %discount factor
t = .2; %choice temperature

Q = zeros(n^2,4); %Q-values, initialized to zero
s = randi([1 n^2],1); %starting state
figure(1) %open graphics
while ishandle(1) %run until user closes figure
    a = find(rand<cumsum(exp(Q(s,:)/t))/sum(exp(Q(s,:)/t)),1,'first'); %softmax action selection
    snew = find(rand<cumsum(squeeze(T(s,a,:))),1,'first'); %new state sampled according to T
    Q(s,a) = (1-e)*Q(s,a) + e*(R(s)+g*max(Q(snew,:))); %learn
    s = snew;
    %draw state space and state values
    figure(1),clf,hold on
    set(gca,'XTick',[],'YTick',[],'Xlim',[0 n],'YLim',[0 n])
    axis square; box
    for i=1:n-1
        plot([0 n],[i i],'k')
        plot([i i],[0 n],'k')
    end
    for i=1:n, for j=1:n
            greyscale = (nthroot(Q((i-1)*n+j,1)*(1-g),3)+1)/2;
            fill([i-1 i-.6 i-.6 i-1 i-1],[n-j n-j+.4 n-j+.6 n-j+1 n-j],[1 1 1]*greyscale) %left action
            greyscale = (nthroot(Q((i-1)*n+j,2)*(1-g),3)+1)/2;
            fill([i-.4 i i i-.4 i-.4],[n-j+.4 n-j n-j+1 n-j+.6 n-j+.4],[1 1 1]*greyscale) %right action
            greyscale = (nthroot(Q((i-1)*n+j,3)*(1-g),3)+1)/2; 
            fill([i-.6 i-.4 i i-1 i-.6],[n-j+.6 n-j+.6 n-j+1 n-j+1 n-j+.6],[1 1 1]*greyscale) %up action
            greyscale = (nthroot(Q((i-1)*n+j,4)*(1-g),3)+1)/2; 
            fill([i-1 i i-.4 i-.6 i-1],[n-j n-j n-j+.4 n-j+.4 n-j],[1 1 1]*greyscale) %down action
    end, end
    fill(ceil(snew/n)-[.6 .5 .4 .5 .6],n-mod(snew-1,n)-[.5 .6 .5 .4 .5],'b')
    drawnow
end
