lesson-09.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Connect Intensive - Machine Learning Nanodegree\n",
    "# Lesson 09: Reinforcement Learning - Gridworld!\n",
    "\n",
    "## Objectives\n",
    "  - Solve the state value function $v(s)$ for the gridworld problem iteratively.\n",
    "  - Use the greedy policy to improve the random policy for gridworld actions.\n",
    "  - Visualize the state value function and greedy policy after each iteration.\n",
    "  \n",
    "## Prerequisites\n",
    "  - You should have the following python packages installed:\n",
    "    - [matplotlib](http://matplotlib.org/index.html)\n",
    "    - [numpy](http://www.scipy.org/scipylib/download.html)\n",
    "\n",
    "    \n",
    "## Acknowledgements\n",
    "The code from this notebook is loosely inspired by [Shangtong Zhang's python code](https://github.com/ShangtongZhang/reinforcement-learning-an-introduction), particularly [the GridWorld.py code](https://github.com/ShangtongZhang/reinforcement-learning-an-introduction/blob/master/chapter03/GridWorld.py) that accompanies Chapter 3 of Sutton and Barto's textbook [Reinforcement Learning: an Introduction](https://webdocs.cs.ualberta.ca/~sutton/book/the-book.html). Shangtong Zhang's copyright notice is in the cell below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#######################################################################\n",
    "# Copyright (C) 2016 Shangtong Zhang(zhangshangtong.cpp@gmail.com)    #\n",
    "# Permission given to modify the code as long as you keep this        #\n",
    "# declaration at the top                                              #\n",
    "#######################################################################"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Agents in Gridworld\n",
    "\n",
    "This week's lesson notebook is pretty brief, but it should illustrate how the [Bellman equation](https://en.wikipedia.org/wiki/Bellman_equation) can be solved iteratively for a prescribed policy. The prescribed policy for our reinforcement learning agent is random movement: the agent is navigating a 2-d world and at every timestep can move north, south, east, or west in an effort to reach a terminal state (or a goal) in the gridworld. To hurry the agent along, there is a negative reward of -1 at each timestep until the agent is in the goal.\n",
    "\n",
    "To evaluate how well the random policy works in the gridworld, we can iteratively solve the Bellman equation until the state value function converges for each state. The `Grid` class below has methods to do so, and to visualize our progress along the way:\n",
    "  - `iterate_value_function()`: Performs one step of the iteration toward solving the Bellman equation.\n",
    "  - `iterate_to_convergence()`: Iterates until the state value function converges to within a specified tolerance.\n",
    "  - `show_greedy_policy()`: Shows the state value function for gridworld and the corresponding greedy policy at the current iteration.\n",
    "  \n",
    "Read through, then **run** the cell below so that we can use the `Grid` class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.patches as patches\n",
    "\n",
    "class Grid:\n",
    "    \n",
    "    def __init__(self, nrows=5, ncols=5, goals = [(4,4)], discount = 1.0):\n",
    "    # @nrows: (integer) The number of rows in the gridworld\n",
    "    # @ncols: (integer) The number of columns in the gridworld\n",
    "    # @goals: (list of tuples) The cell or cells that are terminal state(s in gridworld\n",
    "    # @discount: (float) The discount factor, gamma, for computing state value function\n",
    "    \n",
    "        # Set the dimensions of gridworld\n",
    "        self.nrows = nrows\n",
    "        self.ncols = ncols\n",
    "        \n",
    "        # Counter for number of iterations in solution of Bellman equation\n",
    "        # for the state value function v(s)\n",
    "        self.k = 0\n",
    "        \n",
    "        # Initialize the state value function to be zero for all states\n",
    "        self.values = np.zeros((nrows,ncols))\n",
    "        \n",
    "        # RL agent is finding a goal or terminal state as fast as possible,\n",
    "        # so reward is -1 for all actions to incentivize fast routes to the goal\n",
    "        self.rewards = - np.ones((nrows,ncols))\n",
    "        \n",
    "        # Terminal states for the gridworld\n",
    "        self.goals = goals\n",
    "        \n",
    "        # The discount factor, gamma\n",
    "        self.discount = discount\n",
    "        \n",
    "        # Policy is random -- RL agent randomly selects one of four actions\n",
    "        self.prob = {'N':0.25,'S':0.25,'E':0.25,'W':0.25}\n",
    "        \n",
    "        # Summed difference between new and old state value functions\n",
    "        # for one interation of Bellman equation (check for convergence)\n",
    "        self.diff = np.inf\n",
    "\n",
    "        # Tolerance for convergence of the iterative solution of Bellman equation\n",
    "        self.tolerance = 1e-4\n",
    "            \n",
    "    def iterate_value_function(self):\n",
    "        # Perform one iteration of iterative Bellman equation to determine\n",
    "        # the state value function for the random policy in the gridworld\n",
    "        \n",
    "        # Increment the number of iterations\n",
    "        self.k += 1\n",
    "        \n",
    "        # Initialize a matrix for the new state value function after one iteration\n",
    "        new_values = np.zeros((self.nrows,self.ncols))\n",
    "        \n",
    "        # We are going to iterate over each cell within the gridworld\n",
    "        for row in np.arange(self.nrows):\n",
    "            for col in np.arange(self.ncols):\n",
    "                # Only update non-terminal states\n",
    "                if (row,col) not in self.goals:\n",
    "                    # When solving the Bellman equations, you add two things:\n",
    "                    #  1. the immediate reward upon taking the action, and\n",
    "                    #  2. the discounted value of the state to which you moved\n",
    "                    \n",
    "                    # Action: Move North\n",
    "                    if row == 0:\n",
    "                        # Cannot move North in the top row of gridworld\n",
    "                        north = self.rewards[( row , col )] + self.discount * self.values[( row , col )]\n",
    "                    else:\n",
    "                        north = self.rewards[(row-1, col )] + self.discount * self.values[(row-1, col )]\n",
    "                        \n",
    "                    # Action: Move South\n",
    "                    if row == self.nrows-1:\n",
    "                        # Cannot move South in the bottom row of gridworld\n",
    "                        south = self.rewards[( row , col )] + self.discount * self.values[( row , col )]\n",
    "                    else:\n",
    "                        south = self.rewards[(row+1, col )] + self.discount * self.values[(row+1, col )]\n",
    "                        \n",
    "                    # Action: Move West\n",
    "                    if col == 0:\n",
    "                        # Cannot move West in the leftmost column of gridworld\n",
    "                        west = self.rewards[( row , col )] + self.discount * self.values[( row , col )]\n",
    "                    else:\n",
    "                        west = self.rewards[( row ,col-1)] + self.discount * self.values[( row ,col-1)]\n",
    "                        \n",
    "                    # Action: Move East\n",
    "                    if col == self.ncols-1:\n",
    "                        # Cannot move East in the rightmost column of gridworld\n",
    "                        east = self.rewards[( row , col )] + self.discount * self.values[( row , col )]\n",
    "                    else:\n",
    "                        east = self.rewards[( row ,col+1)] + self.discount * self.values[( row ,col+1)]\n",
    "                        \n",
    "                    # Add up all contributions to the new state value function\n",
    "                    new_values[(row,col)] = self.prob['N'] * north + \\\n",
    "                                            self.prob['S'] * south + \\\n",
    "                                            self.prob['E'] * east  + \\\n",
    "                                            self.prob['W'] * west\n",
    "                else:\n",
    "                    # Value of terminal state is unchanged\n",
    "                    new_values[(row,col)] = self.values[(row,col)]\n",
    "                    \n",
    "        # Find difference between old and new state value function\n",
    "        self.diff = np.sum(np.abs(self.values-new_values))\n",
    "        \n",
    "        # Save new values\n",
    "        self.values = new_values\n",
    "        \n",
    "    def iterate_to_convergence(self):\n",
    "        # Iteratively solve the Bellman equation until convergence\n",
    "        while self.diff > self.tolerance:\n",
    "            self.iterate_value_function()\n",
    "        print(\"Converged to tolerance of {}!\".format(self.tolerance))\n",
    "            \n",
    "    def show_greedy_policy(self):\n",
    "        # Create two subplots\n",
    "        # Subplot 1: Current state value function v_k(s), after k iterations\n",
    "        # Subplot 2: Improvement: greedy policy using v_k(s), after k iterations\n",
    "        fig, (ax1, ax2) = plt.subplots(1,2,figsize=(10, 5))\n",
    "        \n",
    "        # Set x and y limits of both subplots based on gridworld dimensions\n",
    "        ax1.set_xlim([0,self.ncols])\n",
    "        ax1.set_ylim([-self.nrows,0])\n",
    "        ax2.set_xlim([0,self.ncols])\n",
    "        ax2.set_ylim([-self.nrows,0])\n",
    "        \n",
    "        # Iterate through each cell of the gridworld\n",
    "        for row in np.arange(self.nrows):\n",
    "            for col in np.arange(self.ncols):\n",
    "                # Terminal states or goal cells will be gold\n",
    "                # Non-terminal states will be empty\n",
    "                if (row,col) in self.goals:\n",
    "                    color_fill = \"gold\"\n",
    "                else:\n",
    "                    color_fill = \"none\"\n",
    "                # Add a rectangle patch for the current cell of gridworld\n",
    "                ax1.add_patch(\n",
    "                    patches.Rectangle(\n",
    "                    (col, -(row+1)),\n",
    "                    1.0,\n",
    "                    1.0,\n",
    "                    facecolor=color_fill,\n",
    "                    linewidth=3\n",
    "                    )\n",
    "                )\n",
    "                ax2.add_patch(\n",
    "                    patches.Rectangle(\n",
    "                    (col, -(row+1)),\n",
    "                    1.0,\n",
    "                    1.0,\n",
    "                    facecolor=color_fill,\n",
    "                    linewidth=3\n",
    "                    )\n",
    "                )\n",
    "                # For subplot 1, add text for the state value function\n",
    "                # at the current iteration of the Bellman equation\n",
    "                ax1.text(col+0.5, -(row+1)+0.5,'{:.2f}'.format(self.values[(row,col)]),\n",
    "                        horizontalalignment='center',\n",
    "                        verticalalignment='center',\n",
    "                        fontsize=12)\n",
    "                # For subplot 2, draw arrows representing the \n",
    "                # improved (greedy) policy for all non-terminal states\n",
    "                if (row,col) not in self.goals:\n",
    "                    # max_dir will contain a list of the greedy policy actions\n",
    "                    # for the current state (row, col)\n",
    "                    max_dir = []\n",
    "                    # Initialize the maximum value to -infinity\n",
    "                    max_val = -np.inf\n",
    "                    # check state value function for \"north\" action\n",
    "                    if not row == 0:\n",
    "                        if np.abs(max_val - self.values[(row-1,col)]) < self.tolerance:\n",
    "                            max_dir.append('N')\n",
    "                        elif self.values[(row-1,col)] > max_val:\n",
    "                            max_val = self.values[(row-1,col)]\n",
    "                            max_dir = ['N']\n",
    "                    elif np.abs(max_val - self.values[(row,col)]) < self.tolerance:\n",
    "                        max_dir.append('N')\n",
    "                    elif self.values[(row,col)] > max_val:\n",
    "                        max_val = self.values[(row,col)]\n",
    "                        max_dir = ['N']\n",
    "\n",
    "                        \n",
    "                    # check state value function for \"south\" action\n",
    "                    if not row == self.nrows-1:\n",
    "                        if np.abs(max_val - self.values[(row+1,col)]) < self.tolerance:\n",
    "                            max_dir.append('S')\n",
    "                        elif self.values[(row+1,col)] > max_val:\n",
    "                            max_val = self.values[(row+1,col)]\n",
    "                            max_dir = ['S']\n",
    "                    elif np.abs(max_val - self.values[(row,col)]) < self.tolerance:\n",
    "                        max_dir.append('S')\n",
    "                    elif self.values[(row,col)] > max_val:\n",
    "                        max_val = self.values[(row,col)]\n",
    "                        max_dir = ['S']\n",
    "\n",
    "                        \n",
    "                    # check state value function for \"west\" action\n",
    "                    if not col == 0:\n",
    "                        if np.abs(max_val - self.values[(row,col-1)]) < self.tolerance:\n",
    "                            max_dir.append('W')\n",
    "                        elif self.values[(row,col-1)] > max_val:\n",
    "                            max_val = self.values[(row,col-1)]\n",
    "                            max_dir = ['W']\n",
    "                    elif np.abs(max_val - self.values[(row,col)]) < self.tolerance:\n",
    "                        max_dir.append('W')\n",
    "                    elif self.values[(row,col)] > max_val:\n",
    "                        max_val = self.values[(row,col)]\n",
    "                        max_dir = ['W']\n",
    "\n",
    "                        \n",
    "                    # check state value function for \"east\" action\n",
    "                    if not col == self.ncols-1:\n",
    "                        if np.abs(max_val - self.values[(row,col+1)]) < self.tolerance:\n",
    "                            max_dir.append('E')\n",
    "                        elif self.values[(row,col+1)] > max_val:\n",
    "                            max_val = self.values[(row,col+1)]\n",
    "                            max_dir = ['E']\n",
    "                    elif np.abs(max_val - self.values[(row,col)]) < self.tolerance:\n",
    "                        max_dir.append('E')\n",
    "                    elif self.values[(row,col)] > max_val:\n",
    "                        max_val = self.values[(row,col)]\n",
    "                        max_dir = ['E']\n",
    "\n",
    "                        \n",
    "                    # Arrows are centered in each gridworld cell\n",
    "                    x = col+0.5\n",
    "                    y = -(row+1)+0.5\n",
    "                    \n",
    "                    # For each direction in the greedy policy, draw the arrow\n",
    "                    for direc in max_dir:\n",
    "                        if direc == 'N':\n",
    "                            dx = 0\n",
    "                            dy = 0.25\n",
    "                        elif direc == 'S':\n",
    "                            dx = 0\n",
    "                            dy = -0.25\n",
    "                        if direc == 'E':\n",
    "                            dx = 0.25\n",
    "                            dy = 0\n",
    "                        elif direc == 'W':\n",
    "                            dx = -0.25\n",
    "                            dy = 0\n",
    "                        ax2.arrow(x,y,dx,dy,head_width=0.1, fc='k', ec='k')\n",
    "        \n",
    "        # Remove tick marks on both subplots, and title/label the subplots\n",
    "        ax1.axes.get_xaxis().set_ticks([])\n",
    "        ax1.axes.get_yaxis().set_ticks([])\n",
    "        ax1.set_title('State Value Function $v_k(s)$\\nRandom Policy, $k={}$ Iteration(s)'.format(self.k))\n",
    "        ax1.set_xlabel('Discount Factor $\\gamma={:.1f}$'.format(self.discount),fontsize=12)\n",
    "        ax2.set_axis_off()\n",
    "        ax2.set_title('Greedy Policy\\nwith respect to $v_k(s)$')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's create a gridworld! Let's make it $4\\times4$ with a single terminal state in the southeast corner, cell (3,3). Let's be far-sighted and set the discount factor $\\gamma$ for the [Markov Decision Process](https://en.wikipedia.org/wiki/Markov_decision_process) to 1.0. This way, we're not discounting long-term rewards. **Run** the cell below to do so. You will see that the state value function $v_k(s)$ has been initialized to zero for all states, and if we were to create a greedy policy from this initial value, no action is preferred over any other just yet. *Note:* actions that would move the agent outside the gridworld (for example, the *North* action on the topmost edge of gridworld) leave the agent where it is for the current timestep."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "g1 = Grid(nrows=4, ncols=4, goals = [(3,3)], discount = 1.0)\n",
    "g1.show_greedy_policy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now **run** the cell below to perform one iteration in the solution of the Bellman equation to find $v_1(s)$. You should find that all but the goal state get a negative reward of -1. Does that influence the resulting greedy policy?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "g1.iterate_value_function()\n",
    "g1.show_greedy_policy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now **run** the cell below to perform a second iteration in the solution of the Bellman equation to find $v_1(s)$. Things start to get interesting. The information from the goal state starts to propagate further into the gridworld, and the cells farther from the gridworld goal continue to get penalized the hardest."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "g1.iterate_value_function()\n",
    "g1.show_greedy_policy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Run** the cell below for a third, fourth, fifth, and sixth iteration..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Third iteration\n",
    "g1.iterate_value_function()\n",
    "g1.show_greedy_policy()\n",
    "\n",
    "# Fourth iteration\n",
    "g1.iterate_value_function()\n",
    "g1.show_greedy_policy()\n",
    "\n",
    "# Fifth iteration\n",
    "g1.iterate_value_function()\n",
    "g1.show_greedy_policy()\n",
    "\n",
    "# Sixth iteration\n",
    "g1.iterate_value_function()\n",
    "g1.show_greedy_policy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Rather than iterate one step at a time, we can iterate until the state value function converges to within a specified tolerance using the `iterate_to_convergence` method. **Run** the cell below to do so.  The converged state value function and corresponding greedy policy for this gridworld will be printed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Iterate to convergence\n",
    "g1.iterate_to_convergence()\n",
    "g1.show_greedy_policy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once the state value function $v_k(s)$ has converged, we have **evaluated** our policy $\\pi$ (in this case, our policy $\\pi$ was the random policy). The next step would be to **improve** the policy by replacing $\\pi$ with the greedy policy $\\pi'$. We would then need to solve the Bellman equation again for the new policy $\\pi'$ (policy evaluation), until the state value function $v'_k(s)$ for the new policy converges. This back-and-forth between policy evaluation and policy improvement is called **policy iteration.** Eventually, this process leads to an optimal policy.\n",
    "\n",
    "As always, feel free to play around with the code and try your own examples -- create a larger gridworld, add multiple goals, change the discount factor. If you'd like to see another example of policy evaluation, **run** the cells below to see how the state value function converges for a gridworld with two terminal states."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Initialize gridworld with two terminal states\n",
    "g2 = Grid(nrows=4, ncols=4, goals = [(0,0),(3,3)], discount = 1.0)\n",
    "g2.show_greedy_policy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# First iteration\n",
    "g2.iterate_value_function()\n",
    "g2.show_greedy_policy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Second iteration\n",
    "g2.iterate_value_function()\n",
    "g2.show_greedy_policy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Iterate to convergence\n",
    "g2.iterate_to_convergence()\n",
    "g2.show_greedy_policy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}