POMDPs-jl-demo.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[32m\u001b[1m  Updating\u001b[22m\u001b[39m registry at `C:\\Users\\rmoss\\.julia\\registries\\General`\n",
      "\u001b[32m\u001b[1m  Updating\u001b[22m\u001b[39m git-repo `https://github.com/JuliaRegistries/General.git`\n",
      "\u001b[?25l\u001b[2K\u001b[?25h\u001b[32m\u001b[1m  Updating\u001b[22m\u001b[39m registry at `C:\\Users\\rmoss\\.julia\\registries\\JuliaPOMDP`\n",
      "\u001b[32m\u001b[1m  Updating\u001b[22m\u001b[39m git-repo `https://github.com/JuliaPOMDP/Registry`\n",
      "\u001b[?25l\u001b[2K\u001b[?25h\u001b[32m\u001b[1m Resolving\u001b[22m\u001b[39m package versions...\n",
      "\u001b[32m\u001b[1m  Updating\u001b[22m\u001b[39m `C:\\Users\\rmoss\\Code\\sisl\\aa228\\aa228-notebook\\Project.toml`\n",
      "\u001b[90m [no changes]\u001b[39m\n",
      "\u001b[32m\u001b[1m  Updating\u001b[22m\u001b[39m `C:\\Users\\rmoss\\Code\\sisl\\aa228\\aa228-notebook\\Manifest.toml`\n",
      "\u001b[90m [no changes]\u001b[39m\n",
      "\u001b[32m\u001b[1m Resolving\u001b[22m\u001b[39m package versions...\n",
      "\u001b[32m\u001b[1m  Updating\u001b[22m\u001b[39m `C:\\Users\\rmoss\\Code\\sisl\\aa228\\aa228-notebook\\Project.toml`\n",
      "\u001b[90m [no changes]\u001b[39m\n",
      "\u001b[32m\u001b[1m  Updating\u001b[22m\u001b[39m `C:\\Users\\rmoss\\Code\\sisl\\aa228\\aa228-notebook\\Manifest.toml`\n",
      "\u001b[90m [no changes]\u001b[39m\n",
      "\u001b[32m\u001b[1m Resolving\u001b[22m\u001b[39m package versions...\n",
      "\u001b[32m\u001b[1m  Updating\u001b[22m\u001b[39m `C:\\Users\\rmoss\\Code\\sisl\\aa228\\aa228-notebook\\Project.toml`\n",
      "\u001b[90m [no changes]\u001b[39m\n",
      "\u001b[32m\u001b[1m  Updating\u001b[22m\u001b[39m `C:\\Users\\rmoss\\Code\\sisl\\aa228\\aa228-notebook\\Manifest.toml`\n",
      "\u001b[90m [no changes]\u001b[39m\n",
      "\u001b[32m\u001b[1m Resolving\u001b[22m\u001b[39m package versions...\n",
      "\u001b[32m\u001b[1m  Updating\u001b[22m\u001b[39m `C:\\Users\\rmoss\\Code\\sisl\\aa228\\aa228-notebook\\Project.toml`\n",
      "\u001b[90m [no changes]\u001b[39m\n",
      "\u001b[32m\u001b[1m  Updating\u001b[22m\u001b[39m `C:\\Users\\rmoss\\Code\\sisl\\aa228\\aa228-notebook\\Manifest.toml`\n",
      "\u001b[90m [no changes]\u001b[39m\n"
     ]
    }
   ],
   "source": [
    "import Pkg\n",
    "Pkg.add(\"POMDPs\")\n",
    "Pkg.add(\"POMDPPolicies\")\n",
    "Pkg.add(\"POMDPModelTools\")\n",
    "Pkg.add(\"POMDPSimulators\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "using POMDPs\n",
    "using POMDPPolicies\n",
    "using POMDPModelTools\n",
    "using POMDPSimulators"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "include(\"helpers.jl\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chain MDP\n",
    "\n",
    "State space: Integer representing which node we are in.\n",
    "\n",
    "Action space: Try to move left or right (represented by symbols `:left` and `:right`.\n",
    "\n",
    "Transitions: Move in action direction with probability `p_success`, otherwise move in the opposite direction.\n",
    "\n",
    "Reward: +1 for being in rightmost state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "image/svg+xml": [
       "<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n",
       "<svg xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\" width=\"103.074pt\" height=\"18.207pt\" viewBox=\"0 0 103.074 18.207\" version=\"1.1\">\n",
       "<defs>\n",
       "<g>\n",
       "<symbol overflow=\"visible\" id=\"glyph-1567442442225000-0-0\">\n",
       "<path style=\"stroke:none;\" d=\"\"/>\n",
       "</symbol>\n",
       "<symbol overflow=\"visible\" id=\"glyph-1567442442225000-0-1\">\n",
       "<path style=\"stroke:none;\" d=\"M 4.34375 -0.15625 L 4.34375 -0.640625 L 3.859375 -0.640625 C 2.953125 -0.640625 3.09375 -0.578125 3.09375 -0.953125 L 3.09375 -6.546875 C 3.09375 -6.78125 2.9375 -6.96875 2.625 -6.96875 C 1.953125 -6.28125 1.203125 -6.328125 0.703125 -6.328125 L 0.703125 -5.6875 C 1.09375 -5.6875 1.796875 -5.734375 2.015625 -5.859375 L 2.015625 -0.953125 C 2.015625 -0.59375 2.15625 -0.640625 1.265625 -0.640625 L 0.765625 -0.640625 L 0.765625 0.015625 C 1.296875 -0.03125 2.15625 -0.03125 2.5625 -0.03125 C 2.953125 -0.03125 3.828125 -0.03125 4.34375 0.015625 Z M 4.34375 -0.15625 \"/>\n",
       "</symbol>\n",
       "<symbol overflow=\"visible\" id=\"glyph-1567442442225000-0-2\">\n",
       "<path style=\"stroke:none;\" d=\"M 4.46875 -2.0625 L 4.078125 -2.0625 C 4 -1.59375 3.984375 -1.28125 3.875 -1.125 C 3.8125 -1.046875 3.28125 -1.09375 3.0625 -1.09375 L 1.265625 -1.09375 L 1.390625 -0.8125 L 2.4375 -1.84375 C 4 -3.21875 4.640625 -3.875 4.640625 -4.859375 C 4.640625 -6 3.578125 -6.96875 2.359375 -6.96875 C 1.234375 -6.96875 0.3125 -5.875 0.3125 -5 C 0.3125 -4.4375 1 -4.28125 1.03125 -4.28125 C 1.203125 -4.28125 1.71875 -4.5625 1.71875 -4.96875 C 1.71875 -5.21875 1.359375 -5.65625 1.015625 -5.65625 C 0.9375 -5.65625 0.921875 -5.65625 1.140625 -5.734375 C 1.28125 -6.125 1.65625 -6.328125 2.234375 -6.328125 C 3.140625 -6.328125 3.390625 -5.6875 3.390625 -4.859375 C 3.390625 -4.0625 2.9375 -3.40625 2.390625 -2.78125 L 0.484375 -0.640625 C 0.375 -0.53125 0.3125 -0.40625 0.3125 0 L 4.34375 0 L 4.671875 -2.0625 Z M 4.46875 -2.0625 \"/>\n",
       "</symbol>\n",
       "<symbol overflow=\"visible\" id=\"glyph-1567442442225000-0-3\">\n",
       "<path style=\"stroke:none;\" d=\"M 4.71875 -1.859375 C 4.71875 -2.6875 3.921875 -3.625 2.890625 -3.84375 L 2.890625 -3.515625 C 3.703125 -3.78125 4.453125 -4.640625 4.453125 -5.421875 C 4.453125 -6.234375 3.40625 -6.96875 2.453125 -6.96875 C 1.453125 -6.96875 0.515625 -6.203125 0.515625 -5.4375 C 0.515625 -5.109375 0.90625 -4.765625 1.203125 -4.765625 C 1.5 -4.765625 1.875 -5.140625 1.875 -5.4375 C 1.875 -5.9375 1.234375 -6.09375 1.390625 -6.09375 C 1.515625 -6.3125 2.046875 -6.390625 2.40625 -6.390625 C 2.828125 -6.390625 3.1875 -6.328125 3.1875 -5.4375 C 3.1875 -5.3125 3.21875 -4.859375 2.953125 -4.421875 C 2.65625 -3.9375 2.453125 -3.953125 2.203125 -3.953125 C 2.125 -3.9375 1.890625 -3.921875 1.8125 -3.921875 C 1.734375 -3.90625 1.484375 -3.734375 1.484375 -3.625 C 1.484375 -3.515625 1.734375 -3.359375 1.90625 -3.359375 L 2.34375 -3.359375 C 3.15625 -3.359375 3.34375 -2.84375 3.34375 -1.859375 C 3.34375 -0.515625 2.84375 -0.390625 2.40625 -0.390625 C 1.96875 -0.390625 1.34375 -0.515625 1.203125 -0.765625 C 1.21875 -0.765625 1.703125 -1.140625 1.703125 -1.53125 C 1.703125 -1.890625 1.265625 -2.25 0.984375 -2.25 C 0.734375 -2.25 0.234375 -1.9375 0.234375 -1.5 C 0.234375 -0.59375 1.34375 0.21875 2.4375 0.21875 C 3.65625 0.21875 4.71875 -0.84375 4.71875 -1.859375 Z M 4.71875 -1.859375 \"/>\n",
       "</symbol>\n",
       "<symbol overflow=\"visible\" id=\"glyph-1567442442225000-0-4\">\n",
       "<path style=\"stroke:none;\" d=\"M 4.859375 -1.796875 L 4.859375 -2.28125 L 3.875 -2.28125 L 3.875 -6.65625 C 3.875 -6.84375 3.703125 -7.078125 3.53125 -7.078125 C 3.453125 -7.078125 3.296875 -7.03125 3.203125 -6.90625 L 0.09375 -2.15625 L 0.09375 -1.640625 L 2.75 -1.640625 L 2.75 -0.9375 C 2.75 -0.578125 2.90625 -0.640625 2.171875 -0.640625 L 1.78125 -0.640625 L 1.78125 0.015625 C 2.375 -0.03125 2.890625 -0.03125 3.3125 -0.03125 C 3.734375 -0.03125 4.25 -0.03125 4.828125 0.015625 L 4.828125 -0.640625 L 4.453125 -0.640625 C 3.71875 -0.640625 3.875 -0.578125 3.875 -0.9375 L 3.875 -1.640625 L 4.859375 -1.640625 Z M 2.984375 -2.28125 L 0.875 -2.28125 L 3.109375 -5.71875 L 2.8125 -5.828125 L 2.8125 -2.28125 Z M 2.984375 -2.28125 \"/>\n",
       "</symbol>\n",
       "</g>\n",
       "<clipPath id=\"clip-1567442442225000-1\">\n",
       "  <path d=\"M 0 0 L 24 0 L 24 18.207031 L 0 18.207031 Z M 0 0 \"/>\n",
       "</clipPath>\n",
       "<clipPath id=\"clip-1567442442225000-2\">\n",
       "  <path d=\"M 22 0 L 52 0 L 52 18.207031 L 22 18.207031 Z M 22 0 \"/>\n",
       "</clipPath>\n",
       "<clipPath id=\"clip-1567442442225000-3\">\n",
       "  <path d=\"M 56 0 L 75 0 L 75 18.207031 L 56 18.207031 Z M 56 0 \"/>\n",
       "</clipPath>\n",
       "<clipPath id=\"clip-1567442442225000-4\">\n",
       "  <path d=\"M 51 0 L 81 0 L 81 18.207031 L 51 18.207031 Z M 51 0 \"/>\n",
       "</clipPath>\n",
       "<clipPath id=\"clip-1567442442225000-5\">\n",
       "  <path d=\"M 79 0 L 103.074219 0 L 103.074219 18.207031 L 79 18.207031 Z M 79 0 \"/>\n",
       "</clipPath>\n",
       "</defs>\n",
       "<g id=\"surface1\">\n",
       "<path style=\"fill:none;stroke-width:0.3985;stroke-linecap:butt;stroke-linejoin:miter;stroke:rgb(0%,0%,0%);stroke-opacity:1;stroke-miterlimit:10;\" d=\"M 28.345906 -0.00146875 L 113.388875 -0.00146875 \" transform=\"matrix(1,0,0,-1,-19.342,9.104)\"/>\n",
       "<path style=\" stroke:none;fill-rule:nonzero;fill:rgb(100%,100%,100%);fill-opacity:1;\" d=\"M 17.808594 9.105469 C 17.808594 4.242188 13.867188 0.300781 9.003906 0.300781 C 4.140625 0.300781 0.199219 4.242188 0.199219 9.105469 C 0.199219 13.96875 4.140625 17.910156 9.003906 17.910156 C 13.867188 17.910156 17.808594 13.96875 17.808594 9.105469 Z M 17.808594 9.105469 \"/>\n",
       "<g clip-path=\"url(#clip-1567442442225000-1)\" clip-rule=\"nonzero\">\n",
       "<path style=\"fill:none;stroke-width:0.3985;stroke-linecap:butt;stroke-linejoin:miter;stroke:rgb(0%,0%,0%);stroke-opacity:1;stroke-miterlimit:10;\" d=\"M 37.150594 -0.00146875 C 37.150594 4.861813 33.209187 8.803219 28.345906 8.803219 C 23.482625 8.803219 19.541219 4.861813 19.541219 -0.00146875 C 19.541219 -4.86475 23.482625 -8.806156 28.345906 -8.806156 C 33.209187 -8.806156 37.150594 -4.86475 37.150594 -0.00146875 Z M 37.150594 -0.00146875 \" transform=\"matrix(1,0,0,-1,-19.342,9.104)\"/>\n",
       "</g>\n",
       "<g style=\"fill:rgb(0%,0%,0%);fill-opacity:1;\">\n",
       "  <use xlink:href=\"#glyph-1567442442225000-0-1\" x=\"6.513\" y=\"12.421\"/>\n",
       "</g>\n",
       "<path style=\" stroke:none;fill-rule:nonzero;fill:rgb(100%,100%,100%);fill-opacity:1;\" d=\"M 46.15625 9.105469 C 46.15625 4.242188 42.214844 0.300781 37.351563 0.300781 C 32.488281 0.300781 28.546875 4.242188 28.546875 9.105469 C 28.546875 13.96875 32.488281 17.910156 37.351563 17.910156 C 42.214844 17.910156 46.15625 13.96875 46.15625 9.105469 Z M 46.15625 9.105469 \"/>\n",
       "<g clip-path=\"url(#clip-1567442442225000-2)\" clip-rule=\"nonzero\">\n",
       "<path style=\"fill:none;stroke-width:0.3985;stroke-linecap:butt;stroke-linejoin:miter;stroke:rgb(0%,0%,0%);stroke-opacity:1;stroke-miterlimit:10;\" d=\"M 65.49825 -0.00146875 C 65.49825 4.861813 61.556844 8.803219 56.693562 8.803219 C 51.830281 8.803219 47.888875 4.861813 47.888875 -0.00146875 C 47.888875 -4.86475 51.830281 -8.806156 56.693562 -8.806156 C 61.556844 -8.806156 65.49825 -4.86475 65.49825 -0.00146875 Z M 65.49825 -0.00146875 \" transform=\"matrix(1,0,0,-1,-19.342,9.104)\"/>\n",
       "</g>\n",
       "<g style=\"fill:rgb(0%,0%,0%);fill-opacity:1;\">\n",
       "  <use xlink:href=\"#glyph-1567442442225000-0-2\" x=\"34.86\" y=\"12.421\"/>\n",
       "</g>\n",
       "<g clip-path=\"url(#clip-1567442442225000-3)\" clip-rule=\"nonzero\">\n",
       "<path style=\" stroke:none;fill-rule:nonzero;fill:rgb(100%,100%,100%);fill-opacity:1;\" d=\"M 74.601563 9.105469 C 74.601563 4.1875 70.617188 0.199219 65.699219 0.199219 C 60.78125 0.199219 56.792969 4.1875 56.792969 9.105469 C 56.792969 14.023438 60.78125 18.007813 65.699219 18.007813 C 70.617188 18.007813 74.601563 14.023438 74.601563 9.105469 Z M 74.601563 9.105469 \"/>\n",
       "</g>\n",
       "<g clip-path=\"url(#clip-1567442442225000-4)\" clip-rule=\"nonzero\">\n",
       "<path style=\"fill:none;stroke-width:0.3985;stroke-linecap:butt;stroke-linejoin:miter;stroke:rgb(0%,0%,0%);stroke-opacity:1;stroke-miterlimit:10;\" d=\"M 93.943562 -0.00146875 C 93.943562 4.9165 89.959187 8.904781 85.041219 8.904781 C 80.12325 8.904781 76.134969 4.9165 76.134969 -0.00146875 C 76.134969 -4.919437 80.12325 -8.903812 85.041219 -8.903812 C 89.959187 -8.903812 93.943562 -4.919437 93.943562 -0.00146875 Z M 93.943562 -0.00146875 \" transform=\"matrix(1,0,0,-1,-19.342,9.104)\"/>\n",
       "</g>\n",
       "<g style=\"fill:rgb(0%,0%,0%);fill-opacity:1;\">\n",
       "  <use xlink:href=\"#glyph-1567442442225000-0-3\" x=\"63.206\" y=\"12.312\"/>\n",
       "</g>\n",
       "<path style=\" stroke:none;fill-rule:nonzero;fill:rgb(0%,100%,0%);fill-opacity:1;\" d=\"M 102.878906 9.105469 C 102.878906 4.226563 98.921875 0.273438 94.046875 0.273438 C 89.167969 0.273438 85.214844 4.226563 85.214844 9.105469 C 85.214844 13.980469 89.167969 17.9375 94.046875 17.9375 C 98.921875 17.9375 102.878906 13.980469 102.878906 9.105469 Z M 102.878906 9.105469 \"/>\n",
       "<g clip-path=\"url(#clip-1567442442225000-5)\" clip-rule=\"nonzero\">\n",
       "<path style=\"fill:none;stroke-width:0.3985;stroke-linecap:butt;stroke-linejoin:miter;stroke:rgb(0%,0%,0%);stroke-opacity:1;stroke-miterlimit:10;\" d=\"M 122.220906 -0.00146875 C 122.220906 4.877438 118.263875 8.830563 113.388875 8.830563 C 108.509969 8.830563 104.556844 4.877438 104.556844 -0.00146875 C 104.556844 -4.876469 108.509969 -8.8335 113.388875 -8.8335 C 118.263875 -8.8335 122.220906 -4.876469 122.220906 -0.00146875 Z M 122.220906 -0.00146875 \" transform=\"matrix(1,0,0,-1,-19.342,9.104)\"/>\n",
       "</g>\n",
       "<g style=\"fill:rgb(0%,0%,0%);fill-opacity:1;\">\n",
       "  <use xlink:href=\"#glyph-1567442442225000-0-4\" x=\"91.553\" y=\"12.476\"/>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n",
       "\n"
      ],
      "text/plain": [
       "TikzPicture(\"\\\\draw (1cm, 0cm) node[draw=black,circle,fill=white]{1} -- (2cm, 0cm) node[draw=black,circle,fill=white]{2} -- (3cm, 0cm) node[draw=black,circle,fill=white]{3} -- (4cm, 0cm) node[draw=black,circle,fill=green]{4};\", \"\", \"\", true)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "plot_chain(4, fill=Dict(4=>\"green\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, create a new type that represents the MDP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "struct ChainMDP <: MDP{Int, Symbol}\n",
    "    len::Int\n",
    "    p_success::Float64\n",
    "    discount::Float64\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We define the behavior of the MDP by imlementing functions from the POMDPs.jl API (http://juliapomdp.github.io/POMDPs.jl/latest/api/). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "search: \u001b[0m\u001b[1mS\u001b[22m\u001b[0m\u001b[1mp\u001b[22m\u001b[0m\u001b[1ma\u001b[22m\u001b[0m\u001b[1mr\u001b[22m\u001b[0m\u001b[1ms\u001b[22m\u001b[0m\u001b[1me\u001b[22m\u001b[0m\u001b[1mC\u001b[22m\u001b[0m\u001b[1ma\u001b[22m\u001b[0m\u001b[1mt\u001b[22m\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/latex": [
       "\\begin{verbatim}\n",
       "SparseCat(values, probabilities)\n",
       "\\end{verbatim}\n",
       "Create a sparse categorical distribution.\n",
       "\n",
       "\\texttt{values} is an iterable object containing the possible values (can be of any type) in the distribution that have nonzero probability. \\texttt{probabilities} is an iterable object that contains the associated probabilities.\n",
       "\n",
       "This is optimized for value iteration with a fast implementation of \\texttt{weighted\\_iterator}. Both \\texttt{pdf} and \\texttt{rand} are order n.\n",
       "\n"
      ],
      "text/markdown": [
       "```\n",
       "SparseCat(values, probabilities)\n",
       "```\n",
       "\n",
       "Create a sparse categorical distribution.\n",
       "\n",
       "`values` is an iterable object containing the possible values (can be of any type) in the distribution that have nonzero probability. `probabilities` is an iterable object that contains the associated probabilities.\n",
       "\n",
       "This is optimized for value iteration with a fast implementation of `weighted_iterator`. Both `pdf` and `rand` are order n.\n"
      ],
      "text/plain": [
       "\u001b[36m  SparseCat(values, probabilities)\u001b[39m\n",
       "\n",
       "  Create a sparse categorical distribution.\n",
       "\n",
       "  \u001b[36mvalues\u001b[39m is an iterable object containing the possible values (can be of any\n",
       "  type) in the distribution that have nonzero probability. \u001b[36mprobabilities\u001b[39m is an\n",
       "  iterable object that contains the associated probabilities.\n",
       "\n",
       "  This is optimized for value iteration with a fast implementation of\n",
       "  \u001b[36mweighted_iterator\u001b[39m. Both \u001b[36mpdf\u001b[39m and \u001b[36mrand\u001b[39m are order n."
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "?SparseCat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# the transition function returns the distribution of sp given that\n",
    "# action a is taken in state s\n",
    "function POMDPs.transition(mdp::ChainMDP, s::Int, a::Symbol)\n",
    "    if a == :right\n",
    "        success = min(s+1, mdp.len)\n",
    "        failure = max(s-1, 1)\n",
    "    else # a == :left\n",
    "        success = max(s-1, 1)\n",
    "        failure = min(s+1, mdp.len)\n",
    "    end\n",
    "    return SparseCat([success, failure], [mdp.p_success, 1.0-mdp.p_success])\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note: this allocates arrays in the return statement - it could be sped up by using [StaticArrays.jl](https://github.com/JuliaArrays/StaticArrays.jl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "function POMDPs.reward(mdp::ChainMDP, s::Int, a::Symbol)\n",
    "    if s == mdp.len\n",
    "        return 1.0\n",
    "    else \n",
    "        return 0.0\n",
    "    end\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "POMDPs.discount(mdp::ChainMDP) = mdp.discount"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "POMDPs.stateindex(::ChainMDP, s::Int64) = s"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And that's all we need to define the MDP (for simulation at least)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ChainMDP(4, 0.7, 0.9)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mdp = ChainMDP(4, 0.7, 0.9)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\u001b[1m              SparseCat distribution\u001b[22m\n",
       "\u001b[90m     ┌                                        ┐\u001b[39m \n",
       "   \u001b[0m2\u001b[90m ┤\u001b[39m\u001b[32m■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■\u001b[39m\u001b[0m 0.7 \u001b[90m \u001b[39m \n",
       "   \u001b[0m1\u001b[90m ┤\u001b[39m\u001b[32m■■■■■■■■■■■■■■■\u001b[39m\u001b[0m 0.30000000000000004     \u001b[90m \u001b[39m \n",
       "\u001b[90m     └                                        ┘\u001b[39m "
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "transition(mdp, 1, :right)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.0"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "reward(mdp, 1, :right)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.0"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "reward(mdp, 4, :left)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Simulations\n",
    "\n",
    "There are three basic ways to simulate a POMDP."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "\n",
    "\n",
    "If we want to try out simple policies we can construct a `Policy` type and implement the `POMDPs.action` function which defines the policy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SimplePolicy()"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mutable struct SimplePolicy <: Policy end\n",
    "POMDPs.action(::SimplePolicy, s::Int64) = (s == 2 ? :left : :right)\n",
    "p = SimplePolicy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[32mSimulating...100%|██████████████████████████████████████| Time: 0:00:01\u001b[39m\n"
     ]
    }
   ],
   "source": [
    "hr = HistoryRecorder(max_steps=10, show_progress=true)\n",
    "history = simulate(hr, mdp, p, 1);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "State:  1    Action:  right   Reward: 0\n",
      "State:  1    Action:  right   Reward: 0\n",
      "State:  2    Action:   left   Reward: 0\n",
      "State:  1    Action:  right   Reward: 0\n",
      "State:  2    Action:   left   Reward: 0\n",
      "State:  1    Action:  right   Reward: 0\n",
      "State:  1    Action:  right   Reward: 0\n",
      "State:  1    Action:  right   Reward: 0\n",
      "State:  2    Action:   left   Reward: 0\n",
      "State:  1    Action:  right   Reward: 0\n"
     ]
    }
   ],
   "source": [
    "for (s, a, r) in eachstep(history, \"s,a,r\")\n",
    "    @printf(\"State: %2d    Action: %6s   Reward: %1d\\n\", s, a, r)\n",
    "end    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also specify a policy object, for example using a vector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "policy_vector = fill(:right, 4) = Symbol[:right, :right, :right, :right]\n"
     ]
    }
   ],
   "source": [
    "@show policy_vector = fill(:right, 4);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "policy = VectorPolicy(mdp, policy_vector);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "scrolled": true
   },
   "source": [
    "and step through a simulation as it is being performed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "State:  1    Action:  right   Reward: 0\n",
      "State:  2    Action:  right   Reward: 0\n",
      "State:  1    Action:  right   Reward: 0\n",
      "State:  2    Action:  right   Reward: 0\n",
      "State:  3    Action:  right   Reward: 0\n",
      "State:  4    Action:  right   Reward: 1\n",
      "State:  3    Action:  right   Reward: 0\n",
      "State:  4    Action:  right   Reward: 1\n",
      "State:  4    Action:  right   Reward: 1\n",
      "State:  4    Action:  right   Reward: 1\n"
     ]
    }
   ],
   "source": [
    "for (s, a, r) in stepthrough(mdp, policy, 1, \"s,a,r\", max_steps=10)\n",
    "    @printf(\"State: %2d    Action: %6s   Reward: %1d\\n\", s, a, r)\n",
    "    sleep(1.0)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "And also use a simulator object and the `simulate()` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "state_hist(history) = [1, 2, 3, 2, 3, 4, 3, 4, 3, 4, 4]\n"
     ]
    }
   ],
   "source": [
    "simulator = HistoryRecorder(max_steps=10, show_progress=true)\n",
    "history = simulate(simulator, mdp, policy, 1)\n",
    "@show state_hist(history);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Solving with MCTS\n",
    "\n",
    "\n",
    "The required API functions depend on the solver. Suppose we want to use the MCTS solver from MCTS.jl:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Pkg.add(\"MCTS\")\n",
    "using MCTS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO: POMDPs.jl requirements for \u001b[34maction(::AbstractMCTSPlanner, ::Any)\u001b[39m and dependencies. ([✔] = implemented correctly; [X] = missing)\n",
      "\n",
      "For \u001b[34maction(::AbstractMCTSPlanner, ::Any)\u001b[39m:\n",
      "  [No additional requirements]\n",
      "For \u001b[34msimulate(::AbstractMCTSPlanner, ::Any, ::Int64)\u001b[39m (in action(::AbstractMCTSPlanner, ::Any)):\n",
      "\u001b[32m  [✔] discount(::ChainMDP)\u001b[39m\n",
      "\u001b[32m  [✔] isterminal(::ChainMDP, ::Int64)\u001b[39m\n",
      "\u001b[32m  [✔] generate_sr(::ChainMDP, ::Int64, ::Symbol, ::MersenneTwister)\u001b[39m\n",
      "\u001b[32m  [✔] isequal(::Int64, ::Int64)\u001b[39m\n",
      "\u001b[32m  [✔] hash(::Int64)\u001b[39m\n",
      "For \u001b[34minsert_node!(::MCTSPlanner, ::Int64)\u001b[39m (in simulate(::AbstractMCTSPlanner, ::Any, ::Int64)):\n",
      "  [No requirements specified]\n",
      "For \u001b[34mestimate_value(::MCTS.SolvedRolloutEstimator, ::MDP, ::Any, ::Int64)\u001b[39m (in simulate(::AbstractMCTSPlanner, ::Any, ::Int64)):\n",
      "  [No additional requirements]\n",
      "For \u001b[34mrollout(::MCTS.SolvedRolloutEstimator, ::MDP, ::Any, ::Int64)\u001b[39m (in estimate_value(::MCTS.SolvedRolloutEstimator, ::MDP, ::Any, ::Int64)):\n",
      "  [No additional requirements]\n",
      "For \u001b[34msimulate(::RolloutSimulator, ::MDP, ::Policy, ::Any)\u001b[39m (in rollout(::MCTS.SolvedRolloutEstimator, ::MDP, ::Any, ::Int64)):\n",
      "\u001b[32m  [✔] action(::RandomPolicy, ::Int64)\u001b[39m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "true"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "solver = MCTSSolver()\n",
    "mdp = ChainMDP(4, 0.8, 0.95)\n",
    "\n",
    "@requirements_info solver mdp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "search: \u001b[0m\u001b[1ma\u001b[22m\u001b[0m\u001b[1mc\u001b[22m\u001b[0m\u001b[1mt\u001b[22m\u001b[0m\u001b[1mi\u001b[22m\u001b[0m\u001b[1mo\u001b[22m\u001b[0m\u001b[1mn\u001b[22m\u001b[0m\u001b[1ms\u001b[22m n_\u001b[0m\u001b[1ma\u001b[22m\u001b[0m\u001b[1mc\u001b[22m\u001b[0m\u001b[1mt\u001b[22m\u001b[0m\u001b[1mi\u001b[22m\u001b[0m\u001b[1mo\u001b[22m\u001b[0m\u001b[1mn\u001b[22m\u001b[0m\u001b[1ms\u001b[22m \u001b[0m\u001b[1ma\u001b[22m\u001b[0m\u001b[1mc\u001b[22m\u001b[0m\u001b[1mt\u001b[22m\u001b[0m\u001b[1mi\u001b[22m\u001b[0m\u001b[1mo\u001b[22m\u001b[0m\u001b[1mn\u001b[22m_hi\u001b[0m\u001b[1ms\u001b[22mt \u001b[0m\u001b[1ma\u001b[22m\u001b[0m\u001b[1mc\u001b[22m\u001b[0m\u001b[1mt\u001b[22m\u001b[0m\u001b[1mi\u001b[22m\u001b[0m\u001b[1mo\u001b[22m\u001b[0m\u001b[1mn\u001b[22mvalue\u001b[0m\u001b[1ms\u001b[22m ordered_\u001b[0m\u001b[1ma\u001b[22m\u001b[0m\u001b[1mc\u001b[22m\u001b[0m\u001b[1mt\u001b[22m\u001b[0m\u001b[1mi\u001b[22m\u001b[0m\u001b[1mo\u001b[22m\u001b[0m\u001b[1mn\u001b[22m\u001b[0m\u001b[1ms\u001b[22m\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/latex": [
       "\\begin{verbatim}\n",
       "actions(problem::POMDP)\n",
       "actions(problem::MDP)\n",
       "\\end{verbatim}\n",
       "Returns the entire action space of a POMDP.\n",
       "\n",
       "\\rule{\\textwidth}{1pt}\n",
       "\\begin{verbatim}\n",
       "actions{S,A,O}(problem::POMDP{S,A,O}, state::S)\n",
       "actions{S,A}(problem::MDP{S,A}, state::S)\n",
       "\\end{verbatim}\n",
       "Return the action space accessible from the given state.\n",
       "\n",
       "\\rule{\\textwidth}{1pt}\n",
       "\\begin{verbatim}\n",
       "actions(problem::POMDP, belief)\n",
       "\\end{verbatim}\n",
       "Return the action space accessible from the states with nonzero belief.\n",
       "\n"
      ],
      "text/markdown": [
       "```\n",
       "actions(problem::POMDP)\n",
       "actions(problem::MDP)\n",
       "```\n",
       "\n",
       "Returns the entire action space of a POMDP.\n",
       "\n",
       "---\n",
       "\n",
       "```\n",
       "actions{S,A,O}(problem::POMDP{S,A,O}, state::S)\n",
       "actions{S,A}(problem::MDP{S,A}, state::S)\n",
       "```\n",
       "\n",
       "Return the action space accessible from the given state.\n",
       "\n",
       "---\n",
       "\n",
       "```\n",
       "actions(problem::POMDP, belief)\n",
       "```\n",
       "\n",
       "Return the action space accessible from the states with nonzero belief.\n"
      ],
      "text/plain": [
       "\u001b[36m  actions(problem::POMDP)\u001b[39m\n",
       "\u001b[36m  actions(problem::MDP)\u001b[39m\n",
       "\n",
       "  Returns the entire action space of a POMDP.\n",
       "\n",
       "  ────────────────────────────────────────────────────────────────────────────\n",
       "\n",
       "\u001b[36m  actions{S,A,O}(problem::POMDP{S,A,O}, state::S)\u001b[39m\n",
       "\u001b[36m  actions{S,A}(problem::MDP{S,A}, state::S)\u001b[39m\n",
       "\n",
       "  Return the action space accessible from the given state.\n",
       "\n",
       "  ────────────────────────────────────────────────────────────────────────────\n",
       "\n",
       "\u001b[36m  actions(problem::POMDP, belief)\u001b[39m\n",
       "\n",
       "  Return the action space accessible from the states with nonzero belief."
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "?actions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "POMDPs.actions(::ChainMDP) = [:left, :right]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO: POMDPs.jl requirements for \u001b[34maction(::AbstractMCTSPlanner, ::Any)\u001b[39m and dependencies. ([✔] = implemented correctly; [X] = missing)\n",
      "\n",
      "For \u001b[34maction(::AbstractMCTSPlanner, ::Any)\u001b[39m:\n",
      "  [No additional requirements]\n",
      "For \u001b[34msimulate(::AbstractMCTSPlanner, ::Any, ::Int64)\u001b[39m (in action(::AbstractMCTSPlanner, ::Any)):\n",
      "\u001b[32m  [✔] discount(::ChainMDP)\u001b[39m\n",
      "\u001b[32m  [✔] isterminal(::ChainMDP, ::Int64)\u001b[39m\n",
      "\u001b[32m  [✔] generate_sr(::ChainMDP, ::Int64, ::Symbol, ::MersenneTwister)\u001b[39m\n",
      "\u001b[32m  [✔] isequal(::Int64, ::Int64)\u001b[39m\n",
      "\u001b[32m  [✔] hash(::Int64)\u001b[39m\n",
      "For \u001b[34minsert_node!(::MCTSPlanner, ::Int64)\u001b[39m (in simulate(::AbstractMCTSPlanner, ::Any, ::Int64)):\n",
      "  [No requirements specified]\n",
      "For \u001b[34mestimate_value(::MCTS.SolvedRolloutEstimator, ::MDP, ::Any, ::Int64)\u001b[39m (in simulate(::AbstractMCTSPlanner, ::Any, ::Int64)):\n",
      "  [No additional requirements]\n",
      "For \u001b[34mrollout(::MCTS.SolvedRolloutEstimator, ::MDP, ::Any, ::Int64)\u001b[39m (in estimate_value(::MCTS.SolvedRolloutEstimator, ::MDP, ::Any, ::Int64)):\n",
      "  [No additional requirements]\n",
      "For \u001b[34msimulate(::RolloutSimulator, ::MDP, ::Policy, ::Any)\u001b[39m (in rollout(::MCTS.SolvedRolloutEstimator, ::MDP, ::Any, ::Int64)):\n",
      "\u001b[32m  [✔] action(::RandomPolicy, ::Int64)\u001b[39m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "true"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@requirements_info solver mdp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "solver = MCTSSolver(enable_tree_vis=true)\n",
    "planner = solve(solver, mdp);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "MDPHistory{Int64,Symbol}([1, 2, 3, 4, 4, 3, 4, 4, 3, 2, 3], Symbol[:right, :right, :right, :right, :right, :right, :right, :right, :right, :right], [0.0, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0, 1.0, 0.0, 0.0], Any[nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing], Any[nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing], 0.95, nothing, nothing)"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "simulate(simulator, mdp, planner, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "    <!DOCTYPE html>\n",
       "    <html>\n",
       "    <head>\n",
       "        <meta charset=\"UTF-8\">\n",
       "        <title>MCTS tree</title>\n",
       "    </head>\n",
       "    <body>\n",
       "    <div id=\"treevisSqgwcYO9\">\n",
       "    <style>\n",
       "        \n",
       ".node {\n",
       "    cursor: pointer;\n",
       "}\n",
       "\n",
       ".node circle {\n",
       "\t  fill: #fff;\n",
       "\t  stroke: steelblue;\n",
       "\t  stroke-width: 3px;\n",
       "      r: 10;\n",
       "}\n",
       "\n",
       ".node rect {\n",
       "\t  fill: #fff;\n",
       "\t  stroke: steelblue;\n",
       "\t  stroke-width: 3px;\n",
       "      x: -10;\n",
       "      y: -10;\n",
       "      width: 20;\n",
       "      height: 20;\n",
       "}\n",
       "\n",
       ".node path {\n",
       "    fill: #ffffff;\n",
       "    /* stroke: steelblue; */\n",
       "    stroke-width: 2px;\n",
       "}\n",
       ".node text {\n",
       "    font-family: Arial, Helvetica, sans-serif;\n",
       "    font-style: normal;\n",
       "    font-size: 12px;\n",
       "}\n",
       ".link /* path */ {\n",
       "    fill: none;\n",
       "    stroke: #cccccc;\n",
       "    stroke-width: 2px;\n",
       "}\n",
       ".link text {\n",
       "    font-family: Arial, Helvetica, sans-serif;\n",
       "    font-style: normal;\n",
       "    font-size: 12px;\n",
       "}\n",
       ".hover text{\n",
       "    font-family: Arial, Helvetica, sans-serif;\n",
       "    font-style: normal;\n",
       "    font-size: 12px;\n",
       "}\n",
       "\n",
       "/*\n",
       ".svg-container {\n",
       "    display: inline-block;\n",
       "    position: relative;\n",
       "    width: 100%;\n",
       "    padding-bottom: 100%;\n",
       "    vertical-align: top;\n",
       "    overflow: hidden;\n",
       "}\n",
       ".svg-content-responsive {\n",
       "    display: inline-block;\n",
       "    position: absolute;\n",
       "    top: 10px;\n",
       "    left: 0;\n",
       "}\n",
       "*/\n",
       "\n",
       "    </style>\n",
       "    <script>\n",
       "       (function(){\n",
       "        var treeData = {\"children\":[[19,20],[17,18],[23,24],[17,18],[17,18],[21,22],[19,20],[23,24],[19,20],[17,18],[21,22],[23,24],[21,22],[23,24],[19,20],[19,20],[9],[7,11],[10,15],[5,16],[3,4],[2,8],[13,14],[6,12]],\"text\":[\"                        1\\nN:     53\\n\",\"                        2\\nN:    179\\n\",\"                        4\\nN:    511\\n\",\"                        2\\nN:    179\\n\",\"                        2\\nN:    179\\n\",\"                        3\\nN:    231\\n\",\"                        1\\nN:     53\\n\",\"                        4\\nN:    511\\n\",\"                        1\\nN:     53\\n\",\"                        2\\nN:    179\\n\",\"                        3\\nN:    231\\n\",\"                        4\\nN:    511\\n\",\"                        3\\nN:    231\\n\",\"                        4\\nN:    511\\n\",\"                        1\\nN:     53\\n\",\"                        1\\nN:     53\\n\",\"                     left\\nQ:   0.32\\nN:      4\\n\",\"                    right\\nQ:   3.01\\nN:    175\\n\",\"                     left\\nQ:   0.00\\nN:      7\\n\",\"                    right\\nQ:   1.98\\nN:     46\\n\",\"                     left\\nQ:   0.86\\nN:      8\\n\",\"                    right\\nQ:   2.88\\nN:    223\\n\",\"                     left\\nQ:   2.57\\nN:     29\\n\",\"                    right\\nQ:   3.26\\nN:    482\\n\"],\"tooltip\":[\"1\\nN: 53\\n\",\"2\\nN: 179\\n\",\"4\\nN: 511\\n\",\"2\\nN: 179\\n\",\"2\\nN: 179\\n\",\"3\\nN: 231\\n\",\"1\\nN: 53\\n\",\"4\\nN: 511\\n\",\"1\\nN: 53\\n\",\"2\\nN: 179\\n\",\"3\\nN: 231\\n\",\"4\\nN: 511\\n\",\"3\\nN: 231\\n\",\"4\\nN: 511\\n\",\"1\\nN: 53\\n\",\"1\\nN: 53\\n\",\"left\\nQ: 0.3234174602534179\\nN: 4\\n\",\"right\\nQ: 3.014329455610871\\nN: 175\\n\",\"left\\nQ: 0.0\\nN: 7\\n\",\"right\\nQ: 1.9784042519327438\\nN: 46\\n\",\"left\\nQ: 0.855078634765625\\nN: 8\\n\",\"right\\nQ: 2.877043660350687\\nN: 223\\n\",\"left\\nQ: 2.5664803894665944\\nN: 29\\n\",\"right\\nQ: 3.256413469722508\\nN: 482\\n\"],\"style\":[\"\",\"\",\"\",\"\",\"\",\"\",\"\",\"\",\"\",\"\",\"\",\"\",\"\",\"\",\"\",\"\",\"stroke:#E60D00\",\"stroke:#137600\",\"stroke:#FF0000\",\"stroke:#644E00\",\"stroke:#BC2200\",\"stroke:#1E7100\",\"stroke:#366500\",\"stroke:#008000\"],\"link_style\":[\"stroke-width:20.0px\",\"stroke-width:8.782371253997939px\",\"stroke-width:12.24744871391589px\",\"stroke-width:15.811388300841898px\",\"stroke-width:17.693034738587656px\",\"stroke-width:8.87908634229102px\",\"stroke-width:9.319717960171479px\",\"stroke-width:17.968582447064392px\",\"stroke-width:20.0px\",\"stroke-width:10.690449676496975px\",\"stroke-width:17.695842933945165px\",\"stroke-width:17.920988413760583px\",\"stroke-width:16.609095970747994px\",\"stroke-width:11.141720290623113px\",\"stroke-width:16.903085094570333px\",\"stroke-width:9.325048082403137px\",\"stroke-width:2.9897403710077435px\",\"stroke-width:19.77527376584118px\",\"stroke-width:7.26843784311631px\",\"stroke-width:18.632493427363787px\",\"stroke-width:3.721936841593884px\",\"stroke-width:19.650628136199263px\",\"stroke-width:4.764513310312894px\",\"stroke-width:19.42419658353599px\"],\"title\":\"MCTS tree\",\"options\":{\"tooltip\":[\"1\\nN: 53\\n\",\"2\\nN: 179\\n\",\"4\\nN: 511\\n\",\"2\\nN: 179\\n\",\"2\\nN: 179\\n\",\"3\\nN: 231\\n\",\"1\\nN: 53\\n\",\"4\\nN: 511\\n\",\"1\\nN: 53\\n\",\"2\\nN: 179\\n\",\"3\\nN: 231\\n\",\"4\\nN: 511\\n\",\"3\\nN: 231\\n\",\"4\\nN: 511\\n\",\"1\\nN: 53\\n\",\"1\\nN: 53\\n\",\"left\\nQ: 0.3234174602534179\\nN: 4\\n\",\"right\\nQ: 3.014329455610871\\nN: 175\\n\",\"left\\nQ: 0.0\\nN: 7\\n\",\"right\\nQ: 1.9784042519327438\\nN: 46\\n\",\"left\\nQ: 0.855078634765625\\nN: 8\\n\",\"right\\nQ: 2.877043660350687\\nN: 223\\n\",\"left\\nQ: 2.5664803894665944\\nN: 29\\n\",\"right\\nQ: 3.256413469722508\\nN: 482\\n\"],\"link_style\":[\"stroke-width:20.0px\",\"stroke-width:8.782371253997939px\",\"stroke-width:12.24744871391589px\",\"stroke-width:15.811388300841898px\",\"stroke-width:17.693034738587656px\",\"stroke-width:8.87908634229102px\",\"stroke-width:9.319717960171479px\",\"stroke-width:17.968582447064392px\",\"stroke-width:20.0px\",\"stroke-width:10.690449676496975px\",\"stroke-width:17.695842933945165px\",\"stroke-width:17.920988413760583px\",\"stroke-width:16.609095970747994px\",\"stroke-width:11.141720290623113px\",\"stroke-width:16.903085094570333px\",\"stroke-width:9.325048082403137px\",\"stroke-width:2.9897403710077435px\",\"stroke-width:19.77527376584118px\",\"stroke-width:7.26843784311631px\",\"stroke-width:18.632493427363787px\",\"stroke-width:3.721936841593884px\",\"stroke-width:19.650628136199263px\",\"stroke-width:4.764513310312894px\",\"stroke-width:19.42419658353599px\"],\"title\":\"MCTS tree\",\"text\":[\"                        1\\nN:     53\\n\",\"                        2\\nN:    179\\n\",\"                        4\\nN:    511\\n\",\"                        2\\nN:    179\\n\",\"                        2\\nN:    179\\n\",\"                        3\\nN:    231\\n\",\"                        1\\nN:     53\\n\",\"                        4\\nN:    511\\n\",\"                        1\\nN:     53\\n\",\"                        2\\nN:    179\\n\",\"                        3\\nN:    231\\n\",\"                        4\\nN:    511\\n\",\"                        3\\nN:    231\\n\",\"                        4\\nN:    511\\n\",\"                        1\\nN:     53\\n\",\"                        1\\nN:     53\\n\",\"                     left\\nQ:   0.32\\nN:      4\\n\",\"                    right\\nQ:   3.01\\nN:    175\\n\",\"                     left\\nQ:   0.00\\nN:      7\\n\",\"                    right\\nQ:   1.98\\nN:     46\\n\",\"                     left\\nQ:   0.86\\nN:      8\\n\",\"                    right\\nQ:   2.88\\nN:    223\\n\",\"                     left\\nQ:   2.57\\nN:     29\\n\",\"                    right\\nQ:   3.26\\nN:    482\\n\"],\"style\":[\"\",\"\",\"\",\"\",\"\",\"\",\"\",\"\",\"\",\"\",\"\",\"\",\"\",\"\",\"\",\"\",\"stroke:#E60D00\",\"stroke:#137600\",\"stroke:#FF0000\",\"stroke:#644E00\",\"stroke:#BC2200\",\"stroke:#1E7100\",\"stroke:#366500\",\"stroke:#008000\"]}};\n",
       "        var rootID = 1-1;\n",
       "        var div = \"treevisSqgwcYO9\";\n",
       "        var initExpand = 0\n",
       "        var initDuration = 750\n",
       "        var svgHeight = 600\n",
       "        if (typeof $ === 'undefined') {\n",
       "    loadScript(\"https://code.jquery.com/jquery-3.1.1.min.js\", run);\n",
       "} else {\n",
       "    run();\n",
       "}\n",
       "\n",
       "function run() {\n",
       "    if (typeof d3 === 'undefined') {\n",
       "        loadScript(\"https://d3js.org/d3.v3.js\", showTree);\n",
       "    } else {\n",
       "        showTree();\n",
       "    }\n",
       "}\n",
       "\n",
       "function loadScript(url, callback)\n",
       "{\n",
       "    console.log(\"starting script load...\")\n",
       "    // Adding the script tag to the head as suggested before\n",
       "    var head = document.getElementsByTagName('head')[0];\n",
       "    var script = document.createElement('script');\n",
       "    script.type = 'text/javascript';\n",
       "    script.src = url;\n",
       "\n",
       "    // Then bind the event to the callback function.\n",
       "    // There are several events for cross browser compatibility.\n",
       "    script.onreadystatechange = callback;\n",
       "    script.onload = callback;\n",
       "\n",
       "    // Fire the loading\n",
       "    head.appendChild(script);\n",
       "}\n",
       "\n",
       "function showTree() {\n",
       "        \n",
       "    // var margin = {top: 20, right: 120, bottom: 20, left: 120},\n",
       "    var margin = {top: 20, right: 120, bottom: 80, left: 120},\n",
       "        width = $(\"#\"+div).width() - margin.right - margin.left,\n",
       "        height = svgHeight - margin.top - margin.bottom;\n",
       "        // TODO make height a parameter of TreeVisualizer\n",
       "\n",
       "    var i = 0,\n",
       "        root;\n",
       "\n",
       "    var tree = d3.layout.tree()\n",
       "        .size([width, height]);\n",
       "\n",
       "    var diagonal = d3.svg.diagonal();\n",
       "        //.projection(function(d) { return [d.y, d.x]; });\n",
       "        // uncomment above to make the tree go horizontally\n",
       "\n",
       "    // see http://stackoverflow.com/questions/16265123/resize-svg-when-window-is-resized-in-d3-js\n",
       "    if (d3.select(\"#\"+div+\"_svg\").empty()) {\n",
       "        $(\".d3twarn\").remove();\n",
       "        d3.select(\"#\"+div).append(\"svg\")\n",
       "            .attr(\"id\", div+\"_svg\")\n",
       "            .attr(\"width\", width + margin.right + margin.left)\n",
       "            .attr(\"height\", height + margin.top + margin.bottom);\n",
       "    }\n",
       "\n",
       "    d3.select(\"#\"+div+\"_svg\").selectAll(\"*\").remove();\n",
       "\n",
       "    var svg = d3.select(\"#\"+div+\"_svg\")\n",
       "        .append(\"g\")\n",
       "            .attr(\"transform\", \"translate(\" + margin.left + \",\" + margin.top + \")\");\n",
       "\n",
       "    // console.log(\"tree data:\");\n",
       "    // console.log(treeData[rootID]);\n",
       "    root = createDisplayNode(rootID, initExpand);\n",
       "    root.x0 = width / 2;\n",
       "    root.y0 = 0;\n",
       "\n",
       "    update(root, initDuration);\n",
       "    console.log(\"tree should appear\");\n",
       "\n",
       "    function createDisplayNode(id, expandLevel) {\n",
       "      var dnode = {\"dataID\":id,\n",
       "                   \"children\": null,\n",
       "                   \"_children\":null};\n",
       "      if (expandLevel > 0) {\n",
       "        initializeChildren(dnode, expandLevel);\n",
       "      }\n",
       "      return dnode;\n",
       "    }\n",
       "\n",
       "    function initializeChildren(d, expandLevel) {\n",
       "      // create children\n",
       "      var children = treeData.children[d.dataID];\n",
       "      d.children = [];\n",
       "      if (children) {\n",
       "        for (var i = 0; i < children.length; i++) {\n",
       "          var cid = children[i]-1;\n",
       "          d.children.push(createDisplayNode(cid, expandLevel-1));\n",
       "        }\n",
       "      }\n",
       "    }\n",
       "\n",
       "    /**\n",
       "     * Recursively called to update each node in the tree.\n",
       "     * \n",
       "     * source is a d3 node that has position, etc.\n",
       "    */\n",
       "    function update(source, duration) {\n",
       "\n",
       "      width = $(\"#\"+div).width() - margin.right - margin.left,\n",
       "      height = $(\"#\"+div).height() - margin.top - margin.bottom;\n",
       "\n",
       "      tree.size([width,height]);\n",
       "      d3.select(\"#\"+div).attr(\"width\", width + margin.right + margin.left)\n",
       "            .attr(\"height\", height + margin.top + margin.bottom);\n",
       "      d3.select(\"#\"+div+\"_svg\").attr(\"width\", width + margin.right + margin.left)\n",
       "             .attr(\"height\", height + margin.top + margin.bottom);\n",
       "\n",
       "\n",
       "      // Compute the new tree layout.\n",
       "      var nodes = tree.nodes(root).reverse(),\n",
       "          links = tree.links(nodes);\n",
       "\n",
       "\n",
       "      // Update the nodes…\n",
       "      var node = svg.selectAll(\"g.node\")\n",
       "          .data(nodes, function(d) { return d.id || (d.id = ++i); });\n",
       "\n",
       "      // Enter any new nodes at the parent's previous position.\n",
       "      var nodeEnter = node.enter().append(\"g\")\n",
       "          .attr(\"class\", \"node\")\n",
       "          .attr(\"transform\", function(d) { return \"translate(\" + source.x0 + \",\" + source.y0 + \")\"; })\n",
       "          .on(\"click\", click)\n",
       "\n",
       "      nodeEnter.append(\"circle\")\n",
       "          .attr(\"r\", \"10px\")\n",
       "          .attr(\"style\", function(d) { return treeData.style[d.dataID]; } )\n",
       "\n",
       "      var tbox = nodeEnter.append(\"text\")\n",
       "          .attr(\"y\", 25)\n",
       "          .attr(\"text-anchor\", \"middle\")\n",
       "          //.text( function(d) { return treeData.text[d.dataID]; } )\n",
       "          .style(\"fill-opacity\", 1e-6);\n",
       "\n",
       "      tbox.each( function(d) {\n",
       "          var el = d3.select(this)\n",
       "          var text = treeData.text[d.dataID];\n",
       "          var lines = text.split('\\n');\n",
       "          for (var i = 0; i < lines.length; i++) {\n",
       "              var tspan = el.append(\"tspan\").text(lines[i]);\n",
       "              if (i > 0) {\n",
       "                  tspan.attr(\"x\", 0).attr(\"dy\", \"1.2em\");\n",
       "              }\n",
       "          }\n",
       "      });\n",
       "\n",
       "      //tooltip\n",
       "      nodeEnter.append(\"title\").text( function(d) { return treeData.tooltip[d.dataID];} );\n",
       "\n",
       "      // Transition nodes to their new position.\n",
       "      var nodeUpdate = node.transition()\n",
       "          .duration(duration)\n",
       "          .attr(\"transform\", function(d) { return \"translate(\" + d.x + \",\" + d.y + \")\"; });\n",
       "\n",
       "      nodeUpdate.select(\"text\")\n",
       "          .style(\"fill-opacity\", 1);\n",
       "\n",
       "      // Transition exiting nodes to the parent's new position.\n",
       "      var nodeExit = node.exit().transition()\n",
       "          .duration(duration)\n",
       "          .attr(\"transform\", function(d) { return \"translate(\" + source.x + \",\" + source.y + \")\"; })\n",
       "          .remove();\n",
       "\n",
       "      nodeExit.select(\"text\")\n",
       "          .style(\"fill-opacity\", 1e-6);\n",
       "\n",
       "      // Update the links…\n",
       "      var link = svg.selectAll(\"path.link\")\n",
       "          .data(links, function(d) { return d.target.id; });\n",
       "\n",
       "      // Enter any new links at the parent's previous position.\n",
       "      // XXX link width should be based on transition data, not node data\n",
       "      link.enter().insert(\"path\", \"g\")\n",
       "          .attr(\"class\", \"link\")\n",
       "          .attr(\"style\", function(d) {\n",
       "              var ls = treeData.link_style[d.target.dataID];\n",
       "              return ls;\n",
       "          } )\n",
       "          .attr(\"d\", function(d) {\n",
       "            var o = {x: source.x0, y: source.y0};\n",
       "            return diagonal({source: o, target: o});\n",
       "          });\n",
       "\n",
       "      // Transition links to their new position.\n",
       "      link.transition()\n",
       "          .duration(duration)\n",
       "          .attr(\"d\", diagonal);\n",
       "\n",
       "      // Transition exiting nodes to the parent's new position.\n",
       "      link.exit().transition()\n",
       "          .duration(duration)\n",
       "          .attr(\"d\", function(d) {\n",
       "            var o = {x: source.x, y: source.y};\n",
       "            return diagonal({source: o, target: o});\n",
       "          })\n",
       "          .remove();\n",
       "\n",
       "      // Stash the old positions for transition.\n",
       "      nodes.forEach(function(d) {\n",
       "        d.x0 = d.x;\n",
       "        d.y0 = d.y;\n",
       "      });\n",
       "    }\n",
       "\n",
       "    // Toggle children on click.\n",
       "    function click(d) {\n",
       "      if (d.children) {\n",
       "        d._children = d.children;\n",
       "        d.children = null;\n",
       "      } else if (d._children) {\n",
       "        d.children = d._children;\n",
       "        d._children = null;\n",
       "      } else {\n",
       "        initializeChildren(d, 1);\n",
       "      }\n",
       "      update(d, 750);\n",
       "    }\n",
       "\n",
       "}\n",
       "\n",
       "        })();\n",
       "    </script>\n",
       "    <p class=\"d3twarn\">\n",
       "    Attempting to display the tree. If the tree is large, this may take some time.\n",
       "    </p>\n",
       "    <p class=\"d3twarn\">\n",
       "    Note: D3Trees.jl requires an internet connection. If no tree appears, please check your connection. To help fix this, please see <a href=\"https://github.com/sisl/D3Trees.jl/issues/10\">this issue</a>. You may also diagnose errors with the javascript console (Ctrl-Shift-J in chrome).\n",
       "    </p>\n",
       "    </div>\n",
       "    </body>\n",
       "    </html>\n",
       "\n"
      ],
      "text/plain": [
       "                        1\n",
       "N:     53\n",
       "\n",
       "├──                     left\n",
       "│  Q:   0.00\n",
       "│  N:      7\n",
       "\n",
       "│  ├──                        2\n",
       "│  │  N:    179\n",
       "\n",
       "│  │  ├──                     left\n",
       "│  │  │  Q:   0.32\n",
       "│  │  │  N:      4\n",
       " (1 children)\n",
       "│  │  └──                    right\n",
       "│  │     Q:   3.01\n",
       "│  │     N:    175\n",
       " (2 children)\n",
       "│  └──                        1\n",
       "│     N:     53\n",
       "\n",
       "│     ├──                     left\n",
       "│     │  Q:   0.00\n",
       "│     │  N:      7\n",
       " (2 children)\n",
       "│     └──                    right\n",
       "│        Q:   1.98\n",
       "│        N:     46\n",
       " (2 children)\n",
       "└──                    right\n",
       "   Q:   1.98\n",
       "   N:     46\n",
       "\n",
       "   ├──                        2\n",
       "   │  N:    179\n",
       "\n",
       "   │  ├──                     left\n",
       "   │  │  Q:   0.32\n",
       "   │  │  N:      4\n",
       " (1 children)\n",
       "   │  └──                    right\n",
       "   │     Q:   3.01\n",
       "   │     N:    175\n",
       " (2 children)\n",
       "   └──                        1\n",
       "      N:     53\n",
       "\n",
       "      ├──                     left\n",
       "      │  Q:   0.00\n",
       "      │  N:      7\n",
       " (2 children)\n",
       "      └──                    right\n",
       "         Q:   1.98\n",
       "         N:     46\n",
       " (2 children)\n"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "using D3Trees\n",
    "D3Tree(planner, 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Solving with value iteration\n",
    "\n",
    "We can also try value iteration from the DiscreteValueIteration package"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "using DiscreteValueIteration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO: POMDPs.jl requirements for \u001b[34msolve(::ValueIterationSolver, ::Union{MDP,POMDP})\u001b[39m and dependencies. ([✔] = implemented correctly; [X] = missing)\n",
      "\n",
      "For \u001b[34msolve(::ValueIterationSolver, ::Union{MDP,POMDP})\u001b[39m:\n",
      "\u001b[32m  [✔] discount(::ChainMDP)\u001b[39m\n",
      "\u001b[31m  [X] n_states(::ChainMDP)\u001b[39m\n",
      "\u001b[31m  [X] n_actions(::ChainMDP)\u001b[39m\n",
      "\u001b[32m  [✔] transition(::ChainMDP, ::Int64, ::Symbol)\u001b[39m\n",
      "\u001b[32m  [✔] reward(::ChainMDP, ::Int64, ::Symbol, ::Int64)\u001b[39m\n",
      "\u001b[32m  [✔] stateindex(::ChainMDP, ::Int64)\u001b[39m\n",
      "\u001b[31m  [X] actionindex(::ChainMDP, ::Symbol)\u001b[39m\n",
      "\u001b[32m  [✔] actions(::ChainMDP, ::Int64)\u001b[39m\n",
      "\u001b[33m  WARNING: Some requirements may not be shown because a MethodError was thrown.\u001b[39m\n",
      "For \u001b[34mordered_states(::Union{MDP,POMDP})\u001b[39m (in solve(::ValueIterationSolver, ::Union{MDP,POMDP})):\n",
      "\u001b[31m  [X] states(::ChainMDP)\u001b[39m\n",
      "\u001b[33m  WARNING: Some requirements may not be shown because a MethodError was thrown.\u001b[39m\n",
      "For \u001b[34mordered_actions(::Union{MDP,POMDP})\u001b[39m (in solve(::ValueIterationSolver, ::Union{MDP,POMDP})):\n",
      "\u001b[32m  [✔] actions(::ChainMDP)\u001b[39m\n",
      "Note: Missing methods are often due to incorrect importing. You must explicitly import POMDPs functions to add new methods.\n",
      "\n",
      "Throwing the first exception (from processing \u001b[34msolve(::ValueIterationSolver, ::Union{MDP,POMDP})\u001b[39m requirements):\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "┌ Warning: ErrorException(\"rethrow(exc) not allowed outside a catch block\")\n",
      "└ @ Main In[31]:6\n"
     ]
    }
   ],
   "source": [
    "using DiscreteValueIteration\n",
    "solver = ValueIterationSolver(verbose=true)\n",
    "try\n",
    "    @requirements_info solver mdp\n",
    "catch err\n",
    "    @warn err\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "POMDPs.n_states(mdp::ChainMDP) = mdp.len\n",
    "POMDPs.states(mdp::ChainMDP) = 1:mdp.len\n",
    "POMDPs.n_actions(::ChainMDP) = 2\n",
    "POMDPs.actionindex(::ChainMDP, a::Symbol) = a == :left ? 1 : 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO: POMDPs.jl requirements for \u001b[34msolve(::ValueIterationSolver, ::Union{MDP,POMDP})\u001b[39m and dependencies. ([✔] = implemented correctly; [X] = missing)\n",
      "\n",
      "For \u001b[34msolve(::ValueIterationSolver, ::Union{MDP,POMDP})\u001b[39m:\n",
      "\u001b[32m  [✔] discount(::ChainMDP)\u001b[39m\n",
      "\u001b[32m  [✔] n_states(::ChainMDP)\u001b[39m\n",
      "\u001b[32m  [✔] n_actions(::ChainMDP)\u001b[39m\n",
      "\u001b[32m  [✔] transition(::ChainMDP, ::Int64, ::Symbol)\u001b[39m\n",
      "\u001b[32m  [✔] reward(::ChainMDP, ::Int64, ::Symbol, ::Int64)\u001b[39m\n",
      "\u001b[32m  [✔] stateindex(::ChainMDP, ::Int64)\u001b[39m\n",
      "\u001b[32m  [✔] actionindex(::ChainMDP, ::Symbol)\u001b[39m\n",
      "\u001b[32m  [✔] actions(::ChainMDP, ::Int64)\u001b[39m\n",
      "\u001b[32m  [✔] support(::SparseCat)\u001b[39m\n",
      "\u001b[32m  [✔] pdf(::SparseCat, ::Int64)\u001b[39m\n",
      "For \u001b[34mordered_states(::Union{MDP,POMDP})\u001b[39m (in solve(::ValueIterationSolver, ::Union{MDP,POMDP})):\n",
      "\u001b[32m  [✔] states(::ChainMDP)\u001b[39m\n",
      "For \u001b[34mordered_actions(::Union{MDP,POMDP})\u001b[39m (in solve(::ValueIterationSolver, ::Union{MDP,POMDP})):\n",
      "\u001b[32m  [✔] actions(::ChainMDP)\u001b[39m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "true"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@requirements_info solver mdp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Iteration 1   ] residual:          1 | iteration runtime:      0.005 ms, (  5.3E-06 s total)\n",
      "[Iteration 2   ] residual:      0.904 | iteration runtime:      0.005 ms, (  1.06E-05 s total)\n",
      "[Iteration 3   ] residual:      0.839 | iteration runtime:      0.002 ms, (  1.26E-05 s total)\n",
      "[Iteration 4   ] residual:      0.783 | iteration runtime:      0.004 ms, (  1.66E-05 s total)\n",
      "[Iteration 5   ] residual:      0.734 | iteration runtime:      0.002 ms, (  1.9E-05 s total)\n",
      "[Iteration 6   ] residual:      0.688 | iteration runtime:      0.002 ms, (  2.1E-05 s total)\n",
      "[Iteration 7   ] residual:      0.646 | iteration runtime:      0.002 ms, (  2.3E-05 s total)\n",
      "[Iteration 8   ] residual:      0.609 | iteration runtime:      0.002 ms, (  2.52E-05 s total)\n",
      "[Iteration 9   ] residual:      0.576 | iteration runtime:      0.003 ms, (  2.78E-05 s total)\n",
      "[Iteration 10  ] residual:      0.543 | iteration runtime:      0.003 ms, (  3.1E-05 s total)\n",
      "[Iteration 11  ] residual:       0.51 | iteration runtime:      0.002 ms, (  3.26E-05 s total)\n",
      "[Iteration 12  ] residual:      0.479 | iteration runtime:      0.002 ms, (  3.5E-05 s total)\n",
      "[Iteration 13  ] residual:      0.449 | iteration runtime:      0.002 ms, (  3.68E-05 s total)\n",
      "[Iteration 14  ] residual:      0.422 | iteration runtime:      0.002 ms, (  3.85E-05 s total)\n",
      "[Iteration 15  ] residual:      0.396 | iteration runtime:      0.002 ms, (  4.04E-05 s total)\n",
      "[Iteration 16  ] residual:      0.371 | iteration runtime:      0.002 ms, (  4.24E-05 s total)\n",
      "[Iteration 17  ] residual:      0.349 | iteration runtime:      0.002 ms, (  4.48E-05 s total)\n",
      "[Iteration 18  ] residual:      0.327 | iteration runtime:      0.002 ms, (  4.67E-05 s total)\n",
      "[Iteration 19  ] residual:      0.307 | iteration runtime:      0.003 ms, (  4.93E-05 s total)\n",
      "[Iteration 20  ] residual:      0.288 | iteration runtime:      0.004 ms, (  5.3E-05 s total)\n",
      "[Iteration 21  ] residual:       0.27 | iteration runtime:      0.003 ms, (  5.56E-05 s total)\n",
      "[Iteration 22  ] residual:      0.254 | iteration runtime:      0.002 ms, (  5.74E-05 s total)\n",
      "[Iteration 23  ] residual:      0.238 | iteration runtime:      0.002 ms, (  5.94E-05 s total)\n",
      "[Iteration 24  ] residual:      0.223 | iteration runtime:      0.002 ms, (  6.17E-05 s total)\n",
      "[Iteration 25  ] residual:       0.21 | iteration runtime:      0.002 ms, (  6.36E-05 s total)\n",
      "[Iteration 26  ] residual:      0.197 | iteration runtime:      0.004 ms, (  6.71E-05 s total)\n",
      "[Iteration 27  ] residual:      0.185 | iteration runtime:      0.002 ms, (  6.91E-05 s total)\n",
      "[Iteration 28  ] residual:      0.173 | iteration runtime:      0.002 ms, (  7.15E-05 s total)\n",
      "[Iteration 29  ] residual:      0.163 | iteration runtime:      0.002 ms, (  7.35E-05 s total)\n",
      "[Iteration 30  ] residual:      0.153 | iteration runtime:      0.003 ms, (  7.63E-05 s total)\n",
      "[Iteration 31  ] residual:      0.143 | iteration runtime:      0.003 ms, (  7.92E-05 s total)\n",
      "[Iteration 32  ] residual:      0.134 | iteration runtime:      0.003 ms, (  8.2E-05 s total)\n",
      "[Iteration 33  ] residual:      0.126 | iteration runtime:      0.002 ms, (  8.43E-05 s total)\n",
      "[Iteration 34  ] residual:      0.118 | iteration runtime:      0.003 ms, (  8.72E-05 s total)\n",
      "[Iteration 35  ] residual:      0.111 | iteration runtime:      0.003 ms, (  8.97E-05 s total)\n",
      "[Iteration 36  ] residual:      0.104 | iteration runtime:      0.003 ms, (  9.25E-05 s total)\n",
      "[Iteration 37  ] residual:     0.0978 | iteration runtime:      0.003 ms, (  9.57E-05 s total)\n",
      "[Iteration 38  ] residual:     0.0917 | iteration runtime:      0.003 ms, (  9.86E-05 s total)\n",
      "[Iteration 39  ] residual:     0.0861 | iteration runtime:      0.004 ms, (  0.000103 s total)\n",
      "[Iteration 40  ] residual:     0.0808 | iteration runtime:      0.003 ms, (  0.000106 s total)\n",
      "[Iteration 41  ] residual:     0.0758 | iteration runtime:      0.006 ms, (  0.000112 s total)\n",
      "[Iteration 42  ] residual:     0.0711 | iteration runtime:      0.002 ms, (  0.000114 s total)\n",
      "[Iteration 43  ] residual:     0.0668 | iteration runtime:      0.004 ms, (  0.000118 s total)\n",
      "[Iteration 44  ] residual:     0.0627 | iteration runtime:      0.003 ms, (  0.000121 s total)\n",
      "[Iteration 45  ] residual:     0.0588 | iteration runtime:      0.003 ms, (  0.000124 s total)\n",
      "[Iteration 46  ] residual:     0.0552 | iteration runtime:      0.003 ms, (  0.000127 s total)\n",
      "[Iteration 47  ] residual:     0.0518 | iteration runtime:      0.002 ms, (  0.000129 s total)\n",
      "[Iteration 48  ] residual:     0.0486 | iteration runtime:      0.002 ms, (  0.000131 s total)\n",
      "[Iteration 49  ] residual:     0.0456 | iteration runtime:      0.002 ms, (  0.000133 s total)\n",
      "[Iteration 50  ] residual:     0.0428 | iteration runtime:      0.003 ms, (  0.000136 s total)\n",
      "[Iteration 51  ] residual:     0.0402 | iteration runtime:      0.004 ms, (   0.00014 s total)\n",
      "[Iteration 52  ] residual:     0.0377 | iteration runtime:      0.003 ms, (  0.000143 s total)\n",
      "[Iteration 53  ] residual:     0.0354 | iteration runtime:      0.004 ms, (  0.000146 s total)\n",
      "[Iteration 54  ] residual:     0.0332 | iteration runtime:      0.004 ms, (  0.000151 s total)\n",
      "[Iteration 55  ] residual:     0.0311 | iteration runtime:      0.005 ms, (  0.000155 s total)\n",
      "[Iteration 56  ] residual:     0.0292 | iteration runtime:      0.006 ms, (  0.000161 s total)\n",
      "[Iteration 57  ] residual:     0.0274 | iteration runtime:      0.005 ms, (  0.000165 s total)\n",
      "[Iteration 58  ] residual:     0.0257 | iteration runtime:      0.003 ms, (  0.000168 s total)\n",
      "[Iteration 59  ] residual:     0.0241 | iteration runtime:      0.003 ms, (  0.000171 s total)\n",
      "[Iteration 60  ] residual:     0.0227 | iteration runtime:      0.002 ms, (  0.000173 s total)\n",
      "[Iteration 61  ] residual:     0.0213 | iteration runtime:      0.002 ms, (  0.000175 s total)\n",
      "[Iteration 62  ] residual:       0.02 | iteration runtime:      0.002 ms, (  0.000177 s total)\n",
      "[Iteration 63  ] residual:     0.0187 | iteration runtime:      0.004 ms, (  0.000181 s total)\n",
      "[Iteration 64  ] residual:     0.0176 | iteration runtime:      0.003 ms, (  0.000184 s total)\n",
      "[Iteration 65  ] residual:     0.0165 | iteration runtime:      0.002 ms, (  0.000185 s total)\n",
      "[Iteration 66  ] residual:     0.0155 | iteration runtime:      0.002 ms, (  0.000188 s total)\n",
      "[Iteration 67  ] residual:     0.0145 | iteration runtime:      0.002 ms, (   0.00019 s total)\n",
      "[Iteration 68  ] residual:     0.0136 | iteration runtime:      0.003 ms, (  0.000193 s total)\n",
      "[Iteration 69  ] residual:     0.0128 | iteration runtime:      0.002 ms, (  0.000195 s total)\n",
      "[Iteration 70  ] residual:      0.012 | iteration runtime:      0.002 ms, (  0.000197 s total)\n",
      "[Iteration 71  ] residual:     0.0113 | iteration runtime:      0.002 ms, (  0.000199 s total)\n",
      "[Iteration 72  ] residual:     0.0106 | iteration runtime:      0.002 ms, (  0.000201 s total)\n",
      "[Iteration 73  ] residual:    0.00992 | iteration runtime:      0.002 ms, (  0.000203 s total)\n",
      "[Iteration 74  ] residual:    0.00931 | iteration runtime:      0.004 ms, (  0.000207 s total)\n",
      "[Iteration 75  ] residual:    0.00873 | iteration runtime:      0.003 ms, (   0.00021 s total)\n",
      "[Iteration 76  ] residual:     0.0082 | iteration runtime:      0.002 ms, (  0.000212 s total)\n",
      "[Iteration 77  ] residual:    0.00769 | iteration runtime:      0.003 ms, (  0.000215 s total)\n",
      "[Iteration 78  ] residual:    0.00722 | iteration runtime:      0.028 ms, (  0.000242 s total)\n",
      "[Iteration 79  ] residual:    0.00677 | iteration runtime:      0.003 ms, (  0.000245 s total)\n",
      "[Iteration 80  ] residual:    0.00636 | iteration runtime:      0.003 ms, (  0.000247 s total)\n",
      "[Iteration 81  ] residual:    0.00596 | iteration runtime:      0.002 ms, (   0.00025 s total)\n",
      "[Iteration 82  ] residual:     0.0056 | iteration runtime:      0.002 ms, (  0.000252 s total)\n",
      "[Iteration 83  ] residual:    0.00525 | iteration runtime:      0.002 ms, (  0.000254 s total)\n",
      "[Iteration 84  ] residual:    0.00493 | iteration runtime:      0.003 ms, (  0.000256 s total)\n",
      "[Iteration 85  ] residual:    0.00463 | iteration runtime:      0.002 ms, (  0.000258 s total)\n",
      "[Iteration 86  ] residual:    0.00434 | iteration runtime:      0.002 ms, (  0.000261 s total)\n",
      "[Iteration 87  ] residual:    0.00407 | iteration runtime:      0.002 ms, (  0.000263 s total)\n",
      "[Iteration 88  ] residual:    0.00382 | iteration runtime:      0.003 ms, (  0.000266 s total)\n",
      "[Iteration 89  ] residual:    0.00359 | iteration runtime:      0.002 ms, (  0.000268 s total)\n",
      "[Iteration 90  ] residual:    0.00337 | iteration runtime:      0.007 ms, (  0.000275 s total)\n",
      "[Iteration 91  ] residual:    0.00316 | iteration runtime:      0.002 ms, (  0.000277 s total)\n",
      "[Iteration 92  ] residual:    0.00296 | iteration runtime:      0.004 ms, (   0.00028 s total)\n",
      "[Iteration 93  ] residual:    0.00278 | iteration runtime:      0.002 ms, (  0.000283 s total)\n",
      "[Iteration 94  ] residual:    0.00261 | iteration runtime:      0.004 ms, (  0.000286 s total)\n",
      "[Iteration 95  ] residual:    0.00245 | iteration runtime:      0.003 ms, (  0.000289 s total)\n",
      "[Iteration 96  ] residual:     0.0023 | iteration runtime:      0.002 ms, (  0.000292 s total)\n",
      "[Iteration 97  ] residual:    0.00216 | iteration runtime:      0.006 ms, (  0.000298 s total)\n",
      "[Iteration 98  ] residual:    0.00202 | iteration runtime:      0.002 ms, (    0.0003 s total)\n",
      "[Iteration 99  ] residual:     0.0019 | iteration runtime:      0.003 ms, (  0.000303 s total)\n",
      "[Iteration 100 ] residual:    0.00178 | iteration runtime:      0.002 ms, (  0.000305 s total)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "ValueIterationPolicy{Float64}:\n",
       " 1 -> :right\n",
       " 2 -> :right\n",
       " 3 -> :right\n",
       " 4 -> :right"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vi_policy = solve(solver, mdp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "action (generic function with 1 method)"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "action(vip::ValueIterationPolicy, s::Int64) = vip.policy[s]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "State: 1 Action: 2\n",
      "State: 2 Action: 2\n",
      "State: 3 Action: 2\n",
      "State: 4 Action: 2\n"
     ]
    }
   ],
   "source": [
    "for s in states(mdp)\n",
    "    println(\"State: $s Action: $(action(vi_policy, s))\")\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Crying Baby POMDP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "using POMDPModels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "BeliefUpdaters.DiscreteUpdater{BabyPOMDP}(BabyPOMDP(-5.0, -10.0, 0.1, 0.8, 0.1, 0.9))"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bp = BabyPOMDP()\n",
    "up = updater(bp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[32m\u001b[1m Resolving\u001b[22m\u001b[39m package versions...\n",
      "\u001b[32m\u001b[1m  Updating\u001b[22m\u001b[39m `C:\\Users\\rmoss\\Code\\sisl\\aa228\\aa228-notebook\\Project.toml`\n",
      "\u001b[90m [no changes]\u001b[39m\n",
      "\u001b[32m\u001b[1m  Updating\u001b[22m\u001b[39m `C:\\Users\\rmoss\\Code\\sisl\\aa228\\aa228-notebook\\Manifest.toml`\n",
      "\u001b[90m [no changes]\u001b[39m\n"
     ]
    }
   ],
   "source": [
    "Pkg.add(\"SARSOP\")\n",
    "using SARSOP # might take a while to compile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\u001b[33mSARSOPSolver\u001b[39m\n",
       "\u001b[32m  fast: \u001b[39m\u001b[33mBool\u001b[39m false\n",
       "\u001b[32m  randomization: \u001b[39m\u001b[33mBool\u001b[39m false\n",
       "\u001b[32m  precision: \u001b[39m\u001b[33mFloat64\u001b[39m 0.001\n",
       "\u001b[32m  timeout: \u001b[39m\u001b[33mNothing\u001b[39m nothing\n",
       "\u001b[32m  memory: \u001b[39m\u001b[33mNothing\u001b[39m nothing\n",
       "\u001b[32m  trial_improvement_factor: \u001b[39m\u001b[33mFloat64\u001b[39m 0.5\n",
       "\u001b[32m  policy_interval: \u001b[39m\u001b[33mNothing\u001b[39m nothing\n",
       "\u001b[32m  verbose: \u001b[39m\u001b[33mBool\u001b[39m true\n",
       "\u001b[32m  policy_filename: \u001b[39mString \"policy.out\"\n",
       "\u001b[32m  pomdp_filename: \u001b[39mString \"model.pomdpx\"\n"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "solver = SARSOPSolver()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generating a pomdpx file: model.pomdpx\n",
      "\n",
      "Loading the model ...\n",
      "  input file   : model.pomdpx\n",
      "  loading time : 0.00s \n",
      "\n",
      "SARSOP initializing ...\n",
      "  initialization time : 0.00s\n",
      "\n",
      "-------------------------------------------------------------------------------\n",
      " Time   |#Trial |#Backup |LBound    |UBound    |Precision  |#Alphas |#Beliefs  \n",
      "-------------------------------------------------------------------------------\n",
      " 0.001   0       0        -47.3685   -16.0713   31.2971     2        1        \n",
      " 0.002   5       51       -17.9323   -16.0492   1.88318     9        6        \n",
      " 0.003   10      103      -16.3978   -16.2159   0.181953    2        8        \n",
      " 0.004   14      151      -16.3124   -16.277    0.0353847   2        13       \n",
      " 0.005   18      207      -16.3058   -16.2969   0.00887227  2        16       \n",
      " 0.005   21      250      -16.3055   -16.3022   0.00329006  7        16       \n",
      " 0.007   25      305      -16.3055   -16.3044   0.00111584  2        18       \n",
      " 0.007   26      323      -16.3055   -16.3045   0.000944675 2        18       \n",
      "-------------------------------------------------------------------------------\n",
      "\n",
      "SARSOP finishing ...\n",
      "  target precision reached\n",
      "  target precision  : 0.001000\n",
      "  precision reached : 0.000945 \n",
      "\n",
      "-------------------------------------------------------------------------------\n",
      " Time   |#Trial |#Backup |LBound    |UBound    |Precision  |#Alphas |#Beliefs  \n",
      "-------------------------------------------------------------------------------\n",
      " 0.007   26      323      -16.3055   -16.3045   0.000944675 2        18       \n",
      "-------------------------------------------------------------------------------\n",
      "\n",
      "Writing out policy ...\n",
      "  output file : policy.out\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "AlphaVectorPolicy{BabyPOMDP,Bool}(BabyPOMDP(-5.0, -10.0, 0.1, 0.8, 0.1, 0.9), Array{Float64,1}[[-19.6749, -29.6749], [-16.3055, -38.2512]], Bool[1, 0])"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "POMDPPolicies.AlphaVectorPolicy(p::POMDP, alphas::Matrix{Float64}, action_map::BitVector) = POMDPPolicies.AlphaVectorPolicy(p, alphas, Array(action_map))\n",
    "p = solve(solver, bp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      "\u001b[32mSimulating... 10%|████                                  |  ETA: 0:00:03\u001b[39m\r",
      "\u001b[32mSimulating...100%|██████████████████████████████████████| Time: 0:00:00\u001b[39m\n"
     ]
    }
   ],
   "source": [
    "history = simulate(simulator, bp, p, up, BoolDistribution(0.0));"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.00 hungry belief => don't feed\n",
      "0.98 hungry belief => don't feed\n",
      "0.97 hungry belief => don't feed\n",
      "0.97 hungry belief => don't feed\n",
      "0.97 hungry belief => don't feed\n",
      "0.97 hungry belief => don't feed\n",
      "0.46 hungry belief => feed\n",
      "1.00 hungry belief => don't feed\n",
      "0.98 hungry belief => don't feed\n",
      "0.47 hungry belief => feed\n"
     ]
    }
   ],
   "source": [
    "for (b, a) in eachstep(history, \"b,a\")\n",
    "    @printf(\"%3.2f hungry belief => %s\\n\",\n",
    "            b.b[1],\n",
    "            a ? \"feed\" : \"don't feed\")\n",
    "end"
   ]
  }
 ],
 "metadata": {
  "@webio": {
   "lastCommId": null,
   "lastKernelId": null
  },
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Julia 1.2.0",
   "language": "julia",
   "name": "julia-1.2"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.2.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}