code/train/train_backtrack.py

"""Backtracking (stage-replay) training loop for the recurrent latent model.

Motivation
----------
Standard curriculum training runs EVERY batch at `num_passes = curriculum.stage`
(the current frontier depth). When the curriculum advances to stage i, the
earlier, shallower readouts (stage 1..i-1, each of which should be produced by
applying the tied recurrent step exactly t times) drift / are forgotten.

Backtracking = interleaved replay with the recurrence depth MATCHED to the
replayed stage:

    stage t  <->  difficulty level t+2  <->  num_passes = t  <->  t active slots

With cand_slot_mode="depth" and passes_per_stage=pps the depth becomes pps*t and
the active-slot count follows num_passes rather than the puzzle level. Replay is
still safe because that mode maps slot j -> stage j: slots [0, pps*t) hold the
right targets no matter which frontier stage the sampler built the batch for.

  * frontier step : level (i+2) puzzles, num_passes = i   -> trains stage-i readout
  * backtrack step: level (t+2) puzzles, num_passes = t   -> re-derives stage-t
                    for a randomly chosen earlier t in {1..i-1}

Each optimisation step is frontier with prob (1 - backtrack_prob) and a backtrack
step with prob backtrack_prob (t uniform over the earlier stages). Because the
recurrent step is weight-tied, replaying "apply f exactly t times = stage t"
keeps f reusable at every depth instead of specialising to the frontier depth.

Everything else (model, data, per-slot BCE + LM CE loss, eval, promotion) is
reused unchanged from the standard pipeline.
"""

import functools
import math
import os

from absl import logging
from clu import metric_writers
from flax import jax_utils
from flax.training import checkpoints
from flax.training import common_utils
import jax
import numpy as np
import tensorflow as tf

from train import data
from train import evaluater
from train import model
from train import trainer


class StageBatchSampler:
    """Draws minibatches of a *specific* difficulty level from the train set.

    Reuses SudokuDataset (which loads the puzzles + staged candidate masks and
    builds a level->indices map) and assembles numpy batches in the exact tuple
    layout the train step expects: (seq, puzzle, start_index, level, cand_targets).
    """

    def __init__(self, config, seed=0, curriculum=None):
        self.ds = data.SudokuDataset(config, train=True, curriculum=curriculum)
        self.rng = np.random.RandomState(seed)
        n_puz = int(getattr(config, "instance_puzzles", 0))
        n_all = len(self.ds)
        self.n_puz = n_all if n_puz <= 0 else min(n_puz, n_all)
        self.available = {}
        for lvl, idx in self.ds.level_index.items():
            keep = np.asarray(idx)
            keep = keep[keep < self.n_puz]
            if len(keep) > 0:
                self.available[lvl] = keep

    def has_level(self, level):
        return level in self.available

    def _gather(self, idxs):
        seqs, puzzles, starts, levels, cands = [], [], [], [], []
        for idx in idxs:
            item = self.ds.__getitem__(int(idx))
            seq, puzzle, start_index, lvl, cand = item[:5]
            seqs.append(seq)
            puzzles.append(puzzle)
            starts.append(start_index)
            levels.append(lvl)
            cands.append(cand)
        return (
            np.stack(seqs).astype(np.int32),
            np.stack(puzzles).astype(np.int32),
            np.stack(starts).astype(np.int32),
            np.stack(levels).astype(np.int32),
            np.stack(cands).astype(np.int32),
        )

    def sample_all(self, bs):
        """Batch drawn uniformly from the whole corpus, matching the standard
        loop's sampler. Difficulty is not a curriculum axis, so every step --
        frontier or repair -- uses this same distribution and differs only in
        the recurrence depth it trains at."""
        return self._gather(self.rng.randint(self.n_puz, size=bs))

    def sample(self, level, bs):
        idx_pool = self.available[level]
        idxs = idx_pool[self.rng.randint(len(idx_pool), size=bs)]
        return self._gather(idxs)

    def sample_mixed(self, levels, bs):
        """Batch drawn uniformly over `levels` (level first, then a puzzle of
        that level) — the exact distribution the standard curriculum sampler
        uses. Frontier steps must use this, NOT single-level batches: training
        exclusively on the frontier level starves every easier level and is a
        different (worse) recipe than the standard loop, not "standard + BT"."""
        levels = [l for l in levels if l in self.available]
        idxs = []
        for _ in range(bs):
            lvl = levels[self.rng.randint(len(levels))]
            pool = self.available[lvl]
            idxs.append(int(pool[self.rng.randint(len(pool))]))
        return self._gather(idxs)


