scripts/report_hpc_clean_results.py

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

import argparse
import csv
import glob
import json
import math
import re
import statistics
import sys
from collections import OrderedDict, defaultdict
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))

from dovla_cil.utils.io import ensure_dir, write_json  # noqa: E402

EVAL_FILENAMES = ("lattice_eval.json", "causalstress.json", "policy_rollout.json")
FALLBACK_FILENAMES = ("metrics.json",)
METRICS = (
    "pairwise_ranking_accuracy",
    "top1_action_selection",
    "selected_success_rate",
    "oracle_success_rate",
    "ndcg_at_k",
    "effect_prediction_mae",
    "selection_regret",
    "potential_edge_mae",
    "policy_rollout_success_rate",
    "policy_rollout_progress",
    "expert_success_rate",
    "policy_oracle_regret",
    "policy_expert_regret",
    "action_mse_to_best",
    "restore_max_error",
)
EXPECTED_BASELINES = (
    "cross_state_negatives",
    "expert_only_bc",
    "label_only_counterfactual",
    "random_negatives",
    "world_model_auxiliary",
    "no_effect_head",
)
UNCLEAN_MARKERS = (
    "pilot",
    "smoke",
    "maniskill_full_k16_n1000_seed0",
    "maniskill_multitask_full_k16_n500",
    "maniskill_scaling_fixed16k",
    "gxk_",
)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description=(
            "Create a contamination-aware summary of clean HPC DoVLA-CIL result roots. "
            "By default, paths containing pilot or known pre-success-contaminated markers are "
            "excluded and recorded in the manifest."
        )
    )
    parser.add_argument(
        "--inputs",
        nargs="+",
        required=True,
        help="Result directories, JSON files, or glob patterns to scan.",
    )
    parser.add_argument("--out", type=Path, required=True, help="Output report directory.")
    parser.add_argument("--name", default="clean_hpc_results", help="Report title.")
    parser.add_argument(
        "--allow-unclean",
        action="store_true",
        help="Include paths that match known pilot/contaminated markers.",
    )
    args = parser.parse_args(argv)

    summary = generate_clean_hpc_report(
        args.inputs,
        args.out,
        name=args.name,
        allow_unclean=args.allow_unclean,
    )
    print(f"report: {summary['report_md']}")
    print(f"aggregate_csv: {summary['aggregate_csv']}")
    print(f"detail_csv: {summary['detail_csv']}")
    print(f"runs: {summary['num_rows']} clean, {summary['num_excluded']} excluded")
    return 0


def generate_clean_hpc_report(
    inputs: list[str | Path],
    out_dir: str | Path,
    *,
    name: str = "clean_hpc_results",
    allow_unclean: bool = False,
) -> dict[str, Any]:
    output_dir = ensure_dir(out_dir)
    paths = discover_result_paths(inputs)
    clean_paths: list[Path] = []
    excluded: list[dict[str, str]] = []
    for path in paths:
        marker = unclean_marker(path)
        if marker and not allow_unclean:
            excluded.append({"path": str(path), "reason": marker})
            continue
        clean_paths.append(path)

    rows = [normalize_result_payload(path, read_json(path)) for path in clean_paths]
    rows = [row for row in rows if row]
    aggregates = aggregate_rows(rows)
    warnings = claim_warnings(aggregates)

    detail_csv = output_dir / "clean_result_rows.csv"
    aggregate_csv = output_dir / "clean_result_summary.csv"
    report_md = output_dir / "clean_result_summary.md"
    manifest_path = output_dir / "clean_result_manifest.json"
    excluded_path = output_dir / "excluded_unclean_paths.txt"

    write_csv(detail_csv, rows, detail_fieldnames(rows))
    write_csv(aggregate_csv, aggregates, aggregate_fieldnames(aggregates))
    report_md.write_text(
        render_markdown_report(
            name=name,
            rows=rows,
            aggregates=aggregates,
            warnings=warnings,
            excluded=excluded,
        ),
        encoding="utf-8",
    )
    excluded_text = "\n".join(f"{item['reason']}\t{item['path']}" for item in excluded)
    excluded_path.write_text(excluded_text + ("\n" if excluded else ""), encoding="utf-8")
    manifest = {
        "name": name,
        "inputs": [str(value) for value in inputs],
        "out_dir": str(output_dir),
        "num_rows": len(rows),
        "num_aggregates": len(aggregates),
        "num_excluded": len(excluded),
        "aggregate_csv": str(aggregate_csv),
        "detail_csv": str(detail_csv),
        "report_md": str(report_md),
        "excluded_paths": str(excluded_path),
        "warnings": warnings,
    }
    write_json(manifest, manifest_path)
    return manifest


