%%%%%%% CS 229 Machine Learning %%%%%%%%%%%
%%%%%%% Programming Assignment 4 %%%%%%%%%%
%%%
%%% Parts of the code (cart and pole dynamics, and the state
%%% discretization) are adapted from code available at the RL repository 
%%% http://www-anw.cs.umass.edu/rlr/domains.html
%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% This file controls the pole-balancing simulation. You need to write
% code in places marked "CODE HERE" only.

% Briefly, the main simulation loop in this file calls cart_pole.m for
% simulating the pole dynamics, get_state.m for discretizing the
% otherwise continuous state space in discrete states, and show_cart.m
% for display. 

% Some useful parameters are listed below.

% NUM_STATES: Number of states in the discretized state space
% You must assume that states are numbered 1 through NUM_STATES. The
% state numbered NUM_STATES (the last one) is a special state that marks
% the state when the pole has been judged to have fallen (or when the
% cart is out of bounds). However, you should NOT treat this state any
% differently in your code. Any distinctions you need to make between
% states should come automatically from your learning algorithm.

% After each simulation cycle, you are supposed to update the transition
% counts and rewards observed. However, you should not change either
% your value function or the transition probability matrix at each
% cycle. 

% Whenever the pole falls, a section of your code below will be
% executed. At this point, you must use the transition counts and reward
% observations that you have gathered to generate a new model for the MDP
% (i.e., transition probabilities and state rewards). After that, you
% must use value iteration to get the optimal value function for this MDP
% model.

% TOLERANCE: Controls the convergence criteria for each value iteration
% run
% In the value iteration, you can assume convergence when the maximum
% absolute change in the value function at any state in an iteration
% becomes lower than TOLERANCE.

% You need to write code that chooses the best action according
% to your current value function, and the current model of the MDP. The
% action must be either 1 or 2 (corresponding to possible directions of
% pushing the cart).

% Finally, we assume that the simulation has converged when
% 'NO_LEARNING_THRESHOLD' consecutive value function computations all 
% converged within one value function iteration. Intuitively, it seems
% like there will be little learning after this, so we end the simulation
% here, and say the overall algorithm has converged. 


% Learning curves can be generated by calling plot_learning_curve.m (it
% assumes that the learning was just executed, and the array
% time_steps_to_failure that records the time for which the pole was
% balanced before each failure are in memory). num_failures is a variable
% that stores the number of failures (pole drops / cart out of bounds)
% till now.

% Other parameters in the code are described below:

% GAMMA: Discount factor to be used

% The following parameters control the simulation display; you dont
% really need to know about them:

% pause_time: Controls the pause between successive frames of the
% display. Higher values make your simulation slower.
% min_trial_length_to_start_display: Allows you to start the display only
% after the pole has been successfully balanced for at least this many
% trials. Setting this to zero starts the display immediately. Choosing a
% reasonably high value (around 100) can allow you to rush through the
% initial learning quickly, and start the display only after the
% performance is reasonable.

clear all; close all;

%%%%%%%%%% Simulation parameters %%%%%%%%%%

pause_time = 0.001*1e-2;
% min_trial_length_to_start_display = 1e3; % Basically, never turn it on.
min_trial_length_to_start_display = 400;
display_started=0;

NUM_STATES = 163;

GAMMA=0.995;

TOLERANCE=0.01;

NO_LEARNING_THRESHOLD = 20;

%%%%%%%%%%  End parameter list   %%%%%%%%%%

% Time cycle of the simulation
time=0;

% These variables perform bookkeeping (how many cycles was the pole
% balanced for before it fell). Useful for plotting learning curves.
time_steps_to_failure=[];
num_failures=0;
time_at_start_of_current_trial=0;

max_failures=500; % You should reach convergence well before this.  

% Starting state is (0 0 0 0)
% x, x_dot, theta, theta_dot represents the actual continuous state vector
x = 0.0; x_dot = 0.0; theta = 0.0; theta_dot = 0.0;

% state is the number given to this state - you only need to consider
% this representation of the state
state = get_state(x, x_dot, theta, theta_dot);

