Skip to content

Repository files navigation

autoMIL

Autonomous Agent-Driven Multiple Instance Learning

Let AI agents run your ML experiments while you sleep.

Python 3.10+ License: Apache 2.0 Release v1.2.1 Tests

Getting Started | How It Works | Training-Script Contract | Documentation

autoMIL is a plug-and-play experiment framework that automates the research iteration, not the hyperparameter. It overlays onto your existing ML project, in any language with any ML library, and lets any coding agent read the codebase, propose and implement source-level changes (architecture + training recipe), run each as an isolated reproducible experiment, and learn from the results across a persistent experiment tree — under a configurable per-cell budget, with every variant reproducible, attributable to its parent, and portable across machines and LLM runtimes. Unlike grid/menu AutoML (Optuna-style search within a fixed pipeline), the unit of automation is the research iteration itself. The original motivation was Multiple Instance Learning in computational pathology (see examples/ovarian_hrd for a 189-experiment autonomous run); the framework itself is generic and runs sklearn-iris end-to-end via the same contract (see examples/sklearn-iris).


The Problem

Manual ML development is a grind: tweak hyperparameters, edit code, run training, check results, repeat. AutoML tools like Optuna search parameter spaces, but they can't invent new architectures, combine techniques creatively, or learn from what failed last time.

The Solution

autoMIL gives coding agents the infrastructure to run experiments autonomously:

What the agent does:

  • Reads your codebase
  • Designs experiments
  • Modifies any file (models, losses, augmentations)
  • Submits experiments via CLI
  • Learns from results
  • Repeats forever

What autoMIL handles:

  • GPU scheduling (best-fit bin packing)
  • Parallel execution (git worktree isolation)
  • Experiment tracking (directed tree, not flat log)
  • Knowledge persistence (learnings.md)
  • Result evaluation (val-firewall + Ladder keep/discard)
  • 3D visualization (live dashboard)

Real result: On ovarian cancer HRD prediction, autoMIL autonomously ran 189 experiments and improved the primary-metric value from 0.814 to 0.851 (+4.5%), discovering techniques like R-Drop, focal loss, gradient clipping, and coordinate positional encoding that human researchers hadn't tried.


Key Features

Feature Description
Plug-and-play Overlays onto any existing ML project. No restructuring needed.
Multi-runtime agents First-class skills for Claude Code; shared AGENTS.md + trajectory capture for Codex, OpenCode, and DeepSeek (routed via opencode/codex). automil init --runtime auto-detects or installs explicitly.
Full-codebase scope Agent edits any file: architectures, losses, augmentations, optimizers.
Git worktree isolation Each experiment runs in a snapshot. Only changed files stored.
Pluggable backends local (default), slurm (submitit, opt-in via [slurm] extra), ray (raw @ray.remote, opt-in via [ray] extra). Same Backend ABC; same cap contract.
Hardware autodetect automil init probes CUDA / ROCm / CPU via LocalBackend.healthcheck() and stamps detected GPU count, VRAM, and concurrency defaults into config.yaml.
Variant registry Architectural changes ship as committed variant modules (automil/variants/<parent>/<name>.py) selected via config. Registry-only path reproduces a node end-to-end via automil verify-repro.
Configurable per-cell budget cap Two-tier state machine (refusing-new at T-buffer, terminating at T) with per-fold checkpoints and a SIGTERM contract. Budget is set as a duration (cap.budget: 6h, also 30m/90s/2d) — set it with uv run automil budget set 6h / inspect with uv run automil budget show, or per-cell via uv run automil submit --budget-seconds N (D-134, honored only on the submit that creates the cell). cap.mode: agent_active consumes Claude Code's native cumulative claude_code.active_time.total metric (CLI + user active seconds; idle excluded), scraped from Claude's localhost Prometheus endpoint and bound to the cell by synchronous session hooks. wall_clock is the portable fallback for runtimes without that observer. The framework fallback is 6h; the frozen preprint campaign pins 12h plus exactly 30 launches. Budget-killed runs reconcile to executed with partial primary_value, never crash.
Validation-firewall Keep/discard selects on validation only; test metrics are sealed at ingest into a quarantined held_out block and revealed exactly once via automil certify. Test never drives search, so final numbers aren't selected on test.
Generalization gate Pre-registered held-out manifest + paired Wilcoxon + bootstrap CI + Bonferroni, ships a candidate node status, manual nomination by default, promotion-rate metric exposed via SSE.
Trajectory recorder Per-submit JSONL using OpenTelemetry gen_ai.* keys with secret redaction (sk-…, hf_…, AWS keys) and bounded rotation (5 MB soft / 50 MB hard).
Multi-GPU orchestrator Background daemon with bin packing, OOM detection, crash recovery, namespaced running/<backend>/.
Experiment tree UCB-inspired scoring balances exploitation and exploration across branches; primary_value-dominance keep/discard gated by the Ladder keep-margin, on a validation-only primary_value scalar (the val-firewall).
3D dashboard Interactive Three.js visualization with live SSE updates (localhost:8420).
Persistent learnings Knowledge accumulates across sessions. Agents don't repeat mistakes.
Setup validation automil check validates protected files, registry purity, backend directives, and env.required before experiments run.

