% the state space
% 1 5
% 234
nstates=5;

% the action space
% up, right, down, left
nactions=4; 

%transition probability; for simplicity make it deterministic
%T[s,a]=new state
T=[
1 1 2 1
1 3 2 2
3 4 3 2
5 4 4 3
5 5 4 5
];

% rewards, again deterministic
% r[s,a]
r=[
  0 0 0 0
101 0 0 0
  0 0 0 -5
100 0 0 0
  0 0 0 0
];

% initialize the action value function q(s,a) randomly
q=rand(nstates, nactions)

% run forever 
s=3 % initial state
t=1; % iteration
epsilon=0.5; % epsilon-greedy parameter
alpha = 0.1; % Q-learning step size
gamma=0.9; % discounting factor
while 1
  
  % epsilon-greedy policy
  if (rand<=epsilon)
     % randomly pick an action
     a = randi(nactions);
  else
     [tmp, a]=max(q(s,:));
  end

  % perform action a in state s
  snew = T(s,a);
  reward = r(s,a);
  fprintf('iteration=%d s=%d a=%d snew=%d r=%d\n', t, s, a, snew, reward)

  % Q learning
  q(s,a) = (1-alpha)*q(s,a) + alpha*(reward + gamma*max(q(snew,:)))

  pause

  s = snew; % move to the new state
  t=t+1;
end

% To see q(2,1), note we immediately get a reward 101, then the optimal policy is to go down and up;
% every two steps we get reward 101 again.  that is,
% q(2,up=1)=101 + gamma^2*101 + gamma^4*101 + ... = 101*(1/1-gamma^2)=531.579
% Similarly, 
% q(4,up=1)=100 + gamma^2*100 + gamma^4*100 + ... = 100*(1/1-gamma^2)=526.316
% But for q(1,down), we need to wait one step then get to q(2,1), namely
% q(1,down=3)=0 + gamma*q(2,up=1)=478.421
% For q(3,left), we incur -5 immedately, and q(2,1) thereafter, thus
% q(3,left=4)=-5 + gamma*q(2,up=1)=473.421
