onf-c2/src/onf/graph/train/loop.py

"""Train the graph retriever (onf.graph.net.gnn.GraphRetrieverNet) -- the curriculum for the piece
onf.graph.run.retrieve runs at deploy time. Design writeup: docs/technical/02-retrieval-head.md.

One stage trains the retriever under onf.graph.train.loss's selected objective, cfg.qbatch queries
per optimiser step. Whichever objective is trained, the epoch kept is the one with the lowest
held-out mse(a_exec) -- the error of the chunk the deployed controller would actually emit,
averaged within perturbation-radius strata so a radius the held-out draw happens to over-sample
cannot decide the checkpoint. MRR is a logged diagnostic and is allowed to fall.

Under the CHUNK objective the retriever is frozen and only onf.blend.alpha.AlphaNet is optimised,
against query sets the onf.graph.train.corrupt sampler has given deliberately wrong retrievals on;
it is written beside the head as g_alpha.npz.

Checkpoints go to g_head.npz stamped with
onf.graph.core.geometry.graph_hash, so a head trained against one graph can never be silently
loaded against another -- a mismatch would still return confident, well-formed logits, just over
the wrong demonstration strands.
"""

from __future__ import annotations

import os
from dataclasses import dataclass, field as dc_field, replace
from pathlib import Path
from typing import Any, Iterator

import numpy as np
import torch

from onf.config import GraphConfig, default_paths
from onf.device import resolve_device
from onf.field.field import ONFField
from onf.graph.core import schema
from onf.graph.core.edges import EdgeSet
from onf.graph.net.gnn import C2_PREFIX, GraphRetrieverNet
from onf.graph.net.language import LanguageHead
from onf.graph.core.nodes import NodeTable
from onf.graph.core.geometry import CleanlinessConfig, load_cleanliness_field
from onf.graph.core.geometry import CONSTANTS_JSON, graph_hash

from onf.graph.train.data import (
    DRIFT_FRAC,
    ENTRY_STATIC_FRAC,
    NegativeMiner,
    SoftTargetBuilder,
    _node_feature_matrix,
    make_queries,
    measure_inter_demo_spacing,
)
from onf.blend.alpha import AlphaNet
from onf.blend.ee_track import POS_BLOCK, ROT_BLOCK, ActionScale
from onf.blend.kinematics import PandaKinematics
from onf.blend.target import ChunkTargetBuilder, ChunkTargets, PolicyActionCache, ee_segment_table
from onf.graph.train.corrupt import (
    KIND_NAMES, OFFSET, POLICY_KIND_NAMES, Corruption, PolicyCorruptor, RetrievalCorruptor,
)
from onf.graph.train.loss import (
    CHUNK, DEFAULT_LOSS_SPEC, ChunkLossComputer, _query_loss_batch,
)
from onf.graph.train.metrics import _gnn_metrics, evaluate
from onf.graph.train.types import EvalQueries, GraphTensors, LossSpec, QueryBatch, QueryData, QuerySpec, StageSpec, TrainGraphSpec

# ---- this stage's own hyper-parameters -- see docs/technical/02-retrieval-head.md's constants table. The
# query-construction, loss and held-out-eval constants live in their own sub-modules.
N_QUERY = 4_096           # queries generated per stage
N_EVAL = 128               # held-out queries used for the per-epoch GNN metrics report
HELDOUT_STRIDE = 5                   # hold out every 5th demo strand -- ONFTrainConfig.heldout_stride's convention
LR, WEIGHT_DECAY = 1e-3, 1e-5

RADIUS_STRATA = 3         # perturbation-radius strata the perturbed held-out rows are split into,
                           # by quantile, alongside a fourth stratum for the clean rows
# Objectives that freeze the retriever and train the AlphaNet instead.

_OFFSET_BINS = 5          # log-spaced displacement bins in AlphaStrata's calibration curve. Five
                           # over the sampler's two decades is roughly one bin per half-decade, which
                           # is the finest split the held-out OFFSET count (~13 rows of 128) supports.
ACTION_SCALE_JSON = "action_scale.json"


def save_checkpoint(net: GraphRetrieverNet, nodes: NodeTable, edges: EdgeSet, path: str | os.PathLike) -> str:
    """Write the head to path, stamped with the graph hash.

    Args:
        net: The trained retriever.
        nodes: Node table it was trained against.
        edges: Its edge set.
        path: Destination, typically <graph_dir>/g_head.npz.

    Returns:
        The path written.
    """
    out = net.state_npz_dict()
    path = os.fspath(path)
    out["graph_hash"] = np.array(graph_hash(
        nodes, edges, constants=Path(path).parent / CONSTANTS_JSON, c2_state=net.c2_state(),
    ))
    Path(path).parent.mkdir(parents=True, exist_ok=True)
    np.savez(path, **out)
    return path


def _prepare_tensors(
    nodes: NodeTable, edges: EdgeSet, dev: Any, *, cfg: GraphConfig | None = None,
    chunk: ChunkTargetBuilder | None = None, language: Any = None,
) -> GraphTensors:
    """Move the graph to dev once and build the phase-bin basis the WHEN-readout loss terms use.

    Args:
        nodes: Node table to prepare.
        edges: Its edge set.
        dev: Target device.
        cfg: Graph configuration; None leaves SoftTargetBuilder on its own defaults.
        chunk: Chunk-target builder, whose action scale and segment length the [V, K, 6]
            end-effector table is built from. None leaves that table unbuilt.
        language: onf.graph.lang.TaskEmbeddings, or None. Its table and the per-node task ids are
            uploaded here because both are graph-scoped and the seed prior needs them every step.

    Returns:
        The bundle threaded through every stage/epoch/batch of one run.
    """
    node_x_t = torch.as_tensor(_node_feature_matrix(nodes), device=dev)
    et = edges.torch(dev)
    src, dst, rel, log_idf = et["src"].long(), et["dst"].long(), et["rel"], et["log_idf"]
    # This is the only place in the codebase that backpropagates through net.propagate at scale, so
    # the only place the relation-grouped view is worth its extra forward cost.
    src_by_rel, dst_by_rel, rel_offsets = et["src_by_rel"].long(), et["dst_by_rel"].long(), et["rel_offsets"]
    phase_t = torch.as_tensor(nodes.phase.astype(np.float32), device=dev)
    # NOT onf.graph.core.binning.phase_bins: float32 (not f64) and no lower clamp. Kept as-is because
    # the trained head was fit against these bins; switching to the f64 path would move them.
    bin_idx_t = torch.clamp((phase_t * schema.NBINS_ALIGN).long(), max=schema.NBINS_ALIGN - 1)
    return GraphTensors(
        nodes=nodes, edges=edges, node_x_t=node_x_t, src=src, dst=dst, rel=rel, log_idf=log_idf,
        phase_t=phase_t, bin_idx_t=bin_idx_t, src_by_rel=src_by_rel, dst_by_rel=dst_by_rel,
        rel_offsets=rel_offsets,
        # The two row-level samplers, bound to this graph once per run rather than rebuilt per row.
        targets=SoftTargetBuilder(
            nodes,
            **({} if cfg is None else
               {"phase_band": cfg.where_phase_band, "move_temp": cfg.where_move_temp}),
        ),
        miner=NegativeMiner(nodes),
        ee_table=None if chunk is None else ee_segment_table(nodes, chunk.chunk_len).to(dev),
        action_scale=None if chunk is None else chunk.scale,
        e_task=(
            None if language is None
            else torch.as_tensor(language.e_task, dtype=torch.float32, device=dev)
        ),
        task_id_t=(
            None if language is None
            else torch.as_tensor(np.asarray(nodes.task_id, dtype=np.int64), device=dev)
        ),
    )


