Autonomous ML Research Agent

Inspiration

Machine learning engineering has a dirty secret: much of the “engineering” is a five-step loop repeated over and over.

Read the data → engineer features → train → evaluate → reflect → repeat

It is mechanical enough to automate, but judgment-heavy enough to be difficult.

When TikTok TechJam’s Autonomous ML Research Agent challenge asked us to build an agent that could run this loop itself, on a real recommender-systems benchmark, against a real published baseline, we wanted to find out how far an LLM could actually go without a human quietly steering from the sidelines.

That became the core principle of our project:

If the agent is supposed to be autonomous, we should build the system so that we can prove when it actually was.

What We Built

We built an autonomous ML research agent that reproduces and iteratively improves on the provided Factorization Machine (FM) baseline for KuaiRand-Pure, a short-video recommendation benchmark.

Each iteration follows a closed research loop:

Read the current best pipeline and the history of previous experiments.
Hypothesise: an LLM proposes one research hypothesis and a complete modified pipeline.
Execute the candidate in an isolated sandbox with a hard timeout.
Evaluate it against the official validation metrics, where the primary score is:

$(\text{Primary} = \frac{\text{GAUC} + \text{nDCG@5}}{2})$
Decide: keep the candidate if it improves the score; otherwise revert.
Record the hypothesis, actual code diff, metrics, and failures in evidence-grade JSONL logs.
Repeat until convergence, the iteration cap, or the wall-clock budget is reached.

The convergence criterion is:

$\quad (\varepsilon = 0.002 \quad \text{over } N=3 \text{ iterations})$

The key design principle is simple:

A failed experiment must never destroy the current best result. A crashed experiment must never crash the research loop.

How We Built It

We were a team of five, and we made one rule non-negotiable:

Agree on the interfaces before anyone opens an editor.

We decomposed the research loop into four core handoff functions:

propose
run
read-metrics
log

We defined these as shared contracts in agent/contracts.py using TypedDicts, then built an integration test that calls every real module and compares its actual return value against the contract.

This caught interface drift twice during development.

The first time, a stale contract with a different return shape had remained in the repository for two days while another module was being built against it. The test caught the mismatch immediately.

After that, every new module was checked against the real contract before integration. This turned our interface from documentation into an executable guarantee.

The rest of the system

LLM interface

Automatic retries
Token metering
Defensive handling of malformed responses

Sandboxed execution

LLM-generated code runs in an isolated working directory
Syntax-checking happens before execution
Hard subprocess timeouts prevent runaway experiments

Independent evaluation

Metrics are read by a component outside the agent's edit surface
Candidates cannot modify the code that grades them

Evidence-grade logging

Every hypothesis, code diff, metric, and failure is recorded
Malformed JSONL entries cannot bring down an overnight run

The Problems We Didn't Expect

Building an autonomous system exposed several failure modes that would have been easy to miss in a conventional ML pipeline.

1. Contracts drift silently

Our shared interface was edited after another module had already been built against a different shape and nobody noticed.

The problem only became obvious when we stopped trusting documentation and wrote a test that actually called the real modules and checked their return values.

Lesson: documentation describes an interface, while tests enforce one.

2. A scratch directory isn't the repository

Candidate code runs in an isolated directory for safety.

That created a less obvious problem: imports such as import data and import evaluate could no longer resolve because the candidate was no longer running from the repository root.

We fixed this by explicitly injecting the repository root into PYTHONPATH, along with the environment variables that other modules had independently converged on.

Lesson: isolation changes the execution environment, not just the filesystem.

3. Windows quietly ate our tracebacks

Our development environment included Windows, where default file encoding is cp1252.

File writes and subprocess handling without an explicit encoding didn't always fail loudly. Instead, they could corrupt output until a teammate was left staring at an unhelpful:

exit code 1

We made encoding explicit with UTF-8 across file and subprocess boundaries.

Lesson: infrastructure bugs can masquerade as model failures.

4. Is an improvement actually an improvement?

This was one of our most important discoveries.

The FM baseline had a standard deviation of approximately 0.0008 across five seeds.

During a real autonomous run, the agent accepted two candidates with improvements of:

