RL-autoTune

Declare a task in one YAML, launch a staged screen → refine → verify search, and get a hyperparameter winner that reproduces before it's accepted.

Staged hyperparameter optimization for RL tasks (Isaac Lab + skrl by default, extensible to other stacks). Goal: maximum reward in minimum wall-clock time, with results you can trust — every candidate is scored by a fair canonical eval (reward shaping can't cheat it), dead runs are killed live from their TensorBoard curves, and the winner is retrained from scratch with the same seed and must reproduce before it counts.

Tasks are data, not code: each task is one tasks/<name>.yaml. Adding a task requires zero code changes.

An AI agent given this repo to optimize a task? CLAUDE.md is your operating contract; docs/NEW_TASK.md is the task-declaration guide.

---

Table of contents

• The pipeline
• Why results are trustworthy
• Quick start
• What you get back
• Termination-cause reports
• Shared or free GPU
• CLI cheatsheet
• Architecture in one picture
• Repository layout
• Documentation
• Requirements

---

The pipeline

flowchart LR
    C["8-12 candidates<br/>(one axis each,<br/>always a baseline)"] --> S["STAGE 1 — SCREEN<br/>all candidates,<br/>short equal budget"]
    S -->|"live monitor kills:<br/>NaN, reward gates,<br/>hung process"| D[("dead runs<br/>stopped early")]
    S --> E1["canonical eval<br/>of every survivor"]
    E1 --> R["STAGE 2 — REFINE<br/>top-k, full budget,<br/>plateau early-stop"]
    R --> E2["canonical eval"]
    E2 --> V["STAGE 3 — VERIFY<br/>retrain the winner<br/>from scratch, SAME seed"]
    V -->|"score within margin"| W["winner.json<br/>reproduced: true"]
    V -->|"score collapsed"| F["flagged:<br/>not reproduced"]

Loading

• Screen — every candidate gets the same short budget, enough for the baseline to show a clear signal. The live monitor tails each trial's tfevents and kills only certainly dead runs (NaN reward at any poll, reward below a catastrophic gate, no events at all).
• Refine — the top-k rankable candidates get the full budget. Once the smoothed reward flattens, plateau detection stops the run early — the trial still counts, the GPU time isn't wasted.
• Verify — the winner is retrained from scratch at the same seed and re-scored. GPU physics isn't bit-deterministic, so a tolerance margin is used; the margin formula is sign-safe (works for tasks whose valid scores are negative).

Every stage appends to results/<study>/trials.jsonl, so a killed or crashed campaign resumes exactly where it stopped — finished trials are never retrained.

Why results are trustworthy

Mechanism	What it prevents
Canonical eval — every checkpoint is re-scored with all reward/rule parameters forced back to the developer defaults declared in the task spec's canonical block; deterministic policy, fixed eval seed	a candidate that shaped its training reward from leaking that shaping into its score
Rank only by eval, never by training reward — a shaped trial whose eval failed is ranked last by design	"my training curve looks great" winning anything
Verify stage — same-seed retrain must land within tolerance × max(|score|, score_scale) of the original	lucky-seed / one-off winners
Tie band + speed — among near-tied scores (band scaled by |best|), the trial that reached 90 % of its final reward soonest wins	picking a 2 % better but 2× slower config
Termination attribution — every eval reports why episodes end, per cause	"the score went up" hiding "the agent now dies twice as often"
Never touch the task repo — trials import the task from a read-only snapshot (stub_roots)	the campaign silently drifting when someone edits the task

Quick start

export RPO_PYTHON=/path/to/sim-rl-venv/bin/python   # e.g. an Isaac Lab venv

# 0. declare the task (data only, no code) — see docs/NEW_TASK.md
$EDITOR tasks/mytask.yaml

# 1. define the search: candidates, budgets, early-stop rules
cp configs/TEMPLATE_search.yaml configs/mystudy_search.yaml && $EDITOR $_

# 2. validate the wiring with no GPU (specs load, resolve and route)
python scripts/selftest.py

# 3. smoke-test on the box: one tiny real trial through the full machinery.
#    Measures steps/s -> budget advice, reports peak VRAM -> num_envs advice.
python scripts/doctor.py configs/mystudy_search.yaml

# 4. launch the campaign FULLY DETACHED (it runs for hours; the babysitter
#    relaunches on death and trials.jsonl makes restarts idempotent)
setsid nohup bash scripts/run_campaign.sh configs/mystudy_search.yaml &
tail -f results/mystudy/orchestrator.log

A study config is small and readable — candidates are hand-picked axes from a bottleneck diagnosis (see docs/METHODOLOGY.md), not a blind grid:

study: mystudy
base:    { task: mytask, seed: 42, num_envs: 4096 }
eval:    { num_envs: 512, max_steps: 600, min_episodes: 1500, seed: 7777 }
stages:
  screen: { timesteps: 8000,  monitor: { min_reward_gates: [{step: 4000, min_reward: 0.02}] } }
  refine: { top_k: 3, timesteps: 24000, monitor: { plateau: {after_frac: 0.5, window_frac: 0.3, min_improve: 0.04} } }
  verify: { tolerance: 0.15 }
candidates:
  - { name: baseline }                                          # REQUIRED anchor
  - { name: epochs8,  agent_cfg: { agent: { learning_epochs: 8 } } }
  - { name: kl012,    agent_cfg: { agent: { learning_rate_scheduler_kwargs: { kl_threshold: 0.012 } } } }
  - { name: hz120,    env_cfg:   { decimation: 1 } }

What you get back

Everything lands in results/<study>/ — the examples below are illustrative.

orchestrator.log — the campaign narrating itself:

[14:02:11] === STAGE 1 screen: 8 candidates × 8000 steps, 3 parallel slots ===
[14:03:44] [screen/baseline/s42]     step=1200 reward=0.184 (1.5 min)
[14:29:30] [screen/lr5e3/s42]   train status=early_stopped NaN reward at step 2400 wall=3.1 min points=12
[14:41:03] screen ranking: epochs8=1.92, baseline=1.71, hz120=1.66, kl012=1.41, ...
[14:41:03] === STAGE 2 refine: top 3 × 24000 steps ===
[15:37:12] [refine/epochs8/s42]   train status=plateau smoothed reward 2.081->2.104 over last 14 points at step 19200
[16:04:19] verify: original=2.113 rerun=2.041 margin=0.317 reproduced=True (tolerance 15%)

report.md — every trial, ranked:

trial	status	steps	train min	final train rew	eval return	scoring rate	t90 (s)	score
refine/epochs8/s42	plateau	19200	21.4	2.284	2.113	71.2 %	512.3	2.113
verify/epochs8/s42	plateau	19800	22.0	2.241	2.041	69.8 %	530.1	2.041
refine/baseline/s42	completed	24000	26.8	2.019	1.902	64.1 %	741.0	1.902
screen/lr5e3/s42	early_stopped	2400	3.1	nan	—	—	—	-inf

winner.json — the verified answer:

{
  "name": "epochs8",
  "candidate": { "name": "epochs8", "agent_cfg": { "agent": { "learning_epochs": 8 } } },
  "seed": 42,
  "timesteps": 24000,
  "refine_score": 2.113,
  "verify_score": 2.041,
  "reproduced": true,
  "tolerance": 0.15,
  "margin": 0.317,
  "total_campaign_hours": 3.9
}

Need one candidate at a custom budget (e.g. a long run to its true plateau)? scripts/extra_refine.py runs it outside the fixed-budget pipeline and appends to the same trials.jsonl / report.md:

python scripts/extra_refine.py configs/mystudy_search.yaml epochs8 --timesteps 120000 --dry-run
python scripts/extra_refine.py configs/mystudy_search.yaml epochs8 --timesteps 120000

Termination-cause reports

Score alone doesn't tell you what the policy does. Every canonical eval also attributes why episodes end — automatically on manager-based Isaac envs (termination-manager term names), and on Direct envs via terminations.probes declared in the task spec (transcribed 1:1 from the env's _get_dones(), see docs/NEW_TASK.md §4). Zero config still gets you the terminated vs time-out split.

