stable_baselines_her.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/stable_baselines_her.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 Baselines - Hindsight Experience Replay on Highway Env\n",
    "\n",
    "Github Repo: [https://github.com/DLR-RM/stable-baselines3](https://github.com/DLR-RM/stable-baselines3)\n",
    "\n",
    "Highway env: [https://github.com/eleurent/highway-env](https://github.com/eleurent/highway-env) \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",
    "## Install Dependencies and Stable Baselines Using Pip\n",
    "\n",
    "\n",
    "```\n",
    "pip install stable-baselines3[extra]\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": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "qPg7pyvK_Emi",
    "outputId": "95466567-e00c-4aab-a6d9-74f4839a83fe"
   },
   "outputs": [],
   "source": [
    "# Install stable-baselines latest version\n",
    "!pip install \"stable-baselines3[extra]>=2.0.0a4\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "4wA6lU254uNl",
    "outputId": "ce145fd4-800d-49bb-877e-a402fbe7677f"
   },
   "outputs": [],
   "source": [
    "# Install highway-env\n",
    "!pip install highway-env"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FtY8FhliLsGm"
   },
   "source": [
    "## Import policy, RL agent, ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "BIedd7Pz9sOs"
   },
   "outputs": [],
   "source": [
    "import gymnasium as gym\n",
    "import highway_env\n",
    "import numpy as np\n",
    "\n",
    "from stable_baselines3 import HerReplayBuffer, SAC, DDPG\n",
    "from stable_baselines3.common.noise import NormalActionNoise"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "RapkYvTXL7Cd"
   },
   "source": [
    "## Create the Gym env and instantiate the agent\n",
    "\n",
    "For this example, we will be using the parking environment from the [highway-env](https://github.com/Farama-Foundation/HighwayEnv) repo by @eleurent.\n",
    "\n",
    "The parking env is a goal-conditioned continuous control task, in which the vehicle must park in a given space with the appropriate heading.\n",
    "\n",
    "\n",
    "![parking-env](https://raw.githubusercontent.com/eleurent/highway-env/gh-media/docs/media/parking-env.gif)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QT5Io2LwAzlp"
   },
   "source": [
    "### Train Soft Actor-Critic (SAC) agent\n",
    "\n",
    "Here, we use HER \"future\" goal sampling strategy, where we create 4 artificial transitions per real transition\n",
    "\n",
    "Note: the hyperparameters (network architecture, discount factor, ...) were tuned for this task"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "TzkK23C2BCKr"
   },
   "outputs": [],
   "source": [
    "env = gym.make(\"parking-v0\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "QDdPRp0f5Bcz",
    "outputId": "b6bb9d63-a88e-4628-cccd-6c182a7b7278"
   },
   "outputs": [],
   "source": [
    "# SAC hyperparams:\n",
    "model = SAC(\n",
    "    \"MultiInputPolicy\",\n",
    "    env,\n",
    "    replay_buffer_class=HerReplayBuffer,\n",
    "    replay_buffer_kwargs=dict(\n",
    "        n_sampled_goal=4,\n",
    "        goal_selection_strategy=\"future\",\n",
    "    ),\n",
    "    verbose=1,\n",
    "    buffer_size=int(1e6),\n",
    "    learning_rate=1e-3,\n",
    "    gamma=0.95,\n",
    "    batch_size=256,\n",
    "    policy_kwargs=dict(net_arch=[256, 256, 256]),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "P82f568g5g0q",
    "outputId": "72afdd05-0e9f-4d83-f4b2-7585e81b9de5"
   },
   "outputs": [],
   "source": [
    "# Train for 1e5 steps\n",
    "model.learn(int(1e5))\n",
    "# Save the trained agent\n",
    "model.save('her_sac_highway')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "nGa3dNn36PTX"
   },
   "outputs": [],
   "source": [
    "# Load saved model\n",
    "model = SAC.load('her_sac_highway', env=env)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Jg_7vKNGA6Hf"
   },
   "source": [
    "#### Evaluate the agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "D3sxhf1Q6NlL",
    "outputId": "b822f5be-6bb1-474d-dcfd-55d124247038"
   },
   "outputs": [],
   "source": [
    "# we use the gym >v.26 API here. Note that you could also wrap the env in a DummyVecEnv\n",
    "# which allows you to use a simplified API\n",
    "obs, _ = env.reset()\n",
    "\n",
    "# Evaluate the agent\n",
    "episode_reward = 0\n",
    "for _ in range(1000):\n",
    "    action, _ = model.predict(obs, deterministic=True)\n",
    "    obs, reward, terminated, truncated, info = env.step(action)\n",
    "    done = truncated or terminated\n",
    "    episode_reward += reward\n",
    "    if done or info.get(\"is_success\", False):\n",
    "        print(\"Reward:\", episode_reward, \"Success?\", info.get(\"is_success\", False))\n",
    "        episode_reward = 0.0\n",
    "        obs, _ = env.reset()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fMjiWlSEi-3n"
   },
   "source": [
    "### Train DDPG agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "FCLW5cGIZa52",
    "outputId": "31e39167-3d5b-41b1-817b-74302cc355b4"
   },
   "outputs": [],
   "source": [
    "# Create the action noise object that will be used for exploration\n",
    "n_actions = env.action_space.shape[0]\n",
    "noise_std = 0.2\n",
    "action_noise = NormalActionNoise(\n",
    "    mean=np.zeros(n_actions), sigma=noise_std * np.ones(n_actions)\n",
    ")\n",
    "\n",
    "model = DDPG(\n",
    "    \"MultiInputPolicy\",\n",
    "    env,\n",
    "    replay_buffer_class=HerReplayBuffer,\n",
    "    replay_buffer_kwargs=dict(\n",
    "        n_sampled_goal=4,\n",
    "        goal_selection_strategy=\"future\",\n",
    "    ),\n",
    "    verbose=1,\n",
    "    buffer_size=int(1e6),\n",
    "    learning_rate=1e-3,\n",
    "    action_noise=action_noise,\n",
    "    gamma=0.95,\n",
    "    batch_size=256,\n",
    "    policy_kwargs=dict(net_arch=[256, 256, 256]),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "PXAQeENljBIe",
    "outputId": "398faa87-1c38-425b-b972-c65437a1dccf"
   },
   "outputs": [],
   "source": [
    "# Train for 2e5 steps\n",
    "model.learn(int(2e5))\n",
    "# Save the trained agent\n",
    "model.save('her_ddpg_highway')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5o9N-V5FO2v2"
   },
   "outputs": [],
   "source": [
    "# Load saved model\n",
    "model = DDPG.load('her_ddpg_highway', env=env)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "B8uJjxQLBxdS"
   },
   "source": [
    "#### Evaluate the agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "XjohypRhjEjS",
    "outputId": "9fda55c1-2e67-470a-bf15-a65144f19350"
   },
   "outputs": [],
   "source": [
    "# we use the gym >v.26 API here. Note that you could also wrap the env in a DummyVecEnv\n",
    "# which allows you to use the old gym API a simplified API\n",
    "obs, _ = env.reset()\n",
    "\n",
    "# Evaluate the agent\n",
    "episode_reward = 0\n",
    "for _ in range(1000):\n",
    "    action, _ = model.predict(obs, deterministic=True)\n",
    "    obs, reward, terminated, truncated, info = env.step(action)\n",
    "    done = truncated or terminated\n",
    "    episode_reward += reward\n",
    "    if done or info.get(\"is_success\", False):\n",
    "        print(\"Reward:\", episode_reward, \"Success?\", info.get(\"is_success\", False))\n",
    "        episode_reward = 0.0\n",
    "        obs, _ = env.reset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "stable_baselines_HER.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"
  },
  "vscode": {
   "interpreter": {
    "hash": "3201c96db5836b171d01fee72ea1be894646622d4b41771abf25c98b548a611d"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}