saving_loading_a2c.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "saving_loading_a2c.ipynb",
      "version": "0.3.2",
      "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/robertmoni/modelbasedrl/blob/master/saving_loading_a2c.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 Baselines, a Fork of OpenAI Baselines - Training, Saving and Loading\n",
        "\n",
        "Github Repo: [https://github.com/hill-a/stable-baselines](https://github.com/hill-a/stable-baselines)\n",
        "\n",
        "Medium article: [https://medium.com/@araffin/stable-baselines-a-fork-of-openai-baselines-df87c4b2fc82](https://medium.com/@araffin/stable-baselines-a-fork-of-openai-baselines-df87c4b2fc82)\n",
        "\n",
        "## Install Dependencies and Stable Baselines Using Pip\n",
        "\n",
        "List of full dependencies can be found in the [README](https://github.com/hill-a/stable-baselines).\n",
        "\n",
        "```\n",
        "\n",
        "sudo apt-get update && sudo apt-get install cmake libopenmpi-dev zlib1g-dev\n",
        "```\n",
        "\n",
        "\n",
        "```\n",
        "\n",
        "pip install stable-baselines\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gWskDE2c9WoN",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!apt install swig cmake libopenmpi-dev zlib1g-dev\n",
        "!pip install stable-baselines==2.5.1 box2d box2d-kengz"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FtY8FhliLsGm",
        "colab_type": "text"
      },
      "source": [
        "## Import policy, RL agent, ..."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BIedd7Pz9sOs",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import gym\n",
        "import numpy as np\n",
        "\n",
        "from stable_baselines.common.policies import MlpPolicy\n",
        "from stable_baselines.common.vec_env import DummyVecEnv\n",
        "from stable_baselines import A2C"
      ],
      "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 Lunar Lander environment.\n",
        "\n",
        "\"Landing outside landing pad is possible. Fuel is infinite, so an agent can learn to fly and then land on its first attempt. Four discrete actions available: do nothing, fire left orientation engine, fire main engine, fire right orientation engine. \"\n",
        "\n",
        "Lunar Lander environment: [https://gym.openai.com/envs/LunarLander-v2/](https://gym.openai.com/envs/LunarLander-v2/)\n",
        "\n",
        "![Lunar Lander](https://cdn-images-1.medium.com/max/960/1*f4VZPKOI0PYNWiwt0la0Rg.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"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pUWGZp3i9wyf",
        "colab_type": "code",
        "outputId": "e305fe9c-b498-4385-f0e1-9b739eac95b7",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        }
      },
      "source": [
        "env = gym.make('LunarLander-v2')\n",
        "# vectorized environments allow to easily multiprocess training\n",
        "# we demonstrate its usefulness in the next examples\n",
        "env = DummyVecEnv([lambda: env])  # The algorithms require a vectorized environment to run\n",
        "\n",
        "model = A2C(MlpPolicy, env, ent_coef=0.1, verbose=0)\n"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "\u001b[33mWARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.\u001b[0m\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "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_steps=1000):\n",
        "  \"\"\"\n",
        "  Evaluate a RL agent\n",
        "  :param model: (BaseRLModel object) the RL Agent\n",
        "  :param num_steps: (int) number of timesteps to evaluate it\n",
        "  :return: (float) Mean reward for the last 100 episodes\n",
        "  \"\"\"\n",
        "  episode_rewards = [0.0]\n",
        "  obs = env.reset()\n",
        "  for i in range(num_steps):\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, rewards, dones, info = env.step(action)\n",
        "      \n",
        "      # Stats\n",
        "      episode_rewards[-1] += rewards[0]\n",
        "      if dones[0]:\n",
        "          obs = env.reset()\n",
        "          episode_rewards.append(0.0)\n",
        "  # Compute mean reward for the last 100 episodes\n",
        "  mean_100ep_reward = round(np.mean(episode_rewards[-100:]), 1)\n",
        "  print(\"Mean reward:\", mean_100ep_reward, \"Num episodes:\", len(episode_rewards))\n",
        "  \n",
        "  return mean_100ep_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": "00a0b8d7-3729-4b6e-9f08-2b0415eb9eab",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        }
      },
      "source": [
        "# Random Agent, before training\n",
        "mean_reward_before_train = evaluate(model, num_steps=10000)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Mean reward: -210.3 Num episodes: 107\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r5UoXTZPNdFE",
        "colab_type": "text"
      },
      "source": [
        "## Train the agent and save it\n",
        "\n",
        "Warning: this may take a while"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "e4cfSXIB-pTF",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Train the agent\n",
        "model.learn(total_timesteps=10000)\n",
        "# Save the agent\n",
        "model.save(\"a2c_lunar\")\n",
        "del model  # delete trained model to demonstrate loading"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "T31dZJYNrJwF",
        "colab_type": "text"
      },
      "source": [
        "## Load the trained agent"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "K1ExgtyZrIA6",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "model = A2C.load(\"a2c_lunar\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ygl_gVmV_QP7",
        "colab_type": "code",
        "outputId": "d32b1690-4f57-4a60-b100-2ea2b18c635b",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        }
      },
      "source": [
        "# Evaluate the trained agent\n",
        "mean_reward = evaluate(model, num_steps=10000)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Mean reward: -310.2 Num episodes: 68\n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
}