play_env_as_human.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/nathanwispinski/meta-rl/blob/main/play_env_as_human.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aVhgAnD5gFGy"
      },
      "source": [
        "# play_env_as_human.ipynb\n",
        "\n",
        "This is a Google Colab notebook to demo the meta-rl multi-armed bandit environment.\n",
        "\n",
        "For more details, see the GitHub repository (https://github.com/nathanwispinski/meta-rl)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mHdmWcLfgRhb"
      },
      "source": [
        "# Colab setup"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Clone GitHub repository.\n",
        "!git clone https://github.com/nathanwispinski/meta-rl"
      ],
      "metadata": {
        "cellView": "form",
        "id": "G_L6U1HvKUkx",
        "outputId": "782fb4fb-657c-4471-e5cd-12128bde0c1e",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Cloning into 'meta-rl'...\n",
            "remote: Enumerating objects: 60, done.\u001b[K\n",
            "remote: Counting objects: 100% (60/60), done.\u001b[K\n",
            "remote: Compressing objects: 100% (47/47), done.\u001b[K\n",
            "remote: Total 60 (delta 23), reused 38 (delta 12), pack-reused 0\u001b[K\n",
            "Unpacking objects: 100% (60/60), 202.89 KiB | 3.22 MiB/s, done.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Change working directory to cloned repository (i.e., /content/meta-rl/).\n",
        "%cd meta-rl"
      ],
      "metadata": {
        "cellView": "form",
        "id": "5-MC8AkuKV-H",
        "outputId": "43ed182a-0687-49cd-f486-57163b9a4c31",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "/content/meta-rl\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Install dependencies from `requirements.txt`.\n",
        "!pip install -r requirements.txt"
      ],
      "metadata": {
        "cellView": "form",
        "id": "a8299gC5KXhA",
        "outputId": "5518bbde-bc15-4689-be5c-c8d877c7d0c4",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        }
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting absl_py==1.3.0\n",
            "  Downloading absl_py-1.3.0-py3-none-any.whl (124 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m124.6/124.6 KB\u001b[0m \u001b[31m2.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting chex==0.1.5\n",
            "  Downloading chex-0.1.5-py3-none-any.whl (85 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m85.3/85.3 KB\u001b[0m \u001b[31m3.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting dm_haiku==0.0.9\n",
            "  Downloading dm_haiku-0.0.9-py3-none-any.whl (352 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m352.1/352.1 KB\u001b[0m \u001b[31m14.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting jax==0.3.25\n",
            "  Downloading jax-0.3.25.tar.gz (1.1 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.1/1.1 MB\u001b[0m \u001b[31m32.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "Collecting jaxlib==0.3.25\n",
            "  Downloading jaxlib-0.3.25-cp38-cp38-manylinux2014_x86_64.whl (71.2 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m71.2/71.2 MB\u001b[0m \u001b[31m6.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting matplotlib==3.6.2\n",
            "  Downloading matplotlib-3.6.2-cp38-cp38-manylinux_2_12_x86_64.manylinux2010_x86_64.whl (9.4 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m9.4/9.4 MB\u001b[0m \u001b[31m18.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting ml_collections==0.1.1\n",
            "  Downloading ml_collections-0.1.1.tar.gz (77 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m77.9/77.9 KB\u001b[0m \u001b[31m3.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "Requirement already satisfied: numpy==1.22.4 in /usr/local/lib/python3.8/dist-packages (from -r requirements.txt (line 8)) (1.22.4)\n",
            "Collecting optax==0.1.4\n",
            "  Downloading optax-0.1.4-py3-none-any.whl (154 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m154.9/154.9 KB\u001b[0m \u001b[31m11.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting rlax==0.1.4\n",
            "  Downloading rlax-0.1.4-py3-none-any.whl (128 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m128.7/128.7 KB\u001b[0m \u001b[31m6.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting scipy==1.9.3\n",
            "  Downloading scipy-1.9.3-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (33.8 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m33.8/33.8 MB\u001b[0m \u001b[31m15.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: toolz>=0.9.0 in /usr/local/lib/python3.8/dist-packages (from chex==0.1.5->-r requirements.txt (line 2)) (0.12.0)\n",
            "Requirement already satisfied: dm-tree>=0.1.5 in /usr/local/lib/python3.8/dist-packages (from chex==0.1.5->-r requirements.txt (line 2)) (0.1.8)\n",
            "Requirement already satisfied: tabulate>=0.8.9 in /usr/local/lib/python3.8/dist-packages (from dm_haiku==0.0.9->-r requirements.txt (line 3)) (0.8.10)\n",
            "Collecting jmp>=0.0.2\n",
            "  Downloading jmp-0.0.4-py3-none-any.whl (18 kB)\n",
            "Requirement already satisfied: opt_einsum in /usr/local/lib/python3.8/dist-packages (from jax==0.3.25->-r requirements.txt (line 4)) (3.3.0)\n",
            "Requirement already satisfied: typing_extensions in /usr/local/lib/python3.8/dist-packages (from jax==0.3.25->-r requirements.txt (line 4)) (4.5.0)\n",
            "Requirement already satisfied: pyparsing>=2.2.1 in /usr/local/lib/python3.8/dist-packages (from matplotlib==3.6.2->-r requirements.txt (line 6)) (3.0.9)\n",
            "Collecting contourpy>=1.0.1\n",
            "  Downloading contourpy-1.0.7-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (300 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m300.0/300.0 KB\u001b[0m \u001b[31m20.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.8/dist-packages (from matplotlib==3.6.2->-r requirements.txt (line 6)) (1.4.4)\n",
            "Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.8/dist-packages (from matplotlib==3.6.2->-r requirements.txt (line 6)) (0.11.0)\n",
            "Requirement already satisfied: fonttools>=4.22.0 in /usr/local/lib/python3.8/dist-packages (from matplotlib==3.6.2->-r requirements.txt (line 6)) (4.38.0)\n",
            "Requirement already satisfied: python-dateutil>=2.7 in /usr/local/lib/python3.8/dist-packages (from matplotlib==3.6.2->-r requirements.txt (line 6)) (2.8.2)\n",
            "Requirement already satisfied: pillow>=6.2.0 in /usr/local/lib/python3.8/dist-packages (from matplotlib==3.6.2->-r requirements.txt (line 6)) (8.4.0)\n",
            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.8/dist-packages (from matplotlib==3.6.2->-r requirements.txt (line 6)) (23.0)\n",
            "Requirement already satisfied: PyYAML in /usr/local/lib/python3.8/dist-packages (from ml_collections==0.1.1->-r requirements.txt (line 7)) (6.0)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.8/dist-packages (from ml_collections==0.1.1->-r requirements.txt (line 7)) (1.15.0)\n",
            "Requirement already satisfied: contextlib2 in /usr/local/lib/python3.8/dist-packages (from ml_collections==0.1.1->-r requirements.txt (line 7)) (0.5.5)\n",
            "Collecting dm-env\n",
            "  Downloading dm_env-1.6-py3-none-any.whl (26 kB)\n",
            "Collecting distrax>=0.0.2\n",
            "  Downloading distrax-0.1.3-py3-none-any.whl (317 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m318.0/318.0 KB\u001b[0m \u001b[31m23.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: tensorflow-probability>=0.15.0 in /usr/local/lib/python3.8/dist-packages (from distrax>=0.0.2->rlax==0.1.4->-r requirements.txt (line 10)) (0.19.0)\n",
            "Requirement already satisfied: decorator in /usr/local/lib/python3.8/dist-packages (from tensorflow-probability>=0.15.0->distrax>=0.0.2->rlax==0.1.4->-r requirements.txt (line 10)) (4.4.2)\n",
            "Requirement already satisfied: cloudpickle>=1.3 in /usr/local/lib/python3.8/dist-packages (from tensorflow-probability>=0.15.0->distrax>=0.0.2->rlax==0.1.4->-r requirements.txt (line 10)) (2.2.1)\n",
            "Requirement already satisfied: gast>=0.3.2 in /usr/local/lib/python3.8/dist-packages (from tensorflow-probability>=0.15.0->distrax>=0.0.2->rlax==0.1.4->-r requirements.txt (line 10)) (0.4.0)\n",
            "Building wheels for collected packages: jax, ml_collections\n",
            "  Building wheel for jax (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for jax: filename=jax-0.3.25-py3-none-any.whl size=1308510 sha256=3b34d767a97c821cba2fab61d2f17135ee647cf8eac4151dc96a79ec89952d2c\n",
            "  Stored in directory: /root/.cache/pip/wheels/ac/cc/1f/327114e0ec22ee90b4db16c01e1fc843883e18501098a7bbec\n",
            "  Building wheel for ml_collections (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for ml_collections: filename=ml_collections-0.1.1-py3-none-any.whl size=94524 sha256=12891ba50358364d5c154e08626e073a8c6b4ca1dbe668d2dd53082237c13e38\n",
            "  Stored in directory: /root/.cache/pip/wheels/6d/9f/a9/9e8309035a5bf09ed9086bbca8c9b74cb6413d3eb203e2bc8c\n",
            "Successfully built jax ml_collections\n",
            "Installing collected packages: scipy, jmp, contourpy, absl_py, ml_collections, matplotlib, jaxlib, jax, dm_haiku, dm-env, chex, optax, distrax, rlax\n",
            "  Attempting uninstall: scipy\n",
            "    Found existing installation: scipy 1.7.3\n",
            "    Uninstalling scipy-1.7.3:\n",
            "      Successfully uninstalled scipy-1.7.3\n",
            "  Attempting uninstall: absl_py\n",
            "    Found existing installation: absl-py 1.4.0\n",
            "    Uninstalling absl-py-1.4.0:\n",
            "      Successfully uninstalled absl-py-1.4.0\n",
            "  Attempting uninstall: matplotlib\n",
            "    Found existing installation: matplotlib 3.5.3\n",
            "    Uninstalling matplotlib-3.5.3:\n",
            "      Successfully uninstalled matplotlib-3.5.3\n",
            "  Attempting uninstall: jaxlib\n",
            "    Found existing installation: jaxlib 0.4.4+cuda11.cudnn82\n",
            "    Uninstalling jaxlib-0.4.4+cuda11.cudnn82:\n",
            "      Successfully uninstalled jaxlib-0.4.4+cuda11.cudnn82\n",
            "  Attempting uninstall: jax\n",
            "    Found existing installation: jax 0.4.4\n",
            "    Uninstalling jax-0.4.4:\n",
            "      Successfully uninstalled jax-0.4.4\n",
            "Successfully installed absl_py-1.3.0 chex-0.1.5 contourpy-1.0.7 distrax-0.1.3 dm-env-1.6 dm_haiku-0.0.9 jax-0.3.25 jaxlib-0.3.25 jmp-0.0.4 matplotlib-3.6.2 ml_collections-0.1.1 optax-0.1.4 rlax-0.1.4 scipy-1.9.3\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.colab-display-data+json": {
              "pip_warning": {
                "packages": [
                  "matplotlib",
                  "mpl_toolkits"
                ]
              }
            }
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Import dependencies"
      ],
      "metadata": {
        "id": "fiL-QQpVKbDS"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Import dependencies after install.\n",
        "import numpy as np\n",
        "import time\n",
        "from IPython.display import display, clear_output\n",
        "import ipywidgets as widgets\n",
        "\n",
        "import modules.envs as envs"
      ],
      "metadata": {
        "cellView": "form",
        "id": "ouCP3wWcKZVG"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Create an environment to play."
      ],
      "metadata": {
        "id": "wtKvKrEZKlcD"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "f_0q4wR2gd6B",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Create a config for the environment.\n",
        "#@markdown Available `reward_structure` configs are: `independent`, and `correlated`.\n",
        "\n",
        "steps_per_episode = 10 #@param {type:\"integer\"}\n",
        "num_arms = 2 #@param {type:\"integer\"}\n",
        "reward_structure = \"correlated\" #@param {type:\"string\"}\n",
        "total_episodes = 10 #@param {type:\"integer\"}\n",
        "\n",
        "env_config = {\n",
        "    \"steps_per_episode\": steps_per_episode,\n",
        "    \"num_arms\": num_arms,\n",
        "    \"reward_structure\": reward_structure,\n",
        "}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "cellView": "form",
        "id": "VakNQGzvlnvc"
      },
      "outputs": [],
      "source": [
        "#@title Create environment.\n",
        "env = envs.create_env(env_config=env_config)\n",
        "valid_action_range = np.arange(env.num_actions)\n",
        "observation = env.reset()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8atWqT0agqqw"
      },
      "source": [
        "# Play"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 512,
          "referenced_widgets": [
            "c0ddd4dd21a74a2995334ef5bbd25386",
            "42d6fef3ce54402e8144dec553efe2c7",
            "d4c935d928584047b2ad313dd68707a8",
            "cbf0d39d270444d19c6d54b82b75a9de",
            "b0e5acda15f8499cb24701cf9b7599d4",
            "0303109117694eb59e48f37eb431a398",
            "c65f51f037a54c6381b9cb46d1080a3d",
            "a763b522a9f24134bb5d5af38cee883c",
            "1e73e58b66354574b2e401a2817e56ea",
            "136413b48a9b4d76b7a8dd903a90a5de",
            "46f9514b0e484b51901fd2110e26ac2b"
          ]
        },
        "id": "Nr-gndzWigKG",
        "outputId": "07657509-f2b1-4980-8825-e6b4311493ed"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Button(description='Arm 1', style=ButtonStyle())"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "c0ddd4dd21a74a2995334ef5bbd25386"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Button(description='Arm 2', style=ButtonStyle())"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "cbf0d39d270444d19c6d54b82b75a9de"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Button(description='Clear history', style=ButtonStyle())"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "c65f51f037a54c6381b9cb46d1080a3d"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Output()"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "136413b48a9b4d76b7a8dd903a90a5de"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "{'vector_input': array([0., 0., 0.])}\n",
            "Step: 0; Action: 0; Reward: 1.0; {'vector_input': array([1., 0., 1.])}\n",
            "Step: 1; Action: 0; Reward: 0.0; {'vector_input': array([1., 0., 0.])}\n",
            "Step: 2; Action: 0; Reward: 1.0; {'vector_input': array([1., 0., 1.])}\n",
            "Step: 3; Action: 1; Reward: 0.0; {'vector_input': array([0., 1., 0.])}\n",
            "Step: 4; Action: 1; Reward: 0.0; {'vector_input': array([0., 1., 0.])}\n",
            "Step: 5; Action: 1; Reward: 1.0; {'vector_input': array([0., 1., 1.])}\n",
            "Step: 6; Action: 0; Reward: 1.0; {'vector_input': array([1., 0., 1.])}\n",
            "Step: 7; Action: 0; Reward: 1.0; {'vector_input': array([1., 0., 1.])}\n",
            "Step: 8; Action: 0; Reward: 1.0; {'vector_input': array([1., 0., 1.])}\n",
            "Step: 9; Action: 0; Reward: 1.0; {'vector_input': array([1., 0., 1.])}\n",
            "Episode done. Arm win probs were: [0.90476954 0.09523046]\n",
            "Step: 0; Action: 0; Reward: 0.0; {'vector_input': array([1., 0., 0.])}\n",
            "Step: 1; Action: 0; Reward: 1.0; {'vector_input': array([1., 0., 1.])}\n",
            "Step: 2; Action: 0; Reward: 1.0; {'vector_input': array([1., 0., 1.])}\n",
            "Step: 3; Action: 1; Reward: 0.0; {'vector_input': array([0., 1., 0.])}\n",
            "Step: 4; Action: 1; Reward: 1.0; {'vector_input': array([0., 1., 1.])}\n",
            "Step: 5; Action: 1; Reward: 0.0; {'vector_input': array([0., 1., 0.])}\n",
            "Step: 6; Action: 0; Reward: 1.0; {'vector_input': array([1., 0., 1.])}\n",
            "Step: 7; Action: 0; Reward: 0.0; {'vector_input': array([1., 0., 0.])}\n",
            "Step: 8; Action: 0; Reward: 0.0; {'vector_input': array([1., 0., 0.])}\n",
            "Step: 9; Action: 0; Reward: 1.0; {'vector_input': array([1., 0., 1.])}\n",
            "Episode done. Arm win probs were: [0.59072432 0.40927568]\n"
          ]
        }
      ],
      "source": [
        "#@title Play environment as a human.\n",
        "\n",
        "output = widgets.Output()\n",
        "\n",
        "# Make one button per bandit arm\n",
        "num_arms = env_config['num_arms']\n",
        "buttons = []\n",
        "click_fns = []\n",
        "for i in range(num_arms):\n",
        "    arm_name = \"Arm \" + str(i + 1)\n",
        "    button = widgets.Button(description=arm_name)\n",
        "    def button_click(a, action=i):\n",
        "        next_observation, reward, done, info = env.step(action)\n",
        "        step = interactive_info['step']\n",
        "        print(f'Step: {step}; Action: {action}; Reward: {reward}; {next_observation}')\n",
        "        if done:\n",
        "            interactive_info['step'] = 0\n",
        "        else:\n",
        "            interactive_info['step'] += 1\n",
        "        interactive_info['observation'] = next_observation\n",
        "        if done:\n",
        "            print(f'Episode done. Arm win probs were: {env._arm_probs}')\n",
        "            env.reset()\n",
        "    buttons.append(button)\n",
        "    click_fns.append(button_click)\n",
        "\n",
        "for (button, click_fn) in zip(buttons, click_fns):\n",
        "    button.on_click(click_fn)\n",
        "\n",
        "# Make a clear history button\n",
        "button_0 = widgets.Button(description=\"Clear history\")\n",
        "def button_0_click(b):\n",
        "    clear_output()\n",
        "    display(*buttons)\n",
        "button_0.on_click(button_0_click)\n",
        "\n",
        "interactive_info = {\n",
        "    'step': 0,\n",
        "    'observation': None,\n",
        "    'next_observation': None,\n",
        "    'reward': None,\n",
        "    'done': None,\n",
        "    'info': None,\n",
        "    'action': None\n",
        "    }\n",
        "\n",
        "buttons.append(button_0)\n",
        "buttons.append(output)\n",
        "\n",
        "display(*buttons)\n",
        "\n",
        "interactive_info['observation'] = env.reset()\n",
        "print(interactive_info['observation'])\n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "include_colab_link": true
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "meta-rl-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.9.16"
    },
    "orig_nbformat": 4,
    "vscode": {
      "interpreter": {
        "hash": "0e4463702277b73b75dad737c443b2ac1ecf2c3f192baf2a92474f48480d9455"
      }
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "c0ddd4dd21a74a2995334ef5bbd25386": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ButtonModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ButtonModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ButtonView",
            "button_style": "",
            "description": "Arm 1",
            "disabled": false,
            "icon": "",
            "layout": "IPY_MODEL_42d6fef3ce54402e8144dec553efe2c7",
            "style": "IPY_MODEL_d4c935d928584047b2ad313dd68707a8",
            "tooltip": ""
          }
        },
        "42d6fef3ce54402e8144dec553efe2c7": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "d4c935d928584047b2ad313dd68707a8": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ButtonStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ButtonStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "button_color": null,
            "font_weight": ""
          }
        },
        "cbf0d39d270444d19c6d54b82b75a9de": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ButtonModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ButtonModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ButtonView",
            "button_style": "",
            "description": "Arm 2",
            "disabled": false,
            "icon": "",
            "layout": "IPY_MODEL_b0e5acda15f8499cb24701cf9b7599d4",
            "style": "IPY_MODEL_0303109117694eb59e48f37eb431a398",
            "tooltip": ""
          }
        },
        "b0e5acda15f8499cb24701cf9b7599d4": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "0303109117694eb59e48f37eb431a398": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ButtonStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ButtonStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "button_color": null,
            "font_weight": ""
          }
        },
        "c65f51f037a54c6381b9cb46d1080a3d": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ButtonModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ButtonModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ButtonView",
            "button_style": "",
            "description": "Clear history",
            "disabled": false,
            "icon": "",
            "layout": "IPY_MODEL_a763b522a9f24134bb5d5af38cee883c",
            "style": "IPY_MODEL_1e73e58b66354574b2e401a2817e56ea",
            "tooltip": ""
          }
        },
        "a763b522a9f24134bb5d5af38cee883c": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "1e73e58b66354574b2e401a2817e56ea": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ButtonStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ButtonStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "button_color": null,
            "font_weight": ""
          }
        },
        "136413b48a9b4d76b7a8dd903a90a5de": {
          "model_module": "@jupyter-widgets/output",
          "model_name": "OutputModel",
          "model_module_version": "1.0.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/output",
            "_model_module_version": "1.0.0",
            "_model_name": "OutputModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/output",
            "_view_module_version": "1.0.0",
            "_view_name": "OutputView",
            "layout": "IPY_MODEL_46f9514b0e484b51901fd2110e26ac2b",
            "msg_id": "",
            "outputs": []
          }
        },
        "46f9514b0e484b51901fd2110e26ac2b": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        }
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}