if min_trial_length_to_start_display==0 || display_started==1
  show_cart(x, x_dot, theta, theta_dot, pause_time);
end

%%% CODE HERE: Perform all your initializations here %%%

% Assume no transitions or rewards have been observed
% Initialize the value function array to small random values (0 to 0.10,
% say)
% Initialize the transition probabilities uniformly (ie, probability of
% transitioning for state x to state y using action a is exactly
% 1/NUM_STATES). Initialize all state rewards to zero.

%------------------------------------------------------------
% Tony Hyun Kim
% CS 229, PS #4, Problem 6
%------------------------------------------------------------
V = 0.1*rand(NUM_STATES,1);

% Basically, I keep running averages for the relevant variables:
R_acc = zeros(NUM_STATES,1); % Accumulator
R_cnt = zeros(NUM_STATES,1); % Count.

% H(i,j,k): the number of times that we transition from state i 
%       to state j by taking action k.
% Ha(i,k) : the number of times we have taken action k in state i
Hij = zeros(NUM_STATES,NUM_STATES,2);
Ha  = zeros(NUM_STATES,2);

% Default state transition probability is uniform.
P = 1/NUM_STATES*ones(NUM_STATES,NUM_STATES,2);

%%%% END YOUR CODE %%%%%%%%%%%%%%%%%%%%%%%%%%%%





%%% CODE HERE (while loop condition) %%%
% This is the criterion to end the simulation
% You should change it to terminate when the previous
% 'NO_LEARNING_THRESHOLD' consecutive value function computations all
% converged within one value function iteration. Intuitively, it seems
% like there will be little learning after this, so end the simulation
% here, and say the overall algorithm has converged. 

% A list of 4-tuples: 
%   (state, action, new_state, R)
% This is a temporary "buffer" for the results of a single run, since
% we are asked to update the MDP model only after the cart has failed.
observations = [];

no_learning_trials = 0;

while ( (num_failures<max_failures) && ... // Safety brake
        (no_learning_trials < NO_LEARNING_THRESHOLD) )

  %%% CODE HERE: Write code to choose action (1 or 2) %%%

  %------------------------------------------------------------
  % THK:
  %     Basically, we find the optimal policy based on our current
  %     estimate of the MDP. The code is very simple because we only
  %     have two available actions
  %------------------------------------------------------------
  
  % Get the current state
  state = get_state(x,x_dot,theta,theta_dot);
  
  % See which action (1 or 2) yields a better value function from
  %     the current state, and choose accordingly
  v1 = P(state,:,1)*V;
  v2 = P(state,:,2)*V;
  if( v1 > v2 );
      action = 1;
  elseif (v2 > v1)
      action = 2;
  else % If neither is clearly better, choose randomly, so that
       % we can explore the MDP model
      action = randi(2);
  end
   
  % This action choice algorithm is just for illustration. It may
  % convince you that reinforcement learning is nice for control
  % problems!  Replace it with your code to choose an action that is
  % optimal according to the current value function, and the current MDP
  % model.
%   if (theta<0)
%     action=1;
%   else
%     action=2;
%   end

