week04_approx_rl/homework_pytorch_main.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1NDjJiqysoT-"
   },
   "source": [
    "# Deep Q-Network implementation.\n",
    "\n",
    "This homework shamelessly demands you to implement DQN — an approximate Q-learning algorithm with experience replay and target networks — and see if it works any better this way.\n",
    "\n",
    "**Papers:**\n",
    "\n",
    "[1] Original paper, 2013: https://arxiv.org/pdf/1312.5602.pdf\n",
    "\n",
    "[2] Extended paper, Nature, 2015: https://web.stanford.edu/class/psych209/Readings/MnihEtAlHassibis15NatureControlDeepRL.pdf\n",
    "\n",
    "[3] Rainbow, 2017: https://arxiv.org/pdf/1710.02298.pdf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BcLhaXMKsoT_"
   },
   "source": [
    "**This notebook is the main homework notebook.**\n",
    "Another notebook is given for debug: (**homework_pytorch_debug**). The debug notebook is not supported anymore, the codes have diverged. However, it can be useful in some cases. The tasks are similar and they used to share most of the code. The main difference is in environments. In main notebook it can take some 2 hours for the agent to start improving so it seems reasonable to launch the algorithm on a simpler env first. In debug one it is CartPole and it will train in several minutes.\n",
    "\n",
    "**About evaluation:** All points are only given for the main notebook.\n",
    "\n",
    "**Plan and evaluation points:**\n",
    "1. Getting to know the environment: most of the code is implemented for you\n",
    "2. DQN as it is (10 points): the main part of DQN implementation\n",
    "3. Main Loop (3 points): the training loop itself. Please, note, it can be really time-consuming, and implementation bugs can arise.\n",
    "4. Interpretation (2 points): calculation of episode statistics and their interpretation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "IVo0UxTWsoT_"
   },
   "outputs": [],
   "source": [
    "import sys, os\n",
    "if 'google.colab' in sys.modules and not os.path.exists('.setup_complete'):\n",
    "    os.makedirs('dqn', exist_ok=True)\n",
    "    os.makedirs('test_td_loss', exist_ok=True)\n",
    "\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/dqn/atari_wrappers.py -P dqn/\n",
    "    !wget -q https://raw.githubusercontent.com/yandexdataschool/Practical_RL/master/week04_approx_rl/dqn/utils.py -P dqn/\n",
    "    !wget -q https://raw.githubusercontent.com/yandexdataschool/Practical_RL/master/week04_approx_rl/dqn/replay_buffer.py -P dqn/\n",
    "    !wget -q https://raw.githubusercontent.com/yandexdataschool/Practical_RL/master/week04_approx_rl/dqn/framebuffer.py -P dqn/\n",
    "    !wget -q https://raw.githubusercontent.com/yandexdataschool/Practical_RL/master/week04_approx_rl/dqn/analysis.py -P dqn/\n",
    "    !wget -q https://raw.githubusercontent.com/yandexdataschool/Practical_RL/master/week04_approx_rl/dqn/logger.py -P dqn/\n",
    "    !wget -q https://raw.githubusercontent.com/yandexdataschool/Practical_RL/master/week04_approx_rl/test_td_loss/compute_td_loss.py -P test_td_loss/\n",
    "    !wget -q https://raw.githubusercontent.com/yandexdataschool/Practical_RL/master/week04_approx_rl/requirements.txt\n",
    "\n",
    "    !pip install -r requirements.txt\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": {
    "id": "KkrBeP7YsoUA"
   },
   "source": [
    "__Frameworks__ - we'll accept this homework in any deep learning framework. This particular notebook was designed for PyTorch, but you find it easy to adapt it to almost any Python-based deep learning framework."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "P8WoWe9DsoUA"
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import gymnasium as gym\n",
    "import ale_py\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "gym.register_envs(ale_py)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6GQBgViKsoUA"
   },
   "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 have tried before."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**These are various versions of Breakout provided by Gynmasium:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_names = list(gym.envs.registry.keys())\n",
    "names_breakout = [name for name in all_names if \"Break\" in name]\n",
    "names_breakout"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "S_zvw_31soUA"
   },
   "outputs": [],
   "source": [
    "# ENV_NAME = \"BreakoutDeterministic-v4\"  # this one is simpler\n",
    "ENV_NAME = \"ALE/Breakout-v5\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you are curious about Atari environments in Gymnasium, please, refer to:\n",
    "- [4] Gymnasium docs: https://gymnasium.farama.org/environments/atari/\n",
    "- [5] Longer paper: https://arxiv.org/abs/1709.06009\n",
    "- [6] Shorter paper: https://www.ijcai.org/Proceedings/2018/0787.pdf\n",
    "\n",
    "For now it's enough to know about the v5 environments:\n",
    "- v5 environments are recommended for use\n",
    "- frame_skip=5, every 5-th frame is shown to the agent and the chosen action is executed for the next 5 moves\n",
    "- randomness comes from repeat_action_probability=0.25: with this probability the previous action is executed instead of the chosen action"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xT9BvasNsoUA"
   },
   "source": [
    "## Getting to know the Environment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iwN8jA0OsoUA"
   },
   "source": [
    "**Let's see what observations look like.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = gym.make(ENV_NAME, render_mode=\"rgb_array\")\n",
    "env.reset()\n",
    "plt.imshow(env.render())\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Some more observations, coming from taking random actions**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rUZHU2HdsoUB"
   },
   "outputs": [],
   "source": [
    "env = gym.make(ENV_NAME, render_mode=\"rgb_array\")\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())\n",
    "        env.step(env.action_space.sample())\n",
    "plt.show()"
   ]
  },
  {
   "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": "markdown",
   "metadata": {},
   "source": [
    "**Let's take a look at action meanings:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.unwrapped.get_action_meanings()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1. NOOP stands for the action of doing nothing\n",
    "2. Right and Left move the platform to the corresponding direction\n",
    "3. Fire releases the ball in the beginning of a life\n",
    "\n",
    "In this assignment we will wrap the environment to execute the \"Fire\" action in the beginning of a life automatically. It will turn the \"FIRE\" action into another \"NOOP\".  \n",
    "Also, we will wrap the environment to make an episode last for 1 life instead of 5 ones.  \n",
    "These transforms are claimed as non-recommended in the paper [5] but it was done in the original paper [1] and it will help the training to converge faster. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hLNt1fbnsoUB"
   },
   "source": [
    "**Let's play a little.**\n",
    "\n",
    "Pay attention to zoom and fps args of play function. Control: A, D, space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "WOIL47azsoUB"
   },
   "outputs": [],
   "source": [
    "# # Does not work in Colab.\n",
    "# Even on a local laptop breaks matplotlib rendering. So it's recommended to restart the notebook after playing.\n",
    "# # Use the Escape button to continue.\n",
    "\n",
    "# from gymnasium.utils.play import play\n",
    "\n",
    "# play(env=gym.make(ENV_NAME, render_mode=\"rgb_array\"), zoom=4, fps=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Wrapping the Environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_basic_env():\n",
    "    return gym.make(ENV_NAME, render_mode=\"rgb_array\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5DPrxQuXsoUB"
   },
   "source": [
    "### Processing game image\n",
    "\n",
    "Let's check the shape and the dtype of the observation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = make_basic_env()\n",
    "obs, *_ = env.reset()\n",
    "obs.shape, obs.dtype"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Raw Atari images are large, 210x160x3 by default. However, we don't need that level of detail in order to learn from them.\n",
    "\n",
    "We can thus save a lot of time by preprocessing game image, including\n",
    "* Resizing to a smaller shape, 64x64 (or 84 x 84, which is used in literature)\n",
    "* Converting to grayscale\n",
    "* Cropping irrelevant image parts (top, bottom and edges) \\[we won't do this\\]\n",
    "\n",
    "The images are of the uint8 dtype.  \n",
    "uint8 stands for the 8-bit unsigned integer type.  \n",
    "We are going to store 10^5 or 10^6 observations in memory (RAM), so let's pay attention to preserving the 8-bit type after our transforms."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def apply_gray_scale_wrap(env):\n",
    "    # With the argument values chosen as below, the gym.wrappers.AtariPreprocessing wrapper\n",
    "    # only converts images to grayscale and downsamples them the screen_size\n",
    "    env = gym.wrappers.AtariPreprocessing(\n",
    "        env,\n",
    "        noop_max=0,  # the default value 30 can be harmful with FireResetEnv and frame_skip=5\n",
    "        frame_skip=1,  # frame_skip has already been set to 5 inside the env\n",
    "        terminal_on_life_loss=False,  # we do this explicitly in the FireResetEnv wrapper\n",
    "        screen_size=84  # please use 84 (which is the standard value) or 64 (which will save some computations and memory)\n",
    "    )\n",
    "    return env\n",
    "\n",
    "\n",
    "env = make_basic_env()\n",
    "env = apply_gray_scale_wrap(env)\n",
    "\n",
    "obs, *_ = env.reset()\n",
    "\n",
    "assert obs.dtype == np.dtype('uint8'), obs_dtype\n",
    "\n",
    "print(obs.shape, obs.dtype)\n",
    "plt.imshow(obs)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Atari specific wrappers\n",
    "\n",
    "We try to make our lives a little easier with the following wrappers:\n",
    "1. EpisodicLifeEnv it makes the signal that dropping a ball is not good more explicit\n",
    "2. Using FireResetEnv: with it the agent doesn't have to perform a special action to fire the ball in the beginning of a life"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dqn.atari_wrappers import FireResetEnv\n",
    "from dqn.atari_wrappers import EpisodicLifeEnv\n",
    "\n",
    "def apply_atary_specific_wrap(env):\n",
    "    env = EpisodicLifeEnv(env)\n",
    "    env = FireResetEnv(env)\n",
    "    return env\n",
    "\n",
    "env = make_basic_env()\n",
    "env = apply_gray_scale_wrap(env)\n",
    "env = apply_atary_specific_wrap(env)\n",
    "\n",
    "obs, *_ = env.reset()\n",
    "\n",
    "print(obs.shape, obs.dtype)\n",
    "plt.imshow(obs)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### FrameStack\n",
    "To make the game playable from a single observation (note the direction of the ball), we stack 4 consecutive frames:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N_FRAMES_STACKED = 4\n",
    "\n",
    "def make_final_env(apply_frame_stack=True):\n",
    "    \"\"\"\n",
    "    Builds the environment with all the wrappers applied.\n",
    "    The environment is meant be used directly as an RL algorithm input.\n",
    "\n",
    "    apply_frame_stack=False can be useful for vecotrized environments, which are not required for this assignment.\n",
    "    \"\"\"\n",
    "    env = make_basic_env()\n",
    "    env = apply_gray_scale_wrap(env)\n",
    "    env = apply_atary_specific_wrap(env)\n",
    "    if apply_frame_stack:\n",
    "        env = # your code. Please, use gym.wrappers.FrameStackObservation\n",
    "    return env\n",
    "\n",
    "\n",
    "env = make_final_env()\n",
    "\n",
    "obs, *_ = env.reset()\n",
    "print(f\"Shape: {obs.shape}, dtype: {obs.dtype}, Python object type: {type(obs)}\")\n",
    "for _ in range(N_FRAMES_STACKED - 1):\n",
    "    obs, *_ = env.step(env.action_space.sample())\n",
    "print()\n",
    "\n",
    "\n",
    "print(\"Frames, left to right: from older to more recent. The ball is dropping.\")\n",
    "_, axes = plt.subplots(figsize=(len(obs) * 3, 4), ncols=len(obs))\n",
    "for ax, frame in zip(axes, obs):\n",
    "    ax.imshow(frame)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**That is the final version of the environment we are going to do RL on.**\n",
    "\n",
    "**Let's discuss the representation of an obsevation.**\\\n",
    "An observation is a 4-frame-stack of grayscale images with reduced resolution.\\\n",
    "Memory (RAM) is a high-demand resource in this task. That's why:\n",
    "1. We use the uint8 dtype instead of float32 the neural network will operate on\n",
    "2. We don't represent them as numpy.ndarrays. **LazyFrames** are used by gym.wrappers.FrameStack instead. 2 consecutive observations share 3 of 4 frames. LazyFrames make use of this fact to save memory.\n",
    "When we feed the observations to neural networks, we should remember to scale them to the \\[-1, 1\\] range. We'll implement scaling as the first layer of a neural network, but that'll be later."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**The ball is dropping, but its hard to notice. Let's define a function to render more human-readable images:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def merge_frame_stack_to_plot(frame_stack_obs: np.ndarray | gym.wrappers.frame_stack.LazyFrames):\n",
    "    \"\"\"\n",
    "    A helper function to plot a frame stack as a single human-interpretable image.\n",
    "\n",
    "    Brighter pixels are more recent, pale pixels are older.\n",
    "    Motions goes from pale to bright.\n",
    "\n",
    "    Note! This function is designed for human vision convenience and it is NOT supposed to be used as part of\n",
    "    data preprocessing for the Reinforcement Learning agent.\n",
    "    \"\"\"\n",
    "    weights = np.ones(frame_stack_obs.shape[0], dtype=float)\n",
    "    weights[-1] += weights.sum()\n",
    "    weights /= weights.sum()\n",
    "    result = (weights[:, None, None] * frame_stack_obs).sum(0)\n",
    "    return result\n",
    "\n",
    "\n",
    "obs_joint = merge_frame_stack_to_plot(obs)\n",
    "plt.imshow(obs_joint)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Hope that's better"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N_ACTIONS = env.action_space.n\n",
    "STATE_SHAPE = env.observation_space.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4iJM3IAwsoUB"
   },
   "source": [
    "**Let's see if the game is still playable after applying the wrappers.**\n",
    "At playing the EpisodicLifeEnv wrapper seems not to work but actually it does (because after when life finishes a new ball is dropped automatically - it means that FireResetEnv wrapper understands that a new episode began).\n",
    "\n",
    "**Not supported for now.**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jhiOKsQvsoUC"
   },
   "source": [
    "## DQN as it is (10 pts)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aspwJFiGsoUC"
   },
   "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 you can find a couple of examples on diagrams below."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "YbZIucfksoUC"
   },
   "source": [
    "**Dueling network:**\\\n",
    "Paper: https://arxiv.org/pdf/1511.06581.pdf\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 advantage stream $A_\\psi$; and $\\theta = \\{\\xi, \\eta, \\psi\\}$ is their concatenation.\n",
    "\n",
    "This is what it looks like:\n",
    "\n",
    "Simple, expects height=width=64\n",
    "\n",
    "<!-- <img src=\"img/dueling_basic.png\" alt=\"dueling_basic\" width=\"500\"/> -->\n",
    "![img](https://github.com/yandexdataschool/Practical_RL/raw/master/week04_approx_rl/img/dueling_basic.png)\n",
    "\n",
    "Nature DQN ([2]), expects height=width=84\n",
    "\n",
    "<!-- <img src=\"img/dueling_nature.png\" alt=\"dueling_nature\" width=\"500\"/> -->\n",
    "![img](https://github.com/yandexdataschool/Practical_RL/raw/master/week04_approx_rl/img/dueling_nature.png)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "SPPmY6wIsoUC",
    "outputId": "717e2355-008e-4994-b5f2-1c8cf98ac445"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "device"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These constants will be useful"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N_ACTIONS, N_FRAMES_STACKED"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ConvBackbone(nn.Sequential):\n",
    "    \"\"\"\n",
    "    The convolutional part of a DQN model.\n",
    "    Please, don't think about input scaling here: it will be implemented below.\n",
    "    \"\"\"\n",
    "    def __init__(self, c_in: int = N_FRAMES_STACKED) -> None:\n",
    "        super().__init__(\n",
    "            nn.Conv2d(...),\n",
    "            ...\n",
    "            # your code,\n",
    "            nn.Flatten(),\n",
    "        )\n",
    "\n",
    "\n",
    "class DuelingDqnHead(nn.Module):\n",
    "    \"\"\"\n",
    "    Implenets the Dueling DQN logic.\n",
    "    Please, don't think about gradient scaling here (if you know what it is about): it will be implemented below.\n",
    "    \"\"\"\n",
    "    def __init__(self, n_actions, inp_size=64 * 7 * 7, hidden_size=512) -> None:\n",
    "        super().__init__()\n",
    "        self.adv_stream = nn.Sequential(\n",
    "            # your code\n",
    "        )\n",
    "        self.value_stream = nn.Sequential(\n",
    "            # your code\n",
    "        )\n",
    "\n",
    "    def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
    "        assert x.ndim == 2, x.shape  # (batch_size, n_features)\n",
    "        # your code\n",
    "        # When calculating the mean advantage, please, remember, x is a batched input!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's make a simple test for the network architecture:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def test_network_part_shapes(backbone, head):\n",
    "    batch_size = 3\n",
    "    env = make_final_env()\n",
    "    s, _ = env.reset()\n",
    "    inp = torch.rand(batch_size, *s.shape)\n",
    "\n",
    "    features = backbone(inp)\n",
    "    qvalues = head(features)\n",
    "\n",
    "    assert features.ndim == 2, features.shape\n",
    "    assert features.shape[0] == batch_size, features.shape\n",
    "    \n",
    "    assert qvalues.ndim == 2, qvalues.shape\n",
    "    assert qvalues.shape[0] == batch_size, qvalues.shape\n",
    "    assert qvalues.shape[1] == N_ACTIONS, qvalues.shape\n",
    "\n",
    "    print(\"Test passed!\")\n",
    "\n",
    "test_network_part_shapes(\n",
    "    backbone=ConvBackbone(N_FRAMES_STACKED),\n",
    "    head=DuelingDqnHead(N_ACTIONS),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Now let's build a full model.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MAX_UINT_8 = 2 ** 8 - 1\n",
    "\n",
    "\n",
    "class InputScaler(nn.Module):\n",
    "    def __init__(self, mult=1 / MAX_UINT_8):\n",
    "        super().__init__()\n",
    "        self.mult = mult\n",
    "\n",
    "    def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
    "        return x * self.mult\n",
    "\n",
    "\n",
    "class GradScalerFunctional(torch.autograd.Function):\n",
    "    \"\"\"\n",
    "    A torch.autograd.Function works as Identity on forward pass\n",
    "    and scales the gradient by scale_factor on backward pass.\n",
    "    \"\"\"\n",
    "    @staticmethod\n",
    "    def forward(ctx, input, scale_factor):\n",
    "        ctx.scale_factor = scale_factor\n",
    "        return input\n",
    "\n",
    "    @staticmethod\n",
    "    def backward(ctx, grad_output):\n",
    "        scale_factor = ctx.scale_factor\n",
    "        grad_input = grad_output * scale_factor\n",
    "        return grad_input, None\n",
    "\n",
    "\n",
    "class GradScaler(nn.Module):\n",
    "    \"\"\"\n",
    "    An nn.Module incapsulating GradScalerFunctional\n",
    "    \"\"\"\n",
    "    def __init__(self, scale_factor: float):\n",
    "        super().__init__()\n",
    "        self.scale_factor = scale_factor\n",
    "\n",
    "    def forward(self, x):\n",
    "        return GradScalerFunctional.apply(x, self.scale_factor)\n",
    "\n",
    "\n",
    "class DQNetworkDueling(nn.Sequential):\n",
    "    def __init__(self, c_in: int, n_actions: int) -> None:\n",
    "        input_scaler = InputScaler()  # the inputs come from the uint8 range\n",
    "        backbone = ConvBackbone(c_in=c_in)  # your code\n",
    "        grad_scaler = GradScaler(1 / 2**0.5)  # Dueling DQN suggests do scale the gradient by 1 / sqrt(2)\n",
    "        head = DuelingDqnHead(n_actions=n_actions)\n",
    "        super().__init__(input_scaler, backbone, grad_scaler, head)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def test_network_shapes(model):\n",
    "    batch_size = 3\n",
    "    env = make_final_env()\n",
    "    s, _ = env.reset()\n",
    "    inp = torch.rand(batch_size, *s.shape)\n",
    "\n",
    "    qvalues = model(inp)\n",
    "    \n",
    "    assert qvalues.ndim == 2, qvalues.shape\n",
    "    assert qvalues.shape[0] == batch_size, qvalues.shape\n",
    "    assert qvalues.shape[1] == N_ACTIONS, qvalues.shape\n",
    "\n",
    "    print(\"Test passed!\")\n",
    "\n",
    "test_network_shapes(model=DQNetworkDueling(N_FRAMES_STACKED, N_ACTIONS))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Now let's wrap our model into an Agent class.**  \n",
    "It will implement epsilon-greedy policy on numpy ndarrays."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DQNAgent(nn.Module):\n",
    "    \"\"\"\n",
    "    Epsilon-greedy policy with a torch.nn.Module Q-value estimator.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, q_network: nn.Module, epsilon=1) -> None:\n",
    "        super().__init__()\n",
    "        self.epsilon = epsilon\n",
    "        self.q_network = q_network\n",
    "\n",
    "    def forward(self, state_t):\n",
    "        \"\"\"\n",
    "        takes agent's observation (tensor), returns qvalues (tensor)\n",
    "        :param state_t: a batch of 4-frame buffers, shape = [batch_size, 4, h, w]\n",
    "        \"\"\"\n",
    "        # Use your network to compute qvalues for given state\n",
    "        qvalues =  # your code\n",
    "        return qvalues\n",
    "\n",
    "    @torch.no_grad()  # we don't need autograd here, so let's save the computations\n",
    "    def get_qvalues(self, states: np.ndarray) -> np.ndarray:\n",
    "        \"\"\"\n",
    "        like forward, but works on numpy arrays, not tensors\n",
    "        \"\"\"\n",
    "        model_device = next(self.parameters()).device\n",
    "        states_pt = torch.tensor(\n",
    "            np.array(states), device=model_device, dtype=torch.float32\n",
    "        )\n",
    "        # Use your network to compute qvalues for given state\n",
    "        qvalues_pt = # your code\n",
    "        qvalues = qvalues_pt.data.cpu().numpy()\n",
    "        return qvalues\n",
    "\n",
    "    def sample_actions_by_qvalues(self, qvalues: np.ndarray, greedy: bool = False) -> np.ndarray:\n",
    "        \"\"\"pick actions given qvalues. Uses epsilon-greedy exploration strategy.\"\"\"\n",
    "        batch_size, n_actions = qvalues.shape\n",
    "        # greedy_actions =  # your code\n",
    "        greedy_actions = qvalues.argmax(axis=-1)  # your code\n",
    "        if greedy:\n",
    "            return greedy_actions\n",
    "\n",
    "        random_actions = np.random.randint(\"your code\")\n",
    "        should_explore = np.random.binomial(\"your code\")\n",
    "        epsilon_greedy_actions = np.where(\n",
    "            \"your code\"\n",
    "        )\n",
    "        return epsilon_greedy_actions\n",
    "\n",
    "    def sample_actions(self, states: np.ndarray) -> np.ndarray:\n",
    "        qvalues = self.get_qvalues(states)\n",
    "        actions = self.sample_actions_by_qvalues(qvalues)\n",
    "        return actions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_network_shapes(\n",
    "    model=DQNAgent(DQNetworkDueling(N_FRAMES_STACKED, N_ACTIONS))\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "BUFMLKX1soUC"
   },
   "outputs": [],
   "source": [
    "agent = DQNAgent(\n",
    "    DQNetworkDueling(N_FRAMES_STACKED, N_ACTIONS),\n",
    "    epsilon=0.5\n",
    ").to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XbsIT2EdsoUC"
   },
   "source": [
    "Now let's try out our agent to see if it raises any errors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "pZR3qE2esoUC"
   },
   "outputs": [],
   "source": [
    "def evaluate(env, agent, n_games=1, greedy=False, t_max=10000, seed=None):\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(seed=seed)\n",
    "        reward = 0\n",
    "        for _ in range(t_max):\n",
    "            action = agent.sample_actions(np.array(s)[None], greedy=greedy)[0]\n",
    "            s, r, terminated, truncated, _ = env.step(action)\n",
    "            reward += r\n",
    "            if terminated or truncated:\n",
    "                break\n",
    "\n",
    "        rewards.append(reward)\n",
    "    return np.mean(rewards)\n",
    "\n",
    "print(evaluate(env, agent, n_games=1, greedy=False))\n",
    "print(evaluate(env, agent, n_games=1, greedy=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2BiRixA-soUC"
   },
   "source": [
    "### Experience replay\n",
    "For this assignment, we provide you with experience replay buffer.\n",
    "\n",
    "<!-- <img src=\"img/exp_replay.png\" alt=\"exp_replay\" width=\"500\"/> -->\n",
    "![img](https://github.com/yandexdataschool/Practical_RL/raw/master/yet_another_week/_resource/exp_replay.png)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jTBZo5BVsoUC"
   },
   "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": {
    "id": "Ydi0KK9LsoUC"
   },
   "outputs": [],
   "source": [
    "from dqn.replay_buffer import ReplayBuffer\n",
    "exp_replay = ReplayBuffer(10)\n",
    "\n",
    "for _ in range(30):\n",
    "    exp_replay.add(env.reset()[0], env.action_space.sample(), 1.0, env.reset()[0], 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": "markdown",
   "metadata": {},
   "source": [
    "**The `play_and_record` function, defined below is the main way the agent will interact with the environment during training.**\n",
    "\n",
    "Previously we used to train RL algorithms on entire episodes.  \n",
    "\n",
    "This time we keep the environment constantly running and will be getting small portions of interactions with it.\n",
    "\n",
    "The agent takes several actions (4 actions in [2] and [3]), the corresponding (s, a, r, s', terminated) tuples are put into the replay buffer.  \n",
    "Whenever an episode finishes (i.e. `truncated or terminated`), the environment is reset and the procedure continues as usually.  \n",
    "\n",
    "To make the first step in a constantly running environment, the agent needs to know the state of the environment. This is the meaning of the `initial_state` argument of the function.\n",
    "\n",
    "It's worth noting, the agent does not train on the fresh tuples immediately. The agent trains on samples which are sampled from the buffer.\n",
    "\n",
    "**Implementation note:**\n",
    "We define an `ActionSampler` protocol. The goal of it is to let the function `play_and_record` accept not only `DQNAgent` class instances, but any object that can sample actions. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "cEXv69KWsoUC"
   },
   "outputs": [],
   "source": [
    "from typing import Protocol\n",
    "\n",
    "\n",
    "class ActionSampler(Protocol):\n",
    "    \"\"\"\n",
    "    A protocol which defines an Callable which samples actions from states\n",
    "    \"\"\"\n",
    "\n",
    "    def __call__(\n",
    "        self, state: gym.wrappers.stateful_observation.FrameStackObservation\n",
    "    ) -> int: ...\n",
    "\n",
    "\n",
    "class RandomActionSampler:\n",
    "    \"\"\"\n",
    "    We will need this guy to fill the buffer with initial 50-200K observations from a random policy.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, action_space) -> None:\n",
    "        self.action_space = action_space\n",
    "\n",
    "    def __call__(\n",
    "        self, state: gym.wrappers.stateful_observation.FrameStackObservation\n",
    "    ) -> int:\n",
    "        action = self.action_space.sample()\n",
    "        return action\n",
    "\n",
    "\n",
    "class DqnActionSampler:\n",
    "    \"\"\"\n",
    "    DQNAgent works on batched np.ndarray inputs.\n",
    "    This class uses a DQNAgent to sample actions from single LazyFrames observations.\n",
    "\n",
    "    This will be an epsilon-greedy sampler.\n",
    "    A greedy sampler can be defined as well, but we won't need it.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, agent: DQNAgent):\n",
    "        self.agent = agent\n",
    "\n",
    "    def __call__(\n",
    "        self, state: gym.wrappers.stateful_observation.FrameStackObservation\n",
    "    ) -> int:\n",
    "        state_batched = np.array(state)[None]\n",
    "        action_batched = # <YOUR CODE>\n",
    "        action = action_batched.item()\n",
    "        return action\n",
    "\n",
    "\n",
    "@torch.no_grad()\n",
    "def play_and_record(\n",
    "    initial_state: gym.wrappers.stateful_observation.FrameStackObservation,\n",
    "    action_sampler: ActionSampler,\n",
    "    env,\n",
    "    exp_replay,\n",
    "    n_steps=1,\n",
    "):\n",
    "    \"\"\"\n",
    "    Play the game for exactly n_steps, record every (s,a,r,s', done) to replay buffer.\n",
    "    Whenever game ends due to termination or truncation, add record with done=terminated and reset the game.\n",
    "    It is guaranteed that env has terminated=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",
    "    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\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_play_and_record(action_sampler):\n",
    "    exp_replay = ReplayBuffer(10_000)\n",
    "\n",
    "    state, _ = env.reset()\n",
    "    \n",
    "    # action_sampler = RandomActionSampler(env.action_space)\n",
    "    # action_sampler = DqnActionSampler(agent)\n",
    "    play_and_record(state, action_sampler, 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, \\\n",
    "        \"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, \\\n",
    "        \"Please make sure you restart the game whenever it is 'done' and \" \\\n",
    "        \"record the is_done correctly into the buffer. Got %f is_done rate over \" \\\n",
    "        \"%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(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], \\\n",
    "            \"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!\")\n",
    "\n",
    "\n",
    "print(\"Random:\")\n",
    "test_play_and_record(RandomActionSampler(env.action_space))\n",
    "print(\"DQN:\")\n",
    "test_play_and_record(DqnActionSampler(agent))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "y5zyryPOsoUF"
   },
   "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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "target_network = DQNetworkDueling(N_FRAMES_STACKED, N_ACTIONS).to(device)\n",
    "target_network.load_state_dict(agent.q_network.state_dict())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2idY8QX0soUF"
   },
   "source": [
    "### Learning with... Q-learning\n",
    "Here we write a function similar to `agent.update` from tabular q-learning."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "k46MPwwwsoUF"
   },
   "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 major source of of bugs in this homework. We tried to cover it with tests, but 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**\n",
    "\n",
    "$$ Q_{reference}(s,a) = r(s, a) + \\gamma \\cdot\n",
    "Q_{target}(s',argmax_{a'}Q_\\theta(s', a')) $$\n",
    "\n",
    "We will use Double DQN for training, but **we ask you to implement both** of the methods to experience the difference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "V02HcUYasoUG",
    "outputId": "7a11e3d3-d030-40be-8f14-59b5481749fb"
   },
   "outputs": [],
   "source": [
    "def compute_td_loss_on_tensors(\n",
    "    states: torch.Tensor,  # (batch_size, *state_shape)\n",
    "    actions: torch.Tensor,  # (batch_size,)\n",
    "    rewards: torch.Tensor,  # (batch_size,)\n",
    "    next_states: torch.Tensor,  # (batch_size, *state_shape)\n",
    "    is_done: torch.Tensor,  # (batch_size,), torch.bool\n",
    "    agent: nn.Module,\n",
    "    target_network: nn.Module,\n",
    "    gamma: float = 0.99,\n",
    "    check_shapes=False,\n",
    "):\n",
    "    predicted_qvalues = agent(states)  # shape: [batch_size, n_actions]\n",
    "    assert is_done.dtype is torch.bool\n",
    "\n",
    "    # compute q-values for all actions in next states\n",
    "    with torch.no_grad():\n",
    "        predicted_next_qvalues_target = # your code;  # shape: [batch_size, n_actions]\n",
    "\n",
    "    # select q-values for chosen actions\n",
    "    predicted_qvalues_for_actions = predicted_qvalues[\n",
    "        range(len(actions)), actions\n",
    "    ]  # shape: [batch_size]\n",
    "\n",
    "    # compute V*(next_states) using predicted next q-values\n",
    "    next_state_values =  # your code\n",
    "\n",
    "    if check_shapes:\n",
    "        assert (\n",
    "            next_state_values.dim() == 1\n",
    "            and next_state_values.shape[0] == states.shape[0]\n",
    "        ), \"must predict one value per state\"\n",
    "        assert not next_state_values.requires_grad\n",
    "\n",
    "    # compute \"target q-values\" for loss - it's what's inside square parentheses in the above formula.\n",
    "    # at the last state use the simplified formula: Q(s,a) = r(s,a) since s' doesn't exist\n",
    "    target_qvalues_for_actions =  # your code\n",
    "\n",
    "    # mean squared error loss to minimize\n",
    "    loss = torch.mean((predicted_qvalues_for_actions - target_qvalues_for_actions) ** 2)\n",
    "\n",
    "    if check_shapes:\n",
    "        assert (\n",
    "            predicted_next_qvalues_target.data.dim() == 2\n",
    "        ), \"make sure you predicted q-values for all actions in next state\"\n",
    "        assert (\n",
    "            next_state_values.data.dim() == 1\n",
    "        ), \"make sure you computed V(s') as maximum over just the actions axis and not all axes\"\n",
    "        assert (\n",
    "            target_qvalues_for_actions.data.dim() == 1\n",
    "        ), \"there's something wrong with target q-values, they must be a vector\"\n",
    "\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from test_td_loss.compute_td_loss import test_is_done_is_used, test_compute_td_loss_vanilla\n",
    "\n",
    "test_compute_td_loss_vanilla(compute_td_loss_on_tensors)\n",
    "print(\"Well done!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_td_loss_on_tensors_double(\n",
    "    states: torch.Tensor,  # (batch_size, *state_shape)\n",
    "    actions: torch.Tensor,  # (batch_size,)\n",
    "    rewards: torch.Tensor,  # (batch_size,)\n",
    "    next_states: torch.Tensor,  # (batch_size, *state_shape)\n",
    "    is_done: torch.Tensor,  # (batch_size,), torch.bool\n",
    "    agent: nn.Module,\n",
    "    target_network: nn.Module,\n",
    "    gamma: float = 0.99,\n",
    "    check_shapes=False,\n",
    "):\n",
    "    predicted_qvalues = agent(states)  # shape: [batch_size, n_actions]\n",
    "    assert is_done.dtype is torch.bool\n",
    "\n",
    "    # compute q-values for all actions in next states\n",
    "    with torch.no_grad():\n",
    "        predicted_next_qvalues_target = # your code;  # shape: [batch_size, n_actions]\n",
    "\n",
    "    # select q-values for chosen actions\n",
    "    predicted_qvalues_for_actions = predicted_qvalues[\n",
    "        range(len(actions)), actions\n",
    "    ]  # shape: [batch_size]\n",
    "\n",
    "    # compute V*(next_states) using predicted next q-values\n",
    "    next_state_values =  # your code\n",
    "\n",
    "    if check_shapes:\n",
    "        assert (\n",
    "            next_state_values.dim() == 1\n",
    "            and next_state_values.shape[0] == states.shape[0]\n",
    "        ), \"must predict one value per state\"\n",
    "        assert not next_state_values.requires_grad\n",
    "\n",
    "    # compute \"target q-values\" for loss - it's what's inside square parentheses in the above formula.\n",
    "    # at the last state use the simplified formula: Q(s,a) = r(s,a) since s' doesn't exist\n",
    "    target_qvalues_for_actions =  # your code\n",
    "\n",
    "    # mean squared error loss to minimize\n",
    "    loss = torch.mean((predicted_qvalues_for_actions - target_qvalues_for_actions) ** 2)\n",
    "\n",
    "    if check_shapes:\n",
    "        assert (\n",
    "            predicted_next_qvalues_target.data.dim() == 2\n",
    "        ), \"make sure you predicted q-values for all actions in next state\"\n",
    "        assert (\n",
    "            next_state_values.data.dim() == 1\n",
    "        ), \"make sure you computed V(s') as maximum over just the actions axis and not all axes\"\n",
    "        assert (\n",
    "            target_qvalues_for_actions.data.dim() == 1\n",
    "        ), \"there's something wrong with target q-values, they must be a vector\"\n",
    "\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from test_td_loss.compute_td_loss import test_compute_td_loss_double\n",
    "\n",
    "test_compute_td_loss_double(compute_td_loss_on_tensors_double)\n",
    "print(\"Well done!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**The following function works on np.ndarrays: it converts its inputs to torch.Tensors and calls the torch-tensor function**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_td_loss(\n",
    "    states,\n",
    "    actions,\n",
    "    rewards,\n",
    "    next_states,\n",
    "    is_done,\n",
    "    agent,\n",
    "    target_network,\n",
    "    gamma=0.99,\n",
    "    check_shapes=False,\n",
    "    device=None,\n",
    "    tensor_loss_evaluator=compute_td_loss_on_tensors_double,\n",
    "):\n",
    "    \"\"\"Compute td loss using torch operations only. Use the formulae above.\"\"\"\n",
    "\n",
    "    if device is None:\n",
    "        device = next(agent.parameters()).device\n",
    "    states = torch.tensor(\n",
    "        states, device=device, dtype=torch.float32\n",
    "    )  # shape: [batch_size, *state_shape]\n",
    "    actions = torch.tensor(\n",
    "        actions, device=device, dtype=torch.int64\n",
    "    )  # shape: [batch_size]\n",
    "    rewards = torch.tensor(\n",
    "        rewards, device=device, dtype=torch.float32\n",
    "    )  # shape: [batch_size]\n",
    "    # shape: [batch_size, *state_shape]\n",
    "    next_states = torch.tensor(next_states, device=device, dtype=torch.float)\n",
    "    is_done = torch.tensor(\n",
    "        is_done, device=device, dtype=torch.bool\n",
    "    )  # shape: [batch_size]\n",
    "\n",
    "    return tensor_loss_evaluator(\n",
    "        states=states,\n",
    "        actions=actions,\n",
    "        rewards=rewards,\n",
    "        next_states=next_states,\n",
    "        is_done=is_done,\n",
    "        agent=agent,\n",
    "        target_network=target_network,\n",
    "        gamma=gamma,\n",
    "        check_shapes=check_shapes,\n",
    "    )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "x8AvquAtsoUG"
   },
   "source": [
    "Sanity checks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5nRoOn30soUG"
   },
   "outputs": [],
   "source": [
    "obs_batch, act_batch, reward_batch, next_obs_batch, is_done_batch = exp_replay.sample(10)\n",
    "\n",
    "loss = compute_td_loss(obs_batch, act_batch, reward_batch, next_obs_batch, is_done_batch,\n",
    "                       agent, target_network,\n",
    "                       gamma=0.99, check_shapes=True)\n",
    "loss.backward()\n",
    "\n",
    "assert loss.requires_grad and tuple(loss.data.size()) == (), \\\n",
    "    \"you must return scalar loss - mean over batch\"\n",
    "assert np.any(next(agent.parameters()).grad.data.cpu().numpy() != 0), \\\n",
    "    \"loss must be differentiable w.r.t. network weights\"\n",
    "assert np.all(next(target_network.parameters()).grad is None), \\\n",
    "    \"target network should not have grads\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KIplL0hSsoUG"
   },
   "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": {
    "id": "-JV-ulB-soUG"
   },
   "outputs": [],
   "source": [
    "from tqdm.auto import trange\n",
    "from IPython.display import clear_output\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "\n",
    "seed = # your favourite seed\n",
    "random.seed(seed)\n",
    "np.random.seed(seed)\n",
    "torch.manual_seed(seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "-eurxA-_soUG",
    "outputId": "4eb73eb0-771c-442e-a086-50bae3ebd9d8"
   },
   "outputs": [],
   "source": [
    "env = make_final_env()\n",
    "\n",
    "state, _ = env.reset(seed=seed)\n",
    "\n",
    "agent = DQNAgent(\n",
    "    DQNetworkDueling(N_FRAMES_STACKED, N_ACTIONS),\n",
    "    epsilon=1\n",
    ").to(device)\n",
    "target_network = DQNetworkDueling(N_FRAMES_STACKED, N_ACTIONS).to(device)\n",
    "target_network.load_state_dict(agent.q_network.state_dict())\n",
    "\n",
    "action_sampler = DqnActionSampler(agent)\n",
    "action_sampler_random = RandomActionSampler(env.action_space)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WZg25kIasoUG"
   },
   "source": [
    "Buffer of size $10^4$ can probably pass the threshold for this assignment.\n",
    "\n",
    "Larger sizes ($10^5$ and $10^6$ are common) can show a much higher score."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "hWyMxfN4soUG",
    "outputId": "a0d4147b-56b2-4f69-802a-0da87ad82bdb",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from dqn.utils import is_enough_ram\n",
    "\n",
    "REPLAY_BUFFER_SIZE = 10**6\n",
    "# INITIAL_BUFFER_FILL = 50_000  # Nature DQN Extended Data Table 1\n",
    "INITIAL_BUFFER_FILL = 200_000  # Rainbow without prioritization\n",
    "_n_steps = 100\n",
    "\n",
    "exp_replay = ReplayBuffer(REPLAY_BUFFER_SIZE)\n",
    "for i in trange(INITIAL_BUFFER_FILL // _n_steps):\n",
    "    if not is_enough_ram(min_available_gb=0.1):\n",
    "        print(\"\"\"\n",
    "            Less than 100 Mb RAM available.\n",
    "            Make sure the buffer size in not too huge.\n",
    "            Also check, maybe other processes consume RAM heavily.\n",
    "            \"\"\"\n",
    "             )\n",
    "        break\n",
    "    play_and_record(state, action_sampler_random, env, exp_replay, n_steps=_n_steps)\n",
    "    if len(exp_replay) >= INITIAL_BUFFER_FILL:\n",
    "        break\n",
    "print(len(exp_replay))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(exp_replay)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_ca9vbW4soUG"
   },
   "outputs": [],
   "source": [
    "update_frequency = 4  # n_steps for play_and_record; Nature DQN Extended Data Table 1 + Rainbow Table 4: Additional hyper-parameters\n",
    "batch_size = 32  # Nature DQN Extended Data Table 1 + Table 4: Additional hyper-parameters\n",
    "total_steps = 10 * 10**6  # this can be long, feel free to stop the training when the target score is reached\n",
    "decay_steps = 10**6  # Nature DQN Extended Data Table 1\n",
    "\n",
    "opt = torch.optim.Adam(agent.parameters(), lr=6.25e-05, eps=1.4e-4)  # Rainbow\n",
    "\n",
    "init_epsilon = 1  # Nature DQN\n",
    "final_epsilon = 0.1  # Nature DQN\n",
    "\n",
    "loss_freq = 100\n",
    "refresh_target_network_freq = 10_000  # Nature DQN\n",
    "eval_freq = 10_000\n",
    "\n",
    "max_grad_norm = 10  # Dueling DQN\n",
    "\n",
    "n_lives = 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "oJWs0q-6soUG"
   },
   "outputs": [],
   "source": [
    "step = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "675-JU0hsoUG"
   },
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "def wait_for_keyboard_interrupt():\n",
    "    try:\n",
    "        while True:\n",
    "            time.sleep(1)\n",
    "    except KeyboardInterrupt:\n",
    "        pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dqn.utils import linear_decay, is_enough_ram\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "from dqn.logger import Logger"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "use_tensorboard = True  # Set to False to use simple matplotlib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logger = Logger(use_tensorboard=use_tensorboard)\n",
    "\n",
    "if use_tensorboard:\n",
    "    %load_ext tensorboard\n",
    "    %tensorboard --logdir runs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "FgQ1vK3CsoUG",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "env = make_final_env()\n",
    "state, _ = env.reset()\n",
    "\n",
    "with trange(step, total_steps + 1) as progress_bar:\n",
    "    for step in progress_bar:\n",
    "        if not is_enough_ram():\n",
    "            print('Less than 100 MB RAM available, freezing.')\n",
    "            print('Ensure everything is okay and use KeyboardInterrupt to continue.')\n",
    "            wait_for_keyboard_interrupt()\n",
    "\n",
    "        agent.epsilon = linear_decay(init_epsilon, final_epsilon, step, decay_steps)\n",
    "\n",
    "        # Play\n",
    "        _, state = play_and_record(state, action_sampler, env, exp_replay, n_steps=update_frequency)\n",
    "\n",
    "        # Train\n",
    "        s, a, r, s_next, done = exp_replay.sample(batch_size)\n",
    "        # <YOUR CODE: sample batch_size of data from experience replay>\n",
    "\n",
    "        loss = compute_td_loss(s, a, r, s_next, done, agent, target_network, device=device)  # Ensure 'device' is defined\n",
    "\n",
    "        loss.backward()\n",
    "        grad_norm = nn.utils.clip_grad_norm_(agent.parameters(), max_grad_norm)\n",
    "        opt.step()\n",
    "        opt.zero_grad()\n",
    "\n",
    "        if step % loss_freq == 0:\n",
    "            loss_value = loss.data.cpu().item()\n",
    "            grad_norm_value = grad_norm.cpu().item()\n",
    "            logger.log_loss(loss_value, step)\n",
    "            logger.log_grad_norm(grad_norm_value, step)\n",
    "\n",
    "        if step % refresh_target_network_freq == 0:\n",
    "            # Load agent weights into target_network\n",
    "            target_network.load_state_dict(agent.q_network.state_dict())\n",
    "            torch.save(agent.state_dict(), \"last_state_dict.pt\")\n",
    "\n",
    "        if step % eval_freq == 0:\n",
    "            mean_reward = evaluate(\n",
    "                make_final_env(), agent, n_games=3 * n_lives, greedy=True, seed=step\n",
    "            )\n",
    "            initial_state_q_values = agent.get_qvalues(\n",
    "                [make_final_env().reset(seed=step)[0]]\n",
    "            )\n",
    "            initial_v = np.max(initial_state_q_values).item()\n",
    "\n",
    "            logger.log_mean_reward(mean_reward, step)\n",
    "            logger.log_initial_state_v(initial_v, step)\n",
    "\n",
    "            clear_output(True)\n",
    "            print(\"Buffer size = %i, Epsilon = %.5f\" % (len(exp_replay), agent.epsilon))\n",
    "\n",
    "            if not use_tensorboard:\n",
    "                # If using matplotlib, plot the metrics\n",
    "                logger.plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "agent = DQNAgent(\n",
    "    DQNetworkDueling(N_FRAMES_STACKED, N_ACTIONS),\n",
    "    epsilon=1\n",
    ").to(device)\n",
    "agent.load_state_dict(torch.load(\"last_state_dict.pt\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZEDQhQrdsoUG"
   },
   "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": {
    "id": "s0jLjYGwsoUG"
   },
   "source": [
    "Final scoring is done on a whole episode with all 5 lives."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "xTGVrwwQsoUG"
   },
   "outputs": [],
   "source": [
    "final_score = evaluate(\n",
    "  make_final_env(),\n",
    "    agent, n_games=30, greedy=True, t_max=10 * 1000, seed=9\n",
    ")\n",
    "print('final score:', final_score)\n",
    "assert final_score >= 3, 'not as cool as DQN can'\n",
    "print('Cool!')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ovaG8N4lsoUH"
   },
   "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",
    "* **grad norm** just shows the intensivity of training. Not ok is growing to values of about 100 (or maybe even 50) though it depends on network architecture.\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",
    "![img](https://github.com/yandexdataschool/Practical_RL/raw/master/yet_another_week/_resource/training.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kVV72AB-soUH"
   },
   "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 it was generated by. 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",
    "* target network update frequency: has something in common with learning rate. Too frequent updates can lead to divergence. Too rare can lead to slow leraning. 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": {
    "id": "Plp8WC_esoUH"
   },
   "source": [
    "### Video"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "DdExc_AssoUH"
   },
   "outputs": [],
   "source": [
    "# record sessions\n",
    "from gymnasium.wrappers import RecordVideo\n",
    "\n",
    "with make_final_env() as env, RecordVideo(\n",
    "    env=env, video_folder=\"./videos\", episode_trigger=lambda episode_number: True\n",
    ") as env_monitor:\n",
    "    sessions = [\n",
    "        evaluate(env_monitor, agent, n_games=n_lives, greedy=True) for _ in range(10)\n",
    "    ]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lt6xg1n_soUH"
   },
   "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": {
    "id": "fLPx2aI7soUH"
   },
   "source": [
    "## Let's have a closer look at this. Interpretation (2 pts).\n",
    "\n",
    "Let's play 5 episodes (note that the game has 5 lives) and log some stats:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dqn.analysis import play_and_log_episode\n",
    "\n",
    "env = make_final_env()\n",
    "stats = play_and_log_episode(env, agent)\n",
    "\n",
    "print(\"Keys:\", list(stats.keys()))\n",
    "print(\"Shapes:\")\n",
    "for key in [\"states\", \"qvalues\", \"actions\", \"rewards\"]:\n",
    "    print(f\"{key}: {stats[key].shape}\")\n",
    "print(\"terminated:\", stats[\"episode_finished\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's plot rewards:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(stats[\"rewards\"])\n",
    "plt.grid()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Your task it to evaluate the following quantities from the logs:\n",
    "1. Discounted returns: $G[t] = \\sum_{t'=t}^T \\gamma ^ {t' - t}r[t]$, where $T$ is the total time of an episode.\n",
    "2. State Values estimated by the agent: $V_{agent}[t] = \\max_{a}Q_{agent}(s[t], a)$. \n",
    "3. Q-spread: $\\Delta Q[t] = \\max_{a}Q_{agent}(s[t], a) - \\min_{a}Q_{agent}(s[t], a)$\n",
    "\n",
    "Create a new env: `env = make_final_env()`, play for 5 episodes (a full game has 5 lives, so it will be 1 full game).  \n",
    "Plot rewards and the evaluated quantites for each of them.  \n",
    "Using the plots, can you find points where the ball hits the wall?  \n",
    "Where the ball hits the platform?  \n",
    "Probably, you won't need all the quantities for that, but it's still good to check the behavior of a model.\n",
    "\n",
    "The `merge_frame_stack_to_plot` function may be useful."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_discounted_returns(rewards: Reversible[float], gamma: float) -> list[float]:\n",
    "    \"\"\"\n",
    "    Calculates G[t] for each t, given rewards and gamma.\n",
    "\n",
    "    Tip: Iterate backwards through rewards and use the following relationship:\n",
    "    G[t] = r[t] + gamma * G[t + 1]\n",
    "    \"\"\"\n",
    "    returns_reverse: list[float] = []\n",
    "    last_return: float = 0\n",
    "    for r in reversed(rewards):\n",
    "        # your code\n",
    "    returns = list(reversed(returns_reverse))\n",
    "    return returns\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "e36bU0u8soUH"
   },
   "source": [
    "## Bonus I (2 pts). Get High Score!\n",
    "\n",
    "Get mean reward per life:\n",
    "1. $\\geq 20$: +1 pt\n",
    "2. $\\geq 40$: +2 pts"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "78e4nRoSsoUH"
   },
   "source": [
    "## Bonus II (2+ pts). Apply modifications to DQN.\n",
    "\n",
    "* [Prioritized experience replay](https://arxiv.org/abs/1511.05952) (2 pts, please check for effective open-source implementations)\n",
    "* [Noisy Nets](https://arxiv.org/abs/1706.10295) (2 pts, please, pay attention to the exploration policy and whether epsilon-greedy policy is used)\n",
    "* [distributional RL](https://arxiv.org/abs/1707.06887)(distributional and distributed stand for different things here) (3 pts)\n",
    "* Other modifications (2+ pts depending on complexity)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "j_3RXboysoUH"
   },
   "source": [
    "## Bonus III (2 pts). Distributed RL.\n",
    "\n",
    "https://gymnasium.farama.org/api/vector/#gymnasium.vector.AsyncVectorEnv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "f9X5aB56soUI"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}