stable_baselines_getting_started.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/Stable-Baselines-Team/rl-colab-notebooks/blob/sb3/stable_baselines_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": {
    "colab_type": "text",
    "id": "hyyN-2qyK_T2"
   },
   "source": [
    "# Stable Baselines3 Tutorial - Getting Started\n",
    "\n",
    "Github repo: https://github.com/araffin/rl-tutorial-jnrr19\n",
    "\n",
    "Stable-Baselines3: https://github.com/DLR-RM/stable-baselines3\n",
    "\n",
    "Documentation: https://stable-baselines.readthedocs.io/en/master/\n",
    "\n",
    "RL Baselines3 zoo: https://github.com/DLR-RM/rl-baselines3-zoo\n",
    "\n",
    "\n",
    "[RL Baselines3 Zoo](https://github.com/DLR-RM/rl-baselines3-zoo) is a training framework for Reinforcement Learning (RL), using Stable Baselines3.\n",
    "\n",
    "It provides scripts for training, evaluating agents, tuning hyperparameters, plotting results and recording videos.\n",
    "\n",
    "## Introduction\n",
    "\n",
    "In this notebook, you will learn the basics for using stable baselines3 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",
   "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": "gWskDE2c9WoN"
   },
   "outputs": [],
   "source": [
    "!apt-get update && apt-get install ffmpeg freeglut3-dev xvfb  # For visualization\n",
    "!pip install \"stable-baselines3[extra]>=2.0.0a4\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "U29X1-B-AIKE"
   },
   "outputs": [],
   "source": [
    "import stable_baselines3\n",
    "\n",
    "print(f\"{stable_baselines3.__version__=}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "FtY8FhliLsGm"
   },
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "gcX8hEcaUpR0"
   },
   "source": [
    "Stable-Baselines works on environments that follow the [gym interface](https://stable-baselines.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-baselines.readthedocs.io/en/master/guide/algos.html)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "BIedd7Pz9sOs"
   },
   "outputs": [],
   "source": [
    "import gymnasium as gym\n",
    "import numpy as np\n",
    "\n",
    "print(f\"{gym.__version__=}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Ae32CtgzTG3R"
   },
   "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",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "R7tKaBFrTR0a"
   },
   "outputs": [],
   "source": [
    "from stable_baselines3 import PPO"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "-0_8OQbOTTNT"
   },
   "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",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ROUJr675TT01"
   },
   "outputs": [],
   "source": [
    "from stable_baselines3.ppo import MlpPolicy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "RapkYvTXL7Cd"
   },
   "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://gymnasium.farama.org/environments/classic_control/cart_pole/](https://gymnasium.farama.org/environments/classic_control/cart_pole/)\n",
    "\n",
    "![Cartpole](https://cdn-images-1.medium.com/max/1143/1*h4WTQNVIsvMXJTCpXm_TAw.gif)\n",
    "\n",
    "Note: vectorized environments allow to easily multiprocess training. In this example, we are using only one process, hence the DummyVecEnv.\n",
    "\n",
    "We chose the MlpPolicy because input of CartPole 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",
    "\n",
    "Here we are using the [Proximal Policy Optimization](https://stable-baselines.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-baselines.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-baselines.readthedocs.io/en/master/modules/sac.html) or [TD3](https://stable-baselines.readthedocs.io/en/master/modules/td3.html), but is much faster regarding wall-clock time.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "pUWGZp3i9wyf"
   },
   "outputs": [],
   "source": [
    "env = gym.make(\"CartPole-v1\", render_mode=\"rgb_array\")\n",
    "model = PPO(MlpPolicy, env, verbose=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "4efFdrQ7MBvl"
   },
   "source": [
    "We create a helper function to evaluate the agent:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "63M8mSKR-6Zt"
   },
   "outputs": [],
   "source": [
    "def evaluate(model, num_episodes=100, deterministic=True):\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",
    "    vec_env = model.get_env()\n",
    "    all_episode_rewards = []\n",
    "    for i in range(num_episodes):\n",
    "        episode_rewards = []\n",
    "        done = False\n",
    "        obs = vec_env.reset()\n",
    "        while not done:\n",
    "            # _states are only useful when using LSTM policies\n",
    "            action, _states = model.predict(obs, deterministic=deterministic)\n",
    "            # here, action, rewards and dones are arrays\n",
    "            # because we are using vectorized env\n",
    "            # also note that the step only returns a 4-tuple, as the env that is returned\n",
    "            # by model.get_env() is an sb3 vecenv that wraps the >v0.26 API\n",
    "            obs, reward, done, info = vec_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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "6hkyafs--gJz"
   },
   "source": [
    "In fact, Stable-Baselines3 already provides you with that helper:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "s6ZNldIR-fce"
   },
   "outputs": [],
   "source": [
    "from stable_baselines3.common.evaluation import evaluate_policy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "zjEVOIY8NVeK"
   },
   "source": [
    "Let's evaluate the un-trained agent, this should be a random agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "xDHLMA6NFk95",
    "outputId": "a7115c9f-1076-48a2-a78b-61d387609c5f"
   },
   "outputs": [],
   "source": [
    "# Use a separate environement for evaluation\n",
    "eval_env = gym.make(\"CartPole-v1\", render_mode=\"rgb_array\")\n",
    "\n",
    "# Random Agent, before training\n",
    "mean_reward, std_reward = evaluate_policy(model, eval_env, n_eval_episodes=100)\n",
    "\n",
    "print(f\"mean_reward:{mean_reward:.2f} +/- {std_reward:.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "r5UoXTZPNdFE"
   },
   "source": [
    "## Train the agent and evaluate it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "e4cfSXIB-pTF",
    "outputId": "8f6fe592-df0a-4731-f447-b7c059b4f500"
   },
   "outputs": [],
   "source": [
    "# Train the agent for 10000 steps\n",
    "model.learn(total_timesteps=10_000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "ygl_gVmV_QP7",
    "outputId": "e2e2a301-bfea-48a4-c046-2912d9a90124"
   },
   "outputs": [],
   "source": [
    "# Evaluate the trained agent\n",
    "mean_reward, std_reward = evaluate_policy(model, eval_env, n_eval_episodes=100)\n",
    "\n",
    "print(f\"mean_reward:{mean_reward:.2f} +/- {std_reward:.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "A00W6yY3NkHG"
   },
   "source": [
    "Apparently the training went well, the mean reward increased a lot ! "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "xVm9QPNVwKXN"
   },
   "source": [
    "### Prepare video recording"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "MPyfQxD5z26J"
   },
   "outputs": [],
   "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'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "SLzXxO8VMD6N"
   },
   "outputs": [],
   "source": [
    "import base64\n",
    "from pathlib import Path\n",
    "\n",
    "from IPython import display as ipythondisplay\n",
    "\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(\n",
    "            \"\"\"<video alt=\"{}\" autoplay \n",
    "                    loop controls style=\"height: 400px;\">\n",
    "                    <source src=\"data:video/mp4;base64,{}\" type=\"video/mp4\" />\n",
    "                </video>\"\"\".format(\n",
    "                mp4, video_b64.decode(\"ascii\")\n",
    "            )\n",
    "        )\n",
    "    ipythondisplay.display(ipythondisplay.HTML(data=\"<br>\".join(html)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "LTRNUfulOGaF"
   },
   "source": [
    "We will record a video using the [VecVideoRecorder](https://stable-baselines.readthedocs.io/en/master/guide/vec_envs.html#vecvideorecorder) wrapper, you will learn about those wrapper in the next notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Trag9dQpOIhx"
   },
   "outputs": [],
   "source": [
    "from stable_baselines3.common.vec_env import VecVideoRecorder, DummyVecEnv\n",
    "\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(\"CartPole-v1\", render_mode=\"rgb_array\")])\n",
    "    # Start the video at step=0 and record 500 steps\n",
    "    eval_env = VecVideoRecorder(\n",
    "        eval_env,\n",
    "        video_folder=video_folder,\n",
    "        record_video_trigger=lambda step: step == 0,\n",
    "        video_length=video_length,\n",
    "        name_prefix=prefix,\n",
    "    )\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()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "KOObbeu5MMlR"
   },
   "source": [
    "### Visualize trained agent\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "iATu7AiyMQW2",
    "outputId": "7a1ffe06-b8c5-42e5-b075-1b04b98c017e"
   },
   "outputs": [],
   "source": [
    "record_video(\"CartPole-v1\", model, video_length=500, prefix=\"ppo-cartpole\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "-n4i-fW3NojZ"
   },
   "outputs": [],
   "source": [
    "show_videos(\"videos\", prefix=\"ppo\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "9Y8zg4V566qD"
   },
   "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-baselines.readthedocs.io/en/master/guide/quickstart.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "iaOPfOrwWEP4"
   },
   "outputs": [],
   "source": [
    "model = PPO('MlpPolicy', \"CartPole-v1\", verbose=1).learn(1000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "FrI6f5fWnzp-"
   },
   "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 ;)\n"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "1_getting_started.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.11.0"
  },
  "vscode": {
   "interpreter": {
    "hash": "3201c96db5836b171d01fee72ea1be894646622d4b41771abf25c98b548a611d"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}