diff --git a/.gitignore b/.gitignore
index 9a4bb620f..af3dab103 100644
--- a/.gitignore
+++ b/.gitignore
@@ -70,3 +70,4 @@ target/
 
 # dotenv
 .env
+.idea
diff --git a/mdp.ipynb b/mdp.ipynb
index ee9b0ba85..e288d1b49 100644
--- a/mdp.ipynb
+++ b/mdp.ipynb
@@ -133,7 +133,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {
     "collapsed": true
    },
@@ -145,7 +145,7 @@
     "        # All possible actions.\n",
     "        actlist = []\n",
     "        for state in transition_matrix.keys():\n",
-    "            actlist.extend(transition_matrix.keys())\n",
+    "            actlist.extend(transition_matrix[state])\n",
     "        actlist = list(set(actlist))\n",
     "\n",
     "        MDP.__init__(self, init, actlist, terminals=terminals, gamma=gamma)\n",
@@ -155,7 +155,10 @@
     "            self.states.add(state)\n",
     "\n",
     "    def T(self, state, action):\n",
-    "        return [(new_state, prob) for new_state, prob in self.t[state][action].items()]"
+    "        if action is None:\n",
+    "            return [(0.0, state)]\n",
+    "        else: \n",
+    "            return [(prob, new_state) for new_state, prob in self.t[state][action].items()]"
    ]
   },
   {
@@ -449,11 +452,7 @@
      ]
     },
     {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "9aed96e7288d4ed59df439f68399dc12"
-      }
-     },
+     "data": {},
      "metadata": {},
      "output_type": "display_data"
     }
