| |
| """Attach frozen APE loss weights to the unchanged Florence-Bridge draw stream.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import os |
|
|
| from omegaconf import OmegaConf |
|
|
| from external_baselines.keyframe_focused_adapted.keyframe_cache import ( |
| KeyframeWeightCache, |
| ) |
|
|
|
|
| def install_adapter(cache_path: str) -> KeyframeWeightCache: |
| from starVLA.dataloader.gr00t_lerobot.datasets import LeRobotMixtureDataset |
| from starVLA.dataloader.infoflow_sampler import ValidAnchorManifest |
|
|
| cache = KeyframeWeightCache(cache_path) |
| if os.environ.get("KEYFRAME_SMOKE_FROZEN_MANIFEST") == "1": |
| original_build_manifest = LeRobotMixtureDataset._build_valid_anchor_manifest |
|
|
| def load_frozen_manifest(self): |
| path = self.data_cfg.get("valid_anchor_manifest_path") |
| if not path: |
| return original_build_manifest(self) |
| manifest = ValidAnchorManifest.load(path) |
| print( |
| "Smoke-only frozen manifest fast path: " |
| f"{path}", |
| flush=True, |
| ) |
| return manifest |
|
|
| LeRobotMixtureDataset._build_valid_anchor_manifest = load_frozen_manifest |
|
|
| original_getitem = LeRobotMixtureDataset.__getitem__ |
|
|
| def keyframe_getitem(self, index): |
| sample = original_getitem(self, index) |
| trace = sample["sampling_trace"] |
| record = cache.lookup( |
| trace["dataset_name"], |
| int(trace["trajectory_id"]), |
| int(trace["base_index"]), |
| ) |
| |
| trace["keyframe_ape"] = record.ape |
| trace["keyframe_weight"] = record.weight |
| trace["keyframe_cache"] = cache.path.name |
| return sample |
|
|
| LeRobotMixtureDataset.__getitem__ = keyframe_getitem |
| return cache |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--config_yaml", required=True) |
| args, clipargs = parser.parse_known_args() |
|
|
| cache = install_adapter(os.environ["KEYFRAME_APE_CACHE"]) |
| print( |
| f"Installed Keyframe-Focused-adapted cache {cache.path}", |
| flush=True, |
| ) |
|
|
| 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) |
| if cfg.datasets.vla_data.temporal_proposal_path is not None: |
| raise ValueError("Keyframe loss weighting requires the uniform temporal sampler") |
| if cfg.trainer.sample_loss_weight_key != "keyframe_weight": |
| raise ValueError("trainer.sample_loss_weight_key must be keyframe_weight") |
| cfg.config_yaml = args.config_yaml |
| train_starvla.main(cfg) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
