code: """
  CS 229 Machine Learning
  Question: Reinforcement Learning - The Inverted Pendulum
  """
  
  from __future__ import division, print_function
  import matplotlib
  matplotlib.use('TkAgg')
  from env import CartPole, Physics
  import matplotlib.pyplot as plt
  import numpy as np
  from scipy.signal import lfilter
  
  """
  Parts of the code (cart and pole dynamics, and the state
  discretization) are inspired from code available at the RL repository
  http://www-anw.cs.umass.edu/rlr/domains.html
  
  Briefly, the cart-pole system is described in `cart_pole.py`. The main
  simulation loop in this file calls the `simulate()` function for
  simulating the pole dynamics, `get_state()` for discretizing the
  otherwise continuous state space in discrete states, and `show_cart()`
  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 0 through `NUM_STATES` - 1. The
  state numbered `NUM_STATES` - 1 (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 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 0 or 1 (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 a code snippet at the end
  (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 is 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.
  """
  
  def initialize_mdp_data(num_states):
      """
      Return a variable that contains all the parameters/state you need for your MDP.
      Feel free to use whatever data type is most convenient for you (custom classes, tuples, dicts, etc)
  
      Assume that 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.
  
      Args:
          num_states: The number of states
  
      Returns: The initial MDP parameters
      """
      transition_counts = np.zeros((num_states, num_states, 2)) # P_ijk = P_{s_i, a_k} (s_j) i.e. s_i -> s_j under a_k
      transition_probs = np.ones((num_states, num_states, 2)) / num_states # initial uniform distribution
      #Index zero is count of rewards being -1 , index 1 is count of total num state is reached
      reward_counts = np.zeros((num_states, 2)) 
      reward = np.zeros(num_states)
      value = np.random.rand(num_states) * 0.1 # - 0.05 - 1/num_states
  
      return {
          'transition_counts': transition_counts,
          'transition_probs': transition_probs,
          'reward_counts': reward_counts,
          'reward': reward,
          'value': value,
          'num_states': num_states,
          'state_action_history': [[]]
      }
  
  def choose_action(state, mdp_data):
      """
      Choose the next action (0 or 1) that is optimal according to your current
      mdp_data. When there is no optimal action, return a random action.
  
      Args:
          state: The current state in the MDP
          mdp_data: The parameters for your MDP. See initialize_mdp_data.
  
      Returns:
          0 or 1 that is optimal according to your current MDP
      """
  
      # *** START CODE HERE ***
      # argmax over actions a of sum_s' P_sa(s') V(s')
      # s, s', a axes. times in axis s', sum over it, argmax over a and then take state s
      expected_state_action_value = (mdp_data['transition_probs'] * mdp_data['value'].reshape(1, -1, 1)).sum(axis=1)
      optimal = expected_state_action_value.argmax(axis=1)  # argmax over the 2 actions
      is_optimal = expected_state_action_value[state][int(optimal[state])] > \
                       expected_state_action_value[state][int(not optimal[state])]
      action = optimal[state] if is_optimal else np.random.randint(2)
      #print(f'state: {state}; optimal action: {action}; expected values: {expected_state_action_value[state]}')
      return optimal[state] if expected_state_action_value[state][int(optimal[state])] > expected_state_action_value[state][int(not optimal[state])] else np.random.randint(2)
      # *** END CODE HERE ***
  
  def update_mdp_transition_counts_reward_counts(mdp_data, state, action, new_state, reward):
      """
      Update the transition count and reward count information in your mdp_data. 
      Do not change the other MDP parameters (those get changed later).
  
      Record the number of times `state, action, new_state` occurs.
      Record the rewards for every `new_state` 
      (since rewards are -1 or 0, you just need to record number of times reward -1 is seen in 'reward_counts' index new_state,0)
      Record the number of time `new_state` was reached (in 'reward_counts' index new_state,1)
  
      Args:
          mdp_data: The parameters of your MDP. See initialize_mdp_data.
          state: The state that was observed at the start.
          action: The action you performed.
          new_state: The state after your action.
          reward: The reward after your action (i.e. reward corresponding to new_state).
  
      Returns:
          Nothing
      """
  
      # *** START CODE HERE ***
      mdp_data['state_action_history'][-1].append([state, action, new_state])
      if reward == -1:
          mdp_data['reward_counts'][new_state][0] += 1
      mdp_data['reward_counts'][new_state][1] += 1  # the number of times reaching this state regardless of reward -1 or 0
  
      mdp_data['transition_counts'][state, new_state, action] += 1
  
      # *** END CODE HERE ***
  
      # This function does not return anything
      return
  
  def update_mdp_transition_probs_reward(mdp_data):
      """
      Update the estimated transition probabilities and reward values in your MDP.
  
      Make sure you account for the case when 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).
      
      Args:
          mdp_data: The data for your MDP. See initialize_mdp_data.
  
      Returns:
          Nothing
  
      """
  
      # *** START CODE HERE ***
      num_states = mdp_data['value'].shape[0]
      transition_counts = mdp_data['transition_counts']
      # i.e. some s' transition occured for s,a pair
      state_actions_occurred = transition_counts.sum(axis=1) > 0
      state_actions_occurred = np.tile(state_actions_occurred, (num_states, 1, 1)).transpose(1,0,2)
      # where s,a pair occured take from counts, otherwise from the previous probs
      transition_probs = np.where(state_actions_occurred, transition_counts, mdp_data['transition_probs'])
      # counts at s' for s,a pair divided by total counts for s,a is our new prob - or same as old prob if no s,a
      # pair occured
      transition_probs = transition_probs/transition_probs.sum(axis=1).reshape(num_states, 1, 2)
      mdp_data['transition_probs'] = transition_probs
  
      # number of -1 rewards divided by total number of rewards (more general would be sum rewards / number times visited)
      # rewards <-> new states, so did we ever reach this new state?
      reward_occured = transition_counts.sum(axis=(0, 2)) > 0
      reward = np.where(reward_occured, - mdp_data['reward_counts'][:, 0] / mdp_data['reward_counts'][:, 1], mdp_data['reward'])
      mdp_data['reward'] = reward
      # *** END CODE HERE ***
  
      # This function does not return anything
      return
  
  def update_mdp_value(mdp_data, tolerance, gamma):
      """
      Update the estimated values in your MDP.
  
  
      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.
  
      Return true if it converges within one iteration.
      
      Args:
          mdp_data: The data for your MDP. See initialize_mdp_data.
          tolerance: The tolerance to use for the convergence criterion.
          gamma: Your discount factor.
  
      Returns:
          True if the value iteration converged in one iteration
  
      """
  
      # *** START CODE HERE ***
      action_values = (mdp_data['transition_probs'] * mdp_data['value'].reshape(1, -1, 1)).sum(axis=1)
      old_value = mdp_data['value']
      mdp_data['value'] = mdp_data['reward'] + gamma * action_values.max(axis=1)
  
      max_diff = np.max(np.abs(mdp_data['value'] - old_value))
      print(f'max state value diff: {max_diff}')
      return max_diff < tolerance
      # *** END CODE HERE ***
  
  def plot_mdp_data(mdp_data):
      plt.figure()
      plt.plot(mdp_data['value'])
  
  def main(plot=True):
      # Seed the randomness of the simulation so this outputs the same thing each time
      np.random.seed(3)
  
      # Simulation parameters
      pause_time = 0.0001
      min_trial_length_to_start_display = 100
      display_started = min_trial_length_to_start_display == 0
  
      NUM_STATES = 163
      GAMMA = 0.995
      TOLERANCE = 0.01
      NO_LEARNING_THRESHOLD = 20
      # TOTAL_MAX_TRIALS might be useful otherwise it takes up to 300 ish iterations to converge sometimes
  
      # 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
  
      # You should reach convergence well before this
      max_failures = 500
  
      # Initialize a cart pole
      cart_pole = CartPole(Physics())
  
      # Starting `state_tuple` is (0, 0, 0, 0)
      # x, x_dot, theta, theta_dot represents the actual continuous state vector
      x, x_dot, theta, theta_dot = 0.0, 0.0, 0.0, 0.0
      state_tuple = (x, x_dot, theta, theta_dot)
  
      # `state` is the number given to this state, you only need to consider
      # this representation of the state
      state = cart_pole.get_state(state_tuple)
      if min_trial_length_to_start_display == 0 or display_started == 1:
          cart_pole.show_cart(state_tuple, pause_time)
  
      mdp_data = initialize_mdp_data(NUM_STATES)
  
      # 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.
  
  
      consecutive_no_learning_trials = 0
      while consecutive_no_learning_trials < NO_LEARNING_THRESHOLD:
  
          action = choose_action(state, mdp_data)
  
          # Get the next state by simulating the dynamics
          state_tuple = cart_pole.simulate(action, state_tuple)
          # x, x_dot, theta, theta_dot = state_tuple
  
          # Increment simulation time
          time = time + 1
  
          # Get the state number corresponding to new state vector
          new_state = cart_pole.get_state(state_tuple)
          # if display_started == 1:
          #     cart_pole.show_cart(state_tuple, pause_time)
  
          #print(f'state transition prob: {state}, {action} -> {new_state}: {mdp_data["transition_probs"][state, new_state, action]}')
  
          # reward function to use - do not change this!
          if new_state == NUM_STATES - 1:
              R = -1
          else:
              R = 0
  
          # add to transition count for s, s', a triple and to reward count for s'
          update_mdp_transition_counts_reward_counts(mdp_data, 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 - 1:
  
              update_mdp_transition_probs_reward(mdp_data)
  
              converged_in_one_iteration = update_mdp_value(mdp_data, TOLERANCE, GAMMA)
  
              if converged_in_one_iteration:
                  consecutive_no_learning_trials = consecutive_no_learning_trials + 1
              else:
                  consecutive_no_learning_trials = 0
  
          # Do NOT 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 - 1:
              num_failures += 1
              if num_failures >= max_failures:
                  break
              print('[INFO] Failure number {}'.format(num_failures))
              # plot_mdp_data(mdp_data)
              time_steps_to_failure.append(time - time_at_start_of_current_trial)
              print(f'time to failure: {time_steps_to_failure[-1]}')
              # print(f'history: {mdp_data["state_action_history"][-1]}')
              mdp_data["state_action_history"].append([])
              # time_steps_to_failure[num_failures] = time - time_at_start_of_current_trial
              time_at_start_of_current_trial = time
  
              if time_steps_to_failure[num_failures - 1] > min_trial_length_to_start_display:
                  display_started = 1
  
              # Reinitialize state
              # x = 0.0
              x = -1.1 + np.random.uniform() * 2.2
              x_dot, theta, theta_dot = 0.0, 0.0, 0.0
              state_tuple = (x, x_dot, theta, theta_dot)
              state = cart_pole.get_state(state_tuple)
          else:
              state = new_state
  
      if plot:
          plt.figure()
          # plot the learning curve (time balanced vs. trial)
          log_tstf = np.log(np.array(time_steps_to_failure))
          plt.plot(np.arange(len(time_steps_to_failure)), log_tstf, 'k')
          window = 30
          w = np.array([1/window for _ in range(window)])
          weights = lfilter(w, 1, log_tstf)
          x = np.arange(window//2, len(log_tstf) - window//2)
          plt.plot(x, weights[window:len(log_tstf)], 'r--')
          plt.xlabel('Num failures')
          plt.ylabel('Log of num steps to failure')
          plt.savefig('./control_seed3.pdf')
  
      return np.array(time_steps_to_failure)
      
  if __name__ == '__main__':
      main()
  
apis: numpy.random.uniform numpy.ones numpy.tile numpy.zeros matplotlib.pyplot.figure matplotlib.pyplot.savefig numpy.random.seed numpy.abs scipy.signal.lfilter matplotlib.pyplot.ylabel numpy.random.rand matplotlib.use matplotlib.pyplot.plot numpy.where numpy.random.randint numpy.array matplotlib.pyplot.xlabel