workspace/scripts/build_action_scale_vector.py

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

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


PROJECT_ROOT = Path(__file__).resolve().parents[1]


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description=(
            "Build a reproducible per-dimension action scale vector from an "
            "action-bound audit. The default uses train split only, so deployment "
            "diagnostics do not fit action conventions on validation/test outcomes."
        )
    )
    parser.add_argument(
        "--audit",
        type=Path,
        default=Path("runs/action_bound_audit_rgb_refs/metrics.json"),
        help="Action-bound audit metrics.json.",
    )
    parser.add_argument(
        "--out-dir",
        type=Path,
        default=Path("runs/action_scale_vector_train_base_branch_max"),
    )
    parser.add_argument(
        "--splits",
        default="train",
        help="Comma-separated audit splits to use. Default is train only.",
    )
    parser.add_argument(
        "--source",
        choices=("action", "base_action", "base_branch"),
        default="base_branch",
        help="Per-dimension audit source used for max-to-unit scaling.",
    )
    parser.add_argument(
        "--floor",
        type=float,
        default=1.0e-6,
        help="Minimum allowed scale value.",
    )
    parser.add_argument(
        "--ceil",
        type=float,
        default=1.0,
        help="Maximum allowed scale value.",
    )
    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)

    if args.floor <= 0.0:
        parser.error("--floor must be positive")
    if args.ceil <= 0.0:
        parser.error("--ceil must be positive")
    if args.floor > args.ceil:
        parser.error("--floor must be <= --ceil")

    audit = json.loads(args.audit.read_text())
    requested_splits = [item.strip() for item in args.splits.split(",") if item.strip()]
    if not requested_splits:
        parser.error("--splits must name at least one split")

    rows_by_split = {str(row.get("split")): row for row in audit.get("rows", [])}
    missing = [split for split in requested_splits if split not in rows_by_split]
    if missing:
        raise SystemExit(f"missing split(s) in audit: {', '.join(missing)}")

    vectors: list[list[float]] = []
    for split in requested_splits:
        per_dim = rows_by_split[split].get("per_dim", {}).get(args.source, {})
        vector = per_dim.get("suggested_per_dim_scale_to_unit_max")
        if not vector:
            raise SystemExit(f"audit split {split!r} has no per-dim scale for {args.source!r}")
        vectors.append([float(value) for value in vector])

    width = len(vectors[0])
    if any(len(vector) != width for vector in vectors):
        raise SystemExit("requested split vectors have different widths")

    # Use the most conservative per-dimension scale across requested splits.
    scale = [
        min(float(args.ceil), max(float(args.floor), min(vector[dim] for vector in vectors)))
        for dim in range(width)
    ]

    out_dir = args.out_dir
    out_dir.mkdir(parents=True, exist_ok=True)
    payload: dict[str, Any] = {
        "report_type": "action_scale_vector",
        "schema_version": 1,
        "audit": str(args.audit),
        "audit_report_type": audit.get("report_type"),
        "chart_root": audit.get("chart_root"),
        "splits": requested_splits,
        "source": args.source,
        "fit_scope": "train_only" if requested_splits == ["train"] else "multi_split_diagnostic",
        "scale": scale,
        "scale_env": ",".join(f"{value:.12g}" for value in scale),
        "data_hashes": {split: audit.get("data_hashes", {}).get(split) for split in requested_splits},
        "split_hashes": {split: audit.get("split_hashes", {}).get(split) for split in requested_splits},
    }

    (out_dir / "vector.json").write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    (out_dir / "vector_env.txt").write_text(payload["scale_env"] + "\n")
    (out_dir / "config.yaml").write_text(
        "\n".join(
            [
                f"audit: {args.audit}",
                f"splits: {args.splits}",
                f"source: {args.source}",
                f"floor: {args.floor}",
                f"ceil: {args.ceil}",
            ]
        )
        + "\n"
    )
    (out_dir / "command.txt").write_text(
        "python scripts/build_action_scale_vector.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(payload["data_hashes"], sort_keys=True) + "\n")
    (out_dir / "split_hash.txt").write_text(json.dumps(payload["split_hashes"], sort_keys=True) + "\n")
    (out_dir / "train.log").write_text(
        "fit per-dimension action scale from action-bound audit rows only\n"
    )
    (out_dir / "eval.log").write_text("no eval; action convention artifact only\n")
    (out_dir / "table.tex").write_text(_table(payload) + "\n")
    _write_markdown_report(out_dir, payload, no_markdown_report=args.no_markdown_report)
    print(json.dumps({"out_dir": str(out_dir), "scale_env": payload["scale_env"]}, indent=2))
    return 0


def _table(payload: dict[str, Any]) -> str:
    values = payload["scale"]
    lines = [
        "% Auto-generated by scripts/build_action_scale_vector.py",
        "\\begin{tabular}{lrrrrrrr}",
        "\\toprule",
        "Source & d0 & d1 & d2 & d3 & d4 & d5 & d6 \\\\",
        "\\midrule",
        (
            f"{_latex_escape(str(payload['source']))} & "
            + " & ".join(f"{float(value):.4f}" for value in values)
            + " \\\\"
        ),
        "\\bottomrule",
        "\\end{tabular}",
    ]
    return "\n".join(lines)


def _report(payload: dict[str, Any]) -> str:
    return "\n".join(
        [
            "# Action Scale Vector",
            "",
            f"Audit: `{payload['audit']}`",
            f"Splits used: `{','.join(payload['splits'])}`",
            f"Source: `{payload['source']}`",
            f"Fit scope: `{payload['fit_scope']}`",
            "",
            "Scale vector:",
            "",
            f"`{payload['scale_env']}`",
            "",
            "This artifact defines an action-convention diagnostic only. It does "
            "not measure collision/contact safety and does not use validation/test "
            "outcomes when `fit_scope=train_only`.",
        ]
    )


def _write_markdown_report(
    out_dir: Path,
    payload: 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(payload) + "\n")


def _latex_escape(value: str) -> str:
    return value.replace("_", "\\_").replace("%", "\\%").replace("&", "\\&")


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