class EpochReporter:
    """Formats and logs one stage's per-epoch line and its final summary, and accumulates history.

    Best-epoch selection is a separate concern; see BestCheckpoint.
    """

    def __init__(self, logger: Any, stage_name: str, epochs: int):
        self.logger = logger
        self.stage_name = stage_name
        self.epochs = epochs
        self.history: list[dict] = []

    def log(self, msg: str) -> None:
        (self.logger.log(msg) if self.logger is not None else print(msg))

    def record(
        self, ep: int, mean_loss: float, metrics: dict, *,
        guard_mrr: float | None = None, chunk_mse: float | None = None,
        static_metrics: dict | None = None,
    ) -> None:
        """Log and record one epoch.

        Args:
            ep: Zero-based epoch index.
            mean_loss: This epoch's mean per-query loss.
            metrics: Its held-out metrics.
            guard_mrr: Its guard-set MRR, or None when no guard set is in use.
            chunk_mse: Its radius-stratified held-out mse(a_exec), or None when no chunk targets
                were built. This is the model-selection score.
            static_metrics: The ENTRY_STATIC acceptance metrics, or None.
        """
        line = (
            f"[graph.train {self.stage_name}] epoch {ep + 1}/{self.epochs} loss={mean_loss:.4f} "
            f"top1={metrics['top1']:.3f} top10={metrics['top10']:.3f} mrr={metrics['mrr']:.3f} "
            f"when_mae={metrics['when_mae_steps']:.2f} where_rmse={metrics['where_rmse']:.4f} "
            f"move_dist={metrics['move_dist_mean']:.4f} "
            f"phase_argmax[corr/mae]={metrics['when_phase_argmax_corr']:.3f}/"
            f"{metrics['when_phase_argmax_mae']:.3f} "
            f"phase_expect[corr/mae]={metrics['when_phase_expect_corr']:.3f}/"
            f"{metrics['when_phase_expect_mae']:.3f} "
            f"crossed_top1={metrics['crossed_top1']:.3f} "
            f"abstain_rate={metrics['abstain_rate']:.3f} "
            f"abstain_false_fire={metrics['abstain_false_fire_rate']:.3f} "
            f"abstain_recall={metrics['abstain_recovery_recall']:.3f} "
            f"abstain_delta_med[abstain/fire]="
            f"{metrics['abstain_delta_median_abstain_class']:.2f}/"
            f"{metrics['abstain_delta_median_fire_class']:.2f} "
            f"abstain_delta_iqr[abstain/fire]="
            f"{metrics['abstain_delta_iqr_abstain_class']:.2f}/"
            f"{metrics['abstain_delta_iqr_fire_class']:.2f}"
        ) + ("" if guard_mrr is None else f" guard_mrr={guard_mrr:.4f}") + (
            "" if chunk_mse is None else f" chunk_mse={chunk_mse:.5f}"
        ) + (
            "" if static_metrics is None
            else f" static_kl={static_metrics['static_kl']:.4f} static_bary_err={static_metrics['static_bary_err']:.4f}"
        )
        if self.logger is not None:
            self.logger.log(line)
            self.logger.metric(**{f"{self.stage_name}_ep{ep}_loss": mean_loss},
                                **{f"{self.stage_name}_ep{ep}_{k}": v for k, v in metrics.items()},
                                **({} if guard_mrr is None else {f"{self.stage_name}_ep{ep}_guard_mrr": guard_mrr}),
                                **({} if chunk_mse is None else {f"{self.stage_name}_ep{ep}_chunk_mse": chunk_mse}),
                                **({} if static_metrics is None else {
                                    f"{self.stage_name}_ep{ep}_static_kl": static_metrics["static_kl"],
                                    f"{self.stage_name}_ep{ep}_static_bary_err": static_metrics["static_bary_err"],
                                }))
        else:
            print(line)
        self.history.append({
            "epoch": ep, "loss": mean_loss, "guard_mrr": guard_mrr, "chunk_mse": chunk_mse,
            "static_kl": None if static_metrics is None else static_metrics["static_kl"],
            "static_bary_err": None if static_metrics is None else static_metrics["static_bary_err"],
            **metrics,
        })

    def record_final(self, final: dict, final_static: dict | None) -> None:
        """Log and record the stage's closing summary line.

        Args:
            final: The restored best epoch's held-out metrics.
            final_static: Its ENTRY_STATIC acceptance metrics, or None.
        """
        summary = (
            f"[graph.train {self.stage_name}] done: top1={final['top1']:.3f} top10={final['top10']:.3f} "
            f"mrr={final['mrr']:.3f} when_mae={final['when_mae_steps']:.2f} "
            f"where_rmse={final['where_rmse']:.4f} move_dist={final['move_dist_mean']:.4f} "
            f"phase_argmax[corr/mae]={final['when_phase_argmax_corr']:.3f}/{final['when_phase_argmax_mae']:.3f} "
            f"phase_expect[corr/mae]={final['when_phase_expect_corr']:.3f}/{final['when_phase_expect_mae']:.3f} "
            f"crossed_top1={final['crossed_top1']:.3f} | "
            f"abstain_rate={final['abstain_rate']:.3f} "
            f"abstain_false_fire={final['abstain_false_fire_rate']:.3f} "
            f"abstain_recall={final['abstain_recovery_recall']:.3f} "
            f"abstain_delta_med[abstain/fire]="
            f"{final['abstain_delta_median_abstain_class']:.2f}/{final['abstain_delta_median_fire_class']:.2f}"
        ) + (
            "" if final_static is None
            else f" static_kl={final_static['static_kl']:.4f} static_bary_err={final_static['static_bary_err']:.4f}"
        )
        if self.logger is not None:
            self.logger.log(summary)
            self.logger.metric(**{f"{self.stage_name}_final_{k}": v for k, v in final.items()})
        else:
            print(summary)
        self.history.append({"epoch": "final", **final})


@dataclass(frozen=True, eq=False)
class RadiusStrata:
    """A partition of one query set's rows by how far its window was perturbed.

    The clean rows are their own stratum because their radius is exactly 0 and quantiles over a
    mixture with a point mass at 0 do not separate them.

    Attributes:
        label: [n] int64 stratum index per row.
        n_strata: Number of strata, clean included.
    """

    label: np.ndarray
    n_strata: int

    @classmethod
    def from_radii(cls, radius: np.ndarray, n_bins: int = RADIUS_STRATA) -> "RadiusStrata":
        """Bin rows into a clean stratum plus n_bins quantile strata of the perturbed ones.

        Args:
            radius: [n] realised perturbation radius per row, rad.
            n_bins: Quantile strata over the perturbed rows.

        Returns:
            The partition.
        """
        r = np.asarray(radius, dtype=np.float64)
        label = np.zeros(len(r), dtype=np.int64)
        perturbed = r > 0.0
        if perturbed.any():
            edges = np.quantile(r[perturbed], np.linspace(0.0, 1.0, n_bins + 1)[1:-1])
            label[perturbed] = 1 + np.searchsorted(edges, r[perturbed], side="right")
        return cls(label=label, n_strata=n_bins + 1)

    def mean(self, values: np.ndarray) -> float:
        """Average within each non-empty stratum, then across strata.

        Args:
            values: [n] per-row quantity.

        Returns:
            The unweighted mean of the per-stratum means, or NaN when there are no rows.
        """
        per_stratum = [
            float(np.asarray(values)[self.label == s].mean())
            for s in range(self.n_strata)
            if np.any(self.label == s)
        ]
        return float(np.mean(per_stratum)) if per_stratum else float("nan")


@dataclass
class BestCheckpoint:
    """Tracks the lowest-scoring epoch of one stage, and its weights.

    The last epoch is not always the best one. The score is the held-out radius-stratified
    mse(a_exec): the error of the chunk the deployed controller would emit. A run with no chunk
    targets admits no epoch at all, so it keeps its final weights -- silently substituting a
    retrieval proxy would select on a quantity the deployed system does not optimise.

    The module is a parameter rather than the retriever: Stage 2 freezes the retriever and trains
    the AlphaNet, so snapshotting the retriever there would restore nothing.

    Attributes:
        score: The winning epoch's score, or None before any epoch is admitted.
        epoch: Its zero-based index, -1 before any epoch is admitted.
        state: Its cloned state dict, or None.
    """

    score: float | None = None
    epoch: int = -1
    state: dict[str, torch.Tensor] | None = None

    def consider(self, net: torch.nn.Module, ep: int, score: float | None) -> None:
        """Snapshot net as the new best if this epoch scores lowest.

        Args:
            net: Module whose weights would be snapshotted -- whichever one this stage trains.
            ep: Zero-based epoch index.
            score: This epoch's held-out chunk MSE, or None when it could not be measured.
        """
        if score is None or not np.isfinite(score):
            return
        if self.score is not None and score >= self.score:
            return
        self.score, self.epoch = float(score), ep
        self.state = {k: v.detach().clone() for k, v in net.state_dict().items()}

    def restore_into(self, net: torch.nn.Module, final_epoch: int) -> bool:
        """Load the winning weights back into net, unless the final epoch already won.

        Args:
            net: Module to load into.
            final_epoch: Zero-based index of the stage's last epoch.

        Returns:
            True if weights were actually restored.
        """
        if self.state is None or self.epoch == final_epoch:
            return False
        net.load_state_dict(self.state)
        return True


