early_rl.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Early Reinforcement Learning\n",
    "\n",
    "With the advances of modern computing power, the study of Reinforcement Learning is having a heyday. Machines are now able to learn complex tasks once thought to be solely in the domain of humans, from controlling the [heating and cooling in massive data centers](https://www.technologyreview.com/s/611902/google-just-gave-control-over-data-center-cooling-to-an-ai/) to beating [grandmasters at Starcraft](https://storage.googleapis.com/deepmind-media/research/alphastar/AlphaStar_unformatted.pdf). As magnificent as it may seem today, it had humble roots many decades ago. Seeing how far it's come, it's a wonder to see how far it will go!\n",
    "\n",
    "Let's take a step back in time to see how these early algorithms developed. Many of these algorithms make sense given the context of when they were created. Challenge yourself and see if you can come up with the same strategies given the right problem. Ok! Time to cozy up for a story.\n",
    "\n",
    "<img src=\"images/hero.jpg\" width=\"488\" height=\"172\">\n",
    "\n",
    "This is the hero of our story, the gumdrop emoji. It was enjoying a cool winter day building a snowman when suddenly, it slipped and fell on a frozen lake of death.\n",
    "\n",
    "\n",
    "<img src=\"images/lake.jpg\" width=\"900\" height=\"680\">\n",
    "\n",
    "The lake can be thought of as a 4 x 4 grid where the gumdrop can move left (0), down (1), right (2) and up (3). Unfortunately, this frozen lake of death has holes of death where if the gumdrop enters that square, it will fall in and meet an untimely demise. To make matters worse, the lake is surrounded by icy boulders that if the gumdrop attempts to climb, will have it slip back into its original position. Thankfully, at the bottom right of the lake is a safe ramp that leads to a nice warm cup of hot cocoa.\n",
    "\n",
    "## Set Up\n",
    "\n",
    "We can try and save the gumdrop ourselves! This is a common game people begin their Reinforcement Learning journey with, and is included in the OpenAI's python package [Gym](https://gym.openai.com/) and is aptly named [FrozenLake-v0](https://gym.openai.com/envs/FrozenLake-v0/) ([code](https://github.com/openai/gym/blob/master/gym/envs/toy_text/frozen_lake.py)). No time to waste, let's get the environment up and running. Run the below to install the needed libraries if they are not installed already."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting tensorflow==2.11\n",
      "  Downloading tensorflow-2.11.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (3.1 kB)\n",
      "Requirement already satisfied: absl-py>=1.0.0 in /opt/conda/lib/python3.10/site-packages (from tensorflow==2.11) (2.0.0)\n",
      "Collecting astunparse>=1.6.0 (from tensorflow==2.11)\n",
      "  Downloading astunparse-1.6.3-py2.py3-none-any.whl (12 kB)\n",
      "Collecting flatbuffers>=2.0 (from tensorflow==2.11)\n",
      "  Downloading flatbuffers-23.5.26-py2.py3-none-any.whl.metadata (850 bytes)\n",
      "Collecting gast<=0.4.0,>=0.2.1 (from tensorflow==2.11)\n",
      "  Downloading gast-0.4.0-py3-none-any.whl (9.8 kB)\n",
      "Collecting google-pasta>=0.1.1 (from tensorflow==2.11)\n",
      "  Downloading google_pasta-0.2.0-py3-none-any.whl (57 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m57.5/57.5 kB\u001b[0m \u001b[31m1.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0ma \u001b[36m0:00:01\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: grpcio<2.0,>=1.24.3 in /opt/conda/lib/python3.10/site-packages (from tensorflow==2.11) (1.60.0)\n",
      "Collecting h5py>=2.9.0 (from tensorflow==2.11)\n",
      "  Downloading h5py-3.10.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (2.5 kB)\n",
      "Collecting keras<2.12,>=2.11.0 (from tensorflow==2.11)\n",
      "  Downloading keras-2.11.0-py2.py3-none-any.whl.metadata (1.4 kB)\n",
      "Collecting libclang>=13.0.0 (from tensorflow==2.11)\n",
      "  Downloading libclang-16.0.6-py2.py3-none-manylinux2010_x86_64.whl.metadata (5.2 kB)\n",
      "Requirement already satisfied: numpy>=1.20 in /opt/conda/lib/python3.10/site-packages (from tensorflow==2.11) (1.25.2)\n",
      "Collecting opt-einsum>=2.3.2 (from tensorflow==2.11)\n",
      "  Downloading opt_einsum-3.3.0-py3-none-any.whl (65 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m65.5/65.5 kB\u001b[0m \u001b[31m5.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: packaging in /opt/conda/lib/python3.10/site-packages (from tensorflow==2.11) (23.2)\n",
      "Collecting protobuf<3.20,>=3.9.2 (from tensorflow==2.11)\n",
      "  Downloading protobuf-3.19.6-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (787 bytes)\n",
      "Requirement already satisfied: setuptools in /opt/conda/lib/python3.10/site-packages (from tensorflow==2.11) (69.0.3)\n",
      "Requirement already satisfied: six>=1.12.0 in /opt/conda/lib/python3.10/site-packages (from tensorflow==2.11) (1.16.0)\n",
      "Collecting tensorboard<2.12,>=2.11 (from tensorflow==2.11)\n",
      "  Downloading tensorboard-2.11.2-py3-none-any.whl.metadata (1.9 kB)\n",
      "Collecting tensorflow-estimator<2.12,>=2.11.0 (from tensorflow==2.11)\n",
      "  Downloading tensorflow_estimator-2.11.0-py2.py3-none-any.whl.metadata (1.3 kB)\n",
      "Collecting termcolor>=1.1.0 (from tensorflow==2.11)\n",
      "  Downloading termcolor-2.4.0-py3-none-any.whl.metadata (6.1 kB)\n",
      "Requirement already satisfied: typing-extensions>=3.6.6 in /opt/conda/lib/python3.10/site-packages (from tensorflow==2.11) (4.9.0)\n",
      "Requirement already satisfied: wrapt>=1.11.0 in /opt/conda/lib/python3.10/site-packages (from tensorflow==2.11) (1.16.0)\n",
      "Collecting tensorflow-io-gcs-filesystem>=0.23.1 (from tensorflow==2.11)\n",
      "  Downloading tensorflow_io_gcs_filesystem-0.36.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (14 kB)\n",
      "Requirement already satisfied: wheel<1.0,>=0.23.0 in /opt/conda/lib/python3.10/site-packages (from astunparse>=1.6.0->tensorflow==2.11) (0.42.0)\n",
      "Requirement already satisfied: google-auth<3,>=1.6.3 in /opt/conda/lib/python3.10/site-packages (from tensorboard<2.12,>=2.11->tensorflow==2.11) (2.26.2)\n",
      "Collecting google-auth-oauthlib<0.5,>=0.4.1 (from tensorboard<2.12,>=2.11->tensorflow==2.11)\n",
      "  Downloading google_auth_oauthlib-0.4.6-py2.py3-none-any.whl (18 kB)\n",
      "Collecting markdown>=2.6.8 (from tensorboard<2.12,>=2.11->tensorflow==2.11)\n",
      "  Downloading Markdown-3.5.2-py3-none-any.whl.metadata (7.0 kB)\n",
      "Requirement already satisfied: requests<3,>=2.21.0 in /opt/conda/lib/python3.10/site-packages (from tensorboard<2.12,>=2.11->tensorflow==2.11) (2.31.0)\n",
      "Collecting tensorboard-data-server<0.7.0,>=0.6.0 (from tensorboard<2.12,>=2.11->tensorflow==2.11)\n",
      "  Downloading tensorboard_data_server-0.6.1-py3-none-manylinux2010_x86_64.whl (4.9 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m4.9/4.9 MB\u001b[0m \u001b[31m27.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hCollecting tensorboard-plugin-wit>=1.6.0 (from tensorboard<2.12,>=2.11->tensorflow==2.11)\n",
      "  Downloading tensorboard_plugin_wit-1.8.1-py3-none-any.whl.metadata (873 bytes)\n",
      "Collecting werkzeug>=1.0.1 (from tensorboard<2.12,>=2.11->tensorflow==2.11)\n",
      "  Downloading werkzeug-3.0.1-py3-none-any.whl.metadata (4.1 kB)\n",
      "Requirement already satisfied: cachetools<6.0,>=2.0.0 in /opt/conda/lib/python3.10/site-packages (from google-auth<3,>=1.6.3->tensorboard<2.12,>=2.11->tensorflow==2.11) (5.3.2)\n",
      "Requirement already satisfied: pyasn1-modules>=0.2.1 in /opt/conda/lib/python3.10/site-packages (from google-auth<3,>=1.6.3->tensorboard<2.12,>=2.11->tensorflow==2.11) (0.3.0)\n",
      "Requirement already satisfied: rsa<5,>=3.1.4 in /opt/conda/lib/python3.10/site-packages (from google-auth<3,>=1.6.3->tensorboard<2.12,>=2.11->tensorflow==2.11) (4.9)\n",
      "Requirement already satisfied: requests-oauthlib>=0.7.0 in /opt/conda/lib/python3.10/site-packages (from google-auth-oauthlib<0.5,>=0.4.1->tensorboard<2.12,>=2.11->tensorflow==2.11) (1.3.1)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /opt/conda/lib/python3.10/site-packages (from requests<3,>=2.21.0->tensorboard<2.12,>=2.11->tensorflow==2.11) (3.3.2)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /opt/conda/lib/python3.10/site-packages (from requests<3,>=2.21.0->tensorboard<2.12,>=2.11->tensorflow==2.11) (3.6)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /opt/conda/lib/python3.10/site-packages (from requests<3,>=2.21.0->tensorboard<2.12,>=2.11->tensorflow==2.11) (1.26.18)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /opt/conda/lib/python3.10/site-packages (from requests<3,>=2.21.0->tensorboard<2.12,>=2.11->tensorflow==2.11) (2023.11.17)\n",
      "Requirement already satisfied: MarkupSafe>=2.1.1 in /opt/conda/lib/python3.10/site-packages (from werkzeug>=1.0.1->tensorboard<2.12,>=2.11->tensorflow==2.11) (2.1.3)\n",
      "Requirement already satisfied: pyasn1<0.6.0,>=0.4.6 in /opt/conda/lib/python3.10/site-packages (from pyasn1-modules>=0.2.1->google-auth<3,>=1.6.3->tensorboard<2.12,>=2.11->tensorflow==2.11) (0.5.1)\n",
      "Requirement already satisfied: oauthlib>=3.0.0 in /opt/conda/lib/python3.10/site-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<0.5,>=0.4.1->tensorboard<2.12,>=2.11->tensorflow==2.11) (3.2.2)\n",
      "Downloading tensorflow-2.11.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (588.3 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m588.3/588.3 MB\u001b[0m \u001b[31m3.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m:00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hDownloading flatbuffers-23.5.26-py2.py3-none-any.whl (26 kB)\n",
      "Downloading h5py-3.10.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (4.8 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m4.8/4.8 MB\u001b[0m \u001b[31m57.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hDownloading keras-2.11.0-py2.py3-none-any.whl (1.7 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.7/1.7 MB\u001b[0m \u001b[31m89.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading libclang-16.0.6-py2.py3-none-manylinux2010_x86_64.whl (22.9 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m22.9/22.9 MB\u001b[0m \u001b[31m48.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hDownloading protobuf-3.19.6-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.1 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.1/1.1 MB\u001b[0m \u001b[31m74.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading tensorboard-2.11.2-py3-none-any.whl (6.0 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m6.0/6.0 MB\u001b[0m \u001b[31m16.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hDownloading tensorflow_estimator-2.11.0-py2.py3-none-any.whl (439 kB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m439.2/439.2 kB\u001b[0m \u001b[31m988.8 kB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hDownloading tensorflow_io_gcs_filesystem-0.36.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (5.1 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m5.1/5.1 MB\u001b[0m \u001b[31m10.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hDownloading termcolor-2.4.0-py3-none-any.whl (7.7 kB)\n",
      "Downloading Markdown-3.5.2-py3-none-any.whl (103 kB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m103.9/103.9 kB\u001b[0m \u001b[31m610.1 kB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hDownloading tensorboard_plugin_wit-1.8.1-py3-none-any.whl (781 kB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m781.3/781.3 kB\u001b[0m \u001b[31m2.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m:00:01\u001b[0m\n",
      "\u001b[?25hDownloading werkzeug-3.0.1-py3-none-any.whl (226 kB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m226.7/226.7 kB\u001b[0m \u001b[31m921.7 kB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hInstalling collected packages: tensorboard-plugin-wit, libclang, flatbuffers, werkzeug, termcolor, tensorflow-io-gcs-filesystem, tensorflow-estimator, tensorboard-data-server, protobuf, opt-einsum, markdown, keras, h5py, google-pasta, gast, astunparse, google-auth-oauthlib, tensorboard, tensorflow\n",
      "\u001b[33m  WARNING: The script markdown_py is installed in '/home/jupyter/.local/bin' which is not on PATH.\n",
      "  Consider adding this directory to PATH or, if you prefer to suppress this warning, use --no-warn-script-location.\u001b[0m\u001b[33m\n",
      "\u001b[0m\u001b[33m  WARNING: The script google-oauthlib-tool is installed in '/home/jupyter/.local/bin' which is not on PATH.\n",
      "  Consider adding this directory to PATH or, if you prefer to suppress this warning, use --no-warn-script-location.\u001b[0m\u001b[33m\n",
      "\u001b[0m\u001b[33m  WARNING: The script tensorboard is installed in '/home/jupyter/.local/bin' which is not on PATH.\n",
      "  Consider adding this directory to PATH or, if you prefer to suppress this warning, use --no-warn-script-location.\u001b[0m\u001b[33m\n",
      "\u001b[0m\u001b[33m  WARNING: The scripts estimator_ckpt_converter, import_pb_to_tensorboard, saved_model_cli, tensorboard, tf_upgrade_v2, tflite_convert, toco and toco_from_protos are installed in '/home/jupyter/.local/bin' which is not on PATH.\n",
      "  Consider adding this directory to PATH or, if you prefer to suppress this warning, use --no-warn-script-location.\u001b[0m\u001b[33m\n",
      "\u001b[0m\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
      "tensorboardx 2.6.2.2 requires protobuf>=3.20, but you have protobuf 3.19.6 which is incompatible.\u001b[0m\u001b[31m\n",
      "\u001b[0mSuccessfully installed astunparse-1.6.3 flatbuffers-23.5.26 gast-0.4.0 google-auth-oauthlib-0.4.6 google-pasta-0.2.0 h5py-3.10.0 keras-2.11.0 libclang-16.0.6 markdown-3.5.2 opt-einsum-3.3.0 protobuf-3.19.6 tensorboard-2.11.2 tensorboard-data-server-0.6.1 tensorboard-plugin-wit-1.8.1 tensorflow-2.11.0 tensorflow-estimator-2.11.0 tensorflow-io-gcs-filesystem-0.36.0 termcolor-2.4.0 werkzeug-3.0.1\n"
     ]
    }
   ],
   "source": [
    "# Replace with the Tensorflow version that you selected while creating the instance.\n",
    "!pip install tensorflow==2.11 --user"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**NOTE**: In the output of the above cell you may ignore any WARNINGS or ERRORS related to the dependency resolver."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you get any related errors mentioned above please rerun the above cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting gym==0.12.5\n",
      "  Downloading gym-0.12.5.tar.gz (1.5 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.5/1.5 MB\u001b[0m \u001b[31m7.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25h  Preparing metadata (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25hRequirement already satisfied: scipy in /opt/conda/lib/python3.10/site-packages (from gym==0.12.5) (1.11.4)\n",
      "Requirement already satisfied: numpy>=1.10.4 in /opt/conda/lib/python3.10/site-packages (from gym==0.12.5) (1.25.2)\n",
      "Requirement already satisfied: six in /opt/conda/lib/python3.10/site-packages (from gym==0.12.5) (1.16.0)\n",
      "Collecting pyglet>=1.2.0 (from gym==0.12.5)\n",
      "  Downloading pyglet-2.0.10-py3-none-any.whl.metadata (8.5 kB)\n",
      "Downloading pyglet-2.0.10-py3-none-any.whl (858 kB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m858.3/858.3 kB\u001b[0m \u001b[31m54.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hBuilding wheels for collected packages: gym\n",
      "  Building wheel for gym (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for gym: filename=gym-0.12.5-py3-none-any.whl size=1613771 sha256=3cd7bb4958af12655663d92694afd88e300f2edd1c390265518ea1f4d353b96d\n",
      "  Stored in directory: /home/jupyter/.cache/pip/wheels/6d/2a/ab/a34d08c34498b85996bd490dd3e4405dfa1b8b0b06aab5f2b9\n",
      "Successfully built gym\n",
      "Installing collected packages: pyglet, gym\n",
      "Successfully installed gym-0.12.5 pyglet-2.0.10\n"
     ]
    }
   ],
   "source": [
    "!pip install gym==0.12.5 --user"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are [four methods from Gym](http://gym.openai.com/docs/) that are going to be useful to us in order to save the gumdrop.\n",
    "* `make` allows us to build the environment or game that we can pass actions to\n",
    "* `reset` will reset an environment to it's starting configuration and return the state of the player\n",
    "* `render` displays the environment for human eyes\n",
    "* `step` takes an action and returns the player's next state.\n",
    "\n",
    "Let's make, reset, and render the game. The output is an ANSI string with the following characters:\n",
    "* `S` for starting point\n",
    "* `F` for frozen\n",
    "* `H` for hole\n",
    "* `G` for goal\n",
    "* A red square indicates the current position"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note**: Restart the kernel if the above libraries needed to be installed\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n"
     ]
    }
   ],
   "source": [
    "import gym\n",
    "import numpy as np\n",
    "import random\n",
    "env = gym.make('FrozenLake-v0', is_slippery=False)\n",
    "state = env.reset()\n",
    "env.render()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we print the state we'll get `0`. This is telling us which square we're in. Each square is labeled from `0` to `15` from left to right, top to bottom, like this:\n",
    "\n",
    "| | | | |\n",
    "|-|-|-|-|\n",
    "|0|1|2|3|\n",
    "|4|5|6|7|\n",
    "|8|9|10|11|\n",
    "|12|13|14|15|\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n"
     ]
    }
   ],
   "source": [
    "print(state)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can make a simple print function to let us know whether it's game won, game over, or game on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def print_state(state, done):\n",
    "    statement = \"Still Alive!\"\n",
    "    if done:\n",
    "        statement = \"Cocoa Time!\" if state == 15 else \"Game Over!\" \n",
    "    print(state, \"-\", statement)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can control the gumdrop ourselves with the `step` method. Run the below cell over and over again trying to move from the starting position to the goal. Good luck!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  (Right)\n",
      "S\u001b[41mF\u001b[0mFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "1 - Still Alive!\n"
     ]
    }
   ],
   "source": [
    "#0 left\n",
    "#1 down\n",
    "#2 right\n",
    "#3 up\n",
    "\n",
    "# Uncomment to reset the game\n",
    "#env.reset()\n",
    "action = 2  # Change me, please!\n",
    "state, _, done, _ = env.step(action)\n",
    "env.render()\n",
    "print_state(state, done)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Were you able to reach the hot chocolate? If so, great job! There are multiple paths through the maze. One solution is `[1, 1, 2, 2, 1, 2]`. Let's loop through our actions in order to get used to interacting with the environment programmatically."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  (Down)\n",
      "SFFF\n",
      "\u001b[41mF\u001b[0mHFH\n",
      "FFFH\n",
      "HFFG\n",
      "4 - Still Alive!\n",
      "  (Down)\n",
      "SFFF\n",
      "FHFH\n",
      "\u001b[41mF\u001b[0mFFH\n",
      "HFFG\n",
      "8 - Still Alive!\n",
      "  (Right)\n",
      "SFFF\n",
      "FHFH\n",
      "F\u001b[41mF\u001b[0mFH\n",
      "HFFG\n",
      "9 - Still Alive!\n",
      "  (Right)\n",
      "SFFF\n",
      "FHFH\n",
      "FF\u001b[41mF\u001b[0mH\n",
      "HFFG\n",
      "10 - Still Alive!\n",
      "  (Down)\n",
      "SFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HF\u001b[41mF\u001b[0mG\n",
      "14 - Still Alive!\n",
      "  (Right)\n",
      "SFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFF\u001b[41mG\u001b[0m\n",
      "15 - Cocoa Time!\n"
     ]
    }
   ],
   "source": [
    "def play_game(actions):\n",
    "    state = env.reset()\n",
    "    step = 0\n",
    "    done = False\n",
    "\n",
    "    while not done and step < len(actions):\n",
    "        action = actions[step]\n",
    "        state, _, done, _ = env.step(action)\n",
    "        env.render()\n",
    "        step += 1\n",
    "        print_state(state, done)\n",
    "        \n",
    "actions = [1, 1, 2, 2, 1, 2]  # Replace with your favorite path.\n",
    "play_game(actions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Nice, so we know how to get through the maze, but how do we teach that to the gumdrop? It's just some bytes in an android phone. It doesn't have our human insight.\n",
    "\n",
    "We could give it our list of actions directly, but then it would be copying us and not really learning. This was a tricky one to the mathematicians and computer scientists originally trying to solve this problem. How do we teach a machine to do this without human insight?\n",
    "\n",
    "## Value Iteration\n",
    "\n",
    "Let's turn the clock back on our time machines to 1957 to meet Mr. [Richard Bellman](https://en.wikipedia.org/wiki/Richard_E._Bellman). Bellman started his academic career in mathematics, but due to World War II, left his postgraduate studies at John Hopkins to teach electronics as part of the war effort (as chronicled by J. J. O'Connor and E. F. Robertson [here](https://www-history.mcs.st-andrews.ac.uk/Biographies/Bellman.html)). When the war was over, and it came time for him to focus on his next area of research, he became fascinated with [Dynamic Programming](https://en.wikipedia.org/wiki/Dynamic_programming): the idea of breaking a problem down into sub-problems and using recursion to solve the larger problem.\n",
    "\n",
    "Eventually, his research landed him on [Markov Decision Processes](https://en.wikipedia.org/wiki/Markov_decision_process). These processes are a graphical way of representing how to make a decision based on a current state. States are connected to other states with positive and negative rewards that can be picked up along the way.\n",
    "\n",
    "Sound familiar at all? Perhaps our Frozen Lake?\n",
    "\n",
    "In the lake case, each cell is a state. The `H`s and the `G` are a special type of state called a \"Terminal State\", meaning they can be entered, but they have no leaving connections. What of rewards? Let's say the value of losing our life is the negative opposite of getting to the goal and staying alive. Thus, we can assign the reward of entering a death hole as -1, and the reward of escaping as +1.\n",
    "\n",
    "Bellman's first breakthrough with this type of problem is now known as Value Iteration ([his original paper](http://www.iumj.indiana.edu/IUMJ/FULLTEXT/1957/6/56038)). He introduced a variable, gamma (γ), to represent discounted future rewards. He also introduced a function of policy (π) that takes a state (s), and outputs corresponding suggested action (a). The goal is to find the value of a state (V), given the rewards that occur when following an action in a particular state (R).\n",
    "\n",
    "Gamma, the discount, is the key ingredient here. If my time steps were in days, and my gamma was .9, `$100` would be worth `$100` to me today, `$90` tomorrow, `$81` the day after, and so on. Putting this all together, we get the Bellman Equation\n",
    "\n",
    "<img src=\"images/bellman_equation.jpg\" width=\"500\">\n",
    "\n",
    "source: [Wikipedia](https://en.wikipedia.org/wiki/Bellman_equation)\n",
    "\n",
    "In other words, the value of our current state, `current_values`, is equal to the discount times the value of the next state, `next_values`, given the policy the agent will follow. For now, we'll have our agent assume a greedy policy: it will move towards the state with the highest calculated value. If you're wondering what P is, don't worry, we'll get to that later.\n",
    "\n",
    "Let's program it out and see it in action! We'll set up an array representing the lake with -1 as the holes, and 1 as the goal. Then, we'll set up an array of zeros to start our iteration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "LAKE = np.array([[0,  0,  0,  0],\n",
    "                 [0, -1,  0, -1],\n",
    "                 [0,  0,  0, -1],\n",
    "                 [-1, 0,  0,  1]])\n",
    "LAKE_WIDTH = len(LAKE[0])\n",
    "LAKE_HEIGHT = len(LAKE)\n",
    "\n",
    "DISCOUNT = .9  # Change me to be a value between 0 and 1.\n",
    "current_values = np.zeros_like(LAKE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Gym environment class has a handy property for finding the number of states in an environment called `observation_space`. In our case, there a 16 integer states, so it will label it as \"Discrete\". Similarly, `action_space` will tell us how many actions are available to the agent.\n",
    "\n",
    "Let's take advantage of these to make our code portable between different lakes sizes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "env.observation_space - Discrete(16)\n",
      "env.observation_space.n - 16\n",
      "env.action_space - Discrete(4)\n",
      "env.action_space.n - 4\n"
     ]
    }
   ],
   "source": [
    "print(\"env.observation_space -\", env.observation_space)\n",
    "print(\"env.observation_space.n -\", env.observation_space.n)\n",
    "print(\"env.action_space -\", env.action_space)\n",
    "print(\"env.action_space.n -\", env.action_space.n)\n",
    "\n",
    "STATE_SPACE = env.observation_space.n\n",
    "ACTION_SPACE = env.action_space.n\n",
    "STATE_RANGE = range(STATE_SPACE)\n",
    "ACTION_RANGE = range(ACTION_SPACE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll need some sort of function to figure out what the best neighboring cell is. The below function take's a cell of the lake, and looks at the current value mapping (to be called with `current_values`, and see's what the value of the adjacent state is corresponding to the given `action`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def get_neighbor_value(state_x, state_y, values, action):\n",
    "    \"\"\"Returns the value of a state's neighbor.\n",
    "    \n",
    "    Args:\n",
    "        state_x (int): The state's horizontal position, 0 is the lake's left.\n",
    "        state_y (int): The state's vertical position, 0 is the lake's top.\n",
    "        values (float array): The current iteration's state values.\n",
    "        policy (int): Which action to check the value for.\n",
    "        \n",
    "    Returns:\n",
    "        The corresponding action's value.\n",
    "    \"\"\"\n",
    "    left = [state_y, state_x-1]\n",
    "    down = [state_y+1, state_x]\n",
    "    right = [state_y, state_x+1]\n",
    "    up = [state_y-1, state_x]\n",
    "    actions = [left, down, right, up]\n",
    "\n",
    "    direction = actions[action]\n",
    "    check_x = direction[1]\n",
    "    check_y = direction[0]\n",
    "        \n",
    "    is_boulder = check_y < 0 or check_y >= LAKE_HEIGHT \\\n",
    "        or check_x < 0 or check_x >= LAKE_WIDTH\n",
    "    \n",
    "    value = values[state_y, state_x]\n",
    "    if not is_boulder:\n",
    "        value = values[check_y, check_x]\n",
    "        \n",
    "    return value"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But this doesn't find the best action, and the gumdrop is going to need that if it wants to greedily get off the lake. The `get_max_neighbor` function we've defined below takes a number corresponding to a cell as `state_number` and the same value mapping as `get_neighbor_value`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def get_state_coordinates(state_number):\n",
    "    state_x = state_number % LAKE_WIDTH\n",
    "    state_y = state_number // LAKE_HEIGHT\n",
    "    return state_x, state_y\n",
    "\n",
    "def get_max_neighbor(state_number, values):\n",
    "    \"\"\"Finds the maximum valued neighbor for a given state.\n",
    "    \n",
    "    Args:\n",
    "        state_number (int): the state to find the max neighbor for\n",
    "        state_values (float array): the respective value of each state for\n",
    "            each cell of the lake.\n",
    "    \n",
    "    Returns:\n",
    "        max_value (float): the value of the maximum neighbor.\n",
    "        policy (int): the action to take to move towards the maximum neighbor.\n",
    "    \"\"\"\n",
    "    state_x, state_y = get_state_coordinates(state_number)\n",
    "    \n",
    "    # No policy or best value yet\n",
    "    best_policy = -1\n",
    "    max_value = -np.inf\n",
    "\n",
    "    # If the cell has something other than 0, it's a terminal state.\n",
    "    if LAKE[state_y, state_x]:\n",
    "        return LAKE[state_y, state_x], best_policy\n",
    "    \n",
    "    for action in ACTION_RANGE:\n",
    "        neighbor_value = get_neighbor_value(state_x, state_y, values, action)\n",
    "        if neighbor_value > max_value:\n",
    "            max_value = neighbor_value\n",
    "            best_policy = action\n",
    "        \n",
    "    return max_value, best_policy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's write our value iteration code. We'll write a function that comes out one step of the iteration by checking each state and finding its maximum neighbor. The values will be reshaped so that it's in the form of the lake, but the policy will stay as a list of ints. This way, when Gym returns a state, all we need to do is look at the corresponding index in the policy list to tell our agent where to go."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def iterate_value(current_values):\n",
    "    \"\"\"Finds the future state values for an array of current states.\n",
    "    \n",
    "    Args:\n",
    "        current_values (int array): the value of current states.\n",
    "\n",
    "    Returns:\n",
    "        next_values (int array): The value of states based on future states.\n",
    "        next_policies (int array): The recommended action to take in a state.\n",
    "    \"\"\"\n",
    "    next_values = []\n",
    "    next_policies = []\n",
    "\n",
    "    for state in STATE_RANGE:\n",
    "        value, policy = get_max_neighbor(state, current_values)\n",
    "        next_values.append(value)\n",
    "        next_policies.append(policy)\n",
    "    \n",
    "    next_values = np.array(next_values).reshape((LAKE_HEIGHT, LAKE_WIDTH))\n",
    "    return next_values, next_policies\n",
    "\n",
    "next_values, next_policies = iterate_value(current_values)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is what our values look like after one step. Right now, it just looks like the lake. That's because we started with an array of zeros for `current_values`, and the terminal states of the lake were loaded in."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 0,  0,  0,  0],\n",
       "       [ 0, -1,  0, -1],\n",
       "       [ 0,  0,  0, -1],\n",
       "       [-1,  0,  0,  1]])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "next_values"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And this is what our policy looks like reshaped into the form of the lake. The `-1`'s are terminal states. Right now, the agent will move left in any non-terminal state, because it sees all of those states as equal. Remember, if the gumdrop is along the leftmost side of the lake, and tries to move left, it will slip on a boulder and return to the same position."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 0,  0,  0,  0],\n",
       "       [ 0, -1,  0, -1],\n",
       "       [ 0,  0,  0, -1],\n",
       "       [-1,  0,  0, -1]])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.array(next_policies).reshape((LAKE_HEIGHT ,LAKE_WIDTH))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There's one last step to apply the Bellman Equation, the `discount`! We'll multiply our next states by the `discount` and set that to our `current_values`. One loop done!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 0. ,  0. ,  0. ,  0. ],\n",
       "       [ 0. , -0.9,  0. , -0.9],\n",
       "       [ 0. ,  0. ,  0. , -0.9],\n",
       "       [-0.9,  0. ,  0. ,  0.9]])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "current_values = DISCOUNT * next_values\n",
    "current_values"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run the below cell over and over again to see how our values change with each iteration. It should be complete after six iterations when the values no longer change. The policy will also change as the values are updated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Value\n",
      "[[ 0.   0.   0.   0. ]\n",
      " [ 0.  -1.   0.  -1. ]\n",
      " [ 0.   0.   0.  -1. ]\n",
      " [-1.   0.   0.9  1. ]]\n",
      "Policy\n",
      "[[ 0  0  0  0]\n",
      " [ 0 -1  1 -1]\n",
      " [ 0  0  0 -1]\n",
      " [-1  1  2 -1]]\n"
     ]
    }
   ],
   "source": [
    "next_values, next_policies = iterate_value(current_values)\n",
    "print(\"Value\")\n",
    "print(next_values)\n",
    "print(\"Policy\")\n",
    "print(np.array(next_policies).reshape((4,4)))\n",
    "current_values = DISCOUNT * next_values"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Have a completed policy? Let's see it in action! We'll update our `play_game` function to instead take our list of policies. That way, we can start in a random position and still get to the end."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Game Over!\n"
     ]
    }
   ],
   "source": [
    "def play_game(policy):\n",
    "    state = env.reset()\n",
    "    step = 0\n",
    "    done = False\n",
    "\n",
    "    while not done:\n",
    "        action = policy[state]  # This line is new.\n",
    "        state, _, done, _ = env.step(action)\n",
    "        env.render()\n",
    "        step += 1\n",
    "        print_state(state, done)\n",
    "\n",
    "play_game(next_policies)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Phew! Good job, team! The gumdrop made it out alive. So what became of our gumdrop hero? Well, the next day, it was making another snowman and fell onto an even more slippery and deadly lake. Doh! Turns out this story is part of a trilogy. Feel free to move onto the next section after your own sip of cocoa, coffee, tea, or poison of choice.\n",
    "\n",
    "## Policy Iteration\n",
    "\n",
    "You may have noticed that the first lake was built with the parameter `is_slippery=False`. This time, we're going to switch it to `True`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n"
     ]
    }
   ],
   "source": [
    "env = gym.make('FrozenLake-v0', is_slippery=True)\n",
    "state = env.reset()\n",
    "env.render()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Hmm, looks the same as before. Let's try applying our old policy and see what happens."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  (Left)\n",
      "SFFF\n",
      "\u001b[41mF\u001b[0mHFH\n",
      "FFFH\n",
      "HFFG\n",
      "4 - Still Alive!\n",
      "  (Left)\n",
      "SFFF\n",
      "FHFH\n",
      "\u001b[41mF\u001b[0mFFH\n",
      "HFFG\n",
      "8 - Still Alive!\n",
      "  (Left)\n",
      "SFFF\n",
      "FHFH\n",
      "FFFH\n",
      "\u001b[41mH\u001b[0mFFG\n",
      "12 - Game Over!\n"
     ]
    }
   ],
   "source": [
    "play_game(next_policies)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Was there a game over? There's a small chance that the gumdrop made it to the end, but it's much more likely that it accidentally slipped and fell into a hole. Oh no! We can try repeatedly testing the above code cell over and over again, but it might take a while. In fact, this is a similar roadblock Bellman and his colleagues faced.\n",
    "\n",
    "How efficient is Value Iteration? On our modern machines, this algorithm ran fairly quickly, but back in 1960, that wasn't the case. Let's say our lake is a long straight line like this:\n",
    "\n",
    "| | | | | | | |\n",
    "|-|-|-|-|-|-|-|\n",
    "|S|F|F|F|F|F|H|\n",
    "\n",
    "This is the worst case scenario for value iteration. In each iteration, we look at every state (s) and each action per state (a), so one step of value iteration is O(s*a). In the case of our lake line, each iteration only updates one cell. In other words, the value iteration step needs to be run `s` times. In total, that's O(s<sup>2</sup>a).\n",
    "\n",
    "Back in 1960, that was computationally heavy, and so [Ronald Howard](https://en.wikipedia.org/wiki/Ronald_A._Howard) developed an alteration of Value Iteration that mildly sacrificed mathematical accuracy for speed.\n",
    "\n",
    "Here's the strategy: it was observed that the optimal policy often converged before value iteration was complete. To take advantage of this, we'll start with random policy. When we iterate over our values, we'll use this policy instead of trying to find the maximum neighbor. This has been coded out in `find_future_values` below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def find_future_values(current_values, current_policies):\n",
    "    \"\"\"Finds the next set of future values based on the current policy.\"\"\"\n",
    "    next_values = []\n",
    "\n",
    "    for state in STATE_RANGE:\n",
    "        current_policy = current_policies[state]\n",
    "        state_x, state_y = get_state_coordinates(state)\n",
    "\n",
    "        # If the cell has something other than 0, it's a terminal state.\n",
    "        value = LAKE[state_y, state_x]\n",
    "        if not value:\n",
    "            value = get_neighbor_value(\n",
    "                state_x, state_y, current_values, current_policy)\n",
    "        next_values.append(value)\n",
    "\n",
    "    return np.array(next_values).reshape((LAKE_HEIGHT, LAKE_WIDTH))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After we've calculated our new values, then we'll update the policy (and not the values) based on the maximum neighbor. If there's no change in the policy, then we're done. The below is very similar to our `get_max_neighbor` function. Can you see the differences?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def find_best_policy(next_values):\n",
    "    \"\"\"Finds the best policy given a value mapping.\"\"\"\n",
    "    next_policies = []\n",
    "    for state in STATE_RANGE:\n",
    "        state_x, state_y = get_state_coordinates(state)\n",
    "\n",
    "        # No policy or best value yet\n",
    "        max_value = -np.inf\n",
    "        best_policy = -1\n",
    "\n",
    "        if not LAKE[state_y, state_x]:\n",
    "            for policy in ACTION_RANGE:\n",
    "                neighbor_value = get_neighbor_value(\n",
    "                    state_x, state_y, next_values, policy)\n",
    "                if neighbor_value > max_value:\n",
    "                    max_value = neighbor_value\n",
    "                    best_policy = policy\n",
    "                \n",
    "        next_policies.append(best_policy)\n",
    "    return next_policies"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To complete the Policy Iteration algorithm, we'll combine the two functions above. Conceptually, we'll be alternating between updating our value function and updating our policy function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def iterate_policy(current_values, current_policies):\n",
    "    \"\"\"Finds the future state values for an array of current states.\n",
    "    \n",
    "    Args:\n",
    "        current_values (int array): the value of current states.\n",
    "        current_policies (int array): a list where each cell is the recommended\n",
    "            action for the state matching its index.\n",
    "\n",
    "    Returns:\n",
    "        next_values (int array): The value of states based on future states.\n",
    "        next_policies (int array): The recommended action to take in a state.\n",
    "    \"\"\"\n",
    "    next_values = find_future_values(current_values, current_policies)\n",
    "    next_policies = find_best_policy(next_values)\n",
    "    return next_values, next_policies"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, let's modify the `get_neighbor_value` function to now include the slippery ice. Remember the `P` in the Bellman Equation above? It stands for the probability of ending up in a new state given the current state and action taken. That is, we'll take a weighted sum of the values of all possible states based on our chances to be in those states.\n",
    "\n",
    "How does the physics of the slippery ice work? For this lake, whenever the gumdrop tries to move in a particular direction, there are three possible positions that it could end up with. It could move where it was intending to go, but it could also end up to the left or right of the direction it was facing. For instance, if it wanted to move right, it could end up on the square above or below it! This is depicted below, with the yellow squares being potential positions after attempting to move right.\n",
    "\n",
    "<img src=\"images/slipping.jpg\" width=\"360\" height=\"270\">\n",
    "\n",
    "Each of these has an equal probability chance of happening. So since there are three outcomes, they each have about a 33% chance to happen. What happens if we slip in the direction of a boulder? No problem, we'll just end up not moving anywhere. Let's make a function to find what our possible locations could be given a policy and state coordinates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def get_locations(state_x, state_y, policy):\n",
    "    left = [state_y, state_x-1]\n",
    "    down = [state_y+1, state_x]\n",
    "    right = [state_y, state_x+1]\n",
    "    up = [state_y-1, state_x]\n",
    "    directions = [left, down, right, up]\n",
    "    num_actions = len(directions)\n",
    "\n",
    "    gumdrop_right = (policy - 1) % num_actions\n",
    "    gumdrop_left = (policy + 1) % num_actions\n",
    "    locations = [gumdrop_left, policy, gumdrop_right]\n",
    "    return [directions[location] for location in locations]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, we can add it to `get_neighbor_value` to find the weighted value of all the possible states the gumdrop can end up in."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def get_neighbor_value(state_x, state_y, values, policy):\n",
    "    \"\"\"Returns the value of a state's neighbor.\n",
    "    \n",
    "    Args:\n",
    "        state_x (int): The state's horizontal position, 0 is the lake's left.\n",
    "        state_y (int): The state's vertical position, 0 is the lake's top.\n",
    "        values (float array): The current iteration's state values.\n",
    "        policy (int): Which action to check the value for.\n",
    "        \n",
    "    Returns:\n",
    "        The corresponding action's value.\n",
    "    \"\"\"\n",
    "    locations = get_locations(state_x, state_y, policy)\n",
    "    location_chance = 1.0 / len(locations)\n",
    "    total_value = 0\n",
    "\n",
    "    for location in locations:\n",
    "        check_x = location[1]\n",
    "        check_y = location[0]\n",
    "\n",
    "        is_boulder = check_y < 0 or check_y >= LAKE_HEIGHT \\\n",
    "            or check_x < 0 or check_x >= LAKE_WIDTH\n",
    "    \n",
    "        value = values[state_y, state_x]\n",
    "        if not is_boulder:\n",
    "            value = values[check_y, check_x]\n",
    "        total_value += location_chance * value\n",
    "\n",
    "    return total_value"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For Policy Iteration, we'll start off with a random policy if only because the Gumdrop doesn't know any better yet. We'll reset our current values while we're at it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[2, 1, 0, 1],\n",
       "       [2, 1, 2, 2],\n",
       "       [2, 3, 3, 2],\n",
       "       [3, 3, 3, 3]])"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "current_values = np.zeros_like(LAKE)\n",
    "policies = np.random.choice(ACTION_RANGE, size=STATE_SPACE)\n",
    "np.array(policies).reshape((4,4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As before with Value Iteration, run the cell below multiple until the policy no longer changes. It should only take 2-3 clicks compared to Value Iteration's 6."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Value\n",
      "[[ 0.  0.  0.  0.]\n",
      " [ 0. -1.  0. -1.]\n",
      " [ 0.  0.  0. -1.]\n",
      " [-1.  0.  0.  1.]]\n",
      "Policy\n",
      "[[ 0  3  0  3]\n",
      " [ 0 -1  0 -1]\n",
      " [ 3  1  0 -1]\n",
      " [-1  2  1 -1]]\n"
     ]
    }
   ],
   "source": [
    "next_values, policies = iterate_policy(current_values, policies)\n",
    "print(\"Value\")\n",
    "print(next_values)\n",
    "print(\"Policy\")\n",
    "print(np.array(policies).reshape((4,4)))\n",
    "current_values = DISCOUNT * next_values"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Hmm, does this work? Let's see! Run the cell below to watch the gumdrop slip its way to victory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  (Left)\n",
      "SFFF\n",
      "\u001b[41mF\u001b[0mHFH\n",
      "FFFH\n",
      "HFFG\n",
      "4 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "SFFF\n",
      "\u001b[41mF\u001b[0mHFH\n",
      "FFFH\n",
      "HFFG\n",
      "4 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "SFFF\n",
      "\u001b[41mF\u001b[0mHFH\n",
      "FFFH\n",
      "HFFG\n",
      "4 - Still Alive!\n",
      "  (Left)\n",
      "SFFF\n",
      "FHFH\n",
      "\u001b[41mF\u001b[0mFFH\n",
      "HFFG\n",
      "8 - Still Alive!\n",
      "  (Up)\n",
      "SFFF\n",
      "FHFH\n",
      "\u001b[41mF\u001b[0mFFH\n",
      "HFFG\n",
      "8 - Still Alive!\n",
      "  (Up)\n",
      "SFFF\n",
      "\u001b[41mF\u001b[0mHFH\n",
      "FFFH\n",
      "HFFG\n",
      "4 - Still Alive!\n",
      "  (Left)\n",
      "SFFF\n",
      "\u001b[41mF\u001b[0mHFH\n",
      "FFFH\n",
      "HFFG\n",
      "4 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "SFFF\n",
      "\u001b[41mF\u001b[0mHFH\n",
      "FFFH\n",
      "HFFG\n",
      "4 - Still Alive!\n",
      "  (Left)\n",
      "SFFF\n",
      "FHFH\n",
      "\u001b[41mF\u001b[0mFFH\n",
      "HFFG\n",
      "8 - Still Alive!\n",
      "  (Up)\n",
      "SFFF\n",
      "\u001b[41mF\u001b[0mHFH\n",
      "FFFH\n",
      "HFFG\n",
      "4 - Still Alive!\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "0 - Still Alive!\n",
      "  (Left)\n",
      "SFFF\n",
      "\u001b[41mF\u001b[0mHFH\n",
      "FFFH\n",
      "HFFG\n",
      "4 - Still Alive!\n",
      "  (Left)\n",
      "SFFF\n",
      "\u001b[41mF\u001b[0mHFH\n",
      "FFFH\n",
      "HFFG\n",
      "4 - Still Alive!\n",
      "  (Left)\n",
      "SFFF\n",
      "FHFH\n",
      "\u001b[41mF\u001b[0mFFH\n",
      "HFFG\n",
      "8 - Still Alive!\n",
      "  (Up)\n",
      "SFFF\n",
      "FHFH\n",
      "F\u001b[41mF\u001b[0mFH\n",
      "HFFG\n",
      "9 - Still Alive!\n",
      "  (Down)\n",
      "SFFF\n",
      "FHFH\n",
      "FFFH\n",
      "H\u001b[41mF\u001b[0mFG\n",
      "13 - Still Alive!\n",
      "  (Right)\n",
      "SFFF\n",
      "FHFH\n",
      "F\u001b[41mF\u001b[0mFH\n",
      "HFFG\n",
      "9 - Still Alive!\n",
      "  (Down)\n",
      "SFFF\n",
      "FHFH\n",
      "FF\u001b[41mF\u001b[0mH\n",
      "HFFG\n",
      "10 - Still Alive!\n",
      "  (Left)\n",
      "SFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HF\u001b[41mF\u001b[0mG\n",
      "14 - Still Alive!\n",
      "  (Down)\n",
      "SFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFF\u001b[41mG\u001b[0m\n",
      "15 - Cocoa Time!\n"
     ]
    }
   ],
   "source": [
    "play_game(policies)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So what was the learned strategy here? The gumdrop learned to hug the left wall of boulders until it was down far enough to make a break for the exit. Instead of heading directly for it though, it took advantage of actions that did not have a hole of death in them. Patience is a virtue!\n",
    "\n",
    "We promised this story was a trilogy, and yes, the next day, the gumdrop fell upon a frozen lake yet again.\n",
    "\n",
    "## Q Learning\n",
    "Value Iteration and Policy Iteration are great techniques, but what if we don't know how big the lake is? With real world problems, not knowing how many potential states are can be a definite possibility.\n",
    "\n",
    "Enter [Chris Watkins](http://www.cs.rhul.ac.uk/~chrisw/). Inspired by how animals learn with delayed rewards, he came up with the idea of [Q Learning](http://www.cs.rhul.ac.uk/~chrisw/new_thesis.pdf) as an evolution of [Richard Sutton's](https://en.wikipedia.org/wiki/Richard_S._Sutton) [Temporal Difference Learning](https://en.wikipedia.org/wiki/Temporal_difference_learning). Watkins noticed that animals learn from positive and negative rewards, and that they often make mistakes in order to optimize a skill.\n",
    "\n",
    "From this emerged the idea of a Q table. In the lake case, it would look something like this.\n",
    "\n",
    "| |Left|Down|Right|Up|\n",
    "|-|-|-|-|-|\n",
    "|0| | | | |\n",
    "|1| | | | |\n",
    "|...| | | | |\n",
    "\n",
    "Here's the strategy: our agent will explore the environment. As the agent observes new states, we'll add more rows to our table. Whenever it moves from one state to the next, we'll update the cell corresponding to the old state based on the Bellman Equation. The agent doesn't need to know what the probabilities are between transitions. It'll learn the value of these as it experiments.\n",
    "\n",
    "For Q learning, this works by looking at the row that corresponds to the agent's current state. Then, we'll select the action with the highest value. There are multiple ways to initialize the Q-table, but for us, we'll start with all zeros. In that case, when selecting the best action, we'll randomly select between tied max values. If we don't, the agent will favor certain actions which will limit its exploration.\n",
    "\n",
    "To be able to handle an unknown number of states, we'll initialize our q_table as one row to represent our initial state. Then, we'll make a dictionary to map new states to rows in the table."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mapping\n",
      "{0: 0}\n",
      "q_table\n",
      "[[0. 0. 0. 0.]]\n"
     ]
    }
   ],
   "source": [
    "new_row = np.zeros((1, env.action_space.n))\n",
    "q_table = np.copy(new_row)\n",
    "q_map = {0: 0}\n",
    "\n",
    "def print_q(q_table, q_map):\n",
    "    print(\"mapping\")\n",
    "    print(q_map)\n",
    "    print(\"q_table\")\n",
    "    print(q_table)\n",
    "\n",
    "print_q(q_table, q_map)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our new `get_action` function will help us read the `q_table` and find the best action.\n",
    "\n",
    "First, we'll give the agent the ability to act randomly as opposed to choosing the best known action. This gives it the ability to explore and find new situations. This is done with a random chance to act randomly. So random!\n",
    "\n",
    "When the Gumdrop chooses not to act randomly, it will instead act based on the best action recorded in the `q_table`. Numpy's [argwhere](https://docs.scipy.org/doc/numpy/reference/generated/numpy.argwhere.html) is used to find the indexes with the maximum value in the q-table row corresponding to our current state. Since numpy is often used with higher dimensional data, each index is returned as a list of ints. Our indexes are really one dimensional since we're just looking within a single row, so we'll use [np.squeeze](https://docs.scipy.org/doc/numpy/reference/generated/numpy.squeeze.html) to remove the extra brackets. To randomly select from the indexes, we'll use [np.random.choice](https://docs.scipy.org/doc/numpy-1.14.1/reference/generated/numpy.random.choice.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def get_action(q_map, q_table, state_row, random_rate):\n",
    "    \"\"\"Find max-valued actions and randomly select from them.\"\"\"\n",
    "    if random.random() < random_rate:\n",
    "        return random.randint(0, ACTION_SPACE-1)\n",
    "\n",
    "    action_values = q_table[state_row]\n",
    "    max_indexes = np.argwhere(action_values == action_values.max())\n",
    "    max_indexes = np.squeeze(max_indexes, axis=-1)\n",
    "    action = np.random.choice(max_indexes)\n",
    "    return action"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, we'll define how the `q_table` gets updated. We'll apply the Bellman Equation as before, but since there is so much luck involved between slipping and random actions, we'll update our `q_table` as a weighted average between the `old_value` we're updating and the `future_value` based on the best action in the next state. That way, there's a little bit of memory between old and new experiences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def update_q(q_table, new_state_row, reward, old_value):\n",
    "    \"\"\"Returns an updated Q-value based on the Bellman Equation.\"\"\"\n",
    "    learning_rate = .1  # Change to be between 0 and 1.\n",
    "    future_value = reward + DISCOUNT * np.max(q_table[new_state_row])\n",
    "    return old_value + learning_rate * (future_value - old_value)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll update our `play_game` function to take our table and mapping, and at the end, we'll return any updates to them. Once we observe new states, we'll check our mapping and add then to the table if space isn't allocated for them already.\n",
    "\n",
    "Finally, for every `state` - `action` - `new-state` transition, we'll update the cell in `q_table` that corresponds to the `state` and `action` with the Bellman Equation.\n",
    "\n",
    "There's a little secret to solving this lake problem, and that's to have a small negative reward when moving between states. Otherwise, the gumdrop will become too afraid of slipping in a death hole to explore out of what is thought to be safe positions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def play_game(q_table, q_map, random_rate, render=False):\n",
    "    state = env.reset()\n",
    "    step = 0\n",
    "    done = False\n",
    "\n",
    "    while not done:\n",
    "        state_row = q_map[state]\n",
    "        action = get_action(q_map, q_table, state_row, random_rate)\n",
    "        new_state, _, done, _ = env.step(action)\n",
    "\n",
    "        #Add new state to table and mapping if it isn't there already.\n",
    "        if new_state not in q_map:\n",
    "            q_map[new_state] = len(q_table)\n",
    "            q_table = np.append(q_table, new_row, axis=0)\n",
    "        new_state_row = q_map[new_state]\n",
    "\n",
    "        reward = -.01  #Encourage exploration.\n",
    "        if done:\n",
    "            reward = 1 if new_state == 15 else -1\n",
    "        current_q = q_table[state_row, action]\n",
    "        q_table[state_row, action] = update_q(\n",
    "            q_table, new_state_row, reward, current_q)\n",
    "\n",
    "        step += 1\n",
    "        if render:\n",
    "            env.render()\n",
    "            print_state(new_state, done)\n",
    "        state = new_state\n",
    "        \n",
    "    return q_table, q_map"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Ok, time to shine, gumdrop emoji! Let's do one simulation and see what happens."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Run to refresh the q_table.\n",
    "random_rate = 1\n",
    "q_table = np.copy(new_row)\n",
    "q_map = {0: 0}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "q_table, q_map = play_game(q_table, q_map, random_rate, render=True)\n",
    "print_q(q_table, q_map)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Unless the gumdrop was incredibly lucky, chances were, it fell in some death water. Q-learning is markedly different from Value Iteration or Policy Iteration in that it attempts to simulate how an animal learns in unknown situations. Since the layout of the lake is unknown to the Gumdrop, it doesn't know which states are death holes, and which ones are safe. Because of this, it's going to make many mistakes before it can start making successes.\n",
    "\n",
    "Feel free to run the above cell multiple times to see how the gumdrop steps through trial and error. When you're ready, run the below cell to have the gumdrop play 1000 times."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for _ in range(1000):\n",
    "    q_table, q_map = play_game(q_table, q_map, random_rate)\n",
    "    random_rate = random_rate * .99\n",
    "print_q(q_table, q_map)\n",
    "random_rate"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Cats have nine lives, our Gumdrop lived a thousand! Moment of truth. Can it get out of the lake now that it matters?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "q_table, q_map = play_game(q_table, q_map, 0, render=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Third time's the charm!\n",
    "\n",
    "Each of these techniques has its pros and cons. For instance, while Value Iteration is the mathematically correct solution, it's not as time efficient at Policy Iteration or as flexible as Q-Learning.\n",
    "\n",
    "| |Value Iteration|Policy Iteration|Q Tables|\n",
    "|-|-|-|-|\n",
    "|Avoids locally optimal routes|✓|x|x|\n",
    "|On-policy (greedy)|✓|✓|x|\n",
    "|Model Free|x|x|✓|\n",
    "|Most time efficient|x|✓|x|\n",
    "\n",
    "Congratulations on making it through to the end. Now if you ever fall on a Frozen Lake, you'll have many different ways to calculate your survival. The gumdrop thank you!\n",
    "\n",
    "<img src=\"images/end.jpg\" width=\"178\" height=\"234\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Copyright 2020 Google Inc.\n",
    "Licensed under the Apache License, Version 2.0 (the \"License\"); you may not use this file except in compliance with the License. You may obtain a copy of the License at\n",
    "http://www.apache.org/licenses/LICENSE-2.0\n",
    "Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License."
   ]
  }
 ],
 "metadata": {
  "environment": {
   "kernel": "conda-root-py",
   "name": "workbench-notebooks.m115",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/workbench-notebooks:m115"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel) (Local)",
   "language": "python",
   "name": "conda-root-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}