| |
| """Compute dense-panel gradient targets and sparse proposal fidelity.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import csv |
| import hashlib |
| import json |
| import math |
| import os |
| import random |
| import time |
| from datetime import datetime, timezone |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| from omegaconf import OmegaConf |
|
|
| from scripts.validate_stratified_gradient_second_moment import ( |
| _flow_loss, |
| _prepare_conditioning, |
| ) |
| from starVLA.dataloader.lerobot_datasets import get_vla_dataset |
| from starVLA.model.framework.base_framework import build_framework |
| from starVLA.model.framework.share_tools import apply_config_compat |
| from starVLA.scoring.flow_residual import ( |
| randomized_antithetic_stratified_flow_samples, |
| ) |
| from starVLA.scoring.partial_gradient import finite_gradient_l2_norm |
| from starVLA.scoring.residual_cache import ( |
| interpolate_grid_scores, |
| percentile_intensity_masses, |
| quantile_anchor_positions, |
| ) |
| from starVLA.scoring.utility_audit import gradient_second_moment |
|
|
|
|
| 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: |
| encoded = json.dumps( |
| payload, sort_keys=True, separators=(",", ":"), ensure_ascii=True |
| ).encode("utf-8") |
| return hashlib.sha256(encoded).hexdigest() |
|
|
|
|
| def atomic_json(path: Path, payload: dict) -> None: |
| path.parent.mkdir(parents=True, exist_ok=True) |
| temporary = path.with_suffix(path.suffix + ".tmp") |
| temporary.write_text(json.dumps(payload, indent=2), encoding="utf-8") |
| os.replace(temporary, path) |
|
|
|
|
| def seed_all(seed: int) -> None: |
| bounded = int(seed) % (2**32) |
| random.seed(bounded) |
| np.random.seed(bounded) |
| torch.manual_seed(bounded) |
| torch.cuda.manual_seed_all(bounded) |
|
|
|
|
| def load_dense_panel(dense_dir: Path) -> tuple[list[dict], dict]: |
| metadata_path = dense_dir / "dense_residual_metadata.json" |
| scores_path = dense_dir / "dense_residual_scores.npz" |
| metadata = json.loads(metadata_path.read_text(encoding="utf-8")) |
| dense = np.load(scores_path, allow_pickle=False) |
| scores = np.asarray(dense["scores"], dtype=np.float64) |
| offsets = np.asarray(dense["offsets"], dtype=np.int64) |
| records: list[dict] = [] |
| for index in range(len(offsets) - 1): |
| dataset_index = int(dense["dataset_index"][index]) |
| episode_index = int(dense["episode_index"][index]) |
| trajectory_index = int(dense["trajectory_index"][index]) |
| trajectory_id = str(dense["trajectory_id"][index]) |
| shard_path = dense_dir / "dense_shards" / ( |
| f"dataset_{dataset_index:02d}_episode_{episode_index:06d}.json" |
| ) |
| payload = json.loads(shard_path.read_text(encoding="utf-8")) |
| start, stop = int(offsets[index]), int(offsets[index + 1]) |
| expected = { |
| "dataset_index": dataset_index, |
| "episode_index": episode_index, |
| "trajectory_index": trajectory_index, |
| "trajectory_id": trajectory_id, |
| "anchor_count": stop - start, |
| } |
| observed = { |
| "dataset_index": int(payload["dataset_index"]), |
| "episode_index": int(payload["episode_index"]), |
| "trajectory_index": int(payload["trajectory_index"]), |
| "trajectory_id": str(payload["trajectory_id"]), |
| "anchor_count": int(payload["anchor_count"]), |
| } |
| if observed != expected: |
| raise RuntimeError(f"dense trajectory identity mismatch: {shard_path}") |
| if not ( |
| len(payload["anchor_positions"]) |
| == len(payload["base_indices"]) |
| == len(payload["transform_seeds"]) |
| == stop - start |
| ): |
| raise RuntimeError(f"dense anchor identity mismatch: {shard_path}") |
| records.append( |
| { |
| **observed, |
| "trajectory_id": payload["trajectory_id"], |
| "dataset_name": str(payload["dataset_name"]), |
| "manifest_sha256": str(payload["manifest_sha256"]), |
| "global_start": start, |
| "global_stop": stop, |
| "anchor_positions": [int(x) for x in payload["anchor_positions"]], |
| "base_indices": [int(x) for x in payload["base_indices"]], |
| "transform_seeds": [int(x) for x in payload["transform_seeds"]], |
| "dense_scores": scores[start:stop].tolist(), |
| } |
| ) |
| records.sort(key=lambda item: (item["dataset_index"], item["trajectory_index"])) |
| global_offset = 0 |
| for record in records: |
| record["global_start"] = global_offset |
| global_offset += record["anchor_count"] |
| record["global_stop"] = global_offset |
| if len(records) != int(metadata["trajectory_count"]): |
| raise RuntimeError("dense trajectory count differs from metadata") |
| if global_offset != int(metadata["valid_anchor_count"]): |
| raise RuntimeError("dense anchor count differs from metadata") |
| manifests = {record["manifest_sha256"] for record in records} |
| if len(manifests) != 1: |
| raise RuntimeError("dense shards contain multiple manifest identities") |
| identity = { |
| "dense_metadata_sha256": sha256_file(metadata_path), |
| "dense_scores_sha256": sha256_file(scores_path), |
| "dense_manifest_sha256": next(iter(manifests)), |
| "trajectory_count": len(records), |
| "anchor_count": global_offset, |
| "trajectory_keys_sha256": canonical_sha256( |
| { |
| "keys": [ |
| [ |
| record["dataset_index"], |
| record["dataset_name"], |
| record["trajectory_index"], |
| record["trajectory_id"], |
| record["episode_index"], |
| ] |
| for record in records |
| ] |
| } |
| ), |
| "anchor_keys_sha256": canonical_sha256( |
| { |
| "keys": [ |
| [ |
| record["dataset_index"], |
| record["trajectory_id"], |
| position, |
| base_index, |
| transform_seed, |
| record["global_start"] + local_index, |
| ] |
| for record in records |
| for local_index, (position, base_index, transform_seed) in enumerate( |
| zip( |
| record["anchor_positions"], |
| record["base_indices"], |
| record["transform_seeds"], |
| ) |
| ) |
| ] |
| } |
| ), |
| } |
| return records, identity |
|
|
|
|
| def load_runtime(config_path: Path, checkpoint_path: Path, output: Path): |
| cfg = apply_config_compat(OmegaConf.load(config_path)) |
| cfg.output_dir = str(output / "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) |
| checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True) |
| model.load_state_dict(checkpoint, strict=True) |
| model.to(device="cuda", dtype=torch.float32) |
| return cfg, dataset, model |
|
|
|
|
| def measure_trajectory( |
| model, |
| source, |
| record: dict, |
| *, |
| estimator_seeds: list[int], |
| samples_per_estimator: int, |
| ) -> tuple[np.ndarray, np.ndarray, float]: |
| anchor_count = record["anchor_count"] |
| estimator_values = np.empty( |
| (anchor_count, len(estimator_seeds)), dtype=np.float64 |
| ) |
| draw_norms = np.empty( |
| (anchor_count, len(estimator_seeds), samples_per_estimator), |
| dtype=np.float64, |
| ) |
| started = time.perf_counter() |
| for local_index, (base_index, transform_seed) in enumerate( |
| zip(record["base_indices"], record["transform_seeds"]) |
| ): |
| seed_all(transform_seed) |
| raw = source.get_step_data(record["trajectory_id"], base_index) |
| transformed = source.transforms(raw) |
| example = source._pack_sample(transformed) |
| global_index = record["global_start"] + local_index |
| for estimator_index, estimator_seed in enumerate(estimator_seeds): |
| model.requires_grad_(False) |
| model.train() |
| model.qwen_vl_interface.eval() |
| seed_all(estimator_seed * 1_000_003 + global_index) |
| conditioning = _prepare_conditioning(model, example) |
| actions = conditioning[1] |
| action_cfg = model.action_model.config |
| noises, times = randomized_antithetic_stratified_flow_samples( |
| actions, |
| sample_count=samples_per_estimator, |
| seed=estimator_seed * 1_000_033 + global_index, |
| beta_alpha=float(action_cfg.noise_beta_alpha), |
| beta_beta=float(action_cfg.noise_beta_beta), |
| noise_s=float(action_cfg.noise_s), |
| ) |
| model.action_model.requires_grad_(True) |
| parameters = list(model.action_model.parameters()) |
| for draw_index in range(samples_per_estimator): |
| seed_all( |
| estimator_seed * 1_000_037 |
| + global_index * 101 |
| + draw_index // 2 |
| ) |
| model.zero_grad(set_to_none=True) |
| loss = _flow_loss( |
| model, conditioning, noises[draw_index], times[draw_index] |
| ) |
| if loss.ndim != 0 or not torch.isfinite(loss): |
| raise FloatingPointError("flow loss must be a finite scalar") |
| loss.backward() |
| draw_norms[local_index, estimator_index, draw_index] = ( |
| finite_gradient_l2_norm(parameters) |
| ) |
| estimator_values[local_index, estimator_index] = math.sqrt( |
| float(np.mean(np.square(draw_norms[local_index, estimator_index]))) |
| ) |
| model.zero_grad(set_to_none=True) |
| torch.cuda.synchronize() |
| return estimator_values, draw_norms, time.perf_counter() - started |
|
|
|
|
| def worker(args: argparse.Namespace) -> None: |
| if not 0 <= args.shard_index < args.num_shards: |
| raise ValueError("invalid shard index") |
| records, dense_identity = load_dense_panel(args.dense_dir) |
| cfg_sha = sha256_file(args.config) |
| checkpoint_sha = sha256_file(args.checkpoint) |
| worker_identity = { |
| "schema_version": SCHEMA_VERSION, |
| "config_sha256": cfg_sha, |
| "checkpoint_sha256": checkpoint_sha, |
| "dense_identity": dense_identity, |
| "target_estimator_seeds": args.target_seeds, |
| "samples_per_estimator": args.samples_per_estimator, |
| "gradient_definition": "sqrt(mean over estimator seeds and draws of squared full action-model gradient L2 norm)", |
| "parameter_mode": "full_action_model", |
| "frozen_backbone_eval": True, |
| "action_model_mode": "train", |
| "model_dtype": "float32", |
| "global_anchor_index_seed_semantics": True, |
| "num_shards": args.num_shards, |
| } |
| identity_sha = canonical_sha256(worker_identity) |
| _, dataset, model = load_runtime(args.config, args.checkpoint, args.output) |
| if dataset.valid_anchor_manifest is None: |
| raise RuntimeError("current dataset has no valid-anchor manifest") |
| if dataset.valid_anchor_manifest.sha256 != dense_identity["dense_manifest_sha256"]: |
| raise RuntimeError("current dataset manifest differs from dense panel") |
| assigned = records[args.shard_index :: args.num_shards] |
| cache_dir = args.output / "dense_gradient_targets" |
| started = time.perf_counter() |
| new_anchors = 0 |
| completed_anchors = 0 |
| for item_index, record in enumerate(assigned): |
| key = ( |
| f"dataset_{record['dataset_index']:02d}_" |
| f"episode_{record['episode_index']:06d}.json" |
| ) |
| path = cache_dir / key |
| if path.exists(): |
| cached = json.loads(path.read_text(encoding="utf-8")) |
| if ( |
| cached.get("identity_sha256") != identity_sha |
| or cached.get("trajectory_identity") |
| != {field: record[field] for field in ( |
| "dataset_index", "dataset_name", "episode_index", |
| "trajectory_index", "trajectory_id", "anchor_count", |
| "global_start", "global_stop", "anchor_positions", |
| "base_indices", "transform_seeds", |
| )} |
| ): |
| raise RuntimeError(f"incompatible existing cache: {path}") |
| completed_anchors += record["anchor_count"] |
| continue |
| source = dataset.datasets[record["dataset_index"]] |
| estimators, draws, elapsed = measure_trajectory( |
| model, |
| source, |
| record, |
| estimator_seeds=args.target_seeds, |
| samples_per_estimator=args.samples_per_estimator, |
| ) |
| target = np.sqrt(np.mean(np.square(estimators), axis=1)) |
| if not np.isfinite(target).all() or np.any(target <= 0): |
| raise FloatingPointError("gradient targets must be positive and finite") |
| trajectory_identity = {field: record[field] for field in ( |
| "dataset_index", "dataset_name", "episode_index", |
| "trajectory_index", "trajectory_id", "anchor_count", |
| "global_start", "global_stop", "anchor_positions", |
| "base_indices", "transform_seeds", |
| )} |
| atomic_json( |
| path, |
| { |
| "schema_version": SCHEMA_VERSION, |
| "status": "complete", |
| "identity_sha256": identity_sha, |
| "worker_identity": worker_identity, |
| "trajectory_identity": trajectory_identity, |
| "target_estimators": estimators.tolist(), |
| "target_draw_norms": draws.tolist(), |
| "target_gradient_rms": target.tolist(), |
| "scoring_seconds": elapsed, |
| "completed_at": datetime.now(timezone.utc).isoformat(), |
| }, |
| ) |
| new_anchors += record["anchor_count"] |
| completed_anchors += record["anchor_count"] |
| print( |
| f"shard={args.shard_index}/{args.num_shards} " |
| f"trajectory={item_index + 1}/{len(assigned)} " |
| f"anchors={completed_anchors} new={new_anchors} seconds={elapsed:.2f}", |
| flush=True, |
| ) |
| atomic_json( |
| args.output / "logs" / f"worker_{args.shard_index:02d}_of_{args.num_shards:02d}.json", |
| { |
| **worker_identity, |
| "identity_sha256": identity_sha, |
| "status": "complete", |
| "shard_index": args.shard_index, |
| "assigned_trajectories": len(assigned), |
| "completed_anchors": completed_anchors, |
| "new_anchors": new_anchors, |
| "wall_clock_seconds": time.perf_counter() - started, |
| "completed_at": datetime.now(timezone.utc).isoformat(), |
| }, |
| ) |
|
|
|
|
| def q_from_scores(scores: np.ndarray, alpha: float, epsilon: float) -> np.ndarray: |
| masses = percentile_intensity_masses(scores, epsilon=epsilon) |
| baseline = np.full(len(scores), 1.0 / len(scores), dtype=np.float64) |
| information = masses / masses.sum() |
| return (1.0 - alpha) * baseline + alpha * information |
|
|
|
|
| def describe(values: list[float]) -> dict: |
| array = np.asarray(values, dtype=np.float64) |
| return { |
| "mean": float(array.mean()), |
| "median": float(np.median(array)), |
| "std": float(array.std(ddof=1)), |
| "p05": float(np.quantile(array, 0.05)), |
| "p25": float(np.quantile(array, 0.25)), |
| "p75": float(np.quantile(array, 0.75)), |
| "p95": float(np.quantile(array, 0.95)), |
| "min": float(array.min()), |
| "max": float(array.max()), |
| } |
|
|
|
|
| def write_csv(path: Path, rows: list[dict]) -> None: |
| path.parent.mkdir(parents=True, exist_ok=True) |
| with path.open("w", newline="", encoding="utf-8") as handle: |
| writer = csv.DictWriter(handle, fieldnames=list(rows[0])) |
| writer.writeheader() |
| writer.writerows(rows) |
|
|
|
|
| def finalize(args: argparse.Namespace) -> None: |
| records, dense_identity = load_dense_panel(args.dense_dir) |
| cache_dir = args.output / "dense_gradient_targets" |
| cache_files = sorted(cache_dir.glob("dataset_*_episode_*.json")) |
| if len(cache_files) != len(records): |
| raise RuntimeError( |
| f"gradient cache file count {len(cache_files)} != {len(records)}" |
| ) |
| per_trajectory: list[dict] = [] |
| all_targets = np.empty(dense_identity["anchor_count"], dtype=np.float64) |
| cache_identity = None |
| q_checks = [] |
| manual_checks = [] |
| for record_index, record in enumerate(records): |
| path = cache_dir / ( |
| f"dataset_{record['dataset_index']:02d}_" |
| f"episode_{record['episode_index']:06d}.json" |
| ) |
| if not path.exists(): |
| raise FileNotFoundError(path) |
| payload = json.loads(path.read_text(encoding="utf-8")) |
| if payload.get("status") != "complete": |
| raise RuntimeError(f"incomplete cache: {path}") |
| if cache_identity is None: |
| cache_identity = payload["identity_sha256"] |
| if payload["identity_sha256"] != cache_identity: |
| raise RuntimeError("gradient caches contain multiple identities") |
| observed = payload["trajectory_identity"] |
| expected = {field: record[field] for field in ( |
| "dataset_index", "dataset_name", "episode_index", |
| "trajectory_index", "trajectory_id", "anchor_count", |
| "global_start", "global_stop", "anchor_positions", |
| "base_indices", "transform_seeds", |
| )} |
| if observed != expected: |
| raise RuntimeError(f"gradient anchor alignment mismatch: {path}") |
| h = np.asarray(payload["target_gradient_rms"], dtype=np.float64) |
| if len(h) != record["anchor_count"]: |
| raise RuntimeError(f"gradient anchor count mismatch: {path}") |
| if not np.isfinite(h).all() or np.any(h <= 0): |
| raise FloatingPointError(f"invalid gradient target: {path}") |
| all_targets[record["global_start"] : record["global_stop"]] = h |
| dense_scores = np.asarray(record["dense_scores"], dtype=np.float64) |
| proposals: dict[str, np.ndarray] = { |
| "p": np.full(len(h), 1.0 / len(h), dtype=np.float64), |
| "dense": q_from_scores(dense_scores, args.alpha, args.epsilon), |
| } |
| for k in (4, 8, 16): |
| positions = quantile_anchor_positions(len(h), k) |
| estimate = interpolate_grid_scores( |
| len(h), positions, dense_scores[positions] |
| ) |
| proposals[str(k)] = q_from_scores(estimate, args.alpha, args.epsilon) |
| for label, proposal in proposals.items(): |
| q_checks.append( |
| abs(float(proposal.sum()) - 1.0) <= 1e-10 |
| and np.isfinite(proposal).all() |
| and np.all(proposal > 0) |
| ) |
| moments = { |
| label: gradient_second_moment(proposal, h) |
| for label, proposal in proposals.items() |
| } |
| if record_index in {0, len(records) // 2, len(records) - 1}: |
| direct = float( |
| np.sum(np.square(proposals["p"] * h) / proposals["dense"]) |
| ) |
| manual_checks.append(abs(direct - moments["dense"]) <= 1e-12) |
| row = { |
| "dataset_index": record["dataset_index"], |
| "dataset_name": record["dataset_name"], |
| "trajectory_index": record["trajectory_index"], |
| "trajectory_id": record["trajectory_id"], |
| "anchor_count": record["anchor_count"], |
| "M_p": moments["p"], |
| "M_q4": moments["4"], |
| "M_q8": moments["8"], |
| "M_q16": moments["16"], |
| "M_qDense": moments["dense"], |
| } |
| for label in ("4", "8", "16"): |
| row[f"R_{label}"] = moments[label] / moments["dense"] |
| row[f"q{label}_over_p"] = moments[label] / moments["p"] |
| row["qDense_over_p"] = moments["dense"] / moments["p"] |
| per_trajectory.append(row) |
| if not np.isfinite(all_targets).all() or np.any(all_targets <= 0): |
| raise FloatingPointError("assembled gradient targets are incomplete") |
| aggregate = { |
| field: float(np.mean([row[field] for row in per_trajectory])) |
| for field in ("M_p", "M_q4", "M_q8", "M_q16", "M_qDense") |
| } |
| ratios = { |
| "R_4": aggregate["M_q4"] / aggregate["M_qDense"], |
| "R_8": aggregate["M_q8"] / aggregate["M_qDense"], |
| "R_16": aggregate["M_q16"] / aggregate["M_qDense"], |
| "q4_over_p": aggregate["M_q4"] / aggregate["M_p"], |
| "q8_over_p": aggregate["M_q8"] / aggregate["M_p"], |
| "q16_over_p": aggregate["M_q16"] / aggregate["M_p"], |
| "qDense_over_p": aggregate["M_qDense"] / aggregate["M_p"], |
| } |
| trajectory_statistics = {} |
| for label in ("4", "8", "16"): |
| r_values = [row[f"R_{label}"] for row in per_trajectory] |
| p_values = [row[f"q{label}_over_p"] for row in per_trajectory] |
| trajectory_statistics[label] = { |
| "relative_to_dense": describe(r_values), |
| "relative_to_uniform": describe(p_values), |
| "better_than_dense": sum(value < 1.0 - args.comparison_tolerance for value in r_values), |
| "equal_to_dense": sum(abs(value - 1.0) <= args.comparison_tolerance for value in r_values), |
| "worse_than_dense": sum(value > 1.0 + args.comparison_tolerance for value in r_values), |
| "better_than_uniform": sum(value < 1.0 - args.comparison_tolerance for value in p_values), |
| "equal_to_uniform": sum(abs(value - 1.0) <= args.comparison_tolerance for value in p_values), |
| "worse_than_uniform": sum(value > 1.0 + args.comparison_tolerance for value in p_values), |
| } |
| dense_self_ratio = aggregate["M_qDense"] / aggregate["M_qDense"] |
| checks = { |
| "trajectory_count_512": len(per_trajectory) == 512, |
| "anchor_count_19375": len(all_targets) == 19375, |
| "cache_file_count_matches_trajectories": len(cache_files) == len(records), |
| "trajectory_and_anchor_identity_aligned": True, |
| "all_anchor_targets_positive_finite": bool( |
| np.isfinite(all_targets).all() and np.all(all_targets > 0) |
| ), |
| "proposal_normalization_full_support_finite": all(q_checks), |
| "dense_self_ratio_one": abs(dense_self_ratio - 1.0) <= 1e-15, |
| "manual_second_moment_checks": all(manual_checks), |
| "all_moments_finite": all(np.isfinite(list(aggregate.values()))), |
| } |
| write_csv(args.output / "optimization_fidelity_per_trajectory.csv", per_trajectory) |
| readback = list( |
| csv.DictReader( |
| (args.output / "optimization_fidelity_per_trajectory.csv").open( |
| encoding="utf-8" |
| ) |
| ) |
| ) |
| reconstructed = { |
| field: float(np.mean([float(row[field]) for row in readback])) |
| for field in aggregate |
| } |
| checks["aggregate_reconstructs_from_per_trajectory_csv"] = all( |
| math.isclose( |
| reconstructed[field], aggregate[field], rel_tol=1e-12, abs_tol=1e-12 |
| ) |
| for field in aggregate |
| ) |
| summary_rows = [] |
| for label, field in (("Uniform", "M_p"), ("K4", "M_q4"), ("K8", "M_q8"), ("K16", "M_q16"), ("Dense", "M_qDense")): |
| value = aggregate[field] |
| summary_rows.append( |
| { |
| "proposal": label, |
| "aggregate_second_moment": value, |
| "M_over_Mp": value / aggregate["M_p"], |
| "improvement_vs_uniform_percent": 100.0 * (1.0 - value / aggregate["M_p"]), |
| "M_over_Mdense": value / aggregate["M_qDense"], |
| "gap_vs_dense_percent": 100.0 * (value / aggregate["M_qDense"] - 1.0), |
| } |
| ) |
| write_csv(args.output / "optimization_fidelity_summary.csv", summary_rows) |
| np.savez_compressed( |
| args.output / "dense_gradient_targets.npz", |
| target_gradient_rms=all_targets, |
| ) |
| metadata = { |
| "schema_version": SCHEMA_VERSION, |
| "status": "complete" if all(checks.values()) else "failed_checks", |
| "started_at": args.started_at, |
| "completed_at": datetime.now(timezone.utc).isoformat(), |
| "git_commit": None, |
| "git_status": "unavailable_server_copy_has_no_dot_git", |
| "code_path": str(Path(__file__).resolve()), |
| "config": str(args.config.resolve()), |
| "config_sha256": sha256_file(args.config), |
| "checkpoint": str(args.checkpoint.resolve()), |
| "checkpoint_sha256": sha256_file(args.checkpoint), |
| "dense_dir": str(args.dense_dir.resolve()), |
| "dense_identity": dense_identity, |
| "gradient_cache_identity_sha256": cache_identity, |
| "trajectory_count": len(per_trajectory), |
| "anchor_count": len(all_targets), |
| "proposal": { |
| "alpha": args.alpha, |
| "epsilon": args.epsilon, |
| "construction": "within-trajectory percentile-rank safe mixture after linear interpolation at deterministic quantile anchors", |
| "K": [4, 8, 16], |
| }, |
| "target_estimator": { |
| "definition": "sqrt(mean over estimator seeds and draws of squared full action-model gradient L2 norm)", |
| "seeds": args.target_seeds, |
| "samples_per_estimator": args.samples_per_estimator, |
| "randomized_beta_stratification": True, |
| "antithetic_gaussian_noise": True, |
| "common_dropout_within_antithetic_pair": True, |
| "frozen_backbone_eval": True, |
| "action_model_mode": "train", |
| "model_dtype": "float32", |
| }, |
| "aggregation": "equal empirical trajectory weight; mean per-trajectory second moment, matching hierarchical_alignment_audit", |
| "aggregate_second_moments": aggregate, |
| "aggregate_ratios": ratios, |
| "trajectory_statistics": trajectory_statistics, |
| "sanity_checks": checks, |
| } |
| atomic_json(args.output / "optimization_fidelity_metadata.json", metadata) |
| if not all(checks.values()): |
| raise RuntimeError(f"optimization fidelity checks failed: {checks}") |
| readme = f"""# Optimization Fidelity on the Dense Temporal Panel |
| |
| - Status: complete |
| - Panel: {len(per_trajectory)} trajectories, {len(all_targets)} valid anchors |
| - Checkpoint: `{args.checkpoint.resolve()}` |
| - Gradient target: full action-model gradient RMS, four estimator seeds and eight draws per seed |
| - Proposal: within-trajectory percentile safe mixture, alpha={args.alpha}, epsilon={args.epsilon} |
| - Aggregation: equal empirical trajectory weight, matching `hierarchical_alignment_audit` |
| |
| This is a 512-trajectory Dense-reference offline audit. It is not a full-manifest |
| training proposal, a continuation run, or a downstream success-rate evaluation. |
| |
| Formal values are stored in `optimization_fidelity_summary.csv` and |
| `optimization_fidelity_metadata.json`. All required sanity checks passed. |
| """ |
| (args.output / "README.md").write_text(readme, encoding="utf-8") |
| print(json.dumps(metadata, indent=2), flush=True) |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--dense-dir", type=Path, required=True) |
| parser.add_argument("--output", type=Path, required=True) |
| parser.add_argument("--config", type=Path, required=True) |
| parser.add_argument("--checkpoint", type=Path, required=True) |
| parser.add_argument( |
| "--target-seeds", type=int, nargs="+", default=[27300, 27301, 27302, 27303] |
| ) |
| parser.add_argument("--samples-per-estimator", type=int, default=8) |
| subparsers = parser.add_subparsers(dest="command", required=True) |
| worker_parser = subparsers.add_parser("worker") |
| worker_parser.add_argument("--shard-index", type=int, required=True) |
| worker_parser.add_argument("--num-shards", type=int, required=True) |
| worker_parser.set_defaults(func=worker) |
| final_parser = subparsers.add_parser("finalize") |
| final_parser.add_argument("--alpha", type=float, default=0.2) |
| final_parser.add_argument("--epsilon", type=float, default=0.1) |
| final_parser.add_argument("--comparison-tolerance", type=float, default=1e-12) |
| final_parser.add_argument("--started-at", type=str, required=True) |
| final_parser.set_defaults(func=finalize) |
| args = parser.parse_args() |
| if len(args.target_seeds) != 4 or len(set(args.target_seeds)) != 4: |
| raise ValueError("exactly four unique target estimator seeds are required") |
| if args.samples_per_estimator != 8: |
| raise ValueError("the Gate42F contract requires eight samples per estimator") |
| args.func(args) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
