pybullet.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/pybullet.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 - PyBullet: Normalizing Features and Reward\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",
        "Pybullet source code: https://github.com/bulletphysics/bullet3/tree/master/examples/pybullet/\n",
        "\n",
        "Gymnasium compatible envs: https://github.com/araffin/pybullet_envs_gymnasium\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": "j35IfrnHGxF9"
      },
      "outputs": [],
      "source": [
        "# for autoformatting\n",
        "# %load_ext jupyter_black"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gWskDE2c9WoN"
      },
      "outputs": [],
      "source": [
        "!pip install pybullet_envs_gymnasium\n",
        "!pip install \"stable-baselines3[extra]>=2.0.0a4\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FtY8FhliLsGm"
      },
      "source": [
        "## Import policy, RL agent, Wrappers"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BIedd7Pz9sOs"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "# Register pybullet envs\n",
        "import pybullet_envs_gymnasium\n",
        "\n",
        "from stable_baselines3 import PPO\n",
        "from stable_baselines3.common.env_util import make_vec_env\n",
        "from stable_baselines3.common.vec_env import VecNormalize, DummyVecEnv"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7c8VHsiXC7dL"
      },
      "source": [
        "## Create and wrap the environment with `VecNormalize`\n",
        "\n",
        "Normalizing input features may be essential to successful training of an RL agent (by default, images are scaled but not other types of input), for instance when training on [PyBullet](https://github.com/bulletphysics/bullet3/) environments. For that, a wrapper exists and will compute a running average and standard deviation of input features (it can do the same for rewards).\n",
        "\n",
        "More information about `VecNormalize`:\n",
        "- [Documentation](https://stable-baselines3.readthedocs.io/en/master/guide/vec_envs.html#stable_baselines3.common.vec_env.VecNormalize)\n",
        "- [Discussion](https://github.com/hill-a/stable-baselines/issues/698)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kmxIq5UeC3Nj"
      },
      "outputs": [],
      "source": [
        "vec_env = make_vec_env(\"HalfCheetahBulletEnv-v0\", n_envs=1)\n",
        "\n",
        "vec_env = VecNormalize(vec_env, norm_obs=True, norm_reward=True, clip_obs=10.0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pxUMGsl5mabF"
      },
      "source": [
        "### Train the agent"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PQmsSZUHKNRG"
      },
      "outputs": [],
      "source": [
        "model = PPO(\"MlpPolicy\", vec_env, verbose=1)\n",
        "model.learn(total_timesteps=2000)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iZYBIVoLmcR4"
      },
      "source": [
        "### Save the agent and the normalization"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SpMDXP0vmezv"
      },
      "outputs": [],
      "source": [
        "# Don't forget to save the VecNormalize statistics when saving the agent\n",
        "log_dir = \"/tmp/\"\n",
        "model.save(log_dir + \"ppo_halfcheetah\")\n",
        "stats_path = os.path.join(log_dir, \"vec_normalize.pkl\")\n",
        "vec_env.save(stats_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eezphIrRmr-Y"
      },
      "source": [
        "### Test model: load the saved agent and normalization"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CQT1k7lWmmTL"
      },
      "outputs": [],
      "source": [
        "# Load the agent\n",
        "model = PPO.load(log_dir + \"ppo_halfcheetah\")\n",
        "\n",
        "# Load the saved statistics\n",
        "vec_env = make_vec_env(\"HalfCheetahBulletEnv-v0\", n_envs=1,)\n",
        "vec_env = VecNormalize.load(stats_path, vec_env)\n",
        "#  do not update them at test time\n",
        "vec_env.training = False\n",
        "# reward normalization is not needed at test time\n",
        "vec_env.norm_reward = False"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "voCxMxfYm2yc"
      },
      "outputs": [],
      "source": [
        "from stable_baselines3.common.evaluation import evaluate_policy"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6zC2IMJUm3Kw"
      },
      "outputs": [],
      "source": [
        "mean_reward, std_reward = evaluate_policy(model, vec_env)\n",
        "\n",
        "print(f\"Mean reward = {mean_reward:.2f} +/- {std_reward:.2f}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_tHTCBIynFhd"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "colab": {
      "name": "pybullet.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": 0
}