onf-c0/scripts/trace_query.py

#!/usr/bin/env python
"""Walk ONE query through the whole t=0 retrieval pipeline, printing what each stage produces.

Reading graph/run/retrieve.py top to bottom tells you what the code does; running this tells you
what the numbers look like while it does it. Every line below corresponds to one stage of the chain:

    q_hist  --cleanliness-->  pooling weights  --observe-->  per-node scores  --retrieve-->  target

Usage:

    OMP_NUM_THREADS=4 PYTHONPATH=src python scripts/trace_query.py [suite] [perturbation_rad]

The perturbation is how far off the demonstration manifold to place the query, in radians.

It prints the SAME query under both readouts, because the difference is the thing most likely to
confuse a reader of retrieve.py:

  entry_band > 0  (the deployed t=0 recipe) -- observe() masks the pool to
                  task lane AND phase <= entry_band AND reachable, and the aggregation over what
                  survives is UNIFORM. So node == -1 and q_star is the barycentre of that task's
                  start-pose cloud. That is by design, not a failure: at t=0 there is no motion
                  history, so the mask decides the answer and the head's weighting within it does not.
  entry_band = 0  -- no band mask; the head's own posterior decides, and the target lands wherever
                  the query most resembles a demo frame, which mid-episode is NOT the start.

node == -1 in both cases: both registered readout arms (euc_raw, basin) synthesise a config
rather than returning a graph node, so -1 is the normal answer, not an error code. abstain is
likewise computed either way, from the observation model alone; the entry path does not act on it.
"""

from __future__ import annotations

import sys
from pathlib import Path

import numpy as np

REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT / "src") not in sys.path:
    sys.path.insert(0, str(REPO_ROOT / "src"))

from onf.config import GraphConfig, default_paths          # noqa: E402
from onf.graph.core.geometry import finite_diff_vel  # noqa: E402
from onf.graph.run.retrieve import GraphRetriever          # noqa: E402


def trace(suite: str = "long", perturb: float = 0.05, node_idx: int = 100) -> None:
    paths = default_paths()
    cfg = GraphConfig(hist=8, entry_band=0.05, device="cpu")
    r = GraphRetriever.load(paths.graph(suite), cfg=cfg, field_dir=paths.fwm(suite), device="cpu")
    n = r.nodes

    print(f"GRAPH    {suite}: V={len(n)} nodes  E={len(r.edges.src)} edges  "
          f"demos={n.n_demos}  tasks={n.n_tasks}  dim={n.dim}")
    print(f"         field: {type(r.field).__name__ if r.field else 'None (uniform weights)'}")

    q0 = n.q[node_idx].astype(np.float64) + perturb
    q_hist = np.repeat(q0[None, :], cfg.hist, axis=0)
    w = r.scorer.weights(q_hist)
    print(f"\n1 QUERY  q_hist {q_hist.shape} {q_hist.dtype} -- at t=0 the window is ONE pose repeated,")
    print(f"         so finite_diff_vel is all-zero: {np.allclose(finite_diff_vel(q_hist), 0)}")
    print(f"         cleanliness w {w.shape} in [{w.min():.4f}, {w.max():.4f}] -- flat here BY "
          f"CONSTRUCTION (identical rows)")

    obs = r.observe(q_hist)
    finite = int(np.isfinite(obs.logits64.numpy()).sum())
    print(f"\n2 OBSERVE one score per node: logits64 {tuple(obs.logits64.shape)}")
    print(f"         {finite}/{len(n)} survive the task/owner/entry-band masks")
    print(f"         abstain_logit {float(obs.abstain_logit):+.3f}  vs  best node {float(obs.node_max):+.3f}"
          f"   -> {'ABSTAIN' if float(obs.abstain_logit) > float(obs.node_max) else 'retrieve'}")

    res = r.retrieve(q_hist)
    print(f"\n3 RESULT (entry_band={cfg.entry_band}, the DEPLOYED t=0 recipe)")
    print(f"         node={res.node}  <- -1 means 'no single node': q_star is the uniform barycentre")
    print(f"                            of the masked pool. This is the design, not a fallback.")
    print(f"         WHEN  depth (phase in [0,1]) = {res.depth:.4f}")
    print(f"         WHERE q_star {res.q_star.shape}, q_seg {res.q_seg.shape} -- a SEGMENT, not a point,")
    print(f"               so the policy gets its velocity back too")
    print(f"         conf={res.conf:.4f}  abstain={res.abstain} (computed, but the entry path ignores it)")
    print(f"         |q_query - q_star| = {np.linalg.norm(q0 - res.q_star):.4f} rad"
          f"  <- what controller.py's PD loop closes")

    cfg2 = GraphConfig(hist=8, entry_band=0.0, device="cpu")
    r2 = GraphRetriever.load(paths.graph(suite), cfg=cfg2, field_dir=paths.fwm(suite), device="cpu")
    res2 = r2.retrieve(q_hist)
    print(f"\n4 SAME QUERY, entry_band=0 (no band mask -- the head's own posterior decides)")
    print(f"         node={res2.node}  owner(demo)={res2.owner}  task_id={res2.task_id}")
    print(f"         depth={res2.depth:.4f}  conf={res2.conf:.4f}  abstain={res2.abstain}")
    print(f"         |q_query - q_star| = {np.linalg.norm(q0 - res2.q_star):.4f} rad")
    print(f"\n         Same query, two different targets. Which one ships is a CONFIG choice"
          f"\n         (GR_ENTRY_BAND), not a code path someone picked at random.")


if __name__ == "__main__":
    trace(sys.argv[1] if len(sys.argv) > 1 else "long",
          float(sys.argv[2]) if len(sys.argv) > 2 else 0.05)