@dataclass(frozen=True)
class QuerySet:
    """A stage's training queries on device, plus the class-balance weights derived from them.

    Attributes:
        qh: Joint positions [n, H, D].
        vh: Joint velocities [n, H, D].
        gh: Gripper states [n, H].
        w: Per-step cleanliness weights [n, H].
        tgt: Target node ids [n].
        src_v: Source node ids [n].
        abstain_is_correct: Abstain consequence label [n] bool.
        abstain_weight: Per-row inverse-frequency class weight [n].
        is_entry_static: ENTRY_STATIC row mask [n] bool.
        is_drift: DRIFT row mask [n] bool.
        break_row: First perturbed window row [n] i32, -1 when the whole window is clean.
        perturb_radius: Realised offset radius at the window's end [n] f32, 0 when clean.
        targets: The chunk-blend inputs and target, or None when they were not built.
        corrupt: Which rows retrieval is deliberately wrong on, or None when none were drawn.
        policy_corrupt: Which rows the POLICY's chunk is deliberately wrong on. Already folded
            into targets.a_policy; carried so the per-stratum alpha table can group by it.
        e_lang: [n, lang_dim] per-row instruction embedding on device, or None. Drawn ONCE with the
            corruption and for the same reason: a row whose instruction is resampled every epoch is
            a noise process, and the only fit consistent with all of its draws is the mean.
    """

    qh: torch.Tensor
    vh: torch.Tensor
    gh: torch.Tensor
    w: torch.Tensor
    tgt: np.ndarray
    src_v: np.ndarray
    abstain_is_correct: np.ndarray
    abstain_weight: np.ndarray
    is_entry_static: np.ndarray
    is_drift: np.ndarray
    break_row: np.ndarray
    perturb_radius: np.ndarray
    targets: ChunkTargets | None = None
    corrupt: Corruption | None = None
    policy_corrupt: Any = None
    e_lang: torch.Tensor | None = None

    @classmethod
    def from_queries(
        cls, queries: Any, dev: Any, builder: ChunkTargetBuilder | None = None,
        corruptor: RetrievalCorruptor | None = None, rng: Any = None,
        language: Any = None, task_id: np.ndarray | None = None,
        policy_corruptor: PolicyCorruptor | None = None,
    ) -> "QuerySet":
        """Upload one make_queries result, derive its abstain weights and build its chunk targets.

        Class-balance insurance: w_c = n / (2 * n_c) from the query set's OWN realised abstain/fire
        split, so an unweighted term cannot let the majority class dominate the loss regardless of
        how the drift generator happens to be calibrated.

        The corruption is drawn ONCE, here, and then only sliced. Redrawing it per epoch would make
        a corrupted row a noise process rather than an example: alpha would see the same query
        paired with a wrong demo on one epoch and its own on the next, and the only fit consistent
        with both is the average.

        Args:
            queries: A query mapping from make_queries. A mapping predating the ENTRY_STATIC or
                DRIFT feature simply has every row treated as ordinary, a strict no-op.
            dev: Target device.
            builder: Chunk-target builder, or None to leave targets unbuilt.
            corruptor: Deliberately-wrong-retrieval sampler, or None to leave every row clean. None
                draws no randomness at all, so a run that passes none has the rng stream it always had.
            rng: Random state the corruption is drawn from; required with corruptor.
            language: onf.graph.lang.TaskEmbeddings, or None to build no instruction embeddings.
                Requires rng and task_id.
            task_id: [V] per-node task id on the host, used to look up each target's task.
            policy_corruptor: Deliberately-wrong-POLICY sampler, or None to leave every row's
                policy chunk clean. Drawn INDEPENDENTLY of corruptor, so the head sees all four
                quadrants of (retrieval right/wrong) x (policy right/wrong); drawing them jointly
                would leave one empty and leave the two indistinguishable. Needs rng.

        Returns:
            The device-resident query set.

        Raises:
            ValueError: A corruptor was given without an rng to draw it from, or a language table
                without the rng and task ids it needs.
        """
        if (corruptor is not None or policy_corruptor is not None) and rng is None:
            raise ValueError("QuerySet.from_queries: a corruptor needs an rng to draw from")
        if language is not None and (rng is None or task_id is None):
            raise ValueError("QuerySet.from_queries: a language table needs an rng and task_id")
        tgt = queries["tgt_node"]
        n = len(tgt)
        zeros = np.zeros(n, dtype=bool)
        # The loss label is the CONSEQUENCE label, not the is_clean provenance flag.
        abstain = np.asarray(queries.get("abstain_is_correct", zeros), dtype=bool)
        n_abstain = int(abstain.sum())
        n_fire = n - n_abstain
        w_abstain = n / (2.0 * n_abstain) if n_abstain else 0.0
        w_fire = n / (2.0 * n_fire) if n_fire else 0.0
        policy_corrupt = (
            None if policy_corruptor is None
            else policy_corruptor.draw(tgt, np.asarray(queries["t_raw_now"], dtype=np.int64), rng)
        )
        return cls(
            qh=torch.as_tensor(queries["q_hist"], device=dev),
            vh=torch.as_tensor(queries["qdot_hist"], device=dev),
            gh=torch.as_tensor(queries["grip_hist"], device=dev),
            w=torch.as_tensor(queries["w_hist"], device=dev),
            tgt=tgt,
            src_v=queries["src_node"],
            abstain_is_correct=abstain,
            abstain_weight=np.where(abstain, w_abstain, w_fire).astype(np.float64),
            is_entry_static=np.asarray(queries.get("is_entry_static", zeros), dtype=bool),
            is_drift=np.asarray(queries.get("is_drift", zeros), dtype=bool),
            break_row=np.asarray(queries.get("break_row", np.full(n, -1)), dtype=np.int32),
            perturb_radius=np.asarray(
                queries.get("perturb_radius", np.zeros(n)), dtype=np.float32
            ),
            targets=None if builder is None else builder.build(queries, policy_corrupt),
            corrupt=None if corruptor is None else corruptor.draw(tgt, rng),
            policy_corrupt=policy_corrupt,
            e_lang=cls._draw_instructions(language, tgt, task_id, rng, dev),
        )

    @staticmethod
    def _draw_instructions(
        language: Any, tgt: np.ndarray, task_id: np.ndarray | None, rng: Any, dev: Any,
    ) -> torch.Tensor | None:
        """One paraphrase surrogate of each row's own task, as an embedding.

        Args:
            language: The TaskEmbeddings table, or None.
            tgt: [n] target node ids.
            task_id: [V] per-node task id on the host.
            rng: Random state.
            dev: Target device.

        Returns:
            [n, lang_dim] float32 on device, or None when no table was given.

        The surrogate is drawn from the target's OWN task, never a wrong one. The corruption sampler
        already supplies wrong-task rows by swapping what retrieval RETURNS; making the instruction
        wrong as well would describe a different situation -- the operator asked for the wrong thing
        -- which is not a failure this system is meant to survive.
        """
        if language is None:
            return None
        bank = np.asarray(language.e_aug if language.e_aug.size else language.e_task[:, None, :])
        rows = np.asarray(task_id, dtype=np.int64)[np.asarray(tgt, dtype=np.int64)]
        # rng.random, not rng.integers: training threads a legacy np.random.RandomState, whose API
        # has randint rather than integers, while corrupt.py is handed a Generator in tests. random
        # is the one uniform draw both expose under the same name.
        which = np.minimum(
            (np.asarray(rng.random(rows.shape[0])) * bank.shape[1]).astype(np.int64),
            bank.shape[1] - 1,
        )
        return torch.as_tensor(bank[rows, which], dtype=torch.float32, device=dev)

    def __len__(self) -> int:
        return len(self.tgt)

    @property
    def strata(self) -> RadiusStrata:
        """This set's perturbation-radius partition."""
        return RadiusStrata.from_radii(self.perturb_radius)

    @property
    def device(self) -> Any:
        """The device every tensor here lives on."""
        return self.qh.device

    def balance_summary(self) -> str:
        """One-line abstain/fire class-balance report, plus the corruption branch counts.

        Returns:
            The log line. Its abstain/fire half matches the pre-existing format exactly; a set with
            no corruption drawn adds nothing to it.
        """
        n = len(self)
        n_abstain = int(self.abstain_is_correct.sum())
        n_fire = n - n_abstain
        w_abstain = float(self.abstain_weight[self.abstain_is_correct].max(initial=0.0))
        w_fire = float(self.abstain_weight[~self.abstain_is_correct].max(initial=0.0))
        corrupt = "" if self.corrupt is None else " | corruption: " + " ".join(
            f"{name}={count}" for name, count in self.corrupt.counts().items()
        )
        corrupt += "" if self.policy_corrupt is None else " | policy: " + " ".join(
            f"{name}={count}" for name, count in self.policy_corrupt.counts().items()
        )
        return (f"class balance: abstain={n_abstain}/{n} ({100 * n_abstain / max(n, 1):.1f}%) "
                f"fire={n_fire}/{n} ({100 * n_fire / max(n, 1):.1f}%) "
                f"-> loss weights {w_abstain:.3f}/{w_fire:.3f}") + corrupt

    def shuffled_batches(self, qbatch: int) -> "Iterator[QueryBatch]":
        """Yield contiguous chunks of one fresh torch.randperm permutation.

        torch.randperm specifically, NOT RandomSampler/BatchSampler: this consumes the
        process-global torch RNG that torch.manual_seed pinned, and switching samplers would draw a
        different permutation from the same seed. test_train_graph_is_deterministic compares two
        runs' weights exactly.

        Args:
            qbatch: Maximum queries per optimiser step. The last chunk of an epoch may be shorter;
                the loss takes its MEAN over however many rows it is given, so it stays weighted right.

        Yields:
            One onf.graph.train.types.QueryBatch per optimiser step.
        """
        n_q = len(self)
        perm = torch.randperm(n_q).numpy()
        for start in range(0, n_q, qbatch):
            yield self.batch(perm[start : start + qbatch])

    def ordered_batches(self, qbatch: int) -> "Iterator[QueryBatch]":
        """Yield the set in row order, drawing no randomness. Used by the held-out evaluation.

        Args:
            qbatch: Maximum queries per batch.

        Yields:
            One onf.graph.train.types.QueryBatch per chunk, covering every row exactly once.
        """
        n_q = len(self)
        for start in range(0, n_q, qbatch):
            yield self.batch(np.arange(start, min(start + qbatch, n_q)))

    def batch(self, idx: np.ndarray) -> QueryBatch:
        """Slice the set down to the given rows.

        Args:
            idx: [b] row indices.

        Returns:
            The device-resident batch.
        """
        idx_t = torch.as_tensor(idx, dtype=torch.long, device=self.device)
        chunk, corrupt = self.targets, self.corrupt
        pol = self.policy_corrupt
        return QueryBatch(
            qh=self.qh[idx_t], vh=self.vh[idx_t], gh=self.gh[idx_t], w=self.w[idx_t],
            tgt=self.tgt[idx], src_v=self.src_v[idx],
            abstain_is_correct=self.abstain_is_correct[idx],
            abstain_weight=self.abstain_weight[idx],
            is_entry_static=self.is_entry_static[idx],
            is_drift=self.is_drift[idx],
            break_row=self.break_row[idx],
            perturb_radius=self.perturb_radius[idx],
            a_policy=None if chunk is None else torch.as_tensor(chunk.a_policy[idx], device=self.device),
            a_corr=None if chunk is None else torch.as_tensor(chunk.a_corr[idx], device=self.device),
            ee_now=None if chunk is None else torch.as_tensor(chunk.ee_now[idx], device=self.device),
            ee_ref=None if chunk is None else torch.as_tensor(chunk.ee_ref[idx], device=self.device),
            corrupt_node=None if corrupt is None else torch.as_tensor(corrupt.node[idx], device=self.device),
            corrupt_offset=(
                None if corrupt is None else torch.as_tensor(corrupt.offset[idx], device=self.device)
            ),
            # int8 stays on the host: nothing reads it in the loss, only the alpha diagnostics.
            corrupt_kind=None if corrupt is None else corrupt.kind[idx],
            policy_kind=None if pol is None else pol.kind[idx],
            e_lang=None if self.e_lang is None else self.e_lang[idx_t],
        )