python scripts/term_report.py configs/mystudy_search.yaml epochs8 --episodes 500

[term] score at this checkpoint: mean_return=1.847 over 512 episodes (canonical rules, deterministic policy)
512 episodes, cause source: probes
  cause              %  count  mean_len  mean_return
  time_out       58.4%    299     600.0        2.612
  fell           31.1%    159     212.4        0.393
  out_of_bounds   8.2%     42     301.7        0.844
  unattributed    2.3%     12     154.0        0.291

Reading it: 31 % of episodes end in a fall at step ~212 with a fraction of the time-out return — that is the bottleneck the next candidates should target. A growing unattributed share is a built-in honesty check: the probes have drifted from the task's actual termination code. The same breakdown is stored in trials.jsonl for every trial of every stage, so an agent (or you) can diff causes across candidates without re-running anything.

Shared or free GPU

Both regimes are first-class (configs/TEMPLATE_search.yaml):

• Sharing the card with a live training run? Set base.device (pins every trial and eval off cuda:0) and base.min_free_vram_mb — a launch waits up to vram_wait_s for free VRAM, then fails loudly instead of OOMing the neighbouring run. doctor.py measures one trial's peak VRAM and suggests a safe num_envs.
• Card to yourself? stages.screen.max_parallel: K screens K candidates concurrently (~K× less screen wall-time; refine/verify stay sequential). Size K so peak VRAM × K fits with headroom; screen.stagger_s spaces the first wave.

