# ---
# 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="view-in-github"
# <a href="https://colab.research.google.com/github/abhisheksuran/Atari_DQN/blob/master/A2C.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>

# + colab={} colab_type="code" id="udtLxTQyiNUz"
import numpy as np
import tensorflow as tf 
import gym
import tensorflow_probability as tfp
import tensorflow.keras.losses as kls

# + colab={"base_uri": "https://localhost:8080/", "height": 102} colab_type="code" id="nfKWiIG6m5Ed" outputId="e6b2bcea-5ddd-49b1-e8f2-306cb41e42fb"

# !pip3 install box2d-py

# + colab={} colab_type="code" id="4X3S4uD_m-e_"
env= gym.make("CartPole-v0")
low = env.observation_space.low
high = env.observation_space.high


# + colab={} colab_type="code" id="pAZ_YRduimE0"
class critic(tf.keras.Model):
  def __init__(self):
    super().__init__()
    self.d1 = tf.keras.layers.Dense(128,activation='relu')
    #self.d2 = tf.keras.layers.Dense(32,activation='relu')
    self.v = tf.keras.layers.Dense(1, activation = None)

  def call(self, input_data):
    x = self.d1(input_data)
    #x = self.d2(x)
    v = self.v(x)
    return v
    

class actor(tf.keras.Model):
  def __init__(self):
    super().__init__()
    self.d1 = tf.keras.layers.Dense(128,activation='relu')
    #self.d2 = tf.keras.layers.Dense(32,activation='relu')
    self.a = tf.keras.layers.Dense(2,activation='softmax')

  def call(self, input_data):
    x = self.d1(input_data)
    #x = self.d2(x)
    a = self.a(x)
    return a


# + colab={} colab_type="code" id="D0-5_oxXmwT9"
class agent():
    def __init__(self, gamma = 0.99):
        self.gamma = gamma
        self.a_opt = tf.keras.optimizers.RMSprop(learning_rate=7e-3)
        self.c_opt = tf.keras.optimizers.RMSprop(learning_rate=7e-3)
        self.actor = actor()
        self.critic = critic()

          
    def act(self,state):
        prob = self.actor(np.array([state]))
        prob = prob.numpy()
        dist = tfp.distributions.Categorical(probs=prob, dtype=tf.float32)
        action = dist.sample()
        return int(action.numpy()[0])
  


    def actor_loss(self, probs, actions, td):
        
        probability = []
        log_probability= []
        for pb,a in zip(probs,actions):
          dist = tfp.distributions.Categorical(probs=pb, dtype=tf.float32)
          log_prob = dist.log_prob(a)
          prob = dist.prob(a)
          probability.append(prob)
          log_probability.append(log_prob)

        # print(probability)
        # print(log_probability)

        p_loss= []
        e_loss = []
        td = td.numpy()
        #print(td)
        for pb, t, lpb in zip(probability, td, log_probability):
                        t =  tf.constant(t)
                        policy_loss = tf.math.multiply(lpb,t)
                        entropy_loss = tf.math.negative(tf.math.multiply(pb,lpb))
                        p_loss.append(policy_loss)
                        e_loss.append(entropy_loss)
        p_loss = tf.stack(p_loss)
        e_loss = tf.stack(e_loss)
        p_loss = tf.reduce_mean(p_loss)
        e_loss = tf.reduce_mean(e_loss)
        # print(p_loss)
        # print(e_loss)
        loss = -p_loss - 0.0001 * e_loss
        #print(loss)
        return loss

    def learn(self, states, actions, discnt_rewards):
        discnt_rewards = tf.reshape(discnt_rewards, (len(discnt_rewards),))
        
        with tf.GradientTape() as tape1, tf.GradientTape() as tape2:
            p = self.actor(states, training=True)
            v =  self.critic(states,training=True)
            v = tf.reshape(v, (len(v),))
            td = tf.math.subtract(discnt_rewards, v)
            # print(discnt_rewards)
            # print(v)
            #print(td.numpy())
            a_loss = self.actor_loss(p, actions, td)
            c_loss = 0.5*kls.mean_squared_error(discnt_rewards, v)
        grads1 = tape1.gradient(a_loss, self.actor.trainable_variables)
        grads2 = tape2.gradient(c_loss, self.critic.trainable_variables)
        self.a_opt.apply_gradients(zip(grads1, self.actor.trainable_variables))
        self.c_opt.apply_gradients(zip(grads2, self.critic.trainable_variables))
        return a_loss, c_loss


# + colab={"base_uri": "https://localhost:8080/", "height": 1000} colab_type="code" id="vk4Phsi1tvS-" outputId="8849de20-a09f-44c1-ac77-550b47d3400b"
def preprocess1(states, actions, rewards, gamma):
    discnt_rewards = []
    sum_reward = 0
    rewards.reverse()
    for r in rewards:
      sum_reward = r + gamma*sum_reward
      discnt_rewards.append(sum_reward)
    discnt_rewards.reverse()
    states = np.array(states, dtype=np.float32)
    actions = np.array(actions, dtype=np.int32)
    discnt_rewards = np.array(discnt_rewards, dtype=np.float32)

    return states, actions, discnt_rewards

tf.random.set_seed(336699)
agentoo7 = agent()
steps = 250
ep_reward = []
total_avgr = []
for s in range(steps):
  
  done = False
  state = env.reset()
  total_reward = 0
  all_aloss = []
  all_closs = []
  rewards = []
  states = []
  actions = []
  
  while not done:
  
    action = agentoo7.act(state)
    next_state, reward, done, _ = env.step(action)
    rewards.append(reward)
    states.append(state)
    #actions.append(tf.one_hot(action, 2, dtype=tf.int32).numpy().tolist())
    actions.append(action)
    state = next_state
    total_reward += reward
    
    if done:
        ep_reward.append(total_reward)
        avg_reward = np.mean(ep_reward[-100:])
        total_avgr.append(avg_reward)
        print("total reward after {} steps is {} and avg reward is {}".format(s, total_reward, avg_reward))
        states, actions, discnt_rewards = preprocess1(states, actions, rewards, 1)
  
        al,cl = agentoo7.learn(states, actions, discnt_rewards) 
        print(f"al{al}") 
        print(f"cl{cl}")      
      

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

# + colab={"base_uri": "https://localhost:8080/", "height": 295} colab_type="code" id="pquOE2v0Snsh" outputId="4862be43-49d5-4799-d43c-ddfd0247f5a6"
ep = [i  for i in range(250)]
plt.plot(ep,total_avgr,'b')
plt.title("avg reward Vs episods")
plt.xlabel("episods")
plt.ylabel("average reward per 100 episods")
plt.grid(True)
plt.show()




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


StarcoderdataJupyterScriptsDedupFiltered