ReTime / scripts /build_residual_temporal_proposal.py
Add Florence-Bridge code and handoff documentation
977b06e verified 7 days ago
29.3 kB
"""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()