Skip to content

Run population-based training

Goal: let good configurations copy themselves mid-run, instead of running every configuration to completion in isolation.

Ordinary tuning treats a trial as atomic: pick a configuration, run it, score it, pick the next. Population-based training breaks that. Trials run concurrently as a population; periodically each is ranked against the others, and an underperformer exploits a winner — copying its hyperparameters and its checkpoint — then explores by perturbing what it copied.

The consequence worth understanding before you use it: the answer PBT gives you is usually a fork of a fork, and the hyperparameters that produced it were never a single fixed configuration. They are a schedule, discovered while training.

Attach the scheduler

let study = Study::builder()
    .parallelism(THREADS)
    .budget(Budget::new_with_trials(TRIALS))
    // Rank every INTERVAL units; Jaderberg defaults for the rest (the bottom
    // 20% exploit a random member of the top 20%, perturbing by 0.8 / 1.2).
    .scheduler(Arc::new(Pbt::try_new(u64::from(INTERVAL))?))
    .create(StudyConfig::new(STUDY).with_seed(SEED))?;

study.optimize(train)?;

// The genealogy, rebuilt from storage: who forked from whom, and with what
// perturbation. PBT persists no scheduler state — the tree *is* the record.
let lineage = study.lineage()?;

study = atune.create_study(
    direction="minimize",
    # Rank every INTERVAL units; Jaderberg defaults for the rest (the bottom
    # 20% exploit a random member of the top 20%, perturbing by 0.8 / 1.2).
    scheduler=atune.schedulers.Pbt(INTERVAL),
    seed=SEED,
    name=STUDY,
)
study.optimize(objective, n_trials=TRIALS, n_jobs=JOBS)

# The genealogy, rebuilt from storage: who forked from whom, and with what
# perturbation. PBT persists no scheduler state — the tree *is* the record.
lineage = study.lineage()

Pbt::try_new(interval)? ranks the population every interval resource units. The rest are Jaderberg's defaults: the bottom fifth exploits a random member of the top fifth, and perturbation multiplies a copied scalar by 0.8 or 1.2.

Nothing else about the study changes. PBT is a Scheduler, so it attaches exactly where a pruner would — the difference is what it may return. A pruner can say stop; PBT can also say fork, and the study loop then materialises a new trial whose parameters and checkpoint come from another one.

The objective has to be resumable

This is the requirement that makes PBT different to configure. The scheduler forks with or without a checkpoint, so nothing enforces it. Exploit copies a trial's state, not just its numbers. If your objective cannot start from another trial's checkpoint, exploiting a winner gives you its learning rate and none of its progress: the child starts cold and spends its budget relearning what the winner already knew.

So an objective under PBT does three things a plain objective does not:

/// The loss this learning rate converges to — what it is *capable* of.
///
/// A quadratic bowl around [`LR_STAR`]: exactly [`FLOOR_MIN`] at the optimum and
/// worse either side of it. Pure arithmetic, in the same order as the Python arm,
/// so both arms compute the same bits.
fn floor_for(lr: f64) -> f64 {
    let off = lr - LR_STAR;
    FLOOR_MIN + FLOOR_SLOPE * (off * off)
}

/// Trains one segment, warm-starting from the parent's checkpoint if there is one.
fn train(ctx: &mut TrialCtx<'_>) -> Result<Outcome> {
    let lr = ctx.suggest_f64("lr", LR_LOW..=LR_HIGH, Scale::Linear)?;

    // The exploit half of PBT: a fork child inherits its parent's checkpoint
    // reference. atune stored the string, not the state — here the state is one
    // number, so the string is `{loss:?}` and `parse` recovers its exact bits.
    // Parsed into an owned `f64` immediately: the borrow of the context ends
    // here, before the loop needs it mutably again.
    let mut loss = match ctx.parent_checkpoint() {
        None => INITIAL_LOSS,
        Some(reference) => reference.parse::<f64>().map_err(|error| {
            Error::Objective(
                format!("checkpoint reference {reference:?} is not a loss: {error}").into(),
            )
        })?,
    };

    let floor = floor_for(lr);
    for step in 1..=STEPS {
        loss = floor + (loss - floor) * DECAY;
        if step % INTERVAL == 0 {
            // Checkpoint *before* reporting: `report` is the scheduler's decision
            // point, and a fork decided there copies the reference recorded here.
            ctx.record_checkpoint(u64::from(step), format!("{loss:?}"))?;
            ctx.report(u64::from(step), &[loss])?;
        }
    }
    Ok(loss.into())
}

