Home Building an ocean AI scientist
Post
Cancel

Building an ocean AI scientist

Recently I wanted to learn more about ocean emulators, so I decided to build my own. Given my lack of experience working with emulators I decided to see how far I could push using AI to generate the code and develop a self-improving process for improving the output. Over a few days of letting an agent run, I ended up with this emulator of a two-layer shallow water double gyre model. The original model is at the top and the emulator, running freely from the start of a period it never saw in training, is at the bottom. To get to this point the agent ran 134 training runs covering 76 different experiments. In this post I share what I learned in this process.

Surface layer thickness is on the left and vorticity is on the right. The emulator drifts away from the true eddy positions within a few months, as any model of a chaotic flow must, but its eddies keep a realistic size and strength, and the energy in the flow (bottom panel) stays inside the range of the original model.

So what was the goal?

My goal was to build a process that could continually improve some kind of benchmark model. In my case I decided that the model would be an ocean emulator.

Why an emulator project?

Firstly, with an emulator project the source data is the output from another ocean model - in this case the Aronnax shallow water model made by Ed Doddridge. With ocean model output there is no ambiguity about the trustworthiness of the data and so no judgement calls for the model to make on which potential outliers to keep or remove, as these judgement calls are a real weakness of contemporary models.

Secondly, with an emulator project we can make a lot of progress with a single evaluation metric - the mean-square error (MSE) between the original model and the emulator output. At the end of every iteration this metric is passed to the LLM and if the number goes down then we’re probably making progress. As we’ll see below, this metric eventually needs help from others that measure physical realism, but at the outset it is a good indicator.

Thirdly, emulator projects build upon a large open literature on developing deep learning models across much wider domains than geophysical fluid mechanics such as computer vision. LLMs have an excellent knowledge of what works from this literature and can make confident hypotheses on how to make progress from the current point.

Problem domain

The ocean model I looked to emulate was a two-layer shallow water model in a double-gyre set-up on a 200 x 100 grid with cells of roughly 10 km.

I chose this relatively simple domain as I was primarily focused on how to build an autodiscovery pipeline rather than a frontier ocean emulator. While this domain is relatively simple compared to emulating a global ocean model it is still not a trivial problem - as we’ll see, the first models the agent built blew up within a few weeks of simulated time.

A classic double-gyre model has a fixed wind field. I wanted to ensure the agent had to learn how to handle some kind of time-varying surface boundary condition so I created a surface wind field that shifts north and south periodically. With a two-layer model there is enough turbulence in play that this pattern cannot simply be memorised. The target ocean model run was 5000 days, saved as 715 weekly snapshots. This was long enough to train an emulator on but short enough to maintain a fast iteration pace in the emulator development pipeline.

The figure below shows what the emulator has to learn from.

Overview of the training data: the shifting wind profile, the mean upper-layer thickness, a snapshot of upper-layer thickness anomaly and vorticity, and the kinetic energy of both layers over the full run with the held-out period shaded.

  • (a) The wind blows east in the north and south of the basin and west in the middle. The whole pattern slides north and south by 40 km every 40 days.
  • (b) The wind drives the classic double gyre: a thin subpolar gyre in the north and a thick subtropical gyre in the south, with the two meeting along a strong current that leaves the western boundary in the middle of the basin.
  • (c, d) The time-average hides most of the action. Any single snapshot is full of eddies, with thickness anomalies of up to 50 m, that are shed where the boundary current separates and then wander across the basin. These eddies are what make the problem chaotic.
  • (e) The model takes about two years to spin up from rest, and then the energy fluctuates around a steady level. The emulator trains on the first 80% of the run and is evaluated by running freely through the last 20% (shaded).

Auto-discovery pipeline

Agent harness

I used Claude Code as the agent harness throughout. The research loop itself is packaged as a Claude Code skill: a markdown file of instructions that the agent follows on every cycle of evaluating results, choosing a hypothesis, writing the code and tests, running the experiment and recording the outcome.

A key instruction in that skill is to start the next cycle immediately and not to stop because a future cycle might use too much time or too many tokens. Agents have a natural tendency to wrap up and report back after each task. For a research loop we want the opposite: the agent should keep going until I tell it to stop. With this instruction, and with the agent told to make its own decisions when I was away rather than waiting for me, it ran for many hours at a time.

