Keep TAO from lowering training accuracy at lambda=0

cpp/CMakeLists.txt

@@ -240,6 +240,7 @@ if (SGTLEARN_BUILD_TESTS)
             tests/test_splitters.cpp
             tests/test_branch_assignment.cpp
             tests/test_shape_assignment_search.cpp
+            tests/test_tao_objective.cpp
             tests/bench_mae_branch_assignment.cpp
     )
 

cpp/src/algorithms/TAO/ClassificationTaoAdapter.cpp

@@ -91,23 +91,13 @@ void ClassificationTaoAdapter::childRewards(
     return;
   }
 
-  // Single-output: 1/x split uniformly across correct children.
+  // Single-output: full credit (1) for every correct child. Routing is
+  // multi-label, so a sample counts once wherever it lands correctly.
   const size_t label = y_(0, col);
-  size_t numCorrect = 0;
   for (size_t c = 0; c < childLeaves.size(); ++c) {
     const auto predictions =
         argMaxClass(classificationTree_.classCounts[childLeaves[c]]);
-    const bool correct = predictions[0] == label;
-    if (correct)
-      ++numCorrect;
-    reward[c] = correct ? 1.0 : 0.0;
-  }
-  if (numCorrect == 0)
-    return;
-  const double correctReward = 1.0 / static_cast<double>(numCorrect);
-  for (size_t c = 0; c < reward.size(); ++c) {
-    if (reward[c] > 0.0)
-      reward[c] = correctReward;
+    reward[c] = predictions[0] == label ? 1.0 : 0.0;
   }
 }
 
