# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     name: python3
# ---

# + [markdown] id="2BxBk84bp4wC" colab_type="text"
# # Crossentropy method
#
# In this notebook, we will imprement crossentropy method to solve the ['Taxi' problem in Open AI Gym](https://gym.openai.com/envs/Taxi-v2/)
#
# First, we have to make sure we are connected to the right **python 3 reutime and using the GPU**. (Click the 'Runtime' tab and choose 'Change runtime type'), then import the required package (all are already installed in Google Colab)

# + id="w4ucVW4Qp4wD" colab_type="code" colab={}
import gym
import numpy as np, pandas as pd

# + [markdown] id="-7adMV64Et-r" colab_type="text"
# Then we create the 'Taxi' enviroment from the Gym. Render the initial state to check it is imported properly

# + id="IPiSeYf5Erqi" colab_type="code" colab={}
env = gym.make("Taxi-v2")
env.reset()
env.render()

# + [markdown] id="Tl3wtYs_GIje" colab_type="text"
# Find out how many states and action we can have in this enviroment.

# + id="6kcwaSI7p4wI" colab_type="code" colab={}
n_states = env.observation_space.n
n_actions = env.action_space.n

print("n_states=%i, n_actions=%i"%(n_states, n_actions))

# + [markdown] id="DogESWKmp4wL" colab_type="text"
# ## Task 1: Create stochastic policy
#
# Let's create a policy, for crossentrupy method with stochastic policy (updated each epoch), it will be a probability distribution of taking action a will result in state s.
#
# ```policy[s,a] = P(take action a | in state s)```
#
# It could be represented in a 2-D array, and we want to initialize it __uniformly__. In another word, at the beginning state,probability of choosing all actions should be equal (and adding up to 1).
#
# With `n_state` and `n_action`:

# + id="79LDIEUmp4wM" colab_type="code" colab={}
policy = np.ones((n_states,n_actions)) * (1/n_actions)

# + id="a-RRIPR3p4wO" colab_type="code" colab={}
assert type(policy) in (np.ndarray,np.matrix)
assert np.allclose(policy,1./n_actions)
assert np.allclose(np.sum(policy,axis=1), 1)


# + [markdown] id="VtB4cVZJp4wQ" colab_type="text"
# ## Task 2: Play the game
#
# Let's play the game with our policy. The following function will 'play the game' i.e. picking actions according to the state we are in and following the policy given. It will keep on 'playing' till the game is over or the step (t) reached the maximum limit. (To avoid endless loop) The action and states will be reconded and returned. Also the performance of the policy (sum of rewards) will also be returned.

# + id="PQdj_dHLp4wQ" colab_type="code" colab={}
def generate_session(policy,t_max=10**4):
    """
    Play game until end or for t_max ticks.
    :param policy: an array of shape [n_states,n_actions] with action probabilities
    :returns: list of states, list of actions and sum of rewards
    """
    states,actions = [],[]
    total_reward = 0.
    
    s = env.reset()
    
    for t in range(t_max):
        
        a = np.random.choice(n_actions,1,p=policy[s,:])[0]
        
        new_s, r, done, info = env.step(a)
        
        #Record state, action and add up reward to states,actions and total_reward accordingly. 
        states.append(s)
        actions.append(a)
        total_reward += r
        
        s = new_s
        if done:
            break
    return states, actions, total_reward
        


# + id="-OAc-QDyp4wS" colab_type="code" colab={}
s,a,r = generate_session(policy)
assert type(s) == type(a) == list
assert len(s) == len(a)
assert type(r) in [float,np.float]

# + [markdown] id="FLet-FrI-E1C" colab_type="text"
# ## Let's see the initial reward distribution
#
# In the following cell we play the game 200 times using the function `generate_session` we implemented above and plot to examine the performance of our policy...

# + id="vCqRGN36p4wU" colab_type="code" colab={}
import matplotlib.pyplot as plt
# %matplotlib inline

sample_rewards = [generate_session(policy,t_max=1000)[-1] for _ in range(200)]

plt.hist(sample_rewards,bins=20);
plt.vlines([np.percentile(sample_rewards, 50)], [0], [100], label="50'th percentile", color='green')
plt.vlines([np.percentile(sample_rewards, 90)], [0], [100], label="90'th percentile", color='red')
plt.legend()


# + [markdown] id="ZYBfD-5t-tS6" colab_type="text"
# You can see that it is not that great :-( So we now use crossentropy method to improve 'train' a better policy.

