🐛 Bug: MCTSScheduler.select uses all hist indices as "leaves" and ignores uncommited_rec_status — picks internal nodes and re-picks in-flight ones

[Fr3ya]

MCTSScheduler.select in rdagent/scenarios/data_science/proposal/exp_gen/trace_scheduler.py#L358-L394 departs from its sibling schedulers' contract in two ways. Both stem from one line (L365):

# Step 2: consider only available leaves (not being expanded)
available_leaves = list(set(range(len(trace.hist))))

The comment says "only available leaves (not being expanded)" but the code uses every index in trace.hist without filtering. Compare to ProbabilisticScheduler.select (L211-L215):

leaves = trace.get_leaves()
available_leaves = [leaf for leaf in leaves if self.uncommited_rec_status[leaf] == 0]
if not available_leaves:
    return None

Two bugs follow:

1. MCTSScheduler expands internal (already-children-having) nodes. It picks any historical index, not actual leaves of the DAG. That contradicts the docstring of select in the base classes ("Selects the next leaf node...") and the standard MCTS leaf-expansion semantics.

2. MCTSScheduler ignores uncommited_rec_status. The other schedulers refuse to return a leaf that another worker is already expanding. MCTS doesn't, so under parallel ExpGen with step_semaphore['coding'] > 1, two workers will redo the same node's expansion in parallel.

To Reproduce

Self-contained — no RD-Agent install needed (the schedulers are pure-Python):

import math, random
from collections import defaultdict


class FakeTrace:
    """Models a small DAG with leaves [2, 3, 4] (internal: 0, 1)."""
    NEW_ROOT = ()
    def __init__(self):
        self.dag_parent = [(), (0,), (0,), (1,), ()]
        self.hist = [(None, None)] * 5
        self.sub_trace_count = 2
    def get_leaves(self):
        parents = set(p for tp in self.dag_parent for p in tp)
        return sorted(set(range(len(self.hist))) - parents)


class MCTSSelect:
    """Port of MCTSScheduler.select (trace_scheduler.py:358-394)."""
    def __init__(self, max_trace_num):
        self.uncommited_rec_status = defaultdict(int)
        self.max_trace_num = max_trace_num
        self.c_puct = 1.0
        self.node_visit_count = {}
        self.node_value_sum = {}
        self.node_prior = {}
        self.global_visit_count = 0

    def _get_q(self, n):
        v = self.node_visit_count.get(n, 0)
        return 0.0 if v <= 0 else self.node_value_sum.get(n, 0.0) / v

    def _get_u(self, n):
        return self.c_puct * self.node_prior.get(n, 0.0) * math.sqrt(max(1, self.global_visit_count)) / (1 + self.node_visit_count.get(n, 0))

    def select(self, trace):
        if trace.sub_trace_count + self.uncommited_rec_status[trace.NEW_ROOT] < self.max_trace_num:
            return trace.NEW_ROOT
        # ← THE BUG: all hist indices, not leaves, no uncommited filter
        available_leaves = list(set(range(len(trace.hist))))
        if not available_leaves:
            return None
        priors = [1.0 / len(available_leaves)] * len(available_leaves)
        for l, p in zip(available_leaves, priors):
            self.node_prior[l] = p
        best, score = None, -float("inf")
        for l in available_leaves:
            s = self._get_q(l) + self._get_u(l)
            if s > score:
                best, score = l, s
        if best is None:
            return None
        self.global_visit_count += 1
        return (best,)


class ProbSelect:
    """Port of ProbabilisticScheduler.select (trace_scheduler.py:201-229)."""
    def __init__(self, max_trace_num):
        self.uncommited_rec_status = defaultdict(int)
        self.max_trace_num = max_trace_num

    def select(self, trace):
        if trace.sub_trace_count + self.uncommited_rec_status[trace.NEW_ROOT] < self.max_trace_num:
            return trace.NEW_ROOT
        leaves = trace.get_leaves()
        avail = [l for l in leaves if self.uncommited_rec_status[l] == 0]
        return (random.choice(avail),) if avail else None


# Bug #1: MCTS picks INTERNAL nodes
trace = FakeTrace()
print(f"True leaves: {trace.get_leaves()}")
m = MCTSSelect(2); p = ProbSelect(2)
m_picks = defaultdict(int); p_picks = defaultdict(int)
for _ in range(50):
    s = m.select(trace)
    if s and s != (): m_picks[s[0]] += 1
    s = p.select(trace)
    if s and s != (): p_picks[s[0]] += 1
print(f"Probabilistic picks: {dict(p_picks)}  (stays in leaves)")
print(f"MCTS picks:          {dict(m_picks)}  (escapes leaves)")

# Bug #2: MCTS re-picks in-flight leaves
m = MCTSSelect(2); p = ProbSelect(2)
for n in range(5):
    if n != 3:
        m.node_visit_count[n] = 100
m.uncommited_rec_status[3] = 1
p.uncommited_rec_status[3] = 1
mp = sum(1 for _ in range(50) if (s := m.select(trace)) and s != () and s[0] == 3)
pp = sum(1 for _ in range(50) if (s := p.select(trace)) and s != () and s[0] == 3)
print(f"Probabilistic returns in-flight leaf 3: {pp} times")
print(f"MCTS returns in-flight leaf 3:          {mp} times")

Output (deterministic given the seeds):

True leaves: [2, 3, 4]
Probabilistic picks: {2: 15, 3: 20, 4: 15}  (stays in leaves)
MCTS picks:          {0: 50}                  (escapes leaves)
Probabilistic returns in-flight leaf 3: 0 times
MCTS returns in-flight leaf 3:          50 times

50/50 MCTS selections fall on the root node (an internal, already-expanded node). And once the visit-count profile makes leaf 3 the best PUCT candidate, MCTS keeps returning it even though another worker is already expanding it (uncommited_rec_status[3] = 1).

Expected Behavior

• MCTSScheduler.select should pick only leaves of the trace DAG (consistent with the docstring "Selects the next leaf node ...").
• It should respect uncommited_rec_status so multiple parallel ExpGen workers don't expand the same leaf.

Both invariants are upheld by RoundRobinScheduler.select (L116-L131) and ProbabilisticScheduler.select (L201-L229). The fix is to mirror their pattern in MCTSScheduler.select.

Suggested fix

Two lines, matching the pattern used by ProbabilisticScheduler.select:

     def select(self, trace: DSTrace) -> tuple[int, ...] | None:
         # Step 1: keep same policy to reach target number of parallel traces
         if trace.sub_trace_count + self.uncommited_rec_status[trace.NEW_ROOT] < self.max_trace_num:
             return trace.NEW_ROOT

-        # Step 2: consider only available leaves (not being expanded)
-        available_leaves = list(set(range(len(trace.hist))))
+        # Step 2: consider only available leaves (not being expanded)
+        leaves = trace.get_leaves()
+        available_leaves = [leaf for leaf in leaves if self.uncommited_rec_status[leaf] == 0]
         if not available_leaves:
             return None

Additional Notes

• Surfaced by a systematic scan for "scheduler select methods that drift from sibling implementations in candidate-set construction."
• Same class of bug (graph scheduler skips its own in-flight markers) is a recurring multi-worker pattern; a unit test that runs each scheduler 100 times with uncommited_rec_status[leaf] = 1 and asserts leaf is never returned would catch this for all current and future schedulers in one shot.