RL_Object_Detection3.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import torch.utils.data as data\n",
    "from torchvision import models, datasets, transforms\n",
    "from torch.autograd import Variable\n",
    "import numpy as np\n",
    "import os\n",
    "import collections\n",
    "import pandas as pd\n",
    "import random\n",
    "from PIL import Image, ImageDraw\n",
    "import matplotlib.image as mpimg\n",
    "import matplotlib.pyplot as plt\n",
    "from IoU import IoU\n",
    "from Network import ActorNetwork, CriticNetwork\n",
    "from replay_buffer import ReplayBuffer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def ReadData(csv_dir, img_dir, image_sets , test = None):\n",
    "    imgs = []\n",
    "    for y, cls in enumerate(image_sets):\n",
    "        if test is not None:\n",
    "            train_filename = csv_dir+'/val_' + cls + '.csv'\n",
    "        else:\n",
    "            train_filename = csv_dir + '/train_' + cls+ '.csv'\n",
    "\n",
    "        data = pd.read_csv(train_filename)\n",
    "        name = data.fname.values\n",
    "        xmin, ymin, xmax, ymax = data.xmin.values, data.ymin.values, data.xmax.values,\\\n",
    "                                        data.ymax.values\n",
    "        for i, entry in enumerate(name):\n",
    "            img = img_dir+'/'+entry\n",
    "            imgs += [(img, y, xmin[i], ymin[i], xmax[i], ymax[i])]      \n",
    "    return imgs\n",
    "\n",
    "root_dir = '/home/minty/dataset/VOCdevkit/VOC2012/'\n",
    "csv_dir = os.path.join(root_dir,'csvs')\n",
    "img_dir = os.path.join(root_dir,'JPEGImages')\n",
    "image_sets =  ['aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat',\\\n",
    "      'chair', 'cow', 'diningtable', 'dog', 'horse', 'motorbike', 'person', \\\n",
    "      'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor']\n",
    "imgs = ReadData(csv_dir,img_dir,image_sets, test = None)\n",
    "#test_imgs, test_box = ReadData(csv_dir, img_dir, image_sets, test = True)\n",
    "\n",
    "input_transform = transforms.Compose([\n",
    "   #transforms.RandomSizedCrop(300),\n",
    "    transforms.RandomHorizontalFlip(),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(mean = [ 0.485, 0.456, 0.406 ],\n",
    "                        std = [ 0.229, 0.224, 0.225 ]),\n",
    "])\n",
    "\n",
    "class Data(data.Dataset):\n",
    "    def __init__(self, imgs, image_sets, input_transform = None, \\\n",
    "                 target_transform = None ,test = None):\n",
    "        self.test = test\n",
    "        self.input_transform = input_transform\n",
    "        self.target_transform = target_transform\n",
    "        self.classes = len(image_sets)\n",
    "        self.imgs = imgs\n",
    "                 \n",
    "    def __getitem__(self, index):\n",
    "        fn, y, xmin, ymin, xmax, ymax = self.imgs[index]\n",
    "        #y_ = torch.LongTensor([[y]])\n",
    "        #print y_\n",
    "        # One hot encoding buffer that you create out of the loop and just keep reusing\n",
    "        #label = torch.LongTensor(1, self.classes)\n",
    "\n",
    "        # In your for loop\n",
    "        #label.zero_()\n",
    "        #label.scatter_(1, y_, 1)\n",
    "        #print label\n",
    "        label = torch.LongTensor([[y]])\n",
    "        img = Image.open(fn).convert('RGB')\n",
    "        box = torch.LongTensor([[xmin, ymin, xmax, ymax]])\n",
    "        #print img.size\n",
    "        #img = img.resize((300, 300))\n",
    "        if self.input_transform is not None:\n",
    "            img = self.input_transform(img)\n",
    "            #print 'transformed'\n",
    "        return img, box, label, fn\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.imgs)\n",
    "\n",
    "img_data = Data(imgs, image_sets,  input_transform = input_transform, target_transform=None\\\n",
    "                , test=None)\n",
    "#test_data = Data(test_imgs, image_sets,  input_transform = input_transform, target_transform=None\\\n",
    "              #  , test=True)\n",
    "print(len(img_data.imgs))\n",
    "#print(len(test_data.imgs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "( 0 , 0 ,.,.) = \n",
      " -0.4054 -0.3541 -0.3541  ...  -1.5528 -1.6213 -1.6213\n",
      " -0.4226 -0.3712 -0.3369  ...  -1.5185 -1.5185 -1.5014\n",
      " -0.4054 -0.3883 -0.3369  ...  -1.3987 -1.4158 -1.3987\n",
      "           ...             ⋱             ...          \n",
      " -0.4739 -0.8335 -0.2856  ...  -0.9363 -1.5528 -1.7754\n",
      " -0.5424 -0.3883 -0.3027  ...  -0.8678 -1.5357 -1.4672\n",
      " -0.5938 -0.5767 -0.3198  ...  -1.0390 -1.3302 -1.0733\n",
      "\n",
      "( 0 , 1 ,.,.) = \n",
      " -0.2850 -0.2850 -0.2850  ...  -1.9482 -1.9832 -1.9832\n",
      " -0.3025 -0.3025 -0.2675  ...  -2.0182 -2.0357 -2.0182\n",
      " -0.2850 -0.2675 -0.2675  ...  -2.0182 -2.0357 -2.0182\n",
      "           ...             ⋱             ...          \n",
      " -1.1604 -1.6331 -1.3354  ...  -0.8627 -1.5280 -1.7731\n",
      " -1.1078 -1.1078 -1.3004  ...  -0.7752 -1.4405 -1.3529\n",
      " -1.2129 -1.1954 -1.1429  ...  -0.9503 -1.2304 -0.9328\n",
      "\n",
      "( 0 , 2 ,.,.) = \n",
      " -0.2358 -0.2184 -0.2184  ...  -1.7347 -1.7347 -1.6999\n",
      " -0.2532 -0.2358 -0.2010  ...  -1.7696 -1.7347 -1.6824\n",
      " -0.2358 -0.2184 -0.2010  ...  -1.7173 -1.6999 -1.6476\n",
      "           ...             ⋱             ...          \n",
      " -1.1944 -1.6302 -1.2467  ...  -0.5844 -1.3164 -1.6476\n",
      " -1.1770 -1.1247 -1.2293  ...  -0.4624 -1.2119 -1.2293\n",
      " -1.2641 -1.2467 -1.1247  ...  -0.6193 -0.9678 -0.7936\n",
      "[torch.FloatTensor of size 1x3x375x500]\n",
      " \n",
      "(0 ,.,.) = \n",
      "   84   89  123  143\n",
      "[torch.LongTensor of size 1x1x4]\n",
      " \n",
      "(0 ,.,.) = \n",
      "  15\n",
      "[torch.LongTensor of size 1x1x1]\n",
      " ('/home/minty/dataset/VOCdevkit/VOC2012/JPEGImages/2011_001466.jpg',)\n"
     ]
    }
   ],
   "source": [
    "img_batch = data.DataLoader(img_data, batch_size=1 ,shuffle=True, num_workers = 2)\n",
    "img, box, label , fn = next(iter(img_batch))\n",
    "print img, box, label, fn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "resnet = torch.load('resnet184classification.pth')\n",
    "\n",
    "\n",
    "model = torch.nn.Sequential(\n",
    "               resnet.conv1,\n",
    "               resnet.bn1,\n",
    "               resnet.relu,\n",
    "               resnet.maxpool,\n",
    "               resnet.layer1,\n",
    "               resnet.layer2)\n",
    "#model = resnet.conv1\n",
    "#print model(Variable(img)).size()\n",
    "Extractor = nn.Sequential(*list(resnet.children())[:-2])\n",
    "#print Extractor(Variable(img))\n",
    "#have the same result\n",
    "#print resnet.conv1(Variable(img)).size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def Transition(box):\n",
    "    x_img = img[:, :, box[0]:box[2]+1, box[1]:box[3]+1]\n",
    "    return x_img\n",
    "#print Transition(box).size()\n",
    "#box_img = Extractor(Variable(Transition(box)))\n",
    "#print box_img.size()\n",
    "\n",
    "def roi_pooling(box_img):\n",
    "    stride = [int(np.floor(box_img.size(2)/2.0)), int(np.floor(box_img.size(3)/2.0))]\n",
    "    kernel_size = [int(np.ceil(box_img.size(2)/2.0)), int(np.ceil(box_img.size(3)/2.0))]\n",
    "    padding = [0,0]\n",
    "    if box_img.size(2) == 1:\n",
    "        padding[0] = 1\n",
    "        stride[0] = 1\n",
    "        kernel_size[0] = 2\n",
    "    if box_img.size(3) == 1:\n",
    "        padding[1] = 1\n",
    "        stride[1] = 1\n",
    "        kernel_size[1] = 2\n",
    "    padding = tuple(padding)\n",
    "    stride = (stride[0], stride[1])\n",
    "    kernel_size = (kernel_size[0], kernel_size[1])\n",
    "    roi_pooling = nn.MaxPool2d(kernel_size=kernel_size, stride=stride,  \\\n",
    "                               padding = padding, dilation=(1, 1))\n",
    "    #print roi_pooling(box_img).size()\n",
    "    return roi_pooling(box_img)\n",
    "#print roi_pooling(box_img).size()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def Cat_State(box_img, img):\n",
    "    state1 = roi_pooling(box_img).view(-1)\n",
    "    state2 = roi_pooling(Extractor(Variable(img))).view(-1)\n",
    "    state = torch.cat((state1, state2)).unsqueeze(0)\n",
    "    return state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#hyper-parameters\n",
    "tau = 0.001\n",
    "actor_lr = 0.01\n",
    "critic_lr = 0.001\n",
    "gamma = 0.99\n",
    "buffer_size = 100\n",
    "minibatch_size = 8\n",
    "max_episodes = 5\n",
    "max_steps = 10\n",
    "#concatenation of the feature vectors of global image(2048) and box image(2048)\n",
    "state_dim = 4096 \n",
    "action_dim = 9\n",
    "Penalty = 0.1\n",
    "random_seed = 1234"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "actor = ActorNetwork(state_dim, action_dim)\n",
    "critic = CriticNetwork(state_dim, action_dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "A_STEP = 20\n",
    "SCALE_STEP = 0.4\n",
    "#step may be a variable?\n",
    "class Action(object):\n",
    "    def __init__(self, box, img):\n",
    "        self.x_step = img.size(3) / A_STEP\n",
    "        self.y_step = img.size(2) / A_STEP\n",
    "        self.enlarge = 1 + SCALE_STEP\n",
    "        self.shrink = 1 - SCALE_STEP\n",
    "        self.xmin, self.ymin, self.xmax, self.ymax = box\n",
    "        self.center_x = (self.xmax + self.xmin) / 2\n",
    "        self.center_y = (self.ymax + self.ymin) / 2\n",
    "        self.width = self.xmax - self.xmin\n",
    "        self.height = self.ymax - self.ymin\n",
    "        self.narrow = self.width / 5\n",
    "        \n",
    "    def Up(self):\n",
    "        ymin = self.ymin - self.y_step\n",
    "        ymax = self.ymax - self.y_step\n",
    "        new_box = torch.LongTensor([[self.xmin, ymin, self.xmax, ymax]])\n",
    "        #print self.new_box\n",
    "        return new_box, None\n",
    "        \n",
    "    def Down(self):\n",
    "        ymin = self.ymin + self.y_step\n",
    "        ymax = self.ymax + self.y_step\n",
    "        new_box = torch.LongTensor([[self.xmin, ymin, self.xmax, ymax]])\n",
    "        return new_box, None\n",
    "    \n",
    "    def Left(self):\n",
    "        xmin = self.xmin - self.x_step\n",
    "        xmax = self.xmax - self.x_step\n",
    "        new_box = torch.LongTensor([[xmin, self.ymin, xmax, self.ymax]])\n",
    "        return new_box, None\n",
    "    \n",
    "    def Right(self):\n",
    "        xmin = self.xmin + self.x_step\n",
    "        xmax = self.xmax + self.x_step\n",
    "        new_box = torch.LongTensor([[xmin, self.ymin, xmax, self.ymax]])\n",
    "        return new_box, None\n",
    "    \n",
    "    def Enlarge(self):\n",
    "        xmin = self.center_x - int(self.enlarge * self.width) / 2\n",
    "        xmax = self.center_x + int(self.enlarge * self.width) / 2\n",
    "        ymin = self.center_y - int(self.enlarge * self.height) / 2\n",
    "        ymax = self.center_y + int(self.enlarge * self.height) /2\n",
    "        new_box = torch.LongTensor([[xmin, ymin, xmax, ymax]])\n",
    "        return new_box, None\n",
    "    \n",
    "    def Shrink(self):\n",
    "        xmin = self.center_x - int(self.shrink * self.width) / 2\n",
    "        xmax = self.center_x + int(self.shrink * self.width) / 2\n",
    "        ymin = self.center_y - int(self.shrink * self.height) / 2\n",
    "        ymax = self.center_y + int(self.shrink * self.height) / 2\n",
    "        new_box = torch.LongTensor([[xmin, ymin, xmax, ymax]])\n",
    "        return new_box, None\n",
    "    \n",
    "    def Narrow(self):\n",
    "        xmin = self.xmin + self.narrow\n",
    "        xmax = self.xmax - self.narrow\n",
    "        new_box = torch.LongTensor([[xmin, self.ymin, xmax, self.ymax]])\n",
    "        return new_box, None\n",
    "    \n",
    "    def Stretch(self):\n",
    "        xmin = self.xmin - self.narrow\n",
    "        xmax = self.xmax + self.narrow\n",
    "        new_box = torch.LongTensor([[xmin, self.ymin, xmax, self.ymax]])\n",
    "        return new_box, None\n",
    "    \n",
    "    def Trigger(self):\n",
    "        new_box,_ = Initial(img)\n",
    "        return new_box, True\n",
    "       "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def Sample(box_, box, trigger, Trigger, Steps,img):\n",
    "    terminal = None\n",
    "    reward = 0    \n",
    "    while True:\n",
    "        width = box[2] - box[0]\n",
    "        height = box[3] - box[1]\n",
    "        area = width *  height\n",
    "        if area <= 20 or height < 5:\n",
    "            box = Action(box,img).Enlarge()[0].squeeze(0)\n",
    "            reward -= Penalty\n",
    "        if width < 5:\n",
    "            box = Action(box, img).Stretch()[0].squeeze(0)\n",
    "            reward -= Penalty\n",
    "        #penalty for surpassing boundary\n",
    "        if box[0] <= action_bound[0]:\n",
    "            box = Action(box, img).Right()[0].squeeze(0)\n",
    "            reward -= Penalty\n",
    "        if box[1] <= action_bound[1]:\n",
    "            box = Action(box, img).Down()[0].squeeze(0)\n",
    "            reward -= Penalty\n",
    "        if box[2] >= action_bound[2]:\n",
    "            box = Action(box, img).Left()[0].squeeze(0)\n",
    "            reward -= Penalty\n",
    "        if box[3] >= action_bound[3]:\n",
    "            box = Action(box, img).Up()[0].squeeze(0)\n",
    "            reward -= Penalty\n",
    "        if area > 20 and box[0] > action_bound[0] and box[1] > action_bound[1] and box[2] < action_bound[2] and box[3] < action_bound[3]:\n",
    "            break\n",
    "        \n",
    "    reward += np.sign(IoU(box, ground_truth) - IoU(box_, ground_truth))\n",
    "    \n",
    "    box_img = Extractor(Variable(Transition(box)))\n",
    "    next_state = Cat_State(box_img,img)\n",
    "    if trigger is not None:\n",
    "        if IoU(box, ground_truth) > 0.6:\n",
    "            reward += 3.0\n",
    "        else:\n",
    "            reward -= 3.0\n",
    "        Trigger += 1\n",
    "    else:\n",
    "        Steps += 1\n",
    "    if Trigger >= 4 or Steps >=10:\n",
    "        terminal = True\n",
    "        Trigger = 0\n",
    "        Steps = 0\n",
    "    else:\n",
    "        terminal = None\n",
    "    return next_state, reward, terminal, Trigger, Steps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def Initial(img):\n",
    "    num = random.randint(0, 5)\n",
    "    box = torch.LongTensor([0,  0, img.size(3), img.size(2)])\n",
    "    center_x = box[2] / 2\n",
    "    center_y = box[3] / 2\n",
    "    width = box[2] - box[0]\n",
    "    height = box[3] - box[1]\n",
    "    if num == 0:#center\n",
    "        box[0] = center_x - width / 4\n",
    "        box[2] = center_x + width / 4\n",
    "        box[1] = center_y - height / 4\n",
    "        box[3] = center_y + height / 4\n",
    "    elif num == 1: #upper left corner\n",
    "        box[0] = width / 20\n",
    "        box[1] = height / 20\n",
    "        box[2] = width / 2\n",
    "        box[3] = height / 2      \n",
    "    elif num == 2: #bottom left corner\n",
    "        box[0] = width / 20\n",
    "        box[3] = box[3] - height / 20\n",
    "        box[2] = width / 2\n",
    "        box[1] = height / 2\n",
    "    elif num == 3: #upper right corner\n",
    "        box[1] = height / 20\n",
    "        box[2] = box[2] - width / 20\n",
    "        box[0] = width / 2\n",
    "        box[3] = height / 2\n",
    "    else: #bottom right corner\n",
    "        box[2] = box[2] - width / 20\n",
    "        box[3] = box[3] - height / 20\n",
    "        box[0] = width / 2\n",
    "        box[1] = height / 2\n",
    "        \n",
    "    box_img = Extractor(Variable(Transition(box)))\n",
    "    state = Cat_State(box_img,img)\n",
    "    return box, state\n",
    "#action_list = ['Up', 'Down', 'Left', 'Right', 'Enlarge' , 'Shrink', 'Narrow', 'Stretch', 'Trigger']\n",
    "def Generate(a, box, img):\n",
    "    if a == 0:\n",
    "        return Action(box, img).Up()\n",
    "    elif a == 1:\n",
    "        return Action(box, img).Down()\n",
    "    elif a == 2:\n",
    "        return Action(box, img).Left()\n",
    "    elif a == 3:\n",
    "        return Action(box, img).Right()\n",
    "    elif a == 4:\n",
    "        return Action(box, img).Enlarge()\n",
    "    elif a == 5:\n",
    "        return Action(box, img).Shrink()\n",
    "    elif a == 6:\n",
    "        return Action(box, img).Narrow()\n",
    "    elif a == 7:\n",
    "        return Action(box, img).Stretch()\n",
    "    else:\n",
    "        return Action(box, img).Trigger()\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def CreateDict(target, net):\n",
    "    new_state_dict = collections.OrderedDict()\n",
    "    #critic_target.load_state_dict( (1 - tau) * critic_target.state_dict() + tau * critic.state_dict())\n",
    "    keys = net.state_dict().keys()\n",
    "    for item in keys:\n",
    "        params1 = net.state_dict()[item]\n",
    "        params2 = target.state_dict()[item]\n",
    "        new_state_dict[item] = (1 - tau) * params2 + tau * params1\n",
    "    return new_state_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def CheckGradientUpdate(target, net):\n",
    "    for param, shared_param in zip(net.parameters(),\n",
    "                                                           target.parameters()):\n",
    "        if shared_param.grad is None:\n",
    "            shared_param = param.grad\n",
    "        else:\n",
    "            shared_param.grad = (1 - tau) * shared_param.grad + tau * param.grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"optimizer = optim.Adam([\\n                {'params': actor.parameters()},\\n                {'params': critic.parameters(), 'lr': critic_lr}\\n            ], lr = actor_lr)\""
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#pipeline\n",
    "replay_buffer = ReplayBuffer(buffer_size, random_seed)\n",
    "optimizer_a = optim.Adam(actor.parameters(), lr = actor_lr)\n",
    "actor_target = ActorNetwork(state_dim, action_dim)\n",
    "critic_target = CriticNetwork(state_dim, action_dim)\n",
    "optimizer_c = optim.Adam(critic.parameters(), lr = critic_lr)\n",
    "loss_func = nn.MSELoss().cuda()\n",
    "\n",
    "actor.cuda()\n",
    "critic.cuda()\n",
    "actor_target.cuda()\n",
    "critic_target.cuda()\n",
    "loss=[]\n",
    "\n",
    "\"\"\"optimizer = optim.Adam([\n",
    "                {'params': actor.parameters()},\n",
    "                {'params': critic.parameters(), 'lr': critic_lr}\n",
    "            ], lr = actor_lr)\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "img, box, label ,fn= next(iter(img_batch))\n",
    "ground_truth = box[0,0,:]\n",
    "Trigger = 0\n",
    "Steps = 0\n",
    "box_, s = Initial(img)\n",
    "action_bound = torch.LongTensor([0,0,img.size(3),img.size(2)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "print box_, img.size(), ground_truth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "a_output = actor(s.cuda())\n",
    "_ , action = torch.max(a_output.data,1)\n",
    "print action\n",
    "box, trigger = Generate(action[0], box_, img)\n",
    "s2, r, t, Trigger, Steps = Sample(box_.squeeze(0), box.squeeze(0), trigger, Trigger, Steps)\n",
    "\n",
    "target = torch.autograd.Variable(torch.randn(1, 1).unsqueeze(0))\n",
    "output = critic(s.cuda(), a_output)\n",
    "print output.data\n",
    "#update target networks\n",
    "critic_target.load_state_dict(CreateDict(critic_target, critic))\n",
    "actor_target.load_state_dict(CreateDict(actor_target, actor))\n",
    "\n",
    "optimizer_a.zero_grad()\n",
    "optimizer_c.zero_grad()\n",
    "#optimizer.zero_grad()\n",
    "\n",
    "loss_ = loss_func(output,target.cuda())\n",
    "loss_.backward(retain_graph = True)\n",
    "optimizer_c.step()\n",
    "optimizer_a.step()\n",
    "#optimizer.step()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "target = torch.autograd.Variable(torch.randn(1, 1).unsqueeze(0))\n",
    "output = critic(s.cuda(), a_output)\n",
    "print output.data\n",
    "#update target networks\n",
    "critic_target.load_state_dict(CreateDict(critic_target, critic))\n",
    "actor_target.load_state_dict(CreateDict(actor_target, actor))\n",
    "\n",
    "optimizer_a.zero_grad()\n",
    "optimizer_c.zero_grad()\n",
    "#optimizer.zero_grad()\n",
    "\n",
    "loss_ = loss_func(output,target.cuda())\n",
    "loss_.backward(retain_graph = True)\n",
    "optimizer_c.step()\n",
    "optimizer_a.step()\n",
    "#optimizer.step()\n",
    "#print('state: {}, action: {}'.format(s, a_output))\n",
    "_ , action = torch.max(actor(s.cuda()).data,1)\n",
    "#print actor.state_dict()['model.0.weight']\n",
    "#print actor.state_dict()['model.0.bias']\n",
    "print 'action:', action\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "print predicted_q_value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for i in xrange(5):\n",
    "    img, box, label ,fn= next(iter(img_batch))\n",
    "    ground_truth = box[0,0,:]\n",
    "    Trigger = 0\n",
    "    Steps = 0\n",
    "    box_, s = Initial(img)\n",
    "    action_bound = torch.LongTensor([0, 0, img.size(3),img.size(2)])\n",
    "    count = 0\n",
    "\n",
    "    while True:\n",
    "        count = count + 1\n",
    "        a_output = actor(s.cuda())\n",
    "        _ , action = torch.max(a_output.data,1)\n",
    "        print action\n",
    "        \n",
    "        box, trigger = Generate(action[0], box_, img)\n",
    "        s2, r, t, Trigger, Steps = Sample(box_.squeeze(0), box.squeeze(0), trigger, Trigger, Steps,img)\n",
    "       # if Trigger != 0:\n",
    "            \n",
    "        replay_buffer.add(s, a_output, r , t, s2)\n",
    "        if replay_buffer.size() > minibatch_size:\n",
    "            s_batch, a_batch, r_batch, t_batch, s2_batch \\\n",
    "            = replay_buffer.sample_batch(minibatch_size)\n",
    "            \n",
    "            r_batch_tensor = torch.FloatTensor(minibatch_size)\n",
    "            y = Variable(torch.FloatTensor(minibatch_size))\n",
    "            \n",
    "            for k in range(minibatch_size):\n",
    "                s2 = s2_batch[k].unsqueeze(0)\n",
    "                target_a = actor_target(s2.cuda())\n",
    "                target_q = critic_target(s2.cuda(), target_a)\n",
    "                target_q = target_q.squeeze(0)\n",
    "                r_batch_tensor[k] = r_batch[k]\n",
    "\n",
    "                if t_batch[k] is not None:\n",
    "                    y[k] = r_batch_tensor[k]\n",
    "                else:\n",
    "                    y[k] = r_batch_tensor[k] + gamma * target_q\n",
    "            \n",
    "            for m in range(minibatch_size):\n",
    "                state = s_batch[m].unsqueeze(0)\n",
    "                state = state.cuda()\n",
    "                #action = a_batch[8 * m: 8 * (m+1)].unsqueeze(0)\n",
    "                action = actor(state)\n",
    "                predicted_q_value = y[m].cuda()\n",
    "                target = torch.autograd.Variable(predicted_q_value.data.unsqueeze(0))\n",
    "\n",
    "                output = critic(state, action)\n",
    "                #update target networks\n",
    "                critic_target.load_state_dict(CreateDict(critic_target, critic))\n",
    "                actor_target.load_state_dict(CreateDict(actor_target, actor))\n",
    "\n",
    "                optimizer_a.zero_grad()\n",
    "                optimizer_c.zero_grad()\n",
    "                #optimizer.zero_grad()\n",
    "\n",
    "                loss_ = loss_func(output,target)\n",
    "                loss_.backward(retain_graph = True)\n",
    "                # check the target networks\n",
    "                CheckGradientUpdate(actor_target, actor)\n",
    "                CheckGradientUpdate(critic_target, critic)\n",
    "                torch.nn.utils.clip_grad_norm(actor.parameters(), 50)\n",
    "                optimizer_c.step()\n",
    "                optimizer_a.step()\n",
    "                #optimizer.step()\n",
    "                #print output.data\n",
    "\n",
    "            loss.append(loss_.data[0])\n",
    "            print('step: {}, loss: {}'.format(count + 1, loss_.data[0]))\n",
    "        if t is True:\n",
    "            print \"one img terminated\"\n",
    "            break\n",
    "        s = s2\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class Net(nn.Module):\n",
    "    def __init__(self, ):\n",
    "        super(Net, self).__init__()\n",
    "        self.fc1 = nn.Linear(state_dim, 10)\n",
    "        self.fc1.weight.data.normal_(0, 0.1)   # initialization\n",
    "        self.out = nn.Linear(10, action_dim)\n",
    "        self.out.weight.data.normal_(0, 0.1)   # initialization\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.fc1(x)\n",
    "        x = F.relu(x)\n",
    "        actions_value = self.out(x)\n",
    "        return actions_value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "print box_,s\n",
    "print img.size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class DQN(object):\n",
    "    def __init__(self):\n",
    "        self.eval_net, self.target_net = Net().cuda(), Net().cuda()\n",
    "\n",
    "        self.learn_step_counter = 0                                     # for target updating\n",
    "        self.memory_counter = 0                                         # for storing memory\n",
    "        self.memory = Variable(torch.zeros(50, state_dim * 2 + 2))    # initialize memory\n",
    "        self.optimizer = torch.optim.Adam(self.eval_net.parameters(), lr=critic_lr)\n",
    "        self.loss_func = nn.MSELoss()\n",
    "\n",
    "    def choose_action(self, x):\n",
    "        #x = Variable(torch.unsqueeze(torch.FloatTensor(x), 0))\n",
    "        # input only one sample\n",
    "        if np.random.uniform() < 0.9:   # greedy\n",
    "            a_output = self.eval_net.forward(x.cuda())\n",
    "            _ , action = torch.max(a_output.data,1)\n",
    "            #action = torch.LongTensor(([action.numpy()[0]]))   # return the argmax\n",
    "        else:   # random\n",
    "            action = np.random.randint(0, action_dim)\n",
    "            action = torch.LongTensor(([action])).cuda()\n",
    "        return action\n",
    "\n",
    "    def store_transition(self, s, a, r, s_):\n",
    "        a = a.cpu().numpy()[0]\n",
    "        transition = torch.cat((s.cuda(), Variable(torch.FloatTensor(([a]))).unsqueeze(0).cuda(), \\\n",
    "                                Variable(torch.FloatTensor(([r])).unsqueeze(0)).cuda()\\\n",
    "                                , s_.cuda()),1)\n",
    "        # replace the old memory with new memory\n",
    "        index = self.memory_counter % 50 #buffer_size\n",
    "        self.memory[index, :] = transition\n",
    "        self.memory_counter += 1\n",
    "\n",
    "    def learn(self):\n",
    "        # target parameter update\n",
    "        if self.learn_step_counter % 10 == 0:\n",
    "            self.target_net.load_state_dict(self.eval_net.state_dict())\n",
    "        self.learn_step_counter += 1\n",
    "\n",
    "        # sample batch transitions\n",
    "        sample_index = np.random.choice(50, 10)\n",
    "        b_memory = self.memory[sample_index, :]\n",
    "        b_s = b_memory[:, :state_dim].cuda()\n",
    "        b_a = b_memory[:, state_dim:state_dim+1].data.numpy().astype(int)\n",
    "        b_a = Variable(torch.from_numpy(b_a)).cuda()\n",
    "        b_r = b_memory[:, state_dim+1:state_dim+2].cuda()\n",
    "        b_s_ = b_memory[:, -state_dim:].cuda()\n",
    "\n",
    "        # q_eval w.r.t the action in experience\n",
    "        q_eval = self.eval_net(b_s).gather(1, b_a)  # shape (batch, 1)\n",
    "        \n",
    "        q_next = self.target_net(b_s_).detach()     # detach from graph, don't backpropagate\n",
    "        q_target = b_r + gamma * q_next.max(1)[0].unsqueeze(1)   # shape (batch, 1)\n",
    "        q_target = torch.autograd.Variable(q_target.data)\n",
    "        loss = self.loss_func(q_eval, q_target)\n",
    "\n",
    "        self.optimizer.zero_grad()\n",
    "        loss.backward(retain_graph = True)\n",
    "        self.optimizer.step()\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "img, box, label ,fn= next(iter(img_batch))\n",
    "ground_truth = box[0,0,:]\n",
    "Trigger = 0\n",
    "Steps = 0\n",
    "box_, s = Initial(img)\n",
    "action_bound = torch.LongTensor([0,0,img.size(2),img.size(3)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#test for DQN\n",
    "#a = dqn.choose_action(s)\n",
    "#test for DDPG\n",
    "a_output = actor(s.cuda())\n",
    "_ , a = torch.max(a_output.data,1)\n",
    "\n",
    "print a\n",
    "# take action\n",
    "box, trigger = Generate(a[0], box_, img)\n",
    "s_, r, t, Trigger, Steps = Sample(box_.squeeze(0), box.squeeze(0), trigger, Trigger, Steps)\n",
    "#DQN\n",
    "#box, trigger = Generate(a.cpu().numpy()[0] , box_, img)\n",
    "#s_, r, terminal, Trigger, Steps = Sample(box_.squeeze(0), box.squeeze(0), trigger, Trigger, Steps)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Collecting experience...\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 8\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "Trigger:0 Steps:0\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "Trigger:0 Steps:0\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 8\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "Trigger:0 Steps:0\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 0\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "Trigger:0 Steps:0\n",
      "\n",
      " 2\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      "Starting training...\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      "Starting training...\n",
      "('Ep: ', 4, '| Ep_r: ', -13.0)\n",
      "Trigger:0 Steps:0\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      "Starting training...\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      "Starting training...\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      "Starting training...\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      "Starting training...\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      "Starting training...\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      "Starting training...\n",
      "\n",
      " 3\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      "Starting training...\n",
      "\n",
      " 5\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      "Starting training...\n",
      "\n",
      " 5\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      "Starting training...\n",
      "\n",
      " 5\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      "Starting training...\n",
      "('Ep: ', 5, '| Ep_r: ', -3.0)\n",
      "Trigger:0 Steps:0\n",
      "\n",
      " 5\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      "Starting training...\n",
      "\n",
      " 5\n",
      "[torch.cuda.LongTensor of size 1 (GPU 0)]\n",
      "\n",
      "\n",
      "Starting training...\n"
     ]
    }
   ],
   "source": [
    "dqn = DQN()\n",
    "print('\\nCollecting experience...')\n",
    "for i_episode in range(10):\n",
    "    img, box, label ,fn= next(iter(img_batch))\n",
    "    ground_truth = box[0,0,:]\n",
    "    Trigger = 0\n",
    "    Steps = 0\n",
    "    box_, s = Initial(img)\n",
    "    action_bound = torch.LongTensor([0,0,img.size(2),img.size(3)])\n",
    "    ep_r = 0\n",
    "    while True:\n",
    "        a = dqn.choose_action(s)\n",
    "        print a\n",
    "        # take action\n",
    "        box, trigger = Generate(a.cpu().numpy()[0] , box_, img)\n",
    "        s_, r, terminal, Trigger, Steps = Sample(box_.squeeze(0),\\\n",
    "                                                 box.squeeze(0), trigger, Trigger, Steps, img)\n",
    "\n",
    "        # modify the reward\n",
    "        \"\"\"x, x_dot, theta, theta_dot = s_\n",
    "        r1 = (env.x_threshold - abs(x)) / env.x_threshold - 0.8\n",
    "        r2 = (env.theta_threshold_radians - abs(theta)) / env.theta_threshold_radians - 0.5\n",
    "        r = r1 + r2\"\"\"\n",
    "\n",
    "        dqn.store_transition(s, a, r, s_)\n",
    "\n",
    "        ep_r += r\n",
    "        if dqn.memory_counter > 50:\n",
    "            print \"\\nStarting training...\"\n",
    "            dqn.learn()\n",
    "            if terminal is True:\n",
    "                print('Ep: ', i_episode,\n",
    "                      '| Ep_r: ', round(ep_r, 2))\n",
    "\n",
    "        if terminal is True:\n",
    "            print (\"Trigger:{} Steps:{}\".format(Trigger, Steps))\n",
    "            break\n",
    "        s = s_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def PlotBox(img, box):\n",
    "    #img = Image.open(img)\n",
    "    draw = ImageDraw.Draw(img)\n",
    "    #plt.imshow(img_e)\n",
    "    draw.rectangle(box, outline = 'red')  \n",
    "    #draw.rectangle((0,0, box[2], box[3]), outline = 'red')  \n",
    "    plt.imshow(img)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "img, box, label ,fn= next(iter(img_batch))\n",
    "img_e = Image.open(fn[0])\n",
    "ground_truth = box[0,0,:]\n",
    "box = (ground_truth[0],ground_truth[1],ground_truth[2],ground_truth[3])\n",
    "PlotBox(img_e ,box)\n",
    "action_bound = torch.LongTensor([0,0,img.size(2),img.size(3)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "initial_box, s = Initial(img)\n",
    "\n",
    "#a = dqn.choose_action(s)\n",
    "#newbox, trigger = Generate(a.cpu().numpy()[0], initial_box, img)\n",
    "a_output = actor(s.cuda())\n",
    "_ , a = torch.max(a_output.data,1)\n",
    "\n",
    "print a\n",
    "# take action\n",
    "newbox, trigger = Generate(a[0], initial_box, img)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "print action_bound"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "s2, r, t, Trigger, Steps = Sample(initial_box.squeeze(0), newbox.squeeze(0), trigger, Trigger, Steps)\n",
    "newbox = newbox.squeeze(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#print newbox\n",
    "#a = dqn.choose_action(s2)\n",
    "a_output = actor(s2.cuda())\n",
    "_ , a = torch.max(a_output.data,1)\n",
    "print a\n",
    "#newbox1, trigger = Generate(a.cpu().numpy()[0], newbox, img)\n",
    "newbox1, trigger = Generate(a[0], newbox, img)\n",
    "s2, r, t, Trigger, Steps = Sample(newbox.squeeze(0), newbox1.squeeze(0), trigger, Trigger, Steps)\n",
    "print r, Steps, s2, newbox1\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "newbox = newbox1.squeeze(0)\n",
    "box = (newbox[0],newbox[1],newbox[2],newbox[3])\n",
    "PlotBox(img_e ,box)\n",
    "print box"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "print a_output\n",
    "_, action = torch.max(a_output.data,1)\n",
    "print action"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "torch.save(actor, 'actor1.pth')\n",
    "torch.save(critic, 'critic1.pth')\n",
    "torch.save(actor_target, 'actor_target1.pth')\n",
    "torch.save(critic_target, 'critic_target1.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def extract_grads(network):\n",
    "    return [p.grad.data.numpy()*2 for p in list(network.parameters())]\n",
    "\n",
    "\n",
    "def update(network, grad_in, grad_out):\n",
    "    grads = extract_grads(network)\n",
    "    grad_out = grads\n",
    "\n",
    "network.register_backward_hook(self.update)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# img, box, label = next(iter(img_batch))\n",
    "box_, s = Initial(img)\n",
    "action,t = Generate(0, box_)\n",
    "print s\n",
    "#b = Action(box_)\n",
    "#print b.a()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# the Ornstein-Uhlenbeck stochastic process for control tasks\n",
    "class OrnsteinUhlenbeckActionNoise:\n",
    "    def __init__(self, mu, sigma = 0.3, theta=.15, dt=1e-2, x0=None):\n",
    "        self.theta = theta\n",
    "        self.mu = mu\n",
    "        self.sigma = sigma\n",
    "        self.dt = dt\n",
    "        self.x0 = x0\n",
    "        self.reset()\n",
    "\n",
    "    def __call__(self):\n",
    "        x = self.x_prev + self.theta * (self.mu - self.x_prev) * self.dt + \\\n",
    "        self.sigma * np.sqrt(self.dt) * np.random.normal(size=self.mu.shape)\n",
    "        self.x_prev = x\n",
    "        return x\n",
    "\n",
    "    def reset(self):\n",
    "        self.x_prev = self.x0 if self.x0 is not None else np.zeros_like(self.mu)\n",
    "\n",
    "    def __repr__(self):\n",
    "        return 'OrnsteinUhlenbeckActionNoise(mu={}, sigma={})'.format(self.mu, self.sigma)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#reference for mutiple branches\n",
    "class mm(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(mm, self).__init__()\n",
    "        self.n = nn.Linear(4,3)\n",
    "        self.m = nn.Linear(3,2)\n",
    "        self.m2 = nn.Linear(3,4)\n",
    "    def forward(self, input, input2):\n",
    "        input_ = self.n(input)\n",
    "        input2_ = self.n(input2)\n",
    "        o1 = self.m(input_)\n",
    "        o2 = self.m2(input2_)\n",
    "        return o1, o2\n",
    "\n",
    "o1, o2 = mm(input)\n",
    "o = o1 + o2\n",
    "# loss\n",
    "\n",
    "## Or you can do\n",
    "\n",
    "l1 = loss(o1, target)\n",
    "l2 = loss2(o2, target2)\n",
    "torch.autograd.backward([l1, l2])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}