pretraining.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/pretraining.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": "ShofuYy-8dLs"
   },
   "source": [
    "## Stable Baselines3 - Pretraining with Behavior Cloning\n",
    "\n",
    "\n",
    "Github repo: https://github.com/araffin/rl-tutorial-jnrr19\n",
    "\n",
    "Stable-Baselines3: https://github.com/DLR-RM/stable-baselines3\n",
    "\n",
    "Documentation: https://stable-baselines.readthedocs.io/en/master/\n",
    "\n",
    "RL Baselines3 zoo: https://github.com/DLR-RM/rl-baselines3-zoo\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",
    "## Introduction\n",
    "\n",
    "In this notebook, you will learn how to record expert data, then pre-train an agent using this data\n",
    "and finally continue training with Stable-Baselines3.\n",
    "\n",
    "\n",
    "## Install Dependencies and Stable Baselines3 Using Pip\n",
    "\n",
    "List of full dependencies can be found in the [README](https://github.com/DLR-RM/stable-baselines3).\n",
    "\n",
    "Notebook originally created by [skervim](https://github.com/skervim)\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": {},
    "colab_type": "code",
    "id": "QwJ1LHlN9E33"
   },
   "outputs": [],
   "source": [
    "# For Box2D env\n",
    "!apt-get update && apt-get install swig\n",
    "!pip install gymnasium[box2d]\n",
    "!pip install \"stable-baselines3[extra]>=2.0.0a4\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Tgx4AMZo8anP"
   },
   "outputs": [],
   "source": [
    "import gymnasium as gym\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "\n",
    "print(f\"{gym.__version__}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "U_1ZNBum8ane"
   },
   "outputs": [],
   "source": [
    "import torch as th\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.optim.lr_scheduler import StepLR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "iaz2Szrl8anx"
   },
   "outputs": [],
   "source": [
    "from stable_baselines3 import PPO, A2C, SAC, TD3\n",
    "from stable_baselines3.common.evaluation import evaluate_policy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "dS9lr70R9eU9"
   },
   "outputs": [],
   "source": [
    "# Example for continuous actions\n",
    "# env_id = \"LunarLanderContinuous-v2\"\n",
    "\n",
    "# Example for discrete actions\n",
    "env_id = \"CartPole-v1\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "eh88d4oR8an6"
   },
   "outputs": [],
   "source": [
    "env = gym.make(env_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "rKCgHCc_8aoB"
   },
   "source": [
    "## Train Expert Model\n",
    "\n",
    "We create an expert RL agent and let it learn to solve a task by interacting with the evironment.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "EkmIST0r8aoC"
   },
   "outputs": [],
   "source": [
    "ppo_expert = PPO(\"MlpPolicy\", env_id, verbose=1)\n",
    "ppo_expert.learn(total_timesteps=3e4)\n",
    "ppo_expert.save(\"ppo_expert\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "hlyTfGAQ_Az1"
   },
   "source": [
    "check the performance of the trained agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "-_rVEjC0_AQa",
    "outputId": "1873bf5d-c5f5-40b8-bf2b-59331a53fb73"
   },
   "outputs": [],
   "source": [
    "mean_reward, std_reward = evaluate_policy(ppo_expert, env, n_eval_episodes=10)\n",
    "\n",
    "print(f\"Mean reward = {mean_reward} +/- {std_reward}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "S-Kv6v_V8aoJ"
   },
   "source": [
    "## Create Student\n",
    "\n",
    "We also create a student RL agent, which will later be trained with the expert dataset\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "fLdLPUeC8aoL"
   },
   "outputs": [],
   "source": [
    "a2c_student = A2C(\"MlpPolicy\", env_id, verbose=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "sdW8_41-OcXn"
   },
   "outputs": [],
   "source": [
    "# only valid for continuous actions\n",
    "# sac_student = SAC('MlpPolicy', env_id, verbose=1, policy_kwargs=dict(net_arch=[64, 64]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "m3GuNxcU8aoT"
   },
   "source": [
    "\n",
    "We now let our expert interact with the environment (except we already have expert data) and store resultant expert observations and actions to build an expert dataset.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "emodyZDW8aoU"
   },
   "outputs": [],
   "source": [
    "num_interactions = int(4e4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "F3I_2s808aoZ"
   },
   "outputs": [],
   "source": [
    "if isinstance(env.action_space, gym.spaces.Box):\n",
    "    expert_observations = np.empty((num_interactions,) + env.observation_space.shape)\n",
    "    expert_actions = np.empty((num_interactions,) + (env.action_space.shape[0],))\n",
    "\n",
    "else:\n",
    "    expert_observations = np.empty((num_interactions,) + env.observation_space.shape)\n",
    "    expert_actions = np.empty((num_interactions,) + env.action_space.shape)\n",
    "\n",
    "obs, _ = env.reset()\n",
    "\n",
    "for i in tqdm(range(num_interactions)):\n",
    "    action, _ = ppo_expert.predict(obs, deterministic=True)\n",
    "    expert_observations[i] = obs\n",
    "    expert_actions[i] = action\n",
    "    obs, reward, terminated, truncated, info = env.step(action)\n",
    "    done = terminated or truncated\n",
    "    if done:\n",
    "        obs, _ = env.reset()\n",
    "\n",
    "np.savez_compressed(\n",
    "    \"expert_data\",\n",
    "    expert_actions=expert_actions,\n",
    "    expert_observations=expert_observations,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "toKEQE9i8aof"
   },
   "source": [
    "\n",
    "\n",
    "- To seamlessly use PyTorch in the training process, we subclass an `ExpertDataset` from PyTorch's base `Dataset`.\n",
    "- Note that we initialize the dataset with the previously generated expert observations and actions.\n",
    "- We further implement Python's `__getitem__` and `__len__` magic functions to allow PyTorch's dataset-handling to access arbitrary rows in the dataset and inform it about the length of the dataset.\n",
    "- For more information about PyTorch's datasets, you can read: https://pytorch.org/docs/stable/data.html.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "qT72bR1i8aog"
   },
   "outputs": [],
   "source": [
    "from torch.utils.data.dataset import Dataset, random_split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "xUetr5vl8aom"
   },
   "outputs": [],
   "source": [
    "class ExpertDataSet(Dataset):\n",
    "    def __init__(self, expert_observations, expert_actions):\n",
    "        self.observations = expert_observations\n",
    "        self.actions = expert_actions\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        return (self.observations[index], self.actions[index])\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.observations)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "K9bNAhXp8aor"
   },
   "source": [
    "\n",
    "\n",
    "We now instantiate the `ExpertDataSet` and split it into training and test datasets.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "lIdT-zMV8aot"
   },
   "outputs": [],
   "source": [
    "expert_dataset = ExpertDataSet(expert_observations, expert_actions)\n",
    "\n",
    "train_size = int(0.8 * len(expert_dataset))\n",
    "\n",
    "test_size = len(expert_dataset) - train_size\n",
    "\n",
    "train_expert_dataset, test_expert_dataset = random_split(\n",
    "    expert_dataset, [train_size, test_size]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 53
    },
    "colab_type": "code",
    "id": "_LgmtFFq8aox",
    "outputId": "543280f3-84ff-440b-c2f1-a6dd3fe0f7e9"
   },
   "outputs": [],
   "source": [
    "print(\"test_expert_dataset: \", len(test_expert_dataset))\n",
    "print(\"train_expert_dataset: \", len(train_expert_dataset))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "0v8PhG2r8ao4"
   },
   "source": [
    "\n",
    "\n",
    "NOTE: The supervised learning section of this code is adapted from: https://github.com/pytorch/examples/blob/master/mnist/main.py\n",
    "1. We extract the policy network of our RL student agent.\n",
    "2. We load the (labeled) expert dataset containing expert observations as inputs and expert actions as targets.\n",
    "3. We perform supervised learning, that is, we adjust the policy network's parameters such that given expert observations as inputs to the network, its outputs match the targets (expert actions).\n",
    "By training the policy network in this way the corresponding RL student agent is taught to behave like the expert agent that was used to created the expert dataset (Behavior Cloning).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "VwUhCTGU8ao5"
   },
   "outputs": [],
   "source": [
    "def pretrain_agent(\n",
    "    student,\n",
    "    batch_size=64,\n",
    "    epochs=1000,\n",
    "    scheduler_gamma=0.7,\n",
    "    learning_rate=1.0,\n",
    "    log_interval=100,\n",
    "    no_cuda=True,\n",
    "    seed=1,\n",
    "    test_batch_size=64,\n",
    "):\n",
    "    use_cuda = not no_cuda and th.cuda.is_available()\n",
    "    th.manual_seed(seed)\n",
    "    device = th.device(\"cuda\" if use_cuda else \"cpu\")\n",
    "    kwargs = {\"num_workers\": 1, \"pin_memory\": True} if use_cuda else {}\n",
    "\n",
    "    if isinstance(env.action_space, gym.spaces.Box):\n",
    "        criterion = nn.MSELoss()\n",
    "    else:\n",
    "        criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "    # Extract initial policy\n",
    "    model = student.policy.to(device)\n",
    "\n",
    "    def train(model, device, train_loader, optimizer):\n",
    "        model.train()\n",
    "\n",
    "        for batch_idx, (data, target) in enumerate(train_loader):\n",
    "            data, target = data.to(device), target.to(device)\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            if isinstance(env.action_space, gym.spaces.Box):\n",
    "                # A2C/PPO policy outputs actions, values, log_prob\n",
    "                # SAC/TD3 policy outputs actions only\n",
    "                if isinstance(student, (A2C, PPO)):\n",
    "                    action, _, _ = model(data)\n",
    "                else:\n",
    "                    # SAC/TD3:\n",
    "                    action = model(data)\n",
    "                action_prediction = action.double()\n",
    "            else:\n",
    "                # Retrieve the logits for A2C/PPO when using discrete actions\n",
    "                dist = model.get_distribution(data)\n",
    "                action_prediction = dist.distribution.logits\n",
    "                target = target.long()\n",
    "\n",
    "            loss = criterion(action_prediction, target)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            if batch_idx % log_interval == 0:\n",
    "                print(\n",
    "                    \"Train Epoch: {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}\".format(\n",
    "                        epoch,\n",
    "                        batch_idx * len(data),\n",
    "                        len(train_loader.dataset),\n",
    "                        100.0 * batch_idx / len(train_loader),\n",
    "                        loss.item(),\n",
    "                    )\n",
    "                )\n",
    "\n",
    "    def test(model, device, test_loader):\n",
    "        model.eval()\n",
    "        test_loss = 0\n",
    "        with th.no_grad():\n",
    "            for data, target in test_loader:\n",
    "                data, target = data.to(device), target.to(device)\n",
    "\n",
    "                if isinstance(env.action_space, gym.spaces.Box):\n",
    "                    # A2C/PPO policy outputs actions, values, log_prob\n",
    "                    # SAC/TD3 policy outputs actions only\n",
    "                    if isinstance(student, (A2C, PPO)):\n",
    "                        action, _, _ = model(data)\n",
    "                    else:\n",
    "                        # SAC/TD3:\n",
    "                        action = model(data)\n",
    "                    action_prediction = action.double()\n",
    "                else:\n",
    "                    # Retrieve the logits for A2C/PPO when using discrete actions\n",
    "                    dist = model.get_distribution(data)\n",
    "                    action_prediction = dist.distribution.logits\n",
    "                    target = target.long()\n",
    "\n",
    "                test_loss = criterion(action_prediction, target)\n",
    "        test_loss /= len(test_loader.dataset)\n",
    "        print(f\"Test set: Average loss: {test_loss:.4f}\")\n",
    "\n",
    "    # Here, we use PyTorch `DataLoader` to our load previously created `ExpertDataset` for training\n",
    "    # and testing\n",
    "    train_loader = th.utils.data.DataLoader(\n",
    "        dataset=train_expert_dataset, batch_size=batch_size, shuffle=True, **kwargs\n",
    "    )\n",
    "    test_loader = th.utils.data.DataLoader(\n",
    "        dataset=test_expert_dataset,\n",
    "        batch_size=test_batch_size,\n",
    "        shuffle=True,\n",
    "        **kwargs,\n",
    "    )\n",
    "\n",
    "    # Define an Optimizer and a learning rate schedule.\n",
    "    optimizer = optim.Adadelta(model.parameters(), lr=learning_rate)\n",
    "    scheduler = StepLR(optimizer, step_size=1, gamma=scheduler_gamma)\n",
    "\n",
    "    # Now we are finally ready to train the policy model.\n",
    "    for epoch in range(1, epochs + 1):\n",
    "        train(model, device, train_loader, optimizer)\n",
    "        test(model, device, test_loader)\n",
    "        scheduler.step()\n",
    "\n",
    "    # Implant the trained policy network back into the RL student agent\n",
    "    a2c_student.policy = model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "nkEP6i0hEu_R"
   },
   "source": [
    "Evaluate the agent before pretraining, it should be random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "_7kvYIneEui8",
    "outputId": "fb450cd6-ec13-4cd6-bb14-abc12c42c775"
   },
   "outputs": [],
   "source": [
    "mean_reward, std_reward = evaluate_policy(a2c_student, env, n_eval_episodes=10)\n",
    "\n",
    "print(f\"Mean reward = {mean_reward} +/- {std_reward}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "SgduZAbF8ao9"
   },
   "source": [
    "\n",
    "\n",
    "Having defined the training procedure we can now run the pretraining!\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "eI1EFFnW8ao-"
   },
   "outputs": [],
   "source": [
    "pretrain_agent(\n",
    "    a2c_student,\n",
    "    epochs=3,\n",
    "    scheduler_gamma=0.7,\n",
    "    learning_rate=1.0,\n",
    "    log_interval=100,\n",
    "    no_cuda=True,\n",
    "    seed=1,\n",
    "    batch_size=64,\n",
    "    test_batch_size=1000,\n",
    ")\n",
    "a2c_student.save(\"a2c_student\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "RK3Q5Jm58apE"
   },
   "source": [
    "\n",
    "\n",
    "Finally, let us test how well our RL agent student learned to mimic the behavior of the expert\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "GKZ8O--m8apF",
    "outputId": "398abb60-b1ac-43b1-cb40-f9045b8b2339"
   },
   "outputs": [],
   "source": [
    "mean_reward, std_reward = evaluate_policy(a2c_student, env, n_eval_episodes=10)\n",
    "\n",
    "print(f\"Mean reward = {mean_reward} +/- {std_reward}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "name": "pretraining.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
}