Hands_On_6_RL.ipynb

{
  "cells": [
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Reinforcement Learning Tutorial\n",
        "Tutorial adapted by Hongfei from [PyTorch Tutorial](https://pytorch.org/tutorials/intermediate/reinforcement_q_learning.html).\n",
        "\n",
        "This tutorial shows how to use PyTorch to train Deep Q-Learning (DQL) agents\n",
        "on the LunarLander task and the CartPole-v1 task from [Gymnasium](https://www.gymnasium.farama.org)."
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Task -- LunarLander\n",
        "\n",
        "State dimension: 8\n",
        "- $x,y$ positions\n",
        "- $x,y$ linear velocities\n",
        "- angle and angular velocity\n",
        "- two True/False values indicating whether left or right leg touches the ground or not\n",
        "\n",
        "Action space dimension: 4\n",
        "- 0: do nothing\n",
        "- 1: left engine fire\n",
        "- 2: main engine fire\n",
        "- 3: right engine fire\n",
        "\n",
        "For details of this task please refer to SPH6004 lecture notes. Environment provided by the [Gymnasium](https://gymnasium.farama.org/environments/box2d/lunar_lander/) project.\n",
        "\n",
        "![LunarLander](Figs/lunar_lander.gif)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "yaicoV3vuwMi"
      },
      "source": [
        "## Task -- CartPole-v1\n",
        "\n",
        "The agent has to decide between two actions - moving the cart left or\n",
        "right - so that the pole attached to it stays upright. You can find more\n",
        "information about the environment at\n",
        "[Gymnasium's website](https://gymnasium.farama.org/environments/classic_control/cart_pole/).\n",
        "\n",
        "![CartPole-v1](Figs/cartpole.gif)\n",
        "\n",
        "As the agent observes the current state of the environment and chooses\n",
        "an action, the environment *transitions* to a new state, and also\n",
        "returns a reward that indicates the consequences of the action. In this\n",
        "task, rewards are +1 for every incremental timestep and the environment\n",
        "terminates if the pole falls over too far or the cart moves more than 2.4\n",
        "units away from center. This means better performing scenarios will run\n",
        "for longer duration, accumulating larger return.\n",
        "\n",
        "The CartPole task is designed so that the inputs to the agent are 4 real\n",
        "values representing the environment state (0: cart position, 1: cart velocity, 2: pole angle (in radians), 3: pole angular velocity).\n",
        "We take these 4 inputs without any scaling and pass them through a \n",
        "small fully-connected network with 2 outputs, one for each action (0: left, 1: right). "
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "----"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4KYBsUWZuwMl",
        "outputId": "06e3506f-659b-4ccc-b8fd-0220da8d7a9e"
      },
      "outputs": [],
      "source": [
        "# %%bash\n",
        "# pip3 install gymnasium[classic_control]\n",
        "# pip3 install gymnasium[box2d]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "-4ryXFF-uwMl"
      },
      "outputs": [],
      "source": [
        "import gymnasium as gym\n",
        "import math\n",
        "import random\n",
        "from collections import namedtuple, deque\n",
        "from itertools import count\n",
        "import numpy as np\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "import torch.nn.functional as F\n",
        "from time import sleep\n",
        "from tqdm.notebook import tqdm\n",
        "import os\n",
        "\n",
        "task = [\"LunarLander-v2\", \"CartPole-v1\"][0]\n",
        "\n",
        "# if gpu is to be used\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Helper function to render environment\n",
        "\n",
        "def render(task=\"LunarLander-v2\", policy_net = None, wind=False):\n",
        "    # if policy_net is None, use randomly sampled policy for action.\n",
        "    if task==\"LunarLander-v2\":\n",
        "        env = gym.make(task,render_mode='human',enable_wind=wind,wind_power=10)\n",
        "    elif task==\"CartPole-v1\":\n",
        "        env = gym.make(task,render_mode='human')\n",
        "\n",
        "    observation, info = env.reset()\n",
        "    total_reward = 0\n",
        "    with torch.no_grad():\n",
        "        for i in count():\n",
        "            if policy_net is None:\n",
        "                action = env.action_space.sample()        \n",
        "            else:\n",
        "                q_values = policy_net(torch.tensor(observation,device=device).unsqueeze(0))\n",
        "                action = q_values.argmax().item()\n",
        "\n",
        "            observation, reward, terminated, truncated, _ = env.step(action)\n",
        "            total_reward += reward\n",
        "            _ = env.render()\n",
        "            if terminated or truncated:\n",
        "                policy_mode = 'random' if policy_net is None else 'policy_net'\n",
        "                print('Poicy is {}, episode duration: {}'.format(policy_mode,i+1))\n",
        "                print('Total reward is {:.2f}'.format(total_reward))\n",
        "                sleep(1)\n",
        "                env.close()\n",
        "                break"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Poicy is random, episode duration: 44\n",
            "Total reward is 44.00\n"
          ]
        }
      ],
      "source": [
        "render(task)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "bKOw5U9buwMn"
      },
      "source": [
        "## Replay Memory\n",
        "\n",
        "Experience relay stores\n",
        "the transitions that the agent observes, allowing us to reuse this data\n",
        "later. By sampling from it randomly, the transitions that build up a\n",
        "batch are decorrelated. It has been shown that this greatly stabilizes\n",
        "and improves the DQL training procedure.\n",
        "\n",
        "For this, we're going to need two classses:\n",
        "\n",
        "-  ``Transition`` - a named tuple representing a single transition in\n",
        "   our environment.\n",
        "-  ``ReplayMemory`` - a cyclic buffer of bounded size that holds the\n",
        "   transitions observed recently. It also implements a ``.sample()``\n",
        "   method for selecting a random batch of transitions for training.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "WyqDNEnouwMn"
      },
      "outputs": [],
      "source": [
        "Transition = namedtuple('Transition',\n",
        "                        ('state', 'action', 'reward', 'next_state'))\n",
        "\n",
        "\n",
        "class ReplayMemory(object):\n",
        "\n",
        "    def __init__(self, capacity):\n",
        "        self.memory = deque([], maxlen=capacity)\n",
        "\n",
        "    def push(self, *args):\n",
        "        \"\"\"Save a transition\"\"\"\n",
        "        self.memory.append(Transition(*args))\n",
        "\n",
        "    def sample(self, batch_size):\n",
        "        return random.sample(self.memory, batch_size)\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.memory)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "JV2aFH3FuwMp"
      },
      "source": [
        "Now, let's define our model.\n",
        "\n",
        "### Q-network\n",
        "\n",
        "We will use a feed forward neural network $N$ to approximate best utility functions.\n",
        "\n",
        "#### LunarLander-v2\n",
        "\n",
        "The input of network $N$ has dimension 8. It has four outputs, representing $Q(s,\\text{do nothing})$, $Q(s,\\text{left engine fire})$, $Q(s,\\text{main engine fire})$, and $Q(s,\\text{right engine fire})$ (where $s$ is the input to the network representing current state). In effect, the network is trying to predict the *optimal total reward* of\n",
        "taking each action given the current input.\n",
        "\n",
        "#### CartPole-v1\n",
        "\n",
        "The network $N$ input has dimension 4. It has two\n",
        "outputs, representing $Q(s, \\mathrm{left})$ and\n",
        "$Q(s, \\mathrm{right})$ (where $s$ is the input to the\n",
        "network representing current state). In effect, the network is trying to predict the *optimal total reward* of\n",
        "taking each action given the current input.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "gfzY6tmZuwMq"
      },
      "outputs": [],
      "source": [
        "class DQN(nn.Module):\n",
        "\n",
        "    def __init__(self, state_space_dim, n_actions):\n",
        "        super(DQN, self).__init__()\n",
        "        self.layer1 = nn.Linear(state_space_dim, 128)\n",
        "        self.layer2 = nn.Linear(128, 128)\n",
        "        self.layer3 = nn.Linear(128, n_actions)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = F.relu(self.layer1(x))\n",
        "        x = F.relu(self.layer2(x))\n",
        "        return self.layer3(x)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## $\\varepsilon$-greedy algorithm\n",
        "\n",
        "The following cell defines a function to implement $\\varepsilon$-greedy algorithm. Here we use a dynamic $\\varepsilon$, such that at the beginning $t=0$ we have a large $\\varepsilon=\\varepsilon_\\text{start}$, while as number of steps $t\\to \\infty$ the threshold $\\varepsilon\\to \\varepsilon_\\text{end}$ for a small value $\\varepsilon_\\text{end}$. Both $\\varepsilon_\\text{start}$ and $\\varepsilon_\\text{end}$ are hyper-parameters defined later.\n",
        "\n",
        "About `tensor.max()` method:\n",
        "Say we have tensor\n",
        "```\n",
        "a = tensor([[-0.5044, -1.5145],\n",
        "            [-0.8489,  0.0904],\n",
        "            [ 0.0348,  2.0637]])\n",
        "```\n",
        "then `a.max(dim=1)` returns two tensors, where the first contains the maximum along `dim=1`, i.e.\n",
        "```\n",
        "tensor([-0.5044,  0.0904,  2.0637])\n",
        "```\n",
        "and the second contains the indices of the maximum elements, i.e.\n",
        "```\n",
        "indices=tensor([0, 1, 1]).\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "mlJJlg2PuwMr"
      },
      "outputs": [],
      "source": [
        "def select_action(state):\n",
        "    global steps_done\n",
        "    sample = random.random()\n",
        "    eps_threshold = EPS_END + (EPS_START - EPS_END) * \\\n",
        "        math.exp(-1. * steps_done / EPS_DECAY)\n",
        "    steps_done += 1\n",
        "    if sample > eps_threshold:\n",
        "        with torch.no_grad():\n",
        "            # we pick action with the larger optimal total rewards\n",
        "            # .max(dim=1)[1] pick up the maximum indices\n",
        "            return policy_net(state).max(dim=1)[1].view(1, 1)\n",
        "    else:\n",
        "        # we randomly select an action from environment's action space\n",
        "        return torch.tensor([[env.action_space.sample()]], device=device, dtype=torch.long)\n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "RqWRuYdFuwMq"
      },
      "source": [
        "## Training"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "3dAvYr2yuwMs"
      },
      "source": [
        "Finally, the code for training our model.\n",
        "\n",
        "Here, you can find an ``optimize_model`` function that performs a\n",
        "single step of the optimization. It performs the following:\n",
        "1. We sample a batch of experience $(s,a,r,s')$ from our experience replay `memory`, concatenates\n",
        "all the tensors into a single one.\n",
        "\n",
        "2. Computes $Q(s, a)$ using our policy network.\n",
        "\n",
        "3. Compute $V(s') = \\max_a' Q(s', a')$ using our target network.\n",
        "\n",
        "4. Combine $Q(s, a)$, $V(s')$ and the current reward $r$ into our loss.\n",
        "\n",
        "5. (We set $V(s') = 0$ if $s$ is a terminal state.)\n",
        "\n",
        "6. Update policy network by stochastic gradient descent.\n",
        "\n",
        "7. [Soft update](https://arxiv.org/pdf/1509.02971.pdf) update target network. Soft update controlled by the hyperparameter ``TAU``, which will be defined later."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "aeTqGCzPuwMs"
      },
      "outputs": [],
      "source": [
        "def optimize_model():\n",
        "    if len(memory) < BATCH_SIZE:\n",
        "        # If we do not have enough experience in experience replay, do nothing\n",
        "        return\n",
        "    \n",
        "    transitions = memory.sample(BATCH_SIZE)\n",
        "    # Transpose the batch (see https://stackoverflow.com/a/19343/3343043 for\n",
        "    # detailed explanation). This converts batch-array of Transitions\n",
        "    # to Transition of batch-arrays.\n",
        "    batch = Transition(*zip(*transitions))\n",
        "\n",
        "    # Compute a mask of non-final states and concatenate the batch elements\n",
        "    # (a final state would've been the one after which simulation ended)\n",
        "    non_final_mask = torch.tensor(tuple(map(lambda s: s is not None,\n",
        "                                          batch.next_state)), device=device, dtype=torch.bool)\n",
        "    non_final_next_states = torch.cat([s for s in batch.next_state\n",
        "                                                if s is not None])\n",
        "    state_batch = torch.cat(batch.state)\n",
        "    action_batch = torch.cat(batch.action)\n",
        "    reward_batch = torch.cat(batch.reward)\n",
        "\n",
        "    # Compute Q(s_t, a) - the model computes Q(s_t), then we select the\n",
        "    # columns of actions taken. These are the actions which would've been taken\n",
        "    # for each batch state according to policy_net\n",
        "    state_action_values = policy_net(state_batch).gather(dim=1, index=action_batch)\n",
        "\n",
        "    # Compute V(s_{t+1}) for all next states.\n",
        "    # Expected values of actions for non_final_next_states are computed based\n",
        "    # on the \"older\" target_net; selecting their best reward with max(1)[0].\n",
        "    # This is merged based on the mask, such that we'll have either the expected\n",
        "    # state value or 0 in case the state was final.\n",
        "    next_state_values = torch.zeros(BATCH_SIZE, device=device)\n",
        "    with torch.no_grad():\n",
        "        next_state_values[non_final_mask] = target_net(non_final_next_states).max(dim=1)[0]\n",
        "    # Compute the expected Q values\n",
        "    expected_state_action_values = (next_state_values * GAMMA) + reward_batch\n",
        "\n",
        "    # Compute Huber loss\n",
        "    criterion = nn.SmoothL1Loss()\n",
        "    loss = criterion(state_action_values, expected_state_action_values.unsqueeze(1))\n",
        "\n",
        "    # Optimize the model\n",
        "    optimizer.zero_grad()\n",
        "    loss.backward()\n",
        "    # In-place gradient clipping\n",
        "    torch.nn.utils.clip_grad_value_(policy_net.parameters(), 100)\n",
        "    optimizer.step()\n",
        "\n",
        "    # Soft update of the target network's weights\n",
        "    # θ′ ← τ θ + (1 −τ )θ′\n",
        "    target_net_state_dict = target_net.state_dict()\n",
        "    policy_net_state_dict = policy_net.state_dict()\n",
        "    for key in policy_net_state_dict:\n",
        "        target_net_state_dict[key] = policy_net_state_dict[key]*TAU + target_net_state_dict[key]*(1-TAU)\n",
        "    target_net.load_state_dict(target_net_state_dict)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {},
      "outputs": [],
      "source": [
        "# ------------------------ Hyper-parameters ------------------\n",
        "# BATCH_SIZE is the number of experience sampled from experience replay (We train on several experience in parallel)\n",
        "# GAMMA is the discount factor for future rewards\n",
        "# EPS_START is the starting value of epsilon (exploration probability)\n",
        "# EPS_END is the final value of epsilon\n",
        "# EPS_DECAY controls the rate of exponential decay of epsilon, higher means a slower decay\n",
        "# TAU is the update rate of the target network\n",
        "# LR is the learning rate of the AdamW optimizer\n",
        "BATCH_SIZE = 128\n",
        "GAMMA = 0.99\n",
        "EPS_START = 0.9\n",
        "EPS_END = 0.1\n",
        "EPS_DECAY = 10000\n",
        "TAU = 0.005\n",
        "LR = 1e-4\n",
        "memory = ReplayMemory(10000)\n",
        "\n",
        "steps_done = 0\n",
        "num_episodes = 1000\n",
        "total_rewards = []\n",
        "best_total_rewards = float('-inf')\n",
        "\n",
        "# ------------------------ Re-create environment ------------------\n",
        "\n",
        "env = gym.make(task)\n",
        "\n",
        "# Get number of actions from gym action space\n",
        "n_actions = env.action_space.n\n",
        "# Get the number of state observations\n",
        "state, info = env.reset()\n",
        "n_observations = len(state)\n",
        "episode_durations = []\n",
        "\n",
        "# ------------------------- Initialize networks ---------------------\n",
        "\n",
        "policy_net = DQN(n_observations, n_actions).to(device)\n",
        "target_net = DQN(n_observations, n_actions).to(device)\n",
        "target_net.load_state_dict(policy_net.state_dict())\n",
        "\n",
        "optimizer = optim.AdamW(policy_net.parameters(), lr=LR, amsgrad=True)\n",
        "\n",
        "\n",
        "model_save_path = './ckpt'\n",
        "os.makedirs(model_save_path,exist_ok=True)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "Agl462hquwMt"
      },
      "source": [
        "Below, you can find the main training loop. At the beginning we reset\n",
        "the environment and obtain the initial ``state`` Tensor. Then, we sample\n",
        "an action, execute it, observe the next state and the reward, and optimize our model once. When the episode ends, we restart the loop.\n",
        "\n",
        "Note that for LunarLander, environment will terminate if total steps exceeds 1000, while for CartPole, environment will terminate if total steps exceeds 500."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 347
        },
        "id": "7x1myiSEuwMt",
        "outputId": "74f38ab5-3cab-4afb-bd2c-2be49af863ba"
      },
      "outputs": [
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "9e332029ccea4bc8b3d87f6e17acd161",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "  0%|          | 0/500 [00:00<?, ?it/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Episode 0, avg episode_durations 56.00, avg total reward 56.00\n",
            "Episode 50, avg episode_durations 19.28, avg total reward 19.28\n",
            "Episode 100, avg episode_durations 17.34, avg total reward 17.34\n",
            "Episode 150, avg episode_durations 48.16, avg total reward 48.16\n",
            "Episode 200, avg episode_durations 99.02, avg total reward 99.02\n",
            "Episode 250, avg episode_durations 270.04, avg total reward 270.04\n",
            "Episode 300, avg episode_durations 208.62, avg total reward 208.62\n",
            "Episode 350, avg episode_durations 225.74, avg total reward 225.74\n",
            "Episode 400, avg episode_durations 303.62, avg total reward 303.62\n",
            "Episode 450, avg episode_durations 495.30, avg total reward 495.30\n",
            "Complete\n"
          ]
        }
      ],
      "source": [
        "# Training loop\n",
        "for i_episode in tqdm(range(num_episodes)):\n",
        "    # Initialize the environment and get it's state\n",
        "    state, info = env.reset()\n",
        "    total_reward = 0\n",
        "    state = torch.tensor(state, dtype=torch.float32, device=device).unsqueeze(0)\n",
        "    for t in count():\n",
        "        action = select_action(state)\n",
        "        observation, reward, terminated, truncated, _ = env.step(action.item())\n",
        "        total_reward += reward\n",
        "        reward = torch.tensor([reward], device=device)\n",
        "        done = terminated or truncated\n",
        "\n",
        "        if terminated:\n",
        "            next_state = None\n",
        "        else:\n",
        "            next_state = torch.tensor(observation, dtype=torch.float32, device=device).unsqueeze(0)\n",
        "\n",
        "        # Store the transition in memory\n",
        "        memory.push(state, action, reward, next_state)\n",
        "\n",
        "        # Move to the next state\n",
        "        state = next_state\n",
        "\n",
        "        # Perform one step of the optimization (on the policy network)\n",
        "        optimize_model()\n",
        "\n",
        "        if done:\n",
        "            episode_durations.append(t + 1)\n",
        "            total_rewards.append(total_reward)\n",
        "            # print and save model every 50 episodes\n",
        "            aw = np.minimum(len(episode_durations),50)\n",
        "            if i_episode%50 == 0:\n",
        "                if best_total_rewards < np.mean(total_rewards[-aw:]):\n",
        "                    # save model if avg performance is good\n",
        "                    torch.save(policy_net.state_dict(),'./ckpt/{}_checkpoint.pt'.format(task))\n",
        "                    best_total_rewards = np.mean(total_rewards[-aw:])\n",
        "\n",
        "                print('Episode {}, avg episode_durations {:.2f}, avg total reward {:.2f}'.format(i_episode,np.mean(episode_durations[-aw:]),np.mean(total_rewards[-aw:])))\n",
        "            break\n",
        "\n",
        "print('Complete')"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "Jw2Lhc_auwMu"
      },
      "source": [
        "Lastly we visualize the trained policy net."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Poicy is policy_net, episode duration: 500\n",
            "Total reward is 500.00\n"
          ]
        }
      ],
      "source": [
        "\n",
        "policy_net.load_state_dict(torch.load('./ckpt/{}_checkpoint.pt'.format(task)))\n",
        "# Set policy_net=None to visualize a random policy.\n",
        "render(task=task, policy_net=policy_net,wind=False)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "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.10.10"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}