# + [markdown] id="VgEK8Fsnp4wX" colab_type="text"
# ## Task 3: Crossentropy method steps
#
# The goal is to select the state-action pairs which perform better than a certain percentile (elite_states, elites_actions). So we know what is a better choice actions in a certain state. You can say it's trail and error, and yes, that is the basic of reinforcement learning.

# + id="c0jacA8Ep4wY" colab_type="code" colab={}
def select_elites(states_batch,actions_batch,rewards_batch,percentile=50):
    """
    Select states and actions from games that have rewards >= percentile
    :param states_batch: list of lists of states, states_batch[session_i][t]
    :param actions_batch: list of lists of actions, actions_batch[session_i][t]
    :param rewards_batch: list of rewards, rewards_batch[session_i]
    
    :returns: elite_states,elite_actions, both 1D lists of states and respective actions from elite sessions
    
    Please return elite states and actions in their original order 
    [i.e. sorted by session number and timestep within session]
    
    If you're confused, see examples below. Please don't assume that states are integers (they'll get different later).
    """
    
    def gen_list(input_list, output_list):
        for indx,val in enumerate(input_list):
            if rewards_batch[indx] >= reward_threshold:
                if isinstance(val, list):
                    for i in range(len(val)):
                        output_list.append(val[i])
                else:
                    output_list.append(val)
    
    reward_threshold = np.percentile(rewards_batch,percentile)
    
    elite_states  = []   
    elite_actions = []
    
    gen_list(states_batch,elite_states)
    gen_list(actions_batch,elite_actions)
    
    return elite_states,elite_actions
    


# + id="puefOHVFp4wa" colab_type="code" colab={}
states_batch = [
    [1,2,3],   #game1
    [4,2,0,2], #game2
    [3,1]      #game3
]

actions_batch = [
    [0,2,4],   #game1
    [3,2,0,1], #game2
    [3,3]      #game3
]
rewards_batch = [
    3,         #game1
    4,         #game2
    5,         #game3
]

test_result_0 = select_elites(states_batch, actions_batch, rewards_batch, percentile=0)
test_result_40 = select_elites(states_batch, actions_batch, rewards_batch, percentile=30)
test_result_90 = select_elites(states_batch, actions_batch, rewards_batch, percentile=90)
test_result_100 = select_elites(states_batch, actions_batch, rewards_batch, percentile=100)

assert np.all(test_result_0[0] == [1, 2, 3, 4, 2, 0, 2, 3, 1])  \
   and np.all(test_result_0[1] == [0, 2, 4, 3, 2, 0, 1, 3, 3]),\
        "For percentile 0 you should return all states and actions in chronological order"
assert np.all(test_result_40[0] == [4, 2, 0, 2, 3, 1]) and \
        np.all(test_result_40[1] ==[3, 2, 0, 1, 3, 3]),\
        "For percentile 30 you should only select states/actions from two first"
assert np.all(test_result_90[0] == [3,1]) and \
        np.all(test_result_90[1] == [3,3]),\
        "For percentile 90 you should only select states/actions from one game"
assert np.all(test_result_100[0] == [3,1]) and\
       np.all(test_result_100[1] == [3,3]),\
        "Please make sure you use >=, not >. Also double-check how you compute percentile."
print("Ok!")


# + [markdown] id="Tm6nZp6cAQEd" colab_type="text"
# ## Task 4: Update policy
#
# Now we have the 'elites' list, we can update our policy so the probability of the state-action pairs are propotional to their occurance of in the 'elites' list. In other words, the state-action which seems to have a higher chance to perform better get selected more.

# + id="lRMzX-7Lp4we" colab_type="code" colab={}
def update_policy(elite_states,elite_actions):
    """
    Given old policy and a list of elite states/actions from select_elites,
    return new updated policy where each action probability is proportional to
    
    policy[s_i,a_i] ~ #[occurences of si and ai in elite states/actions]
    
    Don't forget to normalize policy to get valid probabilities and handle 0/0 case.
    In case you never visited a state, set probabilities for all actions to 1./n_actions
    
    :param elite_states: 1D list of states from elite sessions
    :param elite_actions: 1D list of actions from elite sessions
    
    """
    
    new_policy = np.zeros([n_states,n_actions])
    esp = 0.0000001
    
    for i,s in enumerate(elite_states):
        new_policy[s,elite_actions[i]] += 1
        
    for s in range(n_states):
        if all(new_policy[s,:] == 0):
            new_policy[s,:] = 1./n_actions
        else:
            new_policy[s,:] /= np.sum(new_policy[s,:])
            
    #Don't forget to set 1/n_actions for all actions in unvisited states.
    
    
    return new_policy

