S34_vacuum_perception.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2mkf6TY5k25H",
   "metadata": {},
   "source": [
    "# Perception with Graphical Models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "pRjyTEbKskua",
   "metadata": {
    "colab_type": "text",
    "tags": [
     "no-tex"
    ]
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/gtbook/robotics/blob/main/S34_vacuum_perception.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "E5rsQom9hatQ",
   "metadata": {
    "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 -U -q gtbook\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "OFkNEfdX_CLe",
   "metadata": {
    "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",
    "\n",
    "def pretty(obj):\n",
    "    return gtbook.display.pretty(obj, VARIABLES)\n",
    "\n",
    "def show(obj, **kwargs):\n",
    "    return gtbook.display.show(obj, VARIABLES, **kwargs)\n",
    "\n",
    "try:\n",
    "    import google.colab\n",
    "except:\n",
    "    import plotly.io as pio\n",
    "    pio.renderers.default = \"png\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "gIariOyXd8PE",
   "metadata": {
    "tags": [
     "no-tex",
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "# From section 3.2:\n",
    "wxyz = gtsam.DiscreteBayesNet()\n",
    "W1 = VARIABLES.binary(\"W\")\n",
    "X1 = VARIABLES.binary(\"X\")\n",
    "Y1 = VARIABLES.binary(\"Y\")\n",
    "Z1 = VARIABLES.binary(\"Z\")\n",
    "wxyz.add(W1, [X1, Z1], \"1/1 1/1 1/1 1/1\")\n",
    "wxyz.add(X1, [Y1, Z1], \"1/1 1/1 1/1 1/1\")\n",
    "wxyz.add(Y1, [Z1], \"1/1 1/1\")\n",
    "wxyz.add(Z1, \"1/1\")\n",
    "\n",
    "# From Section 3.3:\n",
    "N = 3\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",
    "Z = VARIABLES.discrete_series(\"Z\", range(1, N+1), vacuum.light_levels)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "p8jYt599NMVO",
   "metadata": {},
   "source": [
    "```{index} perception; hidden Markov models\n",
    "```\n",
    "\n",
    "> Perception for dynamic Bayes nets is equivalent to inference in hidden Markov models (HMMs)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0ff9a5d",
   "metadata": {
    "tags": [
     "no-pdf"
    ]
   },
   "source": [
    "<img src=\"Figures3/S34-iRobot_vacuuming_robot-05.jpg\" alt=\"Splash image with deeply contemplative robot\" width=\"40%\" align=center style=\"vertical-align:middle;margin:10px 0px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "hMpmGJOx4Bez",
   "metadata": {},
   "source": [
    "```{index} factor graphs\n",
    "```\n",
    "Now that we have a solid mathematical framework to work in, in the remainder of this chapter we will tackle perception, decision-making based on this, and learning to improve both. We used the vacuum cleaning robot as an example system that has discrete states, with actions that mediate discrete-time state transitions, and emitting noisy measurements at every step. The key mathematical concept that tied all this together was the language of dynamic Bayes nets. \n",
    "\n",
    "In this section we tackle *perception* in this setting: can we, by remembering which actions we took and what measurements we observed, infer what states we were in and are likely to be in now? \n",
    "We more formally define what we mean by inference, building on the methods introduced in Section 2.4.\n",
    "Just as we did there, we will characterize what we know and don't know using probability distributions, leaning into the Bayesian view even more.\n",
    "However, here we will consider probability distributions over *sequences* of states, rather than just the current state.\n",
    "\n",
    "To this end, we introduce one of the most important tools used throughout this book: **factor graphs**. Bayes nets are great for *modeling*, but for inferring the state of the robot over time *efficiently* we need a better data structure. \n",
    "We first introduce hidden Markov models (HMMs), and highlight their connection with robot localization over time. We then show how to efficiently perform inference by converting an HMM into a factor graph, and performing full posterior inference and MAP estimation in linear time."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "pgxTTdyGLhbZ",
   "metadata": {},
   "source": [
    "```{index} inference\n",
    "```\n",
    "## Inference in Bayes Nets\n",
    "\n",
    "> Inference begets the full posterior or a single maximum a posteriori estimate.\n",
    "\n",
    "**Inference** is the process of determining knowledge about a set of\n",
    "variables $\\mathcal{X}$ given the known values for another set of variables $\\mathcal{Z}$.\n",
    "In this section we will describe inference when the joint\n",
    "distribution is specified using a Bayes net, but we will not take\n",
    "advantage of the specific sparse structure of the network.\n",
    "Hence, the algorithms\n",
    "below are completely general, for any (discrete) joint probability\n",
    "distribution, as long as we can evaluate the joint distribution."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "jVhHGGJafQ9g",
   "metadata": {},
   "source": [
    "```{index} hidden variables, observations\n",
    "```\n",
    "### Full Posterior Inference\n",
    "\n",
    "> Naive full posterior inference is *easy* but probably expensive.\n",
    "\n",
    "The simplest kind of inference occurs when we can *partition* the variables into two\n",
    "sets: the **hidden variables** $\\mathcal{X}$ and the **observations**\n",
    "$\\mathcal{Z}$. First, let us recall what a Bayes net looks like in this case.\n",
    "From Section 3.2, we have that a general Bayes net over any set of variables is defined as a product of conditionals, hence in our case we would have\n",
    "\\begin{equation}\n",
    "P(\\mathcal{X}, \\mathcal{Z}) = \\prod_{i=1}P(X_{i}|\\Pi_{X_i}) \\prod_{j=1}P(Z_{j}|\\Pi_{Z_j})\n",
    "\\end{equation}\n",
    "where there is no restriction on the parent sets $\\Pi$, just that the directed graph is acyclic. Of course, we will be most interested here in dynamic Bayes nets as defined previously, but for now what follows applies to generic Bayes nets.\n",
    "\n",
    "The workhorse in inference is Bayes’ theorem, which we encountered in the perception section of the previous chapter, Section 2.4. \n",
    "Let us define $\\mathfrak{z}$ as the set of observed *values* for all variables in\n",
    "$\\mathcal{Z}$.\n",
    "We then are interested, *given* the values $\\mathfrak{z}$ for the observed values $\\mathcal{Z}$, in calculating the posterior probability distribution $P(\\mathcal{X}|\\mathcal{Z}=\\mathfrak{z})$.\n",
    "To calculate this we apply Bayes’ theorem, but now applied to\n",
    "*sets* of variables, to obtain an expression for the posterior over the\n",
    "hidden variables $\\mathcal{X}$. Using the \"easy\" version of Bayes’ theorem from Section 2.4, \n",
    "we just re-write the posterior as the joint, but with the values $\\mathfrak{z}$ for all variables in\n",
    "$\\mathcal{Z}$ instantiated:\n",
    "\\begin{equation}\n",
    "P(\\mathcal{X}|\\mathcal{Z}=\\mathfrak{z})\\propto P(\\mathcal{X}, \\mathcal{Z}=\\mathfrak{z}).\n",
    "\\end{equation}\n",
    "This seems almost too simple, but it really works, as shown in a small generic Bayes net example below. Let us recall the 4-variable Bayes net example on variables $X$, $Y$, $W$, and $Z$, and let us take $\\mathcal{X}=(X, Y)$ and $\\mathcal{Z}=(W, Z)$ - which we indicate in the figure by rendering $W$ and $Z$ as boxes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "MjZfcKXhjKDK",
   "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=\"128pt\" height=\"260pt\"\n",
       " viewBox=\"0.00 0.00 128.00 260.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 256)\">\n",
       "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-256 124,-256 124,4 -4,4\"/>\n",
       "<!-- var0 -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>var0</title>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"54,-36 0,-36 0,0 54,0 54,-36\"/>\n",
       "<text text-anchor=\"middle\" x=\"27\" y=\"-12.95\" font-family=\"Times,serif\" font-size=\"14.00\">W</text>\n",
       "</g>\n",
       "<!-- var1 -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>var1</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"55\" cy=\"-90\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"55\" y=\"-84.95\" font-family=\"Times,serif\" font-size=\"14.00\">X</text>\n",
       "</g>\n",
       "<!-- var1&#45;&gt;var0 -->\n",
       "<g id=\"edge4\" class=\"edge\">\n",
       "<title>var1&#45;&gt;var0</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M48.36,-72.41C45.32,-64.79 41.62,-55.55 38.16,-46.9\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"41.5,-45.84 34.54,-37.85 35.01,-48.44 41.5,-45.84\"/>\n",
       "</g>\n",
       "<!-- var2 -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>var2</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"93\" cy=\"-162\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"93\" y=\"-156.95\" font-family=\"Times,serif\" font-size=\"14.00\">Y</text>\n",
       "</g>\n",
       "<!-- var2&#45;&gt;var1 -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>var2&#45;&gt;var1</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M84.19,-144.76C79.73,-136.55 74.2,-126.37 69.16,-117.09\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"72.33,-115.6 64.49,-108.48 66.18,-118.94 72.33,-115.6\"/>\n",
       "</g>\n",
       "<!-- var3 -->\n",
       "<g id=\"node4\" class=\"node\">\n",
       "<title>var3</title>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"65,-252 11,-252 11,-216 65,-216 65,-252\"/>\n",
       "<text text-anchor=\"middle\" x=\"38\" y=\"-228.95\" font-family=\"Times,serif\" font-size=\"14.00\">Z</text>\n",
       "</g>\n",
       "<!-- var3&#45;&gt;var0 -->\n",
       "<g id=\"edge5\" class=\"edge\">\n",
       "<title>var3&#45;&gt;var0</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M33.74,-215.65C27.01,-186.2 15.06,-124.48 19,-72 19.6,-63.99 20.71,-55.35 21.91,-47.43\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"25.35,-48.04 23.51,-37.61 18.45,-46.91 25.35,-48.04\"/>\n",
       "</g>\n",
       "<!-- var3&#45;&gt;var1 -->\n",
       "<g id=\"edge3\" class=\"edge\">\n",
       "<title>var3&#45;&gt;var1</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M40.08,-215.59C42.95,-191.61 48.16,-148.14 51.6,-119.42\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"55.06,-119.96 52.77,-109.61 48.11,-119.13 55.06,-119.96\"/>\n",
       "</g>\n",
       "<!-- var3&#45;&gt;var2 -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>var3&#45;&gt;var2</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M51.6,-215.7C58.37,-207.07 66.68,-196.5 74.06,-187.1\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"76.69,-189.42 80.12,-179.39 71.19,-185.09 76.69,-189.42\"/>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "<gtbook.display.show at 0x107e8f250>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#| caption: Bayes net for the 4-variable model, again.\n",
    "#| label: fig:bayesnet_4\n",
    "show(wxyz, boxes={Z1[0], W1[0]})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dhdxP2iM3AJD",
   "metadata": {},
   "source": [
    "Bayes’ theorem as applied above suggests an easy algorithm to calculate the posterior distribution\n",
    "above: simply enumerate all tuples $\\mathcal{X}$ in a table, evaluate\n",
    "$P(\\mathcal{X}, \\mathcal{Z}=\\mathfrak{z})$ for each one, and then\n",
    "normalize. Note we need *only* enumerate the variables in $\\mathcal{X}$, as all variables in $\\mathcal{Z}$ are fixed. \n",
    "\n",
    "Applying this to our example entails evaluating 100 probabilities, if we assume -as before - that each variable can take on 10 different outcomes. We show the resulting table for the case that $\\mathfrak{z}=(2, 7)$, i.e., every entry in the table is a value for $P(X, Y|W=2, Z=7)\\propto P(W=2, X, Y, Z=7)$. Below we show the factorized joint probability, to make a point:\n",
    "\n",
    "|    *x*   |    *y*   |                 *P(W=2, X=x, Y=y, Z=7)*                |\n",
    "|:--------:|:--------:|:---------------------------------------------------:|\n",
    "|     1    |     1    |*P(W=2\\|X=1, Z=7)P(X=1\\|Y=1, Z=7)P(Y=1\\|Z=7)P(Z=7)*   |\n",
    "|     1    |     2    |    *P(W=2\\|X=1, Z=7)P(X=1\\|Y=2, Z=7)P(Y=2\\|Z=7)P(Z=7)*   |\n",
    "| $\\vdots$ | $\\vdots$ |                       $\\vdots$                      |\n",
    "|    10    |     9    |   *P(W=2\\|X=10, Z=7)P(X=10\\|Y=9, Z=7)P(Y=9\\|Z=7)P(Z=7)*  |\n",
    "|    10    |    10    | *P(W=2\\|X=10, Z=7)P(X=10\\|Y=10, Z=7)P(Y=10\\|Z=7)P(Z=7)* |"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "xMueNLA_YtdX",
   "metadata": {},
   "source": [
    "We normalize by calculating $\\sum_{x, y} P(W=2, X=x, Y=y, Z=7)$ by summing over all these entries, and subsequently dividing all entries by the sum.\n",
    "\n",
    "Estimating the full posterior using this approach is, obviously, *not* efficient. \n",
    "In this example the table contains 100 entries, and in\n",
    "general the number of entries is *exponential* in the size of\n",
    "$\\mathcal{X}$. \n",
    "However, when inspecting the entries in the table\n",
    "there are already some obvious ways to save: for example, $P(Z=7)$ is a\n",
    "common factor in all entries, so clearly we need not bother\n",
    "multiplying it in. Below we will discuss methods to fully\n",
    "exploit the structure of the Bayes net to perform efficient inference.\n",
    "\n",
    "### Exercises\n",
    "\n",
    "1. Think more deeply about which other factors get repeatedly (and redundantly) calculated, and how you might organize the computation to avoid this.\n",
    "\n",
    "2. Show that in the example above, if we instead condition on known values for $\\mathcal{Z}=(X,Z)$, the posterior $P(W,Z|X,Y)$ factors, and as a consequence we only have to enumerate two tables of length 10, instead of a large table of size 100."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "u5az7jXrp9tU",
   "metadata": {},
   "source": [
    "```{index} Robot localization, nuisance variables\n",
    "```\n",
    "### Maximum a Posteriori Estimation\n",
    "\n",
    "> MAP saves a bit on compute, but is still expensive when done naively.\n",
    "\n",
    "For the purpose of decision making,\n",
    "it is often the case that we do not require the full posterior distribution.\n",
    "In these cases, we could rely on methods introduced in Section 2.4, \n",
    "such as Maximum Likelihood Estimation (MLE)\n",
    "or Maximum a posteriori (MAP) estimation.\n",
    "Here, we consider MAP estimation.\n",
    "\n",
    "Suppose we are given the values\n",
    "$\\mathfrak{z}$ for $\\mathcal{Z}$.\n",
    "The MAP estimate of the\n",
    "joint assignment to the other variables $\\mathcal{X}$ is given by\n",
    "\\begin{equation}\n",
    "x^*_{MAP} = \\arg \\max_x P(\\mathcal{X}|\\mathcal{Z}=\\mathfrak{z}).\n",
    "\\end{equation}\n",
    "For example, given\n",
    "$\\mathfrak{z}=(2, 7)$, the MAP estimate for $\\mathcal{X}$ could be $x^*=3$ and\n",
    "$y^*=6$. Note that to compute the MAP estimate, we need not bother with\n",
    "normalizing: we can simply find the maximum entry in a table of *unnormalized*\n",
    "posterior values.\n",
    "\n",
    "**Robot localization** is a capability for autonomous mobile robots that can be achieved via MAP estimation.\n",
    "For example, using the \"robot\" dynamic Bayes net example from the last section, let us\n",
    "assume that we are given the value of all observations and actions.\n",
    "Then the MAP estimate would simply be a trajectory of robot states. \n",
    "Hence, if we had an efficient way to do inference, a MAP estimate would be a\n",
    "great way to estimate the trajectory of a robot over time, i.e., robot localization.\n",
    "\n",
    "We can extend the idea of MAP to estimate only a subset of the unknown variables.\n",
    "This can be useful when there are unknown variables that are not relevant\n",
    "for the decision at hand.\n",
    "We will refer to these unknown and irrelevant variables as **nuisance variables**.\n",
    "In this case, we partition the variables into three sets:\n",
    "the variables of interest $\\mathcal{X}$, \n",
    "the nuisance variables $\\mathcal{Y}$, and the observed variables\n",
    "$\\mathcal{Z}$.\n",
    "Now, the posterior $P(\\mathcal{X}|\\mathcal{Z}=\\mathfrak{z})$\n",
    "can be determined by marginalizing (Section 2.4.7) over\n",
    "the nuisance variables:\n",
    "\\begin{equation}\n",
    "P(\\mathcal{X}|\\mathcal{Z}=\\mathfrak{z})=\\sum_{\\mathfrak{y}}P(\\mathcal{X}, \\mathcal{Y}=\\mathfrak{y}|\\mathcal{Z}=\\mathfrak{z})\\propto\\sum_{\\mathfrak{y}}P(\\mathcal{X}, \\mathcal{Y}=\\mathfrak{y}, \\mathcal{Z}=\\mathfrak{z}).\n",
    "\\end{equation}\n",
    "This approach to dealing with nuisance variables\n",
    "increases the computational cost of finding the MAP estimate.\n",
    "In\n",
    "addition to enumerating all possible combinations of $\\mathcal{X}$ and\n",
    "$\\mathcal{Y}$ values, we now need to calculate\n",
    "a possibly large number of sums, each exponential in the size of\n",
    "$\\mathcal{Y}$. In addition, the *number* of sums is\n",
    "exponential in the size of $\\mathcal{X}$. Below we will see that\n",
    "while we can still exploit the Bayes net structure, \n",
    "this computation can still be quite expensive."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "R82pyXgSkBho",
   "metadata": {},
   "source": [
    "### Exercises\n",
    "\n",
    "1.  Calculate the size of the table needed to enumerate the posterior\n",
    "    over the states $S$ the robot dynamic Bayes net from the previous section,\n",
    "    given the value of all observations $Z$ and actions $A$.\n",
    "\n",
    "2.  Show that if we are given the states, inferring the actions is\n",
    "    actually quite efficient, even with the brute force enumeration.\n",
    "    Hint: this is similar to the first exercise above."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "WxIdUENdGWiT",
   "metadata": {},
   "source": [
    "```{index} pair: hidden Markov model; HMM\n",
    "```\n",
    "## Hidden Markov Models\n",
    "\n",
    "> HMMs provide a general framework for perception over time.\n",
    "\n",
    "In the previous section we discussed dynamic Bayesian networks to model how a robot state evolves over time by taking actions, and how measurements correlate to a particular state. In this section we will ask *how we can recover the state of the robot given only the observations*, i.e. without knowing the states: the state is \"hidden\". Here we will consider a general framework to answer this question.\n",
    "\n",
    "A **hidden Markov model** or **HMM** is a dynamic Bayes net that has two\n",
    "types of variables: states $\\mathcal{X}$ and measurements $\\mathcal{Z}$.\n",
    "The states $\\mathcal{X}$ are connected sequentially and satisfy the Markov property: \n",
    "the probability of a state $X_k$ is\n",
    "only dependent on the value of the previous state $X_{k-1}$. As we saw before, we call a sequence of random variables with this property a **Markov chain.** \n",
    "In addition, in an HMM we refer to the states $\\mathcal{X}$ as *hidden*\n",
    "states, as typically we cannot directly observe their values. Instead, \n",
    "they are indirectly observed through the measurements $\\mathcal{Z}$, \n",
    "where we have one measurement per hidden state. When these two\n",
    "properties are satisfied, we call this probabilistic model a hidden\n",
    "Markov model."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "PAAbIsoSv49l",
   "metadata": {},
   "source": [
    "<figure id=\"fig:unrolledHMM\"> \n",
    "<img src=\"https://raw.githubusercontent.com/gtbook/robotics/main/Figures3/hmm-v2.png?raw=1\" style=\"width:14cm\" alt=\"\">\n",
    "<figcaption>An HMM for three time-steps, represented as a Bayes net.</figcaption>\n",
    "</figure>\n",
    "\n",
    "Figure\n",
    "<a href=\"#fig:unrolledHMM\" data-reference-type=\"ref\" data-reference=\"fig:unrolledHMM\">3.9</a>\n",
    "shows an example of an HMM for three time steps, i.e.., \n",
    "$\\mathcal{X}=\\{X_1, X_2, X_3\\}$ and\n",
    "$\\mathcal{Z}=\\{Z_1, Z_2, Z_3\\}$. As discussed above, in a Bayes net\n",
    "each node is associated with a conditional distribution: the Markov\n",
    "chain has the prior $P(X_1)$ and transition probabilities\n",
    "$P(X_2|X_1)$ and $P(X_3|X_2)$, whereas the measurements $Z_k$\n",
    "depend only on the state $X_k$, modeled by measurement models\n",
    "$P(Z_k|X_k)$. In other words, the Bayes net associated with this HMM encodes the following\n",
    "joint distribution $P(\\mathcal{X}, \\mathcal{Z})$:\n",
    "\\begin{equation}\n",
    "P(\\mathcal{X}, \\mathcal{Z})=P(X_1)P(Z_1|X_1)P(X_2|X_1)P(Z_2|X_2)P(X_3|X_2)P(Z_3|X_3)\n",
    "\\end{equation}\n",
    "\n",
    "Note that we can also write this more succinctly as\n",
    "\\begin{equation}\n",
    "P(\\mathcal{X}, \\mathcal{Z})=P(\\mathcal{Z}|\\mathcal{X})P(\\mathcal{X})\n",
    "\\end{equation}\n",
    "where\n",
    "\\begin{equation}\n",
    "P(\\mathcal{X})=P(X_1, X_2, X_3)=P(X_1)P(X_2|X_1)P(X_3|X_2)\n",
    "\\end{equation}\n",
    "is the prior over state *trajectories*."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23zQ9r0BM5v3",
   "metadata": {},
   "source": [
    "### Example: Robot HMM\n",
    "\n",
    "Let us re-create the dynamic Bayes net from the previous section here, with 3 time steps. However, we now also add conditionals on the three measurements $Z_1$, $Z_2$, and $Z_3$. We show the measurement variables as boxes, as they are *given*. The only variables that are unknown or *hidden* are the three robots states $X_1$, $X_2$, and $X_3$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "iOG7fJuri2Wu",
   "metadata": {},
   "outputs": [],
   "source": [
    "dbn = gtsam.DiscreteBayesNet()\n",
    "for k in range(1, N+1):\n",
    "    dbn.add(Z[k], [X[k]], vacuum.sensor_spec)\n",
    "for k in reversed(range(1, N)):\n",
    "    dbn.add(X[k+1], [X[k], A[k]], vacuum.action_spec)\n",
    "dbn.add(X[1], \"1/1/1/1/1\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "YqaPnxPq4Mq2",
   "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=\"206pt\" height=\"188pt\"\n",
       " viewBox=\"0.00 0.00 206.00 188.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 184)\">\n",
       "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-184 202,-184 202,4 -4,4\"/>\n",
       "<!-- var4683743612465315841 -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>var4683743612465315841</title>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"54,-180 0,-180 0,-144 54,-144 54,-180\"/>\n",
       "<text text-anchor=\"middle\" x=\"27\" y=\"-156.95\" font-family=\"Times,serif\" font-size=\"14.00\">A1</text>\n",
       "</g>\n",
       "<!-- var6341068275337658370 -->\n",
       "<g id=\"node4\" class=\"node\">\n",
       "<title>var6341068275337658370</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"99\" cy=\"-90\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"99\" y=\"-84.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,-143.83C54.54,-134.46 66.07,-122.93 76.05,-112.95\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"78.33,-115.62 82.92,-106.08 73.38,-110.67 78.33,-115.62\"/>\n",
       "</g>\n",
       "<!-- var4683743612465315842 -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>var4683743612465315842</title>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"126,-180 72,-180 72,-144 126,-144 126,-180\"/>\n",
       "<text text-anchor=\"middle\" x=\"99\" y=\"-156.95\" font-family=\"Times,serif\" font-size=\"14.00\">A2</text>\n",
       "</g>\n",
       "<!-- var6341068275337658371 -->\n",
       "<g id=\"node5\" class=\"node\">\n",
       "<title>var6341068275337658371</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"171\" cy=\"-90\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"171\" y=\"-84.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,-143.83C126.54,-134.46 138.07,-122.93 148.05,-112.95\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"150.33,-115.62 154.92,-106.08 145.38,-110.67 150.33,-115.62\"/>\n",
       "</g>\n",
       "<!-- var6341068275337658369 -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>var6341068275337658369</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"27\" cy=\"-90\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"27\" y=\"-84.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,-90C56.28,-90 58.38,-90 60.5,-90\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"60.23,-93.5 70.23,-90 60.23,-86.5 60.23,-93.5\"/>\n",
       "</g>\n",
       "<!-- var6485183463413514241 -->\n",
       "<g id=\"node6\" class=\"node\">\n",
       "<title>var6485183463413514241</title>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"54,-36 0,-36 0,0 54,0 54,-36\"/>\n",
       "<text text-anchor=\"middle\" x=\"27\" y=\"-12.95\" font-family=\"Times,serif\" font-size=\"14.00\">Z1</text>\n",
       "</g>\n",
       "<!-- var6341068275337658369&#45;&gt;var6485183463413514241 -->\n",
       "<g id=\"edge7\" class=\"edge\">\n",
       "<title>var6341068275337658369&#45;&gt;var6485183463413514241</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M27,-71.83C27,-64.55 27,-55.98 27,-47.83\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"30.5,-47.93 27,-37.93 23.5,-47.93 30.5,-47.93\"/>\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,-90C128.28,-90 130.38,-90 132.5,-90\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"132.23,-93.5 142.23,-90 132.23,-86.5 132.23,-93.5\"/>\n",
       "</g>\n",
       "<!-- var6485183463413514242 -->\n",
       "<g id=\"node7\" class=\"node\">\n",
       "<title>var6485183463413514242</title>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"126,-36 72,-36 72,0 126,0 126,-36\"/>\n",
       "<text text-anchor=\"middle\" x=\"99\" y=\"-12.95\" font-family=\"Times,serif\" font-size=\"14.00\">Z2</text>\n",
       "</g>\n",
       "<!-- var6341068275337658370&#45;&gt;var6485183463413514242 -->\n",
       "<g id=\"edge6\" class=\"edge\">\n",
       "<title>var6341068275337658370&#45;&gt;var6485183463413514242</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M99,-71.83C99,-64.55 99,-55.98 99,-47.83\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"102.5,-47.93 99,-37.93 95.5,-47.93 102.5,-47.93\"/>\n",
       "</g>\n",
       "<!-- var6485183463413514243 -->\n",
       "<g id=\"node8\" class=\"node\">\n",
       "<title>var6485183463413514243</title>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"198,-36 144,-36 144,0 198,0 198,-36\"/>\n",
       "<text text-anchor=\"middle\" x=\"171\" y=\"-12.95\" font-family=\"Times,serif\" font-size=\"14.00\">Z3</text>\n",
       "</g>\n",
       "<!-- var6341068275337658371&#45;&gt;var6485183463413514243 -->\n",
       "<g id=\"edge5\" class=\"edge\">\n",
       "<title>var6341068275337658371&#45;&gt;var6485183463413514243</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M171,-71.83C171,-64.55 171,-55.98 171,-47.83\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"174.5,-47.93 171,-37.93 167.5,-47.93 174.5,-47.93\"/>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "<gtbook.display.show at 0x137c53ac0>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#| caption: A hidden Markov model.\n",
    "#| label: fig:hmm_vacuum\n",
    "show(dbn, hints={\"A\": 2, \"X\": 1, \"Z\": 0}, boxes={A[k][0] for k in range(1, N)}.union({Z[k][0] for k in range(1, N+1)}))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "Vs24EealLQ8j",
   "metadata": {},
   "source": [
    "```{index} likelihood\n",
    "```\n",
    "## Naive Inference in HMMs\n",
    "\n",
    "> Inference is easy to implement naively, but hopelessly inefficient.\n",
    "\n",
    "As we saw above,\n",
    "one way to perform inference is to apply Bayes’ rule to obtain an expression for the *posterior* probability distribution over\n",
    "the state trajectory $\\mathcal{X}$, given the measurements\n",
    "$\\mathcal{Z}=\\mathfrak{z}$:\n",
    "\\begin{align}\n",
    "P(\\mathcal{X}|\\mathcal{Z}) & \\propto P(\\mathcal{Z}=\\mathfrak{z}|\\mathcal{X})P(\\mathcal{X}) \\\\\n",
    "& =L(\\mathcal{X}; \\mathcal{Z}=\\mathfrak{z})P(\\mathcal{X})\n",
    "\\end{align}\n",
    "where $P(\\mathcal{X})$ is the trajectory prior\n",
    "and the **likelihood** $L(\\mathcal{X}; \\mathcal{Z}=\\mathfrak{z})$ of\n",
    "$\\mathcal{X}$ given $\\mathcal{Z}=\\mathfrak{z}$ is defined as before as a function of $\\mathcal{X}$: \n",
    "\\begin{align}\n",
    "L(\\mathcal{X}; \\mathcal{Z}=\\mathfrak{z}) & \\propto P(\\mathcal{Z}=\\mathfrak{z}|\\mathcal{X})\\\\\n",
    "& =P(z_1|X_1)P(z_2|X_2)P(z_3|X_3)\\\\\n",
    "& \\propto L(X_1; Z_1)L(X_2; Z_2)L(X_3; Z_3)\n",
    "\\end{align}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "XP_FNxSIiYJK",
   "metadata": {},
   "source": [
    "Hence, a naive implementation for finding the MAP estimate\n",
    "for $\\mathcal{X}$ would tabulate all possible\n",
    "trajectories $\\mathcal{X}$ and calculate the posterior $P(\\mathcal{X}|\\mathcal{Z})$ for each one. \n",
    "Unfortunately the number of entries in this giant table is\n",
    "*exponential* in the number of states. Not only is this computationally\n",
    "prohibitive for long trajectories, but intuitively it is clear that for\n",
    "many of these trajectories we are computing the same values over and\n",
    "over again. There are three different approaches to improve on\n",
    "this:\n",
    "\n",
    "1.  Branch & bound\n",
    "\n",
    "2.  Dynamic programming\n",
    "\n",
    "3.  Inference using factor graphs\n",
    "\n",
    "Branch and bound is a powerful technique but will not generalize to\n",
    "continuous variables; the other two approaches will. And, we will\n",
    "see that dynamic programming, which underlies the classical inference\n",
    "algorithms in the HMM literature, is just a special case of the last\n",
    "approach. Hence, here we will dive in and immediately go for the most\n",
    "general approach: inference in factor graphs."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2Y6wLx76OVk",
   "metadata": {},
   "source": [
    "```{index} factors, unary factors, binary factors\n",
    "```\n",
    "## Factor Graphs\n",
    "\n",
    "> Factor graphs are an excellent representation in which to do inference.\n",
    "\n",
    "We first introduce the notion of factors. \n",
    "Again referring to the example from Figure\n",
    "<a href=\"#fig:unrolledHMM\" data-reference-type=\"ref\" data-reference=\"fig:unrolledHMM\">1</a>, \n",
    "let us consider the posterior.\n",
    "Since the measurements $\\mathcal{Z}$ are *known*, the posterior is\n",
    "proportional to the product of six **factors**, three of which derive\n",
    "from the Markov chain, and three of which are likelihood factors as defined\n",
    "above:\n",
    "\\begin{equation}\n",
    "P(\\mathcal{X}|\\mathcal{Z})\\propto P(X_1)L(X_1; z_1)P(X_2|X_1)L(X_2; z_2)P(X_3|X_2)L(X_3; z_3)\n",
    "\\end{equation}\n",
    "Some of these factors are **unary factors**, and some are **binary factors**. \n",
    "By which we mean: some of the factors depend on just *one* hidden variable, \n",
    "for example $L(X_2; z_2)$, whereas others depend on *two* variables, e.g., the\n",
    "transition model $P(X_3|X_2)$. \n",
    "Measurements are not counted here, \n",
    "because once we are *given* the measurements $\\mathcal{Z}$, they merely\n",
    "function as known parameters in the likelihoods $L(X_k; z_k)$, which\n",
    "are seen as functions of *just* the state $X_k$, i.e., they will yield unary factors in the unknowns."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "sh6peg02BFZQ",
   "metadata": {},
   "source": [
    "```{index} factor graph\n",
    "```\n",
    "<figure id=\"fig:HMM-FG\">\n",
    "<img src=\"https://raw.githubusercontent.com/gtbook/robotics/main/Figures3/fg-v2.png?raw=1\" style=\"width:60%\" alt=\"\">\n",
    "<figcaption>An HMM with observed measurements for three time steps, represented as a factor graph.</figcaption>\n",
    "</figure>\n",
    "\n",
    "This motivates a different graphical model, a **factor graph**, in which\n",
    "we only represent the *hidden* variables $X_1$, $X_2$, and $X_3$, \n",
    "connected to factors that encode probabilistic information. For\n",
    "our example with three hidden states, the corresponding factor graph is\n",
    "shown in Figure\n",
    "<a href=\"#fig:HMM-FG\" data-reference-type=\"ref\" data-reference=\"fig:HMM-FG\">2</a> above.\n",
    "It should be clear from the figure that the connectivity of a factor\n",
    "graph encodes, for each factor $\\phi_{i}$, which subset of variables\n",
    "$\\mathcal{X}_{i}$ it depends on. We write:\n",
    "\\begin{equation}\n",
    "\\phi(\\mathcal{X})=\\phi_1(X_1)\\phi_2(X_1)\\phi_3(X_1, X_2)\\phi_4(X_2)\\phi_5(X_2, X_{3})\\phi_6(X_3)\n",
    "\\end{equation}\n",
    "where the factors above are defined to correspond one-to-one to the six factors in the posterior, \n",
    "e.g., \n",
    "\\begin{equation}\n",
    "\\phi_6(X_3)\\doteq L(X_3; z_3).\n",
    "\\end{equation}\n",
    "\n",
    "All measurements are associated with unary factors, whereas the Markov chain is\n",
    "associated mostly with binary factors, with the exception of the unary\n",
    "factor $\\phi_1(X_1)$. Note that in defining the factors we can omit\n",
    "any normalization factors, which in many cases results in computational\n",
    "savings."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5Awn8wWgKGed",
   "metadata": {},
   "source": [
    "```{index} factors, variables\n",
    "```\n",
    "Formally a factor graph is a bipartite graph\n",
    "$F=(\\mathcal{U}, \\mathcal{V}, \\mathcal{E})$ with two types of nodes:\n",
    "**factors** $\\phi_{i}\\in\\mathcal{U}$ and **variables**\n",
    "*$X_{j}\\in\\mathcal{V}$.* Edges $e_{ij}\\in\\mathcal{E}$ are always between\n",
    "factor nodes and variables nodes, never between variables or between factors. \n",
    "The set of random variable nodes\n",
    "adjacent to a factor $\\phi_{i}$ is written as $\\mathcal{X}_{i}$. With\n",
    "these definitions, a factor graph $F$ defines the factorization of a\n",
    "global function $\\phi(\\mathcal{X})$ as\n",
    "\\begin{equation}\n",
    "\\phi(\\mathcal{X})=\\prod_{i}\\phi_{i}(\\mathcal{X}_{i}).\n",
    "\\end{equation}\n",
    "\n",
    "In other words, the independence relationships are encoded by the edges\n",
    "$e_{ij}$ of the factor graph, with each factor $\\phi_{i}$ a function of\n",
    "*only* the variables $\\mathcal{X}_{i}$ in its adjacency set. As example, \n",
    "for the factor graph in Figure\n",
    "<a href=\"#fig:HMM-FG\" data-reference-type=\"ref\" data-reference=\"fig:HMM-FG\">2</a>\n",
    "we have: \n",
    "\\begin{equation}\n",
    "\\begin{aligned}\n",
    "\\mathcal{X}_1 & =\\{X_1\\}\\\\\n",
    "\\mathcal{X}_2 & =\\{X_1\\}\\\\\n",
    "\\mathcal{X}_3 & =\\{X_1, X_2\\}\\\\\n",
    "\\mathcal{X}_4 & =\\{X_2\\}\\\\\n",
    "\\mathcal{X}_5 & =\\{X_2, X_3\\}\\\\\n",
    "\\mathcal{X}_6 & =\\{X_3\\}\\end{aligned}\n",
    "\\end{equation}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "k0_p5zxqeQiq",
   "metadata": {},
   "source": [
    "## Converting Bayes Nets into Factor Graphs.\n",
    "\n",
    "> It is trivial to convert Bayes nets into factor graphs.\n",
    "\n",
    "<figure id=\"fig:conversion\">\n",
    "<img src=\"https://raw.githubusercontent.com/gtbook/robotics/main/Figures3/hmm-v2.png?raw=1\" style=\"width:12cm\" alt=\"\">\n",
    "<figcaption>Bayes net representation of an HMM.</figcaption>\n",
    "</figure>\n",
    "\n",
    "<figure id=\"fig:conversion-2\">\n",
    "<img src=\"https://raw.githubusercontent.com/gtbook/robotics/main/Figures3/fg-v2.png?raw=1\"  style=\"width:14cm\" alt=\"\">\n",
    "<figcaption>Conversion of HMM above to a factor graph, where measurements are known.</figcaption>\n",
    "</figure>\n",
    "\n",
    "Every Bayes net can be trivially converted to a factor graph, as shown above.\n",
    "Recall that every node in a Bayes net denotes a conditional density on the\n",
    "corresponding variable and its parent nodes. Hence, the conversion is\n",
    "quite simple: every Bayes net node maps to *both* a variable node and\n",
    "a factor node in the corresponding factor graph. The factor is connected\n",
    "to the variable node, as well as the variable nodes corresponding to the\n",
    "parent nodes in the Bayes net. If some nodes in the Bayes net are\n",
    "evidence nodes, i.e., they are given as known variables, we omit the\n",
    "corresponding variable nodes: the known variable simply becomes a fixed\n",
    "parameter in the corresponding factor.\n",
    "\n",
    "### Exercise\n",
    "\n",
    "1.  Convert the dynamic Bayes net from the previous section into a factor graph, assuming *no* known variables.\n",
    "\n",
    "1.  Finally, do the same again, but now assume the states are given. Reflect on the remarkable phenomenon that happens.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c-yQxaGbFtaL",
   "metadata": {},
   "source": [
    "### Factor Graphs in GTSAM\n",
    "\n",
    "Let us create the factor graph directly using GTSAM. Before we do, however, we need to instantiate the given actions and measurements, both of which are assumed known:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "HTqIzNKU08fP",
   "metadata": {},
   "outputs": [],
   "source": [
    "actions = VARIABLES.assignment({A[1]: 'R', A[2]: 'U'})\n",
    "measurements = ['dark', 'medium', 'light']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "xJ7FhVen4NyQ",
   "metadata": {},
   "source": [
    "Now we create the factor graph, first adding the prior $\\phi(X_1)=P(X_1)$ on $X_1$, then the binary factors $\\phi(X_k, X_{k+1}) = P(X_{k+1}|X_k, A_k=a_k)$,\n",
    "and finally, the measurement likelihood factors $\\phi(X_k; Z_k=z_k) \\propto P(Z_k=z_k|X_k)$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "oWiM-RaZp5jZ",
   "metadata": {},
   "outputs": [],
   "source": [
    "graph = gtsam.DiscreteFactorGraph()\n",
    "graph.add(X[1], \"1 1 1 1 1\")  # \\phi(X_1) = P(X_1)\n",
    "for k in range(1, N):\n",
    "    conditional = gtsam.DiscreteConditional(X[k + 1], [X[k], A[k]], vacuum.action_spec)\n",
    "    conditional_a_k = conditional.choose(actions)  # \\phi(X,X+) = P(X+|X,A=a)\n",
    "    graph.push_back(conditional_a_k)\n",
    "for k, measurement in enumerate(measurements, start=1):\n",
    "    conditional = gtsam.DiscreteConditional(Z[k], [X[k]], vacuum.sensor_spec)\n",
    "    z_k = vacuum.light_levels.index(measurement)\n",
    "    factor = conditional.likelihood(z_k)  # \\phi(X) = P(Z=z|X)\n",
    "    graph.push_back(factor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "yOZ6NFg0WLp8",
   "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=\"206pt\" height=\"84pt\"\n",
       " viewBox=\"0.00 0.00 206.00 83.60\" 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 79.6)\">\n",
       "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-79.6 202,-79.6 202,4 -4,4\"/>\n",
       "<!-- var6341068275337658369 -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>var6341068275337658369</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"27\" cy=\"-57.6\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"27\" y=\"-52.55\" font-family=\"Times,serif\" font-size=\"14.00\">X1</text>\n",
       "</g>\n",
       "<!-- factor0 -->\n",
       "<g id=\"node4\" class=\"node\">\n",
       "<title>factor0</title>\n",
       "<ellipse fill=\"black\" stroke=\"black\" cx=\"16\" cy=\"-1.8\" rx=\"1.8\" ry=\"1.8\"/>\n",
       "</g>\n",
       "<!-- var6341068275337658369&#45;&#45;factor0 -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>var6341068275337658369&#45;&#45;factor0</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M23.46,-39.28C20.7,-25.77 17.17,-8.54 16.24,-3.96\"/>\n",
       "</g>\n",
       "<!-- factor1 -->\n",
       "<g id=\"node5\" class=\"node\">\n",
       "<title>factor1</title>\n",
       "<ellipse fill=\"black\" stroke=\"black\" cx=\"69\" cy=\"-1.8\" rx=\"1.8\" ry=\"1.8\"/>\n",
       "</g>\n",
       "<!-- var6341068275337658369&#45;&#45;factor1 -->\n",
       "<g id=\"edge3\" class=\"edge\">\n",
       "<title>var6341068275337658369&#45;&#45;factor1</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M39.15,-41.04C50,-27.14 64.72,-8.29 68.23,-3.79\"/>\n",
       "</g>\n",
       "<!-- factor3 -->\n",
       "<g id=\"node7\" class=\"node\">\n",
       "<title>factor3</title>\n",
       "<ellipse fill=\"black\" stroke=\"black\" cx=\"38\" cy=\"-1.8\" rx=\"1.8\" ry=\"1.8\"/>\n",
       "</g>\n",
       "<!-- var6341068275337658369&#45;&#45;factor3 -->\n",
       "<g id=\"edge6\" class=\"edge\">\n",
       "<title>var6341068275337658369&#45;&#45;factor3</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M30.54,-39.28C33.3,-25.77 36.83,-8.54 37.76,-3.96\"/>\n",
       "</g>\n",
       "<!-- var6341068275337658370 -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>var6341068275337658370</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"99\" cy=\"-57.6\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"99\" y=\"-52.55\" font-family=\"Times,serif\" font-size=\"14.00\">X2</text>\n",
       "</g>\n",
       "<!-- var6341068275337658370&#45;&#45;factor1 -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>var6341068275337658370&#45;&#45;factor1</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M89.84,-40.17C82.12,-26.32 71.95,-8.1 69.53,-3.76\"/>\n",
       "</g>\n",
       "<!-- factor2 -->\n",
       "<g id=\"node6\" class=\"node\">\n",
       "<title>factor2</title>\n",
       "<ellipse fill=\"black\" stroke=\"black\" cx=\"135\" cy=\"-1.8\" rx=\"1.8\" ry=\"1.8\"/>\n",
       "</g>\n",
       "<!-- var6341068275337658370&#45;&#45;factor2 -->\n",
       "<g id=\"edge5\" class=\"edge\">\n",
       "<title>var6341068275337658370&#45;&#45;factor2</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M109.61,-40.75C118.9,-26.86 131.37,-8.22 134.34,-3.78\"/>\n",
       "</g>\n",
       "<!-- factor4 -->\n",
       "<g id=\"node8\" class=\"node\">\n",
       "<title>factor4</title>\n",
       "<ellipse fill=\"black\" stroke=\"black\" cx=\"99\" cy=\"-1.8\" rx=\"1.8\" ry=\"1.8\"/>\n",
       "</g>\n",
       "<!-- var6341068275337658370&#45;&#45;factor4 -->\n",
       "<g id=\"edge7\" class=\"edge\">\n",
       "<title>var6341068275337658370&#45;&#45;factor4</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M99,-39.28C99,-25.77 99,-8.54 99,-3.96\"/>\n",
       "</g>\n",
       "<!-- var6341068275337658371 -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>var6341068275337658371</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"171\" cy=\"-57.6\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"171\" y=\"-52.55\" font-family=\"Times,serif\" font-size=\"14.00\">X3</text>\n",
       "</g>\n",
       "<!-- var6341068275337658371&#45;&#45;factor2 -->\n",
       "<g id=\"edge4\" class=\"edge\">\n",
       "<title>var6341068275337658371&#45;&#45;factor2</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M160.39,-40.75C151.1,-26.86 138.63,-8.22 135.66,-3.78\"/>\n",
       "</g>\n",
       "<!-- factor5 -->\n",
       "<g id=\"node9\" class=\"node\">\n",
       "<title>factor5</title>\n",
       "<ellipse fill=\"black\" stroke=\"black\" cx=\"171\" cy=\"-1.8\" rx=\"1.8\" ry=\"1.8\"/>\n",
       "</g>\n",
       "<!-- var6341068275337658371&#45;&#45;factor5 -->\n",
       "<g id=\"edge8\" class=\"edge\">\n",
       "<title>var6341068275337658371&#45;&#45;factor5</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M171,-39.28C171,-25.77 171,-8.54 171,-3.96\"/>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "<gtbook.display.show at 0x137c55550>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#| caption: A three-variable factor graph for when measurements are *given*\n",
    "#| label: fig:factor_graph_vacuum\n",
    "show(graph)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "RjvqP8Tx6qy6",
   "metadata": {},
   "source": [
    "Note that discrete distributions like $P(X_1)$ and conditionals $P(X_{k+1}|X_k)$, above, are perfectly fine factors, and in fact *derive* from the factor type in GTSAM. This is what allows us to add them directly the graph as is. Note that in a real implementation we might not take the detour to first construct the conditionals as above: we did so because they were conveniently available here, but typically we would construct factors directly."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "DouDxFHzbnQu",
   "metadata": {},
   "source": [
    "```{index} optimization\n",
    "```\n",
    "## Computing with Factor Graphs\n",
    "\n",
    "> We can evaluate, optimize, and sample from factor graphs.\n",
    "\n",
    "Once we convert a Bayes net with evidence into a factor graph where the\n",
    "evidence is all implicit in the factors, we can support a number of\n",
    "different computations. First, given any factor graph defining an\n",
    "unnormalized density $\\phi(X)$, we can easily *evaluate* it for any\n",
    "given value, by simply evaluating every factor and multiplying the\n",
    "results. The factor graph represents the unnormalized posterior, i.e., \n",
    "$\\phi(\\mathcal{X})\\propto P(\\mathcal{X}|\\mathcal{Z})$. \n",
    "\n",
    "Evaluation opens up the way to **optimization**, e.g., finding the MAP estimate,\n",
    "as we will do below. \n",
    "The naive optimization method of enumerating and \n",
    "evaluating all possible assignments for $\\mathcal{X}$ is of course always available,\n",
    "but just as inefficient.\n",
    "In the case of discrete variables, graph search methods can be applied, but we will discuss a\n",
    "different approach in the next sub-section."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "kC2F9SbMjG6l",
   "metadata": {},
   "source": [
    "Because our factor graph is so small, it does not hurt to show off how easy it is to implement the naive algorithm to find the MAP estimate. We just loop over all possible state trajectories, and keep track of the one with the highest value:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "0R4fG7gQl24N",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "found MAP solution with value 0.3277:\n"
     ]
    },
    {
     "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>X1</th><td>Hallway</td></tr>\n",
       "    <tr><th>X2</th><td>Dining Room</td></tr>\n",
       "    <tr><th>X3</th><td>Kitchen</td></tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "<gtbook.display.pretty at 0x137c53c10>"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "map_value = 0\n",
    "map_trajectory = None\n",
    "for x1 in vacuum.rooms:\n",
    "    for x2 in vacuum.rooms:\n",
    "        for x3 in vacuum.rooms:\n",
    "            trajectory = VARIABLES.assignment({X[1]: x1, X[2]: x2, X[3]: x3})\n",
    "            value = graph(trajectory)\n",
    "            if value > map_value:\n",
    "                map_value = value\n",
    "                map_trajectory = trajectory\n",
    "print(f\"found MAP solution with value {map_value:.4f}:\")\n",
    "pretty(map_trajectory)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "qS6jd8Hl4DET",
   "metadata": {},
   "source": [
    "Note that this MAP estimate is *for a given action and measurement sequence*. All those fixed values are implicit in the factors that we have added to the factor graph in `graph` above."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "Okr5Yxr_sk9N",
   "metadata": {},
   "source": [
    "```{index} sampling, Gibbs sampling\n",
    "```\n",
    "Finally, while finding the maximum of the posterior (MAP) is often of most\n",
    "interest, **sampling** from a probability distribution can be used to\n",
    "visualize, explore, and compute statistics and expected values\n",
    "associated with the posterior. The ancestral sampling method we\n",
    "discussed earlier only applies to Bayes bets, however.\n",
    "There are more general sampling algorithms that can be used for factor\n",
    "graphs. \n",
    "One such method is **Gibbs sampling**, which proceeds by sampling one variable\n",
    "at a time from its conditional density given all other variables it is\n",
    "connected to via factors. This assumes that this conditional density can\n",
    "be easily obtained, which is true for discrete variables."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "LVQ5e1tevKLf",
   "metadata": {},
   "source": [
    "## The Max-Product Algorithm for HMMs\n",
    "\n",
    "> Max-product on HMMs, also known as the Viterbi algorithm, is a dynamic-programming algorithm for finding a\n",
    "MAP estimate.\n",
    "\n",
    "In this section we discuss an algorithm that is much faster than the naive algorithm to find the MAP estimate.\n",
    "Given an HMM factor graph of size $n$, the **max-product algorithm** is an $O(n)$ algorithm\n",
    "to find the MAP estimate, which is used by GTSAM under the hood.\n",
    "\n",
    "Let us use the example from Figure\n",
    "<a href=\"#fig:HMM-FG\" data-reference-type=\"ref\" data-reference=\"fig:HMM-FG\">2</a>\n",
    "to understand the main idea behind it. To find the MAP estimate for $\\mathcal{X}$ we need to\n",
    "*maximize* the product\n",
    "\\begin{equation}\n",
    "\\phi(X_1, X_2, X_3)=\\prod\\phi_{i}(\\mathcal{X}_{i})\n",
    "\\end{equation}\n",
    "i.e., the value of the factor graph.\n",
    "Because this is a product of factor values, we can compute its maximum recursively - dynamic programming style. We start by writing out the maximization over the product explicitly:\n",
    "\\begin{equation}\n",
    "\\max_{\\mathcal{X}} \\prod\\phi_{i}(\\mathcal{X}_{i}) = \n",
    "\\max_{X_1, X_2, X_3} ~~~\\phi_1(X_1)\\phi_2(X_1)\\phi_3(X_1, X_2) \\phi_4(X_2)\\phi_5(X_2, X_3)\\phi_6(X_3) \n",
    "\\end{equation}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cvSL_rEsKz6B",
   "metadata": {},
   "source": [
    "The key to our recursion will be to consider each variable in turn, starting with $X_1$. In particular, let us group all the factors connected to $X_1$\n",
    "\\begin{equation}\n",
    "\\max_{\\mathcal{X}} \\prod\\phi_{i}(\\mathcal{X}_{i}) = \n",
    "\\max_{X_1, X_2, X_3} ~~~ \\{ \\phi_1(X_1)\\phi_2(X_1)\\phi_3(X_1, X_2) \\} ~~~ \\phi_4(X_2)\\phi_5(X_2, X_3)\\phi_6(X_3)\n",
    "\\end{equation}\n",
    "which allows us to move the *max* operator over $X_1$ inside:\n",
    "\\begin{equation}\n",
    "\\max_{\\mathcal{X}} \\prod\\phi_{i}(\\mathcal{X}_{i}) = \n",
    "\\max_{X_2, X_3} ~~~ \\{ \\max_{X_1} \\phi_1(X_1)\\phi_2(X_1)\\phi_3(X_1, X_2) \\} ~~~ \\phi_4(X_2)\\phi_5(X_2, X_3)\\phi_6(X_3)\n",
    "\\end{equation}\n",
    "The key to the recursion is that we can simply consider the expression inside the curly braces as new factor on $X_2$, defined as\n",
    "\\begin{equation}\n",
    "\\tau(X_2)\\doteq \\max_{X_1} \\phi_1(X_1)\\phi_2(X_1)\\phi_3(X_1, X_2),\n",
    "\\end{equation}\n",
    "which records the maximum value resulting from *only* maximizing $X_1$. The value of the computation depends on the value of $X_2$, because $X_2$ was involved in factor $\\phi_3$. However, crucially, the variable $X_3$ is not involved in the maximization."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "IzF0jpUAVC3O",
   "metadata": {},
   "source": [
    "Substituting the new factor into the maximization, we now need to maximize over a *reduced* factor graph that no longer is a function of $X_1$,\n",
    "\\begin{equation}\n",
    "\\max_{\\mathcal{X}} \\prod\\phi_{i}(\\mathcal{X}_{i}) =\n",
    "\\max_{X_2, X_3} ~~~\\tau(X_2) \\phi_4(X_2)\\phi_5(X_2, X_3)\\phi_6(X_3)\n",
    "\\end{equation}\n",
    "and this allows us to recurse until no more variables are left!"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "GMpyz7fTxaXq",
   "metadata": {},
   "source": [
    "```{index} Viterbi Algorithm, elimination algorithm\n",
    "```\n",
    "For hidden Markov models this dynamics programming approach is known as the **Viterbi Algorithm**. Note that the algorithmic sketch above only yields the *value* of the factor graph at the maximum. A key part of the Viterbi algorithm is adding some book-keeping that after the recursion terminates, we can recover the actual variable assignments that make the factor graph attain its maximum value.\n",
    "Also, the complexity of max-product for HMMs is *linear* in the number of nodes, which is a nice improvement over exponential complexity. The complexity of every elimination step is quadratic in the number of states, because we have to form the product factors and then maximize over them.\n",
    "\n",
    "Because at every step, one variable is eliminated from the maximization, the max-product algorithm is in fact an instance of the **elimination algorithm**, which works for *arbitrary* factor graphs, albeit not necessarily in $O(n)$ time. Indeed, the max-product (and sum-product below) can be applied in more general settings than the linear chains one finds in HMMs."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5pztOK7ALdX-",
   "metadata": {},
   "source": [
    "### Max-product in GTSAM\n",
    "\n",
    "GTSAM's bread and butter is factor graphs, and finding the MAP value is done via a single `optimize` call, which implements the max-product algorithm internally, book-keeping and all:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "n4P_G37SwHIV",
   "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>X1</th><td>Hallway</td></tr>\n",
       "    <tr><th>X2</th><td>Dining Room</td></tr>\n",
       "    <tr><th>X3</th><td>Kitchen</td></tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "<gtbook.display.pretty at 0x137c53c40>"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "map_value = graph.optimize()\n",
    "pretty(map_value)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83_BNSbWHkbQ",
   "metadata": {},
   "source": [
    "The resulting `gtsam.DiscreteValues` instance corresponds to the MAP trajectory given the actions and measurements at every step. And, of course, this generalizes to much larger sequences, scaling only linearly in complexity with the number of time-steps."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "HWrfjpb6h6tm",
   "metadata": {},
   "source": [
    "## The Sum-Product Algorithm for HMMs\n",
    "\n",
    "> Sum-product on HMMs is a dynamic-programming algorithm for doing full posterior inference.\n",
    "\n",
    "The sum-product algorithm for HMMs is a slight tweak on the max-product\n",
    "algorithm that instead calculates the posterior probability $P(\\mathcal{X}|\\mathcal{Z})$. \n",
    "Whereas the max-product results in a single variable assignment, the sum-product produces a Bayes net that represents the *full Bayesian probability distribution*.\n",
    "The fact that we recover this distribution in the form of a Bayes net again is\n",
    "satisfying, because, as we have seen, that is an economical representation of a\n",
    "probability distribution."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "DC05df5SWTRi",
   "metadata": {},
   "source": [
    "One might wonder about the wisdom of all this: we started with a Bayes\n",
    "net, converted to a factor graph, and now end up with a Bayes net again?\n",
    "There are two important differences. \n",
    "First, in many practical cases we do not even bother with the modeling step, but\n",
    "construct the factor graph directly from the measurements.\n",
    "Second, even if we did, this first Bayes represents the joint \n",
    "distribution $P(\\mathcal{X}, \\mathcal{Z})$, useful for modeling and simulation.\n",
    "However, the second, \"posterior Bayes net\" represents just $P(\\mathcal{X}|\\mathcal{Z})$, and only has nodes for the random variables in $\\mathcal{X}$ and is only half the size. Below we show how to compute it."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "T5vUszbCkGEs",
   "metadata": {},
   "source": [
    "Again, the key is that we can compute the posterior recursively from the product of factors, in dynamic programming style.\n",
    "For the example above, this gives\n",
    "\\begin{equation}\n",
    "\\begin{aligned}\n",
    "P(\\mathcal{X}|\\mathcal{Z}) &\\propto \\prod\\phi_{i}(\\mathcal{X}_{i})\n",
    "\\\\&\\propto \\phi_1(X_1)\\phi_2(X_1)\\phi_3(X_1, X_2) \\phi_4(X_2)\\phi_5(X_2, X_{3})\\phi_6(X_3)\n",
    "\\\\&\\propto \\{\\phi_1(X_1)\\phi_2(X_1)\\phi_3(X_1, X_2)\\} ~~ \\phi_4(X_2)\\phi_5(X_2, X_{3})\\phi_6(X_3)\n",
    "\\end{aligned}\n",
    "\\end{equation}\n",
    "where we once again grouped the factors connected to $X_1$ inside the curly braces. Using the chain rule, we can those as the product of a conditional (which will end up in the Bayes net) and a new factor $\\tau(X_2)$:\n",
    "\\begin{equation}\n",
    "\\phi_1(X_1)\\phi_2(X_1)\\phi_3(X_1, X_2) = P(X_1|X_2, \\mathcal{Z}) \\tau(X_2).\n",
    "\\end{equation}\n",
    "Note that above we explicitly indicate the dependence on $\\mathcal{Z}$ in $P(X_1|X_2, \\mathcal{Z})$, which was left implicit in the factor graph. Substituting this back in, we get a conditional multiplied with a smaller factor graph that no longer involves $X_1$, and we can recurse:\n",
    "\\begin{equation}\n",
    "\\begin{aligned}\n",
    "P(\\mathcal{X}|\\mathcal{Z}) &\\propto P(X_1|X_2, \\mathcal{Z}) ~~ \\{ \\tau(X_2) \\phi_4(X_2)\\phi_5(X_2, X_{3})\\phi_6(X_3)\\}\n",
    "\\\\ &\\propto P(X_1|X_2, \\mathcal{Z}) P(X_2 | X_{3}, \\mathcal{Z}) ~~ \\{\\tau(X_{3})\\phi_6(X_3)\\}\n",
    "\\\\ &\\propto P(X_1|X_2, \\mathcal{Z}) P(X_2 | X_{3}, \\mathcal{Z}) P(X_3| \\mathcal{Z}).\n",
    "\\end{aligned}\n",
    "\\end{equation}\n",
    "Each of the three conditionals above was computed by implementing the chain rule, which is a simple calculation in the case of discrete factors."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fGQcx63IexzH",
   "metadata": {},
   "source": [
    "### Sum-Product in GTSAM\n",
    "\n",
    "In GTSAM, once again the sum-product is but a simple call to the `sumProduct` method of a factor graph. It yields a `gtsam.DiscreteBayes` net which encodes the full posterior distribution:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "WKEIa2RlfylD",
   "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=\"206pt\" height=\"44pt\"\n",
       " viewBox=\"0.00 0.00 206.00 44.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 40)\">\n",
       "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-40 202,-40 202,4 -4,4\"/>\n",
       "<!-- var6341068275337658369 -->\n",
       "<g id=\"node1\" 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",
       "<!-- var6341068275337658370 -->\n",
       "<g id=\"node2\" 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",
       "<!-- var6341068275337658370&#45;&gt;var6341068275337658369 -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>var6341068275337658370&#45;&gt;var6341068275337658369</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M71.78,-18C69.72,-18 67.62,-18 65.5,-18\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"65.77,-14.5 55.77,-18 65.77,-21.5 65.77,-14.5\"/>\n",
       "</g>\n",
       "<!-- var6341068275337658371 -->\n",
       "<g id=\"node3\" 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",
       "<!-- var6341068275337658371&#45;&gt;var6341068275337658370 -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>var6341068275337658371&#45;&gt;var6341068275337658370</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M143.78,-18C141.72,-18 139.62,-18 137.5,-18\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"137.77,-14.5 127.77,-18 137.77,-21.5 137.77,-14.5\"/>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "<gtbook.display.show at 0x137c53760>"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#| caption: The full posterior encoded as a Bayes net.\n",
    "#| label: fig:posterior_bn\n",
    "posterior = graph.sumProduct()\n",
    "show(posterior, hints={\"X\": 1})\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "k0dOKqQiJarj",
   "metadata": {},
   "source": [
    "One of the things we can do with this *exact* posterior is sample from it, which is one possible state history conditioned on the available sensor measurements *and* the known action sequence:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "lkM_FCFVcK7S",
   "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>X1</th><td>Dining Room</td></tr>\n",
       "    <tr><th>X2</th><td>Dining Room</td></tr>\n",
       "    <tr><th>X3</th><td>Kitchen</td></tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "<gtbook.display.pretty at 0x137c53e50>"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sample = posterior.sample()\n",
    "pretty(sample)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "mjvB1v40hCkP",
   "metadata": {},
   "source": [
    "## Sampling and Posterior Means*\n",
    "\n",
    "> Bayesian reasoning at its best\n",
    "\n",
    "When we can produce samples $\\mathcal{X}^{(s)}$ from a posterior\n",
    "$P(\\mathcal{X}|\\mathcal{Z})$, we can calculate the empirical mean of any\n",
    "real-valued function $f(\\mathcal{X})$ as follows:\n",
    "\\begin{equation}\n",
    "E_{P(\\mathcal{X}|\\mathcal{Z})}[f(x)]\\approx\\sum f(\\mathcal{X}^{(s)})\n",
    "\\end{equation}\n",
    "For example, we can calculate the posterior mean of how far the robot\n",
    "traveled, either in Euclidean or Manhattan distance, using this approach. \n",
    "Doing this will provide a more reliable estimate than merely \n",
    "calculating the distance for the MAP value,\n",
    "since this approach averages over the entire probability distribution rather\n",
    "than just using a single (albeit most probable) estimate.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "N30uJddzZEfZ",
   "metadata": {},
   "source": [
    "For example, the code below samples 1000 alternate state histories, parallel universes of what *could* have happened: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "rkvXVghCGKbD",
   "metadata": {},
   "outputs": [],
   "source": [
    "counts = np.zeros((3, 5))\n",
    "num_samples = 1000\n",
    "for i in range(num_samples):\n",
    "    sample = posterior.sample()\n",
    "    for k in range(1,3+1):\n",
    "        key = X[k][0]\n",
    "        room_index = sample[key]\n",
    "        counts[k-1][room_index] += 1 # base 0!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "xclnXZFo-H8h",
   "metadata": {},
   "source": [
    "Posterior means gives us a way to summarize these 1000 alternate histories some other way other than just printing them all out. One idea is to summarize, for every time step, what the probability is to be in a particular room. It turns out we can do this with a one-liner, because we kept track of counts in the code above:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "zO9MnRJ_UnsM",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Living Room</th>\n",
       "      <th>Kitchen</th>\n",
       "      <th>Office</th>\n",
       "      <th>Hallway</th>\n",
       "      <th>Dining Room</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>3.3</td>\n",
       "      <td>1.4</td>\n",
       "      <td>3.8</td>\n",
       "      <td>80.6</td>\n",
       "      <td>10.9</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>0.9</td>\n",
       "      <td>3.8</td>\n",
       "      <td>0.8</td>\n",
       "      <td>5.0</td>\n",
       "      <td>89.5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>5.9</td>\n",
       "      <td>90.7</td>\n",
       "      <td>0.8</td>\n",
       "      <td>0.0</td>\n",
       "      <td>2.6</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   Living Room  Kitchen  Office  Hallway  Dining Room\n",
       "1          3.3      1.4     3.8     80.6         10.9\n",
       "2          0.9      3.8     0.8      5.0         89.5\n",
       "3          5.9     90.7     0.8      0.0          2.6"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.DataFrame(data=100*counts/num_samples,\n",
    "             index=range(1, N+1), columns=vacuum.rooms)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aY65ETPtIFFb",
   "metadata": {},
   "source": [
    "These approximate marginals say how probable it is that the robot was in a particular room at a particular time step. This is much richer information that what is available in the MAP value, which is just a point estimate for the trajectory. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dDyBnnk9CINk",
   "metadata": {},
   "source": [
    "### Exercise\n",
    "\n",
    "Execute the code above multiple times and observe that you *do* get different realizations (almost) every time, but that the approximate marginals stay roughly the same."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "upE7eUIYE77N",
   "metadata": {},
   "source": [
    "## GTSAM 101\n",
    "\n",
    "> The GTSAM concepts used in this section, explained."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3tTNjQ6doA0H",
   "metadata": {},
   "source": [
    "We created, for the first time, an instance of the `gtsam.DiscreteFactorGraph` class. The constructor is trivial - takes no arguments.\n",
    "To add factors, we can use the following methods:\n",
    "\n",
    " 1. `add(self, j: Tuple[int, int], spec: str) -> None`\n",
    "\n",
    " 2. `add(self, j: Tuple[int, int], spec: List[float]) -> None`\n",
    "\n",
    " 3. `add(self, keys: List[Tuple[int, int]], spec: str) -> None`\n",
    "\n",
    "These are very similar to the `gtsam.DiscreteBayesNet` methods, but in factor graphs there is a\n",
    "distinction between frontal and parent values, so we just have a key, or a list of keys as in the last method."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5yGf-Fc2RRo",
   "metadata": {},
   "source": [
    "Two key factor graph methods we used above are `optimize` and `sumProduct`:\n",
    "\n",
    "```python\n",
    "- optimize(self) -> gtsam::DiscreteValues\n",
    "- sumProduct(self) -> gtsam.DiscreteBayesNet\n",
    "```\n",
    "\n",
    "The first one returns the MAP value as an assignment to discrete variables, whereas the second returns an entire Bayes net, encoding the posterior."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "include_colab_link": true,
   "name": "S34_vacuum_perception.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.18"
  },
  "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": 5
}