File size: 4,129 Bytes
977b06e 5cd1287 977b06e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | #!/usr/bin/env python3
"""Thin runtime adapter for the locked Florence-Bridge FrameSkip continuation."""
from __future__ import annotations
import argparse
import json
import os
from dataclasses import asdict
from pathlib import Path
import numpy as np
from omegaconf import OmegaConf
def load_cache(path: str) -> tuple[dict[tuple[str, str], np.ndarray], dict]:
payload = json.loads(Path(path).read_text(encoding="utf-8"))
if payload.get("method") != "FrameSkip-adapted (Gaussian-TPI fallback)":
raise RuntimeError("unexpected FrameSkip cache method")
cache = {}
for record in payload["records"]:
key = (record["dataset_name"], str(record["trajectory_id"]))
keep = np.asarray(record["keep_indices"], dtype=np.int64)
if len(keep) == 0 or np.any(np.diff(keep) <= 0):
raise RuntimeError(f"invalid retained indices for {key}")
cache[key] = keep
return cache, payload
def install_adapter(cache_path: str, start_global_draw_id: int) -> None:
from starVLA.dataloader.gr00t_lerobot.datasets import LeRobotMixtureDataset
cache, cache_payload = load_cache(cache_path)
original_getitem = LeRobotMixtureDataset.__getitem__
global_microbatch_size = 8
period = 6
def frameskip_getitem(self, index):
trace = self.sample_trace(index)
offset = trace.global_draw_id - start_global_draw_id
if offset < 0:
raise RuntimeError("FrameSkip draw precedes the certified continuation start")
microbatch = offset // global_microbatch_size
is_pruned = microbatch % period < 5
logical_index = int(trace.base_index)
actual_index = logical_index
if is_pruned:
key = (trace.dataset_name, str(trace.trajectory_id))
if key not in cache:
raise RuntimeError(f"FrameSkip cache is missing required trajectory {key}")
retained = cache[key]
position = int(np.searchsorted(retained, logical_index))
position = int(np.clip(position, 0, len(retained) - 1))
actual_index = int(retained[position])
dataset = self.datasets[trace.dataset_index]
original_get_step_data = dataset.get_step_data
def redirected_get_step_data(trajectory_id, base_index):
if str(trajectory_id) != str(trace.trajectory_id) or int(base_index) != logical_index:
raise RuntimeError("unexpected nested dataset read during FrameSkip remapping")
return original_get_step_data(trajectory_id, actual_index)
dataset.get_step_data = redirected_get_step_data
try:
sample = original_getitem(self, index)
finally:
dataset.get_step_data = original_get_step_data
sample["sampling_trace"]["frameskip"] = {
"method": cache_payload["method"],
"cache_record_count": len(cache_payload["records"]),
"is_pruned_minibatch": is_pruned,
"logical_base_index": logical_index,
"physical_base_index": actual_index,
"retention_ratio": 0.2 if is_pruned else 1.0,
}
return sample
LeRobotMixtureDataset.__getitem__ = frameskip_getitem
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--config_yaml", required=True)
args, clipargs = parser.parse_known_args()
cache_path = os.environ.get("FRAMESKIP_CACHE")
if not cache_path:
raise RuntimeError("FRAMESKIP_CACHE is required")
start_draw = int(os.environ.get("FRAMESKIP_START_GLOBAL_DRAW_ID", "1280000"))
install_adapter(cache_path, start_draw)
from starVLA.model.framework.share_tools import apply_config_compat
from starVLA.training import train_starvla
from starVLA.training.trainer_utils.trainer_tools import normalize_dotlist_args
cfg = OmegaConf.load(args.config_yaml)
cfg = OmegaConf.merge(cfg, OmegaConf.from_dotlist(normalize_dotlist_args(clipargs)))
cfg = apply_config_compat(cfg, strict=True)
cfg.config_yaml = args.config_yaml
train_starvla.main(cfg)
if __name__ == "__main__":
main()
|