Pong_DQN.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "q4AS7njD7iL6"
      },
      "outputs": [],
      "source": [
        "!nvidia-smi"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "-D98CddQuwKG"
      },
      "outputs": [],
      "source": [
        "import gymnasium as gym\n",
        "import cv2\n",
        "\n",
        "import time\n",
        "import json\n",
        "import random\n",
        "import numpy as np\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "import torch.nn.functional as F\n",
        "\n",
        "from collections import deque"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Hyperparameters"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "I8rGMlN2uzgk"
      },
      "outputs": [],
      "source": [
        "ENVIRONMENT = \"PongDeterministic-v4\"\n",
        "\n",
        "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "\n",
        "SAVE_MODELS = True  # Save models to file so you can test later\n",
        "MODEL_PATH = \"./saved_models/pong-cnn-\"  # Models path for saving or loading\n",
        "SAVE_MODEL_INTERVAL = 10  # Save models at every X epoch\n",
        "TRAIN_MODEL = True  # Train model while playing (Make it False when testing a model)\n",
        "\n",
        "LOAD_MODEL_FROM_FILE = True  # Load model from file\n",
        "LOAD_FILE_EPISODE = 860  # Load Xth episode from file\n",
        "\n",
        "BATCH_SIZE = 64  # Minibatch size that select randomly from mem for train nets\n",
        "MAX_EPISODE = 100000  # Max episode\n",
        "MAX_STEP = 100000  # Max step size for one episode\n",
        "\n",
        "MAX_MEMORY_LEN = 50000  # Max memory len\n",
        "MIN_MEMORY_LEN = 40000  # Min memory len before start train\n",
        "\n",
        "GAMMA = 0.97  # Discount rate\n",
        "ALPHA = 0.00025  # Learning rate\n",
        "EPSILON_DECAY = 0.99  # Epsilon decay rate by step\n",
        "\n",
        "RENDER_GAME_WINDOW = False  # Opens a new window to render the game (Won't work on colab default)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Network Architecture Class"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "HxF5-bzUu1q-"
      },
      "outputs": [],
      "source": [
        "class DuelCNN(nn.Module):\n",
        "    \"\"\"\n",
        "    CNN with Duel Algo. https://arxiv.org/abs/1511.06581\n",
        "    \"\"\"\n",
        "    def __init__(self, h, w, output_size):\n",
        "        super(DuelCNN, self).__init__()\n",
        "        self.conv1 = nn.Conv2d(in_channels=4,  out_channels=32, kernel_size=8, stride=4)\n",
        "        self.bn1 = nn.BatchNorm2d(32)\n",
        "        convw, convh = self.conv2d_size_calc(w, h, kernel_size=8, stride=4)\n",
        "        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=4, stride=2)\n",
        "        self.bn2 = nn.BatchNorm2d(64)\n",
        "        convw, convh = self.conv2d_size_calc(convw, convh, kernel_size=4, stride=2)\n",
        "        self.conv3 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1)\n",
        "        self.bn3 = nn.BatchNorm2d(64)\n",
        "        convw, convh = self.conv2d_size_calc(convw, convh, kernel_size=3, stride=1)\n",
        "\n",
        "        linear_input_size = convw * convh * 64  # Last conv layer's out sizes\n",
        "\n",
        "        # Action layer\n",
        "        self.Alinear1 = nn.Linear(in_features=linear_input_size, out_features=128)\n",
        "        self.Alrelu = nn.LeakyReLU()  # Linear 1 activation funct\n",
        "        self.Alinear2 = nn.Linear(in_features=128, out_features=output_size)\n",
        "\n",
        "        # State Value layer\n",
        "        self.Vlinear1 = nn.Linear(in_features=linear_input_size, out_features=128)\n",
        "        self.Vlrelu = nn.LeakyReLU()  # Linear 1 activation funct\n",
        "        self.Vlinear2 = nn.Linear(in_features=128, out_features=1)  # Only 1 node\n",
        "\n",
        "    def conv2d_size_calc(self, w, h, kernel_size=5, stride=2):\n",
        "        \"\"\"\n",
        "        Helper Function: Calcs conv layers output image sizes\n",
        "        \"\"\"\n",
        "        next_w = (w - (kernel_size - 1) - 1) // stride + 1\n",
        "        next_h = (h - (kernel_size - 1) - 1) // stride + 1\n",
        "        return next_w, next_h\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = F.relu(self.bn1(self.conv1(x)))\n",
        "        x = F.relu(self.bn2(self.conv2(x)))\n",
        "        x = F.relu(self.bn3(self.conv3(x)))\n",
        "\n",
        "        x = x.view(x.size(0), -1)  # Flatten every batch\n",
        "\n",
        "        Ax = self.Alrelu(self.Alinear1(x))\n",
        "        Ax = self.Alinear2(Ax)  # No activation on last layer\n",
        "\n",
        "        Vx = self.Vlrelu(self.Vlinear1(x))\n",
        "        Vx = self.Vlinear2(Vx)  # No activation on last layer\n",
        "\n",
        "        q = Vx + (Ax - Ax.mean())\n",
        "\n",
        "        return q"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Define 'Agent' class"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "plT51MPbu5U5"
      },
      "outputs": [],
      "source": [
        "class Agent:\n",
        "    def __init__(self, environment):\n",
        "        \"\"\"\n",
        "        Hyperparameters definition for Agent\n",
        "        \"\"\"\n",
        "        # State size for breakout env. SS images (210, 160, 3). Used as input size in network\n",
        "        self.state_size_h = environment.observation_space.shape[0]\n",
        "        self.state_size_w = environment.observation_space.shape[1]\n",
        "        self.state_size_c = environment.observation_space.shape[2]\n",
        "\n",
        "        # Activation size for breakout env. Used as output size in network\n",
        "        self.action_size = environment.action_space.n\n",
        "\n",
        "        # Image pre process params\n",
        "        self.target_h = 80  # Height after process\n",
        "        self.target_w = 64  # Widht after process\n",
        "\n",
        "        self.crop_dim = [20, self.state_size_h, 0, self.state_size_w]  # Cut 20 px from top to get rid of the score table\n",
        "\n",
        "        # Trust rate to our experiences\n",
        "        self.gamma = GAMMA  # Discount coef for future predictions\n",
        "        self.alpha = ALPHA  # Learning Rate\n",
        "\n",
        "        # After many experinces epsilon will be 0.05\n",
        "        # So we will do less Explore more Exploit\n",
        "        self.epsilon = 1  # Explore or Exploit\n",
        "        self.epsilon_decay = EPSILON_DECAY  # Adaptive Epsilon Decay Rate\n",
        "        self.epsilon_minimum = 0.05  # Minimum for Explore\n",
        "\n",
        "        # Deque holds replay mem.\n",
        "        self.memory = deque(maxlen=MAX_MEMORY_LEN)\n",
        "\n",
        "        # Create two model for DDQN algorithm\n",
        "        self.online_model = DuelCNN(h=self.target_h, w=self.target_w, output_size=self.action_size).to(DEVICE)\n",
        "        self.target_model = DuelCNN(h=self.target_h, w=self.target_w, output_size=self.action_size).to(DEVICE)\n",
        "        self.target_model.load_state_dict(self.online_model.state_dict())\n",
        "        self.target_model.eval()\n",
        "\n",
        "        # Adam used as optimizer\n",
        "        self.optimizer = optim.Adam(self.online_model.parameters(), lr=self.alpha)\n",
        "\n",
        "    def preProcess(self, image):\n",
        "        \"\"\"\n",
        "        Process image crop resize, grayscale and normalize the images\n",
        "        \"\"\"\n",
        "        frame = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)  # To grayscale\n",
        "        frame = frame[self.crop_dim[0]:self.crop_dim[1], self.crop_dim[2]:self.crop_dim[3]]  # Cut 20 px from top\n",
        "        frame = cv2.resize(frame, (self.target_w, self.target_h))  # Resize\n",
        "        frame = frame / 255  # Normalize\n",
        "        return frame\n",
        "\n",
        "    def act(self, state):\n",
        "        \"\"\"\n",
        "        Get state and do action\n",
        "        Two option can be selected: if 'explore' select random action\n",
        "        if 'exploit' ask nnet for action\n",
        "        \"\"\"\n",
        "\n",
        "        act_protocol = 'Explore' if random.uniform(0, 1) <= self.epsilon else 'Exploit'\n",
        "\n",
        "        if act_protocol == 'Explore':\n",
        "            action = random.randrange(self.action_size)\n",
        "        else:\n",
        "            with torch.no_grad():\n",
        "                state = torch.tensor(state, dtype=torch.float, device=DEVICE).unsqueeze(0)\n",
        "                q_values = self.online_model.forward(state)  # (1, action_size)\n",
        "                action = torch.argmax(q_values).item()  # Returns the indices of the maximum value of all elements\n",
        "\n",
        "        return action\n",
        "\n",
        "    def train(self):\n",
        "        \"\"\"\n",
        "        Train neural nets with replay memory\n",
        "        returns loss and max_q val predicted from online_net\n",
        "        \"\"\"\n",
        "        if len(self.memory) < MIN_MEMORY_LEN:\n",
        "            loss, max_q = [0, 0]\n",
        "            return loss, max_q\n",
        "        # We get out minibatch and turn it to numpy array\n",
        "        state, action, reward, next_state, done = zip(*random.sample(self.memory, BATCH_SIZE))\n",
        "\n",
        "        # Concat batches in one array\n",
        "        # (np.arr, np.arr) ==> np.BIGarr\n",
        "        state = np.concatenate(state)\n",
        "        next_state = np.concatenate(next_state)\n",
        "\n",
        "        # Convert them to tensors\n",
        "        state = torch.tensor(state, dtype=torch.float, device=DEVICE)\n",
        "        next_state = torch.tensor(next_state, dtype=torch.float, device=DEVICE)\n",
        "        action = torch.tensor(action, dtype=torch.long, device=DEVICE)\n",
        "        reward = torch.tensor(reward, dtype=torch.float, device=DEVICE)\n",
        "        done = torch.tensor(done, dtype=torch.float, device=DEVICE)\n",
        "\n",
        "        # Make predictions\n",
        "        state_q_values = self.online_model(state)\n",
        "        next_states_q_values = self.online_model(next_state)\n",
        "        next_states_target_q_values = self.target_model(next_state)\n",
        "\n",
        "        # Select Q-value for the action taken\n",
        "        selected_q_value = state_q_values.gather(1, action.unsqueeze(1)).squeeze(1)\n",
        "\n",
        "        # Double Q-Learning update\n",
        "        next_action_indices = next_states_q_values.max(1)[1].unsqueeze(1)\n",
        "        next_states_target_q_value = next_states_target_q_values.gather(1, next_action_indices).squeeze(1)\n",
        "\n",
        "        # Compute the expected Q-value using Bellman Equation\n",
        "        expected_q_value = reward + self.gamma * next_states_target_q_value * (1 - done)\n",
        "\n",
        "\n",
        "        # Calc loss with expected_q_value and q_value\n",
        "        loss = (selected_q_value - expected_q_value.detach()).pow(2).mean()\n",
        "\n",
        "        self.optimizer.zero_grad()\n",
        "        loss.backward()\n",
        "        self.optimizer.step()\n",
        "\n",
        "        return loss, torch.max(state_q_values).item()\n",
        "\n",
        "    def storeResults(self, state, action, reward, nextState, done):\n",
        "        \"\"\"\n",
        "        Store every result to memory\n",
        "        \"\"\"\n",
        "        self.memory.append([state[None, :], action, reward, nextState[None, :], done])\n",
        "\n",
        "    def adaptiveEpsilon(self):\n",
        "        \"\"\"\n",
        "        Adaptive Epsilon means every step\n",
        "        we decrease the epsilon so we do less Explore\n",
        "        \"\"\"\n",
        "        if self.epsilon > self.epsilon_minimum:\n",
        "            self.epsilon *= self.epsilon_decay"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Train the model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "ve4vYDe3bozg"
      },
      "outputs": [],
      "source": [
        "environment = gym.make(ENVIRONMENT, render_mode=\"rgb_array\")  # Get env\n",
        "environment = gym.wrappers.RecordVideo(environment, f\"videos/OpenAIPong_DQN\")\n",
        "agent = Agent(environment)  # Create Agent\n",
        "\n",
        "if LOAD_MODEL_FROM_FILE:\n",
        "    # torch.load with map_location=torch.device('cpu')\n",
        "    agent.online_model.load_state_dict(torch.load(MODEL_PATH+str(LOAD_FILE_EPISODE)+\".pkl\", map_location=torch.device('cpu')))\n",
        "\n",
        "    with open(MODEL_PATH+str(LOAD_FILE_EPISODE)+'.json') as outfile:\n",
        "        param = json.load(outfile)\n",
        "        agent.epsilon = param.get('epsilon')\n",
        "\n",
        "    startEpisode = LOAD_FILE_EPISODE + 1\n",
        "\n",
        "else:\n",
        "    startEpisode = 1\n",
        "\n",
        "last_100_ep_reward = deque(maxlen=100)  # Last 100 episode rewards\n",
        "total_step = 1  # Cumulkative sum of all steps in episodes\n",
        "for episode in range(startEpisode, MAX_EPISODE):\n",
        "\n",
        "    startTime = time.time()  # Keep time\n",
        "    state, _ = environment.reset()  # Reset env\n",
        "    state = agent.preProcess(state)  # Process image\n",
        "    # Stack state . Every state contains 4 time contionusly frames\n",
        "    # We stack frames like 4 channel image\n",
        "    state = np.stack((state, state, state, state))\n",
        "\n",
        "    total_max_q_val = 0  # Total max q vals\n",
        "    total_reward = 0  # Total reward for each episode\n",
        "    total_loss = 0  # Total loss for each episode\n",
        "    for step in range(MAX_STEP):\n",
        "\n",
        "        if RENDER_GAME_WINDOW:\n",
        "            environment.render()  # Show state visually\n",
        "\n",
        "        # Select and perform an action\n",
        "        action = agent.act(state)  # Act\n",
        "        next_state, reward, done, truncate, info = environment.step(action)  # Observe\n",
        "\n",
        "        next_state = agent.preProcess(next_state)  # Process image\n",
        "\n",
        "        # Stack state . Every state contains 4 time contionusly frames\n",
        "        # We stack frames like 4 channel image\n",
        "        next_state = np.stack((next_state, state[0], state[1], state[2]))\n",
        "\n",
        "        # Store the transition in memory\n",
        "        agent.storeResults(state, action, reward, next_state, done)  # Store to mem\n",
        "\n",
        "        # Move to the next state\n",
        "        state = next_state  # Update state\n",
        "\n",
        "        if TRAIN_MODEL:\n",
        "            # Perform one step of the optimization (on the target network)\n",
        "            loss, max_q_val = agent.train()  # Train with random BATCH_SIZE state taken from mem\n",
        "        else:\n",
        "            loss, max_q_val = [0, 0]\n",
        "\n",
        "        total_loss += loss\n",
        "        total_max_q_val += max_q_val\n",
        "        total_reward += reward\n",
        "        total_step += 1\n",
        "        if total_step % 1000 == 0:\n",
        "            agent.adaptiveEpsilon()  # Decrase epsilon\n",
        "\n",
        "        if done:  # Episode completed\n",
        "            currentTime = time.time()  # Keep current time\n",
        "            time_passed = currentTime - startTime  # Find episode duration\n",
        "            current_time_format = time.strftime(\"%H:%M:%S\", time.gmtime())  # Get current dateTime as HH:MM:SS\n",
        "            epsilonDict = {'epsilon': agent.epsilon}  # Create epsilon dict to save model as file\n",
        "\n",
        "            if SAVE_MODELS and episode % SAVE_MODEL_INTERVAL == 0:  # Save model as file\n",
        "                weightsPath = MODEL_PATH + str(episode) + '.pkl'\n",
        "                epsilonPath = MODEL_PATH + str(episode) + '.json'\n",
        "\n",
        "                torch.save(agent.online_model.state_dict(), weightsPath)\n",
        "                with open(epsilonPath, 'w') as outfile:\n",
        "                    json.dump(epsilonDict, outfile)\n",
        "\n",
        "            if TRAIN_MODEL:\n",
        "                agent.target_model.load_state_dict(agent.online_model.state_dict())  # Update target model\n",
        "\n",
        "            last_100_ep_reward.append(total_reward)\n",
        "            avg_max_q_val = total_max_q_val / step\n",
        "\n",
        "            outStr = \"Episode:{} Time:{} Reward:{:.2f} Loss:{:.2f} Last_100_Avg_Rew:{:.3f} Avg_Max_Q:{:.3f} Epsilon:{:.2f} Duration:{:.2f} Step:{} CStep:{}\".format(\n",
        "                episode, current_time_format, total_reward, total_loss, np.mean(last_100_ep_reward), avg_max_q_val, agent.epsilon, time_passed, step, total_step\n",
        "            )\n",
        "\n",
        "            print(outStr)\n",
        "\n",
        "            if SAVE_MODELS:\n",
        "                outputPath = MODEL_PATH + \"out\" + '.txt'  # Save outStr to file\n",
        "                with open(outputPath, 'a') as outfile:\n",
        "                    outfile.write(outStr+\"\\n\")\n",
        "\n",
        "            break\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Visualize CNN Outputs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "environment = gym.make(ENVIRONMENT, render_mode=\"rgb_array\")  # Get env\n",
        "agent = Agent(environment)  # Create Agent\n",
        "agent.online_model.load_state_dict(torch.load(MODEL_PATH+str(LOAD_FILE_EPISODE)+\".pkl\", map_location=torch.device('cpu')))\n",
        "\n",
        "class GradCAM:\n",
        "    def __init__(self, model, target_layer):\n",
        "        self.model = model\n",
        "        self.gradients = None\n",
        "        self.activations = None\n",
        "\n",
        "        target_layer.register_forward_hook(self.forward_hook)\n",
        "        target_layer.register_backward_hook(self.backward_hook)\n",
        "\n",
        "    def forward_hook(self, module, input, output):\n",
        "        self.activations = output.detach()\n",
        "\n",
        "    def backward_hook(self, module, grad_input, grad_output):\n",
        "        self.gradients = grad_output[0].detach()\n",
        "\n",
        "    def __call__(self, input):\n",
        "        self.model.train()\n",
        "        self.model.zero_grad()\n",
        "        \n",
        "        output = self.model(input)\n",
        "        q_values = output  # Assuming output is the Q-values\n",
        "\n",
        "        # Get the action with the highest Q-value\n",
        "        max_q_value, _ = torch.max(q_values, dim=1)\n",
        "        max_q_value.backward()\n",
        "\n",
        "        # Get gradients and activations\n",
        "        gradients = self.gradients[0]  # Remove batch dimension\n",
        "        activations = self.activations[0]\n",
        "\n",
        "        # Compute weights\n",
        "        weights = torch.mean(gradients, dim=(1, 2))\n",
        "\n",
        "        # Compute Grad-CAM\n",
        "        cam = torch.zeros(activations.shape[1:], dtype=torch.float32)\n",
        "\n",
        "        for i, w in enumerate(weights):\n",
        "            cam += w * activations[i]\n",
        "\n",
        "        cam = cam.cpu().numpy()\n",
        "        cam = np.maximum(cam, 0)\n",
        "        cam = cv2.resize(cam, (input.size(-1), input.size(-2)))\n",
        "        cam = cam - np.min(cam)\n",
        "        cam = cam / np.max(cam)\n",
        "\n",
        "        return cam\n",
        "\n",
        "# Instantiate Grad-CAM with the target convolutional layer (e.g., conv3)\n",
        "model = agent.online_model\n",
        "grad_cam = GradCAM(model, model.conv3)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from matplotlib import pyplot as plt\n",
        "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
        "\n",
        "\n",
        "# reset environment and do some random actions\n",
        "state, _ = environment.reset()\n",
        "state = agent.preProcess(state)\n",
        "state = np.stack((state, state, state, state))\n",
        "action_range = random.randint(1, 500)\n",
        "for _ in range(action_range):\n",
        "    # action = environment.action_space.sample()\n",
        "    with torch.no_grad():\n",
        "        state_tensor = torch.tensor(state, dtype=torch.float, device=DEVICE).unsqueeze(0)\n",
        "        q_values = model.forward(state_tensor)  # (1, action_size)\n",
        "        action = torch.argmax(q_values).item() \n",
        "    next_state, _, done, _, _ = environment.step(action)\n",
        "    if done:\n",
        "        next_state, _ = environment.reset()\n",
        "    next_state = agent.preProcess(next_state)\n",
        "    state = np.stack((next_state, state[0], state[1], state[2]))\n",
        "    \n",
        "\n",
        "# Compute Grad-CAM\n",
        "state = torch.tensor(state, dtype=torch.float32).unsqueeze(0)  # Shape: [1, 4, 80, 64]\n",
        "cam = grad_cam(state)\n",
        "\n",
        "# Get the original frame (grayscale) from the state\n",
        "original_frame = state[0, -1].cpu().numpy()  # Take the last frame in the stack; shape: [80, 64]\n",
        "\n",
        "# # Resize cam to match the original frame size\n",
        "# cam_resized = cv2.resize(cam, (original_frame.shape[1], original_frame.shape[0]))  # (width, height)\n",
        "\n",
        "# # Normalize the heatmap\n",
        "# cam = cam - np.min(cam)\n",
        "# cam = cam / np.max(cam)\n",
        "\n",
        "# Convert the grayscale image to RGB\n",
        "original_frame_rgb = np.stack([original_frame, original_frame, original_frame], axis=2)  # Shape: [80, 64, 3]\n",
        "\n",
        "# Apply color map to the heatmap\n",
        "heatmap = cv2.applyColorMap(np.uint8(255 * cam), cv2.COLORMAP_VIRIDIS)\n",
        "heatmap = np.float32(heatmap) / 255  # Normalize to [0, 1]\n",
        "\n",
        "# Overlay the heatmap on the original image\n",
        "alpha = 0.4  # Transparency factor for the heatmap\n",
        "overlay = alpha * heatmap + (1 - alpha) * original_frame_rgb\n",
        "overlay = overlay / np.max(overlay)  # Normalize to [0, 1] for display\n",
        "\n",
        "# Create a figure with 3 subplots\n",
        "fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(18, 6))\n",
        "\n",
        "# 1. Original Image\n",
        "ax1.imshow(original_frame, cmap='gray')\n",
        "ax1.axis('off')\n",
        "ax1.set_title('Original Image')\n",
        "\n",
        "# 2. Heatmap with Colorbar\n",
        "im2 = ax2.imshow(cam, cmap='viridis')\n",
        "ax2.axis('off')\n",
        "ax2.set_title('Grad-CAM Heatmap')\n",
        "\n",
        "# Add colorbar next to the heatmap\n",
        "divider = make_axes_locatable(ax2)\n",
        "cax = divider.append_axes('right', size='5%', pad=0.05)\n",
        "cbar = fig.colorbar(im2, cax=cax, orientation='vertical')\n",
        "cbar.set_label('Importance')\n",
        "\n",
        "# 3. Overlay Image\n",
        "ax3.imshow(overlay)\n",
        "ax3.axis('off')\n",
        "ax3.set_title('Overlay of Heatmap on Original Image')\n",
        "\n",
        "plt.tight_layout()\n",
        "plt.show()"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "OpenAIPong-DQN.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.12"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}