scripts/build_droid_coverage_sample.py

#!/usr/bin/env python
"""Pick a small DROID subset chosen to COVER the dataset's axes of variation.

Not a random draw. A uniform sample of DROID returns mostly TRI-lab pick-and-place
of rigid objects, which is exactly the regime this project already works in and
therefore the least informative thing to look at. This instead stratifies over the
properties that actually change what a geometry pipeline has to do:

* **Object physics class** -- the axis that decides whether a single rigid SE(3)
  pose per frame is even the right model. Compact-rigid and thin-rigid are the
  regime this project handles; symmetric-rigid breaks pose *fitting* (measured in
  this repo: two bit-identical runs recovered a brick's orientation ~178 deg apart
  with equally good silhouette IoU); articulated has no representation beyond the
  "barely rotates" fixture role; deformable has none at all.
* **Task template** -- DROID's ``current_task`` is usually a *category template*
  ("Move object into or out of container (ex: drawer, ...)"), not a description.
  Templates therefore label the interaction type, which is what we want to spread
  over, and the concrete object has to come from the human annotations instead.
* **Lab** -- each site built its own scenes, props and lighting, so lab is the
  cheapest available proxy for scene diversity.
* **Outcome and length** -- failures and very short/long episodes exercise
  different failure modes (aborted grasps, idle frames) than the median clip.

Object identity comes from ``aggregated-annotations-030724.json`` (3 human
instructions per episode) and is only trusted when >= 2 of the 3 annotators name
the same object; everything else comes from each episode's own
``metadata_<uuid>.json``.

Usage:
    PYTHONPATH=src python scripts/build_droid_coverage_sample.py \\
        --n 48 --out-json outputs/droid_coverage/manifest.json
"""

from __future__ import annotations

import argparse
import json
import pathlib
import re
from collections import Counter, defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed

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

#: Object physics classes, in the order they stress this pipeline. The keywords are
#: matched against the human annotations, not against ``current_task``.
OBJECT_CLASSES: dict[str, list[str]] = {
    "rigid_compact": ["block", "box", "can", "cube", "carton"],
    "rigid_thin": ["marker", "pen", "screwdriver", "spoon", "fork", "knife", "brush"],
    "rigid_symmetric": ["cup", "mug", "bottle", "bowl", "jar", "glass"],
    "articulated": ["drawer", "lid", "cabinet", "door", "oven", "microwave"],
    "deformable": ["towel", "cloth", "shirt", "napkin", "rag", "bag", "sock"],
    "flat_thin": ["plate", "book", "paper", "tray", "card"],
    "organic": ["banana", "apple", "orange", "carrot", "fruit", "bread"],
    "transparent": ["glass", "transparent", "clear bottle", "water bottle"],
}

#: Task templates DROID actually ships, keyed by a short label. Matched as a
#: case-insensitive substring of ``current_task``; ``free_form`` is the catch-all
#: for the "do any task" / "Suggested task:" strings this project's own config
#: already filters with ``meta_tokens``.
TASK_TEMPLATES: dict[str, str] = {
    "into_container": "into or out of container",
    "reposition": "new position and orientation",
    "clean_with_cloth": "use cloth to clean",
    "hang": "hang or unhang",
    "press_button": "press button",
    "open_close": "open or close slidable",
    "fold": "fold, spread out, or clump",
    "lid_on_off": "move lid on or off",
}


def classify_object(instructions: list[str]) -> str | None:
    """Object physics class if >=2 of 3 annotators name a member of the same class."""
    joined = [s.lower() for s in instructions if s]
    best: tuple[str, int] | None = None
    for cls, words in OBJECT_CLASSES.items():
        hits = sum(1 for s in joined if any(re.search(r"\b" + w, s) for w in words))
        if hits >= 2 and (best is None or hits > best[1]):
            best = (cls, hits)
    return best[0] if best else None


