onf-c0/scripts/probe_task_belief.py

#!/usr/bin/env python
"""Does the policy's own action chunk identify WHICH TASK it is performing?

One decisive measurement, no training and no rollout. The hypothesis is stated in full in
onf.graph.run.task_belief's docstring; this script is its falsification test. For N raw frames
spread across the demonstration corpus we run the FROZEN retriever, split the posterior into one
expected chunk per task lane, and ask which lane's chunk the cached policy chunk is closest to.

    s_ret    task_marginal(probs)            the joint-space channel the deleted TaskPosterior had
    s_beh    -|a_policy - a_track^(k)|       NEW: does the policy MOVE like lane k's demos?

The reference numbers this is judged against: the deleted joint-space channel scored 79.2% ungated
at k=8 and was confidently wrong on 20.8% of episodes, 44 of its 63 errors being the single pair
LIVING_ROOM_SCENE5 / SCENE6 (task ids 7 and 8, both "put the white mug on the ... plate"). The text
encoder alone is 64.2% on real rewordings. Chance is 10%. s_beh has to beat s_ret to be worth
building, and its errors have to NOT be that same pair -- an accumulating belief can wait out a
tie between two lanes whose demonstrations coincide, but not a channel that is confidently wrong.

The language prior is OFF throughout (the default SeedLanguage()): the whole point is to measure
what is left when text is removed, so a text-seeded posterior would smuggle the answer in.

The extraction stack itself is onf.graph.build.belief_fit.TaskBeliefProbe, which the belief fit
reuses row for row; what is left here is this measurement's own reporting.

Usage:

    OMP_NUM_THREADS=4 PYTHONPATH=src CUDA_VISIBLE_DEVICES=3 python scripts/probe_task_belief.py \
        --artifacts outputs/long/stage2_alpha_v5/artifacts
"""

from __future__ import annotations

import argparse
import math
import sys
from collections import Counter
from pathlib import Path
from typing import Sequence

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                                          # noqa: E402
from onf.graph.build.belief_fit import CHECK_EVERY, CHUNK_K, Channels, TaskBeliefProbe  # noqa: E402

ACCUM_POINTS = (1, 2, 4, 8, 16)    # report accuracy after this many accumulated checks
_EPS = 1e-12

# Report order, and the one-line reason each row is in the table at all. See Channels.
CHANNEL_NOTES: dict[str, str] = {
    "s_ret": "BASELINE: joint-space marginal, all the deleted channel had",
    "s_beh": "NEW: policy chunk vs each lane's action-unit chunk",
    "s_beh_ff": "CONTROL: chunk rows 1: only -- motion, no live-pose term",
    "s_beh_pose": "BUG CONTROL: metres scored against action units",
    "s_beh + s_ret": "FUSION: the two z-summed across tasks per row",
    "s_beh_ff+s_ret": "FUSION of the motion-only control -- does MOTION add to joint space?",
}


def zscore(x: np.ndarray) -> np.ndarray:
    """Standardise each row across the task axis.

    The two channels live in different units -- s_ret is a probability, s_beh a negative distance in
    action units -- so summing them raw would be a fixed, arbitrary weighting. Per-row
    standardisation across tasks is the only reduction that leaves the argmax of each channel alone
    while making their SUM a meaningful vote.

    Args:
        x: [n, n_tasks] scores.

    Returns:
        [n, n_tasks] standardised scores.
    """
    return (x - x.mean(axis=1, keepdims=True)) / np.clip(x.std(axis=1, keepdims=True), _EPS, None)


# ==================================================================================================
# reporting
# ==================================================================================================
def top1(scores: np.ndarray, truth: np.ndarray) -> float:
    """Fraction of rows whose argmax over tasks is the true task.

    Args:
        scores: [n, n_tasks] per-task scores, higher is better.
        truth: [n] true task ids.

    Returns:
        Accuracy in [0, 1].
    """
    return float((scores.argmax(axis=1) == truth).mean())


