verify/_da3_spatial_specificity.py

"""Does the model use SPATIAL CONTENT, or just a constant bias on the bank pathway?

Zero-ablation cannot tell these apart: if the banks carried a learned constant, zeroing it would be
equally catastrophic. This isolates the question. CPU-only -- never touches the training GPUs.

  T1 BANK DIVERSITY   Are the banks actually different per sample? If banks are ~identical across
                      different observations, they ARE a constant and nothing else matters.
  T2 BUILDER SENSITIVITY  Feed sample A's DA3 features vs sample B's -- do the banks change?
  T3 INTERVENTION LOSS  Same batch, same rng, four conditions:
                        correct / shuffled (each sample gets ANOTHER's geometry) / mean-bank
                        (per-sample info removed, magnitude kept) / zeroed.
                        shuffled ~ correct  => model ignores WHICH geometry (generic signal)
                        shuffled ~ zeroed   => model reads per-sample geometry (what we want)
                        mean ~ correct      => a constant component carries everything
"""
import os, sys, dataclasses, json
import numpy as np, jax, jax.numpy as jnp, flax.nnx as nnx
sys.path.insert(0, "scripts")

CKPT = os.environ["CONTRIB_CKPT"]
BS = int(os.environ.get("SPEC_BS", "4"))
NFS = int(os.environ.get("SPEC_NFS", "4"))


