week06_policy_based/a2c-optional.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "if 'google.colab' in sys.modules:\n",
    "    import os\n",
    "\n",
    "    os.system('apt-get install -y xvfb')\n",
    "    os.system('wget https://raw.githubusercontent.com/yandexdataschool/Practical_RL/master/xvfb -O ../xvfb')\n",
    "    os.system('apt-get install -y python-opengl ffmpeg')\n",
    "    os.system('pip install pyglet==1.2.4')\n",
    "\n",
    "    os.system('python -m pip install -U pygame --user')\n",
    "\n",
    "    print('setup complete')\n",
    "\n",
    "# XVFB will be launched if you run on a server\n",
    "import os\n",
    "if type(os.environ.get(\"DISPLAY\")) is not str or len(os.environ.get(\"DISPLAY\")) == 0:\n",
    "    !bash ../xvfb start\n",
    "    os.environ['DISPLAY'] = ':1'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Implementing Advantage-Actor Critic (A2C)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook you will implement Advantage Actor Critic algorithm that trains on a batch of Atari 2600 environments running in parallel. \n",
    "\n",
    "Firstly, we will use environment wrappers implemented in file `atari_wrappers.py`. These wrappers preprocess observations (resize, grayscal, take max between frames, skip frames and stack them together) and rewards. Some of the wrappers help to reset the environment and pass `done` flag equal to `True` when agent dies.\n",
    "File `env_batch.py` includes implementation of `ParallelEnvBatch` class that allows to run multiple environments in parallel. To create an environment we can use `nature_dqn_env` function. Note that if you are using \n",
    "PyTorch and not using `tensorboardX` you will need to implement a wrapper that will log **raw** total rewards that the *unwrapped* environment returns and redefine the implemention of `nature_dqn_env` function here. \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from atari_wrappers import nature_dqn_env, NumpySummaries\n",
    "\n",
    "\n",
    "env = nature_dqn_env(\"SpaceInvadersNoFrameskip-v4\", nenvs=8, summaries='Numpy')\n",
    "obs = env.reset()\n",
    "assert obs.shape == (8, 84, 84, 4)\n",
    "assert obs.dtype == np.uint8"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we will need to implement a model that predicts logits and values. It is suggested that you use the same model as in [Nature DQN paper](https://web.stanford.edu/class/psych209/Readings/MnihEtAlHassibis15NatureControlDeepRL.pdf) with a modification that instead of having a single output layer, it will have two output layers taking as input the output of the last hidden layer. **Note** that this model is different from the model you used in homework where you implemented DQN. You can use your favorite deep learning framework here. We suggest that you use orthogonal initialization with parameter $\\sqrt{2}$ for kernels and initialize biases with zeros. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import tensorflow as torch\n",
    "# import torch as tf\n",
    "\n",
    "<Define your model here>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You will also need to define and use a policy that wraps the model. While the model computes logits for all actions, the policy will sample actions and also compute their log probabilities.  `policy.act` should return a dictionary of all the arrays that are needed to interact with an environment and train the model.\n",
    " Note that actions must be an `np.ndarray` while the other\n",
    "tensors need to have the type determined by your deep learning framework. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Policy:\n",
    "    def __init__(self, model):\n",
    "        self.model = model\n",
    "    \n",
    "    def act(self, inputs):\n",
    "        <Implement policy by calling model, sampling actions and computing their log probs>\n",
    "        # Should return a dict containing keys ['actions', 'logits', 'log_probs', 'values']."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next will pass the environment and policy to a runner that collects partial trajectories from the environment. \n",
    "The class that does is is already implemented for you."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from runners import EnvRunner"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This runner interacts with the environment for a given number of steps and returns a dictionary containing\n",
    "keys \n",
    "\n",
    "* 'observations' \n",
    "* 'rewards' \n",
    "* 'resets'\n",
    "* 'actions'\n",
    "* all other keys that you defined in `Policy`\n",
    "\n",
    "under each of these keys there is a python `list` of interactions with the environment of specified length $T$ &mdash; the size of partial trajectory. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To train the part of the model that predicts state values you will need to compute the value targets. \n",
    "Any callable could be passed to `EnvRunner` to be applied to each partial trajectory after it is collected. \n",
    "Thus, we can implement and use `ComputeValueTargets` callable. \n",
    "The formula for the value targets is simple:\n",
    "\n",
    "$$\n",
    "\\hat v(s_t) = \\left( \\sum_{t'=0}^{T - 1 - t} \\gamma^{t'}r_{t+t'} \\right) + \\gamma^T \\hat{v}(s_{t+T}),\n",
    "$$\n",
    "\n",
    "In implementation, however, do not forget to use \n",
    "`trajectory['resets']` flags to check if you need to add the value targets at the next step when \n",
    "computing value targets for the current step. You can access `trajectory['state']['latest_observation']`\n",
    "to get last observations in partial trajectory &mdash; $s_{t+T}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ComputeValueTargets:\n",
    "    def __init__(self, policy, gamma=0.99):\n",
    "        self.policy = policy\n",
    "    \n",
    "    def __call__(self, trajectory):\n",
    "        # This method should modify trajectory inplace by adding\n",
    "        # an item with key 'value_targets' to it.\n",
    "        <Compute value targets for a given partial trajectory>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After computing value targets we will transform lists of interactions into tensors\n",
    "with the first dimension `batch_size` which is equal to `T * nenvs`, i.e. you essentially need\n",
    "to flatten the first two dimensions. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MergeTimeBatch:\n",
    "    \"\"\" Merges first two axes typically representing time and env batch. \"\"\"\n",
    "    def __call__(self, trajectory):\n",
    "        # Modify trajectory inplace.\n",
    "        <TODO: implement>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = <Create your model here>\n",
    "policy = Policy(model)\n",
    "runner = EnvRunner(\n",
    "    env, policy, nsteps=5,\n",
    "    transforms=[\n",
    "        ComputeValueTargets(),\n",
    "        MergeTimeBatch(),\n",
    "    ])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now is the time to implement the advantage actor critic algorithm itself. You can look into your lecture,\n",
    "[Mnih et al. 2016](https://arxiv.org/abs/1602.01783) paper, and [lecture](https://www.youtube.com/watch?v=Tol_jw5hWnI&list=PLkFD6_40KJIxJMR-j5A1mkxK26gh_qg37&index=20) by Sergey Levine."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class A2C:\n",
    "    def __init__(self,\n",
    "                 policy,\n",
    "                 optimizer,\n",
    "                 value_loss_coef=0.25,\n",
    "                 entropy_coef=0.01,\n",
    "                 max_grad_norm=0.5):\n",
    "        self.policy = policy\n",
    "        self.optimizer = optimizer\n",
    "        self.value_loss_coef = value_loss_coef\n",
    "        self.entropy_coef = entropy_coef\n",
    "        self.max_grad_norm = max_grad_norm\n",
    "    \n",
    "    def policy_loss(self, trajectory):\n",
    "        # You will need to compute advantages here.\n",
    "        <TODO: implement>\n",
    "    \n",
    "    def value_loss(self, trajectory):\n",
    "        <TODO: implement>\n",
    "    \n",
    "    def loss(self, trajectory):\n",
    "        <TODO: implement>\n",
    "      \n",
    "    def step(self, trajectory):\n",
    "        <TODO: implement>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now you can train your model. With reasonable hyperparameters training on a single GTX1080 for 10 million steps across all batched environments (which translates to about 5 hours of wall clock time)\n",
    "it should be possible to achieve *average raw reward over last 100 episodes* (the average is taken over 100 last \n",
    "episodes in each environment in the batch) of about 600. You should plot this quantity with respect to \n",
    "`runner.step_var` &mdash; the number of interactions with all environments. It is highly \n",
    "encouraged to also provide plots of the following quantities (these are useful for debugging as well):\n",
    "\n",
    "* [Coefficient of Determination](https://en.wikipedia.org/wiki/Coefficient_of_determination) between \n",
    "value targets and value predictions\n",
    "* Entropy of the policy $\\pi$\n",
    "* Value loss\n",
    "* Policy loss\n",
    "* Value targets\n",
    "* Value predictions\n",
    "* Gradient norm\n",
    "* Advantages\n",
    "* A2C loss\n",
    "\n",
    "For optimization we suggest you use RMSProp with learning rate starting from 7e-4 and linearly decayed to 0, smoothing constant (alpha in PyTorch and decay in TensorFlow) equal to 0.99 and epsilon equal to 1e-5."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a2c = <Create instance of the algorithm> \n",
    "\n",
    "<Write your training loop>"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}