diff --git a/src/mcts/node.h b/src/mcts/node.h index 217a5ceec2..752b233c33 100644 --- a/src/mcts/node.h +++ b/src/mcts/node.h @@ -396,8 +396,11 @@ class EdgeAndNode { // Returns U = numerator * p / N. // Passed numerator is expected to be equal to (cpuct * sqrt(N[parent])). - float GetU(float numerator) const { - return numerator * GetP() / (1 + GetNStarted()); + float GetU(float numerator, float childVisits, float shift, float slope) const { + const float p = GetP(); + const float p_dcay = p + (1 - p) * + (1 - 1 / std::pow(1 + shift * p * (childVisits - 1), slope)); + return numerator * p_dcay / (1 + GetNStarted()); } int GetVisitsToReachU(float target_score, float numerator, diff --git a/src/mcts/params.cc b/src/mcts/params.cc index 7b239d0ecd..d76c5c0ccf 100644 --- a/src/mcts/params.cc +++ b/src/mcts/params.cc @@ -57,6 +57,12 @@ 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::kPolicyDecayShiftId{ + "policy-decay-shift", "PolicyDecayShift", + "Adjusts how soon policy effect starts to decay."}; +const OptionId SearchParams::kPolicyDecaySlopeId{ + "policy-decay-slope", "PolicyDecaySlope", + "Adjusts how fast policy effect decays to 1."}; const OptionId SearchParams::kCpuctId{ "cpuct", "CPuct", "cpuct_init constant from \"UCT search\" algorithm. Higher values promote " @@ -264,12 +270,14 @@ void SearchParams::Populate(OptionsParser* options) { options->Add(kMiniBatchSizeId, 1, 1024) = 256; options->Add(kMaxPrefetchBatchId, 0, 1024) = 32; options->Add(kLogitQId) = false; + options->Add(kPolicyDecayShiftId, 0.0f, 100.0f) = 0.1f; + options->Add(kPolicyDecaySlopeId, 0.0f, 100.0f) = 0.5f; 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; options->Add(kCpuctBaseAtRootId, 1.0f, 1000000000.0f) = 18368.0f; - options->Add(kCpuctFactorId, 0.0f, 1000.0f) = 2.815f; - options->Add(kCpuctFactorAtRootId, 0.0f, 1000.0f) = 2.815f; + options->Add(kCpuctFactorId, 0.0f, 1000.0f) = 0.0f; + options->Add(kCpuctFactorAtRootId, 0.0f, 1000.0f) = 0.0f; options->Add(kRootHasOwnCpuctParamsId) = true; options->Add(kTemperatureId, 0.0f, 100.0f) = 0.0f; options->Add(kTempDecayMovesId, 0, 100) = 0; @@ -343,6 +351,8 @@ void SearchParams::Populate(OptionsParser* options) { SearchParams::SearchParams(const OptionsDict& options) : options_(options), kLogitQ(options.Get(kLogitQId)), + kPolicyDecayShift(options.Get(kPolicyDecayShiftId)), + kPolicyDecaySlope(options.Get(kPolicyDecaySlopeId)), kCpuct(options.Get(kCpuctId)), kCpuctAtRoot(options.Get( options.Get(kRootHasOwnCpuctParamsId) ? kCpuctAtRootId diff --git a/src/mcts/params.h b/src/mcts/params.h index 0073a6da50..f5de9ff920 100644 --- a/src/mcts/params.h +++ b/src/mcts/params.h @@ -47,6 +47,8 @@ class SearchParams { return options_.Get(kMaxPrefetchBatchId); } bool GetLogitQ() const { return kLogitQ; } + float GetPolicyDecayShift() const { return kPolicyDecayShift; } + float GetPolicyDecaySlope() const { return kPolicyDecaySlope; } float GetCpuct(bool at_root) const { return at_root ? kCpuctAtRoot : kCpuct; } float GetCpuctBase(bool at_root) const { return at_root ? kCpuctBaseAtRoot : kCpuctBase; @@ -117,6 +119,8 @@ class SearchParams { static const OptionId kMiniBatchSizeId; static const OptionId kMaxPrefetchBatchId; static const OptionId kLogitQId; + static const OptionId kPolicyDecayShiftId; + static const OptionId kPolicyDecaySlopeId; static const OptionId kCpuctId; static const OptionId kCpuctAtRootId; static const OptionId kCpuctBaseId; @@ -174,6 +178,8 @@ class SearchParams { // TODO(crem) Some of those parameters can be converted to be dynamic after // trivial search optimizations. const bool kLogitQ; + const float kPolicyDecayShift; + const float kPolicyDecaySlope; const float kCpuct; const float kCpuctAtRoot; const float kCpuctBase; diff --git a/src/mcts/search.cc b/src/mcts/search.cc index c9a6f82f37..d961c2cdaa 100644 --- a/src/mcts/search.cc +++ b/src/mcts/search.cc @@ -260,6 +260,9 @@ std::vector Search::GetVerboseStats(Node* node) const { const float cpuct = ComputeCpuct(params_, node->GetN(), is_root); const float U_coeff = cpuct * std::sqrt(std::max(node->GetChildrenVisits(), 1u)); + const float childVisits = node->GetChildrenVisits(); + const float shift = params_.GetPolicyDecayShift(); + const float slope = params_.GetPolicyDecaySlope(); const bool logit_q = params_.GetLogitQ(); const float m_slope = params_.GetMovesLeftSlope(); const float m_cap = params_.GetMovesLeftMaxEffect(); @@ -275,13 +278,16 @@ std::vector Search::GetVerboseStats(Node* node) const { std::sort( edges.begin(), edges.end(), - [&fpu, &U_coeff, &logit_q, &draw_score](EdgeAndNode a, EdgeAndNode b) { + [&fpu, &U_coeff, &childVisits, &shift, &slope, &logit_q, + &draw_score](EdgeAndNode a, EdgeAndNode b) { return std::forward_as_tuple( a.GetN(), - a.GetQ(fpu, draw_score, logit_q) + a.GetU(U_coeff)) < + a.GetQ(fpu, draw_score, logit_q) + + a.GetU(U_coeff, childVisits, shift, slope)) < std::forward_as_tuple( b.GetN(), - b.GetQ(fpu, draw_score, logit_q) + b.GetU(U_coeff)); + b.GetQ(fpu, draw_score, logit_q) + + b.GetU(U_coeff, childVisits, shift, slope)); }); std::vector infos; @@ -322,11 +328,11 @@ std::vector Search::GetVerboseStats(Node* node) const { oss << "(Q: " << std::setw(8) << std::setprecision(5) << edge.GetQ(fpu, draw_score, /* logit_q= */ false) << ") "; - oss << "(U: " << std::setw(6) << std::setprecision(5) << edge.GetU(U_coeff) - << ") "; + oss << "(U: " << std::setw(6) << std::setprecision(5) + << edge.GetU(U_coeff, childVisits, shift, slope) << ") "; oss << "(S: " << std::setw(8) << std::setprecision(5) - << Q + edge.GetU(U_coeff) + M_effect << ") "; + << Q + edge.GetU(U_coeff, childVisits, shift, slope) + M_effect << ") "; oss << "(V: "; std::optional v; @@ -1056,6 +1062,9 @@ SearchWorker::NodeToProcess SearchWorker::PickNodeToExtend( const float cpuct = ComputeCpuct(params_, node->GetN(), is_root_node); const float puct_mult = cpuct * std::sqrt(std::max(node->GetChildrenVisits(), 1u)); + const float childVisits = node->GetChildrenVisits(); + const float shift = params_.GetPolicyDecayShift(); + const float slope = params_.GetPolicyDecaySlope(); float best = std::numeric_limits::lowest(); float best_without_u = std::numeric_limits::lowest(); float second_best = std::numeric_limits::lowest(); @@ -1105,7 +1114,7 @@ SearchWorker::NodeToProcess SearchWorker::PickNodeToExtend( M *= a + b * std::abs(Q) + c * Q * Q; } - const float score = child.GetU(puct_mult) + Q + M; + const float score = child.GetU(puct_mult, childVisits, shift, slope) + Q + M; if (score > best) { second_best = best; second_best_edge = best_edge; @@ -1334,6 +1343,9 @@ int SearchWorker::PrefetchIntoCache(Node* node, int budget, bool is_odd_depth) { std::vector scores; const float cpuct = ComputeCpuct(params_, node->GetN(), node == search_->root_node_); + const float childVisits = node->GetChildrenVisits(); + const float shift = params_.GetPolicyDecayShift(); + const float slope = params_.GetPolicyDecaySlope(); const float puct_mult = cpuct * std::sqrt(std::max(node->GetChildrenVisits(), 1u)); const float fpu = @@ -1342,7 +1354,7 @@ int SearchWorker::PrefetchIntoCache(Node* node, int budget, bool is_odd_depth) { if (edge.GetP() == 0.0f) continue; // Flip the sign of a score to be able to easily sort. // TODO: should this use logit_q if set?? - scores.emplace_back(-edge.GetU(puct_mult) - + scores.emplace_back(-edge.GetU(puct_mult, childVisits, shift, slope) - edge.GetQ(fpu, draw_score, /* logit_q= */ false), edge); }