#!/usr/bin/env python3 """Build one resumable shard of the official FrameSkip score/pruning cache.""" from __future__ import annotations import argparse import hashlib import json import os import sys import types from pathlib import Path import numpy as np import timm from omegaconf import OmegaConf def stable_shard(dataset_name: str, trajectory_id: int | str, count: int) -> int: key = f"{dataset_name}\0{trajectory_id}".encode() return int.from_bytes(hashlib.sha256(key).digest()[:8], "little") % count def load_completed(path: Path) -> set[tuple[str, str]]: completed: set[tuple[str, str]] = set() if not path.exists(): return completed with path.open(encoding="utf-8") as source: for line_number, line in enumerate(source, 1): line = line.strip() if not line: continue try: record = json.loads(line) except json.JSONDecodeError as exc: raise RuntimeError(f"corrupt cache shard {path}:{line_number}") from exc completed.add((record["dataset_name"], str(record["trajectory_id"]))) return completed def bind_upstream_helpers(dataset, frame_dataset_class, calculator) -> None: for method_name in ( "_extract_action_from_traj", "_get_actions_from_traj_data", "_extract_frames_from_traj", "_load_trajectory_for_importance", ): method = getattr(frame_dataset_class, method_name) setattr(dataset, method_name, types.MethodType(method, dataset)) dataset.importance_calculator = calculator dataset.vac_video_backend = "torchvision_av" dataset.vac_allow_backend_fallback = False dataset.max_vac_frames = 16 dataset._frameskip_runtime_stats = { "vac_frame_extraction_failures": 0, "vac_trajectories_without_frames": 0, } def protected_indices(actions: np.ndarray) -> set[int]: if len(actions) == 0: return set() protected = {0, len(actions) - 1} if len(actions) > 1: if actions.shape[1] >= 7: gripper = actions[:, -1] changes = np.abs(np.diff(gripper, prepend=gripper[0])) protected.update(int(value) for value in np.where(changes > 0.1)[0]) action_delta = np.linalg.norm(np.diff(actions, axis=0), axis=1) threshold = np.percentile(action_delta, 90) protected.update(int(value) for value in np.where(action_delta >= threshold)[0] + 1) return protected def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--project-root", required=True) parser.add_argument("--upstream-root", required=True) parser.add_argument("--config", required=True) parser.add_argument("--required-draws", required=True) parser.add_argument("--visual-checkpoint", required=True) parser.add_argument("--output-dir", required=True) parser.add_argument("--shard-index", type=int, required=True) parser.add_argument("--shard-count", type=int, required=True) parser.add_argument("--device", default="cuda:0") args = parser.parse_args() project_root = Path(args.project_root).resolve() upstream_package = Path(args.upstream_root).resolve() / "starVLA" sys.path.insert(0, str(project_root)) import starVLA # Resolve the official FrameSkip package while retaining the current StarVLA base classes. starVLA.__path__.append(str(upstream_package)) from starVLA.dataloader.lerobot_datasets import get_vla_dataset from starVLA.frameskip.dataloader.frameskip_dataset import FrameSkipSingleDataset from starVLA.frameskip.utils.frame_pruner import FramePruner from starVLA.frameskip.utils.importance_metrics import GripperAwareImportanceCalculator required_payload = json.loads(Path(args.required_draws).read_text(encoding="utf-8")) config = OmegaConf.load(args.config) data_cfg = config.datasets.vla_data data_cfg.traceable_sampling = False data_cfg.valid_anchor_manifest_path = None data_cfg.valid_anchor_manifest_migration_path = None data_cfg.temporal_proposal_path = None mixture = get_vla_dataset( data_cfg=data_cfg, balance_dataset_weights=False, balance_trajectory_weights=False, ) calculator = GripperAwareImportanceCalculator( alpha=0.6, beta=0.2, gamma=0.2, enable_vac=True, visual_encoder_name=None, device=args.device, gripper_boost=2.0, ) # The released scorer preprocesses at 224px, while current timm metadata fixes # DINOv2-S/14 at 518px. Dynamic positional interpolation preserves the released # 224px preprocessing and the official DINOv2 weights. calculator.visual_encoder = timm.create_model( "vit_small_patch14_dinov2.lvd142m", pretrained=False, num_classes=0, checkpoint_path=args.visual_checkpoint, dynamic_img_size=True, ).to(args.device).eval() pruner = FramePruner( compression_ratios=[0.2], preserve_key_stages=True, min_frames_per_trajectory=5, ) datasets = {dataset.dataset_name: dataset for dataset in mixture.datasets} for dataset in datasets.values(): bind_upstream_helpers(dataset, FrameSkipSingleDataset, calculator) output_dir = Path(args.output_dir) output_dir.mkdir(parents=True, exist_ok=True) output_path = output_dir / f"shard_{args.shard_index:02d}_of_{args.shard_count:02d}.jsonl" completed = load_completed(output_path) assignments: list[tuple[str, int | str]] = [] for dataset_name, trajectory_ids in required_payload["datasets"].items(): for trajectory_id in trajectory_ids: if stable_shard(dataset_name, trajectory_id, args.shard_count) == args.shard_index: assignments.append((dataset_name, trajectory_id)) assignments.sort(key=lambda item: (item[0], type(item[1]).__name__, str(item[1]))) written = 0 with output_path.open("a", encoding="utf-8", buffering=1) as destination: for dataset_name, trajectory_id in assignments: key = (dataset_name, str(trajectory_id)) if key in completed: continue dataset = datasets[dataset_name] trajectory = dataset._load_trajectory_for_importance(trajectory_id) actions = np.asarray(trajectory["actions"]) if len(actions) == 0 or not trajectory["observations"]: raise RuntimeError( f"FrameSkip preprocessing has no usable action/video data for {key}" ) scores = calculator.compute_importance(trajectory) if len(scores) != len(actions) or not np.isfinite(scores).all(): raise RuntimeError(f"invalid FrameSkip scores for {key}") result = pruner.prune_trajectory(trajectory, scores, 0.2) keep_indices = np.asarray(result.keep_indices, dtype=np.int64) protected = protected_indices(actions) retained_set = set(int(value) for value in keep_indices) protected_retained = protected.intersection(retained_set) record = { "dataset_name": dataset_name, "trajectory_id": trajectory_id, "trajectory_length": int(len(actions)), "retention_ratio": 0.2, "actual_ratio": float(len(keep_indices) / len(actions)), "keep_indices": keep_indices.tolist(), "importance_scores": np.asarray(scores, dtype=np.float64).tolist(), "protected_indices": sorted(protected), "protected_retained_count": len(protected_retained), "protected_total_count": len(protected), "protected_retention_rate": ( len(protected_retained) / len(protected) if protected else 1.0 ), "vac_frames": len(trajectory["observations"]), } destination.write(json.dumps(record, separators=(",", ":")) + "\n") written += 1 if written % 25 == 0: print( json.dumps({ "shard": args.shard_index, "new": written, "completed": len(completed) + written, "total": len(assignments), }), flush=True, ) stats = { "schema_version": 1, "shard_index": args.shard_index, "shard_count": args.shard_count, "assigned": len(assignments), "completed_before": len(completed), "written": written, "vac_runtime_stats": { name: dataset._frameskip_runtime_stats for name, dataset in datasets.items() }, } stats_path = output_dir / f"shard_{args.shard_index:02d}_of_{args.shard_count:02d}.status.json" stats_path.write_text(json.dumps(stats, indent=2) + "\n", encoding="utf-8") print(json.dumps(stats, indent=2)) if __name__ == "__main__": main()