{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Info\n",
    "\n",
    "https://github.com/shivaverma/OpenAIGym/blob/master/cart-pole/CartPole-v0.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'2.2.0'"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import gym\n",
    "import random\n",
    "from collections import deque\n",
    "import matplotlib.pyplot as plt\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "\n",
    "# disable eager execution to speed up\n",
    "tf.compat.v1.disable_eager_execution()\n",
    "\n",
    "tf.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ReplayBuffer(object):\n",
    "    def __init__(self, max_size, input_shape, n_actions, discrete=False):\n",
    "        \"\"\"Also called 'Replay Memory' it's designed to store previous\n",
    "            observation, action, reward, next-observation values.\n",
    "         The DQN needs experience replay to reduce correlations in \n",
    "             the sequence of observations, which otherwise might drive \n",
    "             the network into a local minimum.\n",
    "        \"\"\"\n",
    "        # when buffer gets full it goes back and re-writes again\n",
    "        self.mem_size = max_size\n",
    "        # position of first unsaved memory\n",
    "        self.mem_cntr = 0\n",
    "        self.discrete = discrete\n",
    "        self.observation_memory = np.zeros((self.mem_size, input_shape))\n",
    "        self.new_observation_memory = np.zeros((self.mem_size, input_shape))\n",
    "        dtype = np.int8 if self.discrete else np.float32\n",
    "        self.action_memory = np.zeros((self.mem_size, n_actions), dtype=dtype)\n",
    "        self.reward_memory = np.zeros(self.mem_size)\n",
    "        # keep track of our `done` flags becuase the value of the terminal state is 0.\n",
    "        self.terminal_memory = np.zeros(self.mem_size, dtype=np.float32)\n",
    "\n",
    "    def collect_experience(self, state, action, reward, state_, done):\n",
    "        \"\"\" Populate the replay memory with initial experience.\n",
    "        \"\"\"\n",
    "        index = self.mem_cntr % self.mem_size\n",
    "        self.observation_memory[index] = state\n",
    "        self.new_observation_memory[index] = state_\n",
    "        # store one hot encoding of actions, if appropriate\n",
    "        if self.discrete:\n",
    "            actions = np.zeros(self.action_memory.shape[1])\n",
    "            actions[action] = 1.0\n",
    "            self.action_memory[index] = actions\n",
    "        else:\n",
    "            self.action_memory[index] = action\n",
    "        self.reward_memory[index] = reward\n",
    "        # when episode done is 'True'. Use this to multiple the reward.\n",
    "        #  If reached temrinal state we don't give any reward (reward 0)\n",
    "        self.terminal_memory[index] = 1 - done\n",
    "        self.mem_cntr += 1\n",
    "\n",
    "    def sample_buffer(self, batch_size):\n",
    "        \"\"\"Sample memory buffer.\"\"\"\n",
    "        max_mem = min(self.mem_cntr, self.mem_size)\n",
    "        # replace will detemrine if reuse examples if already sampled\n",
    "        batch = np.random.choice(max_mem, batch_size, replace=False)\n",
    "\n",
    "        observations = self.observation_memory[batch]\n",
    "        actions = self.action_memory[batch]\n",
    "        rewards = self.reward_memory[batch]\n",
    "        new_observations = self.new_observation_memory[batch]\n",
    "        terminal = self.terminal_memory[batch]\n",
    "\n",
    "        return observations, actions, rewards, new_observations, terminal\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DQN:\n",
    "\n",
    "    \"\"\" Implementation of deep q learning algorithm \"\"\"\n",
    "\n",
    "    def __init__(self, action_space, state_space):\n",
    "\n",
    "        self.action_space = action_space\n",
    "        self.state_space = state_space\n",
    "        self.epsilon = 1\n",
    "        self.gamma = .95\n",
    "        self.batch_size = 64\n",
    "        self.epsilon_min = .01\n",
    "        self.epsilon_decay = .995\n",
    "        self.learning_rate = 0.001\n",
    "#         self.memory = deque(maxlen=10000)\n",
    "        self.memory = ReplayBuffer(10000, state_space, action_space,\n",
    "                                   discrete=True)\n",
    "        self.model = self.build_model()\n",
    "\n",
    "    def build_model(self):\n",
    "\n",
    "        model = tf.keras.Sequential()\n",
    "        model.add(tf.keras.layers.Dense(24, input_shape=(self.state_space,), activation='relu'))\n",
    "        model.add(tf.keras.layers.Dense(24, activation='relu'))\n",
    "        model.add(tf.keras.layers.Dense(self.action_space, activation='linear'))\n",
    "        model.compile(loss='mse', optimizer=tf.keras.optimizers.Adam(lr=self.learning_rate))\n",
    "        return model\n",
    "\n",
    "    def remember(self, state, action, reward, next_state, done):\n",
    "        self.memory.collect_experience(state, action, reward, next_state, done)\n",
    "#         self.memory.append((state, action, reward, next_state, done))\n",
    "\n",
    "    def act(self, state, use_exploration=True):\n",
    "\n",
    "        if use_exploration and (np.random.rand() <= self.epsilon):\n",
    "            return random.randrange(self.action_space)\n",
    "        act_values = self.model.predict(state)\n",
    "        return np.argmax(act_values[0])\n",
    "\n",
    "    def replay(self):\n",
    "\n",
    "        if self.memory.mem_cntr < self.batch_size:\n",
    "            return\n",
    "        \n",
    "        observation, action, reward, new_observation, done = \\\n",
    "                                      self.memory.sample_buffer(self.batch_size)\n",
    "\n",
    "        action_values = np.array(list(range(self.action_space)), dtype=np.int8)\n",
    "        # from one-hot to indeces\n",
    "        action_indices = np.dot(action, action_values)\n",
    "\n",
    "        # use target model as policy\n",
    "        q_eval = self.model.predict(observation)\n",
    "        # make copy in order to change it later\n",
    "        q_target = q_eval.copy()\n",
    "        # get Q-values for new observation\n",
    "        q_next = self.model.predict(new_observation)\n",
    "\n",
    "        batch_index = np.arange(self.batch_size, dtype=np.int32)\n",
    "        # create Q-values. No reward for ending states.\n",
    "        q_target[batch_index, action_indices] = reward + \\\n",
    "                              self.gamma*np.max(q_next, axis=1)*done\n",
    "        # fit batch to online-model\n",
    "        _ = self.model.fit(observation, q_target, verbose=0)\n",
    "        \n",
    "        self.epsilon = self.epsilon*self.epsilon_decay if self.epsilon > self.epsilon_min else self.epsilon_min\n",
    "#         if self.epsilon > self.epsilon_min:\n",
    "#             self.epsilon *= self.epsilon_decay\n",
    "\n",
    "\n",
    "def train_dqn(episode):\n",
    "\n",
    "    loss = []\n",
    "    test_loss = []\n",
    "    agent = DQN(env.action_space.n, env.observation_space.shape[0])\n",
    "    test_env = gym.make('CartPole-v0')\n",
    "    test_env.seed(0)\n",
    "    np.random.seed(0)\n",
    "\n",
    "    for e in range(episode):\n",
    "        state = env.reset()\n",
    "        state = np.reshape(state, (1, 4))\n",
    "        score = 0\n",
    "        max_steps = 1000\n",
    "        for i in range(max_steps):\n",
    "            action = agent.act(state)\n",
    "            next_state, reward, done, _ = env.step(action)\n",
    "            score += reward\n",
    "            next_state = np.reshape(next_state, (1, 4))\n",
    "            agent.remember(state, action, reward, next_state, done)\n",
    "            state = next_state\n",
    "            agent.replay()\n",
    "            if done:\n",
    "                print(\"episode: {}/{}, score: {}, epsilon {:.3f}\".format(e, episode, score, agent.epsilon))\n",
    "                break\n",
    "        test_loss.append(test_dqn(test_env, 10, agent, False))\n",
    "        loss.append(score)\n",
    "    return loss, test_loss, agent\n",
    "\n",
    "def test_dqn(env, episode, agent, verbose=True):\n",
    "\n",
    "    loss = []\n",
    "    for e in range(episode):\n",
    "        state = env.reset()\n",
    "        state = np.reshape(state, (1, 4))\n",
    "        score = 0\n",
    "        max_steps = 1000\n",
    "        for i in range(max_steps):\n",
    "            action = agent.act(state, use_exploration=False)\n",
    "            next_state, reward, done, _ = env.step(action)\n",
    "            score += reward\n",
    "            next_state = np.reshape(next_state, (1, 4))\n",
    "            state = next_state\n",
    "            if done and verbose:\n",
    "                print(\"episode: {}/{}, score: {}\".format(e, episode, score))\n",
    "                break\n",
    "        loss.append(score)\n",
    "    return loss\n",
    "\n",
    "def random_policy(episode, step):\n",
    "\n",
    "    for i_episode in range(episode):\n",
    "        env.reset()\n",
    "        for t in range(step):\n",
    "            env.render()\n",
    "            action = env.action_space.sample()\n",
    "            state, reward, done, info = env.step(action)\n",
    "            if done:\n",
    "                print(\"Episode finished after {} timesteps\".format(t+1))\n",
    "                break\n",
    "            print(\"Starting next episode\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "episode: 0/100, score: 22.0, epsilon 1.000\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/users/yby1/.local/lib/python3.7/site-packages/gym/logger.py:30: UserWarning: \u001b[33mWARN: You are calling 'step()' even though this environment has already returned done = True. You should always call 'reset()' once you receive 'done = True' -- any further steps are undefined behavior.\u001b[0m\n",
      "  warnings.warn(colorize('%s: %s'%('WARN', msg % args), 'yellow'))\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "episode: 1/100, score: 28.0, epsilon 1.000\n",
      "episode: 2/100, score: 13.0, epsilon 1.000\n",
      "episode: 3/100, score: 9.0, epsilon 0.956\n",
      "episode: 4/100, score: 10.0, epsilon 0.909\n",
      "episode: 5/100, score: 22.0, epsilon 0.814\n",
      "episode: 6/100, score: 17.0, epsilon 0.748\n",
      "episode: 7/100, score: 16.0, epsilon 0.690\n",
      "episode: 8/100, score: 10.0, epsilon 0.656\n",
      "episode: 9/100, score: 29.0, epsilon 0.568\n",
      "episode: 10/100, score: 14.0, epsilon 0.529\n",
      "episode: 11/100, score: 11.0, epsilon 0.501\n",
      "episode: 12/100, score: 17.0, epsilon 0.460\n",
      "episode: 13/100, score: 12.0, epsilon 0.433\n",
      "episode: 14/100, score: 9.0, epsilon 0.414\n",
      "episode: 15/100, score: 14.0, epsilon 0.386\n",
      "episode: 16/100, score: 12.0, epsilon 0.363\n",
      "episode: 17/100, score: 23.0, epsilon 0.324\n",
      "episode: 18/100, score: 12.0, epsilon 0.305\n",
      "episode: 19/100, score: 9.0, epsilon 0.291\n",
      "episode: 20/100, score: 10.0, epsilon 0.277\n",
      "episode: 21/100, score: 21.0, epsilon 0.249\n",
      "episode: 22/100, score: 32.0, epsilon 0.212\n",
      "episode: 23/100, score: 44.0, epsilon 0.170\n",
      "episode: 24/100, score: 35.0, epsilon 0.143\n",
      "episode: 25/100, score: 58.0, epsilon 0.107\n",
      "episode: 26/100, score: 49.0, epsilon 0.084\n",
      "episode: 27/100, score: 26.0, epsilon 0.073\n",
      "episode: 28/100, score: 45.0, epsilon 0.059\n",
      "episode: 29/100, score: 51.0, epsilon 0.045\n",
      "episode: 30/100, score: 60.0, epsilon 0.034\n",
      "episode: 31/100, score: 83.0, epsilon 0.022\n",
      "episode: 32/100, score: 42.0, epsilon 0.018\n",
      "episode: 33/100, score: 37.0, epsilon 0.015\n",
      "episode: 34/100, score: 50.0, epsilon 0.012\n",
      "episode: 35/100, score: 59.0, epsilon 0.010\n",
      "episode: 36/100, score: 106.0, epsilon 0.010\n",
      "episode: 37/100, score: 79.0, epsilon 0.010\n",
      "episode: 38/100, score: 200.0, epsilon 0.010\n",
      "episode: 39/100, score: 113.0, epsilon 0.010\n",
      "episode: 40/100, score: 196.0, epsilon 0.010\n",
      "episode: 41/100, score: 200.0, epsilon 0.010\n",
      "episode: 42/100, score: 191.0, epsilon 0.010\n",
      "episode: 43/100, score: 158.0, epsilon 0.010\n",
      "episode: 44/100, score: 184.0, epsilon 0.010\n",
      "episode: 45/100, score: 187.0, epsilon 0.010\n",
      "episode: 46/100, score: 200.0, epsilon 0.010\n",
      "episode: 47/100, score: 159.0, epsilon 0.010\n",
      "episode: 48/100, score: 200.0, epsilon 0.010\n",
      "episode: 49/100, score: 157.0, epsilon 0.010\n",
      "episode: 50/100, score: 164.0, epsilon 0.010\n",
      "episode: 51/100, score: 173.0, epsilon 0.010\n",
      "episode: 52/100, score: 200.0, epsilon 0.010\n",
      "episode: 53/100, score: 165.0, epsilon 0.010\n",
      "episode: 54/100, score: 151.0, epsilon 0.010\n",
      "episode: 55/100, score: 188.0, epsilon 0.010\n",
      "episode: 56/100, score: 176.0, epsilon 0.010\n",
      "episode: 57/100, score: 200.0, epsilon 0.010\n",
      "episode: 58/100, score: 181.0, epsilon 0.010\n",
      "episode: 59/100, score: 195.0, epsilon 0.010\n",
      "episode: 60/100, score: 200.0, epsilon 0.010\n",
      "episode: 61/100, score: 193.0, epsilon 0.010\n",
      "episode: 62/100, score: 200.0, epsilon 0.010\n",
      "episode: 63/100, score: 200.0, epsilon 0.010\n",
      "episode: 64/100, score: 200.0, epsilon 0.010\n",
      "episode: 65/100, score: 200.0, epsilon 0.010\n",
      "episode: 66/100, score: 177.0, epsilon 0.010\n",
      "episode: 67/100, score: 173.0, epsilon 0.010\n",
      "episode: 68/100, score: 200.0, epsilon 0.010\n",
      "episode: 69/100, score: 200.0, epsilon 0.010\n",
      "episode: 70/100, score: 192.0, epsilon 0.010\n",
      "episode: 71/100, score: 200.0, epsilon 0.010\n",
      "episode: 72/100, score: 200.0, epsilon 0.010\n",
      "episode: 73/100, score: 200.0, epsilon 0.010\n",
      "episode: 74/100, score: 200.0, epsilon 0.010\n",
      "episode: 75/100, score: 200.0, epsilon 0.010\n",
      "episode: 76/100, score: 182.0, epsilon 0.010\n",
      "episode: 77/100, score: 200.0, epsilon 0.010\n",
      "episode: 78/100, score: 186.0, epsilon 0.010\n",
      "episode: 79/100, score: 156.0, epsilon 0.010\n",
      "episode: 80/100, score: 157.0, epsilon 0.010\n",
      "episode: 81/100, score: 164.0, epsilon 0.010\n",
      "episode: 82/100, score: 177.0, epsilon 0.010\n",
      "episode: 83/100, score: 200.0, epsilon 0.010\n",
      "episode: 84/100, score: 200.0, epsilon 0.010\n",
      "episode: 85/100, score: 185.0, epsilon 0.010\n",
      "episode: 86/100, score: 200.0, epsilon 0.010\n",
      "episode: 87/100, score: 151.0, epsilon 0.010\n",
      "episode: 88/100, score: 147.0, epsilon 0.010\n",
      "episode: 89/100, score: 169.0, epsilon 0.010\n",
      "episode: 90/100, score: 138.0, epsilon 0.010\n",
      "episode: 91/100, score: 137.0, epsilon 0.010\n",
      "episode: 92/100, score: 144.0, epsilon 0.010\n",
      "episode: 93/100, score: 119.0, epsilon 0.010\n",
      "episode: 94/100, score: 117.0, epsilon 0.010\n",
      "episode: 95/100, score: 13.0, epsilon 0.010\n",
      "episode: 96/100, score: 11.0, epsilon 0.010\n",
      "episode: 97/100, score: 116.0, epsilon 0.010\n",
      "episode: 98/100, score: 119.0, epsilon 0.010\n",
      "episode: 99/100, score: 114.0, epsilon 0.010\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 17824 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "env = gym.make('CartPole-v0')\n",
    "env.seed(0)\n",
    "np.random.seed(0)\n",
    "\n",
    "agent = DQN(env.action_space.n, env.observation_space.shape[0])\n",
    "\n",
    "\n",
    "ep = 100\n",
    "loss, test_loss, agent = train_dqn(ep)\n",
    "plt.plot([i+1 for i in range(0, ep, 2)], loss[::2])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 35840 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot([i+1 for i in range(0, ep, 2)], test_loss[::2])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = gym.make('CartPole-v0')\n",
    "env.seed(0)\n",
    "np.random.seed(0)\n",
    "\n",
    "ep = 20\n",
    "loss = test_dqn(env, ep, agent)\n",
    "plt.plot([i+1 for i in range(0, ep, 2)], loss[::2])\n",
    "plt.show()\n",
    "\n",
    "np.mean(loss)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
