"""Estimate Theorem-3 gradient second moments with disjoint stratified draws.""" from __future__ import annotations import argparse import hashlib import json 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.flow_residual import ( randomized_antithetic_stratified_flow_samples, ) from starVLA.scoring.partial_gradient import ( activate_partial_gradient_group, define_partial_gradient_groups, finite_gradient_l2_norm, ) 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 _set_random_seed(seed: int) -> None: bounded_seed = int(seed) % (2**32) random.seed(bounded_seed) np.random.seed(bounded_seed) torch.manual_seed(bounded_seed) torch.cuda.manual_seed_all(bounded_seed) def _trace_identity(trace: dict) -> dict: return {field: trace.get(field) for field in TRACE_IDENTITY_FIELDS} def _load_reference_artifact( path: Path, max_samples: int, ) -> 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) < max_samples: raise ValueError("reference artifact does not contain enough anchors") anchors = anchors[:max_samples] return { "path": str(path.resolve()), "sha256": _sha256_file(path), "checkpoint_sha256": payload.get("checkpoint_sha256"), "old_gradient_seeds": payload.get("seed_separation", {}).get( "gradient_seeds" ), }, 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 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 _prepare_conditioning(model, example: dict): images = [example["image"]] instructions = [example["lang"]] raw_actions = [example["action"]] raw_states = [example["state"]] if "state" in example else None with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): vlm_inputs = model.qwen_vl_interface.build_qwenvl_inputs( images=images, instructions=instructions, ) vlm_outputs = model.qwen_vl_interface( **vlm_inputs, output_attentions=False, output_hidden_states=True, return_dict=True, ) hidden_states = vlm_outputs.hidden_states[-1].detach() actions = torch.as_tensor( np.asarray(raw_actions), device=hidden_states.device, dtype=hidden_states.dtype, )[:, -model.action_horizon :, :] states = ( torch.as_tensor( np.asarray(raw_states), device=hidden_states.device, dtype=hidden_states.dtype, ) if raw_states is not None else None ) return hidden_states, actions, states def _flow_loss( model, conditioning, noise: torch.Tensor, time_value: torch.Tensor, ) -> torch.Tensor: hidden_states, actions, states = conditioning time = time_value.reshape(1) expanded_time = time[:, None, None] noisy_trajectory = (1.0 - expanded_time) * noise + expanded_time * actions target_velocity = actions - noise with torch.autocast("cuda", dtype=torch.float32): prediction = model.action_model.predict_velocity( hidden_states, noisy_trajectory, time, state=states, ) return (prediction - target_velocity).square().mean() def _measure_gradient_second_moment( model, examples: list[dict], *, estimator_seeds: list[int], samples_per_estimator: int, parameter_mode: str, frozen_backbone_eval: bool, ) -> tuple[np.ndarray, np.ndarray, dict]: action_config = model.action_model.config groups = define_partial_gradient_groups(model.action_model) estimator_values = np.empty( (len(examples), len(estimator_seeds)), dtype=np.float64, ) draw_norms = np.empty( (len(examples), len(estimator_seeds), samples_per_estimator), dtype=np.float64, ) torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats() started = time.perf_counter() parameters = [] for anchor_index, example in enumerate(examples): for estimator_index, estimator_seed in enumerate(estimator_seeds): model.requires_grad_(False) model.train() if frozen_backbone_eval: model.qwen_vl_interface.eval() conditioning_seed = estimator_seed * 1_000_003 + anchor_index _set_random_seed(conditioning_seed) conditioning = _prepare_conditioning(model, example) actions = conditioning[1] noises, times = randomized_antithetic_stratified_flow_samples( actions, sample_count=samples_per_estimator, seed=estimator_seed * 1_000_033 + anchor_index, beta_alpha=float(action_config.noise_beta_alpha), beta_beta=float(action_config.noise_beta_beta), noise_s=float(action_config.noise_s), ) if parameter_mode == "full": model.action_model.requires_grad_(True) parameters = list(model.action_model.parameters()) elif parameter_mode == "decoder": parameters = activate_partial_gradient_group( model.action_model, groups["decoder"], ) else: raise ValueError(f"unknown parameter mode: {parameter_mode}") for draw_index in range(samples_per_estimator): # Common dropout within an antithetic pair preserves each # marginal draw while strengthening cancellation. dropout_seed = ( estimator_seed * 1_000_037 + anchor_index * 101 + draw_index // 2 ) _set_random_seed(dropout_seed) model.zero_grad(set_to_none=True) loss = _flow_loss( model, conditioning, noises[draw_index], times[draw_index], ) if loss.ndim != 0 or not torch.isfinite(loss): raise FloatingPointError("flow loss must be a finite scalar") loss.backward() draw_norms[ anchor_index, estimator_index, draw_index, ] = finite_gradient_l2_norm(parameters) estimator_values[anchor_index, estimator_index] = np.sqrt( np.mean( np.square( draw_norms[anchor_index, estimator_index] ) ) ) print( f"mode={parameter_mode} anchor={anchor_index + 1}/{len(examples)} " f"h={np.sqrt(np.mean(np.square(draw_norms[anchor_index]))):.6f}", flush=True, ) model.zero_grad(set_to_none=True) torch.cuda.synchronize() elapsed_seconds = time.perf_counter() - started return estimator_values, draw_norms, { "elapsed_seconds": elapsed_seconds, "seconds_per_anchor": elapsed_seconds / len(examples), "seconds_per_draw": ( elapsed_seconds / len(examples) / len(estimator_seeds) / samples_per_estimator ), "peak_allocated_bytes": torch.cuda.max_memory_allocated(), "peak_reserved_bytes": torch.cuda.max_memory_reserved(), "parameter_mode": parameter_mode, "frozen_backbone_eval": frozen_backbone_eval, "parameter_count": sum( parameter.numel() for parameter in parameters ), } def _repeatability(estimator_values: np.ndarray) -> dict: if estimator_values.ndim != 2 or estimator_values.shape[1] < 4: raise ValueError("at least four estimator columns are required") if estimator_values.shape[1] % 2: raise ValueError("estimator count must be even") half = estimator_values.shape[1] // 2 first = np.sqrt(np.mean(np.square(estimator_values[:, :half]), axis=1)) second = np.sqrt(np.mean(np.square(estimator_values[:, half:]), axis=1)) correlation = spearmanr(first, second) concordance = kendalltau(first, second) return { "spearman": float(correlation.statistic), "spearman_pvalue_asymptotic": float(correlation.pvalue), "kendall": float(concordance.statistic), "kendall_pvalue_asymptotic": float(concordance.pvalue), "mean_anchor_cv": float( np.mean( estimator_values.std(axis=1, ddof=1) / estimator_values.mean(axis=1) ) ), } def _gate( sample_count: int, proxy_repeatability: dict, target_repeatability: dict, alignment: dict, robustness: dict, ) -> dict: checks = { "sample_count_at_least_64": sample_count >= 64, "target_repeatability_at_least_0_80": ( target_repeatability["spearman"] >= 0.80 ), "proxy_repeatability_at_least_0_90": ( proxy_repeatability["spearman"] >= 0.90 ), "proxy_target_spearman_at_least_0_70": ( alignment["spearman_partial_gradient_vs_gradient_rms"] >= 0.70 ), "spearman_permutation_p_at_most_0_01": ( alignment["random_permutation_control"][ "one_sided_pvalue_rho_at_least_observed" ] <= 0.01 ), "top25_overlap_at_least_0_625": alignment["top25_overlap"] >= 0.625, "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("--reference-artifact", type=Path, required=True) parser.add_argument("--output", type=Path, required=True) parser.add_argument("--max-samples", type=int, default=8) parser.add_argument("--samples-per-estimator", type=int, default=8) parser.add_argument( "--proxy-seeds", type=int, nargs="+", default=[12300, 12301, 12302, 12303], ) parser.add_argument( "--target-seeds", type=int, nargs="+", default=[13300, 13301, 13302, 13303], ) 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=14100) parser.add_argument("--frozen-backbone-eval", action="store_true") args = parser.parse_args() if args.max_samples < 4: raise ValueError("max-samples must be at least four") for name, seeds in ( ("proxy-seeds", args.proxy_seeds), ("target-seeds", args.target_seeds), ): if len(seeds) < 4 or len(seeds) % 2 or len(set(seeds)) != len(seeds): raise ValueError(f"{name} must contain an even count of unique seeds") if set(args.proxy_seeds) & set(args.target_seeds): raise ValueError("proxy and target seeds must be disjoint") if args.samples_per_estimator < 2 or args.samples_per_estimator % 2: raise ValueError("samples-per-estimator must be a positive even integer") reference, reference_anchors = _load_reference_artifact( args.reference_artifact, args.max_samples, ) checkpoint_sha256 = _sha256_file(args.checkpoint) if ( reference["checkpoint_sha256"] is not None and reference["checkpoint_sha256"] != checkpoint_sha256 ): raise RuntimeError("checkpoint hash differs from reference artifact") all_old_seeds = set(reference["old_gradient_seeds"] or []) if all_old_seeds & (set(args.proxy_seeds) | set(args.target_seeds)): raise ValueError("new seeds overlap the old gradient audit") cfg = apply_config_compat(OmegaConf.load(args.config)) cfg.output_dir = str(args.output.parent / "stratified_gradient_runtime") examples = _load_aligned_examples(cfg, reference_anchors) 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") proxy_estimators, proxy_draws, proxy_runtime = ( _measure_gradient_second_moment( model, examples, estimator_seeds=args.proxy_seeds, samples_per_estimator=args.samples_per_estimator, parameter_mode="decoder", frozen_backbone_eval=args.frozen_backbone_eval, ) ) target_estimators, target_draws, target_runtime = ( _measure_gradient_second_moment( model, examples, estimator_seeds=args.target_seeds, samples_per_estimator=args.samples_per_estimator, parameter_mode="full", frozen_backbone_eval=args.frozen_backbone_eval, ) ) proxy_scores = np.sqrt(np.mean(np.square(proxy_estimators), axis=1)) target_gradient_rms = np.sqrt( np.mean(np.square(target_estimators), axis=1) ) old_gradient_rms = np.asarray( [anchor["gradient_rms"] for anchor in reference_anchors], dtype=np.float64, ) mean_losses = np.asarray( [anchor["mean_gradient_loss"] for anchor in reference_anchors], dtype=np.float64, ) proxy_repeatability = _repeatability(proxy_estimators) target_repeatability = _repeatability(target_estimators) alignment = alignment_audit( proxy_scores, target_gradient_rms, mean_losses, alpha=args.alpha, percentile_epsilon=args.percentile_epsilon, permutation_count=args.permutations, permutation_seed=args.permutation_seed, score_label="partial_gradient", ) robustness = leave_one_out_robustness( proxy_scores, target_gradient_rms, alpha=args.alpha, percentile_epsilon=args.percentile_epsilon, ) gate = _gate( len(examples), proxy_repeatability, target_repeatability, alignment, robustness, ) old_target_correlation = spearmanr( target_gradient_rms, old_gradient_rms, ) payload = { "schema_version": 1, "status": ( "formal_64_anchor_stratified_gradient_gate" if len(examples) >= 64 else "stratified_gradient_smoke_not_a_gate_certificate" ), "config": str(args.config.resolve()), "checkpoint": str(args.checkpoint.resolve()), "checkpoint_sha256": checkpoint_sha256, "reference_artifact": reference, "sample_count": len(examples), "estimator": { "definition": "sqrt(mean over xi of squared per-draw gradient L2 norm)", "randomized_beta_stratification": True, "antithetic_gaussian_noise": True, "common_dropout_within_antithetic_pair": True, "samples_per_estimator": args.samples_per_estimator, "proxy_seeds": args.proxy_seeds, "target_seeds": args.target_seeds, "seeds_disjoint": True, "frozen_backbone_eval": args.frozen_backbone_eval, }, "proxy_runtime": proxy_runtime, "target_runtime": target_runtime, "proxy_repeatability": proxy_repeatability, "target_repeatability": target_repeatability, "new_vs_old_target_spearman": float(old_target_correlation.statistic), "new_vs_old_target_spearman_pvalue": float(old_target_correlation.pvalue), "alignment": alignment, "leave_one_out_robustness": robustness, "gate": gate, "eligible_for_500_step_ablation": gate[ "eligible_for_500_step_ablation" ], "anchors": [ { "anchor_position": index, "sampling_trace": reference_anchors[index]["sampling_trace"], "proxy_estimators": proxy_estimators[index].tolist(), "proxy_draw_norms": proxy_draws[index].tolist(), "proxy_gradient_rms": float(proxy_scores[index]), "target_estimators": target_estimators[index].tolist(), "target_draw_norms": target_draws[index].tolist(), "target_gradient_rms": float(target_gradient_rms[index]), "old_four_draw_gradient_rms": float(old_gradient_rms[index]), } for index in range(len(examples)) ], } 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() def _measure_gradient_second_moment_draw_shard( model, examples: list[dict], *, estimator_seeds: list[int], samples_per_estimator: int, draw_indices: list[int], parameter_mode: str, frozen_backbone_eval: bool, ) -> tuple[np.ndarray, np.ndarray, dict]: action_config = model.action_model.config groups = define_partial_gradient_groups(model.action_model) values = np.empty((len(examples), len(estimator_seeds)), dtype=np.float64) draw_norms = np.empty( (len(examples), len(estimator_seeds), len(draw_indices)), dtype=np.float64, ) torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats() started = time.perf_counter() parameters = [] for anchor_index, example in enumerate(examples): for estimator_index, estimator_seed in enumerate(estimator_seeds): model.requires_grad_(False) model.train() if frozen_backbone_eval: model.qwen_vl_interface.eval() conditioning_seed = estimator_seed * 1_000_003 + anchor_index _set_random_seed(conditioning_seed) conditioning = _prepare_conditioning(model, example) actions = conditioning[1] noises, times = randomized_antithetic_stratified_flow_samples( actions, sample_count=samples_per_estimator, seed=estimator_seed * 1_000_033 + anchor_index, beta_alpha=float(action_config.noise_beta_alpha), beta_beta=float(action_config.noise_beta_beta), noise_s=float(action_config.noise_s), ) if parameter_mode != "full": raise ValueError(f"unknown parameter mode: {parameter_mode}") model.action_model.requires_grad_(True) parameters = list(model.action_model.parameters()) for output_index, draw_index in enumerate(draw_indices): dropout_seed = ( estimator_seed * 1_000_037 + anchor_index * 101 + draw_index // 2 ) _set_random_seed(dropout_seed) model.zero_grad(set_to_none=True) loss = _flow_loss( model, conditioning, noises[draw_index], times[draw_index], ) if loss.ndim != 0 or not torch.isfinite(loss): raise FloatingPointError("flow loss must be a finite scalar") loss.backward() draw_norms[anchor_index, estimator_index, output_index] = finite_gradient_l2_norm(parameters) values[anchor_index, estimator_index] = np.sqrt( np.mean(np.square(draw_norms[anchor_index, estimator_index])) ) print( f"mode={parameter_mode} anchor={anchor_index + 1}/{len(examples)} " f"h={np.sqrt(np.mean(np.square(draw_norms[anchor_index]))):.6f}", flush=True, ) model.zero_grad(set_to_none=True) torch.cuda.synchronize() elapsed_seconds = time.perf_counter() - started return values, draw_norms, { "elapsed_seconds": elapsed_seconds, "seconds_per_anchor": elapsed_seconds / len(examples), "seconds_per_draw": elapsed_seconds / len(examples) / len(estimator_seeds) / len(draw_indices), "peak_allocated_bytes": torch.cuda.max_memory_allocated(), "peak_reserved_bytes": torch.cuda.max_memory_reserved(), "parameter_mode": parameter_mode, "frozen_backbone_eval": frozen_backbone_eval, "parameter_count": sum(parameter.numel() for parameter in parameters), }