Reinforcement Learning in Nengo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# INRC Tutorial: RL in Nengo\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### 1. Theoretical background on PES rule\n",
    "\n",
    "- $E = 1/2 \\int (x-\\hat{x})^2 dx$\n",
    "\n",
    "\n",
    "- $\\delta E/\\delta d_i = (x-\\hat{x})a_i$ (as usual for finding decoders)\n",
    "\n",
    "\n",
    "- So, to move down the gradient:\n",
    "    - $\\Delta d_i = -\\kappa (x - \\hat{x})a_i$ (NEF notation)\n",
    "\n",
    "\n",
    "- How do we make it realistic?\n",
    "    - Need weights\n",
    "    \n",
    "    \n",
    "- The NEF tells us:\n",
    "    - $\\omega_{ij} =  d_i \\cdot e_j$\n",
    "    - $\\Delta \\omega_{ij} = \\kappa a_i E \\cdot e_j $ \n",
    "\n",
    "\n",
    "- What's $ E \\cdot e_j$?\n",
    "    - That's the current that this neuron would get if it had $E$ as an input\n",
    "    - But we don't want this current to drive the neuron\n",
    "    - Rather, we want it to change the weight\n",
    "    - It's a *modulatory* input\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "- This is the \"Prescribed Error Sensitivity\" PES rule\n",
    "    - Any model in the NEF could use this instead of computing decoders\n",
    "    - Requires some other neural group computing the error $E$\n",
    "    - Used in Spaun for Q-value learning (reinforcement task)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### 2. A simple grid world example\n",
    "\n",
    "Some preliminary details:\n",
    "\n",
    "- the agent has four actions (go up, go right, go left, go down)\n",
    "- state is encoded as 25D vector, Q-values are a 4D function of this state \n",
    "- reward is +1 for getting to goal state, -1 for hitting a wall, 0 elsewhere\n",
    "- the error signal is $\\Delta Q(s,a) = \\alpha (R + Q(s', a') - Q(s, a))$\n",
    "- connections into error signal ensemble use fast tau for $Q(s', a')$, slow tau for $Q(s, a)$.\n",
    "\n",
    "First, we'll define some parameters for models, the environment, and some helper functions for managing the environment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "import td_grid\n",
    "import nengo\n",
    "import numpy as np\n",
    "\n",
    "from utils import weight_init\n",
    "\n",
    "# high-level configuration\n",
    "stepsize = 10 # milleseconds between each action\n",
    "n_neurons = 2500\n",
    "n_actions = 4\n",
    "\n",
    "fast_tau = 0\n",
    "slow_tau = 0.01\n",
    "\n",
    "# build the world and add the agent\n",
    "env_map = \"\"\"\n",
    "#######\n",
    "#     #\n",
    "#     #\n",
    "#  G  #\n",
    "#     #\n",
    "#     #\n",
    "#######\n",
    "\"\"\"\n",
    "\n",
    "agent = td_grid.ContinuousAgent()\n",
    "\n",
    "environment = td_grid.World(td_grid.GridCell, map=env_map, directions=4)\n",
    "environment.add(agent, x=2, y=3, dir=2)\n",
    "\n",
    "env_update = td_grid.EnvironmentInterface(agent, n_actions=4, epsilon=0.1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we'll define some helper functions to manage our environment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sensor(t):\n",
    "    '''Return current x,y coordinates of agent as one hot representation'''\n",
    "    data = np.zeros(25)\n",
    "    idx = 5 * (agent.x - 1) + (agent.y - 1)\n",
    "    data[idx] = 1\n",
    "\n",
    "    return data\n",
    "\n",
    "def reward(t):\n",
    "    '''Call to get current reward signal provided to agent'''\n",
    "    return agent.reward"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "with nengo.Network(seed=2) as model:\n",
    "    env = td_grid.GridNode(environment, dt=0.001)\n",
    "    \n",
    "    # define nodes for plotting data, managing agent's interface with environment\n",
    "    reward_node = nengo.Node(reward, size_out=1, label='reward')\n",
    "    sensor_node = nengo.Node(sensor, size_out=25, label='sensor')\n",
    "    update_node = nengo.Node(env_update.step, size_in=4, size_out=12, label='env')\n",
    "    qvalue_node = nengo.Node(size_in=4)\n",
    "\n",
    "    # define neurons to encode state representations\n",
    "    state = nengo.Ensemble(n_neurons=n_neurons, dimensions=25, \n",
    "                           intercepts=nengo.dists.Choice([0.15]), radius=2)\n",
    "    \n",
    "    # define neurons that compute the learning signal\n",
    "    learn_signal = nengo.Ensemble(n_neurons=1000, dimensions=4)\n",
    "                 \n",
    "    # connect the sensor to state ensemble\n",
    "    nengo.Connection(sensor_node, state, synapse=None)\n",
    "    reward_probe = nengo.Probe(reward_node, synapse=fast_tau)\n",
    "\n",
    "    # connect state representation to environment interface\n",
    "    q_conn = nengo.Connection(state.neurons, update_node,\n",
    "                              transform=weight_init(shape=(n_actions, n_neurons)), \n",
    "                              learning_rule_type=nengo.PES(1e-3, pre_tau=slow_tau),\n",
    "                              synapse=fast_tau)\n",
    "    \n",
    "    # connect update node to error signal ensemble w/ fast, slow conns to compute prediction error\n",
    "    nengo.Connection(update_node[0:n_actions], learn_signal, transform=-1, synapse=slow_tau)\n",
    "    nengo.Connection(update_node[n_actions:2*n_actions], learn_signal, transform=1, synapse=fast_tau)\n",
    "    \n",
    "    # connect the learning signal to the learning rule\n",
    "    nengo.Connection(learn_signal, q_conn.learning_rule, transform=-1, synapse=fast_tau)\n",
    "\n",
    "    \n",
    "    # for plotting and visualization purposes\n",
    "    nengo.Connection(update_node[2*n_actions:], qvalue_node, synapse=fast_tau)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "                <div id=\"70bbc289-1bbc-4cfc-bd1f-4bc3db75e00f\">\n",
       "                    <iframe\n",
       "                        src=\"http://localhost:53191/?token=8eb3f65da8a76e6cf31820ef8613cda506f3eeefb3632442\"\n",
       "                        width=\"100%\"\n",
       "                        height=\"600\"\n",
       "                        frameborder=\"0\"\n",
       "                        class=\"cell\"\n",
       "                        style=\"border: 1px solid #eee;\"\n",
       "                        allowfullscreen></iframe>\n",
       "                </div>\n",
       "            "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from nengo_gui.ipython import IPythonViz\n",
    "IPythonViz(model,'configs/default.py.cfg')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### 3. A continuous state grid world example\n",
    "\n",
    "In this model:\n",
    "- the agent always moves randomly, it's not *using* what\n",
    "it learns to change its movement (it is just trying to anticipate future rewards)\n",
    "- the agent is given a reward whenever it is in the green square, and a \n",
    "punishment (negative reward) whenever it is in the red square  \n",
    "- it learns to anticipate the reward/punishment as shown in the value graph\n",
    "- we convert the error rule into the continuous domain by using a long time constant for s-1\n",
    "and a short time constant for s as before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import grid\n",
    "\n",
    "mymap=\"\"\"\n",
    "#######\n",
    "#     #\n",
    "# # # #\n",
    "# # # #\n",
    "#G   R#\n",
    "#######\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "world = grid.World(grid.GridCell, map=mymap, directions=4)\n",
    "body = grid.ContinuousAgent()\n",
    "world.add(body, x=1, y=2, dir=2)\n",
    "\n",
    "tau=0.1\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def move(t, x):\n",
    "    '''Defines a continuous action policy for the agent'''\n",
    "    speed, rotation = x\n",
    "    dt = 0.001\n",
    "    max_speed = 20.0\n",
    "    max_rotate = 10.0\n",
    "    body.turn(rotation * dt * max_rotate)\n",
    "    body.go_forward(speed * dt * max_speed)\n",
    "    \n",
    "    if int(body.x) == 1:\n",
    "        world.grid[4][4].wall = True\n",
    "        world.grid[4][2].wall = False\n",
    "    if int(body.x) == 4:\n",
    "        world.grid[4][2].wall = True\n",
    "        world.grid[4][4].wall = False\n",
    "\n",
    "def sensor(t):\n",
    "    '''Obtain environment state using sensors'''\n",
    "    angles = (np.linspace(-0.5, 0.5, 3) + body.dir) % world.directions\n",
    "    return [body.detect(d, max_distance=4)[0] for d in angles]\n",
    "\n",
    "def braiten(x):\n",
    "    '''Compute input to movement function based on sensor'''\n",
    "    turn = x[2] - x[0]\n",
    "    spd = x[1] - 0.5\n",
    "    return spd, turn\n",
    "\n",
    "def position_func(t):\n",
    "    '''Create unit normalized state representation of grid world'''\n",
    "    return body.x / world.width * 2 - 1, 1 - body.y / world.height * 2, body.dir / world.directions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "with nengo.Network(seed=2) as model:\n",
    "    env = grid.GridNode(world, dt=0.005)\n",
    "\n",
    "    # define nodes and ensembles for managing action policy\n",
    "    movement = nengo.Node(move, size_in=2)\n",
    "    stim_radar = nengo.Node(sensor)\n",
    "    \n",
    "    radar = nengo.Ensemble(n_neurons=50, dimensions=3, radius=4, seed=2,\n",
    "                           noise=nengo.processes.WhiteSignal(10, 0.1, rms=1))\n",
    "    \n",
    "    nengo.Connection(stim_radar, radar)\n",
    "    nengo.Connection(radar, movement, function=braiten)  \n",
    "    \n",
    "    # encode state information in ensemble of neurons\n",
    "    position = nengo.Node(position_func)\n",
    "    state = nengo.Ensemble(100, 3)\n",
    "    \n",
    "    nengo.Connection(position, state, synapse=None)\n",
    "    \n",
    "    reward = nengo.Node(lambda t: body.cell.reward)\n",
    "        \n",
    "    value = nengo.Ensemble(n_neurons=50, dimensions=1)\n",
    "\n",
    "    # this sets up learning on our connection between state and value encodings\n",
    "    learn_conn = nengo.Connection(state, value, function=lambda x: 0,\n",
    "                                  learning_rule_type=nengo.PES(learning_rate=1e-4, pre_tau=tau))\n",
    "    \n",
    "    # this connection adds the reward to the learning signal\n",
    "    nengo.Connection(reward, learn_conn.learning_rule, transform=-1, synapse=tau)\n",
    "    \n",
    "    # this connection adds the observed observed value\n",
    "    nengo.Connection(value, learn_conn.learning_rule, transform=-0.9, synapse=0.01)\n",
    "    \n",
    "    # this connection substracts the predicted value\n",
    "    nengo.Connection(value, learn_conn.learning_rule, transform=1, synapse=tau)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "                <div id=\"bf33c426-004c-4dac-a9bc-93ae7b00137b\">\n",
       "                    <iframe\n",
       "                        src=\"http://localhost:53264/?token=3cd1018bbdcfa030679f51c0bde3a6a54ac1001f9101ec8c\"\n",
       "                        width=\"100%\"\n",
       "                        height=\"600\"\n",
       "                        frameborder=\"0\"\n",
       "                        class=\"cell\"\n",
       "                        style=\"border: 1px solid #eee;\"\n",
       "                        allowfullscreen></iframe>\n",
       "                </div>\n",
       "            "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from nengo_gui.ipython import IPythonViz\n",
    "IPythonViz(model,'configs/learning6-value.py.cfg')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.7"
  },
  "livereveal": {
   "scroll": true,
   "start_slideshow_at": "selected"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}