workspace/scripts/check_tangent_reconstruction.py

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

import argparse
import json
import math
import subprocess
import sys
from pathlib import Path
from typing import Any

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

import numpy as np  # noqa: E402

from scripts.export_cil_charts import _spline_tangent_code  # noqa: E402


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="Verify exported spline tangent codes are deterministic summaries of delta_action."
    )
    parser.add_argument(
        "--chart-index",
        type=Path,
        action="append",
        default=[],
        help="Chart index to check. May be repeated.",
    )
    parser.add_argument("--chart-root", type=Path, default=Path("data/cil_charts"))
    parser.add_argument("--out-dir", type=Path, default=Path("runs/tangent_reconstruction"))
    parser.add_argument(
        "--no-markdown-report",
        action="store_true",
        help="Do not write report.md; persistent prose is consolidated in README.md.",
    )
    args = parser.parse_args(argv)

    indexes = args.chart_index or sorted(args.chart_root.glob("*/index.json"))
    rows = []
    for index_path in indexes:
        rows.extend(_check_index(index_path))
    if not rows:
        raise SystemExit("no tangent rows checked")

    max_abs = max(row["max_abs_error"] for row in rows)
    rmse_values = [row["rmse"] for row in rows]
    summary = {
        "report_type": "tangent_reconstruction_check",
        "num_rows": len(rows),
        "num_failed": sum(1 for row in rows if row["max_abs_error"] > 1.0e-6),
        "max_abs_error": max_abs,
        "mean_rmse": sum(rmse_values) / len(rmse_values),
        "data_hash": _hash_manifest(indexes, "content_hash"),
        "split_hash": _hash_manifest(indexes, "split_hash"),
        "note": "The 21D spline_tangent_code is a deterministic keyframe summary, not a lossless action decoder.",
    }
    out_dir = args.out_dir
    out_dir.mkdir(parents=True, exist_ok=True)
    _write_provenance(out_dir, args, indexes)
    (out_dir / "metrics.json").write_text(json.dumps(summary | {"rows": rows[:100]}, indent=2, sort_keys=True) + "\n")
    (out_dir / "metrics_by_task.json").write_text(json.dumps(_by(rows, "task_id"), indent=2, sort_keys=True) + "\n")
    (out_dir / "metrics_by_seed.json").write_text(json.dumps(_by(rows, "seed"), indent=2, sort_keys=True) + "\n")
    (out_dir / "train.log").write_text("no training; checked exported tangent summaries\n")
    (out_dir / "eval.log").write_text(f"checked {len(rows)} rows across {len(indexes)} indexes\n")
    (out_dir / "table.tex").write_text(_table(summary) + "\n")
    _write_markdown_report(out_dir, summary, no_markdown_report=args.no_markdown_report)
    print(json.dumps({"out_dir": str(out_dir), **summary}, indent=2))
    return 0


def _check_index(index_path: Path) -> list[dict[str, Any]]:
    index = json.loads(index_path.read_text())
    rows = []
    for shard in index.get("shards", []):
        shard_path = index_path.parent / shard["path"]
        with np.load(shard_path, allow_pickle=False) as data:
            delta_actions = data["delta_action"]
            action_shapes = data["action_shape"]
            codes = data["spline_tangent_code"]
            chart_ids = data["chart_id"]
            task_ids = data["task_id"]
            seeds = data["seed"]
            for row in range(delta_actions.shape[0]):
                horizon, action_dim = [int(value) for value in action_shapes[row]]
                delta_action = delta_actions[row, : horizon * action_dim].reshape(horizon, action_dim)
                recomputed = np.asarray(_spline_tangent_code(delta_action.tolist()), dtype=np.float32)
                error = recomputed - codes[row]
                rows.append(
                    {
                        "split": index.get("split"),
                        "chart_id": str(chart_ids[row]),
                        "task_id": str(task_ids[row]),
                        "seed": str(seeds[row]),
                        "max_abs_error": float(np.max(np.abs(error))),
                        "rmse": float(math.sqrt(float(np.mean(error * error)))),
                    }
                )
    return rows


def _by(rows: list[dict[str, Any]], key: str) -> dict[str, dict[str, float]]:
    grouped: dict[str, list[dict[str, Any]]] = {}
    for row in rows:
        grouped.setdefault(str(row[key]), []).append(row)
    return {
        group: {
            "num_rows": len(items),
            "max_abs_error": max(item["max_abs_error"] for item in items),
            "mean_rmse": sum(item["rmse"] for item in items) / len(items),
        }
        for group, items in sorted(grouped.items())
    }


def _write_provenance(out_dir: Path, args: argparse.Namespace, indexes: list[Path]) -> None:
    payload = vars(args) | {"chart_index": [str(path) for path in indexes]}
    (out_dir / "config.yaml").write_text(
        "\n".join(f"{key}: {value}" for key, value in sorted(payload.items())) + "\n"
    )
    (out_dir / "command.txt").write_text("python scripts/check_tangent_reconstruction.py " + " ".join(sys.argv[1:]) + "\n")
    (out_dir / "git_hash.txt").write_text(_run(["git", "rev-parse", "HEAD"]) + "\n")
    (out_dir / "data_hash.txt").write_text(
        json.dumps(_hash_manifest(indexes, "content_hash"), indent=2, sort_keys=True) + "\n"
    )
    (out_dir / "split_hash.txt").write_text(
        json.dumps(_hash_manifest(indexes, "split_hash"), indent=2, sort_keys=True) + "\n"
    )


def _hash_manifest(indexes: list[Path], field: str) -> dict[str, str]:
    manifest: dict[str, str] = {}
    for path in indexes:
        payload = json.loads(path.read_text())
        split = str(payload.get("split", path.parent.name))
        manifest[split] = str(payload.get(field, ""))
    return manifest


def _table(summary: dict[str, Any]) -> str:
    return "\n".join(
        [
            "% Auto-generated by scripts/check_tangent_reconstruction.py",
            "\\begin{tabular}{lrrr}",
            "\\toprule",
            "Rows & Failed & Max abs error & Mean RMSE \\\\",
            "\\midrule",
            (
                f"{summary['num_rows']} & {summary['num_failed']} & "
                f"{summary['max_abs_error']:.2e} & {summary['mean_rmse']:.2e} \\\\"
            ),
            "\\bottomrule",
            "\\end{tabular}",
        ]
    )


def _report(summary: dict[str, Any]) -> str:
    return "\n".join(
        [
            "# Tangent Reconstruction Check",
            "",
            f"Rows checked: `{summary['num_rows']}`",
            f"Rows failed: `{summary['num_failed']}`",
            f"Max abs error: `{summary['max_abs_error']:.8e}`",
            f"Mean RMSE: `{summary['mean_rmse']:.8e}`",
            "",
            summary["note"],
        ]
    )


def _write_markdown_report(
    out_dir: Path,
    summary: dict[str, Any],
    *,
    no_markdown_report: bool,
) -> None:
    report_path = out_dir / "report.md"
    if no_markdown_report:
        report_path.unlink(missing_ok=True)
        return
    report_path.write_text(_report(summary) + "\n")


def _run(command: list[str]) -> str:
    try:
        return subprocess.check_output(command, cwd=PROJECT_ROOT, text=True).strip()
    except (subprocess.CalledProcessError, FileNotFoundError):
        return ""


if __name__ == "__main__":
    raise SystemExit(main())