概要

深層強化学習によってリバーシAIを実装し、GPUで約20分の学習で対ランダム勝率約95%を達成しました。

こんな方におすすめ

深層強化学習の実践的な実装例を探している方
ゲームAIに興味がある方

本記事のポイント

DQNのふんわり解説
リバーシ特有の実装上の注意点と解決方法
サンプルとして複数モデルの実装と性能比較

必要な知識：PythonとPyTorch（基礎程度）

※ 強化学習についてはまだまだ学習中なので、間違いを含む可能性があります。

参考：PyTorch DQN Tutorial
ソースコード：GitHub

DQNについて

DQNのお気持ち

学習ではある状態sについてある行動aで得られる報酬をQ(s,a)として学習することを目的とします。
例えばリバーシでは状態が盤面の様子、行動がどこに石を置くかに相当します。

深層強化学習ではQをニューラルネットワークで近似します。
リバーシの場合、状態sは盤面、行動aは着手可能な位置となり、Q値は最終的な勝敗に基づいて学習されます。
Q値の更新にはベルマン方程式を用います：
状態sから行動aによって報酬rが得られて次の状態s'に移るとき

Q(s,a) = r + \gamma \max_{a'}Q(s',a')

ガンマは割引率と呼ばれています。1より少し小さい値が用いられます。

maxQ(s',a')は次の状態において得られそうな報酬なので、「今の状態で得られそうな報酬は今すぐ得られる報酬と未来に得られそうな報酬を少し割り引いたものである」みたいな気持ちだと理解しています。
割り引く理由は将来の不確実性や報酬の収束性、即時報酬の重視が挙げられます。
将来のことは正確には分からない、報酬が無限にもらえるような自体を回避する、めちゃくちゃ後に報酬をもらえるのは本当に良いのか、みたいな感じだと思います。
AlphaZeroでは割引率を1にして割り引かないでやっているみたいです。

DQNでは右辺と左辺の差を損失として学習します。

今回のリバーシでのDQN

今回は即時報酬を勝敗がついたときのみ与えています。

勝ったときに1、負けたときに0、引き分けで0.5にしました。
それ以外の即時報酬は0にしています。

なぜこれでうまくいくのかというとゲームの後半のQから学習され、割と正しくなった後半のQで中盤のQが学習されるということが起こるからだと思います。

リバーシのDQNの注意点

誰向きの評価なのか注意する

モデルへの入力は自分の石の配置、相手の石の配置を並べたものにしています。
ここで次の状態s'では手番が逆転しているのでmaxQ(s',a')の価値はsの手番の逆の人からみた価値となっています。
したがってモデルの出力を逆転させる必要があります。
これは手番が逆転するゲームだからこそ起こることでPytorchのtutorialにはなく何度もここでハマりました。

パスを1つの行動に含める

リバーシはどこに石を置くかという選択で行動が決まるので行動の数が将棋やチェスなどと比べてとても少ないです。また置ける状態でパスというのは合法ではないことが多いと思います。(今回はそうしています。)

そこで行動数は64にしたくなるのですが、パスを含めた65にする必要があります。
これはパスしかないときの評価値をモデルに出力してしてもらう必要があるためです。

自分はここでバグらせて損失が1e7になっていました。

実際のコード (v_nsはmaxQ(s',a')のことです。)

agents/__init__.py

q_ns: torch.Tensor = self.target_net(next_states_t)
# 64th element is pass
# pass is only legal when the player has no legal moves
legal_actions: torch.Tensor = torch.tensor(
    [ns.get_legal_moves_tf() + [ns.is_pass()] for ns in next_states],
    dtype=torch.bool,
    device=self.config["device"]
)
q_ns = q_ns.masked_fill(~legal_actions, -1e9)
v_ns: torch.Tensor = q_ns.max(1).values
# The value of the next state is the value of the opponent (1 - value is the value of the player)
v_ns = 1.0 - v_ns
game_overs = torch.tensor([ns.is_game_over() for ns in next_states], dtype=torch.bool, device=self.config["device"])
v_ns = v_ns.masked_fill(game_overs, 0.0)    # If the game is over, the value of the next state is 0 and the reward is the final reward

実際の実装

ゲームエンジンについて

強化学習では学習に時間がかかる場合が多くリバーシの効率的な実装は時間短縮の面で重要です。
今回は自分がRustで実装したリバーシをPythonから呼び出すことで高速？なゲームの実行を実現しています。
少なくともPythonの普通の実装よりは早いと思います。

pip install rust-reversi

上のようにpipなどでインストールできるのでぜひ使ってください。

rust-reversi

モデル

モデルには10層のResNetを使用しています。
ボードゲームAIでは畳み込み層を用いたモデルが使用されることが多いようです。

またDueling Networkと呼ばれる構造も採用しました。
最後に状態の価値v(s)とそれぞれの行動の良さa(s,a)を表してほしいという願いを込めて全結合層を分離すると性能が向上するそうです。不思議です。

models/resnet.py

class ResNet10(torch.nn.Module):
    def __init__(self, num_channels, fc_hidden_size):
        super(ResNet10, self).__init__()
        self.conv1 = torch.nn.Conv2d(2, num_channels, kernel_size=3, padding=1, bias=False)
        self.bn1 = torch.nn.BatchNorm2d(num_channels)
        self.res_blocks = torch.nn.ModuleList([ResBlock(num_channels) for _ in range(10)])
        self.fc1 = torch.nn.Linear(num_channels * 8 * 8, fc_hidden_size)
        self.fc2_adv = torch.nn.Linear(fc_hidden_size, OUTPUT_SIZE)
        self.fc2_val = torch.nn.Linear(fc_hidden_size, 1)
        self.num_channels = num_channels
        self.relu = torch.nn.ReLU()
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if len(x.shape) == 3:
            x = x.unsqueeze(0)
        x = self.bn1(self.conv1(x))
        x = self.relu(x)
        for res_block in self.res_blocks:
            x = res_block(x)
        x = x.view(-1, self.num_channels * 8 * 8)
        x = self.fc1(x)
        x = self.relu(x)
        adv = self.fc2_adv(x)
        val = self.fc2_val(x).expand(-1, OUTPUT_SIZE)
        return val + adv - adv.mean(1, keepdim=True).expand(-1, OUTPUT_SIZE)

Prioritized Experience Replay

Pytorchのtutorialでは経験をメモリから取り出す部分では完全にランダムに取り出しています。
ここの確率に重みをつけてサンプリングすることで特定の経験を優先して学習させることができます。
今回は損失が大きいものほど高い確率で再生されるようにしています。
モデルの苦手な部分を何度も復習させてやろうという作戦です。

優先度付きサンプリングの効率化

優先度を計算してサンプリングするのは結構時間がかかってしまいます。
具体的にはそれぞれの損失の値をp_iとして

P(i) = \frac{p_i}{\sum_{j} p_j}

を計算する必要があります。

この問題はsum-treeと呼ばれるものを使用する事により解決できます。
実際以下のように学習時間を改善できました。
約20分(PERなし)->約40分(PERありsum-treeなし)->約20分(PERありsum-treeあり)

ゲームの同時実行

モデルに対して1つの状態について推論させるのを10回繰り返すより、10の状態について1回推論させるほうがはやく計算できます。
そのためゲームを複数同時実行してそれぞれのゲームについてどこに置くかをモデルに一度に推論させるようにしました。
私の環境では一度に240ゲームを実行させています。

同時実行の実装

agents/batch_board.py

from rust_reversi import Board

class BatchBoard:
    def __init__(self, batch_size: int):
        self.batch_size = batch_size
        self.boards: list[Board] = [Board() for _ in range(batch_size)]

    def get_boards(self) -> list[Board]:
        return [board.clone() for board in self.boards]

    def do_move(self, moves: list[int]) -> tuple[list[Board], list[float]]:
        # list for self.boards that are not game over
        new_boards = []
        # return new_boards, rewards
        next_boards = []
        rewards = []
        for board, move in zip(self.boards, moves):
            if move == 64:
                board.do_pass()
            else:
                board.do_move(move)
            next_boards.append(board.clone())
            if board.is_game_over():
                if board.is_win():
                    # turn swapped in do_move, so is_win menas the player that just moved lost
                    rewards.append(0.0)
                elif board.is_lose():
                    # turn swapped in do_move, so is_lose menas the player that just moved won
                    rewards.append(1.0)
                else:
                    # draw
                    rewards.append(0.5)
            else:
                rewards.append(0.0)
                # only append if not game over
                new_boards.append(board.clone())
        self.boards = new_boards
        return next_boards, rewards

    def is_game_over(self) -> bool:
        return len(self.boards) == 0

学習の様子

より簡単なモデルについても学習の様子を載せておきます。
自分の環境は13700+4070Tiです。
ハイパーパラメーターについては適当に決定しているためモデルごとにさらに適切なものがある可能性があります。

ResNet10

学習には23分38秒かかりました。

損失

対ランダム勝率の推移

Episode 0: Win rate vs random = 0.478
Episode 12000: Win rate vs random = 0.643
Episode 24000: Win rate vs random = 0.698
Episode 36000: Win rate vs random = 0.781
Episode 48000: Win rate vs random = 0.741
Episode 60000: Win rate vs random = 0.843
Episode 72000: Win rate vs random = 0.882
Episode 84000: Win rate vs random = 0.918
Episode 96000: Win rate vs random = 0.926
Episode 108000: Win rate vs random = 0.949
Training finished
Win rate vs random = 0.946

config

{'batch_size': 512,
 'board_batch_size': 240,
 'device': device(type='cuda'),
 'episodes_per_optimize': 2,
 'episodes_per_target_update': 4,
 'eps_decay': 10,
 'eps_end': 0.05,
 'eps_start': 1.0,
 'fc_hidden_size': 256,
 'gamma': 0.99,
 'gradient_clip': 1.0,
 'lr': 1e-05,
 'memory_config': {'alpha': 0.5,
                   'beta': 0.5,
                   'memory_size': 24000,
                   'memory_type': <MemoryType.PROPORTIONAL: 0>},
 'model_path': 'cnn_agent.pth',
 'n_episodes': 120000,
 'num_channels': 64,
 'verbose': True}

Conv5Dueling

学習には15分31秒かかりました。
ResNetより性能が良くなっています。
こちらのほうが良いのかもしれません。
ただもっと学習時間を増やせばResNetのほうが性能が良くなる可能性はあると思います。
ResNetで1回しか試していませんがResNetで3時間学習したときの勝率は98.3%を記録したことがあります。(ResNet擁護派)

損失

対ランダム勝率の推移

Episode 0: Win rate vs random = 0.603
Episode 12000: Win rate vs random = 0.585
Episode 24000: Win rate vs random = 0.77
Episode 36000: Win rate vs random = 0.85
Episode 48000: Win rate vs random = 0.919
Episode 60000: Win rate vs random = 0.942
Episode 72000: Win rate vs random = 0.948
Episode 84000: Win rate vs random = 0.948
Episode 96000: Win rate vs random = 0.97
Episode 108000: Win rate vs random = 0.972
Training finished
Win rate vs random = 0.974

config

ResNetと同じです。

モデルの実装

models/cnn_dueling.py

import torch

DROPOUT = 0.0
# 8x8 board + 1 for pass
OUTPUT_SIZE = 65

class ConvLayer(torch.nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, padding):
        super(ConvLayer, self).__init__()
        self.conv = torch.nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, padding=padding)
        self.bn = torch.nn.BatchNorm2d(out_channels)
        self.relu = torch.nn.ReLU()
        self.dropout = torch.nn.Dropout(DROPOUT)
    
    def forward(self, x):
        x = self.bn(self.conv(x))
        x = self.relu(x)
        x = self.dropout(x)
        return x

