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