scripts/aggregate_fuse_paper_table.py

#!/usr/bin/env python3
from __future__ import annotations

import csv
import json
import math
import os
from pathlib import Path
from statistics import mean, stdev

ROOT = Path("/workspace/fcl-vla")
RESULT_ROOT = ROOT / "results/fuse_paper_table_v1"
METHODS = ("sequential", "er", "packnet", "fuse_m1000")
TRAIN_SEEDS = (10000, 20000, 30000)
EVAL_SEEDS = (10000, 20000, 30000, 40000)


def atomic_text(path: Path, value: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(value, encoding="utf-8")
    os.replace(temporary, path)


def wilson(successes: int, total: int, z: float = 1.959963984540054) -> tuple[float, float]:
    if total == 0:
        return float("nan"), float("nan")
    p = successes / total
    denominator = 1 + z * z / total
    center = (p + z * z / (2 * total)) / denominator
    radius = z * math.sqrt(p * (1 - p) / total + z * z / (4 * total * total)) / denominator
    return center - radius, center + radius


def load_cell(method: str, train_seed: int, eval_seed: int) -> dict | None:
    path = RESULT_ROOT / f"train_seed_{train_seed}" / method / f"eval_seed_{eval_seed}.json"
    if not path.exists():
        return None
    try:
        value = json.loads(path.read_text())
        row0 = value["rows"]["after_task0"]
        row1 = value["rows"]["after_task1"]
        assert value["protocol"] == "cfm_two_task_canary_v2_paired_task_seed"
        assert row1["eval_seed_base"] == eval_seed
        assert row1["n_eval_per_task"] == 10
        assert row0["n_eval_per_task"] == 10
        assert len(row1["success"]) == 2
        assert len(row0["success"]) == 1
    except Exception:
        return None
    return {
        "path": str(path),
        "run_dir": value["run_dir"],
        "task0_checkpoint_sha256": row0["checkpoint_sha256"],
        "task1_checkpoint_sha256": row1["checkpoint_sha256"],
        "task0_before_successes": round(row0["success"][0] * 10),
        "task0_after_successes": round(row1["success"][0] * 10),
        "task1_after_successes": round(row1["success"][1] * 10),
    }


def main() -> None:
    RESULT_ROOT.mkdir(parents=True, exist_ok=True)
    seed_rows = []
    provenance = []
    missing = []
    invalid = []
    complete_cells: dict[tuple[str, int], list[dict]] = {}
    for method in METHODS:
        for train_seed in TRAIN_SEEDS:
            cells = []
            for eval_seed in EVAL_SEEDS:
                cell = load_cell(method, train_seed, eval_seed)
                if cell is None:
                    missing.append({"method": method, "train_seed": train_seed, "eval_seed": eval_seed})
                else:
                    cells.append(cell)
                    provenance.append({"method": method, "train_seed": train_seed, "eval_seed": eval_seed, **cell})
            if len(cells) != len(EVAL_SEEDS):
                continue
            complete_cells[(method, train_seed)] = cells

    # FUSE must start from the exact formal Sequential task-0 checkpoint for
    # the same training seed. This also gates the provisional seed-10000 reuse.
    for train_seed in TRAIN_SEEDS:
        sequential_cells = complete_cells.get(("sequential", train_seed))
        fuse_cells = complete_cells.get(("fuse_m1000", train_seed))
        if sequential_cells and fuse_cells:
            sequential_hash = sequential_cells[0]["task0_checkpoint_sha256"]
            fuse_hash = fuse_cells[0]["task0_checkpoint_sha256"]
            if sequential_hash != fuse_hash:
                invalid.append({
                    "method": "fuse_m1000",
                    "train_seed": train_seed,
                    "reason": "task0_checkpoint_hash_mismatch_with_formal_sequential",
                    "sequential_sha256": sequential_hash,
                    "fuse_sha256": fuse_hash,
                })
                del complete_cells[("fuse_m1000", train_seed)]

    for method in METHODS:
        for train_seed in TRAIN_SEEDS:
            cells = complete_cells.get((method, train_seed))
            if cells is None:
                continue
            before = sum(cell["task0_before_successes"] for cell in cells)
            old = sum(cell["task0_after_successes"] for cell in cells)
            new = sum(cell["task1_after_successes"] for cell in cells)
            seed_rows.append({
                "method": method,
                "train_seed": train_seed,
                "task0_before_successes": before,
                "task0_before_total": 40,
                "old_successes": old,
                "old_total": 40,
                "new_successes": new,
                "new_total": 40,
                "task0_before_rate": before / 40,
                "old_rate": old / 40,
                "new_rate": new / 40,
                "final_average": (old + new) / 80,
                "forgetting": (before - old) / 40,
            })

    summaries = []
    for method in METHODS:
        rows = [row for row in seed_rows if row["method"] == method]
        if not rows:
            continue
        old_successes = sum(row["old_successes"] for row in rows)
        new_successes = sum(row["new_successes"] for row in rows)
        old_total = sum(row["old_total"] for row in rows)
        new_total = sum(row["new_total"] for row in rows)
        old_ci = wilson(old_successes, old_total)
        new_ci = wilson(new_successes, new_total)
        summary = {
            "method": method,
            "completed_training_seeds": len(rows),
            "training_seeds": [row["train_seed"] for row in rows],
            "old_pooled_successes": old_successes,
            "old_pooled_total": old_total,
            "new_pooled_successes": new_successes,
            "new_pooled_total": new_total,
            "old_mean": mean(row["old_rate"] for row in rows),
            "new_mean": mean(row["new_rate"] for row in rows),
            "final_average_mean": mean(row["final_average"] for row in rows),
            "forgetting_mean": mean(row["forgetting"] for row in rows),
            "old_sd": stdev(row["old_rate"] for row in rows) if len(rows) > 1 else None,
            "new_sd": stdev(row["new_rate"] for row in rows) if len(rows) > 1 else None,
            "final_average_sd": stdev(row["final_average"] for row in rows) if len(rows) > 1 else None,
            "forgetting_sd": stdev(row["forgetting"] for row in rows) if len(rows) > 1 else None,
            "old_pooled_wilson95": old_ci,
            "new_pooled_wilson95": new_ci,
        }
        summaries.append(summary)

    payload = {
        "schema": "fuse_paper_table_v1",
        "status": "COMPLETE" if not missing and not invalid else "PARTIAL",
        "training_seed_is_statistical_unit": True,
        "train_seeds": TRAIN_SEEDS,
        "eval_seed_bases": EVAL_SEEDS,
        "n_eval_per_cell": 10,
        "seed_rows": seed_rows,
        "method_summaries": summaries,
        "missing_cells": missing,
        "invalid_training_runs": invalid,
        "provenance": provenance,
    }
    atomic_text(RESULT_ROOT / "paper_table.json", json.dumps(payload, indent=2, sort_keys=True) + "\n")

    csv_path = RESULT_ROOT / "paper_table_seed_rows.csv"
    temporary = csv_path.with_suffix(csv_path.suffix + ".tmp")
    with temporary.open("w", newline="", encoding="utf-8") as stream:
        fields = list(seed_rows[0].keys()) if seed_rows else ["method", "train_seed"]
        writer = csv.DictWriter(stream, fieldnames=fields)
        writer.writeheader()
        writer.writerows(seed_rows)
    os.replace(temporary, csv_path)

    lines = [
        "# FUSE paper evidence table v1",
        "",
        f"Status: **{payload['status']}**. Training seed is the statistical unit; pooled rollout intervals are descriptive.",
        "",
        "| Method | Train seeds | Old retention | New plasticity | Final average | Forgetting |",
        "|---|---:|---:|---:|---:|---:|",
    ]
    for row in summaries:
        def fmt(name: str) -> str:
            value = row[name]
            sd = row[name.replace("_mean", "_sd")]
            return f"{100*value:.2f}%" if sd is None else f"{100*value:.2f} ± {100*sd:.2f}%"
        lines.append(
            f"| {row['method']} | {row['completed_training_seeds']}/3 | {fmt('old_mean')} | "
            f"{fmt('new_mean')} | {fmt('final_average_mean')} | {fmt('forgetting_mean')} |"
        )
    lines.extend(["", "## Per-training-seed rows", "", "| Method | Seed | Before T0 | Final old | Final new | Final avg | Forgetting |", "|---|---:|---:|---:|---:|---:|---:|"])
    for row in seed_rows:
        lines.append(
            f"| {row['method']} | {row['train_seed']} | {100*row['task0_before_rate']:.1f}% | "
            f"{100*row['old_rate']:.1f}% | {100*row['new_rate']:.1f}% | "
            f"{100*row['final_average']:.1f}% | {100*row['forgetting']:.1f}% |"
        )
    lines.extend(["", f"Missing atomic cells: {len(missing)}."])
    atomic_text(RESULT_ROOT / "paper_table.md", "\n".join(lines) + "\n")

    ledger = [
        "# Claim-evidence ledger",
        "",
        "| Claim | Evidence required | Current status | Risk |",
        "|---|---|---|---|",
        "| FUSE improves two-task retention over Sequential | Three full training seeds, paired rollout streams | " + ("measured" if payload["status"] == "COMPLETE" else "partial") + " | Training variance |",
        "| FUSE replaces part of raw ER | Equal nominal memory and matched training/evaluation | " + ("measured locally" if payload["status"] == "COMPLETE" else "partial") + " | Generated support is large and contains simulator frames/actions |",
        "| FUSE is federated and privacy preserving | Non-IID multi-client transfer, communication, leakage audit | missing | Central paper claim remains blocked |",
        "| FUSE is architecture independent | Same interface on BC-Transformer, Diffusion Policy, and VLA | missing | Only need has been shown on Diffusion |",
    ]
    atomic_text(RESULT_ROOT / "claim_evidence_ledger.md", "\n".join(ledger) + "\n")


if __name__ == "__main__":
    main()