2_gym_wrappers_saving_loading.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/2_gym_wrappers_saving_loading.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": "3ezJ3Y7XRUnj"
   },
   "source": [
    "# Stable Baselines3 Tutorial - Gym wrappers, saving and loading models\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 *Gym Wrappers* which allow to do monitoring, normalization, limit the number of steps, feature augmentation, ...\n",
    "\n",
    "\n",
    "You will also see the *loading* and *saving* functions, and how to read the outputted files for possible exporting.\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": "YFdlFByORUnl"
   },
   "outputs": [],
   "source": [
    "!pip install swig\n",
    "!pip install \"stable-baselines3[extra]>=2.0.0a4\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "grXe85G9RUnp"
   },
   "outputs": [],
   "source": [
    "import gymnasium as gym\n",
    "from stable_baselines3 import A2C, SAC, PPO, TD3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "hMPAn1SRd32f"
   },
   "source": [
    "# Saving and loading\n",
    "\n",
    "Saving and loading stable-baselines models is straightforward: you can directly call `.save()` and `.load()` on the models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "vBNFnN4Gd32g"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Create save dir\n",
    "save_dir = \"/tmp/gym/\"\n",
    "os.makedirs(save_dir, exist_ok=True)\n",
    "\n",
    "model = PPO(\"MlpPolicy\", \"Pendulum-v1\", verbose=0).learn(8_000)\n",
    "# The model will be saved under PPO_tutorial.zip\n",
    "model.save(f\"{save_dir}/PPO_tutorial\")\n",
    "\n",
    "# sample an observation from the environment\n",
    "obs = model.env.observation_space.sample()\n",
    "\n",
    "# Check prediction before saving\n",
    "print(\"pre saved\", model.predict(obs, deterministic=True))\n",
    "\n",
    "del model  # delete trained model to demonstrate loading\n",
    "\n",
    "loaded_model = PPO.load(f\"{save_dir}/PPO_tutorial\")\n",
    "# Check that the prediction is the same after loading (for the same observation)\n",
    "print(\"loaded\", loaded_model.predict(obs, deterministic=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "gXWPrVqId32o"
   },
   "source": [
    "Saving in stable-baselines is quite powerful, as you save the training hyperparameters, with the current weights. This means in practice, you can simply load a custom model, without redefining the parameters, and continue learning.\n",
    "\n",
    "The loading function can also update the model's class variables when loading."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "LCtxrAbXd32q"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from stable_baselines3.common.vec_env import DummyVecEnv\n",
    "\n",
    "# Create save dir\n",
    "save_dir = \"/tmp/gym/\"\n",
    "os.makedirs(save_dir, exist_ok=True)\n",
    "\n",
    "model = A2C(\"MlpPolicy\", \"Pendulum-v1\", verbose=0, gamma=0.9, n_steps=20).learn(8000)\n",
    "# The model will be saved under A2C_tutorial.zip\n",
    "model.save(f\"{save_dir}/A2C_tutorial\")\n",
    "\n",
    "del model  # delete trained model to demonstrate loading\n",
    "\n",
    "# load the model, and when loading set verbose to 1\n",
    "loaded_model = A2C.load(f\"{save_dir}/A2C_tutorial\", verbose=1)\n",
    "\n",
    "# show the save hyperparameters\n",
    "print(f\"loaded: gamma={loaded_model.gamma}, n_steps={loaded_model.n_steps}\")\n",
    "\n",
    "# as the environment is not serializable, we need to set a new instance of the environment\n",
    "loaded_model.set_env(DummyVecEnv([lambda: gym.make(\"Pendulum-v1\")]))\n",
    "# and continue training\n",
    "loaded_model.learn(8_000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "hKwupU-Jgxjm"
   },
   "source": [
    "# Gym and VecEnv wrappers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ds4AAfmISQIA"
   },
   "source": [
    "## Anatomy of a gym wrapper"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "gnTS9e9hTzZZ"
   },
   "source": [
    "A gym wrapper follows the [gym](https://stable-baselines.readthedocs.io/en/master/guide/custom_env.html) interface: it has a `reset()` and `step()` method.\n",
    "\n",
    "Because a wrapper is *around* an environment, we can access it with `self.env`, this allow to easily interact with it without modifying the original env.\n",
    "There are many wrappers that have been predefined, for a complete list refer to [gym documentation](https://gymnasium.farama.org/api/wrappers/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "hYo0C0TQSL3c"
   },
   "outputs": [],
   "source": [
    "class CustomWrapper(gym.Wrapper):\n",
    "    \"\"\"\n",
    "    :param env: (gym.Env) Gym environment that will be wrapped\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, env):\n",
    "        # Call the parent constructor, so we can access self.env later\n",
    "        super().__init__(env)\n",
    "\n",
    "    def reset(self, **kwargs):\n",
    "        \"\"\"\n",
    "        Reset the environment\n",
    "        \"\"\"\n",
    "        obs, info = self.env.reset(**kwargs)\n",
    "\n",
    "        return obs, info\n",
    "\n",
    "    def step(self, action):\n",
    "        \"\"\"\n",
    "        :param action: ([float] or int) Action taken by the agent\n",
    "        :return: (np.ndarray, float, bool, bool, dict) observation, reward, is this a final state (episode finished),\n",
    "        is the max number of steps reached (episode finished artificially), additional informations\n",
    "        \"\"\"\n",
    "        obs, reward, terminated, truncated, info = self.env.step(action)\n",
    "        return obs, reward, terminated, truncated, info"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "4zeGuyICUN26"
   },
   "source": [
    "## First example: limit the episode length\n",
    "\n",
    "One practical use case of a wrapper is when you want to limit the number of steps by episode, for that you will need to overwrite the `done` signal when the limit is reached. It is also a good practice to pass that information in the `info` dictionary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Eb2U4_K6SNUx"
   },
   "outputs": [],
   "source": [
    "class TimeLimitWrapper(gym.Wrapper):\n",
    "    \"\"\"\n",
    "    :param env: (gym.Env) Gym environment that will be wrapped\n",
    "    :param max_steps: (int) Max number of steps per episode\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, env, max_steps=100):\n",
    "        # Call the parent constructor, so we can access self.env later\n",
    "        super(TimeLimitWrapper, self).__init__(env)\n",
    "        self.max_steps = max_steps\n",
    "        # Counter of steps per episode\n",
    "        self.current_step = 0\n",
    "\n",
    "    def reset(self, **kwargs):\n",
    "        \"\"\"\n",
    "        Reset the environment\n",
    "        \"\"\"\n",
    "        # Reset the counter\n",
    "        self.current_step = 0\n",
    "        return self.env.reset(**kwargs)\n",
    "\n",
    "    def step(self, action):\n",
    "        \"\"\"\n",
    "        :param action: ([float] or int) Action taken by the agent\n",
    "        :return: (np.ndarray, float, bool, bool, dict) observation, reward, is the episode over?, additional informations\n",
    "        \"\"\"\n",
    "        self.current_step += 1\n",
    "        obs, reward, terminated, truncated, info = self.env.step(action)\n",
    "        # Overwrite the truncation signal when when the number of steps reaches the maximum\n",
    "        if self.current_step >= self.max_steps:\n",
    "            truncated = True\n",
    "        return obs, reward, terminated, truncated, info"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "oZufaUJwVM9w"
   },
   "source": [
    "#### Test the wrapper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "szZ43D5PVB07"
   },
   "outputs": [],
   "source": [
    "from gymnasium.envs.classic_control.pendulum import PendulumEnv\n",
    "\n",
    "# Here we create the environment directly because gym.make() already wrap the environment in a TimeLimit wrapper otherwise\n",
    "env = PendulumEnv()\n",
    "# Wrap the environment\n",
    "env = TimeLimitWrapper(env, max_steps=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "cencka9iVg9V"
   },
   "outputs": [],
   "source": [
    "obs, _ = env.reset()\n",
    "done = False\n",
    "n_steps = 0\n",
    "while not done:\n",
    "    # Take random actions\n",
    "    random_action = env.action_space.sample()\n",
    "    obs, reward, terminated, truncated, info = env.step(random_action)\n",
    "    done = terminated or truncated\n",
    "    n_steps += 1\n",
    "\n",
    "print(n_steps, info)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "jkMYA63sV9aA"
   },
   "source": [
    "In practice, `gym` already have a wrapper for that named `TimeLimit` (`gym.wrappers.TimeLimit`) that is used by most environments."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "VIIJbSyQW9R-"
   },
   "source": [
    "## Second example: normalize actions\n",
    "\n",
    "It is usually a good idea to normalize observations and actions before giving it to the agent, this prevents this [hard to debug issue](https://github.com/hill-a/stable-baselines/issues/473).\n",
    "\n",
    "In this example, we are going to normalize the action space of *Pendulum-v1* so it lies in [-1, 1] instead of [-2, 2].\n",
    "\n",
    "Note: here we are dealing with continuous actions, hence the `gym.Box` space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "F5E6kZfzW8vy"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "\n",
    "class NormalizeActionWrapper(gym.Wrapper):\n",
    "    \"\"\"\n",
    "    :param env: (gym.Env) Gym environment that will be wrapped\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, env):\n",
    "        # Retrieve the action space\n",
    "        action_space = env.action_space\n",
    "        assert isinstance(\n",
    "            action_space, gym.spaces.Box\n",
    "        ), \"This wrapper only works with continuous action space (spaces.Box)\"\n",
    "        # Retrieve the max/min values\n",
    "        self.low, self.high = action_space.low, action_space.high\n",
    "\n",
    "        # We modify the action space, so all actions will lie in [-1, 1]\n",
    "        env.action_space = gym.spaces.Box(\n",
    "            low=-1, high=1, shape=action_space.shape, dtype=np.float32\n",
    "        )\n",
    "\n",
    "        # Call the parent constructor, so we can access self.env later\n",
    "        super(NormalizeActionWrapper, self).__init__(env)\n",
    "\n",
    "    def rescale_action(self, scaled_action):\n",
    "        \"\"\"\n",
    "        Rescale the action from [-1, 1] to [low, high]\n",
    "        (no need for symmetric action space)\n",
    "        :param scaled_action: (np.ndarray)\n",
    "        :return: (np.ndarray)\n",
    "        \"\"\"\n",
    "        return self.low + (0.5 * (scaled_action + 1.0) * (self.high - self.low))\n",
    "\n",
    "    def reset(self, **kwargs):\n",
    "        \"\"\"\n",
    "        Reset the environment\n",
    "        \"\"\"\n",
    "        return self.env.reset(**kwargs)\n",
    "\n",
    "    def step(self, action):\n",
    "        \"\"\"\n",
    "        :param action: ([float] or int) Action taken by the agent\n",
    "        :return: (np.ndarray, float,bool, bool, dict) observation, reward, final state? truncated?, additional informations\n",
    "        \"\"\"\n",
    "        # Rescale action from [-1, 1] to original [low, high] interval\n",
    "        rescaled_action = self.rescale_action(action)\n",
    "        obs, reward, terminated, truncated, info = self.env.step(rescaled_action)\n",
    "        return obs, reward, terminated, truncated, info"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "TmJ0eahNaR6K"
   },
   "source": [
    "#### Test before rescaling actions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "UEnjBwisaQIx"
   },
   "outputs": [],
   "source": [
    "original_env = gym.make(\"Pendulum-v1\")\n",
    "\n",
    "print(original_env.action_space.low)\n",
    "for _ in range(10):\n",
    "    print(original_env.action_space.sample())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "jvcll2L3afVd"
   },
   "source": [
    "#### Test the NormalizeAction wrapper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "WsCM9AUGaeBN"
   },
   "outputs": [],
   "source": [
    "env = NormalizeActionWrapper(gym.make(\"Pendulum-v1\"))\n",
    "\n",
    "print(env.action_space.low)\n",
    "\n",
    "for _ in range(10):\n",
    "    print(env.action_space.sample())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "V5h5kk2mbGNs"
   },
   "source": [
    "#### Test with a RL algorithm\n",
    "\n",
    "We are going to use the Monitor wrapper of stable baselines, which allow to monitor training stats (mean episode reward, mean episode length)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "R9FNCN8ybOVU"
   },
   "outputs": [],
   "source": [
    "from stable_baselines3.common.monitor import Monitor\n",
    "from stable_baselines3.common.vec_env import DummyVecEnv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 53
    },
    "colab_type": "code",
    "id": "wutM3c1GbfGP",
    "outputId": "eda3b489-ab0f-45cd-8acd-c36835f063df"
   },
   "outputs": [],
   "source": [
    "env = Monitor(gym.make(\"Pendulum-v1\"))\n",
    "env = DummyVecEnv([lambda: env])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "8cxnE5bdaQ_3"
   },
   "outputs": [],
   "source": [
    "model = A2C(\"MlpPolicy\", env, verbose=1).learn(int(1000))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "EJFSM-Drb3Wc"
   },
   "source": [
    "With the action wrapper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "GszFZthob2wM"
   },
   "outputs": [],
   "source": [
    "normalized_env = Monitor(gym.make(\"Pendulum-v1\"))\n",
    "# Note that we can use multiple wrappers\n",
    "normalized_env = NormalizeActionWrapper(normalized_env)\n",
    "normalized_env = DummyVecEnv([lambda: normalized_env])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "wrKJEO4NcIMd"
   },
   "outputs": [],
   "source": [
    "model_2 = A2C(\"MlpPolicy\", normalized_env, verbose=1).learn(int(1000))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "5BxqXd_6dpJx"
   },
   "source": [
    "## Additional wrappers: VecEnvWrappers\n",
    "\n",
    "In the same vein as gym wrappers, stable baselines provide wrappers for `VecEnv`. Among the different wrappers that exist (and you can create your own), you should know: \n",
    "\n",
    "- VecNormalize: it computes a running mean and standard deviation to normalize observation and returns\n",
    "- VecFrameStack: it stacks several consecutive observations (useful to integrate time in the observation, e.g. successive frame of an atari game)\n",
    "\n",
    "More info in the [documentation](https://stable-baselines3.readthedocs.io/en/master/guide/vec_envs.html#wrappers)\n",
    "\n",
    "Note: when using `VecNormalize` wrapper, you must save the running mean and std along with the model, otherwise you will not get proper results when loading the agent again. If you use the [rl zoo](https://github.com/DLR-RM/rl-baselines3-zoo), this is done automatically"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "zuIcbfv3g9dd"
   },
   "outputs": [],
   "source": [
    "from stable_baselines3.common.vec_env import VecNormalize, VecFrameStack\n",
    "\n",
    "env = DummyVecEnv([lambda: gym.make(\"Pendulum-v1\")])\n",
    "normalized_vec_env = VecNormalize(env)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "-PAbu21pg90A"
   },
   "outputs": [],
   "source": [
    "obs = normalized_vec_env.reset()\n",
    "for _ in range(10):\n",
    "    action = [normalized_vec_env.action_space.sample()]\n",
    "    obs, reward, _, _ = normalized_vec_env.step(action)\n",
    "    print(obs, reward)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "UEpTys28Wz05"
   },
   "source": [
    "## Exercise: code you own monitor wrapper\n",
    "\n",
    "Now that you know how does a wrapper work and what you can do with it, it's time to experiment.\n",
    "\n",
    "The goal here is to create a wrapper that will monitor the training progress, storing both the episode reward (sum of reward for one episode) and episode length (number of steps in for the last episode).\n",
    "\n",
    "You will return those values using the `info` dict after each end of episode."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "8FWeDRd5W7hO"
   },
   "outputs": [],
   "source": [
    "class MyMonitorWrapper(gym.Wrapper):\n",
    "    \"\"\"\n",
    "    :param env: (gym.Env) Gym environment that will be wrapped\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, env):\n",
    "        # Call the parent constructor, so we can access self.env later\n",
    "        super().__init__(env)\n",
    "        # === YOUR CODE HERE ===#\n",
    "        # Initialize the variables that will be used\n",
    "        # to store the episode length and episode reward\n",
    "\n",
    "        # ====================== #\n",
    "\n",
    "    def reset(self, **kwargs):\n",
    "        \"\"\"\n",
    "        Reset the environment\n",
    "        \"\"\"\n",
    "        obs = self.env.reset(**kwargs)\n",
    "        # === YOUR CODE HERE ===#\n",
    "        # Reset the variables\n",
    "\n",
    "        # ====================== #\n",
    "        return obs\n",
    "\n",
    "    def step(self, action):\n",
    "        \"\"\"\n",
    "        :param action: ([float] or int) Action taken by the agent\n",
    "        :return: (np.ndarray, float, bool, bool, dict)\n",
    "            observation, reward, is the episode over?, is the episode truncated?, additional information\n",
    "        \"\"\"\n",
    "        obs, reward, terminated, truncated, info = self.env.step(action)\n",
    "        # === YOUR CODE HERE ===#\n",
    "        # Update the current episode reward and episode length\n",
    "\n",
    "        # ====================== #\n",
    "\n",
    "        if terminated or truncated:\n",
    "            # === YOUR CODE HERE ===#\n",
    "            # Store the episode length and episode reward in the info dict\n",
    "            pass\n",
    "\n",
    "            # ====================== #\n",
    "        return obs, reward, terminated, truncated, info"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "d4fY4QwWXNFK"
   },
   "source": [
    "#### Test your wrapper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "bJbUG-A_liYt"
   },
   "outputs": [],
   "source": [
    "# To use LunarLander, you need to install box2d box2d-kengz (pip) and swig (apt-get)\n",
    "!pip install box2d-py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "oWZp1olSXMUg"
   },
   "outputs": [],
   "source": [
    "env = gym.make(\"LunarLander-v2\")\n",
    "# === YOUR CODE HERE ===#\n",
    "# Wrap the environment\n",
    "\n",
    "# Reset the environment\n",
    "\n",
    "# Take random actions in the environment and check\n",
    "# that it returns the correct values after the end of each episode\n",
    "\n",
    "# ====================== #"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "dJ2IqSM2eOt8"
   },
   "source": [
    " # Conclusion\n",
    " \n",
    " In this notebook, we have seen:\n",
    " - how to easily save and load a model\n",
    " - what is wrapper and what we can do with it\n",
    " - how to create your own wrapper"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "qhWB_bHpSkas"
   },
   "source": [
    "## Wrapper Bonus: changing the observation space: a wrapper for episode of fixed length"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "bBlS9YxYSpJn"
   },
   "outputs": [],
   "source": [
    "from gym.wrappers import TimeLimit\n",
    "\n",
    "\n",
    "class TimeFeatureWrapper(gym.Wrapper):\n",
    "    \"\"\"\n",
    "    Add remaining time to observation space for fixed length episodes.\n",
    "    See https://arxiv.org/abs/1712.00378 and https://github.com/aravindr93/mjrl/issues/13.\n",
    "\n",
    "    :param env: (gym.Env)\n",
    "    :param max_steps: (int) Max number of steps of an episode\n",
    "        if it is not wrapped in a TimeLimit object.\n",
    "    :param test_mode: (bool) In test mode, the time feature is constant,\n",
    "        equal to zero. This allow to check that the agent did not overfit this feature,\n",
    "        learning a deterministic pre-defined sequence of actions.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, env, max_steps=1000, test_mode=False):\n",
    "        assert isinstance(env.observation_space, gym.spaces.Box)\n",
    "        # Add a time feature to the observation\n",
    "        low, high = env.observation_space.low, env.observation_space.high\n",
    "        low, high = np.concatenate((low, [0])), np.concatenate((high, [1.0]))\n",
    "        env.observation_space = gym.spaces.Box(low=low, high=high, dtype=np.float32)\n",
    "\n",
    "        super().__init__(env)\n",
    "\n",
    "        if isinstance(env, TimeLimit):\n",
    "            self._max_steps = env._max_episode_steps\n",
    "        else:\n",
    "            self._max_steps = max_steps\n",
    "        self._current_step = 0\n",
    "        self._test_mode = test_mode\n",
    "\n",
    "    def reset(self, **kwargs):\n",
    "        self._current_step = 0\n",
    "        obs, info = self.env.reset(**kwargs)\n",
    "        return self._get_obs(obs), info\n",
    "\n",
    "    def step(self, action):\n",
    "        self._current_step += 1\n",
    "        obs, reward, terminated, truncated, info = self.env.step(action)\n",
    "        return self._get_obs(obs), reward, terminated, truncated, info\n",
    "\n",
    "    def _get_obs(self, obs):\n",
    "        \"\"\"\n",
    "        Concatenate the time feature to the current observation.\n",
    "\n",
    "        :param obs: (np.ndarray)\n",
    "        :return: (np.ndarray)\n",
    "        \"\"\"\n",
    "        # Remaining time is more general\n",
    "        time_feature = 1 - (self._current_step / self._max_steps)\n",
    "        if self._test_mode:\n",
    "            time_feature = 1.0\n",
    "        # Optionally: concatenate [time_feature, time_feature ** 2]\n",
    "        return np.concatenate((obs, [time_feature]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "z-vWgkZzd4F1"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Ojn4nvNNRUoT"
   },
   "source": [
    "## Going further - Saving format \n",
    "\n",
    "The format for saving and loading models is a zip-archived JSON dump and NumPy zip archive of the arrays:\n",
    "```\n",
    "saved_model.zip/\n",
    "├── data              JSON file of class-parameters (dictionary)\n",
    "├── parameter_list    JSON file of model parameters and their ordering (list)\n",
    "├── parameters        Bytes from numpy.savez (a zip file of the numpy arrays). ...\n",
    "    ├── ...           Being a zip-archive itself, this object can also be opened ...\n",
    "        ├── ...       as a zip-archive and browsed.\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "QWAcc8RFRUoU"
   },
   "source": [
    "## Save and find "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "4tcQxzSCRUoV"
   },
   "outputs": [],
   "source": [
    "# Create save dir\n",
    "save_dir = \"/tmp/gym/\"\n",
    "os.makedirs(save_dir, exist_ok=True)\n",
    "\n",
    "model = PPO(\"MlpPolicy\", \"Pendulum-v1\", verbose=0).learn(8000)\n",
    "model.save(save_dir + \"/PPO_tutorial\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "rGaMNz4HRUoX"
   },
   "outputs": [],
   "source": [
    "!ls /tmp/gym/PPO_tutorial*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "gYY3nQyyRUoa"
   },
   "outputs": [],
   "source": [
    "import zipfile\n",
    "\n",
    "archive = zipfile.ZipFile(\"/tmp/gym/PPO_tutorial.zip\", \"r\")\n",
    "for f in archive.filelist:\n",
    "    print(f.filename)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "cPKkkTvjRUo2"
   },
   "source": [
    "## Exporting saved models\n",
    "\n",
    "And finally some further reading for those who want to export to tensorflowJS or Java.\n",
    "\n",
    "https://stable-baselines3.readthedocs.io/en/master/guide/export.html"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "include_colab_link": true,
   "name": "2_gym_wrappers_saving_loading.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
}