def floor_for(lr: float) -> float:
    """The loss this learning rate converges to — what it is *capable* of.

    A quadratic bowl around ``LR_STAR``: exactly ``FLOOR_MIN`` at the optimum and
    worse either side of it. Pure arithmetic, in the same order as the Rust arm,
    so both arms compute the same bits.
    """
    off = lr - LR_STAR
    return FLOOR_MIN + FLOOR_SLOPE * (off * off)


def objective(trial: atune.Trial) -> float:
    """Trains one segment, warm-starting from the parent's checkpoint if any."""
    lr = trial.suggest_float("lr", LR_LOW, LR_HIGH)

    # The exploit half of PBT: a fork child inherits its parent's checkpoint
    # reference. atune stored the string, not the state — here the state is one
    # number, so the string is `repr(loss)` and `float` recovers its exact bits.
    reference = trial.parent_checkpoint
    loss = INITIAL_LOSS if reference is None else float(reference)

    floor = floor_for(lr)
    for step in range(1, STEPS + 1):
        loss = floor + (loss - floor) * DECAY
        if step % INTERVAL == 0:
            # Checkpoint *before* reporting: `report` is the scheduler's decision
            # point, and a fork decided there copies the reference recorded here.
            trial.record_checkpoint(step, repr(loss))
            trial.report(loss, step)
    return loss

It looks for an inherited checkpoint and warm-starts from it; it reports at intervals, so the scheduler has something to rank; and it writes a checkpoint the next generation can inherit. In a real trainer that reference is a path to model weights. In the example it is deliberately small enough to read, so the mechanism is visible rather than buried in a framework.

Read the genealogy

A flat list of trials is the wrong shape for a PBT study — it hides the only structure that matters. study.lineage() rebuilds the tree from storage: who forked or restarted from whom, at which checkpoint step, and — on the Rust side — with which perturbation; the Python view carries each node's parent, fork and restart flags, restart mode, fork step and value.

When a fork or restart child starts, the study loop writes its provenance — the donor, the step of the checkpoint it starts from, and any restart mode and seed — as an atune.lineage record beside the trial, in the same write that makes it running. Trials written before that record existed are still rebuilt from their parent links, as plain forks.

PBT persists no scheduler state of its own. The tree is the record, which is also why a PBT study resumes correctly: reopen the storage and the population's history is all there.

Export the schedule that produced a trial

study.recipe(leaf) turns one trial's ancestry into a replayable schedule: the segments from the start of its training to the leaf, each with the trial that trained it, the step range it covered, its full configuration, and the fork or restart that began it. The walk stops at a trial that inherited no checkpoint — a root, a fresh restart, or a child whose donor had not checkpointed yet — because nothing before it reached the leaf's weights.

Python's study.recipe(trial_number) returns the same schedule as a dict, and the CLI writes it as JSON:

$ atune recipe --study pbt.atj --trial 17 --out recipe.json

The JSON is versioned ("schema": "atune.recipe", "version": 1) and its bytes are deterministic for one study state. A recipe is built from the lineage records only, never guessed from reports: a fork on the path that has no atune.lineage record, or a donor missing from the study, is an error that names the trial.

Make the population actually pay

Two ways a PBT study can run, report success, and demonstrate nothing.

The segment must be short enough that a warm start still matters. If every trial reaches its floor before the first ranking, exploiting a winner inherits a finished job, and a fork becomes indistinguishable from a fresh trial with the same parameters. Measured on the example above: at eight steps with decay 0.6, about 1.7% of the initial gap survives to the ranking — which is exactly what makes a fork measurably better than a cold trial that drew the same learning rate. Longer segments erase the effect, not because PBT stopped working but because there was nothing left to inherit.