def report_channels(ch: Channels) -> dict[str, np.ndarray]:
    """The measured channels plus the fused one, in report order.

    Args:
        ch: The measured channels.

    Returns:
        Name -> [n, n_tasks] scores.

    Raises:
        KeyError: A channel the fusion needs is missing.
    """
    ret = zscore(ch.scores["s_ret"])
    return {
        **ch.scores,
        "s_beh + s_ret": zscore(ch.scores["s_beh"]) + ret,
        "s_beh_ff+s_ret": zscore(ch.scores["s_beh_ff"]) + ret,
    }


def mcnemar(baseline: np.ndarray, candidate: np.ndarray) -> tuple[int, int, float]:
    """Exact two-sided McNemar test on two per-frame correctness vectors.

    The channels are scored on the SAME frames, so the comparison is paired and only the discordant
    frames carry information. An unpaired proportion test on 600 frames would call a 4pp gap
    ambiguous that the pairing resolves cleanly.

    Args:
        baseline: [n] bool, whether the baseline channel was right on each frame.
        candidate: [n] bool, the same for the candidate.

    Returns:
        (frames the candidate wins, frames it loses, two-sided exact p; nan when none disagree).
    """
    wins = int(np.sum(candidate & ~baseline))
    losses = int(np.sum(baseline & ~candidate))
    n = wins + losses
    if n == 0:
        return wins, losses, float("nan")
    k = min(wins, losses)
    tail = sum(math.comb(n, i) for i in range(k + 1)) / 2.0**n
    return wins, losses, float(min(1.0, 2.0 * tail))


def single_check_table(ch: Channels) -> str:
    """The headline table: one accuracy per channel over independent frames.

    Args:
        ch: The measured channels.

    Returns:
        The rendered table.
    """
    channels = report_channels(ch)
    base = channels["s_ret"].argmax(axis=1) == ch.truth
    out = [f"{'channel':<16}{'top-1':>8}{'n':>6}{'vs s_ret':>22}   note", "-" * 100]
    for name, scores in channels.items():
        correct = scores.argmax(axis=1) == ch.truth
        wins, losses, p = mcnemar(base, correct)
        verdict = "--" if name == "s_ret" else f"+{wins}/-{losses}  p={p:.2g}"
        out.append(
            f"{name:<16}{100 * correct.mean():>7.1f}%{len(ch.truth):>6}{verdict:>22}   "
            f"{CHANNEL_NOTES[name]}"
        )
    out += [
        "-" * 100,
        f"{'chance':<16}{100.0 / ch.n_tasks:>7.1f}%",
        "reference: deleted joint-space-only channel 79.2% ungated at k=8 (confidently wrong",
        "           on 20.8%) | text encoder alone 64.2% on real rewordings | chance 10%",
    ]
    return "\n".join(out)


def accumulation_table(ch: Channels, n_demos: int, n_checks: int) -> str:
    """Accuracy as a function of how many checks have been summed into the belief.

    The task does not change during an episode, so the belief's transition is the identity and
    evidence simply adds. Summing raw per-check scores is exactly that -- a log-space product under
    a flat prior for s_ret, and a total deviation for s_beh. Check 1 sits at the very start of the
    episode, which is where the two channels differ most: the arm has barely moved, so joint space
    cannot yet tell the lanes apart.

    Args:
        ch: Channels measured on accumulation_frames, in row-major [demo, check] order.
        n_demos: Number of demos walked.
        n_checks: Checks per demo.

    Returns:
        The rendered table.
    """
    per_check = {k: v.reshape(n_demos, n_checks, -1) for k, v in report_channels(ch).items()}
    truth = ch.truth.reshape(n_demos, n_checks)[:, 0]

    names = list(per_check)
    out = [f"{'checks':>7}" + "".join(f"{n:>16}" for n in names), "-" * (7 + 16 * len(names))]
    for m in ACCUM_POINTS:
        if m > n_checks:
            continue
        out.append(
            f"{m:>7}"
            + "".join(f"{100 * top1(per_check[n][:, :m].sum(axis=1), truth):>15.1f}%" for n in names)
        )
    out.append("-" * (7 + 16 * len(names)))
    out.append(
        f"n_demos={n_demos}, one check every {CHECK_EVERY} raw frames, check 1 at the episode start"
    )
    return "\n".join(out)