def discover_result_paths(inputs: list[str | Path]) -> list[Path]:
    raw_paths: list[Path] = []
    for value in inputs:
        matches = [Path(match) for match in glob.glob(str(value))]
        raw_paths.extend(matches or [Path(value)])

    discovered: OrderedDict[str, Path] = OrderedDict()
    for path in raw_paths:
        if path.is_dir():
            evaluation_parents: set[Path] = set()
            for filename in EVAL_FILENAMES:
                for found in sorted(path.glob(f"**/{filename}")):
                    discovered[str(found)] = found
                    evaluation_parents.add(found.parent)
            for filename in FALLBACK_FILENAMES:
                for found in sorted(path.glob(f"**/{filename}")):
                    if found.parent not in evaluation_parents:
                        discovered[str(found)] = found
        elif path.exists():
            discovered[str(path)] = path
    return list(discovered.values())


def unclean_marker(path: Path) -> str:
    lowered = str(path).lower()
    for marker in UNCLEAN_MARKERS:
        if marker in lowered:
            return marker
    return ""


def normalize_result_payload(path: Path, payload: dict[str, Any]) -> dict[str, Any]:
    experiment = infer_experiment(path)
    baseline = infer_baseline(path, payload)
    row: dict[str, Any] = {
        "experiment": experiment,
        "evaluation_kind": evaluation_kind(path),
        "run_name": str(payload.get("run_name") or path.parent.name),
        "objective": str(payload.get("objective") or ""),
        "baseline": baseline,
        "observation_mode": str(payload.get("observation_mode") or ""),
        "backbone_type": str(payload.get("backbone_type") or ""),
        "training_k": first_number(payload.get("training_k"), payload.get("k"), payload.get("K")),
        "evaluation_k": first_number(
            payload.get("evaluation_k"), payload.get("k"), payload.get("K")
        ),
        "seed": first_number(payload.get("seed"), infer_seed(path)),
        "num_groups": first_number(payload.get("num_groups")),
        "num_records": first_number(payload.get("num_records")),
        "num_pairs": first_number(payload.get("num_pairs")),
        "dataset": str(payload.get("dataset") or ""),
        "source_path": str(path),
    }
    for metric in METRICS:
        row[metric] = first_number(payload.get(metric))
    return row


def infer_experiment(path: Path) -> str:
    text = str(path)
    if "maniskill_presuccess_scaling_fixed14k" in text:
        return "scaling_fixed14k_pick_common_eval"
    if "maniskill_presuccess_transfer_leave_stack/clip_actionfix" in text:
        return "transfer_leave_stack_rgb_clip_actionfix"
    if "maniskill_presuccess_transfer_leave_stack/state_actionfix" in text:
        return "transfer_leave_stack_state_actionfix"
    if "maniskill_presuccess_transfer_leave_stack" in text:
        return "transfer_leave_stack_state"
    if "maniskill_presuccess_six_task_clip_actionfix" in text:
        return "six_task_rgb_clip_actionfix"
    if "maniskill_presuccess_six_task_rgb_actionfix" in text:
        return "six_task_rgb_actionfix"
    if "maniskill_presuccess_six_task_actionfix" in text:
        return "six_task_state_actionfix"
    if "maniskill_presuccess_six_task_visual_fieldpref" in text:
        return "six_task_rgb_fieldpref"
    if "maniskill_presuccess_six_task_fieldpref" in text:
        return "six_task_state_fieldpref"
    if "maniskill_presuccess_six_task_visual_runs" in text:
        return "six_task_rgb"
    if "maniskill_presuccess_six_task_runs" in text:
        return "six_task_state"
    if "maniskill_presuccess_baseline_runs" in text:
        return "six_task_baseline"
    if "maniskill_presuccess_random_baseline_runs" in text:
        return "six_task_baseline"
    if "maniskill_presuccess_full_runs" in text:
        return "pick_state"
    return path.parent.parent.name if path.parent.parent != path.parent else path.parent.name


