dqn_vs_dqn_pong.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RFqFkfdMdbKq",
        "outputId": "79f02184-1fcb-45d6-cdac-d4c91865690f"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "using device: cuda\n"
          ]
        }
      ],
      "source": [
        "import os\n",
        "import time\n",
        "import random\n",
        "from IPython.display import display, clear_output\n",
        "import PIL.Image\n",
        "import torch\n",
        "from torch import nn, optim\n",
        "import torch.nn.functional as F\n",
        "from collections import deque\n",
        "import imageio\n",
        "import numpy as np\n",
        "from PIL import Image\n",
        "import matplotlib.pyplot as plt\n",
        "import matplotlib.animation as animation\n",
        "\n",
        "import gym\n",
        "import retro\n",
        "import time\n",
        "import json\n",
        "import cv2\n",
        "\n",
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "print(f'using device: {device}')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DFcAtSkW3Kuw"
      },
      "outputs": [],
      "source": [
        "class DuelCNNPlayerModel(nn.Module):\n",
        "    def __init__(self, height, width, output_size):\n",
        "        super(DuelCNNPlayerModel, self).__init__()\n",
        "\n",
        "        # Convolutional neural network\n",
        "        self.conv1 = nn.Conv2d(in_channels = 4, out_channels = 32, kernel_size = 8, stride = 4)\n",
        "        self.bn1 = nn.BatchNorm2d(32)\n",
        "        conv_height, conv_width = self.calculate_updated_size(height, width, 8, 4)\n",
        "        self.conv2 = nn.Conv2d(in_channels = 32, out_channels = 64, kernel_size = 4, stride = 2)\n",
        "        self.bn2 = nn.BatchNorm2d(64)\n",
        "        conv_height, conv_width = self.calculate_updated_size(conv_height, conv_width, 4, 2)\n",
        "        self.conv3 = nn.Conv2d(in_channels = 64, out_channels = 64, kernel_size = 3, stride = 1)\n",
        "        self.bn3 = nn.BatchNorm2d(64)\n",
        "        conv_height, conv_width = self.calculate_updated_size(conv_height, conv_width, 3, 1)\n",
        "\n",
        "        input_size = conv_height * conv_width * 64\n",
        "\n",
        "        # Dueling action layer\n",
        "        self.action_linear_1 = nn.Linear(in_features = input_size, out_features = 128)\n",
        "        self.action_relu = nn.LeakyReLU()\n",
        "        self.action_linear_2 = nn.Linear(in_features = 128, out_features = output_size)\n",
        "\n",
        "        # Dueling state layer\n",
        "        self.state_linear_1 = nn.Linear(in_features = input_size, out_features = 128)\n",
        "        self.state_relu = nn.LeakyReLU()\n",
        "        self.state_linear_2 = nn.Linear(in_features = 128, out_features = 1)\n",
        "\n",
        "    def calculate_updated_size(self, height, width, kernel_size, stride):\n",
        "        return (\n",
        "            ((height - kernel_size) // stride + 1),\n",
        "            ((width - kernel_size) // stride + 1)\n",
        "        )\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)\n",
        "\n",
        "        action_x = self.action_relu(self.action_linear_1(x))\n",
        "        action_x = self.action_linear_2(action_x)\n",
        "\n",
        "        state_x = self.state_relu(self.state_linear_1(x))\n",
        "        state_x = self.state_linear_2(state_x)\n",
        "\n",
        "        q = state_x + (action_x - action_x.mean())\n",
        "\n",
        "        return q\n",
        "\n",
        "# Stub class\n",
        "class DQNPlayerAgent:\n",
        "    def __init__(self, env):\n",
        "\n",
        "        self.state_size_h = env.observation_space.shape[0]\n",
        "        self.state_size_w = env.observation_space.shape[1]\n",
        "        self.state_size_c = env.observation_space.shape[2]\n",
        "\n",
        "        self.height = 80\n",
        "        self.width = 64\n",
        "\n",
        "        self.action_space = 4\n",
        "\n",
        "        self.targetNN = None\n",
        "        self.onlineNN = None\n",
        "\n",
        "        self.targetNN = DuelCNNPlayerModel(self.height, self.width, self.action_space).to(device)\n",
        "        self.onlineNN = DuelCNNPlayerModel(self.height, self.width, self.action_space).to(device)\n",
        "        self.targetNN.load_state_dict(self.onlineNN.state_dict())\n",
        "        self.targetNN.eval()\n",
        "\n",
        "        # Note: if we seem to have a lot of extra RAM, we might want to increase this\n",
        "        self.memory = deque(maxlen=50000)\n",
        "        self.optimizer = optim.Adam(self.onlineNN.parameters(), lr=0.00025)\n",
        "\n",
        "        self.epsilon_min = 0.1\n",
        "        self.epsilon_decay = 0.92\n",
        "        self.epsilon = 1\n",
        "\n",
        "    def predict_action(self, state):\n",
        "        # Greedy algo\n",
        "        if random.random() < self.epsilon:\n",
        "            action = random.choice(range(self.action_space))\n",
        "        else:\n",
        "            with torch.no_grad():\n",
        "                state = torch.tensor(state, dtype=torch.float, device=device).unsqueeze(0)\n",
        "                q_values = self.onlineNN.forward(state)  # (1, action_space)\n",
        "                action = torch.argmax(q_values).item()\n",
        "\n",
        "        return action\n",
        "\n",
        "    def decay_epsilon(self):\n",
        "        if self.epsilon > self.epsilon_min:\n",
        "            self.epsilon *= self.epsilon_decay\n",
        "            self.epsilon = max(self.epsilon, self.epsilon_min)\n",
        "\n",
        "\n",
        "    def train(self):\n",
        "        # Note: if we seem to have a lot of extra RAM, we might want to increase this\n",
        "        if len(self.memory) < 40000:\n",
        "            return 0, 0\n",
        "        # Sample a minibatch and convert to tensors to pass into pytorch\n",
        "        state, action, reward, next_state, done = zip(*random.sample(self.memory, 64))\n",
        "\n",
        "        state = np.concatenate(state)\n",
        "        next_state = np.concatenate(next_state)\n",
        "\n",
        "        state = torch.tensor(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",
        "        next_state = torch.tensor(next_state, dtype=torch.float, device=device)\n",
        "        done = torch.tensor(done, dtype=torch.float, device=device)\n",
        "\n",
        "        state_q_vals = self.onlineNN(state)\n",
        "        next_state_q_vals = self.onlineNN(next_state)\n",
        "        target_q_vals = self.targetNN(next_state)\n",
        "\n",
        "        selected_qvals = state_q_vals.gather(1, action.unsqueeze(1)).squeeze(1)\n",
        "        target_qvals = target_q_vals.gather(1, next_state_q_vals.max(1)[1].unsqueeze(1)).squeeze(1)\n",
        "        expected_qvals = reward + (0.99 * target_qvals * (1 - done))\n",
        "\n",
        "        loss = (selected_qvals - expected_qvals.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_vals).item()\n",
        "\n",
        "    def update_target_to_online(self):\n",
        "        self.targetNN.load_state_dict(self.onlineNN.state_dict())\n",
        "\n",
        "    def store_results(self, data):\n",
        "        self.memory.append([data[0][None, :], data[1], data[2], data[3][None, :], data[4]])\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VNugsO-A-h2u"
      },
      "outputs": [],
      "source": [
        "def preprocess_frame(frame):\n",
        "    frame_cropped = frame[35:195:2, ::2, :]\n",
        "    frame_gray = np.dot(frame_cropped[..., :3], [0.299, 0.587, 0.114])\n",
        "    frame_normalized = frame_gray / 255.0\n",
        "    return np.expand_dims(frame_normalized, axis=0)\n",
        "\n",
        "def preProcess(image):\n",
        "      \"\"\"\n",
        "      Process image crop resize, grayscale and normalize the images\n",
        "      \"\"\"\n",
        "      frame = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)  # To grayscale\n",
        "      crop_dim = [20, 210, 0, 160]\n",
        "      frame = frame[crop_dim[0]:crop_dim[1], crop_dim[2]:crop_dim[3]]  # Cut 20 px from top\n",
        "      frame = cv2.resize(frame, (80, 64))  # Resize\n",
        "      frame = frame.reshape(80, 64) / 255  # Normalize\n",
        "\n",
        "      return frame\n",
        "\n",
        "def train_dqn(model, minibatch, gamma):\n",
        "    # Extract data from the minibatch\n",
        "    states = np.array([sample[0] for sample in minibatch])\n",
        "    actions = np.array([sample[1] for sample in minibatch])\n",
        "    rewards = np.array([sample[2] for sample in minibatch])\n",
        "    next_states = np.array([sample[3] for sample in minibatch])\n",
        "    dones = np.array([sample[4] for sample in minibatch])\n",
        "\n",
        "    # Predict Q-values for the current and next states\n",
        "    q_values = model.predict(states, verbose=0)\n",
        "    next_q_values = model.predict(next_states, verbose=0)\n",
        "\n",
        "    # Update Q-values using the Bellman equation\n",
        "    for i in range(len(minibatch)):\n",
        "        target = rewards[i]\n",
        "        if not dones[i]:\n",
        "            target += gamma * np.max(next_q_values[i])\n",
        "\n",
        "        # Handle multi-dimensional or one-hot encoded actions\n",
        "        if isinstance(actions[i], (list, np.ndarray)):\n",
        "            # Assume one-hot encoding for actions\n",
        "            action_index = np.argmax(actions[i])\n",
        "        else:\n",
        "            # Direct integer index\n",
        "            action_index = int(actions[i])\n",
        "\n",
        "        # Update the target Q-value for the chosen action\n",
        "        q_values[i][action_index] = target\n",
        "\n",
        "    # Train the model on the updated Q-values\n",
        "    model.fit(states, q_values, verbose=0, batch_size=len(minibatch))\n",
        "\n",
        "\n",
        "def player1_action_sample():\n",
        "    action = [0.0] * 16\n",
        "    player1_action = random.choice([0, 4, 5])\n",
        "    action[player1_action] = 1\n",
        "    return action\n",
        "\n",
        "def player2_action_sample():\n",
        "    action = [0.0] * 16\n",
        "    player2_action = random.choice([6, 7, 15])\n",
        "    action[player2_action] = 1\n",
        "    return action\n",
        "\n",
        "def plot_paddle_bounces(episode_numbers, paddle_bounces):\n",
        "    plt.figure(figsize=(12, 6))\n",
        "    plt.plot(episode_numbers, paddle_bounces, linewidth=2)\n",
        "    plt.xlabel(\"Episode Number\", fontsize=14)\n",
        "    plt.ylabel(\"Paddle Bounces per Point\", fontsize=14)\n",
        "    plt.title(\"Number of Paddle Bounces per Episode\", fontsize=16)\n",
        "    plt.grid(True)\n",
        "    plt.show()\n",
        "\n",
        "def plot_average_serving_times(episode_numbers, average_serving_times):\n",
        "    plt.figure(figsize=(12, 6))\n",
        "    plt.plot(episode_numbers, average_serving_times, linewidth=2)\n",
        "    plt.xlabel(\"Episode number\", fontsize=14)\n",
        "    plt.ylabel(\"Average time between ball serve and hit\", fontsize=14)\n",
        "    plt.title(\"Average serving time per episode\", fontsize=16)\n",
        "    plt.grid(True)\n",
        "    plt.show()\n",
        "\n",
        "def plot_rewards(episode_numbers, rewards_p1, rewards_p2):\n",
        "    plt.figure(figsize=(12, 6))\n",
        "    plt.plot(episode_numbers, rewards_p1, label=\"Player 1 Cumulative Rewards\", linewidth=2)\n",
        "    plt.plot(episode_numbers, rewards_p2, label=\"Player 2 Cumulative Rewards\", linewidth=2)\n",
        "    plt.xlabel(\"Episode Number\", fontsize=14)\n",
        "    plt.ylabel(\"Cumulative Rewards\", fontsize=14)\n",
        "    plt.title(\"Cumulative Rewards over Episodes\", fontsize=16)\n",
        "    plt.legend(fontsize=12)\n",
        "    plt.grid(True)\n",
        "    plt.show()\n",
        "\n",
        "def plot_step_rewards(step_rewards_p1, step_rewards_p2, steps_to_display=100):\n",
        "    min_length = min(len(step_rewards_p1), len(step_rewards_p2))\n",
        "    step_rewards_p1 = step_rewards_p1[-min_length:]\n",
        "    step_rewards_p2 = step_rewards_p2[-min_length:]\n",
        "    start_index = max(0, min_length - steps_to_display)\n",
        "    steps = range(start_index, min_length)\n",
        "    step_rewards_p1 = step_rewards_p1[start_index:]\n",
        "    step_rewards_p2 = step_rewards_p2[start_index:]\n",
        "    plt.figure(figsize=(12, 6))\n",
        "    plt.plot(steps, step_rewards_p1, label=\"Player 1 Step Rewards\", marker='o', linestyle='-', linewidth=2)\n",
        "    plt.plot(steps, step_rewards_p2, label=\"Player 2 Step Rewards\", marker='s', linestyle='-', linewidth=2)\n",
        "    plt.xlabel(\"Step Number\", fontsize=14)\n",
        "    plt.ylabel(\"Reward\", fontsize=14)\n",
        "    plt.title(f\"Rewards Throughout the Last {steps_to_display} Steps\", fontsize=16)\n",
        "    plt.legend(fontsize=12)\n",
        "    plt.grid(True)\n",
        "    plt.show()\n",
        "\n",
        "def plot_ball_pos(ball_pos):\n",
        "    steps = range(len(ball_pos))\n",
        "    plt.figure(figsize=(12, 6))\n",
        "    plt.plot(steps, ball_pos, label=\"Ball Position in X axis\", marker='o', linestyle='-', linewidth=2)\n",
        "    plt.xlabel(\"Step Number\", fontsize=14)\n",
        "    plt.ylabel(\"Ball Position X-axis\", fontsize=14)\n",
        "    plt.title(\"Ball Position in X-axis Throughout an Episode\", fontsize=16)\n",
        "    plt.legend(fontsize=12)\n",
        "    plt.grid(True)\n",
        "    plt.show()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YVWi-RVZKnoB",
        "scrolled": true
      },
      "outputs": [],
      "source": [
        "# Run this cell to load the graphs from the autosave data\n",
        "with open('./models/autosave-data.json', 'r') as f:\n",
        "  json_data = json.loads(f.read())\n",
        "  cumulative_rewards_p1 = json_data['cumulative_rewards_p1']\n",
        "  cumulative_rewards_p2 = json_data['cumulative_rewards_p2']\n",
        "  episode_numbers = json_data['episode_numbers']\n",
        "  paddle_bounces_per_point = json_data['paddle_bounces_per_point']\n",
        "  average_serving_times = json_data['average_serving_times']\n",
        "\n",
        "plot_rewards(episode_numbers, cumulative_rewards_p1, cumulative_rewards_p2)\n",
        "plot_paddle_bounces(episode_numbers, paddle_bounces_per_point)\n",
        "plot_average_serving_times(episode_numbers, average_serving_times)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3qz9xEu9KOL6"
      },
      "outputs": [],
      "source": [
        "env = retro.make(game='Pong-Atari2600', players=2, render_mode='rgb_array')\n",
        "env.multi_rewards = True"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "editable": true,
        "id": "jQWw7ZMi3TFH",
        "scrolled": true,
        "tags": []
      },
      "outputs": [],
      "source": [
        "# Set this to a non-zero value to load an existing episode number from\n",
        "# the filesystem.\n",
        "load_episode_number = 820\n",
        "\n",
        "input_shape = (4, 80, 80)\n",
        "action_space = env.action_space.n\n",
        "\n",
        "player1_agent = DQNPlayerAgent(env)\n",
        "player2_agent = DQNPlayerAgent(env)\n",
        "\n",
        "start_episode = load_episode_number\n",
        "episode_save_interval = 20\n",
        "max_episodes = 99999\n",
        "batch_size = 32\n",
        "gamma = 0.99\n",
        "preview_frames = []\n",
        "frame_count = 0\n",
        "# How many frames to include in output gif\n",
        "frame_sampling_rate = 8\n",
        "\n",
        "# These is to keep track for graphing purposes\n",
        "cumulative_rewards_p1 = []\n",
        "cumulative_rewards_p2 = []\n",
        "episode_numbers = []\n",
        "paddle_bounces_per_point = []\n",
        "average_serving_times = []\n",
        "p1_losses = []\n",
        "p1_avg_max_q_val = []\n",
        "p2_losses = []\n",
        "p2_avg_max_q_val = []\n",
        "\n",
        "if load_episode_number != 0:\n",
        "  player1_agent.onlineNN.load_state_dict(torch.load(f'./models/p1-online-{load_episode_number}.pkl'))\n",
        "  player1_agent.targetNN.load_state_dict(torch.load(f'./models/p1-target-{load_episode_number}.pkl'))\n",
        "  player2_agent.onlineNN.load_state_dict(torch.load(f'./models/p2-online-{load_episode_number}.pkl'))\n",
        "  player2_agent.targetNN.load_state_dict(torch.load(f'./models/p2-online-{load_episode_number}.pkl'))\n",
        "  with open(f'./models/autosave-data-{load_episode_number}.json', 'r') as f:\n",
        "    json_data = json.loads(f.read())\n",
        "    cumulative_rewards_p1 = json_data['cumulative_rewards_p1'][:start_episode]\n",
        "    cumulative_rewards_p2 = json_data['cumulative_rewards_p2'][:start_episode]\n",
        "    episode_numbers = json_data['episode_numbers'][:start_episode]\n",
        "    paddle_bounces_per_point = json_data['paddle_bounces_per_point'][:start_episode]\n",
        "    average_serving_times = json_data['average_serving_times'][:start_episode]\n",
        "    p1_losses = json_data['p1_losses'][:start_episode]\n",
        "    p1_avg_max_q_val = json_data['p1_avg_max_q_val'][:start_episode]\n",
        "    p2_losses = json_data['p2_losses'][:start_episode]\n",
        "    p2_avg_max_q_val = json_data['p2_avg_max_q_val'][:start_episode]\n",
        "  player1_agent.epsilon = max(player1_agent.epsilon_decay ** start_episode, player1_agent.epsilon_min)\n",
        "  player2_agent.epsilon = max(player2_agent.epsilon_decay ** start_episode, player2_agent.epsilon_min)\n",
        "\n",
        "# Map of the actions able to be performed by each player to the button\n",
        "# index for the input in the order [SERVE, UP, DOWN, NOOP]\n",
        "PLAYER_ONE_ACTIONS = [0, 4, 5, -1]\n",
        "PLAYER_TWO_ACTIONS = [15, 6, 7, -1]\n",
        "\n",
        "show_output = False\n",
        "\n",
        "\n",
        "def show_output_print(*args):\n",
        "  if show_output:\n",
        "    print(*args)\n",
        "\n",
        "\n",
        "for episode in range(start_episode, max_episodes):\n",
        "  print(f\"Starting episode {episode}/{max_episodes}\")\n",
        "  episode_start_time = time.time()\n",
        "  state, _ = env.reset()\n",
        "  state = preProcess(state)\n",
        "  # Fill state with four of itself for the initial buffer\n",
        "  state = np.stack((state, state, state, state))\n",
        "\n",
        "  done = False\n",
        "  score_reward_p1 = 0\n",
        "  score_reward_p2 = 0\n",
        "\n",
        "  should_reward_p1 = False\n",
        "  should_reward_p2 = False\n",
        "  has_rewarded = True\n",
        "\n",
        "  step_count = 0\n",
        "  train_interval = 100\n",
        "  ram = env.get_ram()\n",
        "\n",
        "  p1_pos = 181\n",
        "  p2_pos = 181\n",
        "  ball_pos = []\n",
        "\n",
        "  step_rewards_p1 = []\n",
        "  step_rewards_p2 = []\n",
        "\n",
        "  prev_frames = []\n",
        "\n",
        "  penalized = False\n",
        "\n",
        "  paddle_bounces = 0\n",
        "  total_points = 0\n",
        "  serving_times = []\n",
        "  last_serve_frame_count = 0\n",
        "\n",
        "  p1_total_loss = 0\n",
        "  p1_total_max_q_val = 0\n",
        "  p2_total_loss = 0\n",
        "  p2_total_max_q_val = 0\n",
        "\n",
        "  while not done:\n",
        "    # -- new --\n",
        "    player1_action = player1_agent.predict_action(state)\n",
        "    player2_action = player2_agent.predict_action(state)\n",
        "\n",
        "    player1_action_index = PLAYER_ONE_ACTIONS[player1_action]\n",
        "    player2_action_index = PLAYER_TWO_ACTIONS[player2_action]\n",
        "\n",
        "    actions = [0.0] * 16\n",
        "    if player1_action_index != -1:\n",
        "      if p1_pos >= 200:\n",
        "        actions[4] = 1.0\n",
        "      elif p1_pos < 40:\n",
        "        actions[5] = 1.0\n",
        "      else:\n",
        "        actions[player1_action_index] = 1.0\n",
        "    if player2_action_index != -1:\n",
        "      if p2_pos >= 200:\n",
        "        actions[6] = 1.0\n",
        "      elif p2_pos < 40:\n",
        "        actions[7] = 1.0\n",
        "      else:\n",
        "        actions[player2_action_index] = 1.0\n",
        "\n",
        "    next_state, rewards, done, _, info = env.step(actions)\n",
        "    next_state = preProcess(next_state)\n",
        "    next_state = np.stack((next_state, state[0], state[1], state[2]))\n",
        "\n",
        "    ball_pos.append(info['ball_x'])\n",
        "\n",
        "    if len(ball_pos) >= 3:\n",
        "\n",
        "      # If the ball was reset, then there would be a big jump from the edge of the border back to \"ball_x: 125\" position, so I said anything less than 5 is safe enough to assume the ball didn't get reset\n",
        "      if abs(ball_pos[-1] - ball_pos[-2]) < 5:\n",
        "\n",
        "        # Assuming that the ball didn't get reset, let's try and observe if there was a change in the ball's direction\n",
        "\n",
        "        # 1st condition: If the third last position is greater than the ball's second last recent position, then this means the ball is moving to the left\n",
        "        # 2nd condition: If ball_pos[-1] > ball_pos[-2] this means that the ball has properly bounced off and is now moving to the right side\n",
        "        if ball_pos[-3] > ball_pos[-2] and ball_pos[-1] > ball_pos[-2]:\n",
        "          rewards[1] = 1.0\n",
        "          rewards[0] = 0.0\n",
        "          paddle_bounces += 1\n",
        "          #print(\"player 2 will be rewarded\")\n",
        "        elif ball_pos[-3] < ball_pos[-2] and ball_pos[-1] < ball_pos[-2]:\n",
        "          rewards[1] = 0.0\n",
        "          rewards[0] = 1.0\n",
        "          paddle_bounces += 1\n",
        "          #print(\"player 1 will be rewarded\")\n",
        "        else:\n",
        "          rewards[0] = 0.0\n",
        "          rewards[1] = 0.0\n",
        "      else:\n",
        "        # print(f'jump {ball_pos[-1] - ball_pos[-2]} >= 5')\n",
        "        rewards[0] = 0.0\n",
        "        rewards[1] = 0.0\n",
        "        ball_pos = []\n",
        "        serving_times.append(frame_count - last_serve_frame_count)\n",
        "        last_serve_frame_count = frame_count\n",
        "    else:\n",
        "      rewards[0] = 0.0\n",
        "      rewards[1] = 0.0\n",
        "\n",
        "    if len(ball_pos) > 10:  # Keep only the last 10 positions\n",
        "      ball_pos.pop(0)\n",
        "\n",
        "    if len(step_rewards_p1) > 10:\n",
        "      step_rewards_p1.pop(0)\n",
        "\n",
        "    if len(step_rewards_p2) > 10:\n",
        "      step_rewards_p2.pop(0)\n",
        "\n",
        "    # Let's check if either player1 or player2 got scored on\n",
        "    show_output_print(\"penalized: \", penalized)\n",
        "\n",
        "    # If the ball makes it past 68, which is when the slider hits the ball\n",
        "    if not penalized and info['ball_x'] < 64:\n",
        "      penalized = True\n",
        "      rewards[0] = 0\n",
        "      rewards[1] = -1\n",
        "\n",
        "    elif not penalized and info['ball_x'] > 192:\n",
        "      penalized = True\n",
        "      rewards[0] = -1\n",
        "      rewards[1] = 0\n",
        "\n",
        "    # If the ball got reset, then we can penalize them again\n",
        "    elif info['ball_x'] > 120 and info['ball_x'] < 130:\n",
        "      penalized = False\n",
        "\n",
        "    show_output_print(\"This is rewards: \", rewards)\n",
        "\n",
        "    player1_agent.store_results((state, player1_action, rewards[0], next_state, done))\n",
        "    player2_agent.store_results((state, player2_action, rewards[1], next_state, done))\n",
        "\n",
        "    p1_pos = info['p1_pos']\n",
        "    p2_pos = info['p2_pos']\n",
        "    score_reward_p1 += rewards[0]\n",
        "    score_reward_p2 += rewards[1]\n",
        "\n",
        "    step_rewards_p1.append(rewards[0])\n",
        "    step_rewards_p2.append(rewards[1])\n",
        "\n",
        "    if step_count % 50 == 0:\n",
        "        p1_loss, p1_max_q_val = player1_agent.train()\n",
        "        p2_loss, p2_max_q_val = player2_agent.train()\n",
        "\n",
        "        p1_total_loss += p1_loss if type(p1_loss) == int else p1_loss.item()\n",
        "        p2_total_loss += p2_loss if type(p2_loss) == int else p2_loss.item()\n",
        "\n",
        "        p1_total_max_q_val += p1_max_q_val\n",
        "        p2_total_max_q_val += p2_max_q_val\n",
        "\n",
        "    state = next_state\n",
        "\n",
        "    step_count += 1\n",
        "\n",
        "    if show_output:\n",
        "      frame = env.render()\n",
        "      frame_image = PIL.Image.fromarray(frame)\n",
        "      #clear_output()\n",
        "      display(frame_image)\n",
        "      plot_step_rewards(step_rewards_p1, step_rewards_p2)\n",
        "      plot_rewards(episode_numbers, cumulative_rewards_p1, cumulative_rewards_p2)\n",
        "      plot_ball_pos(ball_pos)\n",
        "\n",
        "    if (episode % episode_save_interval) == 0 and frame_count % frame_sampling_rate == 0:\n",
        "      frame = env.render()\n",
        "      preview_frames.append(frame)\n",
        "    frame_count += 1\n",
        "\n",
        "    if done:\n",
        "      # score 1 and score 2 are flipped around for some reason\n",
        "      score_p1 = info['score2']\n",
        "      score_p2 = info['score1']\n",
        "      total_points = score_p1 + score_p2\n",
        "      print(f\"Episode {episode} ended. Rewards - Player 1: {score_reward_p1}, \" +\n",
        "          f\"Player 2: {score_reward_p2}, Scores - P1: {score_p1}, P2: {score_p2}\")\n",
        "      state, _ = env.reset()\n",
        "      player1_agent.update_target_to_online()\n",
        "      player2_agent.update_target_to_online()\n",
        "\n",
        "  cumulative_rewards_p1.append(score_reward_p1)\n",
        "  cumulative_rewards_p2.append(score_reward_p2)\n",
        "  if total_points != 0:\n",
        "    paddle_bounces_per_point.append(paddle_bounces / total_points)\n",
        "  episode_numbers.append(episode)\n",
        "  average_serving_times.append(sum(serving_times) / len(serving_times))\n",
        "  p1_losses.append(p1_total_loss)\n",
        "  p2_losses.append(p2_total_loss)\n",
        "  p1_avg_max_q_val.append(p1_total_max_q_val / (1 + (frame_count // 50)))\n",
        "  p2_avg_max_q_val.append(p2_total_max_q_val / (1 + (frame_count // 50)))\n",
        "\n",
        "  player1_agent.decay_epsilon()\n",
        "  player2_agent.decay_epsilon()\n",
        "\n",
        "  if (episode) % episode_save_interval == 0:\n",
        "    imageio.mimsave(f'./episodes/episode-{episode}.gif', preview_frames, fps=60)\n",
        "    torch.save(player1_agent.onlineNN.state_dict(), f'./models/p1-online-{episode}.pkl')\n",
        "    torch.save(player1_agent.targetNN.state_dict(), f'./models/p1-target-{episode}.pkl')\n",
        "    torch.save(player2_agent.onlineNN.state_dict(), f'./models/p2-online-{episode}.pkl')\n",
        "    torch.save(player2_agent.targetNN.state_dict(), f'./models/p2-target-{episode}.pkl')\n",
        "    with open(f'./models/autosave-data-{episode}.json', 'w+') as f:\n",
        "      f.write(json.dumps({\n",
        "        'cumulative_rewards_p1': cumulative_rewards_p1,\n",
        "        'cumulative_rewards_p2': cumulative_rewards_p2,\n",
        "        'episode_numbers': episode_numbers,\n",
        "        'paddle_bounces_per_point': paddle_bounces_per_point,\n",
        "        'average_serving_times': average_serving_times,\n",
        "        'p1_losses': p1_losses,\n",
        "        'p2_losses': p2_losses,\n",
        "        'p1_avg_max_q_val': p1_avg_max_q_val,\n",
        "        'p2_avg_max_q_val': p2_avg_max_q_val\n",
        "      }))\n",
        "  preview_frames = []\n",
        "\n",
        "  episode_elapsed_time = int(time.time() - episode_start_time)\n",
        "  print(f'Took {episode_elapsed_time}s, New epsilon: {player1_agent.epsilon}, Frames: {frame_count}')\n",
        "  frame_count = 0"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "editable": true,
        "id": "IA0ccJNAVlJ8",
        "tags": []
      },
      "outputs": [],
      "source": [
        "env.close()\n",
        "player1_model.save('models/player1.keras')\n",
        "player2_model.save('models/player2.keras')"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "L4",
      "provenance": []
    },
    "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.14"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}