@@ -160,10 +150,15 @@ NodeCareSet ClassificationTaoAdapter::buildCareSet(
   size_t pos = 0;
   for (size_t i = 0; i < care.size(); ++i) {
     const float wi = w_(care.careCols[i]);
+    // Single-output: split the sample's weight across its correct children
+    // (soft labels), so it adds total mass wi to the router fit. childCounts
+    // keeps full weight: dummyChild is correct for the most care weight.
+    const float wRow =
+        nOutputs_ == 1 ? wi / static_cast<float>(goodChildren[i].size()) : wi;
     for (size_t child : goodChildren[i]) {
       care.Xexp.col(pos) = X_.col(care.careCols[i]);
       care.yexp(pos) = child;
-      care.wexp(pos) = wi;
+      care.wexp(pos) = wRow;
       childCounts[child] += static_cast<double>(wi);
       ++pos;
     }

cpp/src/algorithms/TAO/ClassificationTaoAdapter.h

@@ -7,11 +7,14 @@
  * Care-set logic (classification variant):
  *
  * - For each sample at an internal node, walk each child subtree to its leaf
- *   and score child ``c`` with reward ``1/x`` if the leaf's majority class equals
- *   the sample label, where ``x`` is the number of correct children; else ``0``.
+ *   and score child ``c`` with reward ``1`` if the leaf's majority class equals
+ *   the sample label, else ``0``. Routing to any correct child is fully
+ *   correct (multi-label), so rewards are the true accuracy contribution.
  * - Samples where all children tie are excluded from the care set.
  * - Samples with multiple equally good children contribute one expanded router-
- *   training row per good child (multi-to-single pseudolabel expansion).
+ *   training row per good child (multi-to-single pseudolabel expansion). In
+ *   single-output, each row carries ``w / x`` (``x`` = number of correct
+ *   children), so the sample adds total weight ``w`` to the router fit.
  *
  * Leaf refresh recomputes weighted class histograms at every node from
  * ``nodeSamples``. The routing discretizer uses the tree's own criterion
@@ -59,9 +62,8 @@ class ClassificationTaoAdapter final : public ShapeGeneralizedTaoAdapter {
   /**
    * Per-child correctness rewards for one sample.
    *
-   * Single-output: ``1/x`` split uniformly across correct children (``x`` =
-   * count of correct), else ``0``. Multi-output: fraction of outputs correctly
-   * classified by each child's leaf.
+   * Single-output: ``1`` for each correct child, else ``0``. Multi-output:
+   * fraction of outputs correctly classified by each child's leaf.
    */
   void childRewards(const std::vector<size_t> &childLeaves, arma::uword col,
                     std::vector<double> &reward) const;

cpp/src/algorithms/TAO/TaoObjective.cpp

@@ -5,11 +5,13 @@
 #include <cstddef>
 #include "algorithms/TAO/TaoObjective.h"
 
+#include "algorithms/missing_values.h"
+
 #include <algorithm>
 #include <cmath>
 #include <limits>
 #include <stdexcept>
-#include <string>
+#include <utility>
 
 namespace tao {
 
@@ -51,17 +53,15 @@ double TaoObjective::penalizedScore(double rewardSum,
          complexityScale * lambda_ * nodeSampleCount_ / totalCareWeight_;
 }
 
-double TaoObjective::rewardSumForDiscretizer(
-    const ClassificationDiscretizer &disc,
+double TaoObjective::rewardSumForBins(
+    const arma::Row<size_t> &bins,
     const std::vector<size_t> &binToPartition) const {
-  arma::Row<size_t> bins;
-  disc.transform(X_, bins);
   double rewardSum = 0.0;
   for (size_t i = 0; i < nCare_; ++i) {
     const size_t bin = bins(care_.careCols[i]);
     if (bin >= binToPartition.size())
       throw std::runtime_error(
-          "TaoObjective::rewardSumForDiscretizer: bin out of range");
+          "TaoObjective::rewardSumForBins: bin out of range");
     const size_t child = binToPartition[bin];
     rewardSum += careWeight(i) * care_.careRewards[i][child];
   }
@@ -93,31 +93,109 @@ double TaoObjective::scoreDiscretizer(
   if (disc.numLeaves() < 1)
     return -std::numeric_limits<double>::infinity();
 
-  const std::vector<std::vector<std::vector<double>>> &leafStats =
-      disc.leafStats();
-  if (leafStats.empty())
+  // Bin count includes a univariate discretizer's trailing NaN bin.
+  const size_t nBins = disc.leafStats().size();
+  if (nBins == 0)
     throw std::runtime_error(
         "TaoObjective::scoreDiscretizer: discretizer has no leaf stats");
 
-  binToPartitionOut.resize(leafStats.size());
+  arma::Row<size_t> bins;
+  disc.transform(X_, bins);
 
-  // TAO trains the inner discretizer on scalar care-set pseudo-labels (child
-  // partition indices 0..k-1), so each bin's stats are [1][k] not multi-output.
-  for (size_t b = 0; b < leafStats.size(); ++b) {
-    if (leafStats[b].size() != 1)
-      throw std::runtime_error(
-          "TaoObjective::scoreDiscretizer: bin " + std::to_string(b) +
-          " expected 1 output row of class counts, got " +
-          std::to_string(leafStats[b].size()));
-    if (leafStats[b][0].empty())
+  // Relabel each bin by the true care reward, not the fitted histogram: the
+  // router is fit on split weights, whose argmax can miss the best child.
+  const size_t k = nCare_ == 0 ? 0 : care_.careRewards.front().size();
+  std::vector<std::vector<double>> binReward(nBins,
+                                             std::vector<double>(k, 0.0));
+  std::vector<double> binWeight(nBins, 0.0);
+  for (size_t i = 0; i < nCare_; ++i) {
+    const size_t bin = bins(care_.careCols[i]);
+    if (bin >= nBins)
       throw std::runtime_error(
-          "TaoObjective::scoreDiscretizer: bin " + std::to_string(b) +
-          " class histogram is empty");
-    binToPartitionOut[b] = argMax(leafStats[b][0]);
+          "TaoObjective::scoreDiscretizer: bin out of range");
+    const double w = careWeight(i);
+    binWeight[bin] += w;
+    for (size_t c = 0; c < k; ++c)
+      binReward[bin][c] += w * care_.careRewards[i][c];
   }
+  binToPartitionOut.resize(nBins);
+  for (size_t b = 0; b < nBins; ++b)
+    binToPartitionOut[b] =
+        binWeight[b] > 0.0 ? argMax(binReward[b]) : care_.dummyChild;
 
-  return penalizedScore(
-      rewardSumForDiscretizer(disc, binToPartitionOut), complexityScale);
+  return penalizedScore(rewardSumForBins(bins, binToPartitionOut),
+                        complexityScale);
+}
+
+double TaoObjective::scoreBinAssignment(
+    const arma::Row<size_t> &bins, const std::vector<size_t> &binToPartition,
+    double complexityScale) const {
+  return penalizedScore(rewardSumForBins(bins, binToPartition),
+                        complexityScale);
+}
+
+std::optional<ThresholdCut>
+TaoObjective::bestThresholdCut(size_t feature, size_t minLeafSize) const {
+  if (nCare_ == 0)
+    return std::nullopt;
+  if (care_.careRewards.front().size() != 2)
+    throw std::invalid_argument(
+        "TaoObjective::bestThresholdCut: node must have exactly two children");
+
+  std::vector<std::pair<float, size_t>> finite; // (value, care index)
+  finite.reserve(nCare_);
+  double total0 = 0.0, total1 = 0.0;
+  double nan0 = 0.0, nan1 = 0.0, nanWeight = 0.0;
+  for (size_t i = 0; i < nCare_; ++i) {
+    const float value =
+        X_(static_cast<arma::uword>(feature), care_.careCols[i]);
+    const double w = careWeight(i);
+    const double r0 = w * care_.careRewards[i][0];
+    const double r1 = w * care_.careRewards[i][1];
+    if (missing_values::is_finite(value)) {
+      finite.emplace_back(value, i);
+      total0 += r0;
+      total1 += r1;
+    } else {
+      nan0 += r0;
+      nan1 += r1;
+      nanWeight += w;
+    }
+  }
+
+  const size_t minSide = std::max<size_t>(minLeafSize, 1);
+  const size_t n = finite.size();
+  if (n < 2 * minSide)
+    return std::nullopt;
+  std::sort(finite.begin(), finite.end());
+
+  const size_t nanChild =
+      nanWeight > 0.0 ? (nan1 > nan0 ? 1 : 0) : care_.dummyChild;
+  const double nanReward = nanChild == 0 ? nan0 : nan1;
+
+  std::optional<ThresholdCut> best;
+  double left0 = 0.0, left1 = 0.0;
+  for (size_t j = 1; j < n; ++j) {
+    const auto [prevValue, prevIdx] = finite[j - 1];
+    const double w = careWeight(prevIdx);
+    left0 += w * care_.careRewards[prevIdx][0];
+    left1 += w * care_.careRewards[prevIdx][1];
+    if (j < minSide)
+      continue;
+    if (n - j < minSide)
+      break;
+    // Same tie rule as Splitter::findBestSplit, so a depth-1 discretizer
+    // trained on the two sides reproduces this cut.
+    if (finite[j].first <= prevValue + static_cast<float>(1e-7))
+      continue;
+    const double leftTo0 = left0 + (total1 - left1);
+    const double leftTo1 = left1 + (total0 - left0);
+    const size_t leftChild = leftTo1 > leftTo0 ? 1 : 0;
+    const double reward = std::max(leftTo0, leftTo1) + nanReward;
+    if (!best || reward > best->rewardSum)
+      best = ThresholdCut{prevValue, leftChild, nanChild, reward};
+  }
+  return best;
 }
 
 } // namespace tao

cpp/src/algorithms/TAO/TaoObjective.h

@@ -15,10 +15,27 @@
 #include "algorithms/TAO/TaoAdapter.h"
 
 #include <cstddef>
+#include <limits>
+#include <optional>
 #include <vector>
 
 namespace tao {
 
+/**
+ * One univariate threshold routing for a two-child node: finite care values
+ * ``x <= leftMax`` go to ``leftChild``, larger ones to ``1 - leftChild``, and
+ * non-finite ones to ``nanChild``. The installed router's threshold lies
+ * between ``leftMax`` and the next finite care value.
+ */
+struct ThresholdCut {
+  /** Largest finite care value on the left side; the cut lies above it. */
+  float leftMax = 0.0f;
+  size_t leftChild = 0;
+  size_t nanChild = 0;
+  /** Weighted care reward of this routing (no complexity penalty). */
+  double rewardSum = -std::numeric_limits<double>::infinity();
+};
+
 /**
  * Mean care-set reward under a routing rule, with optional split penalty.
  *
@@ -48,21 +65,47 @@ class TaoObjective {
   /**
    * Extract routing from a trained discretizer and score it on the care set.
    *
-   * Bins are mapped to child partitions by argmax over discretizer leaf stats.
-   * Writes the induced bin-to-partition map to the out-param.
+   * Each bin goes to the child with the largest summed care reward over the
+   * care samples in it (ties: lowest child index), not to the argmax of the
+   * discretizer's fitted histogram. Bins with zero care weight, including an
+   * unused NaN bin, go to ``dummyChild``. Writes the bin-to-partition map to
+   * the out-param.
    *
    * @returns Penalized mean reward, or ``-infinity`` when ``disc`` has no bins.
    */
   double scoreDiscretizer(ClassificationDiscretizer &disc,
                           std::vector<size_t> &binToPartitionOut,
                           double complexityScale) const;
 
+  /**
+   * Penalized mean reward when care sample ``i`` routes to
+   * ``binToPartition[bins(careCols[i])]``.
+   *
+   * @param bins Bin of every column of ``X`` (as from ``transform(X, bins)``).
+   */
+  double scoreBinAssignment(const arma::Row<size_t> &bins,
+                            const std::vector<size_t> &binToPartition,
+                            double complexityScale) const;
+
+  /**
+   * Exact best threshold routing on raw ``feature`` for a two-child node.
+   *
+   * Cuts lie between consecutive distinct finite care values (using the inner
+   * splitter's ``1e-7`` tie rule) with at least ``minLeafSize`` finite care
+   * samples per side; both orientations are tried. Non-finite care samples go
+   * to the child with the larger summed reward over them (``dummyChild`` when
+   * they carry no weight).
+   *
+   * @returns The best cut, or ``std::nullopt`` when no cut is feasible.
+   */
+  std::optional<ThresholdCut> bestThresholdCut(size_t feature,
+                                               size_t minLeafSize) const;
+
 private:
   static size_t argMax(const std::vector<double> &counts);
 
-  double rewardSumForDiscretizer(
-      const ClassificationDiscretizer &disc,
-      const std::vector<size_t> &binToPartition) const;
+  double rewardSumForBins(const arma::Row<size_t> &bins,
+                          const std::vector<size_t> &binToPartition) const;
 
   double meanReward(double rewardSum) const;
   double penalizedScore(double rewardSum, double complexityScale) const;