#!/usr/bin/env python3 """Run online sliding speed-set ablations end-to-end. This runner matches ONLINE_SLIDING_ABLATION_WORKFLOW.md. It does not build offline speed datasets. For each selected experiment it: 1. Writes source 1.0x LeRobot parquet norm stats to the experiment asset id. The online-sliding speed set is used for training samples, not for changing normalization statistics. 2. Starts PI05 PyTorch training with online sliding chunks and the experiment's speed set. Examples: uv run python scripts/run_online_sliding_ablations.py \\ --data-root /path/to/libero_data \\ --pi05-base /path/to/pi05_base \\ --only range_base,range_mid \\ --dry-run uv run python scripts/run_online_sliding_ablations.py \\ --data-root /path/to/libero_data \\ --pi05-base /path/to/pi05_base \\ --batch-size 64 --num-workers 0 """ from __future__ import annotations import argparse import dataclasses import json import os import shlex import subprocess import sys from pathlib import Path @dataclasses.dataclass(frozen=True) class Experiment: name: str speeds: tuple[float, ...] asset_id: str exp_name: str EXPERIMENTS: tuple[Experiment, ...] = ( Experiment( "range_base", (1.0,), "online_sliding_range_base_pi05", "pi05_online_sliding_range_base_text", ), Experiment( "range_mid_pilot", (0.75, 1.0, 1.25), "online_sliding_range_mid_pilot_pi05", "pi05_online_sliding_range_mid_pilot_text", ), Experiment( "range_mid", (0.75, 1.0, 1.25, 1.5), "online_sliding_range_mid_pi05", "pi05_online_sliding_range_mid_text", ), Experiment( "range_wide", (0.5, 0.75, 1.0, 1.25, 1.5, 1.75, 2.0), "online_sliding_range_wide_pi05", "pi05_online_sliding_range_wide_text", ), Experiment( "stride_0p5", (0.75, 1.0, 1.5), "online_sliding_stride_0p5_pi05", "pi05_online_sliding_stride_0p5_text", ), Experiment( "stride_0p125", (0.75, 0.875, 1.0, 1.125, 1.25, 1.375, 1.5), "online_sliding_stride_0p125_pi05", "pi05_online_sliding_stride_0p125_text", ), ) SHARED_WIDE_NORM = Experiment( "shared_wide_norm", (0.5, 0.75, 1.0, 1.25, 1.5, 1.75, 2.0), "online_sliding_shared_wide_norm_pi05", "unused_shared_wide_norm", ) SHARED_NORM_CONTROLS: tuple[Experiment, ...] = ( Experiment( "range_base_shared_norm", (1.0,), SHARED_WIDE_NORM.asset_id, "pi05_online_sliding_range_base_shared_norm_text", ), Experiment( "range_mid_shared_norm", (0.75, 1.0, 1.25, 1.5), SHARED_WIDE_NORM.asset_id, "pi05_online_sliding_range_mid_shared_norm_text", ), Experiment( "range_wide_shared_norm", (0.5, 0.75, 1.0, 1.25, 1.5, 1.75, 2.0), SHARED_WIDE_NORM.asset_id, "pi05_online_sliding_range_wide_shared_norm_text", ), ) def _speed_args(speeds: tuple[float, ...]) -> list[str]: return [f"{speed:g}" for speed in speeds] def _shell(cmd: list[str]) -> str: return " ".join(shlex.quote(part) for part in cmd) def _run(cmd: list[str], *, cwd: Path, env: dict[str, str], dry_run: bool) -> None: print(f"$ {_shell(cmd)}") if dry_run: return subprocess.run(cmd, cwd=cwd, env=env, check=True) def _parse_name_list(raw: str | None) -> set[str] | None: if raw is None or raw.strip() == "": return None return {item.strip() for item in raw.split(",") if item.strip()} def _select_experiments(args: argparse.Namespace) -> list[Experiment]: experiments = list(EXPERIMENTS) if args.include_shared_norm_controls: experiments.extend(SHARED_NORM_CONTROLS) only = _parse_name_list(args.only) exclude = _parse_name_list(args.exclude) or set() names = {exp.name for exp in experiments} if only: unknown = only - names if unknown: raise SystemExit(f"Unknown --only experiments: {sorted(unknown)}. Available: {sorted(names)}") experiments = [exp for exp in experiments if exp.name in only] if exclude: unknown = exclude - names if unknown: raise SystemExit(f"Unknown --exclude experiments: {sorted(unknown)}. Available: {sorted(names)}") experiments = [exp for exp in experiments if exp.name not in exclude] if not experiments: raise SystemExit("No experiments selected.") return experiments def _norm_stats_path(project_root: Path, train_config: str, asset_id: str) -> Path: return project_root / "assets" / train_config / asset_id / "norm_stats.json" def _checkpoint_dir(project_root: Path, train_config: str, exp_name: str) -> Path: return project_root / "checkpoints" / train_config / exp_name def _norm_cmd(args: argparse.Namespace, exp: Experiment) -> list[str]: return [ sys.executable, "scripts/compute_norm_stats.py", "--config-name", args.train_config, "--repo-id", str(args.data_root), "--asset-id", exp.asset_id, "--online-sliding-chunks", "--online-sliding-speeds", *_speed_args(exp.speeds), ] def _train_cmd(args: argparse.Namespace, exp: Experiment, log_dir: Path) -> list[str]: cmd = [ sys.executable, "-m", "torch.distributed.run", "--standalone", "--nnodes=1", f"--nproc_per_node={args.num_gpus}", "--log-dir", str(log_dir / "torchrun" / exp.exp_name), "--redirects", "3", "--tee", "3", "scripts/train_pytorch.py", args.train_config, "--exp-name", exp.exp_name, "--pytorch-weight-path", str(args.pi05_base), "--batch-size", str(args.batch_size), "--num-workers", str(args.num_workers), "--num-train-steps", str(args.num_train_steps), "--log-interval", str(args.log_interval), "--save-interval", str(args.save_interval), "--eval-speed-set", *_speed_args(exp.speeds), "--data.repo-id", str(args.data_root), "--data.assets.asset-id", exp.asset_id, "--data.online-sliding-chunks", "--data.online-sliding-speeds", *_speed_args(exp.speeds), "--data.speed-integration", "text", "--model.pytorch-compile-mode", args.compile_mode, ] if args.train_mode == "overwrite": cmd.append("--overwrite") elif args.train_mode == "resume": cmd.append("--resume") if args.no_wandb: cmd.append("--no-wandb-enabled") cmd.extend(args.extra_train_arg) return cmd def _base_env(args: argparse.Namespace) -> dict[str, str]: env = os.environ.copy() if not args.keep_wandb_env: env.pop("WANDB_API_KEY", None) env.pop("WANDB_API_KEY_FILE", None) env.setdefault("WANDB__SERVICE_WAIT", str(args.wandb_service_wait)) return env def _train_env(args: argparse.Namespace) -> dict[str, str]: env = _base_env(args) env["CUDA_VISIBLE_DEVICES"] = args.cuda_devices or ",".join(str(i) for i in range(args.num_gpus)) return env def _write_manifest( project_root: Path, log_dir: Path, args: argparse.Namespace, experiments: list[Experiment], ) -> None: if args.dry_run: return manifest = { "train_config": args.train_config, "data_root": str(args.data_root), "pi05_base": str(args.pi05_base) if args.pi05_base else None, "num_gpus": args.num_gpus, "cuda_devices": args.cuda_devices or ",".join(str(i) for i in range(args.num_gpus)), "batch_size": args.batch_size, "num_workers": args.num_workers, "num_train_steps": args.num_train_steps, "train_mode": args.train_mode, "experiments": [dataclasses.asdict(exp) for exp in experiments], } log_dir.mkdir(parents=True, exist_ok=True) (log_dir / "online_sliding_ablation_manifest.json").write_text(json.dumps(manifest, indent=2)) def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--project-root", type=Path, default=Path.cwd(), help="VLAwithVariousSpeed repo root.") parser.add_argument("--data-root", type=Path, required=True, help="Source LeRobot/LIBERO dataset root.") parser.add_argument("--pi05-base", type=Path, default=None, help="Path to PI05 base weights directory.") parser.add_argument("--train-config", default="pi05_libero_various_speed_all") parser.add_argument("--only", default=None, help="Comma-separated experiment names to run.") parser.add_argument("--exclude", default=None, help="Comma-separated experiment names to skip.") parser.add_argument("--include-shared-norm-controls", action="store_true") parser.add_argument("--stage", choices=("all", "norm", "train"), default="all") parser.add_argument("--force-norm", action="store_true", help="Recompute norm stats even if norm_stats.json exists.") parser.add_argument( "--train-mode", choices=("overwrite", "resume", "skip-existing", "fail-if-exists"), default="overwrite", help="How to handle existing checkpoint directories.", ) parser.add_argument("--dry-run", action="store_true") parser.add_argument("--num-gpus", type=int, default=8) parser.add_argument("--cuda-devices", default=None, help="CUDA_VISIBLE_DEVICES value. Default: 0..num_gpus-1.") parser.add_argument("--batch-size", type=int, default=64) parser.add_argument("--num-workers", type=int, default=0) parser.add_argument("--num-train-steps", type=int, default=30_000) parser.add_argument("--log-interval", type=int, default=100) parser.add_argument("--save-interval", type=int, default=1000) parser.add_argument("--compile-mode", default="None") parser.add_argument("--log-dir", type=Path, default=Path("logs/online_sliding_ablations")) parser.add_argument("--no-wandb", action="store_true") parser.add_argument("--keep-wandb-env", action="store_true") parser.add_argument("--wandb-service-wait", type=int, default=300) parser.add_argument( "--extra-train-arg", action="append", default=[], help="Extra argument appended to train_pytorch.py. Repeat for multiple args.", ) return parser.parse_args() def main() -> None: args = parse_args() project_root = args.project_root.resolve() args.project_root = project_root args.data_root = args.data_root.resolve() if args.pi05_base is not None: args.pi05_base = args.pi05_base.resolve() args.log_dir = (project_root / args.log_dir).resolve() if not args.log_dir.is_absolute() else args.log_dir.resolve() if not (project_root / "scripts" / "train_pytorch.py").exists(): raise SystemExit(f"project root does not look valid: {project_root}") if not args.data_root.exists(): raise SystemExit(f"data root does not exist: {args.data_root}") if args.stage in ("all", "train") and args.pi05_base is None: raise SystemExit("--pi05-base is required for training stages.") if args.stage in ("all", "train") and args.pi05_base is not None and not args.pi05_base.exists(): raise SystemExit(f"pi05 base path does not exist: {args.pi05_base}") if args.batch_size % args.num_gpus != 0: raise SystemExit(f"--batch-size ({args.batch_size}) must be divisible by --num-gpus ({args.num_gpus}).") experiments = _select_experiments(args) if args.include_shared_norm_controls and args.stage in ("all", "norm"): # The shared-control training experiments all point to this shared norm. # Add it only to the norm stage, not to training. norm_experiments = [SHARED_WIDE_NORM, *experiments] else: norm_experiments = experiments print("Online sliding ablation runner") print(f" project_root = {project_root}") print(f" data_root = {args.data_root}") print(f" train_config = {args.train_config}") print(f" stage = {args.stage}") print(f" train_mode = {args.train_mode}") print(f" experiments = {[exp.name for exp in experiments]}") if not args.keep_wandb_env: print(" wandb env = WANDB_API_KEY/WANDB_API_KEY_FILE will be ignored") print() _write_manifest(project_root, args.log_dir, args, experiments) if args.stage in ("all", "norm"): seen_assets: set[str] = set() for exp in norm_experiments: if exp.asset_id in seen_assets: continue seen_assets.add(exp.asset_id) stats_path = _norm_stats_path(project_root, args.train_config, exp.asset_id) if stats_path.exists() and not args.force_norm: print(f"[skip norm] {exp.name}: {stats_path}") continue print(f"\n========== norm: {exp.name} speeds={exp.speeds} asset={exp.asset_id} ==========") _run(_norm_cmd(args, exp), cwd=project_root, env=_base_env(args), dry_run=args.dry_run) if args.stage in ("all", "train"): for exp in experiments: ckpt_dir = _checkpoint_dir(project_root, args.train_config, exp.exp_name) if args.train_mode == "skip-existing" and ckpt_dir.exists(): print(f"[skip train] {exp.name}: checkpoint dir exists at {ckpt_dir}") continue if args.train_mode == "fail-if-exists" and ckpt_dir.exists(): raise SystemExit(f"checkpoint dir already exists for {exp.name}: {ckpt_dir}") stats_path = _norm_stats_path(project_root, args.train_config, exp.asset_id) if not stats_path.exists() and not args.dry_run: raise SystemExit(f"missing norm stats for {exp.name}: {stats_path}") print(f"\n========== train: {exp.name} speeds={exp.speeds} exp={exp.exp_name} ==========") _run(_train_cmd(args, exp, args.log_dir), cwd=project_root, env=_train_env(args), dry_run=args.dry_run) if __name__ == "__main__": main()