@dataclass(frozen=True)
class AlphaStrata:
    """Held-out mean blend weight split by corruption branch -- Stage 2's go/no-go gate.

    The design predicts a LOW weight on wrong_task and offset, where the retrieved demonstration is
    invalid and following it is worse than doing nothing, and a HIGHER one on clean, where following
    it is the entire point. A run whose branches do not separate here has not learned a
    compatibility function, only a displacement monotone, and is not worth a GPU rollout.

    The position/rotation split is the sharpest of these numbers. On the offset branch the object
    has moved, so the demo's position guidance is wrong while its wrist orientation is still right:
    position alpha should collapse there and rotation alpha should not.

    Attributes:
        mean_alpha: Branch name -> mean weight over that branch's rows, NaN when it has none.
        mean_pos: The same over the position block alone (action dims 0:3).
        mean_rot: The same over the rotation block alone (action dims 3:6).
        counts: Branch name -> row count.
        mean_policy: POLICY-branch name -> mean weight, the SECOND direction. The prediction here
            is the opposite sign to mean_alpha's: alpha should be HIGHER on lag/policy_task/gain
            than on clean, because there the demonstration is the only operand still right. A run
            that separates the retrieval branches downward but leaves the policy branches flat has
            learned "disagreement means distrust the demo", which is the fitted-with-the-wrong-sign
            failure the branch exists to remove.
        policy_counts: POLICY-branch name -> row count.
        row_first: Mean weight on chunk row 0, the feedback row.
        row_rest: Mean weight over chunk rows 1..K-1, the demo feedforward, or NaN when K == 1.
        offset_curve: (centre displacement in metres, mean POSITION weight, row count) per
            log-spaced bin of the OFFSET branch -- the calibration curve.

            The branch mean above is a single number and it turned out not to predict anything: the
            deployed v4 head scored 0.504 on offset against 0.639 clean, a 1.3x separation where its
            predecessor managed 6.8x, and v4 is the one that won the axis. A single number cannot
            distinguish "insensitive to the corruption" from "correctly near-indifferent because the
            corruption is small", and with the old one-band uniform sampler those were the same
            measurement. Binned against a log-uniform magnitude the two separate: what the axis needs
            is alpha DECREASING in displacement, and a curve shows that where a mean cannot.
    """

    mean_alpha: dict[str, float]
    mean_pos: dict[str, float]
    mean_rot: dict[str, float]
    counts: dict[str, int]
    row_first: float
    row_rest: float
    offset_curve: tuple[tuple[float, float, int], ...] = ()
    mean_policy: dict[str, float] = dc_field(default_factory=dict)
    policy_counts: dict[str, int] = dc_field(default_factory=dict)

    @classmethod
    def measure(
        cls, computer: ChunkLossComputer, gt: GraphTensors, queries: "QuerySet", qbatch: int,
        curve_set: "QuerySet | None" = None,
    ) -> "AlphaStrata":
        """Run the CHUNK forward pass over a whole query set and stratify its weights.

        Args:
            computer: The chunk-loss computer holding the retriever and the alpha head.
            gt: The device-resident graph bundle.
            queries: The set the branch means are measured on, normally the held-out one.
            qbatch: Rows per forward pass.
            curve_set: The set the calibration curve is measured on instead, when the held-out set
                is too small to bin. The branch means are a generalisation claim and stay held out;
                the curve is a description of what the head learned, so a larger in-sample set is
                the better instrument -- the held-out 128 rows carry only ~13 OFFSET rows, which is
                two per bin.

        Returns:
            The measurement. A set with no corruption drawn reports every row as clean.
        """
        alpha = cls._alpha_over(computer, gt, queries, qbatch)            # [n, K, 6]
        kind = cls._kinds(queries)

        def by_branch(block: slice) -> dict[str, float]:
            per_row = alpha[..., block].mean(axis=(1, 2))
            return {
                name: float(per_row[kind == k].mean()) if np.any(kind == k) else float("nan")
                for k, name in enumerate(KIND_NAMES)
            }

        return cls(
            mean_alpha=by_branch(slice(None)),
            mean_pos=by_branch(POS_BLOCK),
            mean_rot=by_branch(ROT_BLOCK),
            counts={name: int((kind == k).sum()) for k, name in enumerate(KIND_NAMES)},
            mean_policy=cls._by_policy(alpha, queries),
            policy_counts=cls._policy_counts(queries),
            row_first=float(alpha[:, 0, :].mean()),
            row_rest=float(alpha[:, 1:, :].mean()) if alpha.shape[1] > 1 else float("nan"),
            offset_curve=(
                cls._offset_curve(alpha, kind, queries) if curve_set is None
                else cls._offset_curve(
                    cls._alpha_over(computer, gt, curve_set, qbatch), cls._kinds(curve_set),
                    curve_set,
                )
            ),
        )

    @staticmethod
    def _policy_kinds(queries: "QuerySet") -> np.ndarray | None:
        """This set's per-row POLICY-corruption branch, or None when none was drawn."""
        pol = getattr(queries, "policy_corrupt", None)
        return None if pol is None else np.asarray(pol.kind, dtype=np.int8)

    @classmethod
    def _by_policy(cls, alpha: np.ndarray, queries: "QuerySet") -> dict[str, float]:
        """Mean weight per policy-corruption branch.

        Args:
            alpha: [n, K, 6] weights in the set's own order.
            queries: The set they came from.

        Returns:
            {branch name: mean weight}, empty when the set drew no policy corruption.
        """
        kind = cls._policy_kinds(queries)
        if kind is None:
            return {}
        per_row = alpha.mean(axis=(1, 2))
        return {
            name: float(per_row[kind == k].mean()) if np.any(kind == k) else float("nan")
            for k, name in enumerate(POLICY_KIND_NAMES)
        }

    @classmethod
    def _policy_counts(cls, queries: "QuerySet") -> dict[str, int]:
        """Rows per policy-corruption branch, empty when none was drawn."""
        kind = cls._policy_kinds(queries)
        if kind is None:
            return {}
        return {name: int((kind == k).sum()) for k, name in enumerate(POLICY_KIND_NAMES)}

    @staticmethod
    def _alpha_over(
        computer: ChunkLossComputer, gt: GraphTensors, queries: "QuerySet", qbatch: int,
    ) -> np.ndarray:
        """Every row's blend weight, in the set's own order.

        Args:
            computer: The chunk-loss computer.
            gt: The device-resident graph bundle.
            queries: The set to run.
            qbatch: Rows per forward pass.

        Returns:
            [n, K, 6] weights on the host.
        """
        with torch.no_grad():
            return torch.cat(
                [computer.chunk_outputs(gt, qb).alpha for qb in queries.ordered_batches(qbatch)]
            ).detach().cpu().numpy()

    @staticmethod
    def _kinds(queries: "QuerySet") -> np.ndarray:
        """This set's per-row corruption branch, all CLEAN when none was drawn.

        Args:
            queries: The set.

        Returns:
            [n] int8 branch ids.
        """
        return (
            np.zeros(len(queries), dtype=np.int8) if queries.corrupt is None
            else np.asarray(queries.corrupt.kind)
        )

    @staticmethod
    def _offset_curve(
        alpha: np.ndarray, kind: np.ndarray, queries: "QuerySet",
    ) -> tuple[tuple[float, float, int], ...]:
        """Mean position weight against displacement, over log-spaced bins of the OFFSET branch.

        Args:
            alpha: [n, K, 6] weights over the whole query set.
            kind: [n] int8 branch ids.
            queries: The set, for its realised displacement vectors.

        Returns:
            (bin centre in metres, mean position-block weight, row count) per non-empty bin, or ()
            when the set has no OFFSET rows. The magnitude is recomputed from the stored
            displacement rather than carried alongside it, so the curve can never disagree with the
            offset the loss actually applied.
        """
        if queries.corrupt is None:
            return ()
        rows = np.flatnonzero(kind == OFFSET)
        if rows.size == 0:
            return ()
        magnitude = np.linalg.norm(np.asarray(queries.corrupt.offset)[rows], axis=1)
        positive = magnitude > 0.0
        if not np.any(positive):
            return ()
        rows, magnitude = rows[positive], magnitude[positive]
        edges = np.geomspace(magnitude.min(), magnitude.max() * (1.0 + 1e-9), _OFFSET_BINS + 1)
        which = np.clip(np.searchsorted(edges, magnitude, side="right") - 1, 0, _OFFSET_BINS - 1)
        weight = alpha[rows][..., POS_BLOCK].mean(axis=(1, 2))
        return tuple(
            (float(np.sqrt(edges[b] * edges[b + 1])), float(weight[which == b].mean()),
             int((which == b).sum()))
            for b in range(_OFFSET_BINS) if np.any(which == b)
        )

    def lines(self) -> list[str]:
        """The gate, as log lines: one per branch, then the compact summaries and the row split.

        Returns:
            The lines, without any stage prefix.
        """
        return [
            f"alpha[{name}] n={self.counts[name]} mean={self.mean_alpha[name]:.4f} "
            f"pos={self.mean_pos[name]:.4f} rot={self.mean_rot[name]:.4f}"
            for name in KIND_NAMES
        ] + [
            "alpha strata: " + " / ".join(f"{n} {self.mean_alpha[n]:.4f}" for n in KIND_NAMES),
            "alpha strata pos/rot: " + " / ".join(
                f"{n} {self.mean_pos[n]:.4f}/{self.mean_rot[n]:.4f}" for n in KIND_NAMES
            ),
            f"alpha rows: row0={self.row_first:.4f} rows1+={self.row_rest:.4f}",
        ] + ([
            # The OPPOSITE sign to the line above: here alpha must go UP, because the demonstration
            # is the only operand still right. Flat means the head reads disagreement one way only.
            "alpha policy strata: " + " / ".join(
                f"{n} {self.mean_policy[n]:.4f}(n={self.policy_counts.get(n, 0)})"
                for n in POLICY_KIND_NAMES
            ),
        ] if self.mean_policy else []) + ([
            "alpha vs offset, in-sample (cm -> pos alpha): " + "  ".join(
                f"{100.0 * centre:.1f}cm={value:.3f}(n={n})"
                for centre, value, n in self.offset_curve
            ),
        ] if self.offset_curve else [])


