File size: 4,471 Bytes
81ae663 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | """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()
|