@@ -524,7 +523,7 @@
     "022a5fdfc8e44fb09b21c4bd5b67a0db": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -555,7 +554,7 @@
     "0675230fb92f4539bc257b768fb4cd10": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -571,7 +570,7 @@
     "0783e74a8c2b40cc9b0f5706271192f4": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -599,7 +598,7 @@
     "098f12158d844cdf89b29a4cd568fda0": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -624,7 +623,7 @@
     "0b65fb781274495ab498ad518bc274d4": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -733,7 +732,7 @@
     "1af711fe8e4f43f084cef6c89eec40ae": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -749,7 +748,7 @@
     "1c5c913acbde4e87a163abb2e24e6e38": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -774,7 +773,7 @@
     "200e3ebead3d4858a47e2f6d345ca395": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -892,7 +891,7 @@
     "2d3acd8872c342eab3484302cac2cb05": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -902,7 +901,7 @@
     "2e1351ad05384d058c90e594bc6143c1": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -915,7 +914,7 @@
     "2f5438f1b34046a597a467effd43df11": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -952,7 +951,7 @@
     "319425ba805346f5ba366c42e220f9c6": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -971,7 +970,7 @@
     "332a89c03bfb49c2bb291051d172b735": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -1020,7 +1019,7 @@
     "388571e8e0314dfab8e935b7578ba7f9": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -1042,7 +1041,7 @@
     "3a21291c8e7249e3b04417d31b0447cf": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -1055,7 +1054,7 @@
     "3b22d68709b046e09fe70f381a3944cd": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -1065,7 +1064,7 @@
     "3c1b2ec10a9041be8a3fad9da78ff9f6": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -1090,7 +1089,7 @@
     "3e5b9fd779574270bf58101002c152ce": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -1100,7 +1099,7 @@
     "3e8bb05434cb4a0291383144e4523840": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -1149,7 +1148,7 @@
     "428e42f04a1e4347a1f548379c68f91b": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -1165,7 +1164,7 @@
     "4379175239b34553bf45c8ef9443ac55": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -1178,7 +1177,7 @@
     "4421c121414d464bb3bf1b5f0e86c37b": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -1209,7 +1208,7 @@
     "4731208453424514b471f862804d9bb8": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -1258,7 +1257,7 @@
     "4d281cda33fa489d86228370e627a5b0": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -1277,7 +1276,7 @@
     "4ec035cba73647358d416615cf4096ee": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -1302,7 +1301,7 @@
     "5141ae07149b46909426208a30e2861e": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -1339,7 +1338,7 @@
     "55a1b0b794f44ac796bc75616f65a2a1": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -1400,7 +1399,7 @@
     "595c537ed2514006ac823b4090cf3b4b": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -1461,7 +1460,7 @@
     "5f823979d2ce4c34ba18b4ca674724e4": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -1501,14 +1500,14 @@
     "644dcff39d7c47b7b8b729d01f59bee5": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
     "6455faf9dbc6477f8692528e6eb90c9a": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -1521,7 +1520,7 @@
     "665ed2b201144d78a5a1f57894c2267c": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -1564,7 +1563,7 @@
     "6a28f605a5d14589907dba7440ede2fc": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -1589,7 +1588,7 @@
     "6d7effd6bc4c40a4b17bf9e136c5814c": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -1638,7 +1637,7 @@
     "72dfe79a3e52429da1cf4382e78b2144": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -1669,7 +1668,7 @@
     "75e344508b0b45d1a9ae440549d95b1a": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -1727,7 +1726,7 @@
     "7f2f98bbffc0412dbb31c387407a9fed": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -1758,7 +1757,7 @@
     "82e2820c147a4dff85a01bcddbad8645": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -1861,21 +1860,21 @@
     "8cffde5bdb3d4f7597131b048a013929": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
     "8db2abcad8bc44df812d6ccf2d2d713c": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
     "8dd5216b361c44359ba1233ee93683a4": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -1921,7 +1920,7 @@
     "933904217b6045c1b654b7e5749203f5": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -1949,7 +1948,7 @@
     "94f2b877a79142839622a61a3a081c03": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -1971,7 +1970,7 @@
     "97207358fc65430aa196a7ed78b252f0": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -1984,7 +1983,7 @@
     "986c6c4e92964759903d6eb7f153df8a": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -2027,14 +2026,14 @@
     "9d5e9658af264ad795f6a5f3d8c3c30f": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
     "9d7aa65511b6482d9587609ad7898f54": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -2053,7 +2052,7 @@
     "9efb46d2bb0648f6b109189986f4f102": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -2069,7 +2068,7 @@
     "9f43f85a0fb9464e9b7a25a85f6dba9c": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -2082,7 +2081,7 @@
     "9faa50b44e1842e0acac301f93a129c4": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -2107,7 +2106,7 @@
     "a1840ca22d834df2b145151baf6d8241": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -2144,7 +2143,7 @@
     "a39cfb47679c4d2895cda12c6d9d2975": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -2175,7 +2174,7 @@
     "a87c651448f14ce4958d73c2f1e413e1": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -2284,7 +2283,7 @@
     "b7e4c497ff5c4173961ffdc3bd3821a9": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -2309,7 +2308,7 @@
     "b9c138598fce460692cc12650375ee52": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -2328,7 +2327,7 @@
     "bbe5dea9d57d466ba4e964fce9af13cf": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -2362,7 +2361,7 @@
     "beb0c9b29d8d4d69b3147af666fa298b": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -2429,7 +2428,7 @@
     "c74bbd55a8644defa3fcef473002a626": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -2496,7 +2495,7 @@
     "ce3a0e82e80d48b9b2658e0c52196644": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -2506,7 +2505,7 @@
     "ce8d3cd3535b459c823da2f49f3cc526": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -2576,7 +2575,7 @@
     "d83329fe36014f85bb5d0247d3ae4472": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -2610,7 +2609,7 @@
     "dc7376a2272e44179f237e5a1c7f6a49": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -2707,7 +2706,7 @@
     "e4e5dd3dc28d4aa3ab8f8f7c4a475115": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -2723,7 +2722,7 @@
     "e64ab85e80184b70b69d01a9c6851943": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -2820,7 +2819,7 @@
     "f262055f3f1b48029f9e2089f752b0b8": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -2851,7 +2850,7 @@
     "f3df35ce53e0466e81a48234b36a1430": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },
@@ -2930,7 +2929,7 @@
     "f9458080ed534d25856c67ce8f93d5a1": {
      "views": [
       {
-       "cell_index": 27
+       "cell_index": 27.0
       }
      ]
     },