TwinEdgeAI

↑ Course home

← Previous: Week 5a, Multi-Agent Systems  ·  Next: Week 6a, Terminal-Bench →

Engineering AI Agents · Week 5b

Optimization

When it pays to optimize an agent, and how: changing what the model reads (few-shot demos, instructions, textual feedback), changing what the model is (finetuning on rollouts, reinforcement learning, distillation), letting an agent optimize the system itself, and choosing between prompts and weights.

Author: Grok Bot Version: 1.0, October 29, 2026 Study time: about 15 hours

CreditThis coursework is adapted from, and inspired by, Stanford CS329Z: Engineering AI Agents (Diyi Yang, Michael Ryan, John Yang), Week 5, "Optimization" (Wednesday, October 21). The topic list and the readings come from the course's public schedule [31]. I checked the topic order against the public lecture slides [32]. I used the slides for coverage only. The notes, lab, quiz and assignment are original writing, and the lab uses its own examples and code. This pack does not reproduce the slides. It is not an official Stanford document. I added extra sources of my own. Each is marked "added" in the references.

1. Week overview and learning objectives

You now know how to build an agent, give it tools and memory, and put several of them in a team. This week asks a different question: how do you make the system better in a measurable way? The lecture's framing is that three pillars hold up an agent system (data, evals, and the system itself), and that optimization is the "last mile" of adaptation that uses the first two to improve the third [32].

The week has five parts. First, when optimization pays off, and what you must have before you start. Second, context optimization: change what the model reads, through demonstrations, instructions and textual feedback. The methods are bootstrapped few-shot, OPRO, MIPROv2, TextGrad and GEPA. Third, weight optimization: change the model itself, by supervised finetuning on agent rollouts, by reinforcement learning with PPO and GRPO, and by on-policy distillation. Fourth, agents as optimizers: the Darwin Gödel Machine, which edits its own code, and the autoresearch loop, which edits a training script. Fifth, how to choose between prompts and weights, and when to use both. The sixth part is a lab where you build a toy version of every method.

The lab uses the standard library only. Every model is a mock, and every task is a toy. The mocks follow rules that I wrote, so the lab shows how each method works. It says nothing about how well a method does on a real model. The numbers in the cost calculator are made up. I flag this again where it matters.

Learning objectives

After this week you should be able to:

1. Say what you need before you optimize (data and an eval), and list the benefits that optimization can bring. [2.1]
2. Explain why you split data into a part that guides the search and a part that judges it, and how a rollout budget is counted. [2.2]
3. Describe bootstrapped few-shot with a metric filter, and say why demo order and demo quality matter. [2.3]
4. Compare APE, OPRO and MIPROv2 in what they propose, how they score, and how they search. [2.4]
5. Explain how TextGrad and GEPA use text feedback, and how GEPA's Pareto selection differs from keeping the best mean score. [2.5]
6. Describe SFT on filtered teacher rollouts and name what can go wrong with it. [2.6]
7. Write the clipped PPO objective, the GRPO group advantage and the GRPO KL term, and say which models each method keeps in memory. [2.7]
8. Explain on-policy distillation as per-token reverse KL, and why it differs from SFT on teacher text and from RL. [2.8]
9. Explain how the Darwin Gödel Machine picks parents, and why an autoresearch loop needs a locked eval. [2.9]
10. Use a simple decision guide to choose between context optimization, weight optimization and both, and say what evidence the choice needs. [2.10]
11. Do the small calculations of section 2.11 by hand: group advantages, a clipped ratio, and a break-even number of calls. [2.11]
12. Build and test the methods in the lab: bootstrap, a MIPRO-like search, OPRO, a TextGrad-style step, GEPA-style evolution with Pareto selection, REINFORCE, PPO, GRPO, SFT on filtered rollouts, reverse-KL distillation, a DGM-style archive, an eval-guarded autoresearch loop, and a cost calculator. [Lab]

How to use this pack

Time	What to do
3.5 hours	Read the notes (section 2). Skim the figures first.
4 hours	Read the three required papers with the reading guide questions. Skim the optional ones that match your project.
4 hours	Do the lab and at least five exercises.
1.5 hours	Discussion and the quiz.
2 hours	The graded assignment.

2. Lecture notes

2.1 Why optimize, and when it pays

Figure 1 is the map for the week. The lecture starts from three pillars of an agent system: data (examples, traces, feedback), evals (a metric that says what good means), and the system (prompts, tools, model, code). Optimization is the "last mile" of adaptation. It takes the data and the eval and uses them to improve the system [32]. If you lack the first two, there is nothing to optimize against.

The lecture gives conditions under which optimization is worth the work [32]. I restate them in my own words:

• You have data and evals for your domain. Without a trusted metric you cannot tell a gain from noise, or from gaming.
• Examples and feedback express a requirement that is hard to write down. A rule such as "use the house style for dates" is easier to show with ten examples than to describe in a prompt.
• Evals convey intent. A good metric tells the optimizer what you care about, so that it does not have to guess.
• The system has many complex parts. Hand-tuning four prompts that interact is slow. A search can do it.

The lecture lists five benefits: reliability, adaptation to a domain, lower cost and latency, learning from experience, and adaptability when models or needs change [32]. The first one deserves a note. Reliability is often the real prize. An agent that is right 90% of the time and wrong in a random way can be less useful than one that is right 85% of the time and wrong in a predictable way, because you can plan around the second one. That remark is my own, not the lecture's.

2.2 Optimization as search: objective, budget, held-out data

Every method in this week has the same skeleton, and it helps to name the parts before you meet the methods. This framing is mine. The lecture does not use these exact words, but each method below fits it.

1. A candidate space. The things you may change: the text of an instruction, a set of demos, the weights of a model, the code of an agent.
2. An objective. A score on examples, from a metric. It may be exact match, a unit test, a rubric grade by another model, or a measured cost.
3. A proposer. Something that suggests new candidates. It can be random, a language model that reads feedback, or a gradient.
4. A selection rule. How you decide what to keep, and which candidate to extend next.
5. A budget. The number of rollouts you can afford. A rollout is one run of your program on one example. It costs tokens, time and sometimes money. The GEPA paper counts its budget this way, and compares methods at equal or at very unequal rollout counts [1].

Hold some data out

If you pick the best candidate by its score on the same examples that you used to write it, the score is too high. Three of this week's papers handle this in different ways:

• GEPA splits its training data into one part for reflection (a small feedback set, where minibatches of 3 are drawn) and one part for Pareto scoring and selection. It returns the candidate with the best average score on the second part [1].
• MIPROv2 scores combinations on small minibatches because that is cheap, and every so often runs a full evaluation on a larger validation set to check the best ones [2].
• The DSPy documentation says the same in practice: the training set is used for reflective updates, and the validation set is used to score and select [10].

Noise is the second reason to be careful. Here is my own arithmetic. Say a candidate's true accuracy is 0.50 and you score it on 24 examples. The standard error of the score is √(0.5 × 0.5 / 24), about 0.10. If you try 40 candidates and keep the best on those 24 examples, the best one will look good partly by luck. OPRO's famous result has this flavor. The paper's Table 4 lists, for each pair of scorer and optimizer, the instruction with the highest test accuracy among those found, so the headline instruction was chosen among several by looking at test data [5]. This does not make the finding false. It means the headline number is a best-of-several, and that you should hold out a fresh set before you report a result.

2.3 Few-shot demos: ordering, rejection sampling, bootstrapping

The oldest and cheapest form of context optimization is to choose the demonstrations in the prompt. The lecture covers three ideas [32].

Order matters

Lu et al. showed that the order of few-shot examples can make the difference between near state-of-the-art and random-guess accuracy on a task. The effect is present across model sizes, and a good order for one model does not transfer to another [13]. Their remedy is to have the model generate a small artificial probing set, and to pick orders by an entropy statistic on that set. They report a 13% relative improvement for GPT-family models across eleven text classification tasks [13]. I read only the abstract of this paper. The lesson for an agent builder is that a demo set is a hyperparameter, so you should evaluate it and not trust it.

Rejection sampling: keep the traces that pass

STaR (the Self-Taught Reasoner) has a simple loop: generate a rationale for many questions, prompted with a few examples. If the final answer is wrong, try again with the correct answer as a hint. Fine-tune on all the rationales that ended in a correct answer, and repeat [12]. The idea of keeping only the traces that pass the metric is the core of rejection sampling. It returns in section 2.6 for weights. Here it is used to build demonstrations.

DSPy: bootstrap, then search

DSPy treats a language-model pipeline as a program with modules, and its compiler (the optimizer) can build demonstrations for each module [9]. The bootstrapping recipe the lecture shows is: run your program on training examples, keep the traces whose final answer passes the metric, and use the intermediate steps of those traces as demos for each module. Then search over sets of demos (random search is the simple version) and score each set on a dev set [32] [9]. The paper reports that a few lines of DSPy let GPT-3.5 and llama2-13b-chat bootstrap pipelines that beat standard few-shot prompting, and pipelines with expert-written demos in several cases. These are the authors' results on their own tasks [9]. I read the abstract only.

Notice what the metric filter does. A wrong trace used as a demo teaches the wrong thing. The lab shows this in a toy way: demos built from every trace (right and wrong) leave the mock model worse than no demos at all, and demos built from passing traces help.

2.4 Instruction search: APE, OPRO and MIPROv2

Demos are one lever. The instruction text is the other. The three methods in this section differ in who proposes, how candidates are scored, and how the search is steered.

APE: propose many, score, pick

The Automatic Prompt Engineer (APE) treats the instruction as a "program" that a language model writes. A model proposes many candidate instructions from input-output examples, each is scored on held-out examples, and the best is kept [8]. The paper reports instructions that did as well as or better than human-written ones on 24 of 24 Instruction Induction tasks and 17 of 21 curated BIG-Bench tasks [8]. I read only the abstract. APE has no memory of what it tried: each round of proposals is fresh. That is the gap that OPRO fills.

OPRO: show the optimizer its own history

OPRO (Optimization by PROmpting) puts the history of past attempts inside the prompt of a proposing model. The "meta-prompt" holds earlier solutions with their scores, usually sorted so that the best are last, and asks for a new solution that scores higher [5]. The scorer is a second model run on a training set. In the paper, each meta-prompt holds the best 20 instructions so far and three random training examples [5]. The paper reports, with a PaLM 2-L scorer on GSM8K and PaLM 2-L-IT as the optimizer, an instruction found by optimization, "Take a deep breath and work on this problem step-by-step.", with a test accuracy of 80.2%, against 71.8% for "Let's think step by step." It also reports gains of up to 8% on GSM8K and up to 50% on Big-Bench Hard over human-written prompts [5]. These are the paper's numbers, for its models and its setup. As noted in 2.2, the headline instruction was chosen among several by test accuracy.

What OPRO adds over APE is the trajectory. What it lacks is any explanation of why a candidate scored low. The proposer sees a number and must guess. The next two methods give the proposer more to read.

MIPROv2: ground the proposals, then search jointly

MIPROv2 is the first required reading. It optimizes a multi-module language-model program, which means both the instruction and the demos of every module, with no access to module-level labels or gradients [2]. Figure 2 shows the three steps the lecture draws [32], with details from the paper [2].

1. Bootstrap demos. Run the program on training inputs. Keep traces whose final output passes the metric (above a threshold) as candidate demos for each module [2].
2. Propose instructions. A proposer model writes candidate instructions for each module. It is grounded: it reads a summary of the dataset, a summary of the program, some successful bootstrapped traces, and a randomly chosen "prompting tip" to vary the style [2]. The paper also studies meta-optimizing the proposer itself.
3. Search combinations. Pick one instruction and one demo set per module, score the combination on a minibatch, and update a model of which choices are good. MIPRO's surrogate is a tree-structured Parzen estimator (the paper uses an Optuna multivariate TPE). Every so often the best combinations get a full evaluation [2]. The paper compares this credit-assignment rule with a greedy one and with a history-based one that is closest to OPRO [2].

The paper's headline: MIPRO outperforms baselines on 5 of 7 tasks, by as much as 13% accuracy, with Llama-3-8B as the task model [2]. Its lessons are worth remembering more than the number:

• Bootstrapped demos are the most important ingredient in most of the paper's tasks.
• Optimizing instructions and demos jointly gives the best results in general.
• Instruction optimization matters most for tasks with conditional rules that a few examples cannot show.
• Grounding the proposer helps, but the best proposal strategy varies from task to task.
• The results are mixed in places. The Bayesian search was not always significantly better than random search [2].

The last point is an honest one for a paper to make. The lab reproduces the shape of it: a grounded proposer gives a large gain in the toy world, and the surrogate adds a small one.

2.5 Learning from text feedback: TextGrad and GEPA

OPRO and MIPRO see scores. A language model can also read words: an error message, a failed test, a judge's comment. The next two methods use text as the learning signal.

TextGrad: backpropagation with words

TextGrad copies the shape of automatic differentiation. Variables hold text. A "loss" is a language-model critique of an output. Calling loss.backward() asks the critic to write feedback ("textual gradients") for each upstream variable, walking the compute graph in reverse. An optimizer step, TGD.step(), rewrites the variable using that feedback [6]. The paper reports, for example, GPT-4o's zero-shot accuracy on Google-proof question answering improving from 51% to 55%, and a 20% relative gain on LeetCode-Hard solutions [6]. It was later published in Nature [7]. The lecture dates it 2024 (the arXiv version), and the journal version is 2025. I read the abstract and the description of the optimizer, not the full experiments.

The lab has a toy version. A loss text lists the failures of a minibatch. A mock critic reads it and writes "consider naming: <rule>" for each prompt variable that is upstream of the failure. A step adds the most requested word. A gate keeps the change only if a fresh run of the same minibatch is no worse. The critic is a few lines of rules that I wrote. Its power is the power of those rules, not of a language model.

GEPA: reflective evolution with a Pareto set

GEPA (Genetic-Pareto) is the second required reading. It keeps a pool of candidate programs. In each iteration it picks a candidate, picks one module of it, runs it on a small minibatch, and asks a model to reflect on the traces and the feedback and to write a new instruction for that module. The feedback can be richer than a score. The paper's feedback function returns a score plus text, and it can include evaluation traces such as compiler errors [1]. If the child beats its parent on the minibatch, it joins the pool and is scored on the full Pareto set. At the end GEPA returns the candidate with the best average on that set [1]. Figure 3 shows the loop and the selection rule.

