Ex4-Oculoenv.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from skimage.transform import resize\n",
    "import gym\n",
    "import gym.spaces\n",
    "import gym_oculoenv\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torchvision.models as models\n",
    "from torch.autograd import Variable\n",
    "from torch.distributions import Categorical\n",
    "import torch.optim as optim\n",
    "\n",
    "from esncell import ESNCell\n",
    "import utils\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "from itertools import count\n",
    "from collections import namedtuple"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Dataset params\n",
    "sample_length = 1000\n",
    "n_samples = 40\n",
    "batch_size = 1\n",
    "\n",
    "n_obs = 2048\n",
    "n_action = 9\n",
    "oh_mat = np.eye(n_action)\n",
    "\n",
    "# ESN properties\n",
    "input_dim = n_obs + n_action\n",
    "n_hidden = 200\n",
    "w_sparsity=0.1\n",
    "\n",
    "n_iterations = 50\n",
    "\n",
    "SavedAction = namedtuple('SavedAction', ['log_prob', 'value'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class ActorCritic(nn.Module):\n",
    "    def __init__(self, n_action, input_dim, hidden_dim=2000):\n",
    "        super(ActorCritic, self).__init__()\n",
    "        self.affine_a = nn.Linear(input_dim, hidden_dim) # for actor\n",
    "        self.affine_v = nn.Linear(input_dim, hidden_dim) # for critic\n",
    "        self.action_head = nn.Linear(hidden_dim, n_action)\n",
    "        self.value_head = nn.Linear(hidden_dim, 1)\n",
    "        self.saved_actions = []\n",
    "        self.rewards = []\n",
    "\n",
    "    def forward(self, x):\n",
    "        x1 = F.relu(self.affine_a(x))\n",
    "        x2 = F.relu(self.affine_v(x))\n",
    "        action_scores = self.action_head(x1)\n",
    "        state_values = self.value_head(x2)\n",
    "        return F.softmax(action_scores, dim=-1), state_values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "resnet = models.resnet50(pretrained=True)\n",
    "resnet = nn.Sequential(*list(resnet.children())[:-1])\n",
    "resevior = ESNCell(input_dim, n_hidden, batch_size, spectral_radius=0.5, input_scaling=0.3, w_sparsity=w_sparsity)\n",
    "ac_model = ActorCritic(n_action, n_hidden)\n",
    "\n",
    "optimizer = optim.Adam(ac_model.parameters(), lr=1e-3)\n",
    "eps = np.finfo(np.float32).eps.item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def select_action(state):\n",
    "    state = torch.from_numpy(state).float()\n",
    "    state = state.requires_grad_(requires_grad=True)\n",
    "    probs, state_value = ac_model(state)\n",
    "#     print(probs)\n",
    "    m = Categorical(probs)\n",
    "    action = m.sample()\n",
    "    ac_model.saved_actions.append(SavedAction(m.log_prob(action), state_value))\n",
    "    return action.item()\n",
    "\n",
    "\n",
    "def finish_episode(gamma=0.99):\n",
    "    R = 0\n",
    "    saved_actions = ac_model.saved_actions\n",
    "    policy_losses = []\n",
    "    value_losses = []\n",
    "    rewards = []\n",
    "    for r in ac_model.rewards[::-1]:\n",
    "        R = r + gamma * R\n",
    "        rewards.insert(0, R)\n",
    "    rewards = torch.tensor(rewards)\n",
    "    rewards = (rewards - rewards.mean()) / (rewards.std() + eps)\n",
    "    for (log_prob, value), r in zip(saved_actions, rewards):\n",
    "        reward = r - value.item()\n",
    "        policy_losses.append(-log_prob * reward)\n",
    "        value_losses.append(F.smooth_l1_loss(value.view(1), torch.tensor([r])))\n",
    "    optimizer.zero_grad()\n",
    "    loss = torch.stack(policy_losses).sum() + torch.stack(value_losses).sum()\n",
    "    loss.backward()\n",
    "#     retain_graph=True\n",
    "    optimizer.step()\n",
    "    print(loss.item(), np.sum(ac_model.rewards))\n",
    "    del ac_model.rewards[:]\n",
    "    del ac_model.saved_actions[:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NoneType 0.05817764173314431 9 False False\n"
     ]
    }
   ],
   "source": [
    "env = gym.make(\"RedCursor-v0\")\n",
    "obs = env.reset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5.901492595672607 15.71751820531188\n",
      "75.90277099609375 11.676625603246116\n",
      "4.729850769042969 12.26033076200861\n",
      "9.254830360412598 9.256434719774237\n",
      "16.24832534790039 7.577581708653683\n",
      "-0.8216276168823242 7.198015772708286\n",
      "9.305319786071777 5.679752028926697\n",
      "12.164582252502441 5.679752028926697\n",
      "13.040145874023438 5.679752028926697\n",
      "11.194570541381836 5.679752028926697\n",
      "Episode 0\tLast length:   199\tAverage length: 19.62\n",
      "9.128890037536621 14.014790373931564\n",
      "9.151163101196289 13.193009343192527\n",
      "10.161458015441895 13.193009343192527\n",
      "10.600664138793945 13.193009343192527\n",
      "10.031092643737793"
     ]
    }
   ],
   "source": [
    "running_reward = 10\n",
    "hidden_x = resevior.init_hidden(batch_size)\n",
    "for i_episode in range(10000):\n",
    "    obs = env.reset()\n",
    "    action_arr = np.array([0.] * n_action).reshape(1, -1)\n",
    "    for t in range(200):  # Don't infinite loop while learning\n",
    "        obs = obs/255.\n",
    "        obs = resize(obs, (224, 224))\n",
    "        obs = torch.tensor(np.expand_dims(obs.transpose((2, 0, 1)), axis=0).astype(\"float32\"))\n",
    "        obs = resnet(obs).view(1, -1)\n",
    "        obs = torch.cat([obs, torch.tensor(action_arr.astype(\"float32\")).requires_grad_(requires_grad=False)], dim=-1)\n",
    "        hidden_x = resevior(obs, hidden_x=hidden_x)\n",
    "#         print(hidden_x)\n",
    "        state = hidden_x.detach().numpy().reshape(1, -1)\n",
    "        action = select_action(state)\n",
    "        action_arr = np.expand_dims(oh_mat[action], axis=0)\n",
    "        time_step = env.step(action)\n",
    "        obs, reward, done, _ = env.step(action)\n",
    "        ac_model.rewards.append(reward)\n",
    "        if done or (t+1)%20 == 0:\n",
    "            running_reward = running_reward * 0.99 + t * 0.01\n",
    "            finish_episode()\n",
    "            if done:\n",
    "                print('Episode {}\\tLast length: {:5d}\\tAverage length: {:.2f}'.format(i_episode, t, running_reward))\n",
    "                break\n",
    "        if t+1 == 200:\n",
    "            print('Episode {}\\tLast length: {:5d}\\tAverage length: {:.2f}'.format(i_episode, t, running_reward))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.5.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}