12_a2c.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 5,
  "metadata": {
    "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.8.5"
    },
    "colab": {
      "name": "12_a2c.ipynb",
      "provenance": []
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fantastic-department"
      },
      "source": [
        "# The Actor-Critic Method\n",
        "<table align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/matyama/deep-rl-hands-on/blob/main/12_a2c.ipynb\">\n",
        "    <img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />\n",
        "        Run in Google Colab\n",
        "    </a>\n",
        "  </td>\n",
        "</table>"
      ],
      "id": "fantastic-department"
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "buried-kennedy",
        "outputId": "647555a6-581d-447e-e7cf-b53bf9a7249d"
      },
      "source": [
        "%%bash\n",
        "!(stat -t /usr/local/lib/*/dist-packages/google/colab > /dev/null 2>&1) && exit\n",
        "\n",
        "echo \"Running on Google Colab, therefore installing dependencies...\"\n",
        "pip install ptan>=0.7 tensorboardX"
      ],
      "id": "buried-kennedy",
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Running on Google Colab, therefore installing dependencies...\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "accepted-papua"
      },
      "source": [
        "## Variance Reduction\n",
        "Let's start by recalling the policy gradient defined by the *Policy Gradients (PG)* method:\n",
        "$$\n",
        "\\nabla J \\approx \\mathbb{E}[Q(s, a) \\nabla \\log(\\pi(a|s))]\n",
        "$$\n",
        "\n",
        "One of the weak points of the PG method is that the gradient scales $Q(s, a)$ may experience quite significant variance* which does not help the training at all. We fixed this issue by introducing a fixed *baseline* value (e.g. mean reward) that was subtracted from the gradient scales Q.\n",
        "\n",
        "\\* Recall formal defintion: $\\mathbb{V}[X] = \\mathbb{E}[(X - \\mathbb{E}[X])^2]$\n",
        "\n",
        "Let's illustrate this problem and solution on simple example:\n",
        "* Assume there are three actions with $Q_1$, $Q_2$ some small positive values and $Q_3$ being large negative\n",
        "* In this case there will be small positive gradient towards fist two actions and large negative one repelling the policy from the third one\n",
        "* Now imagine $Q_1$ and $Q_2$ were large positive values instead. Then $Q_3$ would become small but positive value. The gradient would still push the policy towards fist two actions but it would direct the gradient towards the trird one a bit as well (instead of pushing it away from it)!\n",
        "\n",
        "Now it's a bit more clear why subtracting a constant value that we called the *baseline* helps."
      ],
      "id": "accepted-papua"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "opposed-leone"
      },
      "source": [
        "## Advantage Actor-Critic (A2C)\n",
        "*Advantage Actor-Critic (A2C)* method can be viewed as a combination of PG and DQN with a simple idea extending the variance reduction theme we discussed above. Until now we treated the *baseline* value as single constant that we subtracted from all $Q(s, a)$ values. A2C pushes this further and uses different baselines for each state $s$.\n",
        "\n",
        "If one recalls the *Duelling DQN* which exploited the fact that $Q(s, a) = V(s) + A(s, a)$ - i.e. state-action values are composed of a *baseline* state values $V(s)$ and action advantages $A(s, a)$ in these states, it is quite straightforward to figure out which values A2C uses as state baselines - the state values $V(s)$!\n",
        "\n",
        "The *Advantage* Actor-Critic name then comes from the fact that our gradient scales turn to action advantages after subtracting state values:\n",
        "$$\n",
        "\\mathbb{E}[Q(s, a) \\nabla \\log(\\pi(a|s))] \\to \\mathbb{E}[A(s, a) \\nabla \\log(\\pi(a|s))]\n",
        "$$\n",
        "\n",
        "Finally, the question is how do we obtain $V(s)$? Here comes the second part which is the combination with the DQN approach - we simply train a DQN alongside our PGN.\n",
        "\n",
        "*Notes*:\n",
        "* *There'll actually be just single NN that will learn both the policy and state values (discussed below)*\n",
        "* *Improvements from both methods are still applicable (also metioned and shown in following sections)*"
      ],
      "id": "opposed-leone"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "assigned-cisco"
      },
      "source": [
        "### Common Imports"
      ],
      "id": "assigned-cisco"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mechanical-expert"
      },
      "source": [
        "# flake8: noqa: E402,I001\n",
        "\n",
        "import time\n",
        "from typing import Any, List, Sequence, Tuple\n",
        "\n",
        "import gym\n",
        "import numpy as np\n",
        "import ptan\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from ptan.experience import ExperienceFirstLast\n",
        "from tensorboardX import SummaryWriter"
      ],
      "id": "mechanical-expert",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "further-tomorrow"
      },
      "source": [
        "### Reward Tracker"
      ],
      "id": "further-tomorrow"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "checked-bones"
      },
      "source": [
        "class RewardTracker:\n",
        "    def __init__(\n",
        "        self,\n",
        "        writer: SummaryWriter,\n",
        "        stop_reward: float,\n",
        "        window_size: int = 100,\n",
        "    ) -> None:\n",
        "        self.writer = writer\n",
        "        self.stop_reward = stop_reward\n",
        "        self.window_size = window_size\n",
        "        self.best_mean_reward = float(\"-inf\")\n",
        "\n",
        "    def __enter__(self) -> \"RewardTracker\":\n",
        "        self.ts = time.time()\n",
        "        self.ts_frame = 0\n",
        "        self.total_rewards = []\n",
        "        return self\n",
        "\n",
        "    def __exit__(self, *args: Any) -> None:\n",
        "        self.writer.close()\n",
        "\n",
        "    def add_reward(self, reward: float, frame: int) -> bool:\n",
        "        \"\"\"\n",
        "        Returns an indication of whether a termination contition was reached.\n",
        "        \"\"\"\n",
        "\n",
        "        self.total_rewards.append(reward)\n",
        "\n",
        "        fps = (frame - self.ts_frame) / (time.time() - self.ts)\n",
        "\n",
        "        self.ts_frame = frame\n",
        "        self.ts = time.time()\n",
        "\n",
        "        mean_reward = np.mean(self.total_rewards[-self.window_size :])\n",
        "\n",
        "        if mean_reward > self.best_mean_reward:\n",
        "            self.best_mean_reward = mean_reward\n",
        "            print(\n",
        "                f\"{frame}: done {len(self.total_rewards)} games, \"\n",
        "                f\"mean reward {mean_reward:.3f}, speed {fps:.2f} fps\"\n",
        "            )\n",
        "\n",
        "        self.writer.add_scalar(\"fps\", fps, frame)\n",
        "        self.writer.add_scalar(\"reward_100\", mean_reward, frame)\n",
        "        self.writer.add_scalar(\"reward\", reward, frame)\n",
        "\n",
        "        return mean_reward > self.stop_reward"
      ],
      "id": "checked-bones",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "australian-brake"
      },
      "source": [
        "### Atari A2C PG Network\n",
        "This NN is quite similar to the *Dueling DQN* architecture but with an important difference. In the Dueling DQN we have also two parts\n",
        "1. Part for the state values $V(s)$\n",
        "1. Part for the action advantages $A(s, a)$\n",
        "\n",
        "But as with any other DQN we did still output $Q(s, a) = V(s) + A(s, a)$. Here we have two separate outputs with common base network:\n",
        "1. Policy network that outputs action logits - basically policy $\\pi(a|s)$ when one converts them to probabilities using softmax\n",
        "1. Value network which computes $V(s)$"
      ],
      "id": "australian-brake"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "killing-lloyd"
      },
      "source": [
        "class AtariA2C(nn.Module):\n",
        "    \"\"\"\n",
        "    A2C network with 2D convolutional base for Atari envs. and two heads:\n",
        "    1. Policy - dense network that outputs action logits\n",
        "    2. Value - dence network that models state values `V(s)`\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, input_shape: Tuple[int, ...], n_actions: int) -> None:\n",
        "        super().__init__()\n",
        "\n",
        "        # 2D conv. base network common to both heads\n",
        "        #  - This way both nets share commonly learned basic features\n",
        "        #  - Also helps with convergence (compared to having two separate NNs)\n",
        "        self.conv = nn.Sequential(\n",
        "            nn.Conv2d(input_shape[0], 32, kernel_size=8, stride=4),\n",
        "            nn.ReLU(),\n",
        "            nn.Conv2d(32, 64, kernel_size=4, stride=2),\n",
        "            nn.ReLU(),\n",
        "            nn.Conv2d(64, 64, kernel_size=3, stride=1),\n",
        "            nn.ReLU(),\n",
        "        )\n",
        "\n",
        "        conv_out_size = self._get_conv_out(input_shape)\n",
        "\n",
        "        # Policy NN - outputs action logits\n",
        "        self.policy = nn.Sequential(\n",
        "            nn.Linear(conv_out_size, 512),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(512, n_actions),\n",
        "        )\n",
        "\n",
        "        # Value NN - outputs state value\n",
        "        self.value = nn.Sequential(\n",
        "            nn.Linear(conv_out_size, 512),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(512, 1),\n",
        "        )\n",
        "\n",
        "    def _get_conv_out(self, shape: Tuple[int, ...]) -> int:\n",
        "        o = self.conv(torch.zeros(1, *shape))\n",
        "        return int(np.prod(o.size()))\n",
        "\n",
        "    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:\n",
        "        \"\"\"\n",
        "        For an input batch of states, returns pair of tensors\n",
        "          1. Policy logits for all actions in these states\n",
        "          2. Values of these states\n",
        "        \"\"\"\n",
        "        inputs = x.float() / 256\n",
        "        conv_out = self.conv(inputs).view(inputs.size()[0], -1)\n",
        "        return self.policy(conv_out), self.value(conv_out)"
      ],
      "id": "killing-lloyd",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EkrZlVejr5Od"
      },
      "source": [
        "### Batch Unpacking\n",
        "Similarly to implementations of other methods, we will use a function that unpacks an experience batch into its components.\n",
        "\n",
        "One important difference here is that we'll use the A2C NN to evaluate final states from the batch to get $V(s_N)$ - here we assume that we make $N$ steps ahead in the environmet/environments.\n",
        "\n",
        "Using all the experienced rewards $r_i$ from the batch we can compute target Q values as\n",
        "$$\n",
        "Q(s, a) = \\sum_{i = 0}^{N - 1} \\gamma^i r_i + \\gamma^T V(s_N)\n",
        "$$\n",
        "* First part (the sum) is the total discounted reward from all but last steps\n",
        "* The second is the discounted future reward from the $N$-th step (predicted by the NN)"
      ],
      "id": "EkrZlVejr5Od"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "parliamentary-posting"
      },
      "source": [
        "def unpack_batch(\n",
        "    batch: Sequence[ExperienceFirstLast],\n",
        "    net: AtariA2C,\n",
        "    gamma: float,\n",
        "    reward_steps: int,\n",
        "    device: str = \"cpu\",\n",
        ") -> Tuple[torch.FloatTensor, torch.LongTensor, torch.FloatTensor]:\n",
        "    \"\"\"\n",
        "    Convert batch into training tensors\n",
        "\n",
        "    :param batch: Experiences from environment(s)\n",
        "    :param net: A2C network that can approximate state values\n",
        "    :returns: states, actions, target Q values as tensors\n",
        "    \"\"\"\n",
        "    states, actions, rewards, not_done_exps, last_states = [], [], [], [], []\n",
        "\n",
        "    # Unwrap each transition from the batch\n",
        "    #  - And mark entries from unfinished episodes\n",
        "    for i, exp in enumerate(batch):\n",
        "\n",
        "        states.append(np.array(exp.state, copy=False))\n",
        "        actions.append(int(exp.action))\n",
        "        rewards.append(exp.reward)\n",
        "\n",
        "        if exp.last_state is not None:\n",
        "            not_done_exps.append(i)\n",
        "            last_states.append(np.array(exp.last_state, copy=False))\n",
        "\n",
        "    # Note: Wrapping states into np array is to fix PyTorch performance issue\n",
        "\n",
        "    # Convert states and actions to tensors\n",
        "    states = torch.FloatTensor(np.array(states, copy=False)).to(device)\n",
        "    actions = torch.LongTensor(actions).to(device)\n",
        "\n",
        "    # Compute target state values V(s) for training the net\n",
        "    #  - Uses given A2C NN to predit the state values\n",
        "\n",
        "    # Init target Q(s, a) to (discounted) rewards\n",
        "    #  - This will be the final value at the end of an episode\n",
        "    target_values = np.array(rewards, dtype=np.float32)\n",
        "\n",
        "    if not_done_exps:\n",
        "\n",
        "        # Convert next states to a tensor\n",
        "        last_states = torch.FloatTensor(np.array(last_states, copy=False)).to(\n",
        "            device\n",
        "        )\n",
        "\n",
        "        # Use given A2C net to predict future values\n",
        "        _, last_state_values = net(last_states)\n",
        "        last_state_values = last_state_values.data.cpu().numpy()[:, 0]\n",
        "\n",
        "        # Add future values to the discounted rewards if episode is not done\n",
        "        last_state_values *= gamma ** reward_steps\n",
        "        target_values[not_done_exps] += last_state_values\n",
        "\n",
        "    # Convert target values to a tensor\n",
        "    target_values = torch.FloatTensor(target_values).to(device)\n",
        "\n",
        "    return states, actions, target_values"
      ],
      "id": "parliamentary-posting",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bXh-fSpB2WQR"
      },
      "source": [
        "### A2C Training *(Atari Pong)*\n",
        "The training loop build on and extends the PG example from previous chapter.\n",
        "\n",
        "The loss function now has three components:\n",
        "1. **Policy loss** - similar to the PG method but with the gradient scales $A(s, a) = Q(s, a) - V(s)$ where $Q(s, a)$ are obtained from the batch unpacking described above (i.e. from experienced rewards and net's $V(s')$) and $V(s)$ is net's prediction for current state(s)\n",
        "1. **Value loss** - simply a MSE between current values $V(s)$ and TD targets (the same values we used for policy loss from the experience batch)\n",
        "1. **Entropy bonus** - the same techinque used for PG that adds an entropy component $\\mathcal{L}_H = \\beta \\sum_i \\pi(s_i) \\log(\\pi(s_i))$ that pushes the policy more towards uniform distribution that favours exploration\n",
        "\n",
        "Finally, we'll use multiple copies of the same environment to sample our experience batch from. This is the same techinque to break correlations that was used for the vanilla PG method. Thre are two basic variants of the *Actor-Critic* method:\n",
        "* A2C which uses parallel environments with synchronized policy gradient updates\n",
        "* A3C which stands for *Asynchronous Advantage Actor-Critic* and will be described in the next chapter"
      ],
      "id": "bXh-fSpB2WQR"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nb78PWxxGjcO"
      },
      "source": [
        "%load_ext tensorboard\n",
        "%tensorboard --logdir runs"
      ],
      "id": "nb78PWxxGjcO",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HcCcZdSOyvMl"
      },
      "source": [
        "# Hyperparameters\n",
        "GAMMA = 0.99\n",
        "LEARNING_RATE = 0.001\n",
        "ADAM_EPS = 1e-3\n",
        "ENTROPY_BETA = 0.01\n",
        "BATCH_SIZE = 128\n",
        "NUM_ENVS = 50\n",
        "REWARD_STEPS = 4\n",
        "CLIP_GRAD = 0.1\n",
        "STOP_REWARD = 18\n",
        "SEED = 42\n",
        "\n",
        "# Set RNG state\n",
        "np.random.seed(SEED)\n",
        "torch.manual_seed(SEED)\n",
        "if torch.cuda.is_available():\n",
        "    torch.cuda.manual_seed(SEED)\n",
        "\n",
        "# Determine where the computations will take place\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "\n",
        "# Make multiple instances of the Atari Pong environment\n",
        "\n",
        "\n",
        "def make_pong_env(i: int, seed: int) -> gym.Env:\n",
        "    env = ptan.common.wrappers.wrap_dqn(gym.make(\"PongNoFrameskip-v4\"))\n",
        "    env.seed(seed + i)\n",
        "    return env\n",
        "\n",
        "\n",
        "envs = [make_pong_env(i, seed=SEED) for i in range(NUM_ENVS)]\n",
        "\n",
        "# Create the A2C network for Atari environments\n",
        "net = AtariA2C(\n",
        "    input_shape=envs[0].observation_space.shape,\n",
        "    n_actions=envs[0].action_space.n,\n",
        ").to(device)\n",
        "print(net)\n",
        "\n",
        "# Initialize the policy agent\n",
        "# - Instead of passing the whole NN, we use a callback over it returning just\n",
        "#   the action logits\n",
        "agent = ptan.agent.PolicyAgent(\n",
        "    model=lambda x: net(x)[0],\n",
        "    apply_softmax=True,\n",
        "    device=device,\n",
        ")\n",
        "\n",
        "# Create `REWARD_STEPS`-ahead experience source over all environments\n",
        "exp_source = ptan.experience.ExperienceSourceFirstLast(\n",
        "    env=envs,\n",
        "    agent=agent,\n",
        "    gamma=GAMMA,\n",
        "    steps_count=REWARD_STEPS,\n",
        ")\n",
        "\n",
        "# Create Adam optimizer\n",
        "#  - Note: We use larger epsilon to make the training converge\n",
        "optimizer = torch.optim.Adam(net.parameters(), lr=LEARNING_RATE, eps=ADAM_EPS)\n",
        "\n",
        "# Create TensorBoard writer for metrics collection\n",
        "writer = SummaryWriter(comment=\"-pong-a2c\")\n",
        "\n",
        "# Create TensorBoard trackers\n",
        "with RewardTracker(writer, stop_reward=STOP_REWARD) as tracker:\n",
        "    with ptan.common.utils.TBMeanTracker(writer, batch_size=10) as tb_tracker:\n",
        "\n",
        "        batch = []\n",
        "\n",
        "        # Run the training loop consuming experiences from the source\n",
        "        for i, exp in enumerate(exp_source):\n",
        "\n",
        "            # Add new experience to current batch\n",
        "            batch.append(exp)\n",
        "\n",
        "            new_rewards = exp_source.pop_total_rewards()\n",
        "            if new_rewards:\n",
        "\n",
        "                # Record new reward and check for termination\n",
        "                solved = tracker.add_reward(reward=new_rewards[0], frame=i)\n",
        "\n",
        "                # Stop if the mean reward was good enough\n",
        "                if solved:\n",
        "                    print(f\"Solved in {i} steps!\")\n",
        "                    break\n",
        "\n",
        "            # Let the batch fill up\n",
        "            if len(batch) < BATCH_SIZE:\n",
        "                continue\n",
        "\n",
        "            # Unpack and clear current batch\n",
        "            states, actions, target_values = unpack_batch(\n",
        "                batch=batch,\n",
        "                net=net,\n",
        "                gamma=GAMMA,\n",
        "                reward_steps=REWARD_STEPS,\n",
        "                device=device,\n",
        "            )\n",
        "            batch.clear()\n",
        "\n",
        "            # Clear gradients\n",
        "            optimizer.zero_grad()\n",
        "\n",
        "            # Compute action logits and values of current states\n",
        "            action_logits, values = net(states)\n",
        "\n",
        "            # Compute V(s) part of the loss function\n",
        "            value_loss = nn.functional.mse_loss(\n",
        "                input=values.squeeze(-1),\n",
        "                target=target_values,\n",
        "            )\n",
        "\n",
        "            # Compute the policy part of the loss function\n",
        "            #  - We use A(s, a) = Q(s, a) - V(s) as the gradient scales\n",
        "            #  - Note: We detach values from the autograph to stop grad flow\n",
        "            log_action_prob = nn.functional.log_softmax(action_logits, dim=1)\n",
        "            advantage = target_values - values.detach()\n",
        "            scaled_log_action_prob = (\n",
        "                advantage * log_action_prob[range(BATCH_SIZE), actions]\n",
        "            )\n",
        "            policy_loss = -scaled_log_action_prob.mean()\n",
        "\n",
        "            # Compute entropy bonus to the loss function\n",
        "            action_prob = nn.functional.softmax(action_logits, dim=1)\n",
        "            entropy_loss = (\n",
        "                ENTROPY_BETA\n",
        "                * (action_prob * log_action_prob).sum(dim=1).mean()\n",
        "            )\n",
        "\n",
        "            # First calculate policy gradients only\n",
        "            policy_loss.backward(retain_graph=True)\n",
        "            grads = np.concatenate(\n",
        "                [\n",
        "                    param.grad.data.cpu().numpy().flatten()\n",
        "                    for param in net.parameters()\n",
        "                    if param.grad is not None\n",
        "                ]\n",
        "            )\n",
        "\n",
        "            # Apply entropy and value gradients\n",
        "            loss = entropy_loss + value_loss\n",
        "            loss.backward()\n",
        "\n",
        "            # Use gradient clipping (by l2 norm) before making next step\n",
        "            nn.utils.clip_grad_norm_(net.parameters(), CLIP_GRAD)\n",
        "            optimizer.step()\n",
        "\n",
        "            # Get total loss for tracking\n",
        "            loss += policy_loss\n",
        "\n",
        "            # Track metrics\n",
        "            tb_tracker.track(\"advantage\", advantage, i)\n",
        "            tb_tracker.track(\"values\", values, i)\n",
        "            tb_tracker.track(\"batch_rewards\", target_values, i)\n",
        "            tb_tracker.track(\"loss_entropy\", entropy_loss, i)\n",
        "            tb_tracker.track(\"loss_policy\", policy_loss, i)\n",
        "            tb_tracker.track(\"loss_value\", value_loss, i)\n",
        "            tb_tracker.track(\"loss_total\", loss, i)\n",
        "            tb_tracker.track(\"grad_l2\", np.sqrt(np.mean(np.square(grads))), i)\n",
        "            tb_tracker.track(\"grad_max\", np.max(np.abs(grads)), i)\n",
        "            tb_tracker.track(\"grad_var\", np.var(grads), i)"
      ],
      "id": "HcCcZdSOyvMl",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MU2Lpp0aHwc1"
      },
      "source": [
        ""
      ],
      "id": "MU2Lpp0aHwc1",
      "execution_count": null,
      "outputs": []
    }
  ]
}