{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "99333bfc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import gym\n",
    "import torch \n",
    "import torch.nn as nn\n",
    "from sklearn.decomposition import PCA\n",
    "from sklearn.manifold import TSNE\n",
    "from random import sample\n",
    "from tqdm import tqdm\n",
    "from time import sleep\n",
    "from model import DQN_Agent\n",
    "import numpy as np      \n",
    "from sklearn.neighbors import KNeighborsClassifier\n",
    "from collections import Counter\n",
    "from sklearn.cluster import KMeans, DBSCAN, OPTICS\n",
    "from scipy.spatial.distance import euclidean\n",
    "from copy import deepcopy\n",
    "from sklearn.metrics import silhouette_score\n",
    "from torch.utils.data import TensorDataset, DataLoader\n",
    "\n",
    "from timm.models.layers import trunc_normal_, DropPath\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "b0aa2dd6",
   "metadata": {},
   "outputs": [],
   "source": [
    "NUM_CLASSES = 2\n",
    "NUM_PROTOTYPES = 2\n",
    "LATENT_SIZE = 64\n",
    "BATCH_SIZE = 32\n",
    "NUM_EPOCHS = 5\n",
    "DEVICE = 'cpu'\n",
    "\n",
    "MAX_SAMPLES = 100000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "e8844f94",
   "metadata": {},
   "outputs": [],
   "source": [
    "def viper_sample(obss, acts, qs, xs, is_reweight=True):\n",
    "\t\n",
    "\t\"\"\"\n",
    "\tFunction taken from: https://github.com/obastani/viper\n",
    "\tobservations\n",
    "\tlatent x activations\n",
    "\tq values\n",
    "\tMax num observations\n",
    "\tUniform or sampling\n",
    "\t\"\"\"\n",
    "\t\n",
    "\t# Step 1: Compute probabilities\n",
    "\tps = np.max(qs, axis=1) - np.min(qs, axis=1)\n",
    "\tps = ps / np.sum(ps)\n",
    "\n",
    "\t# Step 2: Sample points\n",
    "\tif is_reweight:\n",
    "\t\t# According to p(s)\n",
    "\t\tidx = np.random.choice(len(obss), size=min(MAX_SAMPLES, np.sum(ps > 0)), p=ps)\n",
    "\telse:\n",
    "\t\t# Uniformly (without replacement)\n",
    "\t\tidx = np.random.choice(len(obss), size=min(MAX_SAMPLES, np.sum(ps > 0)), replace=False)    \n",
    "\n",
    "\t# Step 3: Obtain sampled indices\n",
    "\treturn obss[idx], acts[idx], qs[idx], xs[idx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "96dd86d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Environment and Agent\n",
    "env = gym.make(\"CartPole-v1\").unwrapped\n",
    "input_dim = env.observation_space.shape[0]\n",
    "output_dim = env.action_space.n\n",
    "exp_replay_size = 256\n",
    "agent = DQN_Agent(seed=1423, layer_sizes=[input_dim, 64, output_dim], lr=1e-3,\n",
    "                  sync_freq=5, exp_replay_size=exp_replay_size)\n",
    "agent.load_pretrained_model(\"weights/cartpole-dqn.pth\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "e7205d9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# q_train.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "8d544585",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train = np.load('data/X_train.npy', ).reshape(-1, LATENT_SIZE)\n",
    "a_train = np.load('data/a_train.npy', ).flatten()\n",
    "obs_train = np.load('data/obs_train.npy', ).reshape(-1, 4)\n",
    "q_train = np.load('data/q_train.npy').reshape(-1, 2)\n",
    "\n",
    "\n",
    "\n",
    "obs_train, a_train, q_train, X_train = viper_sample(obs_train,\n",
    "                                                     a_train,\n",
    "                                                     q_train,\n",
    "                                                     X_train,\n",
    "                                                     is_reweight=True)\n",
    "\n",
    "\n",
    "tensor_x = torch.Tensor(X_train)\n",
    "tensor_y = torch.tensor(a_train)\n",
    "\n",
    "train_dataset = TensorDataset(tensor_x, tensor_y)\n",
    "train_loader = DataLoader(train_dataset, shuffle=True, batch_size=BATCH_SIZE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "13c7caae",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(100000, 64)"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X_train.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "096db306",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(100000, 2)"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "q_train.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "0fdc0eff",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_loader(model, loader, cce_loss):\n",
    "    model.eval()\n",
    "    total_correct = 0\n",
    "    total_loss = 0\n",
    "    total = 0\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        for i, data in enumerate(loader):\n",
    "            imgs, labels = data\n",
    "            imgs, labels = imgs.to(DEVICE), labels.to(DEVICE)\n",
    "            logits, _ = model(imgs)\n",
    "            loss = cce_loss(logits, labels)\n",
    "            preds = torch.argmax(logits, dim=1)\n",
    "            total_correct += sum(preds == labels).item()\n",
    "            total += len(preds)\n",
    "            total_loss += loss.item()\n",
    "            \n",
    "    model.train()\n",
    "    \n",
    "    return (total_correct / total) * 100, (total_loss / len(loader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "797966df",
   "metadata": {},
   "outputs": [],
   "source": [
    "def clust_loss(x, y, p, criterion):\n",
    "            \n",
    "    for idx, i in enumerate(Counter(y.cpu().numpy()).keys()):\n",
    "        x_sub = x[y==i]\n",
    "        target = p[i].repeat(len(x_sub), 1) \n",
    "        \n",
    "        if idx == 0:\n",
    "            loss = criterion(x_sub, target) \n",
    "        else:\n",
    "            loss += criterion(x_sub, target)\n",
    "        \n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "902e0bd4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def sep_loss(x, y, p, criterion):\n",
    "        \n",
    "    for idx, i in enumerate(Counter(y.cpu().numpy()).keys()):\n",
    "        x_sub = x[y==i]\n",
    "        target = p[i].repeat(len(x_sub), 1) \n",
    "        \n",
    "        if idx == 0:\n",
    "            loss = criterion(x_sub, target) \n",
    "        else:\n",
    "            loss += criterion(x_sub, target)\n",
    "    \n",
    "    return -loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "id": "35811975",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Wrapper(nn.Module):\n",
    "\n",
    "    def __init__(self):\n",
    "        super(Wrapper, self).__init__()\n",
    "        protos = (torch.randn([NUM_PROTOTYPES, LATENT_SIZE]) )\n",
    "        self.prototypes = nn.Parameter(protos, requires_grad=True)\n",
    "        self.epsilon = 1e-10\n",
    "        self.maxpool = torch.nn.MaxPool1d(NUM_PROTOTYPES)\n",
    "        self.linear = nn.Linear(NUM_PROTOTYPES, NUM_CLASSES, bias=False) \n",
    "        self.__make_linear_weights()\n",
    "        self.main = nn.Sequential(\n",
    "            nn.Linear(LATENT_SIZE, LATENT_SIZE),\n",
    "            nn.BatchNorm1d(LATENT_SIZE),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(0.5),\n",
    "            nn.Linear(LATENT_SIZE, LATENT_SIZE),\n",
    "            nn.BatchNorm1d(LATENT_SIZE),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(0.5),\n",
    "            nn.Linear(LATENT_SIZE, LATENT_SIZE),\n",
    "        )\n",
    "\n",
    "    def __proto_layer_l2(self, x, b_size):\n",
    "        output = list()\n",
    "        p = self.prototypes.T.view(1, LATENT_SIZE, NUM_PROTOTYPES).tile(b_size, 1, 1).to(DEVICE) \n",
    "        c = x.view(b_size, LATENT_SIZE, 1).tile(1, 1, NUM_PROTOTYPES).to(DEVICE)            \n",
    "        l2s = ( (c - p)**2 ).sum(axis=1).to(DEVICE)        \n",
    "        act = torch.log( (l2s + 1. ) / (l2s + self.epsilon) ).to(DEVICE)   \n",
    "        return act\n",
    "    \n",
    "    def __make_linear_weights(self):\n",
    "        \n",
    "        prototype_class_identity = torch.zeros(NUM_PROTOTYPES, NUM_CLASSES)\n",
    "        num_prototypes_per_class = NUM_PROTOTYPES // NUM_CLASSES\n",
    "        \n",
    "        for j in range(NUM_PROTOTYPES):\n",
    "            prototype_class_identity[j, j // num_prototypes_per_class] = 1\n",
    "            \n",
    "        positive_one_weights_locations = torch.t(prototype_class_identity)\n",
    "        negative_one_weights_locations = 1 - positive_one_weights_locations\n",
    "\n",
    "        incorrect_strength = -0.5\n",
    "        correct_class_connection = 1\n",
    "        incorrect_class_connection = incorrect_strength\n",
    "        self.linear.weight.data.copy_(\n",
    "            correct_class_connection * positive_one_weights_locations\n",
    "            + incorrect_class_connection * negative_one_weights_locations)      \n",
    "            \n",
    "    def forward(self, x): \n",
    "\n",
    "        b_size = x.shape[0]\n",
    "\n",
    "        # Transform\n",
    "        x = self.main(x)\n",
    "\n",
    "        # Prototype layer\n",
    "        ps = self.__proto_layer_l2(x, b_size)\n",
    "        \n",
    "        # Linear layer\n",
    "        logits = self.linear(ps)\n",
    "\n",
    "        return logits, x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a4cdc7b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a694e492",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "id": "73d72ada",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Wrapper()\n",
    "model.drop_path_bool = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "id": "6bbc78f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "instances, labels = next(iter(train_loader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "id": "598ed4f2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([32, 64])"
      ]
     },
     "execution_count": 90,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "instances.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31d69be8",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "id": "92ed0f93",
   "metadata": {},
   "outputs": [],
   "source": [
    "# model(instances)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc571ca2",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "id": "9f619af4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# model.linear.weight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "id": "f57b8a52",
   "metadata": {},
   "outputs": [],
   "source": [
    "# model.linear.requires_grad_(False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "id": "e690f320",
   "metadata": {},
   "outputs": [],
   "source": [
    "# model.linear.weight.requires_grad = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "id": "3c72b217",
   "metadata": {},
   "outputs": [],
   "source": [
    "# model.linear.weight.requires_grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "id": "5135d1d0",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Accuracy: 47.745 %     Train Loss: 0.6932839231681823\n",
      "CCE Loss: 0.3438516201019287\n",
      "Train Accuracy: 84.574 %     Train Loss: 0.26175077248811723\n",
      "CCE Loss: 0.3120229428625107\n",
      "Train Accuracy: 83.478 %     Train Loss: 0.2768279695367813\n",
      "CCE Loss: 0.30981194732427597\n",
      "Train Accuracy: 86.168 %     Train Loss: 0.26253392823696137\n",
      "CCE Loss: 0.30679888998508453\n",
      "Train Accuracy: 85.79299999999999 %     Train Loss: 0.2823270068383217\n",
      "CCE Loss: 0.30230977653503416\n",
      "Train Accuracy: 84.77900000000001 %     Train Loss: 0.2695656241822243\n",
      "CCE Loss: 0.3013905094122887\n",
      "Train Accuracy: 84.795 %     Train Loss: 0.26217432779312133\n",
      "CCE Loss: 0.2974299555468559\n",
      "Train Accuracy: 84.615 %     Train Loss: 0.2603268819665909\n",
      "CCE Loss: 0.2992459752345085\n",
      "Train Accuracy: 84.68299999999999 %     Train Loss: 0.2503831456899643\n",
      "CCE Loss: 0.2969539540195465\n",
      "Train Accuracy: 84.528 %     Train Loss: 0.26323004781246184\n",
      "CCE Loss: 0.29479945239543914\n",
      "Train Accuracy: 84.711 %     Train Loss: 0.2513738033127785\n",
      "CCE Loss: 0.29544061876296995\n",
      "Train Accuracy: 84.649 %     Train Loss: 0.2560209237885475\n",
      "CCE Loss: 0.29446362636327744\n",
      "Train Accuracy: 84.817 %     Train Loss: 0.2516536581873894\n",
      "CCE Loss: 0.29314588712692263\n",
      "Train Accuracy: 84.639 %     Train Loss: 0.2541048629236221\n",
      "CCE Loss: 0.2907930896043777\n",
      "Train Accuracy: 84.68299999999999 %     Train Loss: 0.25521054176568986\n",
      "CCE Loss: 0.29186522459983827\n",
      "Train Accuracy: 85.039 %     Train Loss: 0.2527415077114105\n",
      "CCE Loss: 0.29108982365846636\n",
      "Train Accuracy: 85.15700000000001 %     Train Loss: 0.24974270120978356\n",
      "CCE Loss: 0.289241896841526\n",
      "Train Accuracy: 85.71600000000001 %     Train Loss: 0.2649822499036789\n",
      "CCE Loss: 0.28969356070995333\n",
      "Train Accuracy: 85.197 %     Train Loss: 0.25282575188159945\n",
      "CCE Loss: 0.2876601495695114\n",
      "Train Accuracy: 85.202 %     Train Loss: 0.24864196002960204\n",
      "CCE Loss: 0.28621785912036896\n"
     ]
    }
   ],
   "source": [
    "cce_loss = nn.CrossEntropyLoss()\n",
    "mse_loss = nn.MSELoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.01)\n",
    "scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.9)\n",
    "best_acc = 0\n",
    "\n",
    "# Freeze prototype learning and use full model (inc. transformation)\n",
    "model.train()\n",
    "model.prototypes.requires_grad = True\n",
    "model.linear.weight.requires_grad = False\n",
    "prototypes = model.prototypes\n",
    "\n",
    "\n",
    "for epoch in range(20):\n",
    "    \n",
    "    running_loss = 0\n",
    "    \n",
    "    model.eval()\n",
    "    train_acc, train_loss = evaluate_loader(model, train_loader, cce_loss)\n",
    "    model.train()\n",
    "    print(\"Train Accuracy:\", train_acc, \"%\", \"    Train Loss:\", train_loss)\n",
    "    \n",
    "    if train_acc > best_acc:\n",
    "        torch.save(model.state_dict(), 'weights/wrapper_pre_projection.pth')\n",
    "        best_acc = train_acc\n",
    "    \n",
    "    for instances, labels in train_loader:\n",
    "        \n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        instances, labels = instances.to(DEVICE), labels.to(DEVICE)\n",
    "        logits, x = model(instances)\n",
    "        \n",
    "        loss = cce_loss(logits, labels) \n",
    "                \n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        running_loss += loss.item()\n",
    "        \n",
    "        \n",
    "    print(\"CCE Loss:\", running_loss / len(train_loader))\n",
    "    scheduler.step()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b87db3cd",
   "metadata": {},
   "source": [
    "## Project Prototypes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "id": "79aa58e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# model.linear.weight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "id": "0aa944ef",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 98,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = Wrapper().eval()\n",
    "model.drop_path_bool = False\n",
    "model.load_state_dict(torch.load('weights/wrapper_pre_projection.pth'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "id": "a1c2548c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(54.432, 0.6923979721069335)"
      ]
     },
     "execution_count": 99,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Accuracy before projection\n",
    "evaluate_loader(model, train_loader, cce_loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "id": "aeb2806b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████| 100000/100000 [00:13<00:00, 7682.35it/s]\n"
     ]
    }
   ],
   "source": [
    "trans_x = list()\n",
    "model.eval()\n",
    "\n",
    "with torch.no_grad():\n",
    "    for i in tqdm(range(len(X_train))):\n",
    "        img = X_train[i]\n",
    "        _, x = model(  torch.tensor(img, dtype=torch.float32).view(1, -1)  )\n",
    "        trans_x.append(x[0].tolist())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "id": "598381de",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(100000, 64)"
      ]
     },
     "execution_count": 101,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trans_x = np.array(trans_x)\n",
    "trans_x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "id": "64d1bfbd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.neighbors import KNeighborsClassifier\n",
    "\n",
    "knn = KNeighborsClassifier(algorithm='brute')\n",
    "knn.fit(trans_x, list(range(len(trans_x))))\n",
    "dist, p_idxs = knn.kneighbors(X=model.prototypes.detach().numpy(), n_neighbors=1, return_distance=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "id": "f64ecb8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "p_idxs = np.array(p_idxs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "id": "ed18afd8",
   "metadata": {},
   "outputs": [],
   "source": [
    "projected_prototypes = trans_x[p_idxs.flatten()]\n",
    "proto_images = obs_train[p_idxs.flatten()]\n",
    "proto_before_trans = X_train[p_idxs.flatten()]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 105,
   "id": "151177c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.prototypes = torch.nn.Parameter(torch.tensor(projected_prototypes, dtype=torch.float32))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 106,
   "id": "e0d06415",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(55.306, 0.5885979642438889)"
      ]
     },
     "execution_count": 106,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Accuracy after projection\n",
    "evaluate_loader(model, train_loader, cce_loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 107,
   "id": "ba85265f",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model.state_dict(), 'weights/wrapper_post_projection.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 108,
   "id": "1bdc8d43",
   "metadata": {},
   "outputs": [],
   "source": [
    "# dist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "id": "eb970dca",
   "metadata": {},
   "outputs": [],
   "source": [
    "# proto_images"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5c348e9",
   "metadata": {},
   "source": [
    "\n",
    "    | Num | Observation           | Min                 | Max               |\n",
    "    |-----|-----------------------|---------------------|-------------------|\n",
    "    | 0   | Cart Position         | -4.8                | 4.8               |\n",
    "    | 1   | Cart Velocity         | -Inf                | Inf               |\n",
    "    | 2   | Pole Angle            | ~ -0.418 rad (-24°) | ~ 0.418 rad (24°) |\n",
    "    | 3   | Pole Angular Velocity | -Inf                | Inf               |\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "484a46fa",
   "metadata": {},
   "source": [
    "## Now Train Transformation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 112,
   "id": "4d9be2b5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 112,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = Wrapper().eval()\n",
    "model.drop_path_bool = False\n",
    "model.load_state_dict(torch.load('weights/wrapper_post_projection.pth'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4da98cb4",
   "metadata": {},
   "source": [
    "# Evaluate Model\n",
    "* Plot\n",
    "* Sillhouette score\n",
    "* Distance to clusters\n",
    "* Performace on 1000 simuluations\n",
    "\n",
    "### Variables\n",
    "* Transformer component: This helps A LOT\n",
    "* Num prototypes: Not that important so far, less is better in general\n",
    "* Extra loss functions: Don't help at all, CCE does the job 100%\n",
    "* Training data sampling: Todo\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "483185bc",
   "metadata": {},
   "source": [
    "## Plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 113,
   "id": "87975526",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(55.306, 0.5885979639339447)"
      ]
     },
     "execution_count": 113,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Accuracy after final train\n",
    "evaluate_loader(model, train_loader, cce_loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 114,
   "id": "e7950f14",
   "metadata": {},
   "outputs": [],
   "source": [
    "# model.linear.weight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 115,
   "id": "a82a535f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████| 100000/100000 [00:13<00:00, 7656.50it/s]\n"
     ]
    }
   ],
   "source": [
    "trans_x = list()\n",
    "model.eval()\n",
    "\n",
    "with torch.no_grad():\n",
    "    for i in tqdm(range(len(X_train))):\n",
    "        img = X_train[i]\n",
    "        _, x = model(  torch.tensor(img, dtype=torch.float32).view(1, -1)  )\n",
    "        trans_x.append(x[0].tolist())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 116,
   "id": "13d84b8c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(100000, 64)"
      ]
     },
     "execution_count": 116,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trans_x = np.array(trans_x)\n",
    "trans_x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 117,
   "id": "c846b8a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "rand_idxs = np.random.randint(0, len(a_train), 3000)\n",
    "temp_y = a_train[rand_idxs]\n",
    "temp_x = trans_x[rand_idxs]\n",
    "# temp_x = X_train[rand_idxs]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 118,
   "id": "1644496a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[2, 3]"
      ]
     },
     "execution_count": 118,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[i+NUM_CLASSES for i in range(NUM_PROTOTYPES)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 119,
   "id": "99d5a250",
   "metadata": {},
   "outputs": [],
   "source": [
    "temp_x = np.append(temp_x, model.prototypes.clone().detach().numpy(), axis=0)\n",
    "temp_y = np.append(temp_y, np.array( [i+NUM_CLASSES for i in range(NUM_PROTOTYPES)] ), axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 120,
   "id": "defccfa1",
   "metadata": {},
   "outputs": [],
   "source": [
    "tsne = PCA(n_components=2)\n",
    "emb_x = tsne.fit_transform(temp_x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 121,
   "id": "de165383",
   "metadata": {},
   "outputs": [],
   "source": [
    "# tsne = TSNE(n_components=2)\n",
    "# emb_x = tsne.fit_transform(temp_x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 122,
   "id": "793949cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 123,
   "id": "b925d304",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.DataFrame()\n",
    "df['a'] = temp_y\n",
    "df['x'] = emb_x.T[0]\n",
    "df['y'] = emb_x.T[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 124,
   "id": "0678da3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "df.a = df.a.astype('category')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 125,
   "id": "3145f9f8",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/eoinkenny/opt/anaconda3/envs/rl_env/lib/python3.9/site-packages/seaborn/rcmod.py:82: DeprecationWarning: distutils Version classes are deprecated. Use packaging.version instead.\n",
      "  if LooseVersion(mpl.__version__) >= \"3.0\":\n",
      "/Users/eoinkenny/opt/anaconda3/envs/rl_env/lib/python3.9/site-packages/setuptools/_distutils/version.py:351: DeprecationWarning: distutils Version classes are deprecated. Use packaging.version instead.\n",
      "  other = LooseVersion(other)\n"
     ]
    }
   ],
   "source": [
    "import seaborn as sns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 126,
   "id": "f8c0d423",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<AxesSubplot:xlabel='x', ylabel='y'>"
      ]
     },
     "execution_count": 126,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 25004 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "sns.scatterplot(x='x', y='y', hue='a', data=df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 127,
   "id": "4df025da",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<AxesSubplot:xlabel='x', ylabel='y'>"
      ]
     },
     "execution_count": 127,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 25004 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "sns.scatterplot(x='x', y='y', hue='a', data=df)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4eb0f59b",
   "metadata": {},
   "source": [
    "## Sillhouete Score\n",
    "Re-label dataset with each $x$ nearest prototype"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 128,
   "id": "0789bfd3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "KNeighborsClassifier(algorithm='brute', n_neighbors=1)"
      ]
     },
     "execution_count": 128,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from sklearn.neighbors import KNeighborsClassifier\n",
    "\n",
    "knn = KNeighborsClassifier(algorithm='brute', n_neighbors=1)\n",
    "knn.fit(model.prototypes.detach().numpy(), list(range(NUM_PROTOTYPES)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 129,
   "id": "37727a71",
   "metadata": {},
   "outputs": [],
   "source": [
    "labels = knn.predict(temp_x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 130,
   "id": "51387df4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.5583284946541863"
      ]
     },
     "execution_count": 130,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Transformation\n",
    "silhouette_score(temp_x[:-NUM_PROTOTYPES], labels[:-NUM_PROTOTYPES])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6747bfbe",
   "metadata": {},
   "source": [
    "## Simulation Score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 131,
   "id": "364541c9",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████| 1000/1000 [00:02<00:00, 376.07it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "average reward per episode : 12.01\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "reward_arr = []\n",
    "for i in tqdm(range(1000)):\n",
    "    obs, done, rew = env.reset(), False, 0\n",
    "    count = 0\n",
    "    while not done and count < 200:\n",
    "        _, latent_x, _ = agent.get_action(obs, env.action_space.n, epsilon=0)\n",
    "        A, _ = model( latent_x.view(1,-1) )\n",
    "        A = torch.argmax(A).item()\n",
    "        obs, reward, done, info = env.step(A)\n",
    "        rew += reward\n",
    "        count += 1\n",
    "    reward_arr.append(count)\n",
    "\n",
    "print(\"average reward per episode :\", sum(reward_arr) / len(reward_arr))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abb4e57b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68e935d7",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91fd9f53",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9ff43d0",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05f9b988",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47b4c47b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6d91870",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "0a8dfd47",
   "metadata": {},
   "source": [
    "## Contribution of Each Feature"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "46a3c82c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2098755, 4)"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "obs_train.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "5bb5aa31",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2098755, 64)"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_data_x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "0387db2e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2098755,)"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a_train.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "862921c3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([64, 4])"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "weights = agent.q_net.fc1.weight\n",
    "weights.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "a0806be0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.2815, grad_fn=<DivBackward0>)\n",
      "tensor(0.5814, grad_fn=<DivBackward0>)\n",
      "tensor(0.7222, grad_fn=<DivBackward0>)\n",
      "tensor(0.4661, grad_fn=<DivBackward0>)\n"
     ]
    }
   ],
   "source": [
    "for i in range(4):\n",
    "    print(  sum(abs(weights.T[i])) / len(weights) )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "65082c58",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.1278256763537884\n",
      "0.142202865062279\n",
      "0.017667636993578387\n",
      "0.20191600174775762\n"
     ]
    }
   ],
   "source": [
    "for i in range(4):\n",
    "    print(  sum(abs(obs_train.T[i])) / len(obs_train) )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "a19c4fac",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total Contribution Average: tensor(0.0792, grad_fn=<MulBackward0>)\n",
      "Total Contribution Average: tensor(0.3380, grad_fn=<MulBackward0>)\n",
      "Total Contribution Average: tensor(0.5216, grad_fn=<MulBackward0>)\n",
      "Total Contribution Average: tensor(0.2172, grad_fn=<MulBackward0>)\n"
     ]
    }
   ],
   "source": [
    "data = torch.zeros(4)\n",
    "\n",
    "for i in range(4):\n",
    "    w = sum(abs(weights.T[i])) / len(weights)\n",
    "    x = sum(abs(weights.T[i])) / len(weights)\n",
    "    print(\"Total Contribution Average:\", w*x)\n",
    "    data[i] = w*x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "272659c7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0.2000, 0.2591, 0.3113, 0.2296], grad_fn=<SoftmaxBackward0>)"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.softmax(data, dim=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "b9bf0cba",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2, 4)"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "proto_images.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "7390302b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([64, 4])"
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "weights.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "66b4fb7e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([-0.04302227,  0.11650063, -0.15574128, -0.22569513, -0.16889343,\n",
       "       -0.10001172, -0.01422224,  0.01087987,  0.13577486, -0.03560081,\n",
       "        0.02575449, -0.05889751, -0.10716026, -0.07144366, -0.1550035 ,\n",
       "       -0.19094365, -0.08684268,  0.24042794, -0.14758203, -0.18864988,\n",
       "        0.08164098,  0.05242967, -0.08050377,  0.12700218,  0.00764638,\n",
       "       -0.05004192,  0.01025908, -0.10316804,  0.09904955,  0.16939828,\n",
       "       -0.08770164, -0.00985877,  0.1919707 , -0.05791862, -0.15131058,\n",
       "       -0.00931485, -0.06945472, -0.13663107, -0.05459756, -0.00058078,\n",
       "        0.03391863, -0.05229604,  0.00731769,  0.08461743,  0.01973061,\n",
       "       -0.02384995,  0.00583763, -0.15357726, -0.0069335 , -0.01094437,\n",
       "        0.01309674,  0.08829818, -0.04161307, -0.06451988,  0.05138354,\n",
       "       -0.17514603,  0.086413  ,  0.01803717, -0.05200949,  0.02989575,\n",
       "        0.03192458, -0.01387842,  0.14125147, -0.05025253])"
      ]
     },
     "execution_count": 59,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.matmul(weights.cpu().detach().numpy(), proto_images[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "13323de3",
   "metadata": {},
   "outputs": [],
   "source": [
    "feature_weights = torch.zeros(4)\n",
    "\n",
    "for i in range(4):\n",
    "    w = sum(abs(weights.T[i])) / len(weights)\n",
    "    feature_weights[i] = w"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "id": "5e513a1f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[-0.06225981 -0.00532342 -0.00131571 -0.01817009]\n",
      "[0.01905259 0.00123403 0.00058565 0.01037095]\n"
     ]
    }
   ],
   "source": [
    "for i in range(len(proto_images)):\n",
    "    print( proto_images[i] * feature_weights.detach().numpy() )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "969ccf07",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "rl_env",
   "language": "python",
   "name": "rl_env"
  },
  "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
