03A_RL.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction to Reinforcement Learning (RL) in TensorFlow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Recap of Supervised Learning\n",
    "\n",
    "In previous days, we've primarily concerned ourselves with supervised learning problems (mostly classification).\n",
    "In supervised learning, we are given some sort of training data consisting of input/output pairs, with the goal being to be able to predict the output given some new inputs after learning the model.\n",
    "For example, yesterday's homework was CNN classification model for MNIST; given a training set of 55000 digit images and corresponding digit labels (e.g. '5'), we learned a model that was capable of predicting the digit label of new MNIST images.\n",
    "In order words, something like (but not exactly) this:\n",
    "<img src=\"Figures/mnist_cnn_ex.png\" alt=\"mnist_cnn_ex\" style=\"width: 700px;\"/>\n",
    "<center>Figure from *[Getting started with PyTorch for Deep Learning (Part 3: Neural Network basics)]( https://codetolight.wordpress.com/2017/11/29/getting-started-with-pytorch-for-deep-learning-part-3-neural-network-basics/)*</center>\n",
    "\n",
    "What if we want to learn how to perform more complex behaviors, where data collection can be expensive? \n",
    "How do you teach a robot to walk? \n",
    "Self-driving cars? \n",
    "How do you defeat human champions in the game of Go? "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Reinforcement Learning\n",
    "Enter Reinforcement Learning. \n",
    "In Reinforcement Learning, our model (commonly referred to as an *agent* in this context) interacts with an *environment* by taking *actions* $a$ and receives some sort of feedback from the environment in the form of a *reward* $r$. \n",
    "In this sense, reinforcement learning algorithms learn by experience.\n",
    "We call the trajectory of going from start to finish of a task an *episode*, and often our agent will learn by undergoing many episodes.\n",
    "<img src=\"Figures/RL.PNG\" alt=\"RL\" style=\"width: 400px;\"/>\n",
    "\n",
    "#### Markov Decision Processes (MDPs)\n",
    "Many reinforcement learning algorithms are modeled as Markov Decision Processes (MDPs). \n",
    "In these settings, we have a concept of a *state* $s$, which encapsulates the situation of the agent (e.g. location, velocity).\n",
    "From each state $s_t$, the agent takes an action $a_t$, which results in a transition from one state $s_t$ to another $s_{t+1}$.\n",
    "In many settings, there is stochasticity in this transition, meaning that there's is a distribution over $s_{t+1}$ conditioned on $s_t$ and $a_t$.\n",
    "Often, several of these states are considered episode ending, after which the agent can no longer make any transitions or collect any more reward.\n",
    "These correspond to states such as reaching the final goal, a game concluding, or falling of a cliff.\n",
    "In the end, our goal is to learn a *policy* $\\pi$, or a mapping from states to actions.\n",
    "\n",
    "#### Partially Observable Markov Decision Processes (POMDPs)\n",
    "In an MDP, we assume that we can always tell which state $s_t$ our agents is in.\n",
    "However, this isn't always the case.\n",
    "Sometimes, all we have access to are observations $o_t$ that provide information the state $s_t$, but enough to precisely pinpoint the exact one.\n",
    "We call such settings Partially Observable Markov Decision Processes (POMDPs).\n",
    "Imagine for example a [Roomba](https://en.wikipedia.org/wiki/Roomba) being trained to navigate a living room with RL.\n",
    "From its infrared and mechanical \"bump\" sensors, it receives partial information ($o_t$) as to where it might be, but not a definitive location ($s_t$).\n",
    "Operating as a POMDP adds a whole layer of complexity to RL algorithms.\n",
    "For the rest of day though, we'll focus on MDPs, as their much simpler and easier to use to teach basic concepts.\n",
    "\n",
    "\n",
    "#### A simple MDP example\n",
    "\n",
    "<img src=\"Figures/MDP.png\" alt=\"MDP\" style=\"width: 400px;\"/>\n",
    "<center>Figure from *[Markov decision process](https://en.wikipedia.org/wiki/Markov_decision_process)*</center>\n",
    "\n",
    "In the above example, we can see the 3 possible states for the agent as $s_0$, $s_1$, and $s_2$, with 2 actions $a_0$ and $a_1$ available from each state.\n",
    "We can see that the each action doesn't lead to a determinstic transition to the next state, as shown by multiple paths from each action. \n",
    "Note that each of the outcomes of an action are labeled with a small black number between 0 and 1.\n",
    "This denotes the probability of that outcome (which state we end up at) given the the action; as these are probabilites, the sum of the probabilities of arriving at each of the next states $s_{t+1}$ given a previous state $s_t$ and selected action $a_t$ is 1.\n",
    "\n",
    "#### Objective\n",
    "\n",
    "The goal of the agent is to maximize the total reward $R$ it can receive over some number of steps.\n",
    "It is important to ensure the reward actually captures the true goal we want the agent to achieve.\n",
    "The agent will dutifully attempt maximize the objective it is given, without any considerations to any implicit objectives that a human may desire.\n",
    "There are more than a few (amusing) anecdotes of RL agents learning undesirable behaviors by exploiting some aspect of the reward function.\n",
    "As such, defining this reward requires special care.\n",
    "\n",
    "One countermeasure commonly deployed by RL researchers is the concept of *discounted* rewards.\n",
    "This is done with a multiplicative term $\\gamma$: a reward $T$ steps in the future is discounted as $\\gamma^T r_T$.\n",
    "Using discounting encourages the agent to finish the task sooner rather than later, a common implicit criterion.\n",
    "With discounting then, the RL agent's goal is to maximize:\n",
    "\\begin{align*}\n",
    "\\mathbb{E}[\\sum_{t=0}^{\\infty} \\gamma^t r_t]\n",
    "\\end{align*}\n",
    "This is far from the complete solution to making our rewards accurately capture our desired objectives, but achieving higher rewards sooner rather than later is an almost universal preference, so we almost always add it.\n",
    "Designing a good reward function can be an art is highly dependent on the task.\n",
    "\n",
    "#### Reinforcement Learning as Supervised Learning?\n",
    "At first, this doesn't seem too different from the supervised methods we've looked at before, and some natural questions might arise:\n",
    "- Why can't we just treat RL as a supervised task? \n",
    "Why can't we use the reward (or rather, the negative of the reward) as our supervised loss?\n",
    "\n",
    "Unlike in supervised learning, in reinforcement learning, we often don't have a pre-apportioned dataset to learn from.\n",
    "In some problems set-ups, we may have examples of other agents (oftentimes humans) performing the desired task, but these aren't necessarily optimal examples of how to maximize the reward, which is what we want to learn.\n",
    "In most RL settings, we don't have any examples of state-action trajectories beyond what our agent experiences through trial-and-error, which are even more suboptimal. \n",
    "Learning from the reward signal itself is also not the way to do it, as such an approach would be myopic: it would maximize reward that can be earned now instead of maximizing the overall reward over a trajectory.\n",
    "\n",
    "\n",
    "### Open AI Gym\n",
    "Before we dive any deeper into implementing reinforcement learning models, first we need an environment.\n",
    "Remember, the goal is to learn an agent that can interact with an environment in the way we want, so we need something that our agent can interact with and receive rewards from.\n",
    "In robotics, this is often the real world (or some set-up in the real world).\n",
    "However, it is oftentimes cheaper and quicker to first test our algorithms in simulated settings.\n",
    "There are a number of tasks that are popular benchmarks for the reinforcement learning community, such as [cart pole](https://en.wikipedia.org/wiki/Inverted_pendulum), [mountain car](https://en.wikipedia.org/wiki/Mountain_car_problem), or [Atari 2600 games](https://gym.openai.com/envs/#atari). \n",
    "In the spirit of accelerating progress and promoting openness in the research community, Open AI has very nicely coded up [Open AI Gym](https://gym.openai.com/), which has implementations of many of these environments for public use.\n",
    "We will be using these environments, as it allows us to focus on the algorithms themselves, instead of worrying about implementing each problem setting ourselves.\n",
    "\n",
    "To use it, we first need to download and install it. Make sure you're in your TensorFlow environment first!\n",
    "\n",
    "```Shell\n",
    "# If you environment isn't currently active, activate it (include the word \"source\", without [] if on Windows):\n",
    "# [source] activate tensorflow\n",
    "\n",
    "pip install gym\n",
    "```\n",
    "\n",
    "Once it's installed, we can import it like any other Python module:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### FrozenLake (a Grid World)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's start with a simple environment: FrozenLake.\n",
    "Here's the official description from OpenAI gym:\n",
    "\n",
    "> *Winter is here. You and your friends were tossing around a frisbee at the park when you made a wild throw that left the frisbee out in the middle of the lake. The water is mostly frozen, but there are a few holes where the ice has melted. If you step into one of those holes, you'll fall into the freezing water. At this time, there's an international frisbee shortage, so it's absolutely imperative that you navigate across the lake and retrieve the disc. However, the ice is slippery, so you won't always move in the direction you intend.*\n",
    "\n",
    "A visualization of FrozenLake as a grid world:\n",
    "\n",
    "<img src=\"Figures/FrozenLake.PNG\" alt=\"FrozenLake\" style=\"width: 400px;\"/>\n",
    "<center>Figure from *[Frozen Lake: Beginners Guide To Reinforcement Learning With OpenAI Gym](https://analyticsindiamag.com/openai-gym-frozen-lake-beginners-guide-reinforcement-learning/)*</center>\n",
    "\n",
    "At the start of an episode, we begin in the upper left corner (S).\n",
    "Our goal is to move ourself to the lower right corner (G), avoiding falling into the holes (H).\n",
    "Icy water is cold.\n",
    "\n",
    "In reinforcement learning terms, each of the 16 locations on the grid are a state, and an action is attempting to move in one of four directions (left, down, right, up).\n",
    "Each move will result in the agent's state changing from $s_t$ to $s_{t+1}$ as it changes location, unless it attempts to move in the direction of a wall, which results in the agent's state not changing (the agent doesn't move).\n",
    "We receive a positive reward of \"+1\" for reaching the goal (G), discounted according to how long it took.\n",
    "While there is not a negative reward for falling into a hole (H), the agent still pays a penalty in the sense that falling into the hole is episode-ending and therefore prevents it from receiving any reward. \n",
    "We want to learn a policy $\\pi$ that takes us from our starting location (S) to the goal (G) in as few steps as possible.\n",
    "\n",
    "To really establish what we are trying to accomplish here, it's worth debunking a few common initial misconceptions:\n",
    "\n",
    "- **Knowledge of the states and transition probabilties:** From the top-down view, your first thought might be to plot out a path from the start to the finish, just as you would with a maze.\n",
    "However, this view is provided to us the algorithm designers so we can visualize the problem at hand.\n",
    "The agent learning the task does *not* get this prior knowledge; all we are about to tell it is that there are going to be 16 states and 4 possible actions from each state.\n",
    "A more proper analogy would be if I blindfolded you and dropped you in the middle of a frozen lake, and told you your state (location) every time you decided to take a step in one of four directions, then set off fireworks when you stepped on the frisbee.\n",
    "\n",
    "- **Knowledge of the goal (reward):** In OpenAI's official description of the environment, you (the agent) know what you're hoping to accomplish: You want to retrieve the frisbee, while avoiding falling through the ice.\n",
    "The agent does *not* know this.\n",
    "Rather, it learns the goal by experiencing rewards (or penalties), and the algorithm updates its policy such that it will be more (or less) likely to do those actions again.\n",
    "Note that this means that if an agent never experiences certain rewards, it won't know they exist.\n",
    "\n",
    "- **Prior knowledge of pathfinding, physics, etc.:** As a human, even if you haven't solved this task before, you still bring a tremondous amount of prior knowledge to this problem.\n",
    "For example, you know the shortest path to a destination is a line.\n",
    "You know that North, South, East, and West, are directions, and that going North and then South brings you back to where you already were.\n",
    "You know ice is slippery.\n",
    "You know icy water is cold.\n",
    "You know being in icy cold water is bad.\n",
    "It's important to keep in mind that our agent will begin knowing none of these things; it's initial policy is essentially picking actions completely at random.\n",
    "By the end of the training, it still won't know what abstract concepts like \"North/South,\" \"cold,\" or \"slippery\" mean\\*, but it will have (hopefully) learned a good policy that allows it to complete the goal.\n",
    "\n",
    "\\* *Apologies. Didn't mean to get philosophical.*\n",
    "\n",
    "#### Interacting with FrozenLake\n",
    "This example is simple enough that we could code the environment and its interface ourselves fairly easily, but OpenAI has already done it, and we'd like to focus on the algorithm of solving it as much as possible.\n",
    "We can create an instantiation of FrozenLake in a single line of code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Open AI Gym environments provide a mechanism to observe the state of the environment, and since FrozenLake is an MDP (as opposed to POMDP), the observation is the state itself.\n",
    "For FrozenLake, there are 16 grid locations on the map, meaning we have 16 states.\n",
    "We can confirm this by looking at the size of the `observation_space` attribute for the environment we just created."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our agent will interact with this environment causing its state to change.\n",
    "For FrozenLake, we have 4 options, each corresponding to attempting to step in a particular direction: `[Left, Down, Right, Up]`.\n",
    "We can confirm this by looking at the size of the `action_space` of our environment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before interacting with the environment, we have to first reset it to initialize it. \n",
    "Resetting also returns an observation of the first state after resetting.\n",
    "In FrozenLake, we always start in the upper left corner, which corresponds to state 0.\n",
    "As such, we see the `reset()` command returning `0`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can visualize the FrozenLake environment by calling `render()`. \n",
    "In more complex tasks this will actually add frames to a video showing our agent's progress, but for FrozenLake, it just prints out a text representation, with the highlighted character showing our agent's current location.\n",
    "We can see that we started in the upper-left corner, on the \"S,\" as promised."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's try moving.\n",
    "On thing to keep in mind is that the original FrozenLake environment is \"slippery.\" \n",
    "Because of the ice, if you try to go in one direction, you end up with a 1/3 chance of going in the direction you meant and the two adjacent directions each. \n",
    "For example, if we try going right, we have equal probabilities of slipping and going up and down instead.\n",
    "This makes things a little more complicated, so for now, let's first turn off the stochasticity and make this a deterministic transition instead.\n",
    "We do this by registering a new type of environment, and then instantiating a copy of said environment, making sure to reset it first."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Non-slippery version\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We advance time in an OpenAI environment with the `step()` method, which takes as argument an `action`.\n",
    "Let's trying moving right, which corresponds to action `2`.\n",
    "Notice that the output is a tuple of four elements: the next observation (`object`), the reward (`float`), whether or not the episode is done (`boolean`), and a dictionary of information (`dict`) that may be useful for debugging (this dict shouldn't be used in the final algorithm itself)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, let's `render()` to visualize what happened.\n",
    "Observe that this particular environment prints out the action we took in parentheses up top, in this case \"(Right)\", and then shows the result of that action.\n",
    "Notice that while most of the time, we succeed in going in the direction we want to, occasionally we slip on the ice and go in a direction we didn't intend."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can keep doing this as many times as we want.\n",
    "Since we're in Jupyter, we can just keep running the same cell (making small edits to change our action).\n",
    "\n",
    "Notice that once we fall into a hole, the episode is over, and we can no longer do anything.\n",
    "The same is true after reaching the goal."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before we get into any RL, let's see how random actions perform in this environment:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Hm. \n",
    "Not great. \n",
    "Alright, so clearly picking random steps isn't very likely to take us to the goal.\n",
    "It's apparent just from looking at the map that there're much better policies that we can learn.\n",
    "How are we going to do so?\n",
    "\n",
    "#### Q-learning\n",
    "There are many algorithms that we can use, but let's choose Q-learning, which we covered earlier today.\n",
    "Remember, in Q-learning (and SARSA, it turns out), we're trying learn the Q values for the states in our system.\n",
    "\n",
    "The Q value for a policy $\\pi$ is a function of the state $s$ and action $a$ and is defined as the following:\n",
    "\\begin{equation}\n",
    "Q_\\pi (s,a) = \\mathbb{E}\\big[\\sum_{t=0}^{\\infty} \\gamma^t r_t \\big|\\pi, s_0 = s, a_0 = a\\big]\n",
    "\\end{equation}\n",
    "Intuitively, the Q value is the total reward (including discounting) that the agent will gain if it takes action $a$ from state $s$ and then follows policy $\\pi$ for the rest of the episode.\n",
    "As one might expect, if Q is known exactly, the agent will attain the highest reward from $s$ if the policy $\\pi$ is to pick the $a$ with the highest Q value.\n",
    "\n",
    "Okay, so if we know the Q values for the system, then we can trivially find the optimal policy.\n",
    "So what are the Q values of the system?\n",
    "Well, at the beginning, we don't know, but we can try to learn them through experience.\n",
    "This is where Q-learning comes in.\n",
    "Q-learning iteratively updates the Q values in the following way:\n",
    "\\begin{equation}\n",
    "Q_\\pi (s_t, a_t) \\leftarrow (1 - \\alpha) \\cdot Q_\\pi(s_t, a_t) + \\alpha \\cdot \\big(r_t + \\gamma \\max_a Q_\\pi(s_{t+1}, a)\\big)\n",
    "\\end{equation}\n",
    "Notice that Q-learning is an *off-policy* method, in the sense that you don't actually learn from the trajectory you actually took (otherwise it'd be SARSA).\n",
    "Instead, we learn from the *greedy* transition, i.e. the best action we know how to take.\n",
    "\n",
    "And that's it! \n",
    "We run our agent through many episodes, experiencing many $s_t \\rightarrow a_t \\rightarrow s_{t+1}$ transitions and rewards, and just like that, we eventually learn a good Q function (and thus a good policy).\n",
    "Now of course, there are a bunch of small details and tweaks to make this work in practice, but we'll get to those later.\n",
    "\n",
    "#### Q-learning in FrozenLake\n",
    "FrozenLake is a very simple setting, one that we would call a toy problem.\n",
    "With only 16 states and 4 actions, there are only 64 state-action pairs possible (16x4=64), less if we account for the goal and the holes being episode ending (for simplicity though, we won't). \n",
    "With this few state-action pairs, we can actually solve this problem tabularly.\n",
    "Let's set up a Q table, and initialize the Q-values for all state-action pairs to zeros.\n",
    "Note that we're actually not going to need TensorFlow in this example; we'll use a Numpy array to store the Q table."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "#Initialize table with all zeros to be uniform\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A few hyperparameters we're going to set:\n",
    "- `alpha`: learning rate for the Q function\n",
    "- `gamma`: discount rate for future rewards\n",
    "- `num_episodes`: number of episodes (trajectories from start to goal/hole) our agent will learn from\n",
    "\n",
    "We're also going to store our rewards in an array called `rs`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Learning parameters\n",
    "\n",
    "\n",
    "# array of reward for each episode\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now for the bulk of algorithm itself.\n",
    "Notice that we're going to loop through the process `num_episodes` times, resetting the environment each time.\n",
    "At each step, we take the action with the highest Q value for our current state, with some randomness added in (especially at the beginning) to encourage exploration.\n",
    "After each action, we update our Q table greedily based on the reward experienced and the next best action.\n",
    "We also make sure to update our state, rinse, and repeat. \n",
    "We continue taking actions in an episode until it is `done`, storing the final total reward for the episode."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set total reward and time to zero, done to False\n",
    "\n",
    "\n",
    "# Reset environment and get first new observation\n",
    "\n",
    "    \n",
    "# Choose an action by greedily (with noise) from Q table\n",
    "\n",
    "\n",
    "# Get new state and reward from environment\n",
    "\n",
    "\n",
    "# Update Q-Table with new knowledge\n",
    "\n",
    "\n",
    "# Add reward to episode total\n",
    "\n",
    "\n",
    "# Update state and time\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "How did we do?\n",
    "Let's take a look at the rewards that we saved.\n",
    "We can plot the reward versus the episode number, and hopefully we'll see some sort of an increase over time.\n",
    "RL performance can be extremely noisy, so let's instead plot a moving average."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Plot reward vs episodes\n",
    "\n",
    "\n",
    "# Sliding window average\n",
    "\n",
    "\n",
    "# Plot\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Pretty good.\n",
    "We might also be interested in how often our agent actually reached the goal.\n",
    "This won't account for how quickly the agent got there (which might also of interest), but let's ignore that for now.\n",
    "To prevent us from being overwhelmed by data points, let's bucket the values into 10 stages, printing out how many episodes of each stage resulted in finding the goal."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Print number of times the goal was reached\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our RL agent does a really good job at navigating the FrozenLake when its moves are deterministic, but after all, this is\n",
    "supposed to be *Frozen*Lake, so where's the fun without the slipperiness?\n",
    "Let's go back to the original environment and see how the agent does."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Slippery version"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### TensorFlow in RL\n",
    "Hey, not bad.\n",
    "However, while the previous example was fun and simple, it was noticeably lacking any hint of TensorFlow.\n",
    "After all, this is supposed to be a TensorFlow session, so we should probably remedy that.\n",
    "\n",
    "We could have used a TensorFlow `Variable` to store the Q table and added the appropriate update ops to our computation graph, but that's a rather contrived way to force TensorFlow into the problem.\n",
    "TensorFlow's true utility comes from building neural networks and calculating/applying gradients automatically, which learning the Q table didn't need. \n",
    "\n",
    "#### Continuous domains\n",
    "In our previous example, we mentioned that with only 16 discrete states and 4 actions/state, the Q table only needed to hold 64 values, which is very manageable.\n",
    "However, what if the state or action space is continuous?\n",
    "You could discretize it, but then you have to pick a resolution, and your state-action space could explode exponentially.\n",
    "Treating these binned states or actions as completely different states is also ignoring that two consecutive bins are likely very similar in the needed policy.\n",
    "You can learn these relationships, but doing so is horriby sample inefficient.\n",
    "\n",
    "Instead of learning a Q table then, perhaps a Q function would be more appropriate.\n",
    "This function would take in a state and action as an input and return a Q value as an output.\n",
    "The Q function may be very complex, but as we've learned over the past few days, neural networks are very flexible and good for approximating arbitrary functions.\n",
    "Such an approach is called a [Deep Q Network](https://deepmind.com/research/dqn/)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Cart Pole"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's look at the cart pole problem next. \n",
    "In this setting, we have a pole attached to a hinge on a cart, with the goal being to keep the pole as vertical as possible.\n",
    "Because of gravity, the pole will fall unless the cart is exactly beneath the the pole's center of gravity.\n",
    "To prevent the pole from falling, the agent can apply a force of +1 or -1 to the cart to move it left and right along a track.\n",
    "The agent receives a reward of +1 for every timestamp the pole remains vertical; the game ends when the pole fall past 15 degrees from vertical or the cart moves more than 2.4 units away from the center.\n",
    "\n",
    "<img src=\"Figures/polebalance.gif\" alt=\"polebalance\" style=\"width: 400px;\"/>\n",
    "\n",
    "First, let's create an instance of the cart pole environment:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again, we can look at the `observation_space` for this environment.\n",
    "Also similar to FrozenLake, since this version of cart pole is an MDP (as opposed to POMDP), the observation is the state itself.\n",
    "We can see that the states for cart pole have 4 dimensions, which correspond to `[cart position, cart velocity, pole angle, pole angular velocity]`.\n",
    "Importantly, notice these states are *continuous* values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can look at the `action_space` again too.\n",
    "In cart pole, there are two actions available to the agent: `[apply force left, apply force right]`.\n",
    "We can see this by examining the `action_space` attribute:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Resetting the environment return our first observations, which we can see has 4 values corresponding to the 4 previously mentioned state variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before we get into any reinforcement learning, let's see how we perform actions within the environment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Okay, so clearly choosing a random action at every time step doesn't really achieve our goal of keeping the pole vertical.\n",
    "Let's close that rendering window."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Cart pole is actually a fairly simple problem (it's very low dimensional), and so there are simpler ways to do this, but since we've been having so much fun with neural networks these past few days, let's use a neural network.\n",
    "Specifically, let's build a DQN that uses Q-learning to learn how to balance the pole.\n",
    "There are a lot of small details that go into making these models work well, so instead of coding it live, I've included the full code below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Based on: https://gym.openai.com/evaluations/eval_EIcM1ZBnQW2LBaFN6FY65g/\n",
    "\n",
    "import random\n",
    "import gym\n",
    "import math\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from collections import deque\n",
    "\n",
    "class DQNCartPoleSolver():\n",
    "    def __init__(self, n_episodes=1000, n_win_ticks=195, max_env_steps=None, gamma=1.0, epsilon=1.0, epsilon_min=0.01, epsilon_log_decay=0.995, alpha=0.01, alpha_decay=0.01, batch_size=64, monitor=False, quiet=False):\n",
    "        self.memory = deque(maxlen=100000)\n",
    "        self.env = gym.make('CartPole-v0')\n",
    "        if monitor: self.env = gym.wrappers.Monitor(self.env, '../data/cartpole-1', force=True)\n",
    "        self.gamma = gamma\n",
    "        self.epsilon = epsilon\n",
    "        self.epsilon_min = epsilon_min\n",
    "        self.epsilon_decay = epsilon_log_decay\n",
    "        self.alpha = alpha\n",
    "        self.alpha_decay = alpha_decay\n",
    "        self.n_episodes = n_episodes\n",
    "        self.n_win_ticks = n_win_ticks\n",
    "        self.batch_size = batch_size\n",
    "        self.quiet = quiet\n",
    "        if max_env_steps is not None: self.env._max_episode_steps = max_env_steps\n",
    "\n",
    "        # Init model\n",
    "        self.state_ = tf.placeholder(tf.float32, shape=[None, 4])\n",
    "        h = tf.layers.dense(self.state_, units=24, activation=tf.nn.tanh)\n",
    "        h = tf.layers.dense(h, units=48, activation=tf.nn.tanh)\n",
    "        self.Q = tf.layers.dense(h, units=2)\n",
    "        \n",
    "        self.Q_ = tf.placeholder(tf.float32, shape=[None, 2])\n",
    "        loss = tf.losses.mean_squared_error(self.Q_, self.Q)\n",
    "        self.global_step = tf.Variable(0, name='global_step', trainable=False)\n",
    "        lr = tf.train.exponential_decay(0.01, self.global_step, 0.995, 1)\n",
    "        self.train_step = tf.train.AdamOptimizer(lr).minimize(loss, global_step=self.global_step)\n",
    "        \n",
    "        self.sess = tf.Session()\n",
    "        self.sess.run(tf.global_variables_initializer())\n",
    "\n",
    "    def remember(self, state, action, reward, next_state, done):\n",
    "        self.memory.append((state, action, reward, next_state, done))\n",
    "\n",
    "    def choose_action(self, state, epsilon):\n",
    "        return self.env.action_space.sample() if (np.random.random() <= epsilon) else np.argmax(self.sess.run(self.Q, feed_dict={self.state_: state}))\n",
    "\n",
    "    def get_epsilon(self, t):\n",
    "        return max(self.epsilon_min, min(self.epsilon, 1.0 - math.log10((t + 1) * self.epsilon_decay)))\n",
    "\n",
    "    def preprocess_state(self, state):\n",
    "        return np.reshape(state, [1, 4])\n",
    "\n",
    "    def replay(self, batch_size):\n",
    "        x_batch, y_batch = [], []\n",
    "        minibatch = random.sample(\n",
    "            self.memory, min(len(self.memory), batch_size))\n",
    "        for state, action, reward, next_state, done in minibatch:\n",
    "            y_target = self.sess.run(self.Q, feed_dict={self.state_: state})\n",
    "            y_target[0][action] = reward if done else reward + self.gamma * np.max(self.sess.run(self.Q, feed_dict={self.state_: next_state})[0])\n",
    "            x_batch.append(state[0])\n",
    "            y_batch.append(y_target[0])\n",
    "        \n",
    "        self.sess.run(self.train_step, feed_dict={self.state_: np.array(x_batch), self.Q_: np.array(y_batch)})\n",
    "\n",
    "        if self.epsilon > self.epsilon_min:\n",
    "            self.epsilon *= self.epsilon_decay\n",
    "\n",
    "    def run(self):\n",
    "        scores = deque(maxlen=100)\n",
    "\n",
    "        for e in range(self.n_episodes):\n",
    "            state = self.preprocess_state(self.env.reset())\n",
    "            done = False\n",
    "            i = 0\n",
    "            while not done:\n",
    "                if e % 100 == 0 and not self.quiet:\n",
    "                    self.env.render()\n",
    "                action = self.choose_action(state, self.get_epsilon(e))\n",
    "                next_state, reward, done, _ = self.env.step(action)\n",
    "                next_state = self.preprocess_state(next_state)\n",
    "                self.remember(state, action, reward, next_state, done)\n",
    "                state = next_state\n",
    "                i += 1\n",
    "\n",
    "            scores.append(i)\n",
    "            mean_score = np.mean(scores)\n",
    "            if mean_score >= self.n_win_ticks and e >= 100:\n",
    "                if not self.quiet: print('Ran {} episodes. Solved after {} trials ✔'.format(e, e - 100))\n",
    "                return e - 100\n",
    "            if e % 100 == 0 and not self.quiet:\n",
    "                print('[Episode {}] - Mean survival time over last 100 episodes was {} ticks.'.format(e, mean_score))\n",
    "\n",
    "            self.replay(self.batch_size)\n",
    "        \n",
    "        if not self.quiet: print('Did not solve after {} episodes 😞'.format(e))\n",
    "        return e\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    agent = DQNCartPoleSolver()\n",
    "    agent.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Other materials:\n",
    "Reinforcement Learning can easily be a full course (or 2) on its own at most universities.\n",
    "Given limited time, we can only touch on a couple points.\n",
    "If you're interested in *exploring* this field further, a few recommendations to get you started:\n",
    "\n",
    "- [The definitive textbook on Reinforcement Learning](http://incompleteideas.net/book/the-book-2nd.html)\n",
    "- [Video lectures by David Silver](https://www.youtube.com/playlist?list=PLqYmG7hTraZDM-OYHWgPebj2MfCFzFObQ)\n",
    "- [Deep Q-Networks (DQN) for Atari 2600 games](https://deepmind.com/research/dqn/) \n",
    "- [Popular blog post on using policy gradients to learn to play Pong from pixels](http://karpathy.github.io/2016/05/31/rl/)\n",
    "- [An amusing video of DeepMind AI learning to walk](https://www.youtube.com/watch?v=gn4nRCC9TwQ) "
   ]
  }
 ],
 "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}