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

# #### The project uses reinforcement learning to balance the frictionless cart on the pole by moving the cart according to the position of the pole. The model that will be used is the OpenAI gym. I will be using Neural network for training.

# ### loading Libraries

import sys
assert sys.version_info >=(3,5)
import gym
gym.envs.registry.all()
from tensorflow import keras

import numpy as np
import tensorflow as tf
from tensorflow import keras
assert tf.__version__ >= "2.0"
import sklearn
assert sklearn.__version__ >= "0.20"

# #### setting seed 

np.random.seed(42)
tf.random.set_seed(42)
env = gym.make("CartPole-v1")
env.seed(42)

import matplotlib.pyplot as plt
import matplotlib as mpl

# %matplotlib inline
mpl.rc('axes', labelsize=14)
mpl.rc('xtick', labelsize=12)
mpl.rc('ytick', labelsize=12)

# ###### environmetn

gym.envs.registry.all()
# env = gym.make("Acrobot-v1")

# +
# gym.envs.registry.all()
# -

# #### reset method to return an observation to start over again

obs = env.reset()
print(obs)

# #### attempting the piece of the game action

print(env.action_space)

action = 1
obs, reward, done, info = env.step(action)

print(obs)
print(reward)
print(done)
print(info)

# ### Hardcoding the strategy to see how it will work

# +
env.seed(42)

def basic_policy(obs):
    angle = obs[2]
    if angle < 0:
        return 0
    else:
        return 1


# -

# ###### Now let us play 500 episodes of the game, each episode with 200 steps. For each step we call the basic_policy to get the action, and perform that step with that action. Let us calculate the rewards for each episode and finally see how many minimum, maximum, and mean steps that our basic_policy is able to keep the pole up:
#
#

totals = []
for episode in range(500):
    episode_rewards = 0
    obs = env.reset()
    for step in range(200):
        action = basic_policy(obs)
        obs, reward, done, info = env.step(action)
        episode_rewards += reward
        if done:
            break
    totals.append(episode_rewards)

# ##### alculating the mean, start deviation, minimum, and maximum number of steps that the pole was upright.

print(np.mean(totals), np.std(totals), np.min(totals), np.max(totals))

# ##### Building with Neural Network

# #### clearing session to removes all nodes left from previous models

keras.backend.clear_session()

# #### setting input shape to 4, since we are feeding all values of obs

n_inputs = 4

# #### Building the neural network

model = tf.keras.Sequential([
    keras.layers.Dense(5, activation="elu", input_shape=[n_inputs]),
    keras.layers.Dense(1, activation="sigmoid"),
])
                

model.summary()


# #### function for the untrain network policy
# #### we shall set randomly thrashold to allow the  network decide the set of actions to take that will be more better in achieveing it purpose

def basic_policy_untrained(obs):
    left_proba = model.predict(obs.reshape(1, -1))
    action = int(np.random.rand() > left_proba)
    return action


# #### Runing this for 50 episodes and 200 steps. we shall call the basic_policy_untrained for each step of each episode:

totals = []
for episode in range(50):
    episode_reward = 0
    obs = env.reset()
    for step in range(200):
        action = basic_policy_untrained(obs)
        obs, reward, done, info = env.step(action)
        episode_rewards += reward
        if done:
            break
    totals.append(episode_rewards)
np.mean(totals), np.std(totals), np.min(totals), np.max(totals)

# ##### from the above scores we can see that neural network perform badly

# #### Training Neural Network for better performance 

# ##### setting the number of environment to 50

n_environments = 50

# #### setting numbers of iteration to 5000

n_iterations = 5000

# ###### initializing cartpole environmen

envs = [gym.make("CartPole-v1") for _ in range(n_environments)]

# ##### Setting different seeds to each environment with their respective indices as per the above list

for index, env in enumerate(envs):
    env.seed(index)

# #### Reset all the environments to get the observation for all the environment

observations = [env.reset() for env in envs]

# ##### Initialize the optimizer to keras optimization

optimizer = keras.optimizers.RMSprop()

# #### initializing loss function

loss_fn = keras.losses.binary_crossentropy

# #### Setting the target probabilities of the action. If angle < 0, we want proba(left) = 1., or else proba(left) = 0. Then, we shall fit the observations to the model and calculate the gradients based on the loss with respect to the targets. We shall use the probabilities predicted by the model and determined the actions and proceed for the next step, and the iterations continue.

# +

for iteration in range(n_iterations):
    # if angle < 0, we want proba(left) = 1., or else proba(left) = 0.
    target_probas = np.array([([1.] if obs[2] < 0 else [0.])
                              for obs in observations])

    with tf.GradientTape() as tape:
        left_probas = model(np.array(observations))
        loss = tf.reduce_mean(loss_fn(target_probas, left_probas))
    print("\rIteration: {}, Loss: {:.3f}".format(iteration, loss.numpy()), end="")
    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))

    actions = (np.random.rand(n_environments, 1) > left_probas.numpy()).astype(np.int32)
    for env_index, env in enumerate(envs):
        obs, reward, done, info = env.step(actions[env_index][0])
        observations[env_index] = obs if not done else env.reset()