def evaluation_kind(path: Path) -> str:
    if path.name == "policy_rollout.json":
        return "policy_rollout"
    if path.name == "causalstress.json":
        return "causalstress"
    if path.name == "lattice_eval.json":
        return "lattice"
    return "metrics"


def infer_seed(path: Path) -> int | str:
    match = re.search(r"(?:^|[/_])seed[_-]?(\d+)(?:/|$)", str(path))
    return int(match.group(1)) if match else ""


def infer_baseline(path: Path, payload: dict[str, Any]) -> str:
    if payload.get("baseline"):
        return str(payload["baseline"])
    parts = set(path.parts)
    for name in (
        "cross_state_negatives",
        "expert_only_bc",
        "label_only_counterfactual",
        "no_effect_head",
        "no_rank_regret",
        "random_negatives",
        "world_model_auxiliary",
    ):
        if name in parts:
            return name
    return ""


def aggregate_rows(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
    grouped: dict[tuple[Any, ...], list[dict[str, Any]]] = defaultdict(list)
    for row in rows:
        key = (
            row.get("experiment", ""),
            row.get("evaluation_kind", ""),
            row.get("objective", ""),
            row.get("baseline", ""),
            row.get("observation_mode", ""),
            row.get("training_k", ""),
        )
        grouped[key].append(row)

    aggregates: list[dict[str, Any]] = []
    for key, group_rows in sorted(grouped.items(), key=aggregate_group_sort_key):
        experiment, eval_kind, objective, baseline, observation_mode, training_k = key
        aggregate: dict[str, Any] = {
            "experiment": experiment,
            "evaluation_kind": eval_kind,
            "objective": objective,
            "baseline": baseline,
            "observation_mode": observation_mode,
            "backbone_type": str(group_rows[0].get("backbone_type") or ""),
            "training_k": training_k,
            "n": len(group_rows),
            "seeds": ",".join(
                str(int(row["seed"])) for row in group_rows if is_number(row.get("seed"))
            ),
            "mean_num_groups": mean_value(row.get("num_groups") for row in group_rows),
        }
        for metric in METRICS:
            values = [float(row[metric]) for row in group_rows if is_number(row.get(metric))]
            aggregate[f"{metric}_mean"] = statistics.mean(values) if values else ""
            if len(values) > 1:
                aggregate[f"{metric}_std"] = statistics.pstdev(values)
            else:
                aggregate[f"{metric}_std"] = 0.0 if values else ""
        aggregates.append(aggregate)
    return aggregates


def claim_warnings(aggregates: list[dict[str, Any]]) -> list[str]:
    warnings: list[str] = []
    lattice_aggregates = [
        row for row in aggregates if row.get("evaluation_kind") in ("", "lattice")
    ]
    scaling = [
        row
        for row in lattice_aggregates
        if row.get("experiment") == "scaling_fixed14k_pick_common_eval"
        and is_number(row.get("training_k"))
        and is_number(row.get("pairwise_ranking_accuracy_mean"))
    ]
    if scaling:
        ordered = sorted(scaling, key=lambda row: float(row["training_k"]))
        if float(ordered[-1]["pairwise_ranking_accuracy_mean"]) <= float(
            ordered[0]["pairwise_ranking_accuracy_mean"]
        ):
            warnings.append(
                "Scaling ranking accuracy does not improve from smallest K to largest K."
            )
        beta = log_k_beta(ordered, "pairwise_ranking_accuracy_mean")
        if beta <= 0:
            warnings.append("Scaling beta_log_k for ranking accuracy is non-positive.")
    else:
        warnings.append("No clean scaling rows found.")

    state_rows = [
        row
        for row in lattice_aggregates
        if row.get("experiment") == "six_task_state" and row.get("baseline") == ""
    ]
    fieldpref_rows = [
        row
        for row in lattice_aggregates
        if row.get("experiment") == "six_task_state_fieldpref" and row.get("baseline") == ""
    ]
    actionfix_rows = [
        row
        for row in lattice_aggregates
        if row.get("experiment") == "six_task_state_actionfix" and row.get("baseline") == ""
    ]
    actionfix_lattice = best_matching(actionfix_rows, objective="lattice_field")
    fieldpref_lattice = best_matching(fieldpref_rows, objective="lattice_field")
    standard_lattice = best_matching(state_rows, objective="lattice_field")
    lattice = actionfix_lattice or fieldpref_lattice or standard_lattice
    if actionfix_lattice:
        lattice_label = "Action-vector-corrected IAF"
    elif fieldpref_lattice:
        lattice_label = "Field-preference IAF"
    else:
        lattice_label = "Six-task IAF"
    legacy = best_matching(state_rows, objective="legacy")
    if lattice and legacy:
        if float(lattice.get("selected_success_rate_mean") or 0.0) <= float(
            legacy.get("selected_success_rate_mean") or 0.0
        ):
            warnings.append(f"{lattice_label} selected success does not beat legacy.")
        if float(lattice.get("pairwise_ranking_accuracy_mean") or 0.0) <= float(
            legacy.get("pairwise_ranking_accuracy_mean") or 0.0
        ):
            warnings.append(f"{lattice_label} pairwise ranking does not beat legacy.")
    else:
        warnings.append("Missing six-task IAF or legacy aggregate.")

    cross = best_matching(lattice_aggregates, baseline="cross_state_negatives")
    label = best_matching(lattice_aggregates, baseline="label_only_counterfactual")
    if cross and lattice:
        if float(lattice.get("pairwise_ranking_accuracy_mean") or 0.0) <= float(
            cross.get("pairwise_ranking_accuracy_mean") or 0.0
        ):
            warnings.append("Same-state IAF ranking does not beat cross-state baseline.")
    if label and lattice:
        if float(lattice.get("pairwise_ranking_accuracy_mean") or 0.0) <= float(
            label.get("pairwise_ranking_accuracy_mean") or 0.0
        ):
            warnings.append("Same-state IAF ranking does not beat label-only baseline.")
    present_baselines = {
        str(row.get("baseline")) for row in lattice_aggregates if row.get("baseline")
    }
    for baseline in EXPECTED_BASELINES:
        if baseline not in present_baselines:
            warnings.append(f"Missing expected baseline aggregate: {baseline}.")

    transfer_rows = [
        row
        for row in lattice_aggregates
        if str(row.get("experiment", "")).startswith("transfer_leave_stack")
        and is_number(row.get("selected_success_rate_mean"))
    ]
    if transfer_rows and max(
        float(row["selected_success_rate_mean"]) for row in transfer_rows
    ) < 0.10:
        warnings.append(
            "Held-out Stack selected success is below 10%; do not claim broad OOD task transfer."
        )
    return warnings


def render_markdown_report(
    *,
    name: str,
    rows: list[dict[str, Any]],
    aggregates: list[dict[str, Any]],
    warnings: list[str],
    excluded: list[dict[str, str]],
) -> str:
    fields = [
        "experiment",
        "evaluation_kind",
        "objective",
        "baseline",
        "observation_mode",
        "backbone_type",
        "training_k",
        "n",
        "pairwise_ranking_accuracy_mean",
        "top1_action_selection_mean",
        "selected_success_rate_mean",
        "oracle_success_rate_mean",
        "ndcg_at_k_mean",
        "effect_prediction_mae_mean",
        "selection_regret_mean",
        "policy_rollout_success_rate_mean",
        "policy_rollout_progress_mean",
        "expert_success_rate_mean",
        "policy_oracle_regret_mean",
        "action_mse_to_best_mean",
    ]
    lines = [
        f"# {name}",
        "",
        "## Scope",
        "",
        f"- clean result files: {len(rows)}",
        f"- aggregate rows: {len(aggregates)}",
        f"- excluded unclean files: {len(excluded)}",
        "",
        "## Aggregate Results",
        "",
        markdown_table(aggregates, fields),
        "",
        "## Claim Warnings",
        "",
    ]
    if warnings:
        lines.extend(f"- {warning}" for warning in warnings)
    else:
        lines.append("- No automatic claim warnings.")
    if excluded:
        lines.extend(["", "## Excluded Paths", ""])
        for item in excluded[:40]:
            lines.append(f"- `{item['reason']}`: `{item['path']}`")
        if len(excluded) > 40:
            lines.append(f"- ... {len(excluded) - 40} more")
    return "\n".join(lines).rstrip() + "\n"


def markdown_table(rows: list[dict[str, Any]], fieldnames: list[str]) -> str:
    lines = ["| " + " | ".join(fieldnames) + " |"]
    lines.append("| " + " | ".join("---" for _ in fieldnames) + " |")
    if not rows:
        lines.append("| " + " | ".join("_n/a_" for _ in fieldnames) + " |")
    for row in rows:
        cells = " | ".join(format_cell(row.get(field, "")) for field in fieldnames)
        lines.append(f"| {cells} |")
    return "\n".join(lines)


def write_csv(path: Path, rows: list[dict[str, Any]], fieldnames: list[str]) -> None:
    ensure_dir(path.parent)
    with path.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=fieldnames)
        writer.writeheader()
        for row in rows:
            writer.writerow({field: format_cell(row.get(field, "")) for field in fieldnames})


