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