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