REINFORCE.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "authorship_tag": "ABX9TyM+jtEfUv40FN0BnEyoV0jB"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "!sudo apt-get install -y build-essential swig libopenmpi-dev"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "collapsed": true,
        "id": "1d_MG7nNsN1s",
        "outputId": "6b00eb35-cfce-4568-85c2-4b33ff91e271"
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Reading package lists... Done\n",
            "Building dependency tree... Done\n",
            "Reading state information... Done\n",
            "build-essential is already the newest version (12.9ubuntu3).\n",
            "libopenmpi-dev is already the newest version (4.1.2-2ubuntu1).\n",
            "The following additional packages will be installed:\n",
            "  swig4.0\n",
            "Suggested packages:\n",
            "  swig-doc swig-examples swig4.0-examples swig4.0-doc\n",
            "The following NEW packages will be installed:\n",
            "  swig swig4.0\n",
            "0 upgraded, 2 newly installed, 0 to remove and 29 not upgraded.\n",
            "Need to get 1,116 kB of archives.\n",
            "After this operation, 5,542 kB of additional disk space will be used.\n",
            "Get:1 http://archive.ubuntu.com/ubuntu jammy/universe amd64 swig4.0 amd64 4.0.2-1ubuntu1 [1,110 kB]\n",
            "Get:2 http://archive.ubuntu.com/ubuntu jammy/universe amd64 swig all 4.0.2-1ubuntu1 [5,632 B]\n",
            "Fetched 1,116 kB in 1s (808 kB/s)\n",
            "debconf: unable to initialize frontend: Dialog\n",
            "debconf: (No usable dialog-like program is installed, so the dialog based frontend cannot be used. at /usr/share/perl5/Debconf/FrontEnd/Dialog.pm line 78, <> line 2.)\n",
            "debconf: falling back to frontend: Readline\n",
            "debconf: unable to initialize frontend: Readline\n",
            "debconf: (This frontend requires a controlling tty.)\n",
            "debconf: falling back to frontend: Teletype\n",
            "dpkg-preconfigure: unable to re-open stdin: \n",
            "Selecting previously unselected package swig4.0.\n",
            "(Reading database ... 124947 files and directories currently installed.)\n",
            "Preparing to unpack .../swig4.0_4.0.2-1ubuntu1_amd64.deb ...\n",
            "Unpacking swig4.0 (4.0.2-1ubuntu1) ...\n",
            "Selecting previously unselected package swig.\n",
            "Preparing to unpack .../swig_4.0.2-1ubuntu1_all.deb ...\n",
            "Unpacking swig (4.0.2-1ubuntu1) ...\n",
            "Setting up swig4.0 (4.0.2-1ubuntu1) ...\n",
            "Setting up swig (4.0.2-1ubuntu1) ...\n",
            "Processing triggers for man-db (2.10.2-1) ...\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install gymnasium[box2d]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "EkLI8KQnsAO4",
        "outputId": "0136dd89-419c-444e-aa68-1a569fa6b596"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: gymnasium[box2d] in /usr/local/lib/python3.11/dist-packages (1.0.0)\n",
            "Requirement already satisfied: numpy>=1.21.0 in /usr/local/lib/python3.11/dist-packages (from gymnasium[box2d]) (1.26.4)\n",
            "Requirement already satisfied: cloudpickle>=1.2.0 in /usr/local/lib/python3.11/dist-packages (from gymnasium[box2d]) (3.1.1)\n",
            "Requirement already satisfied: typing-extensions>=4.3.0 in /usr/local/lib/python3.11/dist-packages (from gymnasium[box2d]) (4.12.2)\n",
            "Requirement already satisfied: farama-notifications>=0.0.1 in /usr/local/lib/python3.11/dist-packages (from gymnasium[box2d]) (0.0.4)\n",
            "Collecting box2d-py==2.3.5 (from gymnasium[box2d])\n",
            "  Downloading box2d-py-2.3.5.tar.gz (374 kB)\n",
            "\u001b[?25l     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.0/374.4 kB\u001b[0m \u001b[31m?\u001b[0m eta \u001b[36m-:--:--\u001b[0m\r\u001b[2K     \u001b[91m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[90m╺\u001b[0m \u001b[32m368.6/374.4 kB\u001b[0m \u001b[31m12.9 MB/s\u001b[0m eta \u001b[36m0:00:01\u001b[0m\r\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m374.4/374.4 kB\u001b[0m \u001b[31m9.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "Requirement already satisfied: pygame>=2.1.3 in /usr/local/lib/python3.11/dist-packages (from gymnasium[box2d]) (2.6.1)\n",
            "Collecting swig==4.* (from gymnasium[box2d])\n",
            "  Downloading swig-4.3.0-py2.py3-none-manylinux_2_5_x86_64.manylinux1_x86_64.whl.metadata (3.5 kB)\n",
            "Downloading swig-4.3.0-py2.py3-none-manylinux_2_5_x86_64.manylinux1_x86_64.whl (1.9 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.9/1.9 MB\u001b[0m \u001b[31m58.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hBuilding wheels for collected packages: box2d-py\n",
            "  Building wheel for box2d-py (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for box2d-py: filename=box2d_py-2.3.5-cp311-cp311-linux_x86_64.whl size=2351230 sha256=2b5d358e195f9eaa55ab6dd2073bbf4a0df885b63e610a9b20497d83d9c36e88\n",
            "  Stored in directory: /root/.cache/pip/wheels/ab/f1/0c/d56f4a2bdd12bae0a0693ec33f2f0daadb5eb9753c78fa5308\n",
            "Successfully built box2d-py\n",
            "Installing collected packages: swig, box2d-py\n",
            "Successfully installed box2d-py-2.3.5 swig-4.3.0\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "WfY9g8yHr9e-",
        "outputId": "74815dd4-2ecf-4282-acc3-54e4ae39fc04"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Episode 1: Score = -119.29343324065972, Avg Score = -119.29\n",
            "Episode 11: Score = -141.1762166414236, Avg Score = -156.45\n",
            "Episode 21: Score = -248.61103874304916, Avg Score = -186.31\n",
            "Episode 31: Score = -246.10220665389983, Avg Score = -169.02\n",
            "Episode 41: Score = -403.32881514134453, Avg Score = -188.92\n",
            "Episode 51: Score = -149.19830651571058, Avg Score = -184.56\n",
            "Episode 61: Score = -31.932345993517544, Avg Score = -190.04\n",
            "Episode 71: Score = -83.19592660737484, Avg Score = -178.35\n",
            "Episode 81: Score = -54.57273820874984, Avg Score = -171.67\n",
            "Episode 91: Score = -81.99677183204506, Avg Score = -168.85\n",
            "Episode 101: Score = -199.3870056588955, Avg Score = -167.00\n",
            "Episode 111: Score = -155.12409953340443, Avg Score = -160.42\n",
            "Episode 121: Score = -186.66137387224376, Avg Score = -156.89\n",
            "Episode 131: Score = -80.93929446665268, Avg Score = -159.00\n",
            "Episode 141: Score = -143.01664897626458, Avg Score = -146.57\n",
            "Episode 151: Score = -98.23742098207072, Avg Score = -139.86\n",
            "Episode 161: Score = -75.41364433560273, Avg Score = -127.90\n",
            "Episode 171: Score = -58.75022568784989, Avg Score = -131.67\n",
            "Episode 181: Score = -21.626527610562675, Avg Score = -125.76\n",
            "Episode 191: Score = -59.670876816191885, Avg Score = -125.62\n",
            "Episode 201: Score = -96.87853481773921, Avg Score = -125.92\n",
            "Episode 211: Score = 29.525598405595826, Avg Score = -123.39\n",
            "Episode 221: Score = -27.57211293321184, Avg Score = -108.89\n",
            "Episode 231: Score = -89.99422598085592, Avg Score = -103.10\n",
            "Episode 241: Score = -386.9875869153624, Avg Score = -105.31\n",
            "Episode 251: Score = -104.9785411502358, Avg Score = -113.29\n",
            "Episode 261: Score = -121.01299642292402, Avg Score = -117.77\n",
            "Episode 271: Score = -310.5807580491253, Avg Score = -132.43\n",
            "Episode 281: Score = -17.3289257723224, Avg Score = -149.13\n",
            "Episode 291: Score = -260.2651762147656, Avg Score = -156.02\n",
            "Episode 301: Score = -23.227883199352206, Avg Score = -146.96\n",
            "Episode 311: Score = -26.300368432672073, Avg Score = -147.47\n",
            "Episode 321: Score = -2.4091913574459625, Avg Score = -144.32\n",
            "Episode 331: Score = -5.755042479003592, Avg Score = -143.46\n",
            "Episode 341: Score = -162.34715995126777, Avg Score = -138.83\n",
            "Episode 351: Score = -191.78760724010664, Avg Score = -136.37\n",
            "Episode 361: Score = -155.985389643097, Avg Score = -131.40\n",
            "Episode 371: Score = -25.752568446121703, Avg Score = -116.93\n",
            "Episode 381: Score = -20.463149733834456, Avg Score = -104.24\n",
            "Episode 391: Score = -126.70481104462168, Avg Score = -88.51\n",
            "Episode 401: Score = -115.73874614982606, Avg Score = -93.42\n",
            "Episode 411: Score = -71.30984003912239, Avg Score = -95.94\n",
            "Episode 421: Score = -52.48676419133157, Avg Score = -101.00\n",
            "Episode 431: Score = -30.015326747206256, Avg Score = -103.08\n",
            "Episode 441: Score = -20.37049555500532, Avg Score = -95.08\n",
            "Episode 451: Score = -3.3306588776422394, Avg Score = -81.52\n",
            "Episode 461: Score = -32.92436129514365, Avg Score = -74.85\n",
            "Episode 471: Score = 44.462690801854244, Avg Score = -59.35\n",
            "Episode 481: Score = -8.850063937235058, Avg Score = -48.66\n",
            "Episode 491: Score = -187.84173497525197, Avg Score = -45.14\n",
            "Episode 501: Score = -31.749635187462303, Avg Score = -34.71\n",
            "Episode 511: Score = 16.15770389138484, Avg Score = -25.22\n",
            "Episode 521: Score = 18.69744215973087, Avg Score = -18.27\n",
            "Episode 531: Score = -9.601460704682239, Avg Score = -8.36\n",
            "Episode 541: Score = -17.940286403842194, Avg Score = -9.40\n",
            "Episode 551: Score = -35.24752992665047, Avg Score = -11.31\n",
            "Episode 561: Score = -59.1889313473156, Avg Score = -12.23\n",
            "Episode 571: Score = -10.141438118494733, Avg Score = -17.22\n",
            "Episode 581: Score = 17.803191424104455, Avg Score = -18.82\n",
            "Episode 591: Score = -8.521207443328137, Avg Score = -17.96\n",
            "Episode 601: Score = -37.833244327288156, Avg Score = -18.18\n",
            "Episode 611: Score = -12.385524450906019, Avg Score = -21.74\n",
            "Episode 621: Score = -55.37149422938777, Avg Score = -25.73\n",
            "Episode 631: Score = 7.059279513140538, Avg Score = -25.20\n",
            "Episode 641: Score = 38.20854114354739, Avg Score = -19.27\n",
            "Episode 651: Score = -104.56316095850582, Avg Score = -17.01\n",
            "Episode 661: Score = 20.11125341829644, Avg Score = -13.05\n",
            "Episode 671: Score = 58.6904851869878, Avg Score = -15.16\n",
            "Episode 681: Score = 33.94774340963677, Avg Score = -13.28\n",
            "Episode 691: Score = -262.60383367114486, Avg Score = -13.80\n",
            "Episode 701: Score = -44.03477380542823, Avg Score = -11.14\n",
            "Episode 711: Score = -163.37843645305455, Avg Score = -5.43\n",
            "Episode 721: Score = 17.353415703970455, Avg Score = -3.89\n",
            "Episode 731: Score = 16.420826019390915, Avg Score = -2.42\n",
            "Episode 741: Score = 79.89873856625287, Avg Score = -2.54\n",
            "Episode 751: Score = 71.35183812639121, Avg Score = 3.06\n",
            "Episode 761: Score = -15.010332559545365, Avg Score = 6.39\n",
            "Episode 771: Score = 97.14524506535561, Avg Score = 17.73\n",
            "Episode 781: Score = -10.787376036318324, Avg Score = 21.02\n",
            "Episode 791: Score = 39.59169376376451, Avg Score = 28.03\n",
            "Episode 801: Score = -38.24222858894733, Avg Score = 27.39\n",
            "Episode 811: Score = -3.2031155171148624, Avg Score = 27.80\n",
            "Episode 821: Score = 43.455867189193164, Avg Score = 31.91\n",
            "Episode 831: Score = 67.01816005216853, Avg Score = 34.92\n",
            "Episode 841: Score = 73.92092919340762, Avg Score = 35.86\n",
            "Episode 851: Score = 100.84203243690202, Avg Score = 39.69\n",
            "Episode 861: Score = 80.0575260577081, Avg Score = 45.33\n",
            "Episode 871: Score = 141.92995853310734, Avg Score = 45.63\n",
            "Episode 881: Score = -15.928055588536807, Avg Score = 46.96\n",
            "Episode 891: Score = 0.3069586530661752, Avg Score = 43.62\n",
            "Episode 901: Score = 118.29031427569929, Avg Score = 40.09\n",
            "Episode 911: Score = -12.991536993014535, Avg Score = 38.97\n",
            "Episode 921: Score = 128.2698843169083, Avg Score = 47.20\n",
            "Episode 931: Score = 112.7622605157013, Avg Score = 51.58\n",
            "Episode 941: Score = 111.76340982044664, Avg Score = 60.06\n",
            "Episode 951: Score = 130.96667481898461, Avg Score = 63.97\n",
            "Episode 961: Score = -185.57726162038367, Avg Score = 64.36\n",
            "Episode 971: Score = -131.53893782521578, Avg Score = 67.62\n",
            "Episode 981: Score = 64.6636182316823, Avg Score = 69.20\n",
            "Episode 991: Score = -154.96131780382197, Avg Score = 73.76\n",
            "Episode 1001: Score = 103.61116535882739, Avg Score = 79.64\n",
            "Episode 1011: Score = 231.45072515540403, Avg Score = 92.48\n",
            "Episode 1021: Score = -275.1991023491338, Avg Score = 84.46\n",
            "Episode 1031: Score = -17.389492549026627, Avg Score = 86.79\n",
            "Episode 1041: Score = -182.21785707808584, Avg Score = 81.36\n",
            "Episode 1051: Score = 270.0486917092732, Avg Score = 92.43\n",
            "Episode 1061: Score = -4.0305659045397135, Avg Score = 101.41\n",
            "Episode 1071: Score = 34.74520293426585, Avg Score = 112.07\n",
            "Episode 1081: Score = 207.72589258094007, Avg Score = 124.72\n",
            "Episode 1091: Score = 226.29360218183226, Avg Score = 132.18\n",
            "Episode 1101: Score = 166.23606210371707, Avg Score = 147.21\n",
            "Episode 1111: Score = 54.46627453234339, Avg Score = 153.16\n",
            "Episode 1121: Score = 210.43382387168782, Avg Score = 164.77\n",
            "Episode 1131: Score = -30.42661330778668, Avg Score = 169.31\n",
            "Episode 1141: Score = 17.337906610168915, Avg Score = 173.56\n",
            "Episode 1151: Score = 242.28645768387418, Avg Score = 161.45\n",
            "Episode 1161: Score = 250.98692354952414, Avg Score = 158.91\n",
            "Episode 1171: Score = 288.918795036762, Avg Score = 154.03\n",
            "Episode 1181: Score = 229.10859330199492, Avg Score = 143.82\n",
            "Episode 1191: Score = 4.915537723702968, Avg Score = 146.22\n",
            "Episode 1201: Score = 23.190574113030237, Avg Score = 140.27\n",
            "Episode 1211: Score = -16.952597252601777, Avg Score = 141.09\n",
            "Episode 1221: Score = 244.1009085489999, Avg Score = 148.57\n",
            "Episode 1231: Score = 281.08781968081104, Avg Score = 153.70\n",
            "Episode 1241: Score = 287.4632102153039, Avg Score = 165.71\n",
            "Episode 1251: Score = 291.2230431340448, Avg Score = 179.10\n",
            "Episode 1261: Score = 255.71320047572712, Avg Score = 179.72\n",
            "Episode 1271: Score = -6.033742872038275, Avg Score = 188.88\n",
            "Environment solved!\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "WARNING:imageio_ffmpeg:IMAGEIO FFMPEG_WRITER WARNING: input image is not divisible by macro_block_size=16, resizing from (600, 400) to (608, 400) to ensure video compatibility with most codecs and players. To prevent resizing, make your input image divisible by the macro_block_size or set the macro_block_size to 1 (risking incompatibility).\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Video saved as lunar_lander.mp4\n"
          ]
        }
      ],
      "source": [
        "import gymnasium as gym\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "import numpy as np\n",
        "from collections import deque\n",
        "import imageio\n",
        "\n",
        "# Policy Network\n",
        "class PolicyNetwork(nn.Module):\n",
        "    def __init__(self, state_dim, action_dim, hidden_size=128):\n",
        "        super(PolicyNetwork, self).__init__()\n",
        "        self.fc1 = nn.Linear(state_dim, hidden_size)\n",
        "        self.fc2 = nn.Linear(hidden_size, hidden_size)\n",
        "        self.fc3 = nn.Linear(hidden_size, action_dim)\n",
        "        self.softmax = nn.Softmax(dim=-1)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = torch.relu(self.fc1(x))\n",
        "        x = torch.relu(self.fc2(x))\n",
        "        x = self.fc3(x)\n",
        "        return self.softmax(x)\n",
        "\n",
        "# REINFORCE Algorithm\n",
        "class REINFORCEAgent:\n",
        "    def __init__(self, state_dim, action_dim, lr=1e-3, gamma=0.995):\n",
        "        self.policy = PolicyNetwork(state_dim, action_dim)\n",
        "        self.optimizer = optim.Adam(self.policy.parameters(), lr=lr)\n",
        "        self.gamma = gamma\n",
        "        self.baseline = 0\n",
        "\n",
        "    def select_action(self, state):\n",
        "        state = torch.tensor(state, dtype=torch.float32)\n",
        "        action_probs = self.policy(state)\n",
        "        action_dist = torch.distributions.Categorical(action_probs)\n",
        "        action = action_dist.sample()\n",
        "        return action.item(), action_dist.log_prob(action)\n",
        "\n",
        "    def normalize_rewards(self, rewards):\n",
        "        rewards = np.array(rewards)\n",
        "        rewards -= rewards.mean()\n",
        "        rewards /= (rewards.std() + 1e-6)\n",
        "        return rewards\n",
        "\n",
        "    def update_policy(self, rewards, log_probs):\n",
        "        discounted_rewards = []\n",
        "        G = 0\n",
        "        for r in reversed(rewards):\n",
        "            G = r + self.gamma * G\n",
        "            discounted_rewards.insert(0, G)\n",
        "\n",
        "        discounted_rewards = self.normalize_rewards(discounted_rewards)\n",
        "        discounted_rewards = torch.tensor(discounted_rewards, dtype=torch.float32)\n",
        "        self.baseline = 0.9 * self.baseline + 0.1 * discounted_rewards.mean()\n",
        "        advantages = discounted_rewards - self.baseline\n",
        "\n",
        "        # Ensure advantages are detached for the loss calculation\n",
        "        loss = torch.sum(-torch.stack(log_probs) * advantages.detach())\n",
        "\n",
        "        self.optimizer.zero_grad()\n",
        "        loss.backward()\n",
        "        nn.utils.clip_grad_norm_(self.policy.parameters(), max_norm=1.0)\n",
        "        self.optimizer.step()\n",
        "\n",
        "def record_video(agent, filename=\"lunar_lander.mp4\", max_steps=1000):\n",
        "    env = gym.make(\"LunarLander-v3\", render_mode=\"rgb_array\")\n",
        "    state, _ = env.reset()\n",
        "    frames = []\n",
        "\n",
        "    for _ in range(max_steps):\n",
        "        frame = env.render()\n",
        "        frames.append(frame)\n",
        "        action, _ = agent.select_action(state)\n",
        "        state, _, terminated, truncated, _ = env.step(action)\n",
        "        if terminated or truncated:\n",
        "            break\n",
        "\n",
        "    env.close()\n",
        "    imageio.mimsave(filename, frames, fps=30)\n",
        "    print(f\"Video saved as {filename}\")\n",
        "\n",
        "# Training loop\n",
        "def train_agent(episodes=2000, max_steps=1000):\n",
        "    env = gym.make(\"LunarLander-v3\")\n",
        "    state_dim = env.observation_space.shape[0]\n",
        "    action_dim = env.action_space.n\n",
        "    agent = REINFORCEAgent(state_dim, action_dim)\n",
        "\n",
        "    scores = deque(maxlen=100)\n",
        "\n",
        "    for episode in range(episodes):\n",
        "        state, _ = env.reset()\n",
        "        log_probs = []\n",
        "        rewards = []\n",
        "        score = 0\n",
        "\n",
        "        for _ in range(max_steps):\n",
        "            action, log_prob = agent.select_action(state)\n",
        "            next_state, reward, terminated, truncated, _ = env.step(action)\n",
        "\n",
        "            log_probs.append(log_prob)\n",
        "            rewards.append(reward)\n",
        "            score += reward\n",
        "            state = next_state\n",
        "\n",
        "            if terminated or truncated:\n",
        "                break\n",
        "\n",
        "        agent.update_policy(rewards, log_probs)\n",
        "        scores.append(score)\n",
        "        avg_score = np.mean(scores)\n",
        "\n",
        "        if episode % 10 == 0:\n",
        "            print(f\"Episode {episode + 1}: Score = {score}, Avg Score = {avg_score:.2f}\")\n",
        "\n",
        "        if avg_score >= 200:\n",
        "            print(\"Environment solved!\")\n",
        "            break\n",
        "\n",
        "    env.close()\n",
        "    record_video(agent=agent)\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    train_agent()\n"
      ]
    }
  ]
}