Project set-up

The project used an approach known as autoresearch. In this paradigm:

  • we define an overall metric that the model is evaluated on
  • we keep a markdown file where the agent keeps tracks of the results to date
  • each step of the process is an experiment associated with a given hypothesis that is generated by the agent
  • every experiment is associated with a git branch and its own config file
  • at the end of every experiment the agent writes the key results to a central markdown file
  • if a branch provides a new benchmark result then its results are merged back into the main branch to replace the existing benchmark

A key aim of this process is to ensure that any results arrived at along the way are fully reproducible. This autoresearch approach is in contrast to a simpler alternative for persistent tasks where we use built-in functionality such as /loop in Claude Code. In that case the agent manages the experiments (rather than git and the markdown) and it is easy to end up with earlier results that can no longer be reproduced later.

Setting a proper benchmark

An agent that is rewarded for making a number go down will exploit any route to a lower number, including routes we didn’t intend. In an emulator project the obvious risk is that information about the period the emulator is being evaluated on leaks into the emulator. So rather than relying on asking nicely, I set out ground rules in the project instructions that the agent cannot negotiate:

  • the evaluation code lives in one file the agent is never allowed to edit
  • the train/evaluation split is a fixed chronological 80/20 split
  • every model is scored by the same rollout function, which hands the emulator only past truth and the known wind forcing
  • the emulator can’t get at evaluation-period data any other way, such as through global variables or by reading the data file during a rollout

I also had the agent compute two simple baselines before any training. Repeating the start state (persistence) scores 297.7. Predicting the long-term average of the training period (climatology) scores 136.8.

Hypothesis generation

Left to itself, an agent has a strong bias towards making very small configuration adjustments (e.g. changing the size of a layer or a learning rate), which quickly leads to dead ends.

Successful experimentation requires a balance between making bigger jumps in architecture and then exploring within each of these architectures to see if promising results can be developed. This process is familiar to any scientist - while we want to try different approaches we also know that results are generally not instantaneous and that persistence is required with any hypothesis to evaluate if there is performance that can be unlocked.

To enforce this balance, the skill makes the agent classify every experiment as tuning, a medium step (such as a new loss function or training horizon) or a bigger jump (such as a new architecture). It then sets rules for the mix: no more than two tuning experiments in a row unless one produced a new benchmark, a medium step or bigger jump after two failed tuning experiments, and a bigger jump after four experiments without a new benchmark. On every cycle the agent writes out three to five candidate hypotheses with their mechanism, acceptance criteria and failure modes before picking one.

Results

The chart below shows the final benchmark score of every training run, excluding a few runs that blew up. Red dots are runs of models that were promoted to become the benchmark. The final benchmark MSE of every training run for the two-layer double gyre simulation. Red dots are runs of models that were promoted to become the benchmark.

The chart looks surprisingly flat after the first few promotions, and the reasons why turned out to be the most interesting part of the project.

First challenge: getting a stable model

The early results were poor. A small CNN predicting the next state blew up within weeks. Predicting the change at each step and conserving water volume exactly still diverged, and training on 8-step rollouts only slowed the divergence down. Treating land cells as fixed helped, but the first model that stayed stable across the whole evaluation period came from adding noise to the starting state during training. That teaches the model to damp its own errors rather than amplify them.

That first stable model scored 215.7, much worse than simply predicting the average. Moving to a U-Net, keeping an exponential moving average of the weights (EMA) and training with a probabilistic loss (CRPS) brought this down to roughly 145.

Second challenge: setting the right metric

At this point it became clear that there was a problem with the benchmark itself. The two-layer model is chaotic, so after about 26 weeks no emulator can know where the eddies are. A physically realistic emulator that keeps producing eddies will have them in the wrong places, and its error ends up at around twice the climate variance. A blurry emulator that relaxes towards the average scores better. The benchmark is dominated by this chaotic tail, so on its own it rewards smoothing, and that is the opposite of what I want from an emulator.

