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