Blame for scripts/sparse_temporal_scoring_ablation.py

#!/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()