scripts/sample_appearance_episodes.py

#!/usr/bin/env python
"""Pick a lab-stratified episode set for the appearance run, and fetch it.

Why stratify rather than take the pool's head: the pool is sorted by uuid, so
its head is one lab, and the 63 episodes already on disk are effectively one lab
too (62 of them share the camera pair 22008760/24400334). A LoRA trained on one
lab learns that lab's bench, lighting and gripper, which is the opposite of what
an *appearance* fine-tune is for -- the robot is the constant across DROID and
everything else is the nuisance variable. Proportional-with-a-floor sampling
keeps the big labs dominant (they really are most of DROID) while guaranteeing
the small ones are represented at all: RAD has 512 episodes against TRI's 12300,
and pure proportional sampling at n=400 would give it 4.

Intrinsics are per physical camera, and the appearance pipeline falls back to a
median-ZED guess for a serial it has never seen (principal point off by up to
29 px at 720). Episodes whose serials are absent from
``configs/droid_camera_intrinsics.json`` are therefore reported, not silently
included -- pass ``--allow-unknown-serials`` to take them anyway.

Usage:
    PYTHONPATH=src /home/quang/miniconda3/envs/fpgm/bin/python \\
        scripts/sample_appearance_episodes.py --n 400 --out configs/appearance_400.txt
    ... then --download to fetch mp4 + trajectory.h5 + metadata.json for each.
"""

from __future__ import annotations

import argparse
import collections
import json
import random
import subprocess
import sys
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parents[1]


def stratified(uuids: list[str], n: int, floor: int, seed: int) -> list[str]:
    """``n`` uuids, proportional per lab but never fewer than ``floor`` per lab.

    Labs with fewer than ``floor`` episodes contribute all of them. The residual
    from rounding is handed to the largest labs, so the total is exactly ``n``
    (or the pool size, if that is smaller).
    """
    by_lab: dict[str, list[str]] = collections.defaultdict(list)
    for u in uuids:
        by_lab[u.split("+")[0]].append(u)
    rng = random.Random(seed)
    for v in by_lab.values():
        rng.shuffle(v)

    total = sum(len(v) for v in by_lab.values())
    quota = {lab: max(floor, round(n * len(v) / total)) for lab, v in by_lab.items()}
    quota = {lab: min(q, len(by_lab[lab])) for lab, q in quota.items()}

    # Reconcile to exactly n by walking labs largest-first.
    order = sorted(by_lab, key=lambda k: -len(by_lab[k]))
    while sum(quota.values()) > n:
        for lab in order:
            if quota[lab] > floor and sum(quota.values()) > n:
                quota[lab] -= 1
    while sum(quota.values()) < min(n, total):
        for lab in order:
            if quota[lab] < len(by_lab[lab]) and sum(quota.values()) < n:
                quota[lab] += 1

    picked: list[str] = []
    for lab in order:
        picked.extend(by_lab[lab][: quota[lab]])
    return sorted(picked)


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--pool", type=Path, default=REPO_ROOT / "configs/appearance_uuid_pool.txt")
    ap.add_argument("--cameras-dir", type=Path,
                    default=REPO_ROOT / "data/pointworld/droid/cameras")
    ap.add_argument("--intrinsics", type=Path,
                    default=REPO_ROOT / "configs/droid_camera_intrinsics.json")
    ap.add_argument("--episodes-root", type=Path, default=REPO_ROOT / "data/droid_raw")
    ap.add_argument("--n", type=int, default=400)
    ap.add_argument("--floor", type=int, default=8, help="minimum episodes per lab")
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--out", type=Path, required=True)
    ap.add_argument("--allow-unknown-serials", action="store_true")
    ap.add_argument("--download", action="store_true",
                    help="fetch the sampled episodes with download_droid_episodes.py")
    ap.add_argument("--download-batch", type=int, default=25)
    args = ap.parse_args()

    pool = [u for u in args.pool.read_text().split() if u]
    known = set(json.loads(args.intrinsics.read_text())["cameras"])
    have = {p.name for p in args.episodes_root.iterdir() if p.is_dir()} \
        if args.episodes_root.exists() else set()

    picked = stratified(pool, args.n, args.floor, args.seed)

    kept, unknown = [], []
    for u in picked:
        cam = args.cameras_dir / f"{u}_cameras.json"
        serials = [k for k in json.loads(cam.read_text()) if k.isdigit()]
        (kept if all(s in known for s in serials) else unknown).append(u)
    if args.allow_unknown_serials:
        kept += unknown

    args.out.write_text("\n".join(sorted(kept)) + "\n")
    by_lab = collections.Counter(u.split("+")[0] for u in kept)
    print(f"{len(kept)} episodes -> {args.out}  ({len(unknown)} dropped for unknown "
          f"camera serials{'; kept anyway' if args.allow_unknown_serials else ''})")
    for lab, c in by_lab.most_common():
        print(f"  {lab:12s} {c}")
    todo = [u for u in kept if u not in have]
    print(f"{len(have & set(kept))} already on disk, {len(todo)} to download "
          f"(~{len(todo) * 13 / 1024:.1f} GB)")

    if args.download and todo:
        script = REPO_ROOT / "scripts/download_droid_episodes.py"
        for i in range(0, len(todo), args.download_batch):
            batch = todo[i:i + args.download_batch]
            print(f"  downloading {i + 1}-{i + len(batch)} / {len(todo)}", flush=True)
            # Batched rather than one big argv: the downloader exits non-zero if
            # ANY episode fails, and a batch bounds how much is re-run on a retry.
            subprocess.run([sys.executable, str(script),
                            "--data-dir", str(args.episodes_root),
                            "--cameras-dir", str(args.cameras_dir),
                            "--uuids", *batch], check=False)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())