notebooks/rl-final-project.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-12-21T14:53:11.897618Z",
     "iopub.status.busy": "2024-12-21T14:53:11.896681Z",
     "iopub.status.idle": "2024-12-21T14:53:45.257197Z",
     "shell.execute_reply": "2024-12-21T14:53:45.256261Z",
     "shell.execute_reply.started": "2024-12-21T14:53:11.897555Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "# !pip install -q git+https://github.com/Farama-Foundation/MAgent2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-12-21T14:53:45.259389Z",
     "iopub.status.busy": "2024-12-21T14:53:45.259124Z",
     "iopub.status.idle": "2024-12-21T14:53:47.156959Z",
     "shell.execute_reply": "2024-12-21T14:53:47.156034Z",
     "shell.execute_reply.started": "2024-12-21T14:53:45.259351Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F\n",
    "import numpy as np\n",
    "from collections import deque\n",
    "import random\n",
    "import torch.nn as nn\n",
    "from magent2.environments import battle_v4\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from time import time "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Q Networks\n",
    "- Phần này chứa cài đặt của các mạng Q \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-12-21T14:53:47.158495Z",
     "iopub.status.busy": "2024-12-21T14:53:47.158091Z",
     "iopub.status.idle": "2024-12-21T14:53:47.163331Z",
     "shell.execute_reply": "2024-12-21T14:53:47.162433Z",
     "shell.execute_reply.started": "2024-12-21T14:53:47.158466Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "def kaiming_init(m):\n",
    "    \"\"\"\n",
    "    Khởi tạo tham số của lớp theo Kaiming Initialization.\n",
    "    \"\"\"\n",
    "    if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):\n",
    "        nn.init.kaiming_normal_(m.weight, nonlinearity=\"relu\")  \n",
    "        if m.bias is not None:\n",
    "            nn.init.zeros_(m.bias)  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pretrained Networks\n",
    "- Trong phần này, chúng tôi cài đặt lại các mạng được sử dụng trong mô hình pretrained và final_pretraiend"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-12-21T14:53:47.165405Z",
     "iopub.status.busy": "2024-12-21T14:53:47.165063Z",
     "iopub.status.idle": "2024-12-21T14:53:47.175580Z",
     "shell.execute_reply": "2024-12-21T14:53:47.174857Z",
     "shell.execute_reply.started": "2024-12-21T14:53:47.165361Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Đây là kiến trúc pretrained được sử dụng cho red.pt \n",
    "\"\"\"\n",
    "\n",
    "\n",
    "class PretrainedQNetwork(nn.Module):\n",
    "    def __init__(self, observation_shape, action_shape):\n",
    "        super().__init__()\n",
    "        self.cnn = nn.Sequential(\n",
    "            nn.Conv2d(observation_shape[-1], observation_shape[-1], 3),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(observation_shape[-1], observation_shape[-1], 3),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "        dummy_input = torch.randn(observation_shape).permute(2, 0, 1)\n",
    "        dummy_output = self.cnn(dummy_input)\n",
    "        flatten_dim = dummy_output.view(-1).shape[0]\n",
    "        self.network = nn.Sequential(\n",
    "            nn.Linear(flatten_dim, 120),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(120, 84),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(84, action_shape),\n",
    "        )\n",
    "       # self.apply(kaiming_init)\n",
    "\n",
    "    def forward(self, x):\n",
    "        assert len(x.shape) >= 3\n",
    "        x = self.cnn(x)\n",
    "        if len(x.shape) == 3:\n",
    "            batchsize = 1\n",
    "        else:\n",
    "            batchsize = x.shape[0]\n",
    "        x = x.reshape(batchsize, -1)\n",
    "        return self.network(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-12-21T14:53:47.176553Z",
     "iopub.status.busy": "2024-12-21T14:53:47.176332Z",
     "iopub.status.idle": "2024-12-21T14:53:47.187007Z",
     "shell.execute_reply": "2024-12-21T14:53:47.186316Z",
     "shell.execute_reply.started": "2024-12-21T14:53:47.176531Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Đây là kiến trúc được cài đặt trong final_red.pt \n",
    "\"\"\"\n",
    "\n",
    "\n",
    "class Final_QNets(nn.Module):\n",
    "    def __init__(self, observation_shape, action_shape):\n",
    "        super().__init__()\n",
    "        self.cnn = nn.Sequential(\n",
    "            nn.Conv2d(observation_shape[-1], observation_shape[-1], 3),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(observation_shape[-1], observation_shape[-1], 3),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "        dummy_input = torch.randn(observation_shape).permute(2, 0, 1)\n",
    "        dummy_output = self.cnn(dummy_input)\n",
    "        flatten_dim = dummy_output.view(-1).shape[0]\n",
    "        self.network = nn.Sequential(\n",
    "            nn.Linear(flatten_dim, 120),\n",
    "            # nn.LayerNorm(120),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(120, 84),\n",
    "            # nn.LayerNorm(84),\n",
    "            nn.Tanh(),\n",
    "        )\n",
    "        self.last_layer = nn.Linear(84, action_shape)\n",
    "\n",
    "    def forward(self, x):\n",
    "        assert len(x.shape) >= 3, \"only support magent input observation\"\n",
    "        x = self.cnn(x)\n",
    "        if len(x.shape) == 3:\n",
    "            batchsize = 1\n",
    "        else:\n",
    "            batchsize = x.shape[0]\n",
    "        x = x.reshape(batchsize, -1)\n",
    "        x = self.network(x)\n",
    "        self.last_latent = x\n",
    "        return self.last_layer(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-12-21T14:53:47.188619Z",
     "iopub.status.busy": "2024-12-21T14:53:47.188240Z",
     "iopub.status.idle": "2024-12-21T14:53:47.201938Z",
     "shell.execute_reply": "2024-12-21T14:53:47.201287Z",
     "shell.execute_reply.started": "2024-12-21T14:53:47.188585Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Đây là kiến trúc mạng sử đổi (được sử dụng trong thí nghiệm 2)\n",
    "\"\"\"\n",
    "\n",
    "class MyQNetwork(nn.Module):\n",
    "    def __init__(self, observation_shape, action_shape):\n",
    "        super().__init__()\n",
    "        self.cnn = nn.Sequential(\n",
    "            nn.Conv2d(observation_shape[-1], observation_shape[-1], 3),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(observation_shape[-1], observation_shape[-1], 3),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "        dummy_input = torch.randn(observation_shape).permute(2, 0, 1)\n",
    "        dummy_output = self.cnn(dummy_input)\n",
    "        flatten_dim = dummy_output.view(-1).shape[0]\n",
    "        self.network = nn.Sequential(\n",
    "            nn.Linear(flatten_dim, 120),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(120, 84),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(84, action_shape),\n",
    "        )\n",
    "       # self.apply(kaiming_init)\n",
    "\n",
    "    def forward(self, x):\n",
    "        assert len(x.shape) >= 3\n",
    "        x = self.cnn(x)\n",
    "        if len(x.shape) == 3:\n",
    "            batchsize = 1\n",
    "        else:\n",
    "            batchsize = x.shape[0]\n",
    "        x = x.reshape(batchsize, -1)\n",
    "        return self.network(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## QMix Networks "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-12-21T14:53:47.203430Z",
     "iopub.status.busy": "2024-12-21T14:53:47.203131Z",
     "iopub.status.idle": "2024-12-21T14:53:47.218352Z",
     "shell.execute_reply": "2024-12-21T14:53:47.217736Z",
     "shell.execute_reply.started": "2024-12-21T14:53:47.203403Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Đây là kiến trúc mạng dùng chung giữa các Agent trong thuật toán QMix - tương tự như mạng Pretrained \n",
    "\"\"\"\n",
    "\n",
    "class SharedQNetwork(nn.Module):\n",
    "    def __init__(self, observation_shape, action_shape):\n",
    "        super(SharedQNetwork, self).__init__()\n",
    "        self.cnn = nn.Sequential(\n",
    "            nn.Conv2d(observation_shape[-1], observation_shape[-1], 3, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(observation_shape[-1], observation_shape[-1], 3, padding=1),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "        \n",
    "        dummy_input = torch.randn(1, observation_shape[-1], observation_shape[0], observation_shape[1])\n",
    "        dummy_output = self.cnn(dummy_input)\n",
    "        flatten_dim = dummy_output.view(-1).shape[0]\n",
    "\n",
    "        self.network = nn.Sequential(\n",
    "            nn.Linear(flatten_dim, 256),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(256, 128),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(128, action_shape)\n",
    "        )\n",
    "        self.apply(kaiming_init)\n",
    "\n",
    "    def forward(self, x):\n",
    "\n",
    "        if len(x.shape) == 4: \n",
    "            x = x.permute(0, 3, 1, 2) \n",
    "        elif len(x.shape) == 3: \n",
    "            x = x.permute(2, 0, 1).unsqueeze(0)  # [H, W, C] -> [1, C, H, W\n",
    "\n",
    "        x = self.cnn(x)\n",
    "        x = x.reshape(x.shape[0], -1)  # Flatten\n",
    "        return self.network(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-12-21T14:53:47.219808Z",
     "iopub.status.busy": "2024-12-21T14:53:47.219479Z",
     "iopub.status.idle": "2024-12-21T14:53:47.233139Z",
     "shell.execute_reply": "2024-12-21T14:53:47.232420Z",
     "shell.execute_reply.started": "2024-12-21T14:53:47.219768Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Đây là cài đặt kiến trúc mạng Mixing trong thuật toán QMix\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "class MixingNetwork(nn.Module):\n",
    "    def __init__(self, num_agents, embed_dim=32, channels=5, height=13, width=13):\n",
    "        super(MixingNetwork, self).__init__()\n",
    "        self.num_agents = num_agents\n",
    "        self.embed_dim = embed_dim\n",
    "        \n",
    "        # CNN feature extractor for states\n",
    "        self.cnn = nn.Sequential(\n",
    "            nn.Conv2d(channels, channels, kernel_size=3, stride=1, padding=1),  # [batch_size*num_agents, 16, height, width]\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(kernel_size=2),  # [batch_size*num_agents, 16, height//2, width//2]\n",
    "            nn.Conv2d(channels, channels, kernel_size=3, stride=1, padding=1),  # [batch_size*num_agents, 32, height//2, width//2]\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(kernel_size=2),  # [batch_size*num_agents, 32, height//4, width//4]\n",
    "        )\n",
    "        cnn_out_dim = (height // 4) * (width // 4) * channels  # Flattened output of CNN\n",
    "        self.fc_state = nn.Linear(cnn_out_dim, 1)  # Reduce state feature to scalar per agent\n",
    "\n",
    "        # Hyper-networks\n",
    "        self.hyper_w1 = nn.Linear(1, num_agents * embed_dim)\n",
    "        self.hyper_b1 = nn.Linear(1, embed_dim)\n",
    "        self.hyper_w2 = nn.Linear(1, embed_dim)\n",
    "        self.hyper_b2 = nn.Linear(1, 1)\n",
    "\n",
    "    def forward(self, agent_qs, states):\n",
    "        batch_size = agent_qs.size(0)\n",
    "        num_agents = agent_qs.size(1)\n",
    "        states = states.view(batch_size * num_agents, states.size(2), states.size(3), states.size(4))\n",
    "        \n",
    "        # Process states with CNN\n",
    "        states = self.cnn(states)  # [batch_size*num_agents, 32, height//4, width//4]\n",
    "        states = states.reshape(batch_size * num_agents, -1)  # Flatten to [batch_size*num_agents, cnn_out_dim]\n",
    "        states = self.fc_state(states)  # [batch_size*num_agents, 1]\n",
    "        states = states.view(batch_size, num_agents, 1)  # Reshape to [batch_size, num_agents, 1]\n",
    "\n",
    "        # Aggregate state to batch dimension\n",
    "        global_state = states.mean(dim=1)  # Mean across agents: [batch_size, 1]\n",
    "\n",
    "        # Get weights and biases from hyper-networks\n",
    "        w1 = torch.abs(self.hyper_w1(global_state)).view(batch_size, num_agents, self.embed_dim)\n",
    "        b1 = self.hyper_b1(global_state).view(batch_size, 1, self.embed_dim)\n",
    "        w2 = torch.abs(self.hyper_w2(global_state)).view(batch_size, self.embed_dim, 1)\n",
    "        b2 = self.hyper_b2(global_state).view(batch_size, 1, 1)\n",
    "\n",
    "        # Compute Mixing Network\n",
    "        hidden = torch.bmm(agent_qs.unsqueeze(1), w1) + b1  # [batch_size, 1, embed_dim]\n",
    "        hidden = F.relu(hidden)\n",
    "        q_total = torch.bmm(hidden, w2) + b2  # [batch_size, 1, 1]\n",
    "        return q_total.squeeze(-1)  # [batch_size, 1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Memory "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-12-21T14:53:47.234222Z",
     "iopub.status.busy": "2024-12-21T14:53:47.233984Z",
     "iopub.status.idle": "2024-12-21T14:53:47.245655Z",
     "shell.execute_reply": "2024-12-21T14:53:47.245036Z",
     "shell.execute_reply.started": "2024-12-21T14:53:47.234199Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Cài đặt replay buffer cho thuật toán Double Q Learning \n",
    "\"\"\"\n",
    "\n",
    "\n",
    "# Replay Buffer\n",
    "from torch.utils.data import Dataset, DataLoader \n",
    "\n",
    "class ReplayBuffer(Dataset):\n",
    "    def __init__(self, capacity):\n",
    "        self.buffer = deque(maxlen=capacity)\n",
    "\n",
    "    def push(self, state, action, reward, next_state, done):\n",
    "        self.buffer.append((state, action, reward, next_state, done))\n",
    "\n",
    "    def sample(self, batch_size):\n",
    "        batch = random.sample(self.buffer, batch_size)\n",
    "        state, action, reward, next_state, done = zip(*batch)\n",
    "        return (np.stack(state), np.array(action), np.array(reward), \n",
    "                np.stack(next_state), np.array(done))\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.buffer)\n",
    "\n",
    "\n",
    "    def __getitem__(self, idx): \n",
    "        state, action, reward, next_state, done = self.buffer[idx]\n",
    "        return (\n",
    "            torch.tensor(state), \n",
    "            torch.tensor(action), \n",
    "            torch.tensor(reward, dtype = torch.float),\n",
    "            torch.tensor(next_state), \n",
    "            torch.tensor(done, dtype = torch.float32)\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-12-21T14:53:47.248149Z",
     "iopub.status.busy": "2024-12-21T14:53:47.247890Z",
     "iopub.status.idle": "2024-12-21T14:53:47.262461Z",
     "shell.execute_reply": "2024-12-21T14:53:47.261642Z",
     "shell.execute_reply.started": "2024-12-21T14:53:47.248125Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Cài đặt StateMemory cho thuật toán QMix\n",
    "\"\"\"\n",
    "\n",
    "class StateMemory(Dataset):\n",
    "    def __init__(self, capacity, num_agents = 162, grouped_agents = 18):\n",
    "        self.capacity = capacity\n",
    "        self.memory = [deque(maxlen=capacity) for _ in range(grouped_agents)]\n",
    "        self.num_agents = num_agents \n",
    "        self.grouped_agents = grouped_agents \n",
    "\n",
    "    def push(self, idx, state, action, reward, next_state, done):\n",
    "        self.memory[idx].append((state, action, reward, next_state, done))\n",
    "\n",
    "    def sample(self, batch_size):\n",
    "        \"\"\"\n",
    "        Lấy ngẫu nhiên một batch thông tin từ bộ nhớ -> batch \n",
    "        batch được sample ra là một chuỗi hành đồng \n",
    "        \"\"\"\n",
    "        \n",
    "        batch = random.sample(self.memory, batch_size)\n",
    "        idx, state, action, reward, next_state, done = zip(*batch)\n",
    "        return (\n",
    "            np.stack(state),\n",
    "            np.array(action),\n",
    "            np.array(reward, dtype=np.float32),\n",
    "            np.stack(next_state),\n",
    "            np.array(done, dtype=np.float32)\n",
    "        )\n",
    "\n",
    "    def ensemble(self): \n",
    "        \"\"\"\n",
    "        mở rộng tất cả deque trong self.memory đến chiều dài tối đa bằng cách thêm các giá trị None vào cuối.\n",
    "        state \n",
    "        \"\"\"\n",
    "        max_len = max([len(agent_memory) for agent_memory in self.memory])\n",
    "        min_len = min([len(agent_memory) for agent_memory in self.memory])\n",
    "\n",
    "        if max_len == min_len: return \n",
    "        \n",
    "        for i in range(self.grouped_agents):\n",
    "            current_len = len(self.memory[i])\n",
    "            while current_len < max_len:\n",
    "                \n",
    "                self.memory[i].append((None, None, None, None, None))\n",
    "                current_len += 1\n",
    "\n",
    "    def __len__(self):\n",
    "        \"\"\"\n",
    "        Trả về độ dài của bộ nhớ.\n",
    "        \"\"\"\n",
    "        return min([len(i) for i in self.memory])\n",
    "\n",
    "    \n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        \"\"\"\n",
    "        Trả về dữ liệu tại một chỉ số cụ thể dưới dạng tensor cho tất cả agents.\n",
    "        \"\"\"\n",
    "        states, actions, rewards, next_states, dones = [], [], [], [], []\n",
    "\n",
    "        for i in range(self.grouped_agents):\n",
    "        \n",
    "            state, action, reward, next_state, done = self.memory[i][idx]\n",
    "\n",
    "            state = state if state is not None else np.full_like(self.memory[0][0][0], fill_value=-1)\n",
    "            action = action if action is not None else -1 \n",
    "            reward = reward if reward is not None else 0.0\n",
    "            next_state = next_state if next_state is not None else np.full_like(self.memory[0][0][0], fill_value=-1)\n",
    "            done = done if done is not None else 1.0 \n",
    "\n",
    "            # Thêm vào danh sách\n",
    "            states.append(state)\n",
    "            actions.append(action)\n",
    "            rewards.append(reward)\n",
    "            next_states.append(next_state)\n",
    "            dones.append(done)\n",
    "\n",
    "        states = np.array(states, dtype=np.float32)\n",
    "        actions = np.array(actions, dtype=np.int64)\n",
    "        rewards = np.array(rewards, dtype=np.float32)\n",
    "        next_states = np.array(next_states, dtype=np.float32)\n",
    "        dones = np.array(dones, dtype=np.float32)\n",
    "\n",
    "   \n",
    "        \n",
    "        return (\n",
    "            torch.tensor(states, dtype=torch.float32),      \n",
    "            torch.tensor(actions, dtype=torch.long),        \n",
    "            torch.tensor(rewards, dtype=torch.float32),      \n",
    "            torch.tensor(next_states, dtype=torch.float32),  \n",
    "            torch.tensor(dones, dtype=torch.float32)        \n",
    "        )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Agent"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Đây là module sử dụng để huấn luyện, thực hiện action dựa trên policy đã được huấn luyện "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-12-21T14:53:47.263713Z",
     "iopub.status.busy": "2024-12-21T14:53:47.263466Z",
     "iopub.status.idle": "2024-12-21T14:53:47.276151Z",
     "shell.execute_reply": "2024-12-21T14:53:47.275421Z",
     "shell.execute_reply.started": "2024-12-21T14:53:47.263689Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "class RandomAgent:\n",
    "    def __init__(self, action_space):\n",
    "        self.n_action = action_space\n",
    "\n",
    "    def get_action(self, observation):\n",
    "        return torch.randint(0, self.n_action, (1,)).item()  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pretrained Agent \n",
    "- cài đặt các agent với tham số được cho trước dùng để eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-12-21T14:53:47.277233Z",
     "iopub.status.busy": "2024-12-21T14:53:47.277030Z",
     "iopub.status.idle": "2024-12-21T14:53:47.287003Z",
     "shell.execute_reply": "2024-12-21T14:53:47.286240Z",
     "shell.execute_reply.started": "2024-12-21T14:53:47.277212Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "class PretrainedAgent:\n",
    "    def __init__(self, n_observation, n_actions, device=\"cpu\"):\n",
    "        self.device = torch.device(device)\n",
    "        self.qnetwork = PretrainedQNetwork(n_observation, n_actions).to(self.device)\n",
    "        self.n_action = n_actions\n",
    "        self.qnetwork.load_state_dict(\n",
    "            torch.load(\"/kaggle/input/pretrained/pytorch/default/1/red.pt\", weights_only=True, map_location=self.device)\n",
    "        )\n",
    "\n",
    "    def get_action(self, observation):\n",
    "        \n",
    "        observation = (\n",
    "            torch.Tensor(observation).float().permute([2, 0, 1]).unsqueeze(0).to(self.device)\n",
    "        )\n",
    "        with torch.no_grad():\n",
    "            q_values = self.qnetwork(observation)\n",
    "        action = torch.argmax(q_values, dim=1).cpu().numpy()[0]\n",
    "\n",
    "        return action"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-12-21T14:53:47.288056Z",
     "iopub.status.busy": "2024-12-21T14:53:47.287850Z",
     "iopub.status.idle": "2024-12-21T14:53:47.300256Z",
     "shell.execute_reply": "2024-12-21T14:53:47.299574Z",
     "shell.execute_reply.started": "2024-12-21T14:53:47.288034Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "class FinalAgent: \n",
    "    def __init__(self, n_observation, n_actions, device = \"cpu\"): \n",
    "        self.device = torch.device(device)\n",
    "\n",
    "        self.final_networks = Final_QNets(n_observation, n_actions).to(self.device)\n",
    "\n",
    "        self.final_networks.load_state_dict(\n",
    "            torch.load(\"/kaggle/input/final_rl/pytorch/default/1/red_final.pt\", weights_only = True, map_location = self.device)\n",
    "        )\n",
    "\n",
    "    def get_action(self, observation): \n",
    "        observation = (\n",
    "            torch.Tensor(observation).float().permute([2, 0, 1]).unsqueeze(0).to(self.device)\n",
    "        )\n",
    "        with torch.no_grad():\n",
    "            q_values = self.final_networks(observation)\n",
    "        action = torch.argmax(q_values, dim=1).cpu().numpy()[0]\n",
    "\n",
    "        return action\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training Agent \n",
    "- Chứ các agent được cài đặt thuật toán Double Q Learning + QMix "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-12-21T14:53:47.301671Z",
     "iopub.status.busy": "2024-12-21T14:53:47.301435Z",
     "iopub.status.idle": "2024-12-21T14:53:47.311832Z",
     "shell.execute_reply": "2024-12-21T14:53:47.310949Z",
     "shell.execute_reply.started": "2024-12-21T14:53:47.301647Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "class DQNAgent:\n",
    "    def __init__(self, observation_shape, action_shape, batch_size=64, lr=1e-3, gamma=0.6, device=\"cpu\"):\n",
    "        self.device = torch.device(device)\n",
    "        self.q_net = PretrainedQNetwork(observation_shape, action_shape).float().to(self.device)\n",
    "        self.target_net = PretrainedQNetwork(observation_shape, action_shape).float().to(self.device)\n",
    "        self.target_net.load_state_dict(self.q_net.state_dict())\n",
    "        self.target_net.eval()\n",
    "\n",
    "        self.optimizer = optim.Adam(self.q_net.parameters(), lr=lr)\n",
    "        self.batch_size = batch_size\n",
    "        self.gamma = gamma\n",
    "        self.action_shape = action_shape\n",
    "        self.epsilon = 1.0\n",
    "        self.epsilon_decay = 0.97\n",
    "        self.epsilon_min = 0.05\n",
    "        self.loss_fn = nn.MSELoss()\n",
    "    \n",
    "\n",
    "    def get_action(self, observation):\n",
    "        if np.random.rand() < self.epsilon:\n",
    "            return np.random.randint(self.action_shape)\n",
    "        else:\n",
    "            state_tensor = torch.FloatTensor(observation).unsqueeze(0).permute(0, 3, 1, 2).to(self.device)\n",
    "            with torch.no_grad():\n",
    "                return self.q_net(state_tensor).argmax().item()\n",
    "\n",
    "    def train(self, dataloader):\n",
    "        \"\"\"\n",
    "            cap nhat lai tham so mo hinh voi input dau vao \n",
    "        \"\"\"\n",
    "        self.q_net.train()\n",
    "        for obs, action, reward, next_obs, done in dataloader: \n",
    "            self.q_net.zero_grad()\n",
    "    \n",
    "            obs = obs.permute(0, 3, 1, 2).to(self.device) \n",
    "            action = action.unsqueeze(1).to(self.device)\n",
    "            reward = reward.unsqueeze(1).to(self.device)\n",
    "            next_obs = next_obs.to(self.device)\n",
    "            next_obs = next_obs.permute(0, 3, 1, 2).to(self.device)\n",
    "            done = done.unsqueeze(1).to(self.device)\n",
    "    \n",
    "            # cap nhat gia tri q \n",
    "            with torch.no_grad(): \n",
    "                target_q_values = reward + self.gamma * (1 - done) * self.target_net(next_obs).max(1, keepdim=True)[0]\n",
    "    \n",
    "            q_values = self.q_net(obs).gather(1, action)\n",
    "    \n",
    "            loss = self.loss_fn(q_values, target_q_values)\n",
    "            loss.backward()\n",
    "            self.optimizer.step()\n",
    "       \n",
    "\n",
    "    def update_target_network(self):\n",
    "        self.target_net.load_state_dict(self.q_net.state_dict())\n",
    "\n",
    "    def decay_epsilon(self):\n",
    "        self.epsilon = max(self.epsilon_min, self.epsilon * self.epsilon_decay)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-12-21T14:53:47.313181Z",
     "iopub.status.busy": "2024-12-21T14:53:47.312872Z",
     "iopub.status.idle": "2024-12-21T14:53:47.330436Z",
     "shell.execute_reply": "2024-12-21T14:53:47.329630Z",
     "shell.execute_reply.started": "2024-12-21T14:53:47.313156Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "class QMIXAgent:\n",
    "    def __init__(self, obs_shape, num_agents, action_dim, lr=1e-3, gamma=0.99, device=\"cpu\"):\n",
    "        self.device = torch.device(device)\n",
    "        self.num_agents = num_agents\n",
    "        self.action_dim = action_dim\n",
    "        \n",
    "        self.q_net = MyQNetwork(obs_shape, action_dim).to(self.device)\n",
    "        self.target_q_net = MyQNetwork(obs_shape, action_dim).to(self.device)\n",
    "        self.target_q_net.load_state_dict(self.q_net.state_dict())\n",
    "\n",
    "        self.mixing_net = MixingNetwork(num_agents).to(self.device)\n",
    "        self.target_mixing_net = MixingNetwork(num_agents).to(self.device)\n",
    "        self.target_mixing_net.load_state_dict(self.mixing_net.state_dict())\n",
    "\n",
    "        self.target_q_net.eval()\n",
    "        self.target_mixing_net.eval()\n",
    "\n",
    "        self.optimizer = optim.Adam(list(self.q_net.parameters()) + list(self.mixing_net.parameters()), lr=lr)\n",
    "        self.gamma = gamma\n",
    "        self.loss_fn = nn.MSELoss()\n",
    "\n",
    "        self.epsilon = 1.0\n",
    "        self.epsilon_decay = 0.98\n",
    "        self.epsilon_min = 0.05\n",
    "\n",
    "    def load_pretrained(self): \n",
    "        pass \n",
    "    \n",
    "\n",
    "    def get_action(self, observation):\n",
    "        if np.random.rand() < self.epsilon:\n",
    "            return np.random.randint(self.action_dim)\n",
    "        else:\n",
    "            state_tensor = torch.FloatTensor(observation).unsqueeze(0).permute(0, 3, 1, 2).to(self.device)\n",
    "            with torch.no_grad():\n",
    "                return self.q_net(state_tensor).argmax().item()\n",
    "\n",
    "\n",
    "    def train(self, dataloader):\n",
    "        self.q_net.train()\n",
    "        self.mixing_net.train()\n",
    "    \n",
    "        for obs, action, reward, next_obs, done in dataloader:\n",
    "            self.q_net.zero_grad()\n",
    "            self.mixing_net.zero_grad()\n",
    "            \n",
    "            # Đổi thứ tự permute và chuyển vào device\n",
    "            obs = obs.permute(0, 1, 4, 2, 3).to(self.device)  # [batch_size, num_agents, channels, height, width]\n",
    "            next_obs = next_obs.permute(0, 1, 4, 2, 3).to(self.device)\n",
    "            action = action.to(self.device)  # [batch_size, num_agents]\n",
    "            reward = reward.to(self.device)  # [batch_size, num_agents, 1]\n",
    "            done = done.to(self.device)  # [batch_size, num_agents, 1]\n",
    "    \n",
    "            # Xác định agent chết và sống\n",
    "            alive_agent_mask = (action != -1).float()  # [batch_size, num_agents]\n",
    "        \n",
    "            # Tính Q-values hiện tại\n",
    "            obs_flat = obs.view(-1, *obs.shape[2:])  # [batch_size * num_agents, channels, height, width]\n",
    "            obs_q_values = self.q_net(obs_flat).view(obs.size(0), obs.size(1), -1)  # [batch_size, num_agents, action_dim]\n",
    "        \n",
    "            current_q_values = obs_q_values.gather(-1, action.unsqueeze(-1).clamp(min=0)).squeeze(-1)  # [batch_size, num_agents]\n",
    "            current_q_values = current_q_values * alive_agent_mask  # Bỏ qua Q-values của agent chết\n",
    "\n",
    "            # Tính Q-values mục tiêu\n",
    "            with torch.no_grad():\n",
    "                next_obs_flat = next_obs.view(-1, *next_obs.shape[2:])  # [batch_size * num_agents, channels, height, width]\n",
    "                next_q_values = self.target_q_net(next_obs_flat).view(next_obs.size(0), next_obs.size(1), -1)  # [batch_size, num_agents, action_dim]\n",
    "    \n",
    "                max_next_q_values = next_q_values.max(dim=-1)[0]  # [batch_size, num_agents]\n",
    "                masked_next_q_values = max_next_q_values * alive_agent_mask  # Bỏ qua Q-values của agent chết\n",
    "                \n",
    "                target_q_totals = self.target_mixing_net(masked_next_q_values, next_obs)  # [batch_size, 1]\n",
    "                targets = reward.mean(dim=1, keepdim=True) + self.gamma * target_q_totals * (1 - done.mean(dim=1, keepdim=True))\n",
    "            \n",
    "            # Tính tổng Q-values hiện tại\n",
    "            current_q_totals = self.mixing_net(current_q_values, obs)  # [batch_size, 1]\n",
    "    \n",
    "            # Tính loss và cập nhật\n",
    "            loss = self.loss_fn(current_q_totals, targets)\n",
    "            loss.backward()\n",
    "            self.optimizer.step()\n",
    "\n",
    "\n",
    "       \n",
    "    def update_target_network(self):\n",
    "        self.target_q_net.load_state_dict(self.q_net.state_dict())\n",
    "        self.target_mixing_net.load_state_dict(self.mixing_net.state_dict())\n",
    "\n",
    "    def decay_epsilon(self):\n",
    "        self.epsilon = max(self.epsilon_min, self.epsilon * self.epsilon_decay)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-12-21T14:53:47.331650Z",
     "iopub.status.busy": "2024-12-21T14:53:47.331404Z",
     "iopub.status.idle": "2024-12-21T14:53:47.343875Z",
     "shell.execute_reply": "2024-12-21T14:53:47.343189Z",
     "shell.execute_reply.started": "2024-12-21T14:53:47.331628Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "# import wandb\n",
    "# from kaggle_secrets import UserSecretsClient\n",
    "\n",
    "# user_secrets = UserSecretsClient()\n",
    "# wandb_key = user_secrets.get_secret(\"wandb-key\")\n",
    "\n",
    "# wandb.login(key = wandb_key)\n",
    "\n",
    "# wandb.init(project=\"RL_TRAINING\", name=\"QMix_27\", \n",
    "#             config={\"epochs_num\": 200, \"opponents\": \"random, training with blue + red data\", \"batch_size\" : 128, \"num_agent\": 81})\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-12-21T14:53:47.345409Z",
     "iopub.status.busy": "2024-12-21T14:53:47.345155Z",
     "iopub.status.idle": "2024-12-21T14:53:47.361175Z",
     "shell.execute_reply": "2024-12-21T14:53:47.360430Z",
     "shell.execute_reply.started": "2024-12-21T14:53:47.345353Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "from time import time \n",
    "\n",
    "class Trainer : \n",
    "    \"\"\"\n",
    "    Sử dụng blue để huấn luyện \n",
    "    \n",
    "    \"\"\"\n",
    "    def __init__(self, env, red_agent, blue_agent, buffer, batch_size, is_self_play = False): \n",
    "        self.red_agent = red_agent\n",
    "        self.blue_agent = blue_agent\n",
    "        self.buffer = buffer \n",
    "        self.batch_size = batch_size \n",
    "        self.env = env \n",
    "        self.is_self_play = is_self_play \n",
    "\n",
    "    def agent_give_action(self, name: str, observation):\n",
    "        if self.is_self_play : \n",
    "            return  self.blue_agent.get_action(observation)\n",
    "        if name == \"blue\": \n",
    "            return  self.blue_agent.get_action(observation)\n",
    "        return self.red_agent.get_action(observation)\n",
    "\n",
    "\n",
    "    \n",
    "    def update_memory(self, is_longterm: bool = False): \n",
    "        \"\"\"\n",
    "        Tạo ra một vòng lặp lưu trữ và cập nhật dữ liệu cho từng agent \n",
    "        \"\"\"\n",
    "        self.env.reset()\n",
    "        prev_obs = {}\n",
    "        prev_actions = {}\n",
    "        red_reward = 0 \n",
    "        blue_reward = 0 \n",
    "\n",
    "        prev_team = \"red\"\n",
    "\n",
    "        n_kills = {\"red\": 0, \"blue\": 0}\n",
    "        # vong lap 1 \n",
    "        for idx, agent in enumerate(self.env.agent_iter()): \n",
    "            prev_ob, reward, termination, truncation, _ = self.env.last()\n",
    "            team = agent.split(\"_\")[0]\n",
    "            n_kills[team] += (reward > 4.5)\n",
    "\n",
    "            if truncation or termination: \n",
    "                prev_action = None\n",
    "            else: \n",
    "                if agent.split(\"_\")[0] == \"red\": \n",
    "                    prev_action =  self.agent_give_action(\"red\", prev_ob)\n",
    "                    red_reward += reward\n",
    "                else: \n",
    "                    prev_action = self.agent_give_action(\"blue\", prev_ob)\n",
    "                    blue_reward += reward \n",
    "    \n",
    "\n",
    "        \n",
    "            prev_obs[agent] = prev_ob \n",
    "            prev_actions[agent] = prev_action \n",
    "            self.env.step(prev_action)\n",
    "\n",
    "            if (idx + 1) % self.env.num_agents == 0: break \n",
    "\n",
    "        # vong lap 2 \n",
    "        for agent in self.env.agent_iter(): \n",
    "\n",
    "            obs, reward, termination, truncation, _ = self.env.last()\n",
    "            team = agent.split(\"_\")[0]\n",
    "            n_kills[team] += (reward > 4.5)\n",
    "            \n",
    "            if truncation or termination: \n",
    "                action = None \n",
    "            else: \n",
    "                if agent.split(\"_\")[0] == \"red\" : \n",
    "                    action = self.agent_give_action(\"red\", obs)\n",
    "                    red_reward += reward \n",
    "                \n",
    "                else: \n",
    "                    action = self.agent_give_action(\"blue\", obs)\n",
    "                    blue_reward += reward\n",
    "                \n",
    "\n",
    "            self.env.step(action)\n",
    "            if isinstance(self.buffer, StateMemory):\n",
    "                if team != prev_team : \n",
    "                    self.buffer.ensemble()  \n",
    "                    prev_team = team\n",
    "                idx = int(agent.split(\"_\")[1]) % self.buffer.grouped_agents\n",
    "                self.buffer.push(\n",
    "                    idx,\n",
    "                    prev_obs[agent], \n",
    "                    prev_actions[agent], \n",
    "                    reward, \n",
    "                    obs, \n",
    "                    termination \n",
    "                )\n",
    "            else: \n",
    "                 self.buffer.push(\n",
    "                    prev_obs[agent], \n",
    "                    prev_actions[agent], \n",
    "                    reward, \n",
    "                    obs, \n",
    "                    termination \n",
    "                )\n",
    "\n",
    "            prev_obs[agent] = obs \n",
    "            prev_actions[agent] = action\n",
    "\n",
    "        return  blue_reward - red_reward,  n_kills, blue_reward # red thắng  \n",
    "\n",
    "\n",
    "    def save_model (self, file_path):\n",
    "        \n",
    "        torch.save(self.blue_agent.q_net.state_dict(), file_path)\n",
    "        print(f\"Model saved to {file_path}\")\n",
    "    \n",
    "    def train(self, episodes=500, target_update_freq=2, is_type = \"dqn\"):\n",
    "        gap_rewards = []\n",
    "\n",
    "\n",
    "        for eps in range(episodes): \n",
    "            start = time()\n",
    "            gap_reward, n_kills, blue_reward = self.update_memory()\n",
    "\n",
    "            \n",
    "            if is_type == \"qmix\": \n",
    "                self.buffer.ensemble()\n",
    "            \n",
    "            dataloader = DataLoader(self.buffer, batch_size = self.batch_size, shuffle = True)\n",
    "            # print(f\"Out of dataloader {len(self.buffer)}\")\n",
    "            self.blue_agent.train(dataloader)\n",
    "            \n",
    "            self.blue_agent.decay_epsilon()\n",
    "            if eps % target_update_freq == 0:\n",
    "                self.blue_agent.update_target_network()\n",
    "    \n",
    "            end = time() - start \n",
    "            \n",
    "            # wandb.log({\n",
    "            #     \"episode\": eps,\n",
    "            #     \"gap_rewards\": gap_reward,\n",
    "            #     \"epsilon\": self.blue_agent.epsilon,\n",
    "            #     \"time\": end,\n",
    "            #     \"red_kill\": n_kills[\"red\"], \n",
    "            #     \"blue_kill\": n_kills[\"blue\"]\n",
    "            # })\n",
    "    \n",
    "            \n",
    "            gap_rewards.append(gap_reward)\n",
    "            print(f\"Episode {eps}, Gap Reward: {gap_reward}, Total Reward: {blue_reward}, Epsilon: {self.blue_agent.epsilon:.2f}, Time: {end}, Kill: {n_kills}\")\n",
    "    \n",
    "        self.env.close()\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training \n",
    "\n",
    "- Với thuật toán Double Deep Q -> khởi tạo đối tượng DQNAgent + ReplayBuffer\n",
    "- Với thuật toán QMix -> khởi tạo QMix Agent + StateMemory\n",
    "- Với việc huấn luyện self-play -> set tham số is_self_play = True\n",
    "  < code đang được huấn luyện cho blue agent> "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-12-21T14:53:47.362266Z",
     "iopub.status.busy": "2024-12-21T14:53:47.362035Z",
     "iopub.status.idle": "2024-12-21T15:22:09.800708Z",
     "shell.execute_reply": "2024-12-21T15:22:09.799800Z",
     "shell.execute_reply.started": "2024-12-21T14:53:47.362243Z"
    },
    "trusted": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 0, Gap Reward: -2.4750006729736924, Total Reward: -2987.9551359387115, Epsilon: 0.97, Time: 4.9963459968566895, Kill: {'red': 1, 'blue': 1}\n",
      "Episode 1, Gap Reward: 129.58500491362065, Total Reward: -2907.700127983466, Epsilon: 0.94, Time: 5.836490154266357, Kill: {'red': 4, 'blue': 2}\n",
      "Episode 2, Gap Reward: 7.53498422075063, Total Reward: -2852.0951310805976, Epsilon: 0.91, Time: 6.948957204818726, Kill: {'red': 0, 'blue': 23}\n",
      "Episode 3, Gap Reward: 174.7499812869355, Total Reward: -2502.8851289544255, Epsilon: 0.89, Time: 8.242162466049194, Kill: {'red': 3, 'blue': 35}\n",
      "Episode 4, Gap Reward: -1030.555068277754, Total Reward: -2239.6701236618683, Epsilon: 0.86, Time: 8.72317910194397, Kill: {'red': 3, 'blue': 72}\n",
      "Episode 5, Gap Reward: -1469.000080970116, Total Reward: -2340.995119580999, Epsilon: 0.83, Time: 10.10652470588684, Kill: {'red': 0, 'blue': 79}\n",
      "Episode 6, Gap Reward: -913.8050667922944, Total Reward: -1948.710114103742, Epsilon: 0.81, Time: 10.806957483291626, Kill: {'red': 3, 'blue': 81}\n",
      "Episode 7, Gap Reward: 116.37998250499368, Total Reward: -2283.695109518245, Epsilon: 0.78, Time: 12.882496356964111, Kill: {'red': 3, 'blue': 43}\n",
      "Episode 8, Gap Reward: -231.74503398314118, Total Reward: -2253.79511221312, Epsilon: 0.76, Time: 14.377876281738281, Kill: {'red': 0, 'blue': 53}\n",
      "Episode 9, Gap Reward: -519.490048231557, Total Reward: -2062.3751082951203, Epsilon: 0.74, Time: 14.90550971031189, Kill: {'red': 1, 'blue': 65}\n",
      "Episode 10, Gap Reward: -1111.3650602223352, Total Reward: -2030.990098644048, Epsilon: 0.72, Time: 16.004969358444214, Kill: {'red': 2, 'blue': 71}\n",
      "Episode 11, Gap Reward: 322.1100016562268, Total Reward: -2240.1350913923234, Epsilon: 0.69, Time: 17.943142652511597, Kill: {'red': 1, 'blue': 39}\n",
      "Episode 12, Gap Reward: -1014.5150646176189, Total Reward: -1891.7951006414369, Epsilon: 0.67, Time: 18.308514833450317, Kill: {'red': 0, 'blue': 77}\n",
      "Episode 13, Gap Reward: -67.8800292443484, Total Reward: -584.9100519781932, Epsilon: 0.65, Time: 9.969611644744873, Kill: {'red': 2, 'blue': 81}\n",
      "Episode 14, Gap Reward: -1418.6900690626353, Total Reward: -1777.8500856151804, Epsilon: 0.63, Time: 19.25377893447876, Kill: {'red': 2, 'blue': 80}\n",
      "Episode 15, Gap Reward: -397.5100276414305, Total Reward: -1898.795083269477, Epsilon: 0.61, Time: 20.99543309211731, Kill: {'red': 0, 'blue': 59}\n",
      "Episode 16, Gap Reward: -1178.2800613306463, Total Reward: -1729.4950832389295, Epsilon: 0.60, Time: 21.39218235015869, Kill: {'red': 0, 'blue': 78}\n",
      "Episode 17, Gap Reward: -635.0000422056764, Total Reward: -1667.1950819920748, Epsilon: 0.58, Time: 22.50596308708191, Kill: {'red': 0, 'blue': 72}\n",
      "Episode 18, Gap Reward: 979.565023906529, Total Reward: -1813.630077086389, Epsilon: 0.56, Time: 23.960315227508545, Kill: {'red': 1, 'blue': 38}\n",
      "Episode 19, Gap Reward: 117.9699789499864, Total Reward: -359.6000402774662, Epsilon: 0.54, Time: 11.192858934402466, Kill: {'red': 0, 'blue': 81}\n",
      "Episode 20, Gap Reward: 588.6450021583587, Total Reward: -1600.370078669861, Epsilon: 0.53, Time: 25.19992470741272, Kill: {'red': 1, 'blue': 55}\n",
      "Episode 21, Gap Reward: 157.18499444704503, Total Reward: -1577.2950691292062, Epsilon: 0.51, Time: 25.710083961486816, Kill: {'red': 1, 'blue': 56}\n",
      "Episode 22, Gap Reward: -1052.1200528834015, Total Reward: -1423.0950682824478, Epsilon: 0.50, Time: 26.0499484539032, Kill: {'red': 0, 'blue': 80}\n",
      "Episode 23, Gap Reward: 686.0450135068968, Total Reward: -1495.3950646826997, Epsilon: 0.48, Time: 27.62678861618042, Kill: {'red': 0, 'blue': 56}\n",
      "Episode 24, Gap Reward: 125.59998981375247, Total Reward: -1469.0950689706951, Epsilon: 0.47, Time: 28.631731033325195, Kill: {'red': 0, 'blue': 66}\n",
      "Episode 25, Gap Reward: -624.8750399304554, Total Reward: -1318.5750674912706, Epsilon: 0.45, Time: 28.830407857894897, Kill: {'red': 1, 'blue': 78}\n",
      "Episode 26, Gap Reward: 363.8299875855446, Total Reward: -168.1450334051624, Epsilon: 0.44, Time: 11.666045427322388, Kill: {'red': 0, 'blue': 81}\n",
      "Episode 27, Gap Reward: 147.29999667871743, Total Reward: -1298.125056937337, Epsilon: 0.43, Time: 31.057491540908813, Kill: {'red': 1, 'blue': 60}\n",
      "Episode 28, Gap Reward: 1321.3900437355042, Total Reward: -1430.7950545670465, Epsilon: 0.41, Time: 32.28914976119995, Kill: {'red': 0, 'blue': 42}\n",
      "Episode 29, Gap Reward: 338.0499880509451, Total Reward: 166.2349794888869, Epsilon: 0.40, Time: 7.6544904708862305, Kill: {'red': 0, 'blue': 81}\n",
      "Episode 30, Gap Reward: -449.0400320990011, Total Reward: -1110.6600578865036, Epsilon: 0.39, Time: 31.36219596862793, Kill: {'red': 1, 'blue': 79}\n",
      "Episode 31, Gap Reward: 489.2949929870665, Total Reward: 151.1749782692641, Epsilon: 0.38, Time: 8.304598808288574, Kill: {'red': 3, 'blue': 81}\n",
      "Episode 32, Gap Reward: 440.5950013138354, Total Reward: -1023.3600525120273, Epsilon: 0.37, Time: 32.68964385986328, Kill: {'red': 1, 'blue': 78}\n",
      "Episode 33, Gap Reward: 782.3600204084069, Total Reward: -1148.7850495772436, Epsilon: 0.36, Time: 33.895774841308594, Kill: {'red': 1, 'blue': 54}\n",
      "Episode 34, Gap Reward: 470.84500304050744, Total Reward: -997.2950506778434, Epsilon: 0.34, Time: 34.62361764907837, Kill: {'red': 0, 'blue': 72}\n",
      "Episode 35, Gap Reward: -223.10002576746047, Total Reward: -439.0550352046266, Epsilon: 0.33, Time: 22.806530952453613, Kill: {'red': 0, 'blue': 81}\n",
      "Episode 36, Gap Reward: 1502.1750550456345, Total Reward: -1174.2950401110575, Epsilon: 0.32, Time: 35.37240386009216, Kill: {'red': 0, 'blue': 34}\n",
      "Episode 37, Gap Reward: 612.4199963137507, Total Reward: 210.55998054891825, Epsilon: 0.31, Time: 9.036760091781616, Kill: {'red': 0, 'blue': 81}\n",
      "Episode 38, Gap Reward: 529.500007333234, Total Reward: -928.2950453525409, Epsilon: 0.30, Time: 36.205015659332275, Kill: {'red': 0, 'blue': 68}\n",
      "Episode 39, Gap Reward: 1162.5450398717076, Total Reward: -1022.0950376586989, Epsilon: 0.30, Time: 36.237213134765625, Kill: {'red': 0, 'blue': 47}\n",
      "Episode 40, Gap Reward: -265.89502014778554, Total Reward: -811.5950408056378, Epsilon: 0.29, Time: 35.734429359436035, Kill: {'red': 0, 'blue': 76}\n",
      "Episode 41, Gap Reward: 397.250007272698, Total Reward: -871.7950383108109, Epsilon: 0.28, Time: 36.600464820861816, Kill: {'red': 0, 'blue': 64}\n",
      "Episode 42, Gap Reward: 374.7950021540746, Total Reward: -768.4950397722423, Epsilon: 0.27, Time: 37.066519021987915, Kill: {'red': 0, 'blue': 74}\n",
      "Episode 43, Gap Reward: 59.70999046880752, Total Reward: -700.1950381100178, Epsilon: 0.26, Time: 36.90523672103882, Kill: {'red': 1, 'blue': 79}\n",
      "Episode 44, Gap Reward: 591.1050113942474, Total Reward: -747.8950370429084, Epsilon: 0.25, Time: 37.62710452079773, Kill: {'red': 0, 'blue': 68}\n",
      "Episode 45, Gap Reward: 1327.060042518191, Total Reward: -782.0950328391045, Epsilon: 0.25, Time: 38.55983018875122, Kill: {'red': 0, 'blue': 58}\n",
      "Episode 46, Gap Reward: 587.384995194152, Total Reward: 397.3849868979305, Epsilon: 0.24, Time: 4.8294477462768555, Kill: {'red': 0, 'blue': 81}\n",
      "Episode 47, Gap Reward: 1862.5550714023411, Total Reward: -1156.0950353778899, Epsilon: 0.23, Time: 39.50113582611084, Kill: {'red': 0, 'blue': 28}\n",
      "Episode 48, Gap Reward: 496.35499208047986, Total Reward: 271.38998313900083, Epsilon: 0.22, Time: 7.837052583694458, Kill: {'red': 0, 'blue': 81}\n",
      "Episode 49, Gap Reward: -61.10001493990421, Total Reward: -873.1950450148433, Epsilon: 0.22, Time: 39.326772689819336, Kill: {'red': 0, 'blue': 80}\n",
      "Episode 50, Gap Reward: 747.5300035197288, Total Reward: 108.31997974961996, Epsilon: 0.21, Time: 15.441362380981445, Kill: {'red': 0, 'blue': 81}\n",
      "Episode 51, Gap Reward: 365.6499915448949, Total Reward: 145.50998246576637, Epsilon: 0.21, Time: 14.83320951461792, Kill: {'red': 0, 'blue': 81}\n",
      "Episode 52, Gap Reward: 1673.9500640658662, Total Reward: -787.6950232917443, Epsilon: 0.20, Time: 40.901952505111694, Kill: {'red': 0, 'blue': 37}\n",
      "Episode 53, Gap Reward: 537.6249941485003, Total Reward: 337.1949857864529, Epsilon: 0.19, Time: 7.979888916015625, Kill: {'red': 0, 'blue': 81}\n",
      "Episode 54, Gap Reward: -209.98001596983522, Total Reward: -578.9950301619247, Epsilon: 0.19, Time: 40.62567377090454, Kill: {'red': 0, 'blue': 77}\n",
      "Episode 55, Gap Reward: 12.769989766180515, Total Reward: -477.6750289462507, Epsilon: 0.18, Time: 40.33002305030823, Kill: {'red': 1, 'blue': 80}\n",
      "Episode 56, Gap Reward: 734.8950037462637, Total Reward: 111.27498066145927, Epsilon: 0.18, Time: 18.207478284835815, Kill: {'red': 0, 'blue': 81}\n",
      "Episode 57, Gap Reward: 2220.7250880459324, Total Reward: -764.8950176630169, Epsilon: 0.17, Time: 42.41295504570007, Kill: {'red': 0, 'blue': 27}\n",
      "Episode 58, Gap Reward: 593.1999955587089, Total Reward: 409.0949879921973, Epsilon: 0.17, Time: 6.6185431480407715, Kill: {'red': 0, 'blue': 81}\n",
      "Episode 59, Gap Reward: 668.5700016766787, Total Reward: -125.57502711191773, Epsilon: 0.16, Time: 20.50230050086975, Kill: {'red': 0, 'blue': 81}\n",
      "Episode 60, Gap Reward: 141.96499657444656, Total Reward: -360.0950224297121, Epsilon: 0.16, Time: 42.002259731292725, Kill: {'red': 0, 'blue': 78}\n",
      "Episode 61, Gap Reward: 692.8100001476705, Total Reward: 313.4999853875488, Epsilon: 0.15, Time: 12.730149745941162, Kill: {'red': 0, 'blue': 81}\n",
      "Episode 62, Gap Reward: 371.12500740773976, Total Reward: -338.92002014350146, Epsilon: 0.15, Time: 41.1352276802063, Kill: {'red': 3, 'blue': 71}\n",
      "Episode 63, Gap Reward: 577.3799946084619, Total Reward: 436.64498848933727, Epsilon: 0.14, Time: 5.496417999267578, Kill: {'red': 0, 'blue': 81}\n",
      "Episode 64, Gap Reward: -4434.450198738836, Total Reward: -6136.495259420946, Epsilon: 0.14, Time: 43.222917795181274, Kill: {'red': 0, 'blue': 54}\n",
      "Episode 65, Gap Reward: 341.8349902834743, Total Reward: 241.1199853438884, Epsilon: 0.13, Time: 14.179447650909424, Kill: {'red': 0, 'blue': 81}\n",
      "Episode 66, Gap Reward: 1081.2100359853357, Total Reward: -403.6950173266232, Epsilon: 0.13, Time: 43.57239603996277, Kill: {'red': 0, 'blue': 62}\n",
      "Episode 67, Gap Reward: 310.6750030172989, Total Reward: -272.4950186153874, Epsilon: 0.13, Time: 44.0355806350708, Kill: {'red': 0, 'blue': 79}\n",
      "Episode 68, Gap Reward: 655.5050177657977, Total Reward: -319.795017356053, Epsilon: 0.12, Time: 44.02384090423584, Kill: {'red': 0, 'blue': 71}\n",
      "Episode 69, Gap Reward: 207.68499968759716, Total Reward: -260.7950176782906, Epsilon: 0.12, Time: 44.45601224899292, Kill: {'red': 0, 'blue': 79}\n"
     ]
    }
   ],
   "source": [
    "env = battle_v4.env(map_size=45, render_mode=\"rgb_array\", attack_opponent_reward=0.5)\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "\n",
    "observation_shape = env.observation_space(\"red_0\").shape\n",
    "action_shape = env.action_space(\"red_0\").n\n",
    "num_agents = 27\n",
    "\n",
    "blue_agent = DQNAgent(observation_shape,action_shape, device=device)\n",
    "\n",
    "# blue_agent = QMIXAgent(observation_shape, num_agents , action_shape, device=device)\n",
    "\n",
    "# blue_agent = PretrainedAgent(n_observation = observation_shape, n_actions = action_shape, device = device)\n",
    "red_agent = RandomAgent(action_shape)\n",
    "buffer = ReplayBuffer(capacity=10000)\n",
    "# buffer = StateMemory(capacity = 10000, grouped_agents = num_agents)\n",
    "\n",
    "trainer = Trainer(env, red_agent, blue_agent, buffer, batch_size = 64, is_self_play=False)\n",
    "trainer.train(episodes = 70)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-12-21T15:22:09.802399Z",
     "iopub.status.busy": "2024-12-21T15:22:09.801989Z",
     "iopub.status.idle": "2024-12-21T15:22:09.810834Z",
     "shell.execute_reply": "2024-12-21T15:22:09.810026Z",
     "shell.execute_reply.started": "2024-12-21T15:22:09.802357Z"
    },
    "trusted": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model saved to my_final3.pt\n"
     ]
    }
   ],
   "source": [
    "trainer.save_model(\"my_final.pt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Eval "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-12-21T15:26:24.147980Z",
     "iopub.status.busy": "2024-12-21T15:26:24.147256Z",
     "iopub.status.idle": "2024-12-21T15:26:24.154023Z",
     "shell.execute_reply": "2024-12-21T15:26:24.153122Z",
     "shell.execute_reply.started": "2024-12-21T15:26:24.147940Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "class TestQAgent: \n",
    "\n",
    "    def __init__(self, n_observation, n_actions, model_path: str): \n",
    "        self.qnetwork = PretrainedQNetwork(n_observation, n_actions)\n",
    "        self.n_action = n_actions\n",
    "        self.qnetwork.load_state_dict(\n",
    "            torch.load(model_path, weights_only=True, map_location=\"cpu\")\n",
    "        ) \n",
    "\n",
    "    def get_action(self, observation):\n",
    "\n",
    "        if np.random.rand() < 0.05:\n",
    "            return np.random.randint(self.n_action)\n",
    "        else:\n",
    "            observation = (\n",
    "                        torch.Tensor(observation).float().permute([2, 0, 1]).unsqueeze(0)\n",
    "                    )\n",
    "            with torch.no_grad():\n",
    "                q_values = self.qnetwork(observation)\n",
    "                action = torch.argmax(q_values, dim=1).numpy()[0]\n",
    "\n",
    "        return action\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-12-21T15:26:27.241278Z",
     "iopub.status.busy": "2024-12-21T15:26:27.240967Z",
     "iopub.status.idle": "2024-12-21T15:26:27.256829Z",
     "shell.execute_reply": "2024-12-21T15:26:27.255955Z",
     "shell.execute_reply.started": "2024-12-21T15:26:27.241250Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "def eval():\n",
    "    max_cycles = 300\n",
    "    env = battle_v4.env(map_size=45, max_cycles=max_cycles)\n",
    "    device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "\n",
    "    def random_policy(env, agent, obs):\n",
    "        return env.action_space(agent).sample()\n",
    "    \n",
    "    \n",
    "    q_network = PretrainedQNetwork(\n",
    "        env.observation_space(\"red_0\").shape, env.action_space(\"red_0\").n\n",
    "    )\n",
    "    q_network.load_state_dict(\n",
    "        torch.load(\"/kaggle/input/pretrained/pytorch/default/1/red.pt\", weights_only=True, map_location=\"cpu\")\n",
    "    )\n",
    "    q_network.to(device)\n",
    "\n",
    "    final_q_network = Final_QNets(\n",
    "        env.observation_space(\"red_0\").shape, env.action_space(\"red_0\").n\n",
    "    )\n",
    "    final_q_network.load_state_dict(\n",
    "        torch.load(\"/kaggle/input/final_rl/pytorch/default/1/red_final.pt\", weights_only=True, map_location=\"cpu\")\n",
    "    )\n",
    "    final_q_network.to(device)\n",
    "\n",
    "    def my_policy(env, agent, obs):\n",
    "        my_agent = TestQAgent(env.observation_space(\"red_0\").shape,  env.action_space(\"red_0\").n, model_path= '/kaggle/working/my_final.pt')\n",
    "        return my_agent.get_action(obs)\n",
    "\n",
    "\n",
    "    def pretrain_policy(env, agent, obs):\n",
    "        observation = (\n",
    "            torch.Tensor(obs).float().permute([2, 0, 1]).unsqueeze(0).to(device)\n",
    "        )\n",
    "        with torch.no_grad():\n",
    "            q_values = q_network(observation)\n",
    "        return torch.argmax(q_values, dim=1).cpu().numpy()[0]\n",
    "\n",
    "    def final_pretrain_policy(env, agent, obs):\n",
    "        observation = (\n",
    "            torch.Tensor(obs).float().permute([2, 0, 1]).unsqueeze(0).to(device)\n",
    "        )\n",
    "        with torch.no_grad():\n",
    "            q_values = final_q_network(observation)\n",
    "        return torch.argmax(q_values, dim=1).cpu().numpy()[0]\n",
    "\n",
    "    def run_eval(env, red_policy, blue_policy, n_episode: int = 100):\n",
    "        red_win, blue_win = [], []\n",
    "        red_tot_rw, blue_tot_rw = [], []\n",
    "        n_agent_each_team = len(env.env.action_spaces) // 2\n",
    "        blue_agents = []\n",
    "        red_agents = []\n",
    "\n",
    "        for _ in tqdm(range(n_episode)):\n",
    "            env.reset()\n",
    "            n_kill = {\"red\": 0, \"blue\": 0}\n",
    "            red_reward, blue_reward = 0, 0\n",
    "\n",
    "            for agent in env.agent_iter():\n",
    "                observation, reward, termination, truncation, info = env.last()\n",
    "                agent_team = agent.split(\"_\")[0]\n",
    "\n",
    "                n_kill[agent_team] += (\n",
    "                    reward > 4.5\n",
    "                )  \n",
    "                if agent_team == \"red\":\n",
    "                    red_reward += reward\n",
    "                else:\n",
    "                    blue_reward += reward\n",
    "\n",
    "                if termination or truncation:\n",
    "                    action = None  \n",
    "                else:\n",
    "                    if agent_team == \"red\":\n",
    "                        action = red_policy(env, agent, observation)\n",
    "                    else:\n",
    "                        action = blue_policy(env, agent, observation)\n",
    "\n",
    "                env.step(action)\n",
    "\n",
    "            who_wins = \"red\" if n_kill[\"red\"] >= n_kill[\"blue\"] + 5 else \"draw\"\n",
    "            who_wins = \"blue\" if n_kill[\"red\"] + 5 <= n_kill[\"blue\"] else who_wins\n",
    "            red_win.append(who_wins == \"red\")\n",
    "            blue_win.append(who_wins == \"blue\")\n",
    "\n",
    "            blue_agents.append(n_kill[\"blue\"])\n",
    "            red_agents.append(n_kill[\"red\"])\n",
    "\n",
    "            red_tot_rw.append(red_reward / n_agent_each_team)\n",
    "            blue_tot_rw.append(blue_reward / n_agent_each_team)\n",
    "\n",
    "        return {\n",
    "            \"winrate_red\": np.mean(red_win),\n",
    "            \"winrate_blue\": np.mean(blue_win),\n",
    "            \"average_rewards_red\": np.mean(red_tot_rw),\n",
    "            \"average_rewards_blue\": np.mean(blue_tot_rw),\n",
    "            \"red_kill\": np.mean(red_agents) / n_agent_each_team,\n",
    "            \"blue_kill\": np.mean(blue_agents) / n_agent_each_team,\n",
    "        }\n",
    "\n",
    "    print(\"=\" * 20)\n",
    "    print(\"Eval with random policy\")\n",
    "    print(\n",
    "        run_eval(\n",
    "            env=env, red_policy=random_policy, blue_policy=my_policy, n_episode=30\n",
    "        )\n",
    "    )\n",
    "    print(\"=\" * 20)\n",
    "\n",
    "    print(\"Eval with trained policy\")\n",
    "    print(\n",
    "        run_eval(\n",
    "            env=env, red_policy=pretrain_policy, blue_policy=my_policy, n_episode=30\n",
    "        )\n",
    "    )\n",
    "    print(\"=\" * 20)\n",
    "\n",
    "    print(\"Eval with final trained policy\")\n",
    "    print(\n",
    "        run_eval(\n",
    "            env=env,\n",
    "            red_policy=final_pretrain_policy,\n",
    "            blue_policy=my_policy,\n",
    "            n_episode=30,\n",
    "        )\n",
    "    )\n",
    "    print(\"=\" * 20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "trusted": true
   },
   "outputs": [],
   "source": [
    "eval()"
   ]
  }
 ],
 "metadata": {
  "kaggle": {
   "accelerator": "none",
   "dataSources": [
    {
     "isSourceIdPinned": true,
     "modelId": 186056,
     "modelInstanceId": 163700,
     "sourceId": 192015,
     "sourceType": "modelInstanceVersion"
    },
    {
     "isSourceIdPinned": true,
     "modelId": 192823,
     "modelInstanceId": 170510,
     "sourceId": 199886,
     "sourceType": "modelInstanceVersion"
    }
   ],
   "dockerImageVersionId": 30805,
   "isGpuEnabled": false,
   "isInternetEnabled": true,
   "language": "python",
   "sourceType": "notebook"
  },
  "kernelspec": {
   "display_name": "rl_env",
   "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}