| |
| """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 |
|
|
| |
| 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, |
| ) |
| |
| |
| |
| 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() |
|
|
