AntBO/Backup SQL.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "936098bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch \n",
    "import pandas as pd\n",
    "from einops import rearrange, reduce, asnumpy, parse_shape\n"
   ]
  },
  {
   "cell_type": "raw",
   "id": "c9c17df6",
   "metadata": {},
   "source": [
    "x = np.arange(-1,1,2./11)\n",
    "# x = x - 1 \n",
    "y = np.tanh(x)\n",
    "plt.plot(x,y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "35cbe07c",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SeqOptimisation:\n",
    "    def __init__(self, seqSize=10, numberVars=20, easy=True):\n",
    "        self.seqSize = seqSize\n",
    "        self.hseqSize = seqSize/2.0\n",
    "        self.numberVars = numberVars\n",
    "        self.easy = easy\n",
    "        if self.easy:\n",
    "            # Search for same number repeated\n",
    "            self.target_num1 = np.random.randint(0, numberVars, 1).item()\n",
    "#             self.target_num2 = np.random.randint(0, numberVars, 1).item()\n",
    "            self.target_seq = np.ones(seqSize)*self.target_num1\n",
    "#             self.target_seq = np.concatenate([np.ones(seqSize//2)*self.target_num1, np.ones(int(self.seqSize-seqSize//2))*self.target_num2])\n",
    "        else:    \n",
    "            #TODO: Need to finish implement this harder task \n",
    "            # Also can add to search for seq with target subpattern\n",
    "            self.target_seq = np.random.randint(0, numberVars, seqSize)\n",
    "        self.target_seq = torch.IntTensor(self.target_seq).squeeze().long()\n",
    "        self.br = None\n",
    "        self.bseq = None\n",
    "        print(f'Target Sequence is {self.target_seq}')\n",
    "        self.evls = 0\n",
    "    def reward(self, seq):\n",
    "        seq = seq.squeeze().long()\n",
    "        r = -((self.target_seq!=seq).sum().item())/(self.seqSize)\n",
    "        r = np.tanh(r)\n",
    "        if self.br is None:\n",
    "            self.br = r\n",
    "            self.bseq = seq.squeeze()\n",
    "        if r>self.br:\n",
    "            if r==0:\n",
    "                print(f'Found Optimal Sequence {seq} at evaluation {self.evls}')\n",
    "            self.br = r\n",
    "            self.bseq = seq.squeeze()\n",
    "        self.evls += 1            \n",
    "        return r, False\n",
    "            \n",
    "    def breward(self, seq):\n",
    "        return -((self.target_seq!=seq).sum(1))/(self.seqSize)\n",
    "    \n",
    "    def get_best(self):\n",
    "        print(f'with hamming distance {(self.target_seq!=self.bseq).sum()} and reward {self.br}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2a71d8a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import random\n",
    "import numpy as np\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "from collections import namedtuple, deque\n",
    "from itertools import count\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "925cf0d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "device = 'cpu'\n",
    "# torch.cuda.set_device('cuda')\n",
    "torch.set_default_tensor_type(torch.FloatTensor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "73c1dad2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.autograd import Variable\n",
    "class LayerNorm(nn.Module):\n",
    "    \"Construct a layernorm module (See citation for details).\"\n",
    "    def __init__(self, features, eps=1e-6):\n",
    "        super(LayerNorm, self).__init__()\n",
    "        self.a_2 = nn.Parameter(torch.ones(features))\n",
    "        self.b_2 = nn.Parameter(torch.zeros(features))\n",
    "        self.eps = eps\n",
    "\n",
    "    def forward(self, x):\n",
    "        mean = x.mean(-1, keepdim=True)\n",
    "        std = x.std(-1, keepdim=True)\n",
    "        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2\n",
    "    \n",
    "    \n",
    "class SublayerConnection(nn.Module):\n",
    "    \"\"\"\n",
    "    A residual connection followed by a layer norm.\n",
    "    Note for code simplicity the norm is first as opposed to last.\n",
    "    \"\"\"\n",
    "    def __init__(self, size, dropout):\n",
    "        super(SublayerConnection, self).__init__()\n",
    "        self.norm = LayerNorm(size)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x, sublayer):\n",
    "        \"Apply residual connection to any sublayer with the same size.\"\n",
    "        return x + self.dropout(sublayer(self.norm(x)))\n",
    "    \n",
    "    \n",
    "class PositionalEncoding(nn.Module):\n",
    "\n",
    "    def __init__(self, d_model: int, max_len: int = 5000):\n",
    "        super().__init__()\n",
    "        position = torch.arange(max_len).unsqueeze(1)\n",
    "        div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))\n",
    "        pe = torch.zeros(max_len, 1, d_model)\n",
    "        pe[:, 0, 0::2] = torch.sin(position * div_term)\n",
    "        pe[:, 0, 1::2] = torch.cos(position * div_term)\n",
    "        self.register_buffer('pe', pe)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            x: Tensor, shape [seq_len, batch_size, embedding_dim]\n",
    "        \"\"\"\n",
    "        x = rearrange(x, 'b s e -> s b e')\n",
    "        x = x + self.pe[:x.size(0)]\n",
    "        return rearrange(x, 's b e -> b s e')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "bc26fc4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def attention(q, k, v, d_k, dropout=None):\n",
    "    \n",
    "    scores = torch.matmul(q, k.transpose(-2, -1)) /  math.sqrt(d_k)\n",
    "    scores = F.softmax(scores, dim=-1)\n",
    "\n",
    "    if dropout is not None:\n",
    "        scores = dropout(scores)\n",
    "        \n",
    "    output = torch.matmul(scores, v)\n",
    "    return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "5197591c",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadAttention(nn.Module):\n",
    "    def __init__(self, d_model, dropout=0.1):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.d_model = d_model\n",
    "        self.d_k = d_model\n",
    "        self.h = 1\n",
    "        \n",
    "        self.q_linear = nn.Linear(d_model, d_model)\n",
    "        self.v_linear = nn.Linear(d_model, d_model)\n",
    "        self.k_linear = nn.Linear(d_model, d_model)\n",
    "        self.out = nn.Linear(d_model, d_model)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "    def forward(self, q, k, v, mask=None):\n",
    "        \n",
    "        bs = q.size(0)\n",
    "        \n",
    "        # perform linear operation and split into h heads\n",
    "        \n",
    "        k = self.k_linear(k).view(bs, -1, 1, self.d_k)\n",
    "        q = self.q_linear(q).view(bs, -1, 1, self.d_k)\n",
    "        v = self.v_linear(v).view(bs, -1, 1, self.d_k)\n",
    "        \n",
    "        # transpose to get dimensions bs * h * sl * d_model\n",
    "       \n",
    "        k = k.transpose(1,2)\n",
    "        q = q.transpose(1,2)\n",
    "        v = v.transpose(1,2)\n",
    "        # calculate attention using function we will define next\n",
    "        scores = attention(q, k, v, self.d_k, self.dropout)\n",
    "        \n",
    "        # concatenate heads and put through final linear layer\n",
    "        concat = scores.transpose(1,2).contiguous()\\\n",
    "        .view(bs, -1, self.d_model)\n",
    "        \n",
    "        output = self.out(concat)\n",
    "    \n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "2bc566f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import gym\n",
    "from gym import spaces\n",
    "import numpy as np\n",
    "import os\n",
    "import subprocess\n",
    "import pandas as pd\n",
    "import csv\n",
    "\n",
    "from itertools import groupby\n",
    "import re\n",
    "import time\n",
    "from torch.distributions import Categorical\n",
    "from itertools import count\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "75bf2370",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TransformerNN(nn.Module):\n",
    "    def __init__(self, numberVars=20, seqSize=11, outputs=1, d_model=32, device='cpu', dropout=0.1, PolicyGradients=False):\n",
    "        super(TransformerNN, self).__init__()\n",
    "#         self.dropout1 = nn.Dropout(dropout)\n",
    "#         self.dropout2 = nn.Dropout(dropout)\n",
    "        if PolicyGradients:\n",
    "            self.saved_log_probs = []\n",
    "            self.rewards = []\n",
    "#             encoder_layer = nn.TransformerEncoderLayer(d_model=32, nhead=8)\n",
    "#             transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=1)\n",
    "            \n",
    "#         else:\n",
    "        self.PEncoder = PositionalEncoding(d_model, max_len=seqSize)\n",
    "        self.Embedder = nn.Embedding(numberVars+1, d_model, max_norm=True)\n",
    "        self.numberVars = numberVars\n",
    "        self.d_model = d_model\n",
    "        self.lin1 = nn.Linear(d_model, d_model)\n",
    "        self.head = nn.Linear(d_model*seqSize, outputs)\n",
    "        self.device = device\n",
    "        self.numberVars = numberVars\n",
    "        self.PolicyGradients = PolicyGradients\n",
    "\n",
    "#             for p in self.parameters():\n",
    "#                 print(p.shape)\n",
    "#                 if p.dim() > 1:\n",
    "#                     nn.init.xavier_uniform(p)\n",
    "\n",
    "        \n",
    "        \n",
    "    def forward(self, x):\n",
    "        assert x.max().item()<=(self.numberVars+1), print('Embedding more variables than we should')\n",
    "        # [Batch, SeqLengthxNumVars]\n",
    "        x = self.Embedder(x).float()\n",
    "        x = self.PEncoder(x)\n",
    "        x = self.lin1(x)\n",
    "        x = F.relu(x)\n",
    "        x = x.flatten(1)\n",
    "        x = self.head(x)\n",
    "        if self.PolicyGradients:\n",
    "            x = F.softmax(x, dim=1)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "ca4d4cde",
   "metadata": {},
   "outputs": [],
   "source": [
    "Transition = namedtuple('Transition',\n",
    "                        ('state', 'action', 'next_state', 'reward'))\n",
    "\n",
    "def val_StrCritic(task, policy_net, NumClass, seqSize):\n",
    "    with torch.no_grad():\n",
    "        criterion = nn.L1Loss()\n",
    "        TX = torch.randint(NumClass,(1000,seqSize))\n",
    "        TestY = task.breward(TX)\n",
    "        preds = policy_net(TX)\n",
    "        loss = criterion(preds, TestY)\n",
    "    return loss\n",
    "\n",
    "\n",
    "CombTransition = namedtuple('CombTransition',\n",
    "                        ('structure', 'reward'))\n",
    "\n",
    "class CombReplayMemory(object):\n",
    "\n",
    "    def __init__(self, capacity):\n",
    "        self.memory = deque([],maxlen=capacity)\n",
    "\n",
    "    def push(self, *args):\n",
    "        \"\"\"Save a Combinatorial Optimisation transition\"\"\"\n",
    "        self.memory.append(CombTransition(*args))\n",
    "\n",
    "    def sample(self, batch_size):\n",
    "        return random.sample(self.memory, batch_size)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.memory)\n",
    "\n",
    "class ReplayMemory(object):\n",
    "\n",
    "    def __init__(self, capacity):\n",
    "        self.memory = deque([],maxlen=capacity)\n",
    "\n",
    "    def push(self, *args):\n",
    "        \"\"\"Save a transition\"\"\"\n",
    "        self.memory.append(Transition(*args))\n",
    "\n",
    "    def sample(self, batch_size):\n",
    "        return random.sample(self.memory, batch_size)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.memory)\n",
    "    \n",
    "def getBestMemory(memory):\n",
    "    transitions = memory.memory\n",
    "    batch = CombTransition(*zip(*transitions))\n",
    "    structure_batch = torch.cat(batch.structure)\n",
    "    reward_batch = torch.cat(batch.reward)\n",
    "    rewards_idmax = reward_batch.squeeze().argmax()\n",
    "    return structure_batch[rewards_idmax][None, :]        \n",
    "    \n",
    "def optimize_Str(BATCH_SIZE, policy_net, optimizer, memory, gradient_steps=2, seqS=11):\n",
    "    if len(memory) < BATCH_SIZE:\n",
    "        return None\n",
    "    for _ in range(gradient_steps):\n",
    "        transitions = memory.sample(BATCH_SIZE)\n",
    "        batch = CombTransition(*zip(*transitions))\n",
    "        structure_batch = torch.cat(batch.structure)\n",
    "        reward_batch = torch.cat(batch.reward)   \n",
    "        # for each batch state according to policy_net\n",
    "        expected_structure_values = policy_net(structure_batch)\n",
    "        # Compute the expected Q values\n",
    "\n",
    "        # Compute Huber loss\n",
    "        criterion = nn.SmoothL1Loss()\n",
    "        loss = criterion(reward_batch.squeeze(), expected_structure_values.squeeze())\n",
    "\n",
    "        # Optimize the model\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        torch.nn.utils.clip_grad_norm_(policy_net.parameters(), 100)\n",
    "\n",
    "        optimizer.step()\n",
    "    return loss.detach().item()\n",
    "    \n",
    "\n",
    "def optimize_SubStr(BATCH_SIZE, policy_net, optimizer, memory, gradient_steps=1, model='QL'):\n",
    "    GAMMA = 0.99\n",
    "    \n",
    "    if len(memory) < BATCH_SIZE:\n",
    "        return\n",
    "    for _ in range(gradient_steps):\n",
    "        transitions = memory.sample(BATCH_SIZE)\n",
    "        # Transpose the batch (see https://stackoverflow.com/a/19343/3343043 for\n",
    "        # detailed explanation). This converts batch-array of Transitions\n",
    "        # to Transition of batch-arrays.\n",
    "        batch = Transition(*zip(*transitions))\n",
    "\n",
    "        # Compute a mask of non-final states and concatenate the batch elements\n",
    "        # (a final state would've been the one after which simulation ended)\n",
    "        non_final_mask = torch.tensor(tuple(map(lambda s: s is not None,\n",
    "                                              batch.next_state)), device=device, dtype=torch.bool)\n",
    "        non_final_next_states = torch.cat([s for s in batch.next_state\n",
    "                                                    if s is not None])\n",
    "        state_batch = torch.cat(batch.state)\n",
    "        action_batch = torch.cat(batch.action)\n",
    "        reward_batch = torch.cat(batch.reward)\n",
    "        \n",
    "\n",
    "        # Compute Q(s_t, a) - the model computes Q(s_t), then we select the\n",
    "        # columns of actions taken. These are the actions which would've been taken\n",
    "        # for each batch state according to policy_net\n",
    "        state_action_values = policy_net(state_batch).gather(1, action_batch)\n",
    "\n",
    "        # Compute V(s_{t+1}) for all next states.\n",
    "        # Expected values of actions for non_final_next_states are computed based\n",
    "        # on the \"older\" target_net; selecting their best reward with max(1)[0].\n",
    "        # This is merged based on the mask, such that we'll have either the expected\n",
    "        # state value or 0 in case the state was final.\n",
    "        next_state_values = torch.zeros(BATCH_SIZE, device=device)\n",
    "        #Double Q Learning\n",
    "    #     next_state_values[non_final_mask] = target_net(non_final_next_states).max(1)[0].detach()\n",
    "        #Single Q Learning\n",
    "        next_state_values[non_final_mask] = policy_net(non_final_next_states).max(1)[0].detach()\n",
    "\n",
    "        # Compute the expected Q values\n",
    "        if model=='QL':\n",
    "            expected_state_action_values = (next_state_values * GAMMA) + reward_batch\n",
    "        else:\n",
    "            expected_state_action_values = reward_batch\n",
    "\n",
    "        # Compute Huber loss\n",
    "        criterion = nn.SmoothL1Loss()\n",
    "        loss = criterion(state_action_values, expected_state_action_values.unsqueeze(1))\n",
    "        # Optimize the model\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "#         grads = []\n",
    "#         for param in policy_net.parameters():\n",
    "#             grads.append(param.grad.clone().detach().mean().numpy())\n",
    "        torch.nn.utils.clip_grad_norm_(policy_net.parameters(), 100)\n",
    "#         mean_g = np.mean(grads)\n",
    "        optimizer.step()\n",
    "    return loss.detach().item()\n",
    "\n",
    "    \n",
    "def reset_state(seqSize, nClass):\n",
    "    # Make new variable for empty token. \n",
    "    return (torch.ones((1,seqSize))*nClass).long()\n",
    "\n",
    "def reset_state_naive():\n",
    "    # Make new variable for empty token. \n",
    "    return torch.zeros((1,1)).long()\n",
    "\n",
    "def val_critic(task, policy_net, NumClass, seqSize):\n",
    "    with torch.no_grad():\n",
    "        criterion = nn.L1Loss()\n",
    "        TX = torch.randint(NumClass,(1000,seqSize))\n",
    "        TestY = task.breward(TX)\n",
    "        LastA = TX[:, -1].long().unsqueeze(1).clone()\n",
    "        TX[:, -1] = NumClass\n",
    "        actions = policy_net(TX)\n",
    "        preds = actions.gather(1, LastA)\n",
    "        loss = criterion(preds, TestY)\n",
    "    return loss\n",
    "\n",
    "class ActionSelection:\n",
    "    def __init__(self, total_eps=10000):\n",
    "        self.steps_done = 0 \n",
    "        self.EPS_START = 1.0\n",
    "        self.EPS_END = 0.05\n",
    "        self.EPS_DECAY = total_eps//10\n",
    "        self.action_samples = 0\n",
    "        self.rejects = 0\n",
    "        \n",
    "    def increment(self):\n",
    "        self.steps_done += 1 \n",
    "        \n",
    "    def select_action_pg(self, state, policy):\n",
    "        probs = policy(state)\n",
    "        m = Categorical(probs)\n",
    "        action = m.sample()\n",
    "        policy.saved_log_probs.append(m.log_prob(action))\n",
    "        return action.item()\n",
    "    \n",
    "    def select_action_pg_greedy(self, state, policy):\n",
    "        with torch.no_grad():\n",
    "            probs = policy(state)\n",
    "            action = probs.argmax()\n",
    "            return action.item()\n",
    "        \n",
    "    def select_recurrent(self, q_vals, Nclass):\n",
    "        steps_done = self.steps_done\n",
    "        sample = random.random()\n",
    "        eps_threshold = self.EPS_END + (self.EPS_START - self.EPS_END) * \\\n",
    "            math.exp(-1. * steps_done / self.EPS_DECAY)\n",
    "        self.increment()\n",
    "        if sample > eps_threshold:\n",
    "            self.action_samples += 1\n",
    "            with torch.no_grad():\n",
    "                assert q_vals.shape[1]==Nclass\n",
    "                return q_vals.max(1)[1].view(1, 1)\n",
    "        else:\n",
    "            self.rejects += 1\n",
    "            return torch.randint(Nclass, (1,1), device=device)\n",
    "\n",
    "    def select_action(self, state, policy_net, Nclass, greedy=False):\n",
    "        steps_done = self.steps_done\n",
    "        sample = random.random()\n",
    "        eps_threshold = self.EPS_END + (self.EPS_START - self.EPS_END) * \\\n",
    "            math.exp(-1. * steps_done / self.EPS_DECAY)\n",
    "        self.increment()\n",
    "        if sample > eps_threshold or greedy:\n",
    "            self.action_samples += 1\n",
    "            with torch.no_grad():\n",
    "                q_vals = policy_net(state)\n",
    "                assert q_vals.shape[1]==Nclass\n",
    "                return policy_net(state).max(1)[1].view(1, 1)\n",
    "        else:\n",
    "            self.rejects += 1\n",
    "            return torch.randint(Nclass, (1,1), device=device)\n",
    "        \n",
    "\n",
    "def greedy_perturbation(memory, SeqSize=11, NumClass=20):\n",
    "    if len(memory)<=0:\n",
    "        return torch.randint(NumClass,(1,SeqSize))\n",
    "    seq = getBestMemory(memory)\n",
    "    idx = torch.randint(SeqSize,(1,))\n",
    "    val = torch.randint(NumClass, (1,))\n",
    "    seq[0,idx]=val \n",
    "    return seq \n",
    "\n",
    "def perturbation(seq, SeqSize=11, NumClass=20):\n",
    "    if seq is None:\n",
    "        return torch.randint(NumClass,(1,SeqSize))\n",
    "    idx = torch.randint(SeqSize,(1,))\n",
    "    val = torch.randint(NumClass, (1,))\n",
    "    seq[0,idx]=val \n",
    "    return seq \n",
    "        \n",
    "class StructureSelection:\n",
    "    def __init__(self, t_init=1.):\n",
    "        # Increase T means more accepts, decrease means more rejects\n",
    "        self.t_init = t_init #1. \n",
    "        self.rejects = 0\n",
    "        self.steps_done = 0\n",
    "        self.action_samples = 0\n",
    "    \n",
    "    def increment(self):\n",
    "        self.steps_done += 1 \n",
    "        \n",
    "    def _criterion(self, structure, nstructure, val1, val2, rew_std=None):\n",
    "        \"\"\"Do we accept struct1 or reject it and keep (or choose struct2) val1 is for chosen structure, val2 is for random structure.\"\"\"\n",
    "        if not torch.is_tensor(val1):\n",
    "            val1 = torch.tensor(np.array([val1*1.0])).float()\n",
    "        if not torch.is_tensor(val2):\n",
    "            val2 = torch.tensor(np.array([val2*1.0])).float()\n",
    "        # TODO: Try both with and without annealing temperature.\n",
    "        if rew_std is not None:\n",
    "            t = rew_std\n",
    "        else:\n",
    "            t = self.t_init/(self.steps_done+1)\n",
    "#             t = self.t_init\n",
    "        diff = val2-val1\n",
    "        acceptP = torch.clip(torch.exp(-diff/t), 0, 1).item()\n",
    "        self.increment()\n",
    "        if diff<0 or torch.rand(1).item() < acceptP:\n",
    "            # Critics Agree\n",
    "            self.action_samples += 1\n",
    "            return structure, True, acceptP\n",
    "        else:\n",
    "            # Difference is too minor, so takes random structure\n",
    "            self.rejects += 1\n",
    "            return nstructure, False, None\n",
    "            \n",
    "    def select_MH(self, structure, nstructure, StrCritic, rew_std=None):\n",
    "        \"\"\"Choose action based on second critic estimates.\"\"\"\n",
    "        if nstructure is None:\n",
    "            return structure, True\n",
    "        with torch.no_grad():\n",
    "            val1 = StrCritic(structure).squeeze()\n",
    "            val2 = StrCritic(nstructure).squeeze()\n",
    "            return self._criterion(structure, nstructure, val1, val2, rew_std)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "71325e36",
   "metadata": {},
   "outputs": [],
   "source": [
    "def backwardStore(SubStrmemory, Structure, rStructure, SeqSize, NumClass):\n",
    "    a_next_state = None\n",
    "    a_state = Structure.clone()\n",
    "    for rstep in range(1, SeqSize+1):\n",
    "        a_action = Structure[:, -rstep].unsqueeze(0)\n",
    "        a_state[:, -rstep] = NumClass\n",
    "        SubStrmemory.push(a_state, a_action, a_next_state, rStructure)\n",
    "        a_next_state = a_state.clone()\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "94096b54",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "class SA:\n",
    "    def __init__(self, SeqSize, NumClass, exploration, num_episodes):\n",
    "        self.curr_x = None\n",
    "        self.curr_val = None \n",
    "        self.SeqSize = SeqSize\n",
    "        self.NumClass = NumClass\n",
    "        self.exploration = exploration\n",
    "        self.counter = 0\n",
    "        self.Strmemory = CombReplayMemory(min(num_episodes, 1e5))\n",
    "        \n",
    "    def suggest(self, method):\n",
    "        if method=='SA':\n",
    "            return self.suggest_sa()\n",
    "        if method=='G':\n",
    "            return self.suggest_greedy()\n",
    "        if method=='RS':\n",
    "            return self.suggest_random()\n",
    "        \n",
    "    def suggest_random(self):\n",
    "        self.counter += 1 \n",
    "        return torch.randint(self.NumClass,(1,self.SeqSize))\n",
    "        \n",
    "    def suggest_greedy(self):\n",
    "        if self.counter>self.exploration:\n",
    "            self.counter += 1 \n",
    "            return greedy_perturbation(self.Strmemory, self.SeqSize, self.NumClass)\n",
    "        else:\n",
    "            self.counter += 1\n",
    "            return torch.randint(self.NumClass,(1, self.SeqSize))\n",
    "        \n",
    "    def suggest_sa(self):\n",
    "        if self.counter>self.exploration:\n",
    "            self.counter += 1\n",
    "            return perturbation(self.curr_x, self.SeqSize, self.NumClass)\n",
    "        else:\n",
    "            self.counter += 1\n",
    "            return torch.randint(self.NumClass,(1, self.SeqSize))\n",
    "    \n",
    "    def update(self, Structure, rStructure, sSelector):\n",
    "        \n",
    "        \"\"\"Maximise function\"\"\"\n",
    "        \n",
    "        # Store\n",
    "        rStructure = torch.tensor([rStructure], device=device)\n",
    "        self.Strmemory.push(Structure, rStructure)\n",
    "            \n",
    "        if self.curr_val is None:\n",
    "            self.curr_val = rStructure\n",
    "            self.curr_x = Structure\n",
    "\n",
    "        else:   \n",
    "            _, criterion, _ = sSelector._criterion(Structure, self.curr_x, rStructure, self.curr_val)\n",
    "            if criterion:\n",
    "                sSelector.action_samples += 1\n",
    "                self.curr_val = rStructure\n",
    "                self.curr_x = Structure\n",
    "\n",
    "\n",
    "def run_random(taskF, num_episodes, SeqSize, NumClass, BATCH_SIZE, method='SA'):\n",
    "    rewards = []\n",
    "    exploration = int(BATCH_SIZE)\n",
    "    t_init = 1.\n",
    "    action_samples = 0\n",
    "    rejects = 0\n",
    "    search = SA(SeqSize, NumClass, exploration, num_episodes)\n",
    "    sSelector = StructureSelection()\n",
    "\n",
    "    for i_episode in range(num_episodes):\n",
    "        Structure = search.suggest(method)\n",
    "        rStructure, OptimalDone = taskF.reward(Structure)\n",
    "        search.update(Structure, rStructure, sSelector)\n",
    "        rewards.append(rStructure)\n",
    "        \n",
    "        if OptimalDone:\n",
    "            taskF.get_best()\n",
    "                        \n",
    "    print(f'accepts {sSelector.action_samples} rejects {sSelector.rejects}')\n",
    "    return np.maximum.accumulate(rewards)\n",
    "\n",
    "def run_multiple_rl_seeds(df_init, nseeds=3, neps=2000, nclass=20, nseq=10, verbose=False, model='QL', BATCH_SIZE=32, gradient_steps=1, baseline=True, easy=False):\n",
    "    \n",
    "    df_model = df_init[df_init.model==model]\n",
    "    if len(df_model)>0:\n",
    "        start_seed = df_model.seed.max() + 1\n",
    "    else:\n",
    "        start_seed = 0\n",
    "    for i in range(nseeds):\n",
    "        seed = int(i + start_seed)\n",
    "        torch.manual_seed(seed)\n",
    "        np.random.seed(seed)\n",
    "        df = pd.DataFrame({'model': [], 'seed': [], 'step': [], 'max_reward': [], 'criterion':[], 'task':[], 'size':[]})\n",
    "        task = SeqOptimisation(nseq, nclass, easy=easy)\n",
    "        if model=='QL' or model=='RQL':\n",
    "            print(model)\n",
    "            max_r = run_ql(task, num_episodes=neps, SeqSize=nseq, NumClass=nclass, BATCH_SIZE=BATCH_SIZE, gradient_steps=gradient_steps, model=model)\n",
    "        elif model=='SQL':\n",
    "            max_r, criterion = run_sql(task, num_episodes=neps, SeqSize=nseq, NumClass=nclass, BATCH_SIZE=BATCH_SIZE, verbose=verbose, gradient_steps=gradient_steps, model=model)\n",
    "            df['criterion'] = criterion\n",
    "        elif model=='PG' or model=='PGB':\n",
    "            max_r = run_reinforce(task, num_episodes=neps, SeqSize=nseq, NumClass=nclass, BATCH_SIZE=BATCH_SIZE, verbose=verbose, gradient_steps=gradient_steps, model=model, baseline=baseline)\n",
    "        elif model in ['SA','G','RS']:\n",
    "            max_r = run_random(task, num_episodes=neps, SeqSize=nseq, NumClass=nclass, BATCH_SIZE=BATCH_SIZE, method=model)\n",
    "        else:\n",
    "            return NotImplementedErrort\n",
    "        df['max_reward'] = -1*max_r\n",
    "        df['step'] = np.arange(0, max_r.shape[0], 1)\n",
    "        df['model'] = model\n",
    "        df['seed'] = seed\n",
    "        df['task'] = 'easy' if easy else 'hard'\n",
    "        df['size'] = nseq\n",
    "        df_init = df_init.append(df)\n",
    "    return df_init"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "ed4c383b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "def run_sql(taskF, num_episodes = 1000, SeqSize=11, NumClass=20, BATCH_SIZE=16, train_freq=1, verbose=False, saveD=False, d_model=32, gradient_steps=1, model='QL'):\n",
    "    # Needs to have task that maximises reward, can flip it after\n",
    "    SubStrCritic = TransformerNN(numberVars=NumClass+1, outputs=NumClass, d_model=d_model, seqSize=SeqSize, PolicyGradients=False).to(device)\n",
    "    SubStroptimizer = optim.Adam(SubStrCritic.parameters(), lr=1e-4)\n",
    "    SubStrmemory = ReplayMemory(min(num_episodes, 1e5))\n",
    "    aSelector = ActionSelection(num_episodes*SeqSize)\n",
    "    \n",
    "    StrCritic = TransformerNN(numberVars=NumClass+1, outputs=1, d_model=d_model, seqSize=SeqSize, PolicyGradients=False).to(device)\n",
    "    Stroptimizer = optim.Adam(StrCritic.parameters(), lr=1e-4)    \n",
    "    Strmemory = CombReplayMemory(min(num_episodes, 1e5))\n",
    "    sSelector = StructureSelection()\n",
    "\n",
    "    if verbose:\n",
    "        val_loss = []\n",
    "        train_loss = []\n",
    "        \n",
    "    rewards = []\n",
    "    criterions = []\n",
    "    total_steps = 0\n",
    "    OptimalDone = False\n",
    "    \n",
    "    for i_episode in range(num_episodes):\n",
    "        if (i_episode%1000)==10:\n",
    "            taskF.get_best()\n",
    "        # Initialize the environment and state\n",
    "        state = reset_state(SeqSize, NumClass)\n",
    "        assert state.max().item()<=(NumClass+1)\n",
    "        tmp_buffer = []\n",
    "        for step in range(SeqSize):\n",
    "            # Select and perform an action\n",
    "            action = aSelector.select_action(state, SubStrCritic, NumClass, greedy=True)\n",
    "            reward = 0.\n",
    "            next_state = state.clone()\n",
    "            next_state[0, step] = action\n",
    "            \n",
    "            if step==(SeqSize-1):\n",
    "                done = True\n",
    "                        \n",
    "            else:\n",
    "                done=False\n",
    "                tmp_buffer.append((state, action, next_state, None))\n",
    "                \n",
    "            \n",
    "            if done:\n",
    "                # Perform n steps of the optimization (on the policy network)\n",
    "                tloss = optimize_SubStr(BATCH_SIZE, SubStrCritic, SubStroptimizer, SubStrmemory, gradient_steps=gradient_steps, model=model)\n",
    "                _ = optimize_Str(BATCH_SIZE, StrCritic, Stroptimizer, Strmemory, gradient_steps=2)\n",
    "                nstructure = greedy_perturbation(Strmemory, SeqSize, NumClass)\n",
    "                rew_std = np.std(rewards)\n",
    "                Structure, criterion, acceptP = sSelector.select_MH(next_state, nstructure, StrCritic, rew_std)\n",
    "                \n",
    "                # If Criterion Was Accepted, Add SubStr Optimisation Result to Memory, if not, add Random Structure Into SubStr Memory Instead \n",
    "                rStructure, OptimalDone = taskF.reward(Structure)\n",
    "                \n",
    "                # Add full structure to structure Critic\n",
    "                rStructure = torch.tensor([rStructure], device=device)\n",
    "                Strmemory.push(Structure, rStructure)\n",
    "                if criterion:\n",
    "                    SubStrmemory.push(state, action, None, rStructure)\n",
    "                    for (a_state, a_action, a_next_state, _) in tmp_buffer:\n",
    "                        SubStrmemory.push(a_state, a_action, a_next_state, rStructure)\n",
    "                    \n",
    "                else:\n",
    "                    # Final structure was not selected by agent.\n",
    "                    backwardStore(SubStrmemory, Structure, rStructure, SeqSize, NumClass)\n",
    "\n",
    "                rewards.append(rStructure.detach().item())\n",
    "#                 criterions.append(acceptP)\n",
    "                criterions.append(sSelector.action_samples/(i_episode+1))\n",
    "                total_steps += 1\n",
    "                break\n",
    "                \n",
    "            # Move to the next state\n",
    "            state = next_state\n",
    "            if OptimalDone:\n",
    "                taskF.get_best()\n",
    "                print(f'OPTIMAL with {i_episode} policy actions accepts {sSelector.action_samples} and rejects {sSelector.rejects}')\n",
    "    \n",
    "    taskF.get_best()\n",
    "    print(f'Complete with {i_episode} policy actions accepts {sSelector.action_samples} and rejects {sSelector.rejects}')\n",
    "    return np.maximum.accumulate(rewards), criterions\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "5d3dd107",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "def optimise_rnn(mem, main_model, BATCH_SIZE, optimizer, criterion, gradient_steps, GAMMA = 0.99):\n",
    "    if len(mem)<BATCH_SIZE:\n",
    "        return None\n",
    "    for _ in range(gradient_steps):\n",
    "        hidden_batch, cell_batch = main_model.init_hidden_states(bsize=BATCH_SIZE)\n",
    "\n",
    "        batch = mem.get_batch(bsize=BATCH_SIZE)\n",
    "\n",
    "        current_states = []\n",
    "        acts = []\n",
    "        rewards = []\n",
    "        next_states = []\n",
    "\n",
    "        for b in batch:\n",
    "            cs,ac,rw,ns = [],[],[],[]\n",
    "            for element in b:\n",
    "                cs.append(element[0])\n",
    "                ac.append(element[1])\n",
    "                rw.append(element[2])\n",
    "                ns.append(element[3])\n",
    "            current_states.append(cs)\n",
    "            acts.append(ac)\n",
    "            rewards.append(rw)\n",
    "            next_states.append(ns)\n",
    "\n",
    "        current_states = np.array(current_states)\n",
    "        acts = np.array(acts)\n",
    "        rewards = np.array(rewards)\n",
    "        next_states = np.array(next_states)\n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        torch_current_states = torch.from_numpy(current_states).long().to(device)\n",
    "        torch_acts = torch.from_numpy(acts).long().to(device)\n",
    "        torch_rewards = torch.from_numpy(rewards).float().to(device).squeeze(-1)\n",
    "        torch_next_states = torch.from_numpy(next_states).long().to(device)\n",
    "\n",
    "\n",
    "        Q_next,_ = main_model.forward(torch_next_states, hidden_state=hidden_batch, cell_state=cell_batch)\n",
    "        Q_next_max,__ = Q_next.clone().detach().max(dim=-1)\n",
    "        target_values = torch_rewards + (GAMMA * Q_next_max)\n",
    "        Q_s, _ = main_model.forward(torch_current_states, hidden_state=hidden_batch, cell_state=cell_batch)\n",
    "        Q_s_a = Q_s.gather(dim=-1, index=torch_acts).squeeze(-1)\n",
    "        # make previous grad zero\n",
    "        loss = criterion(Q_s_a,target_values)    \n",
    "\n",
    "    \n",
    "        # backward\n",
    "        loss.backward()\n",
    "        torch.nn.utils.clip_grad_norm_(main_model.parameters(), 100)\n",
    "        # update params\n",
    "        optimizer.step()\n",
    "\n",
    "    return loss\n",
    "\n",
    "    \n",
    "\n",
    "class RecurrentMemory():\n",
    "    \n",
    "    def __init__(self,memsize):\n",
    "        self.memsize = memsize\n",
    "        self.memory = deque(maxlen=self.memsize)\n",
    "        self.currSize = 0 \n",
    "        \n",
    "    def __len__(self):\n",
    "        return self.currSize\n",
    "    \n",
    "    def add_episode(self,epsiode):\n",
    "        self.memory.append(epsiode)\n",
    "        self.currSize += 1\n",
    "        \n",
    "    def get_batch(self,bsize):\n",
    "        sampled_epsiodes = random.sample(self.memory,bsize)\n",
    "        batch = []\n",
    "        for episode in sampled_epsiodes:\n",
    "            batch.append(episode)\n",
    "        return batch\n",
    "\n",
    "class LSTMQNetwork(nn.Module):\n",
    "    \n",
    "    def __init__(self,input_size=1,d_model=32,out_size=1, numberVars=20):\n",
    "        super(LSTMQNetwork,self).__init__()\n",
    "        self.Embedder = nn.Embedding(numberVars+1, d_model, max_norm=True)\n",
    "        self.input_size = input_size\n",
    "        self.out_size = out_size\n",
    "        self.layers = 1\n",
    "        self.d_model = d_model\n",
    "        self.lstm_layer = nn.LSTM(input_size=d_model,hidden_size=d_model,num_layers=self.layers,batch_first=True)\n",
    "        self.val = nn.Linear(in_features=d_model,out_features=out_size)\n",
    "        \n",
    "    def forward(self,x, hidden_state, cell_state):\n",
    "        x = self.Embedder(x).squeeze(-2)\n",
    "        lstm_out = self.lstm_layer(x,(hidden_state,cell_state))\n",
    "        out = lstm_out[0]\n",
    "        h_n = lstm_out[1][0]\n",
    "        c_n = lstm_out[1][1]\n",
    "        val_out = self.val(out)\n",
    "        return val_out, (h_n,c_n)\n",
    "    \n",
    "    def init_hidden_states(self,bsize):\n",
    "        h = torch.zeros(self.layers,bsize,self.d_model).float().to(device)\n",
    "        c = torch.zeros(self.layers,bsize,self.d_model).float().to(device)\n",
    "        \n",
    "        return h,c\n",
    "\n",
    "def run_ql(taskF, num_episodes = 1000, SeqSize=11, NumClass=20, BATCH_SIZE=16, train_freq=1, verbose=False, saveD=False, d_model=32, gradient_steps=1, model='QL'):\n",
    "    if model=='RQL':\n",
    "        recurrent = True \n",
    "    else:\n",
    "        recurrent = False \n",
    "    if recurrent:\n",
    "        main_model = LSTMQNetwork(1, d_model, NumClass)\n",
    "        optimizer = optim.Adam(main_model.parameters(), lr=1e-4)\n",
    "        memory = RecurrentMemory(min(num_episodes, 1e5))\n",
    "        criterion = nn.SmoothL1Loss()\n",
    "    else:\n",
    "        SubStrCritic = TransformerNN(numberVars=NumClass+1, outputs=NumClass, d_model=d_model, seqSize=SeqSize, PolicyGradients=False).to(device)\n",
    "        SubStroptimizer = optim.Adam(SubStrCritic.parameters(), lr=1e-4)\n",
    "        SubStrmemory = ReplayMemory(min(num_episodes, 1e5))\n",
    "    \n",
    "    aSelector = ActionSelection(num_episodes*SeqSize)\n",
    "        \n",
    "    rewards = []\n",
    "    loss_stat = []\n",
    "\n",
    "    total_steps = 0\n",
    "    OptimalDone = False\n",
    "    \n",
    "    \n",
    "    for i_episode in tqdm(range(num_episodes)):\n",
    "        # Initialize the environment and state [1,1]\n",
    "        local_memory = []\n",
    "        Structure = reset_state(SeqSize, NumClass)\n",
    "        if recurrent:\n",
    "            state = reset_state_naive()\n",
    "            hidden_state, cell_state = main_model.init_hidden_states(bsize=1)\n",
    "\n",
    "        for step in range(SeqSize):\n",
    "            if recurrent:\n",
    "                # Select and perform an action\n",
    "                model_out = main_model.forward(state.long().unsqueeze(0), hidden_state=hidden_state, cell_state=cell_state)\n",
    "                qvals = model_out[0][:,-1,:]\n",
    "                hidden_state = model_out[1][0]\n",
    "                cell_state = model_out[1][1]\n",
    "                action = aSelector.select_recurrent(qvals, NumClass)\n",
    "                next_state = action  \n",
    "            else:\n",
    "                action = aSelector.select_action(Structure, SubStrCritic, NumClass, greedy=False)\n",
    "                \n",
    "            reward = 0\n",
    "                      \n",
    "            if step==(SeqSize-1):\n",
    "                done = True       \n",
    "            else:\n",
    "                if recurrent:\n",
    "                    local_memory.append((state.long().flatten().numpy(), np.array([action]), np.array([reward]), next_state.long().flatten().numpy()))\n",
    "                else:\n",
    "                    nStructure = Structure.clone()\n",
    "                    nStructure[:, step] = action\n",
    "                    rStructure = torch.tensor([reward], device=device)\n",
    "                    SubStrmemory.push(Structure.clone(), action, nStructure, rStructure)\n",
    "                done=False\n",
    "            \n",
    "            if done:\n",
    "                nStructure = Structure.clone()\n",
    "                nStructure[:, step] = action\n",
    "                reward, OptimalDone = taskF.reward(nStructure)\n",
    "                if recurrent:\n",
    "                    local_memory.append((state.float().flatten().numpy(), np.array([action]), np.array([reward]), next_state.long().flatten().numpy()))\n",
    "                    memory.add_episode(local_memory)\n",
    "                else:\n",
    "                    rStructure = torch.tensor([reward], device=device)\n",
    "                    SubStrmemory.push(Structure, action, None, rStructure)\n",
    "                if recurrent:\n",
    "                    loss = optimise_rnn(memory, main_model, BATCH_SIZE, optimizer, criterion, gradient_steps)\n",
    "                else:\n",
    "                    loss = optimize_SubStr(BATCH_SIZE, SubStrCritic, SubStroptimizer, SubStrmemory, gradient_steps=gradient_steps, model=model)\n",
    "                    \n",
    "                rewards.append(reward)\n",
    "                total_steps += 1\n",
    "                break\n",
    "                \n",
    "            # Move to the next state\n",
    "            if recurrent:\n",
    "                state = next_state\n",
    "            \n",
    "            Structure[:, step] = action\n",
    "                \n",
    "            if OptimalDone:\n",
    "                taskF.get_best()\n",
    "                print(f'OPTIMAL with {i_episode}')\n",
    "    \n",
    "    taskF.get_best()\n",
    "    return np.maximum.accumulate(rewards)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86c9db22",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "id": "653482d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "EPS_START = 1.0\n",
    "EPS_END = 0.05\n",
    "EPS_DECAY = 100\n",
    "steps_done = np.arange(0, 1000, 1)\n",
    "eps_threshold = EPS_END + (EPS_START - EPS_END) * np.exp(-1. * steps_done/100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "87c89773",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_eg(eps_threshold, steps_done):\n",
    "    return 1.0-np.cumsum(np.random.binomial(1, eps_threshold))/(steps_done+1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 237,
   "id": "f5178b56",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'1'"
      ]
     },
     "execution_count": 237,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "r'{}'.format(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 342,
   "id": "f78f72e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib as mpl\n",
    "mpl.rcParams.update(mpl.rcParamsDefault)\n",
    "\n",
    "def plot_seaborn(df_runs, colourD, markerD, save=False, easy=True, x_label='Protein Designs', y_label='max_reward', legend=False, numeps=1000, color='red', length=10):\n",
    "    plt.rcParams['text.usetex'] = True\n",
    "    \n",
    "    fsize = 25\n",
    "    name = f'Synthetic {\"Exploitation\" if easy else \"Exploration\"} Task (Length {SeqSize})'\n",
    "    matplotlib.rc('xtick', labelsize=fsize)\n",
    "    matplotlib.rc('ytick', labelsize=fsize)\n",
    "    cols = 1\n",
    "    rows = 1.5\n",
    "    fig, (axtop, axeps, ax) = plt.subplots(nrows=3, ncols=1, figsize=(6.2*cols, 5*rows),  gridspec_kw={'height_ratios': [2.5, 2.5, 10]})\n",
    "    sns.set_style(\"white\")\n",
    "    \n",
    "    df_runs = df_runs.reset_index()\n",
    "    df_runs['max_reward'] *= length\n",
    "    df_runs = df_runs.round({'max_reward': 0})\n",
    "    \n",
    "    print(df_runs.shape)\n",
    "    greenpastel = sns.color_palette('pastel', 12)[2]\n",
    "    redpastel = sns.color_palette('pastel', 12)[3]\n",
    "    dfcrit = df_runs[df_runs.model=='SQL']\n",
    "    if len(dfcrit) != 0:\n",
    "        EPS_START = 1.0\n",
    "        EPS_END = 0.05\n",
    "        EPS_DECAY = numeps//10\n",
    "        steps_done = np.arange(0, numeps, 1)\n",
    "\n",
    "        critval = dfcrit.groupby(['step']).aggregate('mean')['criterion'].to_numpy()\n",
    "\n",
    "        ones = np.ones_like(critval)\n",
    "        xs = np.arange(0, len(ones), 1)\n",
    "        print('Started Plotting Crit')       \n",
    "        axtop.fill_between(xs,critval, color=greenpastel)\n",
    "        axtop.fill_between(xs,critval, ones,where=ones>=critval, color=redpastel)\n",
    "        sns.lineplot(ax=axtop, data=dfcrit, x='step', y='criterion', color=color, linewidth=1, legend=False)        \n",
    "        axtop.set_xlabel('')\n",
    "        axtop.set_ylabel(r\"$\\mathbf{\\pi_{\\mathcal{S}C}}$\"\n",
    "#           \"\\n\"\n",
    "#           r\"$\\textbf{Expolitation}$\"\n",
    "                         , fontsize=20, fontweight=\"bold\", labelpad=15)\n",
    "        axtop.set_xticks([])\n",
    "        axtop.set_title(name, fontsize=fsize, fontweight=\"bold\")\n",
    "        s1 = axtop.spines[\"top\"]\n",
    "        s1.set_visible(False)\n",
    "        s3 = axtop.spines[\"right\"]\n",
    "        s3.set_visible(False)\n",
    "        s2 = axtop.spines[\"bottom\"]\n",
    "        s2.set_visible(False)\n",
    "#         axeps.set_ylim([0, 1])\n",
    "        axtop.set_yticks([])\n",
    "        print('Finished Plotting Crit')\n",
    "        \n",
    "        print('Started Plotting Eps') \n",
    "        eps_threshold = EPS_END + (EPS_START - EPS_END) * np.exp(-1. * steps_done/EPS_DECAY)\n",
    "        dfeps = pd.DataFrame(np.array([get_eg(eps_threshold, steps_done) for _ in range(3)]).T).melt()\n",
    "        dfeps.columns = ['seed', 'criterion']\n",
    "        dfeps['step'] = np.array([steps_done for _ in range(3)]).flatten()\n",
    "        eps_threshold = dfeps.groupby(['step']).aggregate('mean')['criterion'].to_numpy()\n",
    "\n",
    "        sns.lineplot(ax=axeps, data=dfeps, x='step', y='criterion', color=color, linewidth=1, legend=False)        \n",
    "        axeps.fill_between(xs,eps_threshold, color=greenpastel)\n",
    "        axeps.fill_between(xs,eps_threshold, ones,where=ones>=eps_threshold, color=redpastel)\n",
    "        axeps.set_xlabel('')\n",
    "        axeps.set_ylabel(r\"$\\mathbf{\\pi_{\\epsilon}}$\"\n",
    "#           \"\\n\"\n",
    "#           r\"$\\textbf{Expolitation}$\"\n",
    "                         , fontsize=20, fontweight=\"bold\", labelpad=15)\n",
    "#         axeps.set_ylim([0, 1])\n",
    "        axeps.set_xticks([])\n",
    "        axeps.set_yticks([])\n",
    "        s1 = axeps.spines[\"top\"]\n",
    "        s1.set_visible(False)\n",
    "        s3 = axeps.spines[\"right\"]\n",
    "        s3.set_visible(False)\n",
    "        \n",
    "        print('Finished Plotting Eps')\n",
    "        \n",
    "    else:\n",
    "        ax.set_title(name, fontsize=fsize, fontweight=\"bold\")\n",
    "\n",
    "    \n",
    "    ax.set_xlabel(r\"$\\textbf{Protein Designs}$\", fontsize=20, fontweight=\"bold\")\n",
    "    ax.set_ylabel(r\"$\\textbf{Hamming Distance}$\"\n",
    "#                   \"\\n\"\n",
    "#                   r\"$\\textbf{Distance}$\"\n",
    "                  , fontsize=20, fontweight=\"bold\", labelpad=15)\n",
    "#     ax.yaxis.set_label_position(\"right\")\n",
    "    print('Started Plotting Curves')\n",
    "    sns.lineplot(ax=ax, data=df_runs, x='step', y='max_reward', hue=\"model\", legend=legend, linewidth=3, palette=colourD, markers=markerD, markevery=50, markersize=3)\n",
    "    save_friendly = name.replace('\\n','').replace(' ','_')\n",
    "    plt.subplots_adjust(wspace=0, hspace=0)\n",
    "#     s1 = ax.spines[\"top\"]\n",
    "#     s1.set_visible(False)\n",
    "    s3 = ax.spines[\"right\"]\n",
    "    s3.set_visible(False)\n",
    "    fig.text(0, .75, r\"$\\textbf{Expolitation}$\", va='center', rotation='vertical', fontsize=20)\n",
    "#     ax.set_ylim(0, 0.8)\n",
    "    if save:\n",
    "        plt.savefig(save_friendly+f'_legend_{legend}.pdf', bbox_inches=\"tight\")\n",
    "#     plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 343,
   "id": "a34454c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "LabelDict = {'Criterion':'Criterion','RS': 'Random Search','SA': 'Simulated Annealing','AS': 'Active Search','AMS': 'Attention Model Sampling', \n",
    "             'AMG': 'Attention Model Greedy', 'RQL':'Recurrent Q Learning', 'SQL':'Structured Q Learning',\n",
    "             'High Affinity':'High Affinity','Super Affinity':'Super Affinity', 'Super+ Affinity':'Super+ Affinity'}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 344,
   "id": "0fbda6a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "colours = sns.color_palette('dark', len(LabelDict.keys()))\n",
    "LabelDictinv = {v: k for k, v in LabelDict.items()}\n",
    "colourDINV = {}\n",
    "for i,(key,_) in enumerate(LabelDict.items()):\n",
    "    colourDINV[key]=colours[i]\n",
    "markerDINV = {}\n",
    "all_markers = ['.', 'o', 'v', '^', '<', '>', '8', 's', 'p', '*', 'h', 'H', 'D', 'd', 'P', 'X']\n",
    "from matplotlib.lines import Line2D\n",
    "for i,(key,_) in enumerate(LabelDict.items()):\n",
    "    markerDINV[key]=all_markers[i]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 345,
   "id": "1445712b",
   "metadata": {},
   "outputs": [],
   "source": [
    "SeqSize = 10\n",
    "easy=True\n",
    "n_seeds = 5\n",
    "neps = 10000\n",
    "batch_size = 32\n",
    "gradient_steps = 1 # 20 is good\n",
    "save = True\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 346,
   "id": "988690db",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "df_mh = pd.read_csv('AllData.csv', index_col=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3efb92d6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(400000, 8)\n",
      "Started Plotting Crit\n"
     ]
    }
   ],
   "source": [
    "# for SeqSize in [10, 100]:    \n",
    "for SeqSize in [10]: \n",
    "    for easy in [True]:\n",
    "#     for easy in [True, False]:\n",
    "        if easy:\n",
    "            df_task = df_mh[df_mh.task=='easy']\n",
    "        else:\n",
    "            df_task = df_mh[df_mh.task=='hard']\n",
    "        plot_seaborn(df_task, colourDINV, markerDINV, save, easy, legend=True, color=colourDINV['Criterion'], length=SeqSize, numeps=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 288,
   "id": "1eb4a500",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Target Sequence is tensor([3, 3, 3, 3, 3, 3, 3, 3, 3, 3])\n",
      "with hamming distance 10 and reward -0.7615941559557649\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/alexc/anaconda3/lib/python3.7/site-packages/numpy/core/_methods.py:263: RuntimeWarning: Degrees of freedom <= 0 for slice\n",
      "  keepdims=keepdims, where=where)\n",
      "/home/alexc/anaconda3/lib/python3.7/site-packages/numpy/core/_methods.py:223: RuntimeWarning: invalid value encountered in true_divide\n",
      "  subok=False)\n",
      "/home/alexc/anaconda3/lib/python3.7/site-packages/numpy/core/_methods.py:254: RuntimeWarning: invalid value encountered in double_scalars\n",
      "  ret = ret.dtype.type(ret / rcount)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found Optimal Sequence tensor([3, 3, 3, 3, 3, 3, 3, 3, 3, 3]) at evaluation 172\n",
      "with hamming distance 0 and reward 0.0\n",
      "with hamming distance 0 and reward 0.0\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m/tmp/ipykernel_28381/3555571973.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      5\u001b[0m     \u001b[0;32mfor\u001b[0m \u001b[0measy\u001b[0m \u001b[0;32min\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;32mFalse\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      6\u001b[0m \u001b[0;31m#         df_mh = run_multiple_rl_seeds(df_mh, nseeds=n_seeds, neps=neps, BATCH_SIZE=batch_size, model='G', easy=easy, nseq=SeqSize)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 7\u001b[0;31m         \u001b[0mdf_mh\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mrun_multiple_rl_seeds\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdf_mh\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnseeds\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mn_seeds\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mneps\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mneps\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mBATCH_SIZE\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mbatch_size\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgradient_steps\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mgradient_steps\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'SQL'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0measy\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0measy\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnseq\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mSeqSize\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      8\u001b[0m \u001b[0;31m#         df_mh = run_multiple_rl_seeds(df_mh, nseeds=n_seeds, neps=neps, BATCH_SIZE=batch_size, gradient_steps=gradient_steps, model='QL', easy=easy, nseq=SeqSize)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      9\u001b[0m         \u001b[0mdf_mh\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mrun_multiple_rl_seeds\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdf_mh\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnseeds\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mn_seeds\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mneps\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mneps\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mBATCH_SIZE\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbatch_size\u001b[0m\u001b[0;34m//\u001b[0m\u001b[0mSeqSize\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m+\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgradient_steps\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mgradient_steps\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'RQL'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0measy\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0measy\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnseq\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mSeqSize\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/tmp/ipykernel_28381/2512860311.py\u001b[0m in \u001b[0;36mrun_multiple_rl_seeds\u001b[0;34m(df_init, nseeds, neps, nclass, nseq, verbose, model, BATCH_SIZE, gradient_steps, baseline, easy)\u001b[0m\n\u001b[1;32m     95\u001b[0m             \u001b[0mmax_r\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mrun_ql\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtask\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnum_episodes\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mneps\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mSeqSize\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mnseq\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mNumClass\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mnclass\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mBATCH_SIZE\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mBATCH_SIZE\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgradient_steps\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mgradient_steps\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     96\u001b[0m         \u001b[0;32melif\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m==\u001b[0m\u001b[0;34m'SQL'\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 97\u001b[0;31m             \u001b[0mmax_r\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcriterion\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mrun_sql\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtask\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnum_episodes\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mneps\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mSeqSize\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mnseq\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mNumClass\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mnclass\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mBATCH_SIZE\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mBATCH_SIZE\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mverbose\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mverbose\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgradient_steps\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mgradient_steps\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     98\u001b[0m             \u001b[0mdf\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'criterion'\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcriterion\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     99\u001b[0m         \u001b[0;32melif\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m==\u001b[0m\u001b[0;34m'PG'\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m==\u001b[0m\u001b[0;34m'PGB'\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/tmp/ipykernel_28381/333335012.py\u001b[0m in \u001b[0;36mrun_sql\u001b[0;34m(taskF, num_episodes, SeqSize, NumClass, BATCH_SIZE, train_freq, verbose, saveD, d_model, gradient_steps, model)\u001b[0m\n\u001b[1;32m     49\u001b[0m                 \u001b[0m_\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0moptimize_Str\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mBATCH_SIZE\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mStrCritic\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mStroptimizer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mStrmemory\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgradient_steps\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     50\u001b[0m                 \u001b[0mnstructure\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mgreedy_perturbation\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mStrmemory\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mSeqSize\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mNumClass\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 51\u001b[0;31m                 \u001b[0mrew_std\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstd\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mrewards\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     52\u001b[0m                 \u001b[0mStructure\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcriterion\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0macceptP\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0msSelector\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mselect_MH\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnext_state\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnstructure\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mStrCritic\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrew_std\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     53\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<__array_function__ internals>\u001b[0m in \u001b[0;36mstd\u001b[0;34m(*args, **kwargs)\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/lib/python3.7/site-packages/numpy/core/fromnumeric.py\u001b[0m in \u001b[0;36mstd\u001b[0;34m(a, axis, dtype, out, ddof, keepdims, where)\u001b[0m\n\u001b[1;32m   3580\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3581\u001b[0m     return _methods._std(a, axis=axis, dtype=dtype, out=out, ddof=ddof,\n\u001b[0;32m-> 3582\u001b[0;31m                          **kwargs)\n\u001b[0m\u001b[1;32m   3583\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3584\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/lib/python3.7/site-packages/numpy/core/_methods.py\u001b[0m in \u001b[0;36m_std\u001b[0;34m(a, axis, dtype, out, ddof, keepdims, where)\u001b[0m\n\u001b[1;32m    261\u001b[0m          where=True):\n\u001b[1;32m    262\u001b[0m     ret = _var(a, axis=axis, dtype=dtype, out=out, ddof=ddof,\n\u001b[0;32m--> 263\u001b[0;31m                keepdims=keepdims, where=where)\n\u001b[0m\u001b[1;32m    264\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    265\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mret\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmu\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mndarray\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/lib/python3.7/site-packages/numpy/core/_methods.py\u001b[0m in \u001b[0;36m_var\u001b[0;34m(a, axis, dtype, out, ddof, keepdims, where)\u001b[0m\n\u001b[1;32m    195\u001b[0m def _var(a, axis=None, dtype=None, out=None, ddof=0, keepdims=False, *,\n\u001b[1;32m    196\u001b[0m          where=True):\n\u001b[0;32m--> 197\u001b[0;31m     \u001b[0marr\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0masanyarray\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0ma\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    198\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    199\u001b[0m     \u001b[0mrcount\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_count_reduce_items\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0marr\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0maxis\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkeepdims\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mkeepdims\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mwhere\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mwhere\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "# df_mh = pd.read_csv('AllData.csv', index_col=0)\n",
    "# df_mh = pd.DataFrame({'model': [], 'seed': [], 'step': [], 'max_reward': [], 'criterion':[], 'task':[], 'size':[]})\n",
    "for SeqSize in [10, 100]:    \n",
    "# for SeqSize in [10]:    \n",
    "    for easy in [True, False]:\n",
    "#         df_mh = run_multiple_rl_seeds(df_mh, nseeds=n_seeds, neps=neps, BATCH_SIZE=batch_size, model='G', easy=easy, nseq=SeqSize)\n",
    "        df_mh = run_multiple_rl_seeds(df_mh, nseeds=n_seeds, neps=neps, BATCH_SIZE=batch_size, gradient_steps=gradient_steps, model='SQL', easy=easy, nseq=SeqSize)        \n",
    "#         df_mh = run_multiple_rl_seeds(df_mh, nseeds=n_seeds, neps=neps, BATCH_SIZE=batch_size, gradient_steps=gradient_steps, model='QL', easy=easy, nseq=SeqSize)\n",
    "        df_mh = run_multiple_rl_seeds(df_mh, nseeds=n_seeds, neps=neps, BATCH_SIZE=int(batch_size//SeqSize)+1, gradient_steps=gradient_steps, model='RQL', easy=easy, nseq=SeqSize)\n",
    "        df_mh = run_multiple_rl_seeds(df_mh, nseeds=n_seeds, neps=neps, BATCH_SIZE=batch_size, model='RS', easy=easy, nseq=SeqSize)\n",
    "        df_mh = run_multiple_rl_seeds(df_mh, nseeds=n_seeds, neps=neps, BATCH_SIZE=batch_size, model='SA', easy=easy, nseq=SeqSize)\n",
    "        df_task = df_mh[df_mh.task==str(easy)]\n",
    "        df_task = df_task[df_task['size']==SeqSize]\n",
    "        plot_seaborn(df_task, colourDINV, markerDINV, save, f'Synthetic {\"Exploitative\" if easy else \"Explorative\"} Task \\n (Length {SeqSize}) \\n', legend=True, color=colourDINV['Criterion'])\n",
    "#         df_mh.to_csv('AllData.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0fce49aa",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Target Sequence is tensor([3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
      "        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/alexc/anaconda3/lib/python3.7/site-packages/numpy/core/_methods.py:263: RuntimeWarning: Degrees of freedom <= 0 for slice\n",
      "  keepdims=keepdims, where=where)\n",
      "/home/alexc/anaconda3/lib/python3.7/site-packages/numpy/core/_methods.py:223: RuntimeWarning: invalid value encountered in true_divide\n",
      "  subok=False)\n",
      "/home/alexc/anaconda3/lib/python3.7/site-packages/numpy/core/_methods.py:254: RuntimeWarning: invalid value encountered in double_scalars\n",
      "  ret = ret.dtype.type(ret / rcount)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "with hamming distance 674 and reward -0.5876049206498948\n",
      "with hamming distance 653 and reward -0.5736860827264039\n",
      "with hamming distance 624 and reward -0.5539069192039144\n",
      "with hamming distance 603 and reward -0.5391808600618703\n"
     ]
    }
   ],
   "source": [
    "# BATCH_SIZE = 1000\n",
    "for SeqSize in [1000]:    \n",
    "# for SeqSize in [10]:    \n",
    "    for easy in [True, False]:\n",
    "#         df_mh = run_multiple_rl_seeds(df_mh, nseeds=n_seeds, neps=neps, BATCH_SIZE=batch_size, model='G', easy=easy, nseq=SeqSize)\n",
    "        df_mh = run_multiple_rl_seeds(df_mh, nseeds=n_seeds, neps=neps, BATCH_SIZE=batch_size, gradient_steps=gradient_steps, model='SQL', easy=easy, nseq=SeqSize)        \n",
    "#         df_mh = run_multiple_rl_seeds(df_mh, nseeds=n_seeds, neps=neps, BATCH_SIZE=batch_size, gradient_steps=gradient_steps, model='QL', easy=easy, nseq=SeqSize)\n",
    "        df_mh = run_multiple_rl_seeds(df_mh, nseeds=n_seeds, neps=neps, BATCH_SIZE=int(batch_size//SeqSize)+1, gradient_steps=gradient_steps, model='RQL', easy=easy, nseq=SeqSize)\n",
    "        df_mh = run_multiple_rl_seeds(df_mh, nseeds=n_seeds, neps=neps, BATCH_SIZE=batch_size, model='RS', easy=easy, nseq=SeqSize)\n",
    "        df_mh = run_multiple_rl_seeds(df_mh, nseeds=n_seeds, neps=neps, BATCH_SIZE=batch_size, model='SA', easy=easy, nseq=SeqSize)\n",
    "        df_task = df_mh[df_mh.task==str(easy)]\n",
    "        df_task = df_task[df_task['size']==SeqSize]\n",
    "        plot_seaborn(df_task, colourDINV, markerDINV, save, f'Synthetic {\"Exploitative\" if easy else \"Explorative\"} Task \\n (Length {SeqSize}) \\n', legend=True, color=colourDINV['Criterion'])\n",
    "        df_mh.to_csv('AllData.csv')"
   ]
  },
  {
   "cell_type": "raw",
   "id": "4409c127",
   "metadata": {},
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "def calculate_returns(policy, gamma=0.99):\n",
    "    R = 0\n",
    "    returns = []\n",
    "    for r in policy.rewards[::-1]:\n",
    "        R = r + gamma * R\n",
    "        returns.insert(0, R)\n",
    "    return torch.tensor(returns)\n",
    "\n",
    "def policy_loss(policy, returns, baseline_vals=None):\n",
    "    policy_loss = []\n",
    "    for i, (log_prob, R) in enumerate(zip(policy.saved_log_probs, returns)):\n",
    "        if baseline_vals is not None:\n",
    "            b = baseline_vals[i]\n",
    "            policy_loss.append(-log_prob * (R-b))\n",
    "        else:\n",
    "            policy_loss.append(-log_prob * R)\n",
    "    return torch.cat(policy_loss).mean()\n",
    "\n",
    "def get_statistics(returns):\n",
    "    mu = returns.mean()\n",
    "    sigma = returns.std()\n",
    "    return mu,sigma\n",
    "\n",
    "def safe_normalise(returns, mu=None, sigma=None):\n",
    "    eps = np.finfo(np.float32).eps.item()\n",
    "    if (mu is None) and (sigma is None):\n",
    "        mu,sigma = get_statistics(returns)\n",
    "    return (returns - mu) / (sigma + eps)            \n",
    "\n",
    "def finish_episode(policy, baseline_net, optimizer):\n",
    "    returns = calculate_returns(policy_loss, gamma)\n",
    "    returns = safe_normalise(returns)\n",
    "    return policy_loss(policy, returns)\n",
    "\n",
    "def finish_episode_baseline(policy, baseline_net, optimizer):\n",
    "    returns = calculate_returns(policy)\n",
    "    returns_baseline = calculate_returns(baseline_net)\n",
    "    returns_all = torch.cat([returns, returns_baseline])\n",
    "    mu,sigma = get_statistics(returns_all)\n",
    "    returns = safe_normalise(returns, mu, sigma)\n",
    "    returns_baseline = safe_normalise(returns_baseline, mu, sigma)\n",
    "    return policy_loss(policy, returns, returns_baseline)\n",
    "\n",
    "\n",
    "def run_reinforce(taskF, num_episodes = 1000, SeqSize=11, NumClass=20, BATCH_SIZE=16, train_freq=1, verbose=False, saveD=False, d_model=32, gradient_steps=1, model='PG', baseline=True):\n",
    "    TARGET_UPDATE = BATCH_SIZE*5\n",
    "    #Always add 1 to Q learning for padded class. \n",
    "    policy_net = TransformerNN(numberVars=NumClass+1, outputs=NumClass, d_model=d_model, seqSize=SeqSize, PolicyGradients=True).to(device)\n",
    "    optimizer = optim.Adam(policy_net.parameters(), lr=1e-3)\n",
    "    baseline_net = TransformerNN(numberVars=NumClass+1, outputs=NumClass, d_model=d_model, seqSize=SeqSize, PolicyGradients=True).to(device)\n",
    "    rewards = []\n",
    "    if verbose:\n",
    "        val_loss = []\n",
    "        train_loss = []\n",
    "    if saveD:\n",
    "        sequence_designs = []\n",
    "    selector = Action(num_episodes)\n",
    "    total_steps = 0\n",
    "    OptimalDone = False\n",
    "    optimizer.zero_grad()\n",
    "    policy_loss = 0\n",
    "    for i_episode in range(num_episodes//2):\n",
    "        if (i_episode%1000)==10:\n",
    "            taskF.get_best()\n",
    "        # Initialize the environment and state\n",
    "        init_state = reset_state(SeqSize, NumClass)\n",
    "        assert init_state.max().item()<=(NumClass+1)\n",
    "        \n",
    "        state = init_state.clone()\n",
    "        state_greedy = init_state.clone()\n",
    "        \n",
    "        tmp_buffer = []\n",
    "        for step in range(SeqSize):\n",
    "            # Select and perform an action\n",
    "            \n",
    "            action = selector.select_action_pg(state, policy_net)\n",
    "            action_greedy = selector.select_action_pg_greedy(state_greedy, baseline_net)\n",
    "            reward = 0.\n",
    "            \n",
    "            next_state = state.clone()\n",
    "            next_state[0, step] = action\n",
    "            \n",
    "            next_state_greedy = state.clone()\n",
    "            next_state_greedy[0, step] = action_greedy\n",
    "            \n",
    "            if step==(SeqSize-1):\n",
    "                done = True\n",
    "                if saveD:\n",
    "                    sequence_designs.append(next_state)\n",
    "                reward, OptimalDone = taskF.reward(next_state)\n",
    "                reward_greedy, _ = taskF.reward(next_state_greedy)\n",
    "\n",
    "                        \n",
    "            else:\n",
    "                done=False\n",
    "            policy_net.rewards.append(reward)\n",
    "            baseline_net.rewards.append(reward)\n",
    "\n",
    "            # Move to the next state\n",
    "            state = next_state\n",
    "            state_greedy = next_state_greedy\n",
    "            \n",
    "            if done:\n",
    "                \n",
    "                # Perform one step of the optimization (on the policy network)\n",
    "                if baseline:\n",
    "                    policy_loss += finish_episode_baseline(policy_net, baseline_net, optimizer)\n",
    "                else:\n",
    "                    policy_loss += finish_episode(policy_net, baseline_net, optimizer)\n",
    "                del policy_net.rewards[:]\n",
    "                del policy_net.saved_log_probs[:]\n",
    "                del baseline_net.rewards[:]\n",
    "                del baseline_net.saved_log_probs[:]   \n",
    "                \n",
    "                if (i_episode%BATCH_SIZE)==(BATCH_SIZE-1):\n",
    "                    policy_loss /= 1.0*BATCH_SIZE\n",
    "                    policy_loss.backward()\n",
    "                    optimizer.step()\n",
    "                    optimizer.zero_grad()\n",
    "                    policy_loss = 0\n",
    "                if verbose:\n",
    "                    train_loss.append(tloss) \n",
    "                rewards.append(reward)\n",
    "                rewards.append(reward_greedy)\n",
    "\n",
    "                total_steps += 1\n",
    "                break\n",
    "                baseline\n",
    "            if OptimalDone:\n",
    "                taskF.get_best()\n",
    "                print(f'Complete with {selector.action_samples} policy actions')\n",
    "                     \n",
    "        # Update the target network, copying all weights and biases in DQN\n",
    "        if i_episode % TARGET_UPDATE == 0:\n",
    "            baseline_net.load_state_dict(policy_net.state_dict())\n",
    "    taskF.get_best()\n",
    "    optimizer.zero_grad()\n",
    "    print(f'Complete with {i_episode} policy actions accepts {selector.action_samples} and rejects {selector.rejects}')\n",
    "    if verbose:\n",
    "        return np.maximum.accumulate(rewards)\n",
    "#     , np.array(sequence_designs)\n",
    "    else:\n",
    "        return np.maximum.accumulate(rewards)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e842687",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "def batch_zero(batch, Nseq, Nclass, numzeros=1):    \n",
    "    replace = (torch.zeros_like(batch)*Nclass).long()\n",
    "    idx = torch.randint(Nseq,(batch.shape[0],))\n",
    "    oneh = F.one_hot(idx, Nseq)\n",
    "    \n",
    "    if numzeros>1:\n",
    "        for _ in range(numzeros-1):\n",
    "            idx = torch.randint(Nseq,(batch.shape[0],))\n",
    "            oneh += F.one_hot(idx, Nseq)\n",
    "            \n",
    "    neighs = torch.where(oneh==1, replace, batch)\n",
    "    # Remove Duplicates\n",
    "#     neighs = torch.unique(neighs, dim=0)\n",
    "    # Remove Observed\n",
    "    return neighs\n",
    "\n",
    "def batch_perturbation(Nseq, Nsamples, Nclass):\n",
    "    idx = torch.randint(Nseq,(Nsamples,))\n",
    "    replace = torch.randint_like(best_s, Nclass)\n",
    "    oneh = F.one_hot(idx, Nseq)\n",
    "    neighs = torch.where(oneh==1, replace, best_s)\n",
    "    # Remove Duplicates\n",
    "#     neighs = torch.unique(neighs, dim=0)\n",
    "    # Remove Observed\n",
    "    return neighs\n",
    "\n",
    "def optimize_structure(policy_net, memory, Nseq, Nclass, Nsamples=1000, BATCH_SIZE=32):\n",
    "    if len(memory) < BATCH_SIZE:\n",
    "        return torch.randint(Nclass,(1,Nseq)), torch.randint(Nclass,(1,Nseq))\n",
    "    #Greedy Search \n",
    "    transitions = memory.memory\n",
    "    batch = CombTransition(*zip(*transitions))\n",
    "    structure_batch = torch.cat(batch.structure)\n",
    "    rewards_idmax = torch.cat(batch.reward).squeeze().argmax()\n",
    "    best_s_single = structure_batch[rewards_idmax][None, :]\n",
    "    best_s = best_s_single.clone().repeat(Nsamples,1)\n",
    "    neighs = batch_perturbation(memory, Nseq, Nsamples, Nclass)\n",
    "    neighs[(structure_batch==neighs[:, None]).all(2).any(1).logical_not()]\n",
    "    with torch.no_grad():\n",
    "        npredsMAX = policy_net(neighs).squeeze().argmax()\n",
    "    return neighs[npredsMAX][None, :], best_s_single\n",
    "\n",
    "            \n",
    "        \n",
    "\n",
    "def run_sql(taskF, num_episodes = 1000, SeqSize=11, NumClass=20, BATCH_SIZE=32, train_freq=1, verbose=False, saveD=False, d_model=32):\n",
    "    total_accepts = 0\n",
    "    total_rejects = 0\n",
    "    # TARGET_UPDATE = 10\n",
    "    #Always add 1 to Q learning for padded class. \n",
    "    policyNet = TransformerNN(numberVars=NumClass, outputs=1, d_model=d_model, seqSize=SeqSize).to(device)\n",
    "    policyNet2 = None\n",
    "#     policyNet2 = TransformerNN(numberVars=NumClass, outputs=1, d_model=d_model, seqSize=SeqSize).to(device)\n",
    "    optimizer = optim.Adam(policyNet.parameters(), lr=1e-4)\n",
    "#     optimizer2 = optim.Adam(policyNet2.parameters(), lr=1e-4)\n",
    "    memory = CombReplayMemory(min(num_episodes, 1e5))\n",
    "    t_init = 1 \n",
    "    selector = StructureAction(num_episodes)\n",
    "    rewards = []\n",
    "    criterion = []\n",
    "    if verbose:\n",
    "        val_loss = []\n",
    "        train_loss = []\n",
    "    if saveD:\n",
    "        sequence_designs = []\n",
    "#     selector = Action()\n",
    "    total_steps = 0\n",
    "    for i_episode in range(num_episodes):\n",
    "        # Random Structure\n",
    "        structure, bests = optimize_structure(policyNet, memory, SeqSize, NumClass)\n",
    "        nstructure = perturbation(bests, SeqSize, NumClass)\n",
    "#         nstructure = torch.randint(NumClass,(1,SeqSize))\n",
    "        selector.select_action(structure, nstructure, policyNet, policyNet2, NumClass, MH=False)\n",
    "#         structure = torch.randint(NumClass,(1,SeqSize))\n",
    "        if saveD:\n",
    "            sequence_designs.append(structure)\n",
    "        reward, _ = taskF.reward(structure)\n",
    "        # Store the transition in memory\n",
    "        reward = torch.tensor([reward], device=device)\n",
    "        memory.push(structure, reward)\n",
    "#         if (i_episode%SeqSize)==3:\n",
    "        tloss = optimize_model_sql(BATCH_SIZE, policyNet, optimizer, memory)\n",
    "#         _ = optimize_model_sql(BATCH_SIZE, policyNet2, optimizer2, memory)\n",
    "        if verbose:\n",
    "            train_loss.append(tloss) \n",
    "            val_loss.append(val_sql_critic(taskF, policyNet, NumClass, SeqSize))\n",
    "        rewards.append(reward.detach().item())\n",
    "        \n",
    "        \n",
    "        \n",
    "                \n",
    "        if (i_episode%1000)==10:\n",
    "            taskF.get_best()\n",
    "                \n",
    "    taskF.get_best()\n",
    "    print(f'Complete with {i_episode} policy actions accepts {selector.action_samples} and rejects {selector.rejects}')\n",
    "    if verbose:\n",
    "        plot_rewards(rewards, val_loss, train_loss)\n",
    "        return np.maximum.accumulate(rewards)\n",
    "#     , np.array(sequence_designs)\n",
    "    else:\n",
    "        return np.maximum.accumulate(rewards)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a4938fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_zero(structure_batch, 11, 20, 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f6d4429",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "NumClass = 20\n",
    "SeqSize = 11\n",
    "task = SeqOptimisation(SeqSize, NumClass)\n",
    "rews = run_sql(task, num_episodes=2000, SeqSize=11, NumClass=20, BATCH_SIZE=32, verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c5a5ce6",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_all.model.unique()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d0707fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_sql = run_multiple_rl_seeds(df_all, nseeds=10, neps=2000, model='SQL', verbose=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f11a0c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_sql.model.unique()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3d1e55f",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_seaborn(df_sql, True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4925feb2",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_sql[df_sql.model=='SQL'].max()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "762ef490",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_rewards(rewards, val_loss, train_loss):\n",
    "    plt.figure(1)\n",
    "    plt.clf()\n",
    "    rewards_ep = torch.tensor(rewards, dtype=torch.float)\n",
    "    val_loss = torch.tensor(val_loss, dtype=torch.float)\n",
    "    plt.title('Training...')\n",
    "    plt.xlabel('Episode')\n",
    "    plt.ylabel('Max Rewards')\n",
    "    \n",
    "    plt.plot(np.maximum.accumulate(rewards_ep.numpy()), label='Max Reward')\n",
    "    plt.show()\n",
    "    \n",
    "    plt.figure(1)\n",
    "    plt.clf()\n",
    "    plt.title('Training...')\n",
    "    plt.xlabel('Episode')\n",
    "    plt.ylabel('Val Loss')\n",
    "    plt.plot(val_loss.numpy(), label='val loss')\n",
    "    plt.show()\n",
    "    \n",
    "    \n",
    "    plt.figure(1)\n",
    "    plt.clf()\n",
    "    plt.title('Training...')\n",
    "    plt.xlabel('Episode')\n",
    "    plt.ylabel('Train Loss')\n",
    "    plt.plot(np.array(train_loss), label='Train loss')\n",
    "    plt.show()\n",
    "    \n",
    "    plt.figure(1)\n",
    "    plt.clf()\n",
    "    plt.title('Training...')\n",
    "    plt.xlabel('Episode')\n",
    "    plt.ylabel('100 Step Average Reward')\n",
    "    # Take 100 episode averages and plot them too\n",
    "    if len(rewards_ep) >= 100:\n",
    "        means = rewards_ep.unfold(0, 100, 1).mean(1).view(-1)\n",
    "        means = torch.cat((torch.zeros(99), means))\n",
    "        plt.plot(means.numpy(), label='100 steps Reward Average')\n",
    "\n",
    "    plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}