3_multiprocessing.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/araffin/rl-tutorial-jnrr19/blob/sb3/3_multiprocessing.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "KnPeMWYi0vAx"
   },
   "source": [
    "# Stable Baselines3 Tutorial - Multiprocessing of environments\n",
    "\n",
    "Github repo: https://github.com/araffin/rl-tutorial-jnrr19/tree/sb3/\n",
    "\n",
    "Stable-Baselines3: https://github.com/DLR-RM/stable-baselines3\n",
    "\n",
    "Documentation: https://stable-baselines3.readthedocs.io/en/master/\n",
    "\n",
    "SB3-Contrib: https://github.com/Stable-Baselines-Team/stable-baselines3-contrib\n",
    "\n",
    "RL Baselines3 zoo: https://github.com/DLR-RM/rl-baselines3-zoo\n",
    "\n",
    "\n",
    "## Introduction\n",
    "\n",
    "In this notebook, you will learn how to use *Vectorized Environments* (aka multiprocessing) to make training faster. You will also see that this speed up comes at a cost of sample efficiency.\n",
    "\n",
    "## Install Dependencies and Stable Baselines3 Using Pip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for autoformatting\n",
    "# %load_ext jupyter_black"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ClRYNMkVvpUX"
   },
   "outputs": [],
   "source": [
    "!apt install swig\n",
    "!pip install \"stable-baselines3[extra]>=2.0.0a4\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "OQunADhw1EXX"
   },
   "source": [
    "## Vectorized Environments and Imports\n",
    "\n",
    "[Vectorized Environments](https://stable-baselines3.readthedocs.io/en/master/guide/vec_envs.html) are a method for stacking multiple independent environments into a single environment. Instead of training an RL agent on 1 environment per step, it allows us to train it on n environments per step. This provides two benefits:\n",
    "* Agent experience can be collected more quickly\n",
    "* The experience will contain a more diverse range of states, it usually improves exploration\n",
    "\n",
    "Stable-Baselines provides two types of Vectorized Environment:\n",
    "- SubprocVecEnv which run each environment in a separate process\n",
    "- DummyVecEnv which run all environment on the same process\n",
    "\n",
    "In practice, DummyVecEnv is usually faster than SubprocVecEnv because of communication delays that subprocesses have."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "AvO5BGrVv2Rk"
   },
   "outputs": [],
   "source": [
    "import time\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "%matplotlib inline\n",
    "\n",
    "import gymnasium as gym\n",
    "\n",
    "from stable_baselines3.common.vec_env import DummyVecEnv, SubprocVecEnv\n",
    "from stable_baselines3.common.utils import set_random_seed\n",
    "from stable_baselines3 import PPO, A2C"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "JcdG_UZS1-yO"
   },
   "source": [
    "Import evaluate function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "NHslfVkuwALj"
   },
   "outputs": [],
   "source": [
    "from stable_baselines3.common.evaluation import evaluate_policy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "WWsIT2vP2FzB"
   },
   "source": [
    "## Define an environment function\n",
    "\n",
    "The multiprocessing implementation requires a function that can be called inside the process to instantiate a gym env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "6S95WiPGwF6z"
   },
   "outputs": [],
   "source": [
    "def make_env(env_id, rank, seed=0):\n",
    "    \"\"\"\n",
    "    Utility function for multiprocessed env.\n",
    "\n",
    "    :param env_id: (str) the environment ID\n",
    "    :param seed: (int) the inital seed for RNG\n",
    "    :param rank: (int) index of the subprocess\n",
    "    \"\"\"\n",
    "\n",
    "    def _init():\n",
    "        env = gym.make(env_id)\n",
    "        # use a seed for reproducibility\n",
    "        # Important: use a different seed for each environment\n",
    "        # otherwise they would generate the same experiences\n",
    "        env.reset(seed=seed + rank)\n",
    "        return env\n",
    "\n",
    "    set_random_seed(seed)\n",
    "    return _init"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "9-QID4O2bd7c"
   },
   "source": [
    "Stable-Baselines also provides directly an helper to create vectorized environment:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Gk7Ukbqlbl-i"
   },
   "outputs": [],
   "source": [
    "from stable_baselines3.common.env_util import make_vec_env"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "DJUP0PQi2WEE"
   },
   "source": [
    "## Define a few constants (feel free to try out other environments and algorithms)\n",
    "We will be using the Cartpole environment: [https://gym.openai.com/envs/CartPole-v1/](https://gym.openai.com/envs/CartPole-v1/)\n",
    "\n",
    "![Cartpole](https://cdn-images-1.medium.com/max/1143/1*h4WTQNVIsvMXJTCpXm_TAw.gif)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "bmdNV8UVwTht"
   },
   "outputs": [],
   "source": [
    "env_id = \"CartPole-v1\"\n",
    "# The different number of processes that will be used\n",
    "PROCESSES_TO_TEST = [1, 2, 4, 8, 16]\n",
    "NUM_EXPERIMENTS = 3  # RL algorithms can often be unstable, so we run several experiments (see https://arxiv.org/abs/1709.06560)\n",
    "TRAIN_STEPS = 5000\n",
    "# Number of episodes for evaluation\n",
    "EVAL_EPS = 20\n",
    "ALGO = A2C\n",
    "\n",
    "# We will create one environment to evaluate the agent on\n",
    "eval_env = gym.make(env_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "y08bJGxj2ezh"
   },
   "source": [
    "## Iterate through the different numbers of processes\n",
    "\n",
    "For each processes, several experiments are run per process\n",
    "This may take a couple of minutes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 106
    },
    "colab_type": "code",
    "id": "kcYpsA8ExB9T",
    "outputId": "11e28f5c-c3d3-4669-ab4b-acff3e710ac1"
   },
   "outputs": [],
   "source": [
    "reward_averages = []\n",
    "reward_std = []\n",
    "training_times = []\n",
    "total_procs = 0\n",
    "for n_procs in PROCESSES_TO_TEST:\n",
    "    total_procs += n_procs\n",
    "    print(f\"Running for n_procs = {n_procs}\")\n",
    "    if n_procs == 1:\n",
    "        # if there is only one process, there is no need to use multiprocessing\n",
    "        train_env = DummyVecEnv([lambda: gym.make(env_id)])\n",
    "    else:\n",
    "        # Here we use the \"fork\" method for launching the processes, more information is available in the doc\n",
    "        # This is equivalent to make_vec_env(env_id, n_envs=n_procs, vec_env_cls=SubprocVecEnv, vec_env_kwargs=dict(start_method='fork'))\n",
    "        train_env = SubprocVecEnv(\n",
    "            [make_env(env_id, i + total_procs) for i in range(n_procs)],\n",
    "            start_method=\"fork\",\n",
    "        )\n",
    "\n",
    "    rewards = []\n",
    "    times = []\n",
    "\n",
    "    for experiment in range(NUM_EXPERIMENTS):\n",
    "        # it is recommended to run several experiments due to variability in results\n",
    "        train_env.reset()\n",
    "        model = ALGO(\"MlpPolicy\", train_env, verbose=0)\n",
    "        start = time.time()\n",
    "        model.learn(total_timesteps=TRAIN_STEPS)\n",
    "        times.append(time.time() - start)\n",
    "        mean_reward, _ = evaluate_policy(model, eval_env, n_eval_episodes=EVAL_EPS)\n",
    "        rewards.append(mean_reward)\n",
    "    # Important: when using subprocesses, don't forget to close them\n",
    "    # otherwise, you may have memory issues when running a lot of experiments\n",
    "    train_env.close()\n",
    "    reward_averages.append(np.mean(rewards))\n",
    "    reward_std.append(np.std(rewards))\n",
    "    training_times.append(np.mean(times))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "2z5paN1q3AaC"
   },
   "source": [
    "## Plot the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "CGnZ8SccKG4D"
   },
   "outputs": [],
   "source": [
    "def plot_training_results(training_steps_per_second, reward_averages, reward_std):\n",
    "    \"\"\"\n",
    "    Utility function for plotting the results of training\n",
    "\n",
    "    :param training_steps_per_second: List[double]\n",
    "    :param reward_averages: List[double]\n",
    "    :param reward_std: List[double]\n",
    "    \"\"\"\n",
    "    plt.figure(figsize=(9, 4))\n",
    "    plt.subplots_adjust(wspace=0.5)\n",
    "    plt.subplot(1, 2, 1)\n",
    "    plt.errorbar(\n",
    "        PROCESSES_TO_TEST,\n",
    "        reward_averages,\n",
    "        yerr=reward_std,\n",
    "        capsize=2,\n",
    "        c=\"k\",\n",
    "        marker=\"o\",\n",
    "    )\n",
    "    plt.xlabel(\"Processes\")\n",
    "    plt.ylabel(\"Average return\")\n",
    "    plt.subplot(1, 2, 2)\n",
    "    plt.bar(range(len(PROCESSES_TO_TEST)), training_steps_per_second)\n",
    "    plt.xticks(range(len(PROCESSES_TO_TEST)), PROCESSES_TO_TEST)\n",
    "    plt.xlabel(\"Processes\")\n",
    "    plt.ylabel(\"Training steps per second\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 279
    },
    "colab_type": "code",
    "id": "fPWfc96JxT-k",
    "outputId": "df2b74d5-61ea-487b-9364-8ec33b4e0624"
   },
   "outputs": [],
   "source": [
    "training_steps_per_second = [TRAIN_STEPS / t for t in training_times]\n",
    "\n",
    "plot_training_results(training_steps_per_second, reward_averages, reward_std)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "R5xE8EX63PO9"
   },
   "source": [
    "## Sample efficiency vs wall clock time trade-off\n",
    "There is clearly a trade-off between sample efficiency, diverse experience and wall clock time. Let's try getting the best performance in a fixed amount of time, say 10 seconds per experiment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 106
    },
    "outputId": "2cdda2c8-e2f0-401b-a6ea-99c80d91fe8e"
   },
   "outputs": [],
   "source": [
    "SECONDS_PER_EXPERIMENT = 10\n",
    "steps_per_experiment = [int(SECONDS_PER_EXPERIMENT * fps) for fps in training_steps_per_second]\n",
    "reward_averages = []\n",
    "reward_std = []\n",
    "training_times = []\n",
    "\n",
    "for n_procs, train_steps in zip(PROCESSES_TO_TEST, steps_per_experiment):\n",
    "    total_procs += n_procs\n",
    "    print(f\"Running for n_procs = {n_procs} for steps = {train_steps}\")\n",
    "    if n_procs == 1:\n",
    "        # if there is only one process, there is no need to use multiprocessing\n",
    "        train_env = DummyVecEnv([lambda: gym.make(env_id)])\n",
    "    else:\n",
    "        train_env = SubprocVecEnv([make_env(env_id, i+total_procs) for i in range(n_procs)], start_method=\"spawn\")\n",
    "        # Alternatively, you can use a DummyVecEnv if the communication delays is the bottleneck\n",
    "        # train_env = DummyVecEnv([make_env(env_id, i+total_procs) for i in range(n_procs)])\n",
    "\n",
    "    rewards = []\n",
    "    times = []\n",
    "\n",
    "    for experiment in range(NUM_EXPERIMENTS):\n",
    "        # it is recommended to run several experiments due to variability in results\n",
    "        train_env.reset()\n",
    "        model = ALGO(\"MlpPolicy\", train_env, verbose=0)\n",
    "        start = time.time()\n",
    "        model.learn(total_timesteps=train_steps)\n",
    "        times.append(time.time() - start)\n",
    "        mean_reward, _ = evaluate_policy(model, eval_env, n_eval_episodes=EVAL_EPS)\n",
    "        rewards.append(mean_reward)\n",
    "\n",
    "    train_env.close()\n",
    "    reward_averages.append(np.mean(rewards))\n",
    "    reward_std.append(np.std(rewards))\n",
    "    training_times.append(np.mean(times))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "G7a7ZiVw5A11"
   },
   "source": [
    "## Plot the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 297
    },
    "colab_type": "code",
    "id": "EQXJ1hI46DVB",
    "outputId": "d5b47716-3551-47b1-f690-16d726e89a05"
   },
   "outputs": [],
   "source": [
    "training_steps_per_second = [s / t for s,t in zip(steps_per_experiment, training_times)]\n",
    "\n",
    "plot_training_results(training_steps_per_second, reward_averages, reward_std)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "0FcOcVf5rY3C"
   },
   "source": [
    "## DummyVecEnv vs SubprocVecEnv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 106
    },
    "colab_type": "code",
    "id": "MebaTHQvqhoH",
    "outputId": "637e9934-e6b1-4ce3-a401-c20f23437e67"
   },
   "outputs": [],
   "source": [
    "reward_averages = []\n",
    "reward_std = []\n",
    "training_times = []\n",
    "total_procs = 0\n",
    "for n_procs in PROCESSES_TO_TEST:\n",
    "    total_procs += n_procs\n",
    "    print(f'Running for n_procs = {n_procs}'))\n",
    "    # Here we are using only one process even for n_env > 1\n",
    "    # this is equivalent to DummyVecEnv([make_env(env_id, i + total_procs) for i in range(n_procs)])\n",
    "    train_env = make_vec_env(env_id, n_envs=n_procs)\n",
    "\n",
    "    rewards = []\n",
    "    times = []\n",
    "\n",
    "    for experiment in range(NUM_EXPERIMENTS):\n",
    "        # it is recommended to run several experiments due to variability in results\n",
    "        train_env.reset()\n",
    "        model = ALGO(\"MlpPolicy\", train_env, verbose=0)\n",
    "        start = time.time()\n",
    "        model.learn(total_timesteps=TRAIN_STEPS)\n",
    "        times.append(time.time() - start)\n",
    "        mean_reward, _ = evaluate_policy(model, eval_env, n_eval_episodes=EVAL_EPS)\n",
    "        rewards.append(mean_reward)\n",
    "\n",
    "    train_env.close()\n",
    "    reward_averages.append(np.mean(rewards))\n",
    "    reward_std.append(np.std(rewards))\n",
    "    training_times.append(np.mean(times))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 297
    },
    "colab_type": "code",
    "id": "kmMr_c1hqmoi",
    "outputId": "cc174025-ed75-4897-f745-c08944493366"
   },
   "outputs": [],
   "source": [
    "training_steps_per_second = [TRAIN_STEPS / t for t in training_times]\n",
    "\n",
    "plot_training_results(training_steps_per_second, reward_averages, reward_std)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "e9PNbT35spZW"
   },
   "source": [
    "### What's happening?\n",
    "\n",
    "It seems that having only one process for n environments is faster in our case.\n",
    "In practice, the bottleneck does not come from the environment computation, but from synchronisation and communication between processes. To learn more about that problem, you can start [here](https://github.com/hill-a/stable-baselines/issues/322#issuecomment-492202915)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "GlcJPYN-6ebp"
   },
   "source": [
    "## Conclusions\n",
    "This notebook has highlighted some of the pros and cons of multiprocessing. It is worth mentioning that colab notebooks only provide two CPU cores per process, so we do not see a linear scaling of the FPS of the environments. State of the art Deep RL research has scaled parallel processing to tens of thousands of CPU cores, [OpenAI RAPID](https://openai.com/blog/how-to-train-your-openai-five/) [IMPALA](https://arxiv.org/abs/1802.01561).\n",
    "\n",
    "Do you think this direction of research is transferable to real world robots / intelligent agents?\n",
    "\n",
    "Things to try:\n",
    "* Another algorithm / environment.\n",
    "* Increase the number of experiments.\n",
    "* Train for more iterations.\n"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "include_colab_link": true,
   "name": "3_multiprocessing.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.10.9"
  },
  "vscode": {
   "interpreter": {
    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}