FIDLE_DQNfromScratch.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "w_5p3EyVknLC",
      "metadata": {
        "id": "w_5p3EyVknLC"
      },
      "source": [
        "<img width=\"800px\" src=\"../fidle/img/header.svg\"></img>\n",
        "\n",
        "# <!-- TITLE --> [DRL1] - Solving CartPole with DQN\n",
        "<!-- DESC --> Using a a Deep Q-Network to play CartPole - an inverted pendulum problem (PyTorch)\n",
        "<!-- AUTHOR : Nathan Cassereau (IDRIS) and Bertrand Cabot (IDRIS) -->\n",
        "\n",
        "\n",
        "\n",
        "By Nathan Cassereau (IDRIS) and Bertrand Cabot (IDRIS)\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "ucB28wGpmFwi",
      "metadata": {
        "id": "ucB28wGpmFwi"
      },
      "source": [
        "## Objectives\n",
        "\n",
        "* Understand the code behind the DQN algorithm\n",
        "* Visualize the result for fun purposes :)\n",
        "\n",
        "This notebook implements a DQN from scratch and trains it. It is simply a vanilla DQN with a target network (sometimes referred as Double DQN). More sophisticated and recent modifications might help stabilize the training.\n",
        "\n",
        "Considering that we are going to use a tiny network for a simple environment, matrix multiplications are not that time consuming, and using a GPU can be detrimental as communications between CPU and GPU are no longer negligeable compared to forward and backward steps. This notebook will therefore be executed on CPU.\n",
        "\n",
        "The chosen environment will be imported from the gym toolkit (https://gym.openai.com/)."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "fqQsB2Jwm-BP",
      "metadata": {
        "id": "fqQsB2Jwm-BP"
      },
      "source": [
        "## Demonstration steps:\n",
        "\n",
        "- Define numerous hyperparameters\n",
        "- Implement the Q-Network\n",
        "- Implement an agent following the Double DQN algorithm\n",
        "- Train it for a few minutes\n",
        "- Visualize the result"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "nRJmgZ0inpkk",
      "metadata": {
        "id": "nRJmgZ0inpkk"
      },
      "source": [
        "## Installations\n",
        "\n",
        "Gym requires a graphical interface to render a state observation. Xvfb allows to run the notebook headless. This software is not available on Jean Zay's compute node, hence the usage of Google colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "y2Y71JbfgkeU",
      "metadata": {
        "id": "y2Y71JbfgkeU"
      },
      "outputs": [],
      "source": [
        "!pip3 install pyvirtualdisplay\n",
        "!pip install pyglet==1.5.11\n",
        "!apt-get install x11-utils > /dev/null 2>&1 \n",
        "!apt-get install -y xvfb python-opengl > /dev/null 2>&1"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "q6eYfBKnoOJQ",
      "metadata": {
        "id": "q6eYfBKnoOJQ"
      },
      "source": [
        "## Imports\n",
        "\n",
        "I chose to use Pytorch to implement this DQN due to its straightforward API and personal preferences.\n",
        "Gym implements the environment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "0fc91d65-4756-4432-906c-7d315d981775",
      "metadata": {
        "id": "0fc91d65-4756-4432-906c-7d315d981775"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "\n",
        "import gym\n",
        "from gym import wrappers\n",
        "\n",
        "import random\n",
        "from tqdm.notebook import tqdm\n",
        "\n",
        "import functools\n",
        "import matplotlib.pyplot as plt\n",
        "import os\n",
        "import io\n",
        "import base64\n",
        "import glob\n",
        "from IPython.display import display, HTML"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "Hao-RYcdowHn",
      "metadata": {
        "id": "Hao-RYcdowHn"
      },
      "source": [
        "## Hyperparameters\n",
        "\n",
        "The size of the replay buffer does not matter much. In this case, it is big enough to hold every transitions we will have in our training. This choice does have a huge impact on memory though.\n",
        "\n",
        "Warm-up allows the network to gather some information before the training process begins.\n",
        "\n",
        "The target network will only be updated once every 10k steps in order to stabilize the training.\n",
        "\n",
        "The exploration rate is linearly decreasing, although an exponential curve is a sound and common choice as well.\n",
        "\n",
        "As mentioned above, only the CPU will be used, the GPU would be useful for bigger networks, and / or environments which have a torch tensor internal state.\n",
        "\n",
        "Considering this is a simple DQN implementation, its stability leaves a lot to be desired. In order not to rely on luck, a decent seed was chosen."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "6fX1X6y6YHXF",
      "metadata": {
        "id": "6fX1X6y6YHXF"
      },
      "outputs": [],
      "source": [
        "learning_rate = 0.0001\n",
        "buffer_size = 200000\n",
        "warmup_steps = 10000\n",
        "batch_size = 32\n",
        "gamma = 0.99\n",
        "train_freq = 4\n",
        "target_update_interval = 10000\n",
        "exploration_fraction = 0.1\n",
        "exploration_initial_eps = 1.0\n",
        "exploration_final_eps = 0.05\n",
        "device = torch.device(\"cpu\") # torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "\n",
        "seed = 987654321\n",
        "np.random.seed(seed)\n",
        "torch.manual_seed(seed)\n",
        "random.seed(seed)\n",
        "if torch.cuda.is_available():\n",
        "    torch.cuda.manual_seed(seed)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "TofGB-s7qfSH",
      "metadata": {
        "id": "TofGB-s7qfSH"
      },
      "source": [
        "## Q-Network and Agent implementation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "4VhftO9PaE9g",
      "metadata": {
        "id": "4VhftO9PaE9g"
      },
      "outputs": [],
      "source": [
        "class DQN(nn.Module):\n",
        "\n",
        "    def __init__(self):\n",
        "        super(DQN, self).__init__()\n",
        "        self.layer1 = nn.Linear(4, 64)\n",
        "        self.layer2 = nn.Linear(64, 64)\n",
        "        self.layer3 = nn.Linear(64, 2)\n",
        "        self.relu = nn.ReLU()\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.relu(self.layer1(x))\n",
        "        x = self.relu(self.layer2(x))\n",
        "        return self.layer3(x)\n",
        "\n",
        "    def compute_target(self, x, rewards):\n",
        "        with torch.no_grad():\n",
        "            values = torch.zeros(x.shape[0], device=device)\n",
        "            values[rewards != 1] = torch.max(self.forward(x[rewards != 1]), dim=-1)[0]\n",
        "            values = rewards + gamma * values\n",
        "        return values\n",
        "\n",
        "    def predict(self, x):\n",
        "        if len(x.shape) < 2:\n",
        "            x = x[None, :]\n",
        "        with torch.no_grad():\n",
        "            x = torch.argmax(self.forward(x), dim=-1)\n",
        "        if x.device.type == \"cuda\":\n",
        "            x = x.cpu()\n",
        "        return x\n",
        "\n",
        "class Agent:\n",
        "\n",
        "    def __init__(self, env):\n",
        "        self.env = env\n",
        "        self.q_network = DQN().to(device)\n",
        "        self.target_network = DQN().to(device)\n",
        "        self.target_network.eval()\n",
        "        self.synchronize()\n",
        "        self.optimizer = torch.optim.Adam(self.q_network.parameters(), lr=learning_rate)\n",
        "        self.criterion = nn.MSELoss()\n",
        "        self.buffer = []\n",
        "        self.n_updates = 0\n",
        "    \n",
        "    def add_transition(self, state, action, reward, nextState):\n",
        "        self.buffer.append((state, action, reward, nextState))\n",
        "        if len(self.buffer) > buffer_size:\n",
        "            self.buffer.pop(random.randrange(len(self.buffer)))\n",
        "\n",
        "    def sample(self):\n",
        "        transitions = random.sample(self.buffer, batch_size)\n",
        "        states, actions, rewards, nextStates = zip(*transitions)\n",
        "        states = torch.stack(states).to(device)\n",
        "        actions = torch.cat(actions).to(device)\n",
        "        rewards = torch.cat(rewards).to(device)\n",
        "        nextStates = torch.stack(nextStates).to(device)\n",
        "        return states, actions, rewards, nextStates\n",
        "    \n",
        "    def train_step(self, step):\n",
        "        if step % target_update_interval == 0:\n",
        "            self.synchronize()\n",
        "        if step < warmup_steps or step % train_freq != 0:\n",
        "            return 0.\n",
        "\n",
        "        states, actions, rewards, nextStates = self.sample()\n",
        "        output = self.q_network(states)\n",
        "        output = torch.gather(output, 1, actions.unsqueeze(-1)).view(-1)\n",
        "        expectedOutput = self.target_network.compute_target(nextStates, rewards).view(-1)\n",
        "        self.optimizer.zero_grad()\n",
        "        loss = self.criterion(output, expectedOutput)\n",
        "        loss.backward()\n",
        "        torch.nn.utils.clip_grad_norm_(self.q_network.parameters(), 10)\n",
        "        self.optimizer.step()\n",
        "        self.n_updates += 1\n",
        "        return loss.item()\n",
        "\n",
        "    def synchronize(self):\n",
        "        self.target_network.load_state_dict(self.q_network.state_dict())\n",
        "\n",
        "    def play(self, state, exploration_rate=0.):\n",
        "        if random.random() > exploration_rate:\n",
        "            return self.q_network.predict(state.to(device))\n",
        "        else:\n",
        "            shape = (state.shape[0],) if len(state.shape) > 1 else (1,)\n",
        "            return torch.randint(0, 2, size=shape)\n",
        "\n",
        "    @functools.lru_cache(maxsize=None)\n",
        "    def exploration_slope(self, total_steps):\n",
        "        return (exploration_initial_eps - exploration_final_eps) / (exploration_fraction * total_steps)\n",
        "\n",
        "    def exploration(self, step, total_steps):\n",
        "        eps = exploration_initial_eps - step * self.exploration_slope(total_steps)\n",
        "        return max(eps, exploration_final_eps)\n",
        "\n",
        "    def train(self, total_steps):\n",
        "        obs = torch.from_numpy(env.reset()).float()\n",
        "\n",
        "        n_episodes = 0\n",
        "        length_current_episode = 0\n",
        "        lengths = []\n",
        "        avg_reward = 0\n",
        "        loss_backup = 0.\n",
        "        acc_loss = 0.\n",
        "        acc_loss_count = 0\n",
        "        self.rewards = []\n",
        "\n",
        "        with tqdm(range(total_steps), desc=\"Training agent\", unit=\"steps\") as pbar:\n",
        "            for step in pbar:\n",
        "                eps = self.exploration(step, total_steps)\n",
        "\n",
        "                action = self.play(obs, eps)\n",
        "                new_obs, _, done, info = env.step(action.item())\n",
        "                reward = torch.tensor([1.0 if not done else -1.0], dtype=torch.float32)\n",
        "                new_obs = torch.from_numpy(new_obs).float()\n",
        "\n",
        "                self.add_transition(obs, action, reward, new_obs)\n",
        "                loss = self.train_step(step)\n",
        "                if loss != 0:\n",
        "                    acc_loss += loss\n",
        "                    acc_loss_count += 1\n",
        "\n",
        "                if done:\n",
        "                    obs = torch.from_numpy(env.reset()).float()\n",
        "                    n_episodes += 1\n",
        "                    lengths.append(length_current_episode)\n",
        "                    self.rewards.append(length_current_episode)\n",
        "                    length_current_episode = 0\n",
        "                    if len(lengths) >= 25:\n",
        "                        avg_reward = sum(lengths) / len(lengths)\n",
        "                        if acc_loss_count != 0:\n",
        "                            loss_backup = acc_loss / acc_loss_count\n",
        "                        else:\n",
        "                            loss_backup = \"??\"\n",
        "                        acc_loss = 0.\n",
        "                        acc_loss_count = 0\n",
        "                        lengths = []\n",
        "                else:\n",
        "                    obs = new_obs\n",
        "                    length_current_episode += 1\n",
        "\n",
        "                pbar.set_postfix({\n",
        "                    \"episodes\": n_episodes,\n",
        "                    \"avg_reward\": avg_reward,\n",
        "                    \"loss\": loss_backup,\n",
        "                    \"exploration_rate\": eps,\n",
        "                    \"n_updates\": self.n_updates,\n",
        "                })"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "Kne9b7vCql3N",
      "metadata": {
        "id": "Kne9b7vCql3N"
      },
      "source": [
        "## Defining the environment"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "BXw4RmGpFkZm",
      "metadata": {
        "id": "BXw4RmGpFkZm"
      },
      "outputs": [],
      "source": [
        "env = gym.make(\"CartPole-v1\")\n",
        "env.seed(seed+2)\n",
        "env.reset()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "i93WQNsbqo68",
      "metadata": {
        "id": "i93WQNsbqo68"
      },
      "source": [
        "## Training our agent"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "rAm6v_0HiEge",
      "metadata": {
        "id": "rAm6v_0HiEge"
      },
      "outputs": [],
      "source": [
        "agent = Agent(env)\n",
        "agent.train(120000)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "PPT-tl4Rqroj",
      "metadata": {
        "id": "PPT-tl4Rqroj"
      },
      "source": [
        "## Episodes length\n",
        "\n",
        "A very noisy curve. It does reach satisfying levels though."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "IoCnHaZKgHqI",
      "metadata": {
        "id": "IoCnHaZKgHqI"
      },
      "outputs": [],
      "source": [
        "fig = plt.figure(figsize=(20, 12))\n",
        "plt.plot(agent.rewards)\n",
        "plt.xlabel(\"Episodes\")\n",
        "plt.ylabel(\"Episode length\")\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "0fuolKppq1Ak",
      "metadata": {
        "id": "0fuolKppq1Ak"
      },
      "source": [
        "## Result visualisation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "GXT1q5ckh0dG",
      "metadata": {
        "id": "GXT1q5ckh0dG"
      },
      "outputs": [],
      "source": [
        "from pyvirtualdisplay import Display\n",
        "\n",
        "virtual_display = Display(visible=0, size=(1400, 900))\n",
        "virtual_display.start()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "710b8294-4f75-49b5-a54a-777439ce8799",
      "metadata": {
        "id": "710b8294-4f75-49b5-a54a-777439ce8799"
      },
      "outputs": [],
      "source": [
        "env = gym.make(\"CartPole-v1\")\n",
        "env.seed(4)\n",
        "env = wrappers.Monitor(env, \"./CartPole-v1/\", force=True)\n",
        "\n",
        "obs = env.reset()\n",
        "i = 0\n",
        "\n",
        "while True:\n",
        "    action = agent.q_network.predict(torch.from_numpy(obs).float().to(device))\n",
        "    \n",
        "    obs, rewards, done, info = env.step(action.item())\n",
        "    env.render()\n",
        "    if done:\n",
        "        break\n",
        "    else:\n",
        "        i += 1\n",
        "env.close()\n",
        "print(f\"Survived {i} steps\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c7ad6655-02b7-436e-a7ae-93a7222b100e",
      "metadata": {
        "id": "c7ad6655-02b7-436e-a7ae-93a7222b100e"
      },
      "outputs": [],
      "source": [
        "def ipython_show_video(path):\n",
        "    \"\"\"Shamelessly stolen from https://stackoverflow.com/a/51183488/9977878\n",
        "    \"\"\"\n",
        "    if not os.path.isfile(path):\n",
        "        raise NameError(\"Cannot access: {}\".format(path))\n",
        "\n",
        "    video = io.open(path, 'r+b').read()\n",
        "    encoded = base64.b64encode(video)\n",
        "\n",
        "    display(HTML(\n",
        "        data=\"\"\"\n",
        "        <video alt=\"test\" controls>\n",
        "        <source src=\"data:video/mp4;base64,{0}\" type=\"video/mp4\" />\n",
        "        </video>\n",
        "        \"\"\".format(encoded.decode('ascii'))\n",
        "    ))\n",
        "\n",
        "ipython_show_video(glob.glob(\"/content/CartPole-v1/*.mp4\")[0])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "31e6af84-489e-4665-919e-8234462c1f0a",
      "metadata": {
        "id": "31e6af84-489e-4665-919e-8234462c1f0a"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "drl(2).ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "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.9.7"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}