Policy decay

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,

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<IntOption>(kMiniBatchSizeId, 1, 1024) = 256;
   options->Add<IntOption>(kMaxPrefetchBatchId, 0, 1024) = 32;
   options->Add<BoolOption>(kLogitQId) = false;
+  options->Add<FloatOption>(kPolicyDecayShiftId, 0.0f, 100.0f) = 0.1f;
+  options->Add<FloatOption>(kPolicyDecaySlopeId, 0.0f, 100.0f) = 0.5f;
   options->Add<FloatOption>(kCpuctId, 0.0f, 100.0f) = 2.147f;
   options->Add<FloatOption>(kCpuctAtRootId, 0.0f, 100.0f) = 2.147f;
   options->Add<FloatOption>(kCpuctBaseId, 1.0f, 1000000000.0f) = 18368.0f;
   options->Add<FloatOption>(kCpuctBaseAtRootId, 1.0f, 1000000000.0f) = 18368.0f;
-  options->Add<FloatOption>(kCpuctFactorId, 0.0f, 1000.0f) = 2.815f;
-  options->Add<FloatOption>(kCpuctFactorAtRootId, 0.0f, 1000.0f) = 2.815f;
+  options->Add<FloatOption>(kCpuctFactorId, 0.0f, 1000.0f) = 0.0f;
+  options->Add<FloatOption>(kCpuctFactorAtRootId, 0.0f, 1000.0f) = 0.0f;
   options->Add<BoolOption>(kRootHasOwnCpuctParamsId) = true;
   options->Add<FloatOption>(kTemperatureId, 0.0f, 100.0f) = 0.0f;
   options->Add<IntOption>(kTempDecayMovesId, 0, 100) = 0;
@@ -343,6 +351,8 @@ void SearchParams::Populate(OptionsParser* options) {
 SearchParams::SearchParams(const OptionsDict& options)
     : options_(options),
       kLogitQ(options.Get<bool>(kLogitQId)),
+      kPolicyDecayShift(options.Get<float>(kPolicyDecayShiftId)),
+      kPolicyDecaySlope(options.Get<float>(kPolicyDecaySlopeId)),
       kCpuct(options.Get<float>(kCpuctId)),
       kCpuctAtRoot(options.Get<float>(
           options.Get<bool>(kRootHasOwnCpuctParamsId) ? kCpuctAtRootId

src/mcts/params.h

@@ -47,6 +47,8 @@ class SearchParams {
     return options_.Get<int>(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;

src/mcts/search.cc

@@ -260,6 +260,9 @@ std::vector<std::string> 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<std::string> 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<std::string> infos;
@@ -322,11 +328,11 @@ std::vector<std::string> 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<float> 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<float>::lowest();
     float best_without_u = std::numeric_limits<float>::lowest();
     float second_best = std::numeric_limits<float>::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<ScoredEdge> 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);
   }