#!/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())