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