def aggregate_group_sort_key(item: tuple[tuple[Any, ...], list[dict[str, Any]]]) -> tuple[Any, ...]:
    key, _group_rows = item
    experiment, eval_kind, objective, baseline, observation_mode, training_k = key
    return (
        str(experiment),
        str(eval_kind),
        str(objective),
        str(baseline),
        str(observation_mode),
        float(training_k) if is_number(training_k) else math.inf,
        str(training_k),
    )


def detail_fieldnames(rows: list[dict[str, Any]]) -> list[str]:
    defaults = [
        "experiment",
        "evaluation_kind",
        "run_name",
        "objective",
        "baseline",
        "observation_mode",
        "backbone_type",
        "training_k",
        "evaluation_k",
        "seed",
        "num_groups",
        "num_records",
        "num_pairs",
        "dataset",
        "source_path",
    ]
    return defaults + [metric for metric in METRICS if any(metric in row for row in rows)]


def aggregate_fieldnames(rows: list[dict[str, Any]]) -> list[str]:
    defaults = [
        "experiment",
        "evaluation_kind",
        "objective",
        "baseline",
        "observation_mode",
        "backbone_type",
        "training_k",
        "n",
        "seeds",
        "mean_num_groups",
    ]
    metric_fields = [
        f"{metric}_{suffix}"
        for metric in METRICS
        for suffix in ("mean", "std")
        if any(f"{metric}_{suffix}" in row for row in rows)
    ]
    return defaults + metric_fields


