ppo_BipedalWalker_v3_basic.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyPlRcS88aQo6+2zPxAzQPJ6",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/BrutFab/ppo_BipedalWalker_v3/blob/main/ppo_BipedalWalker_v3_basic.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# **1. Setup**\n"
      ],
      "metadata": {
        "id": "6MbMQP0b7zGg"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### **Install Packages**"
      ],
      "metadata": {
        "id": "gg4Vl1FH_yVm"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oyYkz0Vc7xNE"
      },
      "outputs": [],
      "source": [
        "# Install necessary packages\n",
        "!apt install swig cmake ffmpeg xvfb python3-opengl\n",
        "!pip install stable-baselines3==2.0.0a5 gymnasium[box2d] huggingface_sb3 pyvirtualdisplay imageio[ffmpeg]"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "The Next Cell will force the notebook runtime to restart. This is to ensure all the new libraries installed will be used."
      ],
      "metadata": {
        "id": "-v81BfAd7yWc"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "os.kill(os.getpid(), 9)"
      ],
      "metadata": {
        "id": "FaiyY-Fn_EZg"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### **Start Virtual Display**"
      ],
      "metadata": {
        "id": "DAlous3h_OrZ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from pyvirtualdisplay import Display\n",
        "virtual_display = Display(visible=0, size=(1400, 900))\n",
        "virtual_display.start()"
      ],
      "metadata": {
        "id": "xblUYFLM_ZGr",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "697a23dd-c35c-4dac-e676-7cf0cb44c2c4"
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<pyvirtualdisplay.display.Display at 0x7c3f41724040>"
            ]
          },
          "metadata": {},
          "execution_count": 1
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### **Setup Environment**"
      ],
      "metadata": {
        "id": "cG91Uw14_9-W"
      }
    },
    {
      "cell_type": "markdown",
      "source": [],
      "metadata": {
        "id": "ddhvmlKeLh-H"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import gymnasium as gym\n",
        "env = gym.make(\"BipedalWalker-v3\", hardcore=True)\n",
        "env.reset()"
      ],
      "metadata": {
        "id": "FyLRt45mACyY",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "9d6ba416-d6df-4ee1-de8a-4abbd196bd79"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(array([ 2.7473392e-03,  1.5248505e-06, -1.9884571e-04, -1.6000010e-02,\n",
              "         9.2200637e-02,  4.6161897e-04,  8.6010480e-01,  1.0362319e-03,\n",
              "         1.0000000e+00,  3.2592569e-02,  4.6158818e-04,  8.5368299e-01,\n",
              "        -3.1407541e-04,  1.0000000e+00,  4.4081384e-01,  4.4581994e-01,\n",
              "         4.6142259e-01,  4.8954999e-01,  5.3410256e-01,  6.0246080e-01,\n",
              "         7.0914859e-01,  8.8593149e-01,  1.0000000e+00,  1.0000000e+00],\n",
              "       dtype=float32),\n",
              " {})"
            ]
          },
          "metadata": {},
          "execution_count": 2
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### **Observation Space**\n",
        "Observation Space Shape (24,) vector of size 24, where each value contains different information about the walker:\n",
        "\n",
        "- **Hull Angle Speed**: The speed at which the main body of the walker is rotating.\n",
        "- **Angular Velocity**: The rate of change of the angular position of the walker.\n",
        "- **Horizontal Speed**: The speed at which the walker is moving horizontally.\n",
        "- **Vertical Speed**: The speed at which the walker is moving vertically.\n",
        "- **Position of Joints**: The positions (angles) of the walker's joints. Given that the walker has 4 joints, this take up 4 values.\n",
        "- **Joints Angular Speed**: The rate of change of the angular position for each joint. Again, this would be 4 values for the 4 joints.\n",
        "- **Legs Contact with Ground**: Indicating whether each leg is in contact with the ground. Given two legs, this contains 2 values.\n",
        "- **10 Lidar Rangefinder Measurements**: These are distance measurements to detect obstacles or terrain features around the walker. There are 10 of these values.\n"
      ],
      "metadata": {
        "id": "xr8LRvHOLxIx"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "print(\"_____OBSERVATION SPACE_____ \\n\")\n",
        "print(\"Observation Space Shape\", env.observation_space.shape)\n",
        "print(\"Sample observation\", env.observation_space.sample()) # Get a random observation"
      ],
      "metadata": {
        "id": "l23gcNuzLr4h",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "7c3391bd-9acb-421c-9f90-56df81132ee6"
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "_____OBSERVATION SPACE_____ \n",
            "\n",
            "Observation Space Shape (24,)\n",
            "Sample observation [ 3.0999613  -4.234623   -3.6569533   3.0880961  -1.957453   -2.572675\n",
            " -2.2593882   4.4289813   3.45656     1.8115904   4.4997163   1.4293909\n",
            "  2.1744916   1.0393379   0.04452281  0.49756885  0.86104745 -0.9264838\n",
            "  0.8241219  -0.8104876   0.41755676 -0.9043817   0.44482002  0.19353716]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### **Action Space**\n",
        "\n",
        " Actions are motor speed values in the [-1, 1] range for each of the 4 joints at both hips and knees."
      ],
      "metadata": {
        "id": "aS3kPiz8I7OH"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "print(\"\\n _____ACTION SPACE_____ \\n\")\n",
        "print(\"Action Space Shape\", env.action_space.shape)\n",
        "print(\"Action Space Sample\", env.action_space.sample()) # Take a random action"
      ],
      "metadata": {
        "id": "r8-of2O-KXe_",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "2b623dda-0217-43b4-e3fa-e7cccc0224ff"
      },
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            " _____ACTION SPACE_____ \n",
            "\n",
            "Action Space Shape (4,)\n",
            "Action Space Sample [ 0.83445764 -0.5473673   0.01842198 -0.15174645]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### **Vectorized Environment**\n",
        "Create a vectorized environment (a method for stacking multiple independent environments into a single environment) of 16 environments to have more diverse experiences."
      ],
      "metadata": {
        "id": "WFheIWeIMR8F"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from stable_baselines3.common.env_util import make_vec_env\n",
        "env = make_vec_env('BipedalWalker-v3', n_envs=16)"
      ],
      "metadata": {
        "id": "G9LsgYvuMqTB"
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# **2. Building the Model**"
      ],
      "metadata": {
        "id": "Ho3QSAXWNOxF"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from stable_baselines3 import PPO\n",
        "model = PPO(\n",
        "    policy = 'MlpPolicy',\n",
        "    env = env,\n",
        "    n_steps = 2048,\n",
        "    batch_size = 128,\n",
        "    n_epochs = 6,\n",
        "    gamma = 0.999,\n",
        "    gae_lambda = 0.98,\n",
        "    ent_coef = 0.01,\n",
        "    verbose=1)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "n5amWvjFNt93",
        "outputId": "6b426c48-561f-438b-ebd8-e6a93ee7aa55"
      },
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Using cpu device\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 3.**Video Generation**"
      ],
      "metadata": {
        "id": "LyYbio5uOiHb"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from wasabi import Printer\n",
        "import numpy as np\n",
        "from stable_baselines3.common.base_class import BaseAlgorithm\n",
        "from pathlib import Path\n",
        "import tempfile\n",
        "from stable_baselines3.common.monitor import Monitor\n",
        "from stable_baselines3.common.vec_env import (\n",
        "    DummyVecEnv,\n",
        "    VecEnv,\n",
        "    VecVideoRecorder,\n",
        ")"
      ],
      "metadata": {
        "id": "bZpJ_3yHOhce"
      },
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "msg = Printer()"
      ],
      "metadata": {
        "id": "r0XIO_k5PBPc"
      },
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def generate_replay(\n",
        "    model: BaseAlgorithm,\n",
        "    eval_env: VecEnv,\n",
        "    video_length: int,\n",
        "    is_deterministic: bool,\n",
        "    local_path: Path,\n",
        "):\n",
        "    \"\"\"\n",
        "    Generate a replay video of the agent\n",
        "    :param model: trained model\n",
        "    :param eval_env: environment used to evaluate the agent\n",
        "    :param video_length: length of the video (in timesteps)\n",
        "    :param is_deterministic: use deterministic or stochastic actions\n",
        "    :param local_path: path of the local repository\n",
        "    \"\"\"\n",
        "    # This is another temporary directory for video outputs\n",
        "    # SB3 created a -step-0-to-... meta files as well as other\n",
        "    # artifacts which we don't want in the repo.\n",
        "    with tempfile.TemporaryDirectory() as tmpdirname:\n",
        "        # Step 1: Create the VecVideoRecorder\n",
        "        env = VecVideoRecorder(\n",
        "            eval_env,\n",
        "            tmpdirname,\n",
        "            record_video_trigger=lambda x: x == 0,\n",
        "            video_length=video_length,\n",
        "            name_prefix=\"\",\n",
        "        )\n",
        "\n",
        "        obs = env.reset()\n",
        "        lstm_states = None\n",
        "        episode_starts = np.ones((env.num_envs,), dtype=bool)\n",
        "\n",
        "        try:\n",
        "            for _ in range(video_length):\n",
        "                action, lstm_states = model.predict(\n",
        "                    obs,\n",
        "                    state=lstm_states,\n",
        "                    episode_start=episode_starts,\n",
        "                    deterministic=is_deterministic,\n",
        "                )\n",
        "                obs, _, episode_starts, _ = env.step(action)\n",
        "\n",
        "            # Save the video\n",
        "            env.close()\n",
        "\n",
        "            # Convert the video with x264 codec\n",
        "            inp = env.video_recorder.path\n",
        "            out = local_path\n",
        "            os.system(f\"ffmpeg -y -i {inp} -vcodec h264 {out}\".format(inp, out))\n",
        "            print(f\"Video saved to: {out}\")\n",
        "        except KeyboardInterrupt:\n",
        "            pass\n",
        "        except Exception as e:\n",
        "            msg.fail(str(e))\n",
        "            # Add a message for video\n",
        "            msg.fail(\n",
        "                \"We are unable to generate a replay of your agent\"\n",
        "            )"
      ],
      "metadata": {
        "id": "TbXxg6viPC-u"
      },
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# **4. Training, Saving and Record the Videos**"
      ],
      "metadata": {
        "id": "7FSBntjZWBZC"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "model.learn(total_timesteps=1000)\n",
        "\n",
        "# Save the model\n",
        "model_name = \"ppo-BipedalWalker-v3\"\n",
        "model.save(model_name)"
      ],
      "metadata": {
        "id": "V95RlWJ8Staf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "video_dir = \".\"\n",
        "video_name = \"replay.mp4\"\n",
        "env_id = \"BipedalWalker-v3\"\n",
        "\n",
        "generate_replay(\n",
        "    model = model,\n",
        "    eval_env=DummyVecEnv([lambda: Monitor(gym.make(env_id, render_mode=\"rgb_array\"))]),\n",
        "    video_length=1000,\n",
        "    is_deterministic=True,\n",
        "    local_path=os.path.join(video_dir, video_name)\n",
        ")"
      ],
      "metadata": {
        "id": "usp5BKC4dlh8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# **5. Visualize Final Video**"
      ],
      "metadata": {
        "id": "zw4c5jJ0XBWU"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from IPython.display import HTML\n",
        "from base64 import b64encode\n",
        "mp4 = open('replay.mp4','rb').read()\n",
        "data_url = \"data:video/mp4;base64,\" + b64encode(mp4).decode()\n",
        "HTML(\"\"\"\n",
        "<video width=600 controls>\n",
        "      <source src=\"%s\" type=\"video/mp4\">\n",
        "</video>\n",
        "\"\"\" % data_url)"
      ],
      "metadata": {
        "id": "NVFpJPohU2_Z"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# **6. Evaluate the Model**"
      ],
      "metadata": {
        "id": "JvCZEk2DXnPa"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from stable_baselines3.common.evaluation import evaluate_policy"
      ],
      "metadata": {
        "id": "JO0spFR9VEX2"
      },
      "execution_count": 16,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "eval_env = Monitor(gym.make(\"BipedalWalker-v3\"))\n",
        "mean_reward, std_reward = evaluate_policy(model, eval_env, n_eval_episodes=10, deterministic=True)\n",
        "print(f\"mean_reward={mean_reward:.2f} +/- {std_reward}\")"
      ],
      "metadata": {
        "id": "AfCcrkOXXX3Y",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "b7b84a0f-333c-4458-b563-6b0ae9f50985"
      },
      "execution_count": 17,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "mean_reward=-95.23 +/- 3.562129299111482\n"
          ]
        }
      ]
    }
  ]
}