| """Run the implemented exact and neural comparisons from a YAML config.""" |
|
|
| import argparse, json, time, platform |
| from pathlib import Path |
| import numpy as np, pandas as pd, torch, yaml |
| from dooable.graph import Graph, toy_graph, grid_graph, string_graph |
| from dooable.exact import ( |
| solve, |
| uniform_policy, |
| tilted_reference, |
| backward_policy, |
| forward_from_backward, |
| sample, |
| ) |
| from dooable.learning import train |
| from dooable.metrics import graph_metrics |
|
|
|
|
| def main(): |
| p = argparse.ArgumentParser() |
| p.add_argument("--config", required=True) |
| a = p.parse_args() |
| cfg = yaml.safe_load(Path(a.config).read_text()) |
| out = Path(cfg["output"]) |
| out.mkdir(parents=True, exist_ok=True) |
| if cfg.get("graph"): |
| g = Graph.load(cfg["graph"]) |
| elif cfg.get("kind") == "grid": |
| g = grid_graph(cfg.get("width", 5), cfg.get("budget", 6)) |
| elif cfg.get("kind") == "strings": |
| g = string_graph(cfg.get("length", 4), cfg.get("budget", 2)) |
| else: |
| g = toy_graph(cfg.get("multiplicity", 8)) |
| rewards = ( |
| json.loads(Path(cfg["rewards"]).read_text()) |
| if cfg.get("rewards") |
| else {y: 0.0 for y in g.terminals} |
| ) |
| temperature = cfg.get("temperature", 0.7) |
| rows = [] |
| g.save(out / "graph.json") |
| for seed in cfg.get("seeds", [0, 1, 2, 3, 4]): |
| for name in cfg["methods"]: |
| start = time.perf_counter() |
| if name == "exact": |
| forward = solve(g, rewards, temperature).forward |
| elif name == "uniform": |
| forward = uniform_policy(g) |
| elif name == "reference_tilt": |
| forward = tilted_reference(g, rewards) |
| elif name == "zero_cost": |
| from dooable.ablations import zero_cost_policy |
|
|
| forward = zero_cost_policy(g, rewards) |
| elif name == "duplicate_endpoints": |
| from dooable.ablations import duplicate_endpoint_policy |
|
|
| forward = duplicate_endpoint_policy(g, rewards, temperature) |
| elif name in ["dooable", "tb_uniform", "tb_exact", "unnormalized"]: |
| backward = { |
| "dooable": "learned", |
| "tb_uniform": "uniform", |
| "tb_exact": "exact", |
| "unnormalized": "unnormalized", |
| }[name] |
| model, _ = train( |
| g, |
| rewards, |
| temperature, |
| steps=cfg.get("steps", 2000), |
| batch_size=cfg.get("batch_size", 64), |
| seed=seed, |
| output=out / f"{name}_seed{seed}", |
| backward=backward, |
| ) |
| forward = model.probabilities() |
| else: |
| raise ValueError(f"Unimplemented comparator {name}") |
| row = { |
| "method": name, |
| "seed": seed, |
| "seconds": time.perf_counter() - start, |
| "nodes": len(g.nodes), |
| "edges": len(g.edges), |
| "outcomes": len(g.terminals), |
| **graph_metrics(g, forward, rewards, temperature), |
| } |
| if g.metadata.get("kind") == "reaction": |
| from dooable.chemistry import replay |
|
|
| paths = sample(g, forward, cfg.get("samples", 1000), seed) |
| row["replay_fraction"] = np.mean( |
| [replay(r, g.metadata["budget"]) for r in paths] |
| ) |
| row["unique_outcomes"] = len({r["outcome"] for r in paths}) |
| (out / f"{name}_seed{seed}_samples.jsonl").write_text( |
| "".join(json.dumps(r) + "\n" for r in paths) |
| ) |
| rows.append(row) |
| pd.DataFrame(rows).to_csv(out / "metrics.csv", index=False) |
| print(json.dumps(row), flush=True) |
| df = pd.DataFrame(rows) |
| numeric = [c for c in df.select_dtypes("number").columns if c != "seed"] |
| df.groupby("method")[numeric].agg(["mean", "sem"]).to_csv(out / "summary.csv") |
| (out / "run.json").write_text( |
| json.dumps( |
| { |
| "config": cfg, |
| "python": platform.python_version(), |
| "torch": torch.__version__, |
| "numpy": np.__version__, |
| "platform": platform.platform(), |
| }, |
| indent=2, |
| ) |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
