Blame for real_robot/source/retime/cli.py

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