| |
| """Verify that a cache covers an unchanged range of uniform logical draws.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
| from starVLA.dataloader.infoflow_sampler import LogicalDraw, ValidAnchorManifest |
|
|
| from external_baselines.keyframe_focused_adapted.keyframe_cache import ( |
| KeyframeWeightCache, |
| ) |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--manifest", type=Path, required=True) |
| parser.add_argument("--cache", type=Path, required=True) |
| parser.add_argument("--seed", type=int) |
| parser.add_argument("--draw-start", type=int) |
| parser.add_argument("--draw-count", type=int) |
| parser.add_argument("--full-support", action="store_true") |
| args = parser.parse_args() |
|
|
| manifest = ValidAnchorManifest.load(args.manifest) |
| cache = KeyframeWeightCache(args.cache) |
| if args.full_support: |
| arrays = np.load(args.cache, allow_pickle=False) |
| metadata = json.loads(str(arrays["metadata_json"])) |
| counts: dict[str, int] = {} |
| weights: dict[str, int] = {} |
| for dataset_index, dataset in enumerate(manifest.datasets): |
| prefix = f"d{dataset_index}_" |
| trajectory_ids = arrays[prefix + "trajectory_ids"] |
| offsets = arrays[prefix + "offsets"] |
| anchors = arrays[prefix + "anchors"] |
| cached_weights = arrays[prefix + "weights"] |
| expected_trajectory_ids = np.asarray( |
| [int(episode.trajectory_id) for episode in dataset.episodes], |
| dtype=np.int64, |
| ) |
| if not np.array_equal(trajectory_ids, expected_trajectory_ids): |
| raise RuntimeError( |
| f"cache trajectory support does not match {dataset.name}" |
| ) |
| for position, episode in enumerate(dataset.episodes): |
| cached = anchors[int(offsets[position]) : int(offsets[position + 1])] |
| expected = np.asarray(episode.anchors, dtype=np.int32) |
| if not np.array_equal(cached, expected): |
| raise RuntimeError( |
| f"cache anchor support does not match " |
| f"{dataset.name}/{int(episode.trajectory_id)}" |
| ) |
| counts[dataset.name] = int(len(anchors)) |
| values, value_counts = np.unique(cached_weights, return_counts=True) |
| for value, count in zip(values, value_counts, strict=True): |
| label = str(float(value)) |
| weights[label] = weights.get(label, 0) + int(count) |
| print( |
| json.dumps( |
| { |
| "status": "PASS", |
| "mode": "full-support", |
| "cache": cache.path.name, |
| "dataset_anchor_counts": counts, |
| "weight_counts": weights, |
| }, |
| sort_keys=True, |
| ) |
| ) |
| return |
| if args.seed is None or args.draw_start is None or args.draw_count is None: |
| parser.error("draw-range validation requires --seed, --draw-start, and --draw-count") |
| counts: dict[str, int] = {} |
| weights: dict[str, int] = {} |
| for draw_id in range(args.draw_start, args.draw_start + args.draw_count): |
| trace = manifest.sample(args.seed, LogicalDraw(0, draw_id)) |
| record = cache.lookup( |
| trace.dataset_name, trace.trajectory_id, trace.base_index |
| ) |
| counts[trace.dataset_name] = counts.get(trace.dataset_name, 0) + 1 |
| label = str(record.weight) |
| weights[label] = weights.get(label, 0) + 1 |
|
|
| print( |
| json.dumps( |
| { |
| "status": "PASS", |
| "cache": cache.path.name, |
| "draw_start": args.draw_start, |
| "draw_count": args.draw_count, |
| "dataset_counts": counts, |
| "weight_counts": weights, |
| }, |
| sort_keys=True, |
| ) |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