def _run_train_step(step_fn, state, batch_tuple, dropout_rngs):
    """Shard a single numpy batch and run one (pmapped) train step."""
    inputs, _, start_index, levels, cand_targets = batch_tuple
    inputs = common_utils.shard(jax.tree_util.tree_map(np.asarray, inputs))
    start_index = common_utils.shard(jax.tree_util.tree_map(np.asarray, start_index))
    levels = common_utils.shard(jax.tree_util.tree_map(np.asarray, levels))
    cand_targets = np.asarray(cand_targets)
    _nd = jax.local_device_count()
    cand_targets = cand_targets.reshape(
        (_nd, cand_targets.shape[0] // _nd) + cand_targets.shape[1:])
    state, metrics, _ = step_fn(
        state, inputs, start_index, levels, cand_targets, dropout_rng=dropout_rngs)
    return state, metrics


def _prepare_backtrack(config, workdir):
    """Build the model, state, pmapped steps, sampler, and writers shared by the
    probabilistic and adaptive backtracking loops.

    Returns a dict of everything both loops need. `state` is already replicated.
    """
    K = int(config.num_latent_slots)
    assert K > 0, "Backtracking is for the recurrent latent model (K>0)."

    logging.info("Creating datasets (backtracking loop)")
    curriculum = data.CurriculumState(
        stage=int(getattr(config, "curriculum_start_stage", 1)),
        max_stage=int(getattr(config, "curriculum_max_stage", 6)))
    # Per-level train sampler (replaces the uniform-over-unlocked sampler so we
    # can draw a batch of a specific replay stage's level on demand).
    sampler = StageBatchSampler(config, seed=int(config.seed),
                                 curriculum=curriculum)
    eval_data_iter = data.create_iter(config, config.minibatch_size, train=False)

    model_config = model.TransformerConfig(
        dtype=config.dtype, vocab_size=config.vocab_size, seq_len=config.seq_len,
        num_heads=config.num_heads, num_layers=config.num_layers,
        emb_dim=config.emb_dim, qkv_dim=config.qkv_dim, mlp_dim=config.mlp_dim,
        dropout_rate=config.dropout_rate,
        attention_dropout_rate=config.attention_dropout_rate,
        deterministic=False, num_latent_slots=K,
        inject_latents=bool(int(getattr(config, "recurrent_latent", 1))),
    )
    print(str(model_config.__dict__), flush=True)

    rng = jax.random.PRNGKey(config.seed)
    rng, init_rng, dropout_rng = jax.random.split(rng, 3)
    net = model.TransformerLMHeadModel(model_config)
    dummy_latents = jax.numpy.zeros(
        (config.minibatch_size, K, config.emb_dim), model_config.dtype)
    dummy_positions = jax.numpy.zeros((config.minibatch_size, K), jax.numpy.int32)
    dummy_active = jax.numpy.zeros((config.minibatch_size, K), bool)
    _, initial_variables = jax.jit(net.init_with_output)(
        {"params": init_rng, "dropout": dropout_rng},
        jax.numpy.ones((config.minibatch_size, config.seq_len), jax.numpy.int32),
        dummy_latents, dummy_positions, dummy_active)

    state, lr_scheduler_fn = trainer.get_state(config, net, initial_variables)
    start_step = 0
    if config.resume_training:
        state = checkpoints.restore_checkpoint(config.ckpt_loc, state)
        start_step = int(state.step)
        print("----------Restored model from", config.ckpt_loc,
              f"at step {start_step}-----------")

    writer = metric_writers.create_default_writer(
        workdir, asynchronous=False, just_logging=(jax.process_index() > 0))
    tf_summary_writer = tf.summary.create_file_writer(workdir)

    state = jax_utils.replicate(state)
    dropout_rngs = jax.random.split(rng, jax.local_device_count())

    def make_p_train_step(num_passes, backtrack, inst_slot):
        return jax.pmap(
            functools.partial(
                trainer.train_step, config=model_config, hyperparams=config,
                learning_rate_fn=lr_scheduler_fn, num_passes=num_passes,
                backtrack=backtrack, inst_slot=inst_slot),
            axis_name="batch", donate_argnums=(0,))

    # Precompile, for every recurrence depth 1..K, both a FRONTIER step (normal
    # loss: LM CE + all active-slot BCE) and a strict BACKTRACK step (only the
    # last active slot's readout, no CE). Frontier is used when t == curriculum
    # stage i; the strict backtrack step is used by the probabilistic loop for
    # replayed earlier stages t < i.
    # Stage -> recurrence depth. With passes_per_stage>1 a stage advances the
    # chain by more than one slot (pps=2, K=12: stage 6 -> all 12 slots), so the
    # replay depth for stage t is pps*t, not t. Keyed by stage either way.
    pps = int(getattr(config, "passes_per_stage", 1))
    _depth = lambda t: min(pps * t, K)
    # Keyed by stage, so the leak penalty's candidate slot is just t-1 whether
    # this is a frontier step or a replay of an earlier stage.
    p_frontier = {t: make_p_train_step(_depth(t), False, t - 1)
                  for t in range(1, K + 1)}
    p_backtrack = {t: make_p_train_step(_depth(t), True, t - 1)
                   for t in range(1, K + 1)}
    p_eval_step = jax.pmap(functools.partial(
        evaluater.eval_step, config=model_config.replace(deterministic=True)),
        axis_name="batch")

    hooks, report_progress, _ = trainer.get_metrics_report_progress(
        config, workdir, writer)

    return {
        "K": K,
        "start_step": start_step,
        "curriculum": curriculum,
        "sampler": sampler,
        "eval_data_iter": eval_data_iter,
        "model_config": model_config,
        "state": state,
        "dropout_rngs": dropout_rngs,
        "p_frontier": p_frontier,
        "p_backtrack": p_backtrack,
        "p_eval_step": p_eval_step,
        "hooks": hooks,
        "report_progress": report_progress,
        "writer": writer,
        "tf_summary_writer": tf_summary_writer,
        "promote_threshold": float(getattr(config, "promote_acc_threshold", 0.85)),
        "promote_patience": int(getattr(config, "promote_patience_steps", 8000)),
        "min_stage_steps": int(getattr(config, "min_stage_steps", 2000)),
        "ckpt_keep": int(getattr(config, "ckpt_keep", 100)),
        "stage_ckpt_dir": os.path.join(workdir, "stage_ckpts"),
    }


def train_and_evaluate_backtrack(config, workdir):
    """Backtracking curriculum training loop (stage replay with matched depth)."""
    workdir = os.path.abspath(workdir)
    backtrack_prob = float(getattr(config, "backtrack_prob", 0.5))
    frontier_mix = bool(int(getattr(config, "backtrack_frontier_mix", 1)))

    ctx = _prepare_backtrack(config, workdir)
    K = ctx["K"]
    curriculum = ctx["curriculum"]
    sampler = ctx["sampler"]
    eval_data_iter = ctx["eval_data_iter"]
    state = ctx["state"]
    dropout_rngs = ctx["dropout_rngs"]
    p_frontier = ctx["p_frontier"]
    p_backtrack = ctx["p_backtrack"]
    p_eval_step = ctx["p_eval_step"]
    hooks = ctx["hooks"]
    writer = ctx["writer"]
    tf_summary_writer = ctx["tf_summary_writer"]
    promote_threshold = ctx["promote_threshold"]
    promote_patience = ctx["promote_patience"]
    min_stage_steps = ctx["min_stage_steps"]
    ckpt_keep = ctx["ckpt_keep"]
    stage_ckpt_dir = ctx["stage_ckpt_dir"]
    stage_started_at = 0
    sched_rng = np.random.RandomState(int(config.seed) + 1)

    # Bookkeeping: how often each replay depth is actually trained.
    stage_step_counts = {t: 0 for t in range(1, K + 1)}

    def sample_target_stage(i):
        """Frontier stage i with prob (1-p); else a backtrack stage 1..i-1."""
        if i <= 1 or sched_rng.rand() >= backtrack_prob:
            return i
        return int(sched_rng.randint(1, i))   # uniform in {1..i-1}

    with metric_writers.ensure_flushes(writer):
        for step in range(0, config.max_steps):
            if step % 10000 == 0:
                print("Step:", step, flush=True)

            i = curriculum.stage
            t = sample_target_stage(i)
            # Frontier step (t == i): normal loss. Backtrack step (t < i): strict
            # replay of the depth-t readout only. Both draw from the whole
            # corpus: depth, not difficulty, is the replayed axis.
            step_fn = p_frontier[t] if t == i else p_backtrack[t]
            batch = sampler.sample_all(config.minibatch_size)
            state, metrics = _run_train_step(step_fn, state, batch, dropout_rngs)
            stage_step_counts[t] += 1

            for h in hooks:
                h(step)

            if math.isnan(metrics["loss"][0]):
                print("Loss became nan; stopping.", flush=True)
                break

            if step % config.eval_every_steps == 0:
                eval_metrics = evaluater.get_eval_metrics(
                    state, eval_data_iter, p_eval_step, config)
                per_level = eval_metrics.pop("per_level_acc")
                per_depth = eval_metrics.pop("per_slot_acc_changed", {})
                per_depth_all = eval_metrics.pop("per_slot_acc", {})
                if not any(v >= 0 for v in per_depth.values()):
                    per_depth = per_depth_all

                def _m(key):
                    v = eval_metrics.get(key, [])
                    return round(float(np.mean(v)), 4) if len(v) else -1.0

                print(step, "stage", curriculum.stage,
                      "target_t", t,
                      "loss", round(float(metrics["loss"].mean()), 4),
                      "ce", round(float(metrics["ce_loss"].mean()), 4),
                      "aux_bce", round(float(metrics["aux_loss"].mean()), 4),
                      "| val_acc", _m("acc"), "loc_acc", _m("loc_acc"),
                      "val|loc", _m("val_given_loc_acc"),
                      "| cand_bit_acc", _m("cand_bit_acc"),
                      "cand_set_acc", _m("cand_set_acc"),
                      "cand_set_chg", _m("cand_set_acc_changed"),
                      "| replay_counts", dict(stage_step_counts),
                      flush=True)

                with tf_summary_writer.as_default():
                    tf.summary.scalar("loss", metrics["loss"].mean(), step=step)
                    tf.summary.scalar("ce_loss", metrics["ce_loss"].mean(), step=step)
                    tf.summary.scalar("aux_bce_loss", metrics["aux_loss"].mean(), step=step)
                    tf.summary.scalar("curriculum_stage", curriculum.stage, step=step)
                    for key in eval_metrics.keys():
                        v = eval_metrics[key]
                        if isinstance(v, dict) or not v:
                            continue
                        tf.summary.scalar("eval_" + key,
                                          np.array(v).mean(), step=step)
                    for lvl, v in per_level.items():
                        if v >= 0:
                            tf.summary.scalar(f"eval_acc_level_{lvl}", v, step=step)

                # ---- Curriculum promotion (same rule as the standard loop) ----
                if curriculum.stage < curriculum.max_stage:
                    frontier_acc = per_depth.get(curriculum.stage, -1.0)
                    steps_in_stage = step - stage_started_at
                    hit_threshold = frontier_acc >= promote_threshold
                    patience_over = steps_in_stage >= promote_patience
                    if steps_in_stage >= min_stage_steps and (hit_threshold or patience_over):
                        reason = "threshold" if hit_threshold else "patience"
                        curriculum.stage += 1
                        stage_started_at = step
                        print(f"[curriculum] step {step}: promote to stage "
                              f"{curriculum.stage} ({reason}; graduated depth "
                              f"{curriculum.stage - 1} cand-set "
                              f"acc={frontier_acc:.3f}); "
                              f"backtrack pool now depths 1..{curriculum.stage-1}",
                              flush=True)
                        if config.save_checkpoint:
                            unrep = jax_utils.unreplicate(state)
                            checkpoints.save_checkpoint_multiprocess(
                                workdir, unrep, step, keep=ckpt_keep, overwrite=True)
                            checkpoints.save_checkpoint_multiprocess(
                                stage_ckpt_dir, unrep, step, keep=100,
                                overwrite=True, prefix=f"stage{curriculum.stage}_")

            if config.save_checkpoint and step > 0 and step % config.save_every_steps == 0:
                checkpoints.save_checkpoint_multiprocess(
                    workdir, jax_utils.unreplicate(state), step,
                    keep=ckpt_keep, overwrite=True)

        if config.save_checkpoint:
            checkpoints.save_checkpoint_multiprocess(
                workdir, jax_utils.unreplicate(state), config.max_steps,
                keep=ckpt_keep, overwrite=True)


def train_and_evaluate_backtrack_adaptive(config, workdir):
    """Deficit-driven ("adaptive") backtracking loop.

    Difference from the probabilistic loop
    --------------------------------------
    The probabilistic loop replays a *uniformly random* earlier stage every step
    with fixed probability p, regardless of whether that stage needs help. This
    loop instead *watches* each earlier stage's held-out accuracy and only goes
    back to repair a stage when it has actually regressed:

      * Reference: when the curriculum promotes past stage t, we record that
        stage's frontier-level (t+2) val_acc as its "graduation" accuracy.
      * Trigger (relative drop): while at frontier stage i, at every eval we scan
        earlier stages 1..i-1; a stage t is *in deficit* if its current level-(t+2)
        val_acc has fallen below (graduation_acc[t] - margin).
      * Selection (most-deficient first): if any stage is in deficit we switch the
        training target to the single most-regressed stage and train there.
      * Repair step: a FULL frontier-style step at the matched depth (level t+2,
        num_passes=t, LM CE + all-active-slot BCE) -- NOT the strict readout-only
        backtrack step. The metric we are trying to restore is placement accuracy
        (driven by the LM CE), so the repair objective must include it.
      * Exit (recover-or-cap): stay on the stage until its val_acc climbs back
        above (graduation_acc[t] - margin), or until a max-repair-steps cap fires
        (so a stuck stage cannot stall the frontier forever). On exit we re-scan
        and either move to the next most-deficient stage or return to the frontier.

    Frontier promotion is paused while repairing, and the repair time is credited
    back to the frontier stage's patience clock on return.

    Depth-matched eval note: the evaluator scores a level-L puzzle with exactly
    k=L-2 active latent slots, so per_level_acc[t+2] is a faithful measurement of
    the stage-t readout -- a clean trigger signal with no extra instrumentation.
    """
    workdir = os.path.abspath(workdir)
    margin = float(getattr(config, "backtrack_margin", 0.03))
    max_repair_steps = int(getattr(config, "backtrack_max_repair_steps", 6000))
    min_frontier_steps = int(getattr(config, "backtrack_min_frontier_steps", 0))
    frontier_target_acc = float(
        getattr(config, "backtrack_frontier_target_acc", 0.0))
    grad_acc_seed_raw = str(getattr(config, "backtrack_grad_acc_seed", "") or "")
    max_repair_fraction = float(
        getattr(config, "backtrack_max_repair_fraction", 0.0))
    grad_decay = float(getattr(config, "backtrack_grad_decay", 0.0))
    freeze_after_step = int(getattr(config, "backtrack_freeze_after_step", 0))
    frontier_mix = bool(int(getattr(config, "backtrack_frontier_mix", 1)))

    ctx = _prepare_backtrack(config, workdir)
    K = ctx["K"]
    start_step = ctx["start_step"]
    curriculum = ctx["curriculum"]
    sampler = ctx["sampler"]
    eval_data_iter = ctx["eval_data_iter"]
    state = ctx["state"]
    dropout_rngs = ctx["dropout_rngs"]
    p_frontier = ctx["p_frontier"]
    p_eval_step = ctx["p_eval_step"]
    hooks = ctx["hooks"]
    writer = ctx["writer"]
    tf_summary_writer = ctx["tf_summary_writer"]
    promote_threshold = ctx["promote_threshold"]
    promote_loc_threshold = float(getattr(config, "promote_loc_threshold", 0.0))
    promote_spread_threshold = float(
        getattr(config, "promote_spread_threshold", 0.85))
    promote_hbound_threshold = float(
        getattr(config, "promote_hbound_threshold", 0.0))
    promote_ceq_threshold = float(getattr(config, "promote_ceq_threshold", 0.0))
    promote_ceq_mode = str(getattr(config, "promote_ceq_mode", "mean"))
    promote_ceq_margin = float(getattr(config, "promote_ceq_margin", 0.0))
    promote_hinge_threshold = float(
        getattr(config, "promote_hinge_threshold", 0.0))
    promote_hinge_slack = float(getattr(config, "promote_hinge_slack", 0.05))
    instance_mode = bool(getattr(config, "instance_dir", None))
    promote_patience = ctx["promote_patience"]
    min_stage_steps = ctx["min_stage_steps"]
    ckpt_keep = ctx["ckpt_keep"]
    stage_ckpt_dir = ctx["stage_ckpt_dir"]

    # Controller state.
    grad_acc = {}                 # stage t -> level-(t+2) val_acc at graduation
    plateau_steps = int(getattr(config, "plateau_steps", 0))
    plateau_delta = float(getattr(config, "plateau_delta", 0.005))
    stage_best_acc = -1.0
    stage_best_step = start_step
    mode = "frontier"             # "frontier" | "repair"
    repair_stage = None           # stage currently being repaired
    repair_started_at = 0         # step the current repair stage began (cap)
    repair_episode_start = 0      # step we first left the frontier (clock credit)
    stage_started_at = start_step  # step the current frontier stage began
    last_repair_return_step = start_step  # for min-frontier cooldown
    step_counts = {t: 0 for t in range(1, K + 1)}   # steps trained at each depth
    repair_steps_total = 0            # for duty-cycle cap
    print(f"[repair] knobs: margin={margin} max_repair_steps={max_repair_steps} "
          f"min_frontier={min_frontier_steps} frontier_target={frontier_target_acc} "
          f"max_repair_frac={max_repair_fraction} grad_decay={grad_decay} "
          f"freeze_after={freeze_after_step} frontier_mix={frontier_mix}",
          flush=True)
    # When seeding from a mid-curriculum checkpoint (start stage > 1), the
    # earlier stages graduated in the *source* run so we have no reference for
    # them. Prefer an explicit SUDOKU_GRAD_ACC_SEED (true graduation refs);
    # otherwise seed from the first eval (old behavior — can over-trigger).
    seed_start_stage = int(getattr(config, "curriculum_start_stage", 1))
    grad_acc_seeded = seed_start_stage <= 1
    if (not grad_acc_seeded) and grad_acc_seed_raw.strip():
        try:
            vals = [float(x) for x in grad_acc_seed_raw.replace(";", ":").replace(",", ":").split(":") if x.strip()]
            for s, v in enumerate(vals, start=1):
                if s < seed_start_stage:
                    grad_acc[s] = v
            if grad_acc:
                grad_acc_seeded = True
                print(f"[repair] seeded graduation refs from env: "
                      f"{dict((k, round(v, 3)) for k, v in grad_acc.items())}",
                      flush=True)
        except ValueError:
            print(f"[repair] WARNING: bad SUDOKU_GRAD_ACC_SEED="
                  f"{grad_acc_seed_raw!r}; falling back to first-eval seeding",
                  flush=True)

    def compute_deficits(frontier_stage, per_depth):
        """{depth t: graduation_acc[t] - current_acc} over earlier graduated
        depths whose current candidate-set accuracy is measured.

        Keyed on reasoning depth, not difficulty level: depth t's accuracy is
        how well wave snapshot t is predicted, which is exactly what stage t
        taught. A drop there means that propagation block has been forgotten."""
        d = {}
        for t in range(1, frontier_stage):
            if t in grad_acc and per_depth.get(t, -1.0) >= 0:
                d[t] = grad_acc[t] - per_depth.get(t, -1.0)
        return d

    def effective_margin(per_depth, frontier_stage):
        """Widen the repair trigger while the frontier itself is still weak."""
        m = margin
        if frontier_target_acc > 0:
            f_acc = per_depth.get(frontier_stage, -1.0)
            if 0.0 <= f_acc < frontier_target_acc:
                # Only repair clearer regressions until the frontier is good.
                m = max(m, margin + (frontier_target_acc - f_acc))
        return m

    def most_deficient(frontier_stage, per_depth, use_margin=None):
        """Most-regressed depth whose drop exceeds the margin, else None."""
        d = compute_deficits(frontier_stage, per_depth)
        if not d:
            return None
        m = margin if use_margin is None else use_margin
        t = max(d, key=d.get)
        return t if d[t] > m else None

    with metric_writers.ensure_flushes(writer):
        for step in range(start_step, config.max_steps):
            if step % 10000 == 0:
                print("Step:", step, flush=True)

            i = curriculum.stage
            # Target depth: the frontier when training normally, else the depth
            # we are repairing. Difficulty is never gated, so a repair differs
            # from a frontier step only in the recurrence depth it trains at:
            # replaying depth t re-supervises wave snapshots 1..t on the same
            # full-corpus batch distribution.
            t = repair_stage if mode == "repair" else i
            # Repair must train on THAT stage's instances, not the frontier's.
            # Override is 0-based (instance_stage), t is 1-based curriculum.
            if instance_mode:
                sampler.ds.instance_stage_override = (
                    t - 1 if mode == "repair" else None)
            batch = sampler.sample_all(config.minibatch_size)
            state, metrics = _run_train_step(
                p_frontier[t], state, batch, dropout_rngs)
            step_counts[t] += 1
            if mode == "repair":
                repair_steps_total += 1

            for h in hooks:
                h(step)

            if math.isnan(metrics["loss"][0]):
                print("Loss became nan; stopping.", flush=True)
                break

            if step % config.eval_every_steps == 0:
                eval_metrics = evaluater.get_eval_metrics(
                    state, eval_data_iter, p_eval_step, config)
                per_level = eval_metrics.pop("per_level_acc")
                # Depth accuracy drives the controller. Changed-cells-only, so
                # a slot that merely copies its predecessor scores zero credit.
                per_depth = eval_metrics.pop("per_slot_acc_changed", {})
                per_depth_all = eval_metrics.pop("per_slot_acc", {})
                per_stage_inset = eval_metrics.pop("per_stage_inset_acc", {})
                per_stage_vmass_m = eval_metrics.pop(
                    "per_stage_val_mass_multi", {})
                per_stage_vspread = eval_metrics.pop(
                    "per_stage_val_spread", {})
                per_stage_vhb = eval_metrics.pop("per_stage_val_hbound", {})
                per_stage_ceq_ok = eval_metrics.pop(
                    "per_stage_val_ceq_ok", {})
                per_stage_ceq = eval_metrics.pop("per_stage_val_ceq_pop", {})
                per_stage_ceq_margin = eval_metrics.pop(
                    "per_stage_val_ceq_margin", {})
                per_stage_hinge = eval_metrics.pop(
                    "per_stage_val_hinge", {})
                per_stage_hinge_floor = eval_metrics.pop(
                    "per_stage_val_hinge_floor", {})
                eval_metrics.pop("per_stage_val_hinge_inst", {})
                eval_metrics.pop("per_stage_val_hinge_floor_inst", {})
                eval_metrics.pop("per_stage_val_klq", {})
                eval_metrics.pop("per_stage_val_klq_inst", {})
                eval_metrics.pop("per_stage_val_hinge_eps", {})
                eval_metrics.pop("per_stage_val_hinge_eps_inst", {})
                eval_metrics.pop("per_stage_val_ceq_pop_inst", {})
                eval_metrics.pop("per_stage_val_hq_pop_inst", {})
                eval_metrics.pop("per_stage_val_ceq_margin_inst", {})
                eval_metrics.pop("nest_clean", {})
                eval_metrics.pop("nest_n", {})
                eval_metrics.pop("nest_grew", {})
                eval_metrics.pop("nest_step", {})
                eval_metrics.pop("nest_min", {})
                eval_metrics.pop("nest_A", {})
                eval_metrics.pop("nest_S", {})
                eval_metrics.pop("nest_A1", {})
                eval_metrics.pop("nest_S1", {})
                eval_metrics.pop("keep_all", {})
                eval_metrics.pop("keep_cell", {})
                eval_metrics.pop("keep_depth", {})
                eval_metrics.pop("per_stage_val_ceq", {})
                eval_metrics.pop("per_stage_val_logm_pop", {})
                eval_metrics.pop("per_stage_val_hq_pop", {})
                eval_metrics.pop("per_stage_val_hq", {})
                eval_metrics.pop("per_stage_val_ceq_gap", {})
                eval_metrics.pop("per_stage_val_ceq_puz", {})
                eval_metrics.pop("per_stage_val_hgap", {})
                eval_metrics.pop("per_stage_val_hbound_puz", {})
                eval_metrics.pop("per_stage_val_excess_multi", {})
                eval_metrics.pop("per_stage_val_kl_multi", {})
                eval_metrics.pop("per_stage_val_mass", {})
                eval_metrics.pop("per_stage_val_excess", {})
                eval_metrics.pop("per_stage_val_ce", {})
                eval_metrics.pop("per_stage_val_floor", {})
                if instance_mode:
                    # Repair if mass* on that stage's raw candidate set S
                    # falls below the value it graduated at. This is the
                    # same "is the mass inside S?" number the latent-only
                    # run prints as 1 - out*.
                    per_depth = per_stage_vmass_m
                elif not any(v >= 0 for v in per_depth.values()):
                    per_depth = per_depth_all

                def _m(key):
                    v = eval_metrics.get(key, [])
                    return round(float(np.mean(v)), 4) if len(v) else -1.0

                # Seed graduation refs for stages inherited from a checkpoint.
                if not grad_acc_seeded:
                    for s in range(1, seed_start_stage):
                        acc_s = per_depth.get(s, -1.0)
                        if acc_s >= 0:
                            grad_acc[s] = float(acc_s)
                    grad_acc_seeded = True
                    print(f"[repair] step {step}: seeded graduation refs from "
                          f"resume: "
                          f"{dict((k, round(v, 3)) for k, v in grad_acc.items())}",
                          flush=True)

                # Soft graduation refs: forgive chronic mild regression.
                if grad_decay > 0 and grad_acc:
                    for s in list(grad_acc.keys()):
                        cur_s = per_depth.get(s, -1.0)
                        if cur_s >= 0 and cur_s < grad_acc[s]:
                            old = grad_acc[s]
                            grad_acc[s] = (
                                (1.0 - grad_decay) * grad_acc[s]
                                + grad_decay * float(cur_s))
                            if step % (config.eval_every_steps * 5) == 0:
                                print(f"[repair] soft-grad depth {s}: "
                                      f"{old:.3f}->{grad_acc[s]:.3f} "
                                      f"(cur={cur_s:.3f})", flush=True)

                frontier_acc = per_depth.get(curriculum.stage, -1.0)
                spread_now = float(per_stage_vspread.get(curriculum.stage, -1.0))
                hbound_now = float(per_stage_vhb.get(curriculum.stage, -1.0))
                ceq_ok_now = float(per_stage_ceq_ok.get(curriculum.stage, -1.0))
                ceq_margin_now = float(
                    per_stage_ceq_margin.get(curriculum.stage, -99.0))
                hinge_now = float(per_stage_hinge.get(curriculum.stage, -1.0))
                hinge_floor_now = float(
                    per_stage_hinge_floor.get(curriculum.stage, 0.0))
                out_now = (round(1.0 - frontier_acc, 4)
                           if frontier_acc >= 0 else -1.0)
                eff_margin = effective_margin(per_depth, i)
                deficits = compute_deficits(i, per_depth)
                elapsed = max(1, step - start_step + 1)
                repair_frac = repair_steps_total / float(elapsed)
                bt_frozen = (
                    freeze_after_step > 0 and step >= freeze_after_step)
                print(step, "stage", curriculum.stage,
                      "mode", mode,
                      "repair_stage", repair_stage,
                      "loss", round(float(metrics["loss"].mean()), 4),
                      "out*", out_now,
                      "spread*", round(spread_now, 4),
                      "hbound", round(hbound_now, 4),
                      "ceq", round(float(per_stage_ceq.get(curriculum.stage, -1.0)), 4),
                      "margin", round(ceq_margin_now, 4),
                      "ceq_ok", round(ceq_ok_now, 4),
                      "hinge", round(hinge_now, 4),
                      "hfloor", round(hinge_floor_now, 4),
                      "mass*", round(float(frontier_acc), 4),
                      "eff_margin", round(float(eff_margin), 4),
                      "| grad_acc", {k: round(v, 3) for k, v in grad_acc.items()},
                      "deficits", {k: round(v, 3) for k, v in deficits.items()},
                      "repair_frac", round(repair_frac, 3),
                      flush=True)

                with tf_summary_writer.as_default():
                    tf.summary.scalar("loss", metrics["loss"].mean(), step=step)
                    tf.summary.scalar("ce_loss", metrics["ce_loss"].mean(), step=step)
                    tf.summary.scalar("aux_bce_loss", metrics["aux_loss"].mean(), step=step)
                    tf.summary.scalar("curriculum_stage", curriculum.stage, step=step)
                    tf.summary.scalar("repair_mode", 1 if mode == "repair" else 0, step=step)
                    tf.summary.scalar("repair_stage", repair_stage or 0, step=step)
                    if frontier_acc >= 0:
                        tf.summary.scalar("frontier_depth_acc", frontier_acc, step=step)
                    tf.summary.scalar("eff_repair_margin", eff_margin, step=step)
                    for key in eval_metrics.keys():
                        v = eval_metrics[key]
                        if isinstance(v, dict) or not v:
                            continue
                        tf.summary.scalar("eval_" + key,
                                          np.array(v).mean(), step=step)
                    for lvl, v in per_level.items():
                        if v >= 0:
                            tf.summary.scalar(f"eval_acc_level_{lvl}", v, step=step)
                    for s, v in per_depth_all.items():
                        if v >= 0:
                            tf.summary.scalar(f"eval_cand_depth_{s}", v, step=step)

                def _save_stage_ckpt(tag):
                    if config.save_checkpoint:
                        unrep = jax_utils.unreplicate(state)
                        checkpoints.save_checkpoint_multiprocess(
                            workdir, unrep, step, keep=ckpt_keep, overwrite=True)
                        checkpoints.save_checkpoint_multiprocess(
                            stage_ckpt_dir, unrep, step, keep=100,
                            overwrite=True, prefix=f"{tag}_")

                # ---------------- Deficit-driven controller ----------------
                # Freeze / duty-cycle: force frontier-only when budget exhausted.
                duty_ok = (
                    max_repair_fraction <= 0
                    or repair_frac < max_repair_fraction)
                if bt_frozen and mode == "repair":
                    print(f"[repair] step {step}: freeze_after="
                          f"{freeze_after_step}; leaving repair", flush=True)
                    mode = "frontier"
                    repair_stage = None
                    last_repair_return_step = step
                if (not duty_ok) and mode == "repair":
                    print(f"[repair] step {step}: duty-cycle cap "
                          f"(repair_frac={repair_frac:.3f}>="
                          f"{max_repair_fraction}); return to frontier",
                          flush=True)
                    mode = "frontier"
                    repair_stage = None
                    last_repair_return_step = step

                if mode == "frontier":
                    cooldown_ok = (
                        min_frontier_steps <= 0
                        or (step - last_repair_return_step) >= min_frontier_steps)
                    can_repair = (not bt_frozen) and duty_ok and cooldown_ok
                    worst = (most_deficient(i, per_depth, use_margin=eff_margin)
                             if can_repair else None)
                    if worst is not None:
                        mode = "repair"
                        repair_stage = worst
                        repair_started_at = step
                        repair_episode_start = step
                        print(f"[repair] step {step}: enter repair of depth "
                              f"{worst} (snapshot {worst} acc="
                              f"{per_depth.get(worst, -1.0):.3f} < grad "
                              f"{grad_acc.get(worst, -1.0):.3f} - "
                              f"eff_margin {eff_margin:.3f}; "
                              f"frontier_acc={frontier_acc:.3f})",
                              flush=True)
                    elif bt_frozen and most_deficient(
                            i, per_depth, use_margin=eff_margin) is not None:
                        if step % (config.eval_every_steps * 5) == 0:
                            print(f"[repair] step {step}: deficit present but "
                                  f"BT frozen after {freeze_after_step}",
                                  flush=True)
                    elif (not duty_ok) and most_deficient(
                            i, per_depth, use_margin=eff_margin) is not None:
                        if step % (config.eval_every_steps * 5) == 0:
                            print(f"[repair] step {step}: deficit present but "
                                  f"duty-cycle cap "
                                  f"(frac={repair_frac:.3f})",
                                  flush=True)
                    elif (not cooldown_ok) and most_deficient(
                            i, per_depth, use_margin=eff_margin) is not None:
                        print(f"[repair] step {step}: deficit present but "
                              f"frontier cooldown "
                              f"({step - last_repair_return_step}/"
                              f"{min_frontier_steps}); staying on stage {i}",
                              flush=True)
                    elif curriculum.stage < curriculum.max_stage:
                        # Normal promotion (same rule as the standard loop).
                        steps_in_stage = step - stage_started_at
                        # A negative accuracy means "not measured this eval": it
                        # must not reset the plateau tracker nor satisfy the
                        # threshold/plateau rule. Only patience fires unmeasured.
                        measured = frontier_acc >= 0
                        if measured and frontier_acc > stage_best_acc + plateau_delta:
                            stage_best_acc = frontier_acc
                            stage_best_step = step
                        loc_now = _m("loc_wave")
                        loc_ready = (promote_loc_threshold <= 0.0
                                     or loc_now >= promote_loc_threshold)
                        if instance_mode and promote_ceq_threshold > 0.0:
                            # Single criterion; see train_and_evaluate.py.
                            if promote_ceq_mode == "frac":
                                hit_threshold = (
                                    ceq_ok_now >= 0 and loc_ready
                                    and ceq_ok_now >= promote_ceq_threshold)
                            elif promote_ceq_mode == "hinge":
                                hinge_bar = (
                                    promote_hinge_threshold
                                    if promote_hinge_threshold > 0.0
                                    else max(hinge_floor_now, 0.0)
                                    + promote_hinge_slack)
                                hit_threshold = (
                                    hinge_now >= 0.0 and loc_ready
                                    and hinge_now <= hinge_bar)
                            else:
                                hit_threshold = (
                                    ceq_margin_now > -99.0 and loc_ready
                                    and ceq_margin_now >= promote_ceq_margin)
                            stalled = False
                            patience_over = False
                        elif instance_mode:
                            hbound_ready = (
                                promote_hbound_threshold <= 0.0
                                or hbound_now >= promote_hbound_threshold)
                            hit_threshold = (measured and loc_ready
                                             and frontier_acc >= promote_threshold
                                             and spread_now >= promote_spread_threshold
                                             and hbound_ready)
                            stalled = False
                            patience_over = False
                        else:
                            hit_threshold = (measured and loc_ready
                                             and frontier_acc >= promote_threshold)
                            stalled = (measured and loc_ready and plateau_steps > 0
                                       and (step - stage_best_step) >= plateau_steps)
                            patience_over = (promote_patience > 0
                                             and steps_in_stage >= promote_patience)
                        if steps_in_stage >= min_stage_steps and (
                                hit_threshold or stalled or patience_over):
                            reason = ("threshold" if hit_threshold
                                      else "plateau" if stalled else "patience")
                            # Record this stage's graduation mass* BEFORE moving on.
                            grad_acc[curriculum.stage] = float(frontier_acc)
                            curriculum.stage += 1
                            stage_started_at = step
                            stage_best_acc = -1.0
                            stage_best_step = step
                            if instance_mode and promote_ceq_threshold > 0.0:
                                extra = (f" margin={ceq_margin_now:+.4f}"
                                         f" ceq_ok={ceq_ok_now:.3f}"
                                         f" hinge={hinge_now:.4f}"
                                         f" hfloor={hinge_floor_now:.4f}")
                            elif instance_mode:
                                extra = (f" mass*={frontier_acc:.3f}"
                                         f" spread*={spread_now:.3f}"
                                         f" hbound={hbound_now:.3f}")
                            else:
                                extra = f" acc={frontier_acc:.3f}"
                            print(f"[curriculum] step {step}: promote to stage "
                                  f"{curriculum.stage} ({reason}; graduated "
                                  f"depth {curriculum.stage - 1}{extra})",
                                  flush=True)
                            _save_stage_ckpt(f"stage{curriculum.stage}")
                    elif (frontier_target_acc > 0
                          and 0.0 <= frontier_acc < frontier_target_acc
                          and step % (config.eval_every_steps * 5) == 0):
                        print(f"[frontier] step {step}: depth {i} "
                              f"acc={frontier_acc:.3f} "
                              f"< target {frontier_target_acc:.3f}; "
                              f"keeping frontier priority",
                              flush=True)
                else:  # mode == "repair"
                    r = repair_stage
                    cur = per_depth.get(r, -1.0)
                    ref = grad_acc.get(r, -1.0)
                    recovered = cur >= (ref - margin)
                    capped = (step - repair_started_at) >= max_repair_steps
                    if recovered or capped:
                        why = "recovered" if recovered else "cap"
                        print(f"[repair] step {step}: depth {r} done ({why}; "
                              f"snapshot acc={cur:.3f} vs grad {ref:.3f})",
                              flush=True)
                        _save_stage_ckpt(f"repair{r}")
                        # Re-scan: chain to the next most-deficient stage, or
                        # return to the frontier and credit repair time back to
                        # the frontier stage's patience clock.
                        nxt = None
                        if (not bt_frozen) and duty_ok:
                            nxt = most_deficient(
                                i, per_depth, use_margin=eff_margin)
                        if nxt is not None:
                            repair_stage = nxt
                            repair_started_at = step
                            print(f"[repair] step {step}: chain to stage {nxt}",
                                  flush=True)
                        else:
                            mode = "frontier"
                            repair_stage = None
                            last_repair_return_step = step
                            stage_started_at += (step - repair_episode_start)

            if config.save_checkpoint and step > 0 and step % config.save_every_steps == 0:
                checkpoints.save_checkpoint_multiprocess(
                    workdir, jax_utils.unreplicate(state), step,
                    keep=ckpt_keep, overwrite=True)

        if config.save_checkpoint:
            checkpoints.save_checkpoint_multiprocess(
                workdir, jax_utils.unreplicate(state), config.max_steps,
                keep=ckpt_keep, overwrite=True)