aic502-mario-rl-nb.ipynb

{"cells":[{"cell_type":"markdown","metadata":{},"source":["## ![](https://i.imgur.com/dbbaf88.png)"]},{"cell_type":"markdown","metadata":{},"source":["# AIC-502 - Assessment 4\n","\n","# Name: Phasit Thanitkul (Kane)\n","\n","## Update on Nov 26, 2023."]},{"cell_type":"code","execution_count":3,"metadata":{"execution":{"iopub.execute_input":"2023-11-25T13:10:57.491153Z","iopub.status.busy":"2023-11-25T13:10:57.490375Z","iopub.status.idle":"2023-11-25T13:10:57.502090Z","shell.execute_reply":"2023-11-25T13:10:57.501088Z","shell.execute_reply.started":"2023-11-25T13:10:57.491109Z"},"trusted":true},"outputs":[],"source":["# TODO: make a versioning txt file\n","TRAIN_EPISODES = 100\n","SAVE_PATH = 'working/mario_net_last.chkpt'"]},{"cell_type":"code","execution_count":4,"metadata":{"execution":{"iopub.execute_input":"2023-11-25T13:10:57.503445Z","iopub.status.busy":"2023-11-25T13:10:57.503170Z","iopub.status.idle":"2023-11-25T13:11:13.206935Z","shell.execute_reply":"2023-11-25T13:11:13.205807Z","shell.execute_reply.started":"2023-11-25T13:10:57.503420Z"},"id":"ZHH60Lb4ZuoZ","outputId":"d8f39afe-fb06-4823-cc09-1ac40804fb0a","trusted":true},"outputs":[{"name":"stdout","output_type":"stream","text":["Collecting pyglet==1.4.11\n","  Downloading pyglet-1.4.11-py2.py3-none-any.whl (1.0 MB)\n","     ---------------------------------------- 0.0/1.0 MB ? eta -:--:--\n","     -- ------------------------------------- 0.1/1.0 MB 1.7 MB/s eta 0:00:01\n","     ------------------------------------ --- 1.0/1.0 MB 12.1 MB/s eta 0:00:01\n","     ---------------------------------------- 1.0/1.0 MB 11.0 MB/s eta 0:00:00\n","Collecting future (from pyglet==1.4.11)\n","  Downloading future-0.18.3.tar.gz (840 kB)\n","     ---------------------------------------- 0.0/840.9 kB ? eta -:--:--\n","     ------------------------------------- 840.9/840.9 kB 26.0 MB/s eta 0:00:00\n","  Installing build dependencies: started\n","  Installing build dependencies: finished with status 'done'\n","  Getting requirements to build wheel: started\n","  Getting requirements to build wheel: finished with status 'done'\n","  Preparing metadata (pyproject.toml): started\n","  Preparing metadata (pyproject.toml): finished with status 'done'\n","Building wheels for collected packages: future\n","  Building wheel for future (pyproject.toml): started\n","  Building wheel for future (pyproject.toml): finished with status 'done'\n","  Created wheel for future: filename=future-0.18.3-py3-none-any.whl size=492054 sha256=407ba002d9819a1b8a2ac6a42bbbad0a916f3bb4139a8e6d3f8734ecb54bfcb2\n","  Stored in directory: c:\\users\\intel\\appdata\\local\\packages\\pythonsoftwarefoundation.python.3.11_qbz5n2kfra8p0\\localcache\\local\\pip\\cache\\wheels\\da\\19\\ca\\9d8c44cd311a955509d7e13da3f0bea42400c469ef825b580b\n","Successfully built future\n","Installing collected packages: future, pyglet\n","Successfully installed future-0.18.3 pyglet-1.4.11\n","Note: you may need to restart the kernel to use updated packages.\n"]},{"name":"stderr","output_type":"stream","text":["\n","[notice] A new release of pip is available: 23.2.1 -> 23.3.1\n","[notice] To update, run: C:\\Users\\Intel\\AppData\\Local\\Microsoft\\WindowsApps\\PythonSoftwareFoundation.Python.3.11_qbz5n2kfra8p0\\python.exe -m pip install --upgrade pip\n"]}],"source":["pip install pyglet==1.4.11"]},{"cell_type":"code","execution_count":5,"metadata":{"execution":{"iopub.execute_input":"2023-11-25T13:11:13.210891Z","iopub.status.busy":"2023-11-25T13:11:13.210148Z","iopub.status.idle":"2023-11-25T13:11:48.779815Z","shell.execute_reply":"2023-11-25T13:11:48.778538Z","shell.execute_reply.started":"2023-11-25T13:11:13.210862Z"},"id":"jBnGVnpOgsta","outputId":"617ead7c-907f-404f-f098-ade42cf12aaf","trusted":true},"outputs":[{"name":"stdout","output_type":"stream","text":["Collecting PyOpenGLNote: you may need to restart the kernel to use updated packages.\n"]},{"name":"stderr","output_type":"stream","text":["\n","[notice] A new release of pip is available: 23.2.1 -> 23.3.1\n","[notice] To update, run: C:\\Users\\Intel\\AppData\\Local\\Microsoft\\WindowsApps\\PythonSoftwareFoundation.Python.3.11_qbz5n2kfra8p0\\python.exe -m pip install --upgrade pip\n"]},{"name":"stdout","output_type":"stream","text":["\n","  Obtaining dependency information for PyOpenGL from https://files.pythonhosted.org/packages/99/48/00e31747821d3fc56faddd00a4725454d1e694a8b67d715cf20f531506a5/PyOpenGL-3.1.7-py3-none-any.whl.metadata\n","  Downloading PyOpenGL-3.1.7-py3-none-any.whl.metadata (3.2 kB)\n","Collecting PyOpenGL_accelerate\n","  Obtaining dependency information for PyOpenGL_accelerate from https://files.pythonhosted.org/packages/c6/e3/e4c02ccbba4012e16ece90f0e1ccc8e2aa781edd95d6ab193e1d4a9dd63d/PyOpenGL_accelerate-3.1.7-cp311-cp311-win_amd64.whl.metadata\n","  Downloading PyOpenGL_accelerate-3.1.7-cp311-cp311-win_amd64.whl.metadata (984 bytes)\n","Downloading PyOpenGL-3.1.7-py3-none-any.whl (2.4 MB)\n","   ---------------------------------------- 0.0/2.4 MB ? eta -:--:--\n","    --------------------------------------- 0.0/2.4 MB 1.4 MB/s eta 0:00:02\n","   --------- ------------------------------ 0.6/2.4 MB 8.8 MB/s eta 0:00:01\n","   --------------------- ------------------ 1.3/2.4 MB 10.3 MB/s eta 0:00:01\n","   ---------------------------------------  2.4/2.4 MB 14.0 MB/s eta 0:00:01\n","   ---------------------------------------- 2.4/2.4 MB 12.9 MB/s eta 0:00:00\n","Downloading PyOpenGL_accelerate-3.1.7-cp311-cp311-win_amd64.whl (311 kB)\n","   ---------------------------------------- 0.0/311.9 kB ? eta -:--:--\n","   ---------------------------------------- 311.9/311.9 kB ? eta 0:00:00\n","Installing collected packages: PyOpenGL_accelerate, PyOpenGL\n","Successfully installed PyOpenGL-3.1.7 PyOpenGL_accelerate-3.1.7\n"]}],"source":["pip install PyOpenGL PyOpenGL_accelerate"]},{"cell_type":"code","execution_count":14,"metadata":{"execution":{"iopub.execute_input":"2023-11-25T13:11:48.781734Z","iopub.status.busy":"2023-11-25T13:11:48.781423Z","iopub.status.idle":"2023-11-25T13:13:11.171674Z","shell.execute_reply":"2023-11-25T13:13:11.170400Z","shell.execute_reply.started":"2023-11-25T13:11:48.781707Z"},"id":"G5uAVxZ2Cp90","outputId":"c586c01f-a7c9-4fc5-96b2-c600b6036795","trusted":true},"outputs":[{"ename":"SyntaxError","evalue":"invalid syntax (3886764517.py, line 5)","output_type":"error","traceback":["\u001b[1;36m  Cell \u001b[1;32mIn[14], line 5\u001b[1;36m\u001b[0m\n\u001b[1;33m    pip install tensordict torchrl\u001b[0m\n\u001b[1;37m        ^\u001b[0m\n\u001b[1;31mSyntaxError\u001b[0m\u001b[1;31m:\u001b[0m invalid syntax\n"]}],"source":["# !pip3 install -U torch  # If in the future, Colab update its environment, we may try specifically torch version 2.1.0+cu121\n","\n","#pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121\n","\n","#pip install tensordict torchrl\n"]},{"cell_type":"code","execution_count":15,"metadata":{"execution":{"iopub.execute_input":"2023-11-25T13:13:11.173606Z","iopub.status.busy":"2023-11-25T13:13:11.173282Z","iopub.status.idle":"2023-11-25T13:13:11.178411Z","shell.execute_reply":"2023-11-25T13:13:11.177487Z","shell.execute_reply.started":"2023-11-25T13:13:11.173576Z"},"id":"ntQfiLcHm3A-","trusted":true},"outputs":[],"source":["import os\n","# os.kill(os.getpid(), 9)"]},{"cell_type":"code","execution_count":17,"metadata":{"execution":{"iopub.execute_input":"2023-11-25T13:13:11.180111Z","iopub.status.busy":"2023-11-25T13:13:11.179725Z","iopub.status.idle":"2023-11-25T13:13:11.192300Z","shell.execute_reply":"2023-11-25T13:13:11.191555Z","shell.execute_reply.started":"2023-11-25T13:13:11.180063Z"},"id":"EMB7A6D9m63D","trusted":true},"outputs":[],"source":["# For tips on running notebooks in Google Colab, see\n","# https://pytorch.org/tutorials/beginner/colab\n","%matplotlib inline"]},{"cell_type":"code","execution_count":19,"metadata":{"execution":{"iopub.execute_input":"2023-11-25T13:13:11.193584Z","iopub.status.busy":"2023-11-25T13:13:11.193347Z","iopub.status.idle":"2023-11-25T13:13:12.254652Z","shell.execute_reply":"2023-11-25T13:13:12.253705Z","shell.execute_reply.started":"2023-11-25T13:13:11.193563Z"},"id":"yp0VuI_Nm88D","outputId":"f4cc8aac-8558-4901-bc63-42608cd8b5ac","trusted":true},"outputs":[],"source":["# Virtual display\n","#from pyvirtualdisplay import Display\n","\n","#virtual_display = Display(visible=0, size=(1400, 900))\n","#virtual_display.start()\n","# Instead, use a flag to control the rendering\n","render_game = False"]},{"cell_type":"markdown","metadata":{"id":"BdGG2-gPZuod"},"source":["\n","# Train a Mario-playing RL Agent\n","\n","**Credit : Original Authors:** [Yuansong Feng](https://github.com/YuansongFeng)_, [Suraj Subramanian](https://github.com/suraj813)_, [Howard Wang](https://github.com/hw26)_, [Steven Guo](https://github.com/GuoYuzhang)_.\n","\n","**Instruction, Exercises & Videos are added for AIC-502 students. Also Torch versioning is fixed so that Colab is now runable as of Oct 18, 2023**\n","\n","This tutorial walks you through the fundamentals of Deep Reinforcement\n","Learning. At the end, you will implement an AI-powered Mario (using\n","[Double Deep Q-Networks](https://arxiv.org/pdf/1509.06461.pdf)_) that\n","can play the game by itself.\n","\n","Although no prior knowledge of RL is necessary for this tutorial, you\n","can familiarize yourself with these RL\n","[concepts](https://spinningup.openai.com/en/latest/spinningup/rl_intro.html)_,\n","and have this handy\n","[cheatsheet](https://colab.research.google.com/drive/1eN33dPVtdPViiS1njTW_-r-IYCDTFU7N)_\n","as your companion. The full code is available\n","[here](https://github.com/yuansongFeng/MadMario/)_.\n","\n","You can see details of Mario mode & reward function [here](https://github.com/Kautenja/gym-super-mario-bros/tree/master).\n","\n",".. figure:: /_static/img/mario.gif\n","   :alt: mario\n"]},{"cell_type":"code","execution_count":null,"metadata":{"trusted":true},"outputs":[],"source":[]},{"cell_type":"code","execution_count":20,"metadata":{"execution":{"iopub.execute_input":"2023-11-25T13:13:12.256472Z","iopub.status.busy":"2023-11-25T13:13:12.256204Z","iopub.status.idle":"2023-11-25T13:13:32.670189Z","shell.execute_reply":"2023-11-25T13:13:32.669284Z","shell.execute_reply.started":"2023-11-25T13:13:12.256449Z"},"id":"hXBPLmQzZuof","outputId":"527545ff-2efe-4d4b-dd1f-f5b90bde29b0","trusted":true},"outputs":[{"name":"stdout","output_type":"stream","text":["Couldn't find program: 'bash'\n"]}],"source":["%%bash\n","pip install -q gym-super-mario-bros==7.4.0"]},{"cell_type":"code","execution_count":21,"metadata":{"execution":{"iopub.execute_input":"2023-11-25T13:13:32.674058Z","iopub.status.busy":"2023-11-25T13:13:32.673777Z","iopub.status.idle":"2023-11-25T13:13:34.396554Z","shell.execute_reply":"2023-11-25T13:13:34.395576Z","shell.execute_reply.started":"2023-11-25T13:13:32.674034Z"},"id":"Id3Ri8CWDZOK","outputId":"16076570-b76b-465f-d6dc-1300bd7e24b6","trusted":true},"outputs":[{"name":"stdout","output_type":"stream","text":["2.1.1+cu121\n"]}],"source":["import torch\n","print(torch.__version__) # LOG: runnable with version 2.1.0+cu121"]},{"cell_type":"code","execution_count":25,"metadata":{"execution":{"iopub.execute_input":"2023-11-25T13:13:34.398433Z","iopub.status.busy":"2023-11-25T13:13:34.397877Z","iopub.status.idle":"2023-11-25T13:13:39.102945Z","shell.execute_reply":"2023-11-25T13:13:39.102121Z","shell.execute_reply.started":"2023-11-25T13:13:34.398398Z"},"id":"uMs-_9-GZuog","trusted":true},"outputs":[{"ename":"ModuleNotFoundError","evalue":"No module named 'nes_py'","output_type":"error","traceback":["\u001b[1;31m---------------------------------------------------------------------------\u001b[0m","\u001b[1;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)","\u001b[1;32md:\\cmkl\\AIC-502\\mario-rl\\aic502-mario-rl.ipynb Cell 14\u001b[0m line \u001b[0;36m1\n\u001b[0;32m     <a href='vscode-notebook-cell:/d%3A/cmkl/AIC-502/mario-rl/aic502-mario-rl.ipynb#Y142sZmlsZQ%3D%3D?line=11'>12</a>\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mgym\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mwrappers\u001b[39;00m \u001b[39mimport\u001b[39;00m FrameStack\n\u001b[0;32m     <a href='vscode-notebook-cell:/d%3A/cmkl/AIC-502/mario-rl/aic502-mario-rl.ipynb#Y142sZmlsZQ%3D%3D?line=13'>14</a>\u001b[0m \u001b[39m# NES Emulator for OpenAI Gym\u001b[39;00m\n\u001b[1;32m---> <a href='vscode-notebook-cell:/d%3A/cmkl/AIC-502/mario-rl/aic502-mario-rl.ipynb#Y142sZmlsZQ%3D%3D?line=14'>15</a>\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mnes_py\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mwrappers\u001b[39;00m \u001b[39mimport\u001b[39;00m JoypadSpace\n\u001b[0;32m     <a href='vscode-notebook-cell:/d%3A/cmkl/AIC-502/mario-rl/aic502-mario-rl.ipynb#Y142sZmlsZQ%3D%3D?line=16'>17</a>\u001b[0m \u001b[39m# Super Mario environment for OpenAI Gym\u001b[39;00m\n\u001b[0;32m     <a href='vscode-notebook-cell:/d%3A/cmkl/AIC-502/mario-rl/aic502-mario-rl.ipynb#Y142sZmlsZQ%3D%3D?line=17'>18</a>\u001b[0m \u001b[39mimport\u001b[39;00m \u001b[39mgym_super_mario_bros\u001b[39;00m\n","\u001b[1;31mModuleNotFoundError\u001b[0m: No module named 'nes_py'"]}],"source":["\n","from torch import nn\n","from torchvision import transforms as T\n","from PIL import Image\n","import numpy as np\n","from pathlib import Path\n","from collections import deque\n","import random, datetime, os, copy\n","\n","# Gym is an OpenAI toolkit for RL\n","import gym\n","from gym.spaces import Box\n","from gym.wrappers import FrameStack\n","\n","# NES Emulator for OpenAI Gym\n","from nes_py.wrappers import JoypadSpace\n","\n","# Super Mario environment for OpenAI Gym\n","import gym_super_mario_bros\n","\n","from tensordict import TensorDict\n","from torchrl.data import TensorDictReplayBuffer, LazyMemmapStorage"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2023-11-25T13:13:39.104576Z","iopub.status.busy":"2023-11-25T13:13:39.104117Z","iopub.status.idle":"2023-11-25T13:13:39.109590Z","shell.execute_reply":"2023-11-25T13:13:39.108123Z","shell.execute_reply.started":"2023-11-25T13:13:39.104547Z"},"id":"1M4LVnwVoCM_","outputId":"7c87d492-8b08-4540-c667-76d3c2431ad3","trusted":true},"outputs":[],"source":["import matplotlib.pyplot as plt\n","from PIL import Image"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"VYEj9rbdndAB"},"outputs":[],"source":[]},{"cell_type":"markdown","metadata":{"id":"dVgiNdYdZuog"},"source":["## RL Definitions\n","\n","**Environment** The world that an agent interacts with and learns from.\n","\n","**Action** $a$ : How the Agent responds to the Environment. The\n","set of all possible Actions is called *action-space*.\n","\n","**State** $s$ : The current characteristic of the Environment. The\n","set of all possible States the Environment can be in is called\n","*state-space*.\n","\n","**Reward** $r$ : Reward is the key feedback from Environment to\n","Agent. It is what drives the Agent to learn and to change its future\n","action. An aggregation of rewards over multiple time steps is called\n","**Return**.\n","\n","**Optimal Action-Value function** $Q^*(s,a)$ : Gives the expected\n","return if you start in state $s$, take an arbitrary action\n","$a$, and then for each future time step take the action that\n","maximizes returns. $Q$ can be said to stand for the “quality” of\n","the action in a state. We try to approximate this function.\n","\n","\n"]},{"cell_type":"markdown","metadata":{"id":"X1lcK6xpZuoh"},"source":["## Environment\n","\n","### Initialize Environment\n","\n","In Mario, the environment consists of tubes, mushrooms and other\n","components.\n","\n","When Mario makes an action, the environment responds with the changed\n","(next) state, reward and other info.\n","\n","\n"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2023-11-25T13:13:39.111644Z","iopub.status.busy":"2023-11-25T13:13:39.110797Z","iopub.status.idle":"2023-11-25T13:13:39.125794Z","shell.execute_reply":"2023-11-25T13:13:39.125010Z","shell.execute_reply.started":"2023-11-25T13:13:39.111607Z"},"id":"U6GK-RyL91Lg","outputId":"c468fa3b-9763-4d28-fbf8-be3a27c84700","trusted":true},"outputs":[],"source":["print(gym.__version__)"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2023-11-25T13:13:39.127459Z","iopub.status.busy":"2023-11-25T13:13:39.127069Z","iopub.status.idle":"2023-11-25T13:13:39.579735Z","shell.execute_reply":"2023-11-25T13:13:39.578718Z","shell.execute_reply.started":"2023-11-25T13:13:39.127425Z"},"id":"ri6jOM8TZuoi","outputId":"e3583d26-26a6-4679-bc18-d4538fa9ca73","trusted":true},"outputs":[],"source":["# Initialize Super Mario environment (in v0.26 change render mode to 'human' to see results on the screen)\n","if gym.__version__ < '0.26':\n","    env = gym_super_mario_bros.make(\"SuperMarioBros-1-1-v0\", new_step_api=True)\n","    print(1)\n","else:\n","    env = gym_super_mario_bros.make(\"SuperMarioBros-1-1-v0\", render_mode='rgb_array', apply_api_compatibility=True)\n","    print(2)\n","\n","# Limit the action-space to\n","#   0. walk right\n","#   1. jump right\n","env = JoypadSpace(env, [[\"right\"], [\"right\", \"A\"]])\n","\n","env.reset()\n","next_state, reward, done, trunc, info = env.step(action=0)\n","print(f\"{next_state.shape},\\n {reward},\\n {done},\\n {info}\")"]},{"cell_type":"markdown","metadata":{"id":"na2a0jVMZuoj"},"source":["### Preprocess Environment\n","\n","Environment data is returned to the agent in ``next_state``. As you saw\n","above, each state is represented by a ``[3, 240, 256]`` size array.\n","Often that is more information than our agent needs; for instance,\n","Mario’s actions do not depend on the color of the pipes or the sky!\n","\n","We use **Wrappers** to preprocess environment data before sending it to\n","the agent.\n","\n","``GrayScaleObservation`` is a common wrapper to transform an RGB image\n","to grayscale; doing so reduces the size of the state representation\n","without losing useful information. Now the size of each state:\n","``[1, 240, 256]``\n","\n","``ResizeObservation`` downsamples each observation into a square image.\n","New size: ``[1, 84, 84]``\n","\n","``SkipFrame`` is a custom wrapper that inherits from ``gym.Wrapper`` and\n","implements the ``step()`` function. Because consecutive frames don’t\n","vary much, we can skip n-intermediate frames without losing much\n","information. The n-th frame aggregates rewards accumulated over each\n","skipped frame.\n","\n","``FrameStack`` is a wrapper that allows us to squash consecutive frames\n","of the environment into a single observation point to feed to our\n","learning model. This way, we can identify if Mario was landing or\n","jumping based on the direction of his movement in the previous several\n","frames.\n","\n","\n"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2023-11-25T13:13:39.581396Z","iopub.status.busy":"2023-11-25T13:13:39.581012Z","iopub.status.idle":"2023-11-25T13:13:39.600176Z","shell.execute_reply":"2023-11-25T13:13:39.599344Z","shell.execute_reply.started":"2023-11-25T13:13:39.581366Z"},"id":"yioJ53_0Zuok","trusted":true},"outputs":[],"source":["class SkipFrame(gym.Wrapper):\n","    def __init__(self, env, skip):\n","        \"\"\"Return only every `skip`-th frame\"\"\"\n","        super().__init__(env)\n","        self._skip = skip\n","\n","    def step(self, action):\n","        \"\"\"Repeat action, and sum reward\"\"\"\n","        total_reward = 0.0\n","        for i in range(self._skip):\n","            # Accumulate reward and repeat the same action\n","            obs, reward, done, trunk, info = self.env.step(action)\n","            total_reward += reward\n","            if done:\n","                break\n","        return obs, total_reward, done, trunk, info\n","\n","\n","class GrayScaleObservation(gym.ObservationWrapper):\n","    def __init__(self, env):\n","        super().__init__(env)\n","        obs_shape = self.observation_space.shape[:2]\n","        self.observation_space = Box(low=0, high=255, shape=obs_shape, dtype=np.uint8)\n","\n","    def permute_orientation(self, observation):\n","        # permute [H, W, C] array to [C, H, W] tensor\n","        observation = np.transpose(observation, (2, 0, 1))\n","        observation = torch.tensor(observation.copy(), dtype=torch.float)\n","        return observation\n","\n","    def observation(self, observation):\n","        observation = self.permute_orientation(observation)\n","        transform = T.Grayscale()\n","        observation = transform(observation)\n","        return observation\n","\n","\n","class ResizeObservation(gym.ObservationWrapper):\n","    def __init__(self, env, shape):\n","        super().__init__(env)\n","        if isinstance(shape, int):\n","            self.shape = (shape, shape)\n","        else:\n","            self.shape = tuple(shape)\n","\n","        obs_shape = self.shape + self.observation_space.shape[2:]\n","        self.observation_space = Box(low=0, high=255, shape=obs_shape, dtype=np.uint8)\n","\n","    def observation(self, observation):\n","        transforms = T.Compose(\n","            [T.Resize(self.shape), T.Normalize(0, 255)]\n","        )\n","        observation = transforms(observation).squeeze(0)\n","        return observation\n","\n","\n","# Apply Wrappers to environment\n","env = SkipFrame(env, skip=4)\n","env = GrayScaleObservation(env)\n","env = ResizeObservation(env, shape=84)\n","if gym.__version__ < '0.26':\n","    env = FrameStack(env, num_stack=4, new_step_api=True)\n","else:\n","    env = FrameStack(env, num_stack=4)"]},{"cell_type":"markdown","metadata":{"id":"5-gQH1_0Zuol"},"source":["After applying the above wrappers to the environment, the final wrapped\n","state consists of 4 gray-scaled consecutive frames stacked together, as\n","shown above in the image on the left. Each time Mario makes an action,\n","the environment responds with a state of this structure. The structure\n","is represented by a 3-D array of size ``[4, 84, 84]``.\n","\n",".. figure:: /_static/img/mario_env.png\n","   :alt: picture\n","\n","\n","\n"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2023-11-25T13:13:39.601524Z","iopub.status.busy":"2023-11-25T13:13:39.601229Z","iopub.status.idle":"2023-11-25T13:13:39.955683Z","shell.execute_reply":"2023-11-25T13:13:39.954789Z","shell.execute_reply.started":"2023-11-25T13:13:39.601497Z"},"id":"mkAhrlSVoSbI","outputId":"2204568a-b695-4136-d222-4e1532ab7bb4","trusted":true},"outputs":[],"source":["\n","state = env.reset()\n","\n","# if gym.__version__ > '0.26': render got no rgb_array arg\n","# img = env.render(\"rgb_array\").copy() # IMPORTANT: render keep modifying the same object\n","img = env.render().copy() # IMPORTANT: render keep modifying the same object\n","\n","print(img.shape)\n","plt.imshow(img)"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2023-11-25T13:13:39.957244Z","iopub.status.busy":"2023-11-25T13:13:39.956933Z","iopub.status.idle":"2023-11-25T13:13:39.962248Z","shell.execute_reply":"2023-11-25T13:13:39.961312Z","shell.execute_reply.started":"2023-11-25T13:13:39.957217Z"},"trusted":true},"outputs":[],"source":["def first_if_tuple(x):\n","    return x[0] if isinstance(x, tuple) else x"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2023-11-25T13:13:39.963757Z","iopub.status.busy":"2023-11-25T13:13:39.963457Z","iopub.status.idle":"2023-11-25T13:13:39.974720Z","shell.execute_reply":"2023-11-25T13:13:39.973649Z","shell.execute_reply.started":"2023-11-25T13:13:39.963724Z"},"id":"CUM0eF7ipXDV","outputId":"a24c7e34-16be-48c5-fdf3-90967ba4c3fb","trusted":true},"outputs":[],"source":["print(env.action_space)\n","print(env.action_space.n)\n","\n","print(first_if_tuple(state).shape)\n","print(env.action_space.n)"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2023-11-25T13:13:39.976242Z","iopub.status.busy":"2023-11-25T13:13:39.975937Z","iopub.status.idle":"2023-11-25T13:13:40.250341Z","shell.execute_reply":"2023-11-25T13:13:40.249394Z","shell.execute_reply.started":"2023-11-25T13:13:39.976217Z"},"id":"2rxdCQ01snwD","outputId":"a6b8bf42-13dc-4c1d-bc8c-764f34517a7c","trusted":true},"outputs":[],"source":["state = env.reset()\n","next_state, reward, done, trunc, info = env.step(action=1)\n","next_state, reward, done, trunc, info = env.step(action=1)\n","next_state, reward, done, trunc, info = env.step(action=1)\n","print(f\"{first_if_tuple(next_state).shape},\\n {reward},\\n {done},\\n {info}\")\n","img2 = env.render()\n","print(img2.shape)\n","plt.imshow(img2)"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2023-11-25T13:13:40.252054Z","iopub.status.busy":"2023-11-25T13:13:40.251694Z","iopub.status.idle":"2023-11-25T13:13:40.259142Z","shell.execute_reply":"2023-11-25T13:13:40.258200Z","shell.execute_reply.started":"2023-11-25T13:13:40.252021Z"},"id":"Yt-yN5AnDIp3","outputId":"034e43c5-5c78-492c-8209-8a73997cb686","trusted":true},"outputs":[],"source":["np.any(img != img2)"]},{"cell_type":"markdown","metadata":{"id":"pLD18fdAZuol"},"source":["## Agent\n","\n","We create a class ``Mario`` to represent our agent in the game. Mario\n","should be able to:\n","\n","-  **Act** according to the optimal action policy based on the current\n","   state (of the environment).\n","\n","-  **Remember** experiences. Experience = (current state, current\n","   action, reward, next state). Mario *caches* and later *recalls* his\n","   experiences to update his action policy.\n","\n","-  **Learn** a better action policy over time\n","\n","\n"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2023-11-25T13:13:40.260629Z","iopub.status.busy":"2023-11-25T13:13:40.260348Z","iopub.status.idle":"2023-11-25T13:13:40.270488Z","shell.execute_reply":"2023-11-25T13:13:40.269564Z","shell.execute_reply.started":"2023-11-25T13:13:40.260605Z"},"id":"JN8bIfvRZuol","trusted":true},"outputs":[],"source":["class Mario:\n","    def __init__():\n","        pass\n","\n","    def act(self, state):\n","        \"\"\"Given a state, choose an epsilon-greedy action\"\"\"\n","        pass\n","\n","    def cache(self, experience):\n","        \"\"\"Add the experience to memory\"\"\"\n","        pass\n","\n","    def recall(self):\n","        \"\"\"Sample experiences from memory\"\"\"\n","        pass\n","\n","    def learn(self):\n","        \"\"\"Update online action value (Q) function with a batch of experiences\"\"\"\n","        pass"]},{"cell_type":"markdown","metadata":{"id":"agaRElpzZuom"},"source":["In the following sections, we will populate Mario’s parameters and\n","define his functions.\n","\n","\n"]},{"cell_type":"markdown","metadata":{"id":"awTm_yL0Zuom"},"source":["### Act\n","\n","For any given state, an agent can choose to do the most optimal action\n","(**exploit**) or a random action (**explore**).\n","\n","Mario randomly explores with a chance of ``self.exploration_rate``; when\n","he chooses to exploit, he relies on ``MarioNet`` (implemented in\n","``Learn`` section) to provide the most optimal action.\n","\n","\n"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2023-11-25T13:13:40.272324Z","iopub.status.busy":"2023-11-25T13:13:40.271738Z","iopub.status.idle":"2023-11-25T13:13:40.287286Z","shell.execute_reply":"2023-11-25T13:13:40.286486Z","shell.execute_reply.started":"2023-11-25T13:13:40.272292Z"},"id":"8x1EA5XZZuom","trusted":true},"outputs":[],"source":["class Mario:\n","    def __init__(self, state_dim, action_dim, save_dir):\n","        self.state_dim = state_dim\n","        self.action_dim = action_dim\n","        self.save_dir = save_dir\n","\n","        self.device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n","\n","        # Mario's DNN to predict the most optimal action - we implement this in the Learn section\n","        self.net = MarioNet(self.state_dim, self.action_dim).float()\n","        self.net = self.net.to(device=self.device)\n","\n","        self.exploration_rate = 1\n","        self.exploration_rate_decay = 0.99999975\n","        self.exploration_rate_min = 0.1\n","        self.curr_step = 0\n","\n","        self.save_every = 25000  # no. of steps of experiences between saving Mario Net\n","\n","    def act(self, state):\n","        \"\"\"\n","    Given a state, choose an epsilon-greedy action and update value of step.\n","\n","    Inputs:\n","    state(``LazyFrame``): A single observation of the current state, dimension is (state_dim)\n","    Outputs:\n","    ``action_idx`` (``int``): An integer representing which action Mario will perform\n","    \"\"\"\n","        # EXPLORE\n","        if np.random.rand() < self.exploration_rate:\n","            action_idx = np.random.randint(self.action_dim)\n","\n","        # EXPLOIT\n","        else:\n","            state = state[0].__array__() if isinstance(state, tuple) else state.__array__()\n","            state = torch.tensor(state, device=self.device).unsqueeze(0)\n","            action_values = self.net(state, model=\"online\")\n","            action_idx = torch.argmax(action_values, axis=1).item()\n","\n","        # decrease exploration_rate\n","        self.exploration_rate *= self.exploration_rate_decay\n","        self.exploration_rate = max(self.exploration_rate_min, self.exploration_rate)\n","\n","        # increment step\n","        self.curr_step += 1\n","        return action_idx"]},{"cell_type":"markdown","metadata":{"id":"Gf9nyEjpZuom"},"source":["### Cache and Recall\n","\n","These two functions serve as Mario’s “memory” process.\n","\n","``cache()``: Each time Mario performs an action, he stores the\n","``experience`` to his memory. His experience includes the current\n","*state*, *action* performed, *reward* from the action, the *next state*,\n","and whether the game is *done*.\n","\n","``recall()``: Mario randomly samples a batch of experiences from his\n","memory, and uses that to learn the game.\n","\n","Learn more about ReplayBuffer\n","[HERE](https://colab.research.google.com/github/pytorch/rl/blob/gh-pages/_downloads/c64b1e484ba3e0219549719cc3c37479/rb_tutorial.ipynb)\n"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2023-11-25T13:13:40.288709Z","iopub.status.busy":"2023-11-25T13:13:40.288377Z","iopub.status.idle":"2023-11-25T13:13:40.305145Z","shell.execute_reply":"2023-11-25T13:13:40.304262Z","shell.execute_reply.started":"2023-11-25T13:13:40.288683Z"},"id":"CvpADBsKZuon","trusted":true},"outputs":[],"source":["class Mario(Mario):  # subclassing for continuity\n","    def __init__(self, state_dim, action_dim, save_dir):\n","        super().__init__(state_dim, action_dim, save_dir)\n","        # self.memory = deque(maxlen=100000) #TensorDictReplayBuffer(storage=LazyMemmapStorage(100000, device=torch.device(\"cpu\")))\n","        self.memory = TensorDictReplayBuffer(storage=LazyMemmapStorage(100000, device=torch.device(\"cpu\")))\n","        self.batch_size = 32\n","\n","    def cache(self, state, next_state, action, reward, done):\n","        \"\"\"\n","        Store the experience to self.memory (replay buffer)\n","\n","        Inputs:\n","        state (``LazyFrame``),\n","        next_state (``LazyFrame``),\n","        action (``int``),\n","        reward (``float``),\n","        done(``bool``))\n","        \"\"\"\n","        \n","        state = first_if_tuple(state).__array__()\n","        next_state = first_if_tuple(next_state).__array__()\n","\n","        state = torch.tensor(state)\n","        next_state = torch.tensor(next_state)\n","        action = torch.tensor([action])\n","        reward = torch.tensor([reward])\n","        done = torch.tensor([done])\n","\n","        # self.memory.append((state, next_state, action, reward, done,))\n","        self.memory.add(TensorDict({\"state\": state, \"next_state\": next_state, \"action\": action, \"reward\": reward, \"done\": done}, batch_size=[]))\n","\n","    def recall(self):\n","        \"\"\"\n","        Retrieve a batch of experiences from memory\n","        \"\"\"\n","\n","        # batch = random.sample(self.memory, self.batch_size)\n","        # state, next_state, action, reward, done = map(torch.stack, zip(*batch))\n","\n","        batch = self.memory.sample(self.batch_size).to(self.device)\n","        state, next_state, action, reward, done = (batch.get(key) for key in (\"state\", \"next_state\", \"action\", \"reward\", \"done\"))\n","        return state, next_state, action.squeeze(), reward.squeeze(), done.squeeze()"]},{"cell_type":"markdown","metadata":{"id":"FCq3sJOZZuon"},"source":["### Learn\n","\n","Mario uses the [DDQN algorithm](https://arxiv.org/pdf/1509.06461)_\n","under the hood. DDQN uses two ConvNets - $Q_{online}$ and\n","$Q_{target}$ - that independently approximate the optimal\n","action-value function.\n","\n","In our implementation, we share feature generator ``features`` across\n","$Q_{online}$ and $Q_{target}$, but maintain separate FC\n","classifiers for each. $\\theta_{target}$ (the parameters of\n","$Q_{target}$) is frozen to prevent updating by backprop. Instead,\n","it is periodically synced with $\\theta_{online}$ (more on this\n","later).\n","\n","#### Neural Network\n","\n"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2023-11-25T13:13:40.306552Z","iopub.status.busy":"2023-11-25T13:13:40.306235Z","iopub.status.idle":"2023-11-25T13:13:40.321717Z","shell.execute_reply":"2023-11-25T13:13:40.320985Z","shell.execute_reply.started":"2023-11-25T13:13:40.306527Z"},"id":"oN_IqoYfZuon","trusted":true},"outputs":[],"source":["class MarioNet(nn.Module):\n","    \"\"\"mini CNN structure\n","  input -> (conv2d + relu) x 3 -> flatten -> (dense + relu) x 2 -> output\n","  \"\"\"\n","\n","    def __init__(self, input_dim, output_dim):\n","        super().__init__()\n","        c, h, w = input_dim\n","\n","        if h != 84:\n","            raise ValueError(f\"Expecting input height: 84, got: {h}\")\n","        if w != 84:\n","            raise ValueError(f\"Expecting input width: 84, got: {w}\")\n","\n","        self.online = nn.Sequential(\n","            nn.Conv2d(in_channels=c, out_channels=32, kernel_size=8, stride=4),\n","            nn.ReLU(),\n","            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=4, stride=2),\n","            nn.ReLU(),\n","            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1),\n","            nn.ReLU(),\n","            nn.Flatten(),\n","            nn.Linear(3136, 512),\n","            nn.ReLU(),\n","            nn.Linear(512, output_dim),\n","        )\n","\n","        self.target = copy.deepcopy(self.online)\n","\n","        # Q_target parameters are frozen.\n","        for p in self.target.parameters():\n","            p.requires_grad = False\n","\n","    def forward(self, input, model):\n","        if model == \"online\":\n","            return self.online(input)\n","        elif model == \"target\":\n","            return self.target(input)"]},{"cell_type":"markdown","metadata":{"id":"Ljg4ifQVZuon"},"source":["#### TD Estimate & TD Target\n","\n","Two values are involved in learning:\n","\n","**TD Estimate** - the predicted optimal $Q^*$ for a given state\n","$s$\n","\n","\\begin{align}{TD}_e = Q_{online}^*(s,a)\\end{align}\n","\n","**TD Target** - aggregation of current reward and the estimated\n","$Q^*$ in the next state $s'$\n","\n","\\begin{align}a' = argmax_{a} Q_{online}(s', a)\\end{align}\n","\n","\\begin{align}{TD}_t = r + \\gamma Q_{target}^*(s',a')\\end{align}\n","\n","Because we don’t know what next action $a'$ will be, we use the\n","action $a'$ maximizes $Q_{online}$ in the next state\n","$s'$.\n","\n","Notice we use the\n","[@torch.no_grad()](https://pytorch.org/docs/stable/generated/torch.no_grad.html#no-grad)_\n","decorator on ``td_target()`` to disable gradient calculations here\n","(because we don’t need to backpropagate on $\\theta_{target}$).\n","\n","\n"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2023-11-25T13:13:40.323053Z","iopub.status.busy":"2023-11-25T13:13:40.322781Z","iopub.status.idle":"2023-11-25T13:13:40.338257Z","shell.execute_reply":"2023-11-25T13:13:40.337497Z","shell.execute_reply.started":"2023-11-25T13:13:40.323030Z"},"id":"fevvRaCOZuoo","trusted":true},"outputs":[],"source":["class Mario(Mario):\n","    def __init__(self, state_dim, action_dim, save_dir):\n","        super().__init__(state_dim, action_dim, save_dir)\n","        self.gamma = 0.9\n","\n","    def td_estimate(self, state, action):\n","        current_Q = self.net(state, model=\"online\")[\n","            np.arange(0, self.batch_size), action\n","        ]  # Q_online(s,a)\n","        return current_Q\n","\n","    @torch.no_grad()\n","    def td_target(self, reward, next_state, done):\n","        next_state_Q = self.net(next_state, model=\"online\")\n","        best_action = torch.argmax(next_state_Q, axis=1)\n","        next_Q = self.net(next_state, model=\"target\")[\n","            np.arange(0, self.batch_size), best_action\n","        ]\n","        return (reward + (1 - done.float()) * self.gamma * next_Q).float()"]},{"cell_type":"markdown","metadata":{"id":"vV8EiRtQZuoo"},"source":["#### Updating the model\n","\n","As Mario samples inputs from his replay buffer, we compute $TD_t$\n","and $TD_e$ and backpropagate this loss down $Q_{online}$ to\n","update its parameters $\\theta_{online}$ ($\\alpha$ is the\n","learning rate ``lr`` passed to the ``optimizer``)\n","\n","\\begin{align}\\theta_{online} \\leftarrow \\theta_{online} + \\alpha \\nabla(TD_e - TD_t)\\end{align}\n","\n","$\\theta_{target}$ does not update through backpropagation.\n","Instead, we periodically copy $\\theta_{online}$ to\n","$\\theta_{target}$\n","\n","\\begin{align}\\theta_{target} \\leftarrow \\theta_{online}\\end{align}\n","\n","\n","\n"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2023-11-25T13:13:40.339557Z","iopub.status.busy":"2023-11-25T13:13:40.339287Z","iopub.status.idle":"2023-11-25T13:13:40.350646Z","shell.execute_reply":"2023-11-25T13:13:40.349796Z","shell.execute_reply.started":"2023-11-25T13:13:40.339535Z"},"id":"1R4PfyuMZuoo","trusted":true},"outputs":[],"source":["class Mario(Mario):\n","    def __init__(self, state_dim, action_dim, save_dir):\n","        super().__init__(state_dim, action_dim, save_dir)\n","        self.optimizer = torch.optim.Adam(self.net.parameters(), lr=0.00025)\n","        self.loss_fn = torch.nn.SmoothL1Loss()\n","\n","    def update_Q_online(self, td_estimate, td_target):\n","        loss = self.loss_fn(td_estimate, td_target)\n","        self.optimizer.zero_grad()\n","        loss.backward()\n","        self.optimizer.step()\n","        return loss.item()\n","\n","    def sync_Q_target(self):\n","        self.net.target.load_state_dict(self.net.online.state_dict())"]},{"cell_type":"markdown","metadata":{"id":"28EO6cWBZuoo"},"source":["#### Save checkpoint\n","\n","\n"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2023-11-25T13:13:40.352045Z","iopub.status.busy":"2023-11-25T13:13:40.351502Z","iopub.status.idle":"2023-11-25T13:13:40.369689Z","shell.execute_reply":"2023-11-25T13:13:40.368763Z","shell.execute_reply.started":"2023-11-25T13:13:40.352010Z"},"id":"m6SIFwlDZuoo","trusted":true},"outputs":[],"source":["class Mario(Mario):\n","    def save(self, save_path=None):\n","        if save_path is None:\n","            try:\n","                save_path = SAVE_PATH\n","            except:\n","                save_path = (\n","                    self.save_dir / f\"mario_net_{int(self.curr_step // self.save_every)}.chkpt\"\n","                )\n","        \n","        torch.save(\n","            dict(model=self.net.state_dict(),\n","                 optimizer_state_dict=self.optimizer.state_dict(),\n","                 exploration_rate=self.exploration_rate,\n","                 memory=self.memory\n","                ),\n","            save_path,\n","        )\n","        print(f\"MarioNet saved to {save_path} at step {self.curr_step}\")\n","\n","    def load(self, load_path):\n","        # if not load_path.exists():\n","        #     raise ValueError(f\"{load_path} does not exist\")\n","\n","        ckp = torch.load(load_path, map_location=('cuda' if self.use_cuda else 'cpu'))\n","        exploration_rate = ckp.get('exploration_rate')\n","        state_dict = ckp.get('model')\n","\n","        print(f\"Loading model at {load_path} with exploration rate {exploration_rate}\")\n","        self.net.load_state_dict(state_dict)\n","        self.exploration_rate = exploration_rate\n","        self.optimizer.load_state_dict(ckp.get('optimizer_state_dict'))\n","#         self.memory = ckp.get('memory')"]},{"cell_type":"markdown","metadata":{"id":"WjczpJVVZuoo"},"source":["#### Putting it all together\n","\n","\n"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2023-11-25T13:13:40.374953Z","iopub.status.busy":"2023-11-25T13:13:40.374658Z","iopub.status.idle":"2023-11-25T13:13:40.383979Z","shell.execute_reply":"2023-11-25T13:13:40.383099Z","shell.execute_reply.started":"2023-11-25T13:13:40.374930Z"},"id":"6JVB0ClpZuop","trusted":true},"outputs":[],"source":["class Mario(Mario):\n","    def __init__(self, state_dim, action_dim, save_dir, load_path=None):\n","        super().__init__(state_dim, action_dim, save_dir)\n","        self.burnin = 1e4  # min. experiences before training\n","        self.learn_every = 3  # no. of experiences between updates to Q_online\n","        self.sync_every = 1e4  # no. of experiences between Q_target & Q_online sync\n","\n","        if load_path:\n","            self.load(load_path)\n","    def learn(self):\n","        if self.curr_step % self.sync_every == 0:\n","            self.sync_Q_target()\n","\n","        if self.curr_step % self.save_every == 0:\n","            self.save()\n","\n","        if self.curr_step < self.burnin:\n","            return None, None\n","\n","        if self.curr_step % self.learn_every != 0:\n","            return None, None\n","\n","        # Sample from memory\n","        state, next_state, action, reward, done = self.recall()\n","\n","        # Get TD Estimate\n","        td_est = self.td_estimate(state, action)\n","\n","        # Get TD Target\n","        td_tgt = self.td_target(reward, next_state, done)\n","\n","        # Backpropagate loss through Q_online\n","        loss = self.update_Q_online(td_est, td_tgt)\n","\n","        return (td_est.mean().item(), loss)"]},{"cell_type":"markdown","metadata":{"id":"WPUWJ0xDZuop"},"source":["### Logging\n","\n","\n"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2023-11-25T13:13:40.385801Z","iopub.status.busy":"2023-11-25T13:13:40.385327Z","iopub.status.idle":"2023-11-25T13:13:40.405250Z","shell.execute_reply":"2023-11-25T13:13:40.404498Z","shell.execute_reply.started":"2023-11-25T13:13:40.385769Z"},"id":"joXWPz_RZuop","trusted":true},"outputs":[],"source":["import numpy as np\n","import time, datetime\n","import matplotlib.pyplot as plt\n","\n","\n","class MetricLogger:\n","    def __init__(self, save_dir):\n","        self.save_log = save_dir / \"log\"\n","        with open(self.save_log, \"w\") as f:\n","            f.write(\n","                f\"{'Episode':>8}{'Step':>8}{'Epsilon':>10}{'MeanReward':>15}\"\n","                f\"{'MeanLength':>15}{'MeanLoss':>15}{'MeanQValue':>15}\"\n","                f\"{'TimeDelta':>15}{'Time':>20}\\n\"\n","            )\n","        self.ep_rewards_plot = save_dir / \"reward_plot.jpg\"\n","        self.ep_lengths_plot = save_dir / \"length_plot.jpg\"\n","        self.ep_avg_losses_plot = save_dir / \"loss_plot.jpg\"\n","        self.ep_avg_qs_plot = save_dir / \"q_plot.jpg\"\n","\n","        # History metrics\n","        self.ep_rewards = []\n","        self.ep_lengths = []\n","        self.ep_avg_losses = []\n","        self.ep_avg_qs = []\n","\n","        # Moving averages, added for every call to record()\n","        self.moving_avg_ep_rewards = []\n","        self.moving_avg_ep_lengths = []\n","        self.moving_avg_ep_avg_losses = []\n","        self.moving_avg_ep_avg_qs = []\n","\n","        # Current episode metric\n","        self.init_episode()\n","\n","        # Timing\n","        self.record_time = time.time()\n","\n","    def log_step(self, reward, loss, q):\n","        self.curr_ep_reward += reward\n","        self.curr_ep_length += 1\n","        if loss:\n","            self.curr_ep_loss += loss\n","            self.curr_ep_q += q\n","            self.curr_ep_loss_length += 1\n","\n","    def log_episode(self):\n","        \"Mark end of episode\"\n","        self.ep_rewards.append(self.curr_ep_reward)\n","        self.ep_lengths.append(self.curr_ep_length)\n","        if self.curr_ep_loss_length == 0:\n","            ep_avg_loss = 0\n","            ep_avg_q = 0\n","        else:\n","            ep_avg_loss = np.round(self.curr_ep_loss / self.curr_ep_loss_length, 5)\n","            ep_avg_q = np.round(self.curr_ep_q / self.curr_ep_loss_length, 5)\n","        self.ep_avg_losses.append(ep_avg_loss)\n","        self.ep_avg_qs.append(ep_avg_q)\n","\n","        self.init_episode()\n","\n","    def init_episode(self):\n","        self.curr_ep_reward = 0.0\n","        self.curr_ep_length = 0\n","        self.curr_ep_loss = 0.0\n","        self.curr_ep_q = 0.0\n","        self.curr_ep_loss_length = 0\n","\n","    def record(self, episode, epsilon, step):\n","        mean_ep_reward = np.round(np.mean(self.ep_rewards[-100:]), 3)\n","        mean_ep_length = np.round(np.mean(self.ep_lengths[-100:]), 3)\n","        mean_ep_loss = np.round(np.mean(self.ep_avg_losses[-100:]), 3)\n","        mean_ep_q = np.round(np.mean(self.ep_avg_qs[-100:]), 3)\n","        self.moving_avg_ep_rewards.append(mean_ep_reward)\n","        self.moving_avg_ep_lengths.append(mean_ep_length)\n","        self.moving_avg_ep_avg_losses.append(mean_ep_loss)\n","        self.moving_avg_ep_avg_qs.append(mean_ep_q)\n","\n","        last_record_time = self.record_time\n","        self.record_time = time.time()\n","        time_since_last_record = np.round(self.record_time - last_record_time, 3)\n","\n","        print(\n","            f\"Episode {episode} - \"\n","            f\"Step {step} - \"\n","            f\"Epsilon {epsilon} - \"\n","            f\"Mean Reward {mean_ep_reward} - \"\n","            f\"Mean Length {mean_ep_length} - \"\n","            f\"Mean Loss {mean_ep_loss} - \"\n","            f\"Mean Q Value {mean_ep_q} - \"\n","            f\"Time Delta {time_since_last_record} - \"\n","            f\"Time {datetime.datetime.now().strftime('%Y-%m-%dT%H:%M:%S')}\"\n","        )\n","\n","        with open(self.save_log, \"a\") as f:\n","            f.write(\n","                f\"{episode:8d}{step:8d}{epsilon:10.3f}\"\n","                f\"{mean_ep_reward:15.3f}{mean_ep_length:15.3f}{mean_ep_loss:15.3f}{mean_ep_q:15.3f}\"\n","                f\"{time_since_last_record:15.3f}\"\n","                f\"{datetime.datetime.now().strftime('%Y-%m-%dT%H:%M:%S'):>20}\\n\"\n","            )\n","\n","        for metric in [\"ep_lengths\", \"ep_avg_losses\", \"ep_avg_qs\", \"ep_rewards\"]:\n","            plt.clf()\n","            plt.plot(getattr(self, f\"moving_avg_{metric}\"), label=f\"moving_avg_{metric}\")\n","            plt.legend()\n","            plt.savefig(getattr(self, f\"{metric}_plot\"))"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2023-11-25T13:13:40.406913Z","iopub.status.busy":"2023-11-25T13:13:40.406524Z","iopub.status.idle":"2023-11-25T13:13:40.420904Z","shell.execute_reply":"2023-11-25T13:13:40.419941Z","shell.execute_reply.started":"2023-11-25T13:13:40.406877Z"},"id":"f0Ai9AEHCvoX","outputId":"4be9320e-ba31-44b0-d4e6-99047da41352","trusted":true},"outputs":[],"source":["\n","\n","save_dir = Path(\"checkpoints\") #/ datetime.datetime.now().strftime(\"%Y-%m-%dT%H-%M-%S\")\n","if not os.path.exists(save_dir):\n","  save_dir.mkdir(parents=True)\n","print(f'All weights will saved at {save_dir}')\n","\n"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2023-11-25T13:13:40.422365Z","iopub.status.busy":"2023-11-25T13:13:40.422023Z","iopub.status.idle":"2023-11-25T13:13:40.806734Z","shell.execute_reply":"2023-11-25T13:13:40.805898Z","shell.execute_reply.started":"2023-11-25T13:13:40.422334Z"},"trusted":true},"outputs":[],"source":["# Initialize mario, this will re-init the weight\n","mario = Mario(state_dim=(4, 84, 84), action_dim=env.action_space.n, save_dir=save_dir)"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2023-11-25T13:13:40.808224Z","iopub.status.busy":"2023-11-25T13:13:40.807881Z","iopub.status.idle":"2023-11-25T13:13:40.813111Z","shell.execute_reply":"2023-11-25T13:13:40.812114Z","shell.execute_reply.started":"2023-11-25T13:13:40.808197Z"},"trusted":true},"outputs":[],"source":["use_cuda = torch.cuda.is_available()\n","print(f\"Using CUDA: {use_cuda}\")\n","print()"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2023-11-25T13:13:40.814652Z","iopub.status.busy":"2023-11-25T13:13:40.814380Z","iopub.status.idle":"2023-11-25T13:13:41.151494Z","shell.execute_reply":"2023-11-25T13:13:41.150456Z","shell.execute_reply.started":"2023-11-25T13:13:40.814627Z"},"trusted":true},"outputs":[],"source":["## Ask your friends : Aut, Mon' or Pokpong, how to fix this cell so that the notebook can load the previously trained weight\n","\n","try:    \n","    ckp = torch.load('/kaggle/input/38000-ep/mario_net_last.chkpt', map_location=('cuda' if use_cuda else 'cpu'))\n","    print(ckp.get('exploration_rate'))\n","\n","    print(1)\n","    mario.exploration_rate = ckp.get('exploration_rate')\n","    mario.net.load_state_dict(ckp.get('model'))\n","    print(2)\n","    mario.optimizer.load_state_dict(ckp.get('optimizer_state_dict'))\n","#     mario.memory = ckp.get('memory')\n","except Exception as e:\n","    print(e)"]},{"cell_type":"markdown","metadata":{"id":"lfRT-HDSZuop"},"source":["## TRAINING Step\n","\n","In this example we run the training loop for 100 episodes (took around <10 mins on Colab T4 GPU), but for Mario to truly learn the ways of\n","his world, we suggest running the loop for at least few ten thousand episodes!\n","\n","\n"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2023-11-25T13:13:41.153370Z","iopub.status.busy":"2023-11-25T13:13:41.152845Z","iopub.status.idle":"2023-11-25T13:13:41.158015Z","shell.execute_reply":"2023-11-25T13:13:41.157105Z","shell.execute_reply.started":"2023-11-25T13:13:41.153334Z"},"id":"zgwB5k5yFIBQ","outputId":"f73361e9-8ea3-46c8-94c1-5f3c045aa468","trusted":true},"outputs":[],"source":["import warnings\n","warnings.filterwarnings('ignore')"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2023-11-25T13:13:41.159832Z","iopub.status.busy":"2023-11-25T13:13:41.159264Z","iopub.status.idle":"2023-11-25T13:13:41.170459Z","shell.execute_reply":"2023-11-25T13:13:41.169266Z","shell.execute_reply.started":"2023-11-25T13:13:41.159805Z"},"id":"2QdjD1EvZuop","trusted":true},"outputs":[],"source":["\n","\n","logger = MetricLogger(save_dir)\n","\n"]},{"cell_type":"code","execution_count":null,"metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.execute_input":"2023-11-20T09:40:24.453659Z","iopub.status.busy":"2023-11-20T09:40:24.452852Z","iopub.status.idle":"2023-11-20T13:30:32.310736Z","shell.execute_reply":"2023-11-20T13:30:32.308736Z","shell.execute_reply.started":"2023-11-20T09:40:24.453631Z"},"id":"cMZcchFifKwB","outputId":"afcd2bb8-5713-458e-8845-55c833abefa3","trusted":true},"outputs":[],"source":["%%time\n","episodes = TRAIN_EPISODES # 7-8 mins per 100 episodes on T4 without fp16\n","for e in range(episodes):\n","\n","    state = env.reset()\n","\n","    # Play the game!\n","    while True:\n","\n","        # Run agent on the state\n","        action = mario.act(state)\n","\n","        # Agent performs action\n","        next_state, reward, done, trunc, info = env.step(action)\n","\n","        # Remember\n","        mario.cache(state, next_state, action, reward, done)\n","\n","        # Learn\n","        q, loss = mario.learn()\n","\n","        # Logging\n","        logger.log_step(reward, loss, q)\n","\n","        # Update state\n","        state = next_state\n","\n","        # Check if end of game\n","        if done or info[\"flag_get\"]:\n","            break\n","\n","    logger.log_episode()\n","\n","    if e % 20 == 0:\n","        logger.record(episode=e, epsilon=mario.exploration_rate, step=mario.curr_step)"]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":[]},{"cell_type":"markdown","metadata":{"id":"6H_c77IwvuRM"},"source":["## Display and Save Video"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2023-11-25T13:16:28.750371Z","iopub.status.busy":"2023-11-25T13:16:28.749623Z","iopub.status.idle":"2023-11-25T13:16:28.754862Z","shell.execute_reply":"2023-11-25T13:16:28.753966Z","shell.execute_reply.started":"2023-11-25T13:16:28.750328Z"},"id":"uXzCBPXg2r1z","outputId":"d856eb96-1f65-4bb2-d5ce-a5ece6b910a7","trusted":true},"outputs":[],"source":["from ipywidgets import Video\n","VIDEO_NAME = \"replay.webm\" # you SHOULD change the name by adding #episodes that already trained"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2023-11-25T13:16:35.756805Z","iopub.status.busy":"2023-11-25T13:16:35.756390Z","iopub.status.idle":"2023-11-25T13:16:35.764448Z","shell.execute_reply":"2023-11-25T13:16:35.763459Z","shell.execute_reply.started":"2023-11-25T13:16:35.756773Z"},"id":"ECtD4JGBvu7E","trusted":true},"outputs":[],"source":["def animate(imgs, video_name=None, _return=True):\n","    # using cv2 to generate videos\n","    import cv2\n","    import os\n","    import string\n","    import random\n","    video_name = video_name if video_name is not None else ''.join(random.choice(string.ascii_letters) for i in range(18))+'.webm'\n","    height, width, layers = imgs[0].shape\n","    video = cv2.VideoWriter(video_name, cv2.VideoWriter_fourcc(*'VP80'), 20, (width,height))\n","    for img in imgs:\n","        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n","        video.write(img)\n","    video.release()\n","    if _return:\n","        from IPython.display import Video\n","        return Video(video_name)\n","\n"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2023-11-25T13:17:26.845441Z","iopub.status.busy":"2023-11-25T13:17:26.844620Z","iopub.status.idle":"2023-11-25T13:17:38.868920Z","shell.execute_reply":"2023-11-25T13:17:38.868125Z","shell.execute_reply.started":"2023-11-25T13:17:26.845401Z"},"id":"j3BtBMRsweSZ","outputId":"9fbd71de-d410-4421-8ea6-6d92273fe28a","trusted":true},"outputs":[],"source":["import time\n","\n","state = env.reset()\n","# time.sleep(0.03)\n","imgs = [env.render()]\n","# time.sleep(0.03)\n","\n","i=0\n","while True:\n","\n","    action = mario.act(state)\n","    next_state, reward, done, trunc, info = env.step(action)\n","    mario.cache(state, next_state, action, reward, done)\n","\n","    state = next_state\n","\n","    time.sleep(0.03)\n","    img = env.render().copy()\n","    # time.sleep(0.03)\n","    # assert np.any(img != imgs[-1])\n","    imgs += [img]\n","\n","    i+=1\n","    if i%100==0:\n","        print(i)\n","\n","    if done or info['flag_get']:\n","        break"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2023-11-25T13:17:41.960831Z","iopub.status.busy":"2023-11-25T13:17:41.959825Z","iopub.status.idle":"2023-11-25T13:17:43.254581Z","shell.execute_reply":"2023-11-25T13:17:43.253745Z","shell.execute_reply.started":"2023-11-25T13:17:41.960786Z"},"id":"QodTpuzb1591","outputId":"c2a22580-2431-4d20-9c0f-c6e9783150fe","trusted":true},"outputs":[],"source":["print(len(imgs))\n","vid = animate(imgs, video_name=VIDEO_NAME, _return=True)\n"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2023-11-25T13:17:44.855135Z","iopub.status.busy":"2023-11-25T13:17:44.854394Z","iopub.status.idle":"2023-11-25T13:17:44.878753Z","shell.execute_reply":"2023-11-25T13:17:44.877725Z","shell.execute_reply.started":"2023-11-25T13:17:44.855094Z"},"id":"0p_STdsW9o-h","outputId":"7be88dc5-e855-4f45-c6ed-570b3bf51177","trusted":true},"outputs":[],"source":["Video.from_file(VIDEO_NAME) # right-click on the bottom-right panel to download the video"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2023-11-20T13:30:34.559261Z","iopub.status.busy":"2023-11-20T13:30:34.558907Z","iopub.status.idle":"2023-11-20T13:30:34.625521Z","shell.execute_reply":"2023-11-20T13:30:34.624444Z","shell.execute_reply.started":"2023-11-20T13:30:34.559235Z"},"id":"kpLbYANiMsDM","outputId":"5d31a95e-9212-48e9-f543-3e0dbcc7c2f9","trusted":true},"outputs":[],"source":["mario.save(SAVE_PATH)"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2023-11-20T13:30:34.627144Z","iopub.status.busy":"2023-11-20T13:30:34.626776Z","iopub.status.idle":"2023-11-20T13:30:36.887335Z","shell.execute_reply":"2023-11-20T13:30:36.886291Z","shell.execute_reply.started":"2023-11-20T13:30:34.627108Z"},"id":"trRV4uYBOeJx","outputId":"be7f228d-75c6-4b41-a10f-6006f3dd42ab","trusted":true},"outputs":[],"source":["!pwd\n","!ls -sh\n"]},{"cell_type":"code","execution_count":null,"metadata":{"_kg_hide-output":true},"outputs":[],"source":[]}],"metadata":{"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"datasetId":4036417,"sourceId":7019651,"sourceType":"datasetVersion"},{"datasetId":4037512,"sourceId":7021241,"sourceType":"datasetVersion"},{"datasetId":4038402,"sourceId":7022748,"sourceType":"datasetVersion"},{"datasetId":4040210,"sourceId":7025324,"sourceType":"datasetVersion"},{"datasetId":4043133,"sourceId":7029449,"sourceType":"datasetVersion"},{"datasetId":4055024,"sourceId":7046838,"sourceType":"datasetVersion"},{"datasetId":4057505,"sourceId":7050445,"sourceType":"datasetVersion"},{"sourceId":149282146,"sourceType":"kernelVersion"}],"dockerImageVersionId":30558,"isGpuEnabled":true,"isInternetEnabled":true,"language":"python","sourceType":"notebook"},"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.11.0"}},"nbformat":4,"nbformat_minor":4}