{ "cells": [ { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/RatInABox-Lab/RatInABox/blob/main/demos/path_integration_example.ipynb)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Case study: Path Integration\n", "### We use RatInABox to built a network capable of learning to path integrate in 1D using local learning rules" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "This work is closely related to the work done by [Vafidis et al. (2021)](https://www.biorxiv.org/content/10.1101/2021.03.12.435035v1). Although the specific details of our learning rules, parameter settings and architecture are all slightly different we do not claim this work to be novel. Rather include it as a demonstraion of what is possible (even \"easy\") to study using `RatInABox`.\n", "\n", "To avoid rewritting their paper we give only sparse, key details here. If you intend to make perfect sense of , or even build upon this example, we recommend reading the \n", "Vafidis et al. (2021) first, and (optionally but highly recommended too) McNaughton et al. (2006). If you don't intend to do this, this section, and the code below, can still be read, run and referenced as a tutorial for `RatInABox`. \n", "\n", "For now we will briefly describe the problem of path integration and our approach for how to solve it in 1D using `RatInABox`.\n", "\n", "### Path integration with ring attractors and velocity inputs\n", "Animals need to maintain an encoding of their state in the world even in the absence of sensory input. For one dimensional variables (such as the animals head direction, or location along a linear track) it is often assumed this is done using a ring-attractor network. Ring attractors are a type of continuous attractor neural network (CANN). \n", "\n", "In a ring attractor network, the neurons are arranged in a 1D ring. A bump of activity is stabilised by local-excitatory and long-range inhibitatory weights between neighbouring cells (sometime long range inhibition is replaced with global inhibition as in Vafidis et al.). \n", "\n", "Inputs from cells encoding the animals velocity can \"push\" the bump activity around the ring in correspondence with the position of the animal. For each ring attractor neuron, suppose there exists a left and right \"conjunctive velocity neuron\" (conjuctive because the velocity neuron being active is contingent on its corresponding bump attractor neuron being active, this is a small detail your intuitive understanding can probably afford to ignore) which connects assymetrically to the bump attractor neurons. Specifically the left conjuctive velocity cell corresponding to ring attractor cell $i$ might connect in an excitatory manner to cells to the left of cell $i$ and in an inhibitory manner to cells to the right. This way when the animal moves left the velocity neuron becomes active and excites neurons on the left of the bump and inhibits thos on the right, thus pushes the bump leftwards. \n", "\n", "A lot of work in bump attractors and path integration relates to maintaining knowledge of head direction (a ring manifold of neurons, manifesting topologically as a ring within the brain, has been found in drosophila and linked to path integration of head direction (Kim et al. 2017)). However it is not necesasry for the neurons arrangement in physical space to correspond with the manifold (e.g a 1D ring) their activity lies on. This is the running hyptohesis for the role of entorhinal cortex as a neural substrate for path integration and the \"cognitive map\". We recommend McNaughton (2006) for a good summary of path integration as applied to hippocampus/entorhinal cortex. Specifically the activity profile of grid cells is though to lie on a torodial (torus = 2D equivalent of a 1D ring) manifold. Everything about this system seems pretty well set up to perform path integration, there are even velocity-conjuctive grod cells which would do the pushing of the bump. \n", "\n", "\n", "\n", " ### Our approach\n", "\n", " The key problem identified by Vafidis et al. is that the weights (local excitation and long range inhibition for the recurrent connection of the ring-attractor and the asymmetric connectivity profile from the conjunctive velocity cells) need to be fine tuned. It is unliekly these connectivety profiles are hard coded by evolution and, for reasons of flexibility, it would be better (and more flexible) if instead they could be *learned* and better still if the learning rule can be local and Hebbian. \n", "\n", " To do this we define a classes of `Neurons` called `PyramidalNeurons`. Each pyramidal neuron has three compartments (two dendritic and one somatic). The first (basal) dendritic compartment recieves bottom-up sensory input from a set of `PlaceCells` from which they are driven in a one-to-one manner. This will act as the target signal, i.e. a localised bump of activity correponding to the nearby place cells currently active.. A second compartment recieves top-down input from the conjunctive velocity cells and recurrent inputs. Which dendritic compartment drives the soma can be adaptively controlled by a parameter we call `theta`. \n", " \n", " The local hebbian learning rule we used is based on a well studied class of neuronal learning rules where the dendritic compartments attempts to adjust their weights to match the somatic compartment firing rate. This has been term \"dendritic prediction of somatic activity\" and, if the dendritc voltage function is linear, looks like a voltage thresholded Hebbian learning rule (see Urbanczik and Senn (2014)).\n", "\n", " The necessary equations are as follows (note a more biophysically details version of our slimmed down learning rule is implememnted in Vafidis etal. so please see tehre for more details). \n", "\n", " $$\\mathsf{V}^{\\textrm{basal}} = \\sum_{i}\\mathsf{w}^{\\textrm{basal}}_{i}\\mathsf{I}_{i}^{\\textrm{basal}}$$\n", "\n", " $$\\mathsf{V}^{\\textrm{apical}} = \\sum_{i}\\mathsf{w}^{\\textrm{apical}}_{i}\\mathsf{I}_{i}^{\\textrm{apical}}$$\n", "\n", " where $\\mathsf{I}_{i}$ is the unitless input from the layer feeding into theis dendritic compartment (just the firing rate of that layer). ${I}_{i}^{\\textrm{basal}} = \\{ {I}_{i}^{\\textrm{place cells}} \\}$ and ${I}_{i}^{\\textrm{apical}} = \\{ {I}_{i}^{\\textrm{ring attractor}}, {I}_{i}^{\\textrm{conjuctive velocity cells}} \\}$. \n", "\n", " $$\\frac{d\\mathsf{w_{i}^{\\textrm{apical}}}(t)}{dt} \\propto \\big( \\mathsf{V}^{\\textrm{soma}} - \\mathsf{V}^{\\textrm{basal}}\\big)\\mathsf{I}_{i}^{\\textrm{basal}} $$\n", "\n", " where the somatic voltage is given simply by \n", "\n", " $$\\mathsf{V}^{\\textrm{soma}} = (1-\\theta(t))\\mathsf{V}^{\\textrm{basal}} + \\theta(t)\\mathsf{V}^{\\textrm{apical}}$$\n", "\n", " Typically we set theta to be oscillatory. The soma rapidly oscillates between being driven by the top-down inputs from its apical dendritic compartment and bottom-up inputs from its basal compartment. \n", "\n", " \n", "\n", "Green weights are plastic, trained via the local Hebbian learning rule, while blue weights are fixed to sensible values. \n", "\n", "Note the agent, the environment, and three of the four cell layers are all taken directly from `RatInABox`. Only one bespoke `Neurons` cell class need to be defined (this is the `PyramidalNeuron` layer which will become the ring attractor. Infact we chooose to define the dendritic compartments of each pyramidal cell as a `Neurons` classes too)\n", "\n", "Other relevant details include: \n", "\n", "* The conjuctive velocity neurons each recieve an input from one and only one ring attractor cell as well as a velocity neuron (left or right, depending on which population it resides in). This fixed weight from the ring attractor input is tuned such that when the ring attractor neuron fires fully (i.e. the bump is centred on its part of the ring) the activity in the conjunctive cell reaches threshold. Any additional input from the velocity cell now causes activity on this neuron. This means \n", "* The input to the conjuctive velocity cells from the velocity neurons is scaled such that the full sigmoid activity range is covered over the first two standard deviations of the animals likely speed. " ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "!pip install ratinabox\n", "!pip install tqdm " ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Import ratinabox\n", "import ratinabox\n", "from ratinabox.Environment import Environment\n", "from ratinabox.Agent import Agent\n", "from ratinabox.Neurons import Neurons, PlaceCells, VelocityCells, FeedForwardLayer\n", "from ratinabox import utils\n", "\n", "#stylise and automatically save ratinabox plots/animations\n", "ratinabox.stylize_plots(); ratinabox.autosave_plots = True; ratinabox.figure_directory = \"../figures/\"\n", "\n", "import numpy as np\n", "import matplotlib\n", "import matplotlib.pyplot as plt\n", "\n", "%load_ext autoreload\n", "%autoreload 2\n", "\n", "from tqdm import tqdm" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Defining the `PyramidalNeurons` class " ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "class PyramidalNeurons(Neurons):\n", " \"\"\"The PyramidalNeuorn class defines a layer of Neurons() whos firing rates are derived from the firing rates in two DendriticCompartments. They are theta modulated, during early theta phase the apical DendriticCompartment (self.apical_compartment) drives the soma, during late theta phases the basal DendriticCompartment (self.basal_compartment) drives the soma.\n", "\n", " Must be initialised with an Agent and a 'params' dictionary.\n", "\n", " Check that the input layers are all named differently.\n", " List of functions:\n", " • get_state()\n", " • update()\n", " • update_dendritic_compartments()\n", " • update_weights()\n", " • plot_loss()\n", " • plot_rate_map()\n", " \"\"\"\n", "\n", " def __init__(self, Agent, params={}):\n", " \"\"\"Initialises a layer of pyramidal neurons\n", "\n", " Args:\n", " Agent (_type_): _description_\n", " params (dict, optional): _description_. Defaults to {}.\n", " \"\"\"\n", " default_params = {\n", " \"n\": 10,\n", " \"name\": \"PyramidalNeurons\",\n", " # theta params\n", " \"theta_freq\": 5,\n", " \"theta_frac\": 0.5, # -->0 all basal input, -->1 all apical input\n", " }\n", " default_params.update(params)\n", " self.params = default_params\n", " super().__init__(Agent, self.params)\n", "\n", " self.history[\"loss\"] = []\n", " self.error = None\n", "\n", " self.basal_compartment = DendriticCompartment(\n", " self.Agent,\n", " params={\n", " \"soma\": self,\n", " \"name\": f\"{self.name}_basal\",\n", " \"n\": self.n,\n", " \"color\": self.color,\n", " },\n", " )\n", " self.apical_compartment = DendriticCompartment(\n", " self.Agent,\n", " params={\n", " \"soma\": self,\n", " \"name\": f\"{self.name}_apical\",\n", " \"n\": self.n,\n", " \"color\": self.color,\n", " },\n", " )\n", "\n", " def update(self):\n", " \"\"\"Updates the firing rate of the layer. Saves a loss (lpf difference between basal and apical). Also adds noise.\"\"\"\n", " super().update() # this sets and saves self.firingrate\n", "\n", " dt = self.Agent.dt\n", " tau_smooth = 10\n", " # update a smoothed history of the loss\n", " fr_b, fr_a = (\n", " self.basal_compartment.firingrate,\n", " self.apical_compartment.firingrate,\n", " )\n", " error = np.mean(np.abs(fr_b - fr_a))\n", " if self.Agent.t < 2 / self.theta_freq:\n", " self.error = None\n", " else:\n", " # loss_smoothing_timescale = dt\n", " self.error = (dt / tau_smooth) * error + (1 - dt / tau_smooth) * (\n", " self.error or error\n", " )\n", " self.history[\"loss\"].append(self.error)\n", " return\n", "\n", " def update_dendritic_compartments(self):\n", " \"\"\"Individually updates teh basal and apical firing rates.\"\"\"\n", " self.basal_compartment.update()\n", " self.apical_compartment.update()\n", " return\n", "\n", " def get_state(self, evaluate_at=\"last\", **kwargs):\n", " \"\"\"Returns the firing rate of the soma. This depends on the firing rates of the basal and apical compartments and the current theta phase. By default the theta is obtained from self.Agent.t but it can be passed manually as an kwarg to override this.\n", "\n", " theta (or theta_gating) is a number between [0,1] controlling flow of information into soma from the two compartment.s 0 = entirely basal. 1 = entirely apical. Between equals weighted combination. he function theta_gating() takes a time and returns theta.\n", " Args:\n", " evaluate_at (str, optional): 'last','agent','all' or None (in which case pos can be passed directly as a kwarg). Defaults to \"last\".\n", " Returns:\n", " firingrate\n", " \"\"\"\n", " # theta can be passed in manually as a kwarg. If it isn't ithe time from the agent will be used to get theta. Theta determines how much basal and how much apical this neurons uses.\n", " if \"theta\" in kwargs:\n", " theta = kwargs[\"theta\"]\n", " else:\n", " theta = theta_gating(\n", " t=self.Agent.t, freq=self.theta_freq, frac=self.theta_frac\n", " )\n", " fr_basal, fr_apical = 0, 0\n", " # these are special cases, no need to even get their fr's if they aren't used\n", " if theta != 0:\n", " fr_apical = self.apical_compartment.get_state(evaluate_at, **kwargs)\n", " if theta != 1:\n", " fr_basal = self.basal_compartment.get_state(evaluate_at, **kwargs)\n", " firingrate = (1 - theta) * fr_basal + (theta) * fr_apical\n", " return firingrate\n", "\n", " def update_weights(self):\n", " \"\"\"Trains the weights, this function actually defined in the dendrite class.\"\"\"\n", " if self.Agent.t > 2 / self.theta_freq:\n", " self.basal_compartment.update_weights()\n", " self.apical_compartment.update_weights()\n", " return\n", "\n", " def plot_loss(self, fig=None, ax=None):\n", " \"\"\"Plots the loss against time to see if learning working\"\"\"\n", " if fig is None and ax is None:\n", " fig, ax = plt.subplots(figsize=(1.5, 1.5))\n", " ylim = 0\n", " else:\n", " ylim = ax.get_ylim()[1]\n", " t = np.array(self.history[\"t\"]) / 60\n", " loss = self.history[\"loss\"]\n", " ax.plot(t, loss, color=self.color, label=self.name)\n", " ax.set_ylim(\n", " bottom=0, top=max(ylim, np.nanmax(np.array(loss, dtype=np.float64)))\n", " )\n", " ax.set_xlim(left=0)\n", " ax.legend(frameon=False)\n", " ax.set_xlabel(\"Training time / min\")\n", " ax.set_ylabel(\"Loss\")\n", " return fig, ax\n", "\n", " def plot_rate_map(self, route=\"basal\", **kwargs):\n", " \"\"\"This is a wrapper function for the general Neuron class function plot_rate_map. It takes the same arguments as Neurons.plot_rate_map() but, in addition, route can be set to basal or apical in which case theta is set correspondingly and teh soma with take its input from downstream or upstream sources entirely.\n", "\n", " The arguments for the standard plottiong function plot_rate_map() can be passed as usual as kwargs.\n", "\n", " Args:\n", " route (str, optional): _description_. Defaults to 'basal'.\n", " \"\"\"\n", " if route == \"basal\":\n", " theta = 0\n", " elif route == \"apical\":\n", " theta = 1\n", " fig, ax = super().plot_rate_map(**kwargs, theta=theta)\n", " return fig, ax\n", "\n", "\n", "class DendriticCompartment(Neurons):\n", " \"\"\"The DendriticCompartment class defines a layer of Neurons() whos firing rates are an activated linear combination of input layers. This class is a subclass of Neurons() and inherits it properties/plotting functions.\n", "\n", " Must be initialised with an Agent and a 'params' dictionary.\n", " Input params dictionary must contain a list of input_layers which feed into these Neurons. This list looks like [Neurons1, Neurons2,...] where each is a Neurons() class.\n", "\n", " Currently supported activations include 'sigmoid' (paramterised by max_fr, min_fr, mid_x, width), 'relu' (gain, threshold) and 'linear' specified with the \"activation_params\" dictionary in the inout params dictionary. See also activate() for full details.\n", "\n", " Check that the input layers are all named differently.\n", " List of functions:\n", " • get_state()\n", " • add_input()\n", " \"\"\"\n", "\n", " def __init__(self, Agent, params={}):\n", " default_params = {\n", " \"soma\": None,\n", " \"activation_params\": {\n", " \"activation\": \"sigmoid\",\n", " \"max_fr\": 1,\n", " \"min_fr\": 0,\n", " \"mid_x\": 1,\n", " \"width_x\": 2,\n", " },\n", " }\n", " self.Agent = Agent\n", " default_params.update(params)\n", " self.params = default_params\n", " super().__init__(Agent, self.params)\n", "\n", " self.firingrate_temp = None\n", " self.firingrate_prime_temp = None\n", " self.inputs = {}\n", "\n", " def add_input(\n", " self, input_layer, eta=0.001, w_init=0.1, L1=0.0001, L2=0.001, tau_PI=100e-3\n", " ):\n", " \"\"\"Adds an input layer to the class. Each input layer is stored in a dictionary of self.inputs. Each has an associated matrix of weights which are initialised randomly.\n", "\n", " Args:\n", " input_layer (_type_): the layer which feeds into this compartment\n", " eta: learning rate of the weights\n", " w_init: initialisation scale of the weights\n", " L1: how much L1 regularisation\n", " L2: how much L2 regularisation\n", " tau_PI: smoothing timescale of plasticity induction variable\n", " \"\"\"\n", " name = input_layer.name\n", " n_in = input_layer.n\n", " w = np.random.normal(loc=0, scale=w_init / np.sqrt(n_in), size=(self.n, n_in))\n", " I = np.zeros(n_in)\n", " PI = np.zeros(n_in)\n", " if name in self.inputs.keys():\n", " print(\n", " f\"There already exists a layer called {input_layer_name}. Overwriting it now.\"\n", " )\n", " self.inputs[name] = {}\n", " self.inputs[name][\"layer\"] = input_layer\n", " self.inputs[name][\"w\"] = w\n", " self.inputs[name][\"w_init\"] = w.copy()\n", " self.inputs[name][\"I\"] = I # input current\n", " self.inputs[name][\"I_temp\"] = None # input current\n", " self.inputs[name][\"PI\"] = PI # plasticity induction variable\n", " self.inputs[name][\"eta\"] = eta\n", " self.inputs[name][\"L2\"] = L2\n", " self.inputs[name][\"L1\"] = L1\n", " self.inputs[name][\"tau_PI\"] = tau_PI\n", "\n", " def get_state(self, evaluate_at=\"last\", **kwargs):\n", " \"\"\"Returns the \"firing rate\" of the dendritic compartment. By default this layer uses the last saved firingrate from its input layers. Alternatively evaluate_at and kwargs can be set to be anything else which will just be passed to the input layer for evaluation.\n", " Once the firing rate of the inout layers is established these are multiplied by the weight matrices and then activated to obtain the firing rate of this FeedForwardLayer.\n", "\n", " Args:\n", " evaluate_at (str, optional). Defaults to 'last'.\n", " Returns:\n", " firingrate: array of firing rates\n", " \"\"\"\n", " if evaluate_at == \"last\":\n", " V = np.zeros(self.n)\n", " elif evaluate_at == \"all\":\n", " V = np.zeros(\n", " (self.n, self.Agent.Environment.flattened_discrete_coords.shape[0])\n", " )\n", " else:\n", " V = np.zeros((self.n, kwargs[\"pos\"].shape[0]))\n", "\n", " for inputlayer in self.inputs.values():\n", " w = inputlayer[\"w\"]\n", " if evaluate_at == \"last\":\n", " I = inputlayer[\"layer\"].firingrate\n", " else: # kick can down the road let input layer decide how to evaluate the firingrate\n", " I = inputlayer[\"layer\"].get_state(evaluate_at, **kwargs)\n", " inputlayer[\"I_temp\"] = I\n", " V += np.matmul(w, I)\n", " firingrate = utils.activate(V, other_args=self.activation_params)\n", " firingrate_prime = utils.activate(\n", " V, other_args=self.activation_params, deriv=True\n", " )\n", "\n", " self.firingrate_temp = firingrate\n", " self.firingrate_prime_temp = firingrate_prime\n", "\n", " return firingrate\n", "\n", " def update(self):\n", " \"\"\"Updates firingrate of this compartment and saves it to file\"\"\"\n", " self.get_state()\n", " self.firingrate = self.firingrate_temp.reshape(-1)\n", " self.firingrate_deriv = self.firingrate_prime_temp.reshape(-1)\n", " for inputlayer in self.inputs.values():\n", " inputlayer[\"I\"] = inputlayer[\"I_temp\"].reshape(-1)\n", " self.save_to_history()\n", " return\n", "\n", " def update_weights(self):\n", " \"\"\"Implements the weight update: dendritic prediction of somatic activity.\"\"\"\n", " target = self.soma.firingrate\n", " delta = (target - self.firingrate) * (self.firingrate_deriv)\n", " dt = self.Agent.dt\n", " for inputlayer in self.inputs.values():\n", " eta = inputlayer[\"eta\"]\n", " if eta != 0:\n", " tau_PI = inputlayer[\"tau_PI\"]\n", " assert (dt / tau_PI) < 0.2\n", " I = inputlayer[\"I\"]\n", " w = inputlayer[\"w\"]\n", " # first updates plasticity induction variable (smoothed delta error outer product with the input current for this input layer)\n", " PI_old = inputlayer[\"PI\"]\n", " PI_update = np.outer(delta, I)\n", " PI_new = (dt / tau_PI) * PI_update + (1 - dt / tau_PI) * PI_old\n", " inputlayer[\"PI\"] = PI_new\n", " # updates weights\n", " dw = eta * (\n", " PI_new - inputlayer[\"L2\"] * w - inputlayer[\"L1\"] * np.sign(w)\n", " )\n", " inputlayer[\"w\"] = w + dw\n", " return\n", "\n", "\n", "def theta_gating(t, freq=10, frac=0.5):\n", " T = 1 / freq\n", " phase = ((t / T) % 1) % 1\n", " if phase < frac:\n", " return 1\n", " elif phase >= frac:\n", " return 0" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Building the network\n", "This involves initialising the environment, the agent and all of the cells layers shown in the above figure. Then we must set the inputs to each cell layer) and, for those layers where the weights are fixed, set them correctly. \n", "\n", "Basically, the following code amounts to constructing the network shown in the above figure. " ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Initialise the 1D environment\n", "Env = Environment(params={\"dimensionality\": \"1D\", \"boundary_conditions\": \"periodic\"})\n", "\n", "# Put agent (who will move randomly under the ratinabox Ornstein Uhlenbeck random motion policy) inside the environement\n", "Ag = Agent(Env, params={'dt':0.02})\n", "Ag.speed_mean = 0\n", "Ag.speed_std = 0.3\n", "\n", "n_cells = 50\n", "# Place cells provide the target signal\n", "PlaceCells_ = PlaceCells(Ag, params={\"n\": n_cells, \"widths\": 0.1, \"name\": \"PlaceCells\"})\n", "\n", "# The key neuron class: Ring attractor at the centre of the network made from our bespoke, custom-define PyramidalNeurons class.\n", "RingAttractor = PyramidalNeurons(Ag, params={\"n\": n_cells, \"name\": \"RingAttractor\"})\n", "\n", "# Velocity cells encode agent velocity\n", "VelocityCells_ = VelocityCells(Ag, params={\"name\": \"VelocityCells\"})\n", "\n", "# Conjuctive cells\n", "ConjunctiveCells_left = FeedForwardLayer(\n", " Ag,\n", " params={\n", " \"n\": n_cells,\n", " \"name\": \"ConjunctiveCells_left\",\n", " # nb. this tutorial is now quite old so the way that FeedForwardLayer --- define in the main codebase --- activations are set (passing \"activation_function\" at initialisation) no longer matches the way DendriticCompartment --- defined above --- activations are set (setting \"activation_params\" after initialisation). Sorry about this! TODO: update DendriticCompartment to be FeedForwardLayer subclass\n", " \"activation_function\": {\n", " \"activation\": \"relu\",\n", " \"threshold\": 1,\n", " }\n", " },\n", ")\n", "\n", "ConjunctiveCells_right = FeedForwardLayer(\n", " Ag,\n", " params={\n", " \"n\": n_cells,\n", " \"name\": \"ConjunctiveCells_right\",\n", " \"activation_function\": {\n", " \"activation\": \"relu\",\n", " \"threshold\": 1,\n", " }\n", " },\n", ")\n", "\n", "# Set inputs into ring attractor compartments\n", "# Make their activation functions linear\n", "# Set the fixed weights from place celles to Ring attractor to be fixed\n", "RingAttractor.apical_compartment.add_input(RingAttractor)\n", "RingAttractor.apical_compartment.add_input(ConjunctiveCells_left)\n", "RingAttractor.apical_compartment.add_input(ConjunctiveCells_right)\n", "RingAttractor.apical_compartment.activation_params = {\"activation\": \"linear\"}\n", "\n", "RingAttractor.basal_compartment.add_input(PlaceCells_, eta=0) # eta=0, these are fixed\n", "RingAttractor.basal_compartment.inputs[\"PlaceCells\"][\"w\"] = np.identity(n_cells)\n", "RingAttractor.basal_compartment.activation_params = {\"activation\": \"linear\"}\n", "\n", "# Set inputs into the conjuctive cells\n", "# Set the (fixed) weights into the conjunctive cells to be their correct values (identity or just 1's)\n", "ConjunctiveCells_left.add_input(VelocityCells_)\n", "ConjunctiveCells_left.add_input(RingAttractor)\n", "ConjunctiveCells_right.add_input(VelocityCells_)\n", "ConjunctiveCells_right.add_input(RingAttractor)\n", "ConjunctiveCells_left.inputs[\"VelocityCells\"][\"w\"] = np.ones((n_cells, 2)) * np.array(\n", " [1, -1]\n", ") # thus left velocity excites these cells and right velocity shuts them off\n", "ConjunctiveCells_right.inputs[\"VelocityCells\"][\"w\"] = np.ones((n_cells, 2)) * np.array(\n", " [-1, 1]\n", ") # thus right velocity excites these cells and rigleftht velocity shuts them off\n", "ConjunctiveCells_left.inputs[\"RingAttractor\"][\"w\"] = np.identity(n_cells)\n", "ConjunctiveCells_right.inputs[\"RingAttractor\"][\"w\"] = np.identity(n_cells)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Train the network\n", "\n", "Train it for 60 minutes" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "for i in tqdm(range(int(60 * 60 / Ag.dt))):\n", " # update agent\n", " Ag.update()\n", " # update firing rates of all the cell layers\n", " PlaceCells_.update()\n", " VelocityCells_.update()\n", " ConjunctiveCells_left.update()\n", " ConjunctiveCells_right.update()\n", " RingAttractor.update_dendritic_compartments()\n", " RingAttractor.update()\n", " # finally, update the weights\n", " RingAttractor.update_weights()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Analysis \n", "Now training is complete we can see whether it as been successful. \n", "First we plot the loss. This is a (smoothed) trace of the mean absolute difference between the voltage in the basal compartment (the ground truth) and the voltage in the apical compartment (the compartment learnign to replicate the ground truth using recurrent and conjunctive velocity inputs). You can think of this as the prediction error. " ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "fig, ax = RingAttractor.plot_loss()\n", "\n", "save_plots = False\n", "if save_plots == True:\n", " tpl.saveFigure(fig, \"PI_loss\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# pull out the weight as they were at initialisation\n", "w_ccl_init = RingAttractor.apical_compartment.inputs[\"ConjunctiveCells_left\"][\"w_init\"]\n", "w_ccr_init = RingAttractor.apical_compartment.inputs[\"ConjunctiveCells_right\"][\"w_init\"]\n", "w_rec_init = RingAttractor.apical_compartment.inputs[\"RingAttractor\"][\"w_init\"]\n", "\n", "# pull out the weights after training\n", "w_ccl = RingAttractor.apical_compartment.inputs[\"ConjunctiveCells_left\"][\"w\"]\n", "w_ccr = RingAttractor.apical_compartment.inputs[\"ConjunctiveCells_right\"][\"w\"]\n", "w_rec = RingAttractor.apical_compartment.inputs[\"RingAttractor\"][\"w\"]\n", "\n", "# plot them\n", "fig, ax = plt.subplots(1, 3, figsize=(12, 4))\n", "ax[0].imshow(w_rec_init)\n", "ax[1].imshow(w_ccl_init)\n", "ax[2].imshow(w_ccr_init)\n", "ax[0].set_ylabel(\"BEFORE TRAINING\")\n", "ax[0].set_title(\"Ring attractor \\nto ring attractor\")\n", "ax[1].set_title(\"Left conjunctive velocity cells \\nto ring attractor\")\n", "ax[2].set_title(\"Right conjunctive velocity cells \\nto ring attractor\")\n", "\n", "fig1, ax1 = plt.subplots(1, 3, figsize=(12, 4))\n", "ax1[0].imshow(w_rec)\n", "ax1[1].imshow(w_ccl)\n", "ax1[2].imshow(w_ccr)\n", "ax1[0].set_ylabel(\"AFTER TRAINING\")\n", "ax1[0].set_title(\"Ring attractor \\nto ring attractor\")\n", "ax1[1].set_title(\"Left conjunctive velocity cells \\nto ring attractor\")\n", "ax1[2].set_title(\"Right conjunctive velocity cells \\nto ring attractor\")\n", "\n", "if save_plots == True:\n", " tpl.saveFigure(fig, \"PIweights_beforelearning\")\n", " tpl.saveFigure(fig1, \"PIweights_afterlearning\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Further analysis\n", "\n", "In a seperate analysis (of the same network, but code not given here) we produced the following figures. \n", "\n", "The first shows the weight matrices before and after learning, along side a row-averaged cross section of the after-learning matrices. One can see how the theoritical hypothsis for how weights should develop in this system is correctly found by the Hebbian local learning rule (in less than 10 minutes of training!). \n", "\n", "In the movie we summarise an experiment where the animal is left to explore for 60 seconds. For the first 60 seconds sensory *and* top-down inouts are available. For times 30s to 60s the sensory input from the place cells are lesioned and only the velocity integration part (the top-down apical dendritic compartment) drives the soma. As can be seen, although errors accumulate, the network does a decent job at maintaining the position estimate in the total absence of sensory input. Purple shows true position of agent, blue shows position in neural manifold of peak activity. We encourage anyone intersted to replicate this very simple experiment using the above code - perhaps using the decoding_position)_example.ipynb script to train a decoder to decode position fromneural activity and show the decoding is accurate even after sensory input is lesioned.\n", "\n", " * Proof the connectivity structure matches that of a classically define veloicty-driven path integrator CANN.\n", " \n", " \n", " \n", "* Proof this connectivity structure can actually perform path integration\n", "\n", " \n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Disclaimer: \n", "Please note this script is intended as a tutorial, here only to demonstrate how `RatInABox` could be used for a path integration project in continuous time and space and not as a scientific result regarding the role of hippocampus and entorhinal cortex role in path integration. " ] } ], "metadata": { "kernelspec": { "display_name": "riab", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.10.9" }, "orig_nbformat": 4 }, "nbformat": 4, "nbformat_minor": 2 }