Quick Start

1. Install

# Install as a global CLI tool (recommended)
uv tool install git+https://github.com/leoyin1127/autoMIL.git

# Or install from a local clone
git clone https://github.com/leoyin1127/autoMIL.git
cd autoMIL
uv tool install -e .

Environment convention: Development in this repository always uses uv run (for example, uv run automil init). The bare automil commands below assume the recommended global installation was created by uv tool install; they are not an instruction to bypass uv-managed environments.

2. Initialize in your project

cd /path/to/your/project    # any existing git repo
automil init                 # creates automil/ subdirectory

3. Setup

With Claude Code (recommended)
claude
# Type: /automil-setup

The agent scopes your codebase, configures everything, verifies the training contract, and establishes a baseline. Fully autonomous. The setup skill follows a documented idempotency protocol and runs a 1-minute dry-run gate before declaring done.

With another runtime (Codex, OpenCode, DeepSeek)
automil init --runtime codex                   # or opencode, deepseek-via-opencode, deepseek-via-codex
automil show-skill --runtime codex             # render the shared skill to stdout

automil init auto-detects from existing .claude/, .codex/, .opencode/ directories when --runtime is omitted. The canonical skill content lives in _shared/; per-runtime directories carry only non-skill assets (hooks, plugins), so show-skill renders the shared content for every runtime.

Manual setup

Edit automil/config.yaml. Minimum sections you must touch:

run:
  script: "train.py"          # your training script (any name)

files:
  editable: ["train.py", "models/", "losses/*.py"]   # files, dirs, or globs
  readonly: ["evaluate.py"]                          # what must not change

baseline:
  primary_value: 0.814             # your starting performance

env:
  required: []                 # vars that MUST be set before submit
  passthrough: [AUTOMIL_*]     # vars forwarded to experiment subprocesses

scoring:
  # The reducer CR-1b uses to recompute the primary_value from the validation
  # `metrics` block at ingest — the val-firewall does not trust the reported
  # scalar. "" -> "mean" over the metrics values (reproduces the standard
  # primary values exactly); also: max | min | trust_reported (explicit opt-out,
  # weakens the firewall). Arithmetic expressions are rejected by
  # `automil check`; test metrics never belong in the primary_value (they live
  # in the sealed `held_out` block).
  formula: ""

cap:
  # Consumer-supplied; 6h is the generic framework fallback. Durations accept
  # 6h / 30m / 90s / 2d (or a bare number = seconds). Set it with
  # `uv run automil budget set 6h`; inspect with `uv run automil budget show`.
  budget: 6h
  safety_buffer: 30m
  # agent_active = Claude Code's native active-time counter, bound to this cell
  # by synchronous session hooks. wall_clock = time since cell creation.
  mode: agent_active

Ensure your training script honors the training-script contract (writes result.json matching automil/schemas/result.schema.json, exits cleanly on SIGTERM with a partial result). Then validate:

automil check

4. Run

Use tmux to keep the orchestrator and agent running in the background:

# Terminal 1: orchestrator (must stay running)
tmux new -s orchestrator
automil orchestrator start
# Ctrl-b d to detach

# Terminal 2: visualization (optional)
tmux new -s viz
automil viz start            # dashboard at localhost:8420
# Ctrl-b d to detach

# Terminal 3: agent loop
tmux new -s automil
claude --dangerously-skip-permissions   # autonomous mode, no permission prompts
# Type: /automil

5. Watch

automil status               # quick summary
automil rank                 # top proposals
# Open http://localhost:8420  # 3D experiment tree

How It Works

  Your Project (unchanged)          autoMIL Overlay
  ========================          ===============
  src/models/clam.py           -->  automil/config.yaml
  src/train.py                      automil/program.md
  src/data_loader.py                automil/learnings.md
  ...                               automil/orchestrator/
                                       queue/ -> running/ -> archive/
                                       completed/

The experiment cycle:

Agent designs experiment
    |
    v
automil submit --files train.py models/clam.py
    |  (snapshots only changed files)
    v
