diff --git a/examples/advanced/complex_action_dependency.ipynb b/examples/advanced/complex_action_dependency.ipynb
index d3dfb17d..42e21b24 100644
--- a/examples/advanced/complex_action_dependency.ipynb
+++ b/examples/advanced/complex_action_dependency.ipynb
@@ -1,5 +1,12 @@
 {
  "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/infer-actively/pymdp/blob/main/examples/advanced/complex_action_dependency.ipynb)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -13,7 +20,18 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "if \"google.colab\" in sys.modules:\n",
+    "    %pip install \"inferactively-pymdp\" -q"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -37,7 +55,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [
     {
@@ -161,7 +179,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [
     {
@@ -274,25 +292,7 @@
    ]
   }
  ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "pymdp_dev_env",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.12.3"
-  }
- },
+ "metadata": {},
  "nbformat": 4,
  "nbformat_minor": 2
 }
diff --git a/examples/advanced/infer_states_optimization/methods_test.ipynb b/examples/advanced/infer_states_optimization/methods_test.ipynb
index 8364ff9a..36d517b7 100644
--- a/examples/advanced/infer_states_optimization/methods_test.ipynb
+++ b/examples/advanced/infer_states_optimization/methods_test.ipynb
@@ -1,8 +1,28 @@
 {
  "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "798edfd3",
+   "metadata": {},
+   "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/infer-actively/pymdp/blob/main/examples/advanced/infer_states_optimization/methods_test.ipynb)"
+   ]
+  },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
+   "id": "9443d173",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "if \"google.colab\" in sys.modules:\n",
+    "    %pip install \"inferactively-pymdp\" -q"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
    "id": "d329903e-0caa-4ad9-8b85-92ed348f1e0f",
    "metadata": {},
    "outputs": [],
@@ -34,7 +54,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 1,
    "id": "87d1c903-5496-4b12-a653-d6a9d8e6337b",
    "metadata": {},
    "outputs": [
@@ -62,7 +82,7 @@
        "  'dim_sampling_type': 'uniform'}}"
       ]
      },
-     "execution_count": 2,
+     "execution_count": 1,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -90,7 +110,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 2,
    "id": "bc5d11ae-0e03-4ebc-9fb2-e31e51f29b63",
    "metadata": {},
    "outputs": [
@@ -132,7 +152,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "id": "ef656d7c-b9da-4d9b-990d-a6b7a99ca725",
    "metadata": {},
    "outputs": [],
@@ -151,7 +171,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 3,
    "id": "c2b7676f-58a7-41d3-8e66-31d2f58213b0",
    "metadata": {
     "scrolled": true
@@ -262,7 +282,7 @@
        "         [[0.24136765, 0.17386995, 0.2222213 , 0.16080104, 0.2017401 ]]],      dtype=float32)])"
       ]
      },
-     "execution_count": 5,
+     "execution_count": 3,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -282,7 +302,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "id": "dd83c3d3",
    "metadata": {},
    "outputs": [],
@@ -295,7 +315,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": null,
    "id": "6c778c1e-e373-4603-a244-6f3455a282d8",
    "metadata": {},
    "outputs": [],
@@ -312,7 +332,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 4,
    "id": "5398c4aa-77bf-438c-8be8-1a5ab72e3d89",
    "metadata": {
     "scrolled": true
@@ -423,7 +443,7 @@
        "         [[0.24136765, 0.17386995, 0.2222213 , 0.16080104, 0.2017401 ]]],      dtype=float32)])"
       ]
      },
-     "execution_count": 8,
+     "execution_count": 4,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -435,7 +455,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 5,
    "id": "14c74aed",
    "metadata": {},
    "outputs": [
@@ -544,7 +564,7 @@
        "         [[0.24136765, 0.17386995, 0.2222213 , 0.16080104, 0.2017401 ]]],      dtype=float32)])"
       ]
      },
-     "execution_count": 9,
+     "execution_count": 5,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -569,7 +589,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": null,
    "id": "9485d771",
    "metadata": {},
    "outputs": [],
@@ -583,7 +603,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": null,
    "id": "891743ae",
    "metadata": {},
    "outputs": [],
@@ -602,7 +622,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 6,
    "id": "1d52085c",
    "metadata": {},
    "outputs": [
@@ -711,7 +731,7 @@
        "         [[0.24136765, 0.17386995, 0.2222213 , 0.16080104, 0.2017401 ]]],      dtype=float32)])"
       ]
      },
-     "execution_count": 12,
+     "execution_count": 6,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -727,7 +747,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 7,
    "id": "78beb618",
    "metadata": {},
    "outputs": [
@@ -836,7 +856,7 @@
        "         [[0.24136765, 0.17386995, 0.2222213 , 0.16080104, 0.2017401 ]]],      dtype=float32)])"
       ]
      },
