麻雀AIを深層強化学習で作る その１１(学習処理)

前回、自己対局で特徴量を作成する処理を実装した。

今回は、自己対局で生成したデータを読み込んで学習する処理を実装した。

データ読み込み

自己対局プログラムでzlibで圧縮したデータをPythonのzlibで解凍し、Numpyのndarrayとして読み込む。
C++の構造体をnumpyのdtypeとして定義しておくことで、「np.frombuffer(data, StepData)」とするだけで構造化されたデータを読み込むことができる。

Numpyの構造体定義

PublicFeatures = np.dtype((np.float32, (N_CHANNELS_PUBLIC + 4, 9, 4)))
PrivateFeatures = np.dtype((np.float32, (N_CHANNELS_PRIVATE, 9, 4)))
Policy = np.dtype((np.float32, N_ACTIONS))
Hupai = np.dtype((np.float32, 54))
HulePlayer = np.dtype((np.float32, 5))
TajiaTingpai = np.dtype((np.float32, (3, 34)))
Fenpei = np.dtype((np.float32, 4))
StepData = np.dtype(
    [
        ("public_features", PublicFeatures),
        ("private_features", PrivateFeatures),
        ("action", np.dtype((np.int64, (1, )))),
        ("value", np.float32),
        ("logits", Policy),
        ("advantage", np.float32),
        ("hupai", Hupai),
        ("hule_player", HulePlayer),
        ("tajia_tingpai", TajiaTingpai),
        ("fenpei", Fenpei),
    ]
)

データローダ

Numpyで読み込んだデータをPyTorchのデータセットでラップして、PyTorchのデータローダで読み込めるようにする。

データセット定義

class RolloutDataset(Dataset):
    def __init__(self):
        super().__init__()
        self.rollout_data = np.empty(0, StepData)

    def load(self, path):
        logging.info(f"loading {path}")
        with open(path, "rb") as f:
            data = zlib.decompress(f.read())
        tmp = np.frombuffer(data, StepData)
        self.rollout_data = np.concatenate((self.rollout_data, tmp))

    def calc_log_prob(self, batch_size, device):
        self.log_prob = np.empty(len(self.rollout_data), np.float32)
        with torch.inference_mode():
            for i in range(0, len(self.rollout_data), batch_size):
                batch = self.rollout_data[i:i+batch_size]
                actions = torch.from_numpy(batch["action"]).to(device)
                logits = torch.from_numpy(batch["logits"]).to(device)
                log_prob = PolicyValueNetWithAux.log_prob(actions, logits)
                self.log_prob[i:i+batch_size] = log_prob.to("cpu").detach().numpy()

    def __len__(self):
        return len(self.rollout_data)

    def __getitem__(self, idx):
        data = self.rollout_data[idx]
        log_prob = self.log_prob[idx]
        
        return data["public_features"], data["private_features"], data["action"], log_prob, data["advantage"], data["advantage"] + data["value"], data["hupai"]

log probability

PPOでは、挙動方策の選択したアクションの確率と、現在の方策のアクションの確率の比の対数を算出する。
実装上は、それぞれの確率の対数の差を計算する。

自己対局では、ロジットを記録しているため、挙動方策の確率の対数(log probability)をあらかじめ計算しておく。
上記、データセットのcalc_log_probメソッドに実装している。

学習処理

PPOアルゴリズムで学習する処理を実装する。

for epoch in range(args.n_epochs):
    logging.info(f"training epoch {epoch}")
    for public_features, private_features, actions, old_log_prob, advantages, returns, hupai in rollout_buffer:
        values, log_prob, entropy, p_aux1, p_aux2, p_aux3, v_aux = model.evaluate_actions_with_aux(public_features, private_features, actions)

        values = values.flatten()

        # Normalize advantage
        if normalize_advantage:
            advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)

        # ratio between old and new policy, should be one at the first iteration
        ratio = torch.exp(log_prob - old_log_prob)

        # clipped surrogate loss
        policy_loss_1 = advantages * ratio
        policy_loss_2 = advantages * torch.clamp(ratio, 1 - clip_range, 1 + clip_range)
        policy_loss = -torch.min(policy_loss_1, policy_loss_2).mean()

        # Value loss using the TD(gae_lambda) target
        value_loss = F.mse_loss(returns, values)

        # Entropy loss favor exploration
        entropy_loss = -torch.mean(entropy)

        # 補助タスク1 役
        p_aux1_loss = bce_with_logits_loss(p_aux1, hupai)

        loss = policy_loss + ent_coef * entropy_loss + vf_coef * value_loss + p_aux1_loss

        # Optimization step
        optimizer.zero_grad()
        loss.backward()
        # Clip grad norm
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
        optimizer.step()

補助タスク

モデルに定義した補助タスクの内、今回は、和了時の役の予測のみ実装した。

学習結果

自己対局で生成したデータで学習を実行し、自己対局→学習のサイクルを4サイクル回して、訓練損失を確認した。

1サイクルで、約45万局面を生成し、各サイクルで10エポック学習した。
自己対局では、ランダムに選択した役から3向聴前の局面を初期局面とする。

訓練損失 合計

方策損失

価値損失

エントロピー損失

補助タスク(役の予測)の損失

各サイクルでグラフの色が分かれている。

訓練損失の合計は、1サイクル目が一番低く、2サイクル目で上昇してそこから徐々に下がっている。
1サイクル目は方策がランダムに近く、合法手を学習するだけで、方策はランダムに近いため、エントロピー損失が大きく下がったためと考えられる。
2サイクル目からは少しずつ意味のある方策学習され始めている。

方策損失は、4サイクル目が高い値になっているが、補助タスクの学習が先に進んだためと思われる。
価値損失は、4サイクル目が一番低い値となっており、局面の勝率を学習できていそうである。

補助タスクの損失は、4サイクル目が一番低い値となっており、局面から役の予測ができていそうである。
方策と補助タスクは、方策ネットワークのヘッドにあるため、損失はサイクルごとに一方が下がるともう一方が高くなっており、交互に学習が進んでいるようである。

まとめ

麻雀AIの学習処理を実装した。
実際に自己対局でデータを生成して学習するサイクルを4サイクルだけ実行して、うまく学習できているか確認した。
サイクルごとに訓練損失が下がることが確認できた。

次は、学習量を増やして、打牌や予測勝率、役の予測結果をサンプリングして確認してみたい。