So with the agent I built a secondary scorecard alongside the fixed benchmark:

  • predictable error: the error over the first 26 weeks, where skill is actually possible
  • realism: how well the energy spectra and the map of eddy energy match the original model, plus a penalty for spurious grid-scale noise
  • drift: whether the energy in a 25-year free run stays inside the range of the original model

I then set a new promotion rule: a new model becomes the benchmark if it clearly improves predictable error, realism or drift without making the fixed benchmark worse by more than the noise between random seeds. Physical realism takes priority over the score in the chaotic tail. This is why the later red dots in the chart don’t sit lower than the grey ones: those models were promoted for being more skilful in the predictable range and more realistic, not for a better score in the tail.

Seed variability

The difference between two runs of the same configuration with different random seeds was as much as 20 points on the benchmark, larger than most of the improvements the agent was proposing. A single lucky run could look like a breakthrough. Every candidate therefore needed at least three seeds before it could be promoted, which is why so many of the 134 runs are seed replicas.

Bigger jumps

Of the bigger jumps, a latent reduced-order model (tried three times), attention layers, ConvNeXt blocks and two-frame inputs all failed to beat the plain U-Net. However, these ideas remain strongly underdeveloped at this point. The ideas that worked were mostly changes to how the model is trained rather than what the model is:

  • training on its own multi-step rollouts
  • noise injection
  • the probabilistic CRPS loss
  • drawing the random inputs on a coarse grid so they act like large-scale variability rather than pixel noise
  • lightly smoothing each step’s update

The final gain came from simply making the U-Net wider. Widening helped up to 64 channels and stopped there: at 80 channels it appeared to start overfitting (on a single seed, see below), which makes sense with only about 570 weekly snapshots to train on.

The current best model is better than its predecessor on every measure over three seeds. Its predictable error is 15.4, its benchmark score is 140 through the second half of training (149 at the end), and its 25-year free run shows no drift. A variant with an extra spectral loss (the one in the video above) gives clearly more realistic eddies, but costs about 10 points on the benchmark. That trade-off is a judgement call, so the agent flagged it to me rather than deciding itself.

Optimising compute

Experiments started on my Mac’s GPU and then moved to RunPod. The serverless endpoint cost about $0.30 per run. Renting a single GPU by the hour and running two or three jobs on it at once brought this down to around $0.14 per run, since a model this size leaves most of a GPU idle. The agent wrote all the tooling for this: uploading the data, queueing jobs, retrying jobs that ran out of memory, fetching results and always deleting the pod at the end. Some of the most useful fixes came from real failures, such as cheap community GPUs that never exposed an SSH port and a host whose GPU was already filled by someone else’s process.

Key human input

The agent ran for many hours at a time without needing me, but the decisions that shaped the project were mine:

  • whether to favour realism over the benchmark score
  • whether a trade-off was worth promoting
  • how much to spend

The agent also made mistakes that its own records caught. It once launched an experiment from a branch that didn’t have the new code, so a key setting was silently ignored, and the run had to be relabelled and repeated. Because every experiment has a config, a branch and an entry in the research log, these slips were visible and fixable rather than quietly corrupting the results.

A second agent’s opinion on the methodology