# + id="mumQNLK6p4wg" colab_type="code" colab={}

elite_states, elite_actions = ([1, 2, 3, 4, 2, 0, 2, 3, 1], [0, 2, 4, 3, 2, 0, 1, 3, 3])


new_policy = update_policy(elite_states,elite_actions)

assert np.isfinite(new_policy).all(), "Your new policy contains NaNs or +-inf. Make sure you don't divide by zero."
assert np.all(new_policy>=0), "Your new policy can't have negative action probabilities"
assert np.allclose(new_policy.sum(axis=-1),1), "Your new policy should be a valid probability distribution over actions"
reference_answer = np.array([
       [ 1.        ,  0.        ,  0.        ,  0.        ,  0.        ],
       [ 0.5       ,  0.        ,  0.        ,  0.5       ,  0.        ],
       [ 0.        ,  0.33333333,  0.66666667,  0.        ,  0.        ],
       [ 0.        ,  0.        ,  0.        ,  0.5       ,  0.5       ]])
assert np.allclose(new_policy[:4,:5],reference_answer)
print("Ok!")

# + [markdown] id="tCK1uSLTp4wj" colab_type="text"
# ## Let's train it
# Let's put all our work above together: Generate sessions, select N best and fit to those. (Double check you are using the GPU runtime type, training may take a while, be patient)

# + id="GJVx8Ifpp4wk" colab_type="code" colab={}
from IPython.display import clear_output

def show_progress(batch_rewards, log, percentile, reward_range=[-990,+10]):
    """
    A convenience function that displays training progress. 
    No cool math here, just charts.
    """
    
    mean_reward, threshold = np.mean(batch_rewards), np.percentile(batch_rewards, percentile)
    log.append([mean_reward,threshold])

    clear_output(True)
    print("mean reward = %.3f, threshold=%.3f"%(mean_reward, threshold))
    plt.figure(figsize=[8,4])
    plt.subplot(1,2,1)
    plt.plot(list(zip(*log))[0], label='Mean rewards')
    plt.plot(list(zip(*log))[1], label='Reward thresholds')
    plt.legend()
    plt.grid()
    
    plt.subplot(1,2,2)
    plt.hist(batch_rewards,range=reward_range);
    plt.vlines([np.percentile(batch_rewards, percentile)], [0], [100], label="percentile", color='red')
    plt.legend()
    plt.grid()

    plt.show()



# + id="hmoj_BaJp4wo" colab_type="code" colab={}
#reset policy just in case
policy = np.ones([n_states, n_actions]) / n_actions 

# + id="hT5piqEyp4wq" colab_type="code" colab={}
n_sessions = 250  #sample this many sessions
percentile = 50  #take this percent of session with highest rewards
learning_rate = 0.5  #add this thing to all counts for stability

log = []

for i in range(100):
    
    # %time sessions = [generate_session(policy) for i in range(n_sessions)]
    
    batch_states,batch_actions,batch_rewards = zip(*sessions)

    elite_states, elite_actions = select_elites(batch_states,batch_actions,batch_rewards,percentile)
    
    new_policy = update_policy(elite_states,elite_actions)
    
    policy = learning_rate * new_policy + (1-learning_rate) * policy
    
    #display results on chart
    show_progress(batch_rewards, log, percentile)
    if np.mean(batch_rewards) > -20:
        break

# + [markdown] id="ZmQcJYjAp4wt" colab_type="text"
# ## Footnote: Reflecting on results
#
# You may have noticed that the taxi problem quickly converges from <-1000 to a near-optimal score and then descends back into -50/-100. This is in part because the environment has some innate randomness. Namely, the starting points of passenger/driver change from episode to episode.
#
# In case CEM failed to learn how to win from one distinct starting point, it will siply discard it because no sessions from that starting point will make it into the "elites".
#
# To mitigate that problem, you can either reduce the threshold for elite sessions (duct tape way) or  change the way you evaluate strategy (theoretically correct way). You can first sample an action for every possible state and then evaluate this choice of actions by running _several_ games and averaging rewards.

StarcoderdataJupyterScriptsDedupFiltered