%   if num_failures<50
%    if (rand(1) < 0.5)
%      action=1;
%    else
%      action=2;
%    end
%   end


  %%% END YOUR CODE %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
  

  % Get the next state by simulating the dynamics
  [x, x_dot, theta, theta_dot] = cart_pole(action, x, x_dot, theta, theta_dot);

  % Increment simulation time
  time = time + 1;
  
  % Get the state number corresponding to new state vector
  new_state = get_state(x, x_dot, theta, theta_dot);

  if display_started==1
    show_cart(x, x_dot, theta, theta_dot, pause_time);
  end
  
  % Reward function to use - do not change this!
  if (new_state==NUM_STATES)
    R=-1;
  else
    %R=-abs(theta)/2.0;
    R=0;
  end


  %%% CODE HERE: Perform updates %%%%%%%%%

  % A transition from 'state' to 'new_state' has just been made using
  % 'action'. The reward observed in 'new_state' (note) is 'R'.
  % Write code to update your statistics about the MDP - i.e., the
  % information you are storing on the transitions and on the rewards
  % observed. Do not change the actual MDP parameters, except when the
  % pole falls (the next if block)!

    observations = [observations; 
                    state action new_state R];
  
  % Recompute MDP model whenever pole falls
  % Compute the value function V for the new model
  if (new_state==NUM_STATES)

    % Update MDP model using the current accumulated statistics about the
    % MDP - transitions and rewards.
    % Make sure you account for the case when total_count is 0, i.e., a
    % state-action pair has never been tried before, or the state has
    % never been visited before. In that case, you must not change that
    % component (and thus keep it at the initialized uniform distribution).

    %------------------------------------------------------------
    % THK: 
    %   Transfer the recorded data in 'observation' buffer to the 
    %   variables of the MDP model
    %------------------------------------------------------------
    Nobs = length(observations(:,1));
    for i = 1:Nobs
        obs = observations(i,:); % (state action new_state R)
        
        % Increment the state transition probability
        Hij(obs(1),obs(3),obs(2)) = Hij(obs(1),obs(3),obs(2))+1;
        Ha(obs(1),obs(2)) = Ha(obs(1),obs(2))+1;
        
        % Accumulate the reward function (for 'new_state')
        R_acc(obs(3)) = R_acc(obs(3))+obs(4);
        R_cnt(obs(3)) = R_cnt(obs(3))+1;
    end
    observations = [];
    
    % Build the state transition probabilities. Note that if the 
    %   state-action pair has not yet been tried, then we just leave 
    %   the existing (uniform) distribution.
    for k = 1:2
        for s = 1:NUM_STATES
            if(Ha(s,k)~=0)
                P(s,:,k) = Hij(s,:,k)/Ha(s,k);
            end
        end
    end
    
    % Perform value iteration using the new estimated model for the MDP
    % The convergence criterion should be based on TOLERANCE as described
    % at the top of the file.
    % If it converges within one iteration, you may want to update your
    % variable that checks when the whole simulation must end

    
    %------------------------------------------------------------
    % THK:
    %   Begin value iteration
    %------------------------------------------------------------
    
    % Compute the best estimate for reward function as well as
    %   the state transition probability
    Rs = zeros(NUM_STATES,1);
    for s = 1:NUM_STATES
        if(R_cnt(s)~=0)
            Rs(s) = R_acc(s)/R_cnt(s);
        end
    end
    
    dV   = 1;
    vi_iter = 0;
    while( dV > TOLERANCE )
        V_new = zeros(NUM_STATES,1); % Synchronous update
        for s = 1:NUM_STATES
            V_new(s) = Rs(s) + ...
                       GAMMA*max([P(s,:,1)*V P(s,:,2)*V]);
        end
        dV = max(abs(V_new-V));
%         dV = norm(V_new-V);
        V  = V_new;
        vi_iter = vi_iter + 1;
    end
    
    if(vi_iter == 1)
        no_learning_trials = no_learning_trials + 1;
    else
        no_learning_trials = 0;
    end
    % THK: Debug output
    fprintf('trial #%3d: vi_iter=%3d, no_learning_trials=%2d\n',...
            num_failures,vi_iter,no_learning_trials);
    
    %pause(0.2); % You can use this to stop for a while!
    
  end
    

  %%% END YOUR CODE %%%%%%%%%%%%%%%%%%%
  

  % Dont change this code: Controls the simulation, and handles the case
  % when the pole fell and the state must be reinitialized
  if (new_state == NUM_STATES)
    num_failures = num_failures+1;
    time_steps_to_failure(num_failures) = time - time_at_start_of_current_trial;
    time_at_start_of_current_trial = time;

    time_steps_to_failure(num_failures);

    if (time_steps_to_failure(num_failures)> ...
	min_trial_length_to_start_display)
      display_started=1;
    end
    
    % Reinitialize state
    x = -1.1 + rand(1)*2.2;
    %x=0.0;
    x_dot = 0.0; theta = 0.0; theta_dot = 0.0;
    state = get_state(x, x_dot, theta, theta_dot);
  else 
    state=new_state;
  end
end

% Plot the learning curve (time balanced vs trial)
plot_learning_curve;
