advanced_saving_loading.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/advanced_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": {
    "id": "hyyN-2qyK_T2"
   },
   "source": [
    "# Stable Baselines3 - Advanced Saving and Loading\n",
    "\n",
    "Github Repo: [https://github.com/DLR-RM/stable-baselines3](https://github.com/DLR-RM/stable-baselines3)\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",
    "Documentation is available online: [https://stable-baselines3.readthedocs.io/](https://stable-baselines3.readthedocs.io/)\n",
    "\n",
    "## Introduction\n",
    "\n",
    "In this notebook, you will learn how to use some advanced features of stable baselines3 (SB3): how to easily create a test environment for periodic evaluation, use a policy independently from a model (and how to save it, load it) and save/load a replay buffer.\n",
    "\n",
    "## Install Dependencies and Stable Baselines Using Pip\n",
    "\n",
    "\n",
    "```\n",
    "pip install stable-baselines3[extra]\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gWskDE2c9WoN"
   },
   "outputs": [],
   "source": [
    "!pip install \"stable-baselines3[extra]>=2.0.0a4\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FtY8FhliLsGm"
   },
   "source": [
    "## Import policy, RL agent, ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "BIedd7Pz9sOs"
   },
   "outputs": [],
   "source": [
    "from stable_baselines3 import SAC, TD3\n",
    "from stable_baselines3.common.evaluation import evaluate_policy\n",
    "from stable_baselines3.common.callbacks import EvalCallback\n",
    "from stable_baselines3.common.env_util import make_vec_env"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RapkYvTXL7Cd"
   },
   "source": [
    "## Create the Gym env and instantiate the agent\n",
    "\n",
    "For this example, we will use Pendulum environment.\n",
    "\n",
    "\"The inverted pendulum swingup problem is a classic problem in the control literature. In this version of the problem, the pendulum starts in a random position, and the goal is to swing it up so it stays upright.\"\n",
    "\n",
    "Pendulum-v1 environment: [documentation](https://gymnasium.farama.org/environments/classic_control/pendulum/), [source code](https://github.com/Farama-Foundation/Gymnasium/blob/main/gymnasium/envs/classic_control/pendulum.py)\n",
    "\n",
    "![Pendulum](https://gymnasium.farama.org/_images/pendulum.gif)\n",
    "\n",
    "\n",
    "We chose the MlpPolicy because input of Pendulum 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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "eyt_yXx78FvB"
   },
   "source": [
    "### Create the agent and evaluation callback\n",
    "\n",
    "We will use an [EvalCallback](https://stable-baselines3.readthedocs.io/en/master/guide/callbacks.html#evalcallback) to periodically evaluate the agent on a separate env."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 88
    },
    "id": "pUWGZp3i9wyf",
    "outputId": "75553a25-9d37-4d73-9ccb-a05c85ea6927"
   },
   "outputs": [],
   "source": [
    "model = SAC(\"MlpPolicy\", \"Pendulum-v1\", verbose=1, learning_rate=1e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_env = make_vec_env(\"Pendulum-v1\", n_envs=5)\n",
    "\n",
    "n_training_envs = 1 # adjust eval freq depending on the number of training envs\n",
    "# Evaluate the model every 1000 steps on 5 test episodes and save the evaluation to the logs folder\n",
    "eval_callback = EvalCallback(eval_env, eval_freq=1000 // n_training_envs, n_eval_episodes=5, log_path=\"./logs/\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4efFdrQ7MBvl"
   },
   "source": [
    "Train the agent and evaluate it periodically on the test env."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "HfOdzGNo8QQg"
   },
   "outputs": [],
   "source": [
    "model.learn(6000, callback=eval_callback, progress_bar=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sw80kaed3Zba"
   },
   "source": [
    "## Save/Load the replay buffer\n",
    "\n",
    "By default, the replay buffer is not saved when calling `model.save()`, in order to save space on the disk (a replay buffer can be up to several GB when using images).\n",
    "However, SB3 provides a `save_replay_buffer()` and `load_replay_buffer()` method to save it separately."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "id": "T65Bo7-k3dWL",
    "outputId": "896cd86f-5b87-48a3-da0b-5f5543a6f889"
   },
   "outputs": [],
   "source": [
    "# save the model\n",
    "model.save(\"sac_pendulum\")\n",
    "\n",
    "# the saved model does not contain the replay buffer\n",
    "loaded_model = SAC.load(\"sac_pendulum\")\n",
    "print(f\"The loaded_model has {loaded_model.replay_buffer.size()} transitions in its buffer\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "id": "xEpCyKQN4kvw",
    "outputId": "36b920b9-9762-4077-fb2b-6fbbae390488"
   },
   "outputs": [],
   "source": [
    "# now save the replay buffer too\n",
    "model.save_replay_buffer(\"sac_replay_buffer\")\n",
    "\n",
    "# load it into the loaded_model\n",
    "loaded_model.load_replay_buffer(\"sac_replay_buffer\")\n",
    "\n",
    "# now the loaded replay is not empty anymore\n",
    "print(f\"The loaded_model has {loaded_model.replay_buffer.size()} transitions in its buffer\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wELlF5un9edF"
   },
   "source": [
    "## Save the policy only\n",
    "\n",
    "In SB3, you save the policy independently from the model if needed.\n",
    "\n",
    "Note: if you don't save the complete model, you cannot continue training afterward"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PeaVBGuJwK97"
   },
   "outputs": [],
   "source": [
    "policy = model.policy\n",
    "policy.save(\"sac_policy_pendulum.pkl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "1PFEOuOI-S7t"
   },
   "outputs": [],
   "source": [
    "env = model.get_env()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "KTEhTvLD9rQy"
   },
   "outputs": [],
   "source": [
    "# Evaluate the policy\n",
    "mean_reward, std_reward = evaluate_policy(policy, eval_env, n_eval_episodes=10, deterministic=True)\n",
    "\n",
    "print(f\"mean_reward={mean_reward:.2f} +/- {std_reward}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "714wHSKT9iEg"
   },
   "source": [
    "## Load the policy only"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "mogPtE0D90hw"
   },
   "outputs": [],
   "source": [
    "from stable_baselines3.sac.policies import MlpPolicy\n",
    "saved_policy = MlpPolicy.load(\"sac_policy_pendulum.pkl\")\n",
    "# also possible:\n",
    "# saved_policy = SAC.policy_aliases[\"MlpPolicy\"].load(\"sac_policy_pendulum.pkl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ygl_gVmV_QP7"
   },
   "outputs": [],
   "source": [
    "# Evaluate the loaded policy\n",
    "mean_reward, std_reward = evaluate_policy(saved_policy, eval_env, n_eval_episodes=10, deterministic=True)\n",
    "\n",
    "print(f\"mean_reward={mean_reward:.2f} +/- {std_reward}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "aQDZI5VEGnUq"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "advanced_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.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}