diff --git a/src/mcts/node.cc b/src/mcts/node.cc index ea97edd368..9235bbcda9 100644 --- a/src/mcts/node.cc +++ b/src/mcts/node.cc @@ -39,6 +39,7 @@ #include "neural/network.h" #include "utils/exception.h" #include "utils/hashcat.h" +#include "utils/fastmath.h" namespace lczero { @@ -276,11 +277,15 @@ void Node::CancelScoreUpdate(int multivisit) { best_child_cached_ = nullptr; } -void Node::FinalizeScoreUpdate(float v, float d, float m, int multivisit) { +void Node::FinalizeScoreUpdate(float v, float d, float m, int multivisit, + float factor) { // Recompute Q. - wl_ += multivisit * (v - wl_) / (n_ + multivisit); - d_ += multivisit * (d - d_) / (n_ + multivisit); - m_ += multivisit * (m - m_) / (n_ + multivisit); + float delta = (factor > 0.0f + ? FastLog2( 2.0f + factor * n_ ) * multivisit + : multivisit); + wl_ += delta * (v - wl_) / (n_ + delta); + d_ += delta * (d - d_) / (n_ + delta); + m_ += delta * (m - m_) / (n_ + delta); // If first visit, update parent's sum of policies visited at least once. if (n_ == 0 && parent_ != nullptr) { diff --git a/src/mcts/node.h b/src/mcts/node.h index 217a5ceec2..7e94fd9f52 100644 --- a/src/mcts/node.h +++ b/src/mcts/node.h @@ -194,7 +194,7 @@ class Node { // * Q (weighted average of all V in a subtree) // * N (+=1) // * N-in-flight (-=1) - void FinalizeScoreUpdate(float v, float d, float m, int multivisit); + void FinalizeScoreUpdate(float v, float d, float m, int multivisit, float factor); // When search decides to treat one visit as several (in case of collisions // or visiting terminal nodes several times), it amplifies the visit by // incrementing n_in_flight. diff --git a/src/mcts/params.cc b/src/mcts/params.cc index 29f247c53e..9aed58ace6 100644 --- a/src/mcts/params.cc +++ b/src/mcts/params.cc @@ -57,6 +57,9 @@ const OptionId SearchParams::kLogitQId{ "logit-q", "LogitQ", "Apply logit to Q when determining Q+U best child. This makes the U term " "less dominant when Q is near -1 or +1."}; +const OptionId SearchParams::kBackupFactorId{ + "backup-factor", "BackupFactor", + "Parameter for weighted backup. Set 0 for standard search."}; const OptionId SearchParams::kCpuctId{ "cpuct", "CPuct", "cpuct_init constant from \"UCT search\" algorithm. Higher values promote " @@ -259,6 +262,7 @@ void SearchParams::Populate(OptionsParser* options) { options->Add(kMiniBatchSizeId, 1, 1024) = 256; options->Add(kMaxPrefetchBatchId, 0, 1024) = 32; options->Add(kLogitQId) = false; + options->Add(kBackupFactorId, 0.0f, 100.0f) = 0.0f; options->Add(kCpuctId, 0.0f, 100.0f) = 2.147f; options->Add(kCpuctAtRootId, 0.0f, 100.0f) = 2.147f; options->Add(kCpuctBaseId, 1.0f, 1000000000.0f) = 18368.0f; @@ -329,6 +333,7 @@ void SearchParams::Populate(OptionsParser* options) { SearchParams::SearchParams(const OptionsDict& options) : options_(options), kLogitQ(options.Get(kLogitQId)), + kBackupFactor(options.Get(kBackupFactorId)), kCpuct(options.Get(kCpuctId)), kCpuctAtRoot(options.Get( options.Get(kRootHasOwnCpuctParamsId) ? kCpuctAtRootId diff --git a/src/mcts/params.h b/src/mcts/params.h index 4ed15a71f8..6d28f94b5f 100644 --- a/src/mcts/params.h +++ b/src/mcts/params.h @@ -47,6 +47,7 @@ class SearchParams { return options_.Get(kMaxPrefetchBatchId); } bool GetLogitQ() const { return kLogitQ; } + float GetBackupFactor() const { return kBackupFactor; } float GetCpuct(bool at_root) const { return at_root ? kCpuctAtRoot : kCpuct; } float GetCpuctBase(bool at_root) const { return at_root ? kCpuctBaseAtRoot : kCpuctBase; @@ -114,6 +115,7 @@ class SearchParams { static const OptionId kMiniBatchSizeId; static const OptionId kMaxPrefetchBatchId; static const OptionId kLogitQId; + static const OptionId kBackupFactorId; static const OptionId kCpuctId; static const OptionId kCpuctAtRootId; static const OptionId kCpuctBaseId; @@ -170,6 +172,7 @@ class SearchParams { // TODO(crem) Some of those parameters can be converted to be dynamic after // trivial search optimizations. const bool kLogitQ; + const float kBackupFactor; const float kCpuct; const float kCpuctAtRoot; const float kCpuctBase; diff --git a/src/mcts/search.cc b/src/mcts/search.cc index 517090fb29..e4926cd0ba 100644 --- a/src/mcts/search.cc +++ b/src/mcts/search.cc @@ -1477,7 +1477,7 @@ void SearchWorker::DoBackupUpdateSingleNode( m = n->GetM(); } n->FinalizeScoreUpdate(v / (1.0f + params_.GetShortSightedness() * depth), - d, m, node_to_process.multivisit); + d, m, node_to_process.multivisit, params_.GetBackupFactor()); // Nothing left to do without ancestors to update. if (!p) break;