Simple-ARE-Example-1.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "937e7701",
   "metadata": {},
   "source": [
    "# Simple ARE Demo\n",
    "\n",
    "This simple demo shows how to estimate the ARE in the case where reafference is the same across observations. The example is equivalent to observing something like angular velocity in the artificial ape environment (without the moving cube exafference). Unlike in the original experiment, here the platform only rotates in one direction. Exafference is therefore positive making the example slightly more interesting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "43d76ca0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sb\n",
    "\n",
    "def gaussian(mu, sigma):\n",
    "    return lambda n=1: np.random.normal(mu, sigma, size=n)\n",
    "\n",
    "def uniform(x1, x2):\n",
    "    return lambda n=1: x2 - (np.random.rand(n) * (x2 - x1))\n",
    "\n",
    "def constant(x):\n",
    "    return lambda n=1: np.full((n,), x)\n",
    "\n",
    "def categorical(x, p=None):\n",
    "    x = np.array(x)\n",
    "    if p is None:\n",
    "        p = np.ones(x.shape[0]) / x.shape[0]\n",
    "    p = np.array(p)\n",
    "    assert x.shape[0] == p.shape[0]\n",
    "    def _dist(n=1):\n",
    "        return np.random.choice(x, size=n, p=p)\n",
    "    return _dist\n",
    "\n",
    "n = 10000\n",
    "A = categorical([-1,0,1])(n)        # action turn head left/noop/right\n",
    "X0 = uniform(-1,1)(n)               # current rotation (observation)\n",
    "P = categorical([0,1])(n)           # noise term is the platform rotation\n",
    "# SCM\n",
    "# experiment here to see how changing the SCM gives different results. A more complex SCM will need a more flexible model (see next cell)\n",
    "X1 = X0 + A + P "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "df2bfab2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bd1ed42bc4154217bc8c7cc6acfac40e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "from tqdm.auto import tqdm\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "class Module(nn.Module):\n",
    "    \n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        # if the SCM is made more complex, use a more complex model with more layers.\n",
    "        self.layers = nn.Sequential(\n",
    "            nn.Linear(2, 1)\n",
    "        )\n",
    "        self.optim = torch.optim.Adam(self.parameters(), lr=0.0005)\n",
    "        self.criterion = nn.MSELoss()\n",
    "        \n",
    "    def forward(self, x, a):\n",
    "        z = torch.cat([x, a], dim=-1)\n",
    "        return self.layers(z)\n",
    "    \n",
    "    def predict(self, x, a):\n",
    "        pred_total_effect = self.forward(x, a)\n",
    "        noop = torch.zeros_like(a)\n",
    "        pred_exafferent_effect = model(x, noop)\n",
    "        pred_reafferent_effect = pred_total_effect - pred_exafferent_effect.detach()\n",
    "        return pred_total_effect, pred_reafferent_effect, pred_exafferent_effect\n",
    "    \n",
    "    def step(self, x1, x2, a):\n",
    "        self.optim.zero_grad()\n",
    "        pred_total_effect, pred_reafferent_effect, pred_exafferent_effect = self.predict(x1, a)\n",
    "        pred_effect = pred_exafferent_effect + pred_reafferent_effect\n",
    "        total_effect = x2 - x1\n",
    "        loss = self.criterion(pred_effect, total_effect)\n",
    "        loss.backward()\n",
    "        self.optim.step()\n",
    "        return loss.detach()\n",
    "    \n",
    "    def train(self, x1, x2, a, epochs=1000):\n",
    "        pbar = tqdm(range(epochs))\n",
    "        for e in pbar:\n",
    "            x1, x2, a = self.shuffle(x1, x2, a)\n",
    "            loss = self.step(x1, x2, a)\n",
    "            pbar.set_description(f\"loss: {loss.item():.5f}\")\n",
    "    \n",
    "    def shuffle(self, *x):\n",
    "        indx = torch.randperm(x[0].shape[0])\n",
    "        return [z[indx] for z in x]\n",
    "            \n",
    "x0, x1, a = [torch.from_numpy(z).unsqueeze(1).float() for z in (X0, X1, A)]   \n",
    "model = Module()\n",
    "model.train(x0, x1, a, epochs=5000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ebcddc40",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ACTION:  1 | GROUND TRUTH | total effect:  1.500 | reafference:  1.000 | exafference :  0.500\n",
      "ACTION:  1 | ESTIMATE     | total effect:  1.511 | reafference:  1.010 | exafference :  0.501\n",
      "\n",
      "ACTION:  0 | GROUND TRUTH | total effect:  0.500 | reafference:  0.000 | exafference :  0.500\n",
      "ACTION:  0 | ESTIMATE     | total effect:  0.501 | reafference:  0.000 | exafference :  0.501\n",
      "\n",
      "ACTION: -1 | GROUND TRUTH | total effect: -0.500 | reafference: -1.000 | exafference :  0.500\n",
      "ACTION: -1 | ESTIMATE     | total effect: -0.509 | reafference: -1.010 | exafference :  0.501\n",
      "\n"
     ]
    }
   ],
   "source": [
    "a1 = torch.ones_like(a)\n",
    "a2 = torch.zeros_like(a)\n",
    "a3 = -torch.ones_like(a)\n",
    "\n",
    "with torch.no_grad():\n",
    "    t, re, ex = model.predict(x0, a1)\n",
    "    print(f\"ACTION:  1 | GROUND TRUTH | total effect: { 1.5     : .3f} | reafference: { 1.        : .3f} | exafference : { 0.5       : .3f}\")\n",
    "    print(f\"ACTION:  1 | ESTIMATE     | total effect: {t.mean() : .3f} | reafference: { re.mean() : .3f} | exafference : { ex.mean() : .3f}\\n\")\n",
    "    \n",
    "    t, re, ex = model.predict(x0, a2)\n",
    "    print(f\"ACTION:  0 | GROUND TRUTH | total effect: { 0.5     : .3f} | reafference: { 0.        : .3f} | exafference : { 0.5       : .3f}\")\n",
    "    print(f\"ACTION:  0 | ESTIMATE     | total effect: {t.mean() : .3f} | reafference: { re.mean() : .3f} | exafference : { ex.mean() : .3f}\\n\")\n",
    "\n",
    "    t, re, ex = model.predict(x0, a3)\n",
    "    print(f\"ACTION: -1 | GROUND TRUTH | total effect: { -0.5    : .3f} | reafference: { -1.       : .3f} | exafference : { 0.5       : .3f}\")\n",
    "    print(f\"ACTION: -1 | ESTIMATE     | total effect: {t.mean() : .3f} | reafference: { re.mean() : .3f} | exafference : { ex.mean() : .3f}\\n\")\n",
    "    \n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6270085c",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "reaff",
   "language": "python",
   "name": "reaff"
  },
  "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}