"""Audit a low-variance frozen residual against an existing gradient artifact.""" from __future__ import annotations import argparse import hashlib import json import time from pathlib import Path import numpy as np import torch from omegaconf import OmegaConf from scipy.stats import kendalltau, spearmanr from starVLA.dataloader import build_dataloader from starVLA.model.framework.base_framework import build_framework from starVLA.model.framework.share_tools import apply_config_compat from starVLA.scoring.starvla_likelihood import ( score_starvla_flow_matching_residual, ) from starVLA.scoring.utility_audit import ( alignment_audit, leave_one_out_robustness, ) TRACE_IDENTITY_FIELDS = ( "manifest_sha256", "epoch", "global_draw_id", "dataset_index", "dataset_name", "trajectory_index", "trajectory_id", "anchor_index", "base_index", ) def _sha256_file(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as source: for chunk in iter(lambda: source.read(8 * 1024 * 1024), b""): digest.update(chunk) return digest.hexdigest() def _trace_identity(trace: dict) -> dict: return {field: trace.get(field) for field in TRACE_IDENTITY_FIELDS} def _load_gradient_artifact(path: Path) -> tuple[dict, list[dict]]: payload = json.loads(path.read_text(encoding="utf-8")) anchors = payload.get("anchors") if not isinstance(anchors, list) or len(anchors) < 4: raise ValueError("gradient artifact must contain at least four anchors") required = ("gradient_rms", "mean_gradient_loss", "sampling_trace") for index, anchor in enumerate(anchors): if any(field not in anchor for field in required): raise ValueError( f"gradient artifact anchor {index} is missing required fields" ) if not ( np.isfinite(anchor["gradient_rms"]) and anchor["gradient_rms"] > 0.0 and np.isfinite(anchor["mean_gradient_loss"]) ): raise ValueError("gradient artifact contains invalid measurements") metadata = { "path": str(path.resolve()), "sha256": _sha256_file(path), "checkpoint_sha256": payload.get("checkpoint_sha256"), "gradient_runtime": payload.get("gradient_runtime"), "gradient_seeds": payload.get("seed_separation", {}).get( "gradient_seeds" ), } return metadata, anchors def _load_aligned_examples(cfg, anchors: list[dict]) -> list[dict]: dataloader = build_dataloader(cfg, dataset_py=cfg.datasets.vla_data.dataset_py) examples = [] for batch in dataloader: examples.extend(batch) if len(examples) >= len(anchors): break examples = examples[: len(anchors)] if len(examples) != len(anchors): raise RuntimeError("dataloader ended before all audit anchors were recovered") for index, (example, anchor) in enumerate(zip(examples, anchors)): observed = _trace_identity(example.get("sampling_trace") or {}) expected = _trace_identity(anchor["sampling_trace"]) if observed != expected: raise RuntimeError( f"sample trace mismatch at position {index}: " f"observed={observed}, expected={expected}" ) return examples def _repeatability(estimator_matrix: np.ndarray) -> dict: estimator_count, sample_count = estimator_matrix.shape if estimator_count < 4 or estimator_count % 2: raise ValueError("estimator matrix must have an even number >= 4 of rows") half = estimator_count // 2 first_half = estimator_matrix[:half].mean(axis=0) second_half = estimator_matrix[half:].mean(axis=0) correlation = spearmanr(first_half, second_half) rank_concordance = kendalltau(first_half, second_half) top_count = max(1, sample_count // 4) first_top = set(np.argsort(first_half)[-top_count:]) second_top = set(np.argsort(second_half)[-top_count:]) anchor_means = estimator_matrix.mean(axis=0) anchor_cv = estimator_matrix.std(axis=0, ddof=1) / anchor_means return { "spearman": float(correlation.statistic), "spearman_pvalue_asymptotic": float(correlation.pvalue), "kendall": float(rank_concordance.statistic), "kendall_pvalue_asymptotic": float(rank_concordance.pvalue), "mean_absolute_difference": float( np.mean(np.abs(first_half - second_half)) ), "max_absolute_difference": float( np.max(np.abs(first_half - second_half)) ), "top25_overlap": len(first_top & second_top) / top_count, "mean_anchor_cv": float(np.mean(anchor_cv)), "median_anchor_cv": float(np.median(anchor_cv)), } def _gate_decision( *, sample_count: int, repeatability: dict, alignment: dict, robustness: dict, ) -> dict: checks = { "sample_count_at_least_64": sample_count >= 64, "score_repeatability_spearman_at_least_0_90": ( repeatability["spearman"] >= 0.90 ), "score_gradient_spearman_at_least_0_50": ( alignment["spearman_residual_vs_gradient_rms"] >= 0.50 ), "score_gradient_permutation_p_at_most_0_01": ( alignment["random_permutation_control"][ "one_sided_pvalue_rho_at_least_observed" ] <= 0.01 ), "top25_overlap_at_least_0_50": alignment["top25_overlap"] >= 0.50, "second_moment_ratio_at_most_0_98": ( alignment["proposal"]["information_to_baseline_ratio"] <= 0.98 ), "second_moment_permutation_p_at_most_0_01": ( alignment["random_permutation_control"][ "one_sided_pvalue_second_moment_at_most_observed" ] <= 0.01 ), "leave_one_out_spearman_all_positive": robustness[ "all_spearman_positive" ], "leave_one_out_second_moment_all_below_one": robustness[ "all_second_moment_ratios_below_one" ], } return { "checks": checks, "pass": all(checks.values()), "eligible_for_500_step_ablation": all(checks.values()), } def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--config", type=Path, required=True) parser.add_argument("--checkpoint", type=Path, required=True) parser.add_argument("--gradient-artifact", type=Path, required=True) parser.add_argument("--output", type=Path, required=True) parser.add_argument("--estimator-seeds", type=int, nargs="+", default=[9300, 9301, 9302, 9303]) parser.add_argument("--samples-per-estimator", type=int, default=8) parser.add_argument("--alpha", type=float, default=0.2) parser.add_argument("--percentile-epsilon", type=float, default=0.1) parser.add_argument("--permutations", type=int, default=10_000) parser.add_argument("--permutation-seed", type=int, default=9100) args = parser.parse_args() if len(args.estimator_seeds) < 4 or len(args.estimator_seeds) % 2: raise ValueError("estimator-seeds must contain an even count of at least 4") if len(set(args.estimator_seeds)) != len(args.estimator_seeds): raise ValueError("estimator-seeds must be unique") if args.samples_per_estimator < 2 or args.samples_per_estimator % 2: raise ValueError("samples-per-estimator must be a positive even integer") gradient_metadata, gradient_anchors = _load_gradient_artifact( args.gradient_artifact ) checkpoint_sha256 = _sha256_file(args.checkpoint) if ( gradient_metadata["checkpoint_sha256"] is not None and gradient_metadata["checkpoint_sha256"] != checkpoint_sha256 ): raise RuntimeError("checkpoint hash differs from the gradient artifact") gradient_seeds = set(gradient_metadata["gradient_seeds"] or []) if gradient_seeds & set(args.estimator_seeds): raise ValueError("residual estimator seeds overlap gradient seeds") cfg = apply_config_compat(OmegaConf.load(args.config)) cfg.output_dir = str(args.output.parent / "stratified_residual_runtime") examples = _load_aligned_examples(cfg, gradient_anchors) model = build_framework(cfg) checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=True) model.load_state_dict(checkpoint, strict=True) model.requires_grad_(False).eval().to("cuda") torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats() started = time.perf_counter() estimator_matrix = np.empty( (len(args.estimator_seeds), len(examples)), dtype=np.float64, ) score_records = [] time_values = None for anchor_index, example in enumerate(examples): with torch.no_grad(): result = score_starvla_flow_matching_residual( model, [example], estimator_seeds=args.estimator_seeds, samples_per_estimator=args.samples_per_estimator, ) values = result.estimator_residuals[:, 0].float().cpu().numpy() estimator_matrix[:, anchor_index] = values if time_values is None: time_values = result.time_values.float().cpu().tolist() score_records.append( { "anchor_position": anchor_index, "estimator_residuals": values.tolist(), "sample_losses": ( result.sample_losses[:, :, 0].float().cpu().tolist() ), "mean_residual": float(values.mean()), "sampling_trace": gradient_anchors[anchor_index][ "sampling_trace" ], "gradient_rms": float( gradient_anchors[anchor_index]["gradient_rms"] ), "mean_gradient_loss": float( gradient_anchors[anchor_index]["mean_gradient_loss"] ), } ) print( f"anchor={anchor_index + 1}/{len(examples)} " f"residual={values.mean():.6f}", flush=True, ) torch.cuda.synchronize() elapsed_seconds = time.perf_counter() - started scores = estimator_matrix.mean(axis=0) gradient_rms = np.asarray( [anchor["gradient_rms"] for anchor in gradient_anchors], dtype=np.float64, ) mean_losses = np.asarray( [anchor["mean_gradient_loss"] for anchor in gradient_anchors], dtype=np.float64, ) repeatability = _repeatability(estimator_matrix) alignment = alignment_audit( scores, gradient_rms, mean_losses, alpha=args.alpha, percentile_epsilon=args.percentile_epsilon, permutation_count=args.permutations, permutation_seed=args.permutation_seed, score_label="residual", ) robustness = leave_one_out_robustness( scores, gradient_rms, alpha=args.alpha, percentile_epsilon=args.percentile_epsilon, ) gate = _gate_decision( sample_count=len(examples), repeatability=repeatability, alignment=alignment, robustness=robustness, ) payload = { "schema_version": 1, "status": ( "formal_64_anchor_variance_reduced_gate" if len(examples) >= 64 else "variance_reduced_smoke_not_a_gate_certificate" ), "config": str(args.config.resolve()), "checkpoint": str(args.checkpoint.resolve()), "checkpoint_sha256": checkpoint_sha256, "gradient_artifact": gradient_metadata, "sample_count": len(examples), "score_definition": { "name": "frozen_flow_matching_residual", "variance_reduction": [ "deterministic midpoint stratification of Beta flow time", "paired antithetic Gaussian noise", "one shared VLM conditioning encode per anchor", "batched action-head evaluation", ], "model_mode": "eval", "estimator_seeds": args.estimator_seeds, "samples_per_estimator": args.samples_per_estimator, "time_values": time_values, "forbidden_labels": ["likelihood", "entropy", "surprisal"], }, "score_runtime": { "elapsed_seconds": elapsed_seconds, "seconds_per_anchor": elapsed_seconds / len(examples), "peak_allocated_bytes": torch.cuda.max_memory_allocated(), "peak_reserved_bytes": torch.cuda.max_memory_reserved(), }, "repeatability": repeatability, "alignment": alignment, "leave_one_out_robustness": robustness, "gate": gate, "anchors": score_records, } args.output.parent.mkdir(parents=True, exist_ok=True) temporary = args.output.with_suffix(args.output.suffix + ".tmp") temporary.write_text(json.dumps(payload, indent=2), encoding="utf-8") temporary.replace(args.output) print(json.dumps(payload, indent=2)) if __name__ == "__main__": main()