The research log can tell me what the agent did, but it can’t tell me whether the agent’s way of reaching conclusions holds up. So I gave a fresh Claude Code session, with none of the research loop’s context, the project docs and the research log and asked it to evaluate the agent’s experimental methodology. It praised the agent’s habits: one change at a time, baselines before any training, numbers each hypothesis must hit, honest write-ups of failures, and controls for hardware differences and lucky seeds. But it also found real problems, listed here with the most serious first.

  1. Models are chosen on the test set. The benchmark, the predictable error and the realism score are all computed on the held-out 20%, and all 134 runs were judged by those numbers. No separate validation set is held back. With late-stage improvements of 5-15 points against seed noise of about ±20, the winning model’s scores are biased low. Locking the benchmark down stopped evaluation data leaking into the emulator, but it didn’t stop the agent’s choices from overfitting to the evaluation period. The fix is to hold back a validation stretch at the end of the training period, roll out from several start points within it, choose models on that, and touch the test set only when promoting a model.

  2. Each score comes from one start and one random draw. Every rollout starts from the same state, and the stochastic models use a single unseeded noise draw, so each score is one sample. Averaging over several start points or noise draws would reduce this variance more cheaply than training more seeds.

  3. The noise threshold kept moving. Over the project the agent described seed noise as ±20 for a single checkpoint, about 10 for a three-seed mean, and about 5 in the promotion rule. Under that last figure it blocked one candidate that was 14.5 worse on the final checkpoint, then shortly afterwards promoted the 48-channel model at 6.8 worse, describing that as within noise. The measured spread suggests 5 was too strict anyway. The tolerance should be computed once from the data, and candidates compared with the incumbent seed by seed, with a confidence interval rather than an eyeballed allowance. The project docs also disagreed: the README still told the agent to promote on the benchmark alone, while the project instructions had the newer rule.

  4. The agent designed a metric and then trained on it. The realism score is a weighted sum the agent designed, and the spectral-loss runs (including the one in the video) directly optimise its largest component. Part of their best-ever realism is just that circularity. A model trained with a spectral loss should be judged on diagnostics it wasn’t trained on, such as eddy energy maps, decorrelation times, low-frequency variability and long-run drift.

  5. Too many conclusions rest on one seed. To keep costs down, much of the late work ran on a single seed. The log still draws conclusions from gaps smaller than one seed’s spread, including the one above that widening stops helping at 64 channels. A hill-climb over channel width and loss weights on single seeds partly follows noise, and models promoted this way will tend to look worse when rerun.

  6. The search narrowed into tuning. The skill demands a bigger jump after four experiments without a new benchmark. Even so, recent work was widths of 48, 64 and 80 channels and spectral-loss weights of 0.1, 0.3 and 1.0. Seed replicas were also labelled as tuning experiments, which made the counters behind the mix rules meaningless. And most of the bigger jumps were tried before the CRPS loss, coarse noise and smoothing existed, so they never got the training recipe that eventually worked.

None of this undermines the early stability work, where each fix changed the error by orders of magnitude. It matters for the late-stage decisions, where the differences are about the size of the noise. The broader lesson is that locking the benchmark down protects against the agent cheating, but not against the agent fooling itself. The discipline of separate training, validation and test data that we apply to network weights applies just as much to the code an agent writes. Most of the fixes can be built into the pipeline rather than left to the agent’s judgement: a validation split, scores averaged over several starts, a paired statistical test for promotion and a config loader that rejects unknown keys.

Understanding the autoresearch process

There is a near analogy between carrying out an autoresearch process and training a deep learning model. When we train a deep learning model we provide the model with some training data and an objective function that we want to minimise. The model then takes a sample of the inputs, passes them through the network to generate an output. It then uses back-propagation to update its weights.

With autoresearch we are also providing training data and an objective function - in this case the MSE between the rollout and truth. However, in between the inputs and outputs there is a codebase that transforms the inputs into predictions instead of neural network weights. As the model is made from code we cannot update the weights with backpropagation. Instead the updating mechanism is an LLM which mutates the code based on our evaluation metric.

Two aligned loops: deep learning updates network weights with back-propagation, while autoresearch updates a codebase with an LLM.

The analogy also explains the biggest lessons of this project. Like a training loop, the LLM will exploit any route to a lower loss, so the benchmark has to be locked down. And like a training loop, it will optimise exactly the objective it is given, so if that objective rewards blurry emulators, blurry emulators are what you get. And like a model whose hyperparameters are tuned on the test set, a codebase that is improved by checking the same held-out data over and over will overfit to it. Choosing the objective is still the scientist’s job.

What these emulators can’t do yet

The biggest open question is generalisation. This emulator has only ever seen one wind regime, and several of the things that make it stable (giving it the average ocean state as an input, and training it to pull perturbed states back towards the training climate) also anchor it to that climate. The literature finds the same problem at scale. The Samudra global ocean emulator is stable for centuries but struggles to capture the size of forced climate trends, and work on transfer learning across CO2 levels found that adding a small amount of data from the new regime closes most of the gap. The next step for this project is to test the emulator on wind regimes it hasn’t seen.

This post is licensed under CC BY 4.0 by the author.