def main():
    os.environ.update(USE_DA3_FULL="1", DA3_INIT_STD="0.01", DA3_LOGIT_GAIN="1",
                      B1K_ACTIVITIES="clean_up_your_desk", B1K_EXTRACT_DEVICES=os.environ.get("B1K_EXTRACT_DEVICES","cpu"),
                      B1K_2026_ROOT="/work/jack/behavior1k/data/behavior_2026_task29_42",
                      B1K_INIT_PARAMS=CKPT)
    import train_2026 as t26
    from b1k.training.b1k_da3 import create_v3_behavior_da3_loader

    config = dataclasses.replace(t26.build_config(), batch_size=BS, num_workers=0)
    sharding = jax.sharding.SingleDeviceSharding(jax.devices("cpu")[0])
    print(f"[1] building CPU loader (BS={BS})...", flush=True)
    loader = create_v3_behavior_da3_loader(
        config, os.environ["B1K_2026_ROOT"], ["clean_up_your_desk"],
        "/work/jack/behavior1k/task_data.json",
        lang_cache="/work/jack/behavior1k/modernbert_b1k_tasks.pkl",
        sharding=sharding, shuffle=True, num_workers=0)
    obs, actions = next(iter(loader))

    model = config.model.create(jax.random.key(0))
    loaded = config.weight_loader.load(jax.tree.map(np.asarray, nnx.state(model).to_pure_dict()))
    gd, st = nnx.split(model); st.replace_by_pure_dict(loaded); model = nnx.merge(gd, st)
    raw = json.load(open("/work/jack/behavior1k/checkpoints/behavior_50t_checkpoint/assets/"
                         "IliaLarchenko/behavior_224_rgb/norm_stats.json"))["norm_stats"]
    model.load_correlation_matrix({"actions": {k: (np.asarray(v, np.float32) if isinstance(v, list) else v)
                                               for k, v in raw["actions"].items()}})
    print("[2] checkpoint loaded", flush=True)

    banks = model._compute_banks(obs)
    keys = sorted(banks.keys()) if isinstance(banks, dict) else None
    print(f"[3] banks type={type(banks).__name__} keys={keys}", flush=True)

    # ---------- T1: are banks different per sample? ----------
    print("\n=== T1 BANK DIVERSITY (constant check) ===")
    print("    cos~1.000 between different samples => the bank IS effectively a constant")
    for k in keys:
        b = np.asarray(banks[k], np.float32)              # [B, tokens, dim]
        flat = b.reshape(b.shape[0], -1)
        n = flat / (np.linalg.norm(flat, axis=1, keepdims=True) + 1e-9)
        C = n @ n.T
        off = C[~np.eye(C.shape[0], dtype=bool)]
        # variance across samples vs across tokens: if per-sample var << per-token var, it's constant
        var_across_samples = float(b.var(axis=0).mean())
        var_across_tokens = float(b.var(axis=1).mean())
        mean_b = b.mean(axis=0, keepdims=True)
        resid = float(np.linalg.norm(b - mean_b) / (np.linalg.norm(b) + 1e-9))
        print(f"  {k:5s} shape={b.shape}  pairwise cos across samples: mean={off.mean():+.4f} "
              f"max={off.max():+.4f}")
        print(f"        var(across samples)={var_across_samples:.5f}  var(across tokens)={var_across_tokens:.5f}"
              f"   ||b-mean||/||b||={resid:.4f}  (near 0 => constant)")

    # ---------- T2: is the bank builder sensitive to its DA3 input? ----------
    print("\n=== T2 BUILDER INPUT SENSITIVITY ===")
    print("    swap sample0's DA3 features for sample1's; banks must change if geometry is read")
    feats = obs.da3_features
    swapped = feats.at[0].set(feats[1])
    obs_sw = dataclasses.replace(obs, da3_features=swapped)
    banks_sw = model._compute_banks(obs_sw)
    for k in keys:
        a = np.asarray(banks[k], np.float32)[0]
        c = np.asarray(banks_sw[k], np.float32)[0]
        rel = float(np.linalg.norm(c - a) / (np.linalg.norm(a) + 1e-9))
        print(f"  {k:5s} rel change of sample0's bank after feature swap = {rel:.4f} "
              f"{'<-- INSENSITIVE (bank ignores DA3 features!)' if rel < 1e-3 else ''}")

    # ---------- T3: loss under interventions ----------
    print("\n=== T3 INTERVENTION LOSS (same batch, same rng) ===")
    rng = jax.random.key(7)
    orig_fn = model._compute_banks

    def run(tag, make):
        # new _compute_banks signature: (obs, return_aux=False, depth_drop_rng=None) -> banks
        # or (banks, aux) when return_aux. compute_detailed_loss calls it with return_aux=True.
        model._compute_banks = (lambda o, return_aux=False, depth_drop_rng=None:
                                ((make(), None) if return_aux else make()))
        ld = model.compute_detailed_loss(rng, obs, actions, train=False, num_flow_samples=NFS)
        model._compute_banks = orig_fn
        a = float(jnp.mean(ld["action_loss"]))
        print(f"  {tag:22s} action_loss={a:.5f}", flush=True)
        return a

    base = run("correct", lambda: banks)
    shuf = run("shuffled (roll +1)", lambda: {k: jnp.roll(banks[k], 1, axis=0) for k in keys})
    meanb = run("mean-bank", lambda: {k: jnp.broadcast_to(banks[k].mean(axis=0, keepdims=True),
                                                          banks[k].shape) for k in keys})
    zero = run("zeroed", lambda: None)

    print("\n=== VERDICT ===")
    d_sh, d_mn, d_z = shuf - base, meanb - base, zero - base
    print(f"  shuffled  delta = {d_sh:+.5f}  ({100*d_sh/max(base,1e-6):+.1f}%)")
    print(f"  mean-bank delta = {d_mn:+.5f}  ({100*d_mn/max(base,1e-6):+.1f}%)")
    print(f"  zeroed    delta = {d_z:+.5f}  ({100*d_z/max(base,1e-6):+.1f}%)")
    frac = d_sh / max(d_z, 1e-9)
    print(f"  shuffled/zeroed damage ratio = {frac:.3f}")
    if d_sh < 0.05 * max(d_z, 1e-9):
        print("  => banks act as a GENERIC/CONSTANT signal: wrong geometry costs almost nothing.")
    elif frac > 0.3:
        print("  => model reads PER-SAMPLE SPATIAL CONTENT: wrong geometry is nearly as bad as none.")
    else:
        print("  => partial: some per-sample use, but a large generic component.")
    print("SPECIFICITY CHECK DONE", flush=True)


if __name__ == "__main__":
    main()