#!/usr/bin/env python3 """Build a RoboTwin FrameSkip cache with the locked upstream scoring code.""" from __future__ import annotations import argparse import importlib.util import json import sys from pathlib import Path import numpy as np import pyarrow.parquet as parquet import timm from PIL import Image from starVLA.dataloader.gr00t_lerobot.video import get_frames_by_timestamps from starVLA.dataloader.infoflow_sampler import ValidAnchorManifest UPSTREAM_COMMIT = "965becf63e9cd5cc339fbf7fde40d1db47efa67a" METHOD = "FrameSkip-adapted (Gaussian-TPI fallback)" def load_upstream_module(name: str, path: Path): spec = importlib.util.spec_from_file_location(name, path) if spec is None or spec.loader is None: raise RuntimeError(f"cannot load upstream module {path}") module = importlib.util.module_from_spec(spec) sys.modules[name] = module spec.loader.exec_module(module) return module def trajectory_paths(root: Path, dataset_name: str, trajectory_id: int) -> tuple[Path, Path]: dataset = root / dataset_name chunk = trajectory_id // 1000 parquet_path = dataset / "data" / f"chunk-{chunk:03d}" / f"episode_{trajectory_id:06d}.parquet" video_path = ( dataset / "videos" / f"chunk-{chunk:03d}" / "observation.images.cam_high" / f"episode_{trajectory_id:06d}.mp4" ) return parquet_path, video_path def load_trajectory( root: Path, dataset_name: str, trajectory_id: int, max_vac_frames: int, ) -> dict: parquet_path, video_path = trajectory_paths(root, dataset_name, trajectory_id) table = parquet.read_table(parquet_path, columns=["action", "timestamp"]) actions = np.asarray(table["action"].to_pylist(), dtype=np.float32) timestamps = np.asarray(table["timestamp"].to_pylist(), dtype=np.float64).reshape(-1) if actions.ndim != 2 or actions.shape[1] != 14: raise RuntimeError(f"unexpected RoboTwin action shape {actions.shape}: {parquet_path}") if len(actions) != len(timestamps): raise RuntimeError(f"action/timestamp length mismatch: {parquet_path}") sample_indices = np.linspace( 0, len(timestamps) - 1, min(max_vac_frames, len(timestamps)), dtype=int ) frames = get_frames_by_timestamps( video_path.as_posix(), timestamps[sample_indices], video_backend="pyav", video_backend_kwargs={}, ) observations = [] for frame in frames: array = np.asarray(frame) if array.ndim == 3 and array.shape[0] in (1, 3, 4) and array.shape[-1] not in (1, 3, 4): array = np.moveaxis(array, 0, -1) observations.append(Image.fromarray(array.astype(np.uint8))) return {"observations": observations, "actions": actions} def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--manifest", required=True, type=Path) parser.add_argument("--data-root", required=True, type=Path) parser.add_argument("--upstream-root", required=True, type=Path) parser.add_argument("--visual-checkpoint", required=True, type=Path) parser.add_argument("--output", required=True, type=Path) parser.add_argument("--device", default="cuda") parser.add_argument("--max-trajectories", type=int) args = parser.parse_args() candidates = ( args.upstream_root / "starVLA" / "starVLA" / "frameskip" / "utils", args.upstream_root / "starVLA" / "frameskip" / "utils", ) utils = next( (path for path in candidates if (path / "importance_metrics.py").is_file()), None, ) if utils is None: checked = ", ".join(str(path) for path in candidates) raise FileNotFoundError(f"FrameSkip utils not found; checked: {checked}") importance_module = load_upstream_module("frameskip_upstream_importance", utils / "importance_metrics.py") pruner_module = load_upstream_module("frameskip_upstream_pruner", utils / "frame_pruner.py") calculator = importance_module.GripperAwareImportanceCalculator( alpha=0.6, beta=0.2, gamma=0.2, enable_vac=True, visual_encoder_name=None, device=args.device, gripper_boost=2.0, ) # Match the verified Bridge cache builder: retain upstream 224px VAC # preprocessing and interpolate DINOv2 positional embeddings dynamically. calculator.visual_encoder = timm.create_model( "vit_small_patch14_dinov2.lvd142m", pretrained=False, num_classes=0, checkpoint_path=str(args.visual_checkpoint), dynamic_img_size=True, ).to(args.device).eval() pruner = pruner_module.FramePruner( compression_ratios=[0.2], preserve_key_stages=True, min_frames_per_trajectory=5, ) manifest = ValidAnchorManifest.load(args.manifest) records = [] for dataset in manifest.datasets: for episode in dataset.episodes: trajectory_id = int(episode.trajectory_id) trajectory = load_trajectory( args.data_root, dataset.name, trajectory_id, max_vac_frames=16 ) scores = calculator.compute_importance(trajectory) result = pruner.prune_trajectory(trajectory, scores, 0.2) protected = pruner._detect_key_stages(trajectory, scores) retained = set(int(x) for x in result.keep_indices) protected_set = set(int(x) for x in protected) records.append( { "dataset_name": dataset.name, "trajectory_id": trajectory_id, "trajectory_length": int(len(trajectory["actions"])), "retention_ratio": 0.2, "actual_ratio": float(result.actual_ratio), "keep_indices": [int(x) for x in result.keep_indices], "importance_scores": [float(x) for x in scores], "protected_indices": sorted(protected_set), "protected_retained_count": len(protected_set & retained), "protected_total_count": len(protected_set), "protected_retention_rate": ( len(protected_set & retained) / len(protected_set) if protected_set else 1.0 ), "vac_frames": len(trajectory["observations"]), } ) print( f"cached {dataset.name}/{trajectory_id}: " f"{len(result.keep_indices)}/{len(trajectory['actions'])}", flush=True, ) if args.max_trajectories is not None and len(records) >= args.max_trajectories: break if args.max_trajectories is not None and len(records) >= args.max_trajectories: break payload = { "schema_version": 1, "method": METHOD, "upstream_commit": UPSTREAM_COMMIT, "required_trajectory_count": sum( len(dataset.episodes) for dataset in manifest.datasets ), "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, "source_shards": [], } args.output.parent.mkdir(parents=True, exist_ok=True) temporary = args.output.with_suffix(args.output.suffix + ".tmp") temporary.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8") temporary.replace(args.output) print(f"wrote {args.output} records={len(records)}") if __name__ == "__main__": main()