"""optimization_toolkit.py - Week 5b lab: optimizing agents with mocks and toy problems. Standard library only (tested on Python 3.13). Every "model" here is a few lines of Python that act like one. All tasks, skill levels, rewards, prices and GPU hours are MADE UP for teaching. Nothing here measures a real model, and no number here should be quoted as a result about real systems. Parts 1 A toy prompt world: a mock model whose accuracy depends on its instruction and demos 2 Context optimization: bootstrap few-shot, MIPRO-like search, OPRO loop, a TextGrad-style feedback step, GEPA-style reflective mutation with Pareto selection 3 Weight optimization on a tabular toy policy: REINFORCE, PPO-clip, GRPO, SFT from filtered rollouts, and reverse-KL on-policy distillation 4 A DGM-style archive of agent variants (toy landscape) 5 An autoresearch-style loop with an eval guard 6 A prompt-versus-weights cost calculator (made-up prices) """ from __future__ import annotations import hashlib import math import random import re import sys from dataclasses import dataclass, field, replace # --------------------------------------------------------------------------- 0. shared bits def count_tokens(text) -> int: """Rough token proxy: words plus punctuation marks. Not a real tokenizer.""" return len(re.findall(r"\w+|[^\w\s]", str(text))) class BudgetExhausted(RuntimeError): """Raised when an optimizer tries to use more rollouts than it was given.""" @dataclass class Meter: """Counts rollouts (one program run on one example) and tokens. Optional rollout cap.""" max_rollouts: int | None = None rollouts: int = 0 tokens: int = 0 def charge(self, n=1, tokens=0): if self.max_rollouts is not None and self.rollouts + n > self.max_rollouts: raise BudgetExhausted(f"{self.rollouts}+{n} > {self.max_rollouts}") self.rollouts += n self.tokens += tokens @property def left(self): return None if self.max_rollouts is None else self.max_rollouts - self.rollouts def mean(xs): xs = list(xs) return sum(xs) / len(xs) if xs else 0.0 def stdev(xs): """Population standard deviation.""" xs = list(xs) if not xs: return 0.0 m = mean(xs) return math.sqrt(sum((x - m) ** 2 for x in xs) / len(xs)) # --------------------------------------------------------------------------- 1. toy prompt world RULES = {"A": "units", "B": "negatives", "C": "dates", "D": "abbrev", "E": "currency", "F": "ties"} MODULE_OF = {"A": "parse", "B": "parse", "C": "parse", "D": "answer", "E": "answer", "F": "answer"} MODULES = ("parse", "answer") DISTRACTORS = ("carefully", "persona", "stakes", "concise", "friendly", "step") VOCAB = tuple(RULES.values()) + DISTRACTORS KIND_OF_WORD = {w: k for k, w in RULES.items()} # (keyword, kind): naming the keyword in that kind's module instruction costs this much on that kind. CONFLICTS = {("negatives", "C"): 0.30, ("abbrev", "F"): 0.25} N_CHOICES = 4 @dataclass(frozen=True) class ToyParams: """MADE-UP accuracy model for the mock LM.""" base: float = 0.35 # chance of a right answer with no help instr_gain: float = 0.40 # the instruction names the rule for this kind demo_gain: float = 0.10 # per correct demo of this kind (at most 2 count) bad_demo_pen: float = 0.10 # per wrong demo of this kind (at most 2 count) length_pen: float = 0.02 # per instruction word beyond free_words free_words: int = 10 demo_len_pen: float = 0.01 # per demo beyond free_demos (context bloat) free_demos: int = 6 floor: float = 0.02 cap: float = 0.98 @dataclass(frozen=True) class Example: id: int kind: str answer: int @dataclass(frozen=True) class Demo: id: int kind: str answer: int correct: bool = True @dataclass(frozen=True) class Program: """A two-module toy program. instructions: module -> text. demos: a tuple of Demo.""" instructions: dict = field(default_factory=lambda: {m: "Solve the task." for m in MODULES}) demos: tuple = () def with_instruction(self, module, text): if module not in MODULES: raise KeyError(module) ins = dict(self.instructions) ins[module] = text return Program(ins, self.demos) def with_demos(self, demos): return Program(dict(self.instructions), tuple(demos)) def words(self, module): return self.instructions.get(module, "").split() def make_dataset(n, seed=0, weights=None): """n examples with kinds spread over A..F (optionally weighted), answers 0..3.""" rng = random.Random(seed) kinds = sorted(RULES) ws = [1.0] * len(kinds) if weights is None else [weights.get(k, 0.0) for k in kinds] if sum(ws) <= 0: raise ValueError("weights must have a positive sum") return [Example(i, rng.choices(kinds, ws)[0], rng.randrange(N_CHOICES)) for i in range(n)] def p_correct(prog, kind, params=ToyParams()): """The mock model's chance of answering an example of this kind correctly.""" m = MODULE_OF[kind] words = set(prog.words(m)) p = params.base if RULES[kind] in words: p += params.instr_gain good = sum(1 for d in prog.demos if d.kind == kind and d.correct) bad = sum(1 for d in prog.demos if d.kind == kind and not d.correct) p += min(good, 2) * params.demo_gain - min(bad, 2) * params.bad_demo_pen for (kw, k), pen in CONFLICTS.items(): if k == kind and kw in words: p -= pen p -= params.length_pen * max(0, len(prog.words(m)) - params.free_words) p -= params.demo_len_pen * max(0, len(prog.demos) - params.free_demos) return min(params.cap, max(params.floor, p)) @dataclass class Trace: example: Example module: str output: int correct: bool feedback: str rationale: str tokens: int def feedback_text(prog, ex, correct): m = MODULE_OF[ex.kind] if correct: return f"ok: kind {ex.kind} ({RULES[ex.kind]}) was right" if RULES[ex.kind] not in prog.words(m): return f"fail: kind {ex.kind} needs the rule '{RULES[ex.kind]}', and the {m} instruction does not mention it" return f"fail: kind {ex.kind}: the rule '{RULES[ex.kind]}' is stated, so this miss looks like noise" def rollout(prog, ex, rng, params=ToyParams(), meter=None): """Run the program on one example. One rollout = one charge on the meter.""" m = MODULE_OF[ex.kind] ok = rng.random() < p_correct(prog, ex.kind, params) if ok: out = ex.answer else: out = (ex.answer + 1 + rng.randrange(N_CHOICES - 1)) % N_CHOICES toks = len(prog.words(m)) + 20 * len(prog.demos) + 15 + 10 if meter is not None: meter.charge(1, toks) return Trace(ex, m, out, out == ex.answer, feedback_text(prog, ex, ok), f"think about kind {ex.kind}, then answer {out}", toks) def scores_on(prog, examples, rng, params=ToyParams(), meter=None): """Per-example 0/1 scores.""" return [1.0 if rollout(prog, ex, rng, params, meter).correct else 0.0 for ex in examples] def evaluate(prog, examples, rng, params=ToyParams(), meter=None): return mean(scores_on(prog, examples, rng, params, meter)) def exact_accuracy(prog, weights=None, params=ToyParams()): """Expected accuracy with no sampling noise, for a kind distribution (default: uniform).""" ks = sorted(RULES) ws = {k: 1.0 for k in ks} if weights is None else {k: weights.get(k, 0.0) for k in ks} tot = sum(ws.values()) return sum(ws[k] / tot * p_correct(prog, k, params) for k in ks) def per_kind_accuracy(prog, params=ToyParams()): return {k: p_correct(prog, k, params) for k in sorted(RULES)} def best_possible_program(params=ToyParams()): """Brute-force the best instruction (a subset of rule words) per module, no demos. Used only as a ruler in the demo and tests.""" best = {} for m in MODULES: kws = [w for k, w in RULES.items() if MODULE_OF[k] == m] top = None for mask in range(1 << len(kws)): chosen = [kws[i] for i in range(len(kws)) if mask >> i & 1] prog = Program({**{x: "Solve the task." for x in MODULES}, m: " ".join(["Solve"] + chosen)}) acc = sum(p_correct(prog, k, params) for k in RULES if MODULE_OF[k] == m) if top is None or acc > top[0] + 1e-12: top = (acc, " ".join(["Solve"] + chosen)) best[m] = top[1] return Program(best) # --------------------------------------------------------------------------- 2. context optimization def uniform_program(instruction): """One instruction used in every module (the single-prompt setting of OPRO).""" return Program({m: instruction for m in MODULES}) def bootstrap_demos(prog, train, rng, params=ToyParams(), meter=None, max_per_kind=2, keep_all=False): """Bootstrapped few-shot: run the program on training examples and keep traces that pass the metric (exact match). keep_all=True skips the filter, to show what the filter is worth.""" demos, per = [], {} for ex in train: t = rollout(prog, ex, rng, params, meter) if (t.correct or keep_all) and per.get(ex.kind, 0) < max_per_kind: per[ex.kind] = per.get(ex.kind, 0) + 1 demos.append(Demo(ex.id, ex.kind, t.output, t.correct)) return demos def bootstrap_random_search(prog, train, dev, rng, n_sets=8, k=4, params=ToyParams(), meter=None): """Bootstrap demos, then random-search subsets of size k on the dev set (the DSPy recipe). Returns (best program, list of (demo ids, dev score)).""" pool = bootstrap_demos(prog, train, rng, params, meter, max_per_kind=3) sets = [()] for _ in range(n_sets): if pool: sets.append(tuple(sorted(rng.sample(pool, min(k, len(pool))), key=lambda d: d.id))) log, best = [], None for s in sets: cand = prog.with_demos(s) sc = evaluate(cand, dev, rng, params, meter) log.append((tuple(d.id for d in s), sc)) if best is None or sc > best[0]: best = (sc, cand) return best[1], log # ---- MIPRO-like ----------------------------------------------------------------------------- def failure_summary(prog, train, rng, params=ToyParams(), meter=None): """Data-aware grounding: failure rate per kind on the training set.""" fails, seen = {}, {} for ex in train: t = rollout(prog, ex, rng, params, meter) seen[ex.kind] = seen.get(ex.kind, 0) + 1 fails[ex.kind] = fails.get(ex.kind, 0) + (0 if t.correct else 1) return {k: fails[k] / seen[k] for k in seen} def propose_instructions(module, n, rng, fail_rates=None, program_aware=True, tip=False): """A mock proposer. With program_aware it only names rules that belong to the module. With fail_rates (data-aware) it weights rules by how often their kind failed. Without any grounding it draws from the whole vocabulary. tip=True adds a flavor word (does nothing).""" own = [w for k, w in RULES.items() if MODULE_OF[k] == module] out, tries = [], 0 while len(out) < n and tries < 50 * n: tries += 1 pool = own if program_aware else list(VOCAB) if fail_rates is not None and program_aware: ws = [fail_rates.get(KIND_OF_WORD[w], 0.5) + 0.05 for w in pool] else: ws = [1.0] * len(pool) picks = [] for _ in range(rng.randint(1, 3)): picks.append(rng.choices(pool, ws)[0]) if tip: picks.append(rng.choice(DISTRACTORS)) text = " ".join(["Solve"] + list(dict.fromkeys(picks))) if text not in out: out.append(text) return out def ucb_choice(stats, rng, surrogate=True, explore=0.25): """Pick one option from [[sum of scores, pulls], ...]. With the surrogate on, pick the option with the best mean plus an exploration bonus 0.5 / sqrt(pulls + 1) (an untried option has mean 0.5), and with probability `explore` pick uniformly at random. With the surrogate off, always pick uniformly at random.""" if not surrogate or rng.random() < explore: return rng.randrange(len(stats)) best, arg = None, [] for i, (tot, n) in enumerate(stats): mu = tot / n if n else 0.5 u = mu + 0.5 / math.sqrt(n + 1) if best is None or u > best + 1e-12: best, arg = u, [i] elif abs(u - best) <= 1e-12: arg.append(i) return rng.choice(arg) @dataclass class MiproResult: program: Program score: float history: list rollouts: int belief: dict = field(default_factory=dict) full_evals: list = field(default_factory=list) # (trial, dev score, choice) def mipro_search(prog, train, dev, rng, n_instr=5, n_demo_sets=4, trials=24, minibatch=10, full_every=8, grounded=True, surrogate=True, params=ToyParams()): """A MIPRO-like search. Step 1 bootstrap demos. Step 2 propose instructions per module. Step 3 search combinations with minibatch trials, an additive belief per choice, and a full dev evaluation every full_every trials. The belief is a UCB score, my stand-in for the tree-structured Parzen estimator that the real optimizer uses [1]. surrogate=False picks combinations at random with everything else equal.""" meter = Meter() fr = failure_summary(prog, train, rng, params, meter) if grounded else None pool = bootstrap_demos(prog, train, rng, params, meter, max_per_kind=2) demo_sets = [()] + [tuple(rng.sample(pool, min(4, len(pool)))) for _ in range(n_demo_sets)] cands = {} for m in MODULES: cands[m] = [prog.instructions[m]] + propose_instructions(m, n_instr, rng, fr, program_aware=grounded, tip=grounded) sizes = {"parse": len(cands["parse"]), "answer": len(cands["answer"]), "demo": len(demo_sets)} stat = {v: [[0.0, 0] for _ in range(sizes[v])] for v in sizes} history, combos = [], {} def build(choice): p = Program({m: cands[m][choice[m]] for m in MODULES}, demo_sets[choice["demo"]]) return p def pick(v): return ucb_choice(stat[v], rng, surrogate) full = [] for t in range(1, trials + 1): choice = {v: pick(v) for v in sizes} batch = rng.sample(train, min(minibatch, len(train))) sc = evaluate(build(choice), batch, rng, params, meter) for v in sizes: stat[v][choice[v]][0] += sc stat[v][choice[v]][1] += 1 key = tuple(choice[v] for v in sorted(sizes)) combos.setdefault(key, []).append(sc) history.append((t, key, sc)) if t % full_every == 0 or t == trials: top = sorted(combos, key=lambda c: -mean(combos[c]))[:2] for c in top: ch = dict(zip(sorted(sizes), c)) full.append((evaluate(build(ch), dev, rng, params, meter), ch, t)) sc, ch, _ = max(full, key=lambda x: x[0]) return MiproResult(build(ch), sc, history, meter.rollouts, stat, [(t_, sc_, c_) for sc_, c_, t_ in full]) # ---- OPRO --------------------------------------------------------------------------------- def opro_prompt(history, k=6): """The meta-prompt: past instructions with their scores, ascending, so the best is last.""" top = sorted(history, key=lambda h: h[1])[-k:] lines = ["Here are earlier instructions and their scores. Write a new instruction that scores higher."] for text, score in top: lines.append(f"text: {text}\nscore: {score:.2f}") return "\n".join(lines) def parse_opro_prompt(prompt): """Read (text, score) pairs back out of a meta-prompt.""" pairs = re.findall(r"text: (.*)\nscore: ([0-9.]+)", prompt) return [(t, float(s)) for t, s in pairs] def opro_propose(prompt, n, rng, max_words=12): """A mock proposer. It reads the meta-prompt, then edits the best instruction. It cannot reason about why a score was high, so this is selection plus random edits, not insight.""" hist = parse_opro_prompt(prompt) ranked = sorted(hist, key=lambda h: h[1]) best = ranked[-1][0].split() if ranked else ["Solve"] second = ranked[-2][0].split() if len(ranked) > 1 else best out = [] for _ in range(n): r = rng.random() words = list(best if r < 0.85 else second) extra = [w for w in VOCAB if w not in words] if r < 0.6 and extra: words.append(rng.choice(extra)) elif r < 0.85 and len(words) > 1 and extra: words.pop(rng.randrange(1, len(words))) words.append(rng.choice(extra)) elif extra: words.append(rng.choice(extra)) out.append(" ".join(words[:max_words])) return out @dataclass class OproResult: instruction: str score: float history: list rollouts: int def opro(train, rng, steps=12, per_step=4, params=ToyParams(), seed_instruction="Solve the task."): meter = Meter() hist = [(seed_instruction, evaluate(uniform_program(seed_instruction), train, rng, params, meter))] seen = {seed_instruction} for _ in range(steps): prompt = opro_prompt(hist) for text in opro_propose(prompt, per_step, rng): if text in seen: continue seen.add(text) hist.append((text, evaluate(uniform_program(text), train, rng, params, meter))) best = max(hist, key=lambda h: h[1]) return OproResult(best[0], best[1], hist, meter.rollouts) # ---- TextGrad-style --------------------------------------------------------------------------- @dataclass class TGVar: """A node that holds text. requires_grad says whether the optimizer may change it.""" value: str role: str requires_grad: bool = True grads: list = field(default_factory=list) def tg_loss_text(traces): """The 'loss': a judge's written critique of a minibatch (here, the failure feedback lines).""" return "\n".join(t.feedback for t in traces if not t.correct) def tg_critic(var_module, loss_text, downstream): """The mock critic that writes one variable's textual gradient from the loss text and from the gradients of the variables downstream of it. It asks for rules that failures named.""" asks = [] for line in loss_text.splitlines(): m = re.match(r"fail: kind (\w) needs the rule '(\w+)'", line) if m and MODULE_OF[m.group(1)] == var_module: asks.append(m.group(2)) for g in downstream: for w in re.findall(r"consider naming: (\w+)", g): if KIND_OF_WORD.get(w) and MODULE_OF[KIND_OF_WORD[w]] == var_module: asks.append(w) if not asks: return "" counts = {} for w in asks: counts[w] = counts.get(w, 0) + 1 ranked = sorted(counts, key=lambda w: (-counts[w], w)) return "consider naming: " + ", ".join(ranked) def tg_backward(graph, loss_text): """graph: list of (module, TGVar) in forward order. Walks in reverse order (loss first). Returns the order visited. Variables with requires_grad=False get no gradient.""" order = [] downstream = [] for module, var in reversed(graph): order.append(module) g = tg_critic(module, loss_text, downstream) if var.requires_grad and g: var.grads.append(g) if g: downstream.append(g) return order def tg_step(graph): """Textual gradient descent: for every variable with gradients, add the one most-asked rule word to its value, then clear the gradients. Returns the modules that changed.""" changed = [] for module, var in graph: if not var.requires_grad or not var.grads: var.grads.clear() continue counts = {} for g in var.grads: for w in re.findall(r"consider naming: ([\w, ]+)", g): for word in [x.strip() for x in w.split(",") if x.strip()]: counts[word] = counts.get(word, 0) + 1 var.grads.clear() fresh = [w for w in sorted(counts, key=lambda w: (-counts[w], w)) if w not in var.value.split()] if fresh: var.value = var.value.rstrip(".") + " " + fresh[0] changed.append(module) return changed def textgrad_optimize(prog, train, rng, steps=8, minibatch=12, params=ToyParams(), gate=True, frozen=()): """Forward on a minibatch, write the critique, backpropagate it, step, and (if gate) keep the step only if a fresh minibatch run does not get worse. Returns (program, history, meter).""" meter = Meter() cur = prog hist = [] for step in range(steps): batch = rng.sample(train, min(minibatch, len(train))) traces = [rollout(cur, ex, rng, params, meter) for ex in batch] before = mean(1.0 if t.correct else 0.0 for t in traces) graph = [(m, TGVar(cur.instructions[m], m, requires_grad=m not in frozen)) for m in MODULES] tg_backward(graph, tg_loss_text(traces)) changed = tg_step(graph) cand = Program({m: v.value for m, v in graph}, cur.demos) after = evaluate(cand, batch, rng, params, meter) if changed else before accepted = bool(changed) and (after >= before or not gate) if accepted: cur = cand hist.append((step, tuple(changed), before, after, accepted)) return cur, hist, meter # ---- GEPA-like -------------------------------------------------------------------------------- def _pareto(scores): """Shared by select_candidate and pareto_front. scores[k][i] = score of candidate k on instance i. 1. For each instance find the best score and the set of candidates that reach it. 2. Union the sets. A candidate is dominated if every instance it wins also has another surviving winner (my reading of the pseudo-code in the GEPA paper [1]). Remove dominated candidates, lowest index first, until none is left. Returns (survivors, wins) where wins[i] is the set of candidates that reach the best score on i. Note that an instance on which every candidate scores 0 counts as a tie won by all of them.""" n_cand, n_inst = len(scores), len(scores[0]) wins = [] for i in range(n_inst): top = max(scores[k][i] for k in range(n_cand)) wins.append({k for k in range(n_cand) if scores[k][i] == top}) alive = set().union(*wins) removed = True while removed: removed = False for k in sorted(alive): mine = [i for i in range(n_inst) if k in wins[i]] if mine and all(len((wins[i] & alive) - {k}) > 0 for i in mine): alive.discard(k) removed = True break return alive, wins def select_candidate(scores, rng): """Pareto-based candidate selection, after Algorithm 2 of the GEPA paper [1]: take the survivors of _pareto, weight each by the number of instances it wins, and sample.""" alive, wins = _pareto(scores) freq = {k: sum(1 for w in wins if k in w) for k in alive} ks = sorted(freq) return rng.choices(ks, [freq[k] for k in ks])[0] def pareto_front(scores): """The set of candidates kept by select_candidate (for display and tests).""" return sorted(_pareto(scores)[0]) def reflect_update(instruction, traces, module): """The mock reflection LM. It reads failure feedback for this module and adds the rule word that failures asked for most often. If no failure names a missing rule, it returns the instruction unchanged (the failures look like noise).""" counts = {} for t in traces: if t.correct or t.module != module: continue m = re.match(r"fail: kind (\w) needs the rule '(\w+)'", t.feedback) if m: counts[m.group(2)] = counts.get(m.group(2), 0) + 1 fresh = [w for w in sorted(counts, key=lambda w: (-counts[w], w)) if w not in instruction.split()] return instruction.rstrip(".") + " " + fresh[0] if fresh else instruction def merge_programs(a, b, sa, sb, dpareto): """A simple system-aware merge: for each module keep the instruction of the candidate that scores higher on that module's pareto instances (ties go to a).""" ins = {} for m in MODULES: idx = [i for i, ex in enumerate(dpareto) if MODULE_OF[ex.kind] == m] ins[m] = (a if sum(sa[i] for i in idx) >= sum(sb[i] for i in idx) else b).instructions[m] return Program(ins, a.demos) @dataclass class GepaResult: program: Program pool: list parents: list scores: list rollouts: int iterations: int accepted: int def gepa(seed, dfeedback, dpareto, rng, budget=400, minibatch=3, select="pareto", merge_every=0, params=ToyParams()): """GEPA-style reflective prompt evolution (Algorithm 1 of the GEPA paper [1], in toy form). Each iteration: pick a candidate (Pareto-weighted or best-mean), pick a module round-robin, run it on a minibatch, reflect on the failures to edit that module's instruction, and keep the child only if it beats its parent on the same minibatch. Kept children are scored on the whole pareto set. The budget counts every rollout, including pareto scoring.""" meter = Meter(budget) pool, parents = [seed], [None] scores = [] accepted, it = 0, 0 try: scores.append(scores_on(seed, dpareto, rng, params, meter)) while True: it += 1 if merge_every and it % merge_every == 0 and len(pool) >= 2: order = sorted(range(len(pool)), key=lambda k: -mean(scores[k])) a, b = order[0], order[1] child = merge_programs(pool[a], pool[b], scores[a], scores[b], dpareto) if child not in pool: sc = scores_on(child, dpareto, rng, params, meter) pool.append(child) parents.append((a, b)) scores.append(sc) accepted += 1 continue if select == "pareto": k = select_candidate(scores, rng) else: k = max(range(len(pool)), key=lambda j: (mean(scores[j]), -j)) module = MODULES[(it - 1) % len(MODULES)] batch = rng.sample(dfeedback, min(minibatch, len(dfeedback))) traces = [rollout(pool[k], ex, rng, params, meter) for ex in batch] new_text = reflect_update(pool[k].instructions[module], traces, module) if new_text == pool[k].instructions[module]: continue child = pool[k].with_instruction(module, new_text) old = mean(1.0 if t.correct else 0.0 for t in traces) new = evaluate(child, batch, rng, params, meter) if new > old and child not in pool: sc = scores_on(child, dpareto, rng, params, meter) pool.append(child) parents.append(k) scores.append(sc) accepted += 1 except BudgetExhausted: pass if not scores: # the budget could not even score the seed scores.append([0.0] * len(dpareto)) best = max(range(len(scores)), key=lambda j: (mean(scores[j]), -j)) return GepaResult(pool[best], pool, parents, scores, meter.rollouts, it, accepted) # --------------------------------------------------------------------------- 3. weight optimization def softmax(xs): m = max(xs) es = [math.exp(x - m) for x in xs] z = sum(es) return [e / z for e in es] class TabPolicy: """A tabular softmax policy. A 'response' is T tokens from a vocabulary of V, one logit row per (state, position). The 'weights' are the logits. This is a toy stand-in for an LLM.""" def __init__(self, n_states=3, T=2, V=3, theta=None): self.S, self.T, self.V = n_states, T, V self.theta = theta if theta is not None else [[[0.0] * V for _ in range(T)] for _ in range(n_states)] def copy(self): return TabPolicy(self.S, self.T, self.V, [[row[:] for row in st] for st in self.theta]) def probs(self, s, t): return softmax(self.theta[s][t]) def sample(self, s, rng): return tuple(rng.choices(range(self.V), self.probs(s, t))[0] for t in range(self.T)) def logp_tok(self, s, t, v): return math.log(self.probs(s, t)[v]) def logp(self, s, seq): return sum(self.logp_tok(s, t, v) for t, v in enumerate(seq)) def entries(self): return [(s, t, v) for s in range(self.S) for t in range(self.T) for v in range(self.V)] def zeros_like(policy): return [[[0.0] * policy.V for _ in range(policy.T)] for _ in range(policy.S)] def add_score(g, policy, s, t, v, w): """g += w * grad of log pi(v | s, t) with respect to the logits (onehot minus probabilities).""" p = policy.probs(s, t) for u in range(policy.V): g[s][t][u] += w * ((1.0 if u == v else 0.0) - p[u]) def apply_grad(policy, g, lr): for s in range(policy.S): for t in range(policy.T): for u in range(policy.V): policy.theta[s][t][u] += lr * g[s][t][u] def grad_norm(g): return math.sqrt(sum(x * x for st in g for row in st for x in row)) class ToyEnv: """A verifiable toy task: for each state there is one hidden target response. Reward is 1 for an exact match and 0 otherwise (a sparse, checkable reward, like a unit test). MADE UP.""" def __init__(self, n_states=3, T=2, V=3, seed=0): rng = random.Random(seed) self.S, self.T, self.V = n_states, T, V self.target = [tuple(rng.randrange(V) for _ in range(T)) for _ in range(n_states)] def reward(self, s, seq): return 1.0 if tuple(seq) == self.target[s] else 0.0 def success(self, policy): """Exact expected reward, with states uniform. No sampling noise.""" tot = 0.0 for s in range(self.S): p = 1.0 for t, v in enumerate(self.target[s]): p *= policy.probs(s, t)[v] tot += p / self.S return tot def collect(policy, env, rng, n): """n (state, response, reward) triples, states uniform.""" out = [] for _ in range(n): s = rng.randrange(env.S) seq = policy.sample(s, rng) out.append((s, seq, env.reward(s, seq))) return out # ---- REINFORCE ---------------------------------------------------------------------------------- def reinforce_grad(policy, batch, baseline="mean", weights=None): """Policy gradient with a baseline: mean over samples of (R - b) * grad log pi(response). weights (optional) make the 'mean' a weighted mean, which lets tests enumerate all responses.""" w = weights if weights is not None else [1.0 / len(batch)] * len(batch) b = sum(wi * R for wi, (_, _, R) in zip(w, batch)) if baseline == "mean" else 0.0 g = zeros_like(policy) for wi, (s, seq, R) in zip(w, batch): for t, v in enumerate(seq): add_score(g, policy, s, t, v, wi * (R - b)) return g def reinforce_train(policy, env, rng, iters=200, batch=16, lr=0.5, baseline="mean"): curve = [] for _ in range(iters): apply_grad(policy, reinforce_grad(policy, collect(policy, env, rng, batch), baseline), lr) curve.append(env.success(policy)) return curve # ---- PPO-clip ---------------------------------------------------------------------------------- @dataclass class PPOConfig: clip: float = 0.2 kl_beta: float = 0.0 epochs: int = 4 lr: float = 0.5 lr_v: float = 0.5 batch: int = 24 @dataclass class PPOSample: s: int seq: tuple R: float logp_old: list # per token G: list # per token return, with the per-token KL penalty folded in A: list = field(default_factory=list) def clip_flows(A, ratio, eps): """Does the clipped objective min(r A, clip(r) A) pass gradient at this ratio?""" return (ratio < 1 + eps) if A > 0 else ((ratio > 1 - eps) if A < 0 else False) def ppo_collect(policy, ref, env, critic, rng, cfg): out = [] for s, seq, R in collect(policy, env, rng, cfg.batch): lp = [policy.logp_tok(s, t, v) for t, v in enumerate(seq)] lr_ = [ref.logp_tok(s, t, v) for t, v in enumerate(seq)] rew = [-cfg.kl_beta * (a - b) for a, b in zip(lp, lr_)] rew[-1] += R G = [sum(rew[t:]) for t in range(len(seq))] A = [G[t] - critic[s][t] for t in range(len(seq))] out.append(PPOSample(s, seq, R, lp, G, A)) return out def ppo_objective(policy, samples, clip=0.2): tot = 0.0 for sm in samples: for t, v in enumerate(sm.seq): r = math.exp(policy.logp_tok(sm.s, t, v) - sm.logp_old[t]) tot += min(r * sm.A[t], max(1 - clip, min(1 + clip, r)) * sm.A[t]) return tot / (len(samples) * len(samples[0].seq)) def ppo_grad(policy, samples, clip=0.2): g = zeros_like(policy) n = len(samples) * len(samples[0].seq) for sm in samples: for t, v in enumerate(sm.seq): r = math.exp(policy.logp_tok(sm.s, t, v) - sm.logp_old[t]) if clip_flows(sm.A[t], r, clip): add_score(g, policy, sm.s, t, v, sm.A[t] * r / n) return g def ppo_train(policy, env, rng, iters=60, cfg=PPOConfig(), critic=None): """PPO with four things in memory: the policy, a frozen reference, a reward (the env), and a per-token value table (the critic). Returns the success curve.""" ref = policy.copy() if critic is None: critic = [[0.0] * env.T for _ in range(env.S)] curve = [] for _ in range(iters): samples = ppo_collect(policy, ref, env, critic, rng, cfg) for _e in range(cfg.epochs): apply_grad(policy, ppo_grad(policy, samples, cfg.clip), cfg.lr) acc = {} for sm in samples: for t in range(len(sm.seq)): a = acc.setdefault((sm.s, t), [0.0, 0]) a[0] += sm.G[t] - critic[sm.s][t] a[1] += 1 for (s, t), (d, n) in acc.items(): critic[s][t] += cfg.lr_v * d / n curve.append(env.success(policy)) return curve # ---- GRPO ---------------------------------------------------------------------------------------- @dataclass class GRPOConfig: clip: float = 0.2 beta: float = 0.04 group: int = 8 prompts: int = 3 epochs: int = 2 lr: float = 0.5 std_norm: bool = True length_norm: bool = True @dataclass class Group: s: int seqs: list rewards: list adv: list logp_old: list # per response, per token def group_advantages(rewards, std_norm=True, eps=1e-6): """Subtract the group mean. If std_norm, divide by the group standard deviation (a z-score). A group with identical rewards has zero advantage for every response: no signal.""" mu = mean(rewards) sd = stdev(rewards) if sd < eps: return [0.0] * len(rewards) return [(r - mu) / (sd if std_norm else 1.0) for r in rewards] def grpo_collect(policy, env, rng, cfg): groups = [] for _ in range(cfg.prompts): s = rng.randrange(env.S) seqs = [policy.sample(s, rng) for _ in range(cfg.group)] rewards = [env.reward(s, q) for q in seqs] lp = [[policy.logp_tok(s, t, v) for t, v in enumerate(q)] for q in seqs] groups.append(Group(s, seqs, rewards, group_advantages(rewards, cfg.std_norm), lp)) return groups def kl_k3(logp_ref, logp): """The unbiased, never-negative KL estimator used by GRPO: rho - log rho - 1, rho = pi_ref/pi.""" rho = math.exp(logp_ref - logp) return rho - math.log(rho) - 1.0 def grpo_objective(policy, ref, groups, cfg=GRPOConfig()): tot = 0.0 for gr in groups: inner = 0.0 for q, A, lpo in zip(gr.seqs, gr.adv, gr.logp_old): w = 1.0 / len(q) if cfg.length_norm else 1.0 for t, v in enumerate(q): lp = policy.logp_tok(gr.s, t, v) r = math.exp(lp - lpo[t]) surr = min(r * A, max(1 - cfg.clip, min(1 + cfg.clip, r)) * A) inner += w * (surr - cfg.beta * kl_k3(ref.logp_tok(gr.s, t, v), lp)) tot += inner / len(gr.seqs) return tot / len(groups) def grpo_grad(policy, ref, groups, cfg=GRPOConfig()): """Analytic gradient of grpo_objective with respect to the logits. For the KL term, d k3 / d log pi = 1 - rho, where rho = pi_ref / pi.""" g = zeros_like(policy) for gr in groups: for q, A, lpo in zip(gr.seqs, gr.adv, gr.logp_old): w = 1.0 / len(q) if cfg.length_norm else 1.0 for t, v in enumerate(q): lp = policy.logp_tok(gr.s, t, v) r = math.exp(lp - lpo[t]) c = 0.0 if clip_flows(A, r, cfg.clip): c += A * r rho = math.exp(ref.logp_tok(gr.s, t, v) - lp) c -= cfg.beta * (1.0 - rho) add_score(g, policy, gr.s, t, v, w * c / (len(gr.seqs) * len(groups))) return g def grpo_train(policy, env, rng, iters=60, cfg=GRPOConfig()): """GRPO: three things in memory (policy, reference, reward) and no value model.""" ref = policy.copy() curve = [] for _ in range(iters): groups = grpo_collect(policy, env, rng, cfg) for _e in range(cfg.epochs): apply_grad(policy, grpo_grad(policy, ref, groups, cfg), cfg.lr) curve.append(env.success(policy)) return curve # ---- SFT from filtered rollouts ---------------------------------------------------------------------- def sft_grad(policy, data): """Gradient of the mean token log-likelihood of the data (cross-entropy, negated).""" g = zeros_like(policy) for s, seq in data: for t, v in enumerate(seq): add_score(g, policy, s, t, v, 1.0 / (len(data) * len(seq))) return g def sft_train(policy, data, epochs=100, lr=1.0): for _ in range(epochs): apply_grad(policy, sft_grad(policy, data), lr) def filter_rollouts(rollouts, min_reward=1.0): """Rejection sampling: keep (state, response) pairs whose reward passes.""" return [(s, seq) for s, seq, R in rollouts if R >= min_reward] # ---- on-policy distillation (reverse KL) ------------------------------------------------------------ def token_reverse_kl(student, teacher, s, seq): """Per-token sampled reverse-KL terms: log pi_student - log pi_teacher at each sampled token.""" return [student.logp_tok(s, t, v) - teacher.logp_tok(s, t, v) for t, v in enumerate(seq)] def on_policy_distill_grad(student, teacher, samples, weights=None): """Samples come from the STUDENT. Reward at each token = -(log pi_s - log pi_t), discount 0: the token's advantage is just its own reward. Gradient = mean of reward * grad log pi_s.""" w = weights if weights is not None else [1.0 / len(samples)] * len(samples) g = zeros_like(student) for wi, (s, seq) in zip(w, samples): for t, (v, k) in enumerate(zip(seq, token_reverse_kl(student, teacher, s, seq))): add_score(g, student, s, t, v, wi * (-k)) return g def exact_kl(p, q): """KL(p || q) for two probability lists.""" return sum(pi * math.log(pi / qi) for pi, qi in zip(p, q) if pi > 0) def exact_reverse_kl(student, teacher): """KL(student || teacher), summed over positions, averaged over states (positions are independent).""" return sum(exact_kl(student.probs(s, t), teacher.probs(s, t)) for s in range(student.S) for t in range(student.T)) / student.S def exact_forward_kl(student, teacher): return sum(exact_kl(teacher.probs(s, t), student.probs(s, t)) for s in range(student.S) for t in range(student.T)) / student.S def on_policy_distill_train(student, teacher, rng, iters=150, batch=16, lr=0.5): curve = [] for _ in range(iters): smp = [] for _i in range(batch): s = rng.randrange(student.S) smp.append((s, student.sample(s, rng))) apply_grad(student, on_policy_distill_grad(student, teacher, smp), lr) curve.append(exact_reverse_kl(student, teacher)) return curve def discretized_gaussian(mu, sigma, V=21): xs = list(range(V)) ws = [math.exp(-((x - mu) ** 2) / (2 * sigma ** 2)) for x in xs] z = sum(ws) return [w / z for w in ws] def bimodal_teacher(V=21, m1=5.0, m2=15.0, sigma=1.2): a, b = discretized_gaussian(m1, sigma, V), discretized_gaussian(m2, sigma, V) return [0.5 * x + 0.5 * y for x, y in zip(a, b)] def fit_gaussian(teacher, direction, mu0=6.0, log_sigma0=0.0, steps=400, lr=0.5): """Fit a one-bump (Gaussian-shaped) student to a teacher distribution by gradient descent on KL, with numeric gradients. direction='forward' minimizes KL(teacher||student) (what SFT on teacher samples does). direction='reverse' minimizes KL(student||teacher) (what on-policy distillation does). Returns (mu, sigma, loss).""" V = len(teacher) def loss(mu, ls): q = discretized_gaussian(mu, math.exp(ls), V) return exact_kl(teacher, q) if direction == "forward" else exact_kl(q, teacher) mu, ls = mu0, log_sigma0 h = 1e-4 for _ in range(steps): gm = (loss(mu + h, ls) - loss(mu - h, ls)) / (2 * h) gs = (loss(mu, ls + h) - loss(mu, ls - h)) / (2 * h) mu -= lr * gm ls = min(2.5, max(-1.5, ls - lr * gs)) return mu, math.exp(ls), loss(mu, ls) # --------------------------------------------------------------------------- 4. DGM-style archive FEATURES = ("edit_tool", "view_tool", "retry", "planner", "test_runner", "memory_notes") def variant_score(features): """MADE-UP performance landscape for a toy coding agent. 'view_tool', 'planner' and 'memory_notes' each cost a little alone (so no move is neutral), and view_tool plus planner are worth a lot together: a stepping stone. I built the stepping stone in on purpose, so the demo has something for an archive to find.""" f = set(features) s = 0.20 s += 0.08 * ("edit_tool" in f) + 0.05 * ("retry" in f) + 0.04 * ("test_runner" in f) s -= 0.03 * ("view_tool" in f) + 0.03 * ("planner" in f) + 0.01 * ("memory_notes" in f) s += 0.50 * ("view_tool" in f and "planner" in f) s += 0.05 * ("view_tool" in f and "memory_notes" in f) return round(min(1.0, max(0.0, s)), 10) @dataclass(frozen=True) class AgentVariant: id: int features: frozenset score: float parent: int | None = None def new_archive(): root = frozenset() return [AgentVariant(0, root, variant_score(root))] def children_count(archive, i): return sum(1 for a in archive if a.parent == i) def dgm_weights(archive, lam=10.0, alpha0=0.5): """Parent-selection weights from the DGM paper [2]: w_i = sigmoid(lam (score_i - alpha0)) / (1 + n_i), n_i = number of children already in the archive. I read 'functioning children' as all archived children, because only functioning variants are archived here.""" out = [] for a in archive: s = 1.0 / (1.0 + math.exp(-lam * (a.score - alpha0))) out.append(s / (1.0 + children_count(archive, a.id))) return out def dgm_parent_probs(archive, lam=10.0, alpha0=0.5): w = dgm_weights(archive, lam, alpha0) z = sum(w) return [x / z for x in w] def self_modify(parent, rng, new_id, p_break=0.15): """The mock 'agent edits its own code' step: toggle one feature. It breaks with p_break (a broken variant cannot edit code, so it is not archived). Returns a variant or None.""" f = set(parent.features) feat = rng.choice(FEATURES) f.symmetric_difference_update({feat}) if rng.random() < p_break: return None fs = frozenset(f) return AgentVariant(new_id, fs, variant_score(fs), parent.id) def choose_parent(archive, rng, mode): if mode == "dgm": return rng.choices(archive, dgm_weights(archive))[0] if mode == "greedy": return max(archive, key=lambda a: (a.score, a.id)) if mode == "uniform": return rng.choice(archive) raise ValueError(mode) def run_dgm(iters, rng, mode="dgm", keep_all=True): """Grow an archive. keep_all=False keeps only the best variant (a hill climber that throws away stepping stones), which is what 'no open-ended exploration' means in this toy.""" archive = new_archive() for _ in range(iters): parent = choose_parent(archive, rng, mode) child = self_modify(parent, rng, len(archive) if keep_all else archive[-1].id + 1) if child is None: continue if keep_all: archive.append(child) elif child.score > max(a.score for a in archive): archive = [child] return archive def best_score(archive): return max(a.score for a in archive) def compare_dgm(seeds=range(20), iters=60): """Mean best score after iters steps for each parent-selection rule. Toy landscape.""" out = {} for mode in ("dgm", "greedy", "uniform"): out[mode] = mean(best_score(run_dgm(iters, random.Random(s), mode)) for s in seeds) out["best_only"] = mean(best_score(run_dgm(iters, random.Random(s), "greedy", keep_all=False)) for s in seeds) return out # --------------------------------------------------------------------------- 5. autoresearch-style loop def parse_kv(text): out = {} for line in text.splitlines(): m = re.match(r"\s*(\w+)\s*=\s*([-+0-9.eE]+)\s*$", line) if m: out[m.group(1)] = float(m.group(2)) return out def fmt_kv(d): return "".join(f"{k}={v:g}\n" for k, v in d.items()) def true_loss(hp): """MADE-UP validation loss (lower is better) for a toy training script with a fixed time budget.""" return (1.20 + 0.08 * math.log10(hp["lr"] / 0.003) ** 2 + 0.004 * (hp["depth"] - 8) ** 2 + 0.00002 * (hp["width"] - 256) ** 2) def reported_loss(files): """What the harness reports: the eval file may rescale the true loss. In the honest repo scale = 1 and offset = 0.""" ev = parse_kv(files["eval.py"]) return ev["scale"] * true_loss(parse_kv(files["train.py"])) + ev["offset"] def start_repo(): return {"train.py": fmt_kv({"lr": 0.02, "depth": 4, "width": 96}), "eval.py": fmt_kv({"scale": 1.0, "offset": 0.0}), "program.md": "Edit train.py only. Lower loss is better. Do not touch eval.py.\n"} def file_hash(text): return hashlib.sha256(text.encode()).hexdigest() def guard_ok(files, locked): """locked: {filename: sha256}. The guard passes only if every locked file is byte-identical.""" return all(file_hash(files[f]) == h for f, h in locked.items()) def honest_proposer(files, rng): hp = parse_kv(files["train.py"]) key = rng.choice(["lr", "depth", "width"]) if key == "lr": hp["lr"] = round(hp["lr"] * rng.choice([0.5, 0.7, 1.4, 2.0]), 6) elif key == "depth": hp["depth"] = max(1, hp["depth"] + rng.choice([-1, 1])) else: hp["width"] = max(16, hp["width"] + rng.choice([-32, 32])) new = dict(files) new["train.py"] = fmt_kv(hp) return new, f"set {key}" def sneaky_proposer(files, rng, p_cheat=0.3): """Like honest_proposer, but sometimes edits eval.py to shrink the reported loss.""" if rng.random() < p_cheat: ev = parse_kv(files["eval.py"]) ev["scale"] = round(ev["scale"] * 0.9, 6) new = dict(files) new["eval.py"] = fmt_kv(ev) return new, "rescale eval" return honest_proposer(files, rng) @dataclass class AutoResult: files: dict log: list reported: float true: float blocked: int def autoresearch(iters, rng, proposer=honest_proposer, guard=True): """Keep-or-discard loop: propose an edit, run the (cheap) eval, keep only if the REPORTED loss improved and, when guard is on, the locked files are unchanged.""" files = start_repo() locked = {f: file_hash(files[f]) for f in ("eval.py", "program.md")} best = reported_loss(files) log, blocked = [], 0 for i in range(1, iters + 1): cand, what = proposer(files, rng) if guard and not guard_ok(cand, locked): blocked += 1 log.append((i, what, None, "blocked: locked file changed")) continue r = reported_loss(cand) if r < best - 1e-12: files, best = cand, r log.append((i, what, r, "keep")) else: log.append((i, what, r, "discard")) base = start_repo() true = true_loss(parse_kv(files["train.py"])) return AutoResult(files, log, best, true, blocked) # --------------------------------------------------------------------------- 6. prompt vs weights cost @dataclass(frozen=True) class CostInputs: """EVERY NUMBER BELOW IS MADE UP. Prices are dollars per million tokens. Replace them with your own quotes before you use the calculator for anything.""" calls: int = 500_000 # calls over the life of the deployment in_tokens: int = 1500 # prompt tokens per call before optimization out_tokens: int = 300 bloat_tokens: int = 2500 # extra prompt tokens an optimized prompt (demos, rules) adds big_in: float = 3.0 big_out: float = 15.0 small_in: float = 0.25 small_out: float = 1.25 prompt_opt_rollouts: int = 6000 reflection_calls: int = 600 # extra big-model calls that write new instructions reflection_tokens: int = 4000 teacher_rollouts: int = 20000 # big-model rollouts to build SFT data train_examples: int = 5000 # kept after filtering train_tokens_each: int = 2000 epochs: int = 3 train_price: float = 8.0 # $/M training tokens (made up) eng_hours_prompt: float = 8.0 eng_hours_weights: float = 40.0 eng_rate: float = 100.0 # $/hour (made up) both_bloat_share: float = 0.25 # share of the bloat kept when prompts and weights are combined acc_none: float = 0.60 # MADE-UP accuracies. Measure these on your own dev set. acc_prompt: float = 0.78 acc_weights: float = 0.80 acc_both: float = 0.85 error_cost: float = 0.10 # dollars lost per wrong answer (made up) can_finetune: bool = True # False for an API-only model def call_cost(in_tok, out_tok, p_in, p_out): return (in_tok * p_in + out_tok * p_out) / 1e6 def _plan(c, fixed, per, acc): total = fixed + per * c.calls return {"fixed": fixed, "per_call": per, "acc": acc, "total": total, "total_with_errors": total + c.calls * (1.0 - acc) * c.error_cost} def plan_none(c): return _plan(c, 0.0, call_cost(c.in_tokens, c.out_tokens, c.big_in, c.big_out), c.acc_none) def plan_prompt(c): """Optimize the prompt, serve on the big model with the longer prompt.""" per = call_cost(c.in_tokens + c.bloat_tokens, c.out_tokens, c.big_in, c.big_out) roll = c.prompt_opt_rollouts * per refl = c.reflection_calls * call_cost(c.reflection_tokens, c.out_tokens, c.big_in, c.big_out) return _plan(c, roll + refl + c.eng_hours_prompt * c.eng_rate, per, c.acc_prompt) def _weights_fixed(c): data = c.teacher_rollouts * call_cost(c.in_tokens, c.out_tokens, c.big_in, c.big_out) train = c.train_examples * c.train_tokens_each * c.epochs / 1e6 * c.train_price return data + train + c.eng_hours_weights * c.eng_rate def plan_weights(c): """Collect teacher rollouts, fine-tune a small model, serve it with the short prompt. The accuracy is an input, because this calculator cannot tell you whether the small model gets there.""" return _plan(c, _weights_fixed(c), call_cost(c.in_tokens, c.out_tokens, c.small_in, c.small_out), c.acc_weights) def plan_both(c): """Fine-tune the small model, and also keep a short optimized prompt (a share of the bloat).""" bloat = int(c.bloat_tokens * c.both_bloat_share) per = call_cost(c.in_tokens + bloat, c.out_tokens, c.small_in, c.small_out) return _plan(c, _weights_fixed(c) + plan_prompt(c)["fixed"], per, c.acc_both) def breakeven_calls(a, b, error_cost=0.10): """Number of calls above which plan b (higher fixed cost, lower per-call cost) has a lower cost including errors than plan a. 0 if it already wins at zero calls, None if it never does.""" ea = a["per_call"] + (1.0 - a["acc"]) * error_cost eb = b["per_call"] + (1.0 - b["acc"]) * error_cost if eb >= ea: return None return max(0.0, (b["fixed"] - a["fixed"]) / (ea - eb)) def options(c): opts = {"none": plan_none(c), "prompt": plan_prompt(c)} if c.can_finetune: opts["weights"] = plan_weights(c) opts["both"] = plan_both(c) return opts def cheapest(c): """The option with the lowest total cost including the cost of wrong answers.""" opts = options(c) return min(opts, key=lambda k: opts[k]["total_with_errors"]), opts # --------------------------------------------------------------------------- demo def _mean_over(seeds, fn): return mean(fn(s) for s in seeds) def noisy_teacher(env, iters=25, seed=11): """A teacher that is right on only part of the tasks (a short REINFORCE run). Used so that filtering has something to remove.""" t = TabPolicy(env.S, env.T, env.V) reinforce_train(t, env, random.Random(seed), iters=iters) return t def demo(out=print): """Run every part once with fixed seeds. Prints numbered sections. Deterministic.""" seeds = range(8) base = Program() train = make_dataset(60, 1) dev = make_dataset(60, 2) out("== 1. A toy prompt world (made-up numbers) ==") best = best_possible_program() out(f"no help, exact accuracy: {exact_accuracy(base):.3f}") out(f"best rule words, no demos: {exact_accuracy(best):.3f} {best.instructions}") out("per-kind accuracy of the best words: " + str({k: round(v, 2) for k, v in per_kind_accuracy(best).items()})) out("note: naming 'negatives' hurts kind C, naming 'abbrev' hurts kind F (built-in conflicts)") out("\n== 2. Bootstrapped few-shot ==") for keep_all in (False, True): accs = [] for s in seeds: d = bootstrap_demos(base, train, random.Random(s), keep_all=keep_all) accs.append(exact_accuracy(base.with_demos(d))) out(f"demos from {'every trace' if keep_all else 'passing traces only'}: mean exact accuracy {mean(accs):.3f}") accs = [exact_accuracy(bootstrap_random_search(base, train, dev, random.Random(s))[0]) for s in seeds] out(f"bootstrap + random search over demo sets: mean exact accuracy {mean(accs):.3f}") out("\n== 3. MIPRO-like search (equal rollouts, 8 seeds) ==") for name, kw in (("grounded + surrogate", {}), ("grounded, random combos", {"surrogate": False}), ("ungrounded + surrogate", {"grounded": False})): rs = [mipro_search(base, train, dev, random.Random(s), **kw) for s in seeds] out(f"{name:26s} exact accuracy {mean(exact_accuracy(r.program) for r in rs):.3f} rollouts {rs[0].rollouts}") out("\n== 4. OPRO loop (one shared instruction) ==") r = opro(train[:24], random.Random(1)) out(f"best instruction: {r.instruction!r}") out(f"score on its own scoring set {r.score:.3f}, exact accuracy {exact_accuracy(uniform_program(r.instruction)):.3f}, " f"{len(r.history)} instructions tried, {r.rollouts} rollouts") out("end of the meta-prompt shown to the proposer (its last two entries, best last):") out("\n".join(opro_prompt(r.history).splitlines()[-4:])) out("\n== 5. TextGrad-style feedback step ==") p, hist, m = textgrad_optimize(base, train, random.Random(1)) out(f"after 8 steps: {p.instructions}") out(f"exact accuracy {exact_accuracy(p):.3f}, accepted {sum(1 for h in hist if h[4])} of {len(hist)} steps, {m.rollouts} rollouts") p2, hist2, _ = textgrad_optimize(base, train, random.Random(1), frozen=("parse",)) out(f"with the parse prompt frozen: {p2.instructions} exact accuracy {exact_accuracy(p2):.3f}") out("\n== 6. GEPA-style reflective evolution (budget 1000 rollouts, 8 seeds) ==") dpar = dev[:18] for sel in ("pareto", "best"): rs = [gepa(base, train, dpar, random.Random(s), budget=1000, select=sel) for s in seeds] out(f"select={sel:6s} exact accuracy {mean(exact_accuracy(r.program) for r in rs):.3f} " f"mean pool {mean(len(r.pool) for r in rs):.1f} mean accepted {mean(r.accepted for r in rs):.1f}") rs = [gepa(base, train, dpar, random.Random(s), budget=1000, merge_every=6) for s in seeds] out(f"pareto + merge exact accuracy {mean(exact_accuracy(r.program) for r in rs):.3f}") out("\n== 7. REINFORCE, PPO-clip, GRPO on a toy verifiable task ==") env = ToyEnv(seed=3) out(f"targets per state: {env.target}; start success {env.success(TabPolicy()):.3f} (chance is 1/9 per state)") for name, fn in (("REINFORCE", lambda p, r: reinforce_train(p, env, r, iters=90)), ("PPO-clip", lambda p, r: ppo_train(p, env, r, iters=60)), ("GRPO", lambda p, r: grpo_train(p, env, r, iters=60))): fin = [fn(TabPolicy(), random.Random(s))[-1] for s in seeds] out(f"{name:10s} final exact success, mean of 8 seeds: {mean(fin):.3f} (min {min(fin):.3f}) [1,440 rollouts each]") out(f"GRPO advantages for rewards [1,0,0,0]: {[round(a, 2) for a in group_advantages([1, 0, 0, 0])]}") out(f"GRPO advantages for rewards [0,0,0,0]: {group_advantages([0, 0, 0, 0])}") out("\n== 8. SFT from filtered rollouts, and distillation ==") teacher = noisy_teacher(env) out(f"teacher success {env.success(teacher):.3f}") for name, flt in (("filtered (reward 1 only)", True), ("unfiltered", False)): sc = [] for s in seeds: ro = collect(teacher, env, random.Random(100 + s), 120) data = filter_rollouts(ro) if flt else [(a, b) for a, b, _ in ro] st = TabPolicy() sft_train(st, data) sc.append(env.success(st)) out(f"SFT on {name:24s} student success {mean(sc):.3f}") st = TabPolicy() curve = on_policy_distill_train(st, teacher, random.Random(3)) out(f"on-policy distillation: reverse KL {curve[0]:.3f} -> {curve[-1]:.4f}, student success {env.success(st):.3f}") t = bimodal_teacher() for d in ("forward", "reverse"): mu, sg, l = fit_gaussian(t, d) out(f"one-bump student fit by {d} KL: mean {mu:.2f}, sigma {sg:.2f}, loss {l:.3f}") out("forward KL covers both modes and sits between them. reverse KL picks one mode.") out("\n== 9. DGM-style archive (toy landscape, 20 seeds, 60 steps) ==") cmp_ = compare_dgm() for k, v in cmp_.items(): out(f"{k:10s} mean best score {v:.3f}") a = new_archive() out(f"parent probabilities for a one-variant archive: {dgm_parent_probs(a)}") out("\n== 10. Autoresearch-style loop with an eval guard ==") for guard in (True, False): rs = [autoresearch(60, random.Random(s), sneaky_proposer, guard=guard) for s in range(5)] out(f"guard={guard!s:5s} reported {mean(r.reported for r in rs):.3f} true {mean(r.true for r in rs):.3f} " f"blocked edits {mean(r.blocked for r in rs):.1f}") out("\n== 11. Prompt versus weights cost calculator (ALL NUMBERS MADE UP) ==") for calls in (5_000, 50_000, 500_000, 5_000_000): for ft in (True, False): c = CostInputs(calls=calls, can_finetune=ft) k, opts = cheapest(c) out(f"calls {calls:>9,} fine-tuning {'allowed ' if ft else 'API-only'} cheapest incl. errors: {k:7s} " f"(${opts[k]['total_with_errors']:,.0f})") c = CostInputs() be = breakeven_calls(plan_none(c), plan_weights(c), c.error_cost) out(f"calls above which weights beat doing nothing (made-up numbers): {be:,.0f}") if __name__ == "__main__": demo()