def classify_task(current_task: str) -> str:
    t = (current_task or "").lower()
    for label, needle in TASK_TEMPLATES.items():
        if needle in t:
            return label
    if not t or any(m in t for m in ("do any task", "anything", "suggested task")):
        return "free_form"
    return "concrete"  # a lab that wrote a real instruction instead of a template


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--annotations", type=pathlib.Path,
                    default=pathlib.Path("data/droid_annotations/aggregated-annotations-030724.json"))
    ap.add_argument("--cameras-dir", type=pathlib.Path,
                    default=pathlib.Path("data/pointworld/droid/cameras"))
    ap.add_argument("--meta-cache", type=pathlib.Path, default=pathlib.Path("data/droid_meta_survey"))
    ap.add_argument("--out-json", required=True, type=pathlib.Path)
    ap.add_argument("--n", type=int, default=48)
    ap.add_argument("--per-class-pool", type=int, default=14,
                    help="candidates fetched per object class before final selection")
    ap.add_argument("--n-failures", type=int, default=8,
                    help="failure episodes to append, drawn WITHOUT object labels (see below)")
    ap.add_argument("--per-lab-in-class", type=int, default=2,
                    help="cap per lab inside one object class when building the candidate pool. "
                         "Raising this widens lab coverage and, because outcome is only known "
                         "after the metadata fetch, is also the only way to get failure episodes "
                         "into the pool at all -- they are the minority everywhere.")
    ap.add_argument("--workers", type=int, default=12)
    ap.add_argument("--log-level", default="WARNING")
    args = ap.parse_args()
    setup_logging(args.log_level)

    ann = json.loads(args.annotations.read_text())
    covered = {p.name.removesuffix("_cameras.json") for p in args.cameras_dir.glob("*_cameras.json")}
    pool = {k: list(v.values()) for k, v in ann.items() if k in covered}
    covered_annotated = set(pool)
    print(f"annotated {len(ann)} | pointworld-covered {len(covered)} | usable {len(pool)}")

    by_class: dict[str, list[str]] = defaultdict(list)
    for uuid, ins in pool.items():
        cls = classify_object(ins)
        if cls:
            by_class[cls].append(uuid)
    print("object-class pool sizes: " + ", ".join(f"{c}={len(v)}" for c, v in sorted(by_class.items())))

    # Candidates: spread labs inside each class before spending any network call.
    cands: list[tuple[str, str]] = []
    for cls, uuids in by_class.items():
        seen_lab: Counter[str] = Counter()
        for uuid in sorted(uuids, key=lambda u: (u.split("+")[0], u)):
            lab = uuid.split("+")[0]
            if seen_lab[lab] >= args.per_lab_in_class:
                continue
            seen_lab[lab] += 1
            cands.append((cls, uuid))
            if seen_lab.total() >= args.per_class_pool:
                break
    print(f"fetching metadata for {len(cands)} candidates")

    client = DroidRawClient(args.meta_cache)

    def fetch(item: tuple[str, str]) -> dict | None:
        cls, uuid = item
        try:
            m = client.load_episode_metadata(uuid)
        except Exception:  # noqa: BLE001 - a dead uuid must not kill the build
            return None
        return {
            "uuid": uuid, "object_class": cls, "lab": uuid.split("+")[0],
            "scene_id": m.get("scene_id"), "success": bool(m.get("success")),
            "trajectory_length": int(m.get("trajectory_length") or 0),
            "current_task": (m.get("current_task") or "").strip(),
            "task_class": classify_task(m.get("current_task") or ""),
            "ext1_cam_serial": m.get("ext1_cam_serial"),
            "instructions": pool[uuid],
        }

    rows: list[dict] = []
    with ThreadPoolExecutor(max_workers=args.workers) as ex:
        for fut in as_completed([ex.submit(fetch, c) for c in cands]):
            r = fut.result()
            if r:
                rows.append(r)
    print(f"metadata ok for {len(rows)}")

    # Greedy max-coverage: repeatedly take the row adding the rarest still-unseen
    # combination, so every class/task/lab/outcome appears before any is doubled.
    picked: list[dict] = []
    seen = {k: Counter() for k in ("object_class", "task_class", "lab", "success")}

    def gain(r: dict) -> tuple:
        return (
            seen["object_class"][r["object_class"]],
            seen["task_class"][r["task_class"]],
            seen["lab"][r["lab"]],
            seen["success"][r["success"]],
            -r["trajectory_length"],
        )

    remaining = list(rows)
    while remaining and len(picked) < args.n:
        remaining.sort(key=gain)
        r = remaining.pop(0)
        picked.append(r)
        for k in seen:
            seen[k][r[k]] += 1

    # Failure episodes must be drawn separately and WITHOUT an object label.
    # Measured on this box's metadata cache: of 1,153 episodes, all 832 that carry
    # human annotations are successes and all 125 failures carry none -- DROID's
    # annotation pass covered successful episodes only. So any object-identity
    # filter silently selects for success, and the only way to see a failure is to
    # accept not knowing what object is in it.
    if args.n_failures:
        fails: list[dict] = []
        seen_fail_lab: Counter[str] = Counter()
        for path in sorted(args.meta_cache.glob("*/metadata.json")):
            try:
                m = json.loads(path.read_text())
            except (OSError, json.JSONDecodeError):
                continue
            uuid = m.get("uuid")
            if not uuid or uuid in covered_annotated or uuid not in covered or m.get("success"):
                continue
            lab = uuid.split("+")[0]
            if seen_fail_lab[lab] >= 2:
                continue
            seen_fail_lab[lab] += 1
            fails.append({
                "uuid": uuid, "object_class": "unknown_no_annotation", "lab": lab,
                "scene_id": m.get("scene_id"), "success": False,
                "trajectory_length": int(m.get("trajectory_length") or 0),
                "current_task": (m.get("current_task") or "").strip(),
                "task_class": classify_task(m.get("current_task") or ""),
                "ext1_cam_serial": m.get("ext1_cam_serial"), "instructions": [],
            })
            if len(fails) >= args.n_failures:
                break
        picked.extend(fails)
        for r in fails:
            for k in seen:
                seen[k][r[k]] += 1
        print(f"added {len(fails)} failure episodes (no object label available)")

    args.out_json.parent.mkdir(parents=True, exist_ok=True)
    args.out_json.write_text(json.dumps(picked, indent=1))
    print(f"\npicked {len(picked)} -> {args.out_json}")
    for k, c in seen.items():
        print(f"  {k:<14}" + ", ".join(f"{kk}={vv}" for kk, vv in sorted(c.items(), key=lambda x: str(x[0]))))
    lens = sorted(r["trajectory_length"] for r in picked)
    print(f"  length        min={lens[0]} median={lens[len(lens)//2]} max={lens[-1]}")


if __name__ == "__main__":
    main()