dcgan_atari_train.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "## Required Installations"
      ],
      "metadata": {
        "id": "XwjDwtgrv0j-"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "bSr-eSN3-dYP",
        "outputId": "a890dcbf-6a55-4277-ff8e-e253d29bc992"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: gym[atari] in /usr/local/lib/python3.7/dist-packages (0.17.3)\n",
            "Requirement already satisfied: scipy in /usr/local/lib/python3.7/dist-packages (from gym[atari]) (1.4.1)\n",
            "Requirement already satisfied: cloudpickle<1.7.0,>=1.2.0 in /usr/local/lib/python3.7/dist-packages (from gym[atari]) (1.3.0)\n",
            "Requirement already satisfied: numpy>=1.10.4 in /usr/local/lib/python3.7/dist-packages (from gym[atari]) (1.19.5)\n",
            "Requirement already satisfied: pyglet<=1.5.0,>=1.4.0 in /usr/local/lib/python3.7/dist-packages (from gym[atari]) (1.5.0)\n",
            "Requirement already satisfied: Pillow in /usr/local/lib/python3.7/dist-packages (from gym[atari]) (7.1.2)\n",
            "Requirement already satisfied: atari-py~=0.2.0 in /usr/local/lib/python3.7/dist-packages (from gym[atari]) (0.2.9)\n",
            "Requirement already satisfied: opencv-python in /usr/local/lib/python3.7/dist-packages (from gym[atari]) (4.1.2.30)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.7/dist-packages (from atari-py~=0.2.0->gym[atari]) (1.15.0)\n",
            "Requirement already satisfied: future in /usr/local/lib/python3.7/dist-packages (from pyglet<=1.5.0,>=1.4.0->gym[atari]) (0.16.0)\n"
          ]
        }
      ],
      "source": [
        "!pip install gym[atari]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OnQGUte-dJlF",
        "outputId": "614cb392-a9da-4d80-a224-35e4e478087c"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: opencv-python in /usr/local/lib/python3.7/dist-packages (4.1.2.30)\n",
            "Requirement already satisfied: numpy>=1.14.5 in /usr/local/lib/python3.7/dist-packages (from opencv-python) (1.19.5)\n"
          ]
        }
      ],
      "source": [
        "!pip install opencv-python"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4mPmD14dhAwE",
        "outputId": "2bccfc4e-4996-456d-d048-74db2fda8f21"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: unrar in /usr/local/lib/python3.7/dist-packages (0.4)\n",
            "\n",
            "UNRAR 5.50 freeware      Copyright (c) 1993-2017 Alexander Roshal\n",
            "\n",
            "\n",
            "Extracting from Roms.rar\n",
            "\n",
            "Extracting  HC ROMS.zip                                                  \b\b\b\b 36%\b\b\b\b\b  OK \n",
            "Extracting  ROMS.zip                                                     \b\b\b\b 74%\b\b\b\b 99%\b\b\b\b\b  OK \n",
            "All OK\n",
            "mkdir: cannot create directory ‘rars’: File exists\n",
            "copying adventure.bin from HC ROMS/BY ALPHABET (PAL)/A-G/Adventure (PAL).bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/adventure.bin\n",
            "copying air_raid.bin from HC ROMS/BY ALPHABET (PAL)/A-G/Air Raid (PAL).bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/air_raid.bin\n",
            "copying alien.bin from HC ROMS/BY ALPHABET (PAL)/A-G/REMAINING NTSC ORIGINALS/Alien.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/alien.bin\n",
            "copying crazy_climber.bin from HC ROMS/BY ALPHABET (PAL)/A-G/REMAINING NTSC ORIGINALS/Crazy Climber.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/crazy_climber.bin\n",
            "copying elevator_action.bin from HC ROMS/BY ALPHABET (PAL)/A-G/REMAINING NTSC ORIGINALS/Elevator Action (Prototype).bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/elevator_action.bin\n",
            "copying gravitar.bin from HC ROMS/BY ALPHABET (PAL)/A-G/REMAINING NTSC ORIGINALS/Gravitar.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/gravitar.bin\n",
            "copying keystone_kapers.bin from HC ROMS/BY ALPHABET (PAL)/H-R/Keystone Kapers (PAL).bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/keystone_kapers.bin\n",
            "copying king_kong.bin from HC ROMS/BY ALPHABET (PAL)/H-R/King Kong (PAL).bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/king_kong.bin\n",
            "copying laser_gates.bin from HC ROMS/BY ALPHABET (PAL)/H-R/Laser Gates (PAL).bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/laser_gates.bin\n",
            "copying mr_do.bin from HC ROMS/BY ALPHABET (PAL)/H-R/Mr. Do! (PAL).bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/mr_do.bin\n",
            "copying pacman.bin from HC ROMS/BY ALPHABET (PAL)/H-R/Pac-Man (PAL).bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/pacman.bin\n",
            "copying jamesbond.bin from HC ROMS/BY ALPHABET (PAL)/H-R/REMAINING NTSC ORIGINALS/James Bond 007.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/jamesbond.bin\n",
            "copying koolaid.bin from HC ROMS/BY ALPHABET (PAL)/H-R/REMAINING NTSC ORIGINALS/Kool-Aid Man.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/koolaid.bin\n",
            "copying krull.bin from HC ROMS/BY ALPHABET (PAL)/H-R/REMAINING NTSC ORIGINALS/Krull.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/krull.bin\n",
            "copying montezuma_revenge.bin from HC ROMS/BY ALPHABET (PAL)/H-R/REMAINING NTSC ORIGINALS/Montezuma's Revenge - Featuring Panama Joe.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/montezuma_revenge.bin\n",
            "copying star_gunner.bin from HC ROMS/BY ALPHABET (PAL)/S-Z/REMAINING NTSC ORIGINALS/Stargunner.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/star_gunner.bin\n",
            "copying time_pilot.bin from HC ROMS/BY ALPHABET (PAL)/S-Z/REMAINING NTSC ORIGINALS/Time Pilot.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/time_pilot.bin\n",
            "copying up_n_down.bin from HC ROMS/BY ALPHABET (PAL)/S-Z/REMAINING NTSC ORIGINALS/Up 'n Down.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/up_n_down.bin\n",
            "copying sir_lancelot.bin from HC ROMS/BY ALPHABET (PAL)/S-Z/Sir Lancelot (PAL).bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/sir_lancelot.bin\n",
            "copying amidar.bin from HC ROMS/BY ALPHABET/A-G/Amidar.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/amidar.bin\n",
            "copying asteroids.bin from HC ROMS/BY ALPHABET/A-G/Asteroids [no copyright].bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/asteroids.bin\n",
            "copying atlantis.bin from HC ROMS/BY ALPHABET/A-G/Atlantis.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/atlantis.bin\n",
            "copying bank_heist.bin from HC ROMS/BY ALPHABET/A-G/Bank Heist.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/bank_heist.bin\n",
            "copying battle_zone.bin from HC ROMS/BY ALPHABET/A-G/Battlezone.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/battle_zone.bin\n",
            "copying beam_rider.bin from HC ROMS/BY ALPHABET/A-G/Beamrider.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/beam_rider.bin\n",
            "copying berzerk.bin from HC ROMS/BY ALPHABET/A-G/Berzerk.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/berzerk.bin\n",
            "copying bowling.bin from HC ROMS/BY ALPHABET/A-G/Bowling.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/bowling.bin\n",
            "copying boxing.bin from HC ROMS/BY ALPHABET/A-G/Boxing.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/boxing.bin\n",
            "copying breakout.bin from HC ROMS/BY ALPHABET/A-G/Breakout - Breakaway IV.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/breakout.bin\n",
            "copying carnival.bin from HC ROMS/BY ALPHABET/A-G/Carnival.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/carnival.bin\n",
            "copying centipede.bin from HC ROMS/BY ALPHABET/A-G/Centipede.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/centipede.bin\n",
            "copying chopper_command.bin from HC ROMS/BY ALPHABET/A-G/Chopper Command.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/chopper_command.bin\n",
            "copying defender.bin from HC ROMS/BY ALPHABET/A-G/Defender.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/defender.bin\n",
            "copying demon_attack.bin from HC ROMS/BY ALPHABET/A-G/Demon Attack.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/demon_attack.bin\n",
            "copying donkey_kong.bin from HC ROMS/BY ALPHABET/A-G/Donkey Kong.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/donkey_kong.bin\n",
            "copying double_dunk.bin from HC ROMS/BY ALPHABET/A-G/Double Dunk.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/double_dunk.bin\n",
            "copying enduro.bin from HC ROMS/BY ALPHABET/A-G/Enduro.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/enduro.bin\n",
            "copying fishing_derby.bin from HC ROMS/BY ALPHABET/A-G/Fishing Derby.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/fishing_derby.bin\n",
            "copying freeway.bin from HC ROMS/BY ALPHABET/A-G/Freeway.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/freeway.bin\n",
            "copying frogger.bin from HC ROMS/BY ALPHABET/A-G/Frogger.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/frogger.bin\n",
            "copying frostbite.bin from HC ROMS/BY ALPHABET/A-G/Frostbite.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/frostbite.bin\n",
            "copying galaxian.bin from HC ROMS/BY ALPHABET/A-G/Galaxian.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/galaxian.bin\n",
            "copying gopher.bin from HC ROMS/BY ALPHABET/A-G/Gopher.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/gopher.bin\n",
            "copying hero.bin from HC ROMS/BY ALPHABET/H-R/H.E.R.O..bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/hero.bin\n",
            "copying ice_hockey.bin from HC ROMS/BY ALPHABET/H-R/Ice Hockey.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/ice_hockey.bin\n",
            "copying journey_escape.bin from HC ROMS/BY ALPHABET/H-R/Journey Escape.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/journey_escape.bin\n",
            "copying kaboom.bin from HC ROMS/BY ALPHABET/H-R/Kaboom!.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/kaboom.bin\n",
            "copying kangaroo.bin from HC ROMS/BY ALPHABET/H-R/Kangaroo.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/kangaroo.bin\n",
            "copying kung_fu_master.bin from HC ROMS/BY ALPHABET/H-R/Kung-Fu Master.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/kung_fu_master.bin\n",
            "copying lost_luggage.bin from HC ROMS/BY ALPHABET/H-R/Lost Luggage [no opening scene].bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/lost_luggage.bin\n",
            "copying ms_pacman.bin from HC ROMS/BY ALPHABET/H-R/Ms. Pac-Man.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/ms_pacman.bin\n",
            "copying name_this_game.bin from HC ROMS/BY ALPHABET/H-R/Name This Game.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/name_this_game.bin\n",
            "copying phoenix.bin from HC ROMS/BY ALPHABET/H-R/Phoenix.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/phoenix.bin\n",
            "copying pitfall.bin from HC ROMS/BY ALPHABET/H-R/Pitfall! - Pitfall Harry's Jungle Adventure.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/pitfall.bin\n",
            "copying pooyan.bin from HC ROMS/BY ALPHABET/H-R/Pooyan.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/pooyan.bin\n",
            "copying private_eye.bin from HC ROMS/BY ALPHABET/H-R/Private Eye.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/private_eye.bin\n",
            "copying qbert.bin from HC ROMS/BY ALPHABET/H-R/Q-bert.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/qbert.bin\n",
            "copying riverraid.bin from HC ROMS/BY ALPHABET/H-R/River Raid.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/riverraid.bin\n",
            "copying road_runner.bin from patched version of HC ROMS/BY ALPHABET/H-R/Road Runner.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/road_runner.bin\n",
            "copying robotank.bin from HC ROMS/BY ALPHABET/H-R/Robot Tank.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/robotank.bin\n",
            "copying seaquest.bin from HC ROMS/BY ALPHABET/S-Z/Seaquest.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/seaquest.bin\n",
            "copying skiing.bin from HC ROMS/BY ALPHABET/S-Z/Skiing.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/skiing.bin\n",
            "copying solaris.bin from HC ROMS/BY ALPHABET/S-Z/Solaris.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/solaris.bin\n",
            "copying space_invaders.bin from HC ROMS/BY ALPHABET/S-Z/Space Invaders.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/space_invaders.bin\n",
            "copying surround.bin from HC ROMS/BY ALPHABET/S-Z/Surround - Chase.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/surround.bin\n",
            "copying tennis.bin from HC ROMS/BY ALPHABET/S-Z/Tennis.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/tennis.bin\n",
            "copying trondead.bin from HC ROMS/BY ALPHABET/S-Z/TRON - Deadly Discs.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/trondead.bin\n",
            "copying tutankham.bin from HC ROMS/BY ALPHABET/S-Z/Tutankham.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/tutankham.bin\n",
            "copying venture.bin from HC ROMS/BY ALPHABET/S-Z/Venture.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/venture.bin\n",
            "copying pong.bin from HC ROMS/BY ALPHABET/S-Z/Video Olympics - Pong Sports.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/pong.bin\n",
            "copying video_pinball.bin from HC ROMS/BY ALPHABET/S-Z/Video Pinball - Arcade Pinball.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/video_pinball.bin\n",
            "copying wizard_of_wor.bin from HC ROMS/BY ALPHABET/S-Z/Wizard of Wor.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/wizard_of_wor.bin\n",
            "copying yars_revenge.bin from HC ROMS/BY ALPHABET/S-Z/Yars' Revenge.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/yars_revenge.bin\n",
            "copying zaxxon.bin from HC ROMS/BY ALPHABET/S-Z/Zaxxon.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/zaxxon.bin\n",
            "copying assault.bin from HC ROMS/NTSC VERSIONS OF PAL ORIGINALS/Assault (AKA Sky Alien) (1983) (Bomb - Onbase) (CA281).bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/assault.bin\n",
            "copying asterix.bin from ROMS/Asterix (AKA Taz) (07-27-1983) (Atari, Jerome Domurat, Steve Woita) (CX2696) (Prototype).bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/asterix.bin\n"
          ]
        }
      ],
      "source": [
        "import urllib.request\n",
        "urllib.request.urlretrieve('http://www.atarimania.com/roms/Roms.rar','Roms.rar')\n",
        "!pip install unrar\n",
        "!unrar x Roms.rar\n",
        "!mkdir rars\n",
        "!mv HC\\ ROMS.zip   rars\n",
        "!mv ROMS.zip  rars\n",
        "!python -m atari_py.import_roms rars"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OGnE4_jcSea2",
        "outputId": "e5726f02-69c9-40fa-f286-aeff24257e7e"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: tensorboardX in /usr/local/lib/python3.7/dist-packages (2.4.1)\n",
            "Requirement already satisfied: protobuf>=3.8.0 in /usr/local/lib/python3.7/dist-packages (from tensorboardX) (3.17.3)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from tensorboardX) (1.19.5)\n",
            "Requirement already satisfied: six>=1.9 in /usr/local/lib/python3.7/dist-packages (from protobuf>=3.8.0->tensorboardX) (1.15.0)\n"
          ]
        }
      ],
      "source": [
        "!pip install tensorboardX"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "WTe9pGh5u5ha",
        "outputId": "dd56e5b5-21c1-4081-beda-9fafcdfcd4a8"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "/content\n"
          ]
        }
      ],
      "source": [
        "%cd /content/"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XJ6wcQ58aoVA"
      },
      "outputs": [],
      "source": [
        "import gym\n",
        "import numpy as np\n",
        "import cv2\n",
        "import random\n",
        "import torch\n",
        "# import models\n",
        "from torch import nn\n",
        "from tensorboardX import SummaryWriter\n",
        "import torchvision.utils as vutils\n",
        "import os\n",
        "import imageio\n",
        "from torchsummary import summary"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## [Code Base](https://github.com/PacktPublishing/Deep-Reinforcement-Learning-Hands-On-Second-Edition/blob/master/Chapter03/03_atari_gan.py)"
      ],
      "metadata": {
        "id": "ias5LecQwfX5"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QgRBsNtwe8Dk"
      },
      "outputs": [],
      "source": [
        "LATENT_VECTOR_SIZE = 100\n",
        "DISCR_FILTERS = 64\n",
        "GENER_FILTERS = 64\n",
        "BATCH_SIZE = 16\n",
        "\n",
        "# dimension input image will be rescaled\n",
        "IMAGE_SIZE = 64\n",
        "\n",
        "LEARNING_RATE = 0.0001\n",
        "REPORT_EVERY_ITER = 10000\n",
        "SAVE_IMAGE_EVERY_ITER = 10000"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "I5ehImJAatTQ"
      },
      "outputs": [],
      "source": [
        "class InputWrapper(gym.ObservationWrapper):\n",
        "  def __init__(self, *args):\n",
        "    super(InputWrapper, self).__init__(*args)\n",
        "    assert isinstance(self.observation_space, gym.spaces.Box)\n",
        "    prev_space = self.observation_space\n",
        "    self.observation_space = gym.spaces.Box(\n",
        "        self.observation(prev_space.low),\n",
        "        self.observation(prev_space.high),\n",
        "        dtype=np.float32\n",
        "    )\n",
        "\n",
        "  def observation(self, obs):\n",
        "    new_obs = cv2.resize(obs, (IMAGE_SIZE, IMAGE_SIZE))\n",
        "    # transform (210, 160, 3) -> (3, 210, 160)\n",
        "    new_obs = np.moveaxis(new_obs, 2, 0)\n",
        "    return new_obs.astype(np.float32)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Y4LDX6zCdjiW"
      },
      "outputs": [],
      "source": [
        "def generate_samples(envs, batch_size=BATCH_SIZE):\n",
        "  batch = [e.reset() for e in envs]\n",
        "  env_gen = iter(lambda: random.choice(envs), None)\n",
        "\n",
        "  while True:\n",
        "    e = next(env_gen)\n",
        "    obs, reward, is_done, _ = e.step(e.action_space.sample())\n",
        "    if np.mean(obs) > 0.01:\n",
        "      batch.append(obs)\n",
        "    if len(batch) == batch_size:\n",
        "      # Normalizing input between -1 to 1\n",
        "      batch_np = np.array(batch, dtype=np.float32)\n",
        "      batch_np *= 2.0/255.0 - 1.0\n",
        "      yield torch.tensor(batch_np)\n",
        "      batch.clear()\n",
        "    if is_done:\n",
        "      e.reset()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0wFKZD8XVDr6"
      },
      "outputs": [],
      "source": [
        "class Discriminator(nn.Module):\n",
        "    def __init__(self, input_shape):\n",
        "        super(Discriminator, self).__init__()\n",
        "        # this pipe converges image into the single number\n",
        "        self.conv_pipe = nn.Sequential(\n",
        "            nn.Conv2d(in_channels=input_shape[0], out_channels=DISCR_FILTERS,\n",
        "                      kernel_size=4, stride=2, padding=1),\n",
        "            nn.ReLU(),\n",
        "            nn.Conv2d(in_channels=DISCR_FILTERS, out_channels=DISCR_FILTERS*2,\n",
        "                      kernel_size=4, stride=2, padding=1),\n",
        "            nn.BatchNorm2d(DISCR_FILTERS*2),\n",
        "            nn.ReLU(),\n",
        "            nn.Conv2d(in_channels=DISCR_FILTERS * 2, out_channels=DISCR_FILTERS * 4,\n",
        "                      kernel_size=4, stride=2, padding=1),\n",
        "            nn.BatchNorm2d(DISCR_FILTERS * 4),\n",
        "            nn.ReLU(),\n",
        "            nn.Conv2d(in_channels=DISCR_FILTERS * 4, out_channels=DISCR_FILTERS * 8,\n",
        "                      kernel_size=4, stride=2, padding=1),\n",
        "            nn.BatchNorm2d(DISCR_FILTERS * 8),\n",
        "            nn.ReLU(),\n",
        "            nn.Conv2d(in_channels=DISCR_FILTERS * 8, out_channels=1,\n",
        "                      kernel_size=4, stride=1, padding=0),\n",
        "            nn.Sigmoid()\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        for layer in self.conv_pipe:\n",
        "          x = layer(x)\n",
        "          print(x.size())\n",
        "        return x.view(-1, 1).squeeze(dim=1)\n",
        "\n",
        "\n",
        "class Generator(nn.Module):\n",
        "    def __init__(self, output_shape):\n",
        "        super(Generator, self).__init__()\n",
        "        # pipe deconvolves input vector into (3, 64, 64) image\n",
        "        self.pipe = nn.Sequential(\n",
        "            nn.ConvTranspose2d(in_channels=LATENT_VECTOR_SIZE, out_channels=GENER_FILTERS * 8,\n",
        "                               kernel_size=4, stride=1, padding=0),\n",
        "            nn.BatchNorm2d(GENER_FILTERS * 8),\n",
        "            nn.ReLU(),\n",
        "            nn.ConvTranspose2d(in_channels=GENER_FILTERS * 8, out_channels=GENER_FILTERS * 4,\n",
        "                               kernel_size=4, stride=2, padding=1),\n",
        "            nn.BatchNorm2d(GENER_FILTERS * 4),\n",
        "            nn.ReLU(),\n",
        "            nn.ConvTranspose2d(in_channels=GENER_FILTERS * 4, out_channels=GENER_FILTERS * 2,\n",
        "                               kernel_size=4, stride=2, padding=1),\n",
        "            nn.BatchNorm2d(GENER_FILTERS * 2),\n",
        "            nn.ReLU(),\n",
        "            nn.ConvTranspose2d(in_channels=GENER_FILTERS * 2, out_channels=GENER_FILTERS,\n",
        "                               kernel_size=4, stride=2, padding=1),\n",
        "            nn.BatchNorm2d(GENER_FILTERS),\n",
        "            nn.ReLU(),\n",
        "            nn.ConvTranspose2d(in_channels=GENER_FILTERS, out_channels=output_shape[0],\n",
        "                               kernel_size=4, stride=2, padding=1),\n",
        "            nn.Tanh()\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        # return self.pipe(x)\n",
        "        for layer in self.pipe:\n",
        "          x = layer(x)\n",
        "          print(x.size())\n",
        "        return x"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "model = Discriminator(input_shape=(3, 64, 64))\n",
        "print(model)\n",
        "x = torch.randn(1, 3, 64, 64)\n",
        "model(x)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "tYkZ6xLIoOX3",
        "outputId": "770c6e8b-a14b-4b0e-bee4-e764ad702aa9"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Discriminator(\n",
            "  (conv_pipe): Sequential(\n",
            "    (0): Conv2d(3, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
            "    (1): ReLU()\n",
            "    (2): Conv2d(64, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
            "    (3): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "    (4): ReLU()\n",
            "    (5): Conv2d(128, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
            "    (6): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "    (7): ReLU()\n",
            "    (8): Conv2d(256, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
            "    (9): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "    (10): ReLU()\n",
            "    (11): Conv2d(512, 1, kernel_size=(4, 4), stride=(1, 1))\n",
            "    (12): Sigmoid()\n",
            "  )\n",
            ")\n",
            "torch.Size([1, 64, 32, 32])\n",
            "torch.Size([1, 64, 32, 32])\n",
            "torch.Size([1, 128, 16, 16])\n",
            "torch.Size([1, 128, 16, 16])\n",
            "torch.Size([1, 128, 16, 16])\n",
            "torch.Size([1, 256, 8, 8])\n",
            "torch.Size([1, 256, 8, 8])\n",
            "torch.Size([1, 256, 8, 8])\n",
            "torch.Size([1, 512, 4, 4])\n",
            "torch.Size([1, 512, 4, 4])\n",
            "torch.Size([1, 512, 4, 4])\n",
            "torch.Size([1, 1, 1, 1])\n",
            "torch.Size([1, 1, 1, 1])\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([0.5849], grad_fn=<SqueezeBackward1>)"
            ]
          },
          "metadata": {},
          "execution_count": 12
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "model = Generator(output_shape=(3, 64, 64))\n",
        "print(model)\n",
        "x = torch.randn(1, 100, 1, 1)\n",
        "model(x)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4HNth_Qp3pJp",
        "outputId": "c8d59ec5-593b-4ebc-820e-74dc51a05b20"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Generator(\n",
            "  (pipe): Sequential(\n",
            "    (0): ConvTranspose2d(100, 512, kernel_size=(4, 4), stride=(1, 1))\n",
            "    (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "    (2): ReLU()\n",
            "    (3): ConvTranspose2d(512, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
            "    (4): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "    (5): ReLU()\n",
            "    (6): ConvTranspose2d(256, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
            "    (7): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "    (8): ReLU()\n",
            "    (9): ConvTranspose2d(128, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
            "    (10): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "    (11): ReLU()\n",
            "    (12): ConvTranspose2d(64, 3, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
            "    (13): Tanh()\n",
            "  )\n",
            ")\n",
            "torch.Size([1, 512, 4, 4])\n",
            "torch.Size([1, 512, 4, 4])\n",
            "torch.Size([1, 512, 4, 4])\n",
            "torch.Size([1, 256, 8, 8])\n",
            "torch.Size([1, 256, 8, 8])\n",
            "torch.Size([1, 256, 8, 8])\n",
            "torch.Size([1, 128, 16, 16])\n",
            "torch.Size([1, 128, 16, 16])\n",
            "torch.Size([1, 128, 16, 16])\n",
            "torch.Size([1, 64, 32, 32])\n",
            "torch.Size([1, 64, 32, 32])\n",
            "torch.Size([1, 64, 32, 32])\n",
            "torch.Size([1, 3, 64, 64])\n",
            "torch.Size([1, 3, 64, 64])\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[[[-0.4019, -0.3402, -0.3185,  ..., -0.4617, -0.1878, -0.1808],\n",
              "          [ 0.1936,  0.5633, -0.2272,  ..., -0.2690, -0.0699,  0.0929],\n",
              "          [ 0.0487,  0.3593,  0.1587,  ..., -0.0763, -0.5029, -0.1871],\n",
              "          ...,\n",
              "          [ 0.0438,  0.0740, -0.1021,  ...,  0.1895,  0.2004, -0.0818],\n",
              "          [-0.1418, -0.3108, -0.1498,  ..., -0.2686, -0.2964,  0.1426],\n",
              "          [-0.0679,  0.1171, -0.0870,  ..., -0.1997, -0.1087, -0.0620]],\n",
              "\n",
              "         [[-0.1243, -0.2389, -0.0303,  ...,  0.0397,  0.0891,  0.0458],\n",
              "          [-0.0809, -0.6348, -0.1133,  ..., -0.1643, -0.5549, -0.3384],\n",
              "          [-0.1248, -0.2911,  0.7508,  ..., -0.4287, -0.3750, -0.0087],\n",
              "          ...,\n",
              "          [-0.1277, -0.5549,  0.2900,  ..., -0.3760, -0.2701, -0.2142],\n",
              "          [-0.0039, -0.4392,  0.2489,  ..., -0.2774,  0.0533, -0.2155],\n",
              "          [ 0.1072, -0.5004,  0.2728,  ..., -0.3679,  0.0271, -0.2887]],\n",
              "\n",
              "         [[-0.0191,  0.3626, -0.1019,  ...,  0.1597, -0.1338, -0.1160],\n",
              "          [ 0.0735,  0.7870,  0.7080,  ...,  0.4138,  0.3771, -0.1133],\n",
              "          [-0.0352,  0.8163, -0.0618,  ...,  0.6118,  0.3656,  0.2386],\n",
              "          ...,\n",
              "          [-0.0874,  0.4892,  0.0923,  ...,  0.4589,  0.1384,  0.0895],\n",
              "          [ 0.1907,  0.3534,  0.2837,  ...,  0.6340,  0.1458,  0.2869],\n",
              "          [ 0.0314,  0.2036,  0.0602,  ...,  0.0089,  0.0453,  0.3452]]]],\n",
              "       grad_fn=<TanhBackward0>)"
            ]
          },
          "metadata": {},
          "execution_count": 13
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "y0aloEMRfu9I"
      },
      "outputs": [],
      "source": [
        "device = torch.device(\"cuda\")\n",
        "envs = [InputWrapper(gym.make(game)) for game in ('Breakout-v0', 'MsPacman-v0', 'SpaceInvaders-v0')]\n",
        "train_data = generate_samples(envs)\n",
        "noise = torch.randn(BATCH_SIZE, 100, 1, 1)\n",
        "input_shape = envs[0].observation_space.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0RvFrLNHuIxh",
        "outputId": "49bff4ea-d2fc-443c-8eb6-f9d18804a0bb"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Iter 10000: gen_loss=11.015083949191167, dis_loss=0.0050717151021431395\n",
            "Iter 20000: gen_loss=15.72122150363922, dis_loss=1.2060845846772138e-06\n",
            "Iter 30000: gen_loss=20.840003999137878, dis_loss=4.699084856341073e-09\n",
            "Iter 40000: gen_loss=24.618247546958923, dis_loss=6.929041242997869e-11\n",
            "Iter 50000: gen_loss=26.428978769874572, dis_loss=1.4901162970204496e-12\n",
            "Iter 60000: gen_loss=27.05434672050476, dis_loss=0.0\n",
            "Iter 70000: gen_loss=26.78728736190796, dis_loss=1.043081425677883e-11\n",
            "Iter 80000: gen_loss=26.567670729446412, dis_loss=5.215407128389415e-12\n",
            "Iter 90000: gen_loss=26.977887541770937, dis_loss=0.0\n",
            "Iter 100000: gen_loss=27.515558718681337, dis_loss=0.0\n",
            "Iter 110000: gen_loss=27.650638795280457, dis_loss=4.470348713425665e-12\n",
            "Iter 120000: gen_loss=27.69635641441345, dis_loss=0.0\n",
            "Iter 130000: gen_loss=27.93418836631775, dis_loss=0.0\n",
            "Iter 140000: gen_loss=26.922232485198975, dis_loss=9.685756197086447e-12\n",
            "Iter 150000: gen_loss=27.55939024143219, dis_loss=0.0\n"
          ]
        }
      ],
      "source": [
        "# training loops\n",
        "net_discr = Discriminator(input_shape=input_shape).to(device)\n",
        "net_gener = Generator(output_shape=input_shape).to(device)\n",
        "writer = SummaryWriter()\n",
        "\n",
        "objective = nn.BCELoss()\n",
        "\n",
        "gen_optimizer = torch.optim.Adam(net_gener.parameters(), lr=0.0001, betas=(0.5, 0.999))\n",
        "dis_optimizer = torch.optim.Adam(net_discr.parameters(), lr=0.0001, betas=(0.5, 0.999))\n",
        "\n",
        "gen_losses = []\n",
        "dis_losses = []\n",
        "iter_no = 0\n",
        "\n",
        "true_labels_v = torch.ones(BATCH_SIZE, device=device)\n",
        "fake_labels_v = torch.zeros(BATCH_SIZE, device=device)\n",
        "\n",
        "for batch_v in train_data:\n",
        "  gen_input_v = torch.FloatTensor(\n",
        "            BATCH_SIZE, LATENT_VECTOR_SIZE, 1, 1)\n",
        "  gen_input_v.normal_(0, 1)\n",
        "  gen_input_v = gen_input_v.to(device)\n",
        "  batch_v = batch_v.to(device)\n",
        "  gen_output_v = net_gener(gen_input_v)\n",
        "\n",
        "  # train discriminator\n",
        "  dis_optimizer.zero_grad()\n",
        "  dis_output_true_v = net_discr(batch_v)\n",
        "  dis_output_fake_v = net_discr(gen_output_v.detach())\n",
        "  dis_loss = objective(dis_output_true_v, true_labels_v) + \\\n",
        "              objective(dis_output_fake_v, fake_labels_v)\n",
        "  dis_loss.backward()\n",
        "  dis_optimizer.step()\n",
        "  dis_losses.append(dis_loss.item())\n",
        "\n",
        "  # train generator\n",
        "  gen_optimizer.zero_grad()\n",
        "  dis_output_v = net_discr(gen_output_v)\n",
        "  gen_loss_v = objective(dis_output_v, true_labels_v)\n",
        "  gen_loss_v.backward()\n",
        "  gen_optimizer.step()\n",
        "  gen_losses.append(gen_loss_v.item())\n",
        "\n",
        "  iter_no += 1\n",
        "  if iter_no % REPORT_EVERY_ITER == 0:\n",
        "    print(f\"Iter {iter_no}: gen_loss={np.mean(gen_losses)}, dis_loss={np.mean(dis_losses)}\")\n",
        "    writer.add_scalar(\n",
        "        \"gen_loss\", np.mean(gen_losses), iter_no)\n",
        "    writer.add_scalar(\n",
        "        \"dis_loss\", np.mean(dis_losses), iter_no)\n",
        "    gen_losses = []\n",
        "    dis_losses = []\n",
        "    \n",
        "    if not os.path.exists(f'{iter_no}'):\n",
        "        os.makedirs(f'{iter_no}')\n",
        "    torch.save(net_discr.state_dict(), f'{iter_no}/D.pkl')\n",
        "    torch.save(net_gener.state_dict(), f'{iter_no}/G.pkl')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KDWP9w3IGLjX"
      },
      "outputs": [],
      "source": [
        ""
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "dcgan-atari-train",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}