224 lines
6.8 KiB
Python
224 lines
6.8 KiB
Python
"""Minimal AlphaEvolve-like evolutionary loop — stdlib Python.
|
|
|
|
Toy symbolic regression. The "LLM" proposes a small mutation to a candidate
|
|
expression (change a constant, change an operator, add a term). The
|
|
"evaluator" scores the expression on training and held-out test points.
|
|
|
|
MAP-elites grid keeps diverse candidates: cell keyed by (expression depth,
|
|
constant magnitude bucket). Without a held-out split the loop overfits
|
|
aggressively; with one the best candidate generalizes.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import math
|
|
import random
|
|
from dataclasses import dataclass
|
|
|
|
|
|
DEFAULT_SEED = 1
|
|
|
|
|
|
# Target function the loop tries to rediscover.
|
|
def target(x: float) -> float:
|
|
return 2.0 * x * x + 3.0 * x - 1.0
|
|
|
|
|
|
Expr = tuple # recursive: ("num", v) | ("x",) | ("add", a, b) | ("mul", a, b)
|
|
|
|
|
|
def evaluate_expr(e: Expr, x: float) -> float:
|
|
tag = e[0]
|
|
if tag == "num":
|
|
return float(e[1])
|
|
if tag == "x":
|
|
return x
|
|
if tag == "add":
|
|
return evaluate_expr(e[1], x) + evaluate_expr(e[2], x)
|
|
if tag == "mul":
|
|
return evaluate_expr(e[1], x) * evaluate_expr(e[2], x)
|
|
raise ValueError(tag)
|
|
|
|
|
|
def depth(e: Expr) -> int:
|
|
tag = e[0]
|
|
if tag in ("num", "x"):
|
|
return 1
|
|
return 1 + max(depth(e[1]), depth(e[2]))
|
|
|
|
|
|
def max_const(e: Expr) -> float:
|
|
tag = e[0]
|
|
if tag == "num":
|
|
return abs(e[1])
|
|
if tag == "x":
|
|
return 0.0
|
|
return max(max_const(e[1]), max_const(e[2]))
|
|
|
|
|
|
def mutate(e: Expr) -> Expr:
|
|
"""Stand-in for the LLM's targeted edit."""
|
|
choice = random.random()
|
|
if choice < 0.25:
|
|
return random_leaf()
|
|
if choice < 0.5:
|
|
return ("add", e, random_leaf())
|
|
if choice < 0.75:
|
|
return ("mul", e, random_leaf())
|
|
# perturb a constant somewhere
|
|
return perturb(e)
|
|
|
|
|
|
def perturb(e: Expr) -> Expr:
|
|
tag = e[0]
|
|
if tag == "num":
|
|
return ("num", e[1] + random.choice([-1.0, -0.5, 0.5, 1.0]))
|
|
if tag == "x":
|
|
return e
|
|
return (tag, perturb(e[1]), e[2]) if random.random() < 0.5 else (tag, e[1], perturb(e[2]))
|
|
|
|
|
|
def random_leaf() -> Expr:
|
|
if random.random() < 0.5:
|
|
return ("x",)
|
|
return ("num", float(random.choice([-2, -1, 0, 1, 2, 3])))
|
|
|
|
|
|
def render(e: Expr) -> str:
|
|
tag = e[0]
|
|
if tag == "num":
|
|
return f"{e[1]:g}"
|
|
if tag == "x":
|
|
return "x"
|
|
op = "+" if tag == "add" else "*"
|
|
return f"({render(e[1])} {op} {render(e[2])})"
|
|
|
|
|
|
def mse(e: Expr, xs: list[float]) -> float:
|
|
total = 0.0
|
|
for x in xs:
|
|
try:
|
|
y = evaluate_expr(e, x)
|
|
except (OverflowError, ValueError):
|
|
return float("inf")
|
|
total += (y - target(x)) ** 2
|
|
return total / max(1, len(xs))
|
|
|
|
|
|
@dataclass
|
|
class Candidate:
|
|
expr: Expr
|
|
train_score: float
|
|
test_score: float
|
|
generation: int
|
|
|
|
|
|
def cell_key(e: Expr) -> tuple[int, int]:
|
|
d = min(depth(e), 6)
|
|
c = min(int(max_const(e) / 2), 4)
|
|
return (d, c)
|
|
|
|
|
|
def seed_candidate(test_xs: list[float], train_xs: list[float], gen: int) -> Candidate:
|
|
e = random_leaf()
|
|
return Candidate(e, mse(e, train_xs), mse(e, test_xs), gen)
|
|
|
|
|
|
def run_loop(
|
|
generations: int,
|
|
pop: int,
|
|
use_holdout: bool,
|
|
seed: int | None = None,
|
|
) -> tuple[Candidate, list[float], list[float]]:
|
|
if seed is not None:
|
|
random.seed(seed)
|
|
train_xs = [-2.0, -1.0, 0.0, 1.0, 2.0, 3.0]
|
|
test_xs = [-2.5, -1.5, -0.5, 0.5, 1.5, 2.5, 3.5]
|
|
|
|
def signal_of(c: Candidate) -> float:
|
|
return 0.5 * (c.train_score + c.test_score) if use_holdout else c.train_score
|
|
|
|
archive: dict[tuple[int, int], Candidate] = {}
|
|
for _ in range(pop):
|
|
c = seed_candidate(test_xs, train_xs, 0)
|
|
key = cell_key(c.expr)
|
|
incumbent = archive.get(key)
|
|
if incumbent is None or signal_of(c) < signal_of(incumbent):
|
|
archive[key] = c
|
|
|
|
best_trace: list[float] = []
|
|
test_trace: list[float] = []
|
|
for g in range(1, generations + 1):
|
|
parent = random.choice(list(archive.values()))
|
|
child_expr = mutate(parent.expr)
|
|
tr = mse(child_expr, train_xs)
|
|
te = mse(child_expr, test_xs)
|
|
child = Candidate(child_expr, tr, te, g)
|
|
key = cell_key(child_expr)
|
|
incumbent = archive.get(key)
|
|
if incumbent is None or signal_of(child) < signal_of(incumbent):
|
|
archive[key] = child
|
|
|
|
best = min(archive.values(), key=lambda c: c.train_score)
|
|
best_trace.append(best.train_score)
|
|
test_trace.append(best.test_score)
|
|
|
|
# Final selection must use the same signal as the search: using the
|
|
# held-out test here when use_holdout=False would silently leak the
|
|
# holdout back into Run B and mask the overfitting the lesson shows.
|
|
best = min(archive.values(), key=signal_of)
|
|
return best, best_trace, test_trace
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument(
|
|
"--no-holdout",
|
|
action="store_true",
|
|
help="skip the held-out test evaluator (Run B only; forces reward-hacking demo)",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
print("=" * 70)
|
|
print("ALPHAEVOLVE-STYLE LOOP (Phase 15, Lesson 3)")
|
|
print("=" * 70)
|
|
print("target: 2x^2 + 3x - 1")
|
|
|
|
if not args.no_holdout:
|
|
print("\nRun A: held-out test included in evaluator signal")
|
|
best, train_trace, _ = run_loop(
|
|
generations=1500, pop=20, use_holdout=True, seed=DEFAULT_SEED
|
|
)
|
|
print(f" best expr : {render(best.expr)}")
|
|
print(f" train MSE : {best.train_score:.4f}")
|
|
print(f" test MSE : {best.test_score:.4f}")
|
|
print(f" generation: {best.generation}")
|
|
print(" progress : gen 100 train={:.3f} gen 500 train={:.3f} gen 1500 train={:.3f}".format(
|
|
train_trace[99], train_trace[499], train_trace[-1]))
|
|
|
|
print("\nRun B: no held-out test (train-only evaluator -> reward hacking risk)")
|
|
best, _train_trace, _test_trace = run_loop(
|
|
generations=1500, pop=20, use_holdout=False, seed=DEFAULT_SEED
|
|
)
|
|
print(f" best expr : {render(best.expr)}")
|
|
print(f" train MSE : {best.train_score:.4f}")
|
|
print(f" test MSE : {best.test_score:.4f}")
|
|
print(f" generation: {best.generation}")
|
|
gap = best.test_score - best.train_score
|
|
print(f" train-to-test gap: {gap:+.4f} (large gap = overfit/reward hacking proxy)")
|
|
|
|
print()
|
|
print("=" * 70)
|
|
print("HEADLINE: the evaluator is the architecture")
|
|
print("-" * 70)
|
|
print(" Run A converges to low train AND low test MSE.")
|
|
print(" Run B converges to low train MSE; test MSE stays loose or worse.")
|
|
print(" A held-out evaluator is the difference between discovery and")
|
|
print(" reward hacking. AlphaEvolve's wins are in domains where such an")
|
|
print(" evaluator exists. Picking those domains is the hard part.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|