train.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/train.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Xh94QYkVL4Nf"
      },
      "source": [
        "# train.ipynb\n",
        "\n",
        "This is a Google Colab notebook to demo model training of a recurrent neural network in a two-armed bandit task using reinforcement learning.\n",
        "\n",
        "This is a single-threaded version of model training, and may take a while depending on your training settings (around 15 minutes in a CPU Colab instance in my experience).\n",
        "\n",
        "For more details, see the GitHub repository (https://github.com/nathanwispinski/meta-rl)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dIZ_cxlLMDRi"
      },
      "source": [
        "# Colab setup"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Clone GitHub repository.\n",
        "!git clone https://github.com/nathanwispinski/meta-rl"
      ],
      "metadata": {
        "id": "pwMXd0gcwkn5",
        "cellView": "form",
        "outputId": "63d4e656-0ee4-4fc6-89da-e88b3483972a",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Cloning into 'meta-rl'...\n",
            "remote: Enumerating objects: 54, done.\u001b[K\n",
            "remote: Counting objects: 100% (54/54), done.\u001b[K\n",
            "remote: Compressing objects: 100% (39/39), done.\u001b[K\n",
            "remote: Total 54 (delta 20), reused 41 (delta 14), pack-reused 0\u001b[K\n",
            "Unpacking objects: 100% (54/54), 143.75 KiB | 920.00 KiB/s, done.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 53
        },
        "id": "J9alcZG3G3-6",
        "outputId": "a183b450-0016-423e-b8d4-4063e8ad1a02",
        "cellView": "form"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "/content/meta-rl\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'/content/meta-rl'"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 2
        }
      ],
      "source": [
        "#@title Change working directory to cloned repository (i.e., /content/meta-rl/).\n",
        "%cd meta-rl"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "TZ2kq_QOvaTL",
        "cellView": "form",
        "outputId": "bc300b85-cc48-4eef-86d4-ea1f92e0f002",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        }
      },
      "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[31m6.1 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[31m10.4 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[31m32.1 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[31m64.5 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[31m11.0 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[31m76.7 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[31m10.7 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[31m21.1 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[31m17.7 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[31m41.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",
            "Collecting jmp>=0.0.2\n",
            "  Downloading jmp-0.0.4-py3-none-any.whl (18 kB)\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",
            "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",
            "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[31m36.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement 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: 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: 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: 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: 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",
            "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: 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 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[31m29.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting dm-env\n",
            "  Downloading dm_env-1.6-py3-none-any.whl (26 kB)\n",
            "Requirement 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: 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",
            "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",
            "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=eefaffc355ac512d64df19673f99b7af77c34567f5a2b8af94b5051e24c6f106\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=87ef79ac822bafef33781f4221c99ea407c32a84ddb6451dfb257b95a8669ee9\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": {}
        }
      ],
      "source": [
        "# @title Install dependencies from `requirements.txt`.\n",
        "!pip install -r requirements.txt"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Import dependencies"
      ],
      "metadata": {
        "id": "T6TcF0-QxEzm"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "S-A9GZC4HauP",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Import dependencies after install.\n",
        "import json\n",
        "import numpy as np\n",
        "import pickle\n",
        "\n",
        "import modules.envs as envs\n",
        "import modules.agents as agents\n",
        "import modules.loggers as loggers"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Set config for training"
      ],
      "metadata": {
        "id": "JvhJKZhwxW3f"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "j1TxsBVzG6wE",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Import config for training.\n",
        "from configs.bandit_config_train import get_config\n",
        "\n",
        "config = get_config()\n",
        "json_config = json.loads(config.to_json_best_effort())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "8de4ByRpK_1j",
        "cellView": "form",
        "outputId": "89f9a0d2-99c2-4ec2-a79d-1c48d3c70977",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'phase': 'train',\n",
              " 'path': './',\n",
              " 'params_filename': 'train_test',\n",
              " 'random_seed': 42,\n",
              " 'num_workers': 5,\n",
              " 'num_evaluators': 1,\n",
              " 'eval_every_steps': 200000,\n",
              " 'num_eval_episodes': 400,\n",
              " 'log_every_steps': 20000,\n",
              " 'environment': {'env_name': 'bandit',\n",
              "  'steps_per_episode': 100,\n",
              "  'reward_structure': 'correlated'},\n",
              " 'eval_environment': {'steps_per_episode': 100,\n",
              "  'reward_structure': 'correlated'},\n",
              " 'agent': {'total_training_steps': 2000000,\n",
              "  'random_seed': 42,\n",
              "  'num_lstm_units': 48,\n",
              "  'learning_rate_start': 0.0003,\n",
              "  'learning_rate_end': 0.0,\n",
              "  'gamma': 0.9,\n",
              "  'v_loss_coef': 0.05,\n",
              "  'e_loss_coef_start': 0.0,\n",
              "  'e_loss_coef_end': 0.0,\n",
              "  'e_loss_decay_factor': 3,\n",
              "  'max_unroll_steps': 200,\n",
              "  'global_norm_grad_clip': 50.0}}"
            ]
          },
          "metadata": {},
          "execution_count": 6
        }
      ],
      "source": [
        "#@title Print loaded config.\n",
        "json_config"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "Q5NX4ggfLGej",
        "cellView": "form",
        "outputId": "4e9c9329-a1ba-41ae-bbf3-bb5efc6467f4",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'phase': 'train',\n",
              " 'path': './',\n",
              " 'params_filename': 'train_test',\n",
              " 'random_seed': 100,\n",
              " 'num_workers': 5,\n",
              " 'num_evaluators': 1,\n",
              " 'eval_every_steps': 200000,\n",
              " 'num_eval_episodes': 400,\n",
              " 'log_every_steps': 20000,\n",
              " 'environment': {'env_name': 'bandit',\n",
              "  'steps_per_episode': 100,\n",
              "  'reward_structure': 'correlated'},\n",
              " 'eval_environment': {'steps_per_episode': 100,\n",
              "  'reward_structure': 'correlated'},\n",
              " 'agent': {'total_training_steps': 2000000,\n",
              "  'random_seed': 42,\n",
              "  'num_lstm_units': 48,\n",
              "  'learning_rate_start': 0.0003,\n",
              "  'learning_rate_end': 0.0,\n",
              "  'gamma': 0.9,\n",
              "  'v_loss_coef': 0.05,\n",
              "  'e_loss_coef_start': 0.0,\n",
              "  'e_loss_coef_end': 0.0,\n",
              "  'e_loss_decay_factor': 3,\n",
              "  'max_unroll_steps': 200,\n",
              "  'global_norm_grad_clip': 50.0}}"
            ]
          },
          "metadata": {},
          "execution_count": 7
        }
      ],
      "source": [
        "#@title Modify config (optional).\n",
        "#@markdown Add as many lines as needed in the code here.\n",
        "config.update({'random_seed': 100})\n",
        "config.update({'params_filename': 'my_colab_agent'})\n",
        "\n",
        "# Print to see changes.\n",
        "json.loads(config.to_json_best_effort())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2ldXLay6L1FU"
      },
      "source": [
        "# Training setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "2GjvnYCnLOK8",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Unpack config.\n",
        "env_config = config.environment\n",
        "agent_config = config.agent\n",
        "random_seed = config.random_seed\n",
        "total_training_steps = config.agent.total_training_steps\n",
        "log_every_steps = config.log_every_steps\n",
        "params_filename = config.params_filename"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "CxRmJFHGMl-m",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Set random seed in NumPy.\n",
        "np.random.seed(random_seed)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "QAUfl96lMZTT",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Initialize environment.\n",
        "env = envs.create_env(env_config)\n",
        "observation = env.reset()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "V-yfq8qHMgG6",
        "outputId": "fc3b428f-9244-469d-b00b-7a0cab0a6fb2",
        "cellView": "form"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "WARNING:jax._src.lib.xla_bridge:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n"
          ]
        }
      ],
      "source": [
        "#@title Initialize agent.\n",
        "# Note: Jax might complain if there is no GPU/TPU found. You can run on a CPU,\n",
        "# or go to Runtime -> Change Runtime Type in Colab to access a GPU or TPU.\n",
        "agent = agents.create_agent(\n",
        "    observation=observation,\n",
        "    num_actions=env.num_actions,\n",
        "    agent_config=agent_config)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "id": "BC74tFT0NrFJ",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Initialize performance logger.\n",
        "logger = loggers.create_logger(logger_name='bandit', config=config, log_to_console=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "id": "bJ3jDrU8NVbb",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Initialize LSTM recurrent state to zeros.\n",
        "initial_lstm_state = agent.get_initial_lstm_state()\n",
        "lstm_state = initial_lstm_state"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Training"
      ],
      "metadata": {
        "id": "zYmAX4wdx-bC"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "id": "C1vAQdgbNZ5R",
        "cellView": "form",
        "outputId": "8f50da78-6877-4d39-e8d0-f14bf0efb79a",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Global step:\t0\t| Worker step:\t0\t| T:\t151.50\t| Mean Reward:\t0.0001\t| Entropy coef:\t0.0000\t| Loss:\t0.00000\t|\n",
            "Global step:\t20000\t| Worker step:\t20000\t| T:\t6.80\t| Mean Reward:\t0.5082\t| Entropy coef:\t0.0000\t| Loss:\t0.24415\t|\n",
            "Global step:\t40000\t| Worker step:\t40000\t| T:\t7.30\t| Mean Reward:\t0.5055\t| Entropy coef:\t0.0000\t| Loss:\t0.08747\t|\n",
            "Global step:\t60000\t| Worker step:\t60000\t| T:\t5.37\t| Mean Reward:\t0.4942\t| Entropy coef:\t0.0000\t| Loss:\t0.08729\t|\n",
            "Global step:\t80000\t| Worker step:\t80000\t| T:\t7.30\t| Mean Reward:\t0.5107\t| Entropy coef:\t0.0000\t| Loss:\t0.05006\t|\n",
            "Global step:\t100000\t| Worker step:\t100000\t| T:\t5.40\t| Mean Reward:\t0.5024\t| Entropy coef:\t0.0000\t| Loss:\t-0.00209\t|\n",
            "Global step:\t120000\t| Worker step:\t120000\t| T:\t6.57\t| Mean Reward:\t0.5116\t| Entropy coef:\t0.0000\t| Loss:\t0.13439\t|\n",
            "Global step:\t140000\t| Worker step:\t140000\t| T:\t6.17\t| Mean Reward:\t0.5090\t| Entropy coef:\t0.0000\t| Loss:\t-0.01373\t|\n",
            "Global step:\t160000\t| Worker step:\t160000\t| T:\t5.55\t| Mean Reward:\t0.5196\t| Entropy coef:\t0.0000\t| Loss:\t0.03370\t|\n",
            "Global step:\t180000\t| Worker step:\t180000\t| T:\t7.20\t| Mean Reward:\t0.5024\t| Entropy coef:\t0.0000\t| Loss:\t-0.00105\t|\n",
            "Global step:\t200000\t| Worker step:\t200000\t| T:\t5.41\t| Mean Reward:\t0.5212\t| Entropy coef:\t0.0000\t| Loss:\t0.07663\t|\n",
            "Global step:\t220000\t| Worker step:\t220000\t| T:\t7.24\t| Mean Reward:\t0.5172\t| Entropy coef:\t0.0000\t| Loss:\t-0.02531\t|\n",
            "Global step:\t240000\t| Worker step:\t240000\t| T:\t5.41\t| Mean Reward:\t0.5100\t| Entropy coef:\t0.0000\t| Loss:\t0.02592\t|\n",
            "Global step:\t260000\t| Worker step:\t260000\t| T:\t6.87\t| Mean Reward:\t0.5077\t| Entropy coef:\t0.0000\t| Loss:\t0.07864\t|\n",
            "Global step:\t280000\t| Worker step:\t280000\t| T:\t5.88\t| Mean Reward:\t0.5135\t| Entropy coef:\t0.0000\t| Loss:\t-0.06668\t|\n",
            "Global step:\t300000\t| Worker step:\t300000\t| T:\t5.60\t| Mean Reward:\t0.5008\t| Entropy coef:\t0.0000\t| Loss:\t-0.07694\t|\n",
            "Global step:\t320000\t| Worker step:\t320000\t| T:\t7.11\t| Mean Reward:\t0.5283\t| Entropy coef:\t0.0000\t| Loss:\t0.13227\t|\n",
            "Global step:\t340000\t| Worker step:\t340000\t| T:\t5.45\t| Mean Reward:\t0.4935\t| Entropy coef:\t0.0000\t| Loss:\t0.12734\t|\n",
            "Global step:\t360000\t| Worker step:\t360000\t| T:\t7.22\t| Mean Reward:\t0.4905\t| Entropy coef:\t0.0000\t| Loss:\t0.10654\t|\n",
            "Global step:\t380000\t| Worker step:\t380000\t| T:\t5.40\t| Mean Reward:\t0.5366\t| Entropy coef:\t0.0000\t| Loss:\t0.03872\t|\n",
            "Global step:\t400000\t| Worker step:\t400000\t| T:\t7.16\t| Mean Reward:\t0.5265\t| Entropy coef:\t0.0000\t| Loss:\t0.00547\t|\n",
            "Global step:\t420000\t| Worker step:\t420000\t| T:\t5.55\t| Mean Reward:\t0.5159\t| Entropy coef:\t0.0000\t| Loss:\t0.06403\t|\n",
            "Global step:\t440000\t| Worker step:\t440000\t| T:\t5.79\t| Mean Reward:\t0.4789\t| Entropy coef:\t0.0000\t| Loss:\t0.08402\t|\n",
            "Global step:\t460000\t| Worker step:\t460000\t| T:\t6.87\t| Mean Reward:\t0.5199\t| Entropy coef:\t0.0000\t| Loss:\t-0.03743\t|\n",
            "Global step:\t480000\t| Worker step:\t480000\t| T:\t5.45\t| Mean Reward:\t0.5123\t| Entropy coef:\t0.0000\t| Loss:\t0.04865\t|\n",
            "Global step:\t500000\t| Worker step:\t500000\t| T:\t7.23\t| Mean Reward:\t0.5059\t| Entropy coef:\t0.0000\t| Loss:\t0.06763\t|\n",
            "Global step:\t520000\t| Worker step:\t520000\t| T:\t5.46\t| Mean Reward:\t0.5143\t| Entropy coef:\t0.0000\t| Loss:\t0.12484\t|\n",
            "Global step:\t540000\t| Worker step:\t540000\t| T:\t7.31\t| Mean Reward:\t0.4924\t| Entropy coef:\t0.0000\t| Loss:\t0.05413\t|\n",
            "Global step:\t560000\t| Worker step:\t560000\t| T:\t5.41\t| Mean Reward:\t0.5100\t| Entropy coef:\t0.0000\t| Loss:\t0.25031\t|\n",
            "Global step:\t580000\t| Worker step:\t580000\t| T:\t6.18\t| Mean Reward:\t0.5089\t| Entropy coef:\t0.0000\t| Loss:\t0.03182\t|\n",
            "Global step:\t600000\t| Worker step:\t600000\t| T:\t6.62\t| Mean Reward:\t0.5160\t| Entropy coef:\t0.0000\t| Loss:\t0.16210\t|\n",
            "Global step:\t620000\t| Worker step:\t620000\t| T:\t5.35\t| Mean Reward:\t0.4895\t| Entropy coef:\t0.0000\t| Loss:\t0.11279\t|\n",
            "Global step:\t640000\t| Worker step:\t640000\t| T:\t7.32\t| Mean Reward:\t0.5079\t| Entropy coef:\t0.0000\t| Loss:\t0.13359\t|\n",
            "Global step:\t660000\t| Worker step:\t660000\t| T:\t5.45\t| Mean Reward:\t0.4939\t| Entropy coef:\t0.0000\t| Loss:\t0.03856\t|\n",
            "Global step:\t680000\t| Worker step:\t680000\t| T:\t7.31\t| Mean Reward:\t0.5094\t| Entropy coef:\t0.0000\t| Loss:\t0.07443\t|\n",
            "Global step:\t700000\t| Worker step:\t700000\t| T:\t5.38\t| Mean Reward:\t0.5057\t| Entropy coef:\t0.0000\t| Loss:\t0.01897\t|\n",
            "Global step:\t720000\t| Worker step:\t720000\t| T:\t6.37\t| Mean Reward:\t0.5220\t| Entropy coef:\t0.0000\t| Loss:\t-0.00122\t|\n",
            "Global step:\t740000\t| Worker step:\t740000\t| T:\t6.32\t| Mean Reward:\t0.5155\t| Entropy coef:\t0.0000\t| Loss:\t0.04437\t|\n",
            "Global step:\t760000\t| Worker step:\t760000\t| T:\t5.41\t| Mean Reward:\t0.5119\t| Entropy coef:\t0.0000\t| Loss:\t0.12577\t|\n",
            "Global step:\t780000\t| Worker step:\t780000\t| T:\t7.29\t| Mean Reward:\t0.5387\t| Entropy coef:\t0.0000\t| Loss:\t-0.05907\t|\n",
            "Global step:\t800000\t| Worker step:\t800000\t| T:\t5.41\t| Mean Reward:\t0.5268\t| Entropy coef:\t0.0000\t| Loss:\t0.17800\t|\n",
            "Global step:\t820000\t| Worker step:\t820000\t| T:\t7.26\t| Mean Reward:\t0.5400\t| Entropy coef:\t0.0000\t| Loss:\t0.15787\t|\n",
            "Global step:\t840000\t| Worker step:\t840000\t| T:\t5.40\t| Mean Reward:\t0.5368\t| Entropy coef:\t0.0000\t| Loss:\t0.06426\t|\n",
            "Global step:\t860000\t| Worker step:\t860000\t| T:\t6.56\t| Mean Reward:\t0.5393\t| Entropy coef:\t0.0000\t| Loss:\t0.16986\t|\n",
            "Global step:\t880000\t| Worker step:\t880000\t| T:\t6.11\t| Mean Reward:\t0.5436\t| Entropy coef:\t0.0000\t| Loss:\t-0.06447\t|\n",
            "Global step:\t900000\t| Worker step:\t900000\t| T:\t5.40\t| Mean Reward:\t0.5598\t| Entropy coef:\t0.0000\t| Loss:\t0.06149\t|\n",
            "Global step:\t920000\t| Worker step:\t920000\t| T:\t10.16\t| Mean Reward:\t0.5676\t| Entropy coef:\t0.0000\t| Loss:\t0.03477\t|\n",
            "Global step:\t940000\t| Worker step:\t940000\t| T:\t5.52\t| Mean Reward:\t0.5575\t| Entropy coef:\t0.0000\t| Loss:\t0.03521\t|\n",
            "Global step:\t960000\t| Worker step:\t960000\t| T:\t7.54\t| Mean Reward:\t0.5501\t| Entropy coef:\t0.0000\t| Loss:\t0.10623\t|\n",
            "Global step:\t980000\t| Worker step:\t980000\t| T:\t5.55\t| Mean Reward:\t0.5574\t| Entropy coef:\t0.0000\t| Loss:\t0.01987\t|\n",
            "Global step:\t1000000\t| Worker step:\t1000000\t| T:\t6.72\t| Mean Reward:\t0.5530\t| Entropy coef:\t0.0000\t| Loss:\t0.09387\t|\n",
            "Global step:\t1020000\t| Worker step:\t1020000\t| T:\t6.20\t| Mean Reward:\t0.5805\t| Entropy coef:\t0.0000\t| Loss:\t-0.06829\t|\n",
            "Global step:\t1040000\t| Worker step:\t1040000\t| T:\t5.69\t| Mean Reward:\t0.5620\t| Entropy coef:\t0.0000\t| Loss:\t0.07255\t|\n",
            "Global step:\t1060000\t| Worker step:\t1060000\t| T:\t7.24\t| Mean Reward:\t0.5648\t| Entropy coef:\t0.0000\t| Loss:\t0.11031\t|\n",
            "Global step:\t1080000\t| Worker step:\t1080000\t| T:\t5.46\t| Mean Reward:\t0.5667\t| Entropy coef:\t0.0000\t| Loss:\t0.08622\t|\n",
            "Global step:\t1100000\t| Worker step:\t1100000\t| T:\t7.45\t| Mean Reward:\t0.5594\t| Entropy coef:\t0.0000\t| Loss:\t0.11483\t|\n",
            "Global step:\t1120000\t| Worker step:\t1120000\t| T:\t7.42\t| Mean Reward:\t0.5497\t| Entropy coef:\t0.0000\t| Loss:\t0.05013\t|\n",
            "Global step:\t1140000\t| Worker step:\t1140000\t| T:\t7.42\t| Mean Reward:\t0.5654\t| Entropy coef:\t0.0000\t| Loss:\t0.11256\t|\n",
            "Global step:\t1160000\t| Worker step:\t1160000\t| T:\t5.49\t| Mean Reward:\t0.5692\t| Entropy coef:\t0.0000\t| Loss:\t0.02050\t|\n",
            "Global step:\t1180000\t| Worker step:\t1180000\t| T:\t7.37\t| Mean Reward:\t0.5555\t| Entropy coef:\t0.0000\t| Loss:\t-0.16843\t|\n",
            "Global step:\t1200000\t| Worker step:\t1200000\t| T:\t5.49\t| Mean Reward:\t0.5497\t| Entropy coef:\t0.0000\t| Loss:\t0.10340\t|\n",
            "Global step:\t1220000\t| Worker step:\t1220000\t| T:\t6.45\t| Mean Reward:\t0.5699\t| Entropy coef:\t0.0000\t| Loss:\t0.00400\t|\n",
            "Global step:\t1240000\t| Worker step:\t1240000\t| T:\t6.40\t| Mean Reward:\t0.5515\t| Entropy coef:\t0.0000\t| Loss:\t0.12483\t|\n",
            "Global step:\t1260000\t| Worker step:\t1260000\t| T:\t5.41\t| Mean Reward:\t0.5544\t| Entropy coef:\t0.0000\t| Loss:\t0.11050\t|\n",
            "Global step:\t1280000\t| Worker step:\t1280000\t| T:\t7.35\t| Mean Reward:\t0.5774\t| Entropy coef:\t0.0000\t| Loss:\t0.05125\t|\n",
            "Global step:\t1300000\t| Worker step:\t1300000\t| T:\t5.45\t| Mean Reward:\t0.5548\t| Entropy coef:\t0.0000\t| Loss:\t0.04472\t|\n",
            "Global step:\t1320000\t| Worker step:\t1320000\t| T:\t7.43\t| Mean Reward:\t0.5676\t| Entropy coef:\t0.0000\t| Loss:\t0.05170\t|\n",
            "Global step:\t1340000\t| Worker step:\t1340000\t| T:\t5.53\t| Mean Reward:\t0.5774\t| Entropy coef:\t0.0000\t| Loss:\t0.10844\t|\n",
            "Global step:\t1360000\t| Worker step:\t1360000\t| T:\t7.07\t| Mean Reward:\t0.5619\t| Entropy coef:\t0.0000\t| Loss:\t0.06119\t|\n",
            "Global step:\t1380000\t| Worker step:\t1380000\t| T:\t5.78\t| Mean Reward:\t0.5634\t| Entropy coef:\t0.0000\t| Loss:\t0.11340\t|\n",
            "Global step:\t1400000\t| Worker step:\t1400000\t| T:\t5.79\t| Mean Reward:\t0.5521\t| Entropy coef:\t0.0000\t| Loss:\t0.04506\t|\n",
            "Global step:\t1420000\t| Worker step:\t1420000\t| T:\t7.07\t| Mean Reward:\t0.5646\t| Entropy coef:\t0.0000\t| Loss:\t0.06300\t|\n",
            "Global step:\t1440000\t| Worker step:\t1440000\t| T:\t5.53\t| Mean Reward:\t0.5536\t| Entropy coef:\t0.0000\t| Loss:\t0.06573\t|\n",
            "Global step:\t1460000\t| Worker step:\t1460000\t| T:\t7.43\t| Mean Reward:\t0.5611\t| Entropy coef:\t0.0000\t| Loss:\t0.11747\t|\n",
            "Global step:\t1480000\t| Worker step:\t1480000\t| T:\t5.47\t| Mean Reward:\t0.5670\t| Entropy coef:\t0.0000\t| Loss:\t0.03684\t|\n",
            "Global step:\t1500000\t| Worker step:\t1500000\t| T:\t7.37\t| Mean Reward:\t0.5625\t| Entropy coef:\t0.0000\t| Loss:\t0.07390\t|\n",
            "Global step:\t1520000\t| Worker step:\t1520000\t| T:\t5.51\t| Mean Reward:\t0.5698\t| Entropy coef:\t0.0000\t| Loss:\t0.04777\t|\n",
            "Global step:\t1540000\t| Worker step:\t1540000\t| T:\t6.45\t| Mean Reward:\t0.5739\t| Entropy coef:\t0.0000\t| Loss:\t0.07612\t|\n",
            "Global step:\t1560000\t| Worker step:\t1560000\t| T:\t6.39\t| Mean Reward:\t0.5655\t| Entropy coef:\t0.0000\t| Loss:\t0.13095\t|\n",
            "Global step:\t1580000\t| Worker step:\t1580000\t| T:\t5.44\t| Mean Reward:\t0.5385\t| Entropy coef:\t0.0000\t| Loss:\t0.06989\t|\n",
            "Global step:\t1600000\t| Worker step:\t1600000\t| T:\t7.41\t| Mean Reward:\t0.5431\t| Entropy coef:\t0.0000\t| Loss:\t0.17077\t|\n",
            "Global step:\t1620000\t| Worker step:\t1620000\t| T:\t5.47\t| Mean Reward:\t0.5613\t| Entropy coef:\t0.0000\t| Loss:\t0.06908\t|\n",
            "Global step:\t1640000\t| Worker step:\t1640000\t| T:\t7.37\t| Mean Reward:\t0.5606\t| Entropy coef:\t0.0000\t| Loss:\t0.00189\t|\n",
            "Global step:\t1660000\t| Worker step:\t1660000\t| T:\t5.46\t| Mean Reward:\t0.5671\t| Entropy coef:\t0.0000\t| Loss:\t0.09706\t|\n",
            "Global step:\t1680000\t| Worker step:\t1680000\t| T:\t7.05\t| Mean Reward:\t0.5551\t| Entropy coef:\t0.0000\t| Loss:\t0.08931\t|\n",
            "Global step:\t1700000\t| Worker step:\t1700000\t| T:\t5.80\t| Mean Reward:\t0.5524\t| Entropy coef:\t0.0000\t| Loss:\t0.00852\t|\n",
            "Global step:\t1720000\t| Worker step:\t1720000\t| T:\t5.82\t| Mean Reward:\t0.5644\t| Entropy coef:\t0.0000\t| Loss:\t0.06990\t|\n",
            "Global step:\t1740000\t| Worker step:\t1740000\t| T:\t7.04\t| Mean Reward:\t0.5586\t| Entropy coef:\t0.0000\t| Loss:\t0.08369\t|\n",
            "Global step:\t1760000\t| Worker step:\t1760000\t| T:\t5.51\t| Mean Reward:\t0.5568\t| Entropy coef:\t0.0000\t| Loss:\t0.06169\t|\n",
            "Global step:\t1780000\t| Worker step:\t1780000\t| T:\t7.38\t| Mean Reward:\t0.5468\t| Entropy coef:\t0.0000\t| Loss:\t0.12262\t|\n",
            "Global step:\t1800000\t| Worker step:\t1800000\t| T:\t5.51\t| Mean Reward:\t0.5669\t| Entropy coef:\t0.0000\t| Loss:\t0.12607\t|\n",
            "Global step:\t1820000\t| Worker step:\t1820000\t| T:\t7.33\t| Mean Reward:\t0.5635\t| Entropy coef:\t0.0000\t| Loss:\t0.05675\t|\n",
            "Global step:\t1840000\t| Worker step:\t1840000\t| T:\t5.43\t| Mean Reward:\t0.5593\t| Entropy coef:\t0.0000\t| Loss:\t0.09104\t|\n",
            "Global step:\t1860000\t| Worker step:\t1860000\t| T:\t6.27\t| Mean Reward:\t0.5641\t| Entropy coef:\t0.0000\t| Loss:\t0.04012\t|\n",
            "Global step:\t1880000\t| Worker step:\t1880000\t| T:\t6.60\t| Mean Reward:\t0.5785\t| Entropy coef:\t0.0000\t| Loss:\t-0.02131\t|\n",
            "Global step:\t1900000\t| Worker step:\t1900000\t| T:\t5.50\t| Mean Reward:\t0.5750\t| Entropy coef:\t0.0000\t| Loss:\t0.02835\t|\n",
            "Global step:\t1920000\t| Worker step:\t1920000\t| T:\t7.41\t| Mean Reward:\t0.5554\t| Entropy coef:\t0.0000\t| Loss:\t-0.08800\t|\n",
            "Global step:\t1940000\t| Worker step:\t1940000\t| T:\t5.49\t| Mean Reward:\t0.5654\t| Entropy coef:\t0.0000\t| Loss:\t0.09280\t|\n",
            "Global step:\t1960000\t| Worker step:\t1960000\t| T:\t7.38\t| Mean Reward:\t0.5574\t| Entropy coef:\t0.0000\t| Loss:\t0.07881\t|\n",
            "Global step:\t1980000\t| Worker step:\t1980000\t| T:\t5.46\t| Mean Reward:\t0.5585\t| Entropy coef:\t0.0000\t| Loss:\t0.09474\t|\n",
            "Done training!\n"
          ]
        }
      ],
      "source": [
        "#@title Main training loop (Note: this might take a while).\n",
        "\n",
        "step, episode, loss = 0, 0, 0\n",
        "while step < total_training_steps:\n",
        "\n",
        "    # Get an action and step the environment with the agent's action\n",
        "    action, _, v_out, new_lstm_state, _ = agent.get_action(observation, lstm_state)\n",
        "    next_observation, reward, done, info = env.step(action)\n",
        "\n",
        "    # Save experience in a buffer\n",
        "    agent.buffer.append(\n",
        "        obs=observation,\n",
        "        action=action,\n",
        "        reward=reward,\n",
        "        next_obs=next_observation,\n",
        "        done=done,\n",
        "        lstm_state=lstm_state,\n",
        "    )\n",
        "\n",
        "    observation = next_observation\n",
        "    lstm_state = new_lstm_state\n",
        "\n",
        "    # Log performance\n",
        "    logger.log_step(\n",
        "        global_step=step,\n",
        "        worker_step=step,\n",
        "        reward=reward,\n",
        "        info=info,\n",
        "        loss=loss,\n",
        "        entropy_coef=agent.e_loss_coef,\n",
        "    )\n",
        "\n",
        "    # Update agent parameters if an episode is done, or\n",
        "    # if the agent experience buffer == max_unroll_steps\n",
        "    loss, grads, num_steps = agent.update(done, update_params=True)\n",
        "    step += 1\n",
        "\n",
        "    # If done, reset the environment and LSTM state\n",
        "    if done:\n",
        "        episode += 1\n",
        "        done = False\n",
        "        lstm_state = initial_lstm_state\n",
        "        observation = env.reset()\n",
        "\n",
        "print('Done training!')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "id": "766CuUZKP1lk",
        "cellView": "form",
        "outputId": "674e286d-c210-45fd-f4ef-67855b0a2b2c",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Saved parameters.\n"
          ]
        }
      ],
      "source": [
        "#@title Save model after training is complete.\n",
        "results = {\n",
        "    \"params\": agent.params,\n",
        "    \"config\": config.to_dict(),\n",
        "}\n",
        "with open(params_filename + '.pickle', 'wb') as fp:\n",
        "    pickle.dump([results], fp)\n",
        "print(\"Saved parameters.\")"
      ]
    }
  ],
  "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"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}