[doctor] measured ~86.4 policy steps/s (final_train_reward=0.412)
[doctor] budget guide: screen 8k = ~2 min, refine 24k = ~5 min (+~1-1.5 min Isaac boot/trial).
[doctor] peak trial VRAM ≈ 6212 MB at num_envs=256 (≤ 24.27 MB/env — conservative, includes fixed sim overhead)
[doctor] GPU 1 free now: 21480 MB → within an 80% budget (17184 MB) num_envs ≲ 707 by the linear upper bound.

CLI cheatsheet

Command	What it does
python scripts/selftest.py	validate all task specs + study configs, no GPU (includes stubbed smokes of parallel screening and termination attribution)
python scripts/doctor.py <config>	one tiny real trial end-to-end: venv, registration, tfevents, checkpoint, eval, steps/s → budgets, peak VRAM → num_envs; with base.resume also compares the checkpoint's score before/after fine-tuning
python scripts/doctor.py <config> --eval-only --checkpoint <pt>	no training: honest canonical score + termination report of any checkpoint
bash scripts/run_campaign.sh <config>	babysat, resumable campaign (use setsid nohup … &)
python scripts/optimize.py <config> --report-only	regenerate report.md from existing results
python scripts/extra_refine.py <config> <cand> --timesteps N	one candidate at a custom budget (auto plateau-stop on long runs; --full to disable)
python scripts/term_report.py <config> <cand>	replay a finished trial's checkpoint → termination-cause table

Architecture in one picture

flowchart TD
    subgraph orchestrator["orchestrator process — plain python (pyyaml + tensorboard)"]
        O["optimizer/orchestrate.py<br/>staged search, ranking, records"]
        RU["optimizer/runner.py<br/>spawn + live monitor + VRAM guard"]
        SC["optimizer/scoring.py<br/>canonical score, tie-band rank"]
        TB["optimizer/tb.py<br/>tfevents reader"]
        TS["optimizer/taskspec.py<br/>tasks/&lt;name&gt;.yaml loader"]
    end
    subgraph sim["sim subprocess — RPO_PYTHON (Isaac Lab venv)"]
        TT["scripts/trial_train.py"]
        TE["scripts/trial_eval.py /<br/>trial_eval_marl.py"]
        TK["optimizer/termstats.py<br/>pre-reset hook"]
    end
    O -->|"trial.json<br/>(embedded task spec)"| RU
    RU -->|spawn| TT
    RU -->|spawn| TE
    TE --- TK
    TT -->|tfevents| TB
    TB --> RU
    TE -->|eval.json| O
    O -->|append| J[("results/&lt;study&gt;/trials.jsonl<br/>+ report.md + winner.json")]
    TS --> O
    SC --> O

Loading

The orchestrator never imports the simulator: each trial is a subprocess under RPO_PYTHON, fully described by a trial.json that embeds the task spec — the edge scripts are task-agnostic and the two python environments stay decoupled. Details: docs/ARCHITECTURE.md.

Repository layout

optimizer/      task-agnostic core: orchestrate (stages), runner (spawn+monitor),
                scoring, tb (tfevents), taskspec (task loader), termstats (cause attribution)
scripts/        edge scripts (trial_train / trial_eval / trial_eval_marl), doctor,
                selftest, extra_refine, term_report, run_campaign.sh, cfg_merge
tasks/          one <name>.yaml per task + tasks/agents/ base agent configs
configs/        one <study>_search.yaml per search (TEMPLATE_search.yaml to copy)
docs/           NEW_TASK, METHODOLOGY, ARCHITECTURE, CONFIG_REFERENCE
results/        per-study outputs (trials.jsonl, report.md, winner.json committed;
                heavy trial dirs gitignored)
workspace/      machine-local task snapshots (gitignored)

Documentation

Doc	Read it when
docs/NEW_TASK.md	adding a task: the tasks/<name>.yaml schema, snapshotting, canonical LOCK lists, termination probes
docs/METHODOLOGY.md	choosing candidates: bottleneck diagnosis → 8–12 axes, budget & gate calibration
docs/ARCHITECTURE.md	understanding/extending the machinery: components, trial lifecycle, resumability, failure handling
docs/CONFIG_REFERENCE.md	every knob of the study config and task spec, with defaults
CLAUDE.md	operating contract for AI agents driving the tool

Requirements

• Orchestrator side: Python ≥ 3.9, pyyaml, tensorboard. No GPU, no sim.
• Sim side (RPO_PYTHON): any venv that imports your simulator + RL framework — e.g. an Isaac Lab venv with skrl. PPO / AMP / IPPO / MAPPO work out of the box (framework: skrl_ppo / skrl_marl); a project-specific skrl runner via skrl_custom; another RL library (RSL-RL, rl_games) is one new edge-script pair — the orchestrator core is unchanged.

License: BSD-3-Clause.