class StageTrainer:
    """Runs one stage: for each epoch, one _query_loss_batch call and one opt.step per chunk of
    cfg.qbatch queries, then evaluate and keep the best epoch.

    evalq.guard is a clean, stage-independent set whose mrr is logged as a diagnostic; evalq.static
    is the separate entry_static_frac=1.0 set the ENTRY_STATIC acceptance metrics read, kept out of
    evalq.held/evalq.guard so their mrr/top1 never mix in a static-window row. Either being None is
    a strict no-op.
    """

    def __init__(
        self, net: GraphRetrieverNet, opt: torch.optim.Optimizer, gt: GraphTensors,
        queries: QuerySet, evalq: EvalQueries, spec: StageSpec,
        alpha_net: AlphaNet | None = None, lang_head: Any = None,
    ) -> None:
        """
        Args:
            net: The retriever to train.
            opt: Its optimiser.
            gt: The device-resident graph bundle (_prepare_tensors).
            queries: This stage's training queries.
            evalq: Held-out, guard and ENTRY_STATIC evaluation sets.
            spec: Epoch count, rng, device, config, logger, stage name, field, loss spec.
            alpha_net: The blend-weight head. Required under CHUNK, and needed under RETRIEVAL too
                whenever chunk targets exist, since model selection scores the blended chunk.
            lang_head: The instruction prior trained alongside it, or None when this graph has no
                task embeddings.

        Raises:
            ValueError: The CHUNK objective was requested without an alpha_net.
        """
        self.net, self.opt, self.gt, self.queries, self.spec = net, opt, gt, queries, spec
        self.alpha_net = alpha_net
        self.lang_head = lang_head
        self._chunk = spec.loss_spec.objective == CHUNK
        if self._chunk and alpha_net is None:
            raise ValueError(f"the {CHUNK!r} objective needs an alpha_net; none was passed")
        # Under CHUNK the retriever is frozen, so restoring it would restore nothing.
        # Best-epoch selection snapshots this. Under CHUNK it must cover EVERY module the optimiser
        # owns, or restoring the best epoch pairs that epoch's alpha with the last epoch's
        # instruction prior -- two halves of one controller, fitted against each other, taken from
        # different points in training.
        self._trained: torch.nn.Module = (
            torch.nn.ModuleList([alpha_net] + ([] if lang_head is None else [lang_head]))
            if self._chunk else net
        )
        self.qbatch = max(1, int(spec.cfg.qbatch))
        self._held = evalq.held.raw
        self._guard = evalq.guard.raw if evalq.guard is not None else None
        self._static = evalq.static.raw if evalq.static is not None else None
        self._held_set = evalq.held_set
        self._reporter = EpochReporter(spec.logger, spec.stage_name, spec.epochs)
        self._best = BestCheckpoint()

    def _fit_physical_stats(self) -> None:
        """Measure AlphaNet's feature standardisation off the training set, before the first step.

        Only the fifteen physical features need this; the identity blocks are layer-normed. It runs
        once, over the whole training set, and is a no-op under RETRIEVAL.

        The clamp those statistics define is what stops a rollout's out-of-range off-manifold
        distance and chunk disagreement from extrapolating through the head -- the failure that
        made the first trained weight score below both fixed weights bracketing its own held-out
        mean. See onf.blend.alpha.AlphaNet.
        """
        if not self._chunk or self.queries.targets is None:
            return
        computer = self._computer()
        scalars, disagree = zip(*(
            computer.physical_features(self.gt, qb)
            for qb in self.queries.ordered_batches(self.qbatch)
        ))
        self.alpha_net.fit_physical_stats(torch.cat(scalars), torch.cat(disagree))
        mean, std = self.alpha_net.phys_mean, self.alpha_net.phys_std
        self._log(
            f"alpha feature stats over {sum(s.shape[0] for s in scalars)} rows: "
            f"off_manifold {mean[0]:.3f}+-{std[0]:.3f} top1 {mean[1]:.3f}+-{std[1]:.3f} "
            f"entropy {mean[2]:.3f}+-{std[2]:.3f} "
            f"lang_match {mean[3]:.3f}+-{std[3]:.3f} "
            f"disagree_row0 {mean[4:10].mean():.3f} disagree_rows1+ {mean[10:].mean():.3f}"
        )

    def run(self) -> list[dict]:
        """Train for spec.epochs, restore the best epoch, then evaluate it.

        Returns:
            The per-epoch history, with a trailing {"epoch": "final", ...} entry.
        """
        self._log(self.queries.balance_summary())
        if self._held_set is None or self._held_set.targets is None:
            self._log("no held-out chunk targets -- keeping the FINAL epoch, not selecting one")
        self._fit_physical_stats()

        for ep in range(self.spec.epochs):
            mean_loss = self._train_epoch()
            metrics = self._metrics(self._held)
            chunk_mse = self._chunk_mse()
            static_metrics = self._static_metrics()
            self._reporter.record(ep, mean_loss, metrics, guard_mrr=self._guard_mrr(),
                                   chunk_mse=chunk_mse, static_metrics=static_metrics)
            if self.lang_head is not None:
                self._log(
                    f"lang: tau={float(self.lang_head.tau):.4f} "
                    f"grad_norm={getattr(self, '_lang_grad', 0.0):.3e} "
                    f"|proj|={float(self.lang_head.proj.weight.norm()):.4f}"
                )
            self._best.consider(self._trained, ep, chunk_mse)

        if self._best.restore_into(self._trained, self.spec.epochs - 1):
            self._log(f"restored epoch {self._best.epoch + 1}/{self.spec.epochs} (held-out "
                       f"chunk_mse={self._best.score:.5f}) -- the final epoch was not the best")

        final = evaluate(self.net, self.gt.nodes, self.gt.edges, self._held, field=self.spec.field,
                          cfg=self.spec.cfg, device=self.spec.dev)
        final_static = self._static_metrics()
        if final_static is not None:
            final["static_kl"] = final_static["static_kl"]
            final["static_bary_err"] = final_static["static_bary_err"]
        self._reporter.record_final(final, final_static)
        self.report_alpha_strata()
        return self._reporter.history

    def report_alpha_strata(self) -> AlphaStrata | None:
        """Measure and log the held-out blend weight per corruption branch -- Stage 2's gate.

        Named rather than inlined into run so a script can re-run the gate against a restored
        checkpoint without retraining.

        Returns:
            The measurement, or None when this stage trains no alpha head or holds out no chunk
            targets to measure it on.
        """
        held = self._held_set
        if not self._chunk or held is None or held.targets is None:
            return None
        self._set_mode(training=False)
        strata = AlphaStrata.measure(
            self._computer(), self.gt, held, self.qbatch, curve_set=self.queries,
        )
        for line in strata.lines():
            self._log(line)
        return strata

    def _log(self, msg: str) -> None:
        """Emit one stage-prefixed line through the reporter's logger."""
        self._reporter.log(f"[graph.train {self.spec.stage_name}] {msg}")

    def _set_mode(self, *, training: bool) -> None:
        """Put the module this stage trains in train mode and every other one in eval.

        Under CHUNK the retriever is frozen, so it stays in eval throughout: restoring it to train()
        after an evaluation would put a module that never trains back into training mode.

        Args:
            training: Whether the trained module should be in train mode.
        """
        self.net.train(training and not self._chunk)
        if self.alpha_net is not None:
            self.alpha_net.train(training and self._chunk)

    def _computer(self) -> ChunkLossComputer:
        """The chunk-loss computer this stage's alpha diagnostics and model selection read."""
        return ChunkLossComputer(
            self.net, self.alpha_net, self.spec.cfg, self.spec.loss_spec, self.lang_head,
        )

    def _train_epoch(self) -> float:
        """One epoch of shuffled mini-batch updates.

        Returns:
            The mean per-query loss. _query_loss_batch returns the MEAN over its chunk, so the
            division is undone per chunk and redone once at the end -- that keeps this number
            comparable across different qbatch values.
        """
        self._set_mode(training=True)
        total_loss = 0.0
        for qb in self.queries.shuffled_batches(self.qbatch):
            self.opt.zero_grad()
            loss = _query_loss_batch(
                self.net, self.gt, qb, self.spec.rng, self.spec.cfg, self.spec.loss_spec,
                alpha_net=self.alpha_net, lang_head=self.lang_head,
            )
            loss.backward()
            self._lang_grad = self._language_grad_norm()
            self.opt.step()
            total_loss += float(loss.detach()) * len(qb.tgt)
        return total_loss / max(len(self.queries), 1)

    def _language_grad_norm(self) -> float:
        """Total gradient norm reaching the instruction prior on this step.

        Logged per epoch because an inert prior is otherwise invisible. The first run of this
        mechanism trained for 15 epochs, reported "instruction prior ON", and wrote a head whose
        weights were bit-identical to their initialisation -- the training step built its own
        loss computer and never passed the head, so the whole feature was off while every log line
        said it was on. A zero here is now a number on the screen.

        Returns:
            The L2 norm over the prior's parameters, or 0.0 when there is no prior.
        """
        if self.lang_head is None:
            return 0.0
        return float(
            sum(
                float(p.grad.detach().pow(2).sum()) for p in self.lang_head.parameters()
                if p.grad is not None
            ) ** 0.5
        )

    def _chunk_mse(self) -> float | None:
        """The held-out radius-stratified mse(a_exec) -- the model-selection score.

        Returns:
            The mean of the per-stratum mean errors, or None when no chunk targets were built or no
            alpha head exists to blend with.
        """
        held = self._held_set
        if held is None or held.targets is None or self.alpha_net is None:
            return None
        self._set_mode(training=False)
        computer = self._computer()
        with torch.no_grad():
            rows = torch.cat(
                [computer.row_mse(self.gt, qb) for qb in held.ordered_batches(self.qbatch)]
            )
        self._set_mode(training=True)
        return held.strata.mean(rows.detach().cpu().numpy())

    def _metrics(self, raw: Any, *, static: bool = False) -> dict:
        """Held-out metrics for one query set.

        Args:
            raw: The raw query mapping to evaluate on.
            static: Whether to also compute the ENTRY_STATIC acceptance metrics.

        Returns:
            The metric mapping.
        """
        return _gnn_metrics(
            self.net, self.gt.nodes, self.gt.edges, raw, field=self.spec.field, cfg=self.spec.cfg,
            device=self.spec.dev,
            **({"entry_static_pool": self.gt.targets.pool} if static else {}),
        )

    def _guard_mrr(self) -> float | None:
        """This stage's guard-set MRR, or None when no guard set was given."""
        return None if self._guard is None else float(self._metrics(self._guard)["mrr"])

    def _static_metrics(self) -> dict | None:
        """The ENTRY_STATIC acceptance metrics, or None when no static set was given."""
        return None if self._static is None else self._metrics(self._static, static=True)


