"""One-command adapter to the accompanying StarVLA training backend. The CLI prepares temporal proposals and dispatches StarVLA training. The sampler and ordinary importance correction are unchanged. """ from __future__ import annotations import argparse import json import math import os from pathlib import Path import subprocess import sys def parser() -> argparse.ArgumentParser: result = argparse.ArgumentParser(description=__doc__) result.add_argument("--config", type=Path, required=True) result.add_argument("--base-policy", type=Path, required=True, help="Complete Accelerate checkpoint directory, committed single-process model file, or model file with --start weights.") result.add_argument("--start", choices=("resume", "weights"), default="weights", help="resume restores the full state; weights explicitly starts a new optimizer/schedule.") result.add_argument("--alpha", type=float, default=0.2) result.add_argument("--K", "--k", dest="k", type=int, default=8) result.add_argument("--mode", nargs="+", choices=("retime", "uniform", "shuffle"), default=["retime"]) result.add_argument("--manifest", type=Path) result.add_argument("--data-root", type=Path, help="Override the dataset root in the config") result.add_argument("--data-mix", type=str, help="Override the named dataset mixture in the config") result.add_argument("--vlm-path", type=Path, help="Override the Florence/Qwen backbone location") result.add_argument("--statistics-path", type=Path, help="Override precomputed dataset statistics, if required") result.add_argument( "--omit-data-identity-files", action="store_true", help="Use the manifest's computed dataset identity without separate provenance/fingerprint files.", ) result.add_argument("--epsilon", type=float, default=0.1) result.add_argument("--shuffle-seed", type=int, default=12345) result.add_argument("--accelerate-config", type=Path, required=True) result.add_argument("--output", type=Path, required=True) result.add_argument("--enable-wandb", action="store_true") result.add_argument("--backend-root", type=Path, default=Path(__file__).resolve().parents[1], help="ReTime-enabled StarVLA source tree shipped with this artifact.") result.add_argument("--dry-run", action="store_true", help="Validate local inputs and print commands; do not score, train, or write files.") return result def _file(path: Path, label: str) -> Path: path = path.expanduser().resolve() if not path.is_file(): raise ValueError(f"{label} is not a file: {path}") return path def _directory(path: Path, label: str) -> Path: path = path.expanduser().resolve() if not path.is_dir(): raise ValueError(f"{label} is not a directory: {path}") return path def _native_commit_path(model_path: Path) -> Path: for suffix in ("_pytorch_model.pt", "_model.safetensors"): if model_path.name.endswith(suffix): return model_path.with_name( model_path.name[: -len(suffix)] + "_checkpoint.json" ) raise ValueError( "A single-process resume checkpoint must end in " "_pytorch_model.pt or _model.safetensors" ) def _validate_resume_checkpoint(base: Path, world_size: int) -> tuple[Path, dict]: if base.is_dir(): model = _file(base / "pytorch_model.pt", "base-policy model") commit = json.loads( _file(base / "checkpoint.json", "checkpoint commit").read_text( encoding="utf-8" ) ) if int(commit.get("world_size", -1)) != world_size: raise ValueError("Base checkpoint world size differs from the launch contract") entries = commit.get("files", []) if not entries or any( not (base / entry["name"]).is_file() or (base / entry["name"]).stat().st_size != int(entry["bytes"]) for entry in entries ): raise ValueError("Base Accelerate checkpoint is incomplete") return model, commit model = _file(base, "base-policy model") if world_size != 1: raise ValueError( "A committed single-process checkpoint requires a one-process launch" ) commit_path = _native_commit_path(model) commit = json.loads(_file(commit_path, "checkpoint commit").read_text(encoding="utf-8")) files = commit.get("files", {}) for label in ("model", "training_state"): entry = files.get(label, {}) path = model.parent / str(entry.get("name", "")) if not path.is_file() or path.stat().st_size != int(entry.get("bytes", -1)): raise ValueError(f"Base single-process checkpoint has an invalid {label} file") if files["model"]["name"] != model.name: raise ValueError("Checkpoint commit does not describe --base-policy") return model, commit def build_plan(args: argparse.Namespace) -> dict: from omegaconf import OmegaConf if not math.isfinite(args.alpha) or not 0 <= args.alpha < 1: raise ValueError("alpha must satisfy 0 <= alpha < 1") if args.k < 1 or not math.isfinite(args.epsilon) or args.epsilon <= 0: raise ValueError("K and epsilon must be positive") if len(set(args.mode)) != len(args.mode): raise ValueError("Each mode can be requested only once") root = args.backend_root.expanduser().resolve() config = _file(args.config, "config") launch_config = _file(args.accelerate_config, "accelerate config") cfg = OmegaConf.load(config) if not args.enable_wandb: cfg.wandb_project = None launch = OmegaConf.load(launch_config) output = args.output.expanduser().resolve() cfg.run_root_dir = str(output / "runs") cfg.output_dir = str(output / "prepare") if args.data_root: cfg.datasets.vla_data.data_root_dir = str(_directory(args.data_root, "dataset root")) if args.data_mix: cfg.datasets.vla_data.data_mix = args.data_mix if args.vlm_path: cfg.framework.qwenvl.base_vlm = str(_directory(args.vlm_path, "VLM backbone")) if args.statistics_path: cfg.datasets.vla_data.precomputed_statistics_path = str(_file(args.statistics_path, "statistics")) if args.omit_data_identity_files: cfg.datasets.vla_data.dataset_provenance_path = None cfg.datasets.vla_data.dataset_content_fingerprint_path = None # Model-only initialization derives support from the current dataset. if args.start == "weights": cfg.trainer.is_resume = False cfg.trainer.resume_accelerate_checkpoint = None cfg.trainer.resume_checkpoint_integrity_manifest = None if OmegaConf.select(cfg, "trainer.initial_completed_steps") is not None: cfg.trainer.initial_completed_steps = 0 if args.manifest is None: cfg.datasets.vla_data.valid_anchor_manifest_path = None cfg.datasets.vla_data.dataset_provenance_path = None cfg.datasets.vla_data.dataset_content_fingerprint_path = None cfg.datasets.vla_data.temporal_proposal_path = None if args.manifest is not None: cfg.datasets.vla_data.valid_anchor_manifest_path = str(args.manifest.expanduser().resolve()) if OmegaConf.select(cfg, "framework.name") != "QwenGR00T": raise ValueError("This adapter supports the supplied Florence/Qwen GR00T configs only") if not OmegaConf.select(cfg, "datasets.vla_data.traceable_sampling", default=False): raise ValueError("ReTime requires traceable_sampling and a frozen valid-anchor manifest") world_size = int(launch.num_processes) if int(OmegaConf.select(launch, "num_machines", default=1)) != 1: raise ValueError("The one-command launcher currently supports one machine") if world_size != int(cfg.datasets.vla_data.expected_world_size): raise ValueError("Accelerate process count differs from the dataset training contract") accumulation = int(cfg.trainer.gradient_accumulation_steps) ds_accumulation = OmegaConf.select(launch, "deepspeed_config.gradient_accumulation_steps") if ds_accumulation is not None and int(ds_accumulation) != accumulation: raise ValueError("Accelerate and trainer gradient accumulation differ") batch = int(cfg.datasets.vla_data.per_device_batch_size) if min(world_size, accumulation, batch) < 1: raise ValueError("World size, accumulation and batch size must be positive") configured_manifest = OmegaConf.select(cfg, "datasets.vla_data.valid_anchor_manifest_path") manifest = args.manifest or (Path(configured_manifest) if configured_manifest else None) prepare_manifest = manifest is None if prepare_manifest and args.start == "resume": raise ValueError("Exact resume requires the parent run's manifest; use --manifest") manifest = output / "proposal" / "manifest.json" if prepare_manifest else _file(manifest, "manifest") if configured_manifest and Path(configured_manifest).expanduser().resolve() != manifest: raise ValueError("--manifest differs from the config manifest") base = args.base_policy.expanduser().resolve() if args.start == "resume": checkpoint, _ = _validate_resume_checkpoint(base, world_size) else: checkpoint = _file(base / "pytorch_model.pt" if base.is_dir() else base, "base-policy model") cfg.trainer.pretrained_checkpoint = str(checkpoint) if args.start == "weights" else None cfg.trainer.resume_accelerate_checkpoint = str(base) if args.start == "resume" else None cfg.trainer.is_resume = args.start == "resume" modes = list(args.mode) for mode in modes: if (output / "runs" / mode).exists(): raise ValueError(f"Training output already exists: {output / 'runs' / mode}; choose a fresh --output") commands = [] cfg.datasets.vla_data.valid_anchor_manifest_path = str(manifest) # Materialize all paths once so every arm and subprocess uses the same config. config_text = OmegaConf.to_yaml(cfg, resolve=True) config = output / "config.yaml" if prepare_manifest: commands.append([sys.executable, "-m", "retime.prepare", "--config", str(config), "--output", str(manifest)]) def script(name: str, *arguments: str) -> None: module = { "build_residual_temporal_proposal.py": "retime.commands.proposal", "build_shuffle_temporal_proposal.py": "retime.commands.shuffle", }[name] commands.append([sys.executable, "-m", module, *map(str, arguments)]) residual = output / "proposal" / "retime.json" shuffled = output / "proposal" / "shuffle.json" if any(mode != "uniform" for mode in modes): script("build_residual_temporal_proposal.py", "--config", config, "--checkpoint", checkpoint, "--manifest", manifest, "--cache-dir", output / "cache", "--output", residual, "--grid-size", args.k, "--alpha", args.alpha, "--epsilon", args.epsilon) if "shuffle" in modes: script("build_shuffle_temporal_proposal.py", "--manifest", manifest, "--residual-proposal", residual, "--permutation-seed", args.shuffle_seed, "--output", shuffled) training = _file(root / "starVLA" / "training" / "train_starvla.py", "training backend") for mode in modes: proposal = {"uniform": None, "retime": str(residual), "shuffle": str(shuffled)}[mode] overrides = { "run_id": mode, "run_root_dir": str(output / "runs"), "output_dir": str(output / "runs" / mode), "datasets.vla_data.valid_anchor_manifest_path": str(manifest), "datasets.vla_data.temporal_proposal_path": proposal, "trainer.is_resume": args.start == "resume", "trainer.allow_model_only_resume": False, "trainer.resume_accelerate_checkpoint": str(base) if args.start == "resume" else None, "trainer.pretrained_checkpoint": str(checkpoint) if args.start == "weights" else None, } commands.append([ sys.executable, "-m", "accelerate.commands.launch", "--config_file", str(launch_config), "--num_processes", str(world_size), str(training), "--config_yaml", str(config), *[f"--{key}={json.dumps(value)}" for key, value in overrides.items()], ]) return { "backend_root": str(root), "output": str(output), "base_policy": str(base), "start": args.start, "alpha": args.alpha, "K": args.k, "modes": modes, "global_batch_size": batch * world_size * accumulation, "config_text": config_text, "commands": commands, } def main(argv: list[str] | None = None) -> int: arguments = parser() args = arguments.parse_args(argv) try: plan = build_plan(args) except (ValueError, OSError, ImportError, KeyError) as error: arguments.error(str(error)) print(json.dumps({key: value for key, value in plan.items() if key != "config_text"}, indent=2), flush=True) if args.dry_run: return 0 output = Path(plan["output"]) output.mkdir(parents=True, exist_ok=True) (output / "proposal").mkdir(exist_ok=True) (output / "config.yaml").write_text(plan["config_text"], encoding="utf-8") (output / "plan.json").write_text(json.dumps(plan, indent=2) + "\n", encoding="utf-8") env = os.environ.copy() env["PYTHONPATH"] = os.pathsep.join(filter(None, [plan["backend_root"], env.get("PYTHONPATH")])) env.setdefault("WANDB_MODE", "offline") env["PYTHONUNBUFFERED"] = "1" for command in plan["commands"]: subprocess.run(command, cwd=plan["backend_root"], env=env, check=True) return 0 if __name__ == "__main__": raise SystemExit(main())