OPAL: a reactive optimizer with an agent in the loop

TL;DR

OPAL (OPtimizer with an Agent in the Loop) is a PyTorch optimizer wrapper from Bai, Liang, Teoh & Sitzmann (MIT) that lets a coding agent (Claude Code by default; Codex supported) rewrite two small Python programs — updater.py and analysis.py — at scheduled decision points during a training run, based on the losses, gradient statistics, and user-logged metrics the agent is shown. The agent can change the update rule per parameter group, learning-rate schedules, momenta, preconditioning, clipping, and loss-spike guards; it is forbidden from touching the model, data, loss, batch size, step count, or evaluation. Reported wins include SiT-B/2 ImageNet 256 FID 35.7 → 23.6 at 400K updates and LightningDiT-B/1 FID 3.22 → 2.31 at 160 epochs, with Goldfish 125M Swahili perplexity, modded-nanogpt record #36 validation loss, and RAEv2-B all improving over their official recipes. Positioned by the authors as the end of the “pick optimizer + LR schedule before the run and never touch them” default.

Key claims

OPAL splits the optimizer into two agent-editable programs that live in opal_runs/<folder>/workspace/: updater.py turns gradients into parameter updates each step (initialized as a call to the user’s own optimizer and scheduler), and analysis.py measures training statistics the agent chooses — the agent rewrites both at decision points inside a single conversation that persists for the whole run, ending each decision with a rationale of what it saw, changed, and expects [README §How it works].
A new updater.py goes live at the next step on every GPU and runs on a 20-step trial; failures are undone exactly to the previous program if a parameter + optimizer-state snapshot fits in snapshot_limit_gb (default 2 GB) and the failure is on the first step, otherwise the agent is told and iterates from there [README §How it works, “Safe swaps”].
Reported head-to-head gains on five settings over each task’s official recipe, each as a single run (OPAL 0.3 reruns noted): SiT-B/2 ImageNet 256 at 400K updates FID-50K 35.7 → 23.6, RAEv2-B 20 epochs FID 3.39 → 3.10 (efficiency prompt) / 3.19 (OPAL 0.3), Goldfish 125M Swahili FLORES-200 bits 79.37 → 77.60 / 77.25 (OPAL 0.3), modded-nanogpt record #36 validation loss 3.2781 → 3.2749, LightningDiT-B/1 ImageNet 256 FID 3.22 → 2.31 (OPAL 0.3) [README §Results table].
Not always a win: a reported modded-nanogpt run on OPAL 0.2.0 ended at validation loss 3.2878, worse than record #36 at 3.2781, and the authors explicitly flag that the agent can choose differently from one run to the next [README §Results].
Documented emergent behaviors the authors call “superpowers”: detecting that it overshot the loss minimum, resetting momentum and rolling back a few steps; detecting loss spikes, zeroing them out, and watching for them thereafter; designing sophisticated LR schedules of the kind that the authors claim were previously only affordable for DeepSeek-scale runs [tweet thread (2/n)–(4/n); README §How it works “Safe swaps”].
Correctness contract on the un-agented path: with backend="keep" (no agent, no cost), the wrapped run must be bit-for-bit identical to the plain optimizer run step-for-step, and the README specifies the smoke-test protocol (~100 steps each) users should run before paying for an agent [README §Start with your coding agent step 6; §Try it without an agent].
Explicit cost + waiting-time model: opal estimate --total-steps N --sec-per-step S prints both agent-waiting time (also as a multiple of the run’s own training time) and dollar cost; by default a blocking run caps the number of decisions so waiting stays under 4× the run’s own training time, with a 50-decision ceiling, and max_cost_usd provides a hard dollar cap. blocking=False keeps training going while the agent thinks; blocking="auto" only waits when an edit would land too late [README §Options].
The agent sees only: the user’s goal sentence, parameter names/shapes, training losses, its own analysis-program outputs, and read-only copies of the training script + model code (toggleable via share_code=False); user-logged metrics are recorded but hidden by default (shown only via show_agent=True), and held-out validation values arrive via opt.report_dev(...) when validation=N is set [README §How it works “What the agent sees”].
Agent choice: Claude Opus 5.5 at effort="high" is the default; a Codex backend using gpt-6.1-sol is experimental, with the caveat that max_cost_usd does not cap Codex and that the run stops with an error when Codex compacts the conversation (~450K tokens of context on Codex CLI 0.160.1) [README §Experimental Codex backend].

Method

