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

# + [markdown] colab_type="text" id="Wt1OjEjJz8hl"
# Thanks to https://medium.com/coinmonks/landing-a-rocket-with-simple-reinforcement-learning-3a0265f8b58c
#
#
# CEM Building Blocks
# To implement the deep cross entropy method we need to follow 4 steps
#
# 1. Generate Sessions:
# Play through several episodes of the game environment with our current agent and save the actions, states and rewards used for each episode
#
# 2. Retrieve Elite Sessions:
# We want to only learn off episodes that achieved a high score in that batch. We determine our elite threshold by taking some percentile of all the episode rewards for that batch of generated sessions.
#
# 3. Train On Elite Sessions
# Now that we have the top X% of our batch, we train our model on those experiences. We use the recorded states as input and the actions carried out as the target output.
#
# 4. Rinse and Repeat
# We keep repeating this process until our model converges on a successfull policy
#
# Now that we have our blueprint of how to build the agent the last thing we need to cover is applying deep learning to the CEM.

# + colab={"base_uri": "https://localhost:8080/", "height": 187} colab_type="code" id="kpIE1pkYLaE5" outputId="6dc87ac0-c53e-415a-cb53-e8c2e343fc59"
import gym
import numpy as np
import torch 
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

class Net(nn.Module):
    def __init__(self,obs_size,hidden_size,n_actions):
        super(Net,self).__init__()
        self.fc1 = nn.Linear(obs_size,hidden_size)
        self.fc2 = nn.Linear(hidden_size,n_actions)
    def forward(self,x):
        x = F.relu(self.fc1(x))
        return self.fc2(x)


def generate_batch(env,batch_size,t_max = 1000):
    activation = nn.Softmax(dim=1)
    batch_actions,batch_states,batch_rewards = [],[],[]
    for b in range(batch_size):
        states,actions=[],[]   
        total_reward = 0
        s = env.reset()
        for t in range(t_max):
            s_v = torch.FloatTensor([s])
            act_probe_v = activation(net(s_v))
            act_probs = act_probe_v.data.numpy()[0]
            a = np.random.choice(len(act_probs),p=act_probs)

            new_s,r,done,info = env.step(a)

            states.append(s)
            actions.append(a)
            total_reward += r

            s = new_s
            if done:
                batch_actions.append(actions)
                batch_states.append(states)
                batch_rewards.append(total_reward)
                break
    return batch_states,batch_actions,batch_rewards

def filter_batch(states_batch,actions_batch,rewards_batch,percentile=50):
    reward_threshold = np.percentile(rewards_batch,percentile)
    elite_states = []
    elite_actions = []
    for i in range(len(rewards_batch)):
        if rewards_batch[i] > reward_threshold:
            for j in range(len(states_batch[i])):
                elite_states.append(states_batch[i][j])
                elite_actions.append(actions_batch[i][j])
    return elite_states,elite_actions

batch_size = 100

#Bununla oynayabilirsiniz
session_size = 100
percentile = 80
hidden_size = 200
learning_rate = 0.0025
completion_score = 200
env = gym.make("LunarLander-v2")
n_states = env.observation_space.shape[0]
n_actions = env.action_space.n
#neural network
net = Net(n_states, hidden_size, n_actions)
#loss function
objective = nn.CrossEntropyLoss()
#optimisation function
optimizer = optim.Adam(params=net.parameters(), lr=learning_rate)
for i in range(session_size):
    #generate new sessions
    batch_states,batch_actions,batch_rewards = generate_batch(env, batch_size, t_max=5000)
    elite_states, elite_actions = filter_batch(batch_states,batch_actions,batch_rewards,percentile)
    optimizer.zero_grad()
    tensor_states = torch.FloatTensor(elite_states)
    tensor_actions = torch.LongTensor(elite_actions)
    action_scores_v = net(tensor_states)
    loss_v = objective(action_scores_v, tensor_actions)
    loss_v.backward()
    optimizer.step()
    #show results
    mean_reward, threshold = np.mean(batch_rewards),np.percentile(batch_rewards, percentile)
    print("%d: loss=%.3f, reward_mean=%.1f,reward_threshold=%.1f"% (i, loss_v.item(), mean_reward, threshold))
    
    #check if 
    if np.mean(batch_rewards)> completion_score:
        print("Environment has been successfullly completed!")

# + colab={} colab_type="code" id="wokqoRisLZ80"
#Watch the intelligent agent.
import gym.wrappers
env = gym.wrappers.Monitor(gym.make("LunarLander-v2"), directory="videos", force=True)
generate_batch(env, 1, t_max=5000)
env.close()

StarcoderdataJupyterScriptsDedupFiltered