reinforcement_learning.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "source": "**Chapter 18 – Reinforcement Learning**",
   "metadata": {
    "id": "WNi0HDRs-qwQ",
    "cell_id": "f54aa67f-57e4-443c-8eb4-becdde79769c",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "markdown",
   "source": "_This notebook contains all the sample code in chapter 18_.",
   "metadata": {
    "id": "51pDOJ9h-qwb",
    "cell_id": "00001-bf27a464-4f6a-49ef-9705-b0af5e524dda",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "markdown",
   "source": "# Setup\nFirst, let's import a few common modules, ensure MatplotLib plots figures inline and prepare a function to save the figures. We also check that Python 3.5 or later is installed (although Python 2.x may work, it is deprecated so we strongly recommend you use Python 3 instead), as well as Scikit-Learn >= 0.20 and TensorFlow >= 2.0.",
   "metadata": {
    "id": "scSqGWdV-qwf",
    "cell_id": "00002-d18a97c4-20a0-42be-b97a-524c1fd0c915",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 187.5625
   }
  },
  {
   "cell_type": "code",
   "source": "# Python >= 3.5 is required\nimport sys\nassert sys.version_info >= (3, 5)\n\n# Is this notebook running on Colab or Kaggle?\nIS_COLAB = \"google.colab\" in sys.modules\nIS_KAGGLE = \"kaggle_secrets\" in sys.modules\n\nif IS_COLAB or IS_KAGGLE:\n    !apt update && apt install -y libpq-dev libsdl2-dev swig xorg-dev xvfb\n    %pip install -U tf-agents pyvirtualdisplay\n    %pip install -U gym>=0.21.0\n    %pip install -U gym[box2d,atari,accept-rom-license]\n\n# Scikit-Learn >= 0.20 is required\nimport sklearn\nassert sklearn.__version__ >= \"0.20\"\n\n# TensorFlow >= 2.0 is required\nimport tensorflow as tf\nfrom tensorflow import keras\nassert tf.__version__ >= \"2.0\"\n\nif not tf.config.list_physical_devices('GPU'):\n    print(\"No GPU was detected. CNNs can be very slow without a GPU.\")\n    if IS_COLAB:\n        print(\"Go to Runtime > Change runtime and select a GPU hardware accelerator.\")\n    if IS_KAGGLE:\n        print(\"Go to Settings > Accelerator and select GPU.\")\n\n# Common imports\nimport numpy as np\nimport os\n\n# To make this notebook's output stable across runs\nnp.random.seed(42)\ntf.random.set_seed(42)\n\n# To plot pretty figures\n%matplotlib inline\nimport matplotlib as mpl\nimport matplotlib.pyplot as plt\nmpl.rc('axes', labelsize=14)\nmpl.rc('xtick', labelsize=12)\nmpl.rc('ytick', labelsize=12)\n\n# To get smooth animations\nimport matplotlib.animation as animation\nmpl.rc('animation', html='jshtml')\n\n# Where to save the figures\nPROJECT_ROOT_DIR = \".\"\nCHAPTER_ID = \"rl\"\nIMAGES_PATH = os.path.join(PROJECT_ROOT_DIR, \"images\", CHAPTER_ID)\nos.makedirs(IMAGES_PATH, exist_ok=True)\n\ndef save_fig(fig_id, tight_layout=True, fig_extension=\"png\", resolution=300):\n    path = os.path.join(IMAGES_PATH, fig_id + \".\" + fig_extension)\n    print(\"Saving figure\", fig_id)\n    if tight_layout:\n        plt.tight_layout()\n    plt.savefig(path, format=fig_extension, dpi=resolution)",
   "metadata": {
    "id": "BG4AFJ-9-qwh",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "5182c070-ac8e-4753-9997-1204d4678228",
    "cell_id": "00003-cf663e27-e67e-4b0f-8411-7fedbc3ccf56",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 1799
   },
   "execution_count": 1,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "# Introduction to OpenAI gym",
   "metadata": {
    "id": "mFjX6jft-qwo",
    "cell_id": "00004-84a466fd-bb5d-4874-a0fe-4c9ae7efa0bc",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 82
   }
  },
  {
   "cell_type": "markdown",
   "source": "In this notebook we will be using [OpenAI gym](https://gym.openai.com/), a great toolkit for developing and comparing Reinforcement Learning algorithms. It provides many environments for your learning *agents* to interact with. Let's start by importing `gym`:",
   "metadata": {
    "id": "STFy7BPG-qwp",
    "cell_id": "00005-4aa05a26-2dc5-45a0-8189-1060e01aef25",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 97.171875
   }
  },
  {
   "cell_type": "code",
   "source": "import gym",
   "metadata": {
    "id": "hKCRppT8-qwr",
    "cell_id": "00006-8f5a30bc-2a36-4396-b99e-c448286b494b",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 61
   },
   "execution_count": 2,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Let's list all the available environments:",
   "metadata": {
    "id": "deCspMkA-qws",
    "cell_id": "00007-ca7cce8e-41b3-4833-af17-4ce845fc27ee",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "code",
   "source": "gym.envs.registry.all()",
   "metadata": {
    "id": "ejSKsf1V-qwu",
    "outputId": "93f17fde-a89a-449a-9d9c-c508a6b58b86",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "cell_id": "00008-19540f33-3cda-4d47-8351-036dc4ba0dc2",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 183.5625,
    "deepnote_output_heights": [
     null,
     21.1875
    ]
   },
   "execution_count": 3,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "The Cart-Pole is a very simple environment composed of a cart that can move left or right, and pole placed vertically on top of it. The agent must move the cart left or right to keep the pole upright.",
   "metadata": {
    "id": "BHkOFetc-qww",
    "cell_id": "00009-e0cb09e0-ce0b-48fd-ba43-c4ea2c6d49e4",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 74.78125
   }
  },
  {
   "cell_type": "code",
   "source": "env = gym.make('CartPole-v1')",
   "metadata": {
    "id": "MYcIQhVl-qwx",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "bc9a49ed-a5ce-4f27-d16d-11dcbd68ca29",
    "cell_id": "00010-e3e72f12-b87b-4b59-bd0f-920827a29df2",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 171.75
   },
   "execution_count": 4,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Let's initialize the environment by calling is `reset()` method. This returns an observation:",
   "metadata": {
    "id": "l6SpqDft-qwy",
    "cell_id": "00011-be2eb12d-43ca-40bc-941d-cef75e278a33",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "code",
   "source": "env.seed(42)\nobs = env.reset()",
   "metadata": {
    "id": "1giw8eN5-qwz",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "aff24c32-83cf-46db-dd84-a8dac72ad27a",
    "cell_id": "00012-56bca012-9762-4e92-b573-9a7189185eab",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 149.375
   },
   "execution_count": 5,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Observations vary depending on the environment. In this case it is a 1D NumPy array composed of 4 floats: they represent the cart's horizontal position, its velocity, the angle of the pole (0 = vertical), and the angular velocity.",
   "metadata": {
    "id": "S1FhuGry-qw1",
    "cell_id": "00013-138743fd-1490-46d5-a56f-37014ccd221f",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 97.171875
   }
  },
  {
   "cell_type": "code",
   "source": "obs",
   "metadata": {
    "id": "ocKMq2xV-qw2",
    "outputId": "5fa3472a-2efe-4e2e-f22d-37a00468cab1",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "cell_id": "00014-f1afc89c-ba23-45ef-b5b2-d433543984ea",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 113.1875,
    "deepnote_output_heights": [
     21.1875
    ]
   },
   "execution_count": 6,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "An environment can be visualized by calling its `render()` method, and you can pick the rendering mode (the rendering options depend on the environment).",
   "metadata": {
    "id": "k3LgW4EU-qw3",
    "cell_id": "00015-618545c6-ef93-49a6-99ea-17e7566d9045",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 74.78125
   }
  },
  {
   "cell_type": "markdown",
   "source": "**Warning**: some environments (including the Cart-Pole) require access to your display, which opens up a separate window, even if you specify `mode=\"rgb_array\"`. In general you can safely ignore that window. However, if Jupyter is running on a headless server (ie. without a screen) it will raise an exception. One way to avoid this is to install a fake X server like [Xvfb](http://en.wikipedia.org/wiki/Xvfb). On Debian or Ubuntu:\n\n```bash\n$ apt update\n$ apt install -y xvfb\n```\n\nYou can then start Jupyter using the `xvfb-run` command:\n\n```bash\n$ xvfb-run -s \"-screen 0 1400x900x24\" jupyter notebook\n```\n\nAlternatively, you can install the [pyvirtualdisplay](https://github.com/ponty/pyvirtualdisplay) Python library which wraps Xvfb:\n\n```bash\n%pip install -U pyvirtualdisplay\n```\n\nAnd run the following code:",
   "metadata": {
    "id": "sdh0VXZT-qw4",
    "cell_id": "00016-bb9686ed-98c1-4dd0-8a09-afe1306c60d2",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 378.296875
   }
  },
  {
   "cell_type": "code",
   "source": "try:\n    import pyvirtualdisplay\n    display = pyvirtualdisplay.Display(visible=0, size=(1400, 900)).start()\nexcept ImportError:\n    pass",
   "metadata": {
    "id": "zUdkuD8Q-qw4",
    "cell_id": "00017-e7f59ad3-1fb4-4f99-a331-94ba773f5e62",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 133
   },
   "execution_count": 7,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "env.render()",
   "metadata": {
    "id": "7G2mMpqD-qw5",
    "outputId": "186c3495-c90d-47ae-9801-17b4a593c9da",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "cell_id": "00018-3e57d8da-405a-4da3-8303-b658f69d1ec4",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 171.75
   },
   "execution_count": 8,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "In this example we will set `mode=\"rgb_array\"` to get an image of the environment as a NumPy array:",
   "metadata": {
    "id": "YnHlp9oj-qw5",
    "cell_id": "00019-c3b87001-0ca6-4496-b8cb-0ab8827b2308",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "code",
   "source": "img = env.render(mode=\"rgb_array\")\nimg.shape",
   "metadata": {
    "id": "V3xWPXMr-qw6",
    "outputId": "eaecf441-d56a-48df-ccde-bd9fd99ba26f",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "cell_id": "00020-de89165c-41a0-4f49-872c-4abc73753d33",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 221.75,
    "deepnote_output_heights": [
     null,
     21.1875
    ]
   },
   "execution_count": 9,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "def plot_environment(env, figsize=(5,4)):\n    plt.figure(figsize=figsize)\n    img = env.render(mode=\"rgb_array\")\n    plt.imshow(img)\n    plt.axis(\"off\")\n    return img",
   "metadata": {
    "id": "Hrk-6g-H-qw7",
    "cell_id": "00021-84288642-1822-4197-b754-0504a67cf57f",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 151
   },
   "execution_count": 10,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "plot_environment(env)\nplt.show()",
   "metadata": {
    "id": "obb9ygV5-qw7",
    "outputId": "305d830c-daef-4080-ea66-f2852882061b",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 217
    },
    "cell_id": "00022-5bf5aac8-32cf-4a74-9058-e5525d3ccdf6",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 312,
    "deepnote_output_heights": [
     202
    ]
   },
   "execution_count": 11,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Let's see how to interact with an environment. Your agent will need to select an action from an \"action space\" (the set of possible actions). Let's see what this environment's action space looks like:",
   "metadata": {
    "id": "Fve5fP8Q-qw8",
    "cell_id": "00023-7f0fff79-f615-4e8a-9296-c4a7695b2ba2",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 74.78125
   }
  },
  {
   "cell_type": "code",
   "source": "env.action_space",
   "metadata": {
    "id": "tQr9Pc0Z-qw8",
    "outputId": "4705a12d-aa68-4001-bcac-b8630f049497",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "cell_id": "00024-3fda931c-9280-4720-b06d-65f6e7ef07e2",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 113.1875,
    "deepnote_output_heights": [
     21.1875
    ]
   },
   "execution_count": 12,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Yep, just two possible actions: accelerate towards the left or towards the right.",
   "metadata": {
    "id": "2p-SEhfx-qw9",
    "cell_id": "00025-04eccac9-ac29-4382-aae7-058fb1bfe7c3",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "markdown",
   "source": "Since the pole is leaning toward the right (`obs[2] > 0`), let's accelerate the cart toward the right:",
   "metadata": {
    "id": "JJFqJwzZ-qw9",
    "cell_id": "00026-5c273f0e-f24b-40ef-8876-64c08177f8c7",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "code",
   "source": "action = 1  # accelerate right\nobs, reward, done, info = env.step(action)\nobs",
   "metadata": {
    "id": "LpZpfqIL-qw-",
    "outputId": "3037574b-47f0-429a-868e-5e605fc45b51",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "cell_id": "00027-e80c9e3f-1577-4e77-b368-f5530a41f2f8",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 149.1875,
    "deepnote_output_heights": [
     21.1875
    ]
   },
   "execution_count": 13,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Notice that the cart is now moving toward the right (`obs[1] > 0`). The pole is still tilted toward the right (`obs[2] > 0`), but its angular velocity is now negative (`obs[3] < 0`), so it will likely be tilted toward the left after the next step.",
   "metadata": {
    "id": "U5JiZyjL-qw-",
    "cell_id": "00028-5f7a969f-03dc-4a19-8470-57304659bb57",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 97.171875
   }
  },
  {
   "cell_type": "code",
   "source": "plot_environment(env)\nsave_fig(\"cart_pole_plot\")",
   "metadata": {
    "id": "l4IxCKI4-qw_",
    "outputId": "e5a38f70-39ff-44b0-a250-4a43be9a3fa7",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 273
    },
    "cell_id": "00029-411db773-43ae-42d2-bbda-b8f67b312de7",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 397.1875,
    "deepnote_output_heights": [
     null,
     241
    ]
   },
   "execution_count": 14,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Looks like it's doing what we're telling it to do!",
   "metadata": {
    "id": "G94qee38-qxA",
    "cell_id": "00030-feb36b92-fead-443a-9267-1327cbb3b333",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "markdown",
   "source": "The environment also tells the agent how much reward it got during the last step:",
   "metadata": {
    "id": "OGIz8dV5-qxA",
    "cell_id": "00031-0c2a8bcf-4a2a-43d6-89d1-dbbd8947577a",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "code",
   "source": "reward",
   "metadata": {
    "id": "E_fPzxiz-qxA",
    "outputId": "0b180b9b-c0b6-4c97-9ec4-3a5a02ee1436",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "cell_id": "00032-1e5d0780-59ad-4735-9794-8c280b5a4f42",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 113.1875,
    "deepnote_output_heights": [
     21.1875
    ]
   },
   "execution_count": 15,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "When the game is over, the environment returns `done=True`:",
   "metadata": {
    "id": "Tpvv0of6-qxC",
    "cell_id": "00033-b00ad661-1389-4c83-93cc-cdeac9f99339",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "code",
   "source": "done",
   "metadata": {
    "id": "PmngZk7z-qxC",
    "outputId": "c4ed451d-6b09-4ec1-d846-04e1da865e27",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "cell_id": "00034-16b5bc72-79f0-4ec0-919f-58038229a1ce",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 113.1875,
    "deepnote_output_heights": [
     21.1875
    ]
   },
   "execution_count": 16,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Finally, `info` is an environment-specific dictionary that can provide some extra information that you may find useful for debugging or for training. For example, in some games it may indicate how many lives the agent has.",
   "metadata": {
    "id": "ciGi9zTy-qxD",
    "cell_id": "00035-00b5ee13-da01-4c0c-9167-07d1c7c8febf",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 97.171875
   }
  },
  {
   "cell_type": "code",
   "source": "info",
   "metadata": {
    "id": "WNzOucsj-qxD",
    "outputId": "16806b91-5d8e-462b-b8af-278aee8cb4d5",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "cell_id": "00036-d6029243-0fdf-4435-affc-aab943d225f1",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 113.1875,
    "deepnote_output_heights": [
     21.1875
    ]
   },
   "execution_count": 17,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "The sequence of steps between the moment the environment is reset until it is done is called an \"episode\". At the end of an episode (i.e., when `step()` returns `done=True`), you should reset the environment before you continue to use it.",
   "metadata": {
    "id": "jDtUm3H7-qxE",
    "cell_id": "00037-877ed3ce-a312-46f1-bf24-5290864628c9",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 97.171875
   }
  },
  {
   "cell_type": "code",
   "source": "if done:\n    obs = env.reset()",
   "metadata": {
    "id": "3MT7_FPU-qxF",
    "cell_id": "00038-2c9223de-d04b-41cc-9cbe-d15967d4c155",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 79
   },
   "execution_count": 18,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Now how can we make the pole remain upright? We will need to define a _policy_ for that. This is the strategy that the agent will use to select an action at each step. It can use all the past actions and observations to decide what to do.",
   "metadata": {
    "id": "N3-CrjtW-qxG",
    "cell_id": "00039-a44a7267-f08b-48aa-abf7-a974f4b7dece",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 97.171875
   }
  },
  {
   "cell_type": "markdown",
   "source": "# A simple hard-coded policy",
   "metadata": {
    "id": "kiQl5hho-qxG",
    "cell_id": "00040-d2142108-69fa-46e3-a165-da873f70580f",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 82
   }
  },
  {
   "cell_type": "markdown",
   "source": "Let's hard code a simple strategy: if the pole is tilting to the left, then push the cart to the left, and _vice versa_. Let's see if that works:",
   "metadata": {
    "id": "CmGxZ13C-qxG",
    "cell_id": "00041-1c21d086-91af-41e3-8c14-9fa643e64ee4",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 74.78125
   }
  },
  {
   "cell_type": "code",
   "source": "env.seed(42)\n\ndef basic_policy(obs):\n    angle = obs[2]\n    return 0 if angle < 0 else 1\n\ntotals = []\nfor episode in range(500):\n    episode_rewards = 0\n    obs = env.reset()\n    for step in range(200):\n        action = basic_policy(obs)\n        obs, reward, done, info = env.step(action)\n        episode_rewards += reward\n        if done:\n            break\n    totals.append(episode_rewards)",
   "metadata": {
    "id": "ZuONWmvn-qxH",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "4b4faee7-d6f1-4a7b-f92c-412c25543c07",
    "cell_id": "00042-507f617c-51b1-47c0-a844-26177048c4cc",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 419.375
   },
   "execution_count": 19,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "np.mean(totals), np.std(totals), np.min(totals), np.max(totals)",
   "metadata": {
    "id": "nFZ8VEK--qxH",
    "outputId": "8513820e-ecab-44c2-ca63-72243c583dac",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "cell_id": "00043-e6449383-ee1d-4abe-8dad-1874cd20c54c",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 113.1875,
    "deepnote_output_heights": [
     21.1875
    ]
   },
   "execution_count": 20,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Well, as expected, this strategy is a bit too basic: the best it did was to keep the pole up for only 68 steps. This environment is considered solved when the agent keeps the pole up for 200 steps.",
   "metadata": {
    "id": "TGouPN55-qxJ",
    "cell_id": "00044-5cfa740e-dab6-4f99-abff-608e92d77640",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 74.78125
   }
  },
  {
   "cell_type": "markdown",
   "source": "Let's visualize one episode:",
   "metadata": {
    "id": "yj4cgkVJ-qxK",
    "cell_id": "00045-3a2b63dc-10f9-4c6f-a923-6efb35dd1a79",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "code",
   "source": "env.seed(42)\n\nframes = []\n\nobs = env.reset()\nfor step in range(200):\n    img = env.render(mode=\"rgb_array\")\n    frames.append(img)\n    action = basic_policy(obs)\n\n    obs, reward, done, info = env.step(action)\n    if done:\n        break",
   "metadata": {
    "id": "SJlEqJ_S-qxK",
    "cell_id": "00046-b33ee2e7-3176-4041-912d-eabb6abdb4dc",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 277
   },
   "execution_count": 21,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Now show the animation:",
   "metadata": {
    "id": "OBo0JsPp-qxL",
    "cell_id": "00047-658972b7-5176-408f-80b5-316afc43f8aa",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "code",
   "source": "def update_scene(num, frames, patch):\n    patch.set_data(frames[num])\n    return patch,\n\ndef plot_animation(frames, repeat=False, interval=40):\n    fig = plt.figure()\n    patch = plt.imshow(frames[0])\n    plt.axis('off')\n    anim = animation.FuncAnimation(\n        fig, update_scene, fargs=(frames, patch),\n        frames=len(frames), repeat=repeat, interval=interval)\n    plt.close()\n    return anim",
   "metadata": {
    "id": "T9SBlptO-qxL",
    "cell_id": "00048-e138239d-bfb0-4468-8ed4-320c11994df0",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 277
   },
   "execution_count": 22,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "plot_animation(frames)",
   "metadata": {
    "id": "piuXVRWy-qxL",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 389
    },
    "outputId": "5da8db40-a9bd-4232-a192-30fd61460234",
    "cell_id": "00049-1238e4b1-8ea0-48a8-8921-c6a1d1588fe4",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 61
   },
   "execution_count": 23,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Clearly the system is unstable and after just a few wobbles, the pole ends up too tilted: game over. We will need to be smarter than that!",
   "metadata": {
    "id": "SzFlTdNo-qxM",
    "cell_id": "00050-fcb8f07e-91a6-406d-a9c7-bd05c0a724b0",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 74.78125
   }
  },
  {
   "cell_type": "markdown",
   "source": "# Neural Network Policies",
   "metadata": {
    "id": "058mM6GR-qxM",
    "cell_id": "00051-ad062416-1736-490e-b1c4-d68b3e49461d",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 82
   }
  },
  {
   "cell_type": "markdown",
   "source": "Let's create a neural network that will take observations as inputs, and output the probabilities of actions to take for each observation. To choose an action, the network will estimate a probability for each action, then we will select an action randomly according to the estimated probabilities. In the case of the Cart-Pole environment, there are just two possible actions (left or right), so we only need one output neuron: it will output the probability `p` of the action 0 (left), and of course the probability of action 1 (right) will be `1 - p`.",
   "metadata": {
    "id": "S57v1QEA-qxM",
    "cell_id": "00052-d6f73864-064d-4301-a469-3c3ad5030b20",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 164.34375
   }
  },
  {
   "cell_type": "code",
   "source": "keras.backend.clear_session()\ntf.random.set_seed(42)\nnp.random.seed(42)\n\nn_inputs = 4 # == env.observation_space.shape[0]\n\nmodel = keras.models.Sequential([\n    keras.layers.Dense(5, activation=\"elu\", input_shape=[n_inputs]),\n    keras.layers.Dense(1, activation=\"sigmoid\"),\n])",
   "metadata": {
    "id": "F4uzlmyJ-qxM",
    "cell_id": "00053-ff25cdc8-30cb-4a97-9ff5-2e6ed385f0d0",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 223
   },
   "execution_count": 24,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "In this particular environment, the past actions and observations can safely be ignored, since each observation contains the environment's full state. If there were some hidden state then you may need to consider past actions and observations in order to try to infer the hidden state of the environment. For example, if the environment only revealed the position of the cart but not its velocity, you would have to consider not only the current observation but also the previous observation in order to estimate the current velocity. Another example is if the observations are noisy: you may want to use the past few observations to estimate the most likely current state. Our problem is thus as simple as can be: the current observation is noise-free and contains the environment's full state.",
   "metadata": {
    "id": "RX6E9bU3-qxN",
    "cell_id": "00054-0cf04790-9bb0-42c4-9b26-8d26a7d55c14",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 231.515625
   }
  },
  {
   "cell_type": "markdown",
   "source": "You may wonder why we plan to pick a random action based on the probability given by the policy network, rather than just picking the action with the highest probability. This approach lets the agent find the right balance between _exploring_ new actions and _exploiting_ the actions that are known to work well. Here's an analogy: suppose you go to a restaurant for the first time, and all the dishes look equally appealing so you randomly pick one. If it turns out to be good, you can increase the probability to order it next time, but you shouldn't increase that probability to 100%, or else you will never try out the other dishes, some of which may be even better than the one you tried.",
   "metadata": {
    "id": "jijDlOGS-qxN",
    "cell_id": "00055-84a8503a-563b-444e-a094-0a65da0b546c",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 186.734375
   }
  },
  {
   "cell_type": "markdown",
   "source": "Let's write a small function that will run the model to play one episode, and return the frames so we can display an animation:",
   "metadata": {
    "id": "TVc_FN0q-qxO",
    "cell_id": "00056-04c08581-a37f-4df8-b364-bba0e30bb562",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 74.78125
   }
  },
  {
   "cell_type": "code",
   "source": "def render_policy_net(model, n_max_steps=200, seed=42):\n    frames = []\n    env = gym.make(\"CartPole-v1\")\n    env.seed(seed)\n    np.random.seed(seed)\n    obs = env.reset()\n    for step in range(n_max_steps):\n        frames.append(env.render(mode=\"rgb_array\"))\n        left_proba = model.predict(obs.reshape(1, -1))\n        action = int(np.random.rand() > left_proba)\n        obs, reward, done, info = env.step(action)\n        if done:\n            break\n    env.close()\n    return frames",
   "metadata": {
    "id": "VUoY0aRG-qxO",
    "cell_id": "00057-92980458-b317-4693-8517-b3774301a35e",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 313
   },
   "execution_count": 25,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Now let's look at how well this randomly initialized policy network performs:",
   "metadata": {
    "id": "EGnxxoHU-qxP",
    "cell_id": "00058-3febccb1-b3c5-4330-8081-4742379954f7",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "code",
   "source": "frames = render_policy_net(model)\nplot_animation(frames)",
   "metadata": {
    "id": "yQaCSctG-qxQ",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 565
    },
    "outputId": "926521dd-d08c-4f03-a65e-08a334682224",
    "cell_id": "00059-945d01a3-43cb-45f5-9260-97d5723a8858",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 79
   },
   "execution_count": 26,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Yeah... pretty bad. The neural network will have to learn to do better. First let's see if it is capable of learning the basic policy we used earlier: go left if the pole is tilting left, and go right if it is tilting right.",
   "metadata": {
    "id": "Mvwp0arb-qxQ",
    "cell_id": "00060-5c8c09d7-cac4-4095-ba06-705eb1036634",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 97.171875
   }
  },
  {
   "cell_type": "markdown",
   "source": "We can make the same net play in 50 different environments in parallel (this will give us a diverse training batch at each step), and train for 5000 iterations. We also reset environments when they are done. We train the model using a custom training loop so we can easily use the predictions at each training step to advance the environments.",
   "metadata": {
    "id": "8F0Njka5-qxQ",
    "cell_id": "00061-90fc8591-50ab-4cc8-a321-44a5ae049778",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 119.5625
   }
  },
  {
   "cell_type": "code",
   "source": "n_environments = 50\nn_iterations = 5000\n\nenvs = [gym.make(\"CartPole-v1\") for _ in range(n_environments)]\nfor index, env in enumerate(envs):\n    env.seed(index)\nnp.random.seed(42)\nobservations = [env.reset() for env in envs]\noptimizer = keras.optimizers.RMSprop()\nloss_fn = keras.losses.binary_crossentropy\n\nfor iteration in range(n_iterations):\n    # if angle < 0, we want proba(left) = 1., or else proba(left) = 0.\n    target_probas = np.array([([1.] if obs[2] < 0 else [0.])\n                              for obs in observations])\n    with tf.GradientTape() as tape:\n        left_probas = model(np.array(observations))\n        loss = tf.reduce_mean(loss_fn(target_probas, left_probas))\n    print(\"\\rIteration: {}, Loss: {:.3f}\".format(iteration, loss.numpy()), end=\"\")\n    grads = tape.gradient(loss, model.trainable_variables)\n    optimizer.apply_gradients(zip(grads, model.trainable_variables))\n    actions = (np.random.rand(n_environments, 1) > left_probas.numpy()).astype(np.int32)\n    for env_index, env in enumerate(envs):\n        obs, reward, done, info = env.step(actions[env_index][0])\n        observations[env_index] = obs if not done else env.reset()\n\nfor env in envs:\n    env.close()",
   "metadata": {
    "id": "XSVHjkOH-qxR",
    "outputId": "49c8c5da-1faa-425d-8200-a12f86a36ea9",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "cell_id": "00062-6d79160e-cdca-455d-8f10-aa417b56fe1d",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 552
   },
   "execution_count": 27,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "frames = render_policy_net(model)\nplot_animation(frames)",
   "metadata": {
    "id": "wIYg_Vnx-qxR",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 461
    },
    "outputId": "cde2e45a-8154-4403-94ad-5be655661b4a",
    "cell_id": "00063-04a430d3-fcf2-4fa7-84c1-b194bf964d2f",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 79
   },
   "execution_count": 28,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Looks like it learned the policy correctly. Now let's see if it can learn a better policy on its own. One that does not wobble as much.",
   "metadata": {
    "id": "-myEPzH0-qxS",
    "cell_id": "00064-24b4b8f2-0c4e-481b-8d8d-c2a9980730eb",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 74.78125
   }
  },
  {
   "cell_type": "markdown",
   "source": "# Policy Gradients",
   "metadata": {
    "id": "JPGbFieb-qxS",
    "cell_id": "00065-067ab9f4-2be1-4b1a-bf20-d907c013c5a3",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 82
   }
  },
  {
   "cell_type": "markdown",
   "source": "To train this neural network we will need to define the target probabilities `y`. If an action is good we should increase its probability, and conversely if it is bad we should reduce it. But how do we know whether an action is good or bad? The problem is that most actions have delayed effects, so when you win or lose points in an episode, it is not clear which actions contributed to this result: was it just the last action? Or the last 10? Or just one action 50 steps earlier? This is called the _credit assignment problem_.\n\nThe _Policy Gradients_ algorithm tackles this problem by first playing multiple episodes, then making the actions in good episodes slightly more likely, while actions in bad episodes are made slightly less likely. First we play, then we go back and think about what we did.",
   "metadata": {
    "id": "9lny54vH-qxS",
    "cell_id": "00066-d1bb6416-6087-439f-8c37-42a3fc954b88",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 245.515625
   }
  },
  {
   "cell_type": "markdown",
   "source": "Let's start by creating a function to play a single step using the model. We will also pretend for now that whatever action it takes is the right one, so we can compute the loss and its gradients (we will just save these gradients for now, and modify them later depending on how good or bad the action turned out to be):",
   "metadata": {
    "id": "z7BIxr5i-qxS",
    "cell_id": "00067-542dc5f8-1c6f-4758-aac6-6e22359c10a7",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 119.5625
   }
  },
  {
   "cell_type": "code",
   "source": "def play_one_step(env, obs, model, loss_fn):\n    with tf.GradientTape() as tape:\n        left_proba = model(obs[np.newaxis])\n        action = (tf.random.uniform([1, 1]) > left_proba)\n        y_target = tf.constant([[1.]]) - tf.cast(action, tf.float32)\n        loss = tf.reduce_mean(loss_fn(y_target, left_proba))\n    grads = tape.gradient(loss, model.trainable_variables)\n    obs, reward, done, info = env.step(int(action[0, 0].numpy()))\n    return obs, reward, done, grads",
   "metadata": {
    "id": "3O7BA9ML-qxT",
    "cell_id": "00068-82087ff4-7549-4d21-b36c-22ede6bb8652",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 205
   },
   "execution_count": 29,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "If `left_proba` is high, then `action` will most likely be `False` (since a random number uniformally sampled between 0 and 1 will probably not be greater than `left_proba`). And `False` means 0 when you cast it to a number, so `y_target` would be equal to 1 - 0 = 1. In other words, we set the target to 1, meaning we pretend that the probability of going left should have been 100% (so we took the right action).",
   "metadata": {
    "id": "I6FWR-V9-qxT",
    "cell_id": "00069-319aa899-e494-4e8c-877c-622338baf79a",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 141.953125
   }
  },
  {
   "cell_type": "markdown",
   "source": "Now let's create another function that will rely on the `play_one_step()` function to play multiple episodes, returning all the rewards and gradients, for each episode and each step:",
   "metadata": {
    "id": "pGFM1EDm-qxT",
    "cell_id": "00070-4041c362-c325-405a-a8cc-ff68ab7df1a6",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 74.78125
   }
  },
  {
   "cell_type": "code",
   "source": "def play_multiple_episodes(env, n_episodes, n_max_steps, model, loss_fn):\n    all_rewards = []\n    all_grads = []\n    for episode in range(n_episodes):\n        current_rewards = []\n        current_grads = []\n        obs = env.reset()\n        for step in range(n_max_steps):\n            obs, reward, done, grads = play_one_step(env, obs, model, loss_fn)\n            current_rewards.append(reward)\n            current_grads.append(grads)\n            if done:\n                break\n        all_rewards.append(current_rewards)\n        all_grads.append(current_grads)\n    return all_rewards, all_grads",
   "metadata": {
    "id": "xRIL_2EH-qxU",
    "cell_id": "00071-f572e03f-2943-4cdd-8219-43efabfe0a20",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 331
   },
   "execution_count": 30,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "The Policy Gradients algorithm uses the model to play the episode several times (e.g., 10 times), then it goes back and looks at all the rewards, discounts them and normalizes them. So let's create couple functions for that: the first will compute discounted rewards; the second will normalize the discounted rewards across many episodes.",
   "metadata": {
    "id": "AeI9afBA-qxU",
    "cell_id": "00072-d0380350-7136-4b00-bf11-bbeeb6bb1f18",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 119.5625
   }
  },
  {
   "cell_type": "code",
   "source": "def discount_rewards(rewards, discount_rate):\n    discounted = np.array(rewards)\n    for step in range(len(rewards) - 2, -1, -1):\n        discounted[step] += discounted[step + 1] * discount_rate\n    return discounted\n\ndef discount_and_normalize_rewards(all_rewards, discount_rate):\n    all_discounted_rewards = [discount_rewards(rewards, discount_rate)\n                              for rewards in all_rewards]\n    flat_rewards = np.concatenate(all_discounted_rewards)\n    reward_mean = flat_rewards.mean()\n    reward_std = flat_rewards.std()\n    return [(discounted_rewards - reward_mean) / reward_std\n            for discounted_rewards in all_discounted_rewards]",
   "metadata": {
    "id": "B-GuR96x-qxV",
    "cell_id": "00073-a41fb4fe-fc92-43de-abec-cf869ae3c73a",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 295
   },
   "execution_count": 31,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Say there were 3 actions, and after each action there was a reward: first 10, then 0, then -50. If we use a discount factor of 80%, then the 3rd action will get -50 (full credit for the last reward), but the 2nd action will only get -40 (80% credit for the last reward), and the 1st action will get 80% of -40 (-32) plus full credit for the first reward (+10), which leads to a discounted reward of -22:",
   "metadata": {
    "id": "vjFC7KUd-qxW",
    "cell_id": "00074-7e1477c2-eaee-4873-8f72-6910a4c82e58",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 119.5625
   }
  },
  {
   "cell_type": "code",
   "source": "discount_rewards([10, 0, -50], discount_rate=0.8)",
   "metadata": {
    "id": "G7ZGX4Vx-qxY",
    "outputId": "55925e67-679b-4cbf-a852-b693768a97d0",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "cell_id": "00075-d4b9cb0e-18d3-4322-ae3c-fc0fc1ee5e6a",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 61
   },
   "execution_count": 32,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "To normalize all discounted rewards across all episodes, we compute the mean and standard deviation of all the discounted rewards, and we subtract the mean from each discounted reward, and divide by the standard deviation:",
   "metadata": {
    "id": "dB_osGTU-qxZ",
    "cell_id": "00076-259f2832-a56f-4bf0-b8dd-4ea8bd6cee14",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 97.171875
   }
  },
  {
   "cell_type": "code",
   "source": "discount_and_normalize_rewards([[10, 0, -50], [10, 20]], discount_rate=0.8)",
   "metadata": {
    "scrolled": true,
    "id": "y7M-RGax-qxa",
    "outputId": "f702f3ec-994f-492e-96b2-4a1c3e2b06e3",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "cell_id": "00077-853a7164-6abf-432e-a1e3-99e19b3404c0",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 61
   },
   "execution_count": 33,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "n_iterations = 150\nn_episodes_per_update = 10\nn_max_steps = 200\ndiscount_rate = 0.95",
   "metadata": {
    "id": "bxTZpXkc-qxb",
    "cell_id": "00078-36a65fd0-8a39-4000-af04-091f892e1f70",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 115
   },
   "execution_count": 34,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "optimizer = keras.optimizers.Adam(learning_rate=0.01)\nloss_fn = keras.losses.binary_crossentropy",
   "metadata": {
    "id": "SbLhJ9Xz-qxe",
    "cell_id": "00079-f2848cfa-784d-4345-a582-91d81da740fa",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 79
   },
   "execution_count": 35,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "keras.backend.clear_session()\nnp.random.seed(42)\ntf.random.set_seed(42)\n\nmodel = keras.models.Sequential([\n    keras.layers.Dense(5, activation=\"elu\", input_shape=[4]),\n    keras.layers.Dense(1, activation=\"sigmoid\"),\n])",
   "metadata": {
    "id": "Wt-syAOv-qxg",
    "cell_id": "00080-dbb0a43b-0288-445f-8fec-0027fc21c618",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 187
   },
   "execution_count": 36,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "env = gym.make(\"CartPole-v1\")\nenv.seed(42);\n\nfor iteration in range(n_iterations):\n    all_rewards, all_grads = play_multiple_episodes(\n        env, n_episodes_per_update, n_max_steps, model, loss_fn)\n    total_rewards = sum(map(sum, all_rewards))                     # Not shown in the book\n    print(\"\\rIteration: {}, mean rewards: {:.1f}\".format(          # Not shown\n        iteration, total_rewards / n_episodes_per_update), end=\"\") # Not shown\n    all_final_rewards = discount_and_normalize_rewards(all_rewards,\n                                                       discount_rate)\n    all_mean_grads = []\n    for var_index in range(len(model.trainable_variables)):\n        mean_grads = tf.reduce_mean(\n            [final_reward * all_grads[episode_index][step][var_index]\n             for episode_index, final_rewards in enumerate(all_final_rewards)\n                 for step, final_reward in enumerate(final_rewards)], axis=0)\n        all_mean_grads.append(mean_grads)\n    optimizer.apply_gradients(zip(all_mean_grads, model.trainable_variables))\n\nenv.close()",
   "metadata": {
    "id": "DmeJ8XO--qxh",
    "outputId": "8e3475cb-ca0d-489d-df74-386c9dc9159c",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "cell_id": "00081-2f9c94e8-9c90-409a-a416-b4d6bf740e43",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 426
   },
   "execution_count": 37,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "frames = render_policy_net(model)\nplot_animation(frames)",
   "metadata": {
    "id": "l-AWST8N-qxn",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 461
    },
    "outputId": "07055398-a510-4d47-f727-2424118fe27b",
    "cell_id": "00082-385fb825-6690-4e06-b3a0-645cec3cf320",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 79
   },
   "execution_count": 38,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "# Markov Chains",
   "metadata": {
    "id": "bcu5A3m8-qxo",
    "cell_id": "00083-984387a9-ded0-4ffc-99f0-cc5d82dc4374",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 82
   }
  },
  {
   "cell_type": "code",
   "source": "np.random.seed(42)\n\ntransition_probabilities = [ # shape=[s, s']\n        [0.7, 0.2, 0.0, 0.1],  # from s0 to s0, s1, s2, s3\n        [0.0, 0.0, 0.9, 0.1],  # from s1 to ...\n        [0.0, 1.0, 0.0, 0.0],  # from s2 to ...\n        [0.0, 0.0, 0.0, 1.0]]  # from s3 to ...\n\nn_max_steps = 50\n\ndef print_sequence():\n    current_state = 0\n    print(\"States:\", end=\" \")\n    for step in range(n_max_steps):\n        print(current_state, end=\" \")\n        if current_state == 3:\n            break\n        current_state = np.random.choice(range(4), p=transition_probabilities[current_state])\n    else:\n        print(\"...\", end=\"\")\n    print()\n\nfor _ in range(10):\n    print_sequence()",
   "metadata": {
    "id": "hhvwpgaE-qxo",
    "outputId": "75846c0e-147b-4d80-a3b2-0b4a753b3b96",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "cell_id": "00084-822c2a15-89af-43c6-886f-467cd430a657",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 480
   },
   "execution_count": 39,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "# Markov Decision Process",
   "metadata": {
    "id": "iLIE_5vh-qxq",
    "cell_id": "00085-30680fde-0e21-487b-9f2e-46c648302d83",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 82
   }
  },
  {
   "cell_type": "markdown",
   "source": "Let's define some transition probabilities, rewards and possible actions. For example, in state s0, if action a0 is chosen then with proba 0.7 we will go to state s0 with reward +10, with probability 0.3 we will go to state s1 with no reward, and with never go to state s2 (so the transition probabilities are `[0.7, 0.3, 0.0]`, and the rewards are `[+10, 0, 0]`):",
   "metadata": {
    "id": "1Wl7lgxK-qxr",
    "cell_id": "00086-8d9706b1-e034-4f6d-bb6b-96db212aaec0",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 119.5625
   }
  },
  {
   "cell_type": "code",
   "source": "transition_probabilities = [ # shape=[s, a, s']\n        [[0.7, 0.3, 0.0], [1.0, 0.0, 0.0], [0.8, 0.2, 0.0]],\n        [[0.0, 1.0, 0.0], None, [0.0, 0.0, 1.0]],\n        [None, [0.8, 0.1, 0.1], None]]\nrewards = [ # shape=[s, a, s']\n        [[+10, 0, 0], [0, 0, 0], [0, 0, 0]],\n        [[0, 0, 0], [0, 0, 0], [0, 0, -50]],\n        [[0, 0, 0], [+40, 0, 0], [0, 0, 0]]]\npossible_actions = [[0, 1, 2], [0, 2], [1]]",
   "metadata": {
    "id": "FY1Ijaip-qxr",
    "cell_id": "00087-d70adf92-c28e-4a56-b3e7-b5fd21c49852",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 205
   },
   "execution_count": 40,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "# Q-Value Iteration",
   "metadata": {
    "id": "GM4mx7ZO-qxt",
    "cell_id": "00088-0000b168-4c17-4d04-8815-ccb008c8f47e",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 82
   }
  },
  {
   "cell_type": "code",
   "source": "Q_values = np.full((3, 3), -np.inf) # -np.inf for impossible actions\nfor state, actions in enumerate(possible_actions):\n    Q_values[state, actions] = 0.0  # for all possible actions",
   "metadata": {
    "id": "GCOAVQCI-qxt",
    "cell_id": "00089-cbd8ccec-7a25-4549-9e82-7fb0870ef284",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 97
   },
   "execution_count": 41,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "gamma = 0.90  # the discount factor\n\nhistory1 = [] # Not shown in the book (for the figure below)\nfor iteration in range(50):\n    Q_prev = Q_values.copy()\n    history1.append(Q_prev) # Not shown\n    for s in range(3):\n        for a in possible_actions[s]:\n            Q_values[s, a] = np.sum([\n                    transition_probabilities[s][a][sp]\n                    * (rewards[s][a][sp] + gamma * np.max(Q_prev[sp]))\n                for sp in range(3)])\n\nhistory1 = np.array(history1) # Not shown",
   "metadata": {
    "id": "01fp0w1P-qxw",
    "cell_id": "00090-d1a2f281-75f6-4d62-a42a-a49c6ef0d433",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 295
   },
   "execution_count": 42,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "Q_values",
   "metadata": {
    "id": "S6Bw4SJE-qxx",
    "outputId": "4d07ccf9-2e23-451e-fd73-c14dc5ff1a93",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "cell_id": "00091-f624656c-a521-45d7-ac19-564dbf743a61",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 61
   },
   "execution_count": 43,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "np.argmax(Q_values, axis=1)",
   "metadata": {
    "id": "0xiMIrPC-qxy",
    "outputId": "32bf6728-39b5-498b-d1c6-10aef8a7c349",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "cell_id": "00092-a11baba8-8b5c-4507-888c-b13d5a02501f",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 61
   },
   "execution_count": 44,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "The optimal policy for this MDP, when using a discount factor of 0.90, is to choose action a0 when in state s0, and choose action a0 when in state s1, and finally choose action a1 (the only possible action) when in state s2.",
   "metadata": {
    "id": "CHGchyOM-qxz",
    "cell_id": "00093-e690ed38-cd2d-4fde-b7b5-c77270b9d75f",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 97.171875
   }
  },
  {
   "cell_type": "markdown",
   "source": "Let's try again with a discount factor of 0.95:",
   "metadata": {
    "id": "KH4BgQd_-qx0",
    "cell_id": "00094-b80f7b83-a709-4ed1-96ab-08ee8c9a1e77",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "code",
   "source": "Q_values = np.full((3, 3), -np.inf) # -np.inf for impossible actions\nfor state, actions in enumerate(possible_actions):\n    Q_values[state, actions] = 0.0  # for all possible actions",
   "metadata": {
    "id": "Y21cPi9q-qx1",
    "cell_id": "00095-0114821a-3394-46a0-95b8-803454ed28f7",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 97
   },
   "execution_count": 45,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "gamma = 0.95  # the discount factor\n\nfor iteration in range(50):\n    Q_prev = Q_values.copy()\n    for s in range(3):\n        for a in possible_actions[s]:\n            Q_values[s, a] = np.sum([\n                    transition_probabilities[s][a][sp]\n                    * (rewards[s][a][sp] + gamma * np.max(Q_prev[sp]))\n                for sp in range(3)])",
   "metadata": {
    "id": "xHxsJahg-qx1",
    "cell_id": "00096-5b415cf0-ba35-4454-ab1c-a07280538b32",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 223
   },
   "execution_count": 46,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "Q_values",
   "metadata": {
    "id": "gqninxie-qx3",
    "outputId": "59867781-a599-4bd1-ad2b-76b5230a87cb",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "cell_id": "00097-e1550525-381e-443f-8991-03307a5b748d",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 61
   },
   "execution_count": 47,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "np.argmax(Q_values, axis=1)",
   "metadata": {
    "id": "Sf7Z4pUp-qx4",
    "outputId": "9dcafb27-4e78-4fd3-d630-dd4b5f24b4c8",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "cell_id": "00098-3717a4e8-6a32-47f4-8665-06f420a3f137",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 61
   },
   "execution_count": 48,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Now the policy has changed! In state s1, we now prefer to go through the fire (choose action a2). This is because the discount factor is larger so the agent values the future more, and it is therefore ready to pay an immediate penalty in order to get more future rewards.",
   "metadata": {
    "id": "UXSM2wmf-qx6",
    "cell_id": "00099-ce741e56-474c-437e-aa21-2a9297c0208b",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 97.171875
   }
  },
  {
   "cell_type": "markdown",
   "source": "# Q-Learning",
   "metadata": {
    "id": "f-FKdcdr-qx6",
    "cell_id": "00100-538e7aa4-51fc-44bb-816d-ea85c530ca88",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 82
   }
  },
  {
   "cell_type": "markdown",
   "source": "Q-Learning works by watching an agent play (e.g., randomly) and gradually improving its estimates of the Q-Values. Once it has accurate Q-Value estimates (or close enough), then the optimal policy consists in choosing the action that has the highest Q-Value (i.e., the greedy policy).",
   "metadata": {
    "id": "zVkVKq6Y-qx7",
    "cell_id": "00101-a61878c2-44ca-4b8d-9fea-16855b61cfe6",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 97.171875
   }
  },
  {
   "cell_type": "markdown",
   "source": "We will need to simulate an agent moving around in the environment, so let's define a function to perform some action and get the new state and a reward:",
   "metadata": {
    "id": "yysnrZFn-qx8",
    "cell_id": "00102-e9c6bd69-4c7f-47fe-b270-38d014eb7aa2",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 74.78125
   }
  },
  {
   "cell_type": "code",
   "source": "def step(state, action):\n    probas = transition_probabilities[state][action]\n    next_state = np.random.choice([0, 1, 2], p=probas)\n    reward = rewards[state][action][next_state]\n    return next_state, reward",
   "metadata": {
    "id": "dmyB4vIH-qx8",
    "cell_id": "00103-400dbfca-f258-4e95-9366-6d34fb1d07d3",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 133
   },
   "execution_count": 49,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "We also need an exploration policy, which can be any policy, as long as it visits every possible state many times. We will just use a random policy, since the state space is very small:",
   "metadata": {
    "id": "dzqNEILC-qx9",
    "cell_id": "00104-f1f44bec-a919-48ba-a8f9-dd79fe2a719c",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 74.78125
   }
  },
  {
   "cell_type": "code",
   "source": "def exploration_policy(state):\n    return np.random.choice(possible_actions[state])",
   "metadata": {
    "id": "kGrVUN9Z-qx9",
    "cell_id": "00105-41cef208-4ac9-4989-a7a9-83c763c05b29",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 79
   },
   "execution_count": 50,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Now let's initialize the Q-Values like earlier, and run the Q-Learning algorithm:",
   "metadata": {
    "id": "3HbbJQ8i-qx_",
    "cell_id": "00106-ff744ac4-6142-4c20-b14e-b25bc4d0f8f0",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "code",
   "source": "np.random.seed(42)\n\nQ_values = np.full((3, 3), -np.inf)\nfor state, actions in enumerate(possible_actions):\n    Q_values[state][actions] = 0\n\nalpha0 = 0.05 # initial learning rate\ndecay = 0.005 # learning rate decay\ngamma = 0.90 # discount factor\nstate = 0 # initial state\nhistory2 = [] # Not shown in the book\n\nfor iteration in range(10000):\n    history2.append(Q_values.copy()) # Not shown\n    action = exploration_policy(state)\n    next_state, reward = step(state, action)\n    next_value = np.max(Q_values[next_state]) # greedy policy at the next step\n    alpha = alpha0 / (1 + iteration * decay)\n    Q_values[state, action] *= 1 - alpha\n    Q_values[state, action] += alpha * (reward + gamma * next_value)\n    state = next_state\n\nhistory2 = np.array(history2) # Not shown",
   "metadata": {
    "id": "v9DznMcW-qx_",
    "cell_id": "00107-346d96d4-62c0-498b-bd8d-2702ba06df77",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 457
   },
   "execution_count": 51,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "Q_values",
   "metadata": {
    "id": "UshrVxii-qyA",
    "outputId": "53a4d171-8220-4d27-c62b-39cdbbec9728",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "cell_id": "00108-0d6c1c87-1071-48d0-b92a-d925ea283b9b",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 61
   },
   "execution_count": 52,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "np.argmax(Q_values, axis=1) # optimal action for each state",
   "metadata": {
    "id": "Dei8dJiL-qyB",
    "outputId": "47bf2a7e-1de9-4dba-aec2-fa72d3d1e9b2",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "cell_id": "00109-527d8a17-0e71-4085-ad0a-6be39fecff17",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 61
   },
   "execution_count": 53,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "true_Q_value = history1[-1, 0, 0]\n\nfig, axes = plt.subplots(1, 2, figsize=(10, 4), sharey=True)\naxes[0].set_ylabel(\"Q-Value$(s_0, a_0)$\", fontsize=14)\naxes[0].set_title(\"Q-Value Iteration\", fontsize=14)\naxes[1].set_title(\"Q-Learning\", fontsize=14)\nfor ax, width, history in zip(axes, (50, 10000), (history1, history2)):\n    ax.plot([0, width], [true_Q_value, true_Q_value], \"k--\")\n    ax.plot(np.arange(width), history[:, 0, 0], \"b-\", linewidth=2)\n    ax.set_xlabel(\"Iterations\", fontsize=14)\n    ax.axis([0, width, 0, 24])\n\nsave_fig(\"q_value_plot\")",
   "metadata": {
    "id": "1Xf3M9zJ-qyC",
    "outputId": "a74f17c7-4893-4176-a560-3dcca40b7aa0",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 314
    },
    "cell_id": "00110-d669ff84-2000-402b-9cf6-3059f5294163",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 277
   },
   "execution_count": 54,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "# Deep Q-Network",
   "metadata": {
    "id": "eRYfkcjn-qyD",
    "cell_id": "00111-6b30b4c7-c36c-4543-b476-ed33388ddade",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 82
   }
  },
  {
   "cell_type": "markdown",
   "source": "Let's build the DQN. Given a state, it will estimate, for each possible action, the sum of discounted future rewards it can expect after it plays that action (but before it sees its outcome):",
   "metadata": {
    "id": "_4Y9WmQ--qyE",
    "cell_id": "00112-3dba9280-e5e0-474d-9baa-af53f21a3e96",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 74.78125
   }
  },
  {
   "cell_type": "code",
   "source": "keras.backend.clear_session()\ntf.random.set_seed(42)\nnp.random.seed(42)\n\nenv = gym.make(\"CartPole-v1\")\ninput_shape = [4] # == env.observation_space.shape\nn_outputs = 2 # == env.action_space.n\n\nmodel = keras.models.Sequential([\n    keras.layers.Dense(32, activation=\"elu\", input_shape=input_shape),\n    keras.layers.Dense(32, activation=\"elu\"),\n    keras.layers.Dense(n_outputs)\n])",
   "metadata": {
    "id": "aUsw9jYN-qyF",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "8e1accbe-14e0-4774-f118-2fe91733e9ed",
    "cell_id": "00113-8cf05e94-b6ae-4d68-a7f0-77b54e09f87a",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 277
   },
   "execution_count": 55,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "To select an action using this DQN, we just pick the action with the largest predicted Q-value. However, to ensure that the agent explores the environment, we choose a random action with probability `epsilon`.",
   "metadata": {
    "id": "CVCaPbZz-qyF",
    "cell_id": "00114-caf84236-a87c-4dca-8728-ead20eeb2b70",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 97.171875
   }
  },
  {
   "cell_type": "code",
   "source": "def epsilon_greedy_policy(state, epsilon=0):\n    if np.random.rand() < epsilon:\n        return np.random.randint(n_outputs)\n    else:\n        Q_values = model.predict(state[np.newaxis])\n        return np.argmax(Q_values[0])",
   "metadata": {
    "id": "Y0mfxaOA-qyH",
    "cell_id": "00115-193de255-f281-443b-8b3b-ed4846067978",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 151
   },
   "execution_count": 56,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "We will also need a replay memory. It will contain the agent's experiences, in the form of tuples: `(obs, action, reward, next_obs, done)`. We can use the `deque` class for that (but make sure to check out DeepMind's excellent [Reverb library](https://github.com/deepmind/reverb) for a much more robust implementation of experience replay):",
   "metadata": {
    "id": "rFm73nN_-qyM",
    "cell_id": "00116-b770f10b-29fe-48e6-9f63-62762a4dc7e0",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 119.5625
   }
  },
  {
   "cell_type": "code",
   "source": "from collections import deque\n\nreplay_memory = deque(maxlen=2000)",
   "metadata": {
    "id": "D9gfLlXp-qyN",
    "cell_id": "00117-31b309e7-85a8-46ec-9537-4ad836b477fb",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 97
   },
   "execution_count": 57,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "And let's create a function to sample experiences from the replay memory. It will return 5 NumPy arrays: `[obs, actions, rewards, next_obs, dones]`.",
   "metadata": {
    "id": "6rtiNLV0-qyO",
    "cell_id": "00118-8d0a9481-31cb-4c15-9f22-e0eb58d15515",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 74.78125
   }
  },
  {
   "cell_type": "code",
   "source": "def sample_experiences(batch_size):\n    indices = np.random.randint(len(replay_memory), size=batch_size)\n    batch = [replay_memory[index] for index in indices]\n    states, actions, rewards, next_states, dones = [\n        np.array([experience[field_index] for experience in batch])\n        for field_index in range(5)]\n    return states, actions, rewards, next_states, dones",
   "metadata": {
    "id": "xQVOfzV7-qyP",
    "cell_id": "00119-6159461b-c5aa-48d5-8e79-2553c1302744",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 169
   },
   "execution_count": 58,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Now we can create a function that will use the DQN to play one step, and record its experience in the replay memory:",
   "metadata": {
    "id": "jhDP9pLC-qya",
    "cell_id": "00120-bed47486-c09b-4ae7-8754-daa574d9b64d",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 74.78125
   }
  },
  {
   "cell_type": "code",
   "source": "def play_one_step(env, state, epsilon):\n    action = epsilon_greedy_policy(state, epsilon)\n    next_state, reward, done, info = env.step(action)\n    replay_memory.append((state, action, reward, next_state, done))\n    return next_state, reward, done, info",
   "metadata": {
    "id": "ykEMUBmS-qyc",
    "cell_id": "00121-76e3d6af-4a10-4359-b0e4-d47f371f5643",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 133
   },
   "execution_count": 59,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Lastly, let's create a function that will sample some experiences from the replay memory and perform a training step:\n\n**Notes**:\n* The first 3 releases of the 2nd edition were missing the `reshape()` operation which converts `target_Q_values` to a column vector (this is required by the `loss_fn()`).\n* The book uses a learning rate of 1e-3, but in the code below I use 1e-2, as it significantly improves training. I also tuned the learning rates of the DQN variants below.",
   "metadata": {
    "id": "6IhCQQMN-qyc",
    "cell_id": "00122-f83d75f8-c3a3-41cb-9bb6-cd3ca30dec8e",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 231.734375
   }
  },
  {
   "cell_type": "code",
   "source": "batch_size = 32\ndiscount_rate = 0.95\noptimizer = keras.optimizers.Adam(learning_rate=1e-2)\nloss_fn = keras.losses.mean_squared_error\n\ndef training_step(batch_size):\n    experiences = sample_experiences(batch_size)\n    states, actions, rewards, next_states, dones = experiences\n    next_Q_values = model.predict(next_states)\n    max_next_Q_values = np.max(next_Q_values, axis=1)\n    target_Q_values = (rewards +\n                       (1 - dones) * discount_rate * max_next_Q_values)\n    target_Q_values = target_Q_values.reshape(-1, 1)\n    mask = tf.one_hot(actions, n_outputs)\n    with tf.GradientTape() as tape:\n        all_Q_values = model(states)\n        Q_values = tf.reduce_sum(all_Q_values * mask, axis=1, keepdims=True)\n        loss = tf.reduce_mean(loss_fn(target_Q_values, Q_values))\n    grads = tape.gradient(loss, model.trainable_variables)\n    optimizer.apply_gradients(zip(grads, model.trainable_variables))",
   "metadata": {
    "id": "pepXUgvh-qyd",
    "cell_id": "00123-a24b1573-088c-46e9-b953-ab65fec4d3b7",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 403
   },
   "execution_count": 60,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "And now, let's train the model!",
   "metadata": {
    "id": "va-aQT8r-qye",
    "cell_id": "00124-f4dd6306-8ce4-474f-824a-42061c9696cd",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "code",
   "source": "env.seed(42)\nnp.random.seed(42)\ntf.random.set_seed(42)\n\nrewards = [] \nbest_score = 0",
   "metadata": {
    "id": "mTMaffE3-qye",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "93af5333-3710-49a1-ed97-eb168c8b1218",
    "cell_id": "00125-8e2fcd64-ad98-44b5-97b3-ad83ca7022b1",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 151
   },
   "execution_count": 61,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "for episode in range(600):\n    obs = env.reset()    \n    for step in range(200):\n        epsilon = max(1 - episode / 500, 0.01)\n        obs, reward, done, info = play_one_step(env, obs, epsilon)\n        if done:\n            break\n    rewards.append(step) # Not shown in the book\n    if step >= best_score: # Not shown\n        best_weights = model.get_weights() # Not shown\n        best_score = step # Not shown\n    print(\"\\rEpisode: {}, Steps: {}, eps: {:.3f}\".format(episode, step + 1, epsilon), end=\"\") # Not shown\n    if episode > 50:\n        training_step(batch_size)\n\nmodel.set_weights(best_weights)",
   "metadata": {
    "scrolled": true,
    "id": "J62QYrFt-qyf",
    "outputId": "21387eea-4646-4764-be8c-75c1308370de",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "cell_id": "00126-42594699-ba2d-4382-877d-88a21ea020fa",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 336
   },
   "execution_count": 62,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "plt.figure(figsize=(8, 4))\nplt.plot(rewards)\nplt.xlabel(\"Episode\", fontsize=14)\nplt.ylabel(\"Sum of rewards\", fontsize=14)\nsave_fig(\"dqn_rewards_plot\")\nplt.show()",
   "metadata": {
    "id": "w3V1Twy1-qyg",
    "outputId": "1a979349-56c6-4262-9a85-74128e56028e",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 314
    },
    "cell_id": "00127-912bf54c-fca1-47ee-8e1f-dd9b807caf41",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 151
   },
   "execution_count": 63,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "env.seed(42)\nstate = env.reset()\n\nframes = []\n\nfor step in range(200):\n    action = epsilon_greedy_policy(state)\n    state, reward, done, info = env.step(action)\n    if done:\n        break\n    img = env.render(mode=\"rgb_array\")\n    frames.append(img)\n    \nplot_animation(frames)",
   "metadata": {
    "id": "MHLQCcci-qyh",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 461
    },
    "outputId": "e53043df-8a47-42ac-8063-6e12b1c8b3a1",
    "cell_id": "00128-b3697557-3998-401c-9da6-e4d9af9d0371",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 295
   },
   "execution_count": 64,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Not bad at all! 😀",
   "metadata": {
    "id": "G6emyBU_-qyh",
    "cell_id": "00129-08b85014-1c3c-44c9-872f-38fe0e94ddc7",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "markdown",
   "source": "## Double DQN",
   "metadata": {
    "id": "V2l9WnhH-qyi",
    "cell_id": "00130-9996fd67-924e-4273-b907-7cd4b8598ec5",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 70
   }
  },
  {
   "cell_type": "code",
   "source": "keras.backend.clear_session()\ntf.random.set_seed(42)\nnp.random.seed(42)\n\nmodel = keras.models.Sequential([\n    keras.layers.Dense(32, activation=\"elu\", input_shape=[4]),\n    keras.layers.Dense(32, activation=\"elu\"),\n    keras.layers.Dense(n_outputs)\n])\n\ntarget = keras.models.clone_model(model)\ntarget.set_weights(model.get_weights())",
   "metadata": {
    "id": "MZipezVm-qyi",
    "cell_id": "00131-daea2478-c71c-418f-81a6-9d6f53b92ace",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 259
   },
   "execution_count": 65,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "batch_size = 32\ndiscount_rate = 0.95\noptimizer = keras.optimizers.Adam(learning_rate=6e-3)\nloss_fn = keras.losses.Huber()\n\ndef training_step(batch_size):\n    experiences = sample_experiences(batch_size)\n    states, actions, rewards, next_states, dones = experiences\n    next_Q_values = model.predict(next_states)\n    best_next_actions = np.argmax(next_Q_values, axis=1)\n    next_mask = tf.one_hot(best_next_actions, n_outputs).numpy()\n    next_best_Q_values = (target.predict(next_states) * next_mask).sum(axis=1)\n    target_Q_values = (rewards + \n                       (1 - dones) * discount_rate * next_best_Q_values)\n    target_Q_values = target_Q_values.reshape(-1, 1)\n    mask = tf.one_hot(actions, n_outputs)\n    with tf.GradientTape() as tape:\n        all_Q_values = model(states)\n        Q_values = tf.reduce_sum(all_Q_values * mask, axis=1, keepdims=True)\n        loss = tf.reduce_mean(loss_fn(target_Q_values, Q_values))\n    grads = tape.gradient(loss, model.trainable_variables)\n    optimizer.apply_gradients(zip(grads, model.trainable_variables))",
   "metadata": {
    "id": "9HJA7dxO-qyj",
    "cell_id": "00132-41cd9635-0e2b-4191-a6e7-788eb6d76c6a",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 439
   },
   "execution_count": 66,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "replay_memory = deque(maxlen=2000)",
   "metadata": {
    "id": "dZfFDKf3-qyk",
    "cell_id": "00133-9b0a5ce9-e7d2-44aa-80cf-5c8b99e95c96",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 61
   },
   "execution_count": 67,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "env.seed(42)\nnp.random.seed(42)\ntf.random.set_seed(42)\n\nrewards = []\nbest_score = 0\n\nfor episode in range(600):\n    obs = env.reset()    \n    for step in range(200):\n        epsilon = max(1 - episode / 500, 0.01)\n        obs, reward, done, info = play_one_step(env, obs, epsilon)\n        if done:\n            break\n    rewards.append(step)\n    if step >= best_score:\n        best_weights = model.get_weights()\n        best_score = step\n    print(\"\\rEpisode: {}, Steps: {}, eps: {:.3f}\".format(episode, step + 1, epsilon), end=\"\")\n    if episode >= 50:\n        training_step(batch_size)\n        if episode % 50 == 0:\n            target.set_weights(model.get_weights())\n    # Alternatively, you can do soft updates at each step:\n    #if episode >= 50:\n        #target_weights = target.get_weights()\n        #online_weights = model.get_weights()\n        #for index in range(len(target_weights)):\n        #    target_weights[index] = 0.99 * target_weights[index] + 0.01 * online_weights[index]\n        #target.set_weights(target_weights)\n\nmodel.set_weights(best_weights)",
   "metadata": {
    "id": "jH6qFp94-qyk",
    "outputId": "567a53b7-6a44-440b-a5b8-01979e970c30",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "cell_id": "00134-c2697af0-9eea-4a36-9573-2df26e7cfa11",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 624
   },
   "execution_count": 68,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "plt.figure(figsize=(8, 4))\nplt.plot(rewards)\nplt.xlabel(\"Episode\", fontsize=14)\nplt.ylabel(\"Sum of rewards\", fontsize=14)\nsave_fig(\"double_dqn_rewards_plot\")\nplt.show()",
   "metadata": {
    "id": "pl32r4FG-qyl",
    "outputId": "a353bddb-e228-441d-c1d4-d712a27adc01",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 314
    },
    "cell_id": "00135-e03a1264-681e-4de1-a989-2430e7d47656",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 151
   },
   "execution_count": 69,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "env.seed(43)\nstate = env.reset()\n\nframes = []\n\nfor step in range(200):\n    action = epsilon_greedy_policy(state)\n    state, reward, done, info = env.step(action)\n    if done:\n        break\n    img = env.render(mode=\"rgb_array\")\n    frames.append(img)\n   \nplot_animation(frames)",
   "metadata": {
    "scrolled": true,
    "id": "l-u33nkx-qym",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 461
    },
    "outputId": "9cbbec76-2747-490b-8214-ee248fc166cd",
    "cell_id": "00136-784f9957-ad2b-4c94-9448-6295ba0814b0",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 295
   },
   "execution_count": 70,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "# Dueling Double DQN",
   "metadata": {
    "id": "wvdTqqx1-qyn",
    "cell_id": "00137-212fd892-7184-43f1-a568-6638377b5e31",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 82
   }
  },
  {
   "cell_type": "code",
   "source": "keras.backend.clear_session()\ntf.random.set_seed(42)\nnp.random.seed(42)\n\nK = keras.backend\ninput_states = keras.layers.Input(shape=[4])\nhidden1 = keras.layers.Dense(32, activation=\"elu\")(input_states)\nhidden2 = keras.layers.Dense(32, activation=\"elu\")(hidden1)\nstate_values = keras.layers.Dense(1)(hidden2)\nraw_advantages = keras.layers.Dense(n_outputs)(hidden2)\nadvantages = raw_advantages - K.max(raw_advantages, axis=1, keepdims=True)\nQ_values = state_values + advantages\nmodel = keras.models.Model(inputs=[input_states], outputs=[Q_values])\n\ntarget = keras.models.clone_model(model)\ntarget.set_weights(model.get_weights())",
   "metadata": {
    "id": "Ymq-o3Uh-qyo",
    "cell_id": "00138-d6b13be4-90ae-456c-8f3a-ec3c1eb529ce",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 331
   },
   "execution_count": 71,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "batch_size = 32\ndiscount_rate = 0.95\noptimizer = keras.optimizers.Adam(learning_rate=7.5e-3)\nloss_fn = keras.losses.Huber()\n\ndef training_step(batch_size):\n    experiences = sample_experiences(batch_size)\n    states, actions, rewards, next_states, dones = experiences\n    next_Q_values = model.predict(next_states)\n    best_next_actions = np.argmax(next_Q_values, axis=1)\n    next_mask = tf.one_hot(best_next_actions, n_outputs).numpy()\n    next_best_Q_values = (target.predict(next_states) * next_mask).sum(axis=1)\n    target_Q_values = (rewards + \n                       (1 - dones) * discount_rate * next_best_Q_values)\n    target_Q_values = target_Q_values.reshape(-1, 1)\n    mask = tf.one_hot(actions, n_outputs)\n    with tf.GradientTape() as tape:\n        all_Q_values = model(states)\n        Q_values = tf.reduce_sum(all_Q_values * mask, axis=1, keepdims=True)\n        loss = tf.reduce_mean(loss_fn(target_Q_values, Q_values))\n    grads = tape.gradient(loss, model.trainable_variables)\n    optimizer.apply_gradients(zip(grads, model.trainable_variables))",
   "metadata": {
    "id": "NBfNgZrd-qyp",
    "cell_id": "00139-f0ea13d9-adab-40f9-8279-b800cc8b2fb6",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 439
   },
   "execution_count": 72,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "replay_memory = deque(maxlen=2000)",
   "metadata": {
    "id": "uJLi3Yi6-qyq",
    "cell_id": "00140-6d74b8e2-a8d7-4d3d-a287-c4ebe1ab4779",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 61
   },
   "execution_count": 73,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "env.seed(42)\nnp.random.seed(42)\ntf.random.set_seed(42)\n\nrewards = []\nbest_score = 0\n\nfor episode in range(600):\n    obs = env.reset()    \n    for step in range(200):\n        epsilon = max(1 - episode / 500, 0.01)\n        obs, reward, done, info = play_one_step(env, obs, epsilon)\n        if done:\n            break\n    rewards.append(step)\n    if step >= best_score:\n        best_weights = model.get_weights()\n        best_score = step\n    print(\"\\rEpisode: {}, Steps: {}, eps: {:.3f}\".format(episode, step + 1, epsilon), end=\"\")\n    if episode >= 50:\n        training_step(batch_size)\n        if episode % 50 == 0:\n            target.set_weights(model.get_weights())\n\nmodel.set_weights(best_weights)",
   "metadata": {
    "id": "AMEcshWL-qyq",
    "outputId": "39fa5c50-9d5a-4e6d-a5c2-6c80d9835607",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "cell_id": "00141-b2964662-b49b-44ee-abca-f22572f3c75c",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 498
   },
   "execution_count": 74,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "plt.plot(rewards)\nplt.xlabel(\"Episode\")\nplt.ylabel(\"Sum of rewards\")\nplt.show()",
   "metadata": {
    "id": "8Vmg6llE-qyr",
    "outputId": "e8aeeba7-07f2-4272-932e-e299c3e9080e",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 286
    },
    "cell_id": "00142-4395e9a6-630f-4785-a0fa-34dc6ddb2745",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 115
   },
   "execution_count": 75,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "env.seed(42)\nstate = env.reset()\n\nframes = []\n\nfor step in range(200):\n    action = epsilon_greedy_policy(state)\n    state, reward, done, info = env.step(action)\n    if done:\n        break\n    img = env.render(mode=\"rgb_array\")\n    frames.append(img)\n    \nplot_animation(frames)",
   "metadata": {
    "scrolled": true,
    "id": "df-zzMKM-qyr",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 461
    },
    "outputId": "e19af665-8f35-4559-b3db-d3da7a592f90",
    "cell_id": "00143-cb2a3869-fba9-4444-a753-d0cd47f0b950",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 295
   },
   "execution_count": 76,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "This looks like a pretty robust agent!",
   "metadata": {
    "id": "aYIrtqQW-qys",
    "cell_id": "00144-df9f320b-1ca1-401a-af5a-b7044a450962",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "code",
   "source": "env.close()",
   "metadata": {
    "id": "wR1wZqOf-qyv",
    "cell_id": "00145-0e96b6b9-fc65-49fd-92cf-4825204ea18a",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 61
   },
   "execution_count": 77,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "# Using TF-Agents to Beat Breakout",
   "metadata": {
    "id": "gscoRcsG-qyw",
    "cell_id": "00146-ac2df0b3-4006-4631-9624-8102fb9025ba",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 82
   }
  },
  {
   "cell_type": "markdown",
   "source": "Let's use TF-Agents to create an agent that will learn to play Breakout. We will use the Deep Q-Learning algorithm, so you can easily compare the components with the previous implementation, but TF-Agents implements many other (and more sophisticated) algorithms!",
   "metadata": {
    "id": "WRa82mr8-qyw",
    "cell_id": "00147-38a3fbbb-0d1a-4973-9e30-2ad92e3ea937",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 97.171875
   }
  },
  {
   "cell_type": "markdown",
   "source": "## TF-Agents Environments",
   "metadata": {
    "id": "LO4-h_c4-qyw",
    "cell_id": "00148-27170cc9-a84b-4631-9338-172c5c8fc4c6",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 70
   }
  },
  {
   "cell_type": "code",
   "source": "tf.random.set_seed(42)\nnp.random.seed(42)",
   "metadata": {
    "id": "E7lCiLjU-qyx",
    "cell_id": "00149-a9505338-64ba-45c5-b99a-16057e30f75b",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 79
   },
   "execution_count": 78,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "from tf_agents.environments import suite_gym\n\nenv = suite_gym.load(\"Breakout-v4\")\nenv",
   "metadata": {
    "id": "v1_B4dxt-qyx",
    "outputId": "533fbd35-63f1-4e16-d68e-7508cfc87c33",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 363
    },
    "cell_id": "00150-c999aa3f-72d0-494b-885a-e2594b822dcd",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 115
   },
   "execution_count": 81,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "env.gym",
   "metadata": {
    "id": "xtcDWnRs-qyy",
    "cell_id": "00151-65b08bc3-0553-423b-b77c-3aaec52019f3",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 61
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "env.seed(42)\nenv.reset()",
   "metadata": {
    "id": "ShTHmVlb-qyz",
    "cell_id": "00152-72ba4ddb-4c08-4f96-8732-47da96cb51c5",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 79
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "env.step(1) # Fire",
   "metadata": {
    "id": "mn3UtL-Q-qy0",
    "cell_id": "00153-0477547b-aa9f-4724-a8ec-3bbb765a1a77",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 61
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "img = env.render(mode=\"rgb_array\")\n\nplt.figure(figsize=(6, 8))\nplt.imshow(img)\nplt.axis(\"off\")\nsave_fig(\"breakout_plot\")\nplt.show()",
   "metadata": {
    "id": "GamEz7Kk-qy0",
    "cell_id": "00154-dbf24e48-2e91-457b-a644-9f39fb2688df",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 169
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "env.current_time_step()",
   "metadata": {
    "id": "gHxpVrx0-qy1",
    "cell_id": "00155-d8e4c419-d49b-414b-8639-e95ec86c3a59",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 61
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "## Environment Specifications",
   "metadata": {
    "id": "NovrdhBo-qy1",
    "cell_id": "00156-d1b91ead-9fc6-43f3-97b3-c219bde03af4",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 70
   }
  },
  {
   "cell_type": "code",
   "source": "env.observation_spec()",
   "metadata": {
    "id": "9X3YS2F8-qy1",
    "cell_id": "00157-d86aa079-9687-4323-a9c5-17477cac664b",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 61
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "env.action_spec()",
   "metadata": {
    "id": "w4lLs-JK-qy2",
    "cell_id": "00158-80248297-e799-4e22-8818-4e958bc5c8c7",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 61
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "env.time_step_spec()",
   "metadata": {
    "id": "EUCVcxWn-qy2",
    "cell_id": "00159-2e3b6954-86af-4bb8-a9f6-c87a7ab1a7e2",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 61
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "## Environment Wrappers",
   "metadata": {
    "id": "9tD3rvoC-qy3",
    "cell_id": "00160-e3f0b61d-94a2-4748-9f0b-a49c9fc8172e",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 70
   }
  },
  {
   "cell_type": "markdown",
   "source": "You can wrap a TF-Agents environments in a TF-Agents wrapper:",
   "metadata": {
    "id": "QT-p-nSW-qy6",
    "cell_id": "00161-0cf1e274-73bb-4274-a386-db06ca1bd663",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "code",
   "source": "from tf_agents.environments.wrappers import ActionRepeat\n\nrepeating_env = ActionRepeat(env, times=4)\nrepeating_env",
   "metadata": {
    "id": "EjrXXVz1-qy7",
    "cell_id": "00162-88813ab7-2a3e-4a03-a272-ea72f11edf1e",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 115
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "repeating_env.unwrapped",
   "metadata": {
    "id": "-Bj1OdD7-qy7",
    "cell_id": "00163-ba43f399-0132-467e-9646-3893c4a28ab6",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 61
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Here is the list of available wrappers:",
   "metadata": {
    "id": "L42VlkcT-qy8",
    "cell_id": "00164-b0e96b30-5233-4ff2-b0af-2e9b6541d36c",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "code",
   "source": "import tf_agents.environments.wrappers\n\nfor name in dir(tf_agents.environments.wrappers):\n    obj = getattr(tf_agents.environments.wrappers, name)\n    if hasattr(obj, \"__base__\") and issubclass(obj, tf_agents.environments.wrappers.PyEnvironmentBaseWrapper):\n        print(\"{:27s} {}\".format(name, obj.__doc__.split(\"\\n\")[0]))",
   "metadata": {
    "id": "a4_48FXm-qy8",
    "cell_id": "00165-f163f2db-6418-4350-8054-8089439b0b2f",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 156
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "The `suite_gym.load()` function can create an env and wrap it for you, both with TF-Agents environment wrappers and Gym environment wrappers (the latter are applied first).",
   "metadata": {
    "id": "PdCt6SWP-qy9",
    "cell_id": "00166-8f2ac2d7-807a-458a-8bee-1ae49295f296",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 74.78125
   }
  },
  {
   "cell_type": "code",
   "source": "from functools import partial\nfrom gym.wrappers import TimeLimit\n\nlimited_repeating_env = suite_gym.load(\n    \"Breakout-v4\",\n    gym_env_wrappers=[partial(TimeLimit, max_episode_steps=10000)],\n    env_wrappers=[partial(ActionRepeat, times=4)],\n)",
   "metadata": {
    "id": "ITH6KPPm-qy9",
    "cell_id": "00167-04f12346-76a2-4a31-84e1-b326626e5255",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 187
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "limited_repeating_env",
   "metadata": {
    "id": "MzN315Je-qy-",
    "cell_id": "00168-96652900-9b76-4481-a049-95a1b3f1d5d5",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 61
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "limited_repeating_env.unwrapped",
   "metadata": {
    "id": "oVmMLNZe-qy-",
    "cell_id": "00169-5774c21c-bfbf-4b16-8a10-2812ae851ea5",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 61
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Create an Atari Breakout environment, and wrap it to apply the default Atari preprocessing steps:",
   "metadata": {
    "id": "sbTqirC2-qy_",
    "cell_id": "00170-b21a8e22-1ec0-47eb-a1f4-0915176784c5",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "markdown",
   "source": "**Warning**: Breakout requires the player to press the FIRE button at the start of the game and after each life lost. The agent may take a very long time learning this because at first it seems that pressing FIRE just means losing faster. To speed up training considerably, we create and use a subclass of the `AtariPreprocessing` wrapper class called `AtariPreprocessingWithAutoFire` which presses FIRE (i.e., plays action 1) automatically at the start of the game and after each life lost. This is different from the book which uses the regular `AtariPreprocessing` wrapper.",
   "metadata": {
    "id": "U59c5oeC-qy_",
    "cell_id": "00171-aab83439-1521-47d4-8ff8-5158324b1656",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 164.34375
   }
  },
  {
   "cell_type": "code",
   "source": "from tf_agents.environments import suite_atari\nfrom tf_agents.environments.atari_preprocessing import AtariPreprocessing\nfrom tf_agents.environments.atari_wrappers import FrameStack4\n\nmax_episode_steps = 27000 # <=> 108k ALE frames since 1 step = 4 frames\nenvironment_name = \"BreakoutNoFrameskip-v4\"\n\nclass AtariPreprocessingWithAutoFire(AtariPreprocessing):\n    def reset(self, **kwargs):\n        obs = super().reset(**kwargs)\n        super().step(1) # FIRE to start\n        return obs\n    def step(self, action):\n        lives_before_action = self.ale.lives()\n        obs, rewards, done, info = super().step(action)\n        if self.ale.lives() < lives_before_action and not done:\n            super().step(1) # FIRE to start after life lost\n        return obs, rewards, done, info\n\nenv = suite_atari.load(\n    environment_name,\n    max_episode_steps=max_episode_steps,\n    gym_env_wrappers=[AtariPreprocessingWithAutoFire, FrameStack4])",
   "metadata": {
    "id": "CZU8kj95-qzA",
    "cell_id": "00172-e8f84817-83bd-4537-ad9a-89a6c92f5cb0",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 457
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "env",
   "metadata": {
    "id": "l8YMdIuI-qzB",
    "cell_id": "00173-cc8dc28e-2070-40f0-8477-28cef911e4dd",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 61
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Play a few steps just to see what happens:",
   "metadata": {
    "id": "6jgVEdQ0-qzB",
    "cell_id": "00174-09a04a26-be6a-4418-9a32-aeac10f20f7e",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "code",
   "source": "env.seed(42)\nenv.reset()\nfor _ in range(4):\n    time_step = env.step(3) # LEFT",
   "metadata": {
    "id": "vm1dqnt4-qzB",
    "cell_id": "00175-72a112bb-3732-45a7-9d63-242abec29ed4",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 115
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "def plot_observation(obs):\n    # Since there are only 3 color channels, you cannot display 4 frames\n    # with one primary color per frame. So this code computes the delta between\n    # the current frame and the mean of the other frames, and it adds this delta\n    # to the red and blue channels to get a pink color for the current frame.\n    obs = obs.astype(np.float32)\n    img = obs[..., :3]\n    current_frame_delta = np.maximum(obs[..., 3] - obs[..., :3].mean(axis=-1), 0.)\n    img[..., 0] += current_frame_delta\n    img[..., 2] += current_frame_delta\n    img = np.clip(img / 150, 0, 1)\n    plt.imshow(img)\n    plt.axis(\"off\")",
   "metadata": {
    "id": "5EN5uo9u-qzC",
    "cell_id": "00176-6b852d31-f0c2-41ec-ad30-5a9a3c072208",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 282
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "plt.figure(figsize=(6, 6))\nplot_observation(time_step.observation)\nsave_fig(\"preprocessed_breakout_plot\")\nplt.show()",
   "metadata": {
    "id": "PRSe97M_-qzD",
    "cell_id": "00177-3972896a-e350-41e9-abb0-a5e712324e05",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 115
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Convert the Python environment to a TF environment:",
   "metadata": {
    "id": "fbVGqoMl-qzD",
    "cell_id": "00178-62c5ea0c-02b3-4bf8-a91a-7d6f787520f1",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "code",
   "source": "from tf_agents.environments.tf_py_environment import TFPyEnvironment\n\ntf_env = TFPyEnvironment(env)",
   "metadata": {
    "id": "kve4DSqo-qzD",
    "cell_id": "00179-131d53a0-608e-4a51-b15e-e7c35d7c95c1",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 97
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "## Creating the DQN",
   "metadata": {
    "id": "wouZRmlb-qzE",
    "cell_id": "00180-5143fec3-4eb0-4ea6-8aa1-52c081412f53",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 70
   }
  },
  {
   "cell_type": "markdown",
   "source": "Create a small class to normalize the observations. Images are stored using bytes from 0 to 255 to use less RAM, but we want to pass floats from 0.0 to 1.0 to the neural network:",
   "metadata": {
    "id": "h6V7Uj7R-qzE",
    "cell_id": "00181-6eaa49e6-349b-48ee-b85a-f19177d4e73e",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 74.78125
   }
  },
  {
   "cell_type": "markdown",
   "source": "Create the Q-Network:",
   "metadata": {
    "id": "LzmT330p-qzE",
    "cell_id": "00182-b869457d-0264-4f7b-a5c5-7736fb488355",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "code",
   "source": "from tf_agents.networks.q_network import QNetwork\n\npreprocessing_layer = keras.layers.Lambda(\n                          lambda obs: tf.cast(obs, np.float32) / 255.)\nconv_layer_params=[(32, (8, 8), 4), (64, (4, 4), 2), (64, (3, 3), 1)]\nfc_layer_params=[512]\n\nq_net = QNetwork(\n    tf_env.observation_spec(),\n    tf_env.action_spec(),\n    preprocessing_layers=preprocessing_layer,\n    conv_layer_params=conv_layer_params,\n    fc_layer_params=fc_layer_params)",
   "metadata": {
    "id": "2ZnaHyaX-qzE",
    "cell_id": "00183-90377b50-c3b4-4470-8be7-cc4927ea0286",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 277
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Create the DQN Agent:",
   "metadata": {
    "id": "JxaGT1fK-qzF",
    "cell_id": "00184-6350a903-0b2d-4bb6-9dca-291c43306d98",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "code",
   "source": "from tf_agents.agents.dqn.dqn_agent import DqnAgent\n\ntrain_step = tf.Variable(0)\nupdate_period = 4 # run a training step every 4 collect steps\noptimizer = keras.optimizers.RMSprop(learning_rate=2.5e-4, rho=0.95, momentum=0.0,\n                                     epsilon=0.00001, centered=True)\nepsilon_fn = keras.optimizers.schedules.PolynomialDecay(\n    initial_learning_rate=1.0, # initial ε\n    decay_steps=250000 // update_period, # <=> 1,000,000 ALE frames\n    end_learning_rate=0.01) # final ε\nagent = DqnAgent(tf_env.time_step_spec(),\n                 tf_env.action_spec(),\n                 q_network=q_net,\n                 optimizer=optimizer,\n                 target_update_period=2000, # <=> 32,000 ALE frames\n                 td_errors_loss_fn=keras.losses.Huber(reduction=\"none\"),\n                 gamma=0.99, # discount factor\n                 train_step_counter=train_step,\n                 epsilon_greedy=lambda: epsilon_fn(train_step))\nagent.initialize()",
   "metadata": {
    "id": "QD72pu0v-qzG",
    "cell_id": "00185-0a465f49-7498-46f6-bac4-25e597be5421",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 408
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Create the replay buffer (this will use a lot of RAM, so please reduce the buffer size if you get an out-of-memory error):",
   "metadata": {
    "id": "m1JUdCHb-qzH",
    "cell_id": "00186-5b3b66af-18f9-4032-8f46-22d188e31b0e",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 74.78125
   }
  },
  {
   "cell_type": "markdown",
   "source": "**Warning**: we use a replay buffer of size 100,000 instead of 1,000,000 (as used in the book) since many people were getting OOM (Out-Of-Memory) errors.",
   "metadata": {
    "id": "PmaUo6ui-qzI",
    "cell_id": "00187-0507f103-4f99-4c74-966b-689d094c680d",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 74.78125
   }
  },
  {
   "cell_type": "code",
   "source": "from tf_agents.replay_buffers import tf_uniform_replay_buffer\n\nreplay_buffer = tf_uniform_replay_buffer.TFUniformReplayBuffer(\n    data_spec=agent.collect_data_spec,\n    batch_size=tf_env.batch_size,\n    max_length=100000) # reduce if OOM error\n\nreplay_buffer_observer = replay_buffer.add_batch",
   "metadata": {
    "id": "zBCRAHf4-qzI",
    "cell_id": "00188-f3ac6d05-8fe7-43b8-a467-a5b9eabea1bd",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 187
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Create a simple custom observer that counts and displays the number of times it is called (except when it is passed a trajectory that represents the boundary between two episodes, as this does not count as a step):",
   "metadata": {
    "id": "JqEQFlAj-qzJ",
    "cell_id": "00189-56385dfb-a161-468d-937e-53ba5b68eb8f",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 97.171875
   }
  },
  {
   "cell_type": "code",
   "source": "class ShowProgress:\n    def __init__(self, total):\n        self.counter = 0\n        self.total = total\n    def __call__(self, trajectory):\n        if not trajectory.is_boundary():\n            self.counter += 1\n        if self.counter % 100 == 0:\n            print(\"\\r{}/{}\".format(self.counter, self.total), end=\"\")",
   "metadata": {
    "id": "4ztBxTYN-qzK",
    "cell_id": "00190-e49d04eb-80df-4ddf-85f6-e3bd17b305c5",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 205
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Let's add some training metrics:",
   "metadata": {
    "id": "58SaacxG-qzL",
    "cell_id": "00191-5810e78d-2547-48bc-8f77-62fabe84350f",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "code",
   "source": "from tf_agents.metrics import tf_metrics\n\ntrain_metrics = [\n    tf_metrics.NumberOfEpisodes(),\n    tf_metrics.EnvironmentSteps(),\n    tf_metrics.AverageReturnMetric(),\n    tf_metrics.AverageEpisodeLengthMetric(),\n]",
   "metadata": {
    "id": "QNf78FkX-qzL",
    "cell_id": "00192-7f9095ee-511d-4d70-a364-5544e823655b",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 187
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "train_metrics[0].result()",
   "metadata": {
    "id": "qcPWmsQB-qzM",
    "cell_id": "00193-6e574f44-7b95-4995-99eb-e555cb2479d6",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 61
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "from tf_agents.eval.metric_utils import log_metrics\nimport logging\nlogging.getLogger().setLevel(logging.INFO)\nlog_metrics(train_metrics)",
   "metadata": {
    "id": "lTu7Yco8-qzM",
    "cell_id": "00194-04de01d7-8e8e-4a9b-9bb7-3020c5b75f97",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 115
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Create the collect driver:",
   "metadata": {
    "id": "h00aQUHx-qzN",
    "cell_id": "00195-09e5d566-318e-43d9-b9ca-c26212a05045",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "code",
   "source": "from tf_agents.drivers.dynamic_step_driver import DynamicStepDriver\n\ncollect_driver = DynamicStepDriver(\n    tf_env,\n    agent.collect_policy,\n    observers=[replay_buffer_observer] + train_metrics,\n    num_steps=update_period) # collect 4 steps for each training iteration",
   "metadata": {
    "id": "HlukQoEx-qzN",
    "cell_id": "00196-92970000-f8db-492c-bdbe-07cb766ea0db",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 169
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Collect the initial experiences, before training:",
   "metadata": {
    "id": "iUXJv9zN-qzN",
    "cell_id": "00197-87113dee-384a-480d-b8fd-83ba34a5d872",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "code",
   "source": "from tf_agents.policies.random_tf_policy import RandomTFPolicy\n\ninitial_collect_policy = RandomTFPolicy(tf_env.time_step_spec(),\n                                        tf_env.action_spec())\ninit_driver = DynamicStepDriver(\n    tf_env,\n    initial_collect_policy,\n    observers=[replay_buffer.add_batch, ShowProgress(20000)],\n    num_steps=20000) # <=> 80,000 ALE frames\nfinal_time_step, final_policy_state = init_driver.run()",
   "metadata": {
    "id": "folKvZWV-qzN",
    "cell_id": "00198-cdccd878-13c3-468a-9273-eb84bddbcb79",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 223
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Let's sample 2 sub-episodes, with 3 time steps each and display them:",
   "metadata": {
    "id": "nCeHW7JU-qzO",
    "cell_id": "00199-efd9a835-bbf3-4426-9d60-8c334e41db4f",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "markdown",
   "source": "**Note**: `replay_buffer.get_next()` is deprecated. We must use `replay_buffer.as_dataset(..., single_deterministic_pass=False)` instead.",
   "metadata": {
    "id": "246LxLm9-qzO",
    "cell_id": "00200-05df0b41-ec31-456f-8d86-762c6bc45de7",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 74.78125
   }
  },
  {
   "cell_type": "code",
   "source": "tf.random.set_seed(9) # chosen to show an example of trajectory at the end of an episode\n\n#trajectories, buffer_info = replay_buffer.get_next( # get_next() is deprecated\n#    sample_batch_size=2, num_steps=3)\n\ntrajectories, buffer_info = next(iter(replay_buffer.as_dataset(\n    sample_batch_size=2,\n    num_steps=3,\n    single_deterministic_pass=False)))",
   "metadata": {
    "id": "zhSle_nL-qzO",
    "cell_id": "00201-7cb7148e-ee40-408e-a29d-f58da850687b",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 210
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "trajectories._fields",
   "metadata": {
    "id": "qrKj-j4V-qzP",
    "cell_id": "00202-f90172ae-fea5-47d5-83ae-c02cac712fdf",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 61
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "trajectories.observation.shape",
   "metadata": {
    "id": "Sy6BdkO5-qzP",
    "cell_id": "00203-85d11438-e5e2-4160-ab4d-75d7e96f770e",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 61
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "from tf_agents.trajectories.trajectory import to_transition\n\ntime_steps, action_steps, next_time_steps = to_transition(trajectories)\ntime_steps.observation.shape",
   "metadata": {
    "id": "SYV7e6BW-qzQ",
    "cell_id": "00204-a81ce534-a404-4fcf-ac43-1f1486fadc2b",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 115
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "trajectories.step_type.numpy()",
   "metadata": {
    "id": "e6a6emT4-qzQ",
    "cell_id": "00205-ce0d3fe1-d8ac-42c2-b7cf-698ce5287238",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 61
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "plt.figure(figsize=(10, 6.8))\nfor row in range(2):\n    for col in range(3):\n        plt.subplot(2, 3, row * 3 + col + 1)\n        plot_observation(trajectories.observation[row, col].numpy())\nplt.subplots_adjust(left=0, right=1, bottom=0, top=1, hspace=0, wspace=0.02)\nsave_fig(\"sub_episodes_plot\")\nplt.show()",
   "metadata": {
    "id": "0F7culzB-qzQ",
    "cell_id": "00206-ba7a2043-3826-49d3-9ad0-3942290bc20f",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 187
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Now let's create the dataset:",
   "metadata": {
    "id": "Y47E_T-k-qzS",
    "cell_id": "00207-69d68888-efa8-4e25-9021-28bad159af00",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "code",
   "source": "dataset = replay_buffer.as_dataset(\n    sample_batch_size=64,\n    num_steps=2,\n    num_parallel_calls=3).prefetch(3)",
   "metadata": {
    "id": "joTj0B-i-qzT",
    "cell_id": "00208-5f388e6a-f47a-4478-93e1-da549b01f720",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 115
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Convert the main functions to TF Functions for better performance:",
   "metadata": {
    "id": "_GQLlx55-qzT",
    "cell_id": "00209-f5f1c28d-8b48-48a0-ba41-a9362b1c68fb",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "code",
   "source": "from tf_agents.utils.common import function\n\ncollect_driver.run = function(collect_driver.run)\nagent.train = function(agent.train)",
   "metadata": {
    "id": "avp8jT16-qzT",
    "cell_id": "00210-b504bcfb-74dc-47d2-96a9-137741db223a",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 115
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "And now we are ready to run the main loop!",
   "metadata": {
    "id": "RBqAJPyE-qzT",
    "cell_id": "00211-8a386e35-b5e8-4dbf-b58c-d99a1af354ca",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "code",
   "source": "def train_agent(n_iterations):\n    time_step = None\n    policy_state = agent.collect_policy.get_initial_state(tf_env.batch_size)\n    iterator = iter(dataset)\n    for iteration in range(n_iterations):\n        time_step, policy_state = collect_driver.run(time_step, policy_state)\n        trajectories, buffer_info = next(iterator)\n        train_loss = agent.train(trajectories)\n        print(\"\\r{} loss:{:.5f}\".format(\n            iteration, train_loss.loss.numpy()), end=\"\")\n        if iteration % 1000 == 0:\n            log_metrics(train_metrics)",
   "metadata": {
    "id": "UpuquA-j-qzU",
    "cell_id": "00212-6c6d0229-ae19-4687-8a76-962bb17b7576",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 259
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Run the next cell to train the agent for 50,000 steps. Then look at its behavior by running the following cell. You can run these two cells as many times as you wish. The agent will keep improving! It will likely take over 200,000 iterations for the agent to become reasonably good.",
   "metadata": {
    "id": "eo3NUkUY-qzU",
    "cell_id": "00213-437a1bd0-0382-4d1d-9c2f-ff1167889f26",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 97.171875
   }
  },
  {
   "cell_type": "code",
   "source": "train_agent(n_iterations=50000)",
   "metadata": {
    "id": "UGT10bMX-qzU",
    "cell_id": "00214-989453e5-7add-4c73-9ed7-9a82099266e9",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 61
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "frames = []\ndef save_frames(trajectory):\n    global frames\n    frames.append(tf_env.pyenv.envs[0].render(mode=\"rgb_array\"))\n\nwatch_driver = DynamicStepDriver(\n    tf_env,\n    agent.policy,\n    observers=[save_frames, ShowProgress(1000)],\n    num_steps=1000)\nfinal_time_step, final_policy_state = watch_driver.run()\n\nplot_animation(frames)",
   "metadata": {
    "id": "7ar9AXD--qzV",
    "cell_id": "00215-ce7d2492-4b09-4df4-86a9-167ea1f95360",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 277
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "If you want to save an animated GIF to show off your agent to your friends, here's one way to do it:",
   "metadata": {
    "id": "x1Sclp6U-qzV",
    "cell_id": "00216-10b9f18f-d2fd-401d-b03a-1befd76bbc66",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "code",
   "source": "import PIL\n\nimage_path = os.path.join(\"images\", \"rl\", \"breakout.gif\")\nframe_images = [PIL.Image.fromarray(frame) for frame in frames[:150]]\nframe_images[0].save(image_path, format='GIF',\n                     append_images=frame_images[1:],\n                     save_all=True,\n                     duration=30,\n                     loop=0)",
   "metadata": {
    "id": "BpiZIic--qzV",
    "cell_id": "00217-d41ba095-637c-481a-9b78-e50f73a29104",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 205
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "%%html\n<img src=\"images/rl/breakout.gif\" />",
   "metadata": {
    "id": "V7wLbfPc-qzW",
    "cell_id": "00218-9c55a01b-9928-4196-8d12-c72ce2fffee5",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 79
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "# Extra material",
   "metadata": {
    "id": "fsXcF2lr-qzW",
    "cell_id": "00219-696b2e68-5a35-49ff-b89e-6c0efba08b31",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 82
   }
  },
  {
   "cell_type": "markdown",
   "source": "## Deque vs Rotating List",
   "metadata": {
    "id": "IPEbwSJF-qzW",
    "cell_id": "00220-6d75c866-797b-4265-9032-8d4b88fd0e56",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 70
   }
  },
  {
   "cell_type": "markdown",
   "source": "The `deque` class offers fast append, but fairly slow random access (for large replay memories):",
   "metadata": {
    "id": "zl2Nk2UT-qzX",
    "cell_id": "00221-30f7c7fb-7c67-4cdf-bbe2-cc1b74507bb8",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "code",
   "source": "from collections import deque\nnp.random.seed(42)\n\nmem = deque(maxlen=1000000)\nfor i in range(1000000):\n    mem.append(i)\n[mem[i] for i in np.random.randint(1000000, size=5)]",
   "metadata": {
    "id": "x5uTSenq-qzX",
    "cell_id": "00222-1068ce04-885b-409c-98f9-ab929e6343b6",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 169
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "%timeit mem.append(1)",
   "metadata": {
    "id": "wDw0RYDo-qzX",
    "cell_id": "00223-f75e8cb7-a321-4ebb-9d4f-b4f2276f8058",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 61
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "%timeit [mem[i] for i in np.random.randint(1000000, size=5)]",
   "metadata": {
    "id": "1Cf37IVS-qzY",
    "cell_id": "00224-c36dbf32-9179-4899-9115-d03538c92be9",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 61
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Alternatively, you could use a rotating list like this `ReplayMemory` class. This would make random access faster for large replay memories:",
   "metadata": {
    "id": "5mbuOzAS-qzY",
    "cell_id": "00225-13f929a2-7826-4591-bf84-b954617a0feb",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 74.78125
   }
  },
  {
   "cell_type": "code",
   "source": "class ReplayMemory:\n    def __init__(self, max_size):\n        self.buffer = np.empty(max_size, dtype=np.object)\n        self.max_size = max_size\n        self.index = 0\n        self.size = 0\n\n    def append(self, obj):\n        self.buffer[self.index] = obj\n        self.size = min(self.size + 1, self.max_size)\n        self.index = (self.index + 1) % self.max_size\n\n    def sample(self, batch_size):\n        indices = np.random.randint(self.size, size=batch_size)\n        return self.buffer[indices]",
   "metadata": {
    "id": "54h8ici--qzY",
    "cell_id": "00226-47881f26-9ad0-4031-b312-d3aad64e6e5a",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 313
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "mem = ReplayMemory(max_size=1000000)\nfor i in range(1000000):\n    mem.append(i)\nmem.sample(5)",
   "metadata": {
    "id": "nW3bawOt-qzZ",
    "cell_id": "00227-3590477b-514f-4b4d-aa15-e14267881226",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 115
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "%timeit mem.append(1)",
   "metadata": {
    "id": "QLse5-8D-qzZ",
    "cell_id": "00228-67669dab-7985-4ea3-af0d-1a46428dac35",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 61
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "%timeit mem.sample(5)",
   "metadata": {
    "id": "zqaZZkdm-qza",
    "cell_id": "00229-91da113f-146f-49ef-81d4-88cdac70b4e7",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 61
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "## Creating a Custom TF-Agents Environment",
   "metadata": {
    "id": "P0MPgGZs-qza",
    "cell_id": "00230-465ee82d-9c99-4b97-b94d-0b21d983848a",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 70
   }
  },
  {
   "cell_type": "markdown",
   "source": "To create a custom TF-Agent environment, you just need to write a class that inherits from the `PyEnvironment` class and implements a few methods. For example, the following minimal environment represents a simple 4x4 grid. The agent starts in one corner (0,0) and must move to the opposite corner (3,3). The episode is done if the agent reaches the goal (it gets a +10 reward) or if the agent goes out of bounds (-1 reward). The actions are up (0), down (1), left (2) and right (3).",
   "metadata": {
    "id": "ibdOlIk--qzb",
    "cell_id": "00231-ebb96744-4300-4965-a623-d9a5a0684321",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 141.953125
   }
  },
  {
   "cell_type": "code",
   "source": "class MyEnvironment(tf_agents.environments.py_environment.PyEnvironment):\n    def __init__(self, discount=1.0):\n        super().__init__()\n        self._action_spec = tf_agents.specs.BoundedArraySpec(\n            shape=(), dtype=np.int32, name=\"action\", minimum=0, maximum=3)\n        self._observation_spec = tf_agents.specs.BoundedArraySpec(\n            shape=(4, 4), dtype=np.int32, name=\"observation\", minimum=0, maximum=1)\n        self.discount = discount\n\n    def action_spec(self):\n        return self._action_spec\n\n    def observation_spec(self):\n        return self._observation_spec\n\n    def _reset(self):\n        self._state = np.zeros(2, dtype=np.int32)\n        obs = np.zeros((4, 4), dtype=np.int32)\n        obs[self._state[0], self._state[1]] = 1\n        return tf_agents.trajectories.time_step.restart(obs)\n\n    def _step(self, action):\n        self._state += [(-1, 0), (+1, 0), (0, -1), (0, +1)][action]\n        reward = 0\n        obs = np.zeros((4, 4), dtype=np.int32)\n        done = (self._state.min() < 0 or self._state.max() > 3)\n        if not done:\n            obs[self._state[0], self._state[1]] = 1\n        if done or np.all(self._state == np.array([3, 3])):\n            reward = -1 if done else +10\n            return tf_agents.trajectories.time_step.termination(obs, reward)\n        else:\n            return tf_agents.trajectories.time_step.transition(obs, reward,\n                                                               self.discount)",
   "metadata": {
    "id": "UYXQVM17-qzb",
    "cell_id": "00232-af9c9ac7-5460-4448-983b-1812a1217014",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 660
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "The action and observation specs will generally be instances of the `ArraySpec` or `BoundedArraySpec` classes from the `tf_agents.specs` package (check out the other specs in this package as well). Optionally, you can also define a `render()` method, a `close()` method to free resources, as well as a `time_step_spec()` method if you don't want the `reward` and `discount` to be 32-bit float scalars. Note that the base class takes care of keeping track of the current time step, which is why we must implement `_reset()` and `_step()` rather than `reset()` and `step()`.\n",
   "metadata": {
    "id": "NGSOYRL1-qzc",
    "cell_id": "00233-c4f0c3b5-9bec-4619-9af4-df1821d5c8b2",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 164.34375
   }
  },
  {
   "cell_type": "code",
   "source": "my_env = MyEnvironment()\ntime_step = my_env.reset()\ntime_step",
   "metadata": {
    "id": "xczpnc0V-qzd",
    "cell_id": "00234-49209ad3-fbcd-48ca-b191-63220d675540",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 97
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "time_step = my_env.step(1)\ntime_step",
   "metadata": {
    "id": "tNTKChsA-qzd",
    "cell_id": "00235-2d901b11-1f25-460b-907b-65830d3e2295",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 79
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "# Exercise Solutions",
   "metadata": {
    "id": "vl5GEQXV-qzd",
    "cell_id": "00236-c6a2e46d-11a5-4dbf-b3d7-61cbc1220445",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 82
   }
  },
  {
   "cell_type": "markdown",
   "source": "## 1. to 7.\n\nSee Appendix A.",
   "metadata": {
    "id": "1GvVZVB8-qze",
    "cell_id": "00237-ff329047-aa8c-469d-b911-1128fdcfaa4d",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 108.390625
   }
  },
  {
   "cell_type": "markdown",
   "source": "## 8.\n_Exercise: Use policy gradients to solve OpenAI Gym's LunarLander-v2 environment. You will need to install the Box2D dependencies (`%pip install -U gym[box2d]`)._",
   "metadata": {
    "id": "2vUwwtWQ-qze",
    "cell_id": "00238-b4a75407-e4ef-4386-acff-ace1ef3a4ed2",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 130.78125
   }
  },
  {
   "cell_type": "markdown",
   "source": "Let's start by creating a LunarLander-v2 environment:",
   "metadata": {
    "id": "RxT1kMjH-qze",
    "cell_id": "00239-b9976e97-1fed-48b0-8526-95d431a10d83",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "code",
   "source": "env = gym.make(\"LunarLander-v2\")",
   "metadata": {
    "id": "kgCUl0EE-qzh",
    "cell_id": "00240-3b2261d4-4636-446c-a97e-e804d702d181",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 61
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "The inputs are 8-dimensional:",
   "metadata": {
    "id": "zv73hj_5-qzh",
    "cell_id": "00241-eb3e14a9-dfd3-4436-9ba0-5e66db4948b6",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "code",
   "source": "env.observation_space",
   "metadata": {
    "id": "GPZOvOId-qzh",
    "cell_id": "00242-9f5d6a26-0c9c-4404-8092-836f0503fbe7",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 61
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "env.seed(42)\nobs = env.reset()\nobs",
   "metadata": {
    "id": "BOcibCtQ-qzi",
    "cell_id": "00243-0f6df1e6-fa24-4adb-82ea-acd633f599e8",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 97
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "From the [source code](https://github.com/openai/gym/blob/master/gym/envs/box2d/lunar_lander.py), we can see that these each 8D observation (x, y, h, v, a, w, l, r) correspond to:\n* x,y: the coordinates of the spaceship. It starts at a random location near (0, 1.4) and must land near the target at (0, 0).\n* h,v: the horizontal and vertical speed of the spaceship. It starts with a small random speed.\n* a,w: the spaceship's angle and angular velocity.\n* l,r: whether the left or right leg touches the ground (1.0) or not (0.0).",
   "metadata": {
    "id": "yAmK8aiE-qzi",
    "cell_id": "00244-f6e66148-1069-4fb5-9d15-00d72e0c0462",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 201.34375
   }
  },
  {
   "cell_type": "markdown",
   "source": "The action space is discrete, with 4 possible actions:",
   "metadata": {
    "id": "cN0BuNyZ-qzi",
    "cell_id": "00245-5bec4dd1-813b-4e76-b067-4cd85756474a",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "code",
   "source": "env.action_space",
   "metadata": {
    "id": "rkoqqxDu-qzj",
    "cell_id": "00246-850ab031-22e9-487f-ad83-631fc7d92b2e",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 61
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Looking at the [LunarLander-v2's description](https://gym.openai.com/envs/LunarLander-v2/), these actions are:\n* do nothing\n* fire left orientation engine\n* fire main engine\n* fire right orientation engine",
   "metadata": {
    "id": "0YaoPjNd-qzj",
    "cell_id": "00247-ea8889fc-c0c5-4a92-b002-577c97810003",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 178.953125
   }
  },
  {
   "cell_type": "markdown",
   "source": "Let's create a simple policy network with 4 output neurons (one per possible action):",
   "metadata": {
    "id": "gC1UO08n-qzj",
    "cell_id": "00248-39a8dbc9-4c9b-4263-82d2-053a7a8ea02a",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "code",
   "source": "keras.backend.clear_session()\nnp.random.seed(42)\ntf.random.set_seed(42)\n\nn_inputs = env.observation_space.shape[0]\nn_outputs = env.action_space.n\n\nmodel = keras.models.Sequential([\n    keras.layers.Dense(32, activation=\"relu\", input_shape=[n_inputs]),\n    keras.layers.Dense(32, activation=\"relu\"),\n    keras.layers.Dense(n_outputs, activation=\"softmax\"),\n])",
   "metadata": {
    "id": "pXwrMdTS-qzj",
    "cell_id": "00249-1303bfc9-9462-4c61-8e26-16c52e61800d",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 259
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Note that we're using the softmax activation function in the output layer, instead of the sigmoid activation function \nlike we did for the CartPole-v1 environment. This is because we only had two possible actions for the CartPole-v1 environment, so a binary classification model worked fine. However, since we now how more than two possible actions, we need a multiclass classification model.",
   "metadata": {
    "id": "f0j3HOMa-qzk",
    "cell_id": "00250-24f9f2d8-6581-43a6-a636-a1a04b3607a5",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 141.953125
   }
  },
  {
   "cell_type": "markdown",
   "source": "Next, let's reuse the `play_one_step()` and `play_multiple_episodes()` functions we defined for the CartPole-v1 Policy Gradient code above, but we'll just tweak the `play_one_step()` function to account for the fact that the model is now a multiclass classification model rather than a binary classification model. We'll also tweak the `play_multiple_episodes()` function to call our tweaked `play_one_step()` function rather than the original one, and we add a big penalty if the spaceship does not land (or crash) before a maximum number of steps.",
   "metadata": {
    "id": "yU9cmz6D-qzk",
    "cell_id": "00251-2a2e93c1-86a1-4345-8974-c3e1848de7c4",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 164.34375
   }
  },
  {
   "cell_type": "code",
   "source": "def lander_play_one_step(env, obs, model, loss_fn):\n    with tf.GradientTape() as tape:\n        probas = model(obs[np.newaxis])\n        logits = tf.math.log(probas + keras.backend.epsilon())\n        action = tf.random.categorical(logits, num_samples=1)\n        loss = tf.reduce_mean(loss_fn(action, probas))\n    grads = tape.gradient(loss, model.trainable_variables)\n    obs, reward, done, info = env.step(action[0, 0].numpy())\n    return obs, reward, done, grads\n\ndef lander_play_multiple_episodes(env, n_episodes, n_max_steps, model, loss_fn):\n    all_rewards = []\n    all_grads = []\n    for episode in range(n_episodes):\n        current_rewards = []\n        current_grads = []\n        obs = env.reset()\n        for step in range(n_max_steps):\n            obs, reward, done, grads = lander_play_one_step(env, obs, model, loss_fn)\n            current_rewards.append(reward)\n            current_grads.append(grads)\n            if done:\n                break\n        all_rewards.append(current_rewards)\n        all_grads.append(current_grads)\n    return all_rewards, all_grads",
   "metadata": {
    "id": "BkvvRw5x-qzk",
    "cell_id": "00252-7f8fca12-1e41-4126-a8d8-b4e6b01e003d",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 516
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "We'll keep exactly the same `discount_rewards()` and `discount_and_normalize_rewards()` functions as earlier:",
   "metadata": {
    "id": "RrqmOnJC-qzk",
    "cell_id": "00253-fa129c5e-e274-4937-8bc0-03dc24983b5a",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 74.78125
   }
  },
  {
   "cell_type": "code",
   "source": "def discount_rewards(rewards, discount_rate):\n    discounted = np.array(rewards)\n    for step in range(len(rewards) - 2, -1, -1):\n        discounted[step] += discounted[step + 1] * discount_rate\n    return discounted\n\ndef discount_and_normalize_rewards(all_rewards, discount_rate):\n    all_discounted_rewards = [discount_rewards(rewards, discount_rate)\n                              for rewards in all_rewards]\n    flat_rewards = np.concatenate(all_discounted_rewards)\n    reward_mean = flat_rewards.mean()\n    reward_std = flat_rewards.std()\n    return [(discounted_rewards - reward_mean) / reward_std\n            for discounted_rewards in all_discounted_rewards]",
   "metadata": {
    "id": "cK5iaq7j-qzl",
    "cell_id": "00254-f781249e-0843-413a-a5b1-bbe65bbbb990",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 295
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Now let's define some hyperparameters:",
   "metadata": {
    "id": "TjhQdOpg-qzl",
    "cell_id": "00255-0591b6ef-12e6-44a2-8b49-3b19eae9ee00",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "code",
   "source": "n_iterations = 200\nn_episodes_per_update = 16\nn_max_steps = 1000\ndiscount_rate = 0.99",
   "metadata": {
    "id": "x3x8-w0O-qzl",
    "cell_id": "00256-c3053e78-2507-4f23-a492-87af7e604d76",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 115
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Again, since the model is a multiclass classification model, we must use the categorical cross-entropy rather than the binary cross-entropy. Moreover, since the `lander_play_one_step()` function sets the targets as class indices rather than class probabilities, we must use the `sparse_categorical_crossentropy()` loss function:",
   "metadata": {
    "id": "5HpI5xmO-qzm",
    "cell_id": "00257-e53fc6d4-8a46-4b14-9514-f7a2fcbea4e6",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 119.5625
   }
  },
  {
   "cell_type": "code",
   "source": "optimizer = keras.optimizers.Nadam(learning_rate=0.005)\nloss_fn = keras.losses.sparse_categorical_crossentropy",
   "metadata": {
    "id": "eWUjPkcV-qzm",
    "cell_id": "00258-5f80a22e-c9ae-460c-99fd-3d15d8606b7f",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 79
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "We're ready to train the model. Let's go!",
   "metadata": {
    "id": "gSHYkPfI-qzm",
    "cell_id": "00259-5daf78f3-f536-441b-a42d-e328a472f5d7",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "code",
   "source": "env.seed(42)\n\nmean_rewards = []\n\nfor iteration in range(n_iterations):\n    all_rewards, all_grads = lander_play_multiple_episodes(\n        env, n_episodes_per_update, n_max_steps, model, loss_fn)\n    mean_reward = sum(map(sum, all_rewards)) / n_episodes_per_update\n    print(\"\\rIteration: {}/{}, mean reward: {:.1f}  \".format(\n        iteration + 1, n_iterations, mean_reward), end=\"\")\n    mean_rewards.append(mean_reward)\n    all_final_rewards = discount_and_normalize_rewards(all_rewards,\n                                                       discount_rate)\n    all_mean_grads = []\n    for var_index in range(len(model.trainable_variables)):\n        mean_grads = tf.reduce_mean(\n            [final_reward * all_grads[episode_index][step][var_index]\n             for episode_index, final_rewards in enumerate(all_final_rewards)\n                 for step, final_reward in enumerate(final_rewards)], axis=0)\n        all_mean_grads.append(mean_grads)\n    optimizer.apply_gradients(zip(all_mean_grads, model.trainable_variables))",
   "metadata": {
    "id": "LkWjHyQi-qzn",
    "cell_id": "00260-a6204eb2-8716-4def-8d27-8f913bdeaa3c",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 421
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Let's look at the learning curve:",
   "metadata": {
    "id": "uaY4_xUv-qzo",
    "cell_id": "00261-1fd1d64b-dfa1-4be4-b8f3-8f4d27d01bd2",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "code",
   "source": "import matplotlib.pyplot as plt\n\nplt.plot(mean_rewards)\nplt.xlabel(\"Episode\")\nplt.ylabel(\"Mean reward\")\nplt.grid()\nplt.show()",
   "metadata": {
    "id": "1hfY0tOu-qzo",
    "cell_id": "00262-d91ddf76-24d5-4934-9d03-50888b74703c",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 169
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Now let's look at the result!",
   "metadata": {
    "id": "EEyIh2EG-qzp",
    "cell_id": "00263-9084fba3-f62a-47b7-9337-0fd6f8de3dda",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 52.390625
   }
  },
  {
   "cell_type": "code",
   "source": "def lander_render_policy_net(model, n_max_steps=500, seed=42):\n    frames = []\n    env = gym.make(\"LunarLander-v2\")\n    env.seed(seed)\n    tf.random.set_seed(seed)\n    np.random.seed(seed)\n    obs = env.reset()\n    for step in range(n_max_steps):\n        frames.append(env.render(mode=\"rgb_array\"))\n        probas = model(obs[np.newaxis])\n        logits = tf.math.log(probas + keras.backend.epsilon())\n        action = tf.random.categorical(logits, num_samples=1)\n        obs, reward, done, info = env.step(action[0, 0].numpy())\n        if done:\n            break\n    env.close()\n    return frames",
   "metadata": {
    "id": "vpZbiK5O-qzp",
    "cell_id": "00264-7364bfbb-c936-49c7-a17f-935e7944274c",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 349
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": "frames = lander_render_policy_net(model, seed=42)\nplot_animation(frames)",
   "metadata": {
    "id": "DrfjMvkh-qzp",
    "cell_id": "00265-12e44af6-6025-4d3d-bb94-3f223dd931ce",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 79
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "That's pretty good. You can try training it for longer and/or tweaking the hyperparameters to see if you can get it to go over 200.",
   "metadata": {
    "id": "8OH8bBDu-qzq",
    "cell_id": "00266-0c357f11-d33f-41c9-9549-6e6214fbeac1",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 74.78125
   }
  },
  {
   "cell_type": "markdown",
   "source": "## 9.\n_Exercise: Use TF-Agents to train an agent that can achieve a superhuman level at SpaceInvaders-v4 using any of the available algorithms._",
   "metadata": {
    "id": "XZ6M8DkA-qzq",
    "cell_id": "00267-c5acbc51-e513-4f4d-b003-46dc77fc59c3",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 130.78125
   }
  },
  {
   "cell_type": "markdown",
   "source": "Please follow the steps in the [Using TF-Agents to Beat Breakout](http://localhost:8888/notebooks/18_reinforcement_learning.ipynb#Using-TF-Agents-to-Beat-Breakout) section above, replacing `\"Breakout-v4\"` with `\"SpaceInvaders-v4\"`. There will be a few things to tweak, however. For example, the Space Invaders game does not require the user to press FIRE to begin the game. Instead, the player's laser cannon blinks for a few seconds then the game starts automatically. For better performance, you may want to skip this blinking phase (which lasts about 40 steps) at the beginning of each episode and after each life lost. Indeed, it's impossible to do anything at all during this phase, and nothing moves. One way to do this is to use the following custom environment wrapper, instead of the `AtariPreprocessingWithAutoFire` wrapper:",
   "metadata": {
    "id": "j7Bz-VQW-qzq",
    "cell_id": "00268-fff8c873-6f72-492a-a87e-7b3752e75dbb",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 209.125
   }
  },
  {
   "cell_type": "code",
   "source": "class AtariPreprocessingWithSkipStart(AtariPreprocessing):\n    def skip_frames(self, num_skip):\n        for _ in range(num_skip):\n          super().step(0) # NOOP for num_skip steps\n    def reset(self, **kwargs):\n        obs = super().reset(**kwargs)\n        self.skip_frames(40)\n        return obs\n    def step(self, action):\n        lives_before_action = self.ale.lives()\n        obs, rewards, done, info = super().step(action)\n        if self.ale.lives() < lives_before_action and not done:\n            self.skip_frames(40)\n        return obs, rewards, done, info",
   "metadata": {
    "id": "W6X0reqb-qzr",
    "cell_id": "00269-e0522d5f-b66e-4549-b40a-4ee75b590b19",
    "output_cleared": true,
    "deepnote_cell_type": "code",
    "deepnote_cell_height": 295
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "Moreover, you should always ensure that the preprocessed images contain enough information to play the game. For example, the blasts from the laser cannon and from the aliens should still be visible despite the limited resolution. In this particular case, the preprocessing we did for Breakout still works fine for Space Invaders, but that's something you should always check if you want try other games. To do this, you can let the agent play randomly for a while, and record the preprocessed frames, then play the animation and ensure the game still looks playable.\n\nYou will also need to let the agent train for quite a long time to get good performance. Sadly, the DQN algorithm is not able to reach superhuman level on Space Invaders, likely because humans are able to learn efficient long term strategies in this game, whereas DQN can only master fairly short strategies. But there has been a lot of progress over the past few years, and now many other RL algorithms are able to surpass human experts at this game. Check out the [State-of-the-Art for Space Invaders on paperswithcode.com](https://paperswithcode.com/sota/atari-games-on-atari-2600-space-invaders).",
   "metadata": {
    "id": "4RO_QJt1-qzs",
    "cell_id": "00270-33e71192-2337-45b2-98a9-b9a98daeeed7",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 312.6875
   }
  },
  {
   "cell_type": "markdown",
   "source": "## 10.\n_Exercise: If you have about $100 to spare, you can purchase a Raspberry Pi 3 plus some cheap robotics components, install TensorFlow on the Pi, and go wild! For an example, check out this [fun post](https://homl.info/2) by Lukas Biewald, or take a look at GoPiGo or BrickPi. Start with simple goals, like making the robot turn around to find the brightest angle (if it has a light sensor) or the closest object (if it has a sonar sensor), and move in that direction. Then you can start using Deep Learning: for example, if the robot has a camera, you can try to implement an object detection algorithm so it detects people and moves toward them. You can also try to use RL to make the agent learn on its own how to use the motors to achieve that goal. Have fun!_",
   "metadata": {
    "id": "XgVacSes-qzs",
    "cell_id": "00271-c999253b-7017-41a7-a506-69fa93c4cbca",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 265.125
   }
  },
  {
   "cell_type": "markdown",
   "source": "It's your turn now: go crazy, be creative, but most of all, be patient and move forward step by step, you can do it!",
   "metadata": {
    "id": "h5pvB94p-qzs",
    "cell_id": "00272-f59a5b3e-28f1-4f5f-8734-c18627ed2d0e",
    "deepnote_cell_type": "markdown",
    "deepnote_cell_height": 74.78125
   }
  },
  {
   "cell_type": "markdown",
   "source": "<a style='text-decoration:none;line-height:16px;display:flex;color:#5B5B62;padding:10px;justify-content:end;' href='https://deepnote.com?utm_source=created-in-deepnote-cell&projectId=8b8747e0-d5e2-4759-af6a-42b7fe569db7' target=\"_blank\">\n<img alt='Created in deepnote.com' style='display:inline;max-height:16px;margin:0px;margin-right:7.5px;' src='data:image/svg+xml;base64,PD94bWwgdmVyc2lvbj0iMS4wIiBlbmNvZGluZz0iVVRGLTgiPz4KPHN2ZyB3aWR0aD0iODBweCIgaGVpZ2h0PSI4MHB4IiB2aWV3Qm94PSIwIDAgODAgODAiIHZlcnNpb249IjEuMSIgeG1sbnM9Imh0dHA6Ly93d3cudzMub3JnLzIwMDAvc3ZnIiB4bWxuczp4bGluaz0iaHR0cDovL3d3dy53My5vcmcvMTk5OS94bGluayI+CiAgICA8IS0tIEdlbmVyYXRvcjogU2tldGNoIDU0LjEgKDc2NDkwKSAtIGh0dHBzOi8vc2tldGNoYXBwLmNvbSAtLT4KICAgIDx0aXRsZT5Hcm91cCAzPC90aXRsZT4KICAgIDxkZXNjPkNyZWF0ZWQgd2l0aCBTa2V0Y2guPC9kZXNjPgogICAgPGcgaWQ9IkxhbmRpbmciIHN0cm9rZT0ibm9uZSIgc3Ryb2tlLXdpZHRoPSIxIiBmaWxsPSJub25lIiBmaWxsLXJ1bGU9ImV2ZW5vZGQiPgogICAgICAgIDxnIGlkPSJBcnRib2FyZCIgdHJhbnNmb3JtPSJ0cmFuc2xhdGUoLTEyMzUuMDAwMDAwLCAtNzkuMDAwMDAwKSI+CiAgICAgICAgICAgIDxnIGlkPSJHcm91cC0zIiB0cmFuc2Zvcm09InRyYW5zbGF0ZSgxMjM1LjAwMDAwMCwgNzkuMDAwMDAwKSI+CiAgICAgICAgICAgICAgICA8cG9seWdvbiBpZD0iUGF0aC0yMCIgZmlsbD0iIzAyNjVCNCIgcG9pbnRzPSIyLjM3NjIzNzYyIDgwIDM4LjA0NzY2NjcgODAgNTcuODIxNzgyMiA3My44MDU3NTkyIDU3LjgyMTc4MjIgMzIuNzU5MjczOSAzOS4xNDAyMjc4IDMxLjY4MzE2ODMiPjwvcG9seWdvbj4KICAgICAgICAgICAgICAgIDxwYXRoIGQ9Ik0zNS4wMDc3MTgsODAgQzQyLjkwNjIwMDcsNzYuNDU0OTM1OCA0Ny41NjQ5MTY3LDcxLjU0MjI2NzEgNDguOTgzODY2LDY1LjI2MTk5MzkgQzUxLjExMjI4OTksNTUuODQxNTg0MiA0MS42NzcxNzk1LDQ5LjIxMjIyODQgMjUuNjIzOTg0Niw0OS4yMTIyMjg0IEMyNS40ODQ5Mjg5LDQ5LjEyNjg0NDggMjkuODI2MTI5Niw0My4yODM4MjQ4IDM4LjY0NzU4NjksMzEuNjgzMTY4MyBMNzIuODcxMjg3MSwzMi41NTQ0MjUgTDY1LjI4MDk3Myw2Ny42NzYzNDIxIEw1MS4xMTIyODk5LDc3LjM3NjE0NCBMMzUuMDA3NzE4LDgwIFoiIGlkPSJQYXRoLTIyIiBmaWxsPSIjMDAyODY4Ij48L3BhdGg+CiAgICAgICAgICAgICAgICA8cGF0aCBkPSJNMCwzNy43MzA0NDA1IEwyNy4xMTQ1MzcsMC4yNTcxMTE0MzYgQzYyLjM3MTUxMjMsLTEuOTkwNzE3MDEgODAsMTAuNTAwMzkyNyA4MCwzNy43MzA0NDA1IEM4MCw2NC45NjA0ODgyIDY0Ljc3NjUwMzgsNzkuMDUwMzQxNCAzNC4zMjk1MTEzLDgwIEM0Ny4wNTUzNDg5LDc3LjU2NzA4MDggNTMuNDE4MjY3Nyw3MC4zMTM2MTAzIDUzLjQxODI2NzcsNTguMjM5NTg4NSBDNTMuNDE4MjY3Nyw0MC4xMjg1NTU3IDM2LjMwMzk1NDQsMzcuNzMwNDQwNSAyNS4yMjc0MTcsMzcuNzMwNDQwNSBDMTcuODQzMDU4NiwzNy43MzA0NDA1IDkuNDMzOTE5NjYsMzcuNzMwNDQwNSAwLDM3LjczMDQ0MDUgWiIgaWQ9IlBhdGgtMTkiIGZpbGw9IiMzNzkzRUYiPjwvcGF0aD4KICAgICAgICAgICAgPC9nPgogICAgICAgIDwvZz4KICAgIDwvZz4KPC9zdmc+' > </img>\nCreated in <span style='font-weight:600;margin-left:4px;'>Deepnote</span></a>",
   "metadata": {
    "tags": [],
    "created_in_deepnote_cell": true,
    "deepnote_cell_type": "markdown"
   }
  }
 ],
 "nbformat": 4,
 "nbformat_minor": 0,
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.12"
  },
  "colab": {
   "name": "Copy of 18_reinforcement_learning.ipynb",
   "provenance": []
  },
  "accelerator": "GPU",
  "gpuClass": "standard",
  "deepnote_notebook_id": "0e46d105-8e41-4ca4-aada-55e0d39ab6fb",
  "deepnote": {},
  "deepnote_execution_queue": []
 }
}