"""Audit frozen Flow-Matching residual as a gradient-importance proxy.""" from __future__ import annotations import argparse import hashlib import json import math import random 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.utility_audit import ( alignment_audit, leave_one_out_robustness, ) 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 _set_random_seed(seed: int) -> None: random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) def _load_examples(cfg, sample_count: int) -> 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) >= sample_count: break examples = examples[:sample_count] if len(examples) != sample_count: raise RuntimeError( f"dataloader returned {len(examples)} examples, expected {sample_count}" ) if any(example.get("sampling_trace") is None for example in examples): raise RuntimeError("every residual audit example must have a sampling trace") return examples def _measure_frozen_residual( model, examples: list[dict], *, repeat_count: int, seed_base: int, ) -> tuple[np.ndarray, list[dict], dict]: model.requires_grad_(False) model.eval() repeat_seeds = [seed_base + repeat for repeat in range(repeat_count)] residual_matrix = np.empty( (len(examples), repeat_count), dtype=np.float64, ) torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats() started = time.perf_counter() with torch.no_grad(): for anchor_index, example in enumerate(examples): for repeat_index, seed in enumerate(repeat_seeds): _set_random_seed(seed) residual = model([example])["action_loss"] if residual.ndim != 0 or not torch.isfinite(residual): raise FloatingPointError("frozen residual must be a finite scalar") residual_matrix[anchor_index, repeat_index] = float( residual.detach().cpu() ) print( f"score anchor={anchor_index + 1}/{len(examples)} " f"residual={residual_matrix[anchor_index].mean():.6f}", flush=True, ) torch.cuda.synchronize() elapsed_seconds = time.perf_counter() - started half = repeat_count // 2 first_half = residual_matrix[:, :half].mean(axis=1) second_half = residual_matrix[:, half:].mean(axis=1) half_spearman = spearmanr(first_half, second_half) half_kendall = kendalltau(first_half, second_half) top_count = max(1, len(examples) // 4) first_top = set(np.argsort(first_half)[-top_count:]) second_top = set(np.argsort(second_half)[-top_count:]) anchor_results = [] for anchor_index, example in enumerate(examples): values = residual_matrix[anchor_index] anchor_results.append( { "anchor_position": anchor_index, "residuals": values.tolist(), "mean_residual": float(values.mean()), "residual_cv": float( values.std(ddof=1) / values.mean() if values.mean() > 0.0 else 0.0 ), "sampling_trace": example["sampling_trace"], } ) 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(), "repeat_count": repeat_count, "repeat_seeds": repeat_seeds, "model_mode": "eval", "backward": False, } repeatability = { "first_half_seeds": repeat_seeds[:half], "second_half_seeds": repeat_seeds[half:], "spearman": float(half_spearman.statistic), "spearman_pvalue_asymptotic": float(half_spearman.pvalue), "kendall": float(half_kendall.statistic), "kendall_pvalue_asymptotic": float(half_kendall.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([result["residual_cv"] for result in anchor_results]) ), "median_anchor_cv": float( np.median([result["residual_cv"] for result in anchor_results]) ), } return residual_matrix, anchor_results, { "runtime": runtime, "repeatability": repeatability, } def _gradient_squared_norm(parameters: list[torch.nn.Parameter]) -> float: total = 0.0 saw_gradient = False for parameter in parameters: if parameter.grad is None: continue saw_gradient = True gradient = parameter.grad.detach() if not torch.isfinite(gradient).all(): raise FloatingPointError("encountered a non-finite action-model gradient") total += gradient.float().square().sum().item() if not saw_gradient: raise RuntimeError("no action-model parameter received a gradient") return total def _measure_gradient_rms( model, examples: list[dict], *, repeat_count: int, seed_base: int, ) -> tuple[list[dict], dict]: model.requires_grad_(False) model.action_model.requires_grad_(True) model.train() trainable_parameters = [ parameter for parameter in model.action_model.parameters() if parameter.requires_grad ] if not trainable_parameters: raise RuntimeError("action model has no trainable parameters") if any( parameter.requires_grad for parameter in model.qwen_vl_interface.parameters() ): raise RuntimeError("the VLM backbone must remain frozen") repeat_seeds = [seed_base + repeat for repeat in range(repeat_count)] torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats() started = time.perf_counter() anchor_results = [] for anchor_index, example in enumerate(examples): losses = [] gradient_norms = [] for seed in repeat_seeds: _set_random_seed(seed) model.zero_grad(set_to_none=True) loss = model([example])["action_loss"] if loss.ndim != 0 or not torch.isfinite(loss): raise FloatingPointError("training loss must be a finite scalar") loss.backward() losses.append(float(loss.detach().cpu())) gradient_norms.append( math.sqrt(_gradient_squared_norm(trainable_parameters)) ) model.zero_grad(set_to_none=True) gradient_norm_array = np.asarray(gradient_norms, dtype=np.float64) anchor_results.append( { "anchor_position": anchor_index, "losses": losses, "gradient_norms": gradient_norms, "mean_loss": float(np.mean(losses)), "gradient_rms": float( math.sqrt(np.mean(np.square(gradient_norm_array))) ), "gradient_norm_cv": float( gradient_norm_array.std(ddof=1) / gradient_norm_array.mean() if repeat_count > 1 and gradient_norm_array.mean() > 0.0 else 0.0 ), } ) print( f"gradient anchor={anchor_index + 1}/{len(examples)} " f"loss={anchor_results[-1]['mean_loss']:.6f} " f"gradient_rms={anchor_results[-1]['gradient_rms']:.6f}", flush=True, ) torch.cuda.synchronize() elapsed_seconds = time.perf_counter() - started return anchor_results, { "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(), "repeat_count": repeat_count, "repeat_seeds": repeat_seeds, "model_mode": "train", "frozen_backbone_mode": "train", "trainable_action_parameters": sum( parameter.numel() for parameter in trainable_parameters ), } 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("--output", type=Path, required=True) parser.add_argument("--max-samples", type=int, default=16) parser.add_argument("--score-repeats", type=int, default=4) parser.add_argument("--score-seed-base", type=int, default=7300) parser.add_argument("--gradient-repeats", type=int, default=4) parser.add_argument("--gradient-seed-base", type=int, default=8300) 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 args.max_samples < 4: raise ValueError("max-samples must be at least 4") if args.score_repeats < 4 or args.score_repeats % 2: raise ValueError("score-repeats must be an even integer of at least 4") if args.gradient_repeats < 2: raise ValueError("gradient-repeats must be at least 2") if args.permutations <= 0: raise ValueError("permutations must be positive") if not 0.0 <= args.alpha < 1.0: raise ValueError("alpha must satisfy 0 <= alpha < 1") if args.percentile_epsilon <= 0.0: raise ValueError("percentile-epsilon must be positive") score_seeds = set( range(args.score_seed_base, args.score_seed_base + args.score_repeats) ) gradient_seeds = set( range( args.gradient_seed_base, args.gradient_seed_base + args.gradient_repeats, ) ) if score_seeds & gradient_seeds: raise ValueError("score seeds and gradient seeds must be disjoint") cfg = apply_config_compat(OmegaConf.load(args.config)) cfg.output_dir = str(args.output.parent / "residual_gate_runtime") examples = _load_examples(cfg, args.max_samples) model = build_framework(cfg) checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=True) model.load_state_dict(checkpoint, strict=True) model.to("cuda") residual_matrix, residual_results, residual_audit = ( _measure_frozen_residual( model, examples, repeat_count=args.score_repeats, seed_base=args.score_seed_base, ) ) gradient_results, gradient_runtime = _measure_gradient_rms( model, examples, repeat_count=args.gradient_repeats, seed_base=args.gradient_seed_base, ) scores = residual_matrix.mean(axis=1) gradient_rms = np.asarray( [result["gradient_rms"] for result in gradient_results], dtype=np.float64, ) mean_losses = np.asarray( [result["mean_loss"] for result in gradient_results], dtype=np.float64, ) 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=args.max_samples, repeatability=residual_audit["repeatability"], alignment=alignment, robustness=robustness, ) for residual_result, gradient_result in zip( residual_results, gradient_results, ): residual_result.update( { "gradient_losses": gradient_result["losses"], "gradient_norms": gradient_result["gradient_norms"], "mean_gradient_loss": gradient_result["mean_loss"], "gradient_rms": gradient_result["gradient_rms"], "gradient_norm_cv": gradient_result["gradient_norm_cv"], } ) payload = { "schema_version": 1, "status": ( "formal_64_anchor_gate" if args.max_samples >= 64 else "smoke_measurement_not_a_gate_certificate" ), "config": str(args.config.resolve()), "checkpoint": str(args.checkpoint.resolve()), "checkpoint_sha256": _sha256_file(args.checkpoint), "score_definition": { "name": "frozen_flow_matching_residual", "model_mode": "eval", "aggregation": "mean over independent flow time/noise seeds", "forbidden_labels": ["likelihood", "entropy", "surprisal"], }, "seed_separation": { "score_seeds": sorted(score_seeds), "gradient_seeds": sorted(gradient_seeds), "disjoint": not bool(score_seeds & gradient_seeds), }, "sample_count": args.max_samples, "residual": residual_audit, "gradient_runtime": gradient_runtime, "alignment": alignment, "leave_one_out_robustness": robustness, "gate": gate, "anchors": residual_results, } 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()