2_supermario_dqn.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "2_supermario_dqn.ipynb",
      "version": "0.3.2",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "metadata": {
        "id": "t86svuYp_AjF",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 432
        },
        "outputId": "3c59b72b-7c48-4269-c43c-153a6a19680d"
      },
      "cell_type": "code",
      "source": [
        "from IPython.display import clear_output\n",
        "import matplotlib.pyplot as plt\n",
        "%matplotlib inline\n",
        "!pip install gym_super_mario_bros\n",
        "!pip install torch==0.3.1"
      ],
      "execution_count": 18,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Requirement already satisfied: gym_super_mario_bros in /usr/local/lib/python3.6/dist-packages (5.0.0)\n",
            "Requirement already satisfied: pygame>=1.9.3 in /usr/local/lib/python3.6/dist-packages (from gym_super_mario_bros) (1.9.4)\n",
            "Requirement already satisfied: nes-py>=2.0.0 in /usr/local/lib/python3.6/dist-packages (from gym_super_mario_bros) (2.0.0)\n",
            "Requirement already satisfied: numpy>=1.14.2 in /usr/local/lib/python3.6/dist-packages (from gym_super_mario_bros) (1.14.6)\n",
            "Requirement already satisfied: tqdm>=4.19.5 in /usr/local/lib/python3.6/dist-packages (from gym_super_mario_bros) (4.28.1)\n",
            "Requirement already satisfied: pyglet>=1.3.2 in /usr/local/lib/python3.6/dist-packages (from gym_super_mario_bros) (1.3.2)\n",
            "Requirement already satisfied: gym>=0.10.5 in /usr/local/lib/python3.6/dist-packages (from gym_super_mario_bros) (0.10.8)\n",
            "Requirement already satisfied: opencv-python>=3.4.0.12 in /usr/local/lib/python3.6/dist-packages (from gym_super_mario_bros) (3.4.3.18)\n",
            "Requirement already satisfied: matplotlib>=2.0.2 in /usr/local/lib/python3.6/dist-packages (from gym_super_mario_bros) (2.1.2)\n",
            "Requirement already satisfied: future in /usr/local/lib/python3.6/dist-packages (from pyglet>=1.3.2->gym_super_mario_bros) (0.16.0)\n",
            "Requirement already satisfied: requests>=2.0 in /usr/local/lib/python3.6/dist-packages (from gym>=0.10.5->gym_super_mario_bros) (2.18.4)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.6/dist-packages (from gym>=0.10.5->gym_super_mario_bros) (1.11.0)\n",
            "Requirement already satisfied: scipy in /usr/local/lib/python3.6/dist-packages (from gym>=0.10.5->gym_super_mario_bros) (0.19.1)\n",
            "Requirement already satisfied: pytz in /usr/local/lib/python3.6/dist-packages (from matplotlib>=2.0.2->gym_super_mario_bros) (2018.7)\n",
            "Requirement already satisfied: python-dateutil>=2.1 in /usr/local/lib/python3.6/dist-packages (from matplotlib>=2.0.2->gym_super_mario_bros) (2.5.3)\n",
            "Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.6/dist-packages (from matplotlib>=2.0.2->gym_super_mario_bros) (0.10.0)\n",
            "Requirement already satisfied: pyparsing!=2.0.4,!=2.1.2,!=2.1.6,>=2.0.1 in /usr/local/lib/python3.6/dist-packages (from matplotlib>=2.0.2->gym_super_mario_bros) (2.3.0)\n",
            "Requirement already satisfied: idna<2.7,>=2.5 in /usr/local/lib/python3.6/dist-packages (from requests>=2.0->gym>=0.10.5->gym_super_mario_bros) (2.6)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.6/dist-packages (from requests>=2.0->gym>=0.10.5->gym_super_mario_bros) (2018.10.15)\n",
            "Requirement already satisfied: chardet<3.1.0,>=3.0.2 in /usr/local/lib/python3.6/dist-packages (from requests>=2.0->gym>=0.10.5->gym_super_mario_bros) (3.0.4)\n",
            "Requirement already satisfied: urllib3<1.23,>=1.21.1 in /usr/local/lib/python3.6/dist-packages (from requests>=2.0->gym>=0.10.5->gym_super_mario_bros) (1.22)\n",
            "Requirement already satisfied: torch==0.3.1 in /usr/local/lib/python3.6/dist-packages (0.3.1)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from torch==0.3.1) (1.14.6)\n",
            "Requirement already satisfied: pyyaml in /usr/local/lib/python3.6/dist-packages (from torch==0.3.1) (3.13)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "id": "jhR9NIiw_Cs-",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "def plot(frame_idx, rewards, losses):\n",
        "    clear_output(True)\n",
        "    plt.figure(figsize=(20,5))\n",
        "    plt.subplot(131)\n",
        "    plt.title('frame %s. reward: %s' % (frame_idx, np.mean(rewards[-10:])))\n",
        "    plt.plot(rewards)\n",
        "    plt.subplot(132)\n",
        "    plt.title('loss')\n",
        "    plt.plot(losses)\n",
        "    plt.show()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "GmHqX2s0_Cux",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "#import retro\n",
        "\n",
        "import gym_super_mario_bros\n",
        "import logging\n",
        "import gym\n",
        "import gym_super_mario_bros\n",
        "import numpy as np\n",
        "import sys\n",
        "from nes_py.wrappers import BinarySpaceToDiscreteSpaceEnv\n",
        "import random\n",
        "import math, random\n",
        "\n",
        "import gym\n",
        "import numpy as np\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "import torch.autograd as autograd\n",
        "import torch.nn.functional as F\n",
        "\n",
        "\n",
        "\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "naw2P7zJ_Cwo",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "\n",
        "\n",
        "movements = [\n",
        "    ['NOP'],\n",
        "    ['A'],\n",
        "    ['B'],\n",
        "    ['right'],\n",
        "    ['right', 'A'],\n",
        "    ['right', 'B'],\n",
        "    ['right', 'A', 'B'],\n",
        "    ['left'],\n",
        "    ['left', 'A'],\n",
        "    ['left', 'B'],\n",
        "    ['left', 'A', 'B'],\n",
        "    ['down'],\n",
        "    ['up']\n",
        "]\n",
        "\n",
        "\n",
        "\n",
        "from IPython.display import clear_output\n",
        "import matplotlib.pyplot as plt\n",
        "get_ipython().run_line_magic('matplotlib', 'inline')\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "USE_CUDA = torch.cuda.is_available()\n",
        "Variable = lambda *args, **kwargs: autograd.Variable(*args, **kwargs).cuda() if USE_CUDA else autograd.Variable(*args, **kwargs)\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "1qNRZ6rf_C7H",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "outputId": "6bb3316a-5514-4cc3-f9f1-a0887466f5b9"
      },
      "cell_type": "code",
      "source": [
        "\n",
        "\n",
        "\n",
        "from collections import deque\n",
        "import numpy as np\n",
        "from collections import deque\n",
        "import gym\n",
        "from gym import spaces\n",
        "import cv2\n",
        "cv2.ocl.setUseOpenCL(False)\n",
        "\n",
        "class NoopResetEnv(gym.Wrapper):\n",
        "    def __init__(self, env, noop_max=30):\n",
        "        \"\"\"Sample initial states by taking random number of no-ops on reset.\n",
        "        No-op is assumed to be action 0.\n",
        "        \"\"\"\n",
        "        gym.Wrapper.__init__(self, env)\n",
        "        self.noop_max = noop_max\n",
        "        self.override_num_noops = None\n",
        "        self.noop_action = 0\n",
        "        assert env.unwrapped.get_action_meanings()[0] == 'NOOP'\n",
        "\n",
        "    def reset(self, **kwargs):\n",
        "        \"\"\" Do no-op action for a number of steps in [1, noop_max].\"\"\"\n",
        "        self.env.reset(**kwargs)\n",
        "        if self.override_num_noops is not None:\n",
        "            noops = self.override_num_noops\n",
        "        else:\n",
        "            noops = self.unwrapped.np_random.randint(1, self.noop_max + 1) #pylint: disable=E1101\n",
        "        assert noops > 0\n",
        "        obs = None\n",
        "        for _ in range(noops):\n",
        "            obs, _, done, _ = self.env.step(self.noop_action)\n",
        "            if done:\n",
        "                obs = self.env.reset(**kwargs)\n",
        "        return obs\n",
        "\n",
        "    def step(self, ac):\n",
        "        return self.env.step(ac)\n",
        "\n",
        "class FireResetEnv(gym.Wrapper):\n",
        "    def __init__(self, env):\n",
        "        \"\"\"Take action on reset for environments that are fixed until firing.\"\"\"\n",
        "        gym.Wrapper.__init__(self, env)\n",
        "        assert env.unwrapped.get_action_meanings()[1] == 'FIRE'\n",
        "        assert len(env.unwrapped.get_action_meanings()) >= 3\n",
        "\n",
        "    def reset(self, **kwargs):\n",
        "        self.env.reset(**kwargs)\n",
        "        obs, _, done, _ = self.env.step(1)\n",
        "        if done:\n",
        "            self.env.reset(**kwargs)\n",
        "        obs, _, done, _ = self.env.step(2)\n",
        "        if done:\n",
        "            self.env.reset(**kwargs)\n",
        "        return obs\n",
        "\n",
        "    def step(self, ac):\n",
        "        return self.env.step(ac)\n",
        "\n",
        "class EpisodicLifeEnv(gym.Wrapper):\n",
        "    def __init__(self, env):\n",
        "        \"\"\"Make end-of-life == end-of-episode, but only reset on true game over.\n",
        "        Done by DeepMind for the DQN and co. since it helps value estimation.\n",
        "        \"\"\"\n",
        "        gym.Wrapper.__init__(self, env)\n",
        "        self.lives = 0\n",
        "        self.was_real_done  = True\n",
        "\n",
        "    def step(self, action):\n",
        "        obs, reward, done, info = self.env.step(action)\n",
        "        self.was_real_done = done\n",
        "        # check current lives, make loss of life terminal,\n",
        "        # then update lives to handle bonus lives\n",
        "        lives = self.env.unwrapped.ale.lives()\n",
        "        if lives < self.lives and lives > 0:\n",
        "            # for Qbert sometimes we stay in lives == 0 condtion for a few frames\n",
        "            # so its important to keep lives > 0, so that we only reset once\n",
        "            # the environment advertises done.\n",
        "            done = True\n",
        "        self.lives = lives\n",
        "        return obs, reward, done, info\n",
        "\n",
        "    def reset(self, **kwargs):\n",
        "        \"\"\"Reset only when lives are exhausted.\n",
        "        This way all states are still reachable even though lives are episodic,\n",
        "        and the learner need not know about any of this behind-the-scenes.\n",
        "        \"\"\"\n",
        "        if self.was_real_done:\n",
        "            obs = self.env.reset(**kwargs)\n",
        "        else:\n",
        "            # no-op step to advance from terminal/lost life state\n",
        "            obs, _, _, _ = self.env.step(0)\n",
        "        self.lives = self.env.unwrapped.ale.lives()\n",
        "        return obs\n",
        "\n",
        "class MaxAndSkipEnv(gym.Wrapper):\n",
        "    def __init__(self, env, skip=4):\n",
        "        \"\"\"Return only every `skip`-th frame\"\"\"\n",
        "        gym.Wrapper.__init__(self, env)\n",
        "        # most recent raw observations (for max pooling across time steps)\n",
        "        self._obs_buffer = np.zeros((2,)+env.observation_space.shape, dtype=np.uint8)\n",
        "        self._skip       = skip\n",
        "\n",
        "    def reset(self):\n",
        "        return self.env.reset()\n",
        "\n",
        "    def step(self, action):\n",
        "        \"\"\"Repeat action, sum reward, and max over last observations.\"\"\"\n",
        "        total_reward = 0.0\n",
        "        done = None\n",
        "        for i in range(self._skip):\n",
        "            obs, reward, done, info = self.env.step(action)\n",
        "            if i == self._skip - 2: self._obs_buffer[0] = obs\n",
        "            if i == self._skip - 1: self._obs_buffer[1] = obs\n",
        "            total_reward += reward\n",
        "            if done:\n",
        "                break\n",
        "        # Note that the observation on the done=True frame\n",
        "        # doesn't matter\n",
        "        max_frame = self._obs_buffer.max(axis=0)\n",
        "\n",
        "        return max_frame, total_reward, done, info\n",
        "\n",
        "    def reset(self, **kwargs):\n",
        "        return self.env.reset(**kwargs)\n",
        "\n",
        "class ClipRewardEnv(gym.RewardWrapper):\n",
        "    def __init__(self, env):\n",
        "        gym.RewardWrapper.__init__(self, env)\n",
        "\n",
        "    def reward(self, reward):\n",
        "        \"\"\"Bin reward to {+1, 0, -1} by its sign.\"\"\"\n",
        "        return np.sign(reward)\n",
        "\n",
        "class WarpFrame(gym.ObservationWrapper):\n",
        "    def __init__(self, env):\n",
        "        \"\"\"Warp frames to 84x84 as done in the Nature paper and later work.\"\"\"\n",
        "        gym.ObservationWrapper.__init__(self, env)\n",
        "        self.width = 84\n",
        "        self.height = 84\n",
        "        self.observation_space = spaces.Box(low=0, high=255,\n",
        "            shape=(self.height, self.width, 1), dtype=np.uint8)\n",
        "\n",
        "    def observation(self, frame):\n",
        "        frame = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY)\n",
        "        frame = cv2.resize(frame, (self.width, self.height), interpolation=cv2.INTER_AREA)\n",
        "        return frame[:, :, None]\n",
        "\n",
        "class FrameStack(gym.Wrapper):\n",
        "    def __init__(self, env, k):\n",
        "        \"\"\"Stack k last frames.\n",
        "        Returns lazy array, which is much more memory efficient.\n",
        "        See Also\n",
        "        --------\n",
        "        baselines.common.atari_wrappers.LazyFrames\n",
        "        \"\"\"\n",
        "        gym.Wrapper.__init__(self, env)\n",
        "        self.k = k\n",
        "        self.frames = deque([], maxlen=k)\n",
        "        shp = env.observation_space.shape\n",
        "        self.observation_space = spaces.Box(low=0, high=255, shape=(shp[0], shp[1], shp[2] * k), dtype=np.uint8)\n",
        "\n",
        "    def reset(self):\n",
        "        ob = self.env.reset()\n",
        "        for _ in range(self.k):\n",
        "            self.frames.append(ob)\n",
        "        return self._get_ob()\n",
        "\n",
        "    def step(self, action):\n",
        "        ob, reward, done, info = self.env.step(action)\n",
        "        self.frames.append(ob)\n",
        "        return self._get_ob(), reward, done, info\n",
        "\n",
        "    def _get_ob(self):\n",
        "        assert len(self.frames) == self.k\n",
        "        return LazyFrames(list(self.frames))\n",
        "\n",
        "class ScaledFloatFrame(gym.ObservationWrapper):\n",
        "    def __init__(self, env):\n",
        "        gym.ObservationWrapper.__init__(self, env)\n",
        "\n",
        "    def observation(self, observation):\n",
        "        # careful! This undoes the memory optimization, use\n",
        "        # with smaller replay buffers only.\n",
        "        return np.array(observation).astype(np.float32) / 255.0\n",
        "\n",
        "class LazyFrames(object):\n",
        "    def __init__(self, frames):\n",
        "        \"\"\"This object ensures that common frames between the observations are only stored once.\n",
        "        It exists purely to optimize memory usage which can be huge for DQN's 1M frames replay\n",
        "        buffers.\n",
        "        This object should only be converted to numpy array before being passed to the model.\n",
        "        You'd not believe how complex the previous solution was.\"\"\"\n",
        "        self._frames = frames\n",
        "        self._out = None\n",
        "\n",
        "    def _force(self):\n",
        "        if self._out is None:\n",
        "            self._out = np.concatenate(self._frames, axis=2)\n",
        "            self._frames = None\n",
        "        return self._out\n",
        "\n",
        "    def __array__(self, dtype=None):\n",
        "        out = self._force()\n",
        "        if dtype is not None:\n",
        "            out = out.astype(dtype)\n",
        "        return out\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self._force())\n",
        "\n",
        "    def __getitem__(self, i):\n",
        "        return self._force()[i]\n",
        "\n",
        "def make_atari(env_id):\n",
        "    env = gym.make(env_id)\n",
        "    assert 'NoFrameskip' in env.spec.id\n",
        "    env = NoopResetEnv(env, noop_max=30)\n",
        "    env = MaxAndSkipEnv(env, skip=4)\n",
        "    return env\n",
        "\n",
        "def wrap_deepmind(env, episode_life=True, clip_rewards=True, frame_stack=False, scale=False):\n",
        "    \"\"\"Configure environment for DeepMind-style Atari.\n",
        "    \"\"\"\n",
        "    if episode_life:\n",
        "        env = EpisodicLifeEnv(env)\n",
        "    if 'FIRE' in env.unwrapped.get_action_meanings():\n",
        "        env = FireResetEnv(env)\n",
        "    env = WarpFrame(env)\n",
        "    if scale:\n",
        "        env = ScaledFloatFrame(env)\n",
        "    if clip_rewards:\n",
        "        env = ClipRewardEnv(env)\n",
        "    if frame_stack:\n",
        "        env = FrameStack(env, 4)\n",
        "    return env\n",
        "\n",
        "\n",
        "\n",
        "class ImageToPyTorch(gym.ObservationWrapper):\n",
        "    \"\"\"\n",
        "    Image shape to num_channels x weight x height\n",
        "    \"\"\"\n",
        "    def __init__(self, env):\n",
        "        super(ImageToPyTorch, self).__init__(env)\n",
        "        old_shape = self.observation_space.shape\n",
        "        self.observation_space = gym.spaces.Box(low=0.0, high=1.0, shape=(old_shape[-1], old_shape[0], old_shape[1]), dtype=np.uint8)\n",
        "\n",
        "    def observation(self, observation):\n",
        "        return np.swapaxes(observation, 2, 0)\n",
        "    \n",
        "\n",
        "def wrap_pytorch(env):\n",
        "    return ImageToPyTorch(env)\n",
        "print(\"done\")\n"
      ],
      "execution_count": 34,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "done\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "id": "eN7ad71i_Ncu",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "from collections import deque\n",
        "\n",
        "class ReplayBuffer(object):\n",
        "    def __init__(self, capacity):\n",
        "        self.buffer = deque(maxlen=capacity)\n",
        "    \n",
        "    def push(self, state, action, reward, next_state, done):\n",
        "        state      = np.expand_dims(state, 0)\n",
        "        next_state = np.expand_dims(next_state, 0)\n",
        "            \n",
        "        self.buffer.append((state, action, reward, next_state, done))\n",
        "    \n",
        "    def sample(self, batch_size):\n",
        "        state, action, reward, next_state, done = zip(*random.sample(self.buffer, batch_size))\n",
        "        return np.concatenate(state), action, reward, np.concatenate(next_state), done\n",
        "    \n",
        "    def __len__(self):\n",
        "        return len(self.buffer)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "pCzg_kNP_NfL",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "epsilon_start = 1.0\n",
        "epsilon_final = 0.01\n",
        "epsilon_decay = 300000\n",
        "epsilon_by_frame = lambda frame_idx: epsilon_final + (epsilon_start - epsilon_final) * math.exp(-1. * frame_idx / epsilon_decay)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "ZGl8XEow_Nh0",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "class CnnDQN(nn.Module):\n",
        "    def __init__(self, input_shape, num_actions):\n",
        "        super(CnnDQN, self).__init__()\n",
        "        \n",
        "        self.input_shape = input_shape\n",
        "        self.num_actions = num_actions\n",
        "        \n",
        "        self.features = nn.Sequential(\n",
        "            nn.Conv2d(input_shape[0], 32, kernel_size=8, stride=4),\n",
        "            nn.ReLU(),\n",
        "            nn.Conv2d(32, 64, kernel_size=4, stride=2),\n",
        "            nn.ReLU(),\n",
        "            nn.Conv2d(64, 64, kernel_size=3, stride=1),\n",
        "            nn.ReLU()\n",
        "        )\n",
        "        \n",
        "        self.fc = nn.Sequential(\n",
        "            nn.Linear(self.feature_size(), 512),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(512, self.num_actions)\n",
        "        )\n",
        "        \n",
        "    def forward(self, x):\n",
        "        x = self.features(x)\n",
        "        x = x.view(x.size(0), -1)\n",
        "        x = self.fc(x)\n",
        "        return x\n",
        "    \n",
        "    def feature_size(self):\n",
        "        return self.features(autograd.Variable(torch.zeros(1, *self.input_shape))).view(1, -1).size(1)\n",
        "    \n",
        "    def act(self, state, epsilon):\n",
        "        if random.random() > epsilon:\n",
        "            state   = Variable(torch.FloatTensor(np.float32(state)).unsqueeze(0), volatile=True)\n",
        "            q_value = self.forward(state)\n",
        "            action  = q_value.max(1)[1].data[0]\n",
        "        else:\n",
        "            action = random.randrange(env.action_space.n)\n",
        "        return action"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "Ftqozk5VALz0",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "\n",
        "_env = gym_super_mario_bros.make('SuperMarioBros-v0')\n",
        "#_env = gym_super_mario_bros.SuperMarioBrosEnv(frames_per_step=1)\n",
        "env = BinarySpaceToDiscreteSpaceEnv(_env, movements)\n",
        "env    = wrap_pytorch(env)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "jgnYIQ3h_Nj5",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "model = CnnDQN(env.observation_space.shape, env.action_space.n)\n",
        "\n",
        "if USE_CUDA:\n",
        "    model = model.cuda()\n",
        "    \n",
        "optimizer = optim.Adam(model.parameters(), lr=0.00001)\n",
        "\n",
        "replay_initial = 10000\n",
        "replay_buffer = ReplayBuffer(500000)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "QpbR5_cPdV4M",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "def compute_td_loss(batch_size):\n",
        "    state, action, reward, next_state, done = replay_buffer.sample(batch_size)\n",
        "\n",
        "    state      = Variable(torch.FloatTensor(np.float32(state)))\n",
        "    next_state = Variable(torch.FloatTensor(np.float32(next_state)), volatile=True)\n",
        "    action     = Variable(torch.LongTensor(action))\n",
        "    reward     = Variable(torch.FloatTensor(reward))\n",
        "    done       = Variable(torch.FloatTensor(done))\n",
        "\n",
        "    q_values      = model(state)\n",
        "    next_q_values = model(next_state)\n",
        "\n",
        "    q_value          = q_values.gather(1, action.unsqueeze(1)).squeeze(1)\n",
        "    next_q_value     = next_q_values.max(1)[0]\n",
        "    expected_q_value = reward + gamma * next_q_value * (1 - done)\n",
        "    \n",
        "    loss = (q_value - Variable(expected_q_value.data)).pow(2).mean()\n",
        "        \n",
        "    optimizer.zero_grad()\n",
        "    loss.backward()\n",
        "    optimizer.step()\n",
        "    \n",
        "    return loss"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "h7DYpf8W_Nmw",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "num_frames = 300000\n",
        "batch_size = 32\n",
        "gamma      = 0.99\n",
        "\n",
        "losses = []\n",
        "all_rewards = []\n",
        "episode_reward = 0\n",
        "\n",
        "state = env.reset()\n",
        "for frame_idx in range(1, num_frames + 1):\n",
        "    epsilon = epsilon_by_frame(frame_idx)\n",
        "    action = model.act(state, epsilon)\n",
        "    \n",
        "    next_state, reward, done, _ = env.step(action)\n",
        "    replay_buffer.push(state, action, reward, next_state, done)\n",
        "    \n",
        "    state = next_state\n",
        "    episode_reward += reward\n",
        "    \n",
        "    if done:\n",
        "        state = env.reset()\n",
        "        all_rewards.append(episode_reward)\n",
        "        episode_reward = 0\n",
        "        \n",
        "    if len(replay_buffer) > replay_initial:\n",
        "        loss = compute_td_loss(batch_size)\n",
        "        losses.append(loss.data[0])\n",
        "        \n",
        "    if frame_idx % 1000 == 0:\n",
        "        plot(frame_idx, all_rewards, losses)"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}