chore: import upstream snapshot with attribution

labml_nn/rl/ppo/__init__.py

@@ -0,0 +1,206 @@
+"""
+---
+title: Proximal Policy Optimization - PPO
+summary: >
+ An annotated implementation of Proximal Policy Optimization - PPO algorithm in PyTorch.
+---
+
+# Proximal Policy Optimization - PPO
+
+This is a [PyTorch](https://pytorch.org) implementation of
+[Proximal Policy Optimization - PPO](https://arxiv.org/abs/1707.06347).
+
+PPO is a policy gradient method for reinforcement learning.
+Simple policy gradient methods do a single gradient update per sample (or a set of samples).
+Doing multiple gradient steps for a single sample causes problems
+because the policy deviates too much, producing a bad policy.
+PPO lets us do multiple gradient updates per sample by trying to keep the
+policy close to the policy that was used to sample data.
+It does so by clipping gradient flow if the updated policy
+is not close to the policy used to sample the data.
+
+You can find an experiment that uses it [here](experiment.html).
+The experiment uses [Generalized Advantage Estimation](gae.html).
+
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/rl/ppo/experiment.ipynb)
+"""
+
+import torch
+from labml_nn.rl.ppo.gae import GAE
+from torch import nn
+
+
+class ClippedPPOLoss(nn.Module):
+    """
+    ## PPO Loss
+
+    Here's how the PPO update rule is derived.
+
+    We want to maximize policy reward
+     $$\max_\theta J(\pi_\theta) =
+       \mathop{\mathbb{E}}_{\tau \sim \pi_\theta}\Biggl[\sum_{t=0}^\infty \gamma^t r_t \Biggr]$$
+     where $r$ is the reward, $\pi$ is the policy, $\tau$ is a trajectory sampled from policy,
+     and $\gamma$ is the discount factor between $[0, 1]$.
+
+    \begin{align}
+    \mathbb{E}_{\tau \sim \pi_\theta} \Biggl[
+     \sum_{t=0}^\infty \gamma^t A^{\pi_{OLD}}(s_t, a_t)
+    \Biggr] &=
+    \\
+    \mathbb{E}_{\tau \sim \pi_\theta} \Biggl[
+      \sum_{t=0}^\infty \gamma^t \Bigl(
+       Q^{\pi_{OLD}}(s_t, a_t) - V^{\pi_{OLD}}(s_t)
+      \Bigr)
+     \Biggr] &=
+    \\
+    \mathbb{E}_{\tau \sim \pi_\theta} \Biggl[
+      \sum_{t=0}^\infty \gamma^t \Bigl(
+       r_t + V^{\pi_{OLD}}(s_{t+1}) - V^{\pi_{OLD}}(s_t)
+      \Bigr)
+     \Biggr] &=
+    \\
+    \mathbb{E}_{\tau \sim \pi_\theta} \Biggl[
+      \sum_{t=0}^\infty \gamma^t \Bigl(
+       r_t
+      \Bigr)
+     \Biggr]
+     - \mathbb{E}_{\tau \sim \pi_\theta}
+        \Biggl[V^{\pi_{OLD}}(s_0)\Biggr] &=
+    J(\pi_\theta) - J(\pi_{\theta_{OLD}})
+    \end{align}
+
+    So,
+     $$\max_\theta J(\pi_\theta) =
+       \max_\theta \mathbb{E}_{\tau \sim \pi_\theta} \Biggl[
+          \sum_{t=0}^\infty \gamma^t A^{\pi_{OLD}}(s_t, a_t)
+       \Biggr]$$
+
+    Define discounted-future state distribution,
+     $$d^\pi(s) = (1 - \gamma) \sum_{t=0}^\infty \gamma^t P(s_t = s | \pi)$$
+
+    Then,
+
+    \begin{align}
+    J(\pi_\theta) - J(\pi_{\theta_{OLD}})
+    &= \mathbb{E}_{\tau \sim \pi_\theta} \Biggl[
+     \sum_{t=0}^\infty \gamma^t A^{\pi_{OLD}}(s_t, a_t)
+    \Biggr]
+    \\
+    &= \frac{1}{1 - \gamma}
+     \mathbb{E}_{s \sim d^{\pi_\theta}, a \sim \pi_\theta} \Bigl[
+      A^{\pi_{OLD}}(s, a)
+     \Bigr]
+    \end{align}
+
+    Importance sampling $a$ from $\pi_{\theta_{OLD}}$,
+
+    \begin{align}
+    J(\pi_\theta) - J(\pi_{\theta_{OLD}})
+    &= \frac{1}{1 - \gamma}
+     \mathbb{E}_{s \sim d^{\pi_\theta}, a \sim \pi_\theta} \Bigl[
+      A^{\pi_{OLD}}(s, a)
+     \Bigr]
+    \\
+    &= \frac{1}{1 - \gamma}
+     \mathbb{E}_{s \sim d^{\pi_\theta}, a \sim \pi_{\theta_{OLD}}} \Biggl[
+      \frac{\pi_\theta(a|s)}{\pi_{\theta_{OLD}}(a|s)} A^{\pi_{OLD}}(s, a)
+     \Biggr]
+    \end{align}
+
+    Then we assume $d^\pi_\theta(s)$ and  $d^\pi_{\theta_{OLD}}(s)$ are similar.
+    The error we introduce to $J(\pi_\theta) - J(\pi_{\theta_{OLD}})$
+     by this assumption is bound by the KL divergence between
+     $\pi_\theta$ and $\pi_{\theta_{OLD}}$.
+    [Constrained Policy Optimization](https://arxiv.org/abs/1705.10528)
+     shows the proof of this. I haven't read it.
+
+
+    \begin{align}
+    J(\pi_\theta) - J(\pi_{\theta_{OLD}})
+    &= \frac{1}{1 - \gamma}
+     \mathop{\mathbb{E}}_{s \sim d^{\pi_\theta} \atop a \sim \pi_{\theta_{OLD}}} \Biggl[
+      \frac{\pi_\theta(a|s)}{\pi_{\theta_{OLD}}(a|s)} A^{\pi_{OLD}}(s, a)
+     \Biggr]
+    \\
+    &\approx \frac{1}{1 - \gamma}
+     \mathop{\mathbb{E}}_{\textcolor{orange}{s \sim d^{\pi_{\theta_{OLD}}}}
+     \atop a \sim \pi_{\theta_{OLD}}} \Biggl[
+      \frac{\pi_\theta(a|s)}{\pi_{\theta_{OLD}}(a|s)} A^{\pi_{OLD}}(s, a)
+     \Biggr]
+    \\
+    &= \frac{1}{1 - \gamma} \mathcal{L}^{CPI}
+    \end{align}
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, log_pi: torch.Tensor, sampled_log_pi: torch.Tensor,
+                advantage: torch.Tensor, clip: float) -> torch.Tensor:
+        # ratio $r_t(\theta) = \frac{\pi_\theta (a_t|s_t)}{\pi_{\theta_{OLD}} (a_t|s_t)}$;
+        # *this is different from rewards* $r_t$.
+        ratio = torch.exp(log_pi - sampled_log_pi)
+
+        # ### Cliping the policy ratio
+        #
+        # \begin{align}
+        # \mathcal{L}^{CLIP}(\theta) =
+        #  \mathbb{E}_{a_t, s_t \sim \pi_{\theta{OLD}}} \biggl[
+        #    min \Bigl(r_t(\theta) \bar{A_t},
+        #              clip \bigl(
+        #               r_t(\theta), 1 - \epsilon, 1 + \epsilon
+        #              \bigr) \bar{A_t}
+        #    \Bigr)
+        #  \biggr]
+        # \end{align}
+        #
+        # The ratio is clipped to be close to 1.
+        # We take the minimum so that the gradient will only pull
+        # $\pi_\theta$ towards $\pi_{\theta_{OLD}}$ if the ratio is
+        # not between $1 - \epsilon$ and $1 + \epsilon$.
+        # This keeps the KL divergence between $\pi_\theta$
+        #  and $\pi_{\theta_{OLD}}$ constrained.
+        # Large deviation can cause performance collapse;
+        #  where the policy performance drops and doesn't recover because
+        #  we are sampling from a bad policy.
+        #
+        # Using the normalized advantage
+        #  $\bar{A_t} = \frac{\hat{A_t} - \mu(\hat{A_t})}{\sigma(\hat{A_t})}$
+        #  introduces a bias to the policy gradient estimator,
+        #  but it reduces variance a lot.
+        clipped_ratio = ratio.clamp(min=1.0 - clip,
+                                    max=1.0 + clip)
+        policy_reward = torch.min(ratio * advantage,
+                                  clipped_ratio * advantage)
+
+        self.clip_fraction = (abs((ratio - 1.0)) > clip).to(torch.float).mean()
+
+        return -policy_reward.mean()
+
+
+class ClippedValueFunctionLoss(nn.Module):
+    """
+    ## Clipped Value Function Loss
+
+    Similarly we clip the value function update also.
+
+    \begin{align}
+    V^{\pi_\theta}_{CLIP}(s_t)
+     &= clip\Bigl(V^{\pi_\theta}(s_t) - \hat{V_t}, -\epsilon, +\epsilon\Bigr)
+    \\
+    \mathcal{L}^{VF}(\theta)
+     &= \frac{1}{2} \mathbb{E} \biggl[
+      max\Bigl(\bigl(V^{\pi_\theta}(s_t) - R_t\bigr)^2,
+          \bigl(V^{\pi_\theta}_{CLIP}(s_t) - R_t\bigr)^2\Bigr)
+     \biggr]
+    \end{align}
+
+    Clipping makes sure the value function $V_\theta$ doesn't deviate
+     significantly from $V_{\theta_{OLD}}$.
+
+    """
+
+    def forward(self, value: torch.Tensor, sampled_value: torch.Tensor, sampled_return: torch.Tensor, clip: float):
+        clipped_value = sampled_value + (value - sampled_value).clamp(min=-clip, max=clip)
+        vf_loss = torch.max((value - sampled_return) ** 2, (clipped_value - sampled_return) ** 2)
+        return 0.5 * vf_loss.mean()

labml_nn/rl/ppo/experiment.ipynb

@@ -0,0 +1,238 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "AYV_dMVDxyc2"
+   },
+   "source": [
+    "[![Github](https://img.shields.io/github/stars/labmlai/annotated_deep_learning_paper_implementations?style=social)](https://github.com/labmlai/annotated_deep_learning_paper_implementations)\n",
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/rl/ppo/experiment.ipynb)                    \n",
+    "\n",
+    "## Proximal Policy Optimization - PPO\n",
+    "\n",
+    "This is an experiment training an agent to play Atari Breakout game using  Proximal Policy Optimization - PPO"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "AahG_i2y5tY9"
+   },
+   "source": [
+    "Install the `labml-nn` package"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/"
+    },
+    "id": "ZCzmCrAIVg0L",
+    "outputId": "028e759e-0c9f-472e-b4b8-fdcf3e4604ee"
+   },
+   "source": [
+    "!pip install labml-nn"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Add Atari ROMs (Doesn't work without this in Google Colab)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "! wget http://www.atarimania.com/roms/Roms.rar\n",
+    "! mkdir /content/ROM/\n",
+    "! unrar e /content/Roms.rar /content/ROM/\n",
+    "! python -m atari_py.import_roms /content/ROM/"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "SE2VUQ6L5zxI"
+   },
+   "source": [
+    "Imports"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {
+    "id": "0hJXx_g0wS2C"
+   },
+   "source": [
+    "from labml import experiment\n",
+    "from labml.configs import FloatDynamicHyperParam, IntDynamicHyperParam\n",
+    "from labml_nn.rl.ppo.experiment import Trainer"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "Lpggo0wM6qb-"
+   },
+   "source": [
+    "Create an experiment"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {
+    "id": "bFcr9k-l4cAg"
+   },
+   "source": [
+    "experiment.create(name=\"ppo\")"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "-OnHLi626tJt"
+   },
+   "source": [
+    "### Configurations\n",
+    "\n",
+    "`IntDynamicHyperParam` and `FloatDynamicHyperParam` are dynamic hyper parameters\n",
+    "that you can change while the experiment is running."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {
+    "id": "Piz0c5f44hRo"
+   },
+   "source": [
+    "configs = {\n",
+    "    # number of updates\n",
+    "    'updates': 10000,\n",
+    "    # number of epochs to train the model with sampled data\n",
+    "    'epochs': IntDynamicHyperParam(8),\n",
+    "    # number of worker processes\n",
+    "    'n_workers': 8,\n",
+    "    # number of steps to run on each process for a single update\n",
+    "    'worker_steps': 128,\n",
+    "    # number of mini batches\n",
+    "    'batches': 4,\n",
+    "    # Value loss coefficient\n",
+    "    'value_loss_coef': FloatDynamicHyperParam(0.5),\n",
+    "    # Entropy bonus coefficient\n",
+    "    'entropy_bonus_coef': FloatDynamicHyperParam(0.01),\n",
+    "    # Clip range\n",
+    "    'clip_range': FloatDynamicHyperParam(0.1),\n",
+    "    # Learning rate\n",
+    "    'learning_rate': FloatDynamicHyperParam(2.5e-4, (0, 1e-3)),\n",
+    "}"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "wwMzCqpD6vkL"
+   },
+   "source": [
+    "Set experiment configurations"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/",
+     "height": 17
+    },
+    "id": "e6hmQhTw4nks",
+    "outputId": "0e978879-5dcd-4140-ec53-24a3fbd547de"
+   },
+   "source": [
+    "experiment.configs(configs)"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "qYQCFt_JYsjd"
+   },
+   "source": [
+    "Create trainer"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {
+    "id": "8LB7XVViYuPG"
+   },
+   "source": [
+    "trainer = Trainer(**configs)"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "KJZRf8527GxL"
+   },
+   "source": [
+    "Start the experiment and run the training loop."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {
+    "id": "aIAWo7Fw5DR8"
+   },
+   "source": [
+    "with experiment.start():\n",
+    "    trainer.run_training_loop()"
+   ],
+   "outputs": [],
+   "execution_count": null
+  }
+ ],
+ "metadata": {
+  "accelerator": "GPU",
+  "colab": {
+   "collapsed_sections": [],
+   "name": "Proximal Policy Optimization - PPO",
+   "provenance": []
+  },
+  "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.7.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

labml_nn/rl/ppo/experiment.py

@@ -0,0 +1,396 @@
+"""
+---
+title: PPO Experiment with Atari Breakout
+summary: Annotated implementation to train a PPO agent on Atari Breakout game.
+---
+
+# PPO Experiment with Atari Breakout
+
+This experiment trains Proximal Policy Optimization (PPO) agent  Atari Breakout game on OpenAI Gym.
+It runs the [game environments on multiple processes](../game.html) to sample efficiently.
+
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/rl/ppo/experiment.ipynb)
+"""
+
+from typing import Dict
+
+import numpy as np
+import torch
+from torch import nn
+from torch import optim
+from torch.distributions import Categorical
+
+from labml import monit, tracker, logger, experiment
+from labml.configs import FloatDynamicHyperParam, IntDynamicHyperParam
+from labml_nn.rl.game import Worker
+from labml_nn.rl.ppo import ClippedPPOLoss, ClippedValueFunctionLoss
+from labml_nn.rl.ppo.gae import GAE
+
+# Select device
+if torch.cuda.is_available():
+    device = torch.device("cuda:0")
+else:
+    device = torch.device("cpu")
+
+
+class Model(nn.Module):
+    """
+    ## Model
+    """
+
+    def __init__(self):
+        super().__init__()
+
+        # The first convolution layer takes a
+        # 84x84 frame and produces a 20x20 frame
+        self.conv1 = nn.Conv2d(in_channels=4, out_channels=32, kernel_size=8, stride=4)
+
+        # The second convolution layer takes a
+        # 20x20 frame and produces a 9x9 frame
+        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=4, stride=2)
+
+        # The third convolution layer takes a
+        # 9x9 frame and produces a 7x7 frame
+        self.conv3 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1)
+
+        # A fully connected layer takes the flattened
+        # frame from third convolution layer, and outputs
+        # 512 features
+        self.lin = nn.Linear(in_features=7 * 7 * 64, out_features=512)
+
+        # A fully connected layer to get logits for $\pi$
+        self.pi_logits = nn.Linear(in_features=512, out_features=4)
+
+        # A fully connected layer to get value function
+        self.value = nn.Linear(in_features=512, out_features=1)
+
+        #
+        self.activation = nn.ReLU()
+
+    def forward(self, obs: torch.Tensor):
+        h = self.activation(self.conv1(obs))
+        h = self.activation(self.conv2(h))
+        h = self.activation(self.conv3(h))
+        h = h.reshape((-1, 7 * 7 * 64))
+
+        h = self.activation(self.lin(h))
+
+        pi = Categorical(logits=self.pi_logits(h))
+        value = self.value(h).reshape(-1)
+
+        return pi, value
+
+
+def obs_to_torch(obs: np.ndarray) -> torch.Tensor:
+    """Scale observations from `[0, 255]` to `[0, 1]`"""
+    return torch.tensor(obs, dtype=torch.float32, device=device) / 255.
+
+
+class Trainer:
+    """
+    ## Trainer
+    """
+
+    def __init__(self, *,
+                 updates: int, epochs: IntDynamicHyperParam,
+                 n_workers: int, worker_steps: int, batches: int,
+                 value_loss_coef: FloatDynamicHyperParam,
+                 entropy_bonus_coef: FloatDynamicHyperParam,
+                 clip_range: FloatDynamicHyperParam,
+                 learning_rate: FloatDynamicHyperParam,
+                 ):
+        # #### Configurations
+
+        # number of updates
+        self.updates = updates
+        # number of epochs to train the model with sampled data
+        self.epochs = epochs
+        # number of worker processes
+        self.n_workers = n_workers
+        # number of steps to run on each process for a single update
+        self.worker_steps = worker_steps
+        # number of mini batches
+        self.batches = batches
+        # total number of samples for a single update
+        self.batch_size = self.n_workers * self.worker_steps
+        # size of a mini batch
+        self.mini_batch_size = self.batch_size // self.batches
+        assert (self.batch_size % self.batches == 0)
+
+        # Value loss coefficient
+        self.value_loss_coef = value_loss_coef
+        # Entropy bonus coefficient
+        self.entropy_bonus_coef = entropy_bonus_coef
+
+        # Clipping range
+        self.clip_range = clip_range
+        # Learning rate
+        self.learning_rate = learning_rate
+
+        # #### Initialize
+
+        # create workers
+        self.workers = [Worker(47 + i) for i in range(self.n_workers)]
+
+        # initialize tensors for observations
+        self.obs = np.zeros((self.n_workers, 4, 84, 84), dtype=np.uint8)
+        for worker in self.workers:
+            worker.child.send(("reset", None))
+        for i, worker in enumerate(self.workers):
+            self.obs[i] = worker.child.recv()
+
+        # model
+        self.model = Model().to(device)
+
+        # optimizer
+        self.optimizer = optim.Adam(self.model.parameters(), lr=2.5e-4)
+
+        # GAE with $\gamma = 0.99$ and $\lambda = 0.95$
+        self.gae = GAE(self.n_workers, self.worker_steps, 0.99, 0.95)
+
+        # PPO Loss
+        self.ppo_loss = ClippedPPOLoss()
+
+        # Value Loss
+        self.value_loss = ClippedValueFunctionLoss()
+
+    def sample(self) -> Dict[str, torch.Tensor]:
+        """
+        ### Sample data with current policy
+        """
+
+        rewards = np.zeros((self.n_workers, self.worker_steps), dtype=np.float32)
+        actions = np.zeros((self.n_workers, self.worker_steps), dtype=np.int32)
+        done = np.zeros((self.n_workers, self.worker_steps), dtype=np.bool)
+        obs = np.zeros((self.n_workers, self.worker_steps, 4, 84, 84), dtype=np.uint8)
+        log_pis = np.zeros((self.n_workers, self.worker_steps), dtype=np.float32)
+        values = np.zeros((self.n_workers, self.worker_steps + 1), dtype=np.float32)
+
+        with torch.no_grad():
+            # sample `worker_steps` from each worker
+            for t in range(self.worker_steps):
+                # `self.obs` keeps track of the last observation from each worker,
+                #  which is the input for the model to sample the next action
+                obs[:, t] = self.obs
+                # sample actions from $\pi_{\theta_{OLD}}$ for each worker;
+                #  this returns arrays of size `n_workers`
+                pi, v = self.model(obs_to_torch(self.obs))
+                values[:, t] = v.cpu().numpy()
+                a = pi.sample()
+                actions[:, t] = a.cpu().numpy()
+                log_pis[:, t] = pi.log_prob(a).cpu().numpy()
+
+                # run sampled actions on each worker
+                for w, worker in enumerate(self.workers):
+                    worker.child.send(("step", actions[w, t]))
+
+                for w, worker in enumerate(self.workers):
+                    # get results after executing the actions
+                    self.obs[w], rewards[w, t], done[w, t], info = worker.child.recv()
+
+                    # collect episode info, which is available if an episode finished;
+                    #  this includes total reward and length of the episode -
+                    #  look at `Game` to see how it works.
+                    if info:
+                        tracker.add('reward', info['reward'])
+                        tracker.add('length', info['length'])
+
+            # Get value of after the final step
+            _, v = self.model(obs_to_torch(self.obs))
+            values[:, self.worker_steps] = v.cpu().numpy()
+
+        # calculate advantages
+        advantages = self.gae(done, rewards, values)
+
+        #
+        samples = {
+            'obs': obs,
+            'actions': actions,
+            'values': values[:, :-1],
+            'log_pis': log_pis,
+            'advantages': advantages
+        }
+
+        # samples are currently in `[workers, time_step]` table,
+        # we should flatten it for training
+        samples_flat = {}
+        for k, v in samples.items():
+            v = v.reshape(v.shape[0] * v.shape[1], *v.shape[2:])
+            if k == 'obs':
+                samples_flat[k] = obs_to_torch(v)
+            else:
+                samples_flat[k] = torch.tensor(v, device=device)
+
+        return samples_flat
+
+    def train(self, samples: Dict[str, torch.Tensor]):
+        """
+        ### Train the model based on samples
+        """
+
+        # It learns faster with a higher number of epochs,
+        #  but becomes a little unstable; that is,
+        #  the average episode reward does not monotonically increase
+        #  over time.
+        # May be reducing the clipping range might solve it.
+        for _ in range(self.epochs()):
+            # shuffle for each epoch
+            indexes = torch.randperm(self.batch_size)
+
+            # for each mini batch
+            for start in range(0, self.batch_size, self.mini_batch_size):
+                # get mini batch
+                end = start + self.mini_batch_size
+                mini_batch_indexes = indexes[start: end]
+                mini_batch = {}
+                for k, v in samples.items():
+                    mini_batch[k] = v[mini_batch_indexes]
+
+                # train
+                loss = self._calc_loss(mini_batch)
+
+                # Set learning rate
+                for pg in self.optimizer.param_groups:
+                    pg['lr'] = self.learning_rate()
+                # Zero out the previously calculated gradients
+                self.optimizer.zero_grad()
+                # Calculate gradients
+                loss.backward()
+                # Clip gradients
+                torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=0.5)
+                # Update parameters based on gradients
+                self.optimizer.step()
+
+    @staticmethod
+    def _normalize(adv: torch.Tensor):
+        """#### Normalize advantage function"""
+        return (adv - adv.mean()) / (adv.std() + 1e-8)
+
+    def _calc_loss(self, samples: Dict[str, torch.Tensor]) -> torch.Tensor:
+        """
+        ### Calculate total loss
+        """
+
+        # $R_t$ returns sampled from $\pi_{\theta_{OLD}}$
+        sampled_return = samples['values'] + samples['advantages']
+
+        # $\bar{A_t} = \frac{\hat{A_t} - \mu(\hat{A_t})}{\sigma(\hat{A_t})}$,
+        # where $\hat{A_t}$ is advantages sampled from $\pi_{\theta_{OLD}}$.
+        # Refer to sampling function in [Main class](#main) below
+        #  for the calculation of $\hat{A}_t$.
+        sampled_normalized_advantage = self._normalize(samples['advantages'])
+
+        # Sampled observations are fed into the model to get $\pi_\theta(a_t|s_t)$ and $V^{\pi_\theta}(s_t)$;
+        #  we are treating observations as state
+        pi, value = self.model(samples['obs'])
+
+        # $-\log \pi_\theta (a_t|s_t)$, $a_t$ are actions sampled from $\pi_{\theta_{OLD}}$
+        log_pi = pi.log_prob(samples['actions'])
+
+        # Calculate policy loss
+        policy_loss = self.ppo_loss(log_pi, samples['log_pis'], sampled_normalized_advantage, self.clip_range())
+
+        # Calculate Entropy Bonus
+        #
+        # $\mathcal{L}^{EB}(\theta) =
+        #  \mathbb{E}\Bigl[ S\bigl[\pi_\theta\bigr] (s_t) \Bigr]$
+        entropy_bonus = pi.entropy()
+        entropy_bonus = entropy_bonus.mean()
+
+        # Calculate value function loss
+        value_loss = self.value_loss(value, samples['values'], sampled_return, self.clip_range())
+
+        # $\mathcal{L}^{CLIP+VF+EB} (\theta) =
+        #  \mathcal{L}^{CLIP} (\theta) +
+        #  c_1 \mathcal{L}^{VF} (\theta) - c_2 \mathcal{L}^{EB}(\theta)$
+        loss = (policy_loss
+                + self.value_loss_coef() * value_loss
+                - self.entropy_bonus_coef() * entropy_bonus)
+
+        # for monitoring
+        approx_kl_divergence = .5 * ((samples['log_pis'] - log_pi) ** 2).mean()
+
+        # Add to tracker
+        tracker.add({'policy_reward': -policy_loss,
+                     'value_loss': value_loss,
+                     'entropy_bonus': entropy_bonus,
+                     'kl_div': approx_kl_divergence,
+                     'clip_fraction': self.ppo_loss.clip_fraction})
+
+        return loss
+
+    def run_training_loop(self):
+        """
+        ### Run training loop
+        """
+
+        # last 100 episode information
+        tracker.set_queue('reward', 100, True)
+        tracker.set_queue('length', 100, True)
+
+        for update in monit.loop(self.updates):
+            # sample with current policy
+            samples = self.sample()
+
+            # train the model
+            self.train(samples)
+
+            # Save tracked indicators.
+            tracker.save()
+            # Add a new line to the screen periodically
+            if (update + 1) % 1_000 == 0:
+                logger.log()
+
+    def destroy(self):
+        """
+        ### Destroy
+        Stop the workers
+        """
+        for worker in self.workers:
+            worker.child.send(("close", None))
+
+
+def main():
+    # Create the experiment
+    experiment.create(name='ppo')
+    # Configurations
+    configs = {
+        # Number of updates
+        'updates': 10000,
+        # ⚙️ Number of epochs to train the model with sampled data.
+        # You can change this while the experiment is running.
+        'epochs': IntDynamicHyperParam(8),
+        # Number of worker processes
+        'n_workers': 8,
+        # Number of steps to run on each process for a single update
+        'worker_steps': 128,
+        # Number of mini batches
+        'batches': 4,
+        # ⚙️ Value loss coefficient.
+        # You can change this while the experiment is running.
+        'value_loss_coef': FloatDynamicHyperParam(0.5),
+        # ⚙️ Entropy bonus coefficient.
+        # You can change this while the experiment is running.
+        'entropy_bonus_coef': FloatDynamicHyperParam(0.01),
+        # ⚙️ Clip range.
+        'clip_range': FloatDynamicHyperParam(0.1),
+        # You can change this while the experiment is running.
+        # ⚙️ Learning rate.
+        'learning_rate': FloatDynamicHyperParam(1e-3, (0, 1e-3)),
+    }
+
+    experiment.configs(configs)
+
+    # Initialize the trainer
+    m = Trainer(**configs)
+
+    # Run and monitor the experiment
+    with experiment.start():
+        m.run_training_loop()
+    # Stop the workers
+    m.destroy()
+
+
+# ## Run it
+if __name__ == "__main__":
+    main()

