02 Intro to Q-learning and DQN.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Google Colab Setup\n",
    "---\n",
    "\n",
    "Make sure to select GPU in Runtime > Change runtime type > Hardware accelerator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@title << Setup Google Colab by running this cell {display-mode: \"form\"}\n",
    "import sys\n",
    "if 'google.colab' in sys.modules:\n",
    "    # Clone GitHub repository\n",
    "    !git clone https://github.com/pacm/rl-workshop.git\n",
    "        \n",
    "    # Copy files required to run the code\n",
    "    !cp -r \"rl-workshop/agents\" \"rl-workshop/env\" \"rl-workshop/helpers\" \"rl-workshop/videos\" .\n",
    "    \n",
    "    # Install packages via pip\n",
    "    !pip install -r \"rl-workshop/colab-requirements.txt\"\n",
    "    \n",
    "    # Restart Runtime\n",
    "    import os\n",
    "    os.kill(os.getpid(), 9)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Imports\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run env/env.py\n",
    "%run helpers/rl_helpers.py\n",
    "%run agents/dqn.py\n",
    "%run agents/qlearning.py\n",
    "%run agents/random.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# You might want to import other libraries\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Intro to Q-Learning (compass Q-table)\n",
    "---\n",
    "\n",
    "You can find a Q-learning implementation in `agents/`\n",
    "\n",
    "```\n",
    "agents/\n",
    "├── curiosity.py\n",
    "├── dqn.py\n",
    "├── logging.py\n",
    "├── qlearning.py    <-- Q-learning agent\n",
    "└── random.py\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Environment without Skyscrapers + discharge\n",
    "env = CompassQTable(DeliveryDrones())\n",
    "env.env_params.update({'n_drones': 3, 'skyscrapers_factor': 0, 'stations_factor': 0,  'discharge': 0})\n",
    "states = env.reset()\n",
    "\n",
    "print('Observation space:', env.observation_space)\n",
    "print('Initial state:', {drone_index: env.format_state(state) for drone_index, state in states.items()})\n",
    "Image.fromarray(env.render(mode='rgb_array'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the agents\n",
    "agents = {drone.index: RandomAgent(env) for drone in env.drones}\n",
    "agents[0] = QLearningAgent(\n",
    "    env,\n",
    "    gamma=0.95, # Discount factor\n",
    "    alpha=0.1, # Learning rate\n",
    "    # Exploration rate\n",
    "    epsilon_start=1, epsilon_decay=0.99, epsilon_end=0.01\n",
    ")\n",
    "agents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train agents\n",
    "trainer = MultiAgentTrainer(env, agents, reset_agents=True, seed=0)\n",
    "trainer.train(5000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_rolling_rewards(trainer.rewards_log, drones_labels={0: 'Q-learning'})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "agents[0].get_qtable()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(agents[0].gamma**np.arange(100))\n",
    "plt.title('Discount factor: {}'.format(agents[0].gamma))\n",
    "plt.xlabel('Number of steps')\n",
    "plt.ylabel('Discount')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rewards_log = test_agents(env, agents, n_steps=1000, seed=0)\n",
    "plot_cumulative_rewards(rewards_log, drones_labels={0: 'Q-learning'})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = os.path.join('videos', 'ql-compass.mp4')\n",
    "render_video(env, agents, video_path=path, n_steps=120, fps=1, seed=0)\n",
    "ColabVideo(path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Scaling Q-learning (compass + lidar Q-table)\n",
    "---\n",
    "\n",
    "Let's see how Q-learning scales to larger observation spaces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Environment with skyscrapers but without discharge\n",
    "env = LidarCompassQTable(DeliveryDrones())\n",
    "env.env_params.update({'n_drones': 3, 'skyscrapers_factor': 3, 'stations_factor': 0, 'discharge': 0})\n",
    "states = env.reset()\n",
    "\n",
    "print('Observation space:', env.observation_space)\n",
    "print('Sample state:', {drone_index: env.format_state(state) for drone_index, state in states.items()})\n",
    "Image.fromarray(env.render(mode='rgb_array'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the agents\n",
    "agents = {drone.index: RandomAgent(env) for drone in env.drones}\n",
    "agents[0] = QLearningAgent(\n",
    "    env,\n",
    "    gamma=0.95, # Discount factor\n",
    "    alpha=0.1, # Learning rate\n",
    "    # Exploration rate\n",
    "    epsilon_start=1, epsilon_decay=0.99, epsilon_end=0.01\n",
    ")\n",
    "agents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train agents\n",
    "trainer = MultiAgentTrainer(env, agents, reset_agents=True, seed=0)\n",
    "trainer.train(5000)\n",
    "plot_rolling_rewards(trainer.rewards_log, drones_labels={0: 'Q-learning'})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rewards_log = test_agents(env, agents, n_steps=1000, seed=0)\n",
    "plot_cumulative_rewards(rewards_log, drones_labels={0: 'Q-learning'})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = os.path.join('videos', 'ql-compass-lidar-1st-try.mp4')\n",
    "render_video(env, agents, video_path=path, n_steps=120, fps=1, seed=0)\n",
    "ColabVideo(path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Issues with Q-learning\n",
    "---\n",
    "\n",
    "Two issues here\n",
    "\n",
    "* Sparse \"reward signal\": pickup rate is around 1%\n",
    "* No generalization: need to explore entire space!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "q_table = agents[0].get_qtable()\n",
    "print('Q-table:', q_table.shape)\n",
    "q_table.sample(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(agents[0].epsilons)\n",
    "plt.xlabel('Number of episodes')\n",
    "plt.ylabel('Exploration rate (epsilon)')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Possible solutions\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# (1/2) Sparse rewards: Create an intermediate \"pickup\" reward\n",
    "env.env_params.update({\n",
    "    'n_drones': 3, 'pickup_reward': 0.99, 'delivery_reward': 1,\n",
    "    'skyscrapers_factor': 3, 'stations_factor': 0, 'discharge': 0})\n",
    "states = env.reset()\n",
    "\n",
    "# (2/2) Train longer ..\n",
    "agents[0].epsilon = 1\n",
    "agents[0].epsilon_decay = 0.999\n",
    "\n",
    "set_seed(env, seed=0) # Make things deterministic\n",
    "trainer.train(30000)\n",
    "\n",
    "plot_rolling_rewards(\n",
    "    trainer.rewards_log,\n",
    "    events={'pickup': [0.99], 'delivery': [1], 'crash': [-1]},\n",
    "    drones_labels={0: 'Q-learning'})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(agents[0].epsilons)\n",
    "plt.xlabel('Number of episodes')\n",
    "plt.ylabel('Exploration rate (epsilon)')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rewards_log = test_agents(env, agents, n_steps=1000, seed=0)\n",
    "plot_cumulative_rewards(\n",
    "    rewards_log,\n",
    "    events={'pickup': [0.99], 'delivery': [1], 'crash': [-1]},\n",
    "    drones_labels={0: 'Q-learning'}\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Overfitting issues: try with different seeds\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rewards_log = test_agents(env, agents, n_steps=1000, seed=1)\n",
    "plot_cumulative_rewards(\n",
    "    rewards_log,\n",
    "    events={'pickup': [0.99], 'delivery': [1], 'crash': [-1]},\n",
    "    drones_labels={0: 'Q-learning'}\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Pick a good seed for your video ;)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = os.path.join('videos', 'ql-compass-lidar-2nd-try.mp4')\n",
    "render_video(env, agents, video_path=path, n_steps=120, fps=1, seed=1)\n",
    "ColabVideo(path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Q-learning limitations: discrete Q-table!\n",
    "---\n",
    "\n",
    "Let's try Q-learning with the full environment: skyscrapers + charge"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = LidarCompassChargeQTable(DeliveryDrones())\n",
    "env.env_params.update({\n",
    "    'n_drones': 3, 'pickup_reward': 0.99, 'delivery_reward': 1,\n",
    "    'discharge': 10, 'charge': 20, 'charge_reward': -0.1  # (default values)\n",
    "})\n",
    "states = env.reset()\n",
    "\n",
    "print('Observation space:', env.observation_space)\n",
    "print('Sample state:', env.format_state(states[0]))\n",
    "Image.fromarray(env.render(mode='rgb_array'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the agents\n",
    "agents = {drone.index: RandomAgent(env) for drone in env.drones}\n",
    "agents[0] = QLearningAgent(\n",
    "    env, gamma=0.95, alpha=0.1,\n",
    "    epsilon_start=1, epsilon_decay=0.999, epsilon_end=0.01\n",
    ")\n",
    "\n",
    "trainer = MultiAgentTrainer(env, agents, reset_agents=True, seed=0)\n",
    "trainer.train(35000)\n",
    "plot_rolling_rewards(trainer.rewards_log, events={'pickup': [0.99], 'delivery': [1], 'crash': [-1], 'charging': [-0.1]})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "q_table = agents[0].get_qtable()\n",
    "print('Q-table:', q_table.shape)\n",
    "q_table.sample(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Don't forget to test with different seeds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rewards_log = test_agents(env, agents, n_steps=1000, seed=0)\n",
    "plot_cumulative_rewards(\n",
    "    rewards_log,\n",
    "    events={'pickup': [0.99], 'delivery': [1], 'crash': [-1], 'charging': [-0.1]},\n",
    "    drones_labels={0: 'Q-learning'}\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = os.path.join('videos', 'ql-compass-lidar-charge.mp4')\n",
    "render_video(env, agents, video_path=path, n_steps=120, fps=1, seed=0)\n",
    "ColabVideo(path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First tests with deep Q-learning (DQN)\n",
    "---\n",
    "\n",
    "You can find a DQN implementation in `agents/`\n",
    "\n",
    "```\n",
    "agents/\n",
    "├── curiosity.py\n",
    "├── dqn.py          <-- DQN agent\n",
    "├── logging.py\n",
    "├── qlearning.py    <-- Q-learning agent\n",
    "└── random.py\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create environment\n",
    "env = LidarCompassChargeQTable(DeliveryDrones())\n",
    "env.env_params.update({\n",
    "    'n_drones': 3, 'pickup_reward': 0.99, 'delivery_reward': 1\n",
    "})\n",
    "states = env.reset()\n",
    "\n",
    "# Create the agents\n",
    "agents = {drone.index: RandomAgent(env) for drone in env.drones}\n",
    "agents[0] = DQNAgent(\n",
    "    env, DenseQNetworkFactory(env, hidden_layers=[256, 256]),\n",
    "    gamma=0.95, epsilon_start=1, epsilon_decay=0.999, epsilon_end=0.01,\n",
    "    memory_size=10000, batch_size=64, target_update_interval=5\n",
    ")\n",
    "trainer = MultiAgentTrainer(env, agents, reset_agents=True, seed=0)\n",
    "agents[0].qnetwork"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train the agent\n",
    "trainer.train(25000)\n",
    "plot_rolling_rewards(\n",
    "    trainer.rewards_log, drones_labels={0: 'DQN'},\n",
    "    events={'pickup': [0.99], 'delivery': [1], 'crash': [-1], 'charging': [-0.1]})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(agents[0].epsilons)\n",
    "plt.xlabel('Number of episodes')\n",
    "plt.ylabel('Exploration rate (epsilon)')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Don't forget to test with different seeds!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rewards_log = test_agents(env, agents, n_steps=1000, seed=0)\n",
    "plot_cumulative_rewards(\n",
    "    rewards_log, drones_labels={0: 'DQN'},\n",
    "    events={'pickup': [0.99], 'delivery': [1], 'crash': [-1], 'charging': [-0.1]})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Inspect replay memory buffer\n",
    "agents[0].inspect_memory(top_n=10, max_col=80)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Take a moment to play with the different parameters: `memory_size`, `batch_size`, `target_update_interval` and the others"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = os.path.join('videos', 'dqn-compass-lidar-charge.mp4')\n",
    "render_video(env, agents, video_path=path, n_steps=120, fps=1, seed=0)\n",
    "ColabVideo(path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "DQN and WindowedGrid (official environment)\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create environment\n",
    "env = WindowedGridView(DeliveryDrones(), radius=3)\n",
    "env.env_params.update({\n",
    "    'n_drones': 3, 'pickup_reward': 0.99, 'delivery_reward': 1\n",
    "})\n",
    "states = env.reset()\n",
    "\n",
    "# Create the agents\n",
    "agents = {drone.index: RandomAgent(env) for drone in env.drones}\n",
    "agents[0] = my_agent = DQNAgent(\n",
    "    env, ConvQNetworkFactory(env, conv_layers=[\n",
    "        {'out_channels': 32, 'kernel_size': 3, 'stride': 1, 'padding': 1},\n",
    "        {'out_channels': 32, 'kernel_size': 3, 'stride': 1, 'padding': 1},\n",
    "        {'out_channels': 32, 'kernel_size': 3, 'stride': 1, 'padding': 1},\n",
    "        {'out_channels': 64, 'kernel_size': 3, 'stride': 1, 'padding': 1},\n",
    "        {'out_channels': 64, 'kernel_size': 3, 'stride': 1, 'padding': 1},\n",
    "        {'out_channels': 64, 'kernel_size': 3, 'stride': 1, 'padding': 1},\n",
    "    ], dense_layers=[1024, 256]),\n",
    "    gamma=0.95, epsilon_start=1, epsilon_decay=0.99, epsilon_end=0.01,\n",
    "    memory_size=10000, batch_size=64, target_update_interval=5\n",
    ")\n",
    "trainer = MultiAgentTrainer(env, agents, reset_agents=True, seed=0)\n",
    "agents[0].qnetwork"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train the agent\n",
    "for run in range(10):\n",
    "  trainer.train(2500)\n",
    "  plot_rolling_rewards(\n",
    "      trainer.rewards_log, drones_labels={0: 'DQN'},\n",
    "      events={'pickup': [0.99], 'delivery': [1], 'crash': [-1], 'charging': [-0.1]})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = os.path.join('videos', 'dqn-windowed.mp4')\n",
    "render_video(env, agents, video_path=path, n_steps=120, fps=1, seed=0)\n",
    "ColabVideo(path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Share your agent q-network\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = os.path.join('agents', 'dqn-agent.pt')\n",
    "agents[0].save(path)\n",
    "# agents[0].load(path) # Later, load the qnetwork!"
   ]
  }
 ],
 "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.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}