{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "import matplotlib.pyplot as plt\n",
    "import gym"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = gym.make(\"CartPole-v1\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<>:74: DeprecationWarning: invalid escape sequence \\g\n",
      "<>:93: DeprecationWarning: invalid escape sequence \\d\n"
     ]
    }
   ],
   "source": [
    "class REINFORCE:\n",
    "  def __init__(self, env, path=None):\n",
    "    self.env=env #import env\n",
    "    self.state_shape=env.observation_space.shape # the state space\n",
    "    self.action_shape=env.action_space.n # the action space\n",
    "    self.gamma=0.99 # decay rate of past observations\n",
    "    self.alpha=1e-4 # learning rate in the policy gradient\n",
    "    self.learning_rate=0.01 # learning rate in deep learning\n",
    "    \n",
    "    if not path:\n",
    "      self.model=self._create_model() #build model\n",
    "    else:\n",
    "      self.model=self.load_model(path) #import model\n",
    "\n",
    "    # record observations\n",
    "    self.states=[]\n",
    "    self.gradients=[] \n",
    "    self.rewards=[]\n",
    "    self.probs=[]\n",
    "    self.discounted_rewards=[]\n",
    "    self.total_rewards=[]\n",
    "\n",
    "  def hot_encode_action(self, action):\n",
    "    '''encoding the actions into a binary list'''\n",
    "\n",
    "    action_encoded=np.zeros(self.action_shape, np.float32)\n",
    "    action_encoded[action]=1\n",
    "\n",
    "    return action_encoded\n",
    "  \n",
    "  def remember(self, state, action, action_prob, reward):\n",
    "    '''stores observations'''\n",
    "    encoded_action=self.hot_encode_action(action)\n",
    "    self.gradients.append(encoded_action-action_prob)\n",
    "    self.states.append(state)\n",
    "    self.rewards.append(reward)\n",
    "    self.probs.append(action_prob)\n",
    "\n",
    "  def _create_model(self):\n",
    "    ''' builds the model using keras'''\n",
    "    model=keras.Sequential()\n",
    "\n",
    "    # input shape is of observations\n",
    "    model.add(keras.layers.Dense(24, input_shape=self.state_shape, activation=\"relu\"))\n",
    "    # add a relu layer \n",
    "    model.add(keras.layers.Dense(12, activation=\"relu\"))\n",
    "\n",
    "    # output shape is according to the number of action\n",
    "    # The softmax function outputs a probability distribution over the actions\n",
    "    model.add(keras.layers.Dense(self.action_shape, activation=\"softmax\")) \n",
    "    model.compile(loss=\"categorical_crossentropy\",\n",
    "            optimizer=keras.optimizers.Adam(lr=self.learning_rate))\n",
    "        \n",
    "    return model\n",
    "\n",
    "  def get_action(self, state):\n",
    "    '''samples the next action based on the policy probabilty distribution \n",
    "      of the actions'''\n",
    "\n",
    "    # transform state\n",
    "    state=state.reshape([1, state.shape[0]])\n",
    "    # get action probably\n",
    "    action_probability_distribution=self.model.predict(state).flatten()\n",
    "    # norm action probability distribution\n",
    "    action_probability_distribution/=np.sum(action_probability_distribution)\n",
    "    \n",
    "    # sample action\n",
    "    action=np.random.choice(self.action_shape,1,\n",
    "                            p=action_probability_distribution)[0]\n",
    "\n",
    "    return action, action_probability_distribution\n",
    "\n",
    "  def get_discounted_rewards(self, rewards): \n",
    "    '''Use gamma to calculate the total reward discounting for rewards\n",
    "    Following - \\gamma ^ t * Gt'''\n",
    "    \n",
    "    discounted_rewards=[]\n",
    "    cumulative_total_return=0\n",
    "    # iterate the rewards backwards and and calc the total return \n",
    "    for reward in rewards[::-1]:      \n",
    "      cumulative_total_return=(cumulative_total_return*self.gamma)+reward\n",
    "      discounted_rewards.insert(0, cumulative_total_return)\n",
    "\n",
    "    # normalize discounted rewards\n",
    "    mean_rewards=np.mean(discounted_rewards)\n",
    "    std_rewards=np.std(discounted_rewards)\n",
    "    norm_discounted_rewards=(discounted_rewards-\n",
    "                          mean_rewards)/(std_rewards+1e-7) # avoiding zero div\n",
    "    \n",
    "    return norm_discounted_rewards\n",
    "\n",
    "  def update_policy(self):\n",
    "    '''Updates the policy network using the NN model.\n",
    "    This function is used after the MC sampling is done - following\n",
    "    \\delta \\theta = \\alpha * gradient + log pi'''\n",
    "      \n",
    "    # get X\n",
    "    states=np.vstack(self.states)\n",
    "\n",
    "    # get Y\n",
    "    gradients=np.vstack(self.gradients)\n",
    "    rewards=np.vstack(self.rewards)\n",
    "    discounted_rewards=self.get_discounted_rewards(rewards)\n",
    "    gradients*=discounted_rewards\n",
    "    gradients=self.alpha*np.vstack([gradients])+self.probs\n",
    "\n",
    "    history=self.model.train_on_batch(states, gradients)\n",
    "    \n",
    "    self.states, self.probs, self.gradients, self.rewards=[], [], [], []\n",
    "\n",
    "    return history\n",
    "\n",
    "  def train(self, episodes, rollout_n=1, render_n=50):\n",
    "    '''train the model\n",
    "        episodes - number of training iterations \n",
    "        rollout_n- number of episodes between policy update\n",
    "        render_n - number of episodes between env rendering ''' \n",
    "    \n",
    "    env=self.env\n",
    "    total_rewards=np.zeros(episodes)\n",
    "\n",
    "    for episode in range(episodes):\n",
    "      # each episode is a new game env\n",
    "      state=env.reset()\n",
    "      done=False          \n",
    "      episode_reward=0 #record episode reward\n",
    "      \n",
    "      while not done:\n",
    "        # play an action and record the game state & reward per episode\n",
    "        action, prob=self.get_action(state)\n",
    "        next_state, reward, done, _=env.step(action)\n",
    "        self.remember(state, action, prob, reward)\n",
    "        state=next_state\n",
    "        episode_reward+=reward\n",
    "\n",
    "        #if episode%render_n==0: ## render env to visualize.\n",
    "          #env.render()\n",
    "        if done:\n",
    "          # update policy \n",
    "          if episode%rollout_n==0:\n",
    "            history=self.update_policy()\n",
    "\n",
    "      total_rewards[episode]=episode_reward\n",
    "      if episode%10==0:\n",
    "        print(f\"{episode} -> {episode_reward}\")\n",
    "      \n",
    "    self.total_rewards=total_rewards"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\winapp\\Miniconda3\\envs\\py38\\lib\\site-packages\\keras\\optimizer_v2\\adam.py:105: UserWarning: The `lr` argument is deprecated, use `learning_rate` instead.\n",
      "  super(Adam, self).__init__(name, **kwargs)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 -> 39.0\n",
      "10 -> 50.0\n",
      "20 -> 18.0\n",
      "30 -> 15.0\n",
      "40 -> 16.0\n",
      "50 -> 17.0\n",
      "60 -> 20.0\n",
      "70 -> 11.0\n",
      "80 -> 20.0\n",
      "90 -> 40.0\n",
      "100 -> 36.0\n",
      "110 -> 18.0\n",
      "120 -> 79.0\n",
      "130 -> 43.0\n",
      "140 -> 41.0\n",
      "150 -> 37.0\n",
      "160 -> 95.0\n",
      "170 -> 72.0\n",
      "180 -> 138.0\n",
      "190 -> 102.0\n"
     ]
    }
   ],
   "source": [
    "r = REINFORCE(env)\n",
    "\n",
    "r.train(200)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x1f7a517f550>]"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 23520 characters]",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(r.total_rewards)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "16af2a8bbb083ea23e5e41c7f5787656b2ce26968575d8763f2c4b17f9cd711f"
  },
  "kernelspec": {
   "display_name": "Python 3.8.12 ('py38')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.12"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
