File size: 4,066 Bytes
5cd1287 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 | #!/usr/bin/env python3
"""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()
|