Pong_AC.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b20f2cab-1ee5-4c0b-9d63-addece2efa39",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install gymnasium[atari]\n",
    "!pip install gymnasium[accept-rom-license]\n",
    "!pip install moviepy\n",
    "!pip install pysdl2\n",
    "!pip install pyvirtualdisplay"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc32e69c-6cd1-4fae-b0b0-8eb8bc0baef4",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install opencv-python"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b1fe38af-1750-40f7-b33f-74ed02ff1c5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import gymnasium as gym\n",
    "import random \n",
    "import torch\n",
    "import numpy as np\n",
    "from torch import nn\n",
    "from torch.distributions import Categorical\n",
    "\n",
    "from gymnasium import wrappers\n",
    "import matplotlib.pyplot as plt\n",
    "from IPython import display\n",
    "from tqdm.notebook import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "864d50f9-5b76-4e6d-9572-9b774cf39a3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "#gym.envs.registration.registry.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "96160bee-4476-4add-91cf-2e50e43f0b03",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "A.L.E: Arcade Learning Environment (version 0.8.1+53f58b7)\n",
      "[Powered by Stella]\n",
      "/opt/conda/lib/python3.10/site-packages/gymnasium/wrappers/record_video.py:94: UserWarning: \u001b[33mWARN: Overwriting existing videos at /home/chgaw002/DeepLearning/video/pong-base folder (try specifying a different `video_folder` for the `RecordVideo` wrapper if this is not desired)\u001b[0m\n",
      "  logger.warn(\n"
     ]
    }
   ],
   "source": [
    "env = gym.make(\"ALE/Pong-v5\", render_mode=\"rgb_array\")\n",
    "# record the game as as an mp4 file\n",
    "env = wrappers.RecordVideo(env, video_folder='video/pong-base', episode_trigger=lambda n: n%200==0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9077e2fd-1b30-4e5d-a71b-fa8f77823656",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['NOOP', 'FIRE', 'RIGHT', 'LEFT', 'RIGHTFIRE', 'LEFTFIRE']"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "env.unwrapped.get_action_meanings()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "303ee66b-58f1-4436-9096-66901e1d6695",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from collections import deque\n",
    "import cv2\n",
    "cv2.ocl.setUseOpenCL(False)\n",
    "\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",
    "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 = gym.spaces.Box(low=0, high=255, shape=(shp[0], shp[1], shp[2] * k), dtype=env.observation_space.dtype)\n",
    "\n",
    "    def reset(self, **kwargs):\n",
    "        ob, info = self.env.reset(**kwargs)\n",
    "        for _ in range(self.k):\n",
    "            self.frames.append(ob)\n",
    "        return self._get_ob(), info\n",
    "\n",
    "    def step(self, action):\n",
    "        ob, reward, done, truncated, info = self.env.step(action)\n",
    "        self.frames.append(ob)\n",
    "        return self._get_ob(), reward, done, truncated, info\n",
    "\n",
    "    def _get_ob(self):\n",
    "        assert len(self.frames) == self.k\n",
    "        return LazyFrames(list(self.frames))\n",
    "\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 = gym.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 EpisodicLifeEnv(gym.Wrapper):\n",
    "    def __init__(self, env=None):\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",
    "        super(EpisodicLifeEnv, self).__init__(env)\n",
    "        self.lives = 0\n",
    "        self.was_real_done = True\n",
    "        self.was_real_reset = False\n",
    "\n",
    "    def step(self, action):\n",
    "        obs, reward, done, truncated, 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 somtimes 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, truncated, 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, info = self.env.reset(**kwargs)\n",
    "            self.was_real_reset = True\n",
    "        else:\n",
    "            # no-op step to advance from terminal/lost life state\n",
    "            obs, _, _, _, info = self.env.step(0)\n",
    "            self.was_real_reset = False\n",
    "        self.lives = self.env.unwrapped.ale.lives()\n",
    "        return obs, info\n",
    "    \n",
    "class MaxAndSkipEnv(gym.Wrapper):\n",
    "    def __init__(self, env=None, skip=4):\n",
    "        \"\"\"Return only every `skip`-th frame\"\"\"\n",
    "        super(MaxAndSkipEnv, self).__init__(env)\n",
    "        # most recent raw observations (for max pooling across time steps)\n",
    "        self._obs_buffer = 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, truncated, info = self.env.step(action)\n",
    "            self._obs_buffer.append(obs)\n",
    "            total_reward += reward\n",
    "            if done:\n",
    "                break\n",
    "\n",
    "        max_frame = np.max(np.stack(self._obs_buffer), axis=0)\n",
    "\n",
    "        return max_frame, total_reward, done, truncated, info\n",
    "\n",
    "    def reset(self, **kwargs):\n",
    "        \"\"\"Clear past frame buffer and init. to first obs. from inner env.\"\"\"\n",
    "        self._obs_buffer.clear()\n",
    "        obs, info = self.env.reset(**kwargs)\n",
    "        self._obs_buffer.append(obs)\n",
    "        return obs, info\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "fca7cd42-3672-4c19-9444-4f92515983a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_env(env, stack_frames=True, episodic_life=True):\n",
    "    if episodic_life:\n",
    "        env = EpisodicLifeEnv(env)\n",
    "\n",
    "    env = MaxAndSkipEnv(env, skip=4)\n",
    "\n",
    "    env = WarpFrame(env)\n",
    "    if stack_frames:\n",
    "        env = FrameStack(env, 4)\n",
    "  \n",
    "    return env\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "9f601ecc-c1df-426d-8e6a-548ef60983b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "env = make_env(env)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "962e5704-6cd3-4d41-9480-7a7b662d9164",
   "metadata": {},
   "outputs": [],
   "source": [
    "def render(env):\n",
    "    img.set_data(env.render())\n",
    "    display.display(plt.gcf())\n",
    "    display.clear_output(wait=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "e0782329-cd11-4bab-bff2-dbc09240a1d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F\n",
    "import torchvision.transforms as T\n",
    "\n",
    "class ActorCritic(nn.Module):\n",
    " \n",
    "    def __init__(self, in_channels=4, n_actions=14):\n",
    "        \"\"\"\n",
    "        Initialize Actor Network\n",
    "        Args:\n",
    "            in_channels (int): number of input channels\n",
    "            n_actions (int): number of outputs\n",
    "        \"\"\"\n",
    "        super(ActorCritic, self).__init__()\n",
    "        self.actor_features = nn.Sequential(\n",
    "            nn.Conv2d(in_channels, 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",
    "            nn.Flatten(),\n",
    "            nn.Linear(7 * 7 * 64, 256),\n",
    "            nn.ReLU()\n",
    "        )\n",
    "        self.critic_features = nn.Sequential(\n",
    "            nn.Conv2d(in_channels, 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",
    "            nn.Flatten(),\n",
    "            nn.Linear(7 * 7 * 64, 256),\n",
    "            nn.ReLU()\n",
    "        )\n",
    "        self.actor_head = nn.Linear(256, n_actions)\n",
    "        self.critic_head = nn.Linear(256, 1)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = x.float() / 255\n",
    "        x_a = self.actor_features(x)\n",
    "        x_c = self.critic_features(x)\n",
    "        return F.softmax(self.actor_head(x_a), dim=-1), self.critic_head(x_c)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "b88156d3-935d-4f0f-a007-b8cd60c517f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import namedtuple\n",
    "import random\n",
    "\n",
    "SavedAction = namedtuple('SavedAction', \n",
    "                        ('log_prob', 'value'))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "026daa27-f114-4092-94b6-abf8a83a4140",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "class ActorCriticPolicy:\n",
    "    \n",
    "    GAMMA = 0.99\n",
    "    EPS = 1e-7\n",
    "    RENDER = False\n",
    "    # NOOP, UP, DOWN\n",
    "    ACTIONS = [ 0, 2, 3 ]\n",
    "\n",
    "    \n",
    "    def __init__(self, lr=1e-4):\n",
    "        self.n_actions = len(ActorCriticPolicy.ACTIONS)\n",
    "        self.steps_done = 0\n",
    "        self.mean_reward = None\n",
    "        self.model = ActorCritic(n_actions=self.n_actions).to(device)\n",
    "        self.saved_actions = []\n",
    "        self.optimizer = optim.AdamW(self.model.parameters(), lr=lr)\n",
    "        self.postfix = {}\n",
    "        \n",
    "    def get_state(self, obs):\n",
    "        state = np.array(obs)\n",
    "        state = state.transpose((2, 0, 1))\n",
    "        state = torch.from_numpy(state)\n",
    "        return state.unsqueeze(0).to(device)\n",
    "\n",
    "    def __call__(self, observation):\n",
    "        state = self.get_state(observation)\n",
    "        probs, value = self.model(state)\n",
    "        \n",
    "        m = Categorical(probs)\n",
    "        action = m.sample()\n",
    "        \n",
    "        self.saved_actions.append(SavedAction(m.log_prob(action), value))\n",
    "        self.steps_done += 1\n",
    "        \n",
    "        return ActorCriticPolicy.ACTIONS[action.item()]\n",
    "  \n",
    "        \n",
    "    def init_game(self, observation):\n",
    "        self.state = self.get_state(observation)\n",
    "        self.total_reward = 0.0\n",
    "        self.rewards = []\n",
    "    \n",
    "    def update(self, obs, reward, terminated, truncated, info, pbar):\n",
    "        self.total_reward += reward\n",
    "        self.rewards.append(reward)\n",
    "        if not terminated:\n",
    "            self.next_state = self.get_state(obs)\n",
    "        else:\n",
    "            if self.mean_reward is None:\n",
    "                self.mean_reward = self.total_reward\n",
    "            else:\n",
    "                self.mean_reward = self.mean_reward * 0.95 + self.total_reward * (1.0 - 0.95)\n",
    "            self.postfix['total_reward'] = self.total_reward\n",
    "            self.postfix['mean_reward'] = self.mean_reward\n",
    "            self.postfix['steps'] = self.steps_done\n",
    "            pbar.set_postfix(self.postfix)\n",
    "            self.next_state = None\n",
    "            \n",
    "            self.finish_episode()    \n",
    "            \n",
    "        if self.steps_done % 100_000 == 0:\n",
    "            self.save(f'model_{self.steps_done}.pt')\n",
    "\n",
    "    def finish_episode(self):\n",
    "        R = 0\n",
    "        policy_losses = []\n",
    "        value_losses = []\n",
    "        returns = []\n",
    "        \n",
    "        for r in self.rewards[::-1]:\n",
    "            if r != 0:\n",
    "                # Game boundary (Pong specific) !\n",
    "                R = 0\n",
    "            R = r + ActorCriticPolicy.GAMMA * R\n",
    "            returns.insert(0, R)\n",
    "            \n",
    "        returns = torch.tensor(returns)\n",
    "        returns = (returns - returns.mean()) / (returns.std() + ActorCriticPolicy.EPS)\n",
    "        \n",
    "        for (log_prob, value), R in zip(self.saved_actions, returns):\n",
    "            advantage = R - value.item()\n",
    "\n",
    "            # calculate actor (policy) loss\n",
    "            policy_losses.append(-log_prob * advantage)\n",
    "\n",
    "            # calculate critic (value) loss using L1 smooth loss\n",
    "            value_losses.append(F.smooth_l1_loss(value, torch.tensor([R]).unsqueeze(0).to(device)))\n",
    "            \n",
    "        self.optimizer.zero_grad()\n",
    "        # sum up all the values of policy_losses and value_losses\n",
    "        policy_loss = torch.stack(policy_losses).sum()\n",
    "        value_loss = torch.stack(value_losses).sum()\n",
    "        loss = policy_loss + value_loss\n",
    "        \n",
    "        self.postfix['policy_loss'] = policy_loss.item()\n",
    "        self.postfix['value_loss'] = value_loss.item()\n",
    "        self.postfix['loss'] = loss.item()\n",
    "\n",
    "        # perform backprop\n",
    "        loss.backward()\n",
    "        for param in self.model.parameters():\n",
    "            param.grad.data.clamp_(-1, 1)\n",
    "        self.optimizer.step()\n",
    "\n",
    "        # reset rewards and action buffer\n",
    "        del self.rewards[:]\n",
    "        del self.saved_actions[:]\n",
    "\n",
    "        \n",
    "    def load(self, PATH):\n",
    "        checkpoint = torch.load(PATH)\n",
    "        print(checkpoint.keys())\n",
    "        self.model.load_state_dict(checkpoint['model_state_dict'])\n",
    "        self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n",
    "        self.steps_done = checkpoint['steps_done']\n",
    "        if \"mean_reward\" in checkpoint:\n",
    "            self.mean_reward = checkpoint['mean_reward']\n",
    "        \n",
    "    def save(self, PATH):\n",
    "        state = {\n",
    "                    'steps_done': self.steps_done,\n",
    "                    'model_state_dict': self.model.state_dict(),\n",
    "                    'optimizer_state_dict': self.optimizer.state_dict(),\n",
    "                    'mean_reward': self.mean_reward\n",
    "        }\n",
    "        print(state.keys())\n",
    "        torch.save(state, PATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "2e63ce68-4b97-43dc-a817-8694357d2f6d",
   "metadata": {},
   "outputs": [],
   "source": [
    "policy = ActorCriticPolicy(lr=1e-5)\n",
    "#policy.load(\"ac_model.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "c1783180-33be-459a-b839-d1bd100c5ead",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c06da3a7e4184594b7a00410364f826e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/10000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Moviepy - Building video /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-400.mp4.\n",
      "Moviepy - Writing video /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-400.mp4\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "t:   0%|          | 0/889 [00:00<?, ?it/s, now=None]\u001b[A\n",
      "t:  37%|███▋      | 332/889 [00:00<00:00, 3319.73it/s, now=None]\u001b[A\n",
      "t:  75%|███████▍  | 664/889 [00:00<00:00, 3229.32it/s, now=None]\u001b[A\n",
      "                                                                \u001b[A"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Moviepy - Done !\n",
      "Moviepy - video ready /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-400.mp4\n",
      "Moviepy - Building video /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-600.mp4.\n",
      "Moviepy - Writing video /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-600.mp4\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "t:   0%|          | 0/949 [00:00<?, ?it/s, now=None]\u001b[A\n",
      "t:  38%|███▊      | 356/949 [00:00<00:00, 3559.57it/s, now=None]\u001b[A\n",
      "t:  87%|████████▋ | 823/949 [00:00<00:00, 4208.96it/s, now=None]\u001b[A\n",
      "                                                                \u001b[A"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Moviepy - Done !\n",
      "Moviepy - video ready /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-600.mp4\n",
      "dict_keys(['steps_done', 'model_state_dict', 'optimizer_state_dict', 'mean_reward'])\n",
      "Moviepy - Building video /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-800.mp4.\n",
      "Moviepy - Writing video /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-800.mp4\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "t:   0%|          | 0/947 [00:00<?, ?it/s, now=None]\u001b[A\n",
      "t:  23%|██▎       | 221/947 [00:00<00:00, 2205.62it/s, now=None]\u001b[A\n",
      "t:  62%|██████▏   | 589/947 [00:00<00:00, 3070.57it/s, now=None]\u001b[A\n",
      "t:  97%|█████████▋| 918/947 [00:00<00:00, 3170.64it/s, now=None]\u001b[A\n",
      "                                                                \u001b[A"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Moviepy - Done !\n",
      "Moviepy - video ready /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-800.mp4\n",
      "Moviepy - Building video /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-1000.mp4.\n",
      "Moviepy - Writing video /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-1000.mp4\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "t:   0%|          | 0/948 [00:00<?, ?it/s, now=None]\u001b[A\n",
      "t:  27%|██▋       | 254/948 [00:00<00:00, 2536.30it/s, now=None]\u001b[A\n",
      "t:  68%|██████▊   | 644/948 [00:00<00:00, 3335.18it/s, now=None]\u001b[A\n",
      "                                                                \u001b[A"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Moviepy - Done !\n",
      "Moviepy - video ready /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-1000.mp4\n",
      "dict_keys(['steps_done', 'model_state_dict', 'optimizer_state_dict', 'mean_reward'])\n",
      "Moviepy - Building video /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-1200.mp4.\n",
      "Moviepy - Writing video /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-1200.mp4\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "t:   0%|          | 0/1200 [00:00<?, ?it/s, now=None]\u001b[A\n",
      "t:  21%|██        | 254/1200 [00:00<00:00, 2539.65it/s, now=None]\u001b[A\n",
      "t:  53%|█████▎    | 632/1200 [00:00<00:00, 3268.45it/s, now=None]\u001b[A\n",
      "t:  90%|████████▉ | 1075/1200 [00:00<00:00, 3796.84it/s, now=None]\u001b[A\n",
      "                                                                  \u001b[A"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Moviepy - Done !\n",
      "Moviepy - video ready /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-1200.mp4\n",
      "Moviepy - Building video /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-1400.mp4.\n",
      "Moviepy - Writing video /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-1400.mp4\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "t:   0%|          | 0/946 [00:00<?, ?it/s, now=None]\u001b[A\n",
      "t:  37%|███▋      | 352/946 [00:00<00:00, 3517.07it/s, now=None]\u001b[A\n",
      "t:  88%|████████▊ | 829/946 [00:00<00:00, 4251.96it/s, now=None]\u001b[A\n",
      "                                                                \u001b[A"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Moviepy - Done !\n",
      "Moviepy - video ready /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-1400.mp4\n",
      "dict_keys(['steps_done', 'model_state_dict', 'optimizer_state_dict', 'mean_reward'])\n",
      "Moviepy - Building video /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-1600.mp4.\n",
      "Moviepy - Writing video /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-1600.mp4\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "t:   0%|          | 0/1134 [00:00<?, ?it/s, now=None]\u001b[A\n",
      "t:  14%|█▍        | 160/1134 [00:00<00:00, 1599.32it/s, now=None]\u001b[A\n",
      "t:  47%|████▋     | 534/1134 [00:00<00:00, 2857.85it/s, now=None]\u001b[A\n",
      "t:  86%|████████▌ | 976/1134 [00:00<00:00, 3568.08it/s, now=None]\u001b[A\n",
      "                                                                 \u001b[A"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Moviepy - Done !\n",
      "Moviepy - video ready /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-1600.mp4\n",
      "Moviepy - Building video /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-2600.mp4.\n",
      "Moviepy - Writing video /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-2600.mp4\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "t:   0%|          | 0/887 [00:00<?, ?it/s, now=None]\u001b[A\n",
      "t:  27%|██▋       | 236/887 [00:00<00:00, 2353.92it/s, now=None]\u001b[A\n",
      "t:  66%|██████▌   | 587/887 [00:00<00:00, 3032.05it/s, now=None]\u001b[A\n",
      "                                                                \u001b[A"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Moviepy - Done !\n",
      "Moviepy - video ready /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-2600.mp4\n",
      "dict_keys(['steps_done', 'model_state_dict', 'optimizer_state_dict', 'mean_reward'])\n",
      "Moviepy - Building video /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-2800.mp4.\n",
      "Moviepy - Writing video /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-2800.mp4\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "t:   0%|          | 0/1007 [00:00<?, ?it/s, now=None]\u001b[A\n",
      "t:  32%|███▏      | 319/1007 [00:00<00:00, 3185.26it/s, now=None]\u001b[A\n",
      "t:  67%|██████▋   | 675/1007 [00:00<00:00, 3404.78it/s, now=None]\u001b[A\n",
      "                                                                 \u001b[A"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Moviepy - Done !\n",
      "Moviepy - video ready /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-2800.mp4\n",
      "Moviepy - Building video /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-3000.mp4.\n",
      "Moviepy - Writing video /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-3000.mp4\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "t:   0%|          | 0/766 [00:00<?, ?it/s, now=None]\u001b[A\n",
      "t:  32%|███▏      | 246/766 [00:00<00:00, 2454.88it/s, now=None]\u001b[A\n",
      "t:  80%|███████▉  | 610/766 [00:00<00:00, 3150.92it/s, now=None]\u001b[A\n",
      "                                                                \u001b[A"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Moviepy - Done !\n",
      "Moviepy - video ready /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-3000.mp4\n",
      "dict_keys(['steps_done', 'model_state_dict', 'optimizer_state_dict', 'mean_reward'])\n",
      "Moviepy - Building video /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-3200.mp4.\n",
      "Moviepy - Writing video /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-3200.mp4\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "t:   0%|          | 0/886 [00:00<?, ?it/s, now=None]\u001b[A\n",
      "t:  34%|███▍      | 305/886 [00:00<00:00, 3043.91it/s, now=None]\u001b[A\n",
      "t:  73%|███████▎  | 644/886 [00:00<00:00, 3242.86it/s, now=None]\u001b[A\n",
      "                                                                \u001b[A"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Moviepy - Done !\n",
      "Moviepy - video ready /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-3200.mp4\n",
      "Moviepy - Building video /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-3400.mp4.\n",
      "Moviepy - Writing video /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-3400.mp4\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "t:   0%|          | 0/766 [00:00<?, ?it/s, now=None]\u001b[A\n",
      "t:  33%|███▎      | 255/766 [00:00<00:00, 2547.42it/s, now=None]\u001b[A\n",
      "t:  90%|█████████ | 692/766 [00:00<00:00, 3615.69it/s, now=None]\u001b[A\n",
      "                                                                \u001b[A"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Moviepy - Done !\n",
      "Moviepy - video ready /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-3400.mp4\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub message rate exceeded.\n",
      "The Jupyter server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--ServerApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "ServerApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "ServerApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dict_keys(['steps_done', 'model_state_dict', 'optimizer_state_dict', 'mean_reward'])\n",
      "Moviepy - Building video /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-18200.mp4.\n",
      "Moviepy - Writing video /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-18200.mp4\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "t:   0%|          | 0/948 [00:00<?, ?it/s, now=None]\u001b[A\n",
      "t:  31%|███       | 294/948 [00:00<00:00, 2938.22it/s, now=None]\u001b[A\n",
      "t:  69%|██████▊   | 651/948 [00:00<00:00, 3303.67it/s, now=None]\u001b[A\n",
      "                                                                \u001b[A"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Moviepy - Done !\n",
      "Moviepy - video ready /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-18200.mp4\n",
      "Moviepy - Building video /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-18400.mp4.\n",
      "Moviepy - Writing video /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-18400.mp4\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "t:   0%|          | 0/890 [00:00<?, ?it/s, now=None]\u001b[A\n",
      "t:  41%|████      | 366/890 [00:00<00:00, 3658.42it/s, now=None]\u001b[A\n",
      "t:  95%|█████████▍| 845/890 [00:00<00:00, 4319.41it/s, now=None]\u001b[A\n",
      "                                                                \u001b[A"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Moviepy - Done !\n",
      "Moviepy - video ready /home/chgaw002/DeepLearning/video/pong-base/rl-video-episode-18400.mp4\n",
      "dict_keys(['steps_done', 'model_state_dict', 'optimizer_state_dict', 'mean_reward'])\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAS0AAAGFCAYAAACorKVtAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAJfUlEQVR4nO3dTW9cZxnH4ft4XmI7TWjiNq1TkjbqgkiUDRUCJCQWbEAs+QJs2LCgbPkYbLrhM3TFogsWLJEqsUAVXZSqbdqmbUyStsqL7Xk7LIoqkCfx2GPP5G9f1y56zpzccuxfzjyeOdO0bdsWQIiVZQ8AcBCiBUQRLSCKaAFRRAuIIlpAFNECoogWEKU764FN0xz45Gvdpn73w3O1ea5z4McCp89rb97d95iZo3X9mZkP/caZTlNnDv6wuZw7u179Xm/q2r0HD2swHB74nOurq7W2embq2sOdndre2T3wOTnZti+crcH59SM9Z//edq3dvX+k50w0c1J+8/2nDvUXrBz8Am0u1174dj27cXHq2j//9V59fvv2gc/5wnOX6urly1PX3v/44/rgk5sHPicn2+1XrtStV68d6Tkv/eNGXf3rO0d6zkQzR6uz6PocVlO1MuWp7HxvsWymnvPrlZCvC4vVNFUrR7tl3B5ii+YkshEPRBEtIIpoAVFEC4iy4BckAP2vHlb/3vbUtcG5tRp862hfKnHSiBYs2MY7N2vzrfemrt169Vrd/Mn1BU+URbRgwZq2rZXxZPrixN3P92NPC4giWkAU0QKiiBYQxUb8nM70+3X+qbMHftxgMKydweAYJoKTTbTmdPnSs7X57DMHftwnt27Vux/eOIaJ4GQTrTmtHPKd/E3jmTkchp8cIIpoAVFEC4giWkAUG/Fz2h0MajgcTV3r9bp1pt9f8ERwsonWnD7d2qoPb346de3q5ma9fPXKgieCk0205jSZtDUaj6eujSePeCc/cGj2tIAoogVEES0gimgBUWzEz6nX7dba6urUtX7Pl5e9Rqu92nl6+odXjNa8RGY/fqrm9MJzl+r5R9zloXPEH4vOyXD7e1fq7vXLU9cmvc6Cp8kjWnPqdDrV6fhGY3aTXrcmrsIPzaUAEEW0gCiiBUQRLSDKidsNHA5HtbO7O3Vtcsj3Ao7Gjz7nYY1G0+8MwcnQ2R1W7972kZ6zuzM80vOlatq2nelzuP/4i4vHPcuR6Ha7j3ypwXA0OlS4usfwG8LxePzIN1qTb9Tv1qR/tN8zK4NxdQcn+z+71968u+8xJ+5KazQa1VH/s44EhgPqDkZVJzwwy2JPC4giWkCU2Z8eNs0xjgEwm5mj9dPfv36ccwDMZObfHt65c+e4ZwFOuY2NjX2PsacFRBEtIIpoAVFEC4giWkAU0QKiiBYQRbSAKKIFRBEtIIpoAVFEC4giWkAU0QKiiBYQZeabAH758bvHOQdAbWz8eN9jZr4J4Ou/enHugQAe57dv3Nj3mJmvtIYP7801DMBRsKcFRBEtIIpoAVFEC4giWkAU0QKiiBYQRbSAKKIFRBEtIIpoAVFEC4giWkAU0QKiiBYQRbSAKKIFRBEtIIpoAVFEC4giWkAU0QKiiBYQRbSAKKIFRBEtIIpoAVFEC4giWkAU0QKiiBYQRbSAKKIFRBEtIIpoAVFEC4giWkAU0QKiiBYQRbSAKKIFRBEtIIpoAVFEC4giWkAU0QKiiBYQRbSAKKIFRBEtIIpoAVFEC4giWkAU0QKiiBYQRbSAKKIFRBEtIIpoAVFEC4giWkAU0QKiiBYQRbSAKKIFRBEtIIpoAVFEC4giWkAU0QKiiBYQRbSAKKIFRBEtIIpoAVFEC4giWkAU0QKiiBYQRbSAKKIFRBEtIIpoAVFEC4giWkAU0QKiiBYQRbSAKKIFRBEtIIpoAVFEC4giWkAU0QKiiBYQRbSAKKIFRBEtIIpoAVG6yx4AWJ7JSlOT3vQMNJNJrQzH1Sx4pv2IFpxiDzYv1Ec/+261zd40nf38y3rxL29XM2mXMNmjiRacYuN+p7Yvnqta2Rut3sPdJUy0P3taQBTRAqKIFhBFtIAoogVEES0gimgBUUQLiCJaQBTRAqKIFhBFtIAoogVEES0gimgBUUQLiCJaQBTRAqKIFhBFtIAoogVEES0gimgBUUQLiCJaQBSfMA2nWVvVtG217fS1J5FowSl2duurevnPf6+2afas9bYH1UyevHKJFpxivYeDevr9rWWPcSD2tIAoogVEES0gimgBUUQLiCJaQBTRAqKIFhBFtIAoogVEES0gimgBUUQLiCJaQBTRAqKIFhBFtIAoogVEES0gimgBUUQLiCJaQBTRAqKIFhBFtIAoogVEES0gimgBUUQLiCJaQBTRAqKIFhBFtIAoogVEES0gimgBUUQLiCJaQBTRAqKIFhBFtIAoogVEES0gimgBUUQLiCJaQBTRAqKIFhBFtIAoogVEES0gimgBUUQLiCJaQBTRAqKIFhBFtIAoogVEES0gimgBUUQLiCJaQBTRAqKIFhBFtIAoogVEES0gimgBUUQLiCJaQBTRAqKIFhBFtIAoogVEES0gimgBUUQLiCJaQBTRAqKIFhBFtIAoogVEES0gimgBUUQLiCJaQBTRAqKIFhBFtIAoogVEES0gimgBUUQLiCJaQBTRAqKIFhBFtIAoogVEES0gimgBUUQLiCJaQBTRAqKIFhBFtIAoogVEES0gimgBUUQLiCJaQBTRAqKIFhBFtIAoogVEES0gimgBUUQLiCJaQBTRAqKIFhBFtIAoogVEES0gSnfZA0zTVNWLly/X+trqnrW2rfros8/qwfb24gcDlu6JjFY1TT1z8UJdOH9+z9KkbWvr7h3RggVpOt3qdPv//VNbo8HO11cPS/JkRgt4Yrz0o1/W9Z//uqqqhtv3629/+kM9+PcnS5tHtIDHWj2/URvXXqmqqt37X1an19/nEcfLRjwQRbSAKKIFRBEtIIqNeOCx2sm4xoPdqqoaDwdLfblDlWgB+7jx1pt154O3q6pqMh7XgzufLXUe0QIea/uLrdr+YmvZY3zDnhYQRbSAKKIFRBEtIIpoAVFEC4giWkAU0QKiiBYQRbSAKKIFRBEtIIpoAVFEC4giWkCUme+ndek7PzjOOf5Ps9LU+pWXqre+vmetrbYuDteqc+/+wuYBnhxN285279Stzxd7t8Kmaapppq+1k7aWe8NX4Dhcen5z32NmvtJa9gc0/q+ms+wJgGWxpwVEES0gimgBUUQLiCJaQBTRAqKIFhBFtIAoogVEES0gimgBUUQLiCJaQBTRAqKIFhBFtIAoogVEES0gimgBUUQLiCJaQBTRAqKIFhBFtIAoogVEES0gStO2bbvsIQBm5UoLiCJaQBTRAqKIFhBFtIAoogVEES0gimgBUUQLiPIfscw79jvYCZkAAAAASUVORK5CYII=",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "observation, info = env.reset()\n",
    "policy.init_game(observation)\n",
    "\n",
    "plt.ion()\n",
    "plt.axis('off')\n",
    "img = plt.imshow(env.render())\n",
    "\n",
    "with tqdm(total=10000) as pbar:\n",
    "    while True:\n",
    "        try:\n",
    "            action = policy(observation)\n",
    "            observation, reward, terminated, truncated, info = env.step(action)\n",
    "            #render(env)\n",
    "            policy.update(observation, reward, terminated, truncated, info, pbar)\n",
    "\n",
    "            if terminated or truncated:\n",
    "                pbar.update()\n",
    "                observation, info = env.reset()\n",
    "                policy.init_game(observation)\n",
    "                \n",
    "        except KeyboardInterrupt:\n",
    "            break\n",
    "env.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ecfeabc-4fdb-4371-8cde-e1c3a32e7e9b",
   "metadata": {},
   "outputs": [],
   "source": [
    "policy.save(\"ac_model.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ad4b6af-78ab-4edb-a449-31fe87aea7b2",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}