DRL_15_16_17_DQN_Pong.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "DRL_15_16_17_DQN_Pong.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "40Yb47zJQglm",
        "colab_type": "text"
      },
      "source": [
        "DEEP REINFORCEMENT LEARNING EXPLAINED - 15 - 16 - 17\n",
        "# **Deep Q-Network (DQN)**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Q40Fa7qM4_lE",
        "colab_type": "text"
      },
      "source": [
        "OpenAI Pong"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FA1Y5VCv20XZ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import gym\n",
        "import gym.spaces\n",
        "\n",
        "DEFAULT_ENV_NAME = \"PongNoFrameskip-v4\" \n",
        "test_env = gym.make(DEFAULT_ENV_NAME)\n",
        "print(test_env.action_space.n)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8QDaXip14JBv",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "print(test_env.unwrapped.get_action_meanings())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1uzLQLz04z2i",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "print(test_env.observation_space.shape)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZzcdmzIL5EMI",
        "colab_type": "text"
      },
      "source": [
        "\n",
        "Type of hardware accelerator provided by Colab"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VjUM99rEKFNt",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!nvidia-smi "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZhmsqgrHikEl",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import warnings\n",
        "warnings.filterwarnings('ignore')"
      ],
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pRcuJGVSQi6g",
        "colab_type": "text"
      },
      "source": [
        "## OpenAI Gym Wrappers"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nPi1lHINMuSu",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Taken from \n",
        "# https://github.com/PacktPublishing/Deep-Reinforcement-Learning-Hands-On/blob/master/Chapter06/lib/wrappers.py\n",
        "\n",
        "import cv2\n",
        "import numpy as np\n",
        "import collections\n",
        "\n",
        "class FireResetEnv(gym.Wrapper):\n",
        "    def __init__(self, env=None):\n",
        "        super(FireResetEnv, self).__init__(env)\n",
        "        assert env.unwrapped.get_action_meanings()[1] == 'FIRE'\n",
        "        assert len(env.unwrapped.get_action_meanings()) >= 3\n",
        "\n",
        "    def step(self, action):\n",
        "        return self.env.step(action)\n",
        "\n",
        "    def reset(self):\n",
        "        self.env.reset()\n",
        "        obs, _, done, _ = self.env.step(1)\n",
        "        if done:\n",
        "            self.env.reset()\n",
        "        obs, _, done, _ = self.env.step(2)\n",
        "        if done:\n",
        "            self.env.reset()\n",
        "        return obs\n",
        "\n",
        "class MaxAndSkipEnv(gym.Wrapper):\n",
        "    def __init__(self, env=None, skip=4):\n",
        "        super(MaxAndSkipEnv, self).__init__(env)\n",
        "        # most recent raw observations (for max pooling across time steps)\n",
        "        self._obs_buffer = collections.deque(maxlen=2)\n",
        "        self._skip = skip\n",
        "\n",
        "    def step(self, action):\n",
        "        total_reward = 0.0\n",
        "        done = None\n",
        "        for _ in range(self._skip):\n",
        "            obs, reward, done, info = self.env.step(action)\n",
        "            self._obs_buffer.append(obs)\n",
        "            total_reward += reward\n",
        "            if done:\n",
        "                break\n",
        "        max_frame = np.max(np.stack(self._obs_buffer), axis=0)\n",
        "        return max_frame, total_reward, done, info\n",
        "\n",
        "    def reset(self):\n",
        "        self._obs_buffer.clear()\n",
        "        obs = self.env.reset()\n",
        "        self._obs_buffer.append(obs)\n",
        "        return obs\n",
        "\n",
        "\n",
        "class ProcessFrame84(gym.ObservationWrapper):\n",
        "    def __init__(self, env=None):\n",
        "        super(ProcessFrame84, self).__init__(env)\n",
        "        self.observation_space = gym.spaces.Box(low=0, high=255, shape=(84, 84, 1), dtype=np.uint8)\n",
        "\n",
        "    def observation(self, obs):\n",
        "        return ProcessFrame84.process(obs)\n",
        "\n",
        "    @staticmethod\n",
        "    def process(frame):\n",
        "        if frame.size == 210 * 160 * 3:\n",
        "            img = np.reshape(frame, [210, 160, 3]).astype(np.float32)\n",
        "        elif frame.size == 250 * 160 * 3:\n",
        "            img = np.reshape(frame, [250, 160, 3]).astype(np.float32)\n",
        "        else:\n",
        "            assert False, \"Unknown resolution.\"\n",
        "        img = img[:, :, 0] * 0.299 + img[:, :, 1] * 0.587 + img[:, :, 2] * 0.114\n",
        "        resized_screen = cv2.resize(img, (84, 110), interpolation=cv2.INTER_AREA)\n",
        "        x_t = resized_screen[18:102, :]\n",
        "        x_t = np.reshape(x_t, [84, 84, 1])\n",
        "        return x_t.astype(np.uint8)\n",
        "\n",
        "\n",
        "class BufferWrapper(gym.ObservationWrapper):\n",
        "    def __init__(self, env, n_steps, dtype=np.float32):\n",
        "        super(BufferWrapper, self).__init__(env)\n",
        "        self.dtype = dtype\n",
        "        old_space = env.observation_space\n",
        "        self.observation_space = gym.spaces.Box(old_space.low.repeat(n_steps, axis=0),\n",
        "                                                old_space.high.repeat(n_steps, axis=0), dtype=dtype)\n",
        "\n",
        "    def reset(self):\n",
        "        self.buffer = np.zeros_like(self.observation_space.low, dtype=self.dtype)\n",
        "        return self.observation(self.env.reset())\n",
        "\n",
        "    def observation(self, observation):\n",
        "        self.buffer[:-1] = self.buffer[1:]\n",
        "        self.buffer[-1] = observation\n",
        "        return self.buffer\n",
        "\n",
        "\n",
        "class ImageToPyTorch(gym.ObservationWrapper):\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], \n",
        "                                old_shape[0], old_shape[1]), dtype=np.float32)\n",
        "\n",
        "    def observation(self, observation):\n",
        "        return np.moveaxis(observation, 2, 0)\n",
        "\n",
        "\n",
        "class ScaledFloatFrame(gym.ObservationWrapper):\n",
        "    def observation(self, obs):\n",
        "        return np.array(obs).astype(np.float32) / 255.0\n",
        "\n",
        "def make_env(env_name):\n",
        "    env = gym.make(env_name)\n",
        "    env = MaxAndSkipEnv(env)\n",
        "    env = FireResetEnv(env)\n",
        "    env = ProcessFrame84(env)\n",
        "    env = ImageToPyTorch(env)\n",
        "    env = BufferWrapper(env, 4)\n",
        "    return ScaledFloatFrame(env)"
      ],
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Wznv9I1KR_I3",
        "colab_type": "text"
      },
      "source": [
        "## The DQN model\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "h6B8v-Qh5Ykk",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import torch\n",
        "import torch.nn as nn        # Pytorch neural network package\n",
        "import torch.optim as optim  # Pytorch optimization package\n",
        "\n",
        "device = torch.device(\"cuda\")"
      ],
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "N4S1I9xWMkf3",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Taken from \n",
        "# https://github.com/PacktPublishing/Deep-Reinforcement-Learning-Hands-On/blob/master/Chapter06/lib/dqn_model.py\n",
        "\n",
        "import numpy as np\n",
        "\n",
        "class DQN(nn.Module):\n",
        "    def __init__(self, input_shape, n_actions):\n",
        "        super(DQN, self).__init__()\n",
        "\n",
        "        self.conv = 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",
        "        conv_out_size = self._get_conv_out(input_shape)\n",
        "        self.fc = nn.Sequential(\n",
        "            nn.Linear(conv_out_size, 512),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(512, n_actions)\n",
        "        )\n",
        "\n",
        "    def _get_conv_out(self, shape):\n",
        "        o = self.conv(torch.zeros(1, *shape))\n",
        "        return int(np.prod(o.size()))\n",
        "\n",
        "    def forward(self, x):\n",
        "        conv_out = self.conv(x).view(x.size()[0], -1)\n",
        "        return self.fc(conv_out)"
      ],
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "taYi5LZnIOqz",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 263
        },
        "outputId": "c96d3a1d-ebf6-471a-a93f-c96b479cc9fe"
      },
      "source": [
        "test_env = make_env(DEFAULT_ENV_NAME)\n",
        "test_net = DQN(test_env.observation_space.shape, test_env.action_space.n).to(device)\n",
        "print(test_net)"
      ],
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "DQN(\n",
            "  (conv): Sequential(\n",
            "    (0): Conv2d(4, 32, kernel_size=(8, 8), stride=(4, 4))\n",
            "    (1): ReLU()\n",
            "    (2): Conv2d(32, 64, kernel_size=(4, 4), stride=(2, 2))\n",
            "    (3): ReLU()\n",
            "    (4): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1))\n",
            "    (5): ReLU()\n",
            "  )\n",
            "  (fc): Sequential(\n",
            "    (0): Linear(in_features=3136, out_features=512, bias=True)\n",
            "    (1): ReLU()\n",
            "    (2): Linear(in_features=512, out_features=6, bias=True)\n",
            "  )\n",
            ")\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Lhv3Yf-aW7UW",
        "colab_type": "text"
      },
      "source": [
        "## Training"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aPJl73Z1YTa4",
        "colab_type": "text"
      },
      "source": [
        "Load Tensorboard extension"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BCBQhXLfNeUG",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from torch.utils.tensorboard import SummaryWriter\n",
        "%load_ext tensorboard"
      ],
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kb_f_onMXkpb",
        "colab_type": "text"
      },
      "source": [
        "Import required modules and define the hyperparameters"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AGwHC9dyXoPd",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import time\n",
        "import numpy as np\n",
        "import collections\n",
        "\n",
        "\n",
        "MEAN_REWARD_BOUND = 19.0           \n",
        "\n",
        "gamma = 0.99                   \n",
        "batch_size = 32                \n",
        "replay_size = 10000            \n",
        "learning_rate = 1e-4           \n",
        "sync_target_frames = 1000      \n",
        "replay_start_size = 10000      \n",
        "\n",
        "eps_start=1.0\n",
        "eps_decay=.999985\n",
        "eps_min=0.02"
      ],
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FFaMmDKqYmo4",
        "colab_type": "text"
      },
      "source": [
        "Experience replay buffer"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Y79CNYsjY4w0",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "Experience = collections.namedtuple('Experience', field_names=['state', 'action', 'reward', 'done', 'new_state'])\n",
        "\n",
        "class ExperienceReplay:\n",
        "    def __init__(self, capacity):\n",
        "        self.buffer = collections.deque(maxlen=capacity)\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.buffer)\n",
        "\n",
        "    def append(self, experience):\n",
        "        self.buffer.append(experience)\n",
        "\n",
        "    def sample(self, batch_size):\n",
        "        indices = np.random.choice(len(self.buffer), batch_size, replace=False)\n",
        "        states, actions, rewards, dones, next_states = zip(*[self.buffer[idx] for idx in indices])\n",
        "        return np.array(states), np.array(actions), np.array(rewards, dtype=np.float32), \\\n",
        "               np.array(dones, dtype=np.uint8), np.array(next_states)\n"
      ],
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fQDV04ktY3xs",
        "colab_type": "text"
      },
      "source": [
        "Agent"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YdAKFiMWZw90",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Agent:\n",
        "    def __init__(self, env, exp_buffer):\n",
        "        self.env = env\n",
        "        self.exp_buffer = exp_buffer\n",
        "        self._reset()\n",
        "\n",
        "    def _reset(self):\n",
        "        self.state = env.reset()\n",
        "        self.total_reward = 0.0\n",
        "\n",
        "    def play_step(self, net, epsilon=0.0, device=\"cpu\"):\n",
        "\n",
        "        done_reward = None\n",
        "        if np.random.random() < epsilon:\n",
        "            action = env.action_space.sample()\n",
        "        else:\n",
        "            state_a = np.array([self.state], copy=False)\n",
        "            state_v = torch.tensor(state_a).to(device)\n",
        "            q_vals_v = net(state_v)\n",
        "            _, act_v = torch.max(q_vals_v, dim=1)\n",
        "            action = int(act_v.item())\n",
        "\n",
        "        new_state, reward, is_done, _ = self.env.step(action)\n",
        "        self.total_reward += reward\n",
        "\n",
        "        exp = Experience(self.state, action, reward, is_done, new_state)\n",
        "        self.exp_buffer.append(exp)\n",
        "        self.state = new_state\n",
        "        if is_done:\n",
        "            done_reward = self.total_reward\n",
        "            self._reset()\n",
        "        return done_reward\n"
      ],
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ipurwYpa6iKn",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import datetime\n",
        "print(\">>>Training starts at \",datetime.datetime.now())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bgpmAtchZwM_",
        "colab_type": "text"
      },
      "source": [
        "Main training loop"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qEoc2PWmM2mu",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "env = make_env(DEFAULT_ENV_NAME)\n",
        "\n",
        "net = DQN(env.observation_space.shape, env.action_space.n).to(device)\n",
        "target_net = DQN(env.observation_space.shape, env.action_space.n).to(device)\n",
        "writer = SummaryWriter(comment=\"-\" + DEFAULT_ENV_NAME)\n",
        " \n",
        "buffer = ExperienceReplay(replay_size)\n",
        "agent = Agent(env, buffer)\n",
        "\n",
        "epsilon = eps_start\n",
        "\n",
        "optimizer = optim.Adam(net.parameters(), lr=learning_rate)\n",
        "total_rewards = []\n",
        "frame_idx = 0  \n",
        "\n",
        "best_mean_reward = None\n",
        "\n",
        "while True:\n",
        "        frame_idx += 1\n",
        "        epsilon = max(epsilon*eps_decay, eps_min)\n",
        "\n",
        "        reward = agent.play_step(net, epsilon, device=device)\n",
        "        if reward is not None:\n",
        "            total_rewards.append(reward)\n",
        "\n",
        "            mean_reward = np.mean(total_rewards[-100:])\n",
        "\n",
        "            print(\"%d:  %d games, mean reward %.3f, (epsilon %.2f)\" % (\n",
        "                frame_idx, len(total_rewards), mean_reward, epsilon))\n",
        "            \n",
        "            writer.add_scalar(\"epsilon\", epsilon, frame_idx)\n",
        "            writer.add_scalar(\"reward_100\", mean_reward, frame_idx)\n",
        "            writer.add_scalar(\"reward\", reward, frame_idx)\n",
        "\n",
        "            if best_mean_reward is None or best_mean_reward < mean_reward:\n",
        "                torch.save(net.state_dict(), DEFAULT_ENV_NAME + \"-best.dat\")\n",
        "                best_mean_reward = mean_reward\n",
        "                if best_mean_reward is not None:\n",
        "                    print(\"Best mean reward updated %.3f\" % (best_mean_reward))\n",
        "\n",
        "            if mean_reward > MEAN_REWARD_BOUND:\n",
        "                print(\"Solved in %d frames!\" % frame_idx)\n",
        "                break\n",
        "\n",
        "        if len(buffer) < replay_start_size:\n",
        "            continue\n",
        "\n",
        "        batch = buffer.sample(batch_size)\n",
        "        states, actions, rewards, dones, next_states = batch\n",
        "\n",
        "        states_v = torch.tensor(states).to(device)\n",
        "        next_states_v = torch.tensor(next_states).to(device)\n",
        "        actions_v = torch.tensor(actions).to(device)\n",
        "        rewards_v = torch.tensor(rewards).to(device)\n",
        "        done_mask = torch.ByteTensor(dones).to(device)\n",
        "\n",
        "        state_action_values = net(states_v).gather(1, actions_v.unsqueeze(-1)).squeeze(-1)\n",
        "\n",
        "        next_state_values = target_net(next_states_v).max(1)[0]\n",
        "\n",
        "        next_state_values[done_mask] = 0.0\n",
        "\n",
        "        next_state_values = next_state_values.detach()\n",
        "\n",
        "        expected_state_action_values = next_state_values * gamma + rewards_v\n",
        "\n",
        "        loss_t = nn.MSELoss()(state_action_values, expected_state_action_values)\n",
        "\n",
        "        optimizer.zero_grad()\n",
        "        loss_t.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "        if frame_idx % sync_target_frames == 0:\n",
        "            target_net.load_state_dict(net.state_dict())\n",
        "       \n",
        "writer.close()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MZPkszw66cmO",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "print(\">>>Training ends at \",datetime.datetime.now())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qNH2N64k3QRz",
        "colab_type": "text"
      },
      "source": [
        "Performance"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WKbcwfK321Hl",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "tensorboard  --logdir=runs"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6p0jvxoC3m5W",
        "colab_type": "text"
      },
      "source": [
        "## Using the model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TLEfbkKl6AZV",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import gym\n",
        "import time\n",
        "import numpy as np\n",
        "\n",
        "import torch\n",
        "\n",
        "import collections\n",
        "\n",
        "DEFAULT_ENV_NAME = \"PongNoFrameskip-v4\"\n",
        "FPS = 25"
      ],
      "execution_count": 18,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4m0Vm4Yp91ZI",
        "colab_type": "text"
      },
      "source": [
        "Tunning the image rendering in colab\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kgpHXywd5SyZ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Taken from \n",
        "# https://towardsdatascience.com/rendering-openai-gym-envs-on-binder-and-google-colab-536f99391cc7\n",
        "\n",
        "!apt-get install -y xvfb x11-utils\n",
        "\n",
        "!pip install pyvirtualdisplay==0.2.* \\\n",
        "             PyOpenGL==3.1.* \\\n",
        "             PyOpenGL-accelerate==3.1.*\n",
        "\n",
        "!pip install gym[box2d]==0.17.*\n",
        "\n",
        "import pyvirtualdisplay\n",
        "\n",
        "_display = pyvirtualdisplay.Display(visible=False, size=(1400, 900))\n",
        "_ = _display.start()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BvN4S8R53mJI",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Taken (partially) from \n",
        "# https://github.com/PacktPublishing/Deep-Reinforcement-Learning-Hands-On/blob/master/Chapter06/03_dqn_play.py\n",
        "\n",
        "\n",
        "model='PongNoFrameskip-v4-best.dat'\n",
        "record_folder=\"video\"  \n",
        "visualize=True\n",
        "\n",
        "env = make_env(DEFAULT_ENV_NAME)\n",
        "if record_folder:\n",
        "        env = gym.wrappers.Monitor(env, record_folder, force=True)\n",
        "net = DQN(env.observation_space.shape, env.action_space.n)\n",
        "net.load_state_dict(torch.load(model, map_location=lambda storage, loc: storage))\n",
        "\n",
        "state = env.reset()\n",
        "total_reward = 0.0\n",
        "\n",
        "while True:\n",
        "        start_ts = time.time()\n",
        "        if visualize:\n",
        "            env.render()\n",
        "        state_v = torch.tensor(np.array([state], copy=False))\n",
        "        q_vals = net(state_v).data.numpy()[0]\n",
        "        action = np.argmax(q_vals)\n",
        "        \n",
        "        state, reward, done, _ = env.step(action)\n",
        "        total_reward += reward\n",
        "        if done:\n",
        "            break\n",
        "        if visualize:\n",
        "            delta = 1/FPS - (time.time() - start_ts)\n",
        "            if delta > 0:\n",
        "                time.sleep(delta)\n",
        "print(\"Total reward: %.2f\" % total_reward)\n",
        "\n",
        "if record_folder:\n",
        "        env.close()"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}