extra.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "The size of tensor a (15) must match the size of tensor b (10) at non-singleton dimension 0",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[13], line 284\u001b[0m\n\u001b[0;32m    281\u001b[0m x \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mrandn(batch_size, obs_size)\n\u001b[0;32m    282\u001b[0m h \u001b[38;5;241m=\u001b[39m ddrnn\u001b[38;5;241m.\u001b[39minitial_state(batch_size)\n\u001b[1;32m--> 284\u001b[0m out, hidden \u001b[38;5;241m=\u001b[39m \u001b[43mddrnn\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mh\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32mc:\\Users\\Andrew\\AppData\\Local\\miniforge3\\envs\\vrAnalysis\\lib\\site-packages\\torch\\nn\\modules\\module.py:1532\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m   1530\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[0;32m   1531\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m-> 1532\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n",
      "File \u001b[1;32mc:\\Users\\Andrew\\AppData\\Local\\miniforge3\\envs\\vrAnalysis\\lib\\site-packages\\torch\\nn\\modules\\module.py:1541\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m   1536\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[0;32m   1537\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[0;32m   1538\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[0;32m   1539\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[0;32m   1540\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[1;32m-> 1541\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m forward_call(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[0;32m   1543\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m   1544\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "Cell \u001b[1;32mIn[13], line 232\u001b[0m, in \u001b[0;36mDynamicDisRNN.forward\u001b[1;34m(self, observations, prev_latents)\u001b[0m\n\u001b[0;32m    230\u001b[0m \u001b[38;5;66;03m# Add penalty for KL divergence on inputs to each update MLP\u001b[39;00m\n\u001b[0;32m    231\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m mu, sigma \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mzip\u001b[39m(update_mlp_inputs, sigma_to_updates):\n\u001b[1;32m--> 232\u001b[0m     penalty \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mbeta_scale \u001b[38;5;241m*\u001b[39m kl_gaussian(mu, sigma)\n\u001b[0;32m    234\u001b[0m \u001b[38;5;66;03m# Measure the updates weights and value for each latent state\u001b[39;00m\n\u001b[0;32m    235\u001b[0m updates \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mstack([mlp(update_mlp_inputs[:, i]) \u001b[38;5;28;01mfor\u001b[39;00m i, mlp \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28menumerate\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mupdate_mlps)], dim\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1\u001b[39m)\n",
      "\u001b[1;31mRuntimeError\u001b[0m: The size of tensor a (15) must match the size of tensor b (10) at non-singleton dimension 0"
     ]
    }
   ],
   "source": [
    "from typing import Callable\n",
    "from collections.abc import Sequence\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "class DisRNN(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        obs_size=2,\n",
    "        target_size=1,\n",
    "        latent_size=10,\n",
    "        update_mlp_shape=(10, 10, 10),\n",
    "        choice_mlp_shape=(10, 10, 10),\n",
    "        eval_mode=0,\n",
    "        beta_scale=1,\n",
    "        activation=nn.ReLU,\n",
    "    ):\n",
    "        super().__init__()\n",
    "        self.target_size = target_size\n",
    "        self.latent_size = latent_size\n",
    "        self.beta_scale = beta_scale\n",
    "        self.eval_mode = eval_mode\n",
    "        self.activation = activation\n",
    "\n",
    "        mlp_input_size = latent_size + obs_size\n",
    "        \n",
    "        # Initialize update MLP parameters\n",
    "        self.update_mlp_sigmas_unsquashed = nn.Parameter(\n",
    "            torch.empty(mlp_input_size, latent_size).uniform_(-3, -2)\n",
    "        )\n",
    "        self.update_mlp_multipliers = nn.Parameter(\n",
    "            torch.ones(mlp_input_size, latent_size)\n",
    "        )\n",
    "        \n",
    "        # Initialize latent parameters\n",
    "        self.latent_sigmas_unsquashed = nn.Parameter(\n",
    "            torch.empty(latent_size).uniform_(-3, -2)\n",
    "        )\n",
    "        self.latent_inits = nn.Parameter(\n",
    "            torch.empty(latent_size).uniform_(-1, 1)\n",
    "        )\n",
    "\n",
    "        # Create MLPs for each latent\n",
    "        self.update_mlps = nn.ModuleList([\n",
    "            MLP(mlp_input_size, update_mlp_shape, 2, activation)\n",
    "            for _ in range(latent_size)\n",
    "        ])\n",
    "        \n",
    "        # Choice MLP\n",
    "        self.choice_mlp = MLP(latent_size, choice_mlp_shape, target_size, activation)\n",
    "\n",
    "    def forward(self, observations, prev_latents):\n",
    "        batch_size = observations.shape[0]\n",
    "        penalty = torch.zeros(batch_size, device=observations.device)\n",
    "\n",
    "        # Update MLPs\n",
    "        update_mlp_sigmas = 2 * torch.sigmoid(self.update_mlp_sigmas_unsquashed) * (1 - self.eval_mode)\n",
    "        update_mlp_mus_unscaled = torch.cat((observations, prev_latents), dim=1)\n",
    "        update_mlp_mus = update_mlp_mus_unscaled.unsqueeze(2) * self.update_mlp_multipliers\n",
    "        \n",
    "        # Calculate updates for each latent\n",
    "        new_latents = torch.zeros_like(prev_latents)\n",
    "        for i in range(self.latent_size):\n",
    "            # Add noise to inputs\n",
    "            noise = torch.randn_like(update_mlp_mus[:, :, i]) * update_mlp_sigmas[:, i]\n",
    "            update_mlp_inputs = update_mlp_mus[:, :, i] + noise\n",
    "            \n",
    "            # Calculate KL divergence\n",
    "            kl = 0.5 * torch.sum(-torch.log(update_mlp_sigmas[:, i]) - 1.0 + \n",
    "                                update_mlp_sigmas[:, i] + update_mlp_mus[:, :, i].pow(2), dim=1)\n",
    "            penalty += self.beta_scale * kl\n",
    "            \n",
    "            # Calculate update and weight\n",
    "            update = self.update_mlps[i](update_mlp_inputs).squeeze(-1)\n",
    "            w = torch.sigmoid(self.update_weight_mlps[i](update_mlp_inputs)).squeeze(-1)\n",
    "            \n",
    "            # Update latent\n",
    "            new_latent = w * update + (1 - w) * prev_latents[:, i]\n",
    "            new_latents[:, i] = new_latent\n",
    "\n",
    "        # Global bottleneck\n",
    "        latent_sigmas = 2 * torch.sigmoid(self.latent_sigmas_unsquashed) * (1 - self.eval_mode)\n",
    "        noised_up_latents = new_latents + latent_sigmas * torch.randn_like(new_latents)\n",
    "        penalty += torch.sum(-torch.log(latent_sigmas) - 1.0 + \n",
    "                           latent_sigmas + new_latents.pow(2), dim=1)\n",
    "\n",
    "        # Choice MLP\n",
    "        y_hat = self.choice_mlp(noised_up_latents)\n",
    "        \n",
    "        # Append penalty\n",
    "        penalty = penalty.unsqueeze(1) * (1 - self.eval_mode)\n",
    "        output = torch.cat((y_hat, penalty), dim=1)\n",
    "\n",
    "        return output, noised_up_latents\n",
    "\n",
    "    def initial_state(self, batch_size, device=None):\n",
    "        return self.latent_inits.unsqueeze(0).repeat(batch_size, 1).to(device)\n",
    "\n",
    "\n",
    "class MLP(nn.Module):\n",
    "    \"\"\"A simple multi-layer perceptron (MLP) network.\n",
    "    \n",
    "    The MLP is a series of linear layers with an activation function applied after each hidden layer\n",
    "    but not after the final output layer. All layers are nested in a `nn.Sequential` container.\n",
    "    \"\"\"\n",
    "    def __init__(self, input_size: int, hidden_sizes: Sequence[int], output_size: int, activation: Callable=nn.ReLU):\n",
    "        \"\"\"Initialize the MLP.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        input_size : int\n",
    "            The size of the input tensor.\n",
    "        hidden_sizes : Sequence[int]\n",
    "            A sequence of integers representing the size of each hidden layer.\n",
    "        output_size : int\n",
    "            The size of the output tensor.\n",
    "        activation : Callable\n",
    "            The activation function to apply after each hidden layer. Default=nn.ReLU.\n",
    "        with_hooks : bool\n",
    "            Whether to include hooks in the model for inspection of intermediate activations. Default=False.\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.input_size = input_size\n",
    "        self.hidden_sizes = hidden_sizes\n",
    "        self.output_size = output_size\n",
    "        self.activation = activation\n",
    "        \n",
    "        mlp = []\n",
    "        prev_size = input_size\n",
    "        for size in hidden_sizes:\n",
    "            mlp.extend([\n",
    "                nn.Linear(prev_size, size),\n",
    "                self.activation()\n",
    "            ])\n",
    "            prev_size = size\n",
    "        \n",
    "        mlp.append(nn.Linear(prev_size, output_size))\n",
    "        \n",
    "        self.mlp = nn.Sequential(*mlp)\n",
    "\n",
    "    def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
    "        \"\"\"Forward pass of the MLP.\"\"\"\n",
    "        return self.mlp(x)\n",
    "\n",
    "def kl_gaussian(mean: torch.Tensor, var: torch.Tensor, dim: int = -1) -> torch.Tensor:\n",
    "  r\"\"\"Calculate KL divergence between given and standard gaussian distributions.\n",
    "\n",
    "  KL(p, q) = H(p, q) - H(p) = -\\int p(x)log(q(x))dx - -\\int p(x)log(p(x))dx\n",
    "          = 0.5 * [log(|s2|/|s1|) - 1 + tr(s1/s2) + (m1-m2)^2/s2]\n",
    "          = 0.5 * [-log(|s1|) - 1 + tr(s1) + m1^2] (if m2 = 0, s2 = 1)\n",
    "  Args:\n",
    "    mean: mean vector of the first distribution\n",
    "    var: diagonal vector of covariance matrix of the first distribution\n",
    "\n",
    "  Returns:\n",
    "    A scalar representing KL divergence of the two Gaussian distributions.\n",
    "  \"\"\"\n",
    "  return 0.5 * torch.sum(-torch.log(var) - 1.0 + var + torch.square(mean), dim=dim)\n",
    "\n",
    "\n",
    "class DynamicDisRNN(nn.Module):\n",
    "    def __init__(self, obs_size=2, target_size=1, latent_size=10, update_mlp_shape=(10, 10, 10), choice_mlp_shape=(10, 10, 10), eval_mode=0, beta_scale=1, activation=nn.ReLU):\n",
    "        super().__init__()\n",
    "        self.target_size = target_size\n",
    "        self.latent_size = latent_size\n",
    "        self.obs_size = obs_size\n",
    "        self.beta_scale = beta_scale\n",
    "        self.eval_mode = eval_mode\n",
    "        self.activation = activation\n",
    "\n",
    "        mlp_input_size = latent_size + obs_size\n",
    "        \n",
    "        # to_mu_update takes in the observations and latent states and transforms them with a single linear layer\n",
    "        # and a sigmoidal activation function to produce a set of multipliers that determine how much to use each\n",
    "        # potential input to the UpdateMLP for each latent state.\n",
    "        self.to_mu_update = nn.ModuleList([\n",
    "            nn.Sequential(nn.Linear(mlp_input_size, mlp_input_size), nn.Sigmoid()) for _ in range(latent_size)\n",
    "        ])\n",
    "\n",
    "        # to_sigma_update takes in the observations and latent states and transforms them with a single linear layer\n",
    "        # and a sigmoidal activation function to produce a set of sigmas that determine the amount of noise to add to\n",
    "        # the inputs to the UpdateMLP for each latent state.\n",
    "        self.to_sigma_update = nn.ModuleList([nn.Linear(mlp_input_size, mlp_input_size) for _ in range(latent_size)])\n",
    "\n",
    "        # update_mlp takes in the transformed observations and latent states and produces a set of updates for each\n",
    "        # latent state. There is an update_mlp for each latent state.\n",
    "        self.update_mlps = nn.ModuleList([\n",
    "            MLP(mlp_input_size, update_mlp_shape, 2, activation=activation)\n",
    "            for _ in range(latent_size)\n",
    "        ])\n",
    "\n",
    "        # to_mu_latent takes in the observations and latent states and transforms them with a single linear layer\n",
    "        # and a sigmoidal activation function to produce a set of multipliers that determine how much to use each\n",
    "        # previous latent state for the next time step. to_sigma_latent does the same thing for noise on each latent.\n",
    "        self.to_mu_latent = nn.Sequential(nn.Linear(mlp_input_size, latent_size), nn.Sigmoid())\n",
    "        self.to_sigma_latent = nn.Linear(mlp_input_size, latent_size)\n",
    "\n",
    "        # to_mu_choice takes in the latent states and transforms them with a single linear layer and a sigmoidal\n",
    "        # activation function to produce a set of multipliers that determine how much to use each latent state and\n",
    "        # observation for the choice MLP.\n",
    "        self.to_mu_choice = nn.Sequential(nn.Linear(mlp_input_size, mlp_input_size), nn.Sigmoid())\n",
    "\n",
    "        # to_sigma_choice takes in the latent states and transforms them with a single linear layer and a sigmoidal\n",
    "        # activation function to produce a set of sigmas that determine the amount of noise to add to the inputs to\n",
    "        # the choice MLP.\n",
    "        self.to_sigma_choice = nn.Linear(mlp_input_size, mlp_input_size)\n",
    " \n",
    "        # choice_mlp takes in the transformed latent states and observations and produces the output of the model.\n",
    "        self.choice_mlp = MLP(mlp_input_size, choice_mlp_shape, target_size, activation=activation)\n",
    "\n",
    "        # set up an initial latent parameter that is learned\n",
    "        self.latent_inits = nn.Parameter(torch.empty(latent_size).uniform_(-1, 1))\n",
    "\n",
    "    def _step(self, observations, prev_latents):\n",
    "        \"\"\"Forward method for the dynamic disentangled RNN class.\"\"\"\n",
    "        batch_size = observations.shape[0]\n",
    "        penalty = torch.zeros(batch_size, device=observations.device)\n",
    "\n",
    "        # Concatenate observations and latent states, which is the input to every component of the network\n",
    "        obs_plus_latents = torch.cat((observations, prev_latents), dim=1) \n",
    "\n",
    "        # Measure mu and sigma for the inputs to each update MLP\n",
    "        mu_to_updates = torch.stack([mu(obs_plus_latents) for mu in self.to_mu_update], dim=1)\n",
    "        sigma_to_updates = torch.stack([sigma(obs_plus_latents) for sigma in self.to_sigma_update], dim=1)\n",
    "        \n",
    "        # Apply information bottleneck to the inputs to each update MLP\n",
    "        update_mlp_inputs = mu_to_updates * obs_plus_latents.unsqueeze(1) + sigma_to_updates * torch.randn_like(obs_plus_latents.unsqueeze(1))\n",
    "        \n",
    "        # Add penalty for KL divergence on inputs to each update MLP\n",
    "        penalty += self.beta_scale * kl_gaussian(update_mlp_inputs, sigma_to_updates, dim=(1, 2))\n",
    "\n",
    "        # Measure the updates weights and value for each latent state\n",
    "        updates = torch.stack([mlp(update_mlp_inputs[:, i]) for i, mlp in enumerate(self.update_mlps)], dim=1)\n",
    "        target = updates[:, :, 0]\n",
    "        weight = torch.sigmoid(updates[:, :, 1])\n",
    "\n",
    "        # Update latents with weighted updates from UpdateMLPs\n",
    "        new_latents = weight * target + (1 - weight) * prev_latents\n",
    "\n",
    "        # Measure mu and sigma for the latent states\n",
    "        mu_latent = self.to_mu_latent(obs_plus_latents)\n",
    "        sigma_latent = self.to_sigma_latent(obs_plus_latents)\n",
    "\n",
    "        # Apply information bottleneck to the latent states\n",
    "        noised_latents = mu_latent * new_latents + sigma_latent * torch.randn_like(new_latents)\n",
    "        \n",
    "        # Add penalty for KL divergence on latent states\n",
    "        penalty += self.beta_scale * kl_gaussian(noised_latents, sigma_latent)\n",
    "        \n",
    "        # Output of the choice MLP depends on the inputs and the updated bottlenecked latent states\n",
    "        obs_plus_new_latents = torch.cat((observations, noised_latents), dim=1)\n",
    "\n",
    "        # Measure mu and sigma for the inputs to the choice MLP\n",
    "        mu_to_choice = self.to_mu_choice(obs_plus_new_latents)\n",
    "        sigma_to_choice = self.to_sigma_choice(obs_plus_new_latents)\n",
    "\n",
    "        # Apply information bottleneck to the inputs to the choice MLP\n",
    "        choice_mlp_inputs = mu_to_choice * obs_plus_new_latents + sigma_to_choice * torch.randn_like(obs_plus_new_latents)\n",
    "\n",
    "        # Add penalty for KL divergence on inputs to the choice MLP\n",
    "        penalty += self.beta_scale * kl_gaussian(choice_mlp_inputs, sigma_to_choice)\n",
    "\n",
    "        # Measure choice output\n",
    "        choice_output = self.choice_mlp(choice_mlp_inputs)\n",
    "\n",
    "        return choice_output, noised_latents, penalty\n",
    "    \n",
    "    def forward(self, observations, prev_latents):\n",
    "        batch_size = observations.shape[0]\n",
    "        seq_size = observations.shape[1]\n",
    "        outputs = torch.zeros(batch_size, seq_size, self.target_size, device=observations.device)\n",
    "        latents = torch.zeros(batch_size, seq_size, self.latent_size, device=observations.device)\n",
    "        penalties = torch.zeros(batch_size, seq_size, device=observations.device)\n",
    "        for i in range(seq_size):\n",
    "            output, latent, penalty = self._step(observations[:, i], prev_latents)\n",
    "            outputs[:, i] = output\n",
    "            latents[:, i] = latent\n",
    "            penalties[:, i] = penalty\n",
    "            prev_latents = latent\n",
    "        return outputs, latents, penalties\n",
    "\n",
    "    def initial_state(self, batch_size, device=None):\n",
    "        return self.latent_inits.unsqueeze(0).repeat(batch_size, 1).to(device)\n",
    "        \n",
    "\n",
    "obs_size = 2\n",
    "target_size = 1\n",
    "latent_size = 10\n",
    "batch_size = 15\n",
    "\n",
    "ddrnn = DynamicDisRNN(obs_size=obs_size, target_size=target_size, latent_size=latent_size, update_mlp_shape=(latent_size,), choice_mlp_shape=(latent_size,), eval_mode=0, beta_scale=1, activation=nn.ReLU)\n",
    "\n",
    "x = torch.randn(batch_size, obs_size)\n",
    "h = ddrnn.initial_state(batch_size)\n",
    "\n",
    "out, hidden = ddrnn(x, h)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "vrAnalysis",
   "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.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}