"""Build a resumable, grid-scored residual temporal proposal. The expensive frozen-policy residual is evaluated at a fixed quantile grid inside every episode. Scores for the remaining anchors are filled by deterministic linear interpolation, ranked within that episode, and converted to epsilon-floored masses. ``TemporalProposal`` then applies the safe mixture and exact ``p/q`` correction used by training. """ from __future__ import annotations import argparse import hashlib import json import os import random import time from pathlib import Path from typing import Any import numpy as np import torch from omegaconf import OmegaConf from starVLA.dataloader.lerobot_datasets import get_vla_dataset from starVLA.dataloader.temporal_proposal import TemporalProposal from starVLA.model.framework.base_framework import build_framework from starVLA.model.framework.share_tools import apply_config_compat from starVLA.scoring.residual_cache import ( interpolate_grid_scores, percentile_intensity_masses, quantile_anchor_positions, ) from starVLA.scoring.starvla_likelihood import ( score_starvla_flow_matching_residual, ) SCHEMA_VERSION = 1 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 _canonical_sha256(payload: dict[str, Any]) -> str: encoded = json.dumps( payload, sort_keys=True, separators=(",", ":"), ).encode("utf-8") return hashlib.sha256(encoded).hexdigest() def _load_gate_certificate( path: Path, *, checkpoint_sha256: str, manifest_sha256: str, alpha: float, epsilon: float, allow_failed_gate: bool, ) -> dict[str, Any]: payload = json.loads(path.read_text(encoding="utf-8")) gate = payload.get("gate") proposal = payload.get("hierarchical_audit", {}).get("proposal", {}) if (not isinstance(gate, dict) or gate.get("pass") is not True) and not allow_failed_gate: raise ValueError("Proposal eligibility gate did not pass") if gate.get("eligible_for_500_step_ablation") is not True and not allow_failed_gate: raise ValueError("Certificate does not authorize a prospective ablation") status = payload.get("status") if status == "formal_within_trajectory_utility_gate": pass elif status == "formal_independent_target_panel_second_moment_gate": panel = payload.get("panel_audit", {}) if ( int(panel.get("panel_count", 0)) < 8 or int(panel.get("draws_per_panel", 0)) < 4 or panel.get("panels_are_disjoint") is not True or float(panel.get("maximum_individual_ratio", 1.0)) > 0.98 or float(panel.get("minimum_individual_improved_fraction", 0.0)) < 0.75 or int(panel.get("complementary_split_count", 0)) < 70 or float( panel.get("maximum_complementary_split_ratio", 1.0) ) > 0.98 or float(panel.get("minimum_split_spearman", 0.0)) < 0.50 ): raise ValueError( "Independent-panel certificate misses its locked thresholds" ) else: raise ValueError( "Proposal eligibility certificate has the wrong formal status" ) if payload.get("checkpoint_sha256") != checkpoint_sha256: raise ValueError("Certificate belongs to a different checkpoint") if payload.get("manifest_sha256") != manifest_sha256: raise ValueError("Certificate belongs to a different valid-anchor manifest") if not np.isclose(float(proposal.get("alpha", -1.0)), alpha, atol=0.0): raise ValueError("Certificate alpha differs from the proposal alpha") if not np.isclose( float(proposal.get("percentile_epsilon", -1.0)), epsilon, atol=0.0, ): raise ValueError("Certificate epsilon differs from the proposal epsilon") selection = payload.get("selection_protocol", {}) episode_count = int(selection.get("episode_count", -1)) anchors_per_episode = int(selection.get("anchors_per_episode", -1)) anchor_count = int(selection.get("anchor_count", -1)) if ( episode_count < 32 or anchors_per_episode != 8 or anchor_count != episode_count * anchors_per_episode ): raise ValueError( "Certificate is not a valid independent-episode/8-anchor audit" ) return { "override_gate_failure": bool(allow_failed_gate and gate.get("pass") is not True), "path": str(path.resolve()), "sha256": _sha256_file(path), "protocol": status, "status": gate.get("status"), "checks": gate.get("checks"), } def _atomic_json(path: Path, payload: dict[str, Any]) -> None: path.parent.mkdir(parents=True, exist_ok=True) temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp") try: with temporary.open("w", encoding="utf-8") as output: json.dump(payload, output, indent=2, sort_keys=True) output.flush() os.fsync(output.fileno()) os.replace(temporary, path) finally: temporary.unlink(missing_ok=True) 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 _json_scalar(value: Any) -> Any: return value.item() if isinstance(value, np.generic) else value def _shard_path(cache_dir: Path, dataset_index: int, episode_index: int) -> Path: return cache_dir / "shards" / ( f"dataset_{dataset_index:02d}_episode_{episode_index:06d}.json" ) def _planned_grid_count(manifest, grid_size: int) -> int: return sum( min(len(episode.anchors), grid_size) for dataset in manifest.datasets for episode in dataset.episodes ) def _episode_belongs_to_shard( global_episode_index: int, *, shard_index: int, num_shards: int, ) -> bool: if num_shards <= 0: raise ValueError("num_shards must be positive") if not 0 <= shard_index < num_shards: raise ValueError("shard_index must satisfy 0 <= shard_index < num_shards") return int(global_episode_index) % int(num_shards) == int(shard_index) def _worker_progress_path(cache_dir: Path, shard_index: int, num_shards: int) -> Path: return cache_dir / "workers" / f"shard_{shard_index:02d}_of_{num_shards:02d}.json" def _load_valid_shard( path: Path, *, identity_sha256: str, dataset_index: int, episode_index: int, trajectory_id: Any, positions: np.ndarray, base_indices: list[int], estimator_count: int, ) -> dict[str, Any]: payload = json.loads(path.read_text(encoding="utf-8")) expected = { "identity_sha256": identity_sha256, "dataset_index": dataset_index, "episode_index": episode_index, "trajectory_id": _json_scalar(trajectory_id), "anchor_positions": positions.tolist(), "base_indices": base_indices, } for field, value in expected.items(): if payload.get(field) != value: raise ValueError(f"Residual cache shard mismatch for {field}: {path}") estimators = np.asarray(payload.get("estimator_residuals"), dtype=np.float64) if estimators.shape != (len(positions), estimator_count): raise ValueError(f"Residual cache shard has wrong estimator shape: {path}") if not np.isfinite(estimators).all(): raise ValueError(f"Residual cache shard contains non-finite values: {path}") scores = np.asarray(payload.get("grid_scores"), dtype=np.float64) if scores.shape != (len(positions),) or not np.isfinite(scores).all(): raise ValueError(f"Residual cache shard contains invalid grid scores: {path}") return payload def _load_examples( source_dataset, episode, positions: np.ndarray, *, dataset_index: int, episode_index: int, transform_seed_base: int, ) -> tuple[list[dict], list[int], list[int]]: examples = [] base_indices = [] transform_seeds = [] for anchor_position in positions: base_index = int(episode.anchors[int(anchor_position)]) transform_seed = ( transform_seed_base + dataset_index * 10_000_019 + episode_index * 10_007 + int(anchor_position) ) _set_random_seed(transform_seed) raw_data = source_dataset.get_step_data( episode.trajectory_id, base_index, ) transformed = source_dataset.transforms(raw_data) examples.append(source_dataset._pack_sample(transformed)) base_indices.append(base_index) transform_seeds.append(transform_seed) return examples, base_indices, transform_seeds def _score_examples( model, examples: list[dict], *, estimator_seeds: list[int], samples_per_estimator: int, batch_size: int, ) -> np.ndarray: all_estimators = np.empty( (len(examples), len(estimator_seeds)), dtype=np.float64, ) 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=estimator_seeds, samples_per_estimator=samples_per_estimator, ) batch_estimators = ( result.estimator_residuals.transpose(0, 1).float().cpu().numpy() ) all_estimators[start : start + len(batch)] = batch_estimators return all_estimators def _build_identity(args, manifest, gate_certificate) -> dict[str, Any]: checkpoint = args.checkpoint.resolve() identity = { "schema_version": SCHEMA_VERSION, # Runtime locations are intentionally excluded from the scientific # identity; the checkpoint content and explicit scoring parameters bind it. "checkpoint_sha256": _sha256_file(checkpoint), "manifest_sha256": manifest.sha256, "grid_size": int(args.grid_size), "interpolation": "linear_within_episode", "percentile_scope": "within_episode_after_interpolation", "epsilon": float(args.epsilon), "alpha": float(args.alpha), "estimator_seeds": [int(seed) for seed in args.estimator_seeds], "samples_per_estimator": int(args.samples_per_estimator), "transform_seed_base": int(args.transform_seed_base), "model_dtype": args.model_dtype, "gate42f_certificate": gate_certificate, } identity["sha256"] = _canonical_sha256(identity) return identity def _initialize_metadata( path: Path, *, identity: dict[str, Any], episode_count: int, grid_anchor_count: int, ) -> dict[str, Any]: if path.exists(): metadata = json.loads(path.read_text(encoding="utf-8")) if metadata.get("identity") != identity: raise ValueError( "Residual cache metadata belongs to a different scorer protocol" ) return metadata metadata = { "schema_version": SCHEMA_VERSION, "status": "in_progress", "identity": identity, "planned_episode_count": int(episode_count), "planned_grid_anchor_count": int(grid_anchor_count), "completed_episode_count": 0, "completed_grid_anchor_count": 0, "scoring_seconds": 0.0, "proposal": None, } _atomic_json(path, metadata) return metadata def _refresh_progress(metadata: dict[str, Any], cache_dir: Path, manifest) -> None: completed_episodes = 0 completed_anchors = 0 for dataset_index, dataset in enumerate(manifest.datasets): for episode_index, _ in enumerate(dataset.episodes): path = _shard_path(cache_dir, dataset_index, episode_index) if not path.exists(): continue payload = json.loads(path.read_text(encoding="utf-8")) completed_episodes += 1 completed_anchors += len(payload.get("anchor_positions", [])) metadata["completed_episode_count"] = completed_episodes metadata["completed_grid_anchor_count"] = completed_anchors def _finalize_proposal( *, manifest, cache_dir: Path, identity: dict[str, Any], output: Path, alpha: float, epsilon: float, allow_failed_gate: bool, estimator_count: int, gate_certificate: dict[str, Any], ) -> TemporalProposal: masses = [] for dataset_index, dataset in enumerate(manifest.datasets): dataset_masses = [] for episode_index, episode in enumerate(dataset.episodes): positions = quantile_anchor_positions( len(episode.anchors), int(identity["grid_size"]), ) base_indices = [int(episode.anchors[int(index)]) for index in positions] shard = _load_valid_shard( _shard_path(cache_dir, dataset_index, episode_index), identity_sha256=identity["sha256"], dataset_index=dataset_index, episode_index=episode_index, trajectory_id=episode.trajectory_id, positions=positions, base_indices=base_indices, estimator_count=estimator_count, ) full_scores = interpolate_grid_scores( len(episode.anchors), positions, np.asarray(shard["grid_scores"], dtype=np.float64), ) dataset_masses.append( percentile_intensity_masses(full_scores, epsilon=epsilon) ) masses.append(dataset_masses) proposal = TemporalProposal.from_masses( manifest=manifest, alpha=alpha, masses=masses, score_kind="residual", scorer_sha256=identity["sha256"], validation_certificate={ "protocol": "prospective-grid-residual-v1", "identity_sha256": identity["sha256"], "grid_size": identity["grid_size"], "interpolation": identity["interpolation"], "percentile_scope": identity["percentile_scope"], "epsilon": epsilon, "gate42f_certificate": gate_certificate, }, ) proposal.save(output) return proposal 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("--gate-result", type=Path, required=True) parser.add_argument("--allow-failed-gate", action="store_true") parser.add_argument("--cache-dir", type=Path, required=True) parser.add_argument("--output", type=Path, required=True) parser.add_argument("--grid-size", type=int, default=8) parser.add_argument("--alpha", type=float, default=0.2) parser.add_argument("--epsilon", type=float, default=0.1) parser.add_argument( "--estimator-seeds", type=int, nargs="+", default=[16200, 16201, 16202, 16203], ) parser.add_argument("--samples-per-estimator", type=int, default=8) parser.add_argument("--transform-seed-base", type=int, default=16500) parser.add_argument("--batch-size", type=int, default=1) parser.add_argument( "--model-dtype", choices=("float32", "bfloat16"), default="float32", ) parser.add_argument("--num-shards", type=int, default=1) parser.add_argument("--shard-index", type=int, default=0) parser.add_argument( "--max-new-episodes", type=int, default=0, help="Stop cleanly after this many newly scored episodes; 0 means all.", ) parser.add_argument("--plan-only", action="store_true") args = parser.parse_args() if args.grid_size <= 0 or args.batch_size <= 0: raise ValueError("grid-size and batch-size must be positive") if args.num_shards <= 0: raise ValueError("num-shards must be positive") if not 0 <= args.shard_index < args.num_shards: raise ValueError("shard-index must satisfy 0 <= shard-index < num-shards") if not 0.0 <= args.alpha < 1.0: raise ValueError("alpha must satisfy 0 <= alpha < 1") if args.epsilon <= 0.0: raise ValueError("epsilon must be positive") if not args.estimator_seeds: raise ValueError("at least one estimator seed is required") if len(set(args.estimator_seeds)) != len(args.estimator_seeds): raise ValueError("estimator seeds must be unique") if args.samples_per_estimator <= 0: raise ValueError("samples-per-estimator must be positive") cfg = apply_config_compat(OmegaConf.load(args.config)) cfg.output_dir = str(args.cache_dir / "runtime") cfg.datasets.vla_data.temporal_proposal_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)), ) manifest = dataset.valid_anchor_manifest if manifest is None: raise RuntimeError("Residual proposal requires traceable_sampling") checkpoint_sha256 = _sha256_file(args.checkpoint.resolve()) gate_certificate = _load_gate_certificate( args.gate_result.resolve(), checkpoint_sha256=checkpoint_sha256, manifest_sha256=manifest.sha256, alpha=float(args.alpha), epsilon=float(args.epsilon), allow_failed_gate=args.allow_failed_gate, ) identity = _build_identity(args, manifest, gate_certificate) episode_count = sum(len(item.episodes) for item in manifest.datasets) grid_anchor_count = _planned_grid_count(manifest, args.grid_size) metadata_path = args.cache_dir / "metadata.json" metadata = _initialize_metadata( metadata_path, identity=identity, episode_count=episode_count, grid_anchor_count=grid_anchor_count, ) _refresh_progress(metadata, args.cache_dir, manifest) _atomic_json(metadata_path, metadata) print( json.dumps( { "identity_sha256": identity["sha256"], "episodes": episode_count, "grid_anchors": grid_anchor_count, "completed_episodes": metadata["completed_episode_count"], "completed_grid_anchors": metadata["completed_grid_anchor_count"], "estimated_hours_at_0.312_seconds_per_anchor": ( grid_anchor_count * 0.312 / 3600.0 ), }, indent=2, ), flush=True, ) if args.plan_only: return parallel_worker = args.num_shards > 1 assigned_episode_count = 0 assigned_grid_anchor_count = 0 global_episode_index = 0 for dataset_support in manifest.datasets: for episode in dataset_support.episodes: if _episode_belongs_to_shard( global_episode_index, shard_index=args.shard_index, num_shards=args.num_shards, ): assigned_episode_count += 1 assigned_grid_anchor_count += min( len(episode.anchors), args.grid_size, ) global_episode_index += 1 missing_exists = metadata["completed_episode_count"] < episode_count if missing_exists: model = build_framework(cfg) checkpoint = torch.load( args.checkpoint, map_location="cpu", weights_only=True, ) model.load_state_dict(checkpoint, strict=True) model.requires_grad_(False).eval().to( device="cuda", dtype=getattr(torch, args.model_dtype), ) else: model = None new_episodes = 0 completed_assigned_episodes = 0 completed_assigned_anchors = 0 session_scoring_seconds = 0.0 worker_progress = { "schema_version": 1, "status": "in_progress", "identity_sha256": identity["sha256"], "shard_index": args.shard_index, "num_shards": args.num_shards, "assigned_episode_count": assigned_episode_count, "assigned_grid_anchor_count": assigned_grid_anchor_count, "completed_episode_count": 0, "completed_grid_anchor_count": 0, "new_episode_count": 0, "scoring_seconds": 0.0, } worker_path = _worker_progress_path( args.cache_dir, args.shard_index, args.num_shards, ) if parallel_worker: _atomic_json(worker_path, worker_progress) try: global_episode_index = 0 for dataset_index, dataset_support in enumerate(manifest.datasets): source_dataset = dataset.datasets[dataset_index] for episode_index, episode in enumerate(dataset_support.episodes): selected_for_worker = _episode_belongs_to_shard( global_episode_index, shard_index=args.shard_index, num_shards=args.num_shards, ) global_episode_index += 1 if not selected_for_worker: continue shard_path = _shard_path( args.cache_dir, dataset_index, episode_index, ) positions = quantile_anchor_positions( len(episode.anchors), args.grid_size, ) base_indices = [ int(episode.anchors[int(index)]) for index in positions ] if shard_path.exists(): _load_valid_shard( shard_path, identity_sha256=identity["sha256"], dataset_index=dataset_index, episode_index=episode_index, trajectory_id=episode.trajectory_id, positions=positions, base_indices=base_indices, estimator_count=len(args.estimator_seeds), ) completed_assigned_episodes += 1 completed_assigned_anchors += len(positions) continue if model is None: raise RuntimeError("Missing shard discovered after model skip") examples, loaded_base_indices, transform_seeds = _load_examples( source_dataset, episode, positions, dataset_index=dataset_index, episode_index=episode_index, transform_seed_base=args.transform_seed_base, ) if loaded_base_indices != base_indices: raise RuntimeError("Loaded base indices differ from frozen plan") started = time.perf_counter() estimators = _score_examples( model, examples, estimator_seeds=args.estimator_seeds, samples_per_estimator=args.samples_per_estimator, batch_size=args.batch_size, ) torch.cuda.synchronize() elapsed = time.perf_counter() - started session_scoring_seconds += elapsed grid_scores = estimators.mean(axis=1) shard = { "schema_version": SCHEMA_VERSION, "identity_sha256": identity["sha256"], "dataset_index": dataset_index, "dataset_name": dataset_support.name, "episode_index": episode_index, "trajectory_id": _json_scalar(episode.trajectory_id), "anchor_count": len(episode.anchors), "anchor_positions": positions.tolist(), "base_indices": base_indices, "transform_seeds": transform_seeds, "estimator_residuals": estimators.tolist(), "grid_scores": grid_scores.tolist(), "scoring_seconds": elapsed, } _atomic_json(shard_path, shard) new_episodes += 1 completed_assigned_episodes += 1 completed_assigned_anchors += len(positions) if parallel_worker: worker_progress.update( { "completed_episode_count": completed_assigned_episodes, "completed_grid_anchor_count": completed_assigned_anchors, "new_episode_count": new_episodes, "scoring_seconds": session_scoring_seconds, } ) if new_episodes % 25 == 0: _atomic_json(worker_path, worker_progress) print( f"residual-cache shard={args.shard_index}/" f"{args.num_shards} episode=" f"{completed_assigned_episodes}/" f"{assigned_episode_count} seconds={elapsed:.3f}", flush=True, ) else: metadata["completed_episode_count"] += 1 metadata["completed_grid_anchor_count"] += len(positions) metadata["scoring_seconds"] += elapsed _atomic_json(metadata_path, metadata) print( f"residual-cache episode={metadata['completed_episode_count']}/" f"{episode_count} grid_anchors={len(positions)} " f"seconds={elapsed:.3f}", flush=True, ) if ( args.max_new_episodes > 0 and new_episodes >= args.max_new_episodes ): if parallel_worker: worker_progress["status"] = "partial" _atomic_json(worker_path, worker_progress) print("Stopped at max-new-episodes; cache is resumable.", flush=True) return finally: if model is not None: del model torch.cuda.empty_cache() if parallel_worker: worker_progress.update( { "status": "complete", "completed_episode_count": completed_assigned_episodes, "completed_grid_anchor_count": completed_assigned_anchors, "new_episode_count": new_episodes, "scoring_seconds": session_scoring_seconds, } ) _atomic_json(worker_path, worker_progress) if completed_assigned_episodes != assigned_episode_count: raise RuntimeError("Residual cache worker finished with missing episodes") print(json.dumps(worker_progress, indent=2, sort_keys=True), flush=True) return _refresh_progress(metadata, args.cache_dir, manifest) if metadata["completed_episode_count"] != episode_count: metadata["status"] = "in_progress" _atomic_json(metadata_path, metadata) raise RuntimeError("Residual cache is incomplete after scoring loop") proposal = _finalize_proposal( manifest=manifest, cache_dir=args.cache_dir, identity=identity, output=args.output, alpha=args.alpha, epsilon=args.epsilon, allow_failed_gate=args.allow_failed_gate, estimator_count=len(args.estimator_seeds), gate_certificate=gate_certificate, ) metadata["status"] = "complete" metadata["proposal"] = { "path": str(args.output.resolve()), "sha256": proposal.sha256, "importance_weight_min": min( float(np.min(weights)) for dataset_weights in proposal.importance_weights for weights in dataset_weights ), "importance_weight_max": max( float(np.max(weights)) for dataset_weights in proposal.importance_weights for weights in dataset_weights ), } _atomic_json(metadata_path, metadata) print( json.dumps( { "status": "complete", "proposal": metadata["proposal"], "scoring_seconds": metadata["scoring_seconds"], "session_scoring_seconds": session_scoring_seconds, }, indent=2, ), flush=True, ) if __name__ == "__main__": main()