labml_nn/rl/ppo/gae.py

@@ -0,0 +1,82 @@
+"""
+---
+title: Generalized Advantage Estimation (GAE)
+summary: A PyTorch implementation/tutorial of Generalized Advantage Estimation (GAE).
+---
+
+# Generalized Advantage Estimation (GAE)
+
+This is a [PyTorch](https://pytorch.org) implementation of paper
+[Generalized Advantage Estimation](https://arxiv.org/abs/1506.02438).
+
+You can find an experiment that uses it [here](experiment.html).
+"""
+
+import numpy as np
+
+
+class GAE:
+    def __init__(self, n_workers: int, worker_steps: int, gamma: float, lambda_: float):
+        self.lambda_ = lambda_
+        self.gamma = gamma
+        self.worker_steps = worker_steps
+        self.n_workers = n_workers
+
+    def __call__(self, done: np.ndarray, rewards: np.ndarray, values: np.ndarray) -> np.ndarray:
+        """
+        ### Calculate advantages
+
+        \begin{align}
+        \hat{A_t^{(1)}} &= r_t + \gamma V(s_{t+1}) - V(s)
+        \\
+        \hat{A_t^{(2)}} &= r_t + \gamma r_{t+1} +\gamma^2 V(s_{t+2}) - V(s)
+        \\
+        ...
+        \\
+        \hat{A_t^{(\infty)}} &= r_t + \gamma r_{t+1} +\gamma^2 r_{t+2} + ... - V(s)
+        \end{align}
+
+        $\hat{A_t^{(1)}}$ is high bias, low variance, whilst
+        $\hat{A_t^{(\infty)}}$ is unbiased, high variance.
+
+        We take a weighted average of $\hat{A_t^{(k)}}$ to balance bias and variance.
+        This is called Generalized Advantage Estimation.
+        $$\hat{A_t} = \hat{A_t^{GAE}} = \frac{\sum_k w_k \hat{A_t^{(k)}}}{\sum_k w_k}$$
+        We set $w_k = \lambda^{k-1}$, this gives clean calculation for
+        $\hat{A_t}$
+
+        \begin{align}
+        \delta_t &= r_t + \gamma V(s_{t+1}) - V(s_t)
+        \\
+        \hat{A_t} &= \delta_t + \gamma \lambda \delta_{t+1} + ... +
+                             (\gamma \lambda)^{T - t + 1} \delta_{T - 1}
+        \\
+        &= \delta_t + \gamma \lambda \hat{A_{t+1}}
+        \end{align}
+        """
+
+        # advantages table
+        advantages = np.zeros((self.n_workers, self.worker_steps), dtype=np.float32)
+        last_advantage = 0
+
+        # $V(s_{t+1})$
+        last_value = values[:, -1]
+
+        for t in reversed(range(self.worker_steps)):
+            # mask if episode completed after step $t$
+            mask = 1.0 - done[:, t]
+            last_value = last_value * mask
+            last_advantage = last_advantage * mask
+            # $\delta_t$
+            delta = rewards[:, t] + self.gamma * last_value - values[:, t]
+
+            # $\hat{A_t} = \delta_t + \gamma \lambda \hat{A_{t+1}}$
+            last_advantage = delta + self.gamma * self.lambda_ * last_advantage
+
+            #
+            advantages[:, t] = last_advantage
+
+            last_value = values[:, t]
+
+        # $\hat{A_t}$
+        return advantages

labml_nn/rl/ppo/readme.md

@@ -0,0 +1,18 @@
+# [Proximal Policy Optimization - PPO](https://nn.labml.ai/rl/ppo/index.html)
+
+This is a [PyTorch](https://pytorch.org) implementation of
+[Proximal Policy Optimization - PPO](https://arxiv.org/abs/1707.06347).
+
+PPO is a policy gradient method for reinforcement learning.
+Simple policy gradient methods one do a single gradient update per sample (or a set of samples).
+Doing multiple gradient steps for a singe sample causes problems
+because the policy deviates too much producing a bad policy.
+PPO lets us do multiple gradient updates per sample by trying to keep the
+policy close to the policy that was used to sample data.
+It does so by clipping gradient flow if the updated policy
+is not close to the policy used to sample the data.
+
+You can find an experiment that uses it [here](https://nn.labml.ai/rl/ppo/experiment.html).
+The experiment uses [Generalized Advantage Estimation](https://nn.labml.ai/rl/ppo/gae.html).
+
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/rl/ppo/experiment.ipynb)