BioLCNet_CartPole_gradient.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Singular-Brain/DeepBioLCNet/blob/main/BioLCNet_CartPole_gradient.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7fTSvrK3T_GA"
      },
      "source": [
        "#Notebook setups"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "lXtgP_iEPE0G"
      },
      "outputs": [],
      "source": [
        "!pip install -q git+https://github.com/Singular-Brain/DeepBioLCNet"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "KXcXvvsXcOlv",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "a8353553-3724-4287-b6fb-d9d554c774dd"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "fatal: destination path 'DeepBioLCNet' already exists and is not an empty directory.\n"
          ]
        }
      ],
      "source": [
        "!git clone https://github.com/Singular-Brain/DeepBioLCNet\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "### gym and colab compatibility\n",
        "\n",
        "!apt-get install -y xvfb python-opengl > /dev/null 2>&1\n",
        "!pip install gym pyvirtualdisplay > /dev/null 2>&1\n",
        "\n",
        "import gym\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "from IPython import display as ipythondisplay\n",
        "\n",
        "from pyvirtualdisplay import Display\n",
        "display = Display(visible=0, size=(400, 300))\n",
        "display.start()"
      ],
      "metadata": {
        "id": "0QFC6xL2uAaF",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "3cfc98ce-c92a-4fcd-d56c-bb870ab35855"
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<pyvirtualdisplay.display.Display at 0x7fb4ec71e810>"
            ]
          },
          "metadata": {},
          "execution_count": 3
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "K4l3AVRbGS4Q"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "from itertools import count\n",
        "from matplotlib.axes import Axes\n",
        "from matplotlib.image import AxesImage\n",
        "from torch.nn.modules.utils import _pair\n",
        "from matplotlib.collections import PathCollection\n",
        "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
        "from typing import Tuple, List, Optional, Sized, Dict, Union\n",
        "import math\n",
        "import random\n",
        "from torchvision.transforms.functional import crop\n",
        "# from ..utils import reshape_locally_connected_weights, reshape_locally_connected_weights_meh, reshape_conv2d_weights\n",
        "\n",
        "import gym\n",
        "import tkinter\n",
        "\n",
        "from PIL import Image\n",
        "from collections import namedtuple"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "BFGNAecpT-Lj"
      },
      "outputs": [],
      "source": [
        "from bindsnet.network.nodes import Nodes\n",
        "import os\n",
        "### import torch\n",
        "import numpy as np\n",
        "import copy\n",
        "import math\n",
        "import matplotlib.pyplot as plt\n",
        "import collections\n",
        "from torchvision import transforms\n",
        "from tqdm.notebook import tqdm\n",
        "\n",
        "from sklearn.metrics import confusion_matrix\n",
        "import seaborn as sn\n",
        "import torch.nn.functional as fn\n",
        "\n",
        "from abc import ABC, abstractmethod\n",
        "from typing import Union, Tuple, Optional, Sequence\n",
        "from torch.nn.modules.utils import _pair\n",
        "\n",
        "from bindsnet.datasets import MNIST\n",
        "from bindsnet.encoding import PoissonEncoder\n",
        "from bindsnet.network import Network\n",
        "from bindsnet.network.nodes import Input, LIFNodes, AdaptiveLIFNodes, IFNodes\n",
        "from bindsnet.network.topology import Connection, MaxPool2dLocalConnection\n",
        "from bindsnet.network.topology import LocalConnection, LocalConnectionOrig\n",
        "from bindsnet.network.monitors import Monitor, AbstractMonitor, TensorBoardMonitor\n",
        "from bindsnet.learning import PostPre, MSTDP, MSTDPET, WeightDependentPostPre, Hebbian\n",
        "from bindsnet.learning.reward import DynamicDopamineInjection, DopaminergicRPE, AbstractReward\n",
        "from bindsnet.analysis.plotting import plot_locally_connected_weights,plot_locally_connected_weights_meh,plot_spikes,\\\n",
        "plot_LC_timepoint_spikes,plot_locally_connected_weights_meh2,plot_convergence_and_histogram,plot_locally_connected_weights_meh3\n",
        "from bindsnet.analysis.visualization import plot_weights_movie, plot_spike_trains_for_example,summary, plot_voltage\n",
        "from bindsnet.utils import reshape_locally_connected_weights, reshape_locally_connected_weights_meh, reshape_conv2d_weights"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "class RLTasks(AbstractReward):\n",
        "    # language=rst\n",
        "    \"\"\"\n",
        "    Computes the reward for a given RL task in the current state\n",
        "    \"\"\"\n",
        "    def __init__(\n",
        "        self,\n",
        "        env_id,\n",
        "        **kwargs,\n",
        "    ) -> None:\n",
        "\n",
        "        super().__init__()\n",
        "\n",
        "        if env_id == \"CartPole-v0\":\n",
        "            self.compute = self._cartPole_compute\n",
        "        elif env_id == \"MountainCar-v0-v0\":\n",
        "            self.compute = self._mountainCar_compute\n",
        "        elif env_id == \"BreakoutDeterministic-v4\":\n",
        "            self.compute = self._breakout_compute\n",
        "        else:\n",
        "            raise NotImplementedError(\n",
        "                \"This rl environment is not currently supported.\"\n",
        "            )\n",
        "    \n",
        "    def _cartPole_compute(self, **kwargs):\n",
        "        success = kwargs['success']\n",
        "        failure = kwargs['failure']\n",
        "        if kwargs['classic_reward']:\n",
        "            if failure:\n",
        "                return torch.tensor([0.])\n",
        "            else:\n",
        "                return torch.tensor([1.])\n",
        "\n",
        "        env = kwargs['env']\n",
        "        state = env.state\n",
        "        success = kwargs['success']\n",
        "        failure = kwargs['failure']\n",
        "        x, x_dot, theta, theta_dot = state\n",
        "        # r1 = (env.x_threshold - abs(x)) / env.x_threshold - 0.8\n",
        "        # r2 = ((env.theta_threshold_radians - abs(theta)) / env.theta_threshold_radians - 0.5)#*2\n",
        "        #reward = r1 + r2\n",
        "        reward = ((env.theta_threshold_radians - abs(theta)) / env.theta_threshold_radians) - 0.5*abs(theta_dot)\n",
        "        if success:\n",
        "            reward += 5\n",
        "        elif failure:\n",
        "            reward -= 5\n",
        "        reward = torch.tensor([reward])\n",
        "        return reward\n",
        "\n",
        "    def _mountainCar_compute(self):\n",
        "        pass\n",
        "\n",
        "    def _breakout_compute(self):\n",
        "        pass\n",
        "        \n",
        "    def update(self):\n",
        "        pass\n",
        "\n",
        "    def online_compute(self,):\n",
        "        pass"
      ],
      "metadata": {
        "id": "WqFD8yro6WCb"
      },
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ULGGHW43UksI"
      },
      "source": [
        "## Sets up Gpu use and manual seed\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "LiUmFrpcUfmR",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "ca92cf5e-f8fb-4879-aa24-6a0230063554"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Running on Device =  cuda\n"
          ]
        }
      ],
      "source": [
        "if torch.cuda.is_available():\n",
        "    device =  torch.device(\"cuda\")\n",
        "    gpu = True\n",
        "else:\n",
        "    device =  torch.device(\"cpu\")\n",
        "    gpu = False\n",
        "\n",
        "def manual_seed(seed):\n",
        "    random.seed(seed)\n",
        "    torch.manual_seed(seed)\n",
        "    np.random.seed(seed)\n",
        "    if torch.cuda.is_available():\n",
        "        torch.cuda.manual_seed_all(seed)\n",
        "\n",
        "\n",
        "SEED = 2045 # The Singularity is Near!\n",
        "manual_seed(SEED)\n",
        "WANDB = False\n",
        "\n",
        "torch.set_num_threads(os.cpu_count() - 1)\n",
        "print(\"Running on Device = \", device)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MBKedMpIleMr"
      },
      "source": [
        "# Custom Monitors"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6tfqpsr2a1WV"
      },
      "source": [
        "## Reward Monitor"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "M44GJ65GleMs"
      },
      "outputs": [],
      "source": [
        "class RewardMonitor(AbstractMonitor):\n",
        "    # language=rst\n",
        "    \"\"\"\n",
        "    Records state variables of interest.\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(\n",
        "        self,\n",
        "        time: None,\n",
        "        batch_size: int = 1,\n",
        "        device: str = \"cpu\",\n",
        "    ):\n",
        "        # language=rst\n",
        "        \"\"\"\n",
        "        Constructs a ``Monitor`` object.\n",
        "\n",
        "        :param obj: An object to record state variables from during network simulation.\n",
        "        :param state_vars: Iterable of strings indicating names of state variables to record.\n",
        "        :param time: If not ``None``, pre-allocate memory for state variable recording.\n",
        "        :param device: Allow the monitor to be on different device separate from Network device\n",
        "        \"\"\"\n",
        "        super().__init__()\n",
        "\n",
        "        self.time = time\n",
        "        self.batch_size = batch_size\n",
        "        self.device = device\n",
        "\n",
        "        # if time is not specified the monitor variable accumulate the logs\n",
        "        if self.time is None:\n",
        "            self.device = \"cpu\"\n",
        "\n",
        "        self.recording = []\n",
        "        self.reset_state_variables()\n",
        "\n",
        "    def get(self,) -> torch.Tensor:\n",
        "        # language=rst\n",
        "        \"\"\"\n",
        "        Return recording to user.\n",
        "\n",
        "        :return: Tensor of shape ``[time, n_1, ..., n_k]``, where ``[n_1, ..., n_k]`` is the shape of the recorded state\n",
        "        variable.\n",
        "        Note, if time == `None`, get return the logs and empty the monitor variable\n",
        "\n",
        "        \"\"\"\n",
        "        # return_logs = torch.as_tensor(self.recording)\n",
        "        # if self.time is None:\n",
        "        #     self.recording = []\n",
        "        return self.recording\n",
        "\n",
        "    def record(self, **kwargs) -> None:\n",
        "        # language=rst\n",
        "        \"\"\"\n",
        "        Appends the current value of the recorded state variables to the recording.\n",
        "        \"\"\"\n",
        "        if \"reward\" in kwargs:\n",
        "            self.recording.append(kwargs[\"reward\"])\n",
        "        # remove the oldest element (first in the list)\n",
        "        # if self.time is not None:\n",
        "        #     self.recording.pop(0)\n",
        "\n",
        "    def reset_state_variables(self) -> None:\n",
        "        # language=rst\n",
        "        \"\"\"\n",
        "        Resets recordings to empty ``List``s.\n",
        "        \"\"\"\n",
        "        self.recording = []\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8clxN_npa1WY"
      },
      "source": [
        "## Plot Eligibility trace"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "SshGlRwpa1WZ"
      },
      "outputs": [],
      "source": [
        "class PlotET(AbstractMonitor):\n",
        "    # language=rst\n",
        "    \"\"\"\n",
        "    Records state variables of interest.\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(\n",
        "        self,\n",
        "        i,\n",
        "        j,\n",
        "        source,\n",
        "        target,\n",
        "        connection,\n",
        "    ):\n",
        "        # language=rst\n",
        "        \"\"\"\n",
        "        Constructs a ``Monitor`` object.\n",
        "\n",
        "        :param obj: An object to record state variables from during network simulation.\n",
        "        :param state_vars: Iterable of strings indicating names of state variables to record.\n",
        "        :param time: If not ``None``, pre-allocate memory for state variable recording.\n",
        "        :param device: Allow the monitor to be on different device separate from Network device\n",
        "        \"\"\"\n",
        "        super().__init__()\n",
        "        self.i = i\n",
        "        self.j = j\n",
        "        self.source = source\n",
        "        self.target = target\n",
        "        self.connection = connection\n",
        "\n",
        "        self.reset_state_variables()\n",
        "\n",
        "    def get(self,) -> torch.Tensor:\n",
        "        # language=rst\n",
        "        \"\"\"\n",
        "        Return recording to user.\n",
        "\n",
        "        :return: Tensor of shape ``[time, n_1, ..., n_k]``, where ``[n_1, ..., n_k]`` is the shape of the recorded state\n",
        "        variable.\n",
        "        Note, if time == `None`, get return the logs and empty the monitor variable\n",
        "\n",
        "        \"\"\"\n",
        "        # return_logs = torch.as_tensor(self.recording)\n",
        "        # if self.time is None:\n",
        "        #     self.recording = []\n",
        "        return self.recording\n",
        "\n",
        "    def record(self, **kwargs) -> None:\n",
        "        # language=rst\n",
        "        \"\"\"\n",
        "        Appends the current value of the recorded state variables to the recording.\n",
        "        \"\"\"\n",
        "        if hasattr(self.connection.update_rule, 'p_plus'):\n",
        "            self.recording['spikes_i'].append(self.source.s.ravel()[self.i].item())\n",
        "            self.recording['spikes_j'].append(self.target.s.ravel()[self.j].item())\n",
        "            self.recording['p_plus'].append(self.connection.update_rule.p_plus[self.i].item())\n",
        "            self.recording['p_minus'].append(self.connection.update_rule.p_minus[self.j].item())\n",
        "            self.recording['eligibility'].append(self.connection.update_rule.eligibility[self.i,self.j].item())\n",
        "            self.recording['eligibility_trace'].append(self.connection.update_rule.eligibility_trace[self.i,self.j].item())\n",
        "            self.recording['w'].append(self.connection.w[self.i,self.j].item())\n",
        "\n",
        "    def plot(self):\n",
        "\n",
        "        fig, axs  = plt.subplots(7)\n",
        "        fig.set_size_inches(10, 20)\n",
        "        for i, (name, p) in enumerate(self.recording.items()):\n",
        "            axs[i].plot(p[-250:])\n",
        "            axs[i].set_title(name)\n",
        "    \n",
        "        fig.show()\n",
        "\n",
        "    def reset_state_variables(self) -> None:\n",
        "        # language=rst\n",
        "        \"\"\"\n",
        "        Resets recordings to empty ``List``s.\n",
        "        \"\"\"\n",
        "        self.recording = {\n",
        "        'spikes_i': [],\n",
        "        'spikes_j': [],\n",
        "        'p_plus':[],\n",
        "        'p_minus':[],\n",
        "        'eligibility':[],\n",
        "        'eligibility_trace':[],\n",
        "        'w': [],\n",
        "        }\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R_YGE1XjvIkZ"
      },
      "source": [
        "## Kernel "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "-4hp2V46vOUv"
      },
      "outputs": [],
      "source": [
        "class AbstractKernel(ABC):\n",
        "    def __init__(self, kernel_size):\n",
        "        \"\"\"\n",
        "        Base class for generating image filter kernels such as Gabor, DoG, etc. Each subclass should override :attr:`__call__` function.\n",
        "        Instantiates a ``Filter Kernel`` object.\n",
        "        :param window_size : The size of the kernel (int)\n",
        "        \"\"\"\n",
        "        self.window_size = kernel_size\n",
        "\n",
        "    def __call__(self):\n",
        "        pass\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "PL2L6_ABwBH4"
      },
      "outputs": [],
      "source": [
        "class DoGKernel(AbstractKernel):\n",
        "    def __init__(self, kernel_size: Union[int, Tuple[int, int]], sigma1 : float, sigma2 : float):\n",
        "        \"\"\"\n",
        "        Generates DoG filter kernels.\n",
        "        :param kernel_size: Horizontal and vertical size of DOG kernels.(If pass int, we consider it as a square filter) \n",
        "        :param sigma1 : The sigma parameter for the first Gaussian function.\n",
        "        :param sigma2 : The sigma parameter for the second Gaussian function.\n",
        "        \"\"\"\n",
        "        super(DoGKernel, self).__init__(kernel_size)\n",
        "        self.sigma1 = sigma1\n",
        "        self.sigma2 = sigma2\n",
        "        \n",
        "    def __call__(self):\n",
        "        k = self.window_size//2\n",
        "        x, y = np.mgrid[-k:k+1:1, -k:k+1:1]\n",
        "        a = 1.0 / (2 * math.pi)\n",
        "        prod = x*x + y*y\n",
        "        f1 = (1/(self.sigma1*self.sigma1)) * np.exp(-0.5 * (1/(self.sigma1*self.sigma1)) * (prod))\n",
        "        f2 = (1/(self.sigma2*self.sigma2)) * np.exp(-0.5 * (1/(self.sigma2*self.sigma2)) * (prod))\n",
        "        dog = a * (f1-f2)\n",
        "        dog_mean = np.mean(dog)\n",
        "        dog = dog - dog_mean\n",
        "        dog_max = np.max(dog)\n",
        "        dog = dog / dog_max\n",
        "        dog_tensor = torch.from_numpy(dog)\n",
        "        # returns a 2d tensor corresponding to the requested DoG filter\n",
        "        return dog_tensor.float()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "id": "zBUT0IUZDXxW"
      },
      "outputs": [],
      "source": [
        "class Filter():\n",
        "    \"\"\"\n",
        "    Applies a filter transform. Each filter contains a sequence of :attr:`FilterKernel` objects.\n",
        "    The result of each filter kernel will be passed through a given threshold (if not :attr:`None`).\n",
        "    Args:\n",
        "        filter_kernels (sequence of FilterKernels): The sequence of filter kernels.\n",
        "        padding (int, optional): The size of the padding for the convolution of filter kernels. Default: 0\n",
        "        thresholds (sequence of floats, optional): The threshold for each filter kernel. Default: None\n",
        "        use_abs (boolean, optional): To compute the absolute value of the outputs or not. Default: False\n",
        "    .. note::\n",
        "        The size of the compund filter kernel tensor (stack of individual filter kernels) will be equal to the \n",
        "        greatest window size among kernels. All other smaller kernels will be zero-padded with an appropriate \n",
        "        amount.\n",
        "    \"\"\"\n",
        "    # filter_kernels must be a list of filter kernels\n",
        "    # thresholds must be a list of thresholds for each kernel\n",
        "    def __init__(self, filter_kernels, padding=0, thresholds=None, use_abs=False):\n",
        "        tensor_list = []\n",
        "        self.max_window_size = 0\n",
        "        for kernel in filter_kernels:\n",
        "            if isinstance(kernel, torch.Tensor):\n",
        "                tensor_list.append(kernel)\n",
        "                self.max_window_size = max(self.max_window_size, kernel.size(-1))\n",
        "            else:\n",
        "                tensor_list.append(kernel().unsqueeze(0))\n",
        "                self.max_window_size = max(self.max_window_size, kernel.window_size)\n",
        "        for i in range(len(tensor_list)):\n",
        "            p = (self.max_window_size - filter_kernels[i].window_size)//2\n",
        "            tensor_list[i] = fn.pad(tensor_list[i], (p,p,p,p))\n",
        "\n",
        "        self.kernels = torch.stack(tensor_list)\n",
        "        self.number_of_kernels = len(filter_kernels)\n",
        "        self.padding = padding\n",
        "        if isinstance(thresholds, list):\n",
        "            self.thresholds = thresholds.clone().detach()\n",
        "            self.thresholds.unsqueeze_(0).unsqueeze_(2).unsqueeze_(3)\n",
        "        else:\n",
        "            self.thresholds = thresholds\n",
        "        self.use_abs = use_abs\n",
        "\n",
        "    # returns a 4d tensor containing the flitered versions of the input image\n",
        "    # input is a 4d tensor. dim: (minibatch=1, filter_kernels, height, width)\n",
        "    def __call__(self, input):\n",
        "\n",
        "        # if input.dim() == 3:\n",
        "        #     input2 = torch.unsqueeze(input, 0)\n",
        "        input.unsqueeze_(0)\n",
        "        output = fn.conv2d(input, self.kernels, padding = self.padding).float()\n",
        "        if not(self.thresholds is None):\n",
        "            output = torch.where(output < self.thresholds, torch.tensor(0.0, device=output.device), output)\n",
        "        if self.use_abs:\n",
        "            torch.abs_(output)\n",
        "        return output.squeeze(0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sCXSLZZoGS4z"
      },
      "source": [
        "# Viz"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "id": "3DIAdG1mGS4z"
      },
      "outputs": [],
      "source": [
        "def plot_LC_timepoint_spikes(spikes: torch.Tensor,\n",
        "    timepoint: int,\n",
        "    n_filters: int,\n",
        "    in_chans: int,\n",
        "    slice_to_plot: int,\n",
        "    conv_size: Union[int, Tuple[int, int]],\n",
        "    im: Optional[AxesImage] = None,\n",
        "    lines: bool = True,\n",
        "    figsize: Tuple[int, int] = (10, 10),\n",
        "    cmap: str = \"hot_r\",\n",
        "    color: str='r'\n",
        "    ):\n",
        "    # language=rst\n",
        "    \"\"\"\n",
        "    \"\"\"\n",
        "\n",
        "    n_sqrt = int(np.ceil(np.sqrt(n_filters)))\n",
        "    sel_slice = spikes[timepoint].view(in_chans, n_filters, conv_size, conv_size).cpu()\n",
        "    sel_slice = sel_slice[slice_to_plot, ...].view(n_filters, conv_size, conv_size)\n",
        "    spikes_ = np.zeros((n_sqrt*conv_size, n_sqrt*conv_size))\n",
        "    filt_counter = 0\n",
        "\n",
        "    fig, ax = plt.subplots(figsize=figsize)\n",
        "\n",
        "    for n1 in range(n_sqrt):\n",
        "        for n2 in range(n_sqrt):\n",
        "            filter_ = sel_slice[filt_counter, :, :].view(conv_size, conv_size)\n",
        "            spikes_[n1 * conv_size : (n1 + 1) * conv_size, n2 * conv_size : (n2 + 1) * conv_size] = filter_\n",
        "            filt_counter += 1\n",
        "            ax.axhline((n1 + 1) * conv_size, color=\"g\", linestyle=\"-\")\n",
        "            ax.axvline((n2 + 1) * conv_size, color=\"g\", linestyle=\"--\")\n",
        "    ax.imshow(spikes_, cmap='Greys')\n",
        "    return spikes_\n",
        "    \n",
        "def plot_FC_response_map(lc: object,\n",
        "    fc: object,\n",
        "    ind_neuron_in_group: int,\n",
        "    label: int,\n",
        "    n_per_action: int,\n",
        "    input_channel: int = 0,\n",
        "    scale_factor: float = 1.0,\n",
        "    lines: bool = True,\n",
        "    figsize: Tuple[int, int] = (5, 5),\n",
        "    cmap: str = \"hot_r\",\n",
        "    color: str='r'\n",
        "    ) -> AxesImage:\n",
        "    # language=rst\n",
        "    \"\"\"\n",
        "    Plot a connection weight matrix of a :code:`Connection` with `locally connected\n",
        "    structure <http://yann.lecun.com/exdb/publis/pdf/gregor-nips-11.pdf>_.\n",
        "    :param lc: LC connection object of LCNet\n",
        "    :param fc: FC connection object of LCNet\n",
        "    :param input_channel: indicates weights which connected to this channel of input\n",
        "    :param scale_factor: determines intensity of activation map  \n",
        "    :param lines: Whether or not to draw horizontal and vertical lines separating input regions.\n",
        "    :param figsize: Horizontal, vertical figure size in inches.\n",
        "    :param cmap: Matplotlib colormap.\n",
        "    :return: Used for re-drawing the weights plot.\n",
        "    \"\"\"\n",
        "\n",
        "    n_sqrt = int(np.ceil(np.sqrt(lc.out_channels)))\n",
        "\n",
        "    sel_slice = lc.w.view(lc.in_channels, lc.out_channels, lc.conv_size[0], lc.conv_size[1], lc.kernel_size[0], lc.kernel_size[1]).cpu()\n",
        "    sel_slice = sel_slice[input_channel, ...]\n",
        "    input_size = _pair(int(np.sqrt(lc.source.n)))\n",
        "    reshaped = reshape_LC_weights(sel_slice, lc.out_channels, lc.kernel_size, lc.conv_size, input_size)\n",
        "\t\n",
        "    ind_neuron = label * n_per_action + ind_neuron_in_group\n",
        "    w = fc.w[:,ind_neuron].view(reshaped.shape[0]//lc.kernel_size[0],reshaped.shape[1]//lc.kernel_size[1])\n",
        "    w = w.clip(lc.wmin,lc.wmax).repeat_interleave(lc.kernel_size[0], dim=0).repeat_interleave(lc.kernel_size[1], dim=1).cpu()\n",
        "\n",
        "    fig, ax = plt.subplots(figsize=figsize)\n",
        "\n",
        "    im = ax.imshow(reshaped.cpu()*w, cmap=cmap, vmin=lc.wmin, vmax=lc.wmax)\n",
        "    div = make_axes_locatable(ax)\n",
        "    cax = div.append_axes(\"right\", size=\"5%\", pad=0.05)\n",
        "\n",
        "    if lines:\n",
        "        for i in range(\n",
        "            n_sqrt * lc.kernel_size[0],\n",
        "            n_sqrt * lc.conv_size[0] * lc.kernel_size[0],\n",
        "            n_sqrt * lc.kernel_size[0],\n",
        "        ):\n",
        "            ax.axhline(i - 0.5, color=color, linestyle=\"--\")\n",
        "\n",
        "        for i in range(\n",
        "            n_sqrt * lc.kernel_size[1],\n",
        "            n_sqrt * lc.conv_size[1] * lc.kernel_size[1],\n",
        "            n_sqrt * lc.kernel_size[1],\n",
        "        ):\n",
        "            ax.axvline(i - 0.5, color=color, linestyle=\"--\")\n",
        "\n",
        "    ax.set_xticks(())\n",
        "    ax.set_yticks(())\n",
        "    ax.set_aspect(\"auto\")\n",
        "\n",
        "    plt.colorbar(im, cax=cax)\n",
        "    fig.tight_layout()\n",
        "\n",
        "    return im\n",
        "\n",
        "def plot_LC_activation_map(lc : object,\n",
        "    spikes: torch.tensor,\n",
        "    input_channel: int = 0,\n",
        "    scale_factor: float = 1.0,\n",
        "    lines: bool = True,\n",
        "    figsize: Tuple[int, int] = (5, 5),\n",
        "    cmap: str = \"hot_r\",\n",
        "    color: str='r'\n",
        "    ) -> AxesImage:\n",
        "    # language=rst\n",
        "    \"\"\"\n",
        "    Plot an activation map of a :code:`Connection` with `locally connected\n",
        "    structure <http://yann.lecun.com/exdb/publis/pdf/gregor-nips-11.pdf>_.\n",
        "    :param lc: LC connection object of LCNet\n",
        "    :param input_channel: indicates weights which connected to this channel of input \n",
        "    :param scale_factor: determines intensity of activation map \n",
        "    :param lines: Whether or not to draw horizontal and vertical lines separating input regions.\n",
        "    :param figsize: Horizontal, vertical figure size in inches.\n",
        "    :param cmap: Matplotlib colormap.\n",
        "    :return: Used for re-drawing the weights plot.\n",
        "    \"\"\"\n",
        "    spikes = spikes.sum(0).squeeze().view(lc.conv_size[0]*int(np.sqrt(lc.out_channels)),lc.conv_size[1]*int(np.sqrt(lc.out_channels)))\n",
        "    x = scale_factor * spikes / torch.max(spikes)\n",
        "    x = x.clip(lc.wmin,lc.wmax).repeat_interleave(lc.kernel_size[0], dim=0).repeat_interleave(lc.kernel_size[1], dim=1).cpu()\n",
        "    n_sqrt = int(np.ceil(np.sqrt(lc.out_channels)))\n",
        "\n",
        "    sel_slice = lc.w.view(lc.in_channels, lc.out_channels, lc.conv_size[0], lc.conv_size[1], lc.kernel_size[0], lc.kernel_size[1]).cpu()\n",
        "    sel_slice = sel_slice[input_channel, ...]\n",
        "    input_size = _pair(int(np.sqrt(lc.source.n)))\n",
        "    reshaped = reshape_LC_weights(sel_slice, lc.out_channels, lc.kernel_size, lc.conv_size, input_size)\n",
        "\n",
        "    fig, ax = plt.subplots(figsize=figsize)\n",
        "\n",
        "    im = ax.imshow(reshaped.cpu()*x, cmap=cmap, vmin=lc.wmin, vmax=lc.wmax)\n",
        "    div = make_axes_locatable(ax)\n",
        "    cax = div.append_axes(\"right\", size=\"5%\", pad=0.05)\n",
        "\n",
        "    if lines:\n",
        "        for i in range(\n",
        "            n_sqrt * lc.kernel_size[0],\n",
        "            n_sqrt * lc.conv_size[0] * lc.kernel_size[0],\n",
        "            n_sqrt * lc.kernel_size[0],\n",
        "        ):\n",
        "            ax.axhline(i - 0.5, color=color, linestyle=\"--\")\n",
        "\n",
        "        for i in range(\n",
        "            n_sqrt * lc.kernel_size[1],\n",
        "            n_sqrt * lc.conv_size[1] * lc.kernel_size[1],\n",
        "            n_sqrt * lc.kernel_size[1],\n",
        "        ):\n",
        "            ax.axvline(i - 0.5, color=color, linestyle=\"--\")\n",
        "\n",
        "    ax.set_xticks(())\n",
        "    ax.set_yticks(())\n",
        "    ax.set_aspect(\"auto\")\n",
        "\n",
        "    plt.colorbar(im, cax=cax)\n",
        "    fig.tight_layout()\n",
        "\n",
        "    return im\n",
        "\n",
        "\n",
        "def reshape_LC_weights(\n",
        "    w: torch.Tensor,\n",
        "    n_filters: int,\n",
        "    kernel_size: Union[int, Tuple[int, int]],\n",
        "    conv_size: Union[int, Tuple[int, int]],\n",
        "    input_sqrt: Union[int, Tuple[int, int]],\n",
        ") -> torch.Tensor:\n",
        "    # language=rst\n",
        "    \"\"\"\n",
        "    Get the weights from a locally connected layer and reshape them to be two-dimensional and square.\n",
        "    :param w: Weights from a locally connected layer.\n",
        "    :param n_filters: No. of neuron filters.\n",
        "    :param kernel_size: Side length(s) of convolutional kernel.\n",
        "    :param conv_size: Side length(s) of convolution population.\n",
        "    :param input_sqrt: Sides length(s) of input neurons.\n",
        "    :return: Locally connected weights reshaped as a collection of spatially ordered square grids.\n",
        "    \"\"\"\n",
        "    k1, k2 = kernel_size\n",
        "    c1, c2 = conv_size\n",
        "    i1, i2 = input_sqrt\n",
        "    c1sqrt, c2sqrt = int(math.ceil(math.sqrt(c1))), int(math.ceil(math.sqrt(c2)))\n",
        "    fs = int(math.ceil(math.sqrt(n_filters)))\n",
        "\n",
        "    w_ = torch.zeros((n_filters * k1, k2 * c1 * c2))\n",
        "\n",
        "    for n1 in range(c1):\n",
        "        for n2 in range(c2):\n",
        "            for feature in range(n_filters):\n",
        "                n = n1 * c2 + n2\n",
        "                filter_ = w[feature, n1, n2, :, :\n",
        "                ].view(k1, k2)\n",
        "                w_[feature * k1 : (feature + 1) * k1, n * k2 : (n + 1) * k2] = filter_\n",
        "\n",
        "    if c1 == 1 and c2 == 1:\n",
        "        square = torch.zeros((i1 * fs, i2 * fs))\n",
        "\n",
        "        for n in range(n_filters):\n",
        "            square[\n",
        "                (n // fs) * i1 : ((n // fs) + 1) * i2,\n",
        "                (n % fs) * i2 : ((n % fs) + 1) * i2,\n",
        "            ] = w_[n * i1 : (n + 1) * i2]\n",
        "\n",
        "        return square\n",
        "    else:\n",
        "        square = torch.zeros((k1 * fs * c1, k2 * fs * c2))\n",
        "\n",
        "        for n1 in range(c1):\n",
        "            for n2 in range(c2):\n",
        "                for f1 in range(fs):\n",
        "                    for f2 in range(fs):\n",
        "                        if f1 * fs + f2 < n_filters:\n",
        "                            square[\n",
        "                                k1 * (n1 * fs + f1) : k1 * (n1 * fs + f1 + 1),\n",
        "                                k2 * (n2 * fs + f2) : k2 * (n2 * fs + f2 + 1),\n",
        "                            ] = w_[\n",
        "                                (f1 * fs + f2) * k1 : (f1 * fs + f2 + 1) * k1,\n",
        "                                (n1 * c2 + n2) * k2 : (n1 * c2 + n2 + 1) * k2,\n",
        "                            ]\n",
        "\n",
        "        return square\n",
        "\n",
        "def plot_semantic_pooling(lc : object,\n",
        "    input_channel: int = 0,\n",
        "    output_channel: int = None,\n",
        "    lines: bool = True,\n",
        "    figsize: Tuple[int, int] = (5, 5),\n",
        "    cmap: str = \"hot_r\",\n",
        "    color: str='r',\n",
        "    ) -> AxesImage:\n",
        "    # language=rst\n",
        "    \"\"\"\n",
        "    Plot a connection weight matrix of a :code:`Connection` with `locally connected\n",
        "    structure <http://yann.lecun.com/exdb/publis/pdf/gregor-nips-11.pdf>_.\n",
        "    :param lc: LC connection object of LCNet\n",
        "    :param input_channel: indicates weights which connected to this channel of input \n",
        "    :param output_channel: indicates weights of specific channel in the output layer\n",
        "    :param lines: Whether or not to draw horizontal and vertical lines separating input regions.\n",
        "    :param figsize: Horizontal, vertical figure size in inches.\n",
        "    :param cmap: Matplotlib colormap.\n",
        "    :return: Used for re-drawing the weights plot.\n",
        "    \"\"\"\n",
        "\n",
        "    n_sqrt = int(np.ceil(np.sqrt(lc.out_channels)))\n",
        "    sel_slice = lc.w.view(lc.in_channels, lc.out_channels, lc.conv_size[0], lc.conv_size[1], lc.kernel_size[0], lc.kernel_size[1]).cpu()\n",
        "    input_size = _pair(int(np.sqrt(lc.source.n)))\n",
        "    \n",
        "    if output_channel is None:\n",
        "        sel_slice = sel_slice[input_channel, ...]\n",
        "        reshaped = reshape_LC_weights(sel_slice, lc.out_channels, lc.kernel_size, lc.conv_size, input_size)\n",
        "    else:\n",
        "        sel_slice = sel_slice[input_channel, output_channel, ...]\n",
        "        sel_slice = sel_slice.unsqueeze(0)\n",
        "        reshaped = reshape_LC_weights(sel_slice, 1, lc.kernel_size, lc.conv_size, input_size)\n",
        "        print(reshaped.shape)\n",
        "\n",
        "    fig, ax = plt.subplots(figsize=figsize)\n",
        "\n",
        "    im = ax.imshow(reshaped.cpu(), cmap=cmap, vmin=lc.wmin, vmax=lc.wmax)\n",
        "    div = make_axes_locatable(ax)\n",
        "    cax = div.append_axes(\"right\", size=\"5%\", pad=0.05)\n",
        "\n",
        "    if lines and  output_channel is None:\n",
        "        for i in range(\n",
        "            n_sqrt * lc.kernel_size[0],\n",
        "            n_sqrt * lc.conv_size[0] * lc.kernel_size[0],\n",
        "            n_sqrt * lc.kernel_size[0],\n",
        "        ):\n",
        "            ax.axhline(i - 0.5, color=color, linestyle=\"--\")\n",
        "\n",
        "        for i in range(\n",
        "            n_sqrt * lc.kernel_size[1],\n",
        "            n_sqrt * lc.conv_size[1] * lc.kernel_size[1],\n",
        "            n_sqrt * lc.kernel_size[1],\n",
        "        ):\n",
        "            ax.axvline(i - 0.5, color=color, linestyle=\"--\")\n",
        "\n",
        "    ax.set_xticks(())\n",
        "    ax.set_yticks(())\n",
        "    ax.set_aspect(\"auto\")\n",
        "\n",
        "    plt.colorbar(im, cax=cax)\n",
        "    fig.tight_layout()\n",
        "\n",
        "    return im\n",
        "\n",
        "def plot_LC_weights(lc : object,\n",
        "    input_channel: int = 0,\n",
        "    output_channel: int = None,\n",
        "    lines: bool = True,\n",
        "    figsize: Tuple[int, int] = (5, 5),\n",
        "    cmap: str = \"hot_r\",\n",
        "    color: str='r',\n",
        "    ) -> AxesImage:\n",
        "    # language=rst\n",
        "    \"\"\"\n",
        "    Plot a connection weight matrix of a :code:`Connection` with `locally connected\n",
        "    structure <http://yann.lecun.com/exdb/publis/pdf/gregor-nips-11.pdf>_.\n",
        "    :param lc: LC connection object of LCNet\n",
        "    :param input_channel: indicates weights which connected to this channel of input \n",
        "    :param output_channel: indicates weights of specific channel in the output layer\n",
        "    :param lines: Whether or not to draw horizontal and vertical lines separating input regions.\n",
        "    :param figsize: Horizontal, vertical figure size in inches.\n",
        "    :param cmap: Matplotlib colormap.\n",
        "    :return: Used for re-drawing the weights plot.\n",
        "    \"\"\"\n",
        "\n",
        "    n_sqrt = int(np.ceil(np.sqrt(lc.out_channels)))\n",
        "    sel_slice = lc.w.view(lc.in_channels, lc.out_channels, lc.conv_size[0], lc.conv_size[1], lc.kernel_size[0], lc.kernel_size[1]).cpu()\n",
        "    input_size = _pair(int(np.sqrt(lc.source.n)))\n",
        "    \n",
        "    if output_channel is None:\n",
        "        sel_slice = sel_slice[input_channel, ...]\n",
        "        reshaped = reshape_LC_weights(sel_slice, lc.out_channels, lc.kernel_size, lc.conv_size, input_size)\n",
        "    else:\n",
        "        sel_slice = sel_slice[input_channel, output_channel, ...]\n",
        "        sel_slice = sel_slice.unsqueeze(0)\n",
        "        reshaped = reshape_LC_weights(sel_slice, 1, lc.kernel_size, lc.conv_size, input_size)\n",
        "        #print(reshaped.shape)\n",
        "\n",
        "    fig, ax = plt.subplots(figsize=figsize)\n",
        "\n",
        "    im = ax.imshow(reshaped.cpu(), cmap=cmap, vmin=lc.wmin, vmax=lc.wmax)\n",
        "    div = make_axes_locatable(ax)\n",
        "    cax = div.append_axes(\"right\", size=\"5%\", pad=0.05)\n",
        "\n",
        "    if lines and  output_channel is None:\n",
        "        for i in range(\n",
        "            lc.kernel_size[0],#n_sqrt * lc.kernel_size[0],\n",
        "            n_sqrt*lc.conv_size[0] * lc.kernel_size[0],#n_sqrt * lc.conv_size[0] * lc.kernel_size[0],\n",
        "            lc.kernel_size[0],#,n_sqrt * lc.kernel_size[0],\n",
        "        ):\n",
        "            #print(i)\n",
        "            ax.axhline(i, color=color, linestyle=\"--\")\n",
        "\n",
        "        for i in range(\n",
        "            lc.kernel_size[1],#n_sqrt * lc.kernel_size[1],\n",
        "            n_sqrt*lc.conv_size[1] * lc.kernel_size[1],#n_sqrt * lc.conv_size[1] * lc.kernel_size[1],\n",
        "            lc.kernel_size[1],#n_sqrt * lc.kernel_size[1],\n",
        "        ):\n",
        "            ax.axvline(i, color=color, linestyle=\"--\")\n",
        "            \n",
        "        for i in range(\n",
        "            n_sqrt * lc.kernel_size[0],\n",
        "            n_sqrt * lc.conv_size[0] * lc.kernel_size[0],\n",
        "            n_sqrt * lc.kernel_size[0],\n",
        "        ):\n",
        "            #print(i)\n",
        "            ax.axhline(i, color='b', linestyle=\"--\")\n",
        "\n",
        "        for i in range(\n",
        "            n_sqrt * lc.kernel_size[1],\n",
        "            n_sqrt * lc.conv_size[1] * lc.kernel_size[1],\n",
        "            n_sqrt * lc.kernel_size[1],\n",
        "        ):\n",
        "            ax.axvline(i, color='b', linestyle=\"--\")\n",
        "\n",
        "    ax.set_xticks(())\n",
        "    ax.set_yticks(())\n",
        "    ax.set_aspect(\"auto\")\n",
        "\n",
        "    plt.colorbar(im, cax=cax)\n",
        "    fig.tight_layout()\n",
        "\n",
        "    return im\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ywXyWP0I83Au"
      },
      "source": [
        "# Design network"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "id": "PcU9FSsVi4Bw",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "0dca7c35-6fa9-47b7-f765-e8ede5df23d4"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33msingularbrain\u001b[0m (use `wandb login --relogin` to force relogin)\n"
          ]
        }
      ],
      "source": [
        "WANDB = True\n",
        "if WANDB:\n",
        "    !pip install -q wandb\n",
        "    !wandb login\n",
        "    import wandb"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 51,
      "metadata": {
        "id": "8bZpJmlrJDa9"
      },
      "outputs": [],
      "source": [
        "compute_size = lambda inp_size, k, s: int((inp_size-k)/s) + 1\n",
        "def convergence(c):\n",
        "    if c.norm is None:\n",
        "        return 1-torch.mean((c.w-c.wmin)*(c.wmax-c.w))/((c.wmax-c.wmin)/2)**2\n",
        "    else:\n",
        "        mean_norm_factor = c.norm / c.w.shape[-1]\n",
        "        return  1-(torch.mean((c.w-c.wmin)*(c.wmax-c.w))/((c.wmax-c.wmin)/2)**2)\n",
        "\n",
        "        \n",
        "class LCNet(Network):\n",
        "    def __init__(\n",
        "        self,\n",
        "        time: int,\n",
        "        n_actions: int,\n",
        "        neuron_per_action: int,\n",
        "        in_channels : int,\n",
        "        n_channels1: int,\n",
        "        n_channels2: int,\n",
        "        filter_size1: int,\n",
        "        filter_size2: int,\n",
        "        stride1: int,\n",
        "        stride2: int,\n",
        "        maxPool1: bool,\n",
        "        maxPool2: bool,\n",
        "        online: bool,\n",
        "        deep: bool,\n",
        "        reward_fn,\n",
        "        n_neurons: int,\n",
        "        pre_observation: bool,\n",
        "        has_decision_period: bool,\n",
        "        local_rewarding: bool,\n",
        "        nu_LC: Union[float, Tuple[float, float]],\n",
        "        nu_LC2: Union[float, Tuple[float, float]],\n",
        "        nu_Output: float,\n",
        "        dt: float,\n",
        "        crop_size:int ,\n",
        "        nu_inh_LC: float,\n",
        "        nu_inh: float,\n",
        "        inh_type,\n",
        "        inh_LC: bool,\n",
        "        inh_LC2: bool,\n",
        "        inh_factor_LC: float,\n",
        "        inh_factor_LC2: float,\n",
        "        inh_factor:float,\n",
        "        single_output_layer:bool,\n",
        "        NodesType_LC,\n",
        "        NodesType_Output, \n",
        "        update_rule_LC,\n",
        "        update_rule_LC2,\n",
        "        update_rule_Output,\n",
        "        update_rule_inh,\n",
        "        update_rule_inh_LC,\n",
        "        wmin: float,\n",
        "        wmax: float ,\n",
        "        soft_bound,\n",
        "        theta_plus: float,\n",
        "        tc_theta_decay: float,\n",
        "        tc_trace:int,\n",
        "        normal_init:bool,\n",
        "        mu: float,\n",
        "        std:float,\n",
        "        norm_factor_fc,\n",
        "        norm_factor_inh_LC: bool,\n",
        "        norm_factor_LC,\n",
        "        norm_factor_LC2,\n",
        "        norm_factor_out,\n",
        "        norm_factor_inh,\n",
        "        trace_additive,\n",
        "        load_path,\n",
        "        save_path,\n",
        "        LC_weights_path,\n",
        "        LC2_weights_path,\n",
        "        confusion_matrix,\n",
        "        lc_weights_vis,\n",
        "        out_weights_vis,\n",
        "        lc_convergence_vis,\n",
        "        out_convergence_vis,\n",
        "        thresh_LC,\n",
        "        thresh_FC,\n",
        "        num_episodes,\n",
        "        max_steps,\n",
        "        NodeType_FC,\n",
        "        add_layer_fc,\n",
        "        num_neurons_in_fc,\n",
        "        nu_fc,\n",
        "        update_rule_fc,\n",
        "        inh_fc,\n",
        "        inh_factor_fc,\n",
        "        classic_reward,\n",
        "        cartpole_online,\n",
        "        wandb_active = False,\n",
        "        batch_size=1,\n",
        "        save_interval = 100,\n",
        "\n",
        "\n",
        "        **kwargs,\n",
        "    ) -> None:\n",
        "        # language=rst\n",
        "        \"\"\"\n",
        "        Constructor for class ``BioLCNet``.\n",
        "\n",
        "        :param n_inpt: Number of input neurons. Matches the 1D size of the input data.\n",
        "        :param n_neurons: Number of excitatory, inhibitory neurons.\n",
        "        :param exc: Strength of synapse weights from excitatory to inhibitory layer.\n",
        "        :param inh: Strength of synapse weights from inhibitory to excitatory layer.\n",
        "        :param dt: Simulation time step.\n",
        "        :param nu: Single or pair of learning rates for pre- and post-synaptic events,\n",
        "            respectively.\n",
        "        :param reduction: Method for reducing parameter updates along the minibatch\n",
        "            dimension.\n",
        "        :param wmin: Minimum allowed weight on input to excitatory synapses.\n",
        "        :param wmax: Maximum allowed weight on input to excitatory synapses.\n",
        "        :param norm: Input to excitatory layer connection weights normalization\n",
        "            constant.\n",
        "        :param theta_plus: On-spike increment of ``(adaptive)LIFNodes`` membrane\n",
        "            threshold potential.\n",
        "        :param tc_theta_decay: Time constant of ``(adaptive)LIFNodes`` threshold\n",
        "            potential decay.\n",
        "        :param inpt_shape: The dimensionality of the input layer.\n",
        "        \"\"\"\n",
        "        manual_seed(SEED)\n",
        "        super().__init__(dt=dt, reward_fn = None, online=online)\n",
        "        kwargs['single_output_layer'] = single_output_layer\n",
        "        kwargs['dt'] = dt\n",
        "        kwargs['n_labels'] = n_actions\n",
        "        kwargs['neuron_per_actionn'] = neuron_per_action\n",
        "        self.true_label = 0\n",
        "        self.dt = dt\n",
        "        self.cartpole_online = True\n",
        "        self.intensity = kwargs['intensity']\n",
        "        self.reward_fn = reward_fn\n",
        "        self.batch_size = batch_size\n",
        "        self.reward_fn.network = self\n",
        "        self.reward_fn.dt = self.dt\n",
        "        self.n_actions = n_actions\n",
        "        self.neuron_per_action = neuron_per_action\n",
        "        self.n_classes = n_actions\n",
        "        self.classic_reward = classic_reward\n",
        "        self.add_layer_fc = add_layer_fc\n",
        "        self.num_neurons_in_fc = num_neurons_in_fc\n",
        "        self.neuron_per_class = neuron_per_action\n",
        "        self.save_path = save_path\n",
        "        self.load_path = load_path\n",
        "        self.save_interval = save_interval\n",
        "        self.deep = deep\n",
        "        self.maxPool1 = maxPool1\n",
        "        self.maxPool2 = maxPool2\n",
        "        self.time = time\n",
        "        self.crop_size = crop_size\n",
        "        self.filter_size1 = filter_size1\n",
        "        self.filter_size2 = filter_size2\n",
        "        self.clamp_intensity = kwargs.get('clamp_intensity',None)\n",
        "        self.single_output_layer = single_output_layer\n",
        "        self.pre_observation = pre_observation\n",
        "        self.has_decision_period = has_decision_period\n",
        "        self.local_rewarding = local_rewarding\n",
        "        self.soft_bound = soft_bound\n",
        "        self.confusion_matrix = confusion_matrix\n",
        "        self.lc_weights_vis = lc_weights_vis\n",
        "        self.out_weights_vis = out_weights_vis\n",
        "        self.lc_convergence_vis = lc_convergence_vis\n",
        "        self.out_convergence_vis = out_convergence_vis\n",
        "        self.frame_analysis = frame_analysis\n",
        "        self.in_channels = in_channels\n",
        "        self.n_channels1 = n_channels1\n",
        "        self.n_channels2 = n_channels2\n",
        "        self.stride1 = stride1\n",
        "        self.stride2 = stride2\n",
        "        self.convergences = {}\n",
        "        self.norm_factor_LC = norm_factor_LC\n",
        "        self.norm_factor_LC2 = norm_factor_LC2\n",
        "        self.norm_factor_out = norm_factor_out\n",
        "        self.wmin = wmin \n",
        "        self.prev_screen = None\n",
        "        self.screen = None\n",
        "        self.wmax = wmax\n",
        "        self.wandb_active = wandb_active\n",
        "        self.epochs_trained = 0\n",
        "        self.num_episodes = num_episodes\n",
        "        self.max_steps= max_steps\n",
        "        self.rew = 0.0\n",
        "        self.env = gym.make('CartPole-v0')\n",
        "        self.env.reset()\n",
        "\n",
        "        self.time_analysis = kwargs.get('time_analysis', False)\n",
        "        if kwargs['variant'] == 'scalar':\n",
        "            assert self.has_decision_period == True, ''\n",
        "\n",
        "        if self.online == False:\n",
        "            assert self.has_decision_period == True, ''\n",
        "        \n",
        "        if self.has_decision_period == True:\n",
        "            assert self.online == False, \"Decision period is not compatible with online learning.\"\n",
        "            self.observation_period = kwargs['observation_period']\n",
        "            assert self.observation_period >= 0, \"\"\n",
        "            self.decision_period = kwargs['decision_period']\n",
        "            assert self.decision_period > 0, \"\"\n",
        "            self.learning_period = self.time - self.observation_period - self.decision_period\n",
        "\n",
        "        elif self.pre_observation == True:\n",
        "            self.observation_period = kwargs['observation_period']\n",
        "            assert self.observation_period >= 0, \"\"\n",
        "            self.learning_period = self.time - self.observation_period\n",
        "            self.decision_period = self.time - self.observation_period\n",
        "\n",
        "        else:\n",
        "            self.observation_period = 0\n",
        "            self.decision_period = self.time\n",
        "            self.learning_period = self.time\n",
        "\n",
        "        ### nodes\n",
        "        inp = Input(shape= [in_channels,crop_size,crop_size], traces=True, tc_trace=tc_trace,traces_additive = trace_additive)\n",
        "        self.add_layer(inp, name=\"input\")\n",
        "\n",
        "        ## First hidden layer\n",
        "        conv_size1 = compute_size(crop_size, filter_size1, stride1)\n",
        "        main1 = NodesType_LC(shape= [n_channels1, conv_size1, conv_size1], thresh = thresh_LC, traces=True, tc_trace=tc_trace,\n",
        "                             traces_additive = trace_additive,tc_theta_decay = tc_theta_decay, theta_plus = theta_plus)\n",
        "        \n",
        "        self.add_layer(main1, name=\"main1\")\n",
        "\n",
        "        ### connections \n",
        "        LC1 = LocalConnectionOrig(inp, main1, filter_size1, stride1, n_channels1,\\\n",
        "                              nu = _pair(nu_LC), update_rule = update_rule_LC,wmin = wmin, wmax= wmax, norm = norm_factor_LC)\n",
        "\n",
        "        # LC1 = LocalConnection(inp, main1, filter_size1, stride1, in_channels, n_channels1,input_shape=(crop_size,crop_size),\\\n",
        "        #                      nu = _pair(nu_LC), update_rule = update_rule_LC,wmin = wmin, wmax= wmax, soft_bound = soft_bound, norm = norm_factor_LC)\n",
        "\n",
        "\n",
        "        if LC_weights_path:\n",
        "            a = torch.load(LC_weights_path)\n",
        "            LC1.w.data = a['state_dict']['input_to_main1.w']\n",
        "            print(\"Weights loaded ...\")\n",
        "        \n",
        "        elif normal_init:\n",
        "            w_lc_init = torch.normal(mu,std, size = (in_channels, n_channels1 * compute_size(crop_size, filter_size1, stride1)**2, filter_size1**2))\n",
        "            LC1.w.data = w_lc_init\n",
        "       \n",
        "        self.add_connection(LC1, \"input\", \"main1\")\n",
        "        self.convergences['lc1'] = []\n",
        "\n",
        "        if inh_LC:\n",
        "            main_width = compute_size(crop_size, filter_size1, stride1)\n",
        "            w_inh_LC = torch.zeros(n_channels1,main_width,main_width,n_channels1,main_width,main_width)\n",
        "            for c in range(n_channels1):\n",
        "                for w1 in range(main_width):\n",
        "                    for w2 in range(main_width):\n",
        "                        w_inh_LC[c,w1,w2,:,w1,w2] = - inh_factor_LC\n",
        "                        w_inh_LC[c,w1,w2,c,w1,w2] = 0\n",
        "        \n",
        "            w_inh_LC = w_inh_LC.reshape(main1.n,main1.n)\n",
        "                                                             \n",
        "            LC_recurrent_inhibition = Connection(\n",
        "                source=main1,\n",
        "                target=main1,\n",
        "                w=w_inh_LC,\n",
        "            )\n",
        "            self.add_connection(LC_recurrent_inhibition, \"main1\", \"main1\")\n",
        "        \n",
        "        self.final_connection_source_name = 'main1'\n",
        "        self.final_connection_source = main1\n",
        "        \n",
        "        if self.add_layer_fc:\n",
        "            FC1 = NodeType_FC(n= self.num_neurons_in_fc, traces=True,traces_additive = trace_additive, thresh=thresh_FC, tc_trace=tc_trace, tc_theta_decay = tc_theta_decay, theta_plus = theta_plus)\n",
        "\n",
        "            self.add_layer(FC1, \"fc1\")\n",
        "\n",
        "            main_fc = Connection(main1, FC1, nu = nu_fc, update_rule = update_rule_fc, wmin = wmin, wmax= wmax, norm = norm_factor_fc)\n",
        "            \n",
        "            self.add_connection(main_fc, \"main1\", \"fc1\")\n",
        "            self.final_connection_source_name = 'fc1'\n",
        "            self.final_connection_source = FC1\n",
        "\n",
        "            if inh_fc:\n",
        "                w = -inh_factor_fc * (torch.ones(FC1.n, FC1.n) - torch.eye(FC1.n, FC1.n))\n",
        "                fc_recurrent_inhibition = Connection(\n",
        "                    source=FC1,\n",
        "                    target=FC1,\n",
        "                    w=w,\n",
        "                    update_rule = None,\n",
        "                    wmin=-inh_factor_fc,\n",
        "                    wmax=0,\n",
        "                )\n",
        "                self.add_connection(fc_recurrent_inhibition, \"fc1\", \"fc1\")\n",
        "\n",
        "        self.hidden2 = main1\n",
        "        self.hidden2_name = 'main1'\n",
        "        if maxPool1:\n",
        "            maxPool_kernel = 2\n",
        "            maxPool_stride = 2\n",
        "            \n",
        "            conv_size1 =compute_size(conv_size1, maxPool_kernel, maxPool_stride)\n",
        "            self.final_connection_source_name = 'maxpool1'\n",
        "            \n",
        "            maxpool1 = LIFNodes(shape= [self.n_channels1, conv_size1, conv_size1], refrac = 0)\n",
        "            self.add_layer(maxpool1, name=\"maxpool1\")\n",
        "            self.final_connection_source = maxpool1\n",
        "            \n",
        "            maxPoolConnection = MaxPool2dLocalConnection(main1, maxpool1, maxPool_kernel, maxPool_stride)\n",
        "            self.add_connection(maxPoolConnection, \"main1\", 'maxpool1')\n",
        "            \n",
        "            self.hidden2 = maxpool1\n",
        "            self.hidden2_name = 'maxpool1'\n",
        "\n",
        "        if deep:\n",
        "            # # Second hidden layer\n",
        "            conv_size2 = compute_size(conv_size1, filter_size2, stride2)\n",
        "\n",
        "            main2 = NodesType_LC(shape= [n_channels2, conv_size2, conv_size2],traces=True, tc_trace=tc_trace,traces_additive = trace_additive,\n",
        "                                            tc_theta_decay = tc_theta_decay, theta_plus = theta_plus)\n",
        "            \n",
        "            self.add_layer(main2, name=\"main2\")\n",
        "\n",
        "            ### connections \n",
        "            lc2_input_shape = (conv_size1,conv_size1)\n",
        "            LC2 = LocalConnection(self.hidden2, main2, filter_size2, stride2, n_channels1, n_channels2, input_shape= lc2_input_shape,\n",
        "            nu = _pair(nu_LC2), update_rule = update_rule_LC2, wmin = wmin, wmax= wmax, soft_bound = soft_bound, norm = norm_factor_LC2)\n",
        "\n",
        "            self.add_connection(LC2,  self.hidden2_name, \"main2\")\n",
        "            self.convergences['lc2'] = []\n",
        "            if LC2_weights_path:\n",
        "                a = torch.load(LC2_weights_path)\n",
        "                LC2.w.data = a['state_dict']['main1_to_main2.w']\n",
        "                print(\"Weights loaded ...\")\n",
        "                \n",
        "            \n",
        "            elif normal_init:\n",
        "                w_lc_init = torch.normal(mu,std, size = (n_channels1, n_channels2 * compute_size(conv_size1, filter_size2, stride2)**2, filter_size2**2))\n",
        "                LC2.w.data = w_lc_init\n",
        "\n",
        "            self.final_connection_source_name = 'main2'\n",
        "            self.final_connection_source = main2\n",
        "\n",
        "            if inh_LC2:\n",
        "                main_width = conv_size2\n",
        "                w_inh_LC2 = torch.zeros(n_channels2,main_width,main_width,n_channels2,main_width,main_width)\n",
        "                for c in range(n_channels2):\n",
        "                    for w1 in range(main_width):\n",
        "                        for w2 in range(main_width):\n",
        "                            w_inh_LC2[c,w1,w2,:,w1,w2] = - inh_factor_LC2\n",
        "                            w_inh_LC2[c,w1,w2,c,w1,w2] = 0\n",
        "            \n",
        "                w_inh_LC2 = w_inh_LC2.reshape(main2.n,main2.n)\n",
        "                                                                \n",
        "                LC_recurrent_inhibition2 = Connection(\n",
        "                    source=main2,\n",
        "                    target=main2,\n",
        "                    w=w_inh_LC2,\n",
        "                )\n",
        "                self.add_connection(LC_recurrent_inhibition2, \"main2\", \"main2\")\n",
        "\n",
        "\n",
        "            if maxPool2:\n",
        "                maxPool_kernel = 2\n",
        "                maxPool_stride = 2\n",
        "                conv_size2 =compute_size(conv_size2, maxPool_kernel, maxPool_stride)\n",
        "                self.final_connection_source_name = 'maxpool2'\n",
        "                maxpool2 = LIFNodes(shape= [self.n_channels2, conv_size2, conv_size2], refrac = 0)\n",
        "                self.final_connection_source = maxpool2\n",
        "                maxPoolConnection2 = MaxPool2dLocalConnection(main2, maxpool2, maxPool_kernel, maxPool_stride)\n",
        "\n",
        "                self.add_layer(maxpool2, name=\"maxpool2\")\n",
        "                self.add_connection(maxPoolConnection2, \"main2\", 'maxpool2')\n",
        "\n",
        "\n",
        "        ### main2 to output\n",
        "        out = NodesType_Output(n= n_neurons, traces=True,traces_additive = trace_additive, thresh=thresh_FC, tc_trace=tc_trace, tc_theta_decay = tc_theta_decay, theta_plus = theta_plus)\n",
        "\n",
        "        self.add_layer(out, \"output\")\n",
        "\n",
        "        last_main_out = Connection(self.final_connection_source, out, nu = nu_Output, update_rule = update_rule_Output, wmin = wmin, wmax= wmax, norm = norm_factor_out)\n",
        "\n",
        "        if normal_init:\n",
        "            w_last_main_init = torch.normal(mu,std,size = (self.final_connection_source.n,out.n)) \n",
        "            last_main_out.w.data = w_last_main_init\n",
        "\n",
        "        self.add_connection(last_main_out, self.final_connection_source_name, \"output\")\n",
        "        self.convergences['last_main_out'] = []\n",
        "        ### Inhibitory:\n",
        "        if inh_type == 'between_layers':\n",
        "            w = -inh_factor * torch.ones(out.n, out.n)\n",
        "            for c in range(n_actions):\n",
        "                ind = slice(c*neuron_per_action,(c+1)*neuron_per_action)\n",
        "                w[ind, ind] = 0\n",
        "\n",
        "            out_recurrent_inhibition = Connection(\n",
        "                source=out,\n",
        "                target=out,\n",
        "                w=w,\n",
        "                update_rule = update_rule_inh,\n",
        "                wmin=-inh_factor,\n",
        "                wmax=0,\n",
        "                nu = nu_inh,\n",
        "                norm = norm_factor_inh,\n",
        "            )\n",
        "            self.add_connection(out_recurrent_inhibition, \"output\", \"output\")\n",
        "        elif inh_type == 'one_2_all':\n",
        "            w = -inh_factor * (torch.ones(out.n, out.n) - torch.eye(out.n, out.n))\n",
        "            out_recurrent_inhibition = Connection(\n",
        "                source=out,\n",
        "                target=out,\n",
        "                w=w,\n",
        "                update_rule = update_rule_inh,\n",
        "                wmin=-inh_factor,\n",
        "                wmax=0,\n",
        "                nu = nu_inh,\n",
        "                norm = norm_factor_inh,\n",
        "            )\n",
        "            self.add_connection(out_recurrent_inhibition, \"output\", \"output\")\n",
        "        # Diehl and Cook\n",
        "        elif inh_type == 'DC':\n",
        "            raise NotImplementedError('Diehl and cook not implemented yet fo r 10 classes')\n",
        "        elif inh_type == None:\n",
        "            pass\n",
        "        # Directs network to GPU\n",
        "        if gpu:\n",
        "            self.to(\"cuda\")\n",
        "\n",
        "    def frame_process(self, x):\n",
        "        x[x<1.0] = 2.0\n",
        "        x[x==1.0] = 0.0\n",
        "        x[x==2.0] = 1.0\n",
        "        return x\n",
        "\n",
        "\n",
        "    def get_state_spiking(self):\n",
        "        intensity = self.intensity\n",
        "        screen = self.env.render(mode='rgb_array')\n",
        "        screen = screen.transpose((2, 0, 1))\n",
        "        _, screen_height, screen_width = screen.shape\n",
        "        screen = screen[:, int(screen_height*0.4):int(screen_height * 0.8)]\n",
        "        view_width = int(screen_width * 0.6)\n",
        "        world_width = self.env.x_threshold * 2\n",
        "        scale = screen_width / world_width\n",
        "        cart_location = int(self.env.state[0] * scale + screen_width / 2.0)\n",
        "        if cart_location < view_width // 2:\n",
        "            slice_range = slice(view_width)\n",
        "        elif cart_location > (screen_width - view_width // 2):\n",
        "            slice_range = slice(-view_width, None)\n",
        "        else:\n",
        "            slice_range = slice(cart_location - view_width // 2,\n",
        "                                cart_location + view_width // 2)\n",
        "            \n",
        "        # Strip off the edges, so that we have a square image centered on a cart\n",
        "        screen = screen[:, :, slice_range]\n",
        "\n",
        "        # Convert to float, rescale, convert to torch tensor\n",
        "        screen = np.ascontiguousarray(screen, dtype=np.float32) / 255\n",
        "        screen = torch.from_numpy(screen)\n",
        "        h = screen.shape[1]\n",
        "        w = screen.shape[2]\n",
        "        resize = transforms.Compose([\n",
        "                    transforms.ToPILImage(),\n",
        "                    transforms.Resize([80, 180]),\n",
        "                    transforms.Lambda(lambda x: crop(x, 0, 60, 60, 60)),\n",
        "                    transforms.Lambda(lambda x: crop(x, 0, 10, 40, 40)),\n",
        "                    transforms.Grayscale(),\n",
        "                    transforms.ToTensor(),])\n",
        "\n",
        "        process = transforms.Lambda(self.frame_process)\n",
        "        spiking = transforms.Compose([transforms.Lambda(lambda x: x * self.intensity),\n",
        "                    transforms.Lambda(lambda x: PoissonEncoder(time=self.time, dt=1)(x))])\n",
        "        screen = resize(screen)\n",
        "        st = None\n",
        "        if self.prev_screen == None:\n",
        "            st = torch.zeros(*screen.shape)\n",
        "            st = spiking(st)\n",
        "        else:\n",
        "            gradient = screen-self.prev_screen\n",
        "            if torch.max(gradient)-torch.min(gradient) != 0:\n",
        "                normalized = (gradient-torch.min(gradient))/(torch.max(gradient)-torch.min(gradient))\n",
        "            else:\n",
        "                normalized = gradient\n",
        "            normalized = process(normalized)\n",
        "            if (normalized==1).sum() > (normalized==0).sum():\n",
        "                normalized = process(normalized)\n",
        "            st = spiking(normalized)\n",
        "\n",
        "\n",
        "        self.prev_screen = copy.deepcopy(screen)\n",
        "        if self.frame_analysis:\n",
        "            f = st.sum(axis=0)\n",
        "            f = f.to(device)\n",
        "            plt.imshow(f.cpu().numpy().squeeze(), cmap='gray')\n",
        "            plt.show()\n",
        "        \n",
        "        return st\n",
        "\n",
        "    def learn(\n",
        "        self,\n",
        "        hparams = None,\n",
        "        online_validate = True,\n",
        "        running_window_length = 250,\n",
        "        verbose = True,\n",
        "        **kwargs\n",
        "    ):\n",
        "        manual_seed(SEED)\n",
        "        if self.wandb_active:\n",
        "            wandb.watch(self)\n",
        "        self.verbose = verbose\n",
        "\n",
        "        reward_monitor = RewardMonitor(time =self.time)\n",
        "\n",
        "        self.add_monitor(reward_monitor, name=\"reward\")\n",
        "        kwargs['target_name'] = None\n",
        "        reward_hist = []\n",
        "\n",
        "        self.spikes = {}\n",
        "        for layer in set(self.layers):\n",
        "            self.spikes[layer] = Monitor(self.layers[layer], state_vars=[\"s\"], time=None)\n",
        "            self.add_monitor(self.spikes[layer], name=\"%s_spikes\" % layer)\n",
        "            self.dopaminergic_layers = self.layers[\"output\"]\n",
        "       \n",
        "        self.episode = 0\n",
        "        rew = 0.0\n",
        "        tot_rew = 0.0\n",
        "\n",
        "\n",
        "        if self.load_path:\n",
        "            \n",
        "            self.model_params = torch.load(self.load_path)\n",
        "            self.load_state_dict(torch.load(self.load_path)['state_dict'])\n",
        "            self.episode =  self.model_params['episode']\n",
        "            hparams = self.model_params['hparams']\n",
        "            reward_hist = self.model_params['reward_hist']\n",
        "            print(f'Previous model loaded! Resuming training from episode {self.episode}...\\n')\n",
        "        else:\n",
        "            self.env = gym.make('CartPole-v0')\n",
        "            self.env.reset()\n",
        "            print(f'Previous model not found! Training from the beginning...\\n')\n",
        "\n",
        "        pbar = tqdm(total=self.num_episodes)\n",
        "        \n",
        "        if self.time_analysis:\n",
        "            self.sample_spikes = {'input': [], 'main1': [], 'output': []}\n",
        "\n",
        "        for ep in range(self.num_episodes):\n",
        "            # print(f\"episode: {ep+1}\")\n",
        "            self.reset_state_variables()\n",
        "            self.env.reset()\n",
        "            self.prev_screen = None\n",
        "            done = False\n",
        "            tot_rew = 0.0\n",
        "            success = False\n",
        "            failure = False\n",
        "            num_steps = 0\n",
        "            for t in count():\n",
        "                \n",
        "                if t != 0:\n",
        "                    print(\"state (x, x_dot, theta (rad), theta_dot):\", round(x,3), round(x_dot*0.02, 3), round(theta, 3), round(theta_dot*0.02,3), \"output\", sum_spikes, 'selected_action:',  selected_action, 'rew:', round(rew,3))\n",
        "                    if self.time_analysis:\n",
        "                        self.sample_spikes['input'].append(self.spikes['input'].get('s'))\n",
        "                        self.sample_spikes['main1'].append(self.spikes['main1'].get('s'))\n",
        "                        self.sample_spikes['output'].append(self.spikes['output'].get('s').view(self.time, self.batch_size, n_actions, neuron_per_action))\n",
        "\n",
        "                        # w_lc1 = self.connections[('input', 'main1')].w\n",
        "                        # w_last_main_out = self.connections[(self.final_connection_source_name,'output')].w\n",
        "                \n",
        "                image = self.get_state_spiking()\n",
        "                if gpu:\n",
        "                    inputs = {\"input\": image.cuda().view(self.time, self.batch_size, self.in_channels, self.crop_size, self.crop_size)}\n",
        "                else:\n",
        "                    inputs = {\"input\": image.view(self.time, self.batch_size, self.in_channels, self.crop_size, self.crop_size)}\n",
        "\n",
        "\n",
        "                clamp = {}\n",
        "                if self.clamp_intensity is not None:\n",
        "                    encoder = PoissonEncoder(time = self.time, dt = self.dt)\n",
        "                    clamp['output'] = encoder.enc(datum = torch.rand(self.layers['output'].n)*self.clamp_intensity,time = self.time, dt = self.dt)\n",
        "\n",
        "                # if done:\n",
        "                #     failure = True\n",
        "\n",
        "                if t >= self.max_steps:\n",
        "                    success = True\n",
        "                    done = True\n",
        "                elif done:\n",
        "                    failure = True\n",
        "\n",
        "\n",
        "                self.run(inputs=inputs, \n",
        "                        time = self.time,\n",
        "                        one_step=False,\n",
        "                        clamp = clamp,\n",
        "                        env = self.env,\n",
        "                        success = success,\n",
        "                        failure = failure,\n",
        "                        **kwargs,\n",
        "                        )\n",
        "                \n",
        "                rew = float(reward_monitor.get()[0])\n",
        "                tot_rew += rew\n",
        "                \n",
        "                lc_spikes1 = self.spikes['main1'].get('s')\n",
        "                #lc_spikes2 = self.spikes['main2'].get('s')\n",
        "                out_spikes = self.spikes[\"output\"].get(\"s\").view(self.time, self.batch_size, self.n_actions, self.neuron_per_action)\n",
        "                sum_spikes = out_spikes[self.observation_period:self.observation_period+self.decision_period,:,:].sum(0).sum(2)\n",
        "                selected_action = torch.argmax(sum_spikes, dim=1)\n",
        "\n",
        "                self.spikes['main1'].reset_state_variables()\n",
        "                self.spikes[\"output\"].reset_state_variables()\n",
        "                reward_monitor.reset_state_variables()\n",
        "                #self.reset_state_variables()\n",
        "\n",
        "                if done:\n",
        "                    if success == True:\n",
        "                        print(\"output\", sum_spikes, 'selected_action:',  selected_action, 'rew:', rew)\n",
        "                        print('\\Successful episode!')\n",
        "                        num_steps = t+1\n",
        "                    else:\n",
        "                        print(\"output\", sum_spikes, 'selected_action:',  selected_action, 'rew:', rew)\n",
        "                        print('\\nEpisode not successful!')\n",
        "                        num_steps = t+1\n",
        "                    break\n",
        "\n",
        "                obs, reward, done, _ = self.env.step(int(selected_action[0]))     \n",
        "                x = obs[0]\n",
        "                x_dot = obs[1]\n",
        "                theta = obs[2]\n",
        "                theta_dot = obs[3]\n",
        "\n",
        "                # Get voltage recording.\n",
        "                #main_voltage = main_monitor.get(\"v\")\n",
        "\n",
        "                #tensorboard.update(step= i)\n",
        "\n",
        "\n",
        "            if self.lc_weights_vis:\n",
        "                plot_locally_connected_weights(self.connections[('input','main1')].w, self.n_channels1, self.filter_size1,\n",
        "                                                compute_size(self.crop_size, self.filter_size1, self.stride1), self.connections[('input','main1')].locations,\n",
        "                                                self.crop_size ** 2)\n",
        "                plt.show()\n",
        "\n",
        "            if self.wandb_active:\n",
        "                wandb.log({\n",
        "                        **{'reward': tot_rew},\n",
        "                        **{' to '.join(name) + ' std': c.w.std().item() for name, c in self.connections.items() if name[0]!=name[1]},\n",
        "                        #**{name + ' spikes': monitor.get('s').sum().item() for name, monitor in self.spikes.items()},\n",
        "                        **{' to '.join(name) + \" gradients\": wandb.Histogram(c.w.cpu()) for name, c in self.connections.items() if name[0]!=name[1]},\n",
        "                    },\n",
        "                    step = self.episode)\n",
        "\n",
        "\n",
        "            #self.reward_fn.update() \n",
        "            #Plot_et.plot()    \n",
        "            # self.reset_state_variables()  # Reset state variables.\n",
        "            \n",
        "            self.episode += 1\n",
        "            reward_hist.append(tot_rew)\n",
        "            print(f'\\nEpisode {self.episode} lasted for {num_steps} time steps with total reward of {tot_rew}\\n')\n",
        "            if self.episode % self.save_interval == 0 and self.save_path:\n",
        "                model_params = {'state_dict': self.state_dict(), 'hparams': network_hparams, 'episode': self.episode, 'reward_hist': reward_hist}\n",
        "                torch.save(model_params, self.save_path)\n",
        "                print(\"\\n Model saved!\\n\")\n",
        "\n",
        "            pbar.set_description_str(\"Episode: \"+str(self.episode)+ \", Number of steps: \" + str(num_steps) +\", Episode Total Reward: \" + \"{:.2f}\".format(tot_rew))\n",
        "            pbar.update()\n",
        "\n",
        "    \n",
        "    def single_trial(self):\n",
        "        self.reset_state_variables()\n",
        "        \n",
        "        image = self.get_state_spiking()\n",
        "\n",
        "        if gpu:\n",
        "            inputs = {\"input\": image.cuda().view(self.time, 1, self.in_channels, self.crop_size, self.crop_size)}\n",
        "        else:\n",
        "            inputs = {\"input\": image.view(self.time, 1, self.in_channels, self.crop_size, self.crop_size)}\n",
        "\n",
        "        clamp = {}\n",
        "        if self.clamp_intensity is not None:\n",
        "            encoder = PoissonEncoder(time = self.time, dt = self.dt)\n",
        "            clamp['output'] = encoder.enc(datum = torch.rand(self.layers['output'].n)*self.clamp_intensity,time = self.time, dt = self.dt)\n",
        "\n",
        "        self.run(inputs=inputs, \n",
        "                time=self.time, \n",
        "                **reward_hparams,\n",
        "                one_step = False,\n",
        "                mode = 'RL',\n",
        "                env = 'CartPole'\n",
        "                )\n",
        "    "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kCqAFucAUDb8"
      },
      "source": [
        "# Set up hyper-parameters"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 55,
      "metadata": {
        "id": "3TerGeJoFdzg"
      },
      "outputs": [],
      "source": [
        "n_neurons = 2 #100\n",
        "n_actions = 2\n",
        "neuron_per_action = int(n_neurons/n_actions)\n",
        "single_output_layer = True\n",
        "thresh_LC = -52\n",
        "thresh_FC = -52\n",
        "batch_size = 1\n",
        "epochs = 1\n",
        "crop_size = 40\n",
        "intensity = 200\n",
        "frame_analysis = False\n",
        "lc_weights_vis = False\n",
        "\n",
        "obs = 100\n",
        "dec = 200\n",
        "learn = 100\n",
        "time = obs+dec+learn\n",
        "\n",
        "max_steps = 100\n",
        "num_episodes = 10000\n",
        "\n",
        "filter_size1 = 24\n",
        "stride1 = 8\n",
        "n_channels1 = 100\n",
        "\n",
        "save_interval = 100\n",
        "\n",
        "learning_rate = 0.01\n",
        "norm_factor = 0.25\n",
        "\n",
        "add_layer_fc = False\n",
        "num_neurons_in_fc = 10\n",
        "NodeTypeFC = LIFNodes\n",
        "\n",
        "nu_fc = (0.0001,0.01)\n",
        "update_rule_fc = PostPre\n",
        "\n",
        "classic_reward = True\n",
        "inh_type = 'between_layers'\n",
        "\n",
        "clamp_intensity = 40\n",
        "cartpole_online = True\n",
        "\n",
        "frame_analysis = False\n",
        "\n",
        "network_hparams = {\n",
        "    # net structure\n",
        "    'crop_size': crop_size,\n",
        "    'intensity': intensity,\n",
        "    'round_input': False,\n",
        "    'neuron_per_action': neuron_per_action,\n",
        "    'deep': False,\n",
        "    'maxPool1': False,\n",
        "    'maxPool2': False,\n",
        "    'cartpole_online': cartpole_online,\n",
        "    'in_channels':1,\n",
        "    'n_channels1': n_channels1,\n",
        "    'n_channels2': 64,\n",
        "    'filter_size1': filter_size1,\n",
        "    'filter_size2': 5,\n",
        "    'stride1': stride1,\n",
        "    'stride2': 1,\n",
        "    'n_neurons' : n_neurons,\n",
        "    'n_actions': n_actions,\n",
        "    'single_output_layer': single_output_layer,\n",
        "    'batch_size': batch_size,\n",
        "    'epochs': epochs,\n",
        "    'max_steps': max_steps,\n",
        "    'num_episodes': num_episodes,\n",
        "    'save_interval': save_interval,\n",
        "    'add_layer_fc': add_layer_fc,\n",
        "    'num_neurons_in_fc': num_neurons_in_fc,\n",
        "    'classic_reward': classic_reward,\n",
        "    \n",
        "    # time & Phase\n",
        "    'dt' : 1,\n",
        "    'pre_observation': True,\n",
        "    'has_decision_period': True,\n",
        "    'observation_period': obs,\n",
        "    'decision_period': dec,\n",
        "    'time_analysis': False,\n",
        "    'online': False,\n",
        "    'local_rewarding': False,\n",
        "     \n",
        "    # Nodes\n",
        "    'NodesType_LC': AdaptiveLIFNodes,\n",
        "    'NodesType_Output': LIFNodes, \n",
        "    'NodeType_FC': NodeTypeFC,\n",
        "    'theta_plus': 0.05,\n",
        "    'tc_theta_decay': 1e6,\n",
        "    'tc_trace':20,\n",
        "    'trace_additive' : False,\n",
        "    \n",
        "    # Learning\n",
        "    'update_rule_LC': PostPre,\n",
        "    'update_rule_LC2': None,\n",
        "    'update_rule_Output': MSTDPET,\n",
        "    'update_rule_fc': update_rule_fc,\n",
        "    'update_rule_inh': None,\n",
        "    'update_rule_inh_LC' : None,\n",
        "    'nu_LC': (0.0001,0.01),\n",
        "    'nu_LC2': (0.0,0.0),\n",
        "    'nu_Output': learning_rate,\n",
        "    'nu_fc': nu_fc,\n",
        "    'nu_inh': 0.0,\n",
        "    'nu_inh_LC': 0.0,\n",
        "    'soft_bound': True,\n",
        "    'thresh_LC': thresh_LC,\n",
        "    'thresh_FC': thresh_FC,\n",
        "\n",
        "    # weights\n",
        "    'normal_init': False,\n",
        "    'mu' : 0.8,\n",
        "    'std' : 0.02,\n",
        "    'wmin': 0.0,\n",
        "    'wmax': 1.0,\n",
        "    \n",
        "    # Inhibition\n",
        "    'inh_type': inh_type,\n",
        "    'inh_factor': 100,\n",
        "    'inh_LC': True,\n",
        "    'inh_factor_LC': 100,\n",
        "    'inh_LC2': False,\n",
        "    'inh_factor_LC2': 0,\n",
        "    'inh_fc': False,\n",
        "    'inh_factor_fc': 100,\n",
        "    \n",
        "    # Normalization\n",
        "    'norm_factor_LC': norm_factor*filter_size1*filter_size1,\n",
        "    'norm_factor_fc' : None,\n",
        "    'norm_factor_LC2': None,\n",
        "    'norm_factor_out': None,\n",
        "    'norm_factor_inh': None,\n",
        "    'norm_factor_inh_LC': None,\n",
        "    \n",
        "    # clamp\n",
        "    'clamp_intensity': clamp_intensity,#1000,\n",
        "\n",
        "    # Save\n",
        "    'save_path': '/content/drive/My Drive/LCNet/BioLCNet_gradient_cartpole_100_w.pth',\n",
        "    'load_path': None,#'/content/drive/My Drive/LCNet/BioLCNet_cartpole_100tr.pth',#'/content/drive/My Drive/LCNet/BioLCNet_layer1_Shallow_f15_s4_inh100_norm25_ch100_inh25.pth',\n",
        "    'LC_weights_path': None,#'/content/drive/My Drive/LCNet/BioLCNet_cartpole1.pth',#'/content/drive/My Drive/LCNet/BioLCNet_layer1_Shallow_f15_s4_inh100_norm25_ch100_inh25_weights.pth',#'/content/drive/My Drive/LCNet/LCNet_ch81_f13_22_2norm_Adapt_fc_test2.pth',\n",
        "    'LC2_weights_path': None,#'/content/drive/My Drive/LCNet/DeepLCNet_layer2_ch64_f5_s2_norm3.pth',\n",
        "\n",
        "    # Plot:\n",
        "    'confusion_matrix' : False,\n",
        "    'lc_weights_vis': lc_weights_vis,\n",
        "    'out_weights_vis': False,\n",
        "    'lc_convergence_vis': False,\n",
        "    'out_convergence_vis': False,\n",
        "    'frame_analysis': frame_analysis,\n",
        "\n",
        "    ## reward\n",
        "    'n_labels': n_actions,\n",
        "    'neuron_per_action': neuron_per_action,\n",
        "    \n",
        "    'variant': 'scalar',  #true_pred, #pure_per_spike (Just in phase I, online : True) , and #scalar #per_spike\n",
        "    'tc_reward':0,\n",
        "    'dopamine_base': 0.0,\n",
        "    'reward_base': 1.,\n",
        "    'punishment_base': 1.,\n",
        "    \n",
        "\n",
        "    'sub_variant': 'static', #static, #RPE, #pred_decay\n",
        "    'td_nu': 0.0005,  #RPE\n",
        "    'ema_window': 10, #RPE\n",
        "    'tc_dps': 20,     #pred_decay\n",
        "    'dps_factor': 20, #pred_decay, #RPE\n",
        "    }\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SokdidkrV2Z5"
      },
      "source": [
        "# Learning"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 56,
      "metadata": {
        "id": "Venb2KhSYrT_",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "f2866518-9d13-4dcc-fc35-641f90e127fc"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Drive already mounted at /content/drive; to attempt to forcibly remount, call drive.mount(\"/content/drive\", force_remount=True).\n"
          ]
        }
      ],
      "source": [
        "from google.colab import drive\n",
        "if network_hparams['save_path'] or network_hparams['LC_weights_path']:    \n",
        "    drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oThYyYvHJzeP",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 386,
          "referenced_widgets": [
            "7f5866e1fea24fec903f0e7a0035bb88",
            "b4db786096874228b077abbbc975272d",
            "bc5befd433d648569c1c11b6a66ad4ce",
            "63ca4af57e0f4938a56b6c3f9577f081",
            "f2a003372455459ba6e9b73a9b43b459",
            "f8147581481941be8d286bb362b83a0d",
            "a08d221a81264e768308fa420510e500",
            "5b33649bfe4e4154a08e7029a82424dc",
            "46edfe4e973f47ac8c4f6eac7ceb5b66",
            "b20ee9a88fc54a7297584207ab7bd0e3",
            "2953a20627914e08a3a6ce816bbf8126",
            "952bb30941b44297bbc2dbe06f28214b",
            "aa4a131e00d840138fc297ee81ce5c7c",
            "65ee162348d048899e3186a2362e42e8",
            "cef349c76d7348daa4fdab2489b00158",
            "cab39115348645abaf478f8cf7290025",
            "d2107e81c1604285a695b0f56ca7500a",
            "3df4dc6320c8456ab6319eca97727bb4",
            "3ac0cd6b794747d7a14cc2c4825870ac"
          ]
        },
        "outputId": "6ed733d4-7870-463e-f169-ba4074f3ddf9"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "Finishing last run (ID:1hasvymu) before initializing another..."
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<br/>Waiting for W&B process to finish, PID 2862... <strong style=\"color:green\">(success).</strong>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "7f5866e1fea24fec903f0e7a0035bb88",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "VBox(children=(Label(value=' 0.00MB of 0.00MB uploaded (0.00MB deduped)\\r'), FloatProgress(value=1.0, max=1.0)…"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<style>\n",
              "    table.wandb td:nth-child(1) { padding: 0 10px; text-align: right }\n",
              "    .wandb-row { display: flex; flex-direction: row; flex-wrap: wrap; width: 100% }\n",
              "    .wandb-col { display: flex; flex-direction: column; flex-basis: 100%; flex: 1; padding: 10px; }\n",
              "    </style>\n",
              "<div class=\"wandb-row\"><div class=\"wandb-col\">\n",
              "<h3>Run history:</h3><br/><table class=\"wandb\"><tr><td>input to main1 std</td><td>▁▁▁▁▂▂▂▂▄▄▄▄▄▄▄▄▄▅▅▅▅▅▅▅▅▆▆▆▇▇▇▇▇▇▇▇▇▇██</td></tr><tr><td>main1 to output std</td><td>▁▁▁▂▂▂▂▃▃▃▃▄▄▄▄▄▄▅▅▅▅▅▅▅▅▆▆▆▆▆▆▆▆███████</td></tr><tr><td>reward</td><td>▄▁▂▂▃▂▃▆▅▂▁▂▂▁▃▂▁▂▁▁▁▂▂▂▂█▁▄▁▂▁▁▁▇▂▂▂▁▃▂</td></tr></table><br/></div><div class=\"wandb-col\">\n",
              "<h3>Run summary:</h3><br/><table class=\"wandb\"><tr><td>input to main1 std</td><td>0.14902</td></tr><tr><td>main1 to output std</td><td>0.29187</td></tr><tr><td>reward</td><td>10.0</td></tr></table>\n",
              "</div></div>\n",
              "Synced 5 W&B file(s), 0 media file(s), 0 artifact file(s) and 0 other file(s)\n",
              "<br/>Synced <strong style=\"color:#cdcd00\">tough-feather-313</strong>: <a href=\"https://wandb.ai/singularbrain/biolcnet/runs/1hasvymu\" target=\"_blank\">https://wandb.ai/singularbrain/biolcnet/runs/1hasvymu</a><br/>\n",
              "Find logs at: <code>./wandb/run-20220209_181619-1hasvymu/logs</code><br/>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "Successfully finished last run (ID:1hasvymu). Initializing new run:<br/>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "                    Syncing run <strong><a href=\"https://wandb.ai/singularbrain/biolcnet/runs/329hqirx\" target=\"_blank\">effortless-cosmos-314</a></strong> to <a href=\"https://wandb.ai/singularbrain/biolcnet\" target=\"_blank\">Weights & Biases</a> (<a href=\"https://docs.wandb.com/integrations/jupyter.html\" target=\"_blank\">docs</a>).<br/>\n",
              "\n",
              "                "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Previous model not found! Training from the beginning...\n",
            "\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "46edfe4e973f47ac8c4f6eac7ceb5b66",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "  0%|          | 0/10000 [00:00<?, ?it/s]"
            ]
          },
          "metadata": {}
        }
      ],
      "source": [
        "manual_seed(SEED)\n",
        "\n",
        "if WANDB:\n",
        "    wandb.init(project='BioLCNet', entity='singularbrain', config=network_hparams)\n",
        "net = LCNet(time,**network_hparams, reward_fn=RLTasks('CartPole-v0'), wandb_active = WANDB)\n",
        "net.learn(**network_hparams)\n",
        "plot_locally_connected_weights(net.connections[('input','main1')].w, net.n_channels1, net.filter_size1,\n",
        "                                                compute_size(net.crop_size, net.filter_size1, net.stride1), net.connections[('input','main1')].locations,\n",
        "                                                net.crop_size ** 2)\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "plot_locally_connected_weights(net.connections[('input','main1')].w, net.n_channels1, net.filter_size1,\n",
        "                                                compute_size(net.crop_size, net.filter_size1, net.stride1), net.connections[('input','main1')].locations,\n",
        "                                                net.crop_size ** 2)\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "KmwqostVRPJa"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Gym Playground"
      ],
      "metadata": {
        "id": "9M7NvCJlr0E8"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "env = gym.make('CartPole-v0')\n",
        "env.reset()\n",
        "intensity = 400\n",
        "crop_size = 90\n",
        "time = 100+300+50\n",
        "def frame_process(x):\n",
        "        x[x<0.99] = 2.0\n",
        "        x[x<=1.0] = 0.0\n",
        "        x[x==2.0] = 1.0\n",
        "        return x\n",
        "screen = env.render(mode='rgb_array')\n",
        "screen = screen.transpose((2, 0, 1))\n",
        "_, screen_height, screen_width = screen.shape\n",
        "screen = screen[:, int(screen_height*0.4):int(screen_height * 0.8)]\n",
        "view_width = int(screen_width * 0.6)\n",
        "world_width = env.x_threshold * 2\n",
        "scale = screen_width / world_width\n",
        "cart_location = int(env.state[0] * scale + screen_width / 2.0)\n",
        "if cart_location < view_width // 2:\n",
        "    slice_range = slice(view_width)\n",
        "elif cart_location > (screen_width - view_width // 2):\n",
        "    slice_range = slice(-view_width, None)\n",
        "else:\n",
        "    slice_range = slice(cart_location - view_width // 2,\n",
        "                        cart_location + view_width // 2)\n",
        "# Strip off the edges, so that we have a square image centered on a cart\n",
        "screen = screen[:, :, slice_range]\n",
        "# Convert to float, rescale, convert to torch tensor\n",
        "screen = np.ascontiguousarray(screen, dtype=np.float32) / 255\n",
        "screen = torch.from_numpy(screen)\n",
        "print(screen.shape)\n",
        "plt.imshow(screen.cpu().permute(\n",
        "        1, 2, 0).numpy().squeeze(), cmap='gray')\n",
        "h = screen.shape[1]\n",
        "w = screen.shape[2]\n",
        "print(h*0.75)\n",
        "resize = transforms.Compose([\n",
        "            transforms.ToPILImage(),\n",
        "            transforms.Resize([80, 180]),\n",
        "            transforms.Lambda(lambda x: crop(x, 0, 60, 60, 60)),\n",
        "            transforms.Lambda(lambda x: crop(x, 0, 10, 40, 40)),\n",
        "            #transforms.Resize([80, 80]),\n",
        "            #transforms.Lambda(lambda x: crop(x, 0, 0, 80, 80)),\n",
        "            # transforms.Resize([self.crop_size, self.crop_size], interpolation=Image.CUBIC),\n",
        "            #transforms.CenterCrop((crop_size, crop_size)),\n",
        "            transforms.Grayscale(),\n",
        "            transforms.ToTensor(),])\n",
        "            #transforms.Lambda(lambda x: -1.0*x +1.0),\n",
        "            #transforms.Lambda(frame_process),])\n",
        "            #transforms.Lambda(lambda x: 0*x[x<1.0]),\n",
        "            #transforms.Lambda(lambda x: x * intensity),\n",
        "            #transforms.Lambda(lambda x: PoissonEncoder(time=time, dt=1)(x))])\n",
        "device = 'cuda'\n",
        "\n",
        "# Resize, and add a batch dimension (BCHW)\n",
        "screen1 = resize(screen)\n",
        "# print(screen.shape)\n",
        "print(int(h*0.8), int(w*0.8))\n",
        "print(screen1.shape)\n",
        "# screen = screen.to(device)\n",
        "plt.imshow(screen1.cpu().permute(\n",
        "        1, 2, 0).numpy().squeeze(), cmap='gray')\n",
        "\n",
        "# print(a.shape)\n",
        "# a = screen.sum(axis=0)\n",
        "\n",
        "# a = a.to(device)\n",
        "#plt.imshow(a.cpu().numpy().squeeze(), cmap='gray')\n",
        "print(env.state)\n",
        "obs, reward, done, info = env.step(1)\n",
        "# obs, reward, done, info "
      ],
      "metadata": {
        "id": "JDFJB9GeldHl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "intensity = 400\n",
        "crop_size = 90\n",
        "time = 100+300+50\n",
        "\n",
        "screen2 = env.render(mode='rgb_array')\n",
        "screen2 = screen2.transpose((2, 0, 1))\n",
        "_, screen_height, screen_width = screen2.shape\n",
        "screen2 = screen2[:, int(screen_height*0.4):int(screen_height * 0.8)]\n",
        "view_width = int(screen_width * 0.6)\n",
        "world_width = env.x_threshold * 2\n",
        "scale = screen_width / world_width\n",
        "cart_location = int(env.state[0] * scale + screen_width / 2.0)\n",
        "if cart_location < view_width // 2:\n",
        "    slice_range = slice(view_width)\n",
        "elif cart_location > (screen_width - view_width // 2):\n",
        "    slice_range = slice(-view_width, None)\n",
        "else:\n",
        "    slice_range = slice(cart_location - view_width // 2,\n",
        "                        cart_location + view_width // 2)\n",
        "# Strip off the edges, so that we have a square image centered on a cart\n",
        "screen2 = screen2[:, :, slice_range]\n",
        "# Convert to float, rescale, convert to torch tensor\n",
        "screen2 = np.ascontiguousarray(screen2, dtype=np.float32) / 255\n",
        "screen2 = torch.from_numpy(screen2)\n",
        "#print(screen.shape)\n",
        "# plt.imshow(screen2.cpu().permute(\n",
        "#         1, 2, 0).numpy().squeeze(), cmap='gray')\n",
        "h = screen2.shape[1]\n",
        "w = screen2.shape[2]\n",
        "\n",
        "\n",
        "device = 'cuda'\n",
        "\n",
        "# Resize, and add a batch dimension (BCHW)\n",
        "screen2 = resize(screen2)\n",
        "# print(screen.shape)\n",
        "\n",
        "print(screen2.shape)\n",
        "# screen = screen.to(device)\n",
        "plt.imshow(screen2.cpu().permute(\n",
        "        1, 2, 0).numpy().squeeze(), cmap='gray')\n",
        "\n",
        "# print(a.shape)\n",
        "# a = screen.sum(axis=0)\n",
        "\n",
        "# a = a.to(device)\n",
        "#plt.imshow(a.cpu().numpy().squeeze(), cmap='gray')\n",
        "print(env.state)\n",
        "# obs, reward, done, info = env.step(0)\n",
        "# obs, reward, done, info "
      ],
      "metadata": {
        "id": "dPU8lgWZXyze"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def frame_process(x):\n",
        "        x[x<0.99] = 2.0\n",
        "        x[x<=1.0] = 0.0\n",
        "        x[x==2.0] = 1.0\n",
        "        return x"
      ],
      "metadata": {
        "id": "oTTEBtjZczbJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "diff = screen2-screen1\n",
        "normalized = (diff-torch.min(diff))/(torch.max(diff)-torch.min(diff))\n",
        "a = transforms.Compose([\n",
        "            transforms.Lambda(frame_process)])\n",
        "n = a(normalized)\n",
        "if (n==1).sum() > (n==0).sum():\n",
        "    n = a(normalized)\n",
        "f = transforms.Compose([transforms.Lambda(lambda x: x * intensity),\n",
        "            transforms.Lambda(lambda x: PoissonEncoder(time=time, dt=1)(x))])\n",
        "#n = normalized\n",
        "# plt.imshow(n.cpu().permute(\n",
        "#         1, 2, 0).numpy().squeeze(), cmap='gray')\n",
        "print(n.shape)\n",
        "c = f(n)\n",
        "print(c.shape)\n",
        "c2 = c.sum(axis=0)\n",
        "\n",
        "plt.imshow(c2.cpu().numpy().squeeze(), cmap='gray')"
      ],
      "metadata": {
        "id": "lXkLehXpZMwt"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "obs, reward, done, _ = env.step(0)   \n",
        "obs"
      ],
      "metadata": {
        "id": "9VmVGl3wPRBc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        ""
      ],
      "metadata": {
        "id": "zt5WXE2TgMzU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "env.step(1)"
      ],
      "metadata": {
        "id": "fmN6aPHyidbF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "x, x_dot, theta, theta_dot = env.state\n",
        "x, x_dot, theta, theta_dot"
      ],
      "metadata": {
        "id": "1MAMB52EiOMt"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "env.theta_threshold_radians"
      ],
      "metadata": {
        "id": "K83DiktzRLGa"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# BindsNet Breakout"
      ],
      "metadata": {
        "id": "9jFzOzYIFdGd"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "! pip install git+https://github.com/BindsNET/bindsnet.git\n"
      ],
      "metadata": {
        "id": "9Qi4v29mDF3i"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "! wget http://www.atarimania.com/roms/Roms.rar\n",
        "! mkdir /content/ROM/\n",
        "! unrar e /content/Roms.rar /content/ROM/\n",
        "! python -m atari_py.import_roms /content/ROM/"
      ],
      "metadata": {
        "id": "6ROgQ2ZWA-VT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "! pip install matplotlib==2.1.1\n",
        "\n",
        "### Restart your runtime after this"
      ],
      "metadata": {
        "id": "_K7CuZmqFTny"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "### gym and colab compatibility\n",
        "\n",
        "!apt-get install -y xvfb python-opengl > /dev/null 2>&1\n",
        "!pip install gym pyvirtualdisplay > /dev/null 2>&1\n",
        "\n",
        "import gym\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "from IPython import display as ipythondisplay\n",
        "\n",
        "from pyvirtualdisplay import Display\n",
        "display = Display(visible=0, size=(400, 300))\n",
        "display.start()\n",
        "%matplotlib inline"
      ],
      "metadata": {
        "id": "wONJ1L1uEGMA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from bindsnet.encoding import bernoulli\n",
        "from bindsnet.environment import GymEnvironment\n",
        "from bindsnet.learning import MSTDP\n",
        "from bindsnet.network import Network\n",
        "from bindsnet.network.nodes import Input, LIFNodes\n",
        "from bindsnet.network.topology import Connection\n",
        "from bindsnet.pipeline import EnvironmentPipeline\n",
        "from bindsnet.pipeline.action import select_softmax\n",
        "import time\n",
        "\n",
        "# Build network.\n",
        "network = Network(dt=1.0)\n",
        "\n",
        "# Layers of neurons.\n",
        "inpt = Input(n=80 * 80, shape=[1, 1, 1, 80, 80], traces=True)\n",
        "middle = LIFNodes(n=100, traces=True)\n",
        "out = LIFNodes(n=4, refrac=0, traces=True)\n",
        "\n",
        "# Connections between layers.\n",
        "inpt_middle = Connection(source=inpt, target=middle, wmin=0, wmax=1e-1)\n",
        "middle_out = Connection(\n",
        "    source=middle,\n",
        "    target=out,\n",
        "    wmin=0,\n",
        "    wmax=1,\n",
        "    update_rule=MSTDP,\n",
        "    nu=1e-1,\n",
        "    norm=0.5 * middle.n,\n",
        ")\n",
        "\n",
        "# Add all layers and connections to the network.\n",
        "network.add_layer(inpt, name=\"Input Layer\")\n",
        "network.add_layer(middle, name=\"Hidden Layer\")\n",
        "network.add_layer(out, name=\"Output Layer\")\n",
        "network.add_connection(inpt_middle, source=\"Input Layer\", target=\"Hidden Layer\")\n",
        "network.add_connection(middle_out, source=\"Hidden Layer\", target=\"Output Layer\")\n",
        "\n",
        "# Load the Breakout environment.\n",
        "environment = GymEnvironment(\"BreakoutDeterministic-v4\")\n",
        "environment.reset()\n",
        "\n",
        "# Build pipeline from specified components.\n",
        "environment_pipeline = EnvironmentPipeline(\n",
        "    network,\n",
        "    environment,\n",
        "    encoding=bernoulli,\n",
        "    action_function=select_softmax,\n",
        "    output=\"Output Layer\",\n",
        "    time=100,\n",
        "    history_length=1,\n",
        "    delta=1,\n",
        "    plot_interval=None,\n",
        "    render_interval=None,\n",
        ")\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "def run_pipeline(pipeline, episode_count):\n",
        "    for i in range(episode_count):\n",
        "        total_reward = 0\n",
        "        pipeline.reset_state_variables()\n",
        "        is_done = False\n",
        "        render_counter = 0\n",
        "        while not is_done:\n",
        "            # if render_counter % 10 == 0:\n",
        "            #     # plt.title(\"Game image\")\n",
        "            #     # plt.imshow(pipeline.env.render())\n",
        "            #     # plt.show()\n",
        "            #     a = pipeline.env.render()\n",
        "            #     print(a)\n",
        "            #     time.sleep(0.1)\n",
        "                \n",
        "            render_counter += 1\n",
        "            result = pipeline.env_step()\n",
        "            pipeline.step(result)\n",
        "\n",
        "            reward = result[1]\n",
        "            total_reward += reward\n",
        "            \n",
        "            is_done = result[2]\n",
        "        print(f\"Episode {i} total reward:{total_reward}\")\n",
        "    pipeline.env.close()\n",
        "\n",
        "print(\"Training: \")\n",
        "run_pipeline(environment_pipeline, episode_count=100)\n",
        "\n",
        "# stop MSTDP\n",
        "environment_pipeline.network.learning = False\n",
        "\n",
        "print(\"Testing: \")\n",
        "run_pipeline(environment_pipeline, episode_count=100)"
      ],
      "metadata": {
        "id": "y_jx13I5-VTL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        ""
      ],
      "metadata": {
        "id": "3PLVXCK0Djg2"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "7fTSvrK3T_GA",
        "ULGGHW43UksI",
        "MBKedMpIleMr",
        "8clxN_npa1WY",
        "sCXSLZZoGS4z"
      ],
      "machine_shape": "hm",
      "name": "BioLCNet_CartPole_gradient.ipynb",
      "provenance": [],
      "include_colab_link": true
    },
    "interpreter": {
      "hash": "87ae7d1e75b14a98f2d7b99b6b39b40721989d38e6517f9dbec64ca4d8e3011b"
    },
    "kernelspec": {
      "display_name": "Python 3.7.9 64-bit ('base': conda)",
      "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.8.5"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "7f5866e1fea24fec903f0e7a0035bb88": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "VBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "VBoxView",
            "_dom_classes": [],
            "_model_name": "VBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_b4db786096874228b077abbbc975272d",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_bc5befd433d648569c1c11b6a66ad4ce",
              "IPY_MODEL_63ca4af57e0f4938a56b6c3f9577f081"
            ]
          }
        },
        "b4db786096874228b077abbbc975272d": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "bc5befd433d648569c1c11b6a66ad4ce": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "LabelModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "LabelView",
            "style": "IPY_MODEL_f2a003372455459ba6e9b73a9b43b459",
            "_dom_classes": [],
            "description": "",
            "_model_name": "LabelModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 0.08MB of 0.08MB uploaded (0.00MB deduped)\r",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_f8147581481941be8d286bb362b83a0d"
          }
        },
        "63ca4af57e0f4938a56b6c3f9577f081": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_a08d221a81264e768308fa420510e500",
            "_dom_classes": [],
            "description": "",
            "_model_name": "FloatProgressModel",
            "bar_style": "",
            "max": 1,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 1,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_5b33649bfe4e4154a08e7029a82424dc"
          }
        },
        "f2a003372455459ba6e9b73a9b43b459": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "f8147581481941be8d286bb362b83a0d": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "a08d221a81264e768308fa420510e500": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "5b33649bfe4e4154a08e7029a82424dc": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "46edfe4e973f47ac8c4f6eac7ceb5b66": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_b20ee9a88fc54a7297584207ab7bd0e3",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_2953a20627914e08a3a6ce816bbf8126",
              "IPY_MODEL_952bb30941b44297bbc2dbe06f28214b",
              "IPY_MODEL_aa4a131e00d840138fc297ee81ce5c7c"
            ]
          }
        },
        "b20ee9a88fc54a7297584207ab7bd0e3": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "2953a20627914e08a3a6ce816bbf8126": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_65ee162348d048899e3186a2362e42e8",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": "  0%",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_cef349c76d7348daa4fdab2489b00158"
          }
        },
        "952bb30941b44297bbc2dbe06f28214b": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_cab39115348645abaf478f8cf7290025",
            "_dom_classes": [],
            "description": "",
            "_model_name": "FloatProgressModel",
            "bar_style": "",
            "max": 10000,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 0,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_d2107e81c1604285a695b0f56ca7500a"
          }
        },
        "aa4a131e00d840138fc297ee81ce5c7c": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_3df4dc6320c8456ab6319eca97727bb4",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 0/10000 [00:00&lt;?, ?it/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_3ac0cd6b794747d7a14cc2c4825870ac"
          }
        },
        "65ee162348d048899e3186a2362e42e8": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "cef349c76d7348daa4fdab2489b00158": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "cab39115348645abaf478f8cf7290025": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "d2107e81c1604285a695b0f56ca7500a": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "3df4dc6320c8456ab6319eca97727bb4": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "3ac0cd6b794747d7a14cc2c4825870ac": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        }
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}