5_custom_gym_env.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/5_custom_gym_env.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": "AoxOjIlOImwx"
   },
   "source": [
    "# Stable Baselines3 Tutorial - Creating a custom Gym environment\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 your own environment following the OpenAI Gym interface.\n",
    "Once it is done, you can easily use any compatible (depending on the action space) RL algorithm from Stable Baselines on that environment.\n",
    "\n",
    "## Install Dependencies and Stable Baselines3 Using Pip\n",
    "\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": "Sp8rSS4DIhEV"
   },
   "outputs": [],
   "source": [
    "!pip install \"stable-baselines3[extra]>=2.0.0a4\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "rzevZcgmJmhi"
   },
   "source": [
    "## First steps with the gym interface\n",
    "\n",
    "As you have noticed in the previous notebooks, an environment that follows the gym interface is quite simple to use.\n",
    "It provides to this user mainly three methods, which have the following signature (for gym versions > 0.26)\n",
    "- `reset()` called at the beginning of an episode, it returns an observation and a dictionary with additional info (defaults to an empty dict)\n",
    "- `step(action)` called to take an action with the environment, it returns the next observation, the immediate reward, whether new state is a terminal state (episode is finished), whether the max number of timesteps is reached (episode is artificially finished), and additional information\n",
    "- (Optional) `render()` which allow to visualize the agent in action. Note that graphical interface does not work on google colab, so we cannot use it directly (we have to rely on `render_mode='rbg_array'` to retrieve an image of the scene).\n",
    "\n",
    "Under the hood, it also contains two useful properties:\n",
    "- `observation_space` which one of the gym spaces (`Discrete`, `Box`, ...) and describe the type and shape of the observation\n",
    "- `action_space` which is also a gym space object that describes the action space, so the type of action that can be taken\n",
    "\n",
    "The best way to learn about [gym spaces](https://gymnasium.farama.org/api/spaces/) is to look at the [source code](https://github.com/Farama-Foundation/Gymnasium/tree/main/gymnasium/spaces), but you need to know at least the main ones:\n",
    "- `gym.spaces.Box`: A (possibly unbounded) box in $R^n$. Specifically, a Box represents the Cartesian product of n closed intervals. Each interval has the form of one of [a, b], (-oo, b], [a, oo), or (-oo, oo). Example: A 1D-Vector or an image observation can be described with the Box space.\n",
    "```python\n",
    "# Example for using image as input:\n",
    "observation_space = spaces.Box(low=0, high=255, shape=(HEIGHT, WIDTH, N_CHANNELS), dtype=np.uint8)\n",
    "```                                       \n",
    "\n",
    "- `gym.spaces.Discrete`: A discrete space in $\\{ 0, 1, \\dots, n-1 \\}$\n",
    "  Example: if you have two actions (\"left\" and \"right\") you can represent your action space using `Discrete(2)`, the first action will be 0 and the second 1.\n",
    "\n",
    "\n",
    "[Documentation on custom env](https://stable-baselines3.readthedocs.io/en/master/guide/custom_env.html)\n",
    "\n",
    "Also keep in mind that Stabe-baselines internally uses the previous gym API (<0.26), so every VecEnv returns only the observation after resetting and returns a 4-tuple instead of a 5-tuple  (terminated & truncated are already combined to done)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "I98IKKyNJl6K"
   },
   "outputs": [],
   "source": [
    "import gymnasium as gym\n",
    "\n",
    "env = gym.make(\"CartPole-v1\")\n",
    "\n",
    "# Box(4,) means that it is a Vector with 4 components\n",
    "print(\"Observation space:\", env.observation_space)\n",
    "print(\"Shape:\", env.observation_space.shape)\n",
    "# Discrete(2) means that there is two discrete actions\n",
    "print(\"Action space:\", env.action_space)\n",
    "\n",
    "# The reset method is called at the beginning of an episode\n",
    "obs, info = env.reset()\n",
    "# Sample a random action\n",
    "action = env.action_space.sample()\n",
    "print(\"Sampled action:\", action)\n",
    "obs, reward, terminated, truncated, info = env.step(action)\n",
    "# Note the obs is a numpy array\n",
    "# info is an empty dict for now but can contain any debugging info\n",
    "# reward is a scalar\n",
    "print(obs.shape, reward, terminated, truncated, info)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "RqxatIwPOXe_"
   },
   "source": [
    "##  Gym env skeleton\n",
    "\n",
    "In practice this is how a gym environment looks like.\n",
    "Here, we have implemented a simple grid world were the agent must learn to go always left."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "rYzDXA9vJfz1"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import gymnasium as gym\n",
    "from gymnasium import spaces\n",
    "\n",
    "\n",
    "class GoLeftEnv(gym.Env):\n",
    "    \"\"\"\n",
    "    Custom Environment that follows gym interface.\n",
    "    This is a simple env where the agent must learn to go always left.\n",
    "    \"\"\"\n",
    "\n",
    "    # Because of google colab, we cannot implement the GUI ('human' render mode)\n",
    "    metadata = {\"render_modes\": [\"console\"]}\n",
    "\n",
    "    # Define constants for clearer code\n",
    "    LEFT = 0\n",
    "    RIGHT = 1\n",
    "\n",
    "    def __init__(self, grid_size=10, render_mode=\"console\"):\n",
    "        super(GoLeftEnv, self).__init__()\n",
    "        self.render_mode = render_mode\n",
    "\n",
    "        # Size of the 1D-grid\n",
    "        self.grid_size = grid_size\n",
    "        # Initialize the agent at the right of the grid\n",
    "        self.agent_pos = grid_size - 1\n",
    "\n",
    "        # Define action and observation space\n",
    "        # They must be gym.spaces objects\n",
    "        # Example when using discrete actions, we have two: left and right\n",
    "        n_actions = 2\n",
    "        self.action_space = spaces.Discrete(n_actions)\n",
    "        # The observation will be the coordinate of the agent\n",
    "        # this can be described both by Discrete and Box space\n",
    "        self.observation_space = spaces.Box(\n",
    "            low=0, high=self.grid_size, shape=(1,), dtype=np.float32\n",
    "        )\n",
    "\n",
    "    def reset(self, seed=None, options=None):\n",
    "        \"\"\"\n",
    "        Important: the observation must be a numpy array\n",
    "        :return: (np.array)\n",
    "        \"\"\"\n",
    "        super().reset(seed=seed, options=options)\n",
    "        # Initialize the agent at the right of the grid\n",
    "        self.agent_pos = self.grid_size - 1\n",
    "        # here we convert to float32 to make it more general (in case we want to use continuous actions)\n",
    "        return np.array([self.agent_pos]).astype(np.float32), {}  # empty info dict\n",
    "\n",
    "    def step(self, action):\n",
    "        if action == self.LEFT:\n",
    "            self.agent_pos -= 1\n",
    "        elif action == self.RIGHT:\n",
    "            self.agent_pos += 1\n",
    "        else:\n",
    "            raise ValueError(\n",
    "                f\"Received invalid action={action} which is not part of the action space\"\n",
    "            )\n",
    "\n",
    "        # Account for the boundaries of the grid\n",
    "        self.agent_pos = np.clip(self.agent_pos, 0, self.grid_size)\n",
    "\n",
    "        # Are we at the left of the grid?\n",
    "        terminated = bool(self.agent_pos == 0)\n",
    "        truncated = False  # we do not limit the number of steps here\n",
    "\n",
    "        # Null reward everywhere except when reaching the goal (left of the grid)\n",
    "        reward = 1 if self.agent_pos == 0 else 0\n",
    "\n",
    "        # Optionally we can pass additional info, we are not using that for now\n",
    "        info = {}\n",
    "\n",
    "        return (\n",
    "            np.array([self.agent_pos]).astype(np.float32),\n",
    "            reward,\n",
    "            terminated,\n",
    "            truncated,\n",
    "            info,\n",
    "        )\n",
    "\n",
    "    def render(self):\n",
    "        # agent is represented as a cross, rest as a dot\n",
    "        if self.render_mode == \"console\":\n",
    "            print(\".\" * self.agent_pos, end=\"\")\n",
    "            print(\"x\", end=\"\")\n",
    "            print(\".\" * (self.grid_size - self.agent_pos))\n",
    "\n",
    "    def close(self):\n",
    "        pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Zy5mlho1-Ine"
   },
   "source": [
    "### Validate the environment\n",
    "\n",
    "Stable Baselines3 provides a [helper](https://stable-baselines3.readthedocs.io/en/master/common/env_checker.html) to check that your environment follows the Gym interface. It also optionally checks that the environment is compatible with Stable-Baselines (and emits warning if necessary)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "9DOpP_B0-LXm"
   },
   "outputs": [],
   "source": [
    "from stable_baselines3.common.env_checker import check_env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "1CcUVatq-P0l"
   },
   "outputs": [],
   "source": [
    "env = GoLeftEnv()\n",
    "# If the environment don't follow the interface, an error will be thrown\n",
    "check_env(env, warn=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "eJ3khFtkSE0g"
   },
   "source": [
    "### Testing the environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "i62yf2LvSAYY"
   },
   "outputs": [],
   "source": [
    "env = GoLeftEnv(grid_size=10)\n",
    "\n",
    "obs, _ = env.reset()\n",
    "env.render()\n",
    "\n",
    "print(env.observation_space)\n",
    "print(env.action_space)\n",
    "print(env.action_space.sample())\n",
    "\n",
    "GO_LEFT = 0\n",
    "# Hardcoded best agent: always go left!\n",
    "n_steps = 20\n",
    "for step in range(n_steps):\n",
    "    print(f\"Step {step + 1}\")\n",
    "    obs, reward, terminated, truncated, info = env.step(GO_LEFT)\n",
    "    done = terminated or truncated\n",
    "    print(\"obs=\", obs, \"reward=\", reward, \"done=\", done)\n",
    "    env.render()\n",
    "    if done:\n",
    "        print(\"Goal reached!\", \"reward=\", reward)\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Pv1e1qJETfHU"
   },
   "source": [
    "### Try it with Stable-Baselines\n",
    "\n",
    "Once your environment follow the gym interface, it is quite easy to plug in any algorithm from stable-baselines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "PQfLBE28SNDr"
   },
   "outputs": [],
   "source": [
    "from stable_baselines3 import PPO, A2C, DQN\n",
    "from stable_baselines3.common.env_util import make_vec_env\n",
    "\n",
    "# Instantiate the env\n",
    "vec_env = make_vec_env(GoLeftEnv, n_envs=1, env_kwargs=dict(grid_size=10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "zRV4Q7FVUKB6"
   },
   "outputs": [],
   "source": [
    "# Train the agent\n",
    "model = A2C(\"MlpPolicy\", env, verbose=1).learn(5000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "BJbeiF0RUN-p"
   },
   "outputs": [],
   "source": [
    "# Test the trained agent\n",
    "# using the vecenv\n",
    "obs = vec_env.reset()\n",
    "n_steps = 20\n",
    "for step in range(n_steps):\n",
    "    action, _ = model.predict(obs, deterministic=True)\n",
    "    print(f\"Step {step + 1}\")\n",
    "    print(\"Action: \", action)\n",
    "    obs, reward, done, info = vec_env.step(action)\n",
    "    print(\"obs=\", obs, \"reward=\", reward, \"done=\", done)\n",
    "    vec_env.render()\n",
    "    if done:\n",
    "        # Note that the VecEnv resets automatically\n",
    "        # when a done signal is encountered\n",
    "        print(\"Goal reached!\", \"reward=\", reward)\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "jOggIa9sU--b"
   },
   "source": [
    "## It is your turn now, be creative!\n",
    "\n",
    "As an exercise, that's now your turn to build a custom gym environment.\n",
    "There is no constrain about what to do, be creative! (but not too creative, there is not enough time for that)\n",
    "\n",
    "If you don't have any idea, here is is a list of the environment you can implement:\n",
    "- Transform the discrete grid world to a continuous one, you will need to change a bit the logic and the action space\n",
    "- Create a 2D grid world and add walls\n",
    "- Create a tic-tac-toe game\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "lBDp4Pm-Uh4D"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "include_colab_link": true,
   "name": "5.custom_gym_env.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": "3201c96db5836b171d01fee72ea1be894646622d4b41771abf25c98b548a611d"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}