% given a policy determine V_pi
% this is for the 5 by 5 grid world 

numactions = 4; % for north, south, east, and west
xsize = 5; % ranging from 1 to 5
ysize = 5; % ranging from 1 to 5

% let us number the states from 1 
states = [1:xsize*ysize];

% define the probabilistic transition and rewards matrix
% this says which actions taken in which states, lead where.
% a lot of simplifications are made in this code, since
%  - at each cell, all actions are possible
%  - transitions are deterministic

for oldstate=1:size(states,2)
  for action=1:numactions
	% figure out old coordinates
	oldx = floor((oldstate-1)/ysize)+1;
	oldy = oldstate-((oldx-1)*ysize);
	% figure out new coordinates
	if (action == 1) % north
	   if (oldx > 1)
		newx = oldx-1; newy = oldy; rew = 0;
	   else
		newx = oldx; newy = oldy; rew = -1;
	   end
	end
	if (action == 2) % south
	   if (oldx < xsize)
		newx = oldx+1; newy = oldy; rew = 0;
	   else
		newx = oldx; newy = oldy; rew = -1;
	   end
	end
	if (action == 3) % east
	   if (oldy < ysize)
		newx = oldx; newy = oldy+1; rew = 0;
	   else
		newx = oldx; newy = oldy; rew = -1;
	   end
	end
	if (action == 4) % west
	   if (oldy > 1)
		newx = oldx; newy = oldy-1; rew = 0;
	   else
		newx = oldx; newy = oldy; rew = -1;
	   end
	end
	% the following two rules are exceptions to the above
	% rules for directional movement.
	% this silly MATLAB version does not have a continue
	% construct, hence the sloppy overwriting
	if (oldx == 1 & oldy == 2)
		newx = 5; newy = 2; rew = 10;
		% i.e., irresp. of action
	end 
	if (oldx == 1 & oldy == 4)
		newx = 3; newy = 4; rew = 5;
		% i.e., irresp. of action
	end
	newstate = (newx-1)*ysize+newy;
	transfer(oldstate,action) = newstate;
	rewards(oldstate,action) = rew;
  end
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% now, the rest of the code is specific to what you want to do 
% with this RL problem. Here, we are given a policy.
% So, we define the policy matrix first
% this basically gives the probability of taking 
% an action in a given state
% for now, we are defining them to be equiprobable

for i=1:size(states,2)
  for j=1:numactions
    pi(i,j) = 1/numactions;
  end
end

% gamma is the discounting rate
gamma = 0.9;

% to find the V values for the given policy, we first
% setup the set of linear equations to be solved;
% for now, only concentrate on the right hand sides
% the left sides are just equal to the variables
% prepare a nice coefficient matrix

for i=1:size(states,2) % there are this many simultaneous equations
  for j=1:size(states,2)+1 % with this many variables
			   % the +1 is for a constant term! :)
    	coeffs(i,j) = 0;  			   
  end
end

% start entering the coefficient values

for i=1:size(states,2)
  % now, you are writing the right side for V_pi(i)
  rewardcount = 0;
  for j=1:numactions
	% figure out where this action leads you to
	newstate = transfer(i,j);
	coeffs(i,newstate) = coeffs(i,newstate) + pi(i,j)*gamma;
	rewardcount = rewardcount + rewards(i,j)*pi(i,j);
  end
  coeffs(i,size(states,2)+1) = rewardcount;
end

% now solve this linear system
% setup A and b, and do x = A\b

Amatrix = coeffs(:,1:size(states,2));
Amatrix = Amatrix - eye(size(Amatrix));
rightside = -1.*coeffs(:,size(states,2)+1);
Vvalues = Amatrix\rightside;

% arrange the V_pi values nicely on a grid
reshape(Vvalues,xsize,ysize)'