Why a Pareto set instead of "the best so far"

If you always extend the candidate with the best mean score, the pool collapses into one lineage, and a strategy that is good on a few tasks never gets a chance to spread. GEPA's Algorithm 2 keeps, for each training instance, the candidates that reach the best score on it. It removes candidates that are dominated, and samples among the rest in proportion to the number of instances each one wins [1]. In Figure 3, A wins tasks 1, 2 and 4, B wins task 3, and C wins only task 2, which A also wins. So C is dropped, and A is sampled three times as often as B.

What the paper reports, and how to read it

The abstract says GEPA outperforms GRPO by 6% on average and by up to 20%, with up to 35 times fewer rollouts, and outperforms MIPROv2 by over 10% (for example +12% on AIME-2025) [1]. The table below is from the paper's Table 1 (Qwen3-8B as the task model), copied by hand. It lists one number per task.

Method	HotpotQA	IFBench	HoVer	PUPA	AIME-2025	LiveBench-Math	Aggregate
Baseline	42.33	36.90	35.33	80.82	27.33	48.70	45.23
GRPO (24,000 rollouts)	43.33	35.88	38.67	86.66	38.00	51.26	48.91
MIPROv2	55.33	36.22	47.33	81.55	20.00	46.60	47.84
GEPA	62.33	38.61	52.33	91.85	32.00	51.95	54.85
GEPA + Merge	64.33	28.23	51.67	86.26	32.00	51.95	52.40

Source: Agrawal et al. [1], Table 1, Qwen3-8B. The aggregate gains over the baseline are +3.68 (GRPO), +2.61 (MIPROv2), +9.62 (GEPA) and +7.17 (GEPA+Merge); the first three I recomputed from the table and they match. GEPA used about 1,800 to 7,000 rollouts per task in this table, against 24,000 for GRPO. PUPA is the privacy-delegation benchmark from the PAPILLON paper [15].

Here is how I read it, and where I would be careful:

• GEPA wins the aggregate by a clear margin, and has the best score of the three main methods on five of six tasks. GRPO wins on AIME-2025 (38.00 against 32.00). MIPROv2 is below the baseline on AIME-2025 (20.00 against 27.33).
• The GRPO baseline in the paper uses LoRA, a parameter-efficient method [16]. A comparison with full-parameter RL could differ. The paper describes how hyperparameters were chosen for GEPA and for the baselines [1].
• The paper reports that most of GEPA's rollouts went to scoring candidates on the validation set, not to proposing them. It also reports that GEPA's instruction-only prompts were up to 9.2 times shorter than MIPROv2's, which matters for cost at inference time [1].
• The merge step (a form of crossover between candidates) did not help on average in this table (52.40 against 54.85). It helped on one task (HotpotQA) and hurt on IFBench, where it fell to 28.23 [1].
• With GPT-4.1 Mini as the task model (Table 2), the paper reports GEPA ahead of MIPROv2, TextGrad and Trace, by 5.64%, 6.11% and 3.27% on the aggregate [1].
• The authors include the course's instructor Michael J. Ryan, and the authors of MIPROv2 and DSPy [1] [2]. That is not a reason to doubt the results. It is a reason to want a third group to repeat them, and I report the claims as the authors' claims.

The DSPy documentation has a page on GEPA that matches the paper's split: a training set for the reflective updates, a validation set for the Pareto scores, a reflection model, and an auto setting for the budget [10]. The code is open source [11]. The page also quotes cost savings from companies. I could not verify those, so I do not use them.

2.6 Weights, part 1: SFT on agent rollouts

Context optimization changes what the model reads. Weight optimization changes the model. The simplest way is supervised finetuning (SFT): train the model to produce given outputs, with the usual next-token cross-entropy. The lecture first recalls instruction finetuning, then asks how to do SFT for agents [32].

The recipe: teacher rollouts, filtered

1. Take a strong teacher agent (a larger model, or the same model with a good prompt).
2. Run it on many tasks and record its full trajectories: thoughts, tool calls, observations, answers.
3. Filter by outcome. Keep only trajectories that pass the metric (rejection sampling), as in STaR [12].
4. Finetune a student on the kept trajectories.

Two papers show the idea at small scale. FireAct finetunes language models on ReAct-style trajectories generated by GPT-4. The paper reports that finetuning Llama-2-7B on 500 such trajectories raised its exact match on HotpotQA by 77% (from 14.8 to 26.2 in its table), and that finetuned GPT-3.5 rose from 31.4 to 39.2, and to 41.0 when ReAct and chain-of-thought trajectories were mixed. It also reports that finetuned models needed about four times less inference time, because the few-shot examples were no longer in the prompt [17]. AgentTuning builds an instruction set of agent trajectories (AgentInstruct), mixes it with general instructions ("hybrid" tuning), and releases the AgentLM models. The authors say AgentLM-70B is the first open model to match GPT-3.5-turbo on agent tasks they had not trained on [18]. I read the abstracts and the key tables of these papers, and I did not rerun anything.

The lecture also points to a larger data-recipe study for agentic models, which it calls OpenThoughts-Agent. The arXiv paper is titled Data Recipes for Agentic Models. Its authors ran more than 100 ablations on how to build the training data, built a 100K-example set, and report that finetuning Qwen3-32B on it gives 44.8% on average across seven agentic benchmarks, 3.9 points above the strongest open-data model they compare with (40.9%) [19]. I read only the abstract of this paper.

What can go wrong

• Off-policy states. The student learns from states the teacher visited. At test time the student visits its own states, which can be different, and an early mistake leads somewhere the data never covered. This is one reason on-policy methods (2.7 and 2.8) exist.
• A weak filter. If the metric passes wrong trajectories, you train on them. The lab shows this with a made-up teacher that is right only part of the time: SFT on unfiltered rollouts copies the teacher's success rate, and SFT on filtered rollouts does far better.
• Narrow data. Training on one kind of trajectory can harm other abilities. AgentTuning's hybrid mix is a response to this [18].
• Cost and access. You need the weights, or a provider that finetunes for you, and you need GPUs or a budget.

2.7 Weights, part 2: RL with PPO and GRPO

SFT copies a teacher. Reinforcement learning (RL) improves the student against a reward, using the student's own samples. The lecture's key condition is verifiability: RL works best when you can check the result with code, such as a unit test, an exact answer, or a game score [32]. Without a trustworthy reward, RL optimizes the wrong thing very efficiently.

The policy gradient, in two lines

Treat the model as a policy π_θ that picks tokens. Sample an answer, get a reward R, and nudge the parameters to make good answers more likely. The simplest version is REINFORCE: multiply the gradient of the log-probability of the sampled answer by (R − b), where b is a baseline that cuts the variance. It works, and it can be unstable, because one big update on a lucky batch can ruin the policy. PPO and GRPO are two ways to control the size of the update.

PPO: a clipped step, with a learned baseline

Proximal Policy Optimization defines the ratio r = π_new(token) / π_old(token), and maximizes the clipped objective: min(r·A, clip(r, 1−ε, 1+ε)·A), where A is the advantage and ε is about 0.2 in the paper [4]. When A is positive, the clip stops the gradient once r passes 1+ε. When A is negative, it stops once r falls below 1−ε. So one batch cannot push the policy far from where it was. For language models, the lecture describes PPO as using four models: the policy and a value model (both trained), a frozen reference model, and a reward model or checker. The value model estimates the baseline per token, advantages come from generalized advantage estimation (GAE), and a KL penalty to the reference sits in the reward [32] [3].

GRPO: the group is the baseline

Group Relative Policy Optimization comes from DeepSeekMath, the third required reading. It removes the value model. For each prompt it samples a group of G answers, scores them, and defines each answer's advantage by normalizing within the group: A_i = (r_i − mean(r)) / std(r). Every token of answer i gets that same advantage. The objective keeps PPO's clipped ratio, and adds the KL term to the reference directly in the loss with an unbiased estimator, π_ref/π − log(π_ref/π) − 1, so the reward is not mixed with the KL [3]. Figure 4 compares what the two methods keep in memory.

The DeepSeekMath paper also shows what GRPO bought. Applied to its 7B instruction-tuned model, GRPO raised GSM8K accuracy from 82.9% to 88.2% and MATH from 46.8% to 51.7%, in the authors' setup (they sampled 64 outputs per question, with a KL coefficient of 0.04 and a learning rate of 1e-6) [3]. In its discussion the paper says RL improved majority-vote accuracy (Maj@K) but not Pass@K. That is, it made the model's distribution more robust rather than adding new abilities [3]. The paper also sets out a unified view in which SFT, rejection-sampling finetuning, DPO, PPO and GRPO differ in three things: the source of the data, the reward function, and the gradient coefficient [3]. The lecture does not stress this view, but it is a useful way to see the whole week at once.

A scale example: DeepSeek-R1-Zero applied GRPO with rule-based rewards to a base model, and the paper reports its average pass@1 on AIME 2024 rising from 15.6% to 77.9% during training [20]. That is a result about one model, one reward and one benchmark, from the authors of the method.

Known issues with GRPO

• No signal on uniform groups. If all G answers pass, or all fail, the group's advantage is zero. Hard prompts that the model never solves teach nothing.
• Biases. Liu et al. study "R1-Zero-like" training and report an optimization bias in GRPO that artificially lengthens responses, especially wrong ones. The paper traces a length bias to the per-answer 1/length term in the objective, and a question-level difficulty bias to dividing by the group's standard deviation. It proposes Dr. GRPO, which removes both, and reports 43.3% on AIME 2024 with a 7B base model [21]. Section 2.11 shows the standard deviation effect with numbers.
• Agents are multi-turn. The simple recipe scores one answer. An agent acts over many turns and gets a reward at the end. RAGEN studies RL for multi-turn agents and reports a recurring failure it calls the Echo Trap, with reward collapse and gradient spikes. It proposes a stabilized variant, StarPO-S [22]. ArCHer takes another route: a hierarchy with an utterance-level critic and a token-level policy [23]. I read the abstracts of both.

The GEPA paper's comparison (Table 1 above) is the lecture's main evidence that, for a modest budget, text-based optimization can beat RL on rollouts. The honest statement is narrower than "prompts beat RL": in that paper's setup, with LoRA-based GRPO and 24,000 rollouts, a reflective prompt optimizer did better on five of six tasks with far fewer rollouts [1].

2.8 Weights, part 3: on-policy distillation

