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