scripts/survey_droid_tasks.py

#!/usr/bin/env python
"""Sample DROID episode uuids across labs and fetch each one's task instruction.

The 42,935 ``*_cameras.json`` filenames already on disk under
``data/pointworld/droid/cameras`` are the full episode universe -- the uuid is the
filename, so no network call is needed to *enumerate* episodes, only to learn what
each one is about. That "what" lives in the episode's own
``metadata_<uuid>.json`` on the public DROID bucket (1.8 KB, key ``current_task``).

**Sampled per lab, not uniformly at random over all 42,935.** The local corpus is
12,300 TRI / 6,140 AUTOLab / ... down to 512 RAD, and a uniform draw would return
mostly TRI. Object variety in DROID tracks the lab (each site set up its own
scenes and props), so a per-lab quota is what actually buys diverse objects, which
is the point of the survey.

Politeness: a small thread pool and a single pass. ``resolve_episode_url`` may probe
``success`` then ``failure``, so budget up to 2 requests per uuid.

Usage:
    PYTHONPATH=src python scripts/survey_droid_tasks.py \\
        --per-lab 60 --out /tmp/droid_tasks.jsonl [--workers 12]
"""

from __future__ import annotations

import argparse
import json
import random
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

from fpgm.data.droid_raw import DroidRawClient
from fpgm.utils.logging import get_logger, setup_logging

logger = get_logger("survey_droid_tasks")


def episode_uuids(cameras_dir: Path) -> list[str]:
    return sorted(p.name.removesuffix("_cameras.json") for p in cameras_dir.glob("*_cameras.json"))


def sample_per_lab(uuids: list[str], per_lab: int, seed: int) -> list[str]:
    by_lab: dict[str, list[str]] = defaultdict(list)
    for u in uuids:
        by_lab[u.split("+", 1)[0]].append(u)
    rng = random.Random(seed)
    picked: list[str] = []
    for lab in sorted(by_lab):
        pool = by_lab[lab]
        picked.extend(rng.sample(pool, min(per_lab, len(pool))))
    rng.shuffle(picked)
    return picked


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--cameras-dir", type=Path, default=Path("data/pointworld/droid/cameras"))
    ap.add_argument("--cache-dir", type=Path, default=Path("data/droid_meta_survey"))
    ap.add_argument("--out", type=Path, required=True)
    ap.add_argument("--per-lab", type=int, default=60)
    ap.add_argument("--workers", type=int, default=12)
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--log-level", default="WARNING")
    args = ap.parse_args()
    setup_logging(args.log_level)

    uuids = episode_uuids(args.cameras_dir)
    picked = sample_per_lab(uuids, args.per_lab, args.seed)
    print(f"universe {len(uuids)} episodes -> sampling {len(picked)}")

    client = DroidRawClient(args.cache_dir)

    def fetch(uuid: str) -> dict | None:
        try:
            meta = client.load_episode_metadata(uuid)
        except Exception as exc:  # noqa: BLE001 - a dead uuid must not kill the survey
            return {"uuid": uuid, "error": type(exc).__name__, "detail": str(exc)[:160]}
        return {
            "uuid": uuid,
            "lab": uuid.split("+", 1)[0],
            "task": (meta.get("current_task") or "").strip(),
            "ext1_serial": meta.get("ext1_cam_serial"),
            "ext2_serial": meta.get("ext2_cam_serial"),
        }

    rows: list[dict] = []
    ok = err = 0
    with ThreadPoolExecutor(max_workers=args.workers) as pool:
        futures = {pool.submit(fetch, u): u for u in picked}
        for i, fut in enumerate(as_completed(futures), 1):
            row = fut.result()
            if row is None:
                continue
            rows.append(row)
            if "error" in row:
                err += 1
            else:
                ok += 1
            if i % 100 == 0:
                print(f"  {i}/{len(picked)}  ok={ok} err={err}")

    args.out.parent.mkdir(parents=True, exist_ok=True)
    with args.out.open("w") as f:
        for row in rows:
            f.write(json.dumps(row) + "\n")
    print(f"done: ok={ok} err={err} -> {args.out}")


if __name__ == "__main__":
    main()