week04_approx_rl/homework_tf.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Deep Q-Network implementation\n",
    "\n",
    "This notebook shamelessly demands you to implement a DQN - an approximate q-learning algorithm with experience replay and target networks - and see if it works any better this way."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys, os\n",
    "if 'google.colab' in sys.modules:\n",
    "    %tensorflow_version 1.x\n",
    "    \n",
    "    if not os.path.exists('.setup_complete'):\n",
    "        !wget -q https://raw.githubusercontent.com/yandexdataschool/Practical_RL/master/setup_colab.sh -O- | bash\n",
    "\n",
    "        !wget -q https://raw.githubusercontent.com/yandexdataschool/Practical_RL/master/week04_approx_rl/framebuffer.py\n",
    "        !wget -q https://raw.githubusercontent.com/yandexdataschool/Practical_RL/master/week04_approx_rl/replay_buffer.py\n",
    "        !wget -q https://raw.githubusercontent.com/yandexdataschool/Practical_RL/master/week04_approx_rl/atari_wrappers.py\n",
    "        !wget -q https://raw.githubusercontent.com/yandexdataschool/Practical_RL/master/week04_approx_rl/utils.py\n",
    "\n",
    "        !touch .setup_complete\n",
    "\n",
    "# This code creates a virtual display to draw game images on.\n",
    "# It will have no effect if your machine has a monitor.\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": [
    "__Frameworks__ - we'll accept this homework in any deep learning framework. This particular notebook was designed for tensorflow, but you will find it easy to adapt it to almost any python-based deep learning framework."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import numpy as np\n",
    "import utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gym\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Let's play some old videogames\n",
    "![img](https://github.com/yandexdataschool/Practical_RL/raw/master/yet_another_week/_resource/nerd.png)\n",
    "\n",
    "This time we're gonna apply approximate q-learning to an Atari game called Breakout. It's not the hardest thing out there, but it's definitely way more complex than anything we tried before.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ENV_NAME = \"BreakoutNoFrameskip-v4\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preprocessing (3 pts)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see what observations look like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = gym.make(ENV_NAME)\n",
    "env.reset()\n",
    "\n",
    "n_cols = 5\n",
    "n_rows = 2\n",
    "fig = plt.figure(figsize=(16, 9))\n",
    "\n",
    "for row in range(n_rows):\n",
    "    for col in range(n_cols):\n",
    "        ax = fig.add_subplot(n_rows, n_cols, row * n_cols + col + 1)\n",
    "        ax.imshow(env.render('rgb_array'))\n",
    "        env.step(env.action_space.sample())\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Processing game image \n",
    "\n",
    "Raw Atari images are large, 210x160x3 by default. However, we don't need that level of detail in order to learn them.\n",
    "\n",
    "We can thus save a lot of time by preprocessing game image, including\n",
    "* Resizing to a smaller shape, 64 x 64\n",
    "* Converting to grayscale\n",
    "* Cropping irrelevant image parts (top & bottom)\n",
    "\n",
    "Tip: You can implement your own grayscale converter and assign a huge weight to the red channel. This dirty trick is not necessary but it will speed up learning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from gym.core import ObservationWrapper\n",
    "from gym.spaces import Box\n",
    "\n",
    "import cv2\n",
    "\n",
    "\n",
    "class PreprocessAtariObs(ObservationWrapper):\n",
    "    def __init__(self, env):\n",
    "        \"\"\"A gym wrapper that crops, scales image into the desired shapes and optionally grayscales it.\"\"\"\n",
    "        super().__init__(env)\n",
    "\n",
    "        self.img_size = (64, 64, 1)\n",
    "        self.observation_space = Box(0.0, 1.0, self.img_size)\n",
    "\n",
    "    def _to_gray_scale(self, rgb, channel_weights=[0.8, 0.1, 0.1]):\n",
    "        <YOUR CODE>\n",
    "\n",
    "    def observation(self, img):\n",
    "        \"\"\"what happens to each observation\"\"\"\n",
    "\n",
    "        # Here's what you need to do:\n",
    "        #  * Crop image, remove irrelevant parts.\n",
    "        #  * Resize image to self.img_size. Use cv2.resize or any other library you want,\n",
    "        #    e.g. PIL or Keras. Do not use skimage.transform.resize because it is roughly\n",
    "        #    6x slower than cv2.resize.\n",
    "        #  * Cast image to grayscale.\n",
    "        #  * Convert image pixels to (0, 1) range, float32 type.\n",
    "\n",
    "        <YOUR CODE>\n",
    "        return <YOUR CODE>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gym\n",
    "# spawn game instance for tests\n",
    "env = gym.make(ENV_NAME)  # create raw env\n",
    "env = PreprocessAtariObs(env)\n",
    "\n",
    "observation_shape = env.observation_space.shape\n",
    "n_actions = env.action_space.n\n",
    "\n",
    "obs = env.reset()\n",
    "\n",
    "\n",
    "# test observation\n",
    "assert obs.ndim == 3, \"observation must be [height, width, channels] even if there's just one channel\"\n",
    "assert obs.shape == observation_shape, (obs.shape, observation_shape)\n",
    "assert obs.dtype == np.float32\n",
    "assert len(np.unique(obs)) > 2, \"your image must not be binary\"\n",
    "assert 0 <= np.min(obs) and np.max(obs) <= 1, \"convert image pixels to (0,1) range\"\n",
    "assert np.max(obs) >= 0.5, \"It would be easier to see a brighter observation\"\n",
    "assert np.mean(obs) >= 0.1, \"It would be easier to see a brighter observation\"\n",
    "\n",
    "assert np.max(obs) >= 0.5, \"It would be easier to see a brighter observation\"\n",
    "assert np.mean(obs) >= 0.1, \"It would be easier to see a brighter observation\"\n",
    "\n",
    "print(\"Formal tests seem fine. Here's an example of what you'll get.\")\n",
    "\n",
    "\n",
    "n_cols = 5\n",
    "n_rows = 2\n",
    "fig = plt.figure(figsize=(16, 9))\n",
    "obs = env.reset()\n",
    "for row in range(n_rows):\n",
    "    for col in range(n_cols):\n",
    "        ax = fig.add_subplot(n_rows, n_cols, row * n_cols + col + 1)\n",
    "        ax.imshow(obs[:, :, 0], interpolation='none', cmap='gray')\n",
    "        obs, _, _, _ = env.step(env.action_space.sample())\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Wrapping."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**About the game:** You have 5 lives and get points for breaking the wall. Higher bricks cost more than the lower ones. There are 4 actions: start game (should be called at the beginning and after each life is lost), move left, move right and do nothing. There are some common wrappers used for Atari environments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import atari_wrappers\n",
    "\n",
    "def PrimaryAtariWrap(env, clip_rewards=True):\n",
    "    assert 'NoFrameskip' in env.spec.id\n",
    "\n",
    "    # This wrapper holds the same action for <skip> frames and outputs\n",
    "    # the maximal pixel value of 2 last frames (to handle blinking\n",
    "    # in some envs)\n",
    "    env = atari_wrappers.MaxAndSkipEnv(env, skip=4)\n",
    "\n",
    "    # This wrapper sends done=True when each life is lost\n",
    "    # (not all the 5 lives that are givern by the game rules).\n",
    "    # It should make easier for the agent to understand that losing is bad.\n",
    "    env = atari_wrappers.EpisodicLifeEnv(env)\n",
    "\n",
    "    # This wrapper laucnhes the ball when an episode starts.\n",
    "    # Without it the agent has to learn this action, too.\n",
    "    # Actually it can but learning would take longer.\n",
    "    env = atari_wrappers.FireResetEnv(env)\n",
    "\n",
    "    # This wrapper transforms rewards to {-1, 0, 1} according to their sign\n",
    "    if clip_rewards:\n",
    "        env = atari_wrappers.ClipRewardEnv(env)\n",
    "\n",
    "    # This wrapper is yours :)\n",
    "    env = PreprocessAtariObs(env)\n",
    "    return env"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Frame buffer\n",
    "\n",
    "Our agent can only process one observation at a time, so we gotta make sure it contains enough information to fing optimal actions. For instance, agent has to react to moving objects so he must be able to measure object's velocity.\n",
    "\n",
    "To do so, we introduce a buffer that stores 4 last images. This time everything is pre-implemented for you."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from framebuffer import FrameBuffer\n",
    "\n",
    "\n",
    "def make_env(clip_rewards=True, seed=None):\n",
    "    env = gym.make(ENV_NAME)  # create raw env\n",
    "    if seed is not None:\n",
    "        env.seed(seed)\n",
    "    env = PrimaryAtariWrap(env, clip_rewards)\n",
    "    env = FrameBuffer(env, n_frames=4, dim_order='tensorflow')\n",
    "    return env\n",
    "\n",
    "\n",
    "env = make_env()\n",
    "env.reset()\n",
    "n_actions = env.action_space.n\n",
    "state_shape = env.observation_space.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for _ in range(50):\n",
    "    obs, _, _, _ = env.step(env.action_space.sample())\n",
    "\n",
    "\n",
    "plt.figure()\n",
    "plt.title(\"Game image\")\n",
    "plt.imshow(env.render(\"rgb_array\"))\n",
    "plt.show()\n",
    "\n",
    "plt.figure(figsize=[12,10])\n",
    "plt.title(\"Agent observation (4 frames left to right)\")\n",
    "plt.imshow(obs.transpose([0, 2, 1]).reshape([state_shape[0], -1]))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## DQN as it is (4 pts)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Building a network\n",
    "\n",
    "We now need to build a neural network that can map images to state q-values. This network will be called on every agent's step so it better not be resnet-152 unless you have an array of GPUs. Instead, you can use strided convolutions with a small number of features to save time and memory.\n",
    "\n",
    "You can build any architecture you want, but for reference, here's something that will more or less work:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![img](https://github.com/yandexdataschool/Practical_RL/raw/master/yet_another_week/_resource/dqn_arch.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Dueling network: (+2 pts)**\n",
    "$$Q_{\\theta}(s, a) = V_{\\eta}(f_{\\xi}(s)) + A_{\\psi}(f_{\\xi}(s), a) - \\frac{\\sum_{a'}A_{\\psi}(f_{\\xi}(s), a')}{N_{actions}},$$\n",
    "where $\\xi$, $\\eta$, and $\\psi$ are, respectively, the parameters of the\n",
    "shared encoder $f_ξ$ , of the value stream $V_\\eta$ , and of the advan\n",
    "tage stream $A_\\psi$; and $\\theta = \\{\\xi, \\eta, \\psi\\}$ is their concatenation.\n",
    "\n",
    "For the architecture on the image $V$ and $A$ heads can follow the dense layer instead of $Q$. Please don't worry that the model becomes a little bigger."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "tf.reset_default_graph()\n",
    "sess = tf.InteractiveSession()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.keras.layers import Conv2D, Dense, Flatten\n",
    "from tensorflow.keras import models\n",
    "\n",
    "\n",
    "class DQNAgent:\n",
    "    def __init__(self, name, state_shape, n_actions, epsilon=0, reuse=False):\n",
    "        \"\"\"A simple DQN agent\"\"\"\n",
    "        with tf.variable_scope(name, reuse=reuse):\n",
    "\n",
    "            <YOUR CODE: define your network body here. Please make sure you don't use any layers created elsewhere>\n",
    "\n",
    "            # prepare a graph for agent step\n",
    "            self.state_t = tf.placeholder('float32', [None, ] + list(state_shape))\n",
    "            self.qvalues_t = self.get_symbolic_qvalues(self.state_t)\n",
    "\n",
    "        self.weights = tf.get_collection(\n",
    "            tf.GraphKeys.TRAINABLE_VARIABLES, scope=name)\n",
    "        self.epsilon = epsilon\n",
    "\n",
    "    def get_symbolic_qvalues(self, state_t):\n",
    "        \"\"\"takes agent's observation, returns qvalues. Both are tf Tensors\"\"\"\n",
    "        <YOUR CODE: apply your network layers here>\n",
    "        qvalues = <YOUR CODE: symbolic tensor for q-values>\n",
    "\n",
    "        assert tf.is_numeric_tensor(qvalues) and qvalues.shape.ndims == 2, \\\n",
    "            \"please return 2d tf tensor of qvalues [you got %s]\" % repr(qvalues)\n",
    "        assert int(qvalues.shape[1]) == n_actions\n",
    "\n",
    "        return qvalues\n",
    "\n",
    "    def get_qvalues(self, state_t):\n",
    "        \"\"\"Same as symbolic step except it operates on numpy arrays\"\"\"\n",
    "        sess = tf.get_default_session()\n",
    "        return sess.run(self.qvalues_t, {self.state_t: state_t})\n",
    "\n",
    "    def sample_actions(self, qvalues):\n",
    "        \"\"\"pick actions given qvalues. Uses epsilon-greedy exploration strategy. \"\"\"\n",
    "        epsilon = self.epsilon\n",
    "        batch_size, n_actions = qvalues.shape\n",
    "        random_actions = np.random.choice(n_actions, size=batch_size)\n",
    "        best_actions = qvalues.argmax(axis=-1)\n",
    "        should_explore = np.random.choice([0, 1], batch_size, p=[1-epsilon, epsilon])\n",
    "        return np.where(should_explore, random_actions, best_actions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "agent = DQNAgent(\"dqn_agent\", state_shape, n_actions, epsilon=1)\n",
    "sess.run(tf.global_variables_initializer())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's try out our agent to see if it raises any errors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(env, agent, n_games=1, greedy=False, t_max=10000):\n",
    "    \"\"\" Plays n_games full games. If greedy, picks actions as argmax(qvalues). Returns mean reward. \"\"\"\n",
    "    rewards = []\n",
    "    for _ in range(n_games):\n",
    "        s = env.reset()\n",
    "        reward = 0\n",
    "        for _ in range(t_max):\n",
    "            qvalues = agent.get_qvalues([s])\n",
    "            action = qvalues.argmax(axis=-1)[0] if greedy else agent.sample_actions(qvalues)[0]\n",
    "            s, r, done, _ = env.step(action)\n",
    "            reward += r\n",
    "            if done:\n",
    "                break\n",
    "\n",
    "        rewards.append(reward)\n",
    "    return np.mean(rewards)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluate(env, agent, n_games=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Experience replay\n",
    "For this assignment, we provide you with experience replay buffer. If you implemented experience replay buffer in last week's assignment, you can copy-paste it here __to get 2 bonus points__.\n",
    "\n",
    "![img](https://github.com/yandexdataschool/Practical_RL/raw/master/yet_another_week/_resource/exp_replay.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### The interface is fairly simple:\n",
    "* `exp_replay.add(obs, act, rw, next_obs, done)` - saves (s,a,r,s',done) tuple into the buffer\n",
    "* `exp_replay.sample(batch_size)` - returns observations, actions, rewards, next_observations and is_done for `batch_size` random samples.\n",
    "* `len(exp_replay)` - returns number of elements stored in replay buffer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from replay_buffer import ReplayBuffer\n",
    "exp_replay = ReplayBuffer(10)\n",
    "\n",
    "for _ in range(30):\n",
    "    exp_replay.add(env.reset(), env.action_space.sample(), 1.0, env.reset(), done=False)\n",
    "\n",
    "obs_batch, act_batch, reward_batch, next_obs_batch, is_done_batch = exp_replay.sample(5)\n",
    "\n",
    "assert len(exp_replay) == 10, \"experience replay size should be 10 because that's what maximum capacity is\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def play_and_record(initial_state, agent, env, exp_replay, n_steps=1):\n",
    "    \"\"\"\n",
    "    Play the game for exactly n steps, record every (s,a,r,s', done) to replay buffer. \n",
    "    Whenever game ends, add record with done=True and reset the game.\n",
    "    It is guaranteed that env has done=False when passed to this function.\n",
    "\n",
    "    PLEASE DO NOT RESET ENV UNLESS IT IS \"DONE\"\n",
    "\n",
    "    :returns: return sum of rewards over time and the state in which the env stays\n",
    "    \"\"\"\n",
    "    # initial state\n",
    "    s = initial_state\n",
    "    sum_rewards = 0\n",
    "\n",
    "    # Play the game for n_steps as per instructions above\n",
    "    <YOUR CODE>\n",
    "\n",
    "    return sum_rewards, s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# testing your code.\n",
    "exp_replay = ReplayBuffer(2000)\n",
    "\n",
    "state = env.reset()\n",
    "play_and_record(state, agent, env, exp_replay, n_steps=1000)\n",
    "\n",
    "# if you're using your own experience replay buffer, some of those tests may need correction.\n",
    "# just make sure you know what your code does\n",
    "assert len(exp_replay) == 1000, \"play_and_record should have added exactly 1000 steps, \"\\\n",
    "                                 \"but instead added %i\" % len(exp_replay)\n",
    "is_dones = list(zip(*exp_replay._storage))[-1]\n",
    "\n",
    "assert 0 < np.mean(is_dones) < 0.1, \"Please make sure you restart the game whenever it is 'done' and record the is_done correctly into the buffer.\"\\\n",
    "                                    \"Got %f is_done rate over %i steps. [If you think it's your tough luck, just re-run the test]\" % (\n",
    "                                        np.mean(is_dones), len(exp_replay))\n",
    "\n",
    "for _ in range(100):\n",
    "    obs_batch, act_batch, reward_batch, next_obs_batch, is_done_batch = exp_replay.sample(\n",
    "        10)\n",
    "    assert obs_batch.shape == next_obs_batch.shape == (10,) + state_shape\n",
    "    assert act_batch.shape == (10,), \\\n",
    "        \"actions batch should have shape (10,) but is instead %s\" % str(act_batch.shape)\n",
    "    assert reward_batch.shape == (10,), \\\n",
    "        \"rewards batch should have shape (10,) but is instead %s\" % str(reward_batch.shape)\n",
    "    assert is_done_batch.shape == (10,), \\\n",
    "        \"is_done batch should have shape (10,) but is instead %s\" % str(is_done_batch.shape)\n",
    "    assert [int(i) in (0, 1) for i in is_dones], \"is_done should be strictly True or False\"\n",
    "    assert [0 <= a < n_actions for a in act_batch], \"actions should be within [0, n_actions)\"\n",
    "\n",
    "print(\"Well done!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Target networks\n",
    "\n",
    "We also employ the so called \"target network\" - a copy of neural network weights to be used for reference Q-values:\n",
    "\n",
    "The network itself is an exact copy of agent network, but it's parameters are not trained. Instead, they are moved here from agent's actual network every so often.\n",
    "\n",
    "$$ Q_{reference}(s,a) = r + \\gamma \\cdot \\max _{a'} Q_{target}(s',a') $$\n",
    "\n",
    "![img](https://github.com/yandexdataschool/Practical_RL/raw/master/yet_another_week/_resource/target_net.png)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "target_network = DQNAgent(\"target_network\", state_shape, n_actions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_weigths_into_target_network(agent, target_network):\n",
    "    \"\"\" assign target_network.weights variables to their respective agent.weights values. \"\"\"\n",
    "    assigns = []\n",
    "    for w_agent, w_target in zip(agent.weights, target_network.weights):\n",
    "        assigns.append(tf.assign(w_target, w_agent, validate_shape=True))\n",
    "    # tf.get_default_session().run(assigns)\n",
    "    return assigns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create the tf copy graph only once.\n",
    "copy_step = load_weigths_into_target_network(agent, target_network)\n",
    "sess.run(copy_step)\n",
    "# check that it works\n",
    "sess.run([tf.assert_equal(w, w_target) for w, w_target in zip(agent.weights, target_network.weights)])\n",
    "print(\"It works!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Learning with... Q-learning\n",
    "Here we write a function similar to `agent.update` from tabular q-learning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# placeholders that will be fed with exp_replay.sample(batch_size)\n",
    "obs_ph = tf.placeholder(tf.float32, shape=(None,) + state_shape)\n",
    "actions_ph = tf.placeholder(tf.int32, shape=[None])\n",
    "rewards_ph = tf.placeholder(tf.float32, shape=[None])\n",
    "next_obs_ph = tf.placeholder(tf.float32, shape=(None,) + state_shape)\n",
    "is_done_ph = tf.placeholder(tf.float32, shape=[None])\n",
    "\n",
    "is_not_done = 1 - is_done_ph\n",
    "gamma = 0.99"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Take q-values for actions agent just took"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "current_qvalues = agent.get_symbolic_qvalues(obs_ph)\n",
    "current_action_qvalues = tf.reduce_sum(tf.one_hot(actions_ph, n_actions) * current_qvalues, axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compute Q-learning TD error:\n",
    "\n",
    "$$ L = { 1 \\over N} \\sum_i [ Q_{\\theta}(s,a) - Q_{reference}(s,a) ] ^2 $$\n",
    "\n",
    "With Q-reference defined as\n",
    "\n",
    "$$ Q_{reference}(s,a) = r(s,a) + \\gamma \\cdot max_{a'} Q_{target}(s', a') $$\n",
    "\n",
    "Where\n",
    "* $Q_{target}(s',a')$ denotes q-value of next state and next action predicted by __target_network__\n",
    "* $s, a, r, s'$ are current state, action, reward and next state respectively\n",
    "* $\\gamma$ is a discount factor defined two cells above.\n",
    "\n",
    "\n",
    "__Note 1:__ there's an example input below. Feel free to experiment with it before you write the function.\n",
    "\n",
    "__Note 2:__ compute_td_loss is a source of 99% of bugs in this homework. If reward doesn't improve, it often helps to go through it line by line [with a rubber duck](https://rubberduckdebugging.com/).\n",
    "\n",
    "**Double DQN (+2 pts)**\n",
    "\n",
    "$$ Q_{reference}(s,a) = r(s, a) + \\gamma \\cdot\n",
    "Q_{target}(s',argmax_{a'}Q_\\theta(s', a')) $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "next_qvalues_target = <YOUR CODE: compute q-values for NEXT states with target network>\n",
    "next_state_values_target = <YOUR CODE: compute state values by taking max over next_qvalues_target for all actions>\n",
    "reference_qvalues = <YOUR CODE: compute Q_reference(s,a) as per formula above>\n",
    "\n",
    "# Define loss function for sgd.\n",
    "td_loss = (current_action_qvalues - reference_qvalues) ** 2\n",
    "td_loss = tf.reduce_mean(td_loss)\n",
    "\n",
    "train_step = tf.train.AdamOptimizer(1e-3).minimize(td_loss, var_list=agent.weights)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sess.run(tf.global_variables_initializer())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for chk_grad in tf.gradients(reference_qvalues, agent.weights):\n",
    "    error_msg = \"Reference q-values should have no gradient w.r.t. agent weights. Make sure you used target_network qvalues! \"\n",
    "    error_msg += \"If you know what you're doing, ignore this assert.\"\n",
    "    assert chk_grad is None or np.allclose(sess.run(chk_grad), sess.run(chk_grad * 0)), error_msg\n",
    "\n",
    "assert tf.gradients(reference_qvalues, is_not_done)[0] is not None, \"make sure you used is_not_done\"\n",
    "assert tf.gradients(reference_qvalues, rewards_ph)[0] is not None, \"make sure you used rewards\"\n",
    "assert tf.gradients(reference_qvalues, next_obs_ph)[0] is not None, \"make sure you used next states\"\n",
    "assert tf.gradients(reference_qvalues, obs_ph)[0] is None, \"reference qvalues shouldn't depend on current observation!\"  # ignore if you're certain it's ok\n",
    "print(\"Splendid!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Main loop (3 pts)\n",
    "\n",
    "**If deadline is tonight and it has not converged:** It is ok. Send the notebook today and when it converges send it again.\n",
    "If the code is exactly the same points will not be discounted.\n",
    "\n",
    "It's time to put everything together and see if it learns anything."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import trange\n",
    "import pandas as pd\n",
    "from IPython.display import clear_output\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "seed = <YOUR CODE: your favourite random seed>\n",
    "random.seed(seed)\n",
    "np.random.seed(seed)\n",
    "tf.set_random_seed(seed)\n",
    "\n",
    "env = make_env(seed)\n",
    "state_shape = env.observation_space.shape\n",
    "n_actions = env.action_space.n\n",
    "state = env.reset()\n",
    "\n",
    "sess.run(tf.global_variables_initializer())\n",
    "sess.run(copy_step)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "REPLAY_BUFFER_SIZE = 10**5\n",
    "\n",
    "exp_replay = ReplayBuffer(REPLAY_BUFFER_SIZE)\n",
    "_, state = play_and_record(state, agent, env, exp_replay, n_steps=10000)\n",
    "\n",
    "\n",
    "def sample_batch(exp_replay, batch_size):\n",
    "    obs_batch, act_batch, reward_batch, next_obs_batch, is_done_batch = exp_replay.sample(batch_size)\n",
    "    return {\n",
    "        obs_ph: obs_batch,\n",
    "        actions_ph: act_batch,\n",
    "        rewards_ph: reward_batch,\n",
    "        next_obs_ph: next_obs_batch,\n",
    "        is_done_ph: is_done_batch,\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "timesteps_per_epoch = 10\n",
    "batch_size = 64\n",
    "total_steps = 3 * 10**5\n",
    "decay_steps = 10**5\n",
    "\n",
    "init_epsilon = 1\n",
    "final_epsilon = 0.1\n",
    "\n",
    "loss_freq = 5\n",
    "refresh_target_network_freq = 500\n",
    "eval_freq = 500\n",
    "\n",
    "n_lives = 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mean_rw_history = []\n",
    "td_loss_history = []\n",
    "initial_state_v_history = []\n",
    "step = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "state = env.reset()\n",
    "\n",
    "with trange(step, total_steps) as progress_bar:\n",
    "    for step in progress_bar:\n",
    "        agent.epsilon = utils.linear_decay(init_epsilon, final_epsilon, step, decay_steps)\n",
    "        \n",
    "        # play\n",
    "        _, state = play_and_record(state, agent, env, exp_replay, timesteps_per_epoch)\n",
    "\n",
    "        # train\n",
    "        _, loss_t = sess.run([train_step, td_loss], sample_batch(exp_replay, batch_size=batch_size))\n",
    "        \n",
    "        if step % loss_freq == 0:\n",
    "            td_loss_history.append(loss_t)\n",
    "            \n",
    "        if step % refresh_target_network_freq == 0:\n",
    "            # Load agent weights into target_network\n",
    "            sess.run(copy_step)\n",
    "\n",
    "        if step % eval_freq == 0:\n",
    "            # eval the agent\n",
    "            mean_rw_history.append(evaluate(\n",
    "                make_env(clip_rewards=True, seed=step), agent, n_games=3 * n_lives, greedy=True)\n",
    "            )\n",
    "            initial_state_q_values = agent.get_qvalues(\n",
    "                [make_env(seed=step).reset()]\n",
    "            )\n",
    "            initial_state_v_history.append(np.max(initial_state_q_values))\n",
    "\n",
    "            clear_output(True)\n",
    "            print(\"buffer size = %i, epsilon = %.5f\" %\n",
    "                (len(exp_replay), agent.epsilon))\n",
    "\n",
    "            plt.figure(figsize=[16, 9])\n",
    "\n",
    "            plt.subplot(2, 2, 1)\n",
    "            plt.title(\"Mean reward per life\")\n",
    "            plt.plot(mean_rw_history)\n",
    "            plt.grid()\n",
    "\n",
    "            assert not np.isnan(td_loss_history[-1])\n",
    "            plt.subplot(2, 2, 2)\n",
    "            plt.title(\"TD loss history (smoothened)\")\n",
    "            plt.plot(utils.smoothen(td_loss_history))\n",
    "            plt.grid()\n",
    "\n",
    "            plt.subplot(2, 2, 3)\n",
    "            plt.title(\"Initial state V\")\n",
    "            plt.plot(initial_state_v_history)\n",
    "            plt.grid()\n",
    "\n",
    "            plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Agent is evaluated for 1 life, not for a whole episode of 5 lives. Rewards in evaluation are also truncated. Cuz this is what environment the agent is learning in and in this way mean rewards per life can be compared with initial state value\n",
    "\n",
    "**The goal is to get 15 points in the real env**. So 3 or better 4 points in the preprocessed one will probably be enough. You can interrupt learning then."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Final scoring is done on a whole episode with all 5 lives."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "final_score = evaluate(\n",
    "  make_env(clip_rewards=False, seed=9),\n",
    "    agent, n_games=30, greedy=True, t_max=10 * 1000\n",
    ") * n_lives\n",
    "print('final score:', final_score)\n",
    "assert final_score >= 15, 'not as cool as DQN can'\n",
    "print('Cool!')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## How to interpret plots:\n",
    "\n",
    "This aint no supervised learning so don't expect anything to improve monotonously. \n",
    "* **TD loss** is the MSE between agent's current Q-values and target Q-values. It may slowly increase or decrease, it's ok. The \"not ok\" behavior includes going NaN or stayng at exactly zero before agent has perfect performance.\n",
    "* **mean reward** is the expected sum of r(s,a) agent gets over the full game session. It will oscillate, but on average it should get higher over time (after a few thousand iterations...). \n",
    " * In basic q-learning implementation it takes about 40k steps to \"warm up\" agent before it starts to get better.\n",
    "* **Initial state V** is the expected discounted reward for episode in the oppinion of the agent. It should behave more smoothly than **mean reward**. It should get higher over time but sometimes can experience drawdowns because of the agaent's overestimates.\n",
    "* **buffer size** - this one is simple. It should go up and cap at max size.\n",
    "* **epsilon** - agent's willingness to explore. If you see that agent's already at 0.01 epsilon before it's average reward is above 0 - it means you need to increase epsilon. Set it back to some 0.2 - 0.5 and decrease the pace at which it goes down.\n",
    "* Smoothing of plots is done with a gaussian kernel\n",
    "\n",
    "At first your agent will lose quickly. Then it will learn to suck less and at least hit the ball a few times before it loses. Finally it will learn to actually score points.\n",
    "\n",
    "**Training will take time.** A lot of it actually. Probably you will not see any improvment during first **150k** time steps (note that by default in this notebook agent is evaluated every 5000 time steps).\n",
    "\n",
    "But hey, long training time isn't _that_ bad:\n",
    "\n",
    "![img](https://github.com/yandexdataschool/Practical_RL/raw/master/yet_another_week/_resource/training.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## About hyperparameters:\n",
    "\n",
    "The task has something in common with supervised learning: loss is optimized through the buffer (instead of Train dataset). But the distribution of states and actions in the buffer **is not stationary** and depends on the policy that generated it. It can even happen that the mean TD error across the buffer is very low but the performance is extremely poor (imagine the agent collecting data to the buffer always manages to avoid the ball).\n",
    "\n",
    "* Total timesteps and training time: It seems to be so huge, but actually it is normal for RL.\n",
    "\n",
    "* $\\epsilon$ decay shedule was taken from the original paper and is like traditional for epsilon-greedy policies. At the beginning of the training the agent's greedy policy is poor so many random actions should be taken.\n",
    "\n",
    "* Optimizer: In the original paper RMSProp was used (they did not have Adam in 2013) and it can work not worse than Adam. For us Adam was default and it worked.\n",
    "\n",
    "* lr: $10^{-3}$ would probably be too huge\n",
    "\n",
    "* batch size: This one can be very important: if it is too small the agent can fail to learn. Huge batch takes more time to process.\n",
    "\n",
    "* target network update frequency: has something in common with learning rate. Too frequent updates can lead to divergence. Too rare can lead to slow learning. For millions of total timesteps thousands of inner steps seem ok. One iteration of target network updating is an iteration of the (this time approximate) $\\gamma$-compression that stands behind Q-learning. The more inner steps it makes the more accurate is the compression."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Let's have a closer look at this.\n",
    "\n",
    "If average episode score is below 200 using all 5 lives, then probably DQN has not converged fully. But anyway let's make a more complete record of an episode."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_env = make_env(clip_rewards=False)\n",
    "record = utils.play_and_log_episode(eval_env, agent)\n",
    "print('total reward for life:', np.sum(record['rewards']))\n",
    "for key in record:\n",
    "    print(key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(5, 5))\n",
    "ax = fig.add_subplot(1, 1, 1)\n",
    "\n",
    "ax.scatter(record['v_mc'], record['v_agent'])\n",
    "ax.plot(sorted(record['v_mc']), sorted(record['v_mc']),\n",
    "       'black', linestyle='--', label='x=y')\n",
    "\n",
    "ax.grid()\n",
    "ax.legend()\n",
    "ax.set_title('State Value Estimates')\n",
    "ax.set_xlabel('Monte-Carlo')\n",
    "ax.set_ylabel('Agent')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$\\hat V_{Monte-Carlo}(s_t) = \\sum_{\\tau=0}^{episode~end} \\gamma^{\\tau-t}r_t$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Is there a big bias? It's ok, anyway it works."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Video"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Record sessions\n",
    "\n",
    "import gym.wrappers\n",
    "\n",
    "with gym.wrappers.Monitor(make_env(), directory=\"videos\", force=True) as env_monitor:\n",
    "    sessions = [evaluate(env_monitor, agent, n_games=n_lives, greedy=True) for _ in range(10)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show video. This may not work in some setups. If it doesn't\n",
    "# work for you, you can download the videos and view them locally.\n",
    "\n",
    "from pathlib import Path\n",
    "from base64 import b64encode\n",
    "from IPython.display import HTML\n",
    "\n",
    "video_paths = sorted([s for s in Path('videos').iterdir() if s.suffix == '.mp4'])\n",
    "video_path = video_paths[-1]  # You can also try other indices\n",
    "\n",
    "if 'google.colab' in sys.modules:\n",
    "    # https://stackoverflow.com/a/57378660/1214547\n",
    "    with video_path.open('rb') as fp:\n",
    "        mp4 = fp.read()\n",
    "    data_url = 'data:video/mp4;base64,' + b64encode(mp4).decode()\n",
    "else:\n",
    "    data_url = str(video_path)\n",
    "\n",
    "HTML(\"\"\"\n",
    "<video width=\"640\" height=\"480\" controls>\n",
    "  <source src=\"{}\" type=\"video/mp4\">\n",
    "</video>\n",
    "\"\"\".format(data_url))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bonus I (2 pts)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**1.** Plot several (say 3) states with high and low spreads of Q estimate by actions i.e.\n",
    "$$\\max_a \\hat Q(s,a) - \\min_a \\hat Q(s,a)\\$$\n",
    "Please take those states from different episodes to make sure that the states are really different.\n",
    "\n",
    "What should high and low spread mean at least in the world of perfect Q-fucntions?\n",
    "\n",
    "Comment the states you like most.\n",
    "\n",
    "**2.** Plot several (say 3) states with high td-error and several states with high values of\n",
    "$$| \\hat V_{Monte-Carlo}(s) - \\hat V_{agent}(s)|,$$ \n",
    "$$\\hat V_{agent}(s)=\\max_a \\hat Q(s,a).$$ Please take those states from different episodes to make sure that the states are really different. From what part (i.e. beginning, middle, end) of an episode did these states come from?\n",
    "\n",
    "Comment the states you like most."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils import play_and_log_episode, img_by_obs\n",
    "\n",
    "<YOUR CODE>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bonus II (1-5 pts). Get High Score!\n",
    "\n",
    "1 point to you for each 50 points of your agent. Truncated by 5 points. Starting with 50 points, **not** 50 + threshold.\n",
    "\n",
    "One way is to train for several days and use heavier hardware (why not actually).\n",
    "\n",
    "Another way is to apply modifications (see **Bonus III**)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bonus III (2+ pts). Apply modifications to DQN.\n",
    "\n",
    "For inspiration see [Rainbow](https://arxiv.org/abs/1710.02298) - a version of q-learning that combines lots of them.\n",
    "\n",
    "Points for Bonus II and Bonus III fully stack. So if modified agent gets score 250+ you get 5 pts for Bonus II + points for modifications. If the final score is 40 then you get the points for modifications.\n",
    "\n",
    "\n",
    "Some modifications:\n",
    "* [Prioritized experience replay](https://arxiv.org/abs/1511.05952) (5 pts for your own implementation, 3 pts for using a ready one)\n",
    "* [double q-learning](https://arxiv.org/abs/1509.06461) (2 pts)\n",
    "* [dueling q-learning](https://arxiv.org/abs/1511.06581) (2 pts)\n",
    "* multi-step heuristics (see [Rainbow](https://arxiv.org/abs/1710.02298)) (3 pts)\n",
    "* [Noisy Nets](https://arxiv.org/abs/1706.10295) (3 pts)\n",
    "* [distributional RL](https://arxiv.org/abs/1707.06887)(distributional and distributed stand for different things here) (5 pts)\n",
    "* Other modifications (2+ pts depending on complexity)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bonus IV (4+ pts). Distributed RL.\n",
    "\n",
    "Solve the task in a distributed way. It can strongly speed up learning. See [article](https://arxiv.org/pdf/1602.01783.pdf) or some guides."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**As usual bonus points for all the tasks fully stack.**"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "language_info": {
   "name": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}