BitFlip_My.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 116,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import datetime\n",
    "import os\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "import torch\n",
    "from torch import nn\n",
    "import torch.backends.cudnn as cudnn\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.tensorboard import SummaryWriter\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 117,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(42)\n",
    "random.seed(42)\n",
    "torch.manual_seed(42)\n",
    "torch.cuda.manual_seed(42)\n",
    "cudnn.benchmark = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 118,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BitFlip():\n",
    "    def __init__(self, bit_length=50, max_steps=50):\n",
    "        self.bit_length = bit_length\n",
    "        self.max_steps = max_steps\n",
    "        self._reset()\n",
    "\n",
    "    def _terminate(self):\n",
    "        if np.array_equal(self.state, self.goal):\n",
    "            return 1\n",
    "        else:\n",
    "            return 0  \n",
    "\n",
    "    def _reward(self):\n",
    "        if np.array_equal(self.state, self.goal):\n",
    "            return 0\n",
    "        else:\n",
    "            return -1  \n",
    "\n",
    "    def _step(self, action):\n",
    "        self.steps += 1\n",
    "\n",
    "        change = np.zeros(len(self.state))\n",
    "        change[action] = 1\n",
    "        if self.state[action] == 0:\n",
    "            self.state = self.state + change\n",
    "        elif  self.state[action] == 1:\n",
    "            self.state = self.state - change\n",
    "            \n",
    "        return self.state, self._reward(), self._terminate()\n",
    "\n",
    "    def _reset(self):\n",
    "        self.steps = 0\n",
    "        self.state = np.random.randint(2, size=(self.bit_length))\n",
    "\n",
    "        # make sure goal is not the initial state\n",
    "        self.goal = self.state\n",
    "        while np.array_equal(self.state, self.goal):\n",
    "            self.goal = np.random.randint(2, size=(self.bit_length))\n",
    "\n",
    "        return self.state, self.goal\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 119,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DQN(nn.Module):\n",
    "    def __init__(self, input_size, hidden_size, output_size):\n",
    "        super(DQN, self).__init__()\n",
    "\n",
    "        self.hidden = nn.Linear(input_size, hidden_size)\n",
    "        nn.init.kaiming_normal_(self.hidden.weight)\n",
    "        self.hidden.bias.data.zero_()\n",
    "\n",
    "        self.output = nn.Linear(hidden_size, output_size)\n",
    "        nn.init.kaiming_normal_(self.output.weight)\n",
    "        self.output.bias.data.zero_()\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = F.relu(self.hidden(x))\n",
    "        return self.output(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 120,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Agent():\n",
    "    def __init__(self, state_size, action_size, goal_size, model_out_path, clip_target_value=True):\n",
    "        self.state_size = state_size\n",
    "        self.goal_size = goal_size\n",
    "        self.action_size = action_size\n",
    "        self.clip_target_value = clip_target_value\n",
    "\n",
    "        # hyperparameters for the agent and HER\n",
    "        # as same as the paper: https://arxiv.org/pdf/1707.01495.pdf 's Appendix A\n",
    "        self.epsilon = 0.2  # exploration\n",
    "        self.epsilon_min = 0.05  # min exploration\n",
    "        self.epsilon_decay = 0.001\n",
    "        self.tau = 0.95  # target net update weight\n",
    "        self.gamma = 0.98\n",
    "        self.batch_size = 128\n",
    "        self.learning_rate = 1e-3\n",
    "        self.buffer_size = int(1e6)\n",
    "        self.hidden_size = 256\n",
    "\n",
    "        self.model_out_path = model_out_path\n",
    "        self.device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "        self.criterion = None\n",
    "        self.optimizer = None\n",
    "        self.writer = SummaryWriter(self.model_out_path + '/tensorboard')\n",
    "        self.build_model()\n",
    "\n",
    "\n",
    "    # build the double DQN\n",
    "    def build_model(self):  \n",
    "        print('     Building double DQN')\n",
    "        \n",
    "        self.DQN_eval = DQN(input_size=self.state_size + self.goal_size, hidden_size=self.hidden_size, output_size=self.action_size).to(self.device).train()\n",
    "        self.DQN_target = DQN(input_size=self.state_size + self.goal_size, hidden_size=self.hidden_size, output_size=self.action_size).to(self.device).eval()\n",
    "        self.DQN_target.load_state_dict(self.DQN_eval.state_dict())\n",
    "\n",
    "        self.criterion = nn.MSELoss().to(self.device)\n",
    "        self.optimizer = optim.Adam(self.DQN_eval.parameters(), lr=self.learning_rate)\n",
    "        \n",
    "        \n",
    "    def choose_action(self, state, goal):\n",
    "        if np.random.rand() <= self.epsilon:\n",
    "            return np.random.randint(self.action_size)\n",
    "        else:\n",
    "            state = torch.tensor(state)\n",
    "            goal =  torch.tensor(goal)\n",
    "            input = torch.concat((state,goal), 0).type(torch.FloatTensor).to(self.device)\n",
    "            act_values = self.DQN_eval(input) # tensor, len=50\n",
    "            \n",
    "            return torch.argmax(act_values).item() \n",
    "\n",
    "    def gen_batch(self, batchsize):\n",
    "        minibatch = []\n",
    "\n",
    "        for _ in range(batchsize):\n",
    "            state = np.random.randint(2, size=(self.state_size))\n",
    "            \n",
    "            goal = state\n",
    "            nums = np.random.randint(1, self.state_size+1)\n",
    "            indexes = random.sample(range(self.state_size), k=nums)\n",
    "            for index in indexes:\n",
    "                change = np.zeros(self.state_size)\n",
    "                change[index] = 1\n",
    "                if state[index] == 0:\n",
    "                    goal = goal + change\n",
    "                elif state[index] == 1:\n",
    "                    goal = goal - change\n",
    "            # print(np.array_equal(state, goal))\n",
    "            # print(np.sum(np.abs(state-goal)))\n",
    "\n",
    "            action = self.choose_action(state,goal)\n",
    "\n",
    "            change = np.zeros(self.state_size)\n",
    "            change[action] = 1\n",
    "            if state[action] == 0:\n",
    "                next_state = state + change\n",
    "            elif state[action] == 1:\n",
    "                next_state = state - change\n",
    "            # print(np.array_equal(state, next_state))\n",
    "\n",
    "            done = 1 if np.array_equal(next_state, goal) else 0\n",
    "\n",
    "            reward = 0 if np.array_equal(next_state, goal) else -1\n",
    "\n",
    "            minibatch += [(state, action, reward, next_state, done, goal)]\n",
    "\n",
    "        return minibatch\n",
    "\n",
    "    def replay(self, epoch, optimization_steps):\n",
    "        print('\\n===> Replaying experience')\n",
    "\n",
    "        losses = 0\n",
    "        for _ in range(optimization_steps):\n",
    "\n",
    "            # transform the list into multiple ndarrays\n",
    "            minibatch = self.gen_batch(self.batch_size) # list, len=128, element:state, action, reward, next_state, done, goal\n",
    "            ss = np.zeros((self.batch_size, self.state_size)) # state\n",
    "            ass = np.zeros(self.batch_size)\n",
    "            rs = np.zeros(self.batch_size) # reward\n",
    "            nss = np.zeros((self.batch_size, self.state_size)) # next_state\n",
    "            ds = np.zeros(self.batch_size) # done\n",
    "            gs = np.zeros((self.batch_size, self.state_size)) # goal\n",
    "            for i in range(self.batch_size):\n",
    "                ss[i] = minibatch[i][0]\n",
    "                ass[i] = minibatch[i][1]\n",
    "                rs[i] = minibatch[i][2]\n",
    "                nss[i] = minibatch[i][3]\n",
    "                ds[i] = minibatch[i][4]\n",
    "                gs[i] = minibatch[i][5]\n",
    "\n",
    "            # Q_EVAL(st, at)\n",
    "            ass = torch.tensor(ass).type(torch.int64).to(self.device).unsqueeze(1) # torch.Size([128, 1])\n",
    "            q_eval = torch.tensor(np.concatenate((ss, gs), axis=1)).type(torch.FloatTensor).to(self.device)\n",
    "            q_eval = self.DQN_eval(q_eval) # torch.Size([128, 50])\n",
    "            pred =  q_eval.gather(index = ass, dim = 1) # tensor, len=128\n",
    "            if self.clip_target_value and epoch >= 16:\n",
    "                pred = torch.clamp(pred, min=-1 / (1 - self.gamma), max=0) \n",
    "            pred = pred.squeeze(1) # tensor, len=128\n",
    "            # print(pred.shape)\n",
    "\n",
    "            # # Q_EVAL(st, at)\n",
    "            # # ass = torch.tensor(ass).type(torch.int64).to(self.device).unsqueeze(1) # torch.Size([128, 1])\n",
    "            # q_eval = torch.tensor(np.concatenate((ss, gs), axis=1)).type(torch.FloatTensor).to(self.device)\n",
    "            # q_eval = self.DQN_eval(q_eval) # torch.Size([128, 50])\n",
    "            # pred = torch.max(q_eval, dim=1)[0]\n",
    "            # # pred =  q_eval.gather(index = ass, dim = 1) # tensor, len=128\n",
    "            # if self.clip_target_value and epoch >= 16:\n",
    "            #     pred = torch.clamp(pred, min=-1 / (1 - self.gamma), max=0) \n",
    "            # # pred = pred.squeeze(1) # tensor, len=128\n",
    "            # # print(pred)\n",
    "            \n",
    "            # Q_TARGET(st+1, at+1)\n",
    "            q_target_next_state = torch.tensor(np.concatenate((nss, gs), axis=1)).type(torch.FloatTensor).to(self.device)\n",
    "            q_target_next_state = self.DQN_target(q_target_next_state)\n",
    "            q_target_next_state = torch.max(q_target_next_state, dim=1)[0] # tensor, len=128\n",
    "            rs = torch.tensor(rs).type(torch.FloatTensor).to(self.device)\n",
    "            ds = torch.tensor(ds).type(torch.FloatTensor).to(self.device)\n",
    "            target = rs + self.gamma * (1 - ds) * q_target_next_state\n",
    "            if self.clip_target_value and epoch >= 16:\n",
    "                target = torch.clamp(target, min=-1 / (1 - self.gamma), max=0) # tensor, len=128\n",
    "            # print(target)\n",
    "\n",
    "            # train DQN eval\n",
    "            self.optimizer.zero_grad()\n",
    "            mse_loss = self.criterion(pred, target)\n",
    "            mse_loss.backward()\n",
    "            self.optimizer.step()\n",
    "            losses += mse_loss.item()\n",
    "\n",
    "        self.writer.add_scalar(tag=\"loss\", scalar_value=losses / optimization_steps, global_step=epoch)\n",
    "        self.writer.add_scalar(tag=\"lr\", scalar_value=self.optimizer.state_dict()['param_groups'][0]['lr'], global_step=epoch)\n",
    "\n",
    "\n",
    "    def update_target_net(self, epoch, mode='RESET', decay=True):\n",
    "\n",
    "        if mode == 'EMA':\n",
    "            print('\\n===> Updating target DQN by EMA')\n",
    "            new_dict = {}\n",
    "            for key, t, e in zip(self.DQN_target.state_dict().keys(), self.DQN_target.state_dict().values(), self.DQN_eval.state_dict().values()):\n",
    "                new_dict[key] = self.tau * e + (1 - self.tau) * t # tau = 0.95\n",
    "            self.DQN_target.load_state_dict(new_dict)\n",
    "        elif mode == 'RESET':\n",
    "            print('\\n===> Updating target DQN by RESET')\n",
    "            self.DQN_target.load_state_dict(self.DQN_eval.state_dict())\n",
    "\n",
    "        print('\\n===> Saving target DQN')\n",
    "        checkpoint={\n",
    "            'epoch':epoch,\n",
    "            'net_state_dict':self.DQN_target.state_dict(),\n",
    "            'opt_state_dict':self.optimizer.state_dict(),\n",
    "                    }\n",
    "        checkpoint_out_path = self.model_out_path +'/checkpoint/'\n",
    "        torch.save(checkpoint, checkpoint_out_path + str(epoch) + '_checkpoint.pkl')\n",
    "        \n",
    "        if decay:\n",
    "            self.epsilon = max(self.epsilon - self.epsilon_decay, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 121,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "     Building double DQN\n",
      "\n",
      "===> Epoch 1 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 2 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 3 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 4 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 5 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 6 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 7 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 8 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 9 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 10 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 11 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 12 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 13 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 14 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 15 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 16 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 17 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 18 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 19 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 20 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 21 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 22 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 23 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 24 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 25 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 26 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 27 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 28 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 29 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 30 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 31 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 32 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 33 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 34 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 35 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 36 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 37 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 38 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 39 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 40 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 41 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 42 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 43 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 44 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 45 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 46 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 47 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 48 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 49 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 50 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 51 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 52 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 53 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 54 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 55 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 56 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 57 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 58 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 59 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 60 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 61 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 62 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 63 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 64 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 65 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 66 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 67 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 68 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 69 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 70 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 71 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 72 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 73 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 74 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 75 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 76 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 77 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 78 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 79 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 80 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 81 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 82 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 83 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 84 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 85 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 86 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 87 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 88 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 89 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 90 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 91 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 92 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 93 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 94 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 95 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 96 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 97 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 98 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 99 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 100 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 101 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 102 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 103 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 104 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 105 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 106 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 107 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 108 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 109 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 110 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 111 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 112 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 113 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 114 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 115 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 116 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 117 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 118 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 119 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 120 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 121 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 122 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n",
      "\n",
      "===> Epoch 123 starts\n",
      "\n",
      "===> Replaying experience\n",
      "\n",
      "===> Updating target DQN by EMA\n",
      "\n",
      "===> Saving target DQN\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[1;32m/home/guozy/RL_Project/BitFlip_My.ipynb Cell 6\u001b[0m in \u001b[0;36m4\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2B10.250.217.225/home/guozy/RL_Project/BitFlip_My.ipynb#W5sdnNjb2RlLXJlbW90ZQ%3D%3D?line=44'>45</a>\u001b[0m done \u001b[39m=\u001b[39m \u001b[39m0\u001b[39m\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2B10.250.217.225/home/guozy/RL_Project/BitFlip_My.ipynb#W5sdnNjb2RlLXJlbW90ZQ%3D%3D?line=45'>46</a>\u001b[0m \u001b[39mfor\u001b[39;00m t \u001b[39min\u001b[39;00m \u001b[39mrange\u001b[39m(n):\n\u001b[0;32m---> <a href='vscode-notebook-cell://ssh-remote%2B10.250.217.225/home/guozy/RL_Project/BitFlip_My.ipynb#W5sdnNjb2RlLXJlbW90ZQ%3D%3D?line=46'>47</a>\u001b[0m     action \u001b[39m=\u001b[39m agent\u001b[39m.\u001b[39;49mchoose_action(state, goal)\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2B10.250.217.225/home/guozy/RL_Project/BitFlip_My.ipynb#W5sdnNjb2RlLXJlbW90ZQ%3D%3D?line=47'>48</a>\u001b[0m     next_state, reward, done \u001b[39m=\u001b[39m env\u001b[39m.\u001b[39m_step(action) \n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2B10.250.217.225/home/guozy/RL_Project/BitFlip_My.ipynb#W5sdnNjb2RlLXJlbW90ZQ%3D%3D?line=48'>49</a>\u001b[0m     distances \u001b[39m+\u001b[39m\u001b[39m=\u001b[39m np\u001b[39m.\u001b[39msum(np\u001b[39m.\u001b[39mabs(next_state\u001b[39m-\u001b[39mgoal)) \u001b[39m# check the progresss of DQN\u001b[39;00m\n",
      "\u001b[1;32m/home/guozy/RL_Project/BitFlip_My.ipynb Cell 6\u001b[0m in \u001b[0;36m4\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2B10.250.217.225/home/guozy/RL_Project/BitFlip_My.ipynb#W5sdnNjb2RlLXJlbW90ZQ%3D%3D?line=43'>44</a>\u001b[0m state \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39mtensor(state)\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2B10.250.217.225/home/guozy/RL_Project/BitFlip_My.ipynb#W5sdnNjb2RlLXJlbW90ZQ%3D%3D?line=44'>45</a>\u001b[0m goal \u001b[39m=\u001b[39m  torch\u001b[39m.\u001b[39mtensor(goal)\n\u001b[0;32m---> <a href='vscode-notebook-cell://ssh-remote%2B10.250.217.225/home/guozy/RL_Project/BitFlip_My.ipynb#W5sdnNjb2RlLXJlbW90ZQ%3D%3D?line=45'>46</a>\u001b[0m \u001b[39minput\u001b[39m \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39;49mconcat((state,goal), \u001b[39m0\u001b[39;49m)\u001b[39m.\u001b[39mtype(torch\u001b[39m.\u001b[39mFloatTensor)\u001b[39m.\u001b[39mto(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mdevice)\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2B10.250.217.225/home/guozy/RL_Project/BitFlip_My.ipynb#W5sdnNjb2RlLXJlbW90ZQ%3D%3D?line=46'>47</a>\u001b[0m act_values \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mDQN_eval(\u001b[39minput\u001b[39m) \u001b[39m# tensor, len=50\u001b[39;00m\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2B10.250.217.225/home/guozy/RL_Project/BitFlip_My.ipynb#W5sdnNjb2RlLXJlbW90ZQ%3D%3D?line=48'>49</a>\u001b[0m \u001b[39mreturn\u001b[39;00m torch\u001b[39m.\u001b[39margmax(act_values)\u001b[39m.\u001b[39mitem()\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "if __name__ == \"__main__\":\n",
    "\n",
    "    # to store different model information\n",
    "    current_path = os.getcwd()\n",
    "    result_path = current_path + '/result/'\n",
    "    if os.path.exists(result_path) == False:\n",
    "        os.mkdir(result_path)\n",
    "    now = datetime.datetime.now()\n",
    "    now = now.strftime(\"%Y-%m-%d_%H:%M:%S\")\n",
    "    model_out_path =  result_path + now\n",
    "    if os.path.exists(model_out_path) == False:\n",
    "        os.mkdir(model_out_path)\n",
    "    checkpoints_out_path = model_out_path +'/checkpoint/'\n",
    "    if os.path.exists(checkpoints_out_path) == False:\n",
    "        os.mkdir(checkpoints_out_path)\n",
    "    writer_out_path = model_out_path + '/tensorboard'\n",
    "    if os.path.exists(writer_out_path) == False:\n",
    "        os.mkdir(writer_out_path)\n",
    "    \n",
    "    # hyperparameters for the agent and HER\n",
    "    n = 35\n",
    "    num_epochs = 200\n",
    "    num_cycles = 20\n",
    "    num_episodes = 1000\n",
    "    optimization_steps = 40\n",
    "\n",
    "    env = BitFlip(bit_length=n, max_steps=n)\n",
    "    agent = Agent(state_size=n, action_size=n, goal_size=n, model_out_path=model_out_path, clip_target_value=True)\n",
    "\n",
    "\n",
    "    for epoch in range(1,num_epochs*num_cycles+1):\n",
    "        print('\\n===> Epoch {} starts'.format(epoch))\n",
    "\n",
    "        # step 1: train the eval DQN\n",
    "        agent.replay(epoch=epoch, optimization_steps=optimization_steps)\n",
    "\n",
    "        # step 2: update and save the target DQN\n",
    "        agent.update_target_net(epoch=epoch, mode='EMA', decay=False)\n",
    "\n",
    "        # step 3: test the target DQN\n",
    "        successes = 0\n",
    "        distances = 0\n",
    "        for eposode in range(1, num_episodes+1):\n",
    "            state, goal = env._reset() \n",
    "            done = 0\n",
    "            for t in range(n):\n",
    "                action = agent.choose_action(state, goal)\n",
    "                next_state, reward, done = env._step(action) \n",
    "                distances += np.sum(np.abs(next_state-goal)) # check the progresss of DQN\n",
    "                state = next_state\n",
    "                if np.array_equal(state, goal):\n",
    "                    successes += 1\n",
    "                    break\n",
    "\n",
    "        # step 5: store the history\n",
    "        agent.writer.add_scalar(tag=\"distance\",scalar_value=distances / num_episodes / n, global_step=epoch)\n",
    "        agent.writer.add_scalar(tag=\"success rate\",scalar_value=successes / num_episodes, global_step=epoch)\n",
    "        agent.writer.add_scalar(tag=\"epsilon\",scalar_value=agent.epsilon, global_step=epoch)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cyan_torch",
   "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.8.15"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}