{
 "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 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.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.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 len(self.memory) < self.batch_size:\n",
    "            return\n",
    "\n",
    "        minibatch = random.sample(self.memory, self.batch_size)\n",
    "        states = np.array([i[0] for i in minibatch])\n",
    "        actions = np.array([i[1] for i in minibatch])\n",
    "        rewards = np.array([i[2] for i in minibatch])\n",
    "        next_states = np.array([i[3] for i in minibatch])\n",
    "        dones = np.array([i[4] for i in minibatch])\n",
    "\n",
    "        states = np.squeeze(states)\n",
    "        next_states = np.squeeze(next_states)\n",
    "\n",
    "        targets = rewards + self.gamma*(np.amax(self.model.predict_on_batch(next_states), axis=1))*(1-dones)\n",
    "        targets_full = self.model.predict_on_batch(states)\n",
    "\n",
    "        ind = np.array([i for i in range(self.batch_size)])\n",
    "        targets_full[[ind], [actions]] = targets\n",
    "\n",
    "        self.model.fit(states, targets_full, epochs=1, verbose=0)\n",
    "        if self.epsilon > self.epsilon_min:\n",
    "            self.epsilon *= self.epsilon_decay\n",
    "\n",
    "\n",
    "def train_dqn(episode):\n",
    "\n",
    "    loss = []\n",
    "    agent = DQN(env.action_space.n, env.observation_space.shape[0])\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: {}\".format(e, episode, score))\n",
    "                break\n",
    "        loss.append(score)\n",
    "    return loss, agent\n",
    "\n",
    "def test_dqn(env, episode, agent):\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:\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": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /users/yby1/.local/lib/python3.7/site-packages/tensorflow/python/ops/resource_variable_ops.py:1666: calling BaseResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "If using Keras pass *_constraint arguments to layers.\n",
      "episode: 0/100, score: 28.0\n",
      "episode: 1/100, score: 21.0\n",
      "episode: 2/100, score: 51.0\n",
      "episode: 3/100, score: 14.0\n",
      "episode: 4/100, score: 21.0\n",
      "episode: 5/100, score: 14.0\n",
      "episode: 6/100, score: 18.0\n",
      "episode: 7/100, score: 10.0\n",
      "episode: 8/100, score: 10.0\n",
      "episode: 9/100, score: 16.0\n",
      "episode: 10/100, score: 30.0\n",
      "episode: 11/100, score: 13.0\n",
      "episode: 12/100, score: 8.0\n",
      "episode: 13/100, score: 11.0\n",
      "episode: 14/100, score: 11.0\n",
      "episode: 15/100, score: 12.0\n",
      "episode: 16/100, score: 11.0\n",
      "episode: 17/100, score: 10.0\n",
      "episode: 18/100, score: 20.0\n",
      "episode: 19/100, score: 58.0\n",
      "episode: 20/100, score: 152.0\n",
      "episode: 21/100, score: 130.0\n",
      "episode: 22/100, score: 124.0\n",
      "episode: 23/100, score: 105.0\n",
      "episode: 24/100, score: 123.0\n",
      "episode: 25/100, score: 200.0\n",
      "episode: 26/100, score: 200.0\n",
      "episode: 27/100, score: 200.0\n",
      "episode: 28/100, score: 200.0\n",
      "episode: 29/100, score: 200.0\n",
      "episode: 30/100, score: 200.0\n",
      "episode: 31/100, score: 200.0\n",
      "episode: 32/100, score: 200.0\n",
      "episode: 33/100, score: 200.0\n",
      "episode: 34/100, score: 200.0\n",
      "episode: 35/100, score: 200.0\n",
      "episode: 36/100, score: 200.0\n",
      "episode: 37/100, score: 200.0\n",
      "episode: 38/100, score: 200.0\n",
      "episode: 39/100, score: 200.0\n",
      "episode: 40/100, score: 200.0\n",
      "episode: 41/100, score: 200.0\n",
      "episode: 42/100, score: 200.0\n",
      "episode: 43/100, score: 200.0\n",
      "episode: 44/100, score: 200.0\n",
      "episode: 45/100, score: 200.0\n",
      "episode: 46/100, score: 200.0\n",
      "episode: 47/100, score: 200.0\n",
      "episode: 48/100, score: 200.0\n",
      "episode: 49/100, score: 200.0\n",
      "episode: 50/100, score: 169.0\n",
      "episode: 51/100, score: 200.0\n",
      "episode: 52/100, score: 200.0\n",
      "episode: 53/100, score: 200.0\n",
      "episode: 54/100, score: 162.0\n",
      "episode: 55/100, score: 148.0\n",
      "episode: 56/100, score: 200.0\n",
      "episode: 57/100, score: 200.0\n",
      "episode: 58/100, score: 200.0\n",
      "episode: 59/100, score: 200.0\n",
      "episode: 60/100, score: 200.0\n",
      "episode: 61/100, score: 200.0\n",
      "episode: 62/100, score: 200.0\n",
      "episode: 63/100, score: 200.0\n",
      "episode: 64/100, score: 200.0\n",
      "episode: 65/100, score: 200.0\n",
      "episode: 66/100, score: 144.0\n",
      "episode: 67/100, score: 153.0\n",
      "episode: 68/100, score: 149.0\n",
      "episode: 69/100, score: 182.0\n",
      "episode: 70/100, score: 200.0\n",
      "episode: 71/100, score: 200.0\n",
      "episode: 72/100, score: 200.0\n",
      "episode: 73/100, score: 200.0\n",
      "episode: 74/100, score: 200.0\n",
      "episode: 75/100, score: 200.0\n",
      "episode: 76/100, score: 200.0\n",
      "episode: 77/100, score: 200.0\n",
      "episode: 78/100, score: 200.0\n",
      "episode: 79/100, score: 200.0\n",
      "episode: 80/100, score: 200.0\n",
      "episode: 81/100, score: 200.0\n",
      "episode: 82/100, score: 200.0\n",
      "episode: 83/100, score: 200.0\n",
      "episode: 84/100, score: 200.0\n",
      "episode: 85/100, score: 200.0\n",
      "episode: 86/100, score: 138.0\n",
      "episode: 87/100, score: 200.0\n",
      "episode: 88/100, score: 200.0\n",
      "episode: 89/100, score: 128.0\n",
      "episode: 90/100, score: 85.0\n",
      "episode: 91/100, score: 200.0\n",
      "episode: 92/100, score: 200.0\n",
      "episode: 93/100, score: 200.0\n",
      "episode: 94/100, score: 200.0\n",
      "episode: 95/100, score: 173.0\n",
      "episode: 96/100, score: 200.0\n",
      "episode: 97/100, score: 200.0\n",
      "episode: 98/100, score: 182.0\n",
      "episode: 99/100, score: 200.0\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 16700 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",
    "ep = 100\n",
    "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": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "episode: 0/10, score: 197.0\n",
      "episode: 1/10, score: 200.0\n",
      "episode: 2/10, score: 200.0\n",
      "episode: 3/10, score: 187.0\n",
      "episode: 4/10, score: 200.0\n",
      "episode: 5/10, score: 194.0\n",
      "episode: 6/10, score: 185.0\n",
      "episode: 7/10, score: 200.0\n",
      "episode: 8/10, score: 198.0\n",
      "episode: 9/10, score: 197.0\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 14216 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "195.8"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "env = gym.make('CartPole-v0')\n",
    "env.seed(0)\n",
    "np.random.seed(0)\n",
    "\n",
    "ep = 10\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
}
