14_rl/01_Deep_Q_Network.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "machine_shape": "hm",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/machine-perception-robotics-group/MPRGDeepLearningLectureNotebook/blob/master/14_rl/01_Deep_Q_Network.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P9WNOydibxSw"
      },
      "source": [
        "# 目的\n",
        "Deep Q-networkの仕組みを理解し，ゲームタスクを用いて強化学習をおこなう．\\\n",
        "学習後のエージェントの可視化を行い，学習がうまくできているか確認を行う．"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W_x6vE50N8dJ"
      },
      "source": [
        "# Deep Q-Network (DQN)\n",
        "Deep Q-Network (DQN)は，GoogleのDeepmindが2016年に発表した手法で，Q学習におけるQテーブルを用いた行動価値の導出を，DCNNを用いた近似関数で代用するのが主な手法の内容である．\n",
        "\n",
        "Q学習では，全ての状態と行動の組み合わせについて，行動価値をQテーブルに記録している．\n",
        "そのため，状態数と行動数の組み合わせが膨大な環境に対して，膨大なメモリが必要となる問題を抱えていました．\\\n",
        "この問題に対して，DQNはQテーブルそのものをDCNNで代用することにより，上記問題を解決し，Atari2600のゲーム環境で人間を超えるスコアの獲得に成功している．\n",
        "また，DQNにはその他にも強化学習における学習の安定性獲得のために，Experience replay，Target Q-Network，Reward clippingなどの工夫がなされている．\n",
        "\n",
        "DQNのように，行動価値関数（Q値）を用いて行動を決定し，最適行動価値関数になるように更新を行っていく手法は，価値ベースの手法と呼ばれている．\n",
        "派生手法には，DDQN，Ape-X，R2D2などの手法がある．\n",
        "\n",
        "<img src=\"https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2708544/0d5711a1-803f-f142-125e-6b70f84a3c77.png\" width = 50%>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5rQGfxWYK_4O"
      },
      "source": [
        "## 準備\n",
        "\n",
        "### Google Colaboratoryの設定確認・変更\n",
        "本チュートリアルではPyTorchを利用してニューラルネットワークの実装を確認，学習および評価を行います．\n",
        "**GPUを用いて処理を行うために，上部のメニューバーの「ランタイム」→「ランタイムのタイプを変更」からハードウェアアクセラレータをGPUにしてください．**\n",
        "\n",
        "\n",
        "### モジュールの追加インストール\n",
        "下記のプログラムを実行して，実験結果の表示に必要な追加ライブラリやモジュールをインストールする．"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install swig==4.3.0\n",
        "!pip install gymnasium==1.0.0\n",
        "!pip install gym[atari,accept-rom-license]==0.23.1\n",
        "\n",
        "!apt update && apt install xvfb\n",
        "!pip install gym-notebook-wrapper==1.3.3"
      ],
      "metadata": {
        "id": "ZSCuMOnGpZNz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CKpJ1mtDFoDo"
      },
      "source": [
        "## モジュールのインポート\n",
        "はじめに必要なモジュールをインポートする．\n",
        "\n",
        "今回はPyTorchに加えて，Pongを実行するためのシミュレータであるOpenAI Gym（gym）をインポートする．"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "import gym\n",
        "import gym.spaces\n",
        "\n",
        "import time\n",
        "import datetime\n",
        "import random\n",
        "import cv2\n",
        "import matplotlib\n",
        "import matplotlib.pyplot as plt\n",
        "import collections\n",
        "%matplotlib inline\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "import torch.nn.functional as F\n",
        "import torchvision.transforms as T\n",
        "\n",
        "# GPUの確認\n",
        "use_cuda = torch.cuda.is_available()\n",
        "print('Use CUDA:', use_cuda)"
      ],
      "metadata": {
        "id": "x2PRCPPtpa0W"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## シード値の固定"
      ],
      "metadata": {
        "id": "2-BDmHImPZ3T"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "seed = 123\n",
        "# Python random\n",
        "random.seed(seed)\n",
        "# Numpy\n",
        "np.random.seed(seed)\n",
        "# Pytorch\n",
        "torch.manual_seed(seed)\n",
        "torch.cuda.manual_seed(seed)\n",
        "torch.backends.cudnn.deterministic = True\n",
        "torch.use_deterministic_algorithms = True"
      ],
      "metadata": {
        "id": "EtsiNrIlpcVv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Jo5u_p4SFoMa"
      },
      "source": [
        "## OpenAI GymによるPong環境の定義\n",
        " [OpenAI Gym](https://github.com/openai/gym) は，様々な種類の環境を提供しているモジュールです．\n",
        "\n",
        "今回は，gymで利用可能なAtari2600のゲームであるPongを使用します．\n",
        "まず，gym.make関数で実行したい環境を指定します（今回はPong）．\n",
        "その後，reset関数を実行することで，環境を初期化します．\n",
        "\n",
        "Pongは，パドルを操作してボールが自分の陣地に入らないように打ち返すゲームです（いわゆる，スカッシュゲーム）．\\\n",
        "Pongにおいて，現在の状態を確認するためにゲームの画面情報が与えられており，`observation_space`という変数で確認することができます．\n",
        "また，`action_space`という変数で，エージェントが選択可能な行動の数を確認することができます．Atari2600はゲームタスクのため，行動は基本的にコントローラの入力に対応します．\\\n",
        "Pongの行動数は6で，内訳は以下の通りです．\n",
        "\n",
        "|  id  |  action  |\n",
        "| ---- | ---- |\n",
        "|  0  |  NOOP（操作なし）  |\n",
        "|  1  |  FIRE  |\n",
        "|  2  |  RIGHT  |\n",
        "|  3  |  LEFT  |\n",
        "|  4  |  RIGHTFIRE  |\n",
        "|  5  |  LEFTFIRE  |\n",
        "\n",
        "ただし，Pongでの動作はパドルを上下どちらかに移動させるのみのため，\n",
        "FIREはNOOP，RIGHTFIREはRIGHT，LEFTFIREはLEFTと同じ動作となります．\n",
        "そのため，行動の種類としては，パドルを上下どちらかに移動させる行動の2つとなっています．\n",
        "\n",
        "Pongのゲーム概要は，相手の陣地にボールを入れることで得点を獲得し，ボールを自分の陣地に入れられることで得点を取られます．\n",
        "どちらかのプレイヤーが21点取った時点でゲーム終了となります．\\\n",
        "Atari2600のおける報酬は，基本的にゲームのスコアとなっています．そのため，Pongにおけるデフォルトの報酬設計は，以下の通りになっています．\n",
        "\n",
        "* 自陣にボールを入れられる：-1\n",
        "* 相手の陣地にボールを入れる：+1\n",
        "\n",
        "また，gym.make関数の各引数は以下の通りです．\n",
        "\n",
        "<img src=\"https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2708544/218f6254-bd91-9bda-4297-3831020bded0.png\" width=70%>\n",
        "\n",
        "ここで，ALEとは強化学習エージェントを，Atari2600のゲームを用いて評価するためのプラットフォームです．"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# 環境の指定\n",
        "env = gym.make('ALE/Pong-v5', frameskip=1, repeat_action_probability=0.0, full_action_space=False)\n",
        "\n",
        "# 環境の初期化\n",
        "obs = env.reset(seed=seed)\n",
        "print('observation space:', env.observation_space) # 状態空間\n",
        "print('action space:', env.action_space) # 行動空間\n",
        "print('initial observation:', obs.shape)\n",
        "\n",
        "# 行動の決定と決定した行動の入力\n",
        "action = env.action_space.sample()\n",
        "obs, r, done, info = env.step(action)\n",
        "print('next observation:', obs.shape) # 環境から返ってくる次状態\n",
        "print('reward:', r)\n",
        "print('done:', done)"
      ],
      "metadata": {
        "id": "WN_ijS0Jpe0u"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mOgbuO2I1bfq"
      },
      "source": [
        "## ネットワーク構造\n",
        "ネットワークモデルを定義します．\n",
        "ここでは，環境からのゲーム画面情報（画像）を入力し，行動に対するQ値を出力するようなネットワークを定義します．\n",
        "ネットワーク構造は，畳み込み層3層と全結合層2層とします．\n",
        "\n",
        "入力データのサイズを`input_shape`，出力する行動の数を`n_actions`とし，ネットワークの作成時に変更できるようにしておきます．"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "class DQN(nn.Module):\n",
        "    def __init__(self, input_shape, n_actions):\n",
        "        super(DQN, self).__init__()\n",
        "\n",
        "        self.conv = nn.Sequential(\n",
        "            nn.Conv2d(input_shape[0], 32, kernel_size=8, stride=4),\n",
        "            nn.ReLU(),\n",
        "            nn.Conv2d(32, 64, kernel_size=4, stride=2),\n",
        "            nn.ReLU(),\n",
        "            nn.Conv2d(64, 64, kernel_size=3, stride=1),\n",
        "            nn.ReLU()\n",
        "        )\n",
        "\n",
        "        conv_out_size = self._get_conv_out(input_shape)\n",
        "        self.fc = nn.Sequential(\n",
        "            nn.Linear(conv_out_size, 512),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(512, n_actions)\n",
        "        )\n",
        "\n",
        "    def _get_conv_out(self, shape):\n",
        "        o = self.conv(torch.zeros(1, *shape))\n",
        "        return int(np.prod(o.size()))\n",
        "\n",
        "    def forward(self, x):\n",
        "        conv_out = self.conv(x).view(x.size()[0], -1)\n",
        "        return self.fc(conv_out)"
      ],
      "metadata": {
        "id": "OFeV-CcypgVB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fsX8hfGurmBt"
      },
      "source": [
        "## Deep Q-Networkにおける学習工夫の定義\n",
        "\n",
        "Deep Q-Networkでは，学習の促進と安定化の為に，いくつか工夫を施して学習を行っています．\n",
        "代表的な工夫として，Experience Replay, Target Q-Network, Reward Clippingと呼ばれる3つの工夫があります．\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vXY4bzSf3fRA"
      },
      "source": [
        "### Experience Replay\n",
        "DQNでは，獲得した経験を直接使用し学習するのではなく，獲得した経験を一度Replay Bufferと呼ばれるBufferに格納し，学習する際にはBufferから経験をランダムに取得することにより学習を行います．\n",
        "これをExperience Replayと呼び，データの再利用を行うことで，データ効率を高め効率的な学習を行います．\n",
        "\n",
        "Bufferへは，「現在の状態，その時に選択された行動，行動によって遷移した状態（次状態），その時の報酬」の4種類の情報を1つの経験として蓄積します．\\\n",
        "まず，`Experience`という変数を定義します．\n",
        "ここでは，`state`, `action`, `reward`, `done`, `next_state`が1セットとなるようなデータ構造（辞書オブジェクト）を定義します．\\\n",
        "その後，Experience Bufferクラスを定義します．\n",
        "Experience Bufferクラスでは，経験を蓄積する`buffer`（`capacity = buffer`へ格納する経験の数）を定義します．\n",
        "append関数では，メモリへ経験を格納します．\n",
        "また，sample関数では，指定したバッチサイズ (`batch_size`) 分の経験をランダムに選択し，返す関数を定義します．"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sL0Oolhp47OJ"
      },
      "source": [
        "Experience = collections.namedtuple('Experience', field_names=['state', 'action', 'reward', 'done', 'new_state'])\n",
        "\n",
        "class ExperienceBuffer:\n",
        "    def __init__(self, capacity):\n",
        "        self.buffer = collections.deque(maxlen=capacity)\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.buffer)\n",
        "\n",
        "    def append(self, experience):\n",
        "        self.buffer.append(experience)\n",
        "\n",
        "    def sample(self, batch_size):\n",
        "        indices = np.random.choice(len(self.buffer), batch_size, replace=False)\n",
        "        states, actions, rewards, dones, next_states = zip(*[self.buffer[idx] for idx in indices])\n",
        "        return np.array(states), np.array(actions), np.array(rewards, dtype=np.float32), \\\n",
        "               np.array(dones, dtype=bool), np.array(next_states)"
      ],
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "通常のBuffer関数"
      ],
      "metadata": {
        "id": "mBHDtRoAPtNU"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "Experience = collections.namedtuple('Experience', field_names=['state', 'action', 'reward', 'done', 'new_state'])\n",
        "\n",
        "class Buffer:\n",
        "    def __init__(self, capacity):\n",
        "        self.buffer = collections.deque(maxlen=capacity)\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.buffer)\n",
        "\n",
        "    def append(self, experience):\n",
        "        self.buffer.append(experience)\n",
        "\n",
        "    def sample(self, batch_size):\n",
        "        states, actions, rewards, dones, next_states = zip(*[self.buffer[idx] for idx in range(batch_size)])\n",
        "        return np.array(states), np.array(actions), np.array(rewards, dtype=np.float32), \\\n",
        "               np.array(dones, dtype=bool), np.array(next_states)"
      ],
      "metadata": {
        "id": "PSOsjgtjPy97"
      },
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sYQi9_LQ5UiF"
      },
      "source": [
        "### Target Q-Network\n",
        "DQNで誤差を算出する際，目標値として用いる行動価値関数と，現在の行動価値関数をそれぞれ別のネットワークの出力を用いて誤差を算出します．\\\n",
        "目標値として算出に用いるネットワークは，一定周期経過するまで重みを固定したネットワークとし，一定周期で現在のネットワークと同期しながら学習を行います．\n",
        "この工夫をTarget Q-Networkと呼び，学習の安定化を図ります．\n",
        "\n",
        "`target_net`を現在の`net`と同期する`sync_network`を定義します．"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Iod_iK4Yt7tB"
      },
      "source": [
        "def sync_network():\n",
        "    tgt_net.load_state_dict(net.state_dict())"
      ],
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sUL1qa38uol6"
      },
      "source": [
        "### Reward clipping\n",
        "DQNでは，学習に用いる報酬値を，以下の通りにクリッピングします．\n",
        "\n",
        "- 報酬値が正：+1\n",
        "- 報酬値が0：0\n",
        "- 報酬値が負：-1\n",
        "\n",
        "これにより，学習における報酬の外れ値に対する過剰反応を防ぐことができます．\n",
        "\n",
        "Reward clippingは，環境に直接ラップするため，`wrapper`クラスを定義します．\n",
        "reward関数では，環境から受け取った報酬を-1, 0, 1にクリップした報酬値を返します．"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LL7Kkp-svXyu"
      },
      "source": [
        "class ClipRewardEnv(gym.RewardWrapper):\n",
        "    def __init__(self, env):\n",
        "        gym.RewardWrapper.__init__(self, env)\n",
        "\n",
        "    def reward(self, reward):\n",
        "        \"\"\"Bin reward to {+1, 0, -1} by its sign.\"\"\"\n",
        "        return np.sign(reward)"
      ],
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Mk9EDXRAt-1I"
      },
      "source": [
        "## その他の学習に必要な処理\n",
        "gymのatari環境で，学習の安定化と効率的な学習を行うための処理をいくつか行います．\n",
        "\n",
        "* MaxAndSkipEnv：1ステップ実行毎に，４フレームゲームを進める（skip frame）\n",
        "* FireResetEnv：エピソード（ゲーム）開始に，Fireを実行しなければ開始されない環境でのreset関数の設定\n",
        "* ProcessFrame84：210×160のRGB画像を84×84のグレースケール画像に変換\n",
        "* ImageToPyTorch：観測情報（画像）のshapeをHWC(高さ，幅，チャネル)からCHW（チャネル，高さ，幅）に変換\n",
        "* ScaledFloatFrame：画像（0から255）を０．０から1.0の範囲で正規化\n",
        "* BufferWrapper：観測情報を4フレームまとめて返す\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "53nSLW5xwH70"
      },
      "source": [
        "class MaxAndSkipEnv(gym.Wrapper):\n",
        "    def __init__(self, env=None, skip=4):\n",
        "        \"\"\"Return only every `skip`-th frame\"\"\"\n",
        "        super(MaxAndSkipEnv, self).__init__(env)\n",
        "        # most recent raw observations (for max pooling across time steps)\n",
        "        self._obs_buffer = collections.deque(maxlen=2)\n",
        "        self._skip = skip\n",
        "\n",
        "    def step(self, action):\n",
        "        total_reward = 0.0\n",
        "        done = None\n",
        "        for _ in range(self._skip):\n",
        "            obs, reward, done, info = self.env.step(action)\n",
        "            self._obs_buffer.append(obs)\n",
        "            total_reward += reward\n",
        "            if done:\n",
        "                break\n",
        "        max_frame = np.max(np.stack(self._obs_buffer), axis=0)\n",
        "        return max_frame, total_reward, done, info\n",
        "\n",
        "class FireResetEnv(gym.Wrapper):\n",
        "    def __init__(self, env=None):\n",
        "        \"\"\"For environments where the user need to press FIRE for the game to start.\"\"\"\n",
        "        super(FireResetEnv, self).__init__(env)\n",
        "        assert env.unwrapped.get_action_meanings()[1] == 'FIRE'\n",
        "        assert len(env.unwrapped.get_action_meanings()) >= 3\n",
        "\n",
        "    def step(self, action):\n",
        "        return self.env.step(action)\n",
        "\n",
        "    def reset(self, seed=None):\n",
        "        self.env.reset()\n",
        "        obs, _, done, _ = self.env.step(1)\n",
        "        if done:\n",
        "            self.env.reset(seed=seed)\n",
        "        obs, _, done, _ = self.env.step(2)\n",
        "        if done:\n",
        "            self.env.reset(seed=seed)\n",
        "        return obs\n",
        "\n",
        "class ProcessFrame84(gym.ObservationWrapper):\n",
        "    def __init__(self, env=None):\n",
        "        super(ProcessFrame84, self).__init__(env)\n",
        "        self.observation_space = gym.spaces.Box(low=0, high=255, shape=(84, 84, 1), dtype=np.uint8)\n",
        "\n",
        "    def observation(self, obs):\n",
        "        return ProcessFrame84.process(obs)\n",
        "\n",
        "    @staticmethod\n",
        "    def process(frame):\n",
        "        if frame.size == 210 * 160 * 3:\n",
        "            img = np.reshape(frame, [210, 160, 3]).astype(np.float32)\n",
        "        elif frame.size == 250 * 160 * 3:\n",
        "            img = np.reshape(frame, [250, 160, 3]).astype(np.float32)\n",
        "        else:\n",
        "            assert False, \"Unknown resolution.\"\n",
        "        img = img[:, :, 0] * 0.299 + img[:, :, 1] * 0.587 + img[:, :, 2] * 0.114\n",
        "        resized_screen = cv2.resize(img, (84, 110), interpolation=cv2.INTER_AREA)\n",
        "        x_t = resized_screen[18:102, :]\n",
        "        x_t = np.reshape(x_t, [84, 84, 1])\n",
        "        return x_t.astype(np.uint8)\n",
        "\n",
        "class ImageToPyTorch(gym.ObservationWrapper):\n",
        "    def __init__(self, env):\n",
        "        super(ImageToPyTorch, self).__init__(env)\n",
        "        old_shape = self.observation_space.shape\n",
        "        self.observation_space = gym.spaces.Box(low=0.0, high=1.0, shape=(old_shape[-1], old_shape[0], old_shape[1]),\n",
        "                                                dtype=np.float32)\n",
        "\n",
        "    def observation(self, observation):\n",
        "        return np.moveaxis(observation, 2, 0)\n",
        "\n",
        "\n",
        "class ScaledFloatFrame(gym.ObservationWrapper):\n",
        "    def observation(self, obs):\n",
        "        return np.array(obs).astype(np.float32) / 255.0\n",
        "\n",
        "\n",
        "class BufferWrapper(gym.ObservationWrapper):\n",
        "    def __init__(self, env, n_steps, dtype=np.float32):\n",
        "        super(BufferWrapper, self).__init__(env)\n",
        "        self.dtype = dtype\n",
        "        old_space = env.observation_space\n",
        "        self.observation_space = gym.spaces.Box(old_space.low.repeat(n_steps, axis=0),\n",
        "                                                old_space.high.repeat(n_steps, axis=0), dtype=dtype)\n",
        "\n",
        "    def reset(self, seed=None):\n",
        "        self.buffer = np.zeros_like(self.observation_space.low, dtype=self.dtype)\n",
        "        return self.observation(self.env.reset(seed=seed))\n",
        "\n",
        "    def observation(self, observation):\n",
        "        self.buffer[:-1] = self.buffer[1:]\n",
        "        self.buffer[-1] = observation\n",
        "        return self.buffer"
      ],
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Q7C3y7rvo9HP"
      },
      "source": [
        "### 学習に必要な処理の適用\n",
        "環境に対して必要となるそれぞれの処理を作成した環境に対して適用します．"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "E1G4H7WxpSKA"
      },
      "source": [
        "env = gym.make('ALE/Pong-v5', frameskip=1, repeat_action_probability=0.0, full_action_space=False)\n",
        "\n",
        "env = ClipRewardEnv(env)\n",
        "env = MaxAndSkipEnv(env)\n",
        "env = FireResetEnv(env)\n",
        "env = ProcessFrame84(env)\n",
        "env = ImageToPyTorch(env)\n",
        "env = BufferWrapper(env, 4)\n",
        "env = ScaledFloatFrame(env)"
      ],
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "23ixNGvKdAB4"
      },
      "source": [
        "## エージェントの定義\n",
        "エージェントが，環境に対して行動価値に沿った行動を選択し，環境から経験を取得，Experience Bufferへ獲得した経験を記録するようにします．\n",
        "\n",
        "エージェントの環境に対する動きのクラスを定義します．\\\n",
        "play_step関数は，環境に対し行動を決定する関数です．\n",
        "$\\epsilon$-greedy法を用いて，一定の割合でランダムに行動選択を行います．\n",
        "それ以外の場合は，ネットワークへ環境情報（画像）を入力し，行動を決定します．また，確率$\\epsilon$の値は，学習が進むごとに大きくなっていき，学習初期は探索を行い，徐々に最適な行動のみを選択するようになります．"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zaGTP_HhdB6m"
      },
      "source": [
        "class Agent:\n",
        "    def __init__(self, env, exp_buffer, seed=123):\n",
        "        self.env = env\n",
        "        self.exp_buffer = exp_buffer\n",
        "        self.seed = seed\n",
        "        self._reset()\n",
        "\n",
        "    def _reset(self):\n",
        "        self.state = env.reset(seed=self.seed)\n",
        "        self.total_reward = 0.0\n",
        "\n",
        "    def play_step(self, net=None, epsilon=0.0):\n",
        "        done_reward = None\n",
        "\n",
        "        if np.random.random() < epsilon:\n",
        "            action = env.action_space.sample()\n",
        "        else:\n",
        "            state_a = np.array([self.state], copy=False)\n",
        "            state_v = torch.tensor(state_a)\n",
        "            if use_cuda:\n",
        "              state_v = state_v.cuda()\n",
        "            q_vals_v = net(state_v)\n",
        "            _, act_v = torch.max(q_vals_v, dim=1)\n",
        "            action = int(act_v.item())\n",
        "\n",
        "        # do step in the environment\n",
        "        new_state, reward, is_done, _ = self.env.step(action)\n",
        "        self.total_reward += reward\n",
        "\n",
        "        exp = Experience(self.state, action, reward, is_done, new_state)\n",
        "        self.exp_buffer.append(exp)\n",
        "        self.state = new_state\n",
        "        if is_done:\n",
        "            done_reward = self.total_reward\n",
        "            self._reset()\n",
        "        return done_reward\n",
        "\n",
        "    def play_random_step(self):\n",
        "        action = env.action_space.sample()\n",
        "\n",
        "        # do step in the environment\n",
        "        new_state, reward, is_done, _ = self.env.step(action)\n",
        "\n",
        "        exp = Experience(self.state, action, reward, is_done, new_state)\n",
        "        self.exp_buffer.append(exp)\n",
        "        self.state = new_state\n",
        "        if is_done:\n",
        "            self._reset()\n",
        "        return"
      ],
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZJrH9M2xiPey"
      },
      "source": [
        "## TD誤差の計算\n",
        "DQNでは，TD誤差と呼ばれる次状態の推定価値と，実際に選択した行動から得られる価値の差を用いて学習を行います．\n",
        "この時，次状態の推定価値は教師あり学習の教師と同じ役割を持ちます．\\\n",
        "DQNは，Q学習をもとにしているため，現在の行動価値関数を最適行動価値関数になるように更新を行っていきます．\n",
        "\n",
        "TD誤差の計算を行う関数を定義します．\\\n",
        "calc_loss関数では，`replay_buffer`からランダムに取得した`batch`分の経験をもとに，以下のLoss計算を行います．\n",
        "\n",
        "$$\n",
        "L_{\\theta}=(r+\\gamma \\max_{a'}Q_{\\theta_{tgt}}(s',a')-Q_{\\theta}(s,a))^{2}\n",
        "$$\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Jjm67ns-iUYR"
      },
      "source": [
        "def calc_loss(batch, net, tgt_net):\n",
        "    states, actions, rewards, dones, next_states = batch\n",
        "\n",
        "    states_v = torch.tensor(states)\n",
        "    next_states_v = torch.tensor(next_states)\n",
        "    actions_v = torch.tensor(actions)\n",
        "    rewards_v = torch.tensor(rewards)\n",
        "    done_mask = torch.BoolTensor(dones)\n",
        "    if use_cuda:\n",
        "      states_v = states_v.cuda()\n",
        "      next_states_v = next_states_v.cuda()\n",
        "      actions_v = actions_v.cuda()\n",
        "      rewards_v = rewards_v.cuda()\n",
        "      done_mask = done_mask.cuda()\n",
        "\n",
        "    state_action_values = net(states_v).gather(1, actions_v.unsqueeze(-1)).squeeze(-1)\n",
        "    next_state_values = tgt_net(next_states_v).max(1)[0]\n",
        "    next_state_values[done_mask] = 0.0\n",
        "    next_state_values = next_state_values.detach()\n",
        "\n",
        "    expected_state_action_values = next_state_values * GAMMA + rewards_v\n",
        "    return nn.MSELoss()(state_action_values, expected_state_action_values)"
      ],
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KMkQ1vGVxTUQ"
      },
      "source": [
        "## 学習\n",
        "DQNを用いて学習を行います．\n",
        "学習環境は，atari環境のPongゲーム環境を用います．\\\n",
        "各パラメータを定義します．\n",
        "Experience Replayで利用するReplay Bufferのサイズは1万とします．\n",
        "Replay Bufferが貯まりきるまで，ランダムに行動を選択し経験を収集，貯まりきってから学習を行います．\n",
        "\n",
        "学習回数を30万フレーム(`num_frame`)とし，環境の終了条件は，どちらかが21点とったら終了としています．\n",
        "また，最適化手法にはRMSprop利用します．\n",
        "\n",
        "学習を開始します．\n",
        "まず，環境を初期化し，経験をReplayBufferへ蓄積します．\n",
        "十分に蓄積された後，パラメータの更新を行います．\n",
        "また，`SNC_TARGET_FRAMES`で指定した回数ごとに，`target_net`のパラメータを`net`のパラメータと同じになるように同期を行います．\n",
        "\n",
        "`use_experience_replay`，`use_target_q_network`により，Experience replayとTarget Q-networkを使用するか指定が可能です．"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "GAMMA = 0.99\n",
        "BATCH_SIZE = 32\n",
        "REPLAY_SIZE = 10000\n",
        "LEARNING_RATE = 1e-4\n",
        "SYNC_TARGET_FRAMES = 1000\n",
        "REPLAY_START_SIZE = 10000\n",
        "\n",
        "EPSILON_DECAY_LAST_FRAME = 10**5\n",
        "EPSILON_START = 1.0\n",
        "EPSILON_FINAL = 0.02\n",
        "\n",
        "num_frame = 300000\n",
        "\n",
        "use_experience_replay = True\n",
        "use_target_q_network =  True"
      ],
      "metadata": {
        "id": "hsfDJsuB3XE4"
      },
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "学習を開始します．"
      ],
      "metadata": {
        "id": "bc7BfgPQFnen"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "Experience = collections.namedtuple('Experience', field_names=['state', 'action', 'reward', 'done', 'new_state'])\n",
        "net = DQN(env.observation_space.shape, env.action_space.n)\n",
        "tgt_net = DQN(env.observation_space.shape, env.action_space.n)\n",
        "\n",
        "if use_cuda:\n",
        "  net = net.cuda()\n",
        "  tgt_net = tgt_net.cuda()\n",
        "\n",
        "if use_experience_replay:\n",
        "  buffer = ExperienceBuffer(REPLAY_SIZE)\n",
        "else:\n",
        "  buffer = Buffer(REPLAY_SIZE)\n",
        "agent = Agent(env, buffer, seed)\n",
        "epsilon = EPSILON_START\n",
        "\n",
        "optimizer = optim.RMSprop(net.parameters(), lr=LEARNING_RATE)\n",
        "total_rewards = []\n",
        "frame_idx = 0\n",
        "best_mean_reward = None\n",
        "\n",
        "record_reward = []\n",
        "record_step = []\n",
        "\n",
        "print(\"Collect experience\")\n",
        "while len(buffer) < REPLAY_START_SIZE:\n",
        "  agent.play_random_step()\n",
        "agent._reset()\n",
        "print(\"Filled buffer {}\".format(len(buffer)))\n",
        "\n",
        "print(\"Leaning start\")\n",
        "ts = time.time()\n",
        "while frame_idx < num_frame:\n",
        "    frame_idx += 1\n",
        "    epsilon = max(EPSILON_FINAL, EPSILON_START - frame_idx / EPSILON_DECAY_LAST_FRAME)\n",
        "\n",
        "    reward = agent.play_step(net, epsilon)\n",
        "    if reward is not None:\n",
        "        total_rewards.append(reward)\n",
        "        mean_reward = np.mean(total_rewards[-100:])\n",
        "        print(\"Frame {0}/{1}: episode {2}, reward {3}, mean reward {4:.3f}, eps {5:.2f}, time {6}\".format(frame_idx, num_frame, len(total_rewards), reward, mean_reward, epsilon, datetime.timedelta(seconds = time.time() - ts)))\n",
        "        record_reward.append(reward)\n",
        "        record_step.append(frame_idx)\n",
        "\n",
        "        if best_mean_reward is None or best_mean_reward <= mean_reward:\n",
        "            torch.save(net.state_dict(), \"Pong-best.dat\")\n",
        "            if best_mean_reward is not None:\n",
        "                print(\"Best mean reward updated {0:.3f} -> {1:.3f}, model saved\".format(best_mean_reward, mean_reward))\n",
        "            best_mean_reward = mean_reward\n",
        "        if len(total_rewards) == 1000000:\n",
        "            print(\"Solved in {0} frames!\".format(frame_idx))\n",
        "            break\n",
        "\n",
        "    if (frame_idx % SYNC_TARGET_FRAMES == 0) and use_target_q_network:\n",
        "        sync_network()\n",
        "\n",
        "    optimizer.zero_grad()\n",
        "    batch = buffer.sample(BATCH_SIZE)\n",
        "    if use_target_q_network:\n",
        "      loss_t = calc_loss(batch, net, tgt_net)\n",
        "    else:\n",
        "      loss_t = calc_loss(batch, net, net)\n",
        "    loss_t.backward()\n",
        "    optimizer.step()"
      ],
      "metadata": {
        "id": "Y8XuD9g1psAd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 学習時のスコアの推移\n",
        "横軸エピソード数，縦軸スコアとしたグラフを描画してみます．"
      ],
      "metadata": {
        "id": "BSHnuplKQo14"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import matplotlib.pyplot as plt\n",
        "\n",
        "fig = plt.figure()\n",
        "plt.plot(record_step, record_reward, color=\"red\")\n",
        "plt.grid()\n",
        "plt.xlabel(\"step\")\n",
        "plt.ylabel(\"reward\")\n",
        "plt.savefig(\"./dqn_step_per_reward.png\")\n",
        "plt.show()\n",
        "plt.clf()\n",
        "plt.close()"
      ],
      "metadata": {
        "id": "jUwFK70Qptgy"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "グラフのダウンロード"
      ],
      "metadata": {
        "id": "aDQWPHZcQxbF"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import files\n",
        "files.download(\"./dqn_step_per_reward.png\")"
      ],
      "metadata": {
        "id": "ywwlV574oF2h"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ksjONOyoxstC"
      },
      "source": [
        "## 学習結果の可視化\n",
        "学習したネットワーク（エージェント）を確認してみます．\n",
        "\n",
        "ここでは，Gym-Notebook-Wrapper (gnwrapper)を用いて学習済みエージェントによるゲームプレイを可視化します．"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "プログラム演習中に十分な学習時間が確保できない場合は，下記のプログラムを実行して学習済みモデルを用いて動作を確認します．"
      ],
      "metadata": {
        "id": "4W8x6e9tVgFa"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# 学習済みモデルのダウンロード\n",
        "!wget -qq http://www.mprg.cs.chubu.ac.jp/~itaya/share/mprg_colab/DQN/trained_model.zip\n",
        "!unzip trained_model.zip\n",
        "\n",
        "# 学習済みモデルの読み込み\n",
        "net = DQN(env.observation_space.shape, env.action_space.n)\n",
        "net.load_state_dict(torch.load(\"./trained_model/Pong-best.dat\", weights_only=True))\n",
        "if use_cuda:\n",
        "  net = net.cuda()"
      ],
      "metadata": {
        "id": "HAoVDTvApu-C"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import gnwrapper\n",
        "\n",
        "env = gnwrapper.Monitor(env,directory=\"./\") # 環境の作成\n",
        "\n",
        "done = False\n",
        "t = 0\n",
        "state = env.reset(seed=seed)\n",
        "\n",
        "while not done:\n",
        "    state_a = np.array([state], copy=False)\n",
        "    state_v = torch.tensor(state_a)\n",
        "    if use_cuda:\n",
        "      state_v = state_v.cuda()\n",
        "    q_vals_v = net(state_v)\n",
        "    _, act_v = torch.max(q_vals_v, dim=1)\n",
        "    action = int(act_v.item())\n",
        "\n",
        "    new_state, reward, done, info = env.step(action)\n",
        "    state = new_state\n",
        "    t += 1\n",
        "\n",
        "env.display()"
      ],
      "metadata": {
        "id": "NPWqV6M3pwMr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kZDh4S6m2hK0"
      },
      "source": [
        "# 課題\n",
        "\n",
        "1. Experience replayやTarget Q-networkの有無による学習効率の違いを見てみましょう．\n",
        " * `use_experience_replay`，`use_target_q_network`によって使用するか指定可能です．\n",
        " * Pongは，デフォルトの報酬設計が-1，0，1のため，Reward clippingの有無による違いを確認する必要はありません．\n",
        "\n",
        "2. Pong以外のゲームで学習してみましょう．\n",
        " * gym.make()での環境の指定を変更することで，任意の環境で学習できます．\n",
        " * 指定できる環境は，`ALE/Breakout-v5`や`ALE/MsPacman-v5`などがあります．詳しくは[ドキュメント](https://ale.farama.org/environments/)をチェックしてください．\n",
        " * 環境を変更すると学習に必要な時間も変わります．"
      ]
    }
  ]
}