dqn.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e93345a-1e2d-45c8-bc8b-10a6f2b184a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import deque\n",
    "import math\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import random\n",
    "import pickle\n",
    "import glob\n",
    "\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "\n",
    "import wandb\n",
    "\n",
    "from tetris.Environment import TetrisEnv\n",
    "\n",
    "import util.decaying\n",
    "import dqnmodel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "949a4f03-ee16-42eb-9491-ac12f441aa3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "PRO_REPLAY_DIRECTORY = 'pro-replays'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "134f63ef-8006-4b90-b20d-1759ac10dd4c",
   "metadata": {},
   "source": [
    "### Hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9468f036-0d54-4303-9e1a-55f92999006d",
   "metadata": {},
   "outputs": [],
   "source": [
    "MEM_SIZE = 10000\n",
    "MIN_MEM_SIZE = 1000\n",
    "\n",
    "DISCOUNT_START = 0.8\n",
    "DISCOUNT_END = 0.94\n",
    "DISCOUNT_DURATION = 4000\n",
    "\n",
    "EPSILON_START = 0.5\n",
    "EPSILON_END = 0.08\n",
    "EPSILON_DURATION = 3000\n",
    "\n",
    "UPDATE_TARGET_EVERY = 100\n",
    "SIMULATE_EVERY = 5\n",
    "USE_PRO_PLAY_CHANCE = 0.2\n",
    "\n",
    "EPISODES = 6000\n",
    "BATCH_SIZE = 164\n",
    "\n",
    "LEARNING_RATE_START = 3e-3\n",
    "LEARNING_RATE_GAMMA = 0.9\n",
    "LEARNING_RATE_STEP = 300\n",
    "\n",
    "use_pro_replays = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a4ee64b-a680-4610-aa11-4d22082ce021",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14b38626-0551-4689-866f-3b777384bdea",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = dqnmodel.DQNModel(UPDATE_TARGET_EVERY)\n",
    "model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd8766de-e03e-4c04-be72-b560c34d89fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "epsilon = util.decaying.DecayingLinear(EPSILON_START, EPSILON_END, EPSILON_DURATION)\n",
    "discount = util.decaying.DecayingDiscount(DISCOUNT_START, DISCOUNT_END, DISCOUNT_DURATION)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03f70331-d5ba-4924-ae12-4a0c989f638c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_best_state(states, use_epsilon=True):\n",
    "    if not use_epsilon or random.random() > epsilon.get():\n",
    "        # use the q-network (not the target network) for chosing the next state\n",
    "        q_values = model.model(states)\n",
    "        return torch.argmax(q_values)\n",
    "    else:\n",
    "        return random.choice(range(len(states)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a97a3758-3b69-46ee-b570-14d115f925f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "env = TetrisEnv()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f5c250e-694d-4bfd-9896-ea3865931f30",
   "metadata": {},
   "source": [
    "### Fill the replay buffer by playing games"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e8aac1f-362b-4c46-8134-5c946e027b73",
   "metadata": {},
   "outputs": [],
   "source": [
    "replay_buffer = []\n",
    "\n",
    "def to_torch(state):\n",
    "    return torch.from_numpy(states.reshape(states.shape[0], -1)).float()\n",
    "\n",
    "with tqdm(total=MIN_MEM_SIZE/20) as pbar:\n",
    "    while len(replay_buffer) < MIN_MEM_SIZE:\n",
    "        env.reset()\n",
    "        pbar.update(1)\n",
    "\n",
    "        # play moves until game over\n",
    "        while True:\n",
    "            states, scores, clears, dones = env.get_next_states()\n",
    "\n",
    "            chosen_index = get_best_state(torch.from_numpy(states.reshape(-1, 1, 20, 10)).float().to(device))\n",
    "\n",
    "            replay_buffer.append((env.get_current_state(), states[chosen_index], scores[chosen_index], dones[chosen_index]))\n",
    "\n",
    "            if dones[chosen_index]:\n",
    "                break\n",
    "            else:\n",
    "                env.step(states[chosen_index], clears[chosen_index], scores[chosen_index])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d022535-9ad3-4831-bb19-66735b324dde",
   "metadata": {},
   "source": [
    "### Load pro-player replays"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f122d02d-9d11-49e8-a67e-97b77edf3e0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ProReplayDataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, path, train):\n",
    "        self.path = path\n",
    "        self.file_list = glob.glob(f'{path}/*.json')\n",
    "        self.train = train\n",
    "        \n",
    "        def string_to_board(string):\n",
    "            return torch.tensor([int(c) for c in list(string)])\n",
    "        \n",
    "        # to speed things up, load the whole dataset to memory\n",
    "        self.buffer = []\n",
    "        for idx, file in tqdm(enumerate(self.file_list)):\n",
    "            df = pd.read_csv(file)\n",
    "\n",
    "            for index, row in df.iterrows():\n",
    "                self.buffer.append((\n",
    "                    string_to_board(row.current).reshape((20, 10)).to(device),\n",
    "                    string_to_board(row.next).reshape((20, 10)).to(device),\n",
    "                    row.score,\n",
    "                    row.done\n",
    "                ))\n",
    "      \n",
    "    def __len__(self):\n",
    "        return len(self.buffer)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return self.buffer[idx]\n",
    "\n",
    "\n",
    "replay_buffer_dataset = ProReplayDataset(PRO_REPLAY_DIRECTORY, True)\n",
    "replay_buffer_loader = torch.utils.data.DataLoader(dataset=replay_buffer_dataset, batch_size=BATCH_SIZE, shuffle=True)\n",
    "replay_buffer_iter = iter(replay_buffer_loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "637c2fa0-925c-49d9-8986-ffaf8b341fef",
   "metadata": {},
   "source": [
    "### Training loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa651a19-cb69-427d-acbc-77b2f344eec8",
   "metadata": {},
   "outputs": [],
   "source": [
    "wandb.init(project='tetris-dqn', config={\n",
    "    'learning-rate-start': LEARNING_RATE_START,\n",
    "    'learning-rate-gamma': LEARNING_RATE_GAMMA,\n",
    "    'learning-rate-step': LEARNING_RATE_STEP,\n",
    "    \n",
    "    'batch-size': BATCH_SIZE,\n",
    "    \n",
    "    'replay-max-size': MEM_SIZE,\n",
    "    'replay-min-size': MIN_MEM_SIZE,\n",
    "    \n",
    "    'epsilon-start': EPSILON_START,\n",
    "    'epsilon-end': EPSILON_END,\n",
    "    'epsilon-duration': EPSILON_DURATION,\n",
    "    \n",
    "    'discount-start': DISCOUNT_START,\n",
    "    'discount-end': DISCOUNT_END,\n",
    "    'discount-duration': DISCOUNT_DURATION,\n",
    "    \n",
    "    'pro-play-chance': USE_PRO_PLAY_CHANCE,\n",
    "    'simulate-every': SIMULATE_EVERY,\n",
    "    'update-target-every': UPDATE_TARGET_EVERY,\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c95226ce-c7bb-4def-a805-11748c9b7a87",
   "metadata": {},
   "outputs": [],
   "source": [
    "epsilon = util.decaying.DecayingLinear(EPSILON_START, EPSILON_END, EPSILON_DURATION)\n",
    "discount = util.decaying.DecayingDiscount(DISCOUNT_START, DISCOUNT_END, DISCOUNT_DURATION)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "538048ff-be42-42dd-9bac-5b209ab8853e",
   "metadata": {},
   "outputs": [],
   "source": [
    "criterion = nn.MSELoss() # HuberLoss()\n",
    "optimizer = torch.optim.AdamW(model.model.parameters(), lr=LEARNING_RATE_START)\n",
    "scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=LEARNING_RATE_STEP, gamma=LEARNING_RATE_GAMMA, verbose=False)\n",
    "\n",
    "training_loss = []\n",
    "training_scores = []\n",
    "\n",
    "model.model.train()\n",
    "\n",
    "for episode in tqdm(range(EPISODES)):\n",
    "    use_pro_replays = random.random() < USE_PRO_PLAY_CHANCE\n",
    "    \n",
    "    # play another game\n",
    "    if episode % SIMULATE_EVERY == 1:\n",
    "        env.reset()\n",
    "        while True:\n",
    "            states, scores, clears, dones = env.get_next_states()\n",
    "\n",
    "            chosen_index = get_best_state(torch.from_numpy(states.reshape(-1, 1, 20, 10)).float().to(device))\n",
    "\n",
    "            replay_buffer.append((env.get_current_state(), states[chosen_index], scores[chosen_index], dones[chosen_index]))\n",
    "\n",
    "            if dones[chosen_index]:\n",
    "                training_scores.append({'epoch': episode, 'score': env.score})\n",
    "                break\n",
    "            else:\n",
    "                env.step(states[chosen_index], clears[chosen_index], scores[chosen_index])\n",
    "\n",
    "        if len(replay_buffer) > MEM_SIZE:\n",
    "            replay_buffer = replay_buffer[int(MEM_SIZE/10):]\n",
    "            \n",
    "        wandb.log({'game/score': env.score,\n",
    "                   'game/singles': env.clears[0], 'game/doubles': env.clears[1], 'game/triples': env.clears[2], 'game/quads': env.clears[3],\n",
    "                   'game/tspins': env.tspins, 'game/all_clears': env.all_clears, 'game/moves': env.moves })\n",
    "\n",
    "\n",
    "    # get the batch, consisting of (current_state, next_state, score, done), and extract current and next states\n",
    "    if use_pro_replays:\n",
    "        batch = next(replay_buffer_iter, None)\n",
    "        # if we took all entries of the pro-player replays, start over again\n",
    "        if batch is None:\n",
    "            replay_buffer_iter = iter(replay_buffer_loader)\n",
    "            batch = next(replay_buffer_iter, None)\n",
    "            \n",
    "        current_states = batch[0].reshape(-1, 1, 20, 10).float()\n",
    "        next_states = batch[0].reshape(-1, 1, 20, 10).float()\n",
    "    else:\n",
    "        # take sample from replay memory\n",
    "        batch = random.sample(replay_buffer, BATCH_SIZE)\n",
    "\n",
    "        current_states = torch.from_numpy(np.array([s[0] for s in batch])).reshape(-1, 1, 20, 10).float().to(device)\n",
    "        next_states = torch.from_numpy(np.array([s[1] for s in batch])).reshape(-1, 1, 20, 10).float().to(device)\n",
    "\n",
    "    # get the q-values of the current state\n",
    "    y_hat = model.model(current_states)\n",
    "\n",
    "    # calculate expected q-values of the next state using the target-network\n",
    "    next_q_values = model.target_model(next_states)\n",
    "    y = []\n",
    "    if use_pro_replays:\n",
    "        for i in range(batch[3].shape[0]):\n",
    "            done = batch[3][i]\n",
    "            score = batch[2][i].float()\n",
    "            \n",
    "            if not done:\n",
    "                new_q = score + discount.get() * next_q_values[i]\n",
    "            else:\n",
    "                new_q = score\n",
    "            \n",
    "            y.append(new_q)\n",
    "    else:\n",
    "        for i, (_, _, score, done) in enumerate(batch):\n",
    "            if not done:\n",
    "                new_q = score + discount.get() * next_q_values[i]\n",
    "            else:\n",
    "                new_q = score\n",
    "\n",
    "            y.append(new_q)\n",
    "\n",
    "    \n",
    "            \n",
    "    # fit the model to the expected q value\n",
    "    loss = criterion(y_hat, torch.tensor(y).reshape(BATCH_SIZE if not use_pro_replays else batch[3].shape[0], 1).to(device))\n",
    "\n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "    epsilon.step()\n",
    "    discount.step()\n",
    "    scheduler.step()\n",
    "    \n",
    "    model.step()\n",
    "\n",
    "    wandb.log({'training/loss': loss.item()})\n",
    "    training_loss.append({'epoch': episode, 'loss': loss.item()})\n",
    "    \n",
    "    if episode % 100 == 0:\n",
    "        torch.save(model.model.state_dict(), f'models/run-cnn-{math.floor(episode / 100) + 105}.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6c0fec7-86dd-444d-94db-601795e62697",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model.model.state_dict(), 'models/run-cnn-after-18000.pt')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}