class Conv5DuelingNet(torch.nn.Module):
    def __init__(self, num_channels, fc_hidden_size):
        super(Conv5DuelingNet, self).__init__()
        self.conv1 = ConvLayer(2, num_channels, kernel_size=3, padding=1)
        self.conv2 = ConvLayer(num_channels, num_channels, kernel_size=3, padding=1)
        self.conv3 = ConvLayer(num_channels, num_channels, kernel_size=3, padding=1)
        self.conv4 = ConvLayer(num_channels, num_channels, kernel_size=3, padding=1)
        self.conv5 = ConvLayer(num_channels, num_channels, kernel_size=3, padding=1)
        self.fc1 = torch.nn.Linear(num_channels * 8 * 8, fc_hidden_size)
        self.fc2_adv = torch.nn.Linear(fc_hidden_size, OUTPUT_SIZE)
        self.fc2_val = torch.nn.Linear(fc_hidden_size, 1)
        self.relu = torch.nn.ReLU()
        self.dropout = torch.nn.Dropout(DROPOUT)
        self.num_channels = num_channels

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if len(x.shape) == 3:
            x = x.unsqueeze(0)
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        x = self.conv4(x)
        x = self.conv5(x)
        x = x.view(-1, self.num_channels * 8 * 8)
        x = self.fc1(x)
        x = self.relu(x)
        x = self.dropout(x)

        adv = self.fc2_adv(x)
        val = self.fc2_val(x).expand(-1, OUTPUT_SIZE)

        return val + adv - adv.mean(1, keepdim=True).expand(-1, OUTPUT_SIZE)

