MCTSにおける価値の漸進的更新

結論

　MCTSの行動価値を漸進的に更新する実装で、総和を保持して平均化する実装と同程度の性能を達成できた。

背景

　以前、MCTSにおいて行動価値を漸進的に更新する方法について記事を書いたが、性能が悪化してしまった。この記事で述べた通り、原因はおそらくバーチャルロスを探索した回数 $N$ へ直接足しており、正しい探索回数になっていないためだと思われる。よってバーチャルロスを探索した回数 $N$ とは別に保持するように変更した。

実装

置換表エントリ

　各行動についてバーチャルロスを保持する変数を置換表エントリに加える。単純にこの変数を増やしただけだと使用メモリ量が増えてしまうが、行動価値を漸進的に更新する場合、子孫ノードの価値の総和 $W$ を保持する必要がなくなる。結局使用メモリ量はほぼ同じとなる。置換表エントリの主要なデータは次のようになる。std::vectorで保持されているもののサイズは合法手の数と一致する。

struct UctHashEntry {
    int32_t sum_N;                    //この局面から探索した回数
    int32_t virtual_sum_N;            //バーチャルロスの合計
    std::vector<Move> moves;          //この局面の合法手
    std::vector<Index> child_indices; //子局面の置換表におけるインデックス
    std::vector<int32_t> N;           //各行動を探索した回数
    std::vector<int32_t> virtual_N;   //各行動へのバーチャルロス
    std::vector<CalcType> nn_policy;  //ニューラルネットワークの出力Policy
    ValueType value;                  //状態価値:漸進的に更新される
};

行動価値の計算

　行動価値は次状態の価値を参照することで得る。次状態の価値を参照する際に、手番が変わるので反転させる必要があることに注意。

ValueType Searcher::QfromNextValue(const UctHashEntry& node, int32_t i) const {
    return -hash_table_[node.child_indices[i]].value;
}

選択ステップ

　選択ステップでは探索回数 $N$ とバーチャルロスの和を用いてUCB値を計算する。

int32_t Searcher::selectMaxUcbChild(const UctHashEntry& node) {
    constexpr double C_PCUT = 1.0;
    const int32_t sum = node.sum_N + node.virtual_sum_N;

    int32_t max_index = -1;
    double max_value = MIN_SCORE - 1;
    for (int32_t i = 0; i < node.moves.size(); i++) {
        int32_t visit_num = node.N[i] + node.virtual_N[i];

        double Q = QfromNextValue(node, i);
        double U = std::sqrt(sum + 1) / (visit_num + 1);
        double ucb = Q + C_PUCT * node.nn_policy[i] * U;

        if (ucb > max_value) {
            max_value = ucb;
            max_index = i;
        }
    }
    return max_index;
}

バックアップ

　バックアップ時にバーチャルロスを0へと戻していく。実際の探索回数である $N$ やその総和は1足すだけで良い。ここで価値を漸進的に更新する。更新幅は探索回数 + 1の逆数となる。

void SearcherForPlay::backup(std::stack<int32_t>& indices, 
                             std::stack<int32_t>& actions) {
    int32_t leaf = indices.top();
    indices.pop();
    float leaf_value = hash_table_[leaf].value;

    while (!actions.empty()) {
        int32_t index = indices.top();
        indices.pop();
        int32_t action = actions.top();
        actions.pop();
        UctHashEntry& node = hash_table_[index];

        //手番が変わるので価値を反転
        leaf_value = MAX_SCORE + MIN_SCORE - leaf_value;

        //探索回数の更新
        node.N[action]++;
        node.sum_N++;
        node.virtual_sum_N -= node.virtual_N[action];
        node.virtual_N[action] = 0;

        //漸進的更新
        float alpha = 1.0f / (node.sum_N + 1);
        node.value += alpha * (leaf_value - node.value);
    }
}

気づいたこと

　以前の実装のようにバーチャルロスを直接 $N$ に足していると、同じバッチ内で次にそのノードへ到達したとき $N$ が本来よりも大きくなっているので $Q = \frac{W}{N}$ が小さく評価されてしまう。今回のようにバーチャルロスを別に保持することでこの欠点も直すことができる。もっともこの欠点は探索回数が多くなれば無視できるほど小さい影響しかもたらさないだろうし、むしろ同じノードが選ばれにくくなるため良い効果があるかもしれない。

実験結果

　漸進的な更新へと変えた実装の探索部を用いて技巧2(深さ1)と１手１秒で200局対局させた。

探索手法	勝率
総和を保持(従来の実装)	81.0%
漸進的更新(今回の実装)	86.5%

　誤差の範囲内かもしれないが、漸進的更新の方がやや強くなる結果となった。強くなる要因としては、同一局面へ違う進行で合流した際に良い価値を返すことができている可能性が考えられる。従来の実装では保存しておいたニューラルネットワークの出力のValueをそのままバックアップしていくことになるが、今回の実装では漸進的に更新をしているため単なるニューラルネットワークの出力値よりは良い推定値になっていると考えられる。しかしそこまで大幅な性能向上でもなく、弱くなっていないことが確認できたことが重要となる。

補足

　置換表エントリにはバーチャルロスをint32_t型で確保したが、この変数は最大でも探索バッチ数までしか増えないので4Byteも使う必要はない。1Byteでは不安だが、2Byteあれば十分だと思われる。そうすることで従来の実装では価値の総和として4Byteの浮動小数点型を確保する必要があったのに対して多少のメモリ消費量削減となる。

　(個人的に)もっと重要な点は、Categorical分布を保持するときのメモリ消費量が抑えられることである。Miacisの手法では51分割していたため $W$ を保持するために通常のAlphaZeroよりも約51倍のメモリが必要になっていた。これにより秒数制限よりも置換表のサイズ制限の方へ速く到達してしまうような状況だったのだが、今回の実装によって使用メモリ量がだいたい1/50となったため問題なくなった。今までは学習部においてもメモリ消費量の問題からActorの並列数を増やせなかったのだが、今後は並列数を増やすことで多少データ生成速度が上がるのではないかと思われる。