| |
| """Validate and merge FrameSkip cache shards into one portable cache.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--required-draws", required=True) |
| parser.add_argument("--shard-dir", required=True) |
| parser.add_argument("--shard-count", type=int, required=True) |
| parser.add_argument("--output", required=True) |
| args = parser.parse_args() |
|
|
| required = json.loads(Path(args.required_draws).read_text(encoding="utf-8")) |
| expected = { |
| (dataset_name, str(trajectory_id)) |
| for dataset_name, trajectory_ids in required["datasets"].items() |
| for trajectory_id in trajectory_ids |
| } |
| records: dict[tuple[str, str], dict] = {} |
| shard_paths = [] |
| for shard_index in range(args.shard_count): |
| path = Path(args.shard_dir) / f"shard_{shard_index:02d}_of_{args.shard_count:02d}.jsonl" |
| if not path.exists(): |
| raise RuntimeError(f"missing FrameSkip cache shard: {path}") |
| shard_paths.append(path.name) |
| with path.open(encoding="utf-8") as source: |
| for line_number, line in enumerate(source, 1): |
| record = json.loads(line) |
| key = (record["dataset_name"], str(record["trajectory_id"])) |
| if key in records: |
| raise RuntimeError(f"duplicate FrameSkip cache key: {key}") |
| records[key] = record |
| if set(records) != expected: |
| missing = sorted(expected - set(records))[:20] |
| extra = sorted(set(records) - expected)[:20] |
| raise RuntimeError(f"cache coverage mismatch: missing={missing}, extra={extra}") |
|
|
| payload = { |
| "schema_version": 1, |
| "method": "FrameSkip-adapted (Gaussian-TPI fallback)", |
| "upstream_commit": "965becf63e9cd5cc339fbf7fde40d1db47efa67a", |
| "retention_ratio": 0.2, |
| "importance": { |
| "calculator": "GripperAwareImportanceCalculator", |
| "avi_weight": 0.6, |
| "vac_weight": 0.2, |
| "tpi_weight": 0.2, |
| "vac_enabled": True, |
| "max_vac_frames": 16, |
| "visual_encoder": "vit_small_patch14_dinov2.lvd142m", |
| "tpi": "Gaussian sigma_squared=0.2 fallback", |
| "gripper_boost": 2.0, |
| }, |
| "pruning": { |
| "first_last": True, |
| "gripper_transition_threshold": 0.1, |
| "top_action_change_percentile": 90, |
| "min_frames_per_trajectory": 5, |
| "mapping": "numpy.searchsorted to the next retained index, clipped at end", |
| }, |
| "schedule": {"pruned_minibatches": 5, "full_minibatches": 1}, |
| "records": [ |
| records[key] |
| for key in sorted(records, key=lambda item: (item[0], item[1])) |
| ], |
| "source_shards": shard_paths, |
| } |
| output = Path(args.output) |
| output.parent.mkdir(parents=True, exist_ok=True) |
| output.write_text(json.dumps(payload, separators=(",", ":")) + "\n", encoding="utf-8") |
| print(json.dumps({ |
| "output": str(output), |
| "records": len(payload["records"]), |
| "actual_ratio_min": min(r["actual_ratio"] for r in payload["records"]), |
| "actual_ratio_max": max(r["actual_ratio"] for r in payload["records"]), |
| "actual_ratio_mean": sum(r["actual_ratio"] for r in payload["records"]) / len(payload["records"]), |
| }, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
