From c16e52689fd0c56d9b10c618e1ebcdd6191f7341 Mon Sep 17 00:00:00 2001 From: Arun Niranjan Date: Mon, 25 May 2026 14:47:34 +0100 Subject: [PATCH 1/5] Add google colab badge and install cell to model construction notebook --- .../api/model_construction_tutorial.ipynb | 82 +++++++++---------- 1 file changed, 37 insertions(+), 45 deletions(-) diff --git a/examples/api/model_construction_tutorial.ipynb b/examples/api/model_construction_tutorial.ipynb index e3d6236e..926a363b 100644 --- a/examples/api/model_construction_tutorial.ipynb +++ b/examples/api/model_construction_tutorial.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/api/model_construction_tutorial.ipynb)" + ] + }, { "cell_type": "markdown", "metadata": {}, @@ -29,7 +36,18 @@ }, { "cell_type": "code", - "execution_count": 2, + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import sys\n", + "if \"google.colab\" in sys.modules:\n", + " %pip install \"inferactively-pymdp[gpu]\" -q" + ] + }, + { + "cell_type": "code", + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -44,7 +62,7 @@ }, { "cell_type": "code", - "execution_count": 3, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -83,7 +101,7 @@ }, { "cell_type": "code", - "execution_count": 4, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -122,19 +140,9 @@ }, { "cell_type": "code", - "execution_count": 5, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Current belief about position: left\n", - "Goal position: center_left\n", - "Action chosen: move_right\n" - ] - } - ], + "outputs": [], "source": [ "gamma = 10 # deterministic behavior; make gamma smaller for stochastic behavior\n", "\n", @@ -168,25 +176,9 @@ }, { "cell_type": "code", - "execution_count": 6, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Agent 0's current belief about position: left\n", - "Agent 1's current belief about position: center_right\n", - "Agent 2's current belief about position: right\n", - "\n", - "Goal position for all agents: center_left\n", - "\n", - "Agent 0's action chosen: move_right\n", - "Agent 1's action chosen: move_left\n", - "Agent 2's action chosen: move_left\n" - ] - } - ], + "outputs": [], "source": [ "batch_size = 3 # running 3 trials or agents in parallel\n", "gamma = 10 # deterministic behavior; make gamma smaller for stochastic behavior\n", @@ -227,7 +219,7 @@ }, { "cell_type": "code", - "execution_count": 7, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -294,7 +286,7 @@ }, { "cell_type": "code", - "execution_count": 8, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -404,7 +396,7 @@ }, { "cell_type": "code", - "execution_count": 9, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -422,7 +414,7 @@ }, { "cell_type": "code", - "execution_count": 10, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -440,7 +432,7 @@ }, { "cell_type": "code", - "execution_count": 11, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -458,7 +450,7 @@ }, { "cell_type": "code", - "execution_count": 12, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -469,7 +461,7 @@ }, { "cell_type": "code", - "execution_count": 13, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -478,7 +470,7 @@ }, { "cell_type": "code", - "execution_count": 14, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -498,7 +490,7 @@ }, { "cell_type": "code", - "execution_count": 15, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -519,7 +511,7 @@ }, { "cell_type": "code", - "execution_count": 16, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -529,7 +521,7 @@ ], "metadata": { "kernelspec": { - "display_name": ".venv311", + "display_name": "inferactively-pymdp", "language": "python", "name": "python3" }, @@ -543,7 +535,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.11.12" + "version": "3.13.11" } }, "nbformat": 4, From e7cfd15c8f5c396fee614119daba5a270c9e39cc Mon Sep 17 00:00:00 2001 From: Arun Niranjan Date: Mon, 25 May 2026 14:56:16 +0100 Subject: [PATCH 2/5] install vanilla pymdp as colab already ships with JAX --- .../api/model_construction_tutorial.ipynb | 89 +++++++++++++++---- 1 file changed, 70 insertions(+), 19 deletions(-) diff --git a/examples/api/model_construction_tutorial.ipynb b/examples/api/model_construction_tutorial.ipynb index 926a363b..c36b3717 100644 --- a/examples/api/model_construction_tutorial.ipynb +++ b/examples/api/model_construction_tutorial.ipynb @@ -42,12 +42,12 @@ "source": [ "import sys\n", "if \"google.colab\" in sys.modules:\n", - " %pip install \"inferactively-pymdp[gpu]\" -q" + " %pip install \"inferactively-pymdp\" -q" ] }, { "cell_type": "code", - "execution_count": null, + "execution_count": 18, "metadata": {}, "outputs": [], "source": [ @@ -62,7 +62,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 19, "metadata": {}, "outputs": [], "source": [ @@ -101,7 +101,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 20, "metadata": {}, "outputs": [], "source": [ @@ -140,9 +140,27 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 21, "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Current belief about position: left\n", + "Goal position: center_left\n", + "Action chosen: move_right\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/tmp/ipykernel_1062531/1053406980.py:4: UserWarning: A JAX array is being set as static! This can result in unexpected behavior and is usually a mistake to do.\n", + " agent = Agent(**model, gamma=gamma)\n" + ] + } + ], "source": [ "gamma = 10 # deterministic behavior; make gamma smaller for stochastic behavior\n", "\n", @@ -176,9 +194,33 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 22, "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Agent 0's current belief about position: left\n", + "Agent 1's current belief about position: center_right\n", + "Agent 2's current belief about position: right\n", + "\n", + "Goal position for all agents: center_left\n", + "\n", + "Agent 0's action chosen: move_right\n", + "Agent 1's action chosen: move_left\n", + "Agent 2's action chosen: move_left\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/tmp/ipykernel_1062531/1678643109.py:5: UserWarning: A JAX array is being set as static! This can result in unexpected behavior and is usually a mistake to do.\n", + " agent = Agent(**model, batch_size=batch_size, gamma=gamma)\n" + ] + } + ], "source": [ "batch_size = 3 # running 3 trials or agents in parallel\n", "gamma = 10 # deterministic behavior; make gamma smaller for stochastic behavior\n", @@ -219,7 +261,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 23, "metadata": {}, "outputs": [], "source": [ @@ -286,7 +328,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 24, "metadata": {}, "outputs": [], "source": [ @@ -396,9 +438,18 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 25, "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/tmp/ipykernel_1062531/53738707.py:3: UserWarning: A JAX array is being set as static! This can result in unexpected behavior and is usually a mistake to do.\n", + " agent = Agent(**model, learn_A=False, learn_B=False, gamma=gamma, sampling_mode=\"full\")\n" + ] + } + ], "source": [ "gamma = 1.0\n", "\n", @@ -414,7 +465,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 26, "metadata": {}, "outputs": [], "source": [ @@ -432,7 +483,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 27, "metadata": {}, "outputs": [], "source": [ @@ -450,7 +501,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 28, "metadata": {}, "outputs": [], "source": [ @@ -461,7 +512,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 29, "metadata": {}, "outputs": [], "source": [ @@ -470,7 +521,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 30, "metadata": {}, "outputs": [], "source": [ @@ -490,7 +541,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 31, "metadata": {}, "outputs": [], "source": [ @@ -511,7 +562,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 32, "metadata": {}, "outputs": [], "source": [ From 0e342f057de99b6cdea235915002f0381593c939 Mon Sep 17 00:00:00 2001 From: Arun Niranjan Date: Mon, 25 May 2026 15:04:49 +0100 Subject: [PATCH 3/5] Clear outputs --- .../api/model_construction_tutorial.ipynb | 87 ++++--------------- 1 file changed, 18 insertions(+), 69 deletions(-) diff --git a/examples/api/model_construction_tutorial.ipynb b/examples/api/model_construction_tutorial.ipynb index c36b3717..e315ed1c 100644 --- a/examples/api/model_construction_tutorial.ipynb +++ b/examples/api/model_construction_tutorial.ipynb @@ -47,7 +47,7 @@ }, { "cell_type": "code", - "execution_count": 18, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -62,7 +62,7 @@ }, { "cell_type": "code", - "execution_count": 19, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -101,7 +101,7 @@ }, { "cell_type": "code", - "execution_count": 20, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -140,27 +140,9 @@ }, { "cell_type": "code", - "execution_count": 21, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Current belief about position: left\n", - "Goal position: center_left\n", - "Action chosen: move_right\n" - ] - }, - { - "name": "stderr", - "output_type": "stream", - "text": [ - "/tmp/ipykernel_1062531/1053406980.py:4: UserWarning: A JAX array is being set as static! This can result in unexpected behavior and is usually a mistake to do.\n", - " agent = Agent(**model, gamma=gamma)\n" - ] - } - ], + "outputs": [], "source": [ "gamma = 10 # deterministic behavior; make gamma smaller for stochastic behavior\n", "\n", @@ -194,33 +176,9 @@ }, { "cell_type": "code", - "execution_count": 22, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Agent 0's current belief about position: left\n", - "Agent 1's current belief about position: center_right\n", - "Agent 2's current belief about position: right\n", - "\n", - "Goal position for all agents: center_left\n", - "\n", - "Agent 0's action chosen: move_right\n", - "Agent 1's action chosen: move_left\n", - "Agent 2's action chosen: move_left\n" - ] - }, - { - "name": "stderr", - "output_type": "stream", - "text": [ - "/tmp/ipykernel_1062531/1678643109.py:5: UserWarning: A JAX array is being set as static! This can result in unexpected behavior and is usually a mistake to do.\n", - " agent = Agent(**model, batch_size=batch_size, gamma=gamma)\n" - ] - } - ], + "outputs": [], "source": [ "batch_size = 3 # running 3 trials or agents in parallel\n", "gamma = 10 # deterministic behavior; make gamma smaller for stochastic behavior\n", @@ -261,7 +219,7 @@ }, { "cell_type": "code", - "execution_count": 23, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -328,7 +286,7 @@ }, { "cell_type": "code", - "execution_count": 24, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -438,18 +396,9 @@ }, { "cell_type": "code", - "execution_count": 25, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "/tmp/ipykernel_1062531/53738707.py:3: UserWarning: A JAX array is being set as static! This can result in unexpected behavior and is usually a mistake to do.\n", - " agent = Agent(**model, learn_A=False, learn_B=False, gamma=gamma, sampling_mode=\"full\")\n" - ] - } - ], + "outputs": [], "source": [ "gamma = 1.0\n", "\n", @@ -465,7 +414,7 @@ }, { "cell_type": "code", - "execution_count": 26, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -483,7 +432,7 @@ }, { "cell_type": "code", - "execution_count": 27, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -501,7 +450,7 @@ }, { "cell_type": "code", - "execution_count": 28, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -512,7 +461,7 @@ }, { "cell_type": "code", - "execution_count": 29, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -521,7 +470,7 @@ }, { "cell_type": "code", - "execution_count": 30, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -541,7 +490,7 @@ }, { "cell_type": "code", - "execution_count": 31, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -562,7 +511,7 @@ }, { "cell_type": "code", - "execution_count": 32, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ From f9b07f4eb2877e677c8473de810b2896582b5956 Mon Sep 17 00:00:00 2001 From: Arun Niranjan Date: Mon, 25 May 2026 15:06:16 +0100 Subject: [PATCH 4/5] Generate outputs --- .../api/model_construction_tutorial.ipynb | 89 +++++++++++++++---- 1 file changed, 70 insertions(+), 19 deletions(-) diff --git a/examples/api/model_construction_tutorial.ipynb b/examples/api/model_construction_tutorial.ipynb index e315ed1c..762cfe19 100644 --- a/examples/api/model_construction_tutorial.ipynb +++ b/examples/api/model_construction_tutorial.ipynb @@ -36,7 +36,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 33, "metadata": {}, "outputs": [], "source": [ @@ -47,7 +47,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 34, "metadata": {}, "outputs": [], "source": [ @@ -62,7 +62,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 35, "metadata": {}, "outputs": [], "source": [ @@ -101,7 +101,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 36, "metadata": {}, "outputs": [], "source": [ @@ -140,9 +140,27 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 37, "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Current belief about position: left\n", + "Goal position: center_left\n", + "Action chosen: move_right\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/tmp/ipykernel_1062531/1053406980.py:4: UserWarning: A JAX array is being set as static! This can result in unexpected behavior and is usually a mistake to do.\n", + " agent = Agent(**model, gamma=gamma)\n" + ] + } + ], "source": [ "gamma = 10 # deterministic behavior; make gamma smaller for stochastic behavior\n", "\n", @@ -176,9 +194,33 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 38, "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Agent 0's current belief about position: left\n", + "Agent 1's current belief about position: center_right\n", + "Agent 2's current belief about position: right\n", + "\n", + "Goal position for all agents: center_left\n", + "\n", + "Agent 0's action chosen: move_right\n", + "Agent 1's action chosen: move_left\n", + "Agent 2's action chosen: move_left\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/tmp/ipykernel_1062531/1678643109.py:5: UserWarning: A JAX array is being set as static! This can result in unexpected behavior and is usually a mistake to do.\n", + " agent = Agent(**model, batch_size=batch_size, gamma=gamma)\n" + ] + } + ], "source": [ "batch_size = 3 # running 3 trials or agents in parallel\n", "gamma = 10 # deterministic behavior; make gamma smaller for stochastic behavior\n", @@ -219,7 +261,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 39, "metadata": {}, "outputs": [], "source": [ @@ -286,7 +328,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 40, "metadata": {}, "outputs": [], "source": [ @@ -396,9 +438,18 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 41, "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/tmp/ipykernel_1062531/53738707.py:3: UserWarning: A JAX array is being set as static! This can result in unexpected behavior and is usually a mistake to do.\n", + " agent = Agent(**model, learn_A=False, learn_B=False, gamma=gamma, sampling_mode=\"full\")\n" + ] + } + ], "source": [ "gamma = 1.0\n", "\n", @@ -414,7 +465,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 42, "metadata": {}, "outputs": [], "source": [ @@ -432,7 +483,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 43, "metadata": {}, "outputs": [], "source": [ @@ -450,7 +501,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 44, "metadata": {}, "outputs": [], "source": [ @@ -461,7 +512,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 45, "metadata": {}, "outputs": [], "source": [ @@ -470,7 +521,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 46, "metadata": {}, "outputs": [], "source": [ @@ -490,7 +541,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 47, "metadata": {}, "outputs": [], "source": [ @@ -511,7 +562,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 48, "metadata": {}, "outputs": [], "source": [ From 8e21022d63e168f773a881bd7d2dd96bb1856d38 Mon Sep 17 00:00:00 2001 From: Arun Niranjan Date: Tue, 26 May 2026 09:57:42 +0100 Subject: [PATCH 5/5] run precommit to sanitise notebook --- .../api/model_construction_tutorial.ipynb | 52 ++++++------------- 1 file changed, 17 insertions(+), 35 deletions(-) diff --git a/examples/api/model_construction_tutorial.ipynb b/examples/api/model_construction_tutorial.ipynb index 762cfe19..71330818 100644 --- a/examples/api/model_construction_tutorial.ipynb +++ b/examples/api/model_construction_tutorial.ipynb @@ -36,7 +36,7 @@ }, { "cell_type": "code", - "execution_count": 33, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -47,7 +47,7 @@ }, { "cell_type": "code", - "execution_count": 34, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -62,7 +62,7 @@ }, { "cell_type": "code", - "execution_count": 35, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -101,7 +101,7 @@ }, { "cell_type": "code", - "execution_count": 36, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -140,7 +140,7 @@ }, { "cell_type": "code", - "execution_count": 37, + "execution_count": 1, "metadata": {}, "outputs": [ { @@ -194,7 +194,7 @@ }, { "cell_type": "code", - "execution_count": 38, + "execution_count": 2, "metadata": {}, "outputs": [ { @@ -261,7 +261,7 @@ }, { "cell_type": "code", - "execution_count": 39, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -328,7 +328,7 @@ }, { "cell_type": "code", - "execution_count": 40, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -438,7 +438,7 @@ }, { "cell_type": "code", - "execution_count": 41, + "execution_count": 3, "metadata": {}, "outputs": [ { @@ -465,7 +465,7 @@ }, { "cell_type": "code", - "execution_count": 42, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -483,7 +483,7 @@ }, { "cell_type": "code", - "execution_count": 43, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -501,7 +501,7 @@ }, { "cell_type": "code", - "execution_count": 44, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -512,7 +512,7 @@ }, { "cell_type": "code", - "execution_count": 45, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -521,7 +521,7 @@ }, { "cell_type": "code", - "execution_count": 46, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -541,7 +541,7 @@ }, { "cell_type": "code", - "execution_count": 47, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -562,7 +562,7 @@ }, { "cell_type": "code", - "execution_count": 48, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -570,25 +570,7 @@ ] } ], - "metadata": { - "kernelspec": { - "display_name": "inferactively-pymdp", - "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.13.11" - } - }, + "metadata": {}, "nbformat": 4, "nbformat_minor": 2 }