-     "execution_count": 13,
+     "execution_count": 7,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -862,7 +882,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": null,
    "id": "1e91bf02",
    "metadata": {},
    "outputs": [],
@@ -874,7 +894,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": null,
    "id": "96d162ed-b35e-44d3-9106-1c37158aab4d",
    "metadata": {},
    "outputs": [],
@@ -893,7 +913,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": 8,
    "id": "72bc7f44-52e9-4843-aef8-f8255fe785fd",
    "metadata": {
     "scrolled": true
@@ -1004,7 +1024,7 @@
        "         [[0.24136762, 0.17386994, 0.22222117, 0.1608011 , 0.20174012]]],      dtype=float32)])"
       ]
      },
-     "execution_count": 16,
+     "execution_count": 8,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -1017,7 +1037,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 17,
+   "execution_count": 9,
    "id": "5f161096",
    "metadata": {},
    "outputs": [
@@ -1126,7 +1146,7 @@
        "         [[0.24136762, 0.17386994, 0.22222117, 0.1608011 , 0.20174012]]],      dtype=float32)])"
       ]
      },
-     "execution_count": 17,
+     "execution_count": 9,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -1142,25 +1162,7 @@
    ]
   }
  ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "venv",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.10.18"
-  }
- },
+ "metadata": {},
  "nbformat": 4,
  "nbformat_minor": 5
 }
diff --git a/examples/advanced/pymdp_with_neural_encoder.ipynb b/examples/advanced/pymdp_with_neural_encoder.ipynb
index a6c27823..30ccd2a4 100644
--- a/examples/advanced/pymdp_with_neural_encoder.ipynb
+++ b/examples/advanced/pymdp_with_neural_encoder.ipynb
@@ -4,6 +4,14 @@
    "cell_type": "markdown",
    "id": "0",
    "metadata": {},
