% File: value_iteration.m 
% Author: Patric Jensfelt
% Date: 2007-10-07
% 
% Solves the MDP described in Section 17.1 in Russell and Norvigs
% Artificial Intelligence A Modern Approach
%
% Usage value_iteration(rNorm,rGoal,rTrap,gamma)
%
function value_iteration(rNorm, rGoal, rTrap, g)

if nargin < 4
  g = 0.9999;
end

if nargin < 3
  rTrap = -1;
end

if nargin < 2
  rGoal = 1;
end

if nargin < 1
  rNorm = -0.04;
end

% Initialize the utilities to 0
U = zeros(3,4);

% Allocate matrix for optimal policy
PI = zeros(3,4);

global LEFT
global RIGHT 
global UP 
global DOWN
global R14
global R24
global r;
global Pforw
global Pside

LEFT = 1;
DOWN = 2;
RIGHT = 3;
UP = 4;

maxchange = 1e-3;
Pforw = 0.8;
Pside = 0.1;
r = rNorm;
R14 = rGoal;
R24 = rTrap;
gamma = g;

disp(sprintf('gamma=%f',gamma))
disp(sprintf('rNorm=%f',r))
disp(sprintf('rGoal=%f',R14))
disp(sprintf('rTrap=%f',R24))

U(2,4) = R24;
U(1,4) = R14;

run = 1;

display_result(U,PI);
disp('Press ENTER to start')
pause

% Loop until we are done
while run

  Uold = U;
  
  for row = 1:3
    for col = 1:4

      % Skip occupied/terminal squares
      if ((row == 2 && col == 2) || (row == 1 && col == 4) || ...
          (row == 2 && col == 4)) 
        continue; 
      end      

      actions = [LEFT DOWN RIGHT UP];

      for a = 1:length(actions)
        Q(a) = get_q(U, actions(a), gamma, row, col);
      end

      if min(Q) == max(Q)
        U(row,col) = Q(1);
        PI(row,col) = -1;
      else
        [u,a] = max(Q);
        U(row, col) = u;
        PI(row, col) = actions(a);
      end
    end
  end

  display_result(U,PI);

  diff = max(max(abs(Uold - U)));
  if diff < maxchange
    run = 0;
  end

  pause

end

return;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

function display_result(U,PI)

global LEFT
global RIGHT 
global UP 
global DOWN

[qx,qy] = get_quivers(PI);
qx = flipud(qx);
qy = flipud(qy);

figure(1)
clf, hold on
display_field;
for row = 1:3
  for col = 1:4
    if (row == 2 && col == 2)
      continue
    end
        
    text(col-0.3, 4-row, sprintf('%.3f', U(row,col)), 'FontSize',15)
  end
end
axis([0 5 0 4])
axis off

figure(2)
clf, hold on
display_field;
quiver(qx,qy)
axis([0 5 0 4])
axis off

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

function display_field

for row = 1:4
  plot([0.5 4.5], row-[0.5 0.5])
end
for col = 1:5
  plot(col-[0.5 0.5], [0.5 3.5])
end
fill_square(2,2,'k')
fill_square(1,4,'g')
fill_square(2,4,'r')

return;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

function fill_square(row, col, c)
fill(col+[-0.5 0.5 0.5 -0.5 -0.5], 4-row+[-0.5 -0.5 0.5 0.5 -0.5], c)

return;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

function R = get_reward(row, col)

global R14
global R24
global r

R = r;

if (row == 1 && col == 4)
  R = R14;
elseif (row == 2 && col == 4) 
  R = R24;
elseif (row == 2 && col == 2)
  R = 0;
end

return;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

function [qx,qy] = get_quivers(PI)

global LEFT 
global RIGHT
global UP
global DOWN

for row = 1:3
  for col = 1:4
    % Skip occupied/terminal squares
    if ((row == 2 && col == 2) || (row == 1 && col == 4) || ...
        (row == 2 && col == 4)) 
      qx(row,col) = 0;
      qy(row,col) = 0;
    else
      if PI(row,col) == LEFT
        qx(row,col) = -1;
        qy(row,col) = 0;
      elseif PI(row,col) == RIGHT
        qx(row,col) = 1;
        qy(row,col) = 0;
      elseif PI(row,col) == UP
        qx(row,col) = 0;
        qy(row,col) = 1;
      elseif PI(row,col) == DOWN
        qx(row,col) = 0;
        qy(row,col) = -1;
      else
        qx(row,col) = 0;
        qy(row,col) = 0;
      end
    end      
  end
end


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

function Q = get_q(U, A, gamma, row, col)

global LEFT 
global RIGHT
global UP
global DOWN
global Pforw
global Pside

Q = get_reward(row, col);

switch A
  case LEFT,
      Q = Q + gamma * (Pforw * get_utility_of_action(row,col,U,LEFT) + Pside * get_utility_of_action(row, col, U, UP) + Pside * get_utility_of_action(row,col, U, DOWN));
  case RIGHT,
      Q = Q + gamma * (Pforw * get_utility_of_action(row,col,U,RIGHT) + Pside * get_utility_of_action(row, col, U, UP) + Pside * get_utility_of_action(row,col, U, DOWN));
  case UP,
      Q = Q + gamma * (Pforw * get_utility_of_action(row,col,U,UP) + Pside * get_utility_of_action(row, col, U, LEFT) + Pside * get_utility_of_action(row,col, U, RIGHT));
  case DOWN,
      Q = Q + gamma * (Pforw * get_utility_of_action(row,col,U,DOWN) + Pside * get_utility_of_action(row, col, U, LEFT) + Pside * get_utility_of_action(row,col, U, RIGHT));
end

return;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

function u = get_utility_of_action(row, col, U, action)

global LEFT 
global RIGHT
global UP
global DOWN

u = U(row,col);

if action == DOWN
  if (row < 3 && (row ~= 1 || col ~= 2))
    u = U(row+1,col);
  end
elseif action == UP
  if (row > 1 && (row ~= 3 || col ~= 2))
    u = U(row-1,col);
  end
elseif action == RIGHT
  if (col < 4 && (row ~= 2 || col ~= 1))
    u = U(row,col+1);
  end
else
  if (col > 1 && (row ~= 2 || col ~= 3))
    u = U(row,col-1);
  end
end

return; 