S35_vacuum_decision.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github",
    "tags": [
     "no-tex"
    ]
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/gtbook/robotics/blob/main/S35_vacuum_decision.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "JoW4C_OkOMhe",
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip install -q -U gtbook\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "10-snNDwOSuC",
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import gtsam\n",
    "import pandas as pd\n",
    "\n",
    "import gtbook\n",
    "import gtbook.display\n",
    "from gtbook import vacuum\n",
    "from gtbook.discrete import Variables\n",
    "VARIABLES = Variables()\n",
    "def pretty(obj):\n",
    "    return gtbook.display.pretty(obj, VARIABLES)\n",
    "def show(obj, **kwargs):\n",
    "    return gtbook.display.show(obj, VARIABLES, **kwargs)\n",
    "\n",
    "# From section 3.2:\n",
    "N = 5\n",
    "X = VARIABLES.discrete_series(\"X\", range(1, N+1), vacuum.rooms)\n",
    "A = VARIABLES.discrete_series(\"A\", range(1, N), vacuum.action_space)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "nAvx4-UCNzt2"
   },
   "source": [
    "# Markov Decision Processes\n",
    "\n",
    "> For controlled Markov chains, planning is the process of choosing the control inputs. This leads to\n",
    "the concept of Markov decision processes (MDPs).\n",
    "\n",
    "<img src=\"Figures3/S35-iRobot_vacuuming_robot-06.jpg\" alt=\"Splash image with thinking vacuum robot\" width=\"40%\" align=center style=\"vertical-align:middle;margin:10px 0px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Previously in this chapter, we described how conditional probability distributions can be used to model uncertainty\n",
    "in the effects of actions. We defined the belief state $b_{k+1}$ to be the posterior probability distribution\n",
    "for the state at time $k+1$ given the sequence of actions $a_1 \\dots a_k$.\n",
    "In every example, the sequence of actions was predetermined, and we merely calculated probabilities\n",
    "associated with performing these actions from some specified initial state, described\n",
    "by a probability distribution $P(X_1)$.\n",
    "\n",
    "In this section, we consider the problem of choosing which actions to execute.\n",
    "Making these decisions requires that we have quantitative criteria for evaluating actions and their effects.\n",
    "We encode these criteria using *reward functions*.\n",
    "Because the effects of actions are uncertain, it is not possible to know the reward\n",
    "that will be obtained by executing a specific action (or a sequence of actions). Thus, we will again\n",
    "invoke the concept of expectation to compute expected future benefits of applying actions.\n",
    "\n",
    "As we work through these concepts, it will rapidly become apparent that executing a predefined sequence of actions\n",
    "is not the best way to face the future.\n",
    "Suppose, for example, that our vacuuming robot wants to move from the office to the living room.\n",
    "Because it is not possible to know how many times we should \"move right\"\n",
    "to arrive at the hallway, we might construct a sequence of actions\n",
    "to move right many times (to increase the probability of arriving at the dining room),\n",
    "then move up many times (to increase the probability of arriving at the kitchen),\n",
    "then move left many times (to increase the probability of arriving at the living room).\n",
    "\n",
    "This kind of plan might make sense if the robot is unable to know its current location.\n",
    "For example, moving for a long time to the right should eventually bring the robot to the dining room with\n",
    "high probability,\n",
    "but if the robot moves to the right for a short time, it could remain in the office,\n",
    "arrive at the hallway, or arrive at the dining room.\n",
    "\n",
    "In contrast, if the robot is able to know its location (using perception), it can act\n",
    "opportunistically when it reaches the hallway, and immediately move up.\n",
    "The key idea here is that the best action at any moment in time depends\n",
    "on the state in which the action is executed.\n",
    "The recipe of which action to execute in each state is called a *policy*,\n",
    "and defining policies and their associated *value function* is the main goal for this section."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reward Functions\n",
    "\n",
    "> Reward functions provide a quantitative way to evaluate the effects of actions.\n",
    "\n",
    "How should our robot evaluate the effects of actions?\n",
    "One simple approach might be to specify the robot's goal, and then determine\n",
    "how effective an action might be with respect to that goal,\n",
    "e.g., by using the conditional probability tables for the various actions.\n",
    "If, for example, the goal is to arrive to the living room, we could evaluate\n",
    "the probability that a particular action would achieve the goal.\n",
    "\n",
    "This approach has several limitations.  First, if the robot is not currently in a room\n",
    "that is adjacent to the living room, it is not clear how to measure the possible progress\n",
    "toward the living room for a specific action.\n",
    "If, for example, the robot is in the dining room, would it be better to move\n",
    "to the kitchen or to the hallway as an intermediate step toward the living room?\n",
    "Second, this approach does not allow the robot to consider\n",
    "benefits that could arise from visiting other states.\n",
    "For example, even though the living room is the goal, arriving to the kitchen might not be\n",
    "such a bad outcome, if for example, the kitchen floor is also in need of cleaning.\n",
    "Finally, it may be the case that the benefit of executing an action depends\n",
    "not only on the destination, but also on the state in which the action\n",
    "is executed.\n",
    "For example, entering the living room from the hallway might be less desirable\n",
    "than entering from the kitchen (if, for example, guests are likely to arrive\n",
    "in the entrance hallway)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Reward functions provide a very general solution that addresses all of these limitation.\n",
    "Let $\\cal X$ denote the set of states and $\\cal A$ the set of actions.\n",
    "The reward function, $R : {\\cal X} \\times {\\cal A} \\times {\\cal X} \\rightarrow \\mathbb{R}$,\n",
    "assigns a numeric reward to specific state transitions under specific actions.\n",
    "In particular, $R(x_k, a_k, x_{k+1})$ is the reward obtained by arriving to state\n",
    "$x_{k+1}$ from state $x_k$ as a result of executing action $a_k$.\n",
    "\n",
    "In what follows we assume that the reward function is time-invariant, i.e., the benefit of moving up from\n",
    "the hallway to the living room does not change as time passes.\n",
    "This allows us to use the more compact notation\n",
    "$R(x, a, x')$ to represent the reward obtained by arriving to (the next) state $x'$\n",
    "by executing action $a$ in state $x$. We will frequently use this notation in the remainder of the section."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This form of the reward function is very general, allowing us to encode context dependent rewards\n",
    "that depend on the state in which an action is executed.\n",
    "It is also common to specify rewards merely as a function of state.\n",
    "In this case, we denote by $R(x)$ the reward obtained for arriving to state $x$, regardless\n",
    "of the action that was applied, and regardless of the previous state.\n",
    "An example of such a reward function is given below, implemented as a python function.\n",
    "It simply returns a reward of 10 upon entering the living room:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "10.0\n",
      "0.0\n"
     ]
    }
   ],
   "source": [
    "def reward_function(state:int, action:int, next_state:int):\n",
    "    \"\"\"Reward that returns 10 upon entering the living room.\"\"\"\n",
    "    return 10.0 if next_state == \"Living Room\" else 0.0\n",
    "\n",
    "print(reward_function(\"Kitchen\", \"L\", \"Living Room\"))\n",
    "print(reward_function(\"Kitchen\", \"L\", \"Kitchen\"))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Expected Reward\n",
    "\n",
    "> Maximizing immediate expected reward leads to greedy planning.\n",
    "\n",
    "Because the effects of actions are uncertain,\n",
    "there is no way for the robot to deterministically predict the value\n",
    "of the next state, $x'$, for a particular action $a$;\n",
    "therefore, it is not possible to evaluate $R(x,a,x')$.\n",
    "In this case, we can use the **expected reward** to evaluate the stochastic\n",
    "effects of the action.\n",
    "As we have seen in the previous chapter, this expectation will provide a good \n",
    "estimate of the average reward that will be received if action $a$ is executed from state $x$ may times,\n",
    "even if it provides no guarantees about the reward that will be obtained by any specific moment in time.\n",
    "\n",
    "We denote the expected reward for executing action $a$ in state $x$\n",
    "by $\\bar{R}(x,a)$,\n",
    "and its value is obtained by evaluating the expectation\n",
    "\n",
    "$$\n",
    "\\bar{R}(x,a) \\doteq E[R(x,a,X')] = \\sum_{x'} P(x'|x,a) R(x, a, x').\n",
    "$$\n",
    "\n",
    "Note that we use the upper case $X'$ to indicate that the expectation is taken with\n",
    "respect to the random next state.\n",
    "Accordingly, the sum is over all possible next states, and the reward $R(x, a, x')$ for each next\n",
    "state is weighted by the probability $P(x'|x,a)$ of arriving to that state from state $x$ by executing action $a$.\n",
    "\n",
    "<!-- In some cases (e.g., when the transition probabilities are *known*),\n",
    "it can be convenient to work only with the expected reward, instead of dealing wit $R(x,a,x')$\n",
    "which some texts do. However, below we will continue to work with the most general formulation. -->"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To *calculate* the expected reward in code, we first create a multidimensional array $R(x,a,x')$ that contains the rewards for every possible transition $x,a \\rightarrow x'$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([10.,  0.,  0.,  0.,  0.])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "R = np.empty((5, 4, 5), float)\n",
    "for x, room_x in enumerate(vacuum.rooms):\n",
    "    for a, action in enumerate(vacuum.action_space):\n",
    "        for y, room_y in enumerate(vacuum.rooms):\n",
    "            R[x, a, y] = 10.0 if room_y == \"Living Room\" else 0.0\n",
    "\n",
    "# For example, taking action \"L\" in \"Kitchen\":\n",
    "R[vacuum.rooms.index(\"Kitchen\"), vacuum.action_space.index(\"L\")]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The result above is simple to interpret: if the robot starts in the kitchen and goes left, it obtains a reward of 10 *iff* it arrives to the living room (index 0), otherwise 0.\n",
    "\n",
    "Observe that $R$ is a *multidimensional* array, with shape $(5, 4, 5)$, corresponding to the dimensionality of the arguments of the reward function $R(x,a,x')$.\n",
    "We can do the same with the transition probabilities, creating a multidimensional array $T$ with an entry for every possible transition $x \\rightarrow a \\rightarrow x'$ with probability $P(x'|a,x)$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.8, 0.2, 0. , 0. , 0. ])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# First create a conditional with ids 0, 1, and 2 for state, action, and next state:\n",
    "conditional = gtsam.DiscreteConditional((2,5), [(0,5), (1,4)], vacuum.action_spec)\n",
    "\n",
    "# Now use `enumerate` to iterate over all possible assignments to the conditional:\n",
    "T = np.empty((5, 4, 5), float)\n",
    "for assignment, P_y_given_ax in conditional.enumerate():\n",
    "    x, a, y = assignment[0], assignment[1], assignment[2]\n",
    "    T[x, a, y] = P_y_given_ax\n",
    "\n",
    "# For example, taking action \"L\" in \"Kitchen\":\n",
    "T[vacuum.rooms.index(\"Kitchen\"), vacuum.action_space.index(\"L\")]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As expected, if we execute the action L when the robot is in the kitchen,\n",
    "we end up in the living room with 80% chance, while staying in place (the kitchen) 20% of the time.\n",
    "\n",
    "Then, finally, calculating the expected reward $\\bar{R}(x,a)$ for a particular state-action pair $(x,a)$ becomes a simple dot product:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Expected reward (Kitchen, L) = 8.0\n"
     ]
    }
   ],
   "source": [
    "x, a = vacuum.rooms.index(\"Kitchen\"), vacuum.action_space.index(\"L\")\n",
    "print(f\"Expected reward (Kitchen, L) = {np.dot(T[x,a], R[x,a])}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This makes total sense: if we stay in the kitchen, we get zero reward, but *if* we succeed to move to the living room, we receive a reward of 10. This happens only in 80% of the cases, though, so the *expected reward* is only 8."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Note that this computation can also be realized using belief states.  In\n",
    "particular, if encode the reward function as a vector (as above),\n",
    "say $R = [10, 0, 0, 0, 0]^T$,\n",
    "the the expected reward at time $k+1$ is given by\n",
    "\n",
    "$$E[ R(X_k, a, X_{k+t})] = b_{k+1} R = \\left( b_k M_{a}\\right) R\n",
    "$$\n",
    "\n",
    "in which $M_{a}$ denotes the transition probability matrix associated\n",
    "to action $ a$. Note that in this case, we the expected reward is computed for a given\n",
    "prior distribution on $X_k$. If we know with certainty that $X_k =x$, the belief\n",
    "vector would have a $1$ in the entry corresponding to the current room, and zero's for\n",
    "all other entries.\n",
    "Therefore, this form is somewhat more general than that given above. -->\n",
    "\n",
    "Equipped with the definition of expected reward, we can introduce a first, **greedy planning** algorithm:\n",
    "Given the current state $x_k$, execute the action that maximizes the expected reward:\n",
    "\n",
    "$$\n",
    "a_k^* = \\arg  \\max_{a \\in {\\cal A}} E[ R(x_k, a, X_{k+t})]\n",
    "$$\n",
    "\n",
    " <!-- = \\arg \\max_{a \\in {\\cal A}}\\left( b_k M_{A}\\right) R -->"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For example, in the kitchen, we can calculate the expected award for all actions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Expected reward (Kitchen, L) = 8.0\n",
      "Expected reward (Kitchen, R) = 0.0\n",
      "Expected reward (Kitchen, U) = 0.0\n",
      "Expected reward (Kitchen, D) = 0.0\n"
     ]
    }
   ],
   "source": [
    "x = vacuum.rooms.index(\"Kitchen\")\n",
    "for a in range(4):\n",
    "    print(f\"Expected reward ({vacuum.rooms[x]}, {vacuum.action_space[a]}) = {np.dot(T[x,a], R[x,a])}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It turns out that $8.0$ is the best we can do by greedily choosing a single action, as the expected *immediate* reward for any other action is zero. Hence, whenever we find ourselves in the Kitchen, a good course of action is to always go left!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Utility for a Defined Sequence of Actions\n",
    "\n",
    "> Utility is the total discounted reward over finite or infinite horizons.\n",
    "\n",
    "The greedy strategy above focuses on the immediate benefit of applying an action.\n",
    "For our robot, and for most robots operating in the real world, it is important to perform\n",
    "effectively over a prolonged period of time, not merely for the next instant.\n",
    "In order to maximize long-term benefit, instead of considering only the next-stage reward (as the greedy strategy does)\n",
    "we could consider the sum of all rewards achieved in the future.\n",
    "There are two immediate disadvantages to a direct implementation of this approach:\n",
    "\n",
    "- First, because the effects of actions are uncertain, it is likely that the further we look\n",
    "into the future, the more uncertain we will be about the robot's anticipated state.\n",
    "Therefore, it makes sense to *discount* our consideration of rewards that might occur far into the future,\n",
    "say at times $X_{k+l}$, for increasing values of $l$.\n",
    "\n",
    "- Second, it is often convenient to reason with an infinite time horizon, e.g., consider the case when the robot will operate forever.\n",
    "While this is certainly not a realistic expectation,\n",
    "reasoning over an infinite time horizon often simplifies the mathematical complexities of planning into the future.\n",
    "If we merely compute the sum of all future rewards, this sum will diverge to infinity as $l$ approaches infinity."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can deal with both of these disadvantages by multiplying the reward at time step $k + l$ by a discounting\n",
    "factor $\\gamma^l$, with $0 < \\gamma \\leq 1$.\n",
    "We refer to $\\gamma$ as the **discount factor**\n",
    "and to the term $\\gamma^l R(x_{k+l},a_{k+l},x_{k+l+1})$ as a **discounted reward.**\n",
    "Note that for $\\gamma = 1$, there is no discount, and all future rewards are treated with equal weight.\n",
    "\n",
    "Suppose the robot executes a sequence of actions, $a_1, \\dots, a_n$,\n",
    "starting in state  $X_1=x_1$, and passes through\n",
    "state sequence $x_1,x_2,x_3\\dots x_{n+1}$.\n",
    "We define the utility function $U: {\\cal A}^n \\times {\\cal X}^{n+1} \\rightarrow \\mathbb{R}$ as\n",
    "\n",
    "$$\n",
    "U(a_1, \\dots, a_n, x_1, \\dots x_{n+1}) =\n",
    "R(x_1,a_1, x_2) + \\gamma R(x_2, a_2, x_3) + \\dots \\gamma^{n-1} R(x_{n}, a_{n}, x_{n+1})$$\n",
    "\n",
    "We can write this more compactly as the summation\n",
    "\n",
    "$$\n",
    "U(a_1, \\dots, a_n, x_1, \\dots x_{n+1}) =\n",
    "\\sum_{k=1}^{n} = \\gamma^{k-1} R(x_k, a_k, x_{k+1}) \n",
    "$$\n",
    "\n",
    "For $n < \\infty$, we refer to this as a finite-horizon utility function.\n",
    "Note that influence of future rewards decreases exponentially with the time horizon,\n",
    "and that the use of $\\gamma^{k-1}$ ensures that the sum will converge for the infinite horizon case (under mild assumptions about $R$)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The expression above is defined for a specific sequence of actions and a specific sequence of states.\n",
    "When planning, as we have noted above, we are unable to know with certainty the future states.\n",
    "We can, again, deal with this difficulty by computing the *expected* utility for a\n",
    "given sequence of actions, \n",
    "$E[U(a_1, \\dots, a_n, X_1, \\dots X_n)]$.\n",
    "We can now formulate a slightly more sophisticated version of our planning problem. Choose the sequence of actions $a_{1:n}$ that maximizes the expected utility:\n",
    "\n",
    "$$ a_1^*, \\dots a_n^* = \\arg  \\max_{a_1 \\dots a_n \\in {\\cal A}^n} E[U(a_1, \\dots, a_n, X_1, \\dots X_{n+1})]\n",
    "$$\n",
    "\n",
    "As formulated above, this problem could be solved by merely enumerating every possible action sequence,\n",
    "and choosing the sequence that maximizes the expectation.\n",
    "Obviously this is not a computationally tractable approach.\n",
    "Not only does the number of possible action sequences grow exponentially with the time horizon $n$,\n",
    "but the computation of the expectation for a specific action sequence is also computationally heavy.\n",
    "We can, however, approximate this optimization process using the concept of rollouts, as we will now see."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Approximating Expected Utility Using Control Tape Rollouts\n",
    "\n",
    "> Rollin', rollin', rollin'...\n",
    "\n",
    "Consider the computation required to determine the expected utility for a sequence of only two actions:\n",
    "\n",
    "$$\n",
    "E[U(a_1,a_2, X_1, X_2, X_3)] = E[R(X_1, a_1, X_2) + \\gamma R(X_2, a_2, X_3)]\n",
    "$$\n",
    "\n",
    "Computing the expectation requires summing over all combinations of values for states\n",
    "$X_1, X_2, X_3$.  Clearly, as $n$ becomes large, this is not a tractable computation.\n",
    "We can approximate this computation by recalling the relationship between\n",
    "the expected value for a probability distribution and the average value over\n",
    "many realizations of the underlying random process.\n",
    "Namely, as we have seen before, the expected value of a random variable (in this case the value of the utility function)\n",
    "corresponds to what we expect to observe as the average of that random variable over many trials.\n",
    "This immediately suggests an approximation algorithm: \n",
    "\n",
    "- Generate many sample trajectories for the action sequence $a_1, \\dots, a_n$;\n",
    "- compute the average of the discounted rewards over these sample trajectories.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A specific sequence of actions, $a_1, \\dots, a_n$ is sometimes called a **control tape**.\n",
    "The process of computing a sample trajectory and\n",
    "evaluating the corresponding utility is called a **rollout** for the control tape.\n",
    "Each rollout produces one sample trajectory, which gives one corresponding utility value.\n",
    "As an example, suppose the robot starts in the office executes the sequence\n",
    "$a_1, a_2, a_3, a_4$ = R,U,L,L (i.e., the robot executes actions *move right, move up, move left, move left*)\n",
    "in an attempt to reach the living room.\n",
    "\n",
    "Because we took care to specify the Markov chain above in reverse topological order, we can then use the GTSAM method `sample` to do the rollouts for us:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "markovChain = gtsam.DiscreteBayesNet()\n",
    "for k in reversed(range(1,N)):\n",
    "    markovChain.add(X[k+1], [X[k], A[k]], vacuum.action_spec)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/svg+xml": [
       "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
       "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
       " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
       "<!-- Generated by graphviz version 9.0.0 (0)\n",
       " -->\n",
       "<!-- Pages: 1 -->\n",
       "<svg width=\"350pt\" height=\"116pt\"\n",
       " viewBox=\"0.00 0.00 350.00 116.00\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 112)\">\n",
       "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-112 346,-112 346,4 -4,4\"/>\n",
       "<!-- var4683743612465315841 -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>var4683743612465315841</title>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"54,-108 0,-108 0,-72 54,-72 54,-108\"/>\n",
       "<text text-anchor=\"middle\" x=\"27\" y=\"-84.95\" font-family=\"Times,serif\" font-size=\"14.00\">A1</text>\n",
       "</g>\n",
       "<!-- var6341068275337658370 -->\n",
       "<g id=\"node6\" class=\"node\">\n",
       "<title>var6341068275337658370</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"99\" cy=\"-18\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"99\" y=\"-12.95\" font-family=\"Times,serif\" font-size=\"14.00\">X2</text>\n",
       "</g>\n",
       "<!-- var4683743612465315841&#45;&gt;var6341068275337658370 -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>var4683743612465315841&#45;&gt;var6341068275337658370</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M45.17,-71.83C54.54,-62.46 66.07,-50.93 76.05,-40.95\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"78.33,-43.62 82.92,-34.08 73.38,-38.67 78.33,-43.62\"/>\n",
       "</g>\n",
       "<!-- var4683743612465315842 -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>var4683743612465315842</title>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"126,-108 72,-108 72,-72 126,-72 126,-108\"/>\n",
       "<text text-anchor=\"middle\" x=\"99\" y=\"-84.95\" font-family=\"Times,serif\" font-size=\"14.00\">A2</text>\n",
       "</g>\n",
       "<!-- var6341068275337658371 -->\n",
       "<g id=\"node7\" class=\"node\">\n",
       "<title>var6341068275337658371</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"171\" cy=\"-18\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"171\" y=\"-12.95\" font-family=\"Times,serif\" font-size=\"14.00\">X3</text>\n",
       "</g>\n",
       "<!-- var4683743612465315842&#45;&gt;var6341068275337658371 -->\n",
       "<g id=\"edge4\" class=\"edge\">\n",
       "<title>var4683743612465315842&#45;&gt;var6341068275337658371</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M117.17,-71.83C126.54,-62.46 138.07,-50.93 148.05,-40.95\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"150.33,-43.62 154.92,-34.08 145.38,-38.67 150.33,-43.62\"/>\n",
       "</g>\n",
       "<!-- var4683743612465315843 -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>var4683743612465315843</title>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"198,-108 144,-108 144,-72 198,-72 198,-108\"/>\n",
       "<text text-anchor=\"middle\" x=\"171\" y=\"-84.95\" font-family=\"Times,serif\" font-size=\"14.00\">A3</text>\n",
       "</g>\n",
       "<!-- var6341068275337658372 -->\n",
       "<g id=\"node8\" class=\"node\">\n",
       "<title>var6341068275337658372</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"243\" cy=\"-18\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"243\" y=\"-12.95\" font-family=\"Times,serif\" font-size=\"14.00\">X4</text>\n",
       "</g>\n",
       "<!-- var4683743612465315843&#45;&gt;var6341068275337658372 -->\n",
       "<g id=\"edge6\" class=\"edge\">\n",
       "<title>var4683743612465315843&#45;&gt;var6341068275337658372</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M189.17,-71.83C198.54,-62.46 210.07,-50.93 220.05,-40.95\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"222.33,-43.62 226.92,-34.08 217.38,-38.67 222.33,-43.62\"/>\n",
       "</g>\n",
       "<!-- var4683743612465315844 -->\n",
       "<g id=\"node4\" class=\"node\">\n",
       "<title>var4683743612465315844</title>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"270,-108 216,-108 216,-72 270,-72 270,-108\"/>\n",
       "<text text-anchor=\"middle\" x=\"243\" y=\"-84.95\" font-family=\"Times,serif\" font-size=\"14.00\">A4</text>\n",
       "</g>\n",
       "<!-- var6341068275337658373 -->\n",
       "<g id=\"node9\" class=\"node\">\n",
       "<title>var6341068275337658373</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"315\" cy=\"-18\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"315\" y=\"-12.95\" font-family=\"Times,serif\" font-size=\"14.00\">X5</text>\n",
       "</g>\n",
       "<!-- var4683743612465315844&#45;&gt;var6341068275337658373 -->\n",
       "<g id=\"edge8\" class=\"edge\">\n",
       "<title>var4683743612465315844&#45;&gt;var6341068275337658373</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M261.17,-71.83C270.54,-62.46 282.07,-50.93 292.05,-40.95\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"294.33,-43.62 298.92,-34.08 289.38,-38.67 294.33,-43.62\"/>\n",
       "</g>\n",
       "<!-- var6341068275337658369 -->\n",
       "<g id=\"node5\" class=\"node\">\n",
       "<title>var6341068275337658369</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"27\" cy=\"-18\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"27\" y=\"-12.95\" font-family=\"Times,serif\" font-size=\"14.00\">X1</text>\n",
       "</g>\n",
       "<!-- var6341068275337658369&#45;&gt;var6341068275337658370 -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>var6341068275337658369&#45;&gt;var6341068275337658370</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M54.22,-18C56.28,-18 58.38,-18 60.5,-18\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"60.23,-21.5 70.23,-18 60.23,-14.5 60.23,-21.5\"/>\n",
       "</g>\n",
       "<!-- var6341068275337658370&#45;&gt;var6341068275337658371 -->\n",
       "<g id=\"edge3\" class=\"edge\">\n",
       "<title>var6341068275337658370&#45;&gt;var6341068275337658371</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M126.22,-18C128.28,-18 130.38,-18 132.5,-18\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"132.23,-21.5 142.23,-18 132.23,-14.5 132.23,-21.5\"/>\n",
       "</g>\n",
       "<!-- var6341068275337658371&#45;&gt;var6341068275337658372 -->\n",
       "<g id=\"edge5\" class=\"edge\">\n",
       "<title>var6341068275337658371&#45;&gt;var6341068275337658372</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M198.22,-18C200.28,-18 202.38,-18 204.5,-18\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"204.23,-21.5 214.23,-18 204.23,-14.5 204.23,-21.5\"/>\n",
       "</g>\n",
       "<!-- var6341068275337658372&#45;&gt;var6341068275337658373 -->\n",
       "<g id=\"edge7\" class=\"edge\">\n",
       "<title>var6341068275337658372&#45;&gt;var6341068275337658373</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M270.22,-18C272.28,-18 274.38,-18 276.5,-18\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"276.23,-21.5 286.23,-18 276.23,-14.5 276.23,-21.5\"/>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "<gtbook.display.show at 0x115892fd0>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#| caption: A Markov chain with 5 states and 4 actions. The actions are given.\n",
    "#| label: fig:markov_chain_5\n",
    "show(markovChain, hints={\"A\":2, \"X\":1, \"Z\":0}, boxes={A[k][0] for k in range(1,N)})\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def perform_rollout(x1, actions):\n",
    "    \"\"\"Roll out states given actions as a dictionary\"\"\"\n",
    "    dict = actions.copy()\n",
    "    dict[X[1]] = x1\n",
    "    given = VARIABLES.assignment(dict)\n",
    "    return markovChain.sample(given)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To execute a specific rollout for the control tape R,U,L,L, we can use the following code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<table class='DiscreteValues'>\n",
       "  <thead>\n",
       "    <tr><th>Variable</th><th>value</th></tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr><th>A1</th><td>R</td></tr>\n",
       "    <tr><th>A2</th><td>U</td></tr>\n",
       "    <tr><th>A3</th><td>L</td></tr>\n",
       "    <tr><th>A4</th><td>L</td></tr>\n",
       "    <tr><th>X1</th><td>Office</td></tr>\n",
       "    <tr><th>X2</th><td>Office</td></tr>\n",
       "    <tr><th>X3</th><td>Office</td></tr>\n",
       "    <tr><th>X4</th><td>Office</td></tr>\n",
       "    <tr><th>X5</th><td>Office</td></tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "<gtbook.display.pretty at 0x115893760>"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "actions = {A[1]:\"R\", A[2]:\"U\", A[3]:\"L\", A[4]:\"L\"}\n",
    "rollout = perform_rollout(\"Office\", actions)\n",
    "pretty(rollout)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It is important to remember that any individual rollout corresponds to a sample trajectory from a stochastic process.\n",
    "If you execute the above code several times, you should observe that the robot does not always arrive to the living room.\n",
    "\n",
    "The code below executes the rollout and computes the corresponding utility for the sample trajectory.\n",
    "Sample trajectories that do not arrive to the living room will have zero utility."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.0\n"
     ]
    }
   ],
   "source": [
    "def reward(rollout, k):\n",
    "    \"\"\"Return state, action, next_state triple for given rollout at time k.\"\"\"\n",
    "    state = rollout[X[k][0]]\n",
    "    action = rollout[A[k][0]]\n",
    "    next_state = rollout[X[k+1][0]]\n",
    "    return R[state, action, next_state]\n",
    "\n",
    "def rollout_reward(rollout, horizon=N, gamma=1.0):\n",
    "    \"\"\"Calculate reward for a given rollout\"\"\"\n",
    "    discounted_rewards = [gamma**(k-1) * reward(rollout,k) for k in range(1,horizon)]\n",
    "    return sum(discounted_rewards)\n",
    "\n",
    "print(rollout_reward(rollout))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following code executes 20 rollouts for the action sequence R,U,L,L, and prints the utility for each sample trajectory.\n",
    "You can see that in many cases the robot fails to arrive to the living room (thus earning zero utility)!\n",
    "This is because each of the first two actions have a 0.2 probability of failure, and if either of these fail, the robot is unable to reach the living room using this control tape."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.0, 30.0, 30.0, 30.0, 30.0, 30.0, 30.0, 30.0, 0.0, 30.0, 30.0, 0.0, 30.0, 30.0, 30.0, 30.0, 30.0, 30.0, 30.0, 0.0]\n"
     ]
    }
   ],
   "source": [
    "def control_tape_reward(x1, actions):\n",
    "    \"\"\"Calculate reward given a dictionary of actions\"\"\"\n",
    "    rollout = perform_rollout(x1, actions)\n",
    "    return rollout_reward(rollout)\n",
    "\n",
    "print([control_tape_reward(\"Office\", actions) for i in range(20)])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, the expected utility of the action sequence R,U,L,L is approximated\n",
    "by simply averaging over all 20 rollouts:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "19.5"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.mean([control_tape_reward(\"Office\", actions) for i in range(20)])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercises:\n",
    "- Execute both code segments above multiple times and observe the effect.\n",
    "- The rewards above seems to be always either zero or 30. Why is that?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Policies\n",
    "\n",
    "> A policy is a function that specifies which action to take in each state.\n",
    "\n",
    "In the example above, the control tape rollout for the action sequence R,U,L,L failed many times.\n",
    "The reason for this failure may seem obvious: the robot executed the same sequence of actions,\n",
    "regardless of the state trajectory.\n",
    "If the robot had been able to choose its actions based on the current state,\n",
    "it would have chosen to move right until it reached the hallway,\n",
    "at which time it would have chosen to move up until reaching the living room.\n",
    "Clearly, the robot could make better choices if it were allowed to dynamically\n",
    "choose which action to execute based on its current state.\n",
    "\n",
    "A **policy**, $\\pi: {\\cal X} \\rightarrow {\\cal A}$ is a mapping from states to actions.\n",
    "Specifying a policy instead of a control tape has the potential to significantly improve the robot's performance, by adapting the action sequence based on the actual state trajectory that occurs during execution."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The code below defines a fairly intuitive policy. If in the office, move right. If in the dining room or kitchen, move left.\n",
    "If in the hallway or living room, move up.\n",
    "Note we implement policies as a simple list in python, so `pi[0]` is the the action taken in state with index 0 (the Living Room). We return an index into the `action_space` to make the code following it efficient: \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "RIGHT = vacuum.action_space.index(\"R\")\n",
    "LEFT  = vacuum.action_space.index(\"L\")\n",
    "UP    = vacuum.action_space.index(\"U\")\n",
    "DOWN  = vacuum.action_space.index(\"D\")\n",
    "\n",
    "reasonable_policy = [UP, LEFT, RIGHT, UP, LEFT]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once we have a given policy, $\\pi$, we can compute a policy rollout in a manner analogous to computing control tape rollouts described above.\n",
    "In particular, at each state, instead of sampling from the distribution\n",
    "$P(X_{k+1} | a_k, x_k)$ we sample from the distribution\n",
    "$P(X_{k+1} | \\pi(x_k), x_k)$. In other words, instead of simulating a pre-specified action $a_k$, we choose $a_k = \\pi(x_k)$.\n",
    "\n",
    "Here is a function that computes a rollout given a policy, rather than a control tape:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<table class='DiscreteValues'>\n",
       "  <thead>\n",
       "    <tr><th>Variable</th><th>value</th></tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr><th>A1</th><td>R</td></tr>\n",
       "    <tr><th>A2</th><td>U</td></tr>\n",
       "    <tr><th>A3</th><td>U</td></tr>\n",
       "    <tr><th>A4</th><td>U</td></tr>\n",
       "    <tr><th>X1</th><td>Office</td></tr>\n",
       "    <tr><th>X2</th><td>Hallway</td></tr>\n",
       "    <tr><th>X3</th><td>Living Room</td></tr>\n",
       "    <tr><th>X4</th><td>Living Room</td></tr>\n",
       "    <tr><th>X5</th><td>Living Room</td></tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "<gtbook.display.pretty at 0x127f8e9a0>"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def policy_rollout(x1, pi, horizon=N):\n",
    "    \"\"\"Roll out states given a policy pi, for given horizon.\"\"\"\n",
    "    rollout = gtsam.DiscreteValues()\n",
    "    x = x1\n",
    "    for k in range(1, horizon):\n",
    "        a = pi[x]\n",
    "        rollout[X[k][0]] = x\n",
    "        rollout[A[k][0]] = a\n",
    "        next_state_distribution = gtsam.DiscreteDistribution(X[k+1], T[x, a])\n",
    "        x = next_state_distribution.sample()\n",
    "    rollout[X[horizon][0]] = x\n",
    "    return rollout\n",
    "\n",
    "pretty(policy_rollout(vacuum.rooms.index(\"Office\"), reasonable_policy, horizon=5))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our policy-based rollouts seem to be much better!\n",
    "\n",
    "In summary, we proposed two methods to plan an action sequence, a greedy approach that maximizes the expected reward for executing a single action,\n",
    "and an optimization-based method that chooses a fixed action sequence to maximize the\n",
    "expected utility of the action sequence.\n",
    "We described how to implement the computations required for the latter approach using control tape rollouts.\n",
    "\n",
    "With the introduction of policies, planning is reduced to the search for an appropriate\n",
    "policy.\n",
    "In the best case, this policy would maximize the expected utility.\n",
    "Because policies are state-dependent, the combinatorics of enumerating all possible\n",
    "policies precludes any approach that attempts to explicitly enumerate candidate policies.\n",
    "Instead, we will focus on methods that explore the utility associated to\n",
    "a policy.\n",
    "As we will now see, this is somewhat more complicated than exploring the utility\n",
    "for a fixed sequence of actions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Value Function (for a given policy)\n",
    "> The value function $V^\\pi$ measures the expected utility from each state, under a given policy $\\pi$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We would now like to characterize the quality of any given policy.\n",
    "Above, we defined the utility for a specific sequence of $n$ actions as\n",
    "\n",
    "$$\n",
    "U(a_1, \\dots, a_n, x_1, \\dots x_{n+1}) =\n",
    "R(x_1,a_1, x_2) + \\gamma R(x_2, a_2, x_3) + \\dots \\gamma^{n-1} R(x_{n}, a_{n}, x_{n+1})\n",
    "$$\n",
    "\n",
    "and used the expected utility, $E[U(a_1, \\dots, a_n, X_1, \\dots X_{n+1})]$ as a quantitative measure\n",
    "of the efficacy of the specific action sequence.\n",
    "We can apply this same type of reasoning for a policy, $\\pi$.\n",
    "When evaluating policies, it is typical to use a discounted\n",
    "reward over an infinite time horizon.\n",
    "In this case, we define the **value function** $V^\\pi:{\\cal X} \\rightarrow \\mathbb{R}$ for policy $\\pi$ as\n",
    "\n",
    "$$\n",
    "V^\\pi(x_1) \\doteq E [R(x_1, \\pi(x_1), X_2) + \\gamma R(X_2, \\pi(X_2), X_3) + \\gamma^2 R(X_3, \\pi(X_3), X_4) + \\dots]\n",
    "$$\n",
    "\n",
    "where the expectation is taken over the possible values\n",
    "of the states $X_2, X_3 \\dots$.\n",
    "Note that the policy $\\pi$ is deterministic, but that $\\pi(X_i)$ is the random action that results\n",
    "from applying the deterministic policy to the stochastic state $X_i$.\n",
    "\n",
    "<!-- In the definition for $V^\\pi$, the argument is given as $x_1$, the initial state.\n",
    "We can easily generalize this to arbitrary time steps as\n",
    "\n",
    "$$\n",
    "V^\\pi(x_k) \\doteq E [R(x_k, \\pi(x_k), X_{k+1}) + \\gamma R(X_{k+1}, \\pi(X_{k+1}), X_{k+2}) + \\dots]\n",
    "$$ -->"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The value function for a policy can be written in a nice recursive form\n",
    "that can be obtained by the following derivation.\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "V^\\pi(x_1) &=\n",
    "E [R(x_1, \\pi(x_1), X_2) + \\gamma R(X_2, \\pi(X_2), X_3) + \\gamma^2 R(X_3, \\pi(X_3), X_4) + \\dots] \\\\\n",
    "&= \\sum_{x_2} P(x_2|x_1, \\pi(x_1)) R(x_1, \\pi(x_1), x_2) \\\\\n",
    "& + \\sum_{x_2} P(x_2|x_1, \\pi(x_1)) E [\\gamma R(x_2, \\pi(x_2), X_3) + \\gamma^2 R(X_3, \\pi(X_3), X_4) + \\dots]\\\\\n",
    "&= \\sum_{x_2} P(x_2|x_1, \\pi(x_1)) \\{R(x_1, \\pi(x_1), x_2) + \\gamma V^\\pi(x_2)\\}\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "where the second line is obtained by explicitly writing the summation for the expectation\n",
    "taken with respect to the random state $X_2$,\n",
    "and the third line is obtained by noticing that\n",
    "$E [\\gamma R(x_2, \\pi(x_2), X_3) + \\gamma^2 R(X_3, \\pi(X_3), X_4) + \\dots] = \\gamma V^\\pi(x_2)$.\n",
    "\n",
    "We can apply the distributivity property to this expression to obtain\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "V^\\pi(x_1) &= \\sum_{x_2} P(x_2|x_1, \\pi(x_1)) R(x_1, \\pi(x_1), x_2) + \\gamma  \\sum_{x_2} P(x_2|x_1, \\pi(x_1))  V^\\pi(x_2) \\\\\n",
    "&= \\bar{R}(x_1,\\pi(x_1)) + \\gamma \\sum_{x_2} P(x_2|x_1, \\pi(x_1)) V^\\pi(x_2)\n",
    "\\end{aligned}\n",
    "$$\n",
    " \n",
    "in which the term $\\bar{R}(x_1,\\pi(x_1))$ is the expected reward for applying action $a = \\pi(x_1)$ in state $x_1$, which we already know how to compute from the reward function $R$ and transition probabilities $T$.\n",
    "By now substituting $x$ for $x_1$, and $x'$ for $x_{2}$ we can generalize this expression\n",
    "to apply to the state $x$ at any arbitrary time:\n",
    "\n",
    "$$V^\\pi(x) = \\bar{R}(x,\\pi(x)) + \\gamma \\sum_{x'} P(x'|x, \\pi(x)) V^\\pi(x')$$\n",
    "\n",
    "This has a very nice interpretation: the value of a state under a given policy is the expected reward $\\bar{R}(x,\\pi(x))$ under that policy, *plus* the discounted expected value of the next state."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Thought exercise\n",
    "\n",
    " Without a discount factor, we would always be hopeful that if we take one more step, we will find a pot of gold. Reflect on what various values for the discount factor mean in real life."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Approximating the Value Function\n",
    "\n",
    "> We can use rollouts to approximate the value function as well.\n",
    "\n",
    "Just as we approximated the expected utility of an action sequence\n",
    "using control tape rollouts,\n",
    "we can approximate the value function by sampling over a number of policy rollouts.\n",
    "This process provides a *Monte Carlo approximation* of the value function for a given policy. \n",
    "Of course we cannot apply a policy rollout over an infinite time horizon, so we apply\n",
    "the rollout only to some finite number of time steps, say $N_{\\rm{ro}}$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note this is an approximation in *two* ways: \n",
    "1. We approximate the expectation by averaging over $N_{\\rm{Samples}}$ sample trajectories.\n",
    "2. We only roll out for  $N_{\\rm{ro}}$ steps.\n",
    "We can improve on this by increasing the number of samples or by increasing the horizon, but at a cost linear in their product, i.e., \n",
    "$O(N_{\\rm{ro}} N_{\\rm{Samples}})$:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You might worry that evaluating sample paths of length $N_{\\rm{ro}}$ could lead to significant errors\n",
    "in our approximation.\n",
    "In fact, it is easy to determine an upper bound on this error.\n",
    "Since $R(x, a, x') $ is finite,\n",
    "we know that there is some upper bound $R_{\\rm{max}}$ such that $R(x, a, x') < R_{\\rm{max}}$, for all\n",
    "possible $x,a,x'$.\n",
    "We can use this fact to find a bound on $V^\\pi$:\n",
    "\n",
    "$$V^\\pi(x_k)\n",
    "\\leq \\max_{x_k, x_{k+1}, \\dots} \\sum_{i=0}^\\infty \\gamma ^i R(x_{k+i},\\pi(x_{k+i}), x_{k+i+1})\n",
    "\\leq \\sum_{i=0}^\\infty \\gamma ^i R_{\\rm{max}} = \\frac{R_{\\rm{max}}}{1 - \\gamma}$$\n",
    "\n",
    "in which the final term applies for $0 < \\gamma < 1$.\n",
    "This expression can then be used to bound the error \n",
    "\n",
    "$$\n",
    "V^\\pi(x) - \\sum_{i=0}^{N_{\\rm{ro}}}\\gamma ^i R(x_i,\\pi(xi), x_{i+1})$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Because we have functions to sample a policy rollout *and* to calculate the value of a rollout, the code is simple enough:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "V(Living Room) ~ 34.39\n",
      "V(Kitchen) ~ 29.49\n",
      "V(Office) ~ 16.00\n",
      "V(Hallway) ~ 32.39\n",
      "V(Dining Room) ~ 19.25\n"
     ]
    }
   ],
   "source": [
    "def approximate_value_function(x1, pi, nr_samples=10, horizon=N, gamma=0.9):\n",
    "    \"\"\" Approximate the value function by performing `nr_samples` rollouts\n",
    "        starting from x1, and averaging the result.\n",
    "    \"\"\"\n",
    "    rollouts = [policy_rollout(x1, pi, horizon) for _ in range(nr_samples)]\n",
    "    rewards = [rollout_reward(rollout, horizon, gamma) for rollout in rollouts]\n",
    "    return sum(rewards)/nr_samples\n",
    "\n",
    "nr_samples=10\n",
    "for x1, room in enumerate(vacuum.rooms):\n",
    "    V_x1 = approximate_value_function(x1, reasonable_policy, nr_samples)\n",
    "    print(f\"V({room}) ~ {V_x1:.2f}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above calculation was done with a horizon of $N=5$ and 10 samples. Of course, we can use more samples and a longer horizon to obtain a much more accurate estimate:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "V(Living Room) ~ 99.43\n",
      "V(Kitchen) ~ 97.64\n",
      "V(Office) ~ 84.74\n",
      "V(Hallway) ~ 97.37\n",
      "V(Dining Room) ~ 85.89\n"
     ]
    }
   ],
   "source": [
    "nr_samples = 100\n",
    "horizon = 50\n",
    "X = VARIABLES.discrete_series('X', range(1, horizon+1), vacuum.rooms)\n",
    "A = VARIABLES.discrete_series('A', range(1, horizon), vacuum.action_space)\n",
    "for x1, room in enumerate(vacuum.rooms):\n",
    "    V_x1 = approximate_value_function(x1, reasonable_policy, nr_samples, horizon)\n",
    "    print(f\"V({room}) ~ {V_x1:.2f}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Computing The Value Function*\n",
    "\n",
    "> For any *fixed* policy we can compute the exact value of value function.\n",
    "\n",
    "We can compute the exact value of $V^\\pi$ by solving a system of linear equations.\n",
    "Recall our recursive definition of the value function:\n",
    "\n",
    "$$V^\\pi(x) = \\bar{R}(x,\\pi(x)) + \\gamma \\sum_{x'} P(x'|x, \\pi(x)) V^\\pi(x')$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This equation holds for every possible value $x$ for the state.\n",
    "Hence, if there are $n$ possible states, we obtain $n$ linear equations in $n$ unknowns.\n",
    "Each of these $n$ equations is obtained by evaluating $V^\\pi(x)$ for a specific value of $x$.\n",
    "Collecting the unknown $V^\\pi$ terms on the left hand side and the known $\\bar{R}(x,\\pi(x))$ \n",
    "terms on the right hand side, we obtain\n",
    "\n",
    "$$V^\\pi(x) - \\gamma \\sum_{x'} P(x'|x, \\pi(x)) V^\\pi(x') = \\bar{R}(x,\\pi(x)).$$\n",
    "\n",
    " \n",
    "To make this explicit yet concise for our vacuum cleaning robot example,\n",
    "let us define the *scalar* $T^\\pi_{xy}\\doteq P(y|x,\\pi(x))$ as the transition probability from state $x$ to state $y$ under policy $\\pi$. \n",
    "In addition, we use the abbreviations L,K,O,H, and D for the rooms, and use the shorthand $V^\\pi_x\\doteq V^\\pi(x)$ for the value of state $x$ under policy $\\pi$. \n",
    "Using this notation, we can evaluate the above expression.\n",
    "For $x = L$, we obtain\n",
    "\n",
    "$$ \n",
    "V^\\pi_L - \\gamma \\sum_{x'\\in {L,K,O,H,D}} T^\\pi_{Lx'} V^\\pi_{x'} = \\bar{R}(L,\\pi(L))\n",
    "$$\n",
    "\n",
    "or, after some algebra:\n",
    "\n",
    "$$ \n",
    "(1 - \\gamma T^\\pi_{LL}) V^\\pi_{L}\n",
    "- \\gamma T^\\pi_{LK} V^\\pi_{K} \n",
    "- \\gamma T^\\pi_{LO} V^\\pi_{O} \n",
    "- \\gamma T^\\pi_{LH} V^\\pi_{H} \n",
    "- \\gamma T^\\pi_{LD} V^\\pi_{D} \n",
    "= \\bar{R}(L,\\pi(L))\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we apply this same process for each of the five rooms,\n",
    "we obtain the following five equations:\n",
    "\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "(1 - \\gamma T^\\pi_{LL}) V^\\pi_{L}\n",
    "- \\gamma T^\\pi_{LK} V^\\pi_{K} \n",
    "- \\gamma T^\\pi_{LO} V^\\pi_{O} \n",
    "- \\gamma T^\\pi_{LH} V^\\pi_{H} \n",
    "- \\gamma T^\\pi_{LD} V^\\pi_{D} \n",
    "&= \\bar{R}(L,\\pi(L))\n",
    "\\\\\n",
    "- \\gamma T^\\pi_{KL} V^\\pi_{L}\n",
    "+ (1 - \\gamma T^\\pi_{KK}) V^\\pi_{K} \n",
    "- \\gamma T^\\pi_{KO} V^\\pi_{O} \n",
    "- \\gamma T^\\pi_{KH} V^\\pi_{H} \n",
    "- \\gamma T^\\pi_{KD} V^\\pi_{D} \n",
    "&= \\bar{R}(K,\\pi(K))\n",
    "\\\\\n",
    "- \\gamma T^\\pi_{OL} V^\\pi_{L}\n",
    "- \\gamma T^\\pi_{OK} V^\\pi_{K} \n",
    "+ (1 - \\gamma T^\\pi_{OO}) V^\\pi_{O} \n",
    "- \\gamma T^\\pi_{OH} V^\\pi_{H} \n",
    "- \\gamma T^\\pi_{OD} V^\\pi_{D} \n",
    "&= \\bar{R}(O,\\pi(O))\n",
    "\\\\\n",
    "- \\gamma T^\\pi_{HL} V^\\pi_{L}\n",
    "- \\gamma T^\\pi_{HK} V^\\pi_{K} \n",
    "- \\gamma T^\\pi_{HO} V^\\pi_{O} \n",
    "+ (1 - \\gamma T^\\pi_{HH}) V^\\pi_{H} \n",
    "- \\gamma T^\\pi_{HD} V^\\pi_{D} \n",
    "&= \\bar{R}(H,\\pi(H))\n",
    "\\\\\n",
    "- \\gamma T^\\pi_{DL} V^\\pi_{L}\n",
    "- \\gamma T^\\pi_{DK} V^\\pi_{K} \n",
    "- \\gamma T^\\pi_{DO} V^\\pi_{O} \n",
    "- \\gamma T^\\pi_{DH} V^\\pi_{H} \n",
    "+ (1 - \\gamma T^\\pi_{DD}) V^\\pi_{D} \n",
    "&= \\bar{R}(D,\\pi(D))\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "The unknowns in these equations are \n",
    "$V^\\pi_{L}, V^\\pi_{K}, V^\\pi_{O}, V^\\pi_{H}, V^\\pi_{D}$. All of the other terms are\n",
    "either transition probabilities or expected rewards, whose values are either given,\n",
    "or can easily be computed.\n",
    "\n",
    "In code, this becomes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_value_system(pi, gamma=0.9):\n",
    "    \"\"\"Calculate A, b matrix of linear system for value computation.\"\"\"\n",
    "    b = np.empty((5,), float)\n",
    "    AA = np.empty((5,5), float)\n",
    "    for x, room in enumerate(vacuum.rooms):\n",
    "        a = pi[x] # action under policy\n",
    "        b[x] = T[x,a] @ R[x,a] # expected reward under policy pi\n",
    "        AA[x] = -gamma * T[x,a]\n",
    "        AA[x,x] += 1\n",
    "    return AA,b\n",
    "\n",
    "def calculate_value_function(pi, gamma=0.9):\n",
    "    \"\"\"Calculate value function for given policy\"\"\"\n",
    "    AA, b = calculate_value_system(pi, gamma)\n",
    "    return np.linalg.solve(AA,b)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is the linear system for the policy `reasonable_policy` (which is really almost always wrong):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "A =\n",
      "[[ 0.1  -0.   -0.   -0.   -0.  ]\n",
      " [-0.72  0.82 -0.   -0.   -0.  ]\n",
      " [-0.   -0.    0.82 -0.72 -0.  ]\n",
      " [-0.72 -0.   -0.    0.82 -0.  ]\n",
      " [-0.   -0.   -0.   -0.72  0.82]]\n",
      "b = [10.  8.  0.  8.  0.]\n"
     ]
    }
   ],
   "source": [
    "AA, b = calculate_value_system(reasonable_policy)\n",
    "print(f\"A =\\n{AA}\")\n",
    "print(f\"b = {b}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When we calculate the value function $V^\\pi$ under the policy `reasonable_policy` we see that our *exact* value function was well approximated by the Monte Carlo estimate above:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "V(reasonable_policy):\n",
      "  Living Room : 100.00\n",
      "  Kitchen     : 97.56\n",
      "  Office      : 85.66\n",
      "  Hallway     : 97.56\n",
      "  Dining Room : 85.66\n"
     ]
    }
   ],
   "source": [
    "value_for_pi = calculate_value_function(reasonable_policy)\n",
    "print(\"V(reasonable_policy):\")\n",
    "for i,room in enumerate(vacuum.rooms):\n",
    "    print(f\"  {room:12}: {value_for_pi[i]:.2f}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise\n",
    "\n",
    "Why is the value function in the living room $100$ and not the immediate reward $10$?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "Markov Decision Processes or MDPs can be used to model decision making in a stochastic environment, albeit with complete knowledge of the state. It is a rich subject, and we introduced many new concepts in this section:\n",
    "\n",
    "- The reward function, $R : {\\cal X} \\times {\\cal A} \\times {\\cal X} \\rightarrow \\mathbb{R}$.\n",
    "- The expected reward, $\\bar{R}(x,a)$ for executing action $a$ in state $x$, and a corresponding greedy planning algorithm.\n",
    "- The utility function, $U: {\\cal A}^n \\times {\\cal X}^{n+1} \\rightarrow \\mathbb{R}$, as the sum of discounted rewards.\n",
    "- The notion of rollouts to approximate the expected utility of actions.\n",
    "- The policy $\\pi: {\\cal X} \\rightarrow {\\cal A}$, as a mapping from states to actions.\n",
    "- The value function, $V^\\pi:{\\cal X} \\rightarrow \\mathbb{R}$, associated with a given policy $\\pi$.\n",
    "- The use of policy rollouts to approximate the value function $V^\\pi$.\n",
    "- Exact calculation of the value function for a fixed policy.\n",
    "\n",
    "We'll leave the concepts of an optimal policy and how to compute it for the next section. There are two other extensions to MDPs that we did not cover:\n",
    "\n",
    "- Partially Observable MDPs (or POMDPS) are appropriate when we cannot directly observe the state.\n",
    "- Reinforcement learning, a way to learn MDP policies from *data*. This will be covered in the next section as well."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "include_colab_link": true,
   "name": "S35_vacuum_decision.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3.8.12 ('nbdev')",
   "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.9.19"
  },
  "latex_metadata": {
   "affiliation": "Georgia Institute of Technology",
   "author": "Frank Dellaert and Seth Hutchinson",
   "title": "Introduction to Robotics"
  },
  "vscode": {
   "interpreter": {
    "hash": "341996cd3f3db7b5e0d1eaea072c5502d80452314e72e6b77c40445f6e9ba101"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}