scripts/run_online_sliding_ablations.py

#!/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()