# -

# ##### from the above score you can see that it performance much more beteer

# #### creating function to play a single step using the model

def play_one_step(env, obs, model, loss_fn):
    with tf.GradientTape() as tape:
        left_probas = model(obs[np.newaxis])
        action = (tf.random.uniform([1, 1]) > left_probas)
        y_target = tf.constant([[1.]]) - tf.cast(action, tf.float32)
        loss = tf.reduce_mean(loss_fn(y_target, left_probas))
    grads = tape.gradient(loss, model.trainable_variables)
    obs, reward, done, info = env.step(int(action[0,0].numpy()))
    return obs, reward, done, grads


# ##### function that will return all rewards

def play_multiple_episodes(env, n_episodes, n_max_steps, model, loss_fn):
    all_rewards = []
    all_grads = []
    for episode in range(n_episodes):
        current_rewards = []
        current_grads = []
        obs = env.reset()
        for step in range(n_max_steps):
            obs, reward, done, grads = play_one_step(env, obs, model,loss_fn)
            current_rewards.append(reward)
            current_grads.append(grads)
            if done:
                break
            all_rewards.append(current_rewards)
            all_grads.append(current_grads)
        return all_rewards, all_grads


# #### function to define the dicount reward

def discount_rewards(rewards, discount_rate):
    discounted = np.array(rewards)
    for step in range(len(rewards) - 2, -1, -1):
        discounted[step] += discounted[step + 1] * discount_rate
    return discounted


def discount_and_normalize_rewards(all_rewards, discount_rate):
    all_discounted_rewards = [discount_rewards(rewards, discount_rate)
                            for rewards in all_rewards]
    flat_rewards = np.concatenate(all_discounted_rewards)
    reward_mean = flat_rewards.mean()
    reward_std = flat_rewards.std()
    return [(discounted_rewards - reward_mean) / reward_std
            for discounted_rewards in all_discounted_rewards]


# ##### Train with policy Gradients

# ##### seeting the number of environments

# +
tf.random.set_seed(42)
np.random.seed(42)

n_iterations = 150
n_episodes_per_update = 10
n_max_steps = 200
discount_rate = 0.95

optimizer = keras.optimizers.Adam(lr=0.01)
loss_fn = keras.losses.binary_crossentropy

n_inputs = 4 


# +
# n_iteration = 150

# +
# n_max_steps = 200
# -

# #### setting discount rate 

# +
# discount_rate = 0.95

# +
# n_inputs = 4

# +
# loss_fn = keras.losses.binary_crossentropy

# +
# optimizer = keras.optimizers.Adam(lr=0.01)
# -

# ##### applying gradient algorithm while training the neural network

def nn_policy_gradient(model, n_iterations, n_episodes_per_update, n_max_steps, loss_fn):
    env = gym.make("CartPole-v1")
    env.seed(42);

    for iteration in range(n_iterations):
        all_rewards, all_grads = play_multiple_episodes(
            env, n_episodes_per_update, n_max_steps, model, loss_fn)
        total_rewards = sum(map(sum, all_rewards))                     # Not shown in the book
        print("\rIteration: {}, mean rewards: {:.1f}".format(          # Not shown
            iteration, total_rewards / n_episodes_per_update), end="") # Not shown
        all_final_rewards = discount_and_normalize_rewards(all_rewards,
                                                        discount_rate)
        all_mean_grads = []
        for var_index in range(len(model.trainable_variables)):
            mean_grads = tf.reduce_mean(
                [final_reward * all_grads[episode_index][step][var_index]
                for episode_index, final_rewards in enumerate(all_final_rewards)
                    for step, final_reward in enumerate(final_rewards)], axis=0)
            all_mean_grads.append(mean_grads)
        optimizer.apply_gradients(zip(all_mean_grads, model.trainable_variables))

    return model

    env.close()



# #### models

model = keras.models.Sequential([
    keras.layers.Dense(5, activation="elu", input_shape=[n_inputs]),
    keras.layers.Dense(1, activation="sigmoid"),
])

# ##### calling the function

model = nn_policy_gradient(model, n_iterations, n_episodes_per_update, n_max_steps, loss_fn)


# +
totals = []
for episode in range(20):
    print("Episode:",episode)
    episode_rewards = 0
    obs = env.reset()
    for step in range(200):
        action = basic_policy_untrained(obs)
        obs, reward, done, info = env.step(action)
        episode_rewards += reward
        if done:
            break
    totals.append(episode_rewards)

np.mean(totals), np.std(totals), np.min(totals), np.max(totals)
# -



StarcoderdataJupyterScriptsDedupFiltered