%% Reinforcement learning in a maze
clear; close all;

tau = @(x,u) x+heaviside(x-1.5).*heaviside(4.5-x)*u;
rho = @(x,u) (x==2&&u==-1)+2*(x==4&&u==1);
idx = @(u)   (u+3)/2';
gamma=0.5; 

%% Analytic solution for optimal policy
i=0; 
for x=1:5; for u=-1:2:1;  
  i=i+1; R(i)=rho(x,u); 
end; end;

T=zeros(10,10);
T(1,1)=1; T(2,2)=1; T(3,1)=1; T(4,6)=1; T(5,3)=1;
T(6,8)=1; T(7,6)=1; T(8,10)=1; T(9,9)=1; T(10,10)=1;

Q=inv(eye(10)-gamma*T)*R';
Q=reshape(Q,2,5)

%% policy iteration
Q=zeros(5,2); pi=-ones(5,1); 
for iter=1:3;
    for x=1:5; for u=-1:2:1
        Q(x,idx(u))=rho(x,u)+gamma*Q(tau(x,u),idx(pi(tau(x,u))));
    end; end
    for x=1:5; 
        [~,uidx]=max([Q(x,1),Q(x,2)]); pi(x)=2*uidx-3;       
    end
end
disp('Policy iteration'); disp(Q')


%% Q-iteration
Q=zeros(5,2); 
for iter=1:3;
    for x=1:5; for u=-1:2:1
        Qnew(x,idx(u))=rho(x,u)+gamma*max(Q(tau(x,u),1),Q(tau(x,u),2));
    end; end
Q=Qnew;
end
disp('Q-iteration'); disp(Q')

%% SARSA
Q=zeros(5,2); alpha=0.1;  %starting state
for trial=1:1000
    x=3;
    for t=1:5;
        [~,u]=max([Q(x,1),Q(x,2)]);u=2*u-3;
        if rand()<0.1; u=-u; end %epsilon exploration
        Q(x,idx(u))=(1-alpha)*Q(x,idx(u))+alpha*(rho(x,u)+gamma*Q(tau(x,u),idx(u)));
        x=tau(x,u);
    end;
    error(trial)=sum(sum(abs(Q-Qnew)));
end
disp('SARSA'); disp(Q')
plot(error)

return

%% Q-learning
Q=zeros(5,2); alpha=0.1;  %starting state
for trial=1:1000
    x=3;
    for t=1:5;
        [~,u]=max([Q(x,1),Q(x,2)]);u=2*u-3;
        if rand()<0.1; u=-u; end %epsilon exploration
        Q(x,idx(u))=(1-alpha)*Q(x,idx(u))+alpha*(rho(x,u)+gamma*max(Q(tau(x,u),1),Q(tau(x,u),2)));
        x=tau(x,u);
    end;
    error(trial)=sum(sum(abs(Q-Qnew)));
end
disp('Q-learning'); disp(Q')
plot(error)