Conv5

学習には15分03秒かかりました。
またもやResNetより性能が良いです。
ResNet擁護派としては厳しくなってきたかもしれません。

損失

対ランダム勝率の推移

Episode 0: Win rate vs random = 0.473
Episode 12000: Win rate vs random = 0.558
Episode 24000: Win rate vs random = 0.726
Episode 36000: Win rate vs random = 0.702
Episode 48000: Win rate vs random = 0.9
Episode 60000: Win rate vs random = 0.92
Episode 72000: Win rate vs random = 0.945
Episode 84000: Win rate vs random = 0.962
Episode 96000: Win rate vs random = 0.97
Episode 108000: Win rate vs random = 0.965
Training finished
Win rate vs random = 0.974

config

ResNetと同じです。

モデルの実装

models/cnn.py

import torch

# 8x8 board + 1 for pass
OUTPUT_SIZE = 65
DROP_RATE = 0.0

class Conv5Net(torch.nn.Module):
    def __init__(self, num_channels, fc_hidden_size):
        super(Conv5Net, self).__init__()
        self.conv1 = torch.nn.Conv2d(2, num_channels, kernel_size=3, padding=1)
        self.bn1 = torch.nn.BatchNorm2d(num_channels)
        self.conv2 = torch.nn.Conv2d(num_channels, num_channels, kernel_size=3, padding=1)
        self.bn2 = torch.nn.BatchNorm2d(num_channels)
        self.conv3 = torch.nn.Conv2d(num_channels, num_channels, kernel_size=3, padding=1)
        self.bn3 = torch.nn.BatchNorm2d(num_channels)
        self.conv4 = torch.nn.Conv2d(num_channels, num_channels, kernel_size=3, padding=1)
        self.bn4 = torch.nn.BatchNorm2d(num_channels)
        self.conv5 = torch.nn.Conv2d(num_channels, num_channels, kernel_size=3, padding=1)
        self.bn5 = torch.nn.BatchNorm2d(num_channels)
        self.fc1 = torch.nn.Linear(num_channels * 8 * 8, fc_hidden_size)
        self.fc2 = torch.nn.Linear(fc_hidden_size, OUTPUT_SIZE)
        self.relu = torch.nn.ReLU()
        self.dropout = torch.nn.Dropout(DROP_RATE)
        self.num_channels = num_channels

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if len(x.shape) == 3:
            x = x.unsqueeze(0)
        x = self.bn1(self.conv1(x))
        x = self.relu(x)
        x = self.dropout(x)
        x = self.bn2(self.conv2(x))
        x = self.relu(x)
        x = self.dropout(x)
        x = self.bn3(self.conv3(x))
        x = self.relu(x)
        x = self.dropout(x)
        x = self.bn4(self.conv4(x))
        x = self.relu(x)
        x = self.dropout(x)
        x = self.bn5(self.conv5(x))
        x = self.relu(x)
        x = self.dropout(x)
        x = x.view(-1, self.num_channels * 8 * 8)
        x = self.fc1(x)
        x = self.dropout(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

Dense

学習には33分52秒かかりました。
またもやResNetより成績が良いです。
ResNet擁護派はもうだめかも知れません。

損失

対ランダム勝率の推移

Episode 0: Win rate vs random = 0.396
Episode 12000: Win rate vs random = 0.645
Episode 24000: Win rate vs random = 0.743
Episode 36000: Win rate vs random = 0.804
Episode 48000: Win rate vs random = 0.846
Episode 60000: Win rate vs random = 0.873
Episode 72000: Win rate vs random = 0.914
Episode 84000: Win rate vs random = 0.914
Episode 96000: Win rate vs random = 0.943
Episode 108000: Win rate vs random = 0.953
Training finished
Win rate vs random = 0.959

config

{'batch_size': 512,
 'board_batch_size': 240,
 'device': device(type='cuda'),
 'episodes_per_optimize': 2,
 'episodes_per_target_update': 4,
 'eps_decay': 10,
 'eps_end': 0.05,
 'eps_start': 1.0,
 'gamma': 0.99,
 'gradient_clip': 1.0,
 'hidden_size': 256,
 'lr': 1e-05,
 'memory_config': {'alpha': 0.5,
                   'beta': 0.5,
                   'memory_size': 24000,
                   'memory_type': <MemoryType.PROPORTIONAL: 0>},
 'model_path': 'dense_agent.pth',
 'n_episodes': 120000,
 'verbose': True}

モデルの実装

models/dense.py

import torch

INPUT_SIZE = 128
# 8x8 board + 1 for pass
OUTPUT_SIZE = 65

class DenseNet(torch.nn.Module):
    def __init__(self, hidden_size: int):
        super(DenseNet, self).__init__()
        self.fc1 = torch.nn.Linear(INPUT_SIZE, hidden_size)
        self.fc2 = torch.nn.Linear(hidden_size, hidden_size)
        self.fc3 = torch.nn.Linear(hidden_size, OUTPUT_SIZE)
        self.relu = torch.nn.ReLU()

    def forward(self, x) -> torch.Tensor:
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        x = self.relu(x)
        x = self.fc3(x)
        return x

ラッキー

ここまでみてくださってありがとうございます🤗

今後の方針

記事の執筆と同時に学習をしていたらResNetより性能のいいモデルが出てきてしまいました。モデルの検証と、ハイパーパラメーターの検証についてはまだまだなのでやっていきたいです。
ResNetのほうが表現力が高いはずなので時間をかけるとより強くなりやすいことを期待しています。
Attention機構をつけたモデルも試してみたいと思っています。

また強化学習で作成したモデルを生かして蒸留と呼ばれるものをやってみたいと思っています。
最終的にはRustで実装した全結合モデルで高速に探索するところまでやってみたいです。

ReversiはCodingameというサイトに提出することで自分のAIの強さを評価できます。
Rustでの実装はここに提出するためにやろうと思っています。

自分はCondingameを宣伝するほど何もなしていませんが、興味があればぜひ「俺の作った†最強†のAI」を提出しましょう。

まとめ

深層強化学習の実装例としてリバーシを紹介しました。
個人的には3目並べを学習しても「うーん」って感じですがリバーシまでいくと「すごー」ってなるので学習が成功して嬉しいです。
しかし将棋などは実装が難しく状態も複雑なためリバーシの「ちょうど良さ」が自分にとっては良かったのでオススメです。
実装するときの目安に使っていただけたら嬉しいです。

参考

今回のソース
Reinforcement Learning (DQN) Tutorial
rust-reversi
リバーシで深層強化学習
非常に参考になりました。
~~このQiitaの記事いらなくないですか~~
sum tree