The population must be large enough to have a top and a bottom fifth. The population at a decision is every trial that has reported at that step and is still rankable — running and paused trials and finished ones alike — so it accumulates over the study rather than being capped by parallelism (the example runs single-threaded and still forks). What matters is that enough trials report at the same aligned steps, because both quantiles count whole trials, rounded down. Of n trials at a step, the reporter included, the worst floor(0.2 × n) may exploit and the best floor(0.2 × n) may be copied, and while either count is zero every trial continues, so a baseline forms before anything is copied. With the default fractions a step therefore needs five trials before any of them can exploit; four or fewer never do. A smaller fraction, set with Pbt::with_bottom_fraction or with_top_fraction (Python's bottom_fraction and top_fraction), needs a larger population: a nonzero fraction f names nobody until at least 1/f trials have reported at the step, and a zero fraction names nobody at all.

A pruned trial stays in the population. Pruning ends a trial's training, not the reports it already made, and those are real measurements. A trial pruned after it reported at a step — because its own objective stopped it, say — still ranks there, and if it is among the best it can be the donor whose parameters and checkpoint an exploit copies. Pb2 can use a pruned child's latest report as model data under the same conditions as a completed child's. Only failed trials, and waiting ones, which have not reported, are left out of the population.

Declare where trials end. Without a horizon PBT cannot tell a trial's final report from an interior one, so a worst-ranked final report forks a child with nothing left to train. Pbt::with_horizon(final_step) makes a report at or past that step terminal: it is recorded and ranked by the members that report there, but never forks. Python passes Pbt(interval, horizon=final_step).

Ties never count against a trial by default. A population that reports one shared value — a sparse reward still at zero — has no underperformer, so it does not exploit. Pbt::with_tie_break(TieBreak::Seeded) places each reporter at a reproducible random position inside its tie group, so a tied population exploits its bottom fifth in expectation. Python passes tie_break="seeded". Pb2 takes the same with_horizon and with_tie_break.

Restart instead of fork

An exploit normally forks: the child loads the donor's checkpoint and keeps training. On problems where a population collapses onto one lineage, or where inherited weights stop being plastic, it can be better to keep the donor's configuration and start the weights again. Pbt::with_restart(mode) makes every exploit a restart; the choice of who exploits whom, and the explore mutation, are drawn exactly as for a fork, so only what the child does with the donor's state changes. Pb2 takes the same setting. Both fork by default.

  • RestartMode::Fresh — the child gets the donor's mutated configuration and no checkpoint, so it re-initializes from scratch. It has no parent link; its lineage record is what names the donor.
  • RestartMode::try_new_with_shrink_perturb(shrink, noise)? — the child loads the donor's checkpoint like a fork, then shrinks the weights by shrink (in (0, 1]) and adds noise of scale noise (finite, non-negative).

atune never touches weights, so the objective applies the restart. In Rust, ctx.restart() returns the mode and a reproducible restart seed; a subprocess reads ATUNE_RESTART_MODE, ATUNE_RESTART_SEED, and, for shrink-and-perturb, ATUNE_RESTART_SHRINK and ATUNE_RESTART_NOISE (see the environment reference). Plain forks and ordinary trials see none of these. Python passes Pbt(interval, restart="fresh") or Pbt(interval, restart="shrink_perturb", shrink=0.5, noise=0.01), and the child reads trial.restart: a dict with the mode, the seed, and for shrink-and-perturb the shrink and noise. Leaving restart unset forks, and shrink or noise without restart="shrink_perturb" raises.

Evaluate what the population produced

A PBT trial's tuning value scores a configuration schedule that ran under the tuning seeds, but what you take away is either the trained weights or the schedule itself. Three final-stage estimands rank the same top k trials by tuning value, evaluate each under the study's test seeds, and return the same report with its tune, test and gap columns:

  • study.reevaluate(objective, k) — the configuration: train afresh with the trial's own parameters, as for any study.
  • study.evaluate_checkpoints(objective, k) — the artifact: evaluate the trial's latest checkpoint without training. ctx.evaluation() is Evaluation::Checkpoint and ctx.parent_checkpoint() is that checkpoint; a trial that never checkpointed scores NaN.
  • study.replay_recipes(objective, k) — the schedule: train afresh following the trial's recipe. ctx.evaluation() is Evaluation::Recipe and ctx.recipe() is the schedule; the sampled parameters are its first segment's, and recipe.params_at(step) gives the configuration to switch to at each later segment. A trial whose recipe cannot be built fails the call before anything runs.

Python reads the same through trial.evaluation ("checkpoint", "recipe", or None) and trial.recipe (the recipe dict); a subprocess reads ATUNE_EVALUATE, plus ATUNE_CHECKPOINT or ATUNE_RECIPE (the recipe as single-line JSON). Each estimand stores its replicates in an atune.estimate record beside the trial and reuses them only while the test seeds and the evaluated checkpoint or recipe are unchanged. The trial's value and its reevaluate test replicates are never touched, so the three answers sit side by side.

Account for the resource

study.resource_ledger() reports, per trial and in the study's resource unit, what it trained (the amount the fidelity budget charged it), the lineage resource its weights represent (where its recipe ends), its age since initialization or the last shrink-and-perturb restart, and the number of replicates each stored evaluation stage holds. The ledger's total is the study's fidelity charge. Compare estimands and methods at equal trained resource, not equal lineage: a fork's lineage counts the prefix it inherited.

The same charge is a budget. Declare what a reported step counts when the study is created — ResourceUnit::Steps (environment steps, the default), Episodes, Walltime, or a custom unit through StudyConfig::with_resource_unit in Rust, create_study(resource_unit="steps") in Python — and cap a run with Budget::with_fidelity_units(n) in Rust or study.optimize(objective, max_resource=n) in Python. The Rust limit is on the study's total charge; the Python amount is added to what the study had already spent when the call began, as n_trials is. Both are cooperative: no trial starts once the charge reaches the limit, and the trial that crossed it finishes. An objective must report its steps to be charged at all.

PB2, and when to reach for it

Pb2 replaces PBT's random perturbation with one chosen by a time-varying GP-bandit: instead of multiplying by 0.8 or 1.2, it proposes the perturbation its model expects to help most given what the population has done so far. Each decision fits that model to at most the newest 512 improvement rows by report step, the window GP-EI and CARBS use, so the fit's cost stops growing in a long study; gathering the rows still reads every trial. It is behind the pb2 feature, which implies gp — see the feature reference.

Reach for it when perturbation is the bottleneck: a population that exploits correctly but explores badly. It is not a drop-in improvement, because it adds a model whose own assumptions can be wrong.

Six opt-in settings, all off by default, address the ways the default model misreads a population; Python spells them as atune.schedulers.Pb2 keywords. with_estimated_noise fits the GP's observation noise rather than treating rewards as exact. with_round_time counts GP time in intervals, so old data can be forgotten. with_segments reads one improvement per interval of every trial, roots included, instead of one per fork. with_start_fitness, which needs segments, tells the model where each segment started. with_sibling_diversity steers a proposal away from configurations already at the boundary. with_trust_region(r) keeps it near the parent. None is yet measured to help on a real learner, so treat them as experiments; the in-run controller below uses all of them.

Populations inside one run

Everything above treats each member as a trial: its own process, its own environment interaction, and a checkpoint through storage at every fork. When the members can live in one process instead — off-policy learners sharing one replay buffer, PPO members owning disjoint blocks of one vectorized environment, or model-based agents sharing one world model — the population costs gradient compute rather than environment steps, and an exploit is an in-memory weight copy.

atune::pb2::InRunPopulation (the pb2 feature) is the controller for that case. It holds no learners and touches no storage. Your loop asks it for each member's parameters, reports every member's progress, fitness and health at an iteration boundary, calls decide, and applies the returned exploits: copy the donor's learner into the recipient, then set the recipient's new parameters. Explore::Pb2 proposes them with a time-varying GP-UCB that treats every interval as an observation, takes the starting fitness as an input, estimates its noise, keeps one round's proposals apart, and stays within a trust region around the donor (trust_region, half-width 0.2 of each unit axis by default, because GP-UCB over the whole space chases its corners); Explore::Perturb is plain PBT. An unhealthy member is never copied, and is replaced whenever a healthy member can donate. lineage(member) is the schedule that produced a member, and snapshot/restore checkpoint the controller with your population.

Python exposes the same controller as atune.pb2.InRunPopulation(search_space, members, seed=...), with the distribution dicts accepted by create_study. Use explore="pb2" (default) or explore="perturb"; initial maps member indices to complete parameter dicts. Call report(member, progress, fitness, feasible=True) for each member, then decide(): each returned dict has recipient, donor, and params. Copy the donor's learner first, then install those parameters. members and observations are properties; params(member), lineage(member) and best() inspect the schedule. Save the host population alongside snapshot() bytes; restore into a controller with matching settings.

Shared substrate is not free: members sharing replay learn from each other's data, which helps off-policy learners and can mislead on-policy ones, and a shared world model is tuned by none of them. Tune its own hyperparameters on a slower timescale.

IPBT

atune::pb2::Ipbt is a synchronous controller for iterated population-based training (the same pb2 feature). Start with Ipbt::builder(space, members, budget, initial_step).seed(seed).build(). The defaults follow the authors' IPBT configuration: eight target members are typically started in sixteen slots, and the best eight survive the first interval. Set members explicitly; the default oversampling multiplier is two.

For each round(), evaluate every listed slot before training when evaluate_start is true and call report_start(slot, fitness), passing NaN for an infeasible start. Train each slot for round.step using params(slot), then report(slot, fitness, feasible). After every active slot reports, call decide() and apply all returned actions in order: Archive preserves a checkpoint, Copy overwrites a recipient's learner, Reinit computes w = shrink*w + perturb*w_fresh, and Retire stops training a slot. Install the controller's new parameters before the next round. Initial slots need fresh weights supplied by the host. A process-based driver can use the same protocol; shared replay is optional.

Stagnation restarts archive the best checkpoint, reactivate oversampled slots, double the training interval, replace all parameters (half uniform, half restart-level BO), and reset half the weights; the remainder use shrink 0.2 and perturb 0.1. An oversampled round charges the oversampling multiplier times its per-slot interval to the campaign clock. round() returns None when the budget is spent; best() identifies a live member or an archived winner. Save host learners and archives alongside the controller's snapshot() for restore(), taken before decide() or after its actions are applied, since a snapshot holds no pending actions.

Python uses atune.pb2.Ipbt(search_space, members, budget, initial_step, seed=...). Its round() dict contains slots, step, and evaluate_start; report_start, report, and params use the fixed slot indices. decide() returns ordered dicts whose kind is archive (slot/restart), copy (recipient/donor), reinit (slot/shrink/perturb), or retire (slot). Apply them in order, then install the next parameters. best() returns kind="member" with slot/fitness or kind="archived" with restart/fitness; members, slots, restarts, and clock are properties. Both Python controllers return compact UTF-8 JSON bytes containing the complete Rust StateBlob (kind, version, data), and reject foreign or mismatched snapshots with atune.AtuneError. The same IPBT snapshot timing and host-checkpoint rules apply.

The port follows pinned code 827b7c7, including its always-successful trust-region labels and suppression of late restarts. Its deterministic exact GP fits and clipped proposal starts approximate the reference's optimized GPs and gradient search; the Rust module documents every difference. Deliver the winning checkpoint: schedule replay needs its own measurement, since the paper's Humanoid replay lost IQM (9272 to 6939). A trial-level adapter waits with fixed population slots; sibling forks cannot retire oversampled members or express doubled intervals.

What this costs you

  • The result is a schedule, not a configuration. You cannot hand someone "the best hyperparameters" from a PBT study and expect them to reproduce your number — they need the trajectory. Report the lineage alongside the winner.
  • Forking is history-dependent. Which trial sits in the bottom fifth at a ranking depends on what has finished, so a PBT study with several workers is not reproducible trial-for-trial. Determinism states what does still hold.
  • A fidelity budget counts training, not inherited progress. A trial is charged its highest reported step less the step of the checkpoint it started from, so a fork's inherited prefix is charged once, to the trial that trained it, however many children share it. A fresh restart starts from zero and is charged in full. study.resource_ledger() shows the charge per trial.
  • Checkpoints cost storage. A population that forks every interval writes a checkpoint per trial per interval. That is a real disk budget, and the reason checkpoint references are references rather than payloads.

Next