Orchestrator picks up from queue
    |  (creates git worktree at base commit)
    |  (overlays changed files on top)
    v
Runs on a typed CPU/CUDA/ROCm slot in isolation
    |  (framework-owned device masks)
    v
Collects result.json
    |  (Ladder keep-margin on the val primary_value: keep or discard?)
    v
Updates experiment graph
    |  (UCB scoring across branches)
    v
Agent reads results + learnings --> designs next experiment

Each experiment stores only its diff, not the full repo. A worktree provides the complete project context at runtime.


Training Script Contract

The seam between autoMIL and your code is the training-script contract: write a result.json matching automil/schemas/result.schema.json before exiting, honor SIGTERM for partial flush, declare required env vars in automil/config.yaml: env.required. Any language, any ML library qualifies.

The minimum valid payload is:

{"primary_value": 0.912}

primary_value is the single scalar the experiment tree uses for ranking (higher is always better; for loss minimization, negate). Everything else is optional. Under the val-firewall, metrics is validation-only and is what primary_value is computed from; any test metrics go in a sealed held_out block that the orchestrator quarantines under archive/<node>/certify/ and reveals once via automil certify, so test never drives search; held-out-named keys (e.g. test_*) inside metrics fail the node closed at ingest. A full example payload from the autobench consumer:

{
  "status": "completed",
  "metrics": {"val_auc": 0.870, "val_bacc": 0.810},
  "held_out": {"test_auc": 0.872, "test_bacc": 0.830},
  "primary_value": 0.840,
  "elapsed_seconds": 4098,
  "peak_vram_mb": 4500
}

The sklearn-iris consumer writes {"primary_value": <accuracy>} with no metrics dict; both shapes validate against the same schema.

The schema is JSON Schema 2020-12 and is validated at ingest by the orchestrator; malformed payloads transition the node to crashed with a schema-location pointer.

Environment variables available to your script
Variable Value Description
AUTOMIL_ACCELERATOR cpu / cuda / rocm Authoritative execution substrate.
CUDA_VISIBLE_DEVICES Physical CUDA ID, ROCm logical 0, or empty Framework-owned compatibility mask.
ROCR_VISIBLE_DEVICES Physical ROCm ID or empty Primary ROCm/Linux host-device mask.
HIP_VISIBLE_DEVICES / GPU_DEVICE_ORDINAL ROCm logical 0 or empty Post-ROCR logical-device masks.
AUTOMIL_GPU 0 Backward-compatible logical slot; CPU consumers ignore it.
AUTOMIL_NODE_ID node_0042 Experiment identifier.
AUTOMIL_DESC "try focal loss" Experiment description.
AUTOMIL_RUNTIME claude / codex / ... Runtime declared by the agent for trajectory tagging.

The orchestrator owns all device-mask variables; a per-experiment spec.env cannot override them. On ROCm/Linux it first selects the physical host GPU with ROCR_VISIBLE_DEVICES, then exposes that device as logical 0 to HIP-compatible training code. CPU consumers receive empty masks and must branch on AUTOMIL_ACCELERATOR, not on AUTOMIL_GPU.

Vars listed under env.passthrough in config.yaml are forwarded from the orchestrator process to each experiment subprocess. AUTOBENCH_ROOT-style auto-injection was removed in v1.0 (Phase 8 / DEC-01); declare what you need.


CLI Reference

# Project setup + validation
automil init [--runtime <r>] [--no-healthcheck]   Overlay automil/ on current repo
automil check                                     Validate setup (protected files, env.required, backend, registry)
automil show-skill --runtime <r>                  Render merged per-runtime skill file to stdout

# Experiment lifecycle
automil submit --node <id> --desc "..." --mil-model <m> [--files <f>] [--max-time SEC]
               [--budget-seconds N] [--safety-buffer-seconds M]
                                                  Snapshot changed files and queue. --mil-model is
                                                  required unless run.mil_model is set in config.yaml.
                                                  --budget-seconds / --safety-buffer-seconds override
                                                  cap.* for the cell this submit creates (D-134;
                                                  ignored on subsequent submits joining the same cell).
automil cancel <node_id>                          Cancel a running experiment
automil resubmit <node_id>                        Re-queue a terminal experiment as a new node
automil rank                                      Show top-ranked proposals (UCB)
automil propose --parent <id> --kind <k> --desc "..."  Add a proposal (kind: architecture|regularization|hp|data|ensemble)
automil portfolio [--threshold 0.5]               Architecture-vs-HP mix of pending proposals; exits non-zero below target
automil reconcile [--recompute-best]              Sync graph with orchestrator state
automil status                                    Show experiment summary