+0.0005
+0.0007

Both were smaller than the observed noise floor. This meant the agent was technically following the rules, but statistically, it was chasing noise.

We built a standalone multi-seed verification step so that accepted candidates could be re-evaluated before being trusted.

This exposed an important flaw in our original convergence logic:

A convergence threshold and a keep/revert decision cannot meaningfully operate on the same statistic while ignoring its variance.

5. The agent kept rediscovering our dead ends

Our first live iteration proposed adding static user/video features to the FM input.

We had already tested this offline but it provided no meaningful benefit.

The agent didn't know that because our previous experiments from Day 0 weren't being passed into its context.

We added a concise “known dead ends” section to the system prompt.

From then on, the agent could distinguish between:

ideas that had never been tested, and
ideas we had already ruled out.

That small change made the research history part of the agent's memory rather than lost context.

6. The autonomy trap

Our biggest temptation was also the most important test of the project.

Partway through development, it became clear that a hand-designed DeepFM architecture with out-of-fold target encoding was likely to outperform whatever the autonomous loop could discover within the remaining time.

So we built it. But we deliberately did not fold it into the autonomous results.

Instead, we documented it separately as a manual engineering intervention for the final submission.

Why?

Because the challenge is evaluating autonomy. Quietly editing the model until we got a better score would have produced a better number, but a weaker experiment.

We wanted to be able to say exactly:

This result came from the agent. This other result came from us.

What We Achieved

One of the most interesting outcomes came from the agent itself.

In its second real iteration, the autonomous loop independently proposed out-of-fold target encoding as a research direction.

This was striking because it was closely related to the approach we later chose for our hand-designed final architecture.

The agent found the signal before we deliberately built beyond it ourselves.

But score wasn't the only thing we cared about.

We are particularly proud that all five handoff functions follow the same failure-safe principle:

They return a plain dictionary and never raise.

A bad LLM response, generated-code crash, missing metrics file, unavailable API, or malformed log entry becomes:

{"ok": False, "error": "..."}

And the loop continues.

For an overnight autonomous experiment, this matters as much as model performance.

What We Learned

Interfaces aren't real until they're tested

A shared contract written in a file is still just documentation if nothing verifies it.

Our contract test caught integration problems before they reached the full system.

Autonomous and manual improvements are different claims

“The model improved” and “the agent improved the model” are not equivalent statements.

We deliberately separated those results, even when the manual approach produced a better score.

That distinction became one of the most important engineering principles of the project.

Noise is part of the research problem

An autonomous researcher doesn't just need to generate ideas.

It needs to know whether the evidence supports those ideas.

If the noise floor is larger than the improvement being rewarded, the agent can spend its entire compute budget chasing randomness.

Robustness is part of autonomy

An agent that produces a brilliant hypothesis but crashes on experiment #3 isn't autonomous.

For us, autonomy meant the entire loop had to survive bad generations, broken code, missing files, API failures, and statistically weak experiments without requiring a human to restart it.

What's Next

There are three improvements we would build next:

1. Automatic provider fallback

If one LLM provider becomes unavailable or unreliable, the research loop should automatically fail over to another provider.

2. Structured hypothesis planning

Instead of generating one hypothesis at a time, the agent could generate and rank several candidate directions before spending compute on the most promising one.

3. Multi-seed evaluation inside the decision loop

Rather than treating statistical verification as a separate manual step, the keep/revert decision itself should account for variance and require evidence that an improvement is likely to be real.

Results

The autonomous-loop results, including the best validation primary score, iteration count, token usage, and final hidden-test score, are generated directly from our final frozen run using:

python agent/summarise.py

See the Results section of the repository README for the exact numbers.

Final Takeaway

We started with a simple question:

How much of ML research can an LLM actually do by itself?

Our answer isn't that LLMs can replace ML engineers.

It's more interesting than that.

We found that an autonomous research agent can meaningfully explore a real ML problem, but only when the surrounding engineering makes its autonomy observable, reproducible, and fail-safe.

The hardest part wasn't getting an LLM to generate code.

It was building a system where we could trust what happened after it did.