Tree Search Reasoning: Allocating Inference Budget Across Prefixes

This witness runs the same finite prefix tree as the demo for $2(x+3)=14$. In clean mode, the best-first frontier reaches a correct terminal trace. In noisy mode, an invalid shortcut is falsely scored high, and the max backup sends the root toward the wrong terminal.

NODES = {
    "root": {"children": ["A", "B", "C"]},
    "A": {"parent": "root", "children": ["A1"], "clean": -1.2, "noisy": 1.2, "b_clean": -0.3, "b_noisy": 1.1},
    "A1": {"parent": "A", "terminal": True, "correct": False, "clean": -0.4, "noisy": 1.6},
    "B": {"parent": "root", "children": ["B1", "B2"], "clean": 0.8, "noisy": 0.8, "b_clean": 0.9, "b_noisy": 0.9},
    "B1": {"parent": "B", "terminal": True, "correct": True, "clean": 1.0, "noisy": 1.0},
    "B2": {"parent": "B", "terminal": True, "correct": False, "clean": -0.6, "noisy": -0.6},
    "C": {"parent": "root", "children": ["C1", "C2"], "clean": 0.7, "noisy": 0.7, "b_clean": 1.3, "b_noisy": 1.3},
    "C1": {"parent": "C", "children": ["C1a", "C1b"], "clean": 0.9, "noisy": 0.9, "b_clean": 0.8, "b_noisy": 0.8},
    "C1a": {"parent": "C1", "terminal": True, "correct": True, "clean": 0.8, "noisy": 0.8},
    "C1b": {"parent": "C1", "terminal": True, "correct": False, "clean": -0.6, "noisy": -0.6},
    "C2": {"parent": "C", "terminal": True, "correct": False, "clean": -0.8, "noisy": -0.8},
}

ORDER = ["root", "A", "A1", "B", "B1", "B2", "C", "C1", "C1a", "C1b", "C2"]

def score(node_id, mode):
    return NODES[node_id][mode]

def heuristic(node_id, mode):
    return NODES[node_id].get(f"b_{mode}", 0.0)

def cumulative(node_id, mode):
    if node_id == "root":
        return 0.0
    parent = NODES[node_id]["parent"]
    return cumulative(parent, mode) + score(node_id, mode)

def run(mode="clean", budget=2):
    visible = set(["root", *NODES["root"]["children"]])
    expanded = {"root"}

    def is_frontier(node_id):
        node = NODES[node_id]
        return node_id in visible and not node.get("terminal") and node_id not in expanded

    for _ in range(budget):
        frontier = [node_id for node_id in ORDER if is_frontier(node_id)]
        if not frontier:
            break
        chosen = max(frontier, key=lambda node_id: (cumulative(node_id, mode) + heuristic(node_id, mode), node_id))
        expanded.add(chosen)
        visible.update(NODES[chosen].get("children", []))

    def value(node_id):
        node = NODES[node_id]
        if node.get("terminal"):
            return 0.0
        if node_id not in expanded:
            return heuristic(node_id, mode)
        return max(score(child, mode) + value(child) for child in node["children"] if child in visible)

    path = ["root"]
    while path[-1] in expanded and not NODES[path[-1]].get("terminal"):
        children = [child for child in NODES[path[-1]]["children"] if child in visible]
        if not children:
            break
        path.append(max(children, key=lambda child: (score(child, mode) + value(child), child)))

    terminal = path[-1] if NODES[path[-1]].get("terminal") else None
    return {"path": path, "root_value": value("root"), "terminal": terminal}

clean = run("clean", budget=2)
noisy = run("noisy", budget=2)

assert clean["path"] == ["root", "C", "C1", "C1a"]
assert NODES[clean["terminal"]]["correct"] is True
assert noisy["path"] == ["root", "A", "A1"]
assert NODES[noisy["terminal"]]["correct"] is False
print(clean)
print(noisy)