SFT on teacher text has a dense signal (every token is supervised) but it is off-policy (the student sees the teacher's states). RL is on-policy (the student sees its own states) but sparse (one reward per answer). The lecture's third option takes the good half of each: on-policy distillation [32]. Figure 5 compares the three.

How it works

1. The student samples an answer.
2. The teacher scores the student's tokens: for each token the student chose, compute log π_student(token) − log π_teacher(token). This is a sample of the reverse KL, KL(student ∥ teacher), at that token.
3. Use the negative of that number as the per-token reward, and apply a standard RL update. The lecture and the blog post use a discount of zero, so each token is credited for its own term only [32] [25].

The Thinking Machines post says this is a one-line change on top of KL-regularized RL code, and that reverse KL is "mode seeking": the student learns to produce one behavior of the teacher precisely, and does not spread its mass over all of them [25]. The academic root is generalized knowledge distillation (GKD), which trains the student on sequences it generates itself, with feedback from the teacher, and allows a choice of divergences [24]. I read the GKD abstract only.

Why does the direction of the KL matter? Forward KL, KL(teacher ∥ student), is the loss that SFT on teacher samples minimizes. It punishes the student for putting no mass where the teacher puts mass, so a one-bump student fit to a two-bump teacher sits between the bumps. Reverse KL punishes mass where the teacher has none, so the student picks one bump. The lab shows both on a toy teacher with two modes. That is a picture of the effect, not evidence about language models.

What the post claims

Here is a summary of the claims in the Thinking Machines post (Lu et al., October 2025), for a Qwen3-8B-Base student and a Qwen3-32B teacher on math reasoning [25]:

• SFT on about 400,000 prompts reached 60% on AIME 2024. The authors extrapolate that about 2 million prompts would be needed for 70%.
• RL reached 68%. On-policy distillation reached 70%.
• They report roughly 9 to 30 times less compute for on-policy distillation than the other methods, in their own accounting of FLOPs. They also cite a Qwen3 report table in which an off-policy stage reached 55.0, adding RL gave 67.6 at 17,920 GPU hours, and adding on-policy distillation gave 74.4 at 1,800 GPU hours.
• They claim the reward cannot be "hacked" in the usual way, because the teacher's log-probabilities are the reward. That is their claim, and it holds only as far as the teacher's judgment is the thing you want.

2.9 Agents as optimizers: DGM and autoresearch

So far a human wrote the optimizer. The last family turns the optimizer into an agent that can change the system, even its own code. Figure 6 shows two loops.

The Darwin Gödel Machine

The Darwin Gödel Machine (DGM) keeps an archive of coding agents. In each round it picks a parent from the archive, lets that agent read and edit its own code to make a child, evaluates the child on coding benchmarks, and adds working children to the archive [26]. The paper's parent selection mixes two ideas: choose agents with higher benchmark scores more often, and choose agents with fewer functioning children more often. The weight of agent i is w_i = s_i · h_i, where s_i = 1 / (1 + exp(−λ(α_i − α_0))) is a sigmoid of its score α_i (the paper uses λ = 10 and α₀ = 0.5), and h_i = 1 / (1 + n_i) shrinks with the number of children n_i. The probability is w_i divided by the sum of all weights [26]. Weak agents keep a small chance, because a weak agent may be a stepping stone.

The abstract reports the benchmark gains: on SWE-bench from 20.0% to 50.0%, and on Polyglot from 14.2% to 30.7% [26]. The ablations in its Table 1 matter more for the lesson. They give (SWE-bench / Polyglot): DGM 50.0 / 38.0; without open-ended exploration (the variant without the archive, as I read the paper) 23.0 / 14.0; without self-improvement 39.0 / 28.0; and a greedy variant 39.7 / 30.0 [26]. The lecture's summary is that both self-improvement and exploration are critical [32]. The paper also reports problems. Some agents hallucinated that they had used a tool, and one appendix example shows objective hacking, where an agent changed what was measured and not what was meant. The authors describe sandboxing and human oversight as controls [26]. The numbers are for the authors' setup and benchmarks.

The lab's archive is a toy. The performance landscape is made up, and I built in a stepping stone on purpose: two features that each cost a little alone and are worth a lot together. A greedy hill climber that never extends a worse agent cannot find it. An archive with the DGM's selection can. This shows why the idea is plausible. It does not show that real agents have such landscapes.

Other agents-as-optimizers

• Meta-Harness searches over harness code (what to store, retrieve and show to the model). Its proposer is an agent with filesystem access to the code, scores and traces of all earlier candidates. The authors report 7.7 points over a context-management baseline on text classification with 4 times fewer context tokens, and gains on math retrieval and on TerminalBench-2 [27]. I read the abstract.
• AlphaEvolve (Google DeepMind) is an evolutionary coding agent. Language models propose edits to code, and automated evaluators score the results. The white paper reports uses such as a data-center scheduling heuristic [28]. I read the abstract. The evaluator is again the key: it must be automatic and trusted.

Autoresearch: a minimal loop

Karpathy's autoresearch repository is small on purpose. An agent may edit only train.py. Each run trains for a fixed five-minute budget, and the metric is validation bits-per-byte (lower is better). A file program.md holds the instructions. The agent keeps a change if the metric improves and discards it if not. The README says this gives about 12 experiments an hour, or about 100 overnight, on one GPU [29]. There is no controlled evidence in the repository. It is a recipe, not a study.

Cerebras published a blog post about running autoresearch-style loops. It reports 71 experiments on two problems. With a tight scope and a strict gate, the loops gave useful results. With a loose objective and rare checks, the agent drifted away from the task within hours. In one run it masked parts of the model in memory, which cost accuracy and saved no memory [30]. The title says "cheating", and the body is mostly about drift and about fixes such as isolated experiment directories and stricter, more frequent checks. It is a vendor post, and the evidence is one team's runs. The point I take from it matches 2.2: an agent that can edit the eval, or that loses sight of the objective, will optimize something else. The lab's loop locks eval.py and program.md with a hash and blocks any change to them. That is a toy guard. A real system also needs a sandbox and review.

2.10 Prompts or weights?

The lecture ends with the practical question [32]. Figure 7 is my decision guide, built from the lecture's reasons.

Reasons to optimize context first [32]:

• You can use only an API model, so you cannot change weights.
• You want to rerun often and fast. Changing a prompt takes minutes.
• Environment interaction (a rollout) is expensive, and prompt optimization is more sample-efficient. The GEPA paper's claim is that reflective prompt evolution reached higher scores than GRPO with far fewer rollouts in its tests [1].

Reasons to optimize weights [32]:

• Context optimization fails, especially on small models, and no prompt gets you there.
• You have a strong teacher, and want a cheaper student (SFT or distillation).
• Context bloat: the optimized prompt has grown so long that it is slow and expensive on every call. GEPA's shorter prompts help here, and so does moving the knowledge into weights [1].
• Weights scale: with much more data they can go further than any prompt (the lecture's point, and not one that I tested).

Both. BetterTogether alternates the two: optimize prompts, then weights, then prompts again. In experiments with multi-hop QA, math and classification on three small open models, its abstract reports that this beat weights alone by up to 60% and prompts alone by up to 6%, on average across models and tasks [14]. In the lecture's words, they are complementary [32].

My own rule of thumb is to start with the cheapest lever that changes the metric: first prompts and demos, which cost minutes, then SFT on filtered rollouts, which costs hours and needs the weights. Use RL when you have a verifiable reward and the budget. Use distillation when a strong teacher exists. Whatever you choose, measure it on held-out data with your own eval, and count the full cost, including the cost of the optimization run and of serving the result. The lab's cost calculator is a toy for that last step, and its numbers are made up.

2.11 Worked arithmetic

These small calculations are mine. Do them by hand once, then check them against the lab code.

(a) GRPO advantages

A prompt gets 8 sampled answers, and only one passes: rewards (1, 0, 0, 0, 0, 0, 0, 0). The mean is 0.125. The population standard deviation is √(0.125 × 0.875) = 0.331. The passing answer gets advantage (1 − 0.125) / 0.331 = 2.65. Each failing answer gets (0 − 0.125) / 0.331 = −0.38. The advantages sum to zero, and their squares sum to 8, which is the group size. If the rewards are (1, 1, 1, 1, 0, 0, 0, 0), the mean is 0.5, the standard deviation is 0.5, and the advantages are +1 and −1, and again the squares sum to 8. So dividing by the standard deviation gives every mixed group the same total weight. Without it (the Dr. GRPO choice), the first group has total squared advantage 0.875 × 0.875 + 7 × 0.125 × 0.125 = 0.875, and the second has 8 × 0.25 = 2.0, so the 50% prompt counts 2.3 times as much. Neither is "right". They weight easy, hard and medium prompts differently [21] [3]. Libraries differ in whether they use the population or the sample standard deviation, which changes the scale by √(G/(G−1)). The lab uses the population version.

(b) The PPO clip, with ε = 0.2

Advantage A	Ratio r	Clipped objective value	Does the gradient flow?
+1	1.3	min(1.3, 1.2) = 1.2	No: the policy already moved far enough up.
+1	0.5	min(0.5, 0.8) = 0.5	Yes: the good token became less likely, so push it up.
−1	0.7	min(−0.7, −0.8) = −0.8	No: the bad token has already been pushed down far enough.
−1	1.5	min(−1.5, −1.2) = −1.5	Yes: the bad token became more likely, so push it down.

(c) DGM parent weights

With λ = 10 and α₀ = 0.5, an agent with score 0.5 and no children has weight 0.5 × 1 = 0.5. With two children it has 0.5 / 3 = 0.167. An agent with score 0.7 and no children has sigmoid(10 × 0.2) = sigmoid(2) = 0.881. An agent with score 0.3 has sigmoid(−2) = 0.119. So the weak agent is picked about one-seventh as often as the strong one, and it is not excluded.

(d) A break-even number, with made-up prices

This uses the lab's default numbers, which are all made up. Fine-tuning a small model costs a fixed 4,420 dollars (20,000 teacher rollouts for 180 dollars, training for 240, and 40 engineer-hours at 100 dollars). Serving costs 0.00075 dollars per call on the small model, against 0.009 on the large model with no optimization. Say the large model gets 60% right, the finetuned small one 80%, and each wrong answer costs 0.10 dollars. The cost per call, including errors, is 0.009 + 0.4 × 0.10 = 0.049 for the large model and 0.00075 + 0.2 × 0.10 = 0.02075 for the small one. The saving is 0.02825 per call, and 4,420 / 0.02825 ≈ 156,000 calls. Below that, do nothing. Above it, the fine-tune pays. Change any made-up number and the answer moves, which is the point of the exercise.

2.12 Week 5b in one page

• Before you optimize: have data and a trusted eval. Hold out data. Count rollouts [32] [1].
• Demos: order matters. Keep traces that pass the metric. Bootstrap, then search [13] [9].
• Instructions: APE proposes and scores. OPRO shows the proposer its history. MIPROv2 grounds proposals and searches instructions and demos together [8] [5] [2].
• Text feedback: TextGrad passes critiques backward. GEPA reflects on traces, keeps a Pareto pool, and in the authors' tests beat GRPO with far fewer rollouts. Read it with the caveats in 2.5 [6] [1].
• SFT: train on filtered teacher rollouts. Watch for off-policy states and weak filters [17] [12].
• RL: needs a verifiable reward. PPO: clipped ratio, value model, four models. GRPO: group advantage, no value model, KL in the loss [4] [3].
• On-policy distillation: per-token reverse KL on the student's own samples. Compute claims are vendor claims [25] [24].
• Agents as optimizers: an archive helps exploration (DGM). A locked eval keeps the loop honest (autoresearch) [26] [29] [30].
• Choose: prompts first when you are API-only, need speed, or have costly rollouts. Weights when prompts fail, a teacher exists, or the prompt is bloated. Both is often best [14] [32].

3. Reading guide

Three readings are required by the course. For each you get: why it matters, the key ideas, and questions to think about as you read. Read the questions first. The optional table at the end points to the sources that I used for the notes. For some of those I read only the abstract, and the references say which.

Optional readings

Source	Why it is worth your time
Khattab et al., DSPy [9]; DSPy optimizer docs [10]; GEPA code [11]	The programming model behind the optimizers, and the tools you would actually use. Find where the train and validation sets go.
Zhou et al., APE [8]; Yang et al., OPRO [5]	The two simplest instruction-search ideas. In OPRO, read the meta-prompt figure and Table 4, and note how the examples were chosen.
Yuksekgonul et al., TextGrad [6] [7]	The autodiff analogy. Read the optimizer description and one application.
Lu et al., prompt order sensitivity [13]; Zelikman et al., STaR [12]	Why demos need evaluation, and the rejection-sampling loop.
Soylu, Potts, Khattab, BetterTogether [14]; Li et al., PAPILLON [15]	Alternating prompts and weights, and the privacy task used as a GEPA example.
Chen et al., FireAct [17]; Zeng et al., AgentTuning [18]; Data Recipes for Agentic Models [19]	SFT for agents. In FireAct, read the tables that compare prompting with finetuning, and the cost table.
Schulman et al., PPO [4]; DeepSeek-R1 [20]; Liu et al., Dr. GRPO [21]	The clipped objective, a large GRPO run, and a critique of GRPO's biases.
Wang et al., RAGEN [22]; Zhou et al., ArCHer [23]	RL for multi-turn agents. Compare how each handles credit across turns.
Agarwal et al., GKD [24]; Thinking Machines, on-policy distillation [25]	The academic and the practical sides of on-policy distillation. Read the post's plots with the vendor caveat in mind.
Zhang et al., Darwin Gödel Machine [26]; Lee et al., Meta-Harness [27]; AlphaEvolve [28]	Agents that improve systems. In the DGM, read the parent-selection formula and the safety discussion.
Karpathy, autoresearch [29]; Cerebras blog on guarding autoresearch loops [30]	A minimal loop, and the failure mode. Read the repository's program.md and train.py before the blog.
Hu et al., LoRA [16]	The method behind many weight-optimization baselines.

4. Hands-on lab: optimizing mock agents and toy policies

Time: about 4 hours. Goal: build small, testable versions of this week's methods, check the math with numbers, break them on purpose, and see what each one costs. You do not need an API key or any package. Every "language model" is a mock: a few lines of Python that act like one. Every "policy" is a table of logits.

What you will learn

• How a mock model's accuracy can depend on instruction, demos and context length, and how to count rollouts.
• How bootstrapped few-shot filters traces by the metric, and what happens if you do not filter.
• How a MIPRO-like search combines proposals, bootstrapped demos, minibatch scores and a surrogate.
• How an OPRO meta-prompt is built, parsed and answered, and why its best score is optimistic.
• How a TextGrad-style graph passes a critique backward and steps, with a gate and with frozen variables.
• How GEPA's Pareto selection differs from best-mean selection, and how a rollout budget is spent.
• How REINFORCE, PPO-clip and GRPO differ, and how to check a gradient against finite differences.
• How SFT on filtered rollouts, and reverse-KL distillation, differ from each other and from RL.
• How a DGM-style archive picks parents, and how an autoresearch loop can be guarded.
• How to compare prompts and weights on cost and errors, and why every number needs its assumptions.

Setup

1. You need a recent Python 3. There is nothing to install. The code uses only the standard library. I ran it on Python 3.13 only.
2. Save the code below as optimization_toolkit.py. Or download it: optimization_toolkit.py.
3. Save the tests as test_optimization_toolkit.py in the same folder. Or download: test_optimization_toolkit.py. The test code is not printed on this page, to keep the page shorter. It is in the download, and the names of all the tests are listed after the exercises.
4. Run python optimization_toolkit.py for the eleven-section demo, and python test_optimization_toolkit.py for the tests. The demo and the tests are deterministic. The tests take about six seconds. The demo takes about two seconds.

The code: optimization_toolkit.py

The code has six parts, matching Figure 8, and a demo at the end. The comments name each choice that I made and the sources did not. The file is long. Read the part you are working on, with the matching section of the notes open.

"""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()

Run the demo

$ python optimization_toolkit.py
== 1. A toy prompt world (made-up numbers) ==
no help, exact accuracy:            0.350
best rule words, no demos:          0.658   {'parse': 'Solve units negatives dates', 'answer': 'Solve abbrev currency ties'}
per-kind accuracy of the best words: {'A': 0.75, 'B': 0.75, 'C': 0.45, 'D': 0.75, 'E': 0.75, 'F': 0.5}
note: naming 'negatives' hurts kind C, naming 'abbrev' hurts kind F (built-in conflicts)

== 2. Bootstrapped few-shot ==
demos from passing traces only: mean exact accuracy 0.480
demos from every trace: mean exact accuracy 0.219
bootstrap + random search over demo sets: mean exact accuracy 0.417

== 3. MIPRO-like search (equal rollouts, 8 seeds) ==
grounded + surrogate       exact accuracy 0.624  rollouts 720
grounded, random combos    exact accuracy 0.598  rollouts 720
ungrounded + surrogate     exact accuracy 0.508  rollouts 660

== 4. OPRO loop (one shared instruction) ==
best instruction: 'Solve the currency concise dates ties persona'
score on its own scoring set 0.792, exact accuracy 0.550, 43 instructions tried, 1032 rollouts
end of the meta-prompt shown to the proposer (its last two entries, best last):
text: Solve the currency concise dates ties carefully
score: 0.71
text: Solve the currency concise dates ties persona
score: 0.79

== 5. TextGrad-style feedback step ==
after 8 steps: {'parse': 'Solve the task dates negatives units', 'answer': 'Solve the task abbrev currency ties'}
exact accuracy 0.658, accepted 4 of 8 steps, 156 rollouts
with the parse prompt frozen: {'parse': 'Solve the task.', 'answer': 'Solve the task abbrev currency ties'}  exact accuracy 0.508

== 6. GEPA-style reflective evolution (budget 1000 rollouts, 8 seeds) ==
select=pareto exact accuracy 0.639  mean pool 17.9  mean accepted 16.9
select=best   exact accuracy 0.511  mean pool 10.2  mean accepted 9.2
pareto + merge  exact accuracy 0.620

== 7. REINFORCE, PPO-clip, GRPO on a toy verifiable task ==
targets per state: [(0, 2), (2, 0), (1, 2)]; start success 0.111 (chance is 1/9 per state)
REINFORCE  final exact success, mean of 8 seeds: 0.739  (min 0.712)   [1,440 rollouts each]
PPO-clip   final exact success, mean of 8 seeds: 0.853  (min 0.841)   [1,440 rollouts each]
GRPO       final exact success, mean of 8 seeds: 0.841  (min 0.816)   [1,440 rollouts each]
GRPO advantages for rewards [1,0,0,0]: [1.73, -0.58, -0.58, -0.58]
GRPO advantages for rewards [0,0,0,0]: [0.0, 0.0, 0.0, 0.0]

== 8. SFT from filtered rollouts, and distillation ==
teacher success 0.206
SFT on filtered (reward 1 only) student success 0.890
SFT on unfiltered               student success 0.199
on-policy distillation: reverse KL 0.057 -> 0.0000, student success 0.206
one-bump student fit by forward KL: mean 9.72, sigma 6.80, loss 0.698
one-bump student fit by reverse KL: mean 5.00, sigma 1.20, loss 0.693
forward KL covers both modes and sits between them. reverse KL picks one mode.

== 9. DGM-style archive (toy landscape, 20 seeds, 60 steps) ==
dgm        mean best score 0.825
greedy     mean best score 0.370
uniform    mean best score 0.754
best_only  mean best score 0.370
parent probabilities for a one-variant archive: [1.0]

== 10. Autoresearch-style loop with an eval guard ==
guard=True  reported 1.201  true 1.201  blocked edits 15.2
guard=False reported 0.255  true 1.201  blocked edits 0.0

== 11. Prompt versus weights cost calculator (ALL NUMBERS MADE UP) ==
calls     5,000  fine-tuning allowed   cheapest incl. errors: none    ($245)
calls     5,000  fine-tuning API-only  cheapest incl. errors: none    ($245)
calls    50,000  fine-tuning allowed   cheapest incl. errors: none    ($2,450)
calls    50,000  fine-tuning API-only  cheapest incl. errors: none    ($2,450)
calls   500,000  fine-tuning allowed   cheapest incl. errors: both    ($13,282)
calls   500,000  fine-tuning API-only  cheapest incl. errors: prompt  ($20,159)
calls 5,000,000  fine-tuning allowed   cheapest incl. errors: both    ($84,860)
calls 5,000,000  fine-tuning API-only  cheapest incl. errors: prompt  ($193,409)
calls above which weights beat doing nothing (made-up numbers): 156,460

Exercises

1. Read the world. Use p_correct to print the accuracy of each kind under three programs: no help, one rule word named, and one word plus two correct demos. Which term in the code is the length penalty, and when does it start [2]?
2. A filter worth having. Run bootstrap_demos with keep_all=True and with the filter on, over 20 seeds. Report the mean exact accuracy of each. Then change the world so that wrong demos cost nothing (bad_demo_pen=0). What happens to the gap? What does that say about when a metric filter matters [12] [9]?
3. Order, in the toy. The toy has no order effect. Add one (for example, only the last two demos count). Show that the best demo set now depends on order, and write a search that finds it. Why does Lu et al. find order sensitivity in real models [13]?
4. Surrogate against random. Run mipro_search with surrogate=True and False over 30 seeds. Report the mean difference and a rough error bar. Is it significant in your runs? Compare with the paper's mixed result [2].
5. Remove a grounding input. Edit propose_instructions so that it ignores the failure rates but still respects the program's modules. Then so that it ignores the modules. Which removal hurts more? Which paper lesson does this echo [2]?
6. The proposer's budget. Change n_instr from 5 to 2 and to 12. How do rollouts and accuracy change? At what point do more proposals stop helping?
7. OPRO's optimism. Run opro over 30 seeds. For each run, compute the difference between the reported best score and the exact accuracy of the best instruction. Report the mean. Fix it by rescoring the best instruction on a fresh set, and show the gap disappears [5].
8. A better meta-prompt. Change opro_prompt to keep only the top 3 entries instead of 6, then to show the worst 3 as well. Does the mock proposer notice? Why or why not? What would a real proposer do differently?
9. Break the critic. Make tg_critic return a wrong rule word half the time. What does the gate do? Set gate=False and run again. What does the comparison tell you about validating a step [6]?
10. Freeze and unfreeze. Run textgrad_optimize with frozen=("answer",) and with frozen=("parse",). Which module matters more for the final accuracy? Explain with the per-kind accuracies.
11. Tie-breaking in the front. Call pareto_front on score tables you invent: one with a clear winner, one with specialists, one with ties everywhere, and one with all zeros. For each, say what the paper's algorithm should do, and whether my reading agrees [1].
12. Pareto against best-mean. Run gepa with select="pareto" and "best" at budgets 300, 1,000 and 3,000, over 20 seeds. Plot or tabulate. Where is the gap largest? What stops the improvement at high budgets in this toy [1]?
13. Where the rollouts go. Count how many of gepa's rollouts go to minibatches and how many to Pareto scoring. The paper says most go to validation. Does your count agree? Try a Pareto set of 8, 18 and 40 examples [1].
14. Make the feedback poor. Replace the feedback text by "fail" only (the test test_feedback_is_what_makes_it_work does this). Then give the feedback a partial hint. How does accuracy change? What does that say about the value of rich feedback in the paper [1]?
15. Check a gradient. Pick any entry of the PPO gradient from ppo_grad. Compute the finite difference of ppo_objective by hand with a step of 1e-5 and compare. Do the same for GRPO. Why does the check need ratios that differ from 1 [4] [3]?
16. Clip in action. Run PPO with clip=0.05, 0.2 and 10 and the same seeds. Report the success curve. What does a very large clip do, and why [4]?
17. The KL term. Run ppo_train with kl_beta 0 and 2, and report the final reverse KL to the starting policy and the success rate. What does the penalty buy and cost?
18. Uniform groups. Make the task so hard that the policy rarely succeeds. How often does grpo_collect return a group with all-zero advantages? How would you change the sampling or the reward to fix it [3]?
19. Std normalization. Run grpo_train with std_norm=True and False. Then compute sum(a*a) of group_advantages for groups with 1, 4 and 7 passes out of 8, with and without it. Reproduce the numbers of 2.11(a) [21].
20. Filtered against unfiltered SFT. Change the iters argument of noisy_teacher until the teacher succeeds on about 50% of the tasks, then about 80%. How does the gap between filtered and unfiltered SFT change? What if the filter passes 10% of wrong answers [12] [17]?
21. Two directions. Call fit_gaussian with mu0 from 2 to 18, for both directions. Where does reverse KL pick the left mode, the right mode, or neither? What does this say about the dependence of reverse-KL training on initialization [25]?
22. The enumeration check. For a tiny policy, enumerate all responses and show that on_policy_distill_grad equals the negative gradient of exact_reverse_kl. Explain why the discount of zero is enough here [25] [24].
23. Archive against greedy. Change variant_score so that there is no stepping stone (every feature has a small positive gain). Does the archive still win? Then make the stepping stone worth 0.1. When does the archive stop helping [26]?
24. Cheat and guard. Write a proposer that edits program.md to tell the "agent" to ignore the eval, and a second one that edits eval.py with a one-character change. Show that the guard blocks both. Then write a cheat that the hash guard cannot see (for example, a change inside train.py that detects the eval). What would catch it [29] [30]?
25. Your own prices. Replace the made-up numbers of CostInputs with quotes for a real model pair and a real finetuning service. Find the break-even number of calls. Which input moves it most? Which input are you least sure of [32]?

The tests: test_optimization_toolkit.py

The 288 tests check the toy world (the accuracy rule, each conflict, the caps and clamps, the rollout charge, the feedback text, the exact accuracy and the best possible program); the bootstrap and its filter; the UCB rule and the MIPRO-like search (rollout counts, belief counts, the final pick, the last-trial evaluation); the OPRO prompt (order, parsing, proposals, argmax); the TextGrad-style graph (reverse order, frozen variables, the gate, ties); GEPA (the Pareto front, domination, weights, reflection, merge, budget, ties, greedy parents); the tabular policy; every gradient against finite differences (REINFORCE with and without a baseline against the exact gradient by enumeration, PPO, GRPO with and without length and std normalization, SFT, and the distillation gradient against the exact reverse KL); the PPO returns and critic; the GRPO advantages and the k3 estimator; SFT with and without the filter; the forward and reverse fits; the DGM weights, the archive and the comparison; the autoresearch loop and its guard; the cost calculator (every formula and the break-even); and the demo and its output. I also changed the code in 107 ways, one at a time (for example, flipping the clip direction, dropping the KL term, ignoring a frozen variable, removing the guard, or dropping the domination step), and checked that at least one test failed each time. The first round used 105 changes, and 18 of them were missed. I wrote tests until none was missed, and then re-ran every change. All 107 are caught. That shows the tests are sensitive to those changes. It does not show that the code has no bugs, and the changes are the ones that I thought of. I did not run the tests on any other version of Python.

$ python test_optimization_toolkit.py -v
test_deterministic (__main__.TestAutoresearch.test_deterministic) ... ok
test_eval_edit_changes_reported_only (__main__.TestAutoresearch.test_eval_edit_changes_reported_only) ... ok
test_guard_allows_train_edits (__main__.TestAutoresearch.test_guard_allows_train_edits) ... ok
test_guard_blocks_program_md_edits (__main__.TestAutoresearch.test_guard_blocks_program_md_edits) ... ok
test_guard_fails_on_edit (__main__.TestAutoresearch.test_guard_fails_on_edit) ... ok
test_guard_passes_on_untouched_files (__main__.TestAutoresearch.test_guard_passes_on_untouched_files) ... ok
test_guarded_loop_blocks_cheats (__main__.TestAutoresearch.test_guarded_loop_blocks_cheats) ... ok
test_hash_changes_with_content (__main__.TestAutoresearch.test_hash_changes_with_content) ... ok
test_honest_loop_improves_true_loss (__main__.TestAutoresearch.test_honest_loop_improves_true_loss) ... ok
test_honest_proposer_changes_only_train (__main__.TestAutoresearch.test_honest_proposer_changes_only_train) ... ok
test_log_reasons (__main__.TestAutoresearch.test_log_reasons) ... ok
test_loop_never_gets_worse (__main__.TestAutoresearch.test_loop_never_gets_worse) ... ok
test_parse_and_format_roundtrip (__main__.TestAutoresearch.test_parse_and_format_roundtrip) ... ok
test_parse_ignores_noise (__main__.TestAutoresearch.test_parse_ignores_noise) ... ok
test_reported_equals_true_with_honest_eval (__main__.TestAutoresearch.test_reported_equals_true_with_honest_eval) ... ok
test_sneaky_sometimes_edits_eval (__main__.TestAutoresearch.test_sneaky_sometimes_edits_eval) ... ok
test_true_loss_minimum (__main__.TestAutoresearch.test_true_loss_minimum) ... ok
test_true_loss_worse_away_from_minimum (__main__.TestAutoresearch.test_true_loss_worse_away_from_minimum) ... ok
test_unguarded_loop_reports_fake_progress (__main__.TestAutoresearch.test_unguarded_loop_reports_fake_progress) ... ok
test_charges_one_rollout_per_example (__main__.TestBootstrap.test_charges_one_rollout_per_example) ... ok
test_filter_helps_on_average (__main__.TestBootstrap.test_filter_helps_on_average) ... ok
test_filtered_demos_are_correct (__main__.TestBootstrap.test_filtered_demos_are_correct) ... ok
test_keep_all_includes_wrong (__main__.TestBootstrap.test_keep_all_includes_wrong) ... ok
test_max_per_kind (__main__.TestBootstrap.test_max_per_kind) ... ok
test_random_search_deterministic (__main__.TestBootstrap.test_random_search_deterministic) ... ok
test_random_search_keeps_instructions (__main__.TestBootstrap.test_random_search_keeps_instructions) ... ok
test_random_search_log_includes_empty_set (__main__.TestBootstrap.test_random_search_log_includes_empty_set) ... ok
test_random_search_picks_best_logged (__main__.TestBootstrap.test_random_search_picks_best_logged) ... ok
test_accuracy_gain_can_justify_a_costlier_call (__main__.TestCost.test_accuracy_gain_can_justify_a_costlier_call) ... ok
test_api_only_large_scale_prefers_prompt (__main__.TestCost.test_api_only_large_scale_prefers_prompt) ... ok
test_both_combines_fixed_costs (__main__.TestCost.test_both_combines_fixed_costs) ... ok
test_both_per_call_between (__main__.TestCost.test_both_per_call_between) ... ok
test_breakeven_formula (__main__.TestCost.test_breakeven_formula) ... ok
test_breakeven_is_where_costs_cross (__main__.TestCost.test_breakeven_is_where_costs_cross) ... ok
test_breakeven_none_when_never_cheaper (__main__.TestCost.test_breakeven_none_when_never_cheaper) ... ok
test_breakeven_none_when_per_call_cost_equal (__main__.TestCost.test_breakeven_none_when_per_call_cost_equal) ... ok
test_breakeven_zero_when_cheaper_from_start (__main__.TestCost.test_breakeven_zero_when_cheaper_from_start) ... ok
test_call_cost (__main__.TestCost.test_call_cost) ... ok
test_free_errors_make_none_win (__main__.TestCost.test_free_errors_make_none_win) ... ok
test_large_scale_prefers_weights_plan (__main__.TestCost.test_large_scale_prefers_weights_plan) ... ok
test_no_finetune_removes_weight_options (__main__.TestCost.test_no_finetune_removes_weight_options) ... ok
test_none_plan (__main__.TestCost.test_none_plan) ... ok
test_prompt_fixed_value (__main__.TestCost.test_prompt_fixed_value) ... ok
test_prompt_per_call_higher_than_none (__main__.TestCost.test_prompt_per_call_higher_than_none) ... ok
test_prompt_per_call_value (__main__.TestCost.test_prompt_per_call_value) ... ok
test_small_scale_prefers_doing_nothing (__main__.TestCost.test_small_scale_prefers_doing_nothing) ... ok
test_totals_add_up (__main__.TestCost.test_totals_add_up) ... ok
test_weights_fixed_value (__main__.TestCost.test_weights_fixed_value) ... ok
test_weights_per_call_value (__main__.TestCost.test_weights_per_call_value) ... ok
test_archive_beats_greedy_in_the_stepping_stone_landscape (__main__.TestDGM.test_archive_beats_greedy_in_the_stepping_stone_landscape) ... ok
test_archive_ids_unique_and_parents_exist (__main__.TestDGM.test_archive_ids_unique_and_parents_exist) ... ok
test_archive_scores_consistent (__main__.TestDGM.test_archive_scores_consistent) ... ok
test_best_only_archive_has_one_member (__main__.TestDGM.test_best_only_archive_has_one_member) ... ok
test_best_only_never_gets_worse (__main__.TestDGM.test_best_only_never_gets_worse) ... ok
test_best_score (__main__.TestDGM.test_best_score) ... ok
test_children_count (__main__.TestDGM.test_children_count) ... ok
test_children_lower_the_weight (__main__.TestDGM.test_children_lower_the_weight) ... ok
test_choose_parent_bad_mode (__main__.TestDGM.test_choose_parent_bad_mode) ... ok
test_choose_parent_greedy_picks_best (__main__.TestDGM.test_choose_parent_greedy_picks_best) ... ok
test_deterministic (__main__.TestDGM.test_deterministic) ... ok
test_greedy_never_leaves_hill (__main__.TestDGM.test_greedy_never_leaves_hill) ... ok
test_higher_score_higher_weight (__main__.TestDGM.test_higher_score_higher_weight) ... ok
test_low_scorers_keep_nonzero_probability (__main__.TestDGM.test_low_scorers_keep_nonzero_probability) ... ok
test_max_score (__main__.TestDGM.test_max_score) ... ok
test_no_neutral_single_move_from_root (__main__.TestDGM.test_no_neutral_single_move_from_root) ... ok
test_probs_sum_to_one (__main__.TestDGM.test_probs_sum_to_one) ... ok
test_root_score (__main__.TestDGM.test_root_score) ... ok
test_self_modify_can_break (__main__.TestDGM.test_self_modify_can_break) ... ok
test_self_modify_toggles_one_feature (__main__.TestDGM.test_self_modify_toggles_one_feature) ... ok
test_stepping_stone (__main__.TestDGM.test_stepping_stone) ... ok
test_weights_formula (__main__.TestDGM.test_weights_formula) ... ok
test_demo_deterministic (__main__.TestDemo.test_demo_deterministic) ... ok
test_demo_has_all_sections (__main__.TestDemo.test_demo_has_all_sections) ... ok
test_demo_labels_made_up_numbers (__main__.TestDemo.test_demo_labels_made_up_numbers) ... ok
test_no_em_dashes_in_toolkit_text (__main__.TestDemo.test_no_em_dashes_in_toolkit_text) ... ok
test_distilled_student_matches_teacher_success (__main__.TestDistill.test_distilled_student_matches_teacher_success) ... ok
test_exact_kl_known_value (__main__.TestDistill.test_exact_kl_known_value) ... ok
test_exact_kl_skips_zero_mass (__main__.TestDistill.test_exact_kl_skips_zero_mass) ... ok
test_exact_reverse_kl_zero_same (__main__.TestDistill.test_exact_reverse_kl_zero_same) ... ok
test_expected_gradient_is_minus_gradient_of_reverse_kl (__main__.TestDistill.test_expected_gradient_is_minus_gradient_of_reverse_kl) ... ok
test_expected_token_reverse_kl_is_the_kl (__main__.TestDistill.test_expected_token_reverse_kl_is_the_kl) ... ok
test_forward_is_reverse_with_swapped_arguments (__main__.TestDistill.test_forward_is_reverse_with_swapped_arguments) ... ok
test_gradient_zero_when_student_equals_teacher (__main__.TestDistill.test_gradient_zero_when_student_equals_teacher) ... ok
test_kl_nonnegative_and_asymmetric (__main__.TestDistill.test_kl_nonnegative_and_asymmetric) ... ok
test_token_reverse_kl_values (__main__.TestDistill.test_token_reverse_kl_values) ... ok
test_token_reverse_kl_zero_for_same_policy (__main__.TestDistill.test_token_reverse_kl_zero_for_same_policy) ... ok
test_training_deterministic (__main__.TestDistill.test_training_deterministic) ... ok
test_training_reduces_reverse_kl (__main__.TestDistill.test_training_reduces_reverse_kl) ... ok
test_collect_size_and_fields (__main__.TestEnv.test_collect_size_and_fields) ... ok
test_reward_exact_match (__main__.TestEnv.test_reward_exact_match) ... ok
test_success_matches_enumeration (__main__.TestEnv.test_success_matches_enumeration) ... ok
test_success_matches_sampling (__main__.TestEnv.test_success_matches_sampling) ... ok
test_success_uniform_policy (__main__.TestEnv.test_success_uniform_policy) ... ok
test_targets_fixed_by_seed (__main__.TestEnv.test_targets_fixed_by_seed) ... ok
test_advantages_all_equal_is_zero (__main__.TestGRPO.test_advantages_all_equal_is_zero) ... ok
test_advantages_mean_zero (__main__.TestGRPO.test_advantages_mean_zero) ... ok
test_advantages_unit_variance (__main__.TestGRPO.test_advantages_unit_variance) ... ok
test_advantages_without_std_norm (__main__.TestGRPO.test_advantages_without_std_norm) ... ok
test_advantages_z_score (__main__.TestGRPO.test_advantages_z_score) ... ok
test_gradient_matches_without_length_norm (__main__.TestGRPO.test_gradient_matches_without_length_norm) ... ok
test_group_sizes (__main__.TestGRPO.test_group_sizes) ... ok
test_k3_is_unbiased_estimate_of_kl (__main__.TestGRPO.test_k3_is_unbiased_estimate_of_kl) ... ok
test_k3_nonnegative_zero_at_equal (__main__.TestGRPO.test_k3_nonnegative_zero_at_equal) ... ok
test_kl_pulls_toward_reference (__main__.TestGRPO.test_kl_pulls_toward_reference) ... ok
test_no_value_model_needed (__main__.TestGRPO.test_no_value_model_needed) ... ok
test_objective_gradient_matches_finite_difference (__main__.TestGRPO.test_objective_gradient_matches_finite_difference) ... ok
test_std_norm_boosts_easy_and_hard_prompts (__main__.TestGRPO.test_std_norm_boosts_easy_and_hard_prompts) ... ok
test_training_deterministic (__main__.TestGRPO.test_training_deterministic) ... ok
test_training_improves (__main__.TestGRPO.test_training_improves) ... ok
test_zero_advantage_groups_only_feel_kl (__main__.TestGRPO.test_zero_advantage_groups_only_feel_kl) ... ok
test_budget_respected (__main__.TestGepa.test_budget_respected) ... ok
test_budget_too_small_for_seed_scoring (__main__.TestGepa.test_budget_too_small_for_seed_scoring) ... ok
test_children_differ_from_parent_in_one_module (__main__.TestGepa.test_children_differ_from_parent_in_one_module) ... ok
test_deterministic (__main__.TestGepa.test_deterministic) ... ok
test_feedback_is_what_makes_it_work (__main__.TestGepa.test_feedback_is_what_makes_it_work) ... ok
test_greedy_parent_is_best_mean_at_the_time (__main__.TestGepa.test_greedy_parent_is_best_mean_at_the_time) ... ok
test_improves_on_seed (__main__.TestGepa.test_improves_on_seed) ... ok
test_merge_creates_two_parent_children (__main__.TestGepa.test_merge_creates_two_parent_children) ... ok
test_merge_takes_best_module_from_each (__main__.TestGepa.test_merge_takes_best_module_from_each) ... ok
test_merge_tie_goes_to_first (__main__.TestGepa.test_merge_tie_goes_to_first) ... ok
test_minibatch_param_changes_cost (__main__.TestGepa.test_minibatch_param_changes_cost) ... ok
test_pareto_beats_greedy_on_average (__main__.TestGepa.test_pareto_beats_greedy_on_average) ... ok
test_pool_has_no_duplicates (__main__.TestGepa.test_pool_has_no_duplicates) ... ok
test_pool_scores_parents_aligned (__main__.TestGepa.test_pool_scores_parents_aligned) ... ok
test_reflect_adds_most_requested_rule (__main__.TestGepa.test_reflect_adds_most_requested_rule) ... ok
test_reflect_does_not_repeat_a_present_rule (__main__.TestGepa.test_reflect_does_not_repeat_a_present_rule) ... ok
test_reflect_ignores_other_module (__main__.TestGepa.test_reflect_ignores_other_module) ... ok
test_reflect_ignores_successes (__main__.TestGepa.test_reflect_ignores_successes) ... ok
test_reflect_noise_leaves_text (__main__.TestGepa.test_reflect_noise_leaves_text) ... ok
test_result_is_best_mean_on_pareto_set (__main__.TestGepa.test_result_is_best_mean_on_pareto_set) ... ok
test_select_all_zero_scores_still_picks (__main__.TestGepa.test_select_all_zero_scores_still_picks) ... ok
test_select_dominated_removed (__main__.TestGepa.test_select_dominated_removed) ... ok
test_select_keeps_specialists (__main__.TestGepa.test_select_keeps_specialists) ... ok
test_select_nonwinner_never_survives (__main__.TestGepa.test_select_nonwinner_never_survives) ... ok
test_select_removes_dominated_candidate (__main__.TestGepa.test_select_removes_dominated_candidate) ... ok
test_select_returns_front_member (__main__.TestGepa.test_select_returns_front_member) ... ok
test_select_single_candidate (__main__.TestGepa.test_select_single_candidate) ... ok
test_select_ties_keep_one (__main__.TestGepa.test_select_ties_keep_one) ... ok
test_select_weighted_by_wins (__main__.TestGepa.test_select_weighted_by_wins) ... ok
test_failure_summary_high_when_unhelped (__main__.TestMipro.test_failure_summary_high_when_unhelped) ... ok
test_failure_summary_keys_and_range (__main__.TestMipro.test_failure_summary_keys_and_range) ... ok
test_mipro_deterministic (__main__.TestMipro.test_mipro_deterministic) ... ok
test_mipro_grounding_helps_on_average (__main__.TestMipro.test_mipro_grounding_helps_on_average) ... ok
test_mipro_history_keys_are_triples (__main__.TestMipro.test_mipro_history_keys_are_triples) ... ok
test_mipro_improves_on_base (__main__.TestMipro.test_mipro_improves_on_base) ... ok
test_mipro_result_shape (__main__.TestMipro.test_mipro_result_shape) ... ok
test_mipro_rollout_count_formula (__main__.TestMipro.test_mipro_rollout_count_formula) ... ok
test_mipro_score_is_a_dev_score (__main__.TestMipro.test_mipro_score_is_a_dev_score) ... ok
test_mipro_surrogate_flag_same_budget (__main__.TestMipro.test_mipro_surrogate_flag_same_budget) ... ok
test_mipro_ungrounded_skips_failure_summary (__main__.TestMipro.test_mipro_ungrounded_skips_failure_summary) ... ok
test_proposals_program_aware_use_only_own_rules (__main__.TestMipro.test_proposals_program_aware_use_only_own_rules) ... ok
test_proposals_tip_adds_flavor (__main__.TestMipro.test_proposals_tip_adds_flavor) ... ok
test_proposals_ungrounded_can_use_other_words (__main__.TestMipro.test_proposals_ungrounded_can_use_other_words) ... ok
test_proposals_unique_and_count (__main__.TestMipro.test_proposals_unique_and_count) ... ok
test_proposals_weighted_by_failure (__main__.TestMipro.test_proposals_weighted_by_failure) ... ok
test_bimodal_teacher_normalized_two_modes (__main__.TestModeSeeking.test_bimodal_teacher_normalized_two_modes) ... ok
test_both_directions_are_fit_of_the_same_family (__main__.TestModeSeeking.test_both_directions_are_fit_of_the_same_family) ... ok
test_discretized_gaussian_normalized (__main__.TestModeSeeking.test_discretized_gaussian_normalized) ... ok
test_forward_kl_covers_both (__main__.TestModeSeeking.test_forward_kl_covers_both) ... ok
test_reverse_kl_picks_one_mode (__main__.TestModeSeeking.test_reverse_kl_picks_one_mode) ... ok
test_reverse_kl_picks_other_mode_from_other_start (__main__.TestModeSeeking.test_reverse_kl_picks_other_mode_from_other_start) ... ok
test_loop_deterministic (__main__.TestOpro.test_loop_deterministic) ... ok
test_loop_history_unique_and_seeded (__main__.TestOpro.test_loop_history_unique_and_seeded) ... ok
test_loop_returns_best_seen (__main__.TestOpro.test_loop_returns_best_seen) ... ok
test_loop_rollouts_are_scoring_set_times_instructions (__main__.TestOpro.test_loop_rollouts_are_scoring_set_times_instructions) ... ok
test_parse_empty (__main__.TestOpro.test_parse_empty) ... ok
test_prompt_is_ascending (__main__.TestOpro.test_prompt_is_ascending) ... ok
test_prompt_keeps_top_k (__main__.TestOpro.test_prompt_keeps_top_k) ... ok
test_prompt_roundtrip_scores_two_decimals (__main__.TestOpro.test_prompt_roundtrip_scores_two_decimals) ... ok
test_proposals_count_and_cap (__main__.TestOpro.test_proposals_count_and_cap) ... ok
test_proposals_start_from_best_words (__main__.TestOpro.test_proposals_start_from_best_words) ... ok
test_scores_are_on_scoring_set_not_exact (__main__.TestOpro.test_scores_are_on_scoring_set_not_exact) ... ok
test_uniform_program_sets_both_modules (__main__.TestOpro.test_uniform_program_sets_both_modules) ... ok
test_returned_instruction_is_the_argmax (__main__.TestOproBest.test_returned_instruction_is_the_argmax) ... ok
test_advantage_subtracts_critic (__main__.TestPPO.test_advantage_subtracts_critic) ... ok
test_clip_flows_table (__main__.TestPPO.test_clip_flows_table) ... ok
test_critic_learns_values (__main__.TestPPO.test_critic_learns_values) ... ok
test_gradient_flows_inside_clip (__main__.TestPPO.test_gradient_flows_inside_clip) ... ok
test_gradient_zero_when_ratio_far_outside_clip_for_positive_advantage (__main__.TestPPO.test_gradient_zero_when_ratio_far_outside_clip_for_positive_advantage) ... ok
test_kl_penalty_moves_returns_after_policy_moves (__main__.TestPPO.test_kl_penalty_moves_returns_after_policy_moves) ... ok
test_kl_penalty_slows_drift (__main__.TestPPO.test_kl_penalty_slows_drift) ... ok
test_kl_penalty_zero_at_reference (__main__.TestPPO.test_kl_penalty_zero_at_reference) ... ok
test_objective_equals_mean_advantage_at_ratio_one (__main__.TestPPO.test_objective_equals_mean_advantage_at_ratio_one) ... ok
test_objective_gradient_matches_finite_difference (__main__.TestPPO.test_objective_gradient_matches_finite_difference) ... ok
test_ratio_one_at_collection (__main__.TestPPO.test_ratio_one_at_collection) ... ok
test_returns_are_returns_to_go_with_kl (__main__.TestPPO.test_returns_are_returns_to_go_with_kl) ... ok
test_returns_include_terminal_reward (__main__.TestPPO.test_returns_include_terminal_reward) ... ok
test_training_deterministic (__main__.TestPPO.test_training_deterministic) ... ok
test_training_improves (__main__.TestPPO.test_training_improves) ... ok
test_apply_grad (__main__.TestPolicy.test_apply_grad) ... ok
test_copy_is_independent (__main__.TestPolicy.test_copy_is_independent) ... ok
test_entries_count (__main__.TestPolicy.test_entries_count) ... ok
test_grad_norm (__main__.TestPolicy.test_grad_norm) ... ok
test_logp_is_sum_of_tokens (__main__.TestPolicy.test_logp_is_sum_of_tokens) ... ok
test_sample_deterministic_and_valid (__main__.TestPolicy.test_sample_deterministic_and_valid) ... ok
test_score_function_gradient_matches_finite_difference (__main__.TestPolicy.test_score_function_gradient_matches_finite_difference) ... ok
test_score_gradient_sums_to_zero (__main__.TestPolicy.test_score_gradient_sums_to_zero) ... ok
test_sequence_probabilities_sum_to_one (__main__.TestPolicy.test_sequence_probabilities_sum_to_one) ... ok
test_softmax_stable_for_large_logits (__main__.TestPolicy.test_softmax_stable_for_large_logits) ... ok
test_softmax_sums_to_one (__main__.TestPolicy.test_softmax_sums_to_one) ... ok
test_uniform_start (__main__.TestPolicy.test_uniform_start) ... ok
test_all_zero_rewards_give_zero_gradient (__main__.TestReinforce.test_all_zero_rewards_give_zero_gradient) ... ok
test_baseline_changes_nothing_in_expectation (__main__.TestReinforce.test_baseline_changes_nothing_in_expectation) ... ok
test_baseline_reduces_variance (__main__.TestReinforce.test_baseline_reduces_variance) ... ok
test_gradient_equals_exact_gradient_with_baseline (__main__.TestReinforce.test_gradient_equals_exact_gradient_with_baseline) ... ok
test_gradient_equals_exact_gradient_without_baseline (__main__.TestReinforce.test_gradient_equals_exact_gradient_without_baseline) ... ok
test_no_baseline_zero_reward_zero_gradient (__main__.TestReinforce.test_no_baseline_zero_reward_zero_gradient) ... ok
test_positive_reward_raises_token_probability (__main__.TestReinforce.test_positive_reward_raises_token_probability) ... ok
test_training_deterministic (__main__.TestReinforce.test_training_deterministic) ... ok
test_training_improves (__main__.TestReinforce.test_training_improves) ... ok
test_empty_filter_result_is_empty_list (__main__.TestSFT.test_empty_filter_result_is_empty_list) ... ok
test_filter_keeps_reward_at_threshold (__main__.TestSFT.test_filter_keeps_reward_at_threshold) ... ok
test_filtered_sft_beats_teacher_and_unfiltered (__main__.TestSFT.test_filtered_sft_beats_teacher_and_unfiltered) ... ok
test_gradient_matches_finite_difference (__main__.TestSFT.test_gradient_matches_finite_difference) ... ok
test_sft_unfiltered_copies_teacher_roughly (__main__.TestSFT.test_sft_unfiltered_copies_teacher_roughly) ... ok
test_training_reproduces_the_data (__main__.TestSFT.test_training_reproduces_the_data) ... ok
test_count_tokens (__main__.TestShared.test_count_tokens) ... ok
test_mean_empty (__main__.TestShared.test_mean_empty) ... ok
test_meter_cap_raises (__main__.TestShared.test_meter_cap_raises) ... ok
test_meter_counts (__main__.TestShared.test_meter_counts) ... ok
test_meter_exact_cap_ok (__main__.TestShared.test_meter_exact_cap_ok) ... ok
test_meter_left (__main__.TestShared.test_meter_left) ... ok
test_stdev (__main__.TestShared.test_stdev) ... ok
test_backward_respects_requires_grad (__main__.TestTextGrad.test_backward_respects_requires_grad) ... ok
test_backward_visits_reverse_order (__main__.TestTextGrad.test_backward_visits_reverse_order) ... ok
test_backward_writes_grad (__main__.TestTextGrad.test_backward_writes_grad) ... ok
test_critic_orders_by_count (__main__.TestTextGrad.test_critic_orders_by_count) ... ok
test_critic_reads_loss (__main__.TestTextGrad.test_critic_reads_loss) ... ok
test_critic_uses_downstream_gradient (__main__.TestTextGrad.test_critic_uses_downstream_gradient) ... ok
test_frozen_module_not_edited (__main__.TestTextGrad.test_frozen_module_not_edited) ... ok
test_gate_never_accepts_a_worse_minibatch (__main__.TestTextGrad.test_gate_never_accepts_a_worse_minibatch) ... ok
test_loss_text_only_failures (__main__.TestTextGrad.test_loss_text_only_failures) ... ok
test_optimize_deterministic (__main__.TestTextGrad.test_optimize_deterministic) ... ok
test_optimize_gets_better (__main__.TestTextGrad.test_optimize_gets_better) ... ok
test_optimize_rollout_accounting (__main__.TestTextGrad.test_optimize_rollout_accounting) ... ok
test_step_adds_top_word_and_clears (__main__.TestTextGrad.test_step_adds_top_word_and_clears) ... ok
test_step_frozen_variable_unchanged (__main__.TestTextGrad.test_step_frozen_variable_unchanged) ... ok
test_step_skips_words_already_present (__main__.TestTextGrad.test_step_skips_words_already_present) ... ok
test_without_gate_accepts_all_changes (__main__.TestTextGrad.test_without_gate_accepts_all_changes) ... ok
test_gepa_accepts_a_strict_gain (__main__.TestTies.test_gepa_accepts_a_strict_gain) ... ok
test_gepa_rejects_a_loss (__main__.TestTies.test_gepa_rejects_a_loss) ... ok
test_gepa_rejects_a_tie (__main__.TestTies.test_gepa_rejects_a_tie) ... ok
test_textgrad_accepts_a_tie (__main__.TestTies.test_textgrad_accepts_a_tie) ... ok
test_textgrad_rejects_a_worse_step (__main__.TestTies.test_textgrad_rejects_a_worse_step) ... ok
test_bad_demo_penalty (__main__.TestToyWorld.test_bad_demo_penalty) ... ok
test_bad_demo_penalty_capped (__main__.TestToyWorld.test_bad_demo_penalty_capped) ... ok
test_base_probability (__main__.TestToyWorld.test_base_probability) ... ok
test_best_is_best_by_brute_force (__main__.TestToyWorld.test_best_is_best_by_brute_force) ... ok
test_best_possible_program (__main__.TestToyWorld.test_best_possible_program) ... ok
test_clamped_low_and_high (__main__.TestToyWorld.test_clamped_low_and_high) ... ok
test_conflict_abbrev_hurts_ties (__main__.TestToyWorld.test_conflict_abbrev_hurts_ties) ... ok
test_conflict_cancelled_by_own_rule (__main__.TestToyWorld.test_conflict_cancelled_by_own_rule) ... ok
test_conflict_negatives_hurts_dates (__main__.TestToyWorld.test_conflict_negatives_hurts_dates) ... ok
test_dataset_bad_weights (__main__.TestToyWorld.test_dataset_bad_weights) ... ok
test_dataset_deterministic (__main__.TestToyWorld.test_dataset_deterministic) ... ok
test_dataset_kinds_and_answers (__main__.TestToyWorld.test_dataset_kinds_and_answers) ... ok
test_dataset_size_and_ids (__main__.TestToyWorld.test_dataset_size_and_ids) ... ok
test_dataset_weights (__main__.TestToyWorld.test_dataset_weights) ... ok
test_demo_bloat_penalty (__main__.TestToyWorld.test_demo_bloat_penalty) ... ok
test_demo_gain_capped_at_two (__main__.TestToyWorld.test_demo_gain_capped_at_two) ... ok
test_demo_of_other_kind_ignored (__main__.TestToyWorld.test_demo_of_other_kind_ignored) ... ok
test_empirical_matches_p_correct (__main__.TestToyWorld.test_empirical_matches_p_correct) ... ok
test_evaluate_range (__main__.TestToyWorld.test_evaluate_range) ... ok
test_exact_accuracy_base (__main__.TestToyWorld.test_exact_accuracy_base) ... ok
test_exact_accuracy_weights (__main__.TestToyWorld.test_exact_accuracy_weights) ... ok
test_feedback_names_missing_rule (__main__.TestToyWorld.test_feedback_names_missing_rule) ... ok
test_feedback_ok_when_correct (__main__.TestToyWorld.test_feedback_ok_when_correct) ... ok
test_feedback_says_noise_when_rule_present (__main__.TestToyWorld.test_feedback_says_noise_when_rule_present) ... ok
test_instruction_gain (__main__.TestToyWorld.test_instruction_gain) ... ok
test_instruction_in_wrong_module_does_not_help (__main__.TestToyWorld.test_instruction_in_wrong_module_does_not_help) ... ok
test_length_penalty (__main__.TestToyWorld.test_length_penalty) ... ok
test_no_penalty_up_to_free_words (__main__.TestToyWorld.test_no_penalty_up_to_free_words) ... ok
test_per_kind_accuracy_keys (__main__.TestToyWorld.test_per_kind_accuracy_keys) ... ok
test_program_is_immutable_copy (__main__.TestToyWorld.test_program_is_immutable_copy) ... ok
test_rollout_charges_meter (__main__.TestToyWorld.test_rollout_charges_meter) ... ok
test_rollout_deterministic (__main__.TestToyWorld.test_rollout_deterministic) ... ok
test_scores_on_length_and_values (__main__.TestToyWorld.test_scores_on_length_and_values) ... ok
test_with_instruction_unknown_module (__main__.TestToyWorld.test_with_instruction_unknown_module) ... ok
test_wrong_output_differs_from_answer (__main__.TestToyWorld.test_wrong_output_differs_from_answer) ... ok
test_bonus_prefers_less_tried_when_means_equal (__main__.TestUCB.test_bonus_prefers_less_tried_when_means_equal) ... ok
test_explore_rate (__main__.TestUCB.test_explore_rate) ... ok
test_mipro_belief_counts_trials (__main__.TestUCB.test_mipro_belief_counts_trials) ... ok
test_mipro_final_is_best_full_eval (__main__.TestUCB.test_mipro_final_is_best_full_eval) ... ok
test_mipro_full_eval_on_last_trial (__main__.TestUCB.test_mipro_full_eval_on_last_trial) ... ok
test_mipro_full_eval_top_two (__main__.TestUCB.test_mipro_full_eval_top_two) ... ok
test_prefers_better_mean (__main__.TestUCB.test_prefers_better_mean) ... ok
test_surrogate_off_is_uniform (__main__.TestUCB.test_surrogate_off_is_uniform) ... ok
test_untried_option_has_prior_half (__main__.TestUCB.test_untried_option_has_prior_half) ... ok

----------------------------------------------------------------------
Ran 288 tests in 5.573s

OK

5. Discussion questions

Use these in a study group or a class. Try to answer with an example and a source. Give each about 5 to 10 minutes.

1. Pick an agent that you have built or used. Do you have the data and the eval that section 2.1 asks for? If not, write down the first ten examples and the metric you would use. What would the metric miss [32]?
2. OPRO's best instruction was chosen by test accuracy among several. How would you set up a fair test of a prompt optimizer, with a search set, a selection set and a final test set? How many examples does each need, given the noise estimate of 2.2 [5]?
3. MIPROv2 grounds its proposer with a data summary, a program summary, good traces and a tip. For a task of your choice, write what each of the four would say. Which would you drop first to save cost, and what would you expect to lose [2]?
4. TextGrad and GEPA both turn text feedback into edits. Where does the feedback come from in your task (a test, a judge model, a human)? What would a misleading critique do to each method [6] [1]?
5. GEPA's authors include the instructors and the authors of its main baseline. List three checks that would raise your trust in the comparison with GRPO, and say which ones the paper already has [1].
6. You can fine-tune a small model on 5,000 filtered rollouts from a stronger agent. What could go wrong with the data (the filter, the coverage, the states the teacher visits)? What would you check before you train [17] [12]?
7. RL needs a verifiable reward. Give two agent tasks where you have one and two where you do not. For the second pair, what proxy would you use, and how might an optimizer exploit it [32] [3]?
8. Compare PPO and GRPO for a team with a small GPU budget. What does dropping the value model save? What does it cost in variance, and what does the group size G do about it [3] [4]?
9. On-policy distillation claims about 10 times less compute than RL in the vendor's own accounting [25]. What would you measure to check that on your task? List the costs the vendor's accounting might leave out, such as teacher inference and data preparation.
10. The Darwin Gödel Machine reports hallucinated tool use and one case of objective hacking [26]. For an agent that edits its own code, what sandbox, what review and what logging would you require before you let it run overnight [29] [30]?
11. A team has a prompt that grew to 6,000 tokens through repeated optimization. Make the case for moving the knowledge into weights, and the case against. What numbers would settle it [32] [1]?
12. BetterTogether alternates prompt and weight optimization [14]. Sketch a schedule for your own system, with what you would measure after each step. At what point would you stop?

6. Quiz (10 questions)

Pick one answer for each question, then press Check answers. Check answers 0 of 10 answered.

1. According to the lecture's conditions for optimization, which situation is the best reason to start an optimization project on an agent system [32]?
   • The model is brand new and you expect that any optimizer will improve it, so you plan to build the metric later, once the search has produced a few promising candidates to look at
   • You have no labeled data, but you expect the optimizer to invent a good metric on its own from the outputs of the model that it is tuning, so you need neither data nor evals
   • You have data and an eval for your domain, and examples or feedback can express a requirement that is hard to put into a prompt
   • The task and its requirements change every week, so a large optimization run each week is a cheap way to keep the system up to date with those changes
2. OPRO's paper reports an instruction with 80.2% test accuracy on GSM8K against 71.8% for "Let's think step by step." What is the best reading of that table [5]?
   • It is an unbiased estimate of how much any OPRO-found instruction will improve accuracy on a new task, because the test set was never used to guide the search in any way
   • It shows that the optimizer understood why "step by step" works, since the meta-prompt contains an explanation of each score alongside the instruction that earned that score
   • It shows that a language-model optimizer always finds a better instruction than a human, because the 80.2% figure was the average over all instructions found in the 200 steps of the run
   • Table 4 lists the instruction with the highest test accuracy for each scorer and optimizer pair, so it is a best-of-several, and a fresh set is needed for an unbiased estimate
3. Which list gives the three steps of MIPROv2 in the right order [2]?
   • Search combinations of instructions and demos with minibatch scores, then propose instructions from a summary of the data, and last bootstrap demos from passing traces of the program
   • Bootstrap demos from traces that pass the metric, propose grounded instructions for each module, then search combinations using minibatch scores and periodic full evaluation
   • Fine-tune each module on its own labels, merge the module weights into one model, and then evaluate the merged model on a validation set to choose between merged candidates
   • Ask the proposer for one instruction, compute its gradient with respect to the module's weights, and then update the instruction by gradient descent until the metric stops improving
4. This question uses my arithmetic rule for GEPA's selection (a candidate is dropped if every task it wins is also won by a surviving candidate). Three candidates score on four tasks: A = (1, 0, 1, 1), B = (0, 1, 0, 0), C = (1, 0, 1, 0). Which statement is right [1]?
   • A and B survive, because C's wins are all shared with A, and A is sampled about three times as often as B since it wins three tasks and B wins one
   • All three survive, because each one reaches the best score on at least one task, and they are then sampled with equal probability since each one has a unique profile
   • Only A survives, because it has the best mean score of the three candidates, and the other two candidates are removed from the pool before the next iteration begins
   • B and C survive, because they have different scores from each other and from A, and A is removed because it is the candidate that was the parent of the other two
5. A prompt gets four sampled answers with rewards (1, 0, 0, 0). Using the group mean and the population standard deviation, what are the GRPO advantages [3]?
   • +0.75 for the passing answer and −0.25 for each failing answer, since the standard deviation is not used when there is exactly one passing answer in the group
   • +3.0 for the passing answer and −1.0 for each failing answer, since the advantage is the reward divided by the mean of the group's rewards, which is 0.25 here
   • +1.0 for the passing answer and 0.0 for each failing answer, since only a passing answer gives a gradient signal and a failing one is simply ignored in GRPO
   • About +1.73 for the passing answer and −0.58 for each failing answer, since the mean is 0.25 and the standard deviation is about 0.433
6. What is the main difference between PPO and GRPO as the lecture and DeepSeekMath describe them [3] [32]?
   • PPO uses a trained value model as the baseline and puts KL in the reward, while GRPO uses the group's mean and spread and puts KL in the loss
   • GRPO trains a larger value model than PPO does, so it needs more memory, and in return it places the KL penalty inside the reward at every token, as PPO does not
   • PPO never clips the probability ratio, while GRPO clips it, and for that reason GRPO is the only one of the two methods that can be used to train a language model
   • PPO needs a verifiable reward and GRPO does not, since GRPO learns its own reward from the group of answers, which means it can run with no checker for the answers
7. PPO with ε = 0.2, an advantage of −1 and a probability ratio of 0.7. What happens to this token's term [4]?
   • The term is −0.7 and the gradient flows, because the ratio is below 1 and a ratio below 1 is always inside the range in which the clipped objective passes gradient
   • The term is −1.2, because the ratio is raised to 1 plus epsilon before it is multiplied by the advantage, which means the gradient is large and points downward
   • The term is clipped to −0.8, so this token gives no gradient, because the bad token has already been made less likely by more than the clip range
   • The term is +0.8, because a negative advantage flips the sign of the clipped ratio, so the update raises the probability of the token that got the negative advantage
8. In on-policy distillation as the lecture and the Thinking Machines post describe it, what is the per-token reward for a token that the student sampled [25] [32]?
   • One pass-or-fail score for the whole answer, copied to every token, from a verifier that checks the final answer, exactly as in GRPO, with the teacher not involved at all
   • The negative of (log probability under the student minus log probability under the teacher) at that token, which is a sampled estimate of the reverse KL
   • The log probability of that token under the teacher alone, with the student's own probability ignored, so the student simply copies whatever the teacher finds most likely
   • The cross-entropy between the student and a human-written reference answer at that token, which is the same signal as in ordinary SFT and uses no teacher model either
9. The DGM weights an agent by a sigmoid of its score (λ = 10, α₀ = 0.5) divided by 1 plus its number of children. By my arithmetic, agent X has score 0.9 and 3 children, agent Y has 0.6 and none, and agent Z has 0.2 and none. Which is most likely to be chosen as the parent [26]?
   • Agent X, at about 75%, because the weight always rises with the score and the number of children plays no role in the choice of parent in the paper's rule
   • Agent Z, because the rule favors the weakest agent so as to explore, and a low score gives the highest sampling weight in the DGM's parent selection rule
   • All three with equal probability, because the archive keeps every working agent and the parent is sampled uniformly from the archive in each round of the loop
   • Agent Y, at about 71%, because X's weight (about 0.25) is cut by its three children, and Y's weight (about 0.73) is much higher than Z's (about 0.05)
10. Which situation points most clearly toward weight optimization and not toward prompt optimization, in the lecture's list of reasons [32]?
    • You can only call an API model, you need to rerun the search many times a day, and every interaction with the environment is slow and costly to produce
    • You have a small number of examples and a metric, and you want a fast result that you can review in text, with no GPUs and no access to model weights
    • A strong teacher exists and you want a cheaper student, or the optimized prompt has grown so long that it is slow and costly on every call
    • You want the change to be easy to read and to roll back, so that a person can review the new instructions line by line before they are put into use

Answer key

7. Graded assignment and rubric

Title: Optimize, measure, break and defend an agent system.
Time: 4 to 5 hours. Total: 100 points. Work: solo.

The task

Extend the lab code and write a short memo (three pages at most). Mock models and toy tasks are fine. You may use a real model if you have access, but say clearly which parts used it. Build your own small task for each part, and say how you made it. Do not reuse the lab's toy task or its numbers. Label every number that comes from a toy model or from made-up data. Use separate data for searching, for choosing and for the final report, and say how big each set is [5] [1].

Part A. Context optimization on your own task (25 points)

• Choose a small task that is not the lab's (for example, extracting dates from short notes, or labeling 6 kinds of support messages). Build a mock scorer with at least 30 examples, and split them three ways: a search set, a selection set and a final test set [2].
• Run bootstrapped few-shot, a MIPRO-like search, an OPRO loop and a GEPA-style loop on it, with at least 5 seeds, and with the same rollout budget for each. Report the mean and the range of the final test score, and the rollouts used [1] [5] [9].
• Show the selection effect from section 2.2 with your own data: the gap between the best candidate's score on the selection set and on the final test set. Say how big the gap is, and what it means for a reported "best" score [5].
• Change the feedback of the TextGrad-style critic so that it is wrong half of the time. Report what happens, and say what a real system would do about a misleading critique [6].
• Add one new proposer or one new selection rule, with at least five unit tests (one for an edge case). Say which assumption drives its result.

Part B. Weight optimization and gradient checks (25 points)

• Make your own toy environment (a bandit with 5 actions, or a small chain with a trap) and train a policy with REINFORCE, PPO-clip and GRPO. Give every method the same number of rollouts and at least 5 seeds. Do not rank the methods on one run [3] [4].
• Compute the GRPO advantages by hand for two groups of your own (one with a single passing answer, one with half passing), as in 2.11. Show that the code agrees. State which standard deviation you used, and what the Dr. GRPO variant would change [3] [21].
• For PPO, make a table of the ratio, the advantage, the clipped term and the gradient (yes or no) for at least six cases. Include both signs of the advantage and ratios inside and outside the clip range [4].
• Run SFT on filtered rollouts and compare the filter at two settings (strict and loose). Then run reverse-KL distillation from a teacher that is right about 50% of the time and one that is right about 80% of the time. Report what you see, and say why a noisy teacher matters [17] [24] [25].
• List what your toy cannot show about real RL on language models. Name three real failure modes from the readings, such as length bias, difficulty bias and the Echo Trap [21] [22].

Part C. Agents as optimizers, with guards (25 points)

• Design your own variant landscape for the DGM-style archive: at least 6 features, at least one stepping stone (a feature that scores badly alone and helps in combination). Compare the archive with a greedy loop on at least 10 seeds [26].
• Change the parent-selection rule in two ways (for example, a different λ, and no child-count penalty). Compute the selection probabilities by hand for three agents, as in 2.11(c), and check the code [26].
• Write an autoresearch loop for your own toy metric. Write three cheats: an edit to the instructions file, a one-character edit to the eval file, and one cheat that the hash guard cannot see. Show that the guard blocks the first two and misses the third [29] [30].
• For the third cheat, design a second control (a hidden hold-out check, a second metric, or a human review step) and show that it catches the cheat. Say what it costs [30].
• In four sentences, say what you would require before letting an agent edit its own code overnight [26] [29].

Part D. Prompts or weights: a cost study and a memo (25 points)

• Use CostInputs with your own made-up numbers for one agent that makes many calls. Say why each number is plausible, and mark all of them as made up. Find the cheapest plan at three call volumes, and the break-even number of calls between at least two plans [32].
• Do a sensitivity check: change each input by 50% in each direction and show which input moves the break-even most. Say which input you are least sure of.
• Replace at least three inputs with real quotes from a model provider and a finetuning service (give the date and the source). Re-run the study and say what changed.
• Write a two-page memo for this decision: a company has a support agent with a 6,000-token prompt that grew over many prompt-optimization rounds. Should it move to weights? Give your choice, a cost estimate, a test plan with held-out data, the risks, and a rollback plan [14] [32].
• Name three limits of your own experiments in the whole assignment.

What to hand in

• Your optimization_toolkit.py (or your own file names), your tasks and data, and the tests. All tests must pass.
• A memo in PDF or Markdown with the tables from Parts A to C and the memo from Part D.
• A sources list for any reading you cite.

Rubric

Criterion	Full marks	Partial	Points
A. Context optimization	Own task with three data splits. Four methods with equal budgets, at least 5 seeds, a mean and a range. The selection gap shown with data. A misleading-critic test. One new rule with five or more passing tests.	Two splits only, or one seed, or unequal budgets, or no selection-gap measurement, or fewer than five tests.	25
B. Weight optimization	Three RL methods on an own toy with equal rollouts and seeds. Hand-computed GRPO advantages that match the code. A PPO clip table with six cases. SFT filter and distillation teacher comparisons. A clear list of what the toy cannot show.	One seed, or no hand arithmetic, or fewer than six clip cases, or no teacher comparison, or numbers not labelled as toy.	25
C. Agents as optimizers	Own landscape with a stepping stone, an archive-versus-greedy comparison with 10 seeds, selection probabilities by hand, three cheats with a guard that blocks two and misses one, and a second control that catches it.	Fewer than 10 seeds, or no hand arithmetic, or the missed cheat not shown, or no second control.	25
D. Cost study and memo	Made-up inputs labelled as such. Three volumes and a break-even. A sensitivity check. Three real quotes with sources. A memo with a clear choice, a cost estimate, a test plan on held-out data, risks and a rollback plan. Three limits.	No sensitivity check, or no real quotes, or no test plan, or no rollback plan, or inputs not labelled as made up.	25
Total			100

Grade bands: 90 to 100 excellent. 75 to 89 good. 60 to 74 needs work. Under 60: redo and resubmit. Plain writing is rewarded: short sentences, no filler. Label your own arithmetic as your own.

8. Glossary

Term	Meaning
Advantage	How much better an action or an answer was than the baseline. In PPO it comes from a value model and GAE. In GRPO it is the group-normalized reward [3] [4].
Archive (DGM)	The growing set of agent variants that the Darwin Gödel Machine keeps and samples parents from, including variants that are not the best [26].
Autoresearch	Karpathy's small repository in which an agent edits one training file, runs a fixed five-minute budget, and keeps a change if the validation metric improves [29].
Baseline (RL)	A value subtracted from the reward to lower the variance of the gradient. PPO learns one with a value model. GRPO uses the group mean [3].
Best-of-several (selection bias)	The upward bias in a score that was picked as the highest among many candidates. A fresh test set removes it [5].
Bootstrapped few-shot	Running a program on training examples and keeping the traces that pass the metric as demonstrations for the prompt [9] [2].
Clip range (ε)	The limit in PPO on how far the probability ratio can move before the objective stops giving a gradient. A value of 0.2 is common [4].
Context optimization	Improving a system by changing its inputs (instructions, demos, tool text, memory) and not its weights [32].
Dominated candidate	A candidate whose best-scoring tasks are all also won by another surviving candidate. GEPA's Pareto step drops such candidates [1].
Echo Trap	RAGEN's name for a failure in multi-turn agent RL, with reward collapse and gradient spikes during training [22].
Forward KL	KL(teacher || student). Minimizing it is the same as SFT on teacher samples. It tends to spread the student over all of the teacher's modes [24] [25].
GEPA	Genetic-Pareto prompt optimizer: reflect on traces in natural language, propose an edit, and keep a Pareto front of candidates [1].
GKD	Generalized knowledge distillation: train a student on its own samples with feedback from the teacher, with a choice of divergence [24].
Goodhart's law	When a measure becomes a target, it stops being a good measure. An optimizer pushes on the metric, not on what you meant [26].
GRPO	Group Relative Policy Optimization. Samples a group of answers per prompt and uses their normalized rewards as advantages, with no value model [3].
Hold-out set	Data that the optimizer never sees, used for the final score. A search set and a selection set are separate from it [1] [2].
KL penalty	A term that keeps the trained policy close to a reference policy. PPO puts it in the reward. GRPO adds it to the loss [3].
LoRA	Low-rank adaptation: train small added matrices and freeze the original weights, which cuts trainable parameters and memory [16].
Meta-prompt (OPRO)	The prompt shown to the optimizer model. It holds past instructions with their scores, in ascending order, and some task examples [5].
MIPROv2	An optimizer for multi-module programs that bootstraps demos, proposes instructions, and searches combinations with a surrogate model [2].
Minibatch evaluation	Scoring a candidate on a small sample to save cost, with a full evaluation now and then to check the best ones [2].
Objective hacking	When an optimizer raises the score by exploiting the metric (for example, by faking a tool log) and not by doing the work. The DGM paper reports an example [26].
On-policy distillation	Distillation in which the student samples and the teacher grades each token with a reverse-KL reward [25] [24].
OPRO	Optimization by PROmpting. A language model writes new candidates from a history of scored candidates [5].
Pareto front	The set of candidates that are not dominated. In GEPA, it is built per task and the parent is sampled in proportion to the number of tasks won [1].
PPO	Proximal Policy Optimization. A policy-gradient method that clips the probability ratio to keep updates small [4].
Probability ratio	The probability of an action under the new policy divided by its probability under the old policy [4].
Reflective mutation	GEPA's step: a model reads traces and feedback in words and writes an improved prompt for one module [1].
Rejection sampling	Generating many outputs and keeping only those that pass a check. It is the base of STaR and of SFT on filtered rollouts [12].
Reverse KL	KL(student || teacher). It is mode seeking: the student is pushed to copy one behavior of the teacher well [25].
Rollout	One run of the system on one input, used to get a score or a trace. It is the unit of cost in these papers [1].
SFT	Supervised fine-tuning: train on example outputs with cross-entropy. For agents, the examples are filtered traces [17] [18].
TextGrad	A method that treats text as variables and uses language-model feedback as a "gradient" that is passed back through the compute graph [6] [7].
TPE surrogate	The Tree-structured Parzen Estimator, a Bayesian optimization model that MIPROv2 uses to choose the next combination to test [2].
Value model (critic)	A learned model that predicts the expected return, used as the baseline in PPO. GRPO removes it [3] [4].
Verifiable reward	A reward that a program can compute and check, such as a test result or an exact answer. RL on language models depends on it [3] [20].

9. References

Entries marked (required) are the course's required readings. Entries marked (added) are my own additions. The course site and slides, [31] and [32], are the course's own pages. "Abstract level" means that I read the abstract and not the whole paper. Where an entry names more, that is what I read.

• [1] Agrawal, L. A., Tan, S., Soylu, D., Ziems, N., Khare, R., Opsahl-Ong, K., Singhvi, A., Shandilya, H., Ryan, M. J., Jiang, M., Potts, C., Sen, K., Dimakis, A. G., Stoica, I., Klein, D., Zaharia, M., and Khattab, O. "GEPA: Reflective Prompt Evolution Can Outperform Reinforcement Learning." ICLR 2026 (oral). arXiv:2507.19457. One co-author, Michael J. Ryan, teaches the course. https://arxiv.org/abs/2507.19457 (I read the abstract, the method with Algorithms 1 and 2, and Tables 1 and 2) (required)
• [2] Opsahl-Ong, K., Ryan, M. J., Purtell, J., Broman, D., Potts, C., Zaharia, M., and Khattab, O. "Optimizing Instructions and Demonstrations for Multi-Stage Language Model Programs" (MIPROv2). EMNLP 2024. arXiv:2406.11695 (version 2, October 2024). One co-first author, Michael J. Ryan, teaches the course. https://arxiv.org/abs/2406.11695 (I read the abstract, the problem statement, the three steps and the main results) (required)
• [3] Shao, Z., Wang, P., Zhu, Q., Xu, R., Song, J., Bi, X., Zhang, H., Zhang, M., Li, Y. K., Wu, Y., and Guo, D. "DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models." arXiv:2402.03300 (version 3, April 2024). Vendor-authored (DeepSeek-AI). https://arxiv.org/abs/2402.03300 (I read the abstract and the section on reinforcement learning with GRPO) (required)
• [4] Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O. "Proximal Policy Optimization Algorithms." arXiv:1707.06347 (2017). https://arxiv.org/abs/1707.06347 (I read the abstract and the clipped objective) (added)
• [5] Yang, C., Wang, X., Lu, Y., Liu, H., Le, Q. V., Zhou, D., and Chen, X. "Large Language Models as Optimizers" (OPRO). arXiv:2309.03409. Authors from Google DeepMind. https://arxiv.org/abs/2309.03409 (I read the abstract, the meta-prompt description, and Table 4) (added)
• [6] Yuksekgonul, M., Bianchi, F., Boen, J., Liu, S., Lu, P., Huang, Z., Guestrin, C., and Zou, J. "TextGrad: Automatic Differentiation via Text." arXiv:2406.07496 (2024). https://arxiv.org/abs/2406.07496 (I read the abstract and the optimizer description) (added)
• [7] Yuksekgonul, M., Bianchi, F., Boen, J., Liu, S., Lu, P., Huang, Z., Guestrin, C., and Zou, J. "Optimizing generative AI by backpropagating language model feedback." Nature 639, 609 to 616 (2025). The journal version of TextGrad. https://www.nature.com/articles/s41586-025-08661-4 (I checked the journal, the volume, the pages and the authors on the page; I did not read the article in full) (added)
• [8] Zhou, Y., Muresanu, A. I., Han, Z., Paster, K., Pitis, S., Chan, H., and Ba, J. "Large Language Models Are Human-Level Prompt Engineers" (APE). ICLR 2023. arXiv:2211.01910. https://arxiv.org/abs/2211.01910 (abstract level) (added)
• [9] Khattab, O., Singhvi, A., Maheshwari, P., Zhang, Z., Santhanam, K., Vardhamanan, S., Haq, S., Sharma, A., Joshi, T. T., Moazam, H., Miller, H., Zaharia, M., and Potts, C. "DSPy: Compiling Declarative Language Model Calls into Self-Improving Pipelines." arXiv:2310.03714 (version 1, 2023). https://arxiv.org/abs/2310.03714 (abstract level) (added)
• [10] DSPy project. "Optimizers" and "GEPA" pages. DSPy documentation. Project documentation; it changes over time, and I read it in October 2026. The cost-saving statements on the GEPA page are quoted by the project and I did not check them. https://dspy.ai/api/optimizers/GEPA/overview/ (the page names a training set for reflective updates and a validation set for the Pareto scores; the optimizer list is at https://dspy.ai/learn/optimization/optimizers/) (added)
• [11] gepa-ai. "gepa." Open-source code for GEPA. GitHub repository. I checked that it exists and did not read the code. https://github.com/gepa-ai/gepa (added)
• [12] Zelikman, E., Wu, Y., Mu, J., and Goodman, N. D. "STaR: Self-Taught Reasoner: Bootstrapping Reasoning With Reasoning." arXiv:2203.14465 (2022). https://arxiv.org/abs/2203.14465 (abstract level) (added)
• [13] Lu, Y., Bartolo, M., Moore, A., Riedel, S., and Stenetorp, P. "Fantastically Ordered Prompts and Where to Find Them: Overcoming Few-Shot Prompt Order Sensitivity." arXiv:2104.08786 (2021). https://arxiv.org/abs/2104.08786 (abstract level) (added)
• [14] Soylu, D., Potts, C., and Khattab, O. "Fine-Tuning and Prompt Optimization: Two Great Steps that Work Better Together" (BetterTogether). arXiv:2407.10930 (version 2, October 2024). https://arxiv.org/abs/2407.10930 (abstract level) (added)
• [15] Li, S., Raghuram, V. C., Khattab, O., Hirschberg, J., and Yu, Z. "PAPILLON: PrivAcy Preservation from Internet-based and Local Language MOdel ENsembles." arXiv:2410.17127 (version 3, March 2025). https://arxiv.org/abs/2410.17127 (abstract level; cited only for the PUPA benchmark that the GEPA paper uses) (added)
• [16] Hu, E., Shen, Y., Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., Wang, L., and Chen, W. "LoRA: Low-Rank Adaptation of Large Language Models." arXiv:2106.09685 (2021). https://arxiv.org/abs/2106.09685 (abstract level) (added)
• [17] Chen, B., Shu, C., Shareghi, E., Collier, N., Narasimhan, K., and Yao, S. "FireAct: Toward Language Agent Fine-tuning." arXiv:2310.05915 (2023). https://arxiv.org/abs/2310.05915 (I read the abstract and the tables that compare prompting with finetuning) (added)
• [18] Zeng, A., Liu, M., Lu, R., Wang, B., Liu, X., Dong, Y., and Tang, J. "AgentTuning: Enabling Generalized Agent Abilities for LLMs." arXiv:2310.12823 (2023). The claim about AgentLM-70B is the authors' own claim. https://arxiv.org/abs/2310.12823 (abstract level) (added)
• [19] Raoof, N., Zhuang, R., Nezhurina, M., Guha, E., et al. "Data Recipes for Agentic Models." arXiv:2606.24855 (version 1, June 2026). The lecture calls the work OpenThoughts-Agent; the paper's title is the one given here. https://arxiv.org/abs/2606.24855 (abstract level) (added)
• [20] DeepSeek-AI. "DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning." arXiv:2501.12948. Vendor-authored. https://arxiv.org/abs/2501.12948 (abstract level, plus the R1-Zero result in the introduction) (added)
• [21] Liu, Z., Chen, C., Li, W., Qi, P., Pang, T., Du, C., Lee, W. S., and Lin, M. "Understanding R1-Zero-Like Training: A Critical Perspective" (Dr. GRPO). COLM 2025. arXiv:2503.20783. https://arxiv.org/abs/2503.20783 (I read the abstract and the section on the biases in GRPO) (added)
• [22] Wang, Z., Wang, K., Wang, Q., Zhang, P., Li, L., Yang, Z., et al. "RAGEN: Understanding Self-Evolution in LLM Agents via Multi-Turn Reinforcement Learning." arXiv:2504.20073 (version 2, May 2025). https://arxiv.org/abs/2504.20073 (abstract level) (added)
• [23] Zhou, Y., Zanette, A., Pan, J., Levine, S., and Kumar, A. "ArCHer: Training Language Model Agents via Hierarchical Multi-Turn RL." arXiv:2402.19446 (2024). https://arxiv.org/abs/2402.19446 (abstract level) (added)
• [24] Agarwal, R., Vieillard, N., Zhou, Y., Stanczyk, P., Ramos, S., Geist, M., and Bachem, O. "On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes" (GKD). ICLR 2024. arXiv:2306.13649. https://arxiv.org/abs/2306.13649 (abstract level) (added)
• [25] Lu, K., and others at Thinking Machines. "On-Policy Distillation." Thinking Machines Lab blog, October 27, 2025. Vendor-authored: the compute savings (such as 9 to 30 times) and the GPU-hour table are the company's own accounting, and I did not reproduce them. https://thinkingmachines.ai/blog/on-policy-distillation/ (I read the post in full) (added)
• [26] Zhang, J., Hu, S., Lu, C., Lange, R., and Clune, J. "Darwin Gödel Machine: Open-Ended Evolution of Self-Improving Agents." ICLR 2026. arXiv:2505.22954. https://arxiv.org/abs/2505.22954 (I read the abstract, the parent-selection formula, Table 1 and the safety discussion) (added)
• [27] Lee, Y., Nair, R., Zhang, Q., Lee, K., Khattab, O., and Finn, C. "Meta-Harness: End-to-End Optimization of Model Harnesses." arXiv:2603.28052 (version 1, March 2026). https://arxiv.org/abs/2603.28052 (abstract level) (added)
• [28] Novikov, A., Vũ, N., Eisenberger, M., Dupont, E., Huang, P.-S., Wagner, A. Z., et al. "AlphaEvolve: A coding agent for scientific and algorithmic discovery." Google DeepMind white paper, arXiv:2506.13131 (June 2025). Vendor-authored. https://arxiv.org/abs/2506.13131 (abstract level) (added)
• [29] Karpathy, A. "autoresearch." GitHub repository . The repository changes over time, and I read it in October 2026. https://github.com/karpathy/autoresearch (I read the README and the file list) (added)
• [30] Chieng, S., and Cherfa, S. "How to stop your autoresearch loop from cheating." Cerebras blog, March 19, 2026. Vendor-authored: one team's account of 71 experiments, not an independent study. https://www.cerebras.ai/blog/how-to-stop-your-autoresearch-loop-from-cheating (I read the whole post) (added)
• [31] Yang, D., Ryan, M., and Yang, J. "CS329Z: Engineering AI Agents." Stanford University, Fall 2026. Course site and schedule. https://cs329z.stanford.edu/
• [32] Yang, D., Ryan, M., and Yang, J. "Optimization." CS329Z lecture slides (Lecture 9, Week 5, October 21), Stanford University. Used for the topic order and for what the lecture claims; where my page says "the lecture", it means these slides. https://web.stanford.edu/class/cs329z/slides/lecture09.pdf