OPAL wraps a user’s existing torch.optim.* optimizer and scheduler in one line: opt = opal.wrap(opt, model, scheduler=sched, total_steps=steps, goal="...", backend="keep"). The user keeps their regular training loop and passes loss=loss into opt.step(...) so the agent can read losses; optional opt.log(...) records metrics for the user, opt.report_dev(...) surfaces held-out dev-set results to the agent when validation=N is set. The default execution model is single-conversation: a persistent agent process (Claude Code by default) is attached to the training job on rank 0 of the DDP (or single-GPU) setup, periodically reads a report from analysis.py and, at scheduled decision times, may edit updater.py and/or analysis.py. Decisions default to a weighted schedule — a fifth of them fall in the first 5% of the run — with schedule="uniform" as the alternative. The agent is explicitly instructed not to touch the model, dataset, loss, batch size, step count, or evaluation, and this is a rule enforced by prompt and workspace layout rather than a sandbox: the agent’s edits run inside the training process. Safe swaps run each new optimizer program on a 20-step trial; a first-step failure with available snapshot triggers exact rollback to the previous program, later failures stay and the agent is informed. Internet / login constraints on GPU machines are handled by running the agent on a login node connected through a shared mailbox folder. Users set max_cost_usd to bound spend; blocking=False or "auto" trade determinism for speed by overlapping agent reasoning with training.

Results

SiT-B/2 ImageNet 256² (400K updates): FID-50K 35.7 (authors’ evaluator on the official recipe) → 23.6 with the reactive optimizer — the largest absolute FID improvement in the five-row table [README §Results].
RAEv2-B ImageNet 256² (20 epochs): FID-50K 3.39 (official) → 3.10 with an efficiency-prompted agent; 3.46 with a neutral prompt (a regression); 3.19 on an OPAL 0.3 rerun. Prompt choice visibly moves the result in both directions [README §Results].
Goldfish 125M GPT on Swahili (FLORES-200 bits): 79.37 (official) → 77.60; OPAL 0.3 rerun 77.25 [README §Results].
modded-nanogpt record #36 (validation loss): 3.2781 (reference record) → 3.2749. The authors also report an earlier OPAL 0.2.0 run at 3.2878, i.e. worse than the record — making this the only row with an explicit filed negative result [README §Results].
LightningDiT-B/1 ImageNet 256² (160 epochs, with guidance): FID-50K 3.22 (official) → 2.31 on an OPAL 0.3 rerun — the second-largest absolute win behind SiT-B/2 [README §Results].
No independent-replication numbers are included; every row is a single run from the authors’ own benchmarking, with weights and run records released at huggingface.co/datasets/snow-leopard/opal-results [README §Results].
No ablation of the agent’s load-bearing decisions is given — i.e. how much of the gain comes from LR schedule rewrites vs. loss-spike guards vs. momentum resets vs. preconditioning changes — only the aggregated end-of-run metric per task. The opal report HTML page is advertised as containing per-decision rationales + code diffs for individual runs.

Why it’s interesting

OPAL is the first filed instance of the AI-for-AI-research pattern aimed at a running training process’s optimizer, placing it between AIDE²: First Evidence of Recursive Self-Improvement (AIDE² optimizes the autoresearch harness across runs) and the architecture/optimizer stability cluster on Training stability at scale (where humans choose a geometric or spectral stability recipe up front, then the optimizer runs it unchanged). The author team overlaps with the SRG group at MIT, and the first author Xingjian Bai is a student of Phillip Isola — making this the Isola group’s second entry in the agent-driven research-automation line alongside Robot-Use Agents. The explicit “end of pick-a-schedule-before-the-run” framing is the sharpest filed instance of the broader trend seen in Delphi: Scaling Laws That Extrapolate 300× Past the Fit and MuonH (hyperball) sets new modded-NanoGPT speedrun record at 3325 steps — Keller Jordan announcing @wen_kaiyue's result (Hyperball) where the stability knob is removed by construction — OPAL goes further by letting a per-run agent decide what the knob should be at every scheduled decision point. The reported superpowers (loss-spike guards, LR-schedule redesign mid-run) address pathologies that papers like A Unified View of Attention and Residual Sinks: Outlier-Driven Rescaling is Essential for Transformer Training, Mitigate Silent Expert Death In Ultra-Sparse MoE, and Cracks in the Foundation: Seemingly Minor Architectural Choices Impact Long Context Extension all target with architectural priors — raising the question of whether a reactive agent at the optimizer layer substitutes for or composes with those architectural interventions, which no filed paper has yet tested.