![](\"../.images/demos/rl_agent.gif\")
"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Policy Evaluation"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"There are two key components to most model-free RL algorithms. Policy *evaluation* evaluates a policy to learn its \"value function\". Policy *improvement* uses this value function to design a better policy. Iterating between these two often results in convergence towards a near-optimal policy. \n",
"\n",
"To help with these process we define a `ValueNeuron` class. The `ValueNeuron`, which is a subclass of `FeedForwardLayer` (which is a subclass of `Neurons`) recieves feedforward input from a set of input features (these will be `PlaceCells` scattered across the `Environment`) $\\{f_{i}\\}_{i=1}^{N}$. The firing rate of the `ValueNeuron` is given as a function (here, linear) of the input features:\n",
"\n",
"$$ \\mathsf{\\hat{V}}(t;\\mathsf{w}) = \\sigma^{\\textrm{non-lin}} \\big( \\sum_{i=1}^{N}\\mathsf{w}_{i}f_{i}(t) \\big).$$\n",
"\n",
"The goal is to adjust the weights, $\\mathsf{w}_{i}$, such that $\\mathsf{\\hat{V}}(t)$ approximates the \"value\" of the current policy. The value of the current policy $\\pi$ is defined as the decaying sum of expected future rewards\n",
"\n",
"$$\\mathsf{\\hat{V}}(t) \\approx \\mathsf{V}^{\\pi}(t) = \\mathbb{E} \\bigg[ \\int_{t}^{\\infty} e^{-\\frac{t^{\\prime}-t}{\\tau}} \\mathsf{R}(t^{\\prime}) dt^{\\prime}\\bigg] $$\n",
"\n",
"where the expectation is taken over any stochasticity present in the current policy (i.e. how the `Agent` moves) and `Environement` (although in this case the environment will be deterministic). This definition of value is temporally continuous: the key differences, compared to the more common form where value is written as a sum of rewards over discrete future timesteps, is that it is now a continuous integral over a reward *density* function and temporal discounting is done by the exponential decaying over a time period $\\tau$. \n",
"\n",
"We can take the temporal derivative of this and derive a consistency equation (analagous to the Bellman equation) satisfied by this value function. This naturally gives a temporal difference-style update rule which relies on \"bootstrapping\" (the current estimate of the value function is used in lieu of the true value function) to optimize the weights of the value function approximation. A good reference for continuous RL is Doya (2000) if you want to know more about deriving this learning rule.\n",
"\n",
"$$ \\delta \\mathsf{w_i}(t) = \\eta \\bigg( \\mathsf{R}(t) + \\frac{d\\mathsf{\\hat{V}}(t)}{dt} - \\frac{\\mathsf{\\hat{V}(t)}}{\\tau} \\bigg) \\mathsf{e}_i(t) $$\n",
"\n",
"For now it suffices to observe that this learning rule is very similar to the temporally discrete TD-update rule. The first term represents the continuous analog of the temporal difference error (in fact, if you rediscretise with using the euler formula $\\dot{\\mathsf{V}}(t) = \\frac{\\mathsf{V}(t+dt) - \\mathsf{V}(t)}{dt}$ to replace the derivative, and set $dt =$ 1, you will see they are identical). The second term is the `eligibility trace' (using an eligibility trace is optional, and it could just be replaced with $f_i(t)$, however doing so aids stability of the learning) of feature inputs, defined as: \n",
"\n",
"$$\\mathsf{e}_i(t) = \\frac{1}{\\tau_{\\mathsf{e}}} \\int_{-\\infty}^{t} e^{-\\frac{t-t^{\\prime}}{\\tau_{\\mathsf{e}}}} f_i(t^{\\prime}) dt^{\\prime}$$\n",
"\n",
"To learn the value function, therefore, the `ValueNeuron` needs to\n",
"1. Be able to linearly summate its inputs. Already implemented by `FeedForwardLayer.update()`\n",
"2. Store and update the eligibility trace.\n",
"3. Implement the learning rule. This will require access to the reward density function $\\mathsf{R}(t)$, the eligibiilty traces $\\mathsf{e}_i(t)$, the firing rate $\\mathsf{\\hat{V}}(t)$ and the temporal derivitive of the firing rate $\\frac{d\\mathsf{\\hat{V}}(t)}{dt}$\n",
"\n",
"We will use a temporal discounting timescale of $\\tau =$ 10 seconds and an eligibility trace timescale of $\\tau_{\\mathsf{e}} =$ 5 seconds. The input features and rewards will both be provided by layers of `RatInABox` \"`PlaceCells`\". \n",
"\n",
"This figure gives a summary of the network structure and the policy improvement procedure. \n",
"\n",
" ![](\"../.images/demos/rl_networksummary.png\")
\n"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"## Install RatInABox"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Requirement already satisfied: ratinabox in /Users/tomgeorge/miniconda3/envs/riab/lib/python3.10/site-packages (1.11.1)\n",
"Requirement already satisfied: matplotlib in /Users/tomgeorge/miniconda3/envs/riab/lib/python3.10/site-packages (from ratinabox) (3.5.3)\n",
"Requirement already satisfied: shapely in /Users/tomgeorge/miniconda3/envs/riab/lib/python3.10/site-packages (from ratinabox) (2.0.1)\n",
"Requirement already satisfied: scipy in /Users/tomgeorge/miniconda3/envs/riab/lib/python3.10/site-packages (from ratinabox) (1.9.3)\n",
"Requirement already satisfied: numpy in /Users/tomgeorge/miniconda3/envs/riab/lib/python3.10/site-packages (from ratinabox) (1.24.2)\n",
"Requirement already satisfied: python-dateutil>=2.7 in /Users/tomgeorge/miniconda3/envs/riab/lib/python3.10/site-packages (from matplotlib->ratinabox) (2.8.2)\n",
"Requirement already satisfied: packaging>=20.0 in /Users/tomgeorge/miniconda3/envs/riab/lib/python3.10/site-packages (from matplotlib->ratinabox) (23.0)\n",
"Requirement already satisfied: pyparsing>=2.2.1 in /Users/tomgeorge/miniconda3/envs/riab/lib/python3.10/site-packages (from matplotlib->ratinabox) (3.0.9)\n",
"Requirement already satisfied: cycler>=0.10 in /Users/tomgeorge/miniconda3/envs/riab/lib/python3.10/site-packages (from matplotlib->ratinabox) (0.11.0)\n",
"Requirement already satisfied: fonttools>=4.22.0 in /Users/tomgeorge/miniconda3/envs/riab/lib/python3.10/site-packages (from matplotlib->ratinabox) (4.38.0)\n",
"Requirement already satisfied: pillow>=6.2.0 in /Users/tomgeorge/miniconda3/envs/riab/lib/python3.10/site-packages (from matplotlib->ratinabox) (9.4.0)\n",
"Requirement already satisfied: kiwisolver>=1.0.1 in /Users/tomgeorge/miniconda3/envs/riab/lib/python3.10/site-packages (from matplotlib->ratinabox) (1.4.4)\n",
"Requirement already satisfied: six>=1.5 in /Users/tomgeorge/miniconda3/envs/riab/lib/python3.10/site-packages (from python-dateutil>=2.7->matplotlib->ratinabox) (1.16.0)\n"
]
}
],
"source": [
"!pip install ratinabox"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Import RatInABox\n",
"..along with the `ValueNeuron` class"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"import numpy as np\n",
"from tqdm.notebook import tqdm\n",
"from time import time\n",
"import copy\n",
"\n",
"\n",
"import ratinabox\n",
"from ratinabox.Environment import Environment\n",
"from ratinabox.Agent import Agent\n",
"from ratinabox.Neurons import *\n",
"ratinabox.stylize_plots(); ratinabox.autosave_plots=True; ratinabox.figure_directory = \"../figures/\"\n",
"\n",
"from ratinabox.contribs.ValueNeuron import ValueNeuron\n",
"\n",
"%load_ext autoreload\n",
"%autoreload 2"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"## Make the Environment, Agent and Neuron populations \n",
"We now initialise an `Environment` with a wall. We put an `Agent` inside the environment. 400 `PlaceCells` of random size uniformly distributed across the environment are the input features. A single `PlaceCell` hidden behind the wall is the reward and a `ValueNeuron` (discussed above) will learn the value function by TD learning as the agent explores. "
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Figure saved to /Users/tomgeorge/Documents/RatInABox/figures/09_04_24/PlaceCells_ratemaps_1141.png & .svg\n",
"Figure saved to /Users/tomgeorge/Documents/RatInABox/figures/09_04_24/PlaceCells_ratemaps_1141_1.png & .svg\n",
"Figure saved to /Users/tomgeorge/Documents/RatInABox/figures/09_04_24/FeedForwardLayer_ratemaps_1141.png & .svg\n"
]
},
{
"data": {
"text/plain": [
"(\n",
" ![]()
\n",
"
\n",
"\n",
"\n",
"\n"
],
"text/plain": [
"\n",
" \n",
" \n",
" \n",
"
\n",
"\n",
" ![]()
\n",
"
\n",
"\n",
"\n",
"\n"
],
"text/plain": [
"\n",
" \n",
" \n",
" \n",
"
\n",
"![](\"../.images/demos/riab_for_rl_full.gif\")
\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Disclaimer: \n",
"Please note this script is intended as a tutorial, here only to demonstrate of how `RatInABox` could be used for a reinforcement learning project in continuous time and space and not, in any way, as a scientific result regarding the role of hippocampus' role in navigational reinforcement learning. "
]
}
],
"metadata": {
"interpreter": {
"hash": "f3cdbb439150f468ef73eca921e462ba1cf0894f575003150b1b7f3566c8fd4a"
},
"kernelspec": {
"display_name": "Python 3.9.7 ('base')",
"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
}