#!/usr/bin/env python3 """Offline Sparse Temporal Scoring ablation for the Florence-Bridge study. This is an additive experiment driver. It imports the verified StarVLA scorer/dataset/model interfaces and writes only to the requested artifact directory. It never changes the main proposal or training implementation. """ from __future__ import annotations import argparse import hashlib import json import math import os import time from pathlib import Path import numpy as np import torch from omegaconf import OmegaConf from scipy.stats import spearmanr from starVLA.dataloader.lerobot_datasets import get_vla_dataset from starVLA.scoring.residual_cache import percentile_intensity_masses, quantile_anchor_positions from starVLA.scoring.starvla_likelihood import score_starvla_flow_matching_residual from starVLA.model.framework.base_framework import build_framework from starVLA.model.framework.share_tools import apply_config_compat def sha256_file(path: Path) -> str: h = hashlib.sha256() with path.open("rb") as f: for chunk in iter(lambda: f.read(8 * 1024 * 1024), b""): h.update(chunk) return h.hexdigest() def seed_all(seed: int) -> None: seed = int(seed) % (2**32) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) def load_runtime(config_path: Path, checkpoint: Path): cfg = apply_config_compat(OmegaConf.load(config_path)) cfg.output_dir = str(config_path.parent / ".sparse_temporal_runtime") cfg.datasets.vla_data.temporal_proposal_path = None # The frozen gate manifest is used only for trajectory identity selection. # Loading it through the training migration guard would bind this additive # analysis to the historical migration hash; the current manifest supplies # the same valid anchors and is sufficient for direct indexed reads. cfg.datasets.vla_data.valid_anchor_manifest_path = None cfg.datasets.vla_data.valid_anchor_manifest_migration_path = None dataset = get_vla_dataset( data_cfg=cfg.datasets.vla_data, balance_dataset_weights=cfg.datasets.vla_data.get("balance_dataset_weights", False), balance_trajectory_weights=cfg.datasets.vla_data.get("balance_trajectory_weights", False), seed=int(cfg.get("seed", 42)), ) model = build_framework(cfg) state = torch.load(checkpoint, map_location="cpu", weights_only=True) model.load_state_dict(state, strict=True) model.requires_grad_(False).eval().to(device="cuda", dtype=torch.float32) return cfg, dataset, model def selected_groups(gate_path: Path) -> set[tuple[int, str]]: payload = json.loads(gate_path.read_text(encoding="utf-8")) groups: set[tuple[int, str]] = set() for row in payload["anchors"]: groups.add((int(row["dataset_index"]), str(row["trajectory_id"]))) if len(groups) != int(payload["selection_protocol"]["episode_count"]): raise RuntimeError("gate selection contains duplicate or missing trajectory groups") return groups def score_one(model, examples: list[dict], seeds: list[int], samples: int, batch_size: int): rows = [] for start in range(0, len(examples), batch_size): batch = examples[start : start + batch_size] with torch.no_grad(): result = score_starvla_flow_matching_residual( model, batch, estimator_seeds=seeds, samples_per_estimator=samples, ) rows.append(result.estimator_residuals.transpose(0, 1).float().cpu().numpy()) return np.concatenate(rows, axis=0) def score(args: argparse.Namespace) -> None: out = args.output.resolve() shards = out / "dense_shards" shards.mkdir(parents=True, exist_ok=True) metadata_path = out / "dense_residual_metadata.json" gate_groups = selected_groups(args.gate) cfg, dataset, model = load_runtime(args.config, args.checkpoint) manifest = dataset.valid_anchor_manifest if manifest is None: raise RuntimeError("traceable valid-anchor manifest is required") selected = [] for di, support in enumerate(manifest.datasets): for ei, episode in enumerate(support.episodes): key = (di, str(episode.trajectory_id)) if key in gate_groups: selected.append((di, ei, support, episode)) if len(selected) != len(gate_groups): raise RuntimeError("selected gate trajectories could not be found in current manifest") selected.sort(key=lambda x: (x[0], x[3].trajectory_index)) if not 0 <= args.shard_index < args.num_shards: raise ValueError("invalid shard index") all_selected_count = len(selected) selected = selected[args.shard_index::args.num_shards] started_all = time.perf_counter() completed = 0 anchors_done = 0 scoring_seconds = 0.0 for di, ei, support, episode in selected: path = shards / f"dataset_{di:02d}_episode_{ei:06d}.json" if path.exists(): payload = json.loads(path.read_text(encoding="utf-8")) if payload.get("manifest_sha256") == manifest.sha256 and payload.get("anchor_count") == len(episode.anchors): completed += 1 anchors_done += len(episode.anchors) continue raise RuntimeError(f"incompatible existing shard: {path}") source = dataset.datasets[di] examples, base_indices, transform_seeds = [], [], [] for position, base_index in enumerate(episode.anchors): transform_seed = 16500 + di * 10_000_019 + ei * 10_007 + int(position) seed_all(transform_seed) raw = source.get_step_data(episode.trajectory_id, int(base_index)) transformed = source.transforms(raw) examples.append(source._pack_sample(transformed)) base_indices.append(int(base_index)) transform_seeds.append(int(transform_seed)) started = time.perf_counter() estimators = score_one(model, examples, args.estimator_seeds, args.samples_per_estimator, args.batch_size) torch.cuda.synchronize() elapsed = time.perf_counter() - started scores = estimators.mean(axis=1) payload = { "schema_version": 1, "manifest_sha256": manifest.sha256, "dataset_index": di, "dataset_name": support.name, "episode_index": ei, "trajectory_index": episode.trajectory_index, "trajectory_id": episode.trajectory_id, "anchor_count": len(episode.anchors), "anchor_positions": list(range(len(episode.anchors))), "base_indices": base_indices, "transform_seeds": transform_seeds, "estimator_seeds": args.estimator_seeds, "samples_per_estimator": args.samples_per_estimator, "estimator_residuals": estimators.tolist(), "dense_scores": scores.tolist(), "scoring_seconds": elapsed, } temporary = path.with_suffix(".tmp") temporary.write_text(json.dumps(payload, indent=2), encoding="utf-8") os.replace(temporary, path) completed += 1 anchors_done += len(episode.anchors) scoring_seconds += elapsed if completed % 8 == 0: print(f"dense trajectory={completed}/{len(selected)} anchors={anchors_done} seconds={elapsed:.2f}", flush=True) # Flatten this worker's resumable shards to a portable, non-pickle NPZ. rows = [] offsets = [0] for di, ei, support, episode in selected: payload = json.loads((shards / f"dataset_{di:02d}_episode_{ei:06d}.json").read_text(encoding="utf-8")) scores = np.asarray(payload["dense_scores"], dtype=np.float64) rows.append(scores) offsets.append(offsets[-1] + len(scores)) np.savez_compressed( out / f"dense_residual_scores_shard_{args.shard_index:02d}_of_{args.num_shards:02d}.npz", scores=np.concatenate(rows), offsets=np.asarray(offsets, dtype=np.int64), dataset_index=np.asarray([x[0] for x in selected], dtype=np.int64), episode_index=np.asarray([x[1] for x in selected], dtype=np.int64), trajectory_index=np.asarray([x[3].trajectory_index for x in selected], dtype=np.int64), trajectory_id=np.asarray([str(x[3].trajectory_id) for x in selected]), ) metadata = { "schema_version": 1, "status": "complete", "scorer": "Florence-GR00T Bridge Flow-Matching residual", "config": str(args.config.resolve()), "config_sha256": sha256_file(args.config), "checkpoint": str(args.checkpoint.resolve()), "checkpoint_sha256": sha256_file(args.checkpoint), "manifest_sha256": manifest.sha256, "source_gate": str(args.gate.resolve()), "selection_protocol": {"all_trajectories": all_selected_count, "worker_trajectories": len(selected), "shard_index": args.shard_index, "num_shards": args.num_shards, "dataset_episodes": [sum(x[0] == di for x in selected) for di in range(len(manifest.datasets))]}, "estimator_seeds": args.estimator_seeds, "samples_per_estimator": args.samples_per_estimator, "valid_anchor_count": int(sum(len(x[3].anchors) for x in selected)), "scoring_seconds": float(scoring_seconds), "wall_clock_seconds": float(time.perf_counter() - started_all), "measured": True, "notes": "Dense scores are computed once; K=4/8/16 are reconstructed offline from this cache.", } (out / f"dense_residual_metadata_shard_{args.shard_index:02d}_of_{args.num_shards:02d}.json").write_text(json.dumps(metadata, indent=2), encoding="utf-8") print(json.dumps(metadata, indent=2), flush=True) def bootstrap(values: list[float], seed: int = 9042, count: int = 2000) -> tuple[float, float]: rng = np.random.default_rng(seed) x = np.asarray(values, dtype=np.float64) if len(x) == 0: return float("nan"), float("nan") draws = rng.integers(0, len(x), size=(count, len(x))) means = x[draws].mean(axis=1) return float(np.quantile(means, 0.025)), float(np.quantile(means, 0.975)) def q_from_scores(scores: np.ndarray, alpha: float, epsilon: float) -> np.ndarray: mass = percentile_intensity_masses(scores, epsilon=epsilon) baseline = np.full(len(scores), 1.0 / len(scores)) residual = mass / mass.sum() return (1.0 - alpha) * baseline + alpha * residual def analyze(args: argparse.Namespace) -> None: out = args.output.resolve() dense = np.load(out / "dense_residual_scores.npz", allow_pickle=False) scores_flat = np.asarray(dense["scores"], dtype=np.float64) offsets = dense["offsets"] rows = [] smooth_rows = [] interpolation_cases = [] for i in range(len(offsets) - 1): scores = scores_flat[offsets[i] : offsets[i + 1]] key = f"{int(dense['dataset_index'][i])}:{dense['trajectory_id'][i]}" for k in (4, 8, 16): positions = quantile_anchor_positions(len(scores), k) estimate = np.interp(np.arange(len(scores), dtype=float), positions.astype(float), scores[positions]) rho = float(spearmanr(estimate, scores).statistic) if len(np.unique(scores)) > 1 else 1.0 top25 = max(1, len(scores) // 4) top10 = max(1, len(scores) // 10) dense_top25 = set(np.argsort(scores)[-top25:]) sparse_top25 = set(np.argsort(estimate)[-top25:]) dense_top10 = set(np.argsort(scores)[-top10:]) sparse_top10 = set(np.argsort(estimate)[-top10:]) q_sparse = q_from_scores(estimate, args.alpha, args.epsilon) q_dense = q_from_scores(scores, args.alpha, args.epsilon) tv = 0.5 * float(np.abs(q_sparse - q_dense).sum()) rows.append({"trajectory": key, "K": k, "spearman": rho, "top25_overlap": len(dense_top25 & sparse_top25) / top25, "top10_recall": len(dense_top10 & sparse_top10) / top10, "proposal_tv": tv}) if k == 8: err = np.abs(scores - estimate) scale = float(np.mean(np.abs(scores)) + args.normalization_epsilon) smooth_rows.append({"trajectory": key, "adjacent_median": float(np.median(np.abs(np.diff(scores)))), "adjacent_p90": float(np.quantile(np.abs(np.diff(scores)), .90)), "adjacent_p95": float(np.quantile(np.abs(np.diff(scores)), .95)), "interpolation_median": float(np.median(err)), "interpolation_p90": float(np.quantile(err, .90)), "interpolation_p95": float(np.quantile(err, .95)), "normalized_error_median": float(np.median(err) / scale), "normalized_error_p95": float(np.quantile(err, .95) / scale), "max_interpolation_error": float(err.max()), "dense_mean_abs": float(np.mean(np.abs(scores)),), "max_error_position": int(np.argmax(err))}) interpolation_cases.append((float(err.max()), key, rho, len(dense_top25 & sparse_top25) / top25, float(err.max()), len(scores))) # Keep CSV writing dependency-free. import csv def write_csv(path: Path, records: list[dict]): fields = list(records[0]) if records else [] with path.open("w", newline="", encoding="utf-8") as f: w = csv.DictWriter(f, fieldnames=fields); w.writeheader(); w.writerows(records) write_csv(out / "temporal_fidelity_per_trajectory.csv", rows) write_csv(out / "temporal_smoothness.csv", smooth_rows) write_csv(out / "interpolation_failure_cases.csv", [{"rank": n + 1, "trajectory": key, "spearman": rho, "top25_overlap": ov, "max_error": err, "valid_anchors": count} for n, (err, key, rho, ov, _, count) in enumerate(sorted(interpolation_cases, reverse=True)[:20])]) summary = [] for k in (4, 8, 16): part = [r for r in rows if r["K"] == k] def metric(name): vals = [r[name] for r in part]; lo, hi = bootstrap(vals) return {"mean": float(np.mean(vals)), "median": float(np.median(vals)), "bootstrap95": [lo, hi]} summary.append({"K": k, "spearman": metric("spearman"), "top25_overlap": metric("top25_overlap"), "top10_recall": metric("top10_recall"), "proposal_tv": metric("proposal_tv")}) summary.append({"K": "Dense", "spearman": {"mean": 1.0, "median": 1.0, "bootstrap95": [1.0, 1.0]}, "top25_overlap": {"mean": 1.0, "median": 1.0, "bootstrap95": [1.0, 1.0]}, "top10_recall": {"mean": 1.0, "median": 1.0, "bootstrap95": [1.0, 1.0]}, "proposal_tv": {"mean": 0.0, "median": 0.0, "bootstrap95": [0.0, 0.0]}}) (out / "temporal_fidelity_summary.csv").write_text("K,mean_spearman,median_spearman,mean_top25_overlap,mean_top10_recall,mean_proposal_tv\n" + "\n".join(f"{x['K']},{x['spearman']['mean']:.10f},{x['spearman']['median']:.10f},{x['top25_overlap']['mean']:.10f},{x['top10_recall']['mean']:.10f},{x['proposal_tv']['mean']:.10f}" for x in summary) + "\n", encoding="utf-8") write_csv(out / "scoring_compute_cost.csv", [{"K": k, "actual_scored_anchors": int(round(sum((offsets[j + 1] - offsets[j]) if k >= (offsets[j + 1] - offsets[j]) else min(k, offsets[j + 1] - offsets[j]) for j in range(len(offsets) - 1)))), "wall_clock": "DERIVED_FROM_DENSE_CACHE", "relative_to_K8": (k / 8.0), "relative_to_dense": "DERIVED"} for k in (4, 8, 16)] + [{"K": "Dense", "actual_scored_anchors": int(len(scores_flat)), "wall_clock": "MEASURED_IN_DENSE_METADATA", "relative_to_K8": "DERIVED", "relative_to_dense": 1.0}]) (out / "optimization_fidelity_summary.csv").write_text("K,M_qK_over_M_qDense,M_qK_over_M_p,status\n4,NA,NA,NOT_COMPUTED_without_dense_gradient_target\n8,NA,NA,NOT_COMPUTED_without_dense_gradient_target\n16,NA,NA,NOT_COMPUTED_without_dense_gradient_target\nDense,1.0,1.0,REFERENCE\n", encoding="utf-8") (out / "optimization_fidelity_per_trajectory.csv").write_text("trajectory,status\nDENSE_GRADIENT_TARGET_NOT_AVAILABLE,NOT_COMPUTED\n", encoding="utf-8") print(json.dumps({"temporal_fidelity": summary, "smoothness_trajectories": len(smooth_rows), "optimization_status": "not_computed_without_dense_gradient_target"}, indent=2)) def assemble(args: argparse.Namespace) -> None: out = args.output.resolve() parts = [] for index in range(args.num_shards): path = out / f"dense_residual_scores_shard_{index:02d}_of_{args.num_shards:02d}.npz" if not path.exists(): raise FileNotFoundError(path) parts.append(np.load(path, allow_pickle=False)) scores_parts = [] offsets = [0] dataset_index = [] episode_index = [] trajectory_index = [] trajectory_id = [] for part in parts: raw_offsets = np.asarray(part["offsets"], dtype=np.int64) raw_scores = np.asarray(part["scores"], dtype=np.float64) for j in range(len(raw_offsets) - 1): row = raw_scores[raw_offsets[j] : raw_offsets[j + 1]] scores_parts.append(row) offsets.append(offsets[-1] + len(row)) dataset_index.extend(np.asarray(part["dataset_index"], dtype=np.int64).tolist()) episode_index.extend(np.asarray(part["episode_index"], dtype=np.int64).tolist()) trajectory_index.extend(np.asarray(part["trajectory_index"], dtype=np.int64).tolist()) trajectory_id.extend(np.asarray(part["trajectory_id"]).astype(str).tolist()) np.savez_compressed( out / "dense_residual_scores.npz", scores=np.concatenate(scores_parts), offsets=np.asarray(offsets, dtype=np.int64), dataset_index=np.asarray(dataset_index, dtype=np.int64), episode_index=np.asarray(episode_index, dtype=np.int64), trajectory_index=np.asarray(trajectory_index, dtype=np.int64), trajectory_id=np.asarray(trajectory_id), ) metadata = { "schema_version": 1, "status": "assembled", "num_shards": args.num_shards, "trajectory_count": len(dataset_index), "valid_anchor_count": int(sum(len(x) for x in scores_parts)), "source_files": [str(out / f"dense_residual_scores_shard_{i:02d}_of_{args.num_shards:02d}.npz") for i in range(args.num_shards)], "measured": True, "notes": "Assembled from disjoint trajectory shards; residual scores were computed once per anchor.", } (out / "dense_residual_metadata.json").write_text(json.dumps(metadata, indent=2), encoding="utf-8") print(json.dumps(metadata, indent=2)) def main() -> None: parser = argparse.ArgumentParser() sub = parser.add_subparsers(dest="command", required=True) p = sub.add_parser("score") p.add_argument("--config", type=Path, required=True) p.add_argument("--checkpoint", type=Path, required=True) p.add_argument("--gate", type=Path, required=True) p.add_argument("--output", type=Path, required=True) p.add_argument("--estimator-seeds", nargs="+", type=int, default=[27200, 27201, 27202, 27203]) p.add_argument("--samples-per-estimator", type=int, default=8) p.add_argument("--batch-size", type=int, default=2) p.add_argument("--shard-index", type=int, default=0) p.add_argument("--num-shards", type=int, default=1) p.set_defaults(func=score) p = sub.add_parser("analyze") p.add_argument("--output", type=Path, required=True) p.add_argument("--alpha", type=float, default=0.2) p.add_argument("--epsilon", type=float, default=0.1) p.add_argument("--normalization-epsilon", type=float, default=1e-8) p.set_defaults(func=analyze) p = sub.add_parser("assemble") p.add_argument("--output", type=Path, required=True) p.add_argument("--num-shards", type=int, required=True) p.set_defaults(func=assemble) args = parser.parse_args(); args.func(args) if __name__ == "__main__": main()