+   "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/infer-actively/pymdp/blob/main/examples/advanced/pymdp_with_neural_encoder.ipynb)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1",
+   "metadata": {},
    "source": [
     "# Training a neural network to transform continuous into discrete observations for a `pymdp` Agent\n",
     "\n",
@@ -17,7 +25,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "1",
+   "id": "2",
    "metadata": {},
    "source": [
     "## 1. Imports\n",
@@ -29,7 +37,19 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "2",
+   "id": "3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "if \"google.colab\" in sys.modules:\n",
+    "    %pip install \"inferactively-pymdp\" -q"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -51,7 +71,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "3",
+   "id": "5",
    "metadata": {},
    "source": [
     "## 2. Setup\n",
@@ -62,7 +82,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "4",
+   "id": "6",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -120,7 +140,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "5",
+   "id": "7",
    "metadata": {},
    "source": [
     "## 3. Build the Fixed Discrete Agent\n",
@@ -139,7 +159,7 @@
   {
    "cell_type": "code",
    "execution_count": 1,
-   "id": "6",
+   "id": "8",
    "metadata": {},
    "outputs": [
     {
@@ -229,7 +249,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "7",
+   "id": "9",
    "metadata": {},
    "source": [
     "## 4. Generate Offline Continuous Trajectories\n",
@@ -247,7 +267,7 @@
   {
    "cell_type": "code",
    "execution_count": 2,
-   "id": "8",
+   "id": "10",
    "metadata": {},
    "outputs": [
     {
@@ -346,7 +366,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "9",
+   "id": "11",
    "metadata": {},
    "source": [
     "## 5. Define the Differentiable Front-End\n",
@@ -362,7 +382,7 @@
   {
    "cell_type": "code",
    "execution_count": 3,
-   "id": "10",
+   "id": "12",
    "metadata": {},
    "outputs": [
     {
@@ -414,7 +434,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "11",
+   "id": "13",
    "metadata": {},
    "source": [
     "## 6. Define the Training Objective\n",
@@ -435,7 +455,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "12",
+   "id": "14",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -541,7 +561,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "13",
+   "id": "15",
    "metadata": {},
    "source": [
     "## 7. Train the Encoder\n",
@@ -563,7 +583,7 @@
   {
    "cell_type": "code",
    "execution_count": 4,
-   "id": "14",
+   "id": "16",
    "metadata": {},
    "outputs": [
     {
@@ -657,7 +677,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "15",
+   "id": "17",
    "metadata": {},
    "source": [
     "## 8. Evaluate Representation Quality\n",
@@ -674,7 +694,7 @@
   {
    "cell_type": "code",
    "execution_count": 5,
-   "id": "16",
+   "id": "18",
    "metadata": {},
    "outputs": [
     {
@@ -769,7 +789,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "17",
+   "id": "19",
    "metadata": {},
    "source": [
     "## 9. Plan from Continuous Observations\n",
@@ -783,7 +803,7 @@
   {
    "cell_type": "code",
    "execution_count": 6,
-   "id": "18",
+   "id": "20",
    "metadata": {},
    "outputs": [
     {
@@ -862,7 +882,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "19",
+   "id": "21",
    "metadata": {},
    "source": [
     "## 10. Animation: Continuous Density + Inference During Control\n",
@@ -878,7 +898,7 @@
   {
    "cell_type": "code",
    "execution_count": 7,
-   "id": "20",
+   "id": "22",
    "metadata": {},
    "outputs": [
     {
@@ -18871,7 +18891,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "21",
+   "id": "23",
    "metadata": {},
    "source": [
     "## 11. Animation: Encoder Decision Regions + Control Trajectory\n",
@@ -18885,7 +18905,7 @@
   {
    "cell_type": "code",
    "execution_count": 8,
-   "id": "22",
+   "id": "24",
    "metadata": {},
    "outputs": [
     {
@@ -27347,7 +27367,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "23",
+   "id": "25",
    "metadata": {},
    "source": [
     "## 12. Summary\n",
diff --git a/examples/envs/cue_chaining_demo.ipynb b/examples/envs/cue_chaining_demo.ipynb
index 1a309350..65f5d50b 100644
--- a/examples/envs/cue_chaining_demo.ipynb
+++ b/examples/envs/cue_chaining_demo.ipynb
@@ -1,5 +1,13 @@
 {
  "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "e6e8061c",
+   "metadata": {},
+   "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/infer-actively/pymdp/blob/main/examples/envs/cue_chaining_demo.ipynb)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "b803090f",
@@ -22,7 +30,19 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
+   "id": "8210db71",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "if \"google.colab\" in sys.modules:\n",
+    "    %pip install \"inferactively-pymdp\" -q"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
    "id": "b59618c1",
    "metadata": {
     "execution": {
@@ -56,7 +76,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 1,
    "id": "21de5c24",
    "metadata": {
     "execution": {
@@ -110,7 +130,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 2,
    "id": "203e2547",
    "metadata": {
     "execution": {
@@ -180,7 +200,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 3,
    "id": "ea613d7d",
    "metadata": {
     "execution": {
@@ -220,7 +240,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 4,
    "id": "c6a9125f",
    "metadata": {
     "execution": {
@@ -379,7 +399,7 @@
        "10  10  (3, 5)     Null              Null     Cheese   STAY"
       ]
      },
-     "execution_count": 5,
+     "execution_count": 4,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -418,7 +438,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 5,
    "id": "ee345b31",
    "metadata": {
     "execution": {
@@ -494,25 +514,7 @@
    ]
   }
  ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "Python 3",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.11.12"
-  }
- },
+ "metadata": {},
  "nbformat": 4,
  "nbformat_minor": 5
 }
diff --git a/examples/envs/generalized_tmaze_demo.ipynb b/examples/envs/generalized_tmaze_demo.ipynb
index eaaeeb2a..c27651ae 100644
--- a/examples/envs/generalized_tmaze_demo.ipynb
+++ b/examples/envs/generalized_tmaze_demo.ipynb
@@ -1,5 +1,12 @@
 {
  "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/infer-actively/pymdp/blob/main/examples/envs/generalized_tmaze_demo.ipynb)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -31,6 +38,17 @@
     "```"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "if \"google.colab\" in sys.modules:\n",
+    "    %pip install \"inferactively-pymdp[nb]\" -q"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -70,7 +88,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [
     {
@@ -126,7 +144,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -168,7 +186,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [
     {
@@ -258,7 +276,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -269,7 +287,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [
     {
@@ -291,7 +309,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [
     {
@@ -1552,7 +1570,7 @@
        "<IPython.core.display.HTML object>"
       ]
      },
-     "execution_count": 12,
+     "execution_count": 4,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -1593,25 +1611,7 @@
    ]
   }
  ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "pymdp_dev_env",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.12.3"
-  }
- },
+ "metadata": {},
  "nbformat": 4,
  "nbformat_minor": 2
 }
diff --git a/examples/envs/graph_worlds_demo.ipynb b/examples/envs/graph_worlds_demo.ipynb
index 295edc97..a636ca6f 100644
--- a/examples/envs/graph_worlds_demo.ipynb
+++ b/examples/envs/graph_worlds_demo.ipynb
@@ -1,5 +1,12 @@
 {
  "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/infer-actively/pymdp/blob/main/examples/envs/graph_worlds_demo.ipynb)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -11,7 +18,18 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "if \"google.colab\" in sys.modules:\n",
+    "    %pip install \"inferactively-pymdp\" -q"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:30:18.683564Z",
@@ -50,7 +68,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 1,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:30:22.675254Z",
@@ -87,7 +105,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:30:22.728168Z",
@@ -115,7 +133,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 2,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:30:23.548655Z",
@@ -155,7 +173,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:30:24.285830Z",
@@ -179,7 +197,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:30:25.188717Z",
@@ -202,7 +220,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 3,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:30:25.768119Z",
@@ -218,7 +236,7 @@
        "dict_keys(['action', 'empirical_prior', 'env_state', 'neg_efe', 'observation', 'qpi', 'qs'])"
       ]
      },
-     "execution_count": 7,
+     "execution_count": 3,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -236,7 +254,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 4,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:30:25.778621Z",
@@ -269,7 +287,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 5,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:30:25.786826Z",
@@ -327,7 +345,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 6,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:30:26.032916Z",
@@ -469,20 +487,7 @@
    ]
   }
  ],
- "metadata": {
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.11.12"
-  }
- },
+ "metadata": {},
  "nbformat": 4,
  "nbformat_minor": 2
 }
diff --git a/examples/envs/knapsack_demo.ipynb b/examples/envs/knapsack_demo.ipynb
index 03fd8223..1916a91f 100644
--- a/examples/envs/knapsack_demo.ipynb
+++ b/examples/envs/knapsack_demo.ipynb
@@ -1,5 +1,12 @@
 {
  "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/infer-actively/pymdp/blob/main/examples/envs/knapsack_demo.ipynb)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -15,7 +22,18 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "if \"google.colab\" in sys.modules:\n",
+    "    %pip install \"inferactively-pymdp\" -q"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -38,7 +56,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [
     {
@@ -98,7 +116,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [
     {
@@ -119,7 +137,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [
     {
@@ -213,7 +231,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -225,7 +243,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [
     {
@@ -259,7 +277,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [
     {
@@ -290,7 +308,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [
     {
@@ -337,25 +355,7 @@
    ]
   }
  ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "pymdp_dev_env",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.12.3"
-  }
- },
+ "metadata": {},
  "nbformat": 4,
  "nbformat_minor": 2
 }
diff --git a/examples/envs/tmaze_demo.ipynb b/examples/envs/tmaze_demo.ipynb
index b1cb4f64..ce7c9b70 100644
--- a/examples/envs/tmaze_demo.ipynb
+++ b/examples/envs/tmaze_demo.ipynb
@@ -1,5 +1,12 @@
 {
  "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/infer-actively/pymdp/blob/main/examples/envs/tmaze_demo.ipynb)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -36,7 +43,18 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "if \"google.colab\" in sys.modules:\n",
+    "    %pip install \"inferactively-pymdp[nb]\" -q"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:30:30.989078Z",
@@ -134,7 +152,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:30:34.893886Z",
@@ -188,7 +206,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 1,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:30:35.263798Z",
@@ -275,7 +293,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:30:36.074529Z",
@@ -309,7 +327,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 2,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:30:37.067936Z",
@@ -379,7 +397,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 3,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:30:39.107336Z",
@@ -489,7 +507,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:30:57.523036Z",
@@ -521,7 +539,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 4,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:30:57.553800Z",
@@ -589,7 +607,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 5,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:31:00.134138Z",
@@ -653,7 +671,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:31:13.114299Z",
@@ -690,7 +708,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 6,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:31:13.167263Z",
@@ -762,7 +780,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:31:13.736663Z",
@@ -787,7 +805,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 7,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:31:14.974645Z",
@@ -831,7 +849,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 8,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:31:15.010576Z",
@@ -878,7 +896,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": 9,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:31:16.053250Z",
@@ -935,7 +953,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": 10,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:31:18.012533Z",
@@ -979,7 +997,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 17,
+   "execution_count": 11,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:31:18.063914Z",
@@ -1020,7 +1038,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 18,
+   "execution_count": 12,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:31:18.195581Z",
@@ -1066,20 +1084,7 @@
    ]
   }
  ],
- "metadata": {
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.11.12"
-  }
- },
+ "metadata": {},
  "nbformat": 4,
  "nbformat_minor": 2
 }
diff --git a/examples/experimental/sophisticated_inference/mcts_generalized_tmaze.ipynb b/examples/experimental/sophisticated_inference/mcts_generalized_tmaze.ipynb
index ccab8e39..f954eaf9 100644
--- a/examples/experimental/sophisticated_inference/mcts_generalized_tmaze.ipynb
+++ b/examples/experimental/sophisticated_inference/mcts_generalized_tmaze.ipynb
@@ -1,5 +1,12 @@
 {
  "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/infer-actively/pymdp/blob/main/examples/experimental/sophisticated_inference/mcts_generalized_tmaze.ipynb)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -99,6 +106,17 @@
     "\\end{equation}"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "if \"google.colab\" in sys.modules:\n",
+    "    %pip install \"inferactively-pymdp\" -q"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -130,7 +148,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -189,7 +207,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [
     {
@@ -293,7 +311,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -304,7 +322,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [
     {
@@ -313,7 +331,7 @@
        "dict_keys(['action', 'action_weights', 'empirical_prior', 'env_state', 'observation', 'qpi', 'qs', 'search_tree'])"
       ]
      },
-     "execution_count": 15,
+     "execution_count": 2,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -324,7 +342,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [
     {
@@ -350,7 +368,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [
     {
@@ -372,7 +390,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [
     {
@@ -1495,7 +1513,7 @@
        "<IPython.core.display.HTML object>"
       ]
      },
-     "execution_count": 12,
+     "execution_count": 5,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -1506,25 +1524,7 @@
    ]
   }
  ],
- "metadata": {
-  "kernelspec": {
-   "display_name": ".venv311",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.11.12"
-  }
- },
+ "metadata": {},
  "nbformat": 4,
  "nbformat_minor": 2
 }
diff --git a/examples/experimental/sophisticated_inference/mcts_graph_world.ipynb b/examples/experimental/sophisticated_inference/mcts_graph_world.ipynb
index 3c057b3f..47fb6a94 100644
--- a/examples/experimental/sophisticated_inference/mcts_graph_world.ipynb
+++ b/examples/experimental/sophisticated_inference/mcts_graph_world.ipynb
@@ -1,8 +1,26 @@
 {
  "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/infer-actively/pymdp/blob/main/examples/experimental/sophisticated_inference/mcts_graph_world.ipynb)"
+   ]
+  },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "if \"google.colab\" in sys.modules:\n",
+    "    %pip install \"inferactively-pymdp[nb]\" -q"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:31:43.007199Z",
@@ -38,7 +56,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:31:43.765921Z",
@@ -119,7 +137,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 1,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:31:43.771113Z",
@@ -160,7 +178,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 2,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:31:47.658379Z",
@@ -204,7 +222,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:31:48.371274Z",
@@ -224,7 +242,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 3,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:31:48.390928Z",
@@ -1861,7 +1879,7 @@
        "<IPython.core.display.SVG object>"
       ]
      },
-     "execution_count": 6,
+     "execution_count": 3,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -1892,20 +1910,7 @@
    ]
   }
  ],
- "metadata": {
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.11.12"
-  }
- },
+ "metadata": {},
  "nbformat": 4,
  "nbformat_minor": 2
 }
diff --git a/examples/experimental/sophisticated_inference/si_generalized_tmaze.ipynb b/examples/experimental/sophisticated_inference/si_generalized_tmaze.ipynb
index d8bfd968..38b1c8d0 100644
--- a/examples/experimental/sophisticated_inference/si_generalized_tmaze.ipynb
+++ b/examples/experimental/sophisticated_inference/si_generalized_tmaze.ipynb
@@ -1,5 +1,12 @@
 {
  "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/infer-actively/pymdp/blob/main/examples/experimental/sophisticated_inference/si_generalized_tmaze.ipynb)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -11,7 +18,18 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "if \"google.colab\" in sys.modules:\n",
+    "    %pip install \"inferactively-pymdp\" -q"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -32,7 +50,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [
     {
@@ -67,7 +85,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -96,7 +114,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -110,7 +128,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -130,7 +148,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [
     {
@@ -202,7 +220,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -212,7 +230,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [
     {
@@ -243,25 +261,7 @@
    "source": []
   }
  ],
- "metadata": {
-  "kernelspec": {
-   "display_name": ".venv",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.11.11"
-  }
- },
+ "metadata": {},
  "nbformat": 4,
  "nbformat_minor": 2
 }
diff --git a/examples/experimental/sophisticated_inference/si_graph_world.ipynb b/examples/experimental/sophisticated_inference/si_graph_world.ipynb
index 539ebbda..3be1286c 100644
--- a/examples/experimental/sophisticated_inference/si_graph_world.ipynb
+++ b/examples/experimental/sophisticated_inference/si_graph_world.ipynb
@@ -1,5 +1,12 @@
 {
  "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/infer-actively/pymdp/blob/main/examples/experimental/sophisticated_inference/si_graph_world.ipynb)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -18,7 +25,18 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "if \"google.colab\" in sys.modules:\n",
+    "    %pip install \"inferactively-pymdp\" -q"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -30,7 +48,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -44,7 +62,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [
     {
@@ -71,7 +89,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -100,7 +118,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -120,7 +138,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [
     {
@@ -161,25 +179,7 @@
    ]
   }
  ],
- "metadata": {
-  "kernelspec": {
-   "display_name": ".venv",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.11.11"
-  }
- },
+ "metadata": {},
  "nbformat": 4,
  "nbformat_minor": 2
 }
diff --git a/examples/experimental/sophisticated_inference/si_tmaze_SIvalidation.ipynb b/examples/experimental/sophisticated_inference/si_tmaze_SIvalidation.ipynb
index 36142196..e995caeb 100644
--- a/examples/experimental/sophisticated_inference/si_tmaze_SIvalidation.ipynb
+++ b/examples/experimental/sophisticated_inference/si_tmaze_SIvalidation.ipynb
@@ -1,5 +1,12 @@
 {
  "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/infer-actively/pymdp/blob/main/examples/experimental/sophisticated_inference/si_tmaze_SIvalidation.ipynb)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -32,7 +39,18 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "if \"google.colab\" in sys.modules:\n",
+    "    %pip install \"inferactively-pymdp[nb]\" -q"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:31:52.929478Z",
@@ -133,7 +151,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:31:56.833504Z",
@@ -154,7 +172,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:31:56.848285Z",
@@ -183,7 +201,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:31:57.168745Z",
@@ -207,7 +225,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:31:57.219677Z",
@@ -224,7 +242,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 1,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:31:57.257395Z",
@@ -287,7 +305,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:31:57.879961Z",
@@ -304,7 +322,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:31:57.905455Z",
@@ -332,7 +350,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:31:57.915721Z",
@@ -349,7 +367,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:00.063840Z",
@@ -380,7 +398,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 2,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:00.074536Z",
@@ -412,7 +430,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 3,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:00.799469Z",
@@ -451,7 +469,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:01.461295Z",
@@ -531,7 +549,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 4,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:01.475190Z",
@@ -558,7 +576,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": 5,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:01.514071Z",
@@ -596,7 +614,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": 6,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:01.524335Z",
@@ -625,7 +643,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 17,
+   "execution_count": 7,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:02.047026Z",
@@ -663,7 +681,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 18,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:02.258821Z",
@@ -757,7 +775,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 19,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:02.268359Z",
@@ -779,7 +797,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 20,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:02.276198Z",
@@ -809,7 +827,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 21,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:02.470266Z",
@@ -834,7 +852,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 22,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:02.481139Z",
@@ -859,7 +877,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 23,
+   "execution_count": 8,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:02.516898Z",
@@ -924,7 +942,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 24,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:03.125390Z",
@@ -941,7 +959,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 25,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:03.135256Z",
@@ -969,7 +987,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 26,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:03.143984Z",
@@ -986,7 +1004,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 27,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:06.215016Z",
@@ -1018,7 +1036,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 28,
+   "execution_count": 9,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:06.226122Z",
@@ -1050,7 +1068,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 29,
+   "execution_count": 10,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:07.086134Z",
@@ -1089,7 +1107,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 30,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:07.908846Z",
@@ -1185,7 +1203,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 31,
+   "execution_count": 11,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:07.923120Z",
@@ -1214,7 +1232,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 32,
+   "execution_count": 12,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:07.953763Z",
@@ -1252,7 +1270,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 33,
+   "execution_count": 13,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:07.964611Z",
@@ -1282,7 +1300,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 34,
+   "execution_count": 14,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:08.565830Z",
@@ -1311,25 +1329,7 @@
    ]
   }
  ],
- "metadata": {
-  "kernelspec": {
-   "display_name": ".venv",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.11.12"
-  }
- },
+ "metadata": {},
  "nbformat": 4,
  "nbformat_minor": 2
 }
diff --git a/examples/inductive_inference/inductive_inference_example.ipynb b/examples/inductive_inference/inductive_inference_example.ipynb
index 6e3c6afc..8d31b86d 100644
--- a/examples/inductive_inference/inductive_inference_example.ipynb
+++ b/examples/inductive_inference/inductive_inference_example.ipynb
@@ -1,5 +1,12 @@
 {
  "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/infer-actively/pymdp/blob/main/examples/inductive_inference/inductive_inference_example.ipynb)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -7,6 +14,17 @@
     "### Imports"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "if \"google.colab\" in sys.modules:\n",
+    "    %pip install \"inferactively-pymdp\" -q"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -29,7 +47,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -81,7 +99,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -105,7 +123,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -119,25 +137,7 @@
    ]
   }
  ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "pymdp_dev_env",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.11.11"
-  }
- },
+ "metadata": {},
  "nbformat": 4,
  "nbformat_minor": 2
 }
diff --git a/examples/inductive_inference/inductive_inference_gridworld.ipynb b/examples/inductive_inference/inductive_inference_gridworld.ipynb
index 4d45619a..caa04b89 100644
--- a/examples/inductive_inference/inductive_inference_gridworld.ipynb
+++ b/examples/inductive_inference/inductive_inference_gridworld.ipynb
@@ -1,5 +1,12 @@
 {
  "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/infer-actively/pymdp/blob/main/examples/inductive_inference/inductive_inference_gridworld.ipynb)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -9,7 +16,18 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "if \"google.colab\" in sys.modules:\n",
+    "    %pip install \"inferactively-pymdp\" -q"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:10.138587Z",
@@ -41,7 +59,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:14.054405Z",
@@ -77,7 +95,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:14.380265Z",
@@ -108,7 +126,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 1,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:14.446460Z",
@@ -147,7 +165,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:15.110855Z",
@@ -164,7 +182,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 2,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:16.108728Z",
@@ -180,7 +198,7 @@
        "(5, 8)"
       ]
      },
-     "execution_count": 6,
+     "execution_count": 2,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -191,7 +209,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 3,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:16.112658Z",
@@ -231,7 +249,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:16.233977Z",
@@ -261,7 +279,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:16.250817Z",
@@ -282,7 +300,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 4,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:16.254701Z",
@@ -320,7 +338,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:16.383442Z",
@@ -342,7 +360,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 5,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:17.231657Z",
@@ -381,20 +399,7 @@
    ]
   }
  ],
- "metadata": {
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.11.12"
-  }
- },
+ "metadata": {},
  "nbformat": 4,
  "nbformat_minor": 2
 }
diff --git a/examples/inference_and_learning/inference_methods_comparison.ipynb b/examples/inference_and_learning/inference_methods_comparison.ipynb
index fb71c463..b80279ae 100644
--- a/examples/inference_and_learning/inference_methods_comparison.ipynb
+++ b/examples/inference_and_learning/inference_methods_comparison.ipynb
@@ -1,8 +1,26 @@
 {
  "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/infer-actively/pymdp/blob/main/examples/inference_and_learning/inference_methods_comparison.ipynb)"
+   ]
+  },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "if \"google.colab\" in sys.modules:\n",
+    "    %pip install \"inferactively-pymdp\" -q"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:18.347670Z",
@@ -34,7 +52,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:19.685687Z",
@@ -88,7 +106,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 1,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:19.995016Z",
@@ -152,7 +170,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:20.426515Z",
@@ -181,7 +199,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 2,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:21.659932Z",
@@ -213,7 +231,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 3,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:23.748100Z",
@@ -247,7 +265,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 4,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:27.439187Z",
@@ -285,7 +303,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 5,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:27.721291Z",
@@ -330,7 +348,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 6,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:27.822551Z",
@@ -399,7 +417,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 7,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:28.736227Z",
@@ -423,7 +441,7 @@
        "Text(0.5, 0.98, 'VMP smoothed beliefs')"
       ]
      },
-     "execution_count": 10,
+     "execution_count": 7,
      "metadata": {},
      "output_type": "execute_result"
     },
@@ -480,7 +498,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 8,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2026-03-06T15:32:28.889131Z",
@@ -551,25 +569,7 @@
    ]
   }
  ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "pymdp_dev_env",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.11.12"
-  }
- },
+ "metadata": {},
  "nbformat": 4,
  "nbformat_minor": 2
 }
diff --git a/examples/learning/learning_gridworld.ipynb b/examples/learning/learning_gridworld.ipynb
index 7f81924f..b3bcd0dd 100644
--- a/examples/learning/learning_gridworld.ipynb
+++ b/examples/learning/learning_gridworld.ipynb
@@ -1,5 +1,12 @@
 {
  "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/infer-actively/pymdp/blob/main/examples/learning/learning_gridworld.ipynb)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -7,6 +14,17 @@
     "### Imports"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "if \"google.colab\" in sys.modules:\n",
+    "    %pip install \"inferactively-pymdp\" -q"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 1,
diff --git a/examples/model_fitting/fitting_with_pybefit.ipynb b/examples/model_fitting/fitting_with_pybefit.ipynb
index 97ddda73..b50f503f 100644
--- a/examples/model_fitting/fitting_with_pybefit.ipynb
+++ b/examples/model_fitting/fitting_with_pybefit.ipynb
@@ -1,9 +1,17 @@
 {
  "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "0",
+   "metadata": {},
+   "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/infer-actively/pymdp/blob/main/examples/model_fitting/fitting_with_pybefit.ipynb)"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 1,
-   "id": "0",
+   "id": "1",
    "metadata": {},
    "outputs": [
     {
@@ -23,7 +31,19 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "1",
+   "id": "2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "if \"google.colab\" in sys.modules:\n",
+    "    %pip install \"inferactively-pymdp[modelfit]\" -q"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -54,7 +74,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "2",
+   "id": "4",
    "metadata": {},
    "source": [
     "## Fitting the parameters of active inference agents performing in the `T-Maze` environment\n",
@@ -66,7 +86,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "3",
+   "id": "5",
    "metadata": {},
    "source": [
     "<details>\n",
@@ -84,7 +104,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "4",
+   "id": "6",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -117,7 +137,7 @@
   {
    "cell_type": "code",
    "execution_count": 2,
-   "id": "5",
+   "id": "7",
    "metadata": {},
    "outputs": [
     {
@@ -193,7 +213,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "6",
+   "id": "8",
    "metadata": {},
    "source": [
     "### Sample from the TMaze environment and visualize the results from one block, showing the expected behavior (visit the cue, then choose an arm)"
@@ -202,7 +222,7 @@
   {
    "cell_type": "code",
    "execution_count": 3,
-   "id": "7",
+   "id": "9",
    "metadata": {},
    "outputs": [
     {
@@ -233,7 +253,7 @@
   {
    "cell_type": "code",
    "execution_count": 4,
-   "id": "8",
+   "id": "10",
    "metadata": {},
    "outputs": [
     {
@@ -275,7 +295,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "9",
+   "id": "11",
    "metadata": {},
    "source": [
     "### Inference Method 1: HMC with NUTS\n",
@@ -286,7 +306,7 @@
   {
    "cell_type": "code",
    "execution_count": 5,
-   "id": "10",
+   "id": "12",
    "metadata": {},
    "outputs": [
     {
@@ -326,7 +346,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "11",
+   "id": "13",
    "metadata": {},
    "source": [
     "### Plot each ground truth parameter alongside their posterior means (mean taken over parallel HMC samples/chains)"
@@ -335,7 +355,7 @@
   {
    "cell_type": "code",
    "execution_count": 6,
-   "id": "12",
+   "id": "14",
    "metadata": {},
    "outputs": [
     {
@@ -373,7 +393,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "13",
+   "id": "15",
    "metadata": {},
    "source": [
     "### Transform the latent parameter corresponding to the reward probability into probability space and investigate overlap between ground-truth and inferred parameter"
@@ -382,7 +402,7 @@
   {
    "cell_type": "code",
    "execution_count": 7,
-   "id": "14",
+   "id": "16",
    "metadata": {},
    "outputs": [
     {
@@ -437,7 +457,7 @@
   {
    "cell_type": "code",
    "execution_count": 8,
-   "id": "15",
+   "id": "17",
    "metadata": {},
    "outputs": [
     {
@@ -528,7 +548,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "16",
+   "id": "18",
    "metadata": {},
    "source": [
     "### Inference Method 2: Black-Box Stochastic Variational Inference\n",
@@ -539,7 +559,7 @@
   {
    "cell_type": "code",
    "execution_count": 9,
-   "id": "17",
+   "id": "19",
    "metadata": {},
    "outputs": [
     {
@@ -579,7 +599,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "18",
+   "id": "20",
    "metadata": {},
    "source": [
     "### Plot the variational free energy over time (negative ELBO)"
@@ -588,7 +608,7 @@
   {
    "cell_type": "code",
    "execution_count": 10,
-   "id": "19",
+   "id": "21",
    "metadata": {},
    "outputs": [
     {
@@ -618,7 +638,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "20",
+   "id": "22",
    "metadata": {},
    "source": [
     "### Plot each ground truth parameter alongside their posterior means (mean taken over posterior samples from the guide)"
@@ -627,7 +647,7 @@
   {
    "cell_type": "code",
    "execution_count": 11,
-   "id": "21",
+   "id": "23",
    "metadata": {},
    "outputs": [
     {
@@ -665,7 +685,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "22",
+   "id": "24",
    "metadata": {},
    "source": [
     "### Transform the latent parameter corresponding to the reward probability into probability space and investigate overlap between ground-truth and inferred parameter"
@@ -674,7 +694,7 @@
   {
    "cell_type": "code",
    "execution_count": 12,
-   "id": "23",
+   "id": "25",
    "metadata": {},
    "outputs": [
     {
@@ -729,7 +749,7 @@
   {
    "cell_type": "code",
    "execution_count": 13,
-   "id": "24",
+   "id": "26",
    "metadata": {},
    "outputs": [
     {
diff --git a/examples/sparse/sparse_benchmark.ipynb b/examples/sparse/sparse_benchmark.ipynb
index 6d094d2c..164f7b7c 100644
--- a/examples/sparse/sparse_benchmark.ipynb
+++ b/examples/sparse/sparse_benchmark.ipynb
@@ -1,5 +1,12 @@
 {
  "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/infer-actively/pymdp/blob/main/examples/sparse/sparse_benchmark.ipynb)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -17,6 +24,17 @@
     "%autoreload 2"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "if \"google.colab\" in sys.modules:\n",
+    "    %pip install \"inferactively-pymdp\" -q"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,