4_pong_dqn (1).ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "4_pong_dqn.ipynb",
      "version": "0.3.2",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "metadata": {
        "id": "ak-uOKKUzPCH",
        "colab_type": "code",
        "outputId": "8ac3066c-ddae-4a09-f960-b032677caed0",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 328
        }
      },
      "cell_type": "code",
      "source": [
        "from IPython.display import clear_output\n",
        "import matplotlib.pyplot as plt\n",
        "%matplotlib inline\n",
        "!pip install torch==0.3.1\n",
        "!pip install gym[atari]"
      ],
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Requirement already satisfied: torch==0.3.1 in /usr/local/lib/python3.6/dist-packages (0.3.1)\n",
            "Requirement already satisfied: pyyaml in /usr/local/lib/python3.6/dist-packages (from torch==0.3.1) (3.13)\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: gym[atari] in /usr/local/lib/python3.6/dist-packages (0.10.8)\n",
            "Requirement already satisfied: scipy in /usr/local/lib/python3.6/dist-packages (from gym[atari]) (0.19.1)\n",
            "Requirement already satisfied: pyglet>=1.2.0 in /usr/local/lib/python3.6/dist-packages (from gym[atari]) (1.3.2)\n",
            "Requirement already satisfied: numpy>=1.10.4 in /usr/local/lib/python3.6/dist-packages (from gym[atari]) (1.14.6)\n",
            "Requirement already satisfied: requests>=2.0 in /usr/local/lib/python3.6/dist-packages (from gym[atari]) (2.18.4)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.6/dist-packages (from gym[atari]) (1.11.0)\n",
            "Requirement already satisfied: atari-py>=0.1.4; extra == \"atari\" in /usr/local/lib/python3.6/dist-packages (from gym[atari]) (0.1.6)\n",
            "Requirement already satisfied: Pillow; extra == \"atari\" in /usr/local/lib/python3.6/dist-packages (from gym[atari]) (4.0.0)\n",
            "Requirement already satisfied: PyOpenGL; extra == \"atari\" in /usr/local/lib/python3.6/dist-packages (from gym[atari]) (3.1.0)\n",
            "Requirement already satisfied: future in /usr/local/lib/python3.6/dist-packages (from pyglet>=1.2.0->gym[atari]) (0.16.0)\n",
            "Requirement already satisfied: urllib3<1.23,>=1.21.1 in /usr/local/lib/python3.6/dist-packages (from requests>=2.0->gym[atari]) (1.22)\n",
            "Requirement already satisfied: chardet<3.1.0,>=3.0.2 in /usr/local/lib/python3.6/dist-packages (from requests>=2.0->gym[atari]) (3.0.4)\n",
            "Requirement already satisfied: idna<2.7,>=2.5 in /usr/local/lib/python3.6/dist-packages (from requests>=2.0->gym[atari]) (2.6)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.6/dist-packages (from requests>=2.0->gym[atari]) (2018.10.15)\n",
            "Requirement already satisfied: olefile in /usr/local/lib/python3.6/dist-packages (from Pillow; extra == \"atari\"->gym[atari]) (0.46)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "id": "SGAp2VI3zVmd",
        "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": "4NU8fm3XzVpF",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "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"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "vP8CGTJuzVrl",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "\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": "8cYXyDun1-Xy",
        "colab_type": "code",
        "outputId": "ab2dfb34-a873-44db-b7a6-3166a084f506",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        }
      },
      "cell_type": "code",
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')\n"
      ],
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Drive already mounted at /content/drive; to attempt to forcibly remount, call drive.mount(\"/content/drive\", force_remount=True).\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "id": "ln8l71k-3qiE",
        "colab_type": "code",
        "outputId": "e3185859-e8f7-47cb-d49b-b5dfda93b6f4",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        }
      },
      "cell_type": "code",
      "source": [
        "\n",
        "\n",
        "\n",
        "from collections import deque\n",
        "\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": 7,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "done\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "id": "gxVNtpKd2lUo",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "# New Section"
      ]
    },
    {
      "metadata": {
        "id": "_rZTbEhOzVuA",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "\n",
        "epsilon_start = 1.0\n",
        "epsilon_final = 0.01\n",
        "epsilon_decay = 30000\n",
        "\n",
        "epsilon_by_frame = lambda frame_idx: epsilon_final + (epsilon_start - epsilon_final) * math.exp(-1. * frame_idx / epsilon_decay)\n",
        "epsilon_by_frame = lambda frame_idx: epsilon_final + (epsilon_start - epsilon_final) * math.exp(-1. * frame_idx / epsilon_decay)\n",
        "\n",
        "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)\n",
        "\n",
        "\n",
        "\n",
        "env_id = \"CartPole-v0\"\n",
        "env = gym.make(env_id)\n",
        "\n",
        "env_id = \"PongNoFrameskip-v4\"\n",
        "env    = make_atari(env_id)\n",
        "env    = wrap_deepmind(env)\n",
        "env    = wrap_pytorch(env)\n",
        "\n",
        "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\n",
        "\n",
        "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\n",
        "  \n",
        "#로스에 관한 고민\n",
        "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(100000)\n",
        "\n",
        "num_frames = 1000000\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": []
    }
  ]
}