1_getting_started.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "1_getting_started.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/araffin/rl-tutorial-jnrr19/blob/sb3/1_getting_started.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hyyN-2qyK_T2",
        "colab_type": "text"
      },
      "source": [
        "# Stable Baselines3 Tutorial - Getting Started\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",
        "RL Baselines3 zoo: https://github.com/DLR-RM/rl-baselines3-zoo\n",
        "\n",
        "[RL Baselines3 Zoo](https://github.com/DLR-RM/rl-baselines3-zoo) is a collection of pre-trained Reinforcement Learning agents using Stable-Baselines3.\n",
        "\n",
        "It also provides basic scripts for training, evaluating agents, tuning hyperparameters and recording videos.\n",
        "\n",
        "\n",
        "## Introduction\n",
        "\n",
        "In this notebook, you will learn the basics for using stable baselines library: how to create a RL model, train it and evaluate it. Because all algorithms share the same interface, we will see how simple it is to switch from one algorithm to another.\n",
        "\n",
        "\n",
        "## Install Dependencies and Stable Baselines3 Using Pip\n",
        "\n",
        "List of full dependencies can be found in the [README](https://github.com/DLR-RM/stable-baselines3).\n",
        "\n",
        "\n",
        "```\n",
        "pip install stable-baselines3[extra]\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gWskDE2c9WoN",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 784
        },
        "outputId": "03477445-4249-49c3-ddba-4e12df09e98e"
      },
      "source": [
        "!apt-get install ffmpeg freeglut3-dev xvfb  # For visualization\n",
        "!pip install stable-baselines3[extra]"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Reading package lists... Done\n",
            "Building dependency tree       \n",
            "Reading state information... Done\n",
            "freeglut3-dev is already the newest version (2.8.1-3).\n",
            "freeglut3-dev set to manually installed.\n",
            "ffmpeg is already the newest version (7:3.4.6-0ubuntu0.18.04.1).\n",
            "The following NEW packages will be installed:\n",
            "  xvfb\n",
            "0 upgraded, 1 newly installed, 0 to remove and 31 not upgraded.\n",
            "Need to get 784 kB of archives.\n",
            "After this operation, 2,266 kB of additional disk space will be used.\n",
            "Get:1 http://archive.ubuntu.com/ubuntu bionic-updates/universe amd64 xvfb amd64 2:1.19.6-1ubuntu4.4 [784 kB]\n",
            "Fetched 784 kB in 0s (7,570 kB/s)\n",
            "Selecting previously unselected package xvfb.\n",
            "(Reading database ... 144433 files and directories currently installed.)\n",
            "Preparing to unpack .../xvfb_2%3a1.19.6-1ubuntu4.4_amd64.deb ...\n",
            "Unpacking xvfb (2:1.19.6-1ubuntu4.4) ...\n",
            "Setting up xvfb (2:1.19.6-1ubuntu4.4) ...\n",
            "Processing triggers for man-db (2.8.3-2ubuntu0.1) ...\n",
            "Collecting stable-baselines3[extra]\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/bb/2a/85e3540a0e497bec9ca97f18bab3c9af86809362fe68362964a4e711e44f/stable_baselines3-0.6.0a8-py3-none-any.whl (102kB)\n",
            "\u001b[K     |████████████████████████████████| 112kB 6.3MB/s \n",
            "\u001b[?25hRequirement already satisfied: cloudpickle in /usr/local/lib/python3.6/dist-packages (from stable-baselines3[extra]) (1.3.0)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from stable-baselines3[extra]) (1.18.4)\n",
            "Requirement already satisfied: torch>=1.4.0 in /usr/local/lib/python3.6/dist-packages (from stable-baselines3[extra]) (1.5.0+cu101)\n",
            "Requirement already satisfied: matplotlib in /usr/local/lib/python3.6/dist-packages (from stable-baselines3[extra]) (3.2.1)\n",
            "Requirement already satisfied: gym>=0.17 in /usr/local/lib/python3.6/dist-packages (from stable-baselines3[extra]) (0.17.1)\n",
            "Requirement already satisfied: pandas in /usr/local/lib/python3.6/dist-packages (from stable-baselines3[extra]) (1.0.3)\n",
            "Requirement already satisfied: atari-py~=0.2.0; extra == \"extra\" in /usr/local/lib/python3.6/dist-packages (from stable-baselines3[extra]) (0.2.6)\n",
            "Requirement already satisfied: pillow; extra == \"extra\" in /usr/local/lib/python3.6/dist-packages (from stable-baselines3[extra]) (7.0.0)\n",
            "Requirement already satisfied: opencv-python; extra == \"extra\" in /usr/local/lib/python3.6/dist-packages (from stable-baselines3[extra]) (4.1.2.30)\n",
            "Requirement already satisfied: future in /usr/local/lib/python3.6/dist-packages (from torch>=1.4.0->stable-baselines3[extra]) (0.16.0)\n",
            "Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.6/dist-packages (from matplotlib->stable-baselines3[extra]) (0.10.0)\n",
            "Requirement already satisfied: pyparsing!=2.0.4,!=2.1.2,!=2.1.6,>=2.0.1 in /usr/local/lib/python3.6/dist-packages (from matplotlib->stable-baselines3[extra]) (2.4.7)\n",
            "Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.6/dist-packages (from matplotlib->stable-baselines3[extra]) (1.2.0)\n",
            "Requirement already satisfied: python-dateutil>=2.1 in /usr/local/lib/python3.6/dist-packages (from matplotlib->stable-baselines3[extra]) (2.8.1)\n",
            "Requirement already satisfied: pyglet<=1.5.0,>=1.4.0 in /usr/local/lib/python3.6/dist-packages (from gym>=0.17->stable-baselines3[extra]) (1.5.0)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.6/dist-packages (from gym>=0.17->stable-baselines3[extra]) (1.12.0)\n",
            "Requirement already satisfied: scipy in /usr/local/lib/python3.6/dist-packages (from gym>=0.17->stable-baselines3[extra]) (1.4.1)\n",
            "Requirement already satisfied: pytz>=2017.2 in /usr/local/lib/python3.6/dist-packages (from pandas->stable-baselines3[extra]) (2018.9)\n",
            "Installing collected packages: stable-baselines3\n",
            "Successfully installed stable-baselines3-0.6.0a8\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FtY8FhliLsGm",
        "colab_type": "text"
      },
      "source": [
        "## Imports"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gcX8hEcaUpR0",
        "colab_type": "text"
      },
      "source": [
        "Stable-Baselines3 works on environments that follow the [gym interface](https://stable-baselines3.readthedocs.io/en/master/guide/custom_env.html).\n",
        "You can find a list of available environment [here](https://gym.openai.com/envs/#classic_control).\n",
        "\n",
        "It is also recommended to check the [source code](https://github.com/openai/gym) to learn more about the observation and action space of each env, as gym does not have a proper documentation.\n",
        "Not all algorithms can work with all action spaces, you can find more in this [recap table](https://stable-baselines3.readthedocs.io/en/master/guide/algos.html)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BIedd7Pz9sOs",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import gym\n",
        "import numpy as np"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ae32CtgzTG3R",
        "colab_type": "text"
      },
      "source": [
        "The first thing you need to import is the RL model, check the documentation to know what you can use on which problem"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "R7tKaBFrTR0a",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from stable_baselines3 import PPO"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-0_8OQbOTTNT",
        "colab_type": "text"
      },
      "source": [
        "The next thing you need to import is the policy class that will be used to create the networks (for the policy/value functions).\n",
        "This step is optional as you can directly use strings in the constructor: \n",
        "\n",
        "```PPO('MlpPolicy', env)``` instead of ```PPO(MlpPolicy, env)```\n",
        "\n",
        "Note that some algorithms like `SAC` have their own `MlpPolicy`, that's why using string for the policy is the recommened option."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ROUJr675TT01",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from stable_baselines3.ppo.policies import MlpPolicy"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RapkYvTXL7Cd",
        "colab_type": "text"
      },
      "source": [
        "## Create the Gym env and instantiate the agent\n",
        "\n",
        "For this example, we will use CartPole environment, a classic control problem.\n",
        "\n",
        "\"A pole is attached by an un-actuated joint to a cart, which moves along a frictionless track. The system is controlled by applying a force of +1 or -1 to the cart. The pendulum starts upright, and the goal is to prevent it from falling over. A reward of +1 is provided for every timestep that the pole remains upright. \"\n",
        "\n",
        "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",
        "\n",
        "We chose the MlpPolicy because the observation of the CartPole task is a feature vector, not images.\n",
        "\n",
        "The type of action to use (discrete/continuous) will be automatically deduced from the environment action space\n",
        "\n",
        "Here we are using the [Proximal Policy Optimization](https://stable-baselines3.readthedocs.io/en/master/modules/ppo2.html) algorithm, which is an Actor-Critic method: it uses a value function to improve the policy gradient descent (by reducing the variance).\n",
        "\n",
        "It combines ideas from [A2C](https://stable-baselines3.readthedocs.io/en/master/modules/a2c.html) (having multiple workers and using an entropy bonus for exploration) and [TRPO](https://stable-baselines.readthedocs.io/en/master/modules/trpo.html) (it uses a trust region to improve stability and avoid catastrophic drops in performance).\n",
        "\n",
        "PPO is an on-policy algorithm, which means that the trajectories used to update the networks must be collected using the latest policy.\n",
        "It is usually less sample efficient than off-policy alorithms like [DQN](https://stable-baselines.readthedocs.io/en/master/modules/dqn.html), [SAC](https://stable-baselines3.readthedocs.io/en/master/modules/sac.html) or [TD3](https://stable-baselines3.readthedocs.io/en/master/modules/td3.html), but is much faster regarding wall-clock time.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pUWGZp3i9wyf",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "env = gym.make('CartPole-v1')\n",
        "\n",
        "model = PPO(MlpPolicy, env, verbose=0)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4efFdrQ7MBvl",
        "colab_type": "text"
      },
      "source": [
        "We create a helper function to evaluate the agent:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "63M8mSKR-6Zt",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def evaluate(model, num_episodes=100):\n",
        "    \"\"\"\n",
        "    Evaluate a RL agent\n",
        "    :param model: (BaseRLModel object) the RL Agent\n",
        "    :param num_episodes: (int) number of episodes to evaluate it\n",
        "    :return: (float) Mean reward for the last num_episodes\n",
        "    \"\"\"\n",
        "    # This function will only work for a single Environment\n",
        "    env = model.get_env()\n",
        "    all_episode_rewards = []\n",
        "    for i in range(num_episodes):\n",
        "        episode_rewards = []\n",
        "        done = False\n",
        "        obs = env.reset()\n",
        "        while not done:\n",
        "            # _states are only useful when using LSTM policies\n",
        "            action, _states = model.predict(obs)\n",
        "            # here, action, rewards and dones are arrays\n",
        "            # because we are using vectorized env\n",
        "            obs, reward, done, info = env.step(action)\n",
        "            episode_rewards.append(reward)\n",
        "\n",
        "        all_episode_rewards.append(sum(episode_rewards))\n",
        "\n",
        "    mean_episode_reward = np.mean(all_episode_rewards)\n",
        "    print(\"Mean reward:\", mean_episode_reward, \"Num episodes:\", num_episodes)\n",
        "\n",
        "    return mean_episode_reward"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zjEVOIY8NVeK",
        "colab_type": "text"
      },
      "source": [
        "Let's evaluate the un-trained agent, this should be a random agent."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xDHLMA6NFk95",
        "colab_type": "code",
        "outputId": "231b2170-a607-48ed-e9d9-daef596f6384",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        }
      },
      "source": [
        "# Random Agent, before training\n",
        "mean_reward_before_train = evaluate(model, num_episodes=100)"
      ],
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Mean reward: 23.6 Num episodes: 100\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QjjPxrwkYJ2i",
        "colab_type": "text"
      },
      "source": [
        "Stable-Baselines already provides you with that helper:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8z6K9YImYJEx",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from stable_baselines3.common.evaluation import evaluate_policy"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4oPTHjxyZSOL",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "mean_reward, std_reward = evaluate_policy(model, env, n_eval_episodes=100)\n",
        "\n",
        "print(f\"mean_reward:{mean_reward:.2f} +/- {std_reward:.2f}\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r5UoXTZPNdFE",
        "colab_type": "text"
      },
      "source": [
        "## Train the agent and evaluate it"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "e4cfSXIB-pTF",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Train the agent for 10000 steps\n",
        "model.learn(total_timesteps=10000)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ygl_gVmV_QP7",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Evaluate the trained agent\n",
        "mean_reward, std_reward = evaluate_policy(model, env, n_eval_episodes=100)\n",
        "\n",
        "print(f\"mean_reward:{mean_reward:.2f} +/- {std_reward:.2f}\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A00W6yY3NkHG",
        "colab_type": "text"
      },
      "source": [
        "Apparently the training went well, the mean reward increased a lot ! "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xVm9QPNVwKXN",
        "colab_type": "text"
      },
      "source": [
        "### Prepare video recording"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MPyfQxD5z26J",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Set up fake display; otherwise rendering will fail\n",
        "import os\n",
        "os.system(\"Xvfb :1 -screen 0 1024x768x24 &\")\n",
        "os.environ['DISPLAY'] = ':1'"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SLzXxO8VMD6N",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import base64\n",
        "from pathlib import Path\n",
        "\n",
        "from IPython import display as ipythondisplay\n",
        "\n",
        "def show_videos(video_path='', prefix=''):\n",
        "  \"\"\"\n",
        "  Taken from https://github.com/eleurent/highway-env\n",
        "\n",
        "  :param video_path: (str) Path to the folder containing videos\n",
        "  :param prefix: (str) Filter the video, showing only the only starting with this prefix\n",
        "  \"\"\"\n",
        "  html = []\n",
        "  for mp4 in Path(video_path).glob(\"{}*.mp4\".format(prefix)):\n",
        "      video_b64 = base64.b64encode(mp4.read_bytes())\n",
        "      html.append('''<video alt=\"{}\" autoplay \n",
        "                    loop controls style=\"height: 400px;\">\n",
        "                    <source src=\"data:video/mp4;base64,{}\" type=\"video/mp4\" />\n",
        "                </video>'''.format(mp4, video_b64.decode('ascii')))\n",
        "  ipythondisplay.display(ipythondisplay.HTML(data=\"<br>\".join(html)))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LTRNUfulOGaF",
        "colab_type": "text"
      },
      "source": [
        "We will record a video using the [VecVideoRecorder](https://stable-baselines3.readthedocs.io/en/master/guide/vec_envs.html#vecvideorecorder) wrapper, you will learn about those wrapper in the next notebook."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Trag9dQpOIhx",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from stable_baselines3.common.vec_env import VecVideoRecorder, DummyVecEnv\n",
        "\n",
        "def record_video(env_id, model, video_length=500, prefix='', video_folder='videos/'):\n",
        "  \"\"\"\n",
        "  :param env_id: (str)\n",
        "  :param model: (RL model)\n",
        "  :param video_length: (int)\n",
        "  :param prefix: (str)\n",
        "  :param video_folder: (str)\n",
        "  \"\"\"\n",
        "  eval_env = DummyVecEnv([lambda: gym.make(env_id)])\n",
        "  # Start the video at step=0 and record 500 steps\n",
        "  eval_env = VecVideoRecorder(eval_env, video_folder=video_folder,\n",
        "                              record_video_trigger=lambda step: step == 0, video_length=video_length,\n",
        "                              name_prefix=prefix)\n",
        "\n",
        "  obs = eval_env.reset()\n",
        "  for _ in range(video_length):\n",
        "    action, _ = model.predict(obs)\n",
        "    obs, _, _, _ = eval_env.step(action)\n",
        "\n",
        "  # Close the video recorder\n",
        "  eval_env.close()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KOObbeu5MMlR",
        "colab_type": "text"
      },
      "source": [
        "### Visualize trained agent\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iATu7AiyMQW2",
        "colab_type": "code",
        "outputId": "68acb027-6c94-4389-8456-2cfb11494814",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        }
      },
      "source": [
        "record_video('CartPole-v1', model, video_length=500, prefix='ppo2-cartpole')"
      ],
      "execution_count": 17,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Saving video to  /content/videos/ppo2-cartpole-step-0-to-step-500.mp4\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-n4i-fW3NojZ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "show_videos('videos', prefix='ppo2')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9Y8zg4V566qD",
        "colab_type": "text"
      },
      "source": [
        "## Bonus: Train a RL Model in One Line\n",
        "\n",
        "The policy class to use will be inferred and the environment will be automatically created. This works because both are [registered](https://stable-baselines3.readthedocs.io/en/master/guide/quickstart.html)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iaOPfOrwWEP4",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "model = PPO('MlpPolicy', \"CartPole-v1\", verbose=1).learn(1000)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FrI6f5fWnzp-",
        "colab_type": "text"
      },
      "source": [
        "## Conclusion\n",
        "\n",
        "In this notebook we have seen:\n",
        "- how to define and train a RL model using stable baselines3, it takes only one line of code ;)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "73ji3gbNDkf7",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}