"""Run the learning contract on a saved PlasticDynamics checkpoint, one report per baseline mode. uv run python -m scripts.experiments.transfer_contract --model-id phys_mps_3k --out artifacts/experiments/contract-phys_mps_3k Modes: ``frozen`` (no lasting update; fast weights only), ``continued`` (plain gradient steps on the stream), ``in_context`` (the stream prepended at measurement time). The output directory gets one JSON per mode, a ``README.md`` summary table, and ``manifest.json`` with the checkpoint digest and the execution commit. Spec: docs/superpowers/specs/2026-09-23-mechanism-testbed-and-contract.md """ from __future__ import annotations import argparse import hashlib import json import os import subprocess import time from dataclasses import asdict from typing import Any, Callable import torch from plastic.eval.contract import ContractSpec, DynamicsLearner, RetrievalLearner, run_contract from plastic.store import ArtifactStore MODES = ("frozen", "continued", "in_context", "retrieval") def execution_commit(root: str) -> str | None: try: return subprocess.check_output(["git", "-C", root, "rev-parse", "HEAD"], text=True).strip() except Exception: return None def file_digest(path: str) -> str | None: if not os.path.exists(path): return None h = hashlib.sha256() with open(path, "rb") as f: for chunk in iter(lambda: f.read(1 << 20), b""): h.update(chunk) return h.hexdigest() def run_modes( model_factory: Callable[[], torch.nn.Module], *, modes: tuple[str, ...], spec: ContractSpec, seed: int, lr: float, steps: int, device: torch.device, knn: int = 8, ) -> dict[str, dict[str, Any]]: """A fresh model per mode, so no mode's lasting update leaks into another. The retrieval mode uses no model at all.""" reports: dict[str, dict[str, Any]] = {} for mode in modes: if mode == "retrieval": learner: Any = RetrievalLearner(k=knn) reports[mode] = run_contract(learner, spec, seed=seed) reports[mode]["learner"] = {"k": knn, "stored_transitions": learner.stored_transitions()} else: model = model_factory().to(device) learner = DynamicsLearner(model, mode=mode, lr=lr, steps=steps, device=device) reports[mode] = run_contract(learner, spec, seed=seed) reports[mode]["learner"] = {"lr": lr, "steps": steps} reports[mode]["mode"] = mode return reports def _fmt(x: Any) -> str: if x is None: return "n/a" if isinstance(x, bool): return "yes" if x else "no" if isinstance(x, float): return f"{x:+.4f}" if abs(x) < 100 else f"{x:.3g}" return str(x) def summary_table(reports: dict[str, dict[str, Any]]) -> str: """Markdown: one row per mode. MSE deltas are after minus before; negative is improvement.""" policies = list(next(iter(reports.values()))["transfer"].keys()) head = ["mode"] + [f"transfer Δ ({p})" for p in policies] + ["forgetting Δ", "poison harm (vs clean)", "poison harm (vs start)", "corr. residual", "revert ok", "accepted-good", "refused-bad", "tokens consumed", "tokens measured"] lines = ["| " + " | ".join(head) + " |", "|" + "---|" * len(head)] for mode, r in reports.items(): acc = r["acceptance"] row = [mode] row += [_fmt(r["transfer"][p]["delta_mse"]) for p in policies] row += [ _fmt(r["forgetting"]["delta_mse"]), _fmt(r["correction"]["harm"]), _fmt(r["correction"]["harm_vs_before"]), _fmt(r["correction"]["residual"]), _fmt(r["revert"]["ok"]), _fmt(acc["accepted_good"]) + f" (n={acc['n_good']})", _fmt(acc["refused_bad"]) + f" (n={acc['n_bad']})", str(r["compute"]["tokens_consumed"]), str(r["compute"]["tokens_measured"]), ] lines.append("| " + " | ".join(row) + " |") return "\n".join(lines) def before_table(reports: dict[str, dict[str, Any]]) -> str: """The before-stream measurements, which are the same for every mode of one checkpoint.""" r = next(iter(reports.values())) lines = ["| measurement | with adaptation | writes disabled | elements |", "|---|---|---|---|"] for p, row in r["transfer"].items(): lines.append(f"| held-out combos, {p} | {row['before']['adapt']:.4f} | {row['before']['no_adapt']:.4f} | {row['elements']} |") f = r["forgetting"]["before"] lines.append(f"| training distribution | {f['adapt']:.4f} | {f['no_adapt']:.4f} | {r['forgetting']['elements']} |") s = r["speed"]["before"] lines.append(f"| adaptation speed (held-out; adapting error as a fraction of writes-disabled error, by step) | mean {s['area']:.3f}; below one half at step {s['steps_to_half']} | | {s['episodes']} episodes |") return "\n".join(lines) def write_outputs(out: str, reports: dict[str, dict[str, Any]], manifest: dict[str, Any]) -> None: os.makedirs(out, exist_ok=True) for mode, r in reports.items(): with open(os.path.join(out, f"{mode}.json"), "w", encoding="utf-8") as f: json.dump(r, f, indent=1) with open(os.path.join(out, "manifest.json"), "w", encoding="utf-8") as f: json.dump(manifest, f, indent=1) body = [ f"# Learning contract on `{manifest['model_id']}`", "", f"Checkpoint digest `{manifest['checkpoint_digest']}`, execution commit `{manifest['execution_commit']}`, " f"device {manifest['device']}, contract {manifest['contract_version']}, seed {manifest['seed']}.", "", "## Before any stream (identical for every mode)", "", before_table(reports), "", "## After the stream, per mode", "", "MSE deltas are after minus before on identical inputs; negative is improvement. " "`poison harm (vs clean)` is transfer MSE after the poisoned stream minus after the clean stream " "(damage plus the forgone clean gain); `poison harm (vs start)` is minus the poison arm's own " "pre-stream start (damage alone); `corr. residual` is after the corrective stream minus after clean. " "Continued training uses Adam. Everything-in-context prepends the stream to the model's own recurrent carry, " "which is bounded by its forget gate and decay horizon, unlike a transformer's context window; the retrieval " "mode is the model-free lookup baseline (mean delta of the nearest stored transitions). Acceptance is a pair of rates; " "n/a means no decision was recorded, not zero.", "", summary_table(reports), "", "Every measurement is taken from a fresh state with the stream removed, one episode per row. " f"Adaptation window: {manifest.get('adaptation_window')}. " "Spec: docs/superpowers/specs/2026-09-23-mechanism-testbed-and-contract.md", "", ] with open(os.path.join(out, "README.md"), "w", encoding="utf-8") as f: f.write("\n".join(body)) def main(argv: list[str] | None = None) -> int: ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) ap.add_argument("--artifacts-root", default="artifacts") ap.add_argument("--model-id", required=True) ap.add_argument("--out", required=True) ap.add_argument("--modes", nargs="+", default=list(MODES), choices=MODES) ap.add_argument("--device", default="cpu") ap.add_argument("--seed", type=int, default=0) ap.add_argument("--lr", type=float, default=1e-3) ap.add_argument("--steps", type=int, default=10) ap.add_argument("--knn", type=int, default=8, help="neighbours for the retrieval baseline") ap.add_argument("--seq-len", type=int, default=64, help="length of the one episode each scored row holds") ap.add_argument("--probe-steps", type=int, default=None, help="steps of the adaptation curve to report (default: the whole episode)") ap.add_argument("--eval-batch", type=int, default=8) ap.add_argument("--stream-episodes", type=int, default=16) ap.add_argument("--n-heldout", type=int, default=5) ap.add_argument("--split-seed", type=int, default=0) ap.add_argument("--poison-bias", type=float, default=0.5) args = ap.parse_args(argv) device = torch.device(args.device) store = ArtifactStore(args.artifacts_root) cfg, _, record = store.load_checkpoint(args.model_id, device=device) if cfg.domain != "physics": raise SystemExit("the contract runner needs a physics-domain checkpoint") spec = ContractSpec( seq_len=args.seq_len, probe_steps=args.probe_steps, eval_batch=args.eval_batch, stream_episodes=args.stream_episodes, n_heldout=args.n_heldout, split_seed=args.split_seed, poison_bias=args.poison_bias, ) def factory() -> torch.nn.Module: _, model, _ = store.load_checkpoint(args.model_id, device=device) return model t0 = time.time() torch.manual_seed(args.seed) reports = run_modes(factory, modes=tuple(args.modes), spec=spec, seed=args.seed, lr=args.lr, steps=args.steps, device=device, knn=args.knn) root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) manifest = { "model_id": args.model_id, "checkpoint_digest": file_digest(os.path.join(store.model_dir(args.model_id), "checkpoint.pt")), "model_record": {"step": record.get("step"), "extra": record.get("extra")}, "model_signature": store.model_signature(args.model_id), "execution_commit": execution_commit(root), "device": str(device), "contract_version": next(iter(reports.values()))["contract_version"], "adaptation_window": next(iter(reports.values()))["adaptation_window"], "split_id": next(iter(reports.values()))["split"]["id"], "spec": asdict(spec), "seed": args.seed, "modes": list(args.modes), "learner": {"lr": args.lr, "steps": args.steps}, "wall_clock_s": time.time() - t0, } write_outputs(args.out, reports, manifest) print(summary_table(reports)) print(f"[contract] wrote {args.out} in {manifest['wall_clock_s']:.1f}s") return 0 if __name__ == "__main__": raise SystemExit(main())