def read_json(path: Path) -> dict[str, Any]:
    with path.open("r", encoding="utf-8") as handle:
        payload = json.load(handle)
    if not isinstance(payload, dict):
        raise ValueError(f"Expected JSON object in {path}")
    return payload


def first_number(*values: Any) -> float | str:
    for value in values:
        if is_number(value):
            return float(value)
    return ""


def mean_value(values: Any) -> float | str:
    numbers = [float(value) for value in values if is_number(value)]
    return statistics.mean(numbers) if numbers else ""


def best_matching(rows: list[dict[str, Any]], **criteria: str) -> dict[str, Any] | None:
    candidates = [
        row
        for row in rows
        if all(str(row.get(key, "")) == str(value) for key, value in criteria.items())
    ]
    if not candidates:
        return None
    return max(candidates, key=lambda row: float(row.get("pairwise_ranking_accuracy_mean") or 0.0))


def log_k_beta(rows: list[dict[str, Any]], metric: str) -> float:
    pairs = [
        (math.log(float(row["training_k"])), float(row[metric]))
        for row in rows
        if is_number(row.get("training_k"))
        and is_number(row.get(metric))
        and float(row["training_k"]) > 0
    ]
    if len(pairs) < 2:
        return 0.0
    mean_x = statistics.mean(x for x, _ in pairs)
    mean_y = statistics.mean(y for _, y in pairs)
    denom = sum((x - mean_x) ** 2 for x, _ in pairs)
    if denom == 0:
        return 0.0
    return sum((x - mean_x) * (y - mean_y) for x, y in pairs) / denom


def is_number(value: Any) -> bool:
    return (
        isinstance(value, int | float)
        and not isinstance(value, bool)
        and math.isfinite(float(value))
    )


def format_cell(value: Any) -> str:
    if value == "":
        return ""
    if is_number(value):
        number = float(value)
        if number.is_integer() and abs(number) >= 10:
            return str(int(number))
        return f"{number:.6g}"
    return str(value)


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