@dataclass
class TrainResult:
    """Outcome of one GraphTrainer.train call.

    Subscriptable so result["checkpoint_path"] reads in onf.graph.cli and the training tests keep
    working against what used to be a plain dict.

    Attributes:
        net: The trained retriever.
        graph_hash: Hash of the graph it was trained against.
        history: Per-stage epoch history.
        checkpoint_path: Where the head was written, or None when save was off.
        inter_demo_spacing: The corpus-derived threshold (rad) the abstain decision's consequence
            label was trained against, recorded so a run's metrics are self-describing.
    """

    net: GraphRetrieverNet
    graph_hash: str
    history: dict[str, list[dict]] = dc_field(default_factory=dict)
    checkpoint_path: str | None = None
    inter_demo_spacing: float = 0.0

    def __getitem__(self, key: str) -> Any:
        try:
            return getattr(self, key)
        except AttributeError as exc:
            raise KeyError(key) from exc

    def __contains__(self, key: str) -> bool:
        return key in vars(self)

    def keys(self) -> Any:
        """The field names, so dict(result) still round-trips."""
        return vars(self).keys()


class GraphTrainer:
    """One training run, start to finish: load, resolve the field, build queries, run the stage, save.

    Construction resolves device/rng/config and seeds torch; nothing is read from disk until train.
    """

    def __init__(self, graph_dir: str | os.PathLike, spec: TrainGraphSpec | None = None) -> None:
        """
        Args:
            graph_dir: Directory holding g_nodes.npz/g_edges.npz.
            spec: Run options; None is exactly onf.graph.train.types.TrainGraphSpec's default.
        """
        self.graph_dir = Path(graph_dir)
        self.spec = spec if spec is not None else TrainGraphSpec()

        # torch.manual_seed is a process-global side effect, placed here so both sources of
        # randomness for this run are pinned at the same point.
        self.dev = resolve_device(self.spec.device)
        self.rng = np.random.RandomState(self.spec.seed)
        torch.manual_seed(self.spec.seed)
        self.cfg: GraphConfig = self.spec.cfg or GraphConfig()

        self.nodes: NodeTable | None = None
        self.edges: EdgeSet | None = None
        self.field: Any = None
        self.net: GraphRetrieverNet | None = None
        self.alpha: AlphaNet | None = None
        self.opt: torch.optim.Optimizer | None = None
        self.inter_demo_spacing: float = 0.0
        self.chunk: ChunkTargetBuilder | None = None
        self.corruptor: RetrievalCorruptor | None = None
        self.loss_spec: LossSpec = DEFAULT_LOSS_SPEC
        self.language: Any = None                       # onf.graph.lang.TaskEmbeddings, or None
        self.lang_head: LanguageHead | None = None

    # -- public API --------------------------------------------------------------------------------
    def train(self) -> TrainResult:
        """Run the curriculum and, if spec.save, write <graph_dir>/g_head.npz.

        Returns:
            The trained net, its graph hash, the per-stage history, the checkpoint path and the
            corpus-derived inter-demo spacing.

        Raises:
            ValueError: If spec.stages is anything but (1,), or the corpus is too small for the
                stride-HELDOUT_STRIDE by-demo split.
        """
        self._load_graph()
        self.field = self._resolve_field()
        self.chunk = self._resolve_chunk_targets()
        self.language = self._resolve_language()
        self.loss_spec = replace(DEFAULT_LOSS_SPEC, objective=self.spec.objective)

        # Computed once per graph and threaded into every make_queries call below, so every query
        # set is labelled against the identical corpus-derived scale.
        self.inter_demo_spacing = measure_inter_demo_spacing(self.nodes)
        self._log(f"[graph.train] INTER_DEMO_SPACING={self.inter_demo_spacing:.4f} rad "
                   "(median leave-one-demo-out node NN distance)")

        # Only Stage 2 corrupts: the retrieval objective scores node identity, and a row whose
        # retrieval was sabotaged by hand carries no information about it.
        self.corruptor = (
            RetrievalCorruptor(self.nodes) if self.spec.objective == CHUNK else None
        )
        # And the mirror image of it: a policy that is wrong on a query whose retrieval is right.
        # See onf.graph.train.corrupt's "THE OTHER OPERAND" for why one without the other fits the
        # head's dominant feature with the wrong sign on most of the benchmark.
        self.policy_corruptor = (
            PolicyCorruptor(self.nodes, self.corruptor) if self.corruptor is not None else None
        )
        train_owners, heldout_owners = self._split_owners()
        self._build_model()

        # These four make_queries calls share self.rng, so reordering them (or moving the interleaved
        # steps between them) would hand every query set a different draw.
        guard = self._make_queries(heldout_owners, entry_static_frac=0.0)
        static_held = self._make_queries(heldout_owners, entry_frac=0.0, entry_static_frac=1.0)
        gt = _prepare_tensors(
            self.nodes, self.edges, self.dev, cfg=self.cfg, chunk=self.chunk,
            language=self.language,
        )
        if tuple(self.spec.stages) != (1,):
            raise ValueError(
                f"train_graph: stages={self.spec.stages!r} -- only (1,) is legal (stage 2's drift "
                "finetuning was a closed ablation and has been deleted; stage 3's outcome/REINFORCE "
                "was always out of scope, see docs/technical/02-retrieval-head.md)"
            )
        # The one site that turns ENTRY_STATIC and DRIFT on: make_queries defaults both fractions to
        # 0.0 because "entry_frac=0.0 means a pure traversal set" is a contract callers rely on, and
        # the eval sets below stay free of both so their mrr/top1 keep meaning what they did.
        train_q = self._make_queries(train_owners, n=self._n_query,
                                      entry_static_frac=ENTRY_STATIC_FRAC, drift_frac=DRIFT_FRAC)
        held = self._make_queries(heldout_owners, entry_static_frac=0.0)

        stage = StageTrainer(
            net=self.net, opt=self.opt, gt=gt,
            queries=QuerySet.from_queries(
                train_q, self.dev, self.chunk, self.corruptor, self.rng,
                language=self.language, task_id=self._task_id_host(),
                policy_corruptor=self.policy_corruptor,
            ),
            evalq=EvalQueries(
                held=QueryData.from_dict(held),
                guard=QueryData.from_dict(guard),
                static=QueryData.from_dict(static_held),
                # The held-out set is corrupted too: the gate stratifies its alpha by branch, and
                # model selection should reward an epoch that refuses an invalid retrieval.
                held_set=QuerySet.from_queries(
                    held, self.dev, self.chunk, self.corruptor, self.rng,
                    language=self.language, task_id=self._task_id_host(),
                    policy_corruptor=self.policy_corruptor,
                ),
            ),
            spec=StageSpec(epochs=self.spec.epochs, rng=self.rng, dev=self.dev, cfg=self.cfg,
                            logger=self.spec.logger, stage_name="stage1", field=self.field,
                            loss_spec=self.loss_spec),
            alpha_net=self.alpha, lang_head=self.lang_head,
        )

        result = TrainResult(
            net=self.net, graph_hash=self._graph_hash(),
            history={"stage1": stage.run()}, inter_demo_spacing=self.inter_demo_spacing,
        )
        if self.spec.save:
            result.checkpoint_path = self.save()
        return result

    def save(self) -> str:
        """Write this run's checkpoint(s) into graph_dir and log the paths.

        RETRIEVAL writes g_head.npz. CHUNK writes g_alpha.npz and leaves the head alone: it froze
        that file's weights rather than fitting them, and rewriting it would put this run's bytes
        over an artifact other runs are pinned to. The AlphaNet is likewise written only when it was
        trained -- a randomly-initialised g_alpha.npz beside the graph would be loaded at deploy as
        a learned weight.

        Returns:
            The head checkpoint path, written or pre-existing.
        """
        head_path = self.graph_dir / schema.HEAD_NPZ
        if self.spec.objective == CHUNK:
            alpha_path = self.alpha.save(self.graph_dir)
            self._log(f"[graph.train] saved {alpha_path} (head at {head_path} left untouched -- "
                       "this run froze it)")
            if self.lang_head is not None:
                self._log(f"[graph.train] saved {self.lang_head.save(self.graph_dir)} "
                           f"(tau={float(self.lang_head.tau):.4f})")
            return os.fspath(head_path)
        path = save_checkpoint(self.net, self.nodes, self.edges, head_path)
        self._log(f"[graph.train] saved {path}")
        return path

    # -- steps -------------------------------------------------------------------------------------
    def _log(self, msg: str) -> None:
        """Emit one line through spec.logger, or stdout when there is none."""
        (self.spec.logger.log(msg) if self.spec.logger is not None else print(msg))

    def _load_graph(self) -> None:
        """Load and validate <graph_dir>/{g_nodes,g_edges}.npz."""
        self.nodes = NodeTable.load(self.graph_dir)
        self.edges = EdgeSet.load(self.graph_dir / schema.EDGES_NPZ, strict=True)
        self.edges.validate(self.nodes)

    def _resolve_field(self) -> Any:
        """Resolve the onf.field.field.ONFField used for cleanliness weighting.

        A caller-supplied field wins; then ONF_CLEANLINESS=geo's geometric stand-in; otherwise the
        graph dir is tried first (so a co-located field still wins) and then, when spec.suite is
        given, the suite's fwm dir -- the trained field is a pre-existing ONF artifact and does NOT
        live beside the graph. Missing both degrades to uniform window weights, which is precisely
        the failure the design avoids (when recovery is needed the most recent states are the worst
        ones), so whichever path is taken is logged loudly.

        Returns:
            The field, or None for the uniform fallback.

        Raises:
            ValueError: If ONF_CLEANLINESS=geo is set before the graph has been loaded.
        """
        field = self.spec.field
        clean_cfg = CleanlinessConfig.from_env()
        if field is not None:
            self._log("[graph.train] using caller-supplied ONFField for cleanliness weighting")
            return field

        if clean_cfg.use_geometric:
            if self.nodes is None:
                raise ValueError("GraphTrainer._resolve_field: ONF_CLEANLINESS=geo needs the graph's nodes")
            field = load_cleanliness_field(self.nodes, clean_cfg)
            self._log(f"[graph.train] ONF_CLEANLINESS=geo -- GeometricField(scale={field.scale:.4f} rad) "
                       f"over {len(self.nodes)} nodes, no trained ONFField")
            return field

        candidates = [self.graph_dir]
        if self.spec.suite:
            candidates.append(default_paths().fwm(self.spec.suite))
        msg = ""
        for cand in candidates:
            try:
                field = ONFField.load(cand)
                self._log(f"[graph.train] loaded ONFField from {cand} for cleanliness weighting")
                return field
            except (FileNotFoundError, OSError) as exc:
                msg = (f"[graph.train] no trained ONFField at {cand} ({exc}) -- cleanliness weights "
                        "fall back to UNIFORM (see onf.graph.core.geometry.CleanlinessScorer)")
        self._log(msg)
        return None

    def _task_id_host(self) -> np.ndarray | None:
        """Per-node task ids on the host, for the per-row instruction draw.

        Returns:
            [V] int64, or None when this run has no instruction prior to draw for.
        """
        return None if self.language is None else np.asarray(self.nodes.task_id, dtype=np.int64)

    def _resolve_language(self) -> Any:
        """Load this graph's task-instruction embeddings when it has any.

        Returns:
            The onf.graph.lang.TaskEmbeddings, or None. Absent is the ordinary state of a graph
            built before the instruction prior existed, and trains exactly as it did then.

        Raises:
            ValueError: The table exists but was built for different task names.
        """
        from onf.graph.lang import LANG_NPZ, TaskEmbeddings

        if not (self.graph_dir / LANG_NPZ).is_file():
            self._log(f"[graph.train] no {LANG_NPZ} beside the graph -- training WITHOUT the "
                       "instruction prior (build one with scripts/embed_task_names.py)")
            return None
        table = TaskEmbeddings.load(self.graph_dir).check_against(self.nodes.task_names)
        self._log(f"[graph.train] loaded {LANG_NPZ}: {table.n_tasks} tasks, encoder "
                   f"{table.model_id}, {table.e_aug.shape[1] if table.e_aug.size else 0} paraphrase "
                   "surrogates per task")
        return table

    def _resolve_chunk_targets(self) -> ChunkTargetBuilder | None:
        """Build the chunk-blend target builder from the artifacts beside the graph.

        Both artifacts are produced once per suite, by scripts/cache_policy_actions.py and
        scripts/fit_action_scale.py. Missing them is fatal under the CHUNK objective and merely
        disables model selection under RETRIEVAL, so the two cases are separated here rather than
        left to fail deep inside a training step.

        Returns:
            The builder, or None when the artifacts are absent and the objective does not need them.

        Raises:
            FileNotFoundError: The CHUNK objective was requested without the artifacts.
        """
        scale_path = self.graph_dir / ACTION_SCALE_JSON
        try:
            policy = PolicyActionCache.load(self.graph_dir).check_against(self.nodes)
            scale = ActionScale.from_json(scale_path)
        except FileNotFoundError as exc:
            if self.spec.objective == CHUNK:
                raise FileNotFoundError(
                    f"objective={CHUNK!r} needs the chunk-blend artifacts beside the graph "
                    f"({self.graph_dir}), but: {exc}"
                ) from exc
            self._log(f"[graph.train] no chunk-blend targets ({exc}) -- model selection is off, "
                       "the final epoch will be kept")
            return None
        self._log(f"[graph.train] chunk targets: {len(policy.a_pi_raw)} cached policy frames, "
                   f"action scale pos={scale.pos:.6f} m/unit rot={scale.rot:.6f} rad/unit")
        return ChunkTargetBuilder(
            nodes=self.nodes, scale=scale, policy=policy, kinematics=PandaKinematics(),
            chunk_len=self.cfg.seg_k,
        )

    def _split_owners(self) -> tuple[np.ndarray, np.ndarray]:
        """The by-demo held-out split: every HELDOUT_STRIDE-th strand is held out.

        Returns:
            (train_owners, heldout_owners).

        Raises:
            ValueError: If either side of the split comes out empty.
        """
        strands = np.arange(self.nodes.n_demos)
        is_heldout = np.zeros(len(strands), dtype=bool)
        is_heldout[strands[::HELDOUT_STRIDE]] = True
        train_owners, heldout_owners = strands[~is_heldout], strands[is_heldout]
        if len(train_owners) == 0 or len(heldout_owners) == 0:
            raise ValueError(
                f"train_graph: n_demos={self.nodes.n_demos} too small for a stride-{HELDOUT_STRIDE} "
                "by-demo held-out split (need >= 2 demos)"
            )
        return train_owners, heldout_owners

    def _build_model(self) -> None:
        """Construct the retriever and, when chunk targets exist, the AlphaNet, plus the optimiser.

        Under CHUNK the retriever is Stage 1's trained head, frozen, and ONLY the AlphaNet is
        optimised. The two confound each other -- the same loss falls either by retrieving a better
        segment or by trusting the segment less -- and alpha, at four orders of magnitude fewer
        parameters, converges far sooner. Training them apart is what makes a negative result about
        alpha attributable to alpha rather than to a retriever that moved underneath it.

        Raises:
            FileNotFoundError: CHUNK was requested with no trained head beside the graph.
            ValueError: That head was trained against a different graph.
        """
        self.net = (
            self._load_frozen_head() if self.spec.objective == CHUNK
            else self._fresh_retriever()
        )

        if self.chunk is not None:
            # Forked so the AlphaNet's init draws no numbers out of the stream that seeds the
            # retriever and every epoch's shuffle -- a RETRIEVAL run must be unchanged by its
            # presence, and it is only ever trained under CHUNK.
            with torch.random.fork_rng(devices=[]):
                torch.manual_seed(self.spec.seed)
                self.alpha = AlphaNet(
                    embed_dim=self.net.hidden, chunk_len=self.chunk.chunk_len
                ).to(self.dev)

        if self.spec.objective == CHUNK:
            self.net.requires_grad_(False)
            self.net.eval()
            self._build_language_head()
            trained = list(self.alpha.parameters()) + (
                [] if self.lang_head is None else list(self.lang_head.parameters())
            )
            params = trained
            self._log(f"[graph.train] objective={CHUNK!r}: retriever FROZEN "
                       f"({sum(p.numel() for p in self.net.parameters())} params), training the "
                       f"AlphaNet alone ({sum(p.numel() for p in self.alpha.parameters())} params, "
                       f"embed_dim={self.alpha.embed_dim} chunk_len={self.alpha.chunk_len})")
            if self.lang_head is not None:
                self._log(
                    f"[graph.train] instruction prior ON: LanguageHead "
                    f"({sum(p.numel() for p in self.lang_head.parameters())} params, "
                    f"embed_dim={self.lang_head.embed_dim}) trained alongside the AlphaNet"
                )
        else:
            params = self.net.parameters()
        self.opt = torch.optim.AdamW(params, lr=LR, weight_decay=WEIGHT_DECAY)

    def _build_language_head(self) -> None:
        """Construct the instruction prior when this graph shipped task embeddings.

        Trained here rather than in Stage 1 for the same reason the AlphaNet is: it is a
        deploy-time mechanism, and the quantity it should be scaled against is the error of the
        chunk the controller emits, not the retrieval loss.
        """
        if self.language is None:
            self.lang_head = None
            return
        with torch.random.fork_rng(devices=[]):
            torch.manual_seed(self.spec.seed + 1)
            self.lang_head = LanguageHead(
                embed_dim=int(self.language.e_task.shape[1]), hidden=self.net.hidden
            ).to(self.dev)

    def _fresh_retriever(self) -> GraphRetrieverNet:
        """A randomly-initialised retriever on self.dev, normalised against this corpus.

        Returns:
            The model Stage 1 trains from scratch.
        """
        net = GraphRetrieverNet(
            dim=self.nodes.dim, hidden=self.cfg.hidden, layers=self.cfg.layers,
            n_rel=schema.N_RELATIONS, agg=self.cfg.agg, n_psi=schema.PSI_FREQS,
        ).to(self.dev)
        net.set_stats(*self.nodes.feature_stats())
        return net

    def _graph_hash(self, c2_state: Any = None) -> str:
        """graph_hash v2 of this run's graph: geometry, C0's constants and a C2 checkpoint.

        Args:
            c2_state: The C2 checkpoint to stamp against; None uses the live retriever's.

        Returns:
            The hex digest.
        """
        return graph_hash(
            self.nodes, self.edges, constants=self.graph_dir / CONSTANTS_JSON,
            c2_state=self.net.c2_state() if c2_state is None else c2_state,
        )

    def _load_frozen_head(self) -> GraphRetrieverNet:
        """Stage 1's trained head from <graph_dir>/g_head.npz, for Stage 2 to freeze.

        Refusing to fall back on a fresh model is the point: alpha is a compatibility function
        between a query and what retrieval returned, and an untrained retriever returns noise, so
        the whole run would fit a weight for a posterior that will never be deployed.

        Returns:
            The loaded retriever on self.dev.

        Raises:
            FileNotFoundError: No head beside the graph.
            ValueError: The head's stamped graph hash is not this graph's, so its node ids -- which
                are positional and mean nothing across builder runs -- index other strands.
        """
        head_path = self.graph_dir / schema.HEAD_NPZ
        if not head_path.is_file():
            raise FileNotFoundError(
                f"objective={CHUNK!r} freezes Stage 1's retriever, but there is no head at "
                f"{head_path}. Train one first: "
                f"`python -m onf.graph train --suite <suite> --graph-dir {self.graph_dir}`."
            )
        with np.load(os.fspath(head_path)) as raw:
            got = str(raw["graph_hash"]) if "graph_hash" in raw.files else "<unstamped>"
            # The FILE's own C2 tensors, not the live net's: that checkpoint is what is being
            # validated, and under C2 it is part of what the graph means.
            c2 = {k: raw[k] for k in raw.files if k.startswith(C2_PREFIX)}
        want = self._graph_hash(c2 or None)
        if got != want:
            raise ValueError(
                f"{head_path}: graph_hash mismatch -- head={got!r} graph={want!r}. That head was "
                "not trained on this (g_nodes.npz, g_edges.npz) pair."
            )
        self._log(f"[graph.train] loaded the frozen Stage 1 head from {head_path}")
        return GraphRetrieverNet.load_npz(head_path).to(self.dev)

    @property
    def _n_query(self) -> int:
        """Training queries per stage."""
        return N_QUERY if self.spec.n_query is None else self.spec.n_query

    @property
    def _n_eval(self) -> int:
        """Held-out queries per evaluation set."""
        return N_EVAL if self.spec.n_eval is None else self.spec.n_eval

    def _make_queries(self, owners: Any, *, n: int | None = None, **spec_kwargs: float) -> Any:
        """One query set against this run's shared rng, field, window and spacing.

        Args:
            owners: Demo indices to draw from.
            n: Query count; None uses _n_eval (every eval set is the same size).
            **spec_kwargs: The class-mix fractions this set differs by.

        Returns:
            The query mapping from make_queries.
        """
        return make_queries(
            self.nodes, self.rng, self._n_eval if n is None else n, owners,
            QuerySpec(field=self.field, window=self.cfg.hist,
                       inter_demo_spacing=self.inter_demo_spacing, **spec_kwargs),
        )


def train_graph(graph_dir: str | os.PathLike, spec: TrainGraphSpec | None = None) -> TrainResult:
    """Train GraphRetrieverNet on <graph_dir>/{g_nodes,g_edges}.npz and, if spec.save,
    write <graph_dir>/g_head.npz.

    Thin entry point kept for the existing call sites; GraphTrainer is the implementation.

    Args:
        graph_dir: Directory holding the graph artifacts.
        spec: Run options; None is onf.graph.train.types.TrainGraphSpec's default.

    Returns:
        The TrainResult, which is also subscriptable like the dict this used to return.
    """
    return GraphTrainer(graph_dir, spec).train()