Skills › Scientific Agent Skills › arbor › tree.py
tree.py
607 lines · skills/arbor/scripts/tree.py at 5f2bc03
1#!/usr/bin/env python3
2"""
3tree.py -- persistent hypothesis-tree state manager for Arbor-style
4Hypothesis Tree Refinement (HTR).
6The hypothesis tree is the durable research state for an Autonomous
7Optimization (AO) run. This script owns the *mechanical* parts of that state
8-- creating nodes, writing back evidence, propagating insights up the tree,
9pruning falsified branches, recording the held-out merge gate, and rendering
10an "Observe" projection -- so the coordinator (you, the model) can spend its
11judgment on what the evidence *means* rather than on bookkeeping.
13Division of labor:
14 - This script keeps the state consistent and auditable. It never decides
15 which hypothesis to try or whether one is good.
16 - The coordinator reads the projection (`observe`), forms hypotheses,
17 interprets executor reports, and calls the mutating commands to record
18 those decisions.
20State lives in `.arbor/` under the run directory (default: current dir):
21 .arbor/tree.json -- the hypothesis tree (nodes, edges, evidence, insights)
22 .arbor/run.json -- run-level config: objective, evaluators, budget, M_best
24Node fields mirror the paper's research unit n = <h, iota, mu>:
25 hypothesis (h) -- the falsifiable claim this node tests
26 insight (iota) -- distilled, reusable lesson (filled after execution)
27 metadata (mu) -- status, dev_score, test_score, result, branch_ref, depth
29Run `python tree.py --help` or `python tree.py <command> --help`.
30"""
32import argparse
33import json
34import math
35import sys
36import time
37from pathlib import Path
39# ----------------------------------------------------------------------------
40# Storage helpers
41# ----------------------------------------------------------------------------
43VALID_STATUS = {"pending", "running", "executed", "merged", "pruned", "root"}
46def _finite_score(value):
47 """Reject failed-evaluation sentinels instead of persisting invalid scores."""
48 score = float(value)
49 if not math.isfinite(score):
50 raise argparse.ArgumentTypeError("score must be finite (not NaN or infinity)")
51 return score
54def _dir(run_dir):
55 return Path(run_dir) / ".arbor"
58def _tree_path(run_dir):
59 return _dir(run_dir) / "tree.json"
62def _run_path(run_dir):
63 return _dir(run_dir) / "run.json"
66def _load(path, what):
67 if not path.exists():
68 sys.exit(
69 f"error: no {what} found at {path}. Run `tree.py init` first "
70 f"(from the run directory, or pass --run-dir)."
71 )
72 with open(path) as f:
73 return json.load(f)
76def _save(path, data):
77 path.parent.mkdir(parents=True, exist_ok=True)
78 tmp = path.with_suffix(".tmp")
79 with open(tmp, "w") as f:
80 json.dump(data, f, indent=2)
81 tmp.replace(path)
84def _load_tree(run_dir):
85 return _load(_tree_path(run_dir), "hypothesis tree")
88def _load_run(run_dir):
89 return _load(_run_path(run_dir), "run config")
92def _next_id(tree):
93 n = tree.get("_counter", 0) + 1
94 tree["_counter"] = n
95 return f"n{n}"
98def _id_sort_key(node_id):
99 """Order generated ids numerically, then any manually assigned ids by name."""
100 if isinstance(node_id, str) and node_id.startswith("n") and node_id[1:].isdecimal():
101 return (0, int(node_id[1:]), node_id)
102 return (1, 0, str(node_id))
105def _node(tree, node_id):
106 node = tree["nodes"].get(node_id)
107 if node is None:
108 sys.exit(f"error: no node with id '{node_id}'. Run `tree.py status` to list nodes.")
109 return node
112def _children(tree, node_id):
113 return [nid for nid, n in tree["nodes"].items() if n.get("parent") == node_id]
116def _ancestors(tree, node_id):
117 """Path from node up to (and including) root, nearest first."""
118 path = []
119 cur = tree["nodes"][node_id].get("parent")
120 while cur is not None:
121 path.append(cur)
122 cur = tree["nodes"][cur].get("parent")
123 return path
126def _depth(tree, node_id):
127 return len(_ancestors(tree, node_id))
130def _stamp(node):
131 node["updated_at"] = int(time.time())
134# ----------------------------------------------------------------------------
135# Commands
136# ----------------------------------------------------------------------------
139def cmd_init(args):
140 d = _dir(args.run_dir)
141 if _tree_path(args.run_dir).exists() and not args.force:
142 sys.exit(
143 f"error: a tree already exists at {_tree_path(args.run_dir)}. "
144 f"Use --force to overwrite (this erases the current run)."
145 )
146 root = {
147 "id": "n0",
148 "parent": None,
149 "depth": 0,
150 "status": "root",
151 "hypothesis": args.objective,
152 "insight": "", # global insights accumulate here via backpropagation
153 "metadata": {"dev_score": None, "test_score": None, "result": "", "branch_ref": None},
154 "created_at": int(time.time()),
155 "updated_at": int(time.time()),
156 }
157 tree = {"_counter": 0, "root": "n0", "nodes": {"n0": root}}
158 run = {
159 "objective": args.objective,
160 "metric_direction": args.metric_direction,
161 "dev_eval": args.dev_eval,
162 "test_eval": args.test_eval,
163 "material": args.material,
164 "branching": args.branching,
165 "max_depth": args.max_depth,
166 "budget_cycles": args.budget,
167 "cycles_used": 0,
168 "best_node": None, # node id of current M_best
169 "best_test_score": None, # held-out score of M_best
170 "best_branch_ref": None, # git ref / path of M_best artifact
171 "created_at": int(time.time()),
172 }
173 _save(_tree_path(args.run_dir), tree)
174 _save(_run_path(args.run_dir), run)
175 print(f"Initialized Arbor run in {d}")
176 print(f" objective : {args.objective}")
177 print(f" metric direction: {args.metric_direction} (higher-is-better after orienting)")
178 print(f" dev evaluator : {args.dev_eval}")
179 print(f" test evaluator : {args.test_eval}")
180 print(f" budget : {args.budget} cycles, branching {args.branching}, max depth {args.max_depth}")
181 print("\nNext: `tree.py observe` to read the state, then add direction nodes under n0.")
184def cmd_add_node(args):
185 tree = _load_tree(args.run_dir)
186 parent = _node(tree, args.parent)
187 run = _load_run(args.run_dir)
188 nid = _next_id(tree)
189 depth = _depth(tree, args.parent) + 1
190 if depth > run["max_depth"]:
191 print(
192 f"warning: node depth {depth} exceeds max_depth {run['max_depth']}. "
193 f"Deep nodes should be concrete, executable interventions.",
194 file=sys.stderr,
195 )
196 node = {
197 "id": nid,
198 "parent": args.parent,
199 "depth": depth,
200 "status": "pending",
201 "hypothesis": args.hypothesis,
202 "insight": "",
203 "metadata": {"dev_score": None, "test_score": None, "result": "", "branch_ref": None},
204 "created_at": int(time.time()),
205 "updated_at": int(time.time()),
206 }
207 tree["nodes"][nid] = node
208 _save(_tree_path(args.run_dir), tree)
209 kind = "direction" if depth == 1 else "intervention"
210 print(f"Added {kind} node {nid} (depth {depth}) under {args.parent}: {args.hypothesis}")
213def cmd_baseline(args):
214 """Record the evaluated initial artifact once, before any merge decisions."""
215 tree = _load_tree(args.run_dir)
216 run = _load_run(args.run_dir)
217 if run["best_test_score"] is not None or run["best_node"] is not None:
218 sys.exit("error: a baseline or best artifact is already recorded; start a new run to change the baseline.")
219 root = _node(tree, tree["root"])
220 root["metadata"]["test_score"] = args.test_score
221 root["metadata"]["branch_ref"] = args.branch_ref
222 run["best_node"] = tree["root"]
223 run["best_test_score"] = args.test_score
224 run["best_branch_ref"] = args.branch_ref
225 _stamp(root)
226 _save(_tree_path(args.run_dir), tree)
227 _save(_run_path(args.run_dir), run)
228 print(f"Recorded baseline {tree['root']}: test={args.test_score}, ref={args.branch_ref}")
231def cmd_set_status(args):
232 tree = _load_tree(args.run_dir)
233 node = _node(tree, args.node)
234 if args.status not in VALID_STATUS:
235 sys.exit(f"error: status must be one of {sorted(VALID_STATUS)}")
236 node["status"] = args.status
237 _stamp(node)
238 _save(_tree_path(args.run_dir), tree)
239 print(f"{args.node} -> status={args.status}")
242def cmd_set_evidence(args):
243 """Write an executor's report back into its node (the Backpropagate step,
244 leaf part). Insight propagation upward is a separate, deliberate call."""
245 tree = _load_tree(args.run_dir)
246 node = _node(tree, args.node)
247 if args.status is not None and args.status not in VALID_STATUS:
248 sys.exit(f"error: status must be one of {sorted(VALID_STATUS)}")
249 meta = node["metadata"]
250 if args.dev_score is not None:
251 meta["dev_score"] = args.dev_score
252 if args.result is not None:
253 meta["result"] = args.result
254 if args.branch_ref is not None:
255 meta["branch_ref"] = args.branch_ref
256 if args.insight is not None:
257 node["insight"] = args.insight
258 node["status"] = args.status or "executed"
259 _stamp(node)
260 _save(_tree_path(args.run_dir), tree)
261 print(f"Wrote evidence to {args.node}: dev_score={meta['dev_score']} status={node['status']}")
262 if node["insight"]:
263 print(f" insight: {node['insight']}")
264 anc = _ancestors(tree, args.node)
265 if anc:
266 print(
267 "\nReminder: abstract this leaf insight upward. Decide what direction-level "
268 f"lesson it implies for ancestors {anc} and record it with "
269 f"`tree.py propagate --node {args.node} --insight \"...\"` (and update n0 "
270 "global insights if it generalizes)."
271 )
274def cmd_propagate(args):
275 """Backpropagate a distilled, direction-level lesson up the ancestor path.
277 The coordinator decides the abstracted wording; this appends it to the
278 chosen ancestor(s) so later ideation is conditioned on it. By default it
279 updates the immediate parent; --to-root also updates global insights."""
280 tree = _load_tree(args.run_dir)
281 _node(tree, args.node) # validate
282 targets = _ancestors(tree, args.node)
283 if not targets:
284 sys.exit("error: node has no ancestors (is it the root?).")
285 if not args.to_root:
286 targets = targets[:1] # immediate parent only
287 for tid in targets:
288 anc = tree["nodes"][tid]
289 existing = anc.get("insight", "")
290 line = f"[from {args.node}] {args.insight}"
291 anc["insight"] = (existing + "\n" + line).strip() if existing else line
292 _stamp(anc)
293 _save(_tree_path(args.run_dir), tree)
294 print(f"Propagated insight from {args.node} to ancestors {targets}")
297def cmd_prune(args):
298 """Mark a node (and its subtree) pruned. Pruned hypotheses become negative
299 constraints -- record *why* so future ideation avoids the dead end."""
300 tree = _load_tree(args.run_dir)
301 _node(tree, args.node)
302 stack = [args.node]
303 pruned = []
304 while stack:
305 cur = stack.pop()
306 node = tree["nodes"][cur]
307 if node["status"] in ("merged", "root"):
308 continue
309 node["status"] = "pruned"
310 if cur == args.node and args.reason:
311 node["metadata"]["prune_reason"] = args.reason
312 _stamp(node)
313 pruned.append(cur)
314 stack.extend(_children(tree, cur))
315 _save(_tree_path(args.run_dir), tree)
316 print(f"Pruned {pruned}" + (f" -- reason: {args.reason}" if args.reason else ""))
319def cmd_merge(args):
320 """Record a held-out merge gate decision. Only call this AFTER evaluating
321 the candidate on the TEST evaluator in a fresh worktree. Admitting a
322 candidate that only improved dev defeats the purpose of the split."""
323 tree = _load_tree(args.run_dir)
324 run = _load_run(args.run_dir)
325 node = _node(tree, args.node)
327 direction = run["metric_direction"]
328 prev = run["best_test_score"]
329 if prev is None:
330 sys.exit("error: evaluate the initial artifact and run `tree.py baseline` before merging candidates.")
331 if not math.isfinite(prev):
332 sys.exit("error: stored best score is not finite; restore a valid scored run before merging.")
333 if args.node == tree["root"]:
334 sys.exit("error: merge a candidate node; the root stores the initial baseline.")
336 improves = args.test_score > prev if direction == "max" else args.test_score < prev
337 node["metadata"]["test_score"] = args.test_score
338 if args.branch_ref is not None:
339 node["metadata"]["branch_ref"] = args.branch_ref
340 if improves:
341 node["status"] = "merged"
342 run["best_node"] = args.node
343 run["best_test_score"] = args.test_score
344 run["best_branch_ref"] = args.branch_ref or node["metadata"].get("branch_ref")
345 _stamp(node)
346 _save(_tree_path(args.run_dir), tree)
347 _save(_run_path(args.run_dir), run)
348 print(
349 f"MERGE GATE PASSED: {args.node} test={args.test_score} "
350 f"beats previous best={prev}. M_best is now {args.node} "
351 f"(ref: {run['best_branch_ref']})."
352 )
353 else:
354 _stamp(node)
355 _save(_tree_path(args.run_dir), tree)
356 print(
357 f"MERGE GATE REJECTED: {args.node} test={args.test_score} does not beat "
358 f"best={prev} (direction={direction}). This is informative, not a failure: "
359 "a high-dev / low-test gap means the candidate may be exploiting the dev "
360 "signal. Record that lesson and keep M_best unchanged."
361 )
364def cmd_cycle(args):
365 """Increment the cycle counter (call once per Observe->Decide pass)."""
366 run = _load_run(args.run_dir)
367 run["cycles_used"] += 1
368 _save(_run_path(args.run_dir), run)
369 left = run["budget_cycles"] - run["cycles_used"]
370 print(f"Cycle {run['cycles_used']}/{run['budget_cycles']} ({left} remaining).")
371 if left <= 0:
372 print("Budget exhausted -- finish the run: do a final merge-gate check and report.")
375# ----------------------------------------------------------------------------
376# Read-only projections
377# ----------------------------------------------------------------------------
380def _fmt_score(node, run):
381 s = node["metadata"].get("dev_score")
382 t = node["metadata"].get("test_score")
383 bits = []
384 if s is not None:
385 bits.append(f"dev={s}")
386 if t is not None:
387 bits.append(f"test={t}")
388 return (" [" + " ".join(bits) + "]") if bits else ""
391def cmd_observe(args):
392 """The Observe step: a compact projection the coordinator re-grounds on at
393 the start of each cycle, so decisions come from the tree rather than from a
394 lossy conversation history."""
395 tree = _load_tree(args.run_dir)
396 run = _load_run(args.run_dir)
397 nodes = tree["nodes"]
398 root = nodes[tree["root"]]
400 print("=" * 72)
401 print("OBSERVE -- current research state")
402 print("=" * 72)
403 print(f"Objective : {run['objective']}")
404 print(f"Metric direction: {run['metric_direction']}")
405 print(f"Dev evaluator : {run['dev_eval']}")
406 print(f"Test evaluator : {run['test_eval']}")
407 print(
408 f"Budget : cycle {run['cycles_used']}/{run['budget_cycles']}, "
409 f"branching {run['branching']}, max depth {run['max_depth']}"
410 )
411 print(
412 f"Current best : "
413 + (
414 f"{run['best_node']} (test={run['best_test_score']}, ref={run['best_branch_ref']})"
415 if run["best_node"]
416 else "none yet -- M_best is the initial material"
417 )
418 )
420 print("\n-- Global insights (root) --")
421 print(root["insight"].strip() if root["insight"].strip() else " (none yet)")
423 # Active frontier = pending/running leaves
424 frontier = [
425 n
426 for n in nodes.values()
427 if n["status"] in ("pending", "running") and not _children(tree, n["id"])
428 ]
429 print("\n-- Active frontier (selectable hypotheses) --")
430 if not frontier:
431 print(" (empty -- ideate new children under a promising node)")
432 for n in sorted(frontier, key=lambda x: _id_sort_key(x["id"])):
433 anc = _ancestors(tree, n["id"])
434 anc_ins = " | ".join(
435 nodes[a]["insight"].replace("\n", " ")[:80] for a in anc if nodes[a]["insight"].strip()
436 )
437 print(f" {n['id']} (depth {n['depth']}, {n['status']}): {n['hypothesis']}")
438 if anc_ins:
439 print(f" ancestor insights: {anc_ins}")
441 # Validated / executed leaves with evidence
442 executed = [n for n in nodes.values() if n["status"] in ("executed", "merged")]
443 print("\n-- Executed / merged nodes (evidence) --")
444 if not executed:
445 print(" (none yet)")
446 for n in sorted(executed, key=lambda x: _id_sort_key(x["id"])):
447 print(f" {n['id']} [{n['status']}]{_fmt_score(n, run)}: {n['hypothesis']}")
448 if n["insight"].strip():
449 print(f" insight: {n['insight'].splitlines()[0][:120]}")
451 # Pruned lessons (negative constraints)
452 pruned = [n for n in nodes.values() if n["status"] == "pruned"]
453 print("\n-- Pruned lessons (negative constraints -- avoid these) --")
454 if not pruned:
455 print(" (none yet)")
456 for n in sorted(pruned, key=lambda x: _id_sort_key(x["id"])):
457 reason = n["metadata"].get("prune_reason", "")
458 print(f" {n['id']}: {n['hypothesis']}" + (f" -- {reason}" if reason else ""))
460 print("\n" + "=" * 72)
461 print(
462 "Next: Ideate children under a promising node, Select one or more frontier\n"
463 "leaves, Dispatch each to an executor subagent in an isolated worktree."
464 )
467def cmd_status(args):
468 """ASCII tree render -- useful for reports and quick scans."""
469 tree = _load_tree(args.run_dir)
470 run = _load_run(args.run_dir)
471 nodes = tree["nodes"]
473 symbol = {
474 "root": "*",
475 "pending": "o",
476 "running": "~",
477 "executed": "=",
478 "merged": "V",
479 "pruned": "x",
480 }
482 def render(nid, prefix=""):
483 n = nodes[nid]
484 sym = symbol.get(n["status"], "?")
485 best = " <== M_best" if nid == run["best_node"] else ""
486 print(f"{prefix}[{sym}] {nid} {n['hypothesis'][:70]}{_fmt_score(n, run)}{best}")
487 kids = sorted(_children(tree, nid), key=_id_sort_key)
488 for i, c in enumerate(kids):
489 render(c, prefix + " ")
491 print(f"Run: {run['objective']}")
492 print(f"Cycle {run['cycles_used']}/{run['budget_cycles']} | legend: * root o pending ~ running = executed V merged x pruned\n")
493 render(tree["root"])
496def cmd_validate(args):
497 """Check invariants -- catch a corrupted or inconsistent tree early."""
498 tree = _load_tree(args.run_dir)
499 run = _load_run(args.run_dir)
500 nodes = tree["nodes"]
501 problems = []
502 if tree["root"] not in nodes:
503 problems.append("root id not present in nodes")
504 for nid, n in nodes.items():
505 p = n.get("parent")
506 if p is not None and p not in nodes:
507 problems.append(f"{nid}: parent {p} missing")
508 if n["status"] not in VALID_STATUS:
509 problems.append(f"{nid}: invalid status {n['status']}")
510 if run["best_node"] and run["best_node"] not in nodes:
511 problems.append(f"best_node {run['best_node']} missing")
512 merged = [nid for nid, n in nodes.items() if n["status"] == "merged"]
513 if run["best_node"] and run["best_node"] not in merged and run["best_node"] != tree["root"]:
514 problems.append(f"best_node {run['best_node']} is not marked merged")
515 if problems:
516 print("INVALID:")
517 for p in problems:
518 print(f" - {p}")
519 sys.exit(1)
520 print(f"OK -- {len(nodes)} nodes, root={tree['root']}, best={run['best_node']}")
523# ----------------------------------------------------------------------------
524# CLI
525# ----------------------------------------------------------------------------
528def build_parser():
529 p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
530 p.add_argument("--run-dir", default=".", help="Run directory holding .arbor/ (default: current dir)")
531 sub = p.add_subparsers(dest="command", required=True)
533 s = sub.add_parser("init", help="Initialize a new AO run / hypothesis tree")
534 s.add_argument("--objective", required=True, help="Natural-language research objective (root hypothesis)")
535 s.add_argument("--dev-eval", required=True, help="Command/description of the development evaluator")
536 s.add_argument("--test-eval", required=True, help="Command/description of the held-out test evaluator")
537 s.add_argument("--material", default="", help="Path/ref to the initial artifact M_0")
538 s.add_argument("--metric-direction", choices=["max", "min"], default="max", help="Is higher or lower the better score?")
539 s.add_argument("--branching", type=int, default=3, help="Max children proposed per parent (k)")
540 s.add_argument("--max-depth", type=int, default=2, help="Max tree depth (directions at 1, interventions at 2+)")
541 s.add_argument("--budget", type=int, default=20, help="Coordinator cycle budget B")
542 s.add_argument("--force", action="store_true", help="Overwrite an existing run")
543 s.set_defaults(func=cmd_init)
545 s = sub.add_parser("baseline", help="Record the initial artifact's held-out score before candidate merges")
546 s.add_argument("--test-score", type=_finite_score, required=True)
547 s.add_argument("--branch-ref", required=True, help="Immutable git commit or named branch of M_0")
548 s.set_defaults(func=cmd_baseline)
550 s = sub.add_parser("observe", help="Print the research-state projection (start of each cycle)")
551 s.set_defaults(func=cmd_observe)
553 s = sub.add_parser("add-node", help="Add a pending child hypothesis under a parent (Ideate)")
554 s.add_argument("--parent", required=True, help="Parent node id (n0 for a new research direction)")
555 s.add_argument("--hypothesis", required=True, help="Falsifiable claim this node tests")
556 s.set_defaults(func=cmd_add_node)
558 s = sub.add_parser("set-status", help="Set a node's status manually")
559 s.add_argument("--node", required=True)
560 s.add_argument("--status", required=True, help=f"One of {sorted(VALID_STATUS)}")
561 s.set_defaults(func=cmd_set_status)
563 s = sub.add_parser("set-evidence", help="Write an executor report into its node (Backpropagate, leaf)")
564 s.add_argument("--node", required=True)
565 s.add_argument("--dev-score", type=_finite_score, default=None, help="Dev evaluator score returned by the executor")
566 s.add_argument("--result", default=None, help="Factual result summary")
567 s.add_argument("--insight", default=None, help="Distilled, reusable lesson from this experiment")
568 s.add_argument("--branch-ref", default=None, help="Git branch/commit/worktree path of the artifact")
569 s.add_argument("--status", default=None, help=f"Override status (default: executed), one of {sorted(VALID_STATUS)}")
570 s.set_defaults(func=cmd_set_evidence)
572 s = sub.add_parser("propagate", help="Abstract a leaf insight up to ancestors (Backpropagate, upward)")
573 s.add_argument("--node", required=True, help="The leaf the lesson came from")
574 s.add_argument("--insight", required=True, help="Direction-level abstraction of the lesson")
575 s.add_argument("--to-root", action="store_true", help="Also record as a global insight on the root")
576 s.set_defaults(func=cmd_propagate)
578 s = sub.add_parser("prune", help="Prune a falsified node and its subtree (Decide)")
579 s.add_argument("--node", required=True)
580 s.add_argument("--reason", default="", help="Why this direction is a dead end (becomes a negative constraint)")
581 s.set_defaults(func=cmd_prune)
583 s = sub.add_parser("merge", help="Record a held-out merge gate decision (Decide)")
584 s.add_argument("--node", required=True)
585 s.add_argument("--test-score", type=_finite_score, required=True, help="Score on the TEST evaluator in a fresh worktree")
586 s.add_argument("--branch-ref", default=None, help="Artifact ref to promote if it passes")
587 s.set_defaults(func=cmd_merge)
589 s = sub.add_parser("cycle", help="Increment the coordinator cycle counter")
590 s.set_defaults(func=cmd_cycle)
592 s = sub.add_parser("status", help="Render the tree as ASCII (for reports)")
593 s.set_defaults(func=cmd_status)
595 s = sub.add_parser("validate", help="Check tree invariants")
596 s.set_defaults(func=cmd_validate)
598 return p
601def main():
602 args = build_parser().parse_args()
603 args.func(args)
606if __name__ == "__main__":
607 main()