scripts/experiments/transfer_contract.py

"""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())