| |
| """ORACLE-STATE CONTROL: does PNS-World admit a bounded sufficient state? |
| |
| The negative result says learned recurrent states never become causally useful. |
| A reviewer's remaining escape hatch is: "maybe the semantic target or the |
| output architecture is simply not learnable." This control closes it. |
| |
| We hand the model the simulator's *ground-truth bounded* semantic state - |
| the current (entity, attribute) -> value map and the (entity, relation) -> |
| entity map, as at most 120 typed slots, exactly the capacity Sigma has - plus |
| the question, and ask the SAME query-conditioned read + enum head to answer. |
| |
| S_oracle_t + e_t -> enum head -> answer |
| |
| If this scores high, the task admits a bounded sufficient state and the answer |
| heads can exploit one, so the learned models' failure is in CONSTRUCTING or |
| USING state, not in the task or the readout. If it scores low, the negative |
| result is confounded and must be reported as such. |
| |
| The slot CONTENTS are oracle; their embeddings are learned (a symbol must be |
| embedded to be read). Nothing about the future or the gold answer's identity |
| is supplied - only the world state that a perfect bounded summariser would |
| hold at t-. |
| """ |
| import argparse |
| import sys |
| import time |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src")) |
| sys.path.insert(0, str(Path(__file__).resolve().parent)) |
| from pns.common import atomic_write_json, eval_root |
| from pns.model.modules import enum_legal_mask |
| from pns.world.schema import ENUM_ID, ENUM_VOCAB, ET, Fam |
| from pns.world.sim import LifetimeGen |
|
|
| N_SLOT, N_ENT, N_ATTR, N_VAL = 120, 40, 24, len(ENUM_VOCAB) + 1 |
| ATTRS = sorted({a for a in |
| ["location", "team", "status", "phase", "priority", "destination", |
| "owner", "project", "vault"]}) |
| AID = {a: i for i, a in enumerate(ATTRS)} |
|
|
|
|
| def build(seed, length=256): |
| """Replay one lifetime; at every SEM question emit the oracle slot table |
| as it stands BEFORE that question, plus the query and gold.""" |
| g = LifetimeGen(seed, length) |
| g.plan() |
| events, red, _ = g.walk() |
| out = [] |
| |
| |
| live, rel = {}, {} |
| hist = {} |
| clock = 0 |
| for t, ev in enumerate(events): |
| clock += ev.dt |
| p = ev.payload |
| if ev.etype == ET.QUESTION and ev.family in (int(Fam.SEM_LATEST), |
| int(Fam.SEM_2HOP)): |
| tk = g.slots[t] |
| d_ = tk.data if tk is not None else {} |
| slots = [] |
| for (e, a), v in live.items(): |
| if a in AID and v in ENUM_ID: |
| slots.append((e % N_ENT, AID[a], ENUM_ID[v] + 1, 0)) |
| for (e, r), dst in rel.items(): |
| if r in AID: |
| slots.append((e % N_ENT, AID[r], 0, dst % N_ENT + 1)) |
| if ev.family == int(Fam.SEM_LATEST): |
| q_ent, q_rel, q_attr = d_.get("ent", 0) % N_ENT, 0, d_.get("attr") |
| else: |
| q_ent = d_.get("src", 0) % N_ENT |
| q_rel = AID.get(d_.get("rel", ""), 0) + 1 |
| q_attr = d_.get("attr") |
| if q_attr not in AID: |
| continue |
| lm = np.zeros(8, np.uint8) |
| for _e in ev.enum_legal: |
| lm[_e >> 3] |= 1 << (_e & 7) |
| out.append(dict(slots=slots[:N_SLOT], gold=ev.enum_gold, legal=lm, |
| fam=ev.family, q_ent=q_ent, q_rel=q_rel, |
| q_attr=AID[q_attr], delay=ev.meta.get("delay", -1))) |
| |
| if ev.etype == ET.OBSERVATION: |
| if p.get("rel"): |
| rel[(p["ent"], p["attr"])] = p["ent2"] |
| else: |
| k = (p["ent"], p["attr"]) |
| live[k] = p["value"] |
| hist.setdefault(k, []).append((clock, p.get("source", ""), p["value"])) |
| elif ev.etype == ET.CORRECTION: |
| k = (p["ent"], p["key"]) |
| if p["mode"] == "replace": |
| if not p.get("exact") and p.get("value") is not None: |
| live[k] = p["value"] |
| hist.setdefault(k, []).append((clock, "", p["value"])) |
| else: |
| h = hist.get(k, []) |
| tgt = None |
| if p.get("ref_kind") == "source": |
| for c, src, v in h: |
| if src == p.get("ref_source"): |
| tgt = v |
| else: |
| for c, src, v in h: |
| if c == p.get("ref_clock"): |
| tgt = v |
| if tgt is not None: |
| live[k] = tgt |
| hist[k] = [x for x in h if x[0] <= (p.get("ref_clock") or c)] |
| return out |
|
|
|
|
| class OracleReader(nn.Module): |
| """Query-conditioned attention read over oracle slots -> enum head. |
| Same shape of computation as PNSR's read path, same enum output head.""" |
|
|
| def __init__(self, d=256): |
| super().__init__() |
| self.ent = nn.Embedding(N_ENT + 1, d) |
| self.attr = nn.Embedding(N_ATTR + 1, d) |
| self.val = nn.Embedding(N_VAL + 1, d) |
| self.dst = nn.Embedding(N_ENT + 2, d) |
| self.qe = nn.Embedding(N_ENT + 1, d) |
| self.qa = nn.Embedding(N_ATTR + 1, d) |
| self.qr = nn.Embedding(N_ATTR + 2, d) |
| self.rd = nn.MultiheadAttention(d, 8, batch_first=True) |
| self.rd2 = nn.MultiheadAttention(d, 8, batch_first=True) |
| self.mlp = nn.Sequential(nn.LayerNorm(d), nn.Linear(d, d), nn.GELU(), |
| nn.Linear(d, d)) |
| self.enum = nn.Linear(d, len(ENUM_VOCAB)) |
|
|
| def forward(self, se, sa, sv, sd, mask, qe, qa, qr): |
| s = self.ent(se) + self.attr(sa) + self.val(sv) + self.dst(sd) |
| q = (self.qe(qe) + self.qa(qa) + self.qr(qr)).unsqueeze(1) |
| |
| r1, _ = self.rd(q, s, s, key_padding_mask=~mask, need_weights=False) |
| |
| q2 = r1 + self.qa(qa).unsqueeze(1) |
| r2, _ = self.rd2(q2, s, s, key_padding_mask=~mask, need_weights=False) |
| h = r2.squeeze(1) |
| return self.enum(h + self.mlp(h)) |
|
|
|
|
| def pack(rows, dev): |
| B = len(rows) |
| se = torch.zeros(B, N_SLOT, dtype=torch.long) |
| sa = torch.zeros(B, N_SLOT, dtype=torch.long) |
| sv = torch.zeros(B, N_SLOT, dtype=torch.long) |
| sd = torch.zeros(B, N_SLOT, dtype=torch.long) |
| mk = torch.zeros(B, N_SLOT, dtype=torch.bool) |
| for i, r in enumerate(rows): |
| for j, (e, a, v, dd) in enumerate(r["slots"][:N_SLOT]): |
| se[i, j], sa[i, j], sv[i, j], sd[i, j], mk[i, j] = e, a, v, dd, True |
| qe = torch.tensor([r["q_ent"] for r in rows]) |
| qa = torch.tensor([r["q_attr"] for r in rows]) |
| qr = torch.tensor([r["q_rel"] for r in rows]) |
| y = torch.tensor([r["gold"] for r in rows]) |
| lg = torch.tensor(np.stack([r["legal"] for r in rows])) |
| return [x.to(dev) for x in (se, sa, sv, sd, mk, qe, qa, qr, y, lg)] |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--train-lifetimes", type=int, default=3000) |
| ap.add_argument("--val-lifetimes", type=int, default=400) |
| ap.add_argument("--steps", type=int, default=3000) |
| args = ap.parse_args() |
| dev = "cuda" |
| print("building oracle rows ...", flush=True) |
| tr = [r for s in range(args.train_lifetimes) for r in build(s)] |
| va = [r for s in range(1_000_000, 1_000_000 + args.val_lifetimes) for r in build(s)] |
| print(f"train rows {len(tr)} val rows {len(va)}", flush=True) |
| m = OracleReader().to(dev) |
| opt = torch.optim.AdamW(m.parameters(), lr=3e-4, weight_decay=0.01) |
| rng = np.random.default_rng(0) |
| t0 = time.time() |
| for step in range(args.steps): |
| idx = rng.integers(0, len(tr), 256) |
| se, sa, sv, sd, mk, qe, qa, qr, y, lg = pack([tr[i] for i in idx], dev) |
| logits = m(se, sa, sv, sd, mk, qe, qa, qr).masked_fill(~enum_legal_mask(lg), -1e9) |
| loss = F.cross_entropy(logits, y) |
| opt.zero_grad(); loss.backward(); opt.step() |
| if (step + 1) % 500 == 0: |
| m.eval() |
| accs, per = [], {} |
| with torch.no_grad(): |
| for k in range(0, len(va), 512): |
| ch = va[k:k + 512] |
| se, sa, sv, sd, mk, qe, qa, qr, y, lg = pack(ch, dev) |
| p = m(se, sa, sv, sd, mk, qe, qa, qr).masked_fill(~enum_legal_mask(lg), -1e9).argmax(-1) |
| ok = (p == y).float().cpu().numpy() |
| accs.extend(ok) |
| for r, o in zip(ch, ok): |
| per.setdefault(Fam(r["fam"]).name, []).append(o) |
| m.train() |
| print(f"step {step+1} loss {float(loss):.3f} val {np.mean(accs):.4f} " |
| + " ".join(f"{k} {np.mean(v):.3f}" for k, v in sorted(per.items())) |
| + f" ({time.time()-t0:.0f}s)", flush=True) |
| rep = {"val_acc": float(np.mean(accs)), "n_val": len(va), |
| "per_family": {k: float(np.mean(v)) for k, v in per.items()}, |
| "n_train_rows": len(tr)} |
| atomic_write_json(eval_root() / "ORACLE_STATE.json", rep) |
| print(rep) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
