{
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "# And for visualization on Colab install\n",
        "# !apt-get install x11-utils > /dev/null 2>&1 \n",
        "# !pip install pyglet\n",
        "# !apt-get install -y xvfb python-opengl > /dev/null 2>&1\n",
        "# !pip install gym pyvirtualdisplay > /dev/null 2>&1"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "tOhAux1ubEKG",
        "outputId": "727c632b-d755-4ac2-dd1a-9b84c4f1259f"
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: pyglet in /usr/local/lib/python3.7/dist-packages (1.5.0)\n",
            "Requirement already satisfied: future in /usr/local/lib/python3.7/dist-packages (from pyglet) (0.16.0)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "JZV-qP-yay8_"
      },
      "outputs": [],
      "source": [
        "import random\n",
        "import gym\n",
        "#import math\n",
        "import numpy as np\n",
        "from collections import deque\n",
        "import tensorflow as tf\n",
        "from tensorflow.keras.models import Sequential\n",
        "from tensorflow.keras.layers import Dense\n",
        "from tensorflow.keras.optimizers import Adam\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "## Uncomment if working on Colab\n",
        "# from pyvirtualdisplay import Display\n",
        "# display = Display(visible=0, size=(600, 400))\n",
        "# display.start()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "i63z1vW0c4Sp",
        "outputId": "222984a5-6556-4942-a509-c5a1a639c63f"
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<pyvirtualdisplay.display.Display at 0x7f0d25df98d0>"
            ]
          },
          "metadata": {},
          "execution_count": 3
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "ikpmIrLyay9B"
      },
      "outputs": [],
      "source": [
        "EPOCHS = 1000\n",
        "THRESHOLD = 45\n",
        "MONITOR = True"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "trKmD7d2ay9C"
      },
      "outputs": [],
      "source": [
        "class DQN():\n",
        "    def __init__(self, env_string,batch_size=64):\n",
        "        self.memory = deque(maxlen=100000)\n",
        "        self.env = gym.make(env_string)\n",
        "        input_size = self.env.observation_space.shape[0]\n",
        "        action_size = self.env.action_space.n\n",
        "        self.batch_size = batch_size\n",
        "        self.gamma = 1.0\n",
        "        self.epsilon = 1.0\n",
        "        self.epsilon_min = 0.01\n",
        "        self.epsilon_decay = 0.995\n",
        "        \n",
        "        alpha=0.01\n",
        "        alpha_decay=0.01\n",
        "        if MONITOR: self.env = gym.wrappers.Monitor(self.env, 'data/'+env_string, force=True)\n",
        "        \n",
        "        # Init model\n",
        "        self.model = Sequential()\n",
        "        self.model.add(Dense(24, input_dim=input_size, activation='tanh'))\n",
        "        self.model.add(Dense(48, activation='tanh'))\n",
        "        self.model.add(Dense(action_size, activation='linear'))\n",
        "        self.model.compile(loss='mse', optimizer=Adam(lr=alpha, decay=alpha_decay))\n",
        "\n",
        "    def remember(self, state, action, reward, next_state, done):\n",
        "        self.memory.append((state, action, reward, next_state, done))\n",
        "\n",
        "    def choose_action(self, state, epsilon):\n",
        "        if np.random.random() <= epsilon:\n",
        "            return self.env.action_space.sample()\n",
        "        else:\n",
        "            return np.argmax(self.model.predict(state))\n",
        "\n",
        "    def preprocess_state(self, state):\n",
        "        return np.reshape(state, [1, 4])\n",
        "\n",
        "    def replay(self, batch_size):\n",
        "        x_batch, y_batch = [], []\n",
        "        minibatch = random.sample(self.memory, min(len(self.memory), batch_size))\n",
        "        for state, action, reward, next_state, done in minibatch:\n",
        "            y_target = self.model.predict(state)\n",
        "            y_target[0][action] = reward if done else reward + self.gamma * np.max(self.model.predict(next_state)[0])\n",
        "            x_batch.append(state[0])\n",
        "            y_batch.append(y_target[0])\n",
        "        \n",
        "        self.model.fit(np.array(x_batch), np.array(y_batch), batch_size=len(x_batch), verbose=0)\n",
        "        #epsilon = max(epsilon_min, epsilon_decay*epsilon) # decrease epsilon\n",
        "       \n",
        "\n",
        "    def train(self):\n",
        "        scores = deque(maxlen=100)\n",
        "        avg_scores = []\n",
        "        \n",
        "\n",
        "        for e in range(EPOCHS):\n",
        "            state = self.env.reset()\n",
        "            state = self.preprocess_state(state)\n",
        "            done = False\n",
        "            i = 0\n",
        "            while not done:\n",
        "                action = self.choose_action(state,self.epsilon)\n",
        "                next_state, reward, done, _ = self.env.step(action)\n",
        "                next_state = self.preprocess_state(next_state)\n",
        "                self.remember(state, action, reward, next_state, done)\n",
        "                state = next_state\n",
        "                self.epsilon = max(self.epsilon_min, self.epsilon_decay*self.epsilon) # decrease epsilon\n",
        "                i += 1\n",
        "\n",
        "            scores.append(i)\n",
        "            mean_score = np.mean(scores)\n",
        "            avg_scores.append(mean_score)\n",
        "            if mean_score >= THRESHOLD and e >= 100:\n",
        "                print('Ran {} episodes. Solved after {} trials ✔'.format(e, e - 100))\n",
        "                return avg_scores\n",
        "            if e % 100 == 0:\n",
        "                print('[Episode {}] - Mean survival time over last 100 episodes was {} ticks.'.format(e, mean_score))\n",
        "\n",
        "            self.replay(self.batch_size)\n",
        "        \n",
        "        print('Did not solve after {} episodes 😞'.format(e))\n",
        "        return avg_scores\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4STstW_7ay9E",
        "outputId": "b9a26bf3-dd8c-4b4f-c92c-9b6a75333acf"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.7/dist-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"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[Episode 0] - Mean survival time over last 100 episodes was 28.0 ticks.\n",
            "[Episode 100] - Mean survival time over last 100 episodes was 15.71 ticks.\n",
            "[Episode 200] - Mean survival time over last 100 episodes was 27.81 ticks.\n",
            "Ran 259 episodes. Solved after 159 trials ✔\n"
          ]
        }
      ],
      "source": [
        "env_string = 'CartPole-v0'\n",
        "agent = DQN(env_string)\n",
        "scores = agent.train()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 265
        },
        "id": "28iEbGwzay9F",
        "outputId": "e9ab9177-f5eb-472f-dead-64a17be16cf7"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 432x288 with 1 Axes>"
            ],
            "image/png": "[encoded data, 15768 characters]\n"
          },
          "metadata": {
            "needs_background": "light"
          }
        }
      ],
      "source": [
        "import matplotlib.pyplot as plt\n",
        "plt.plot(scores)\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "178ufOPzay9F",
        "outputId": "3836158e-b35a-471d-c59f-a2c0460712b8"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Model: \"sequential\"\n",
            "_________________________________________________________________\n",
            " Layer (type)                Output Shape              Param #   \n",
            "=================================================================\n",
            " dense (Dense)               (None, 24)                120       \n",
            "                                                                 \n",
            " dense_1 (Dense)             (None, 48)                1200      \n",
            "                                                                 \n",
            " dense_2 (Dense)             (None, 2)                 98        \n",
            "                                                                 \n",
            "=================================================================\n",
            "Total params: 1,418\n",
            "Trainable params: 1,418\n",
            "Non-trainable params: 0\n",
            "_________________________________________________________________\n"
          ]
        }
      ],
      "source": [
        "agent.model.summary()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "5E_2klZ3ay9G"
      },
      "outputs": [],
      "source": [
        "agent.env.close()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "b0TrCnMbay9H"
      },
      "outputs": [],
      "source": [
        ""
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "colab": {
      "name": "DQNCartPole.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}