| |
| """Run the Florence-RoboTwin FrameSkip continuation from the 10K state.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import os |
| from pathlib import Path |
|
|
| import numpy as np |
| from omegaconf import OmegaConf |
|
|
|
|
| METHOD = "FrameSkip-adapted (Gaussian-TPI fallback)" |
|
|
|
|
| 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") != METHOD: |
| raise RuntimeError(f"unexpected FrameSkip method: {payload.get('method')!r}") |
| if float(payload.get("retention_ratio", -1.0)) != 0.2: |
| raise RuntimeError("FrameSkip cache must use retention_ratio=0.2") |
|
|
| cache: dict[tuple[str, str], np.ndarray] = {} |
| for record in payload["records"]: |
| key = (record["dataset_name"], str(record["trajectory_id"])) |
| keep = np.asarray(record["keep_indices"], dtype=np.int64) |
| if key in cache: |
| raise RuntimeError(f"duplicate FrameSkip trajectory: {key}") |
| 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, |
| global_microbatch_size: int, |
| ) -> None: |
| from starVLA.dataloader.gr00t_lerobot.datasets import LeRobotMixtureDataset |
|
|
| cache, cache_payload = load_cache(cache_path) |
| original_getitem = LeRobotMixtureDataset.__getitem__ |
| period = 6 |
|
|
| def frameskip_getitem(self, index): |
| trace = self.sample_trace(index) |
| offset = int(trace.global_draw_id) - start_global_draw_id |
| if offset < 0: |
| raise RuntimeError("FrameSkip draw precedes the 10K continuation boundary") |
|
|
| 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)) |
| try: |
| retained = cache[key] |
| except KeyError as exc: |
| raise RuntimeError(f"FrameSkip cache is missing trajectory {key}") from exc |
| position = int(np.searchsorted(retained, logical_index)) |
| actual_index = int(retained[min(position, len(retained) - 1)]) |
|
|
| 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): |
| raise RuntimeError("FrameSkip attempted to change trajectory identity") |
| if 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, overrides = parser.parse_known_args() |
|
|
| install_adapter( |
| cache_path=os.environ["FRAMESKIP_CACHE"], |
| start_global_draw_id=int(os.environ.get("FRAMESKIP_START_GLOBAL_DRAW_ID", "640000")), |
| global_microbatch_size=int(os.environ.get("FRAMESKIP_GLOBAL_MICROBATCH_SIZE", "4")), |
| ) |
|
|
| 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(overrides))) |
| cfg = apply_config_compat(cfg, strict=True) |
| cfg.config_yaml = args.config_yaml |
| train_starvla.main(cfg) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