def confusion(scores: np.ndarray, ch: Channels, names: Sequence[str], top: int = 5) -> str:
    """The ordered (true -> predicted) pairs that dominate one channel's errors.

    Args:
        scores: [n, n_tasks] the channel's scores.
        ch: The measured channels, for the truth column.
        names: Task names, for the readable label.
        top: How many pairs to print.

    Returns:
        The rendered list.
    """
    pred = scores.argmax(axis=1)
    wrong = pred != ch.truth
    pairs = Counter(zip(ch.truth[wrong].tolist(), pred[wrong].tolist()))
    n_err = int(wrong.sum())
    out = [f"top-{top} confused ordered pairs (true -> predicted), {n_err} errors total", "-" * 72]
    for (t, p), count in pairs.most_common(top):
        out.append(
            f"  {t} -> {p}   {count:>4}  ({100 * count / max(n_err, 1):>4.1f}% of errors)   "
            f"{names[t][:28]} -> {names[p][:28]}"
        )
    share = sum(c for (t, p), c in pairs.items() if {t, p} == {7, 8})
    out.append("-" * 72)
    out.append(
        f"  the old channel's failure pair 7<->8 (SCENE5/SCENE6 white mug): {share} errors "
        f"= {100 * share / max(n_err, 1):.1f}% of this channel's errors"
    )
    return "\n".join(out)


def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
    """Parse the command line.

    Args:
        argv: Argument vector; None reads sys.argv.

    Returns:
        The namespace.
    """
    ap = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    ap.add_argument("--artifacts", default="outputs/long/stage2_alpha_v5/artifacts")
    ap.add_argument("--suite", default="long", help="suite whose ONFField weights the windows")
    ap.add_argument("--n-frames", type=int, default=600, help="independent frames for the headline table")
    ap.add_argument("--n-demos", type=int, default=80, help="demos walked for the accumulation table")
    ap.add_argument("--n-checks", type=int, default=max(ACCUM_POINTS), help="checks per demo")
    ap.add_argument("--batch", type=int, default=32, help="forward-pass batch size")
    ap.add_argument("--device", default="cuda", help="torch device")
    ap.add_argument("--seed", type=int, default=0)
    return ap.parse_args(argv)


def main(argv: Sequence[str] | None = None) -> int:
    """Run both passes and print the three tables.

    Args:
        argv: Argument vector; None reads sys.argv.

    Returns:
        Process exit status.
    """
    args = parse_args(argv)
    artifacts = Path(args.artifacts)
    # from_env, not a bare GraphConfig: a bare one would silently reset every other GR_* knob away
    # from whatever the head was trained and is deployed under.
    cfg = GraphConfig.from_env()
    cfg.hist, cfg.device = CHUNK_K, args.device
    # entry_band belongs to the deployed t=0 readout, not to this mid-episode posterior; it is not
    # read by _forward_batch at all, and is named here only so a reader does not go looking.

    probe = TaskBeliefProbe(artifacts, cfg, args.suite, args.device)
    names = list(probe.nodes.task_names)
    print(
        f"[probe] graph {artifacts}: V={len(probe.nodes)} nodes, {probe.nodes.n_demos} demos, "
        f"{probe.n_tasks} tasks, {probe.nodes.n_raw} raw frames",
        file=sys.stderr,
    )

    rng = np.random.RandomState(args.seed)
    frames = probe.spread_frames(args.n_frames, rng)
    single = probe.channels(frames, batch=args.batch)

    grid = probe.accumulation_frames(args.n_demos, args.n_checks, rng)
    accum = probe.channels(grid.reshape(-1), batch=args.batch)

    print("\n=== ONE CHECK, independent frames ===")
    print(single_check_table(single))
    print("\n=== ACCUMULATED over an episode ===")
    print(accumulation_table(accum, grid.shape[0], grid.shape[1]))
    for name in ("s_beh", "s_beh_ff", "s_ret"):
        print(f"\n=== {name} CONFUSION -- {CHANNEL_NOTES[name]} ===")
        print(confusion(single.scores[name], single, names))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())