# Variant registry (Phase 1)
automil port-variant <node_id>                    Convert a node's overlay into a registered variant module
automil promote-variant <variant_id>              Move a gate-passing candidate to canonical
automil refresh-registry                          Regenerate per-kind variants/__init__.py deterministically
automil apply <node_id>                           Apply a node's variant selection to config.yaml
automil revert-baseline                           Reset registry.protected paths to base_commit (mandatory pre-stash)
automil verify-repro <node_id>                    Reproduce a node via the registry path; assert |actual - expected| < tolerance

# Cell budget cap (Phase 4)
automil cell status [<id>] / list                 Inspect cell budget state and consumed seconds

# Generalization gate (Phase 5)
automil nominate <node_id>                        Mark keep-status node as a gate candidate
automil promote <candidate_id>                    Run Stage B gate (paired Wilcoxon + bootstrap CI + Bonferroni)
automil gate register-manifest / retire-manifest / status / stats
                                                  Manage / inspect the gate manifest

# Held-out certification (val-firewall)
automil certify [--node <id>] [--top-k N]         Reveal sealed held-out TEST metrics for val-selected node(s); once, post-search

# Trajectory recorder (Phase 3)
automil trajectory record / export                JSONL trajectory capture and redacted export bundle

# Loop + daemons
automil start-loop / stop-loop                    Control agent loop flag
automil orchestrator start / stop / status        GPU scheduler daemon (best-fit bin packing)
automil viz start / stop / status                 3D visualization dashboard at localhost:8420

Run automil <command> --help for full flag listings.


Project Structure

your-project/                    # your repo (untouched)
  src/
  models/
  train.py
  ...
  automil/                       # added by automil init
    config.yaml                  # project settings (run, files, env, scoring, cap, gate, backend, hardware)
    program.md                   # agent instructions for the loop
    learnings.md                 # accumulated insights
    graph.json                   # experiment tree (gitignored)
    cells/                       # cell budget state (Phase 4)
    variants/                    # registered variant modules (Phase 1)
      <parent>/                  #   one subdir per registered parent
        <name>.py                #   committed code; selected via config
        __init__.py              #   regenerated by `automil refresh-registry`
    orchestrator/
      queue/                     # pending
      running/<backend>/         # per-backend live job specs (Phase 6)
      archive/                   # permanent record
        node_0001/
          train.py               # only changed files
          spec.json              # experiment spec
          run.log                # stdout/stderr (orchestrator-owned, drained from backend.log_iter)
          result.json            # metrics
          trajectory.jsonl       # agent prompt + tool-call events (Phase 3, gitignored by default)
      completed/                 # notifications

Agent Compatibility

Runtime Support Level How to Start
Claude Code First-class automil init --runtime claude then /automil-setup, then /automil
Codex First-class automil init --runtime codex; installs the shared AGENTS.md + CLI-fallback trajectory capture
OpenCode First-class automil init --runtime opencode; installs the shared AGENTS.md + trajectory plugin
DeepSeek First-class (routed) automil init --runtime deepseek-via-opencode (or deepseek-via-codex); DeepSeek is a model accessed through a host runtime
Cursor / Aider / Windsurf Compatible Point the agent at automil/program.md and the contract, any agent that can read files, edit code, and run shell commands works

The canonical skill content lives under _shared/; per-runtime directories carry only non-skill assets (hooks, plugins), and automil show-skill --runtime <r> renders the shared content for every runtime. See the Agent Compatibility Guide.


Examples

Example Task Library Notes Result
sklearn-iris 3-class iris classification scikit-learn Reference second-consumer (~80 LOC, no automil.* imports) demonstrating the training-script contract primary_value ≈ 0.95
ovarian_hrd Binary HRD classification CLAM-MB / H-optimus-1 Pre-v1.0 autonomous run 0.814 → 0.851 (+4.5%, 189 experiments)
clwd 7-class lung subtype classification autobench Skeleton -
placeholder - - Template emitted by automil init -

Documentation

  • Getting Started, full setup, configuration, and usage
  • Training-Script Contract, the seam between framework and consumer (6 contract items + SIGTERM patterns)
  • Agent Compatibility, per-runtime setup, overlay merge model, multi-runtime asset layout
  • Implementation Report, v1.0 architecture, design decisions, and the 9-phase refactor that produced it (point-in-time; the Ladder gate + val-firewall came later — see the contract doc and CHANGELOG)
  • CHANGELOG, release notes through v1.2 plus the Unreleased preprint-campaign integration (preprint-v3 protocol, native active-time metering)

Get Started | View Examples | Report Issues

Apache 2.0 License

About

Autonomous Agent-Driven Multiple Instance Learning

Resources

Stars

4 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages