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

"""Query construction for the graph retriever's training curriculum.

The pretext task (make_queries), hard-negative mining (NegativeMiner), the multi-positive
node-identity target (SoftTargetBuilder), and the ENTRY_STATIC pool cache. Design writeup:
docs/technical/02-retrieval-head.md.
"""

from __future__ import annotations

import warnings
from typing import NamedTuple

import numpy as np

from onf.graph.core import schema
from onf.graph.core.geometry import leave_one_demo_out_floors
from onf.graph.train.drift import PANDA_JOINT_HI, PANDA_JOINT_LO
from onf.graph.core.nodes import NodeTable
from onf.graph.core.geometry import CleanlinessScorer, edge_pad_hist, finite_diff_vel, off_manifold_batch
from onf.graph.train.types import QuerySpec

# ---- curriculum hyper-parameters -- see docs/technical/02-retrieval-head.md's constants table.
N_NEG = 16                  # mined hard negatives per query, split across categories (a)/(b)/(c)
SELFX_K = 16                 # local kNN pool mine_negatives searches for the self-intersection bucket (b)
SELFX_PHASE_GAP = 0.15        # |phase| gap that promotes a near-in-q neighbour into bucket (b)

# ---- episode-start (entry) query classes ------------------------------------------------------------
ENTRY_FRAC = 1.0 / 3.0      # share of make_queries' queries that are episode-start (entry) queries
ENTRY_LO_MULT = 0.25        # entry perturbation radius floor:   0.25 * demo-start noise floor p50
ENTRY_HI_MULT = 50.0        # entry perturbation radius ceiling: 50 * demo-start noise floor p99
ENTRY_LOGUNIFORM = True     # sample the radius log-uniformly -- equal density per decade

# ---- ENTRY_STATIC: trains the deployed t=0 input shape (a single perturbed frame repeated across the
# window, IDENTICALLY ZERO finite-differenced velocity), which the ENTRY class never produces.
ENTRY_STATIC_FRAC = 0.15    # share of ENTRY queries that are ENTRY_STATIC instead
ENTRY_STATIC_BAND = 0.05    # phase band for the static class's soft target pool -- matches the
                             # deployed GR_ENTRY_BAND constant
ENTRY_STATIC_NODE_W = 1.0   # absolute weight of the node-identity term on ENTRY_STATIC rows, replacing
                             # cfg.node_w (the phase terms are trivial: the whole pool sits in bin 0)
ENTRY_STATIC_SRC_AUX_W = 0.0  # zeroed on ENTRY_STATIC rows -- the source-node warm-up would sharpen mass
                             # onto the row's own demo-start node, fighting the uniform-over-pool target

# ---- DRIFT: the only class with a clean-to-perturbed breakpoint INSIDE the window. Every other class
# is uniformly clean (TRAVERSAL) or uniformly perturbed (ENTRY/ENTRY_STATIC), so "which rows of this
# history are still trustworthy" is a question they never pose.
DRIFT_FRAC = 0.15           # share of the rows the static AND entry draws both declined
DRIFT_MIN_HIST = 3          # >= 1 clean row and >= 2 perturbed ones, so the offset always ramps over
                             # a finite-difference step instead of landing as a single-frame impulse


def _node_at_raw(nodes: NodeTable, owner: int, raw_idx: int) -> int:
    """The node on demo owner whose raw span contains raw_idx.

    Last node with t_raw <= raw_idx, clamped to the demo's own node range;
    onf.graph.core.nodes.NodeTable.node_at_raw is the vectorised canonical form.

    Args:
        nodes: Node table to look up in.
        owner: Demo index.
        raw_idx: Raw frame index.

    Returns:
        The node id.
    """
    node_ids = np.asarray(nodes.demo_nodes(int(owner)))
    t_raws = np.asarray(nodes.t_raw)[node_ids]
    pos = int(np.searchsorted(t_raws, int(raw_idx), side="right")) - 1
    pos = min(max(pos, 0), len(node_ids) - 1)
    return int(node_ids[pos])


# There is exactly ONE finite-difference velocity implementation, onf.graph.core.geometry.
# finite_diff_vel, shared by every query builder here AND by GraphRetriever.retrieve: two
# reimplementations with different row-0 conventions once produced two different e_Q for the same
# window depending on train-vs-deploy, one of the divergences behind the abstain-collapse bug.


def entry_perturbation_floor(nodes: NodeTable) -> dict[str, float]:
    """Leave-one-demo-out 1-NN distance (rad) from each demo's first raw frame to every other demo's.

    The same procedure onf.field.build.build_qmanifold_full uses for the Sentinel's noise-floor tau,
    specialised to t=0 points, since that is what an episode-start perturbation is calibrated
    against. Purely a function of the corpus on nodes -- no simulator, no external constant.

    Args:
        nodes: Node table providing raw_ptr and q_raw.

    Returns:
        {"p50": ..., "p90": ..., "p99": ...} in radians.

    Raises:
        ValueError: If the corpus has fewer than two demos.
    """
    raw_ptr = np.asarray(nodes.raw_ptr, dtype=np.int64)
    q_raw = np.asarray(nodes.q_raw, dtype=np.float64)
    j_demos = len(raw_ptr) - 1
    if j_demos < 2:
        raise ValueError("entry_perturbation_floor: need >= 2 demos for a leave-one-demo-out floor")
    starts = q_raw[raw_ptr[:-1]]                                 # [J, D] -- each demo's t=0 config
    owner_raw = np.repeat(np.arange(j_demos, dtype=np.int64), np.diff(raw_ptr))
    own_ids = np.arange(j_demos, dtype=np.int64)

    floors = leave_one_demo_out_floors(starts, own_ids, q_raw, owner_raw)
    return {
        "p50": float(np.percentile(floors, 50)), "p90": float(np.percentile(floors, 90)),
        "p99": float(np.percentile(floors, 99)),
    }


def measure_inter_demo_spacing(nodes: NodeTable) -> float:
    """The corpus's own median leave-one-demo-out nearest-neighbour distance (rad) between nodes.

    This is the abstain decision's consequence-label threshold (docs/technical/02-retrieval-head.md, "The
    abstain decision"): a query whose most recent frame sits closer than this to any node is within
    one demo-spacing of genuine support, so abstaining is correct; farther, recovery is correct.
    Computed once per graph by train_graph. Measured ~0.054 rad on the real long-suite graph.

    Args:
        nodes: Node table providing q and owner.

    Returns:
        The median floor, in radians.

    Raises:
        ValueError: If the corpus has fewer than two demos.
    """
    if nodes.n_demos < 2:
        raise ValueError("measure_inter_demo_spacing: need >= 2 demos for a leave-one-demo-out floor")
    q = np.asarray(nodes.q, dtype=np.float64)
    owner = np.asarray(nodes.owner, dtype=np.int64)
    floors = leave_one_demo_out_floors(q, owner, q, owner)
    return float(np.median(floors))


def _entry_perturb(
    q0: np.ndarray, rng: np.random.RandomState, dim: int, sigma_range: tuple[float, float],
) -> np.ndarray:
    """q0 -> q0 + a random-direction offset of radius drawn from sigma_range, clipped to the limits.

    Args:
        q0: [dim] unperturbed configuration.
        rng: Random generator, consuming a direction draw then a radius draw.
        dim: Degrees of freedom.
        sigma_range: (lo, hi) radius range in radians.

    Returns:
        [dim] perturbed configuration.
    """
    direction = rng.randn(dim)
    n = np.linalg.norm(direction)
    direction = direction / n if n > 1e-12 else direction
    lo_r, hi_r = float(sigma_range[0]), float(sigma_range[1])
    # Log-uniform: over a two-decade range a linear draw puts ~90% of its mass in the top decade and
    # starves the small-offset end, which is exactly the boundary the abstain decision has to learn.
    radius = float(np.exp(rng.uniform(np.log(max(lo_r, 1e-6)), np.log(max(hi_r, 1e-6))))) \
        if ENTRY_LOGUNIFORM and hi_r > lo_r > 0 else float(rng.uniform(lo_r, hi_r))
    lo, hi = PANDA_JOINT_LO[:dim], PANDA_JOINT_HI[:dim]
    return np.clip(np.asarray(q0, np.float64) + radius * direction, lo, hi)


def _node_feature_matrix(nodes: NodeTable) -> np.ndarray:
    """[q, qdot, grip, t_frac] raw columns -- GraphRetrieverNet.node_encode's expected input.

    NOT NodeTable.features, which already applies psi() and feeds the ONF field trunk instead.

    Args:
        nodes: Node table to read.

    Returns:
        [V, 2*dim+2] float32 feature matrix.
    """
    grip = np.asarray(nodes.grip, dtype=np.float32)[:, None]
    t_frac = np.asarray(nodes.t_frac, dtype=np.float32)[:, None]
    return np.concatenate([nodes.q, nodes.qdot, grip, t_frac], axis=1).astype(np.float32)


class _QueryRow(NamedTuple):
    """One row of make_queries output, as returned by each per-class row-builder.

    Attributes:
        q_hist: [window, dim] float32 joint-position window.
        src_owner: Demo the window was drawn from.
        src_node: Node at the window's own end, T.
        tgt_node: Node at T+advance -- the retrieval label.
        t_raw_now: The window's own current RAW frame, T. Not recoverable from src_node, whose
            t_raw is up to coarsen-1 frames behind it.
        is_clean: Provenance: True for TRAVERSAL, False for the perturbed classes.
        is_entry_static: Whether this row is an ENTRY_STATIC row.
        is_drift: Whether this row is a DRIFT row.
        break_row: First perturbed row of the window, or -1 when every row is clean.
        perturb_radius: Realised (post-clip) offset radius at the window's end, rad; 0 when clean.
    """

    q_hist: np.ndarray
    src_owner: int
    src_node: int
    tgt_node: int
    t_raw_now: int
    is_clean: bool
    is_entry_static: bool
    is_drift: bool
    break_row: int
    perturb_radius: float


class _QueryCorpus(NamedTuple):
    """The per-call corpus context shared by the three row-builders.

    Attributes:
        nodes: The node table.
        raw_ptr: [J+1] per-demo boundaries into the raw arrays.
        q_raw: [N_raw, dim] float32 raw joint configs.
        dim: Degrees of freedom.
        window: This call's window length.
    """

    nodes: NodeTable
    raw_ptr: np.ndarray
    q_raw: np.ndarray
    dim: int
    window: int


def _entry_static_row(
    corpus: _QueryCorpus, rng: np.random.RandomState, pool_entry: np.ndarray,
    entry_sigma_range: tuple[float, float],
) -> _QueryRow:
    """Build one ENTRY_STATIC row -- see make_queries.

    Args:
        corpus: The per-call corpus context.
        rng: Random generator; draws rng.randint (pick demo) then _entry_perturb's own draws.
        pool_entry: Demos long enough to serve as entry queries.
        entry_sigma_range: (lo, hi) perturbation radius range.

    Returns:
        The row.
    """
    nodes, raw_ptr, q_raw, dim, window = corpus
    j = int(pool_entry[rng.randint(len(pool_entry))])
    lo = int(raw_ptr[j])
    q0 = q_raw[lo].astype(np.float64)
    q0_pert = _entry_perturb(q0, rng, dim, entry_sigma_range)
    # Every row of this window is a copy of the SAME perturbed frame, so finite_diff_vel yields
    # identically zero velocity -- the exact shape a live t=0 call's repeat(q0, hist) produces.
    q_hist_row = np.repeat(q0_pert[None, :].astype(np.float32), window, axis=0)
    tgt = _node_at_raw(nodes, j, lo)
    return _QueryRow(q_hist=q_hist_row, src_owner=j, src_node=tgt, tgt_node=tgt, t_raw_now=lo,
                      is_clean=False, is_entry_static=True, is_drift=False, break_row=0,
                      perturb_radius=float(np.linalg.norm(q0_pert - q0)))


def _entry_row(
    corpus: _QueryCorpus, rng: np.random.RandomState, pool_entry: np.ndarray, advance: int,
    entry_sigma_range: tuple[float, float],
) -> _QueryRow:
    """Build one ENTRY row -- see make_queries.

    Args:
        corpus: The per-call corpus context.
        rng: Random generator; draws rng.randint (pick demo), _entry_perturb's own draws, then
            rng.randint (window length).
        pool_entry: Demos long enough to serve as entry queries.
        advance: Raw frames from the window end to the target.
        entry_sigma_range: (lo, hi) perturbation radius range.

    Returns:
        The row.
    """
    nodes, raw_ptr, q_raw, dim, window = corpus
    j = int(pool_entry[rng.randint(len(pool_entry))])
    lo = int(raw_ptr[j])
    demo_hi = int(raw_ptr[j + 1])
    q0 = q_raw[lo].astype(np.float64)
    q0_pert = _entry_perturb(q0, rng, dim, entry_sigma_range)
    delta = q0_pert - q0                             # the REALIZED constant offset (post-clip)

    w_use = min(int(rng.randint(1, window + 1)), demo_hi - lo)   # same W distribution as TRAVERSAL
    t_raw_end = lo + w_use - 1                                   # window's OWN (current) end
    raw_window = q_raw[lo : lo + w_use].astype(np.float64) + delta[None, :]
    raw_window = np.clip(raw_window, PANDA_JOINT_LO[:dim], PANDA_JOINT_HI[:dim])
    q_hist_row = edge_pad_hist(raw_window.astype(np.float32), window)[0]
    return _QueryRow(
        q_hist=q_hist_row, src_owner=j, src_node=_node_at_raw(nodes, j, t_raw_end),
        tgt_node=_node_at_raw(nodes, j, t_raw_end + advance), t_raw_now=t_raw_end, is_clean=False,
        is_entry_static=False, is_drift=False, break_row=0,
        perturb_radius=float(np.linalg.norm(delta)),
    )


def _drift_row(
    corpus: _QueryCorpus, rng: np.random.RandomState, pool_trav: np.ndarray, advance: int,
    entry_sigma_range: tuple[float, float],
) -> _QueryRow:
    """Build one DRIFT row -- see make_queries.

    Args:
        corpus: The per-call corpus context.
        rng: Random generator; draws rng.randint (pick demo), rng.randint (T), rng.randint (window
            length), rng.randint (breakpoint), then _entry_perturb's own draws.
        pool_trav: Demos long enough for a full window plus advance.
        advance: Raw frames from the window end to the target.
        entry_sigma_range: (lo, hi) perturbation radius range.

    Returns:
        The row.
    """
    nodes, raw_ptr, q_raw, dim, window = corpus
    j = int(pool_trav[rng.randint(len(pool_trav))])
    lo, hi = int(raw_ptr[j]), int(raw_ptr[j + 1])
    t_raw_end = lo + int(rng.randint(window - 1, (hi - lo) - advance))   # same T draw as TRAVERSAL
    w_use = int(rng.randint(DRIFT_MIN_HIST, window + 1))
    n_clean = int(rng.randint(1, w_use - 1))                             # leaves >= 2 perturbed rows
    raw_window = q_raw[t_raw_end - w_use + 1 : t_raw_end + 1].astype(np.float64)

    delta = _entry_perturb(raw_window[-1], rng, dim, entry_sigma_range) - raw_window[-1]
    # A RAMP to the full radius at T, not a step at the breakpoint: a step adds a one-frame |delta|
    # impulse to the finite-differenced velocity that no deployed window contains, while the ramp
    # leaves a sustained per-row bias -- what an arm executing a wrong action chunk actually does.
    ramp = np.arange(1, w_use - n_clean + 1, dtype=np.float64) / (w_use - n_clean)
    drifted = raw_window.copy()
    drifted[n_clean:] += ramp[:, None] * delta[None, :]
    drifted = np.clip(drifted, PANDA_JOINT_LO[:dim], PANDA_JOINT_HI[:dim])
    q_hist_row = edge_pad_hist(drifted.astype(np.float32), window)[0]
    return _QueryRow(
        q_hist=q_hist_row, src_owner=j, src_node=_node_at_raw(nodes, j, t_raw_end),
        tgt_node=_node_at_raw(nodes, j, t_raw_end + advance), t_raw_now=t_raw_end, is_clean=False,
        is_entry_static=False, is_drift=True,
        # edge_pad_hist right-aligns, so a w_use-row window's own row i sits at window - w_use + i.
        break_row=window - w_use + n_clean, perturb_radius=float(np.linalg.norm(delta)),
    )


def _traversal_row(corpus: _QueryCorpus, rng: np.random.RandomState, pool_trav: np.ndarray, advance: int) -> _QueryRow:
    """Build one TRAVERSAL row -- see make_queries.

    Args:
        corpus: The per-call corpus context.
        rng: Random generator; draws rng.randint (pick demo), rng.randint (T), rng.randint (window
            length).
        pool_trav: Demos long enough for a full window plus advance.
        advance: Raw frames from the window end to the target.

    Returns:
        The row.
    """
    nodes, raw_ptr, q_raw, _dim, window = corpus
    j = int(pool_trav[rng.randint(len(pool_trav))])
    lo, hi = int(raw_ptr[j]), int(raw_ptr[j + 1])
    length = hi - lo
    t_local = int(rng.randint(window - 1, length - advance))   # T, local to this demo
    t_raw_end = lo + t_local
    w_use = int(rng.randint(1, window + 1))                    # variable true history length
    start = t_raw_end - w_use + 1
    raw_window = q_raw[start : t_raw_end + 1]
    q_hist_row = edge_pad_hist(raw_window, window)[0]
    return _QueryRow(
        q_hist=q_hist_row, src_owner=j, src_node=_node_at_raw(nodes, j, t_raw_end),
        tgt_node=_node_at_raw(nodes, j, t_raw_end + advance), t_raw_now=t_raw_end, is_clean=True,
        is_entry_static=False, is_drift=False, break_row=-1, perturb_radius=0.0,
    )


def make_queries(
    nodes: NodeTable, rng: np.random.RandomState, n: int, owners: np.ndarray, spec: QuerySpec | None = None,
) -> dict[str, np.ndarray]:
    """Draw n training queries from demos owned by owners, four classes mixed per spec.

    TRAVERSAL (the rows no other class claims) is the primary pretext task: a window of up
    to spec.window clean raw frames ending at T, labelled with the node advance raw frames further
    along the SAME demo. ENTRY is W real consecutive frames from a demo's own start, offset by a
    single constant perturbation and clipped to the joint limits, labelled advance frames past the
    window's own unperturbed end. ENTRY_STATIC is the deployed t=0 shape: one perturbed t=0 frame
    repeated across the window, so its finite-differenced velocity is identically zero, labelled
    with the demo's own t=0 node. DRIFT is a mid-demo window whose first rows are clean and whose
    remaining rows ramp out to a perturbation of the same log-uniform radius -- the only class with
    a breakpoint inside the window. See docs/technical/02-retrieval-head.md for why each is shaped this way.

    Args:
        nodes: Node table to draw from.
        rng: Random generator driving the class mix and every row builder.
        n: Number of queries.
        owners: Demo indices eligible for this set.
        spec: Keyword options; None is exactly QuerySpec(), which reproduces the original
            bare-default call.

    Returns:
        A dict of length-n arrays: q_hist [n,window,D] f32, qdot_hist [n,window,D] f32
        (onf.graph.core.geometry.finite_diff_vel), grip_hist [n,window] f32 (broadcast from the
        query's own T node -- raw per-frame grip is not persisted), w_hist [n,window] f32
        cleanliness weights, src_owner [n] i32, src_node [n] i64 (the T node), tgt_node [n] i64
        (the T+advance node, the retrieval label; advance=0 makes it the T node itself),
        t_raw_now [n] i64 (the T RAW frame), is_clean
        [n] bool (provenance: which branch built the window), abstain_is_correct [n] bool (the
        abstain LOSS/eval label, by consequence not provenance), is_entry_static [n] bool,
        is_drift [n] bool, break_row [n] i32 (first perturbed window row, -1 when clean),
        perturb_radius [n] f32 (realised offset radius at T, rad).

    Raises:
        ValueError: If owners is empty, no demo is long enough for the traversal queries the
            requested class mix still needs, or DRIFT is requested on a window too short to hold a
            breakpoint.
    """
    spec = spec if spec is not None else QuerySpec()
    field = spec.field
    advance = schema.ADVANCE if spec.advance is None else spec.advance
    window = schema.HIST_H if spec.window is None else spec.window
    entry_frac = ENTRY_FRAC if spec.entry_frac is None else spec.entry_frac
    entry_sigma_range = spec.entry_sigma_range
    inter_demo_spacing = spec.inter_demo_spacing
    entry_static_frac = spec.entry_static_frac
    drift_frac = float(spec.drift_frac)

    owners = np.asarray(owners, dtype=np.int64)
    if len(owners) == 0:
        raise ValueError("make_queries: `owners` is empty")
    raw_ptr = np.asarray(nodes.raw_ptr, dtype=np.int64)
    lengths = np.diff(raw_ptr)

    min_len_trav = window + advance
    pool_trav = owners[lengths[owners] >= min_len_trav]
    min_len_entry = advance + 1
    pool_entry = owners[lengths[owners] >= min_len_entry]

    entry_frac = float(entry_frac)
    entry_static_frac = float(entry_static_frac)
    if len(pool_entry) == 0 and (entry_frac > 0 or entry_static_frac > 0):
        warnings.warn(
            f"make_queries: no demo in `owners` has >= {min_len_entry} raw frames -- entry_frac "
            f"{entry_frac} / entry_static_frac {entry_static_frac} forced to 0 for this call", stacklevel=2,
        )
        entry_frac = 0.0
        entry_static_frac = 0.0
    if len(pool_trav) == 0:
        if entry_frac + entry_static_frac < 1.0:
            raise ValueError(
                f"make_queries: no demo in `owners` has >= {min_len_trav} raw frames (window={window} "
                f"+ advance={advance}), and entry_frac={entry_frac}+entry_static_frac={entry_static_frac} "
                "< 1 still needs traversal queries"
            )
        pool_trav = pool_entry   # unused when entry_frac+entry_static_frac == 1.0, kept non-empty defensively
    if drift_frac > 0 and window < DRIFT_MIN_HIST:
        raise ValueError(
            f"make_queries: drift_frac={drift_frac} needs window >= {DRIFT_MIN_HIST} for a window "
            f"that can hold a clean-to-perturbed breakpoint, got window={window}"
        )

    if (entry_frac > 0 or entry_static_frac > 0 or drift_frac > 0) and entry_sigma_range is None:
        floor = entry_perturbation_floor(nodes)
        entry_sigma_range = (ENTRY_LO_MULT * floor["p50"], ENTRY_HI_MULT * floor["p99"])

    dim = nodes.dim
    q_raw = np.asarray(nodes.q_raw, dtype=np.float32)
    corpus = _QueryCorpus(nodes=nodes, raw_ptr=raw_ptr, q_raw=q_raw, dim=dim, window=window)
    q_hist = np.empty((n, window, dim), dtype=np.float32)
    src_owner = np.empty(n, dtype=np.int32)
    src_node = np.empty(n, dtype=np.int64)
    tgt_node = np.empty(n, dtype=np.int64)
    t_raw_now = np.empty(n, dtype=np.int64)
    is_clean = np.empty(n, dtype=bool)
    is_entry_static = np.zeros(n, dtype=bool)
    is_drift = np.zeros(n, dtype=bool)
    break_row = np.full(n, -1, dtype=np.int32)
    perturb_radius = np.zeros(n, dtype=np.float32)

    for i in range(n):
        # A zero fraction must reproduce the RNG stream of the mix that predates its class, so every
        # check short-circuits its own rng.rand() away rather than drawing and discarding it.
        if entry_static_frac > 0:
            is_static = rng.rand() < entry_static_frac
        else:
            is_static = False
        is_entry = False if is_static else (entry_frac > 0 and rng.rand() < entry_frac)
        drifts = False if (is_static or is_entry) else (drift_frac > 0 and rng.rand() < drift_frac)

        if is_static:
            row = _entry_static_row(corpus, rng, pool_entry, entry_sigma_range)
        elif is_entry:
            row = _entry_row(corpus, rng, pool_entry, advance, entry_sigma_range)
        elif drifts:
            row = _drift_row(corpus, rng, pool_trav, advance, entry_sigma_range)
        else:
            row = _traversal_row(corpus, rng, pool_trav, advance)

        q_hist[i] = row.q_hist
        src_owner[i] = row.src_owner
        src_node[i] = row.src_node
        tgt_node[i] = row.tgt_node
        t_raw_now[i] = row.t_raw_now
        is_clean[i] = row.is_clean
        is_entry_static[i] = row.is_entry_static
        is_drift[i] = row.is_drift
        break_row[i] = row.break_row
        perturb_radius[i] = row.perturb_radius

    qdot_hist = finite_diff_vel(q_hist)
    grip_hist = np.repeat(np.asarray(nodes.grip, dtype=np.float32)[src_node][:, None], window, axis=1)
    w_hist = CleanlinessScorer(field).weights_batch(q_hist)

    spacing = float(inter_demo_spacing) if inter_demo_spacing is not None else measure_inter_demo_spacing(nodes)
    abstain_is_correct = off_manifold_batch(nodes, q_hist[:, -1, :]) < spacing

    return {
        "q_hist": q_hist, "qdot_hist": qdot_hist, "grip_hist": grip_hist, "w_hist": w_hist,
        "src_owner": src_owner, "src_node": src_node, "tgt_node": tgt_node,
        "t_raw_now": t_raw_now, "is_clean": is_clean,
        "abstain_is_correct": abstain_is_correct, "is_entry_static": is_entry_static,
        "is_drift": is_drift, "break_row": break_row, "perturb_radius": perturb_radius,
    }

class NegativeMiner:
    """Hard-negative sampler bound to one fixed node table.

    A class rather than a free function because the loss mines once per row, and the free-function
    form re-ran np.asarray over the [V, D] node arrays on every one of those calls. The sampling
    itself is unchanged, RNG draw for RNG draw.

    Three categories, each falling back to uniform random negatives when it is empty or under quota
    -- most importantly bucket (b), measured at only ~4.1% of frames:
    (a) same t_idx as the positive, different strand;
    (b) among the query's selfx_k nearest nodes in raw q, those whose phase differs from the
        positive's by more than selfx_phase_gap;
    (c) the query's own single nearest node, if it is not the positive.
    """

    def __init__(
        self, nodes: NodeTable, *, selfx_k: int = SELFX_K, selfx_phase_gap: float = SELFX_PHASE_GAP,
    ) -> None:
        """
        Args:
            nodes: The node table to mine from.
            selfx_k: Size of the local kNN pool searched for bucket (b).
            selfx_phase_gap: Phase gap that promotes a near-in-q neighbour into bucket (b).
        """
        self.selfx_k = int(selfx_k)
        self.selfx_phase_gap = float(selfx_phase_gap)
        self.n_nodes = len(nodes)
        self._owner = np.asarray(nodes.owner)
        self._t_idx = np.asarray(nodes.t_idx)
        self._phase = np.asarray(nodes.phase)
        self._q = np.asarray(nodes.q, dtype=np.float64)
        self._all_ids = np.arange(self.n_nodes, dtype=np.int64)

    def mine(
        self, query: np.ndarray, positive: int, rng: np.random.RandomState, n_neg: int,
    ) -> np.ndarray:
        """Mine hard negatives for one query.

        Args:
            query: [D] raw joint configuration, typically the window's most recent frame.
            positive: Node id of the true continuation; never returned.
            rng: Random generator, consumed in the documented bucket order (a, b, c) then the
                uniform top-up, then one final shuffle.
            n_neg: Cap on the number of negatives.

        Returns:
            [<=n_neg] i64 node ids.
        """
        v = self.n_nodes
        positive = int(positive)
        if v <= 1 or n_neg <= 0:
            return np.zeros(0, dtype=np.int64)

        owner, t_idx, phase = self._owner, self._t_idx, self._phase
        query = np.asarray(query, dtype=np.float64).reshape(-1)
        pos_owner, pos_tidx, pos_phase = int(owner[positive]), int(t_idx[positive]), float(phase[positive])

        bucket_a = np.flatnonzero((t_idx == pos_tidx) & (owner != pos_owner))

        d = np.linalg.norm(self._q - query[None, :], axis=1)
        # A full argsort, NOT argpartition: order[0] IS bucket (c), and bucket (b) is a slice of this
        # order, so a different tie-break would change both the negatives drawn and the rng draws spent.
        order = np.argsort(d)
        knn = order[: min(self.selfx_k, v)]
        bucket_b = knn[np.abs(phase[knn] - pos_phase) > self.selfx_phase_gap]

        bucket_c = np.array([int(order[0])], dtype=np.int64)

        quota = max(1, n_neg // 3)
        picks: list[np.ndarray] = []
        for bucket in (bucket_a, bucket_b, bucket_c):
            bucket = np.asarray(bucket, dtype=np.int64)
            bucket = bucket[bucket != positive]
            if len(bucket) == 0:
                continue
            take = min(quota, len(bucket))
            picks.append(rng.choice(bucket, size=take, replace=False))
        picked = np.unique(np.concatenate(picks)) if picks else np.zeros(0, dtype=np.int64)

        if len(picked) < n_neg:
            pool = np.setdiff1d(self._all_ids, np.append(picked, positive))
            need = n_neg - len(picked)
            if len(pool) > 0:
                extra = rng.choice(pool, size=min(need, len(pool)), replace=False)
                picked = np.concatenate([picked, extra])

        rng.shuffle(picked)
        return picked[:n_neg].astype(np.int64)


def mine_negatives(
    nodes: NodeTable, query: np.ndarray, positive: int, rng: np.random.RandomState, n_neg: int,
) -> np.ndarray:
    """One-shot NegativeMiner.mine for a caller that has no miner to hand.

    Args:
        nodes: Node table to mine from.
        query: [D] raw joint configuration.
        positive: Node id of the true continuation.
        rng: Random generator.
        n_neg: Cap on the number of negatives.

    Returns:
        [<=n_neg] i64 node ids.
    """
    return NegativeMiner(nodes).mine(query, positive, rng, n_neg)

class SoftTargetBuilder:
    """Multi-positive node-identity targets for one fixed node table.

    Owns the graph arrays (smooth is called once per training row, and re-deriving [V, D] q per row
    is pure waste), the ENTRY_STATIC task_id -> pool cache, and the log-once empty-pool latch. Both
    target formulas are unchanged; this class only decides where their inputs live.

    Attributes:
        pool: task_id -> node ids with phase <= entry_static_band, the ENTRY_STATIC target pool.
    """

    def __init__(
        self, nodes: NodeTable, *, phase_band: float = schema.WHERE_PHASE_BAND,
        move_temp: float = schema.WHERE_MOVE_TEMP, true_bonus: float = schema.WHERE_TRUE_BONUS,
        entry_static_band: float = ENTRY_STATIC_BAND,
    ) -> None:
        """
        Args:
            nodes: Node table the targets index into.
            phase_band: Half-width of the accepted phase band around the positive.
            move_temp: Temperature of the inverse-movement-cost weighting.
            true_bonus: Multiplier applied to the true continuation before renormalising.
            entry_static_band: Phase ceiling defining the ENTRY_STATIC pool.
        """
        self.phase_band = float(phase_band)
        self.move_temp = float(move_temp)
        self.true_bonus = float(true_bonus)
        self.entry_static_band = float(entry_static_band)
        self._task_id = np.asarray(nodes.task_id)
        self._phase = np.asarray(nodes.phase)
        self._q = np.asarray(nodes.q, dtype=np.float64)
        self.pool = _entry_static_pools(self._task_id, self._phase, self.entry_static_band)
        self._warned = False

    def smooth(
        self, positive: int, q_now: np.ndarray | None = None,
    ) -> tuple[np.ndarray, np.ndarray]:
        """The movement-weighted multi-positive target. See label_smooth_targets.

        Args:
            positive: Node id of the true continuation.
            q_now: [D] current raw joint configuration, or None for the phase-distance proxy.

        Returns:
            (candidates [K] i64, weights [K] f32 summing to 1).
        """
        task_id, phase = self._task_id, self._phase
        positive = int(positive)
        pos_task = int(task_id[positive])
        pos_phase = float(phase[positive])

        cand = np.flatnonzero((task_id == pos_task) & (np.abs(phase - pos_phase) <= self.phase_band))
        if positive not in cand:
            cand = np.append(cand, np.int64(positive))
        cand = cand.astype(np.int64)

        if q_now is not None:
            q_now = np.asarray(q_now, dtype=np.float64).reshape(-1)
            cost = np.linalg.norm(self._q[cand] - q_now[None, :], axis=1)
        else:
            cost = np.abs(phase[cand].astype(np.float64) - pos_phase)

        logw = -cost / max(self.move_temp, 1e-8)
        logw = logw - logw.max()                  # numerical stability only; exp() below is unaffected
        w = np.exp(logw)
        w[cand == positive] *= self.true_bonus
        return cand, (w / w.sum()).astype(np.float32)

    def entry_static(
        self, positive: int, q_now: np.ndarray | None = None,
    ) -> tuple[np.ndarray, np.ndarray]:
        """The ENTRY_STATIC class's target: uniform over the positive's task pool.

        The deployed hand rule this class exists to teach IS "uniform over task lane intersect
        phase<=band", so the label reproduces that rule exactly. An empty pool is rare and worth
        knowing about, so it warns -- once per builder, not once per row.

        Args:
            positive: Node id of the true continuation.
            q_now: [D] current raw joint configuration, used only by the fallback.

        Returns:
            (candidates [K] i64, weights [K] f32 summing to 1).
        """
        task_id = int(self._task_id[int(positive)])
        pool = self.pool.get(task_id)
        if pool is not None and len(pool) > 0:
            return pool, np.full(len(pool), 1.0 / len(pool), dtype=np.float32)

        if not self._warned:
            self._warned = True
            warnings.warn(
                f"SoftTargetBuilder: empty ENTRY_STATIC pool for task_id={task_id} (no node at "
                f"phase<=ENTRY_STATIC_BAND={self.entry_static_band}) -- falling back to the smoothed "
                "target for this row (further occurrences not logged)", stacklevel=2,
            )
        return self.smooth(positive, q_now)


def _entry_static_pools(
    task_id: np.ndarray, phase: np.ndarray, band: float
) -> dict[int, np.ndarray]:
    """task_id -> node ids with phase <= band.

    Args:
        task_id: [V] per-node task ids.
        phase: [V] per-node phase.
        band: Phase ceiling.

    Returns:
        One i64 id array per distinct task id.
    """
    in_band = phase <= float(band)
    return {int(t): np.flatnonzero(in_band & (task_id == t)).astype(np.int64) for t in np.unique(task_id)}


def label_smooth_targets(
    nodes: NodeTable, positive: int, q_now: np.ndarray | None = None, *,
    phase_band: float = schema.WHERE_PHASE_BAND, move_temp: float = schema.WHERE_MOVE_TEMP,
    true_bonus: float = schema.WHERE_TRUE_BONUS,
) -> tuple[np.ndarray, np.ndarray]:
    """Soft multi-positive target: every node sharing positive's task_id within +/- phase_band.

    Weighted by inverse movement cost from q_now: w_v = exp(-||nodes.q[v] - q_now|| / move_temp),
    normalised to sum to 1. q_now=None falls back to a phase-distance proxy so the target still
    decays smoothly within the band rather than collapsing to uniform. The true continuation is
    always the single highest-weighted candidate (cost 0, then multiplied by true_bonus), and is
    always included, so the candidate array is never empty.

    One-shot form of SoftTargetBuilder.smooth; training goes through a long-lived builder.

    Args:
        nodes: Node table the targets index into.
        positive: Node id of the true continuation.
        q_now: [D] current raw joint configuration, or None.
        phase_band: Half-width of the accepted phase band.
        move_temp: Temperature of the inverse-movement-cost weighting.
        true_bonus: Multiplier applied to the true continuation.

    Returns:
        (candidates [K] i64, weights [K] f32 summing to 1).
    """
    return SoftTargetBuilder(
        nodes, phase_band=phase_band, move_temp=move_temp, true_bonus=true_bonus
    ).smooth(positive, q_now)


def build_entry_static_pool_cache(nodes: NodeTable, band: float = ENTRY_STATIC_BAND) -> dict[int, np.ndarray]:
    """task_id -> node ids with phase <= band -- the ENTRY_STATIC class's soft-target pool.

    Training reads the same mapping off SoftTargetBuilder.pool, built once per run.

    Args:
        nodes: Node table to bucket.
        band: Phase ceiling; ENTRY_STATIC_BAND matches the deployed GR_ENTRY_BAND.

    Returns:
        One i64 id array per distinct task id.
    """
    return _entry_static_pools